mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
fix: initialize grad_weight and grad_bias on flce no_grad path (#931)
## Summary tiny fix re: #930 , grad_weight and grad_bias were never set on no_grad path First commit is my fix, second commit was from running `make checkstyle` and a small import change for `qwen3_vl` as the import was broken and tests were not passing? Believe this was introduced in #911 ## Testing Done Repro provided in #930 now passes: ```python import torch import torch.nn as nn from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss vocab_size, hidden_dim, num_tokens = 1000, 512, 256 device = "cuda" if torch.cuda.is_available() else "cpu" linear = nn.Linear(hidden_dim, vocab_size, bias=False).to(device) fused_loss_fn = LigerFusedLinearCrossEntropyLoss() hidden_states = torch.randn(num_tokens, hidden_dim, device=device) labels = torch.randint(0, vocab_size, (num_tokens,), device=device) with torch.no_grad(): loss = fused_loss_fn(linear.weight, hidden_states, labels) print(f"Loss: {loss.item()}") ``` - Hardware Type: 3090 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence > 2110 passed, 255 skipped, 41 warnings, 1 rerun in 276.17s (0:04:36)
This commit is contained in:
@@ -58,6 +58,9 @@ def fused_linear_cross_entropy_forward(
|
||||
else:
|
||||
grad_weight = torch.zeros_like(weight, dtype=accum_dtype, device=device) if weight.requires_grad else None
|
||||
grad_bias = torch.zeros_like(bias, dtype=accum_dtype, device=device) if bias is not None else None
|
||||
else:
|
||||
grad_weight = None
|
||||
grad_bias = None
|
||||
|
||||
loss_1d = torch.zeros(BT, dtype=torch.float32, device=device)
|
||||
z_loss_1d = torch.zeros(BT, dtype=_input.dtype, device=_input.device) if return_z_loss else None
|
||||
|
||||
Reference in New Issue
Block a user