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:
Keaton Elvins
2025-11-06 03:04:05 +08:00
committed by GitHub
parent 1f7a66aeb7
commit 4a1677b84c
@@ -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