fix: Canonicalize LoRA compute dtype - #3638
Open
benthecarman wants to merge 1 commit into
Open
Conversation
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 <benthecarman@live.com>
benthecarman
force-pushed
the
fix/lora-compute-dtype
branch
from
August 24, 2026 02:27
6c5ad4f to
29d32bd
Compare
Contributor
|
/claude review |
Contributor
|
LGTM |
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Fixes LoRA forward failures when adapter, activation, and base-weight dtypes disagree, by canonicalizing the compute dtype on both the Triton and eager paths.
Changelog
TritonLinearLoRA.forward): anchor the kernel's single dtype on the base output dtype (res.dtype); cast the activation and adapter weights to it beforeapply_memory_efficient_lora/ the kernel call, so the fused residual add stays dtype-clean.LinearLoRA.forward): compute the LoRA branch in the adapter dtype (F.linear requires matching dtypes) and cast the addend back to the base output dtype.Function, so gradients flow back in the original dtypes (fp32 adapter grads, bf16 activation grads).Before your PR is "Ready for review"
Pre checks:
Additional Information
lora_dtype=torch.float32):RuntimeError: expected m1 and m2 to have the same dtype, but got: c10::BFloat16 != float.torch.squareunder bf16 autocast, as used by relu虏 activations): kernel dtype mismatch / fp32 output leaking past a bf16 module.test_eager_lora_fp32_adapters_over_bf16_base(CPU, always runs) andTestTritonLoRAMixedDtype(CUDA-gated): plain bf16 input and autocast-promoted fp32 input, checking output dtype, fp32 adapter grads, and bf16 input grads.main: the memory-efficient LoRA path (apply_memory_efficient_lora) is now the default route on both forwards, so the canonicalization feeds it directly. All existing memory-efficiency, DTensor, and compile tests pass unchanged (the casts are no-ops when dtypes already agree); full file: 33 passed, 4 TE-gated skips.馃 Generated with Claude Code