Upload code/faz3_v6/meshai_train/lora_inject.py with huggingface_hub

code/faz3_v6/meshai_train/lora_inject.py

@@ -53,9 +53,18 @@ def inject_lora_into_linears(
             break
     replaced: list[str] = []
     for name, lin in picked:
+        # Keep original device/dtype (Paint often fp16)
+        device = lin.weight.device
+        dtype = lin.weight.dtype
         w = lin.weight.detach().float().cpu()
         lora = FrozenLinearLoRA(w, rank=rank, name=name)
-        # match device/dtype later via .to()
+        lora = lora.to(device=device, dtype=dtype)
+        # Ensure LoRA params stay trainable float32 for Adam stability when base is half
+        if dtype != torch.float32:
+            lora.lora_A.data = lora.lora_A.data.float()
+            lora.lora_B.data = lora.lora_B.data.float()
+            # weight buffer matches module dtype for matmul with activations
+            lora.weight.data = lora.weight.data.to(dtype=dtype)
         _set_module_by_name(root, name, lora)
         replaced.append(name)
     log_fn(f"[faz3-lora] injected {len(replaced)} LoRA layers (rank={rank})")