Fused CE backward: guard scaling=0, drop tensor path, use out-of-place mul#610
Conversation
mmathew23
commented
Apr 24, 2026
mmathew23
commented
- Raise on scaling=0 with non-zero grad_output instead of silently producing NaN grads (saved grad is zeroed by scaling in forward and the unscaled grad cannot be recovered). Return zero grads when grad_output is also 0 (chain rule).
- Collapse tensor scaling into the scalar path via a single .item() at the boundary. All current callers pass a Python float via GradScaler.get_scale(); removes the torch.where branch and the duplicated logging handling.
- Switch the final mul to out-of-place so ctx.saved_tensors is not version-bumped, fixing retain_graph=True. Measured peak memory delta is <3 MB across 14 LoRA / full-FT / MoE / vision configs up to bsz=16, seq=8192.
- Raise on scaling=0 with non-zero grad_output instead of silently producing NaN grads (saved grad is zeroed by scaling in forward and the unscaled grad cannot be recovered). Return zero grads when grad_output is also 0 (chain rule). - Collapse tensor scaling into the scalar path via a single .item() at the boundary. All current callers pass a Python float via GradScaler.get_scale(); removes the torch.where branch and the duplicated logging handling. - Switch the final mul to out-of-place so ctx.saved_tensors is not version-bumped, fixing retain_graph=True. Measured peak memory delta is <3 MB across 14 LoRA / full-FT / MoE / vision configs up to bsz=16, seq=8192.
![gemini-code-assist[bot]](https://avatars.mcszy.xyz/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enhances the UnslothFusedLoss backward pass by implementing explicit handling for zero scaling factors and transitioning to out-of-place multiplication to support retain_graph and double-backward flows. Feedback suggests further optimizing performance by skipping the multiplication operation when the scale factor is exactly 1.0, which prevents redundant memory copies of large gradient tensors.
| grad_inputs = grad_inputs * scale_factor | ||
| if grad_lm_head is not None: grad_lm_head = grad_lm_head * scale_factor | ||
| if grad_lm_head_bias is not None: grad_lm_head_bias = grad_lm_head_bias * scale_factor |
There was a problem hiding this comment.
While switching to out-of-place multiplication correctly avoids version-bumping ctx.saved_tensors (fixing retain_graph=True issues), it unconditionally creates new tensor copies even when scale_factor is 1.0. For large tensors like grad_inputs, this can be expensive in terms of both memory and compute.
Consider adding a check to return the original tensors if scale_factor is 1.0. Since scale_factor is a 0-dim tensor, calling .item() to check its value is a small sync cost that is significantly outweighed by avoiding a multi-gigabyte copy.
| grad_inputs = grad_inputs * scale_factor | |
| if grad_lm_head is not None: grad_lm_head = grad_lm_head * scale_factor | |
| if grad_lm_head_bias is not None: grad_lm_head_bias = grad_lm_head_bias * scale_factor | |
| if scale_factor.item() != 1.0: | |
| grad_inputs = grad_inputs * scale_factor | |
| if grad_lm_head is not None: grad_lm_head = grad_lm_head * scale_factor | |
| if grad_lm_head_bias is not None: grad_lm_head_bias = grad_lm_head_bias * scale_factor |
1 participant
