diff --git a/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index a74919ddf4..d434b1a76c 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/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: diff --git a/tests/unit_tests/_peft/test_lora.py b/tests/unit_tests/_peft/test_lora.py index 39919daf54..905bae2b83 100644 --- a/tests/unit_tests/_peft/test_lora.py +++ b/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