fix: create LoRA parameters on base weight device instead of CPU
When `inject_lora()` replaces Linear layers with LoRALinear after the model has been moved to CUDA, the new lora_A and lora_B parameters were always created on CPU, causing a device mismatch error during the forward pass. Now lora_A and lora_B are created on the same device and dtype as the parent weight, matching the model's current device. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
2c50b3cf37
commit
a5678c9185
@@ -39,8 +39,10 @@ class LoRALinear(nn.Module):
|
||||
|
||||
self.r = r
|
||||
self.scaling = alpha / r
|
||||
self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1]) / r)
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r))
|
||||
device = self.weight.device
|
||||
dtype = self.weight.dtype
|
||||
self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1], device=device, dtype=dtype) / r)
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r, device=device, dtype=dtype))
|
||||
self._merged = False
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
Reference in New Issue
Block a user