diff --git a/experiments/lora/lora_mlp.py b/experiments/lora/lora_mlp.py new file mode 100644 index 0000000..4123b05 --- /dev/null +++ b/experiments/lora/lora_mlp.py @@ -0,0 +1,159 @@ +"""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 matplotlib.pyplot as plt +import torch +import torch.nn.functional as F +from lora.models import BaseMLP, LoRAMLP + +# 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/experiments/lora/lora_pinn.py b/experiments/lora/lora_pinn.py new file mode 100644 index 0000000..c430047 --- /dev/null +++ b/experiments/lora/lora_pinn.py @@ -0,0 +1,170 @@ +"""LoRA applied to Physics-Informed Neural Networks (PINNs) for Burgers' equation.""" + +import numpy as np +import scipy +import torch +import torch.nn as nn +from lora.models import BaseMLP, LoRAMLP +from torch.autograd import grad +from torch.utils.data import DataLoader, TensorDataset + + +def loss_burgers( + model: nn.Module, + xs_pde: torch.Tensor, + nu: float = 0.01 / torch.pi, + x_min: float = 0.0, + x_max: float = 1.0, + t_min: float = 0.0, + t_max: float = 1.0, + n_ic: int = 128, + n_bc: int = 128, +): + """Compute the PINN loss for the Burgers' equation.""" + # xs_pde: (B_pde, 2) with requires_grad=True + xs_pde = xs_pde.requires_grad_(True) + u = model(xs_pde) # (B_pde, 1) + + # grads wrt inputs + g = grad(u, xs_pde, torch.ones_like(u), create_graph=True)[0] + u_x = g[:, 0:1] + u_t = g[:, 1:2] + u_xx = grad(u_x, xs_pde, torch.ones_like(u_x), create_graph=True)[0][:, 0:1] + + r = u_t + u * u_x - nu * u_xx + loss_pde = torch.mean(r**2) + + # sample IC/BC fresh each call (common in PINNs) + device = xs_pde.device + x_ic = x_min + (x_max - x_min) * torch.rand(n_ic, 1, device=device) + t_ic = torch.zeros_like(x_ic) + u_ic = model(torch.cat([x_ic, t_ic], dim=1)) + target_ic = -torch.sin(torch.pi * x_ic) + loss_ic = torch.mean((u_ic - target_ic) ** 2) + + t_bc = t_min + (t_max - t_min) * torch.rand(n_bc, 1, device=device) + x0 = torch.full_like(t_bc, x_min) + x1 = torch.full_like(t_bc, x_max) + u_left = model(torch.cat([x0, t_bc], dim=1)) + u_right = model(torch.cat([x1, t_bc], dim=1)) + loss_bc = torch.mean(u_left**2) + torch.mean(u_right**2) + + return loss_pde + loss_ic + loss_bc + + +if __name__ == "__main__": + filename = "./data/burgers_shock.mat" + data = scipy.io.loadmat(filename) + + nu = 0.01 / torch.pi + x = torch.tensor(data["x"], dtype=torch.float32) + t = torch.tensor(data["t"], dtype=torch.float32) + u_sol = torch.tensor(data["usol"], dtype=torch.float32) + + # Create meshgrid (training points) + X_grid, T_grid = torch.meshgrid(x.flatten(), t.flatten(), indexing="ij") + X_star = torch.stack([X_grid.flatten(), T_grid.flatten()], axis=-1) # (25600, 2) + + # Subsample the training data + n_pde, n_obs = 1000, 1000 + idx_pde = np.random.choice(len(X_star), size=(n_pde,), replace=False) + idx_obs = np.random.choice(len(X_star), size=(n_obs,), replace=False) + xs_pde = X_star[idx_pde] + xs_obs = X_star[idx_obs] + u_obs = u_sol.flatten()[idx_obs].unsqueeze(-1) + + # Base MLP + xs_pde.requires_grad = True + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = BaseMLP(input_dim=2, hidden_dim=10, output_dim=1, num_layers=10).to(device) + + # Training base model + optim = torch.optim.AdamW(params=model.parameters(), lr=1e-3, weight_decay=0.0) + obs_loader = DataLoader(TensorDataset(xs_obs, u_obs), batch_size=64, shuffle=True) + + num_epochs = 1000 + loss_data = nn.MSELoss() + train_loss = [] + for epoch in range(num_epochs): + model.train() + loss_train_epoch = 0.0 + for xs_batch_obs, u_obs_batch in obs_loader: + xs_batch_obs, u_obs_batch = xs_batch_obs.to(device), u_obs_batch.to(device) + optim.zero_grad() + pred = model(xs_batch_obs) + loss = loss_data(pred, u_obs_batch) + loss.backward() + optim.step() + loss_train_epoch += loss.item() + loss_train_epoch /= len(obs_loader) + train_loss.append(loss_train_epoch) + + # Fine-tune with LoRA + pde_loader = DataLoader(TensorDataset(xs_pde), batch_size=128, shuffle=True) + + model_lora = LoRAMLP( + input_dim=2, + hidden_dim=10, + output_dim=1, + num_layers=10, + base_state_dict=model.state_dict(), + r=4, + alpha=8.0, + ).to(device) + + lora_optim = torch.optim.AdamW( + [p for p in model_lora.parameters() if p.requires_grad], + lr=1e-3, + weight_decay=0.0, + ) + + loss_train_lora = [] + model_lora.train() + for _ in range(1000): + loss_epoch = 0.0 + for xs_pde_batch in pde_loader: + xs_pde_batch = xs_pde_batch[0].to(device) + lora_optim.zero_grad() + loss = loss_burgers(model_lora, xs_pde_batch) + loss.backward() + loss_epoch += loss.item() + lora_optim.step() + loss_train_lora.append(loss_epoch / len(pde_loader)) + + # Plot losses + import matplotlib.pyplot as plt + + fig, axs = plt.subplots(1, 2, figsize=(12, 4)) + axs[0].semilogy(train_loss, color="k", ls="-") + axs[1].semilogy(loss_train_lora, color="b", ls="-") + plt.show() + + # Plot predictions + model.eval() + with torch.no_grad(): + pred_base_model = model(X_star).reshape(X_grid.shape).cpu().numpy() + pred_lora_model = model_lora(X_star).reshape(X_grid.shape).cpu().numpy() + u_sol = u_sol.cpu().numpy() + + fig, axs = plt.subplots(2, 2, figsize=(10, 8)) + c2 = axs[0, 0].pcolormesh( + T_grid, X_grid, pred_base_model, shading="auto", cmap="jet" + ) + fig.colorbar(c2, ax=axs[0, 0], label="u(x, t)") + axs[0, 0].set_title("Base MLP solution") + c2 = axs[0, 1].pcolormesh( + T_grid, X_grid, pred_lora_model, shading="auto", cmap="jet" + ) + fig.colorbar(c2, ax=axs[0, 1], label="u(x, t)") + axs[0, 1].set_title("LoRA fine-tuned solution") + c2 = axs[1, 0].pcolormesh( + T_grid, X_grid, np.abs(pred_base_model - u_sol), shading="auto", cmap="jet" + ) + fig.colorbar(c2, ax=axs[1, 0], label="Error") + axs[1, 0].set_title("Abs error") + c2 = axs[1, 1].pcolormesh( + T_grid, X_grid, np.abs(pred_lora_model - u_sol), shading="auto", cmap="jet" + ) + fig.colorbar(c2, ax=axs[1, 1], label="Error") + axs[1, 1].set_title("Abs error") + plt.show() diff --git a/libs/lora/README.md b/libs/lora/README.md new file mode 100644 index 0000000..e69de29 diff --git a/libs/lora/pyproject.toml b/libs/lora/pyproject.toml new file mode 100644 index 0000000..0579e50 --- /dev/null +++ b/libs/lora/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "lora" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +authors = [ + { name = "Brian Staber", email = "brian.staber@gmail.com" } +] +requires-python = ">=3.11" +dependencies = [] + +[build-system] +requires = ["uv_build>=0.8.11,<0.9.0"] +build-backend = "uv_build" diff --git a/libs/lora/src/lora/__init__.py b/libs/lora/src/lora/__init__.py new file mode 100644 index 0000000..8f5a6c3 --- /dev/null +++ b/libs/lora/src/lora/__init__.py @@ -0,0 +1 @@ +"""LoRA module.""" diff --git a/libs/lora/src/lora/linear.py b/libs/lora/src/lora/linear.py new file mode 100644 index 0000000..a2c0927 --- /dev/null +++ b/libs/lora/src/lora/linear.py @@ -0,0 +1,44 @@ +"""LoRA Linear layer implementation.""" + +import math + +import torch.nn as nn + + +class LoRALayer(nn.Module): + """Implement a standalone LoRA layer.""" + + def __init__( + self, in_features: int, out_features: int, r: int = 4, alpha: float = 1.0 + ): + super().__init__() + 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.scaling * self.A(self.B(x)) + + +class LoRALinear(nn.Module): + """Linear layer with LoRA..""" + + def __init__( + self, in_features: int, out_features: int, r: int = 4, alpha: float = 1.0 + ): + super().__init__() + self.base = nn.Linear(in_features, out_features) + self.lora_layer = LoRALayer(in_features, out_features, r, alpha) + + def forward(self, x): + """Forward pass through the LoRA Linear layer.""" + return self.base(x) + self.lora_layer(x) diff --git a/libs/lora/src/lora/models.py b/libs/lora/src/lora/models.py new file mode 100644 index 0000000..3cee5de --- /dev/null +++ b/libs/lora/src/lora/models.py @@ -0,0 +1,70 @@ +"""Simple MLP models for LoRA experiments.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from lora.linear import LoRALinear + + +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.tanh(layer(x)) + x = self.layers[-1](x) + return 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.tanh(layer(x)) + x = self.layers[-1](x) + return x diff --git a/libs/lora/src/lora/py.typed b/libs/lora/src/lora/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/uv.lock b/uv.lock index 9adeeab..633b7e3 100644 --- a/uv.lock +++ b/uv.lock @@ -12,6 +12,7 @@ members = [ "dummy-app", "dummy-lib", "endurance", + "lora", "mini-neural-operators", "mini-nn", "mini-rnn", @@ -1246,6 +1247,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/e9/0d4add7873a73e462aeb45c036a2dead2562b825aa46ba326727b3f31016/kiwisolver-1.4.9-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:fb940820c63a9590d31d88b815e7a3aa5915cad3ce735ab45f0c730b39547de1", size = 73929, upload-time = "2025-08-10T21:27:48.236Z" }, ] +[[package]] +name = "lora" +version = "0.1.0" +source = { editable = "libs/lora" } + [[package]] name = "markdown" version = "3.9"