fix: Canonicalize LoRA compute dtype

nemo_automodel/components/_peft/lora.py

@@ -318,18 +318,23 @@ def forward(self, x):
             if self.dropout_position == "pre":
                 x = F.dropout(x, p=self.dropout_p, training=self.training)
 
+            # F.linear requires the adapter and activation dtypes to match;
+            # compute in the adapter dtype, cast the addend to `res`'s.
+            lora_x = x.to(self.lora_A.weight.dtype)
             # Apply scale before lora_B to keep lora_res as a Partial tensor.
             # This allows both res and lora_res to remain Partial, so only one reduce-scatter is needed after addition.
             # Multiplying after lora_B would convert Partial to Replicate, causing an extra reduce-scatter operation.
-            use_memory_efficient_lora = self._should_use_memory_efficient_lora(x)
+            use_memory_efficient_lora = self._should_use_memory_efficient_lora(lora_x)
             if use_memory_efficient_lora:
                 if self.dropout_position == "pre" or not self.training or self.dropout_p == 0.0:
                     return apply_memory_efficient_lora(
-                        x, self.lora_A.weight, self.lora_B.weight, self.scale, False, res
-                    )
-                lora_res = apply_memory_efficient_lora(x, self.lora_A.weight, self.lora_B.weight, self.scale, False)
+                        lora_x, self.lora_A.weight, self.lora_B.weight, self.scale, False, res
+                    ).to(res.dtype)
+                lora_res = apply_memory_efficient_lora(
+                    lora_x, self.lora_A.weight, self.lora_B.weight, self.scale, False
+                ).to(res.dtype)
             else:
-                lora_res = self.lora_B(self.lora_A(x) * self.scale)
+                lora_res = self.lora_B(self.lora_A(lora_x) * self.scale).to(res.dtype)
             if self.dropout_position == "post":
                 lora_res = F.dropout(lora_res, p=self.dropout_p, training=self.training)
             if use_memory_efficient_lora:
@@ -414,11 +419,22 @@ def forward(self, x):
         if self.dropout_position == "pre":
             x = F.dropout(x, p=self.dropout_p, training=self.training)
         if self.use_memory_efficient_lora:
+            # Canonicalize to the base output dtype: the kernel requires one
+            # dtype across x, adapters, and output, and the addend must match
+            # `res` for the fused addition. The casts are outside the autograd
+            # Function, so gradients flow back in the original dtypes.
+            compute_dtype = res.dtype
+            x = x.to(compute_dtype)
+            lora_A_weight = self.lora_A.weight.to(compute_dtype)
+            lora_B_weight = self.lora_B.weight.to(compute_dtype)
             if self.dropout_position == "pre" or not self.training or self.dropout_p == 0.0:
-                return apply_memory_efficient_lora(x, self.lora_A.weight, self.lora_B.weight, self.scale, True, res)
-            lora_res = apply_memory_efficient_lora(x, self.lora_A.weight, self.lora_B.weight, self.scale, True)
+                return apply_memory_efficient_lora(x, lora_A_weight, lora_B_weight, self.scale, True, res)
+            lora_res = apply_memory_efficient_lora(x, lora_A_weight, lora_B_weight, self.scale, True)
         else:
-            lora_res = self.lora_B(self.lora_A(x) * self.scale)
+            # Module calls need the activation in the adapter dtype; cast the
+            # addend back to the base output dtype.
+            lora_x = x.to(self.lora_A.weight.dtype)
+            lora_res = self.lora_B(self.lora_A(lora_x) * self.scale).to(res.dtype)
         if self.dropout_position == "post":
             lora_res = F.dropout(lora_res, p=self.dropout_p, training=self.training)
         if self.use_memory_efficient_lora:

tests/unit_tests/_peft/test_lora.py

@@ -741,3 +741,62 @@ def test_linear_lora_negative_last_dim_shard_takes_async_shaping(single_rank_pg)
     bmm_spy.assert_not_called()
     assert isinstance(out, DTensor)
     assert torch.allclose(out.full_tensor(), ref, atol=1e-6)
+
+
+def test_eager_lora_fp32_adapters_over_bf16_base():
+    """Eager LoRA with fp32 adapters over a bf16 base runs and keeps grad dtypes.
+
+    F.linear cannot mix an fp32 adapter weight with a bf16 activation: the
+    activation is computed in the adapter dtype and the addend is cast back to
+    the base output dtype. Adapter gradients stay fp32, input gradients stay
+    in the input dtype.
+    """
+    base = nn.Linear(16, 16, dtype=torch.bfloat16)
+    patched = patch_linear_module(base, dim=4, alpha=8, lora_dtype=torch.float32, use_triton=False)
+    x = torch.randn(2, 16, dtype=torch.bfloat16, requires_grad=True)
+    out = patched(x)
+    assert out.dtype == torch.bfloat16
+    out.sum().backward()
+    assert patched.lora_A.weight.dtype == torch.float32
+    assert patched.lora_A.weight.grad is not None
+    assert patched.lora_A.weight.grad.dtype == torch.float32
+    assert patched.lora_B.weight.grad is not None
+    assert patched.lora_B.weight.grad.dtype == torch.float32
+    assert x.grad is not None and x.grad.dtype == torch.bfloat16
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestTritonLoRAMixedDtype:
+    """Triton LoRA with fp32 adapters over a bf16 base, on GPU.
+
+    The kernel requires one dtype across activation, adapters, and output, so
+    the dtype is canonicalized to the base output dtype. Autocast-promoted
+    fp32 activations (e.g. torch.square under bf16 autocast) must not break
+    the kernel or leak an fp32 output.
+    """
+
+    def _build(self):
+        base = nn.Linear(512, 256, bias=False, dtype=torch.bfloat16, device="cuda")
+        return patch_linear_module(base, dim=32, alpha=64, lora_dtype=torch.float32, use_triton=True)
+
+    def test_bf16_input(self):
+        patched = self._build()
+        x = torch.randn(4, 8, 512, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        out = patched(x)
+        assert out.dtype == torch.bfloat16
+        out.float().pow(2).sum().backward()
+        assert patched.lora_A.weight.grad is not None
+        assert patched.lora_A.weight.grad.dtype == torch.float32
+
+    def test_autocast_promoted_fp32_input(self):
+        patched = self._build()
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            xb = torch.randn(4, 8, 512, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+            xp = torch.square(xb)  # autocast promotes square to fp32
+            assert xp.dtype == torch.float32
+            out = patched(xp)
+        assert out.dtype == torch.bfloat16
+        out.float().pow(2).sum().backward()
+        assert patched.lora_A.weight.grad is not None
+        assert patched.lora_A.weight.grad.dtype == torch.float32
+        assert xb.grad is not None and xb.grad.dtype == torch.bfloat16