From 29d32bdfdabbcc2a326ffa243d5a378aedc6e83b Mon Sep 17 00:00:00 2001 From: benthecarman Date: Sun, 23 Aug 2026 18:55:20 -0500 Subject: [PATCH] fix: Canonicalize LoRA compute dtype The Triton LoRA path required the activation, both adapters, and the kernel output to already share one dtype, and the eager path required the adapter dtype to match the activation dtype. Three setups break that: - fp32 master adapters over a bf16 base (wanted for small-lr optimizer precision) fail the Triton dtype assert or F.linear's dtype check; - activations promoted under autocast (torch.square, and ReLUSquaredActivation feeds exactly that) make the Triton kernel compile fp32 operands, which needs 2x shared memory and exceeds the sm_120 limit; - an fp32 LoRA addend then also leaks into `res + lora_res`, promoting the base output. Canonicalize at the boundary: the Triton path computes in the base output dtype, the eager path computes in the adapter dtype and casts the addend back to the base output dtype. The casts sit outside the autograd Function, so gradients flow back through them to the original dtypes (fp32 adapters keep fp32 gradients). Verified on an RTX 5090: bf16 input, autocast-promoted fp32 input, and the eager path all return bf16 outputs with fp32 adapter gradients and bf16 activation gradients; the promoted case previously raised OutOfResources at 9B scale. Signed-off-by: benthecarman --- nemo_automodel/components/_peft/lora.py | 32 ++++++++++---- tests/unit_tests/_peft/test_lora.py | 59 +++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 8 deletions(-) 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