loralib

import torch import torch.nn as nn import loralib as lora class MyModel(nn.Module): def __init__(self): super().__init__() self.linear1 = nn.Linear(10, 20) self.linear2 = nn.Linear(20, 5) def forward(self, x): return self.linear2(self.linear1(x)) # 1. Instantiate the base model base_model = MyModel() # 2. Convert a layer to its LoRA equivalent # Replace nn.Linear with lora.Linear, specifying rank 'r' # Here, we convert linear1 to a LoRA-enabled layer base_model.linear1 = lora.Linear(base_model.linear1.in_features, base_model.linear1.out_features, r=4) # (Optional: Convert more layers) # base_model.linear2 = lora.Linear(base_model.linear2.in_features, base_model.linear2.out_features, r=4) # 3. Mark only LoRA parameters as trainable lora.mark_only_lora_as_trainable(base_model) # Verify trainable parameters print("Trainable parameters after LoRA conversion:") for name, param in base_model.named_parameters(): if param.requires_grad: print(f" {name}: {param.shape}") # Example usage (forward pass) input_tensor = torch.randn(1, 10) output_tensor = base_model(input_tensor) print(f"Output shape: {output_tensor.shape}") # 4. Save only the LoRA-specific state_dict lora_weights = lora.lora_state_dict(base_model) # torch.save(lora_weights, 'my_model_lora.pt')