Skip to content

Commit d8d6630

Browse files
authored
fix: guard save_for_backward on grad_bias not bias in fused linear CE forward (#1157)
## Summary Fixes a bug in `LigerFusedLinearCrossEntropyFunction.forward` where `ctx.save_for_backward` checks `bias is not None` instead of `grad_bias is not None` before calling `.detach()`. When `input_requires_grad=False`, `fused_linear_cross_entropy_forward` sets `grad_bias = None` regardless of whether `bias` is provided. The stale check then calls `None.detach()` and raises `AttributeError`. Fixes #1156 ## Details One-line fix in `src/liger_kernel/ops/fused_linear_cross_entropy.py`: Before: `grad_bias.detach() if bias is not None else None,` After: `grad_bias.detach() if grad_bias is not None else None,` ## Testing Done - Hardware Type: CPU (logic-only fix, no GPU kernel change) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence
1 parent 167741f commit d8d6630

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

src/liger_kernel/ops/fused_linear_cross_entropy.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,7 @@ def forward(
361361
ctx.save_for_backward(
362362
grad_input.detach(),
363363
grad_weight.detach() if grad_weight is not None else None,
364-
grad_bias.detach() if bias is not None else None,
364+
grad_bias.detach() if grad_bias is not None else None,
365365
)
366366
ctx.return_z_loss = return_z_loss
367367
ctx.return_token_accuracy = return_token_accuracy

0 commit comments

Comments
 (0)