From 34e7a0b3141b319ac8e39597c7599c5254bfcdcb Mon Sep 17 00:00:00 2001 From: Brian Staber Date: Wed, 8 Oct 2025 23:03:17 +0200 Subject: [PATCH] :cake: Simple LoRA MLP example --- experiments/lora/mlp.py | 253 ++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + uv.lock | 53 +++++++++ 3 files changed, 307 insertions(+) create mode 100644 experiments/lora/mlp.py diff --git a/experiments/lora/mlp.py b/experiments/lora/mlp.py new file mode 100644 index 0000000..7e07fd5 --- /dev/null +++ b/experiments/lora/mlp.py @@ -0,0 +1,253 @@ +"""LoRA MLP model implementation for experiments. + +We implement a simple MLP model and its LoRA variant for experiments. + +The code includes: +- BaseMLP: A standard MLP model. +- LoRALinear: A linear layer with LoRA adaptation. +- LoRAMLP: An MLP model using LoRA layers. +- Example usage with dummy regression data. + +In the example, we: +1. Train a base MLP on a regression task and plot the predictions (sanity check). +2. Fully fine-tune the base MLP on a new task for comparison. +3. Fine-tune the LoRA MLP on the new task and compare performance. +4. Fully retrain a new base MLP on the new task for comparison. +5. Then, for the three models, we plot the training losses and predictions for comparison. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class BaseMLP(nn.Module): + """Base MLP model used in LoRA experiments.""" + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + layers = nn.ModuleList() + layers.append(nn.Linear(input_dim, hidden_dim)) + for _ in range(num_layers - 2): + layers.append(nn.Linear(hidden_dim, hidden_dim)) + layers.append(nn.Linear(hidden_dim, output_dim)) + self.layers = layers + + def forward(self, x): + """Forward pass through the MLP.""" + for layer in self.layers[:-1]: + x = F.relu(layer(x)) + x = self.layers[-1](x) + return x + + +class LoRALinear(nn.Module): + """LoRA Linear layer.""" + + def __init__(self, in_features, out_features, r=4, alpha=1.0): + super().__init__() + self.base = nn.Linear(in_features, out_features) + self.in_features = in_features + self.out_features = out_features + self.r = r + self.alpha = alpha + + self.B = nn.Linear(in_features, r, bias=False) + self.A = nn.Linear(r, out_features, bias=False) + self.scaling = alpha / r + + nn.init.kaiming_uniform_(self.B.weight, a=math.sqrt(5)) + nn.init.zeros_(self.A.weight) + + def forward(self, x): + """Forward pass through the LoRA Linear layer.""" + return self.base(x) + self.scaling * self.A(self.B(x)) + + +class LoRAMLP(nn.Module): + """MLP model with LoRA layers.""" + + def __init__( + self, + input_dim, + hidden_dim, + output_dim, + num_layers, + base_state_dict, + r=4, + alpha=1.0, + ): + super().__init__() + layers = nn.ModuleList() + layers.append(LoRALinear(input_dim, hidden_dim, r=r, alpha=alpha)) + for _ in range(num_layers - 2): + layers.append(LoRALinear(hidden_dim, hidden_dim, r=r, alpha=alpha)) + layers.append(LoRALinear(hidden_dim, output_dim, r=r, alpha=alpha)) + self.layers = layers + + with torch.no_grad(): + for name, param in base_state_dict.items(): + if "weight" in name: + layer_idx = int(name.split(".")[1]) + self.layers[layer_idx].base.weight.copy_(param) + elif "bias" in name: + layer_idx = int(name.split(".")[1]) + self.layers[layer_idx].base.bias.copy_(param) + + for module in self.modules(): + if hasattr(module, "base") and isinstance(module.base, nn.Linear): + for p in module.base.parameters(): + p.requires_grad_(False) + + def forward(self, x): + """Forward pass through the LoRA MLP.""" + for layer in self.layers[:-1]: + x = F.relu(layer(x)) + x = self.layers[-1](x) + return x + + +if __name__ == "__main__": + """Example usage of LoRAMLP with dummy data.""" + import matplotlib.pyplot as plt + + # Example usage + input_dim = 1 + hidden_dim = 20 + output_dim = input_dim + num_layers = 3 + + # Dummy regression data + x = torch.rand(1000, input_dim) + y = torch.cos(x * 3.14 * 2) + torch.randn_like(x) * 0.1 + x_test = torch.linspace(0, 1, 100).unsqueeze(1).repeat(1, input_dim) + + # Train base MLP + model_first_task = BaseMLP(input_dim, hidden_dim, output_dim, num_layers) + + dataset = torch.utils.data.TensorDataset(x, y) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True) + model_first_task.train() + optimizer = torch.optim.AdamW(model_first_task.parameters(), lr=5e-3) + for _ in range(1000): + optimizer.zero_grad() + loss = F.mse_loss(model_first_task(x), y) + loss.backward() + optimizer.step() + + model_first_task.eval() + with torch.no_grad(): + preds = model_first_task(x_test) + + plt.figure() + plt.scatter(x.numpy(), y.numpy(), label="data", alpha=0.3) + plt.plot(x_test.numpy(), preds.numpy(), color="red", label="model") + plt.legend() + plt.show() + + # Dummy new task data + x = torch.rand(1000, input_dim) + y = torch.sin(x * 3.14 * 4) + torch.randn_like(x) * 0.1 + x_test = torch.linspace(0, 1, 100).unsqueeze(1).repeat(1, input_dim) + dataset = torch.utils.data.TensorDataset(x, y) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True) + + # Train the base model on the new task for comparison + model = BaseMLP(input_dim, hidden_dim, output_dim, num_layers) + optimizer = torch.optim.AdamW(model.parameters(), lr=5e-3) + + loss_train_base = [] + model.train() + for _ in range(50): + loss_epoch = 0.0 + for xb, yb in dataloader: + optimizer.zero_grad() + loss = F.mse_loss(model(xb), yb) + loss.backward() + loss_epoch += loss.item() + optimizer.step() + loss_train_base.append(loss_epoch / len(dataloader)) + + # Fully finetune the first model + model_full_finetuned = BaseMLP(input_dim, hidden_dim, output_dim, num_layers) + model_full_finetuned.load_state_dict(model_first_task.state_dict()) + optimizer = torch.optim.AdamW(model_full_finetuned.parameters(), lr=5e-3) + + loss_train_full_finetuned = [] + model_full_finetuned.train() + for _ in range(50): + loss_epoch = 0.0 + for xb, yb in dataloader: + optimizer.zero_grad() + loss = F.mse_loss(model_full_finetuned(xb), yb) + loss.backward() + loss_epoch += loss.item() + optimizer.step() + loss_train_full_finetuned.append(loss_epoch / len(dataloader)) + + # Fine-tune the LoRA MLP on another task + # Initialize LoRA MLP with the pre-trained base model's state dict + model_lora = LoRAMLP( + input_dim, + hidden_dim, + output_dim, + num_layers, + model_first_task.state_dict(), + r=4, + alpha=8.0, + ) + opt = torch.optim.AdamW( + [p for p in model_lora.parameters() if p.requires_grad], + lr=5e-3, + weight_decay=1e-2, + ) + + loss_train_lora = [] + model_lora.train() + for _ in range(50): + loss_epoch = 0.0 + for xb, yb in dataloader: + opt.zero_grad() + loss = F.mse_loss(model_lora(xb), yb) + loss.backward() + loss_epoch += loss.item() + opt.step() + loss_train_lora.append(loss_epoch / len(dataloader)) + + model_lora.eval() + with torch.no_grad(): + preds_lora = model_lora(x_test) + preds_full_finetune = model_full_finetuned(x_test) + preds_base_retrained = model(x_test) + + plt.figure() + plt.plot(loss_train_lora, label="LoRA model on 2nd task") + plt.plot(loss_train_base, label="Fully trained base model") + plt.plot(loss_train_full_finetuned, label="Fully finetuned 1st base model") + plt.yscale("log") + plt.xlabel("Iteration") + plt.ylabel("MSE Loss") + plt.legend() + plt.show() + + plt.figure() + plt.scatter(x.numpy(), y.numpy(), color="k", label="data", alpha=0.3) + plt.plot( + x_test.numpy(), preds_lora.numpy(), color="red", label="Finetuned LoRA model" + ) + plt.plot( + x_test.numpy(), + preds_full_finetune.numpy(), + color="blue", + label="Fully finetuned 1st base model", + ) + plt.plot( + x_test.numpy(), + preds_base_retrained.numpy(), + color="green", + label="Fully trained base model", + ) + plt.legend() + plt.show() diff --git a/pyproject.toml b/pyproject.toml index 2005d94..8b817f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "hetgpy>=1.0.4", "hydra-core>=1.3.2", "nvidia-physicsnemo>=1.2.0", + "pyqt5>=5.15.11", "pyvista>=0.46.3", "scikit-learn>=1.7.2", "the-well>=1.1.0", diff --git a/uv.lock b/uv.lock index ef0e64c..9872811 100644 --- a/uv.lock +++ b/uv.lock @@ -507,6 +507,7 @@ dependencies = [ { name = "hetgpy" }, { name = "hydra-core" }, { name = "nvidia-physicsnemo" }, + { name = "pyqt5" }, { name = "pyvista" }, { name = "scikit-learn" }, { name = "the-well" }, @@ -536,6 +537,7 @@ requires-dist = [ { name = "hetgpy", specifier = ">=1.0.4" }, { name = "hydra-core", specifier = ">=1.3.2" }, { name = "nvidia-physicsnemo", specifier = ">=1.2.0" }, + { name = "pyqt5", specifier = ">=5.15.11" }, { name = "pyvista", specifier = ">=0.46.3" }, { name = "scikit-learn", specifier = ">=1.7.2" }, { name = "the-well", specifier = ">=1.1.0" }, @@ -2106,6 +2108,57 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/05/e7/df2285f3d08fee213f2d041540fa4fc9ca6c2d44cf36d3a035bf2a8d2bcc/pyparsing-3.2.3-py3-none-any.whl", hash = "sha256:a749938e02d6fd0b59b356ca504a24982314bb090c383e3cf201c95ef7e2bfcf", size = 111120, upload-time = "2025-03-25T05:01:24.908Z" }, ] +[[package]] +name = "pyqt5" +version = "5.15.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyqt5-qt5" }, + { name = "pyqt5-sip" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/07/c9ed0bd428df6f87183fca565a79fee19fa7c88c7f00a7f011ab4379e77a/PyQt5-5.15.11.tar.gz", hash = "sha256:fda45743ebb4a27b4b1a51c6d8ef455c4c1b5d610c90d2934c7802b5c1557c52", size = 3216775, upload-time = "2024-07-19T08:39:57.756Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/64/42ec1b0bd72d87f87bde6ceb6869f444d91a2d601f2e67cd05febc0346a1/PyQt5-5.15.11-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:c8b03dd9380bb13c804f0bdb0f4956067f281785b5e12303d529f0462f9afdc2", size = 6579776, upload-time = "2024-07-19T08:39:19.775Z" }, + { url = "https://files.pythonhosted.org/packages/49/f5/3fb696f4683ea45d68b7e77302eff173493ac81e43d63adb60fa760b9f91/PyQt5-5.15.11-cp38-abi3-macosx_11_0_x86_64.whl", hash = "sha256:6cd75628f6e732b1ffcfe709ab833a0716c0445d7aec8046a48d5843352becb6", size = 7016415, upload-time = "2024-07-19T08:39:32.977Z" }, + { url = "https://files.pythonhosted.org/packages/b4/8c/4065950f9d013c4b2e588fe33cf04e564c2322842d84dbcbce5ba1dc28b0/PyQt5-5.15.11-cp38-abi3-manylinux_2_17_x86_64.whl", hash = "sha256:cd672a6738d1ae33ef7d9efa8e6cb0a1525ecf53ec86da80a9e1b6ec38c8d0f1", size = 8188103, upload-time = "2024-07-19T08:39:40.561Z" }, + { url = "https://files.pythonhosted.org/packages/f3/f0/ae5a5b4f9b826b29ea4be841b2f2d951bcf5ae1d802f3732b145b57c5355/PyQt5-5.15.11-cp38-abi3-win32.whl", hash = "sha256:76be0322ceda5deecd1708a8d628e698089a1cea80d1a49d242a6d579a40babd", size = 5433308, upload-time = "2024-07-19T08:39:46.932Z" }, + { url = "https://files.pythonhosted.org/packages/56/d5/68eb9f3d19ce65df01b6c7b7a577ad3bbc9ab3a5dd3491a4756e71838ec9/PyQt5-5.15.11-cp38-abi3-win_amd64.whl", hash = "sha256:bdde598a3bb95022131a5c9ea62e0a96bd6fb28932cc1619fd7ba211531b7517", size = 6865864, upload-time = "2024-07-19T08:39:53.572Z" }, +] + +[[package]] +name = "pyqt5-qt5" +version = "5.15.17" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d3/f9/accb06e76e23fb23053d48cc24fd78dec6ed14cb4d5cbadb0fd4a0c1b02e/PyQt5_Qt5-5.15.17-py3-none-macosx_10_13_x86_64.whl", hash = "sha256:d8b8094108e748b4bbd315737cfed81291d2d228de43278f0b8bd7d2b808d2b9", size = 39972275, upload-time = "2025-05-24T11:15:42.259Z" }, + { url = "https://files.pythonhosted.org/packages/87/1a/e1601ad6934cc489b8f1e967494f23958465cf1943712f054c5a306e9029/PyQt5_Qt5-5.15.17-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b68628f9b8261156f91d2f72ebc8dfb28697c4b83549245d9a68195bd2d74f0c", size = 37135109, upload-time = "2025-05-24T11:15:59.786Z" }, + { url = "https://files.pythonhosted.org/packages/ac/e1/13d25a9ff2ac236a264b4603abaa39fa8bb9a7aa430519bb5f545c5b008d/PyQt5_Qt5-5.15.17-py3-none-manylinux2014_x86_64.whl", hash = "sha256:b018f75d1cc61146396fa5af14da1db77c5d6318030e5e366f09ffdf7bd358d8", size = 61112954, upload-time = "2025-05-24T11:16:26.036Z" }, +] + +[[package]] +name = "pyqt5-sip" +version = "12.17.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ea/08/88a20c862f40b5c178c517cdc7e93767967dec5ac1b994e226d517991c9b/pyqt5_sip-12.17.1.tar.gz", hash = "sha256:0eab72bcb628f1926bf5b9ac51259d4fa18e8b2a81d199071135458f7d087ea8", size = 104136, upload-time = "2025-10-08T09:04:19.893Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/e4/451e465c75584a7cbd10e10404317b7443af83f56a64e02080b1f3cda5b5/pyqt5_sip-12.17.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:5134d637efadd108a70306bab55b3d7feaa951bf6b8162161a67ae847bea9130", size = 122581, upload-time = "2025-10-08T09:04:13.607Z" }, + { url = "https://files.pythonhosted.org/packages/dc/b2/330f97434b21fbc99ab16f6ce71358ff5ea1bf1f09ed14dfe6b28b5ed8f5/pyqt5_sip-12.17.1-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:155cf755266c8bf64428916e2ff720d5efa1aec003d4ccc40c003b147dbdac03", size = 276844, upload-time = "2025-10-08T09:15:33.713Z" }, + { url = "https://files.pythonhosted.org/packages/3b/fd/53925099d0fc8aaf7adee613b6cebfb3fdfcd1238add64ff9edf6711e5f8/pyqt5_sip-12.17.1-cp311-cp311-win32.whl", hash = "sha256:9dfa7fe4ac93b60004430699c4bf56fef842a356d64dfea7cbc6d580d0427d6d", size = 49099, upload-time = "2025-10-08T09:11:23.928Z" }, + { url = "https://files.pythonhosted.org/packages/33/f8/f47a849c17676557c4220fbce9fcc24e15736af247c4dddcaf9ff0124b57/pyqt5_sip-12.17.1-cp311-cp311-win_amd64.whl", hash = "sha256:2ddd214cf40119b86942a5da2da5a7345334955ab00026d8dcc56326b30e6d3c", size = 58988, upload-time = "2025-10-08T09:08:34.903Z" }, + { url = "https://files.pythonhosted.org/packages/a5/15/291f83f336558300626bebb0c403084ec171bbc8a70683e3376234422eb6/pyqt5_sip-12.17.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:c362606de782d2d46374a38523632786f145c517ee62de246a6069e5f2c5f336", size = 124521, upload-time = "2025-10-08T09:04:15.825Z" }, + { url = "https://files.pythonhosted.org/packages/45/85/ea1ae099260fd1859d71b31f51760b4226abfa778d5796b76d92c8fe6dcd/pyqt5_sip-12.17.1-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:140cc582151456103ebb149fefc678f3cae803e7720733db51212af5219cd45c", size = 282182, upload-time = "2025-10-08T09:15:35.752Z" }, + { url = "https://files.pythonhosted.org/packages/b4/b3/d5b50c721651a0f2ccbef6f8db3dabf3db296b9ec239ba007f5615f57dd7/pyqt5_sip-12.17.1-cp312-cp312-win32.whl", hash = "sha256:9dc1f1525d4d42c080f6cfdfc70d78239f8f67b0a48ea0745497251d8d848b1d", size = 49447, upload-time = "2025-10-08T09:11:24.843Z" }, + { url = "https://files.pythonhosted.org/packages/14/b6/474d8b17763683ab45fb364f3a44f25fdc25d97b47b29ad8819b95a15ac8/pyqt5_sip-12.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:d5e2e9e175559017cd161d661e0ee0b551684f824bb90800c5a8c8a3bea9355e", size = 57946, upload-time = "2025-10-08T09:08:35.775Z" }, + { url = "https://files.pythonhosted.org/packages/1d/58/9ecb688050e79ffe7bbd9fc917aa13f63856a5081ac46bbce87bb11ab971/pyqt5_sip-12.17.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:9ebbd7769ccdaaa6295e9c872553b6cde17f38e171056f17300d8af9a14d1fc8", size = 124485, upload-time = "2025-10-08T09:04:17.473Z" }, + { url = "https://files.pythonhosted.org/packages/b1/9f/ae691360a9f18e3e06fd297e854d7ad175367e35ea184fd2fcf6c79b8c25/pyqt5_sip-12.17.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:b023da906a70af2cf5e6fc1932f441ede07530f3e164dd52c6c2bb5ab7c6f424", size = 281923, upload-time = "2025-10-08T09:15:37.004Z" }, + { url = "https://files.pythonhosted.org/packages/d7/31/491c45423174a359a4b8a8d84a7b541c453f48497ae928cbe4006bcd3e01/pyqt5_sip-12.17.1-cp313-cp313-win32.whl", hash = "sha256:36dbef482bd638786b909f3bda65b7b3d5cbd6cbf16797496de38bae542da307", size = 49400, upload-time = "2025-10-08T09:11:25.769Z" }, + { url = "https://files.pythonhosted.org/packages/64/61/e28681dd5200094f7b2e6671e85c02a4d6693da36d23ad7d39ffbc70b15c/pyqt5_sip-12.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:d04e5551bbc3bcec98acc63b3b0618ddcbf31ff107349225b516fe7e7c0a7c8b", size = 57979, upload-time = "2025-10-08T09:08:37.036Z" }, + { url = "https://files.pythonhosted.org/packages/71/f9/06c09dc94474ffe3f518f80e47fc69d34abf8e4a971ae7e7c667d6ff30a7/pyqt5_sip-12.17.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:c49918287e1ad77956d1589f1d3d432a0be7630c646ea02cf652413a48e14458", size = 124400, upload-time = "2025-10-08T08:38:23.927Z" }, + { url = "https://files.pythonhosted.org/packages/40/ae/be6e338ea427deac5cd81a93f51ae3fb6505d99d6d5e5d5341bcc099327e/pyqt5_sip-12.17.1-cp314-cp314-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:944a4bf1e1ee18ad03a54964c1c6433fb6de582313a1f0b17673e7203e22fc83", size = 282291, upload-time = "2025-10-08T08:38:25.735Z" }, + { url = "https://files.pythonhosted.org/packages/fc/a3/8b758518bd0dd5d1581f7a6d522c9b4d9b58d05087b1d0b4dfaad5376434/pyqt5_sip-12.17.1-cp314-cp314-win32.whl", hash = "sha256:99a2935fd662a67748625b1e6ffa0a2d1f2da068b9df6db04fa59a4a5d4ee613", size = 50578, upload-time = "2025-10-08T08:38:28.72Z" }, + { url = "https://files.pythonhosted.org/packages/40/8c/e96f9877548810b1e537f46fc21ba74552dd4e8c498658114a8353bdf659/pyqt5_sip-12.17.1-cp314-cp314-win_amd64.whl", hash = "sha256:aaa33232cc80793d14fdb3b149b27eec0855612ed66aad480add5ac49b9cee63", size = 59763, upload-time = "2025-10-08T08:38:27.443Z" }, +] + [[package]] name = "pyright" version = "1.1.405"