LoRA - TorchCode

Problem Description

Implement LoRA — parameter-efficient fine-tuning for large models.

$$h = W_0 x + \frac{\alpha}{r} B A x$$

Signature

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, rank, alpha=1.0): ...
    def forward(self, x: Tensor) -> Tensor: ...

Requirements

• self.linear: frozen nn.Linear (weight & bias requires_grad=False)

• self.lora_A: nn.Parameter(rank, in_features) — random init

• self.lora_B: nn.Parameter(out_features, rank) — zero init

• Scaling: alpha / rank

Template

Implement the function below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, rank, alpha=1.0):
        super().__init__()
        pass  # frozen linear + lora_A + lora_B

    def forward(self, x):
        pass  # base + lora

Test Your Implementation

Use this code to debug before submitting.

# 🧪 Debug
layer = LoRALinear(16, 8, rank=4)
x = torch.randn(2, 16)
print('Output:', layer(x).shape)
print('Trainable:', sum(p.numel() for p in layer.parameters() if p.requires_grad))
print('Total:    ', sum(p.numel() for p in layer.parameters()))

Reference Solution

Try solving it yourself first! Click below to reveal the solution.

# ✅ SOLUTION

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, rank, alpha=1.0):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.linear.weight.requires_grad_(False)
        self.linear.bias.requires_grad_(False)
        self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
        self.scaling = alpha / rank

    def forward(self, x):
        return self.linear(x) + (x @ self.lora_A.T @ self.lora_B.T) * self.scaling

Tips

Run Locally

For interactive practice with auto-grading, run TorchCode locally:
pip install torch-judge then use check("lora")

Key Concepts

Low-rank adaptation, frozen base + BA update