ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: create LoRA parameters on base weight device instead of CPU

Browse Source 
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:
ccx1324
2026-07-20 16:11:48 +08:00
co-authored by Claude Opus 4.7
parent 2c50b3cf37
commit a5678c9185
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/model/components/lora.py
+4 -2
View File
@@ -39,8 +39,10 @@ class LoRALinear(nn.Module):
self.r = r self.r = r
self.scaling = alpha / r self.scaling = alpha / r
self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1]) / r) device = self.weight.device
self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r)) 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 self._merged = False
def forward(self, x): def forward(self, x):