ByteDance-Seed / ByteDance-Seed/Triton-distributed
swiglu_backward writes every row's grad_fc1_output to row 0 (missing dA/dB output row-pointer offset) → NaN backward in fused EP-MoE
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 172
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
`swiglu_backward` (used by `TritonDistFusedEpMoeFunction` and available standalone) returns a **corrupt `grad_fc1_output`**: every row's gradient is written to **row 0** of the output. The gate gradient (`dscale`) is correct, which masks the bug because it's a per-row reduction. This produces `NaN`/garbage backward when the fused EP-MoE is used for training.
Both backward kernels in `python/triton_dist/kernels/nvidia/swiglu.py` advance the **input** row pointers (`dC_ptr`, `A_ptr`, `B_ptr`) by the row index but never advance the **output** row pointers (`dA_ptr`, `dB_ptr`). So all Triton programs load the correct input row, compute the correct gradient, then store it into the base pointer (row 0). All rows race into row 0; every row except the race winner is left uninitialized.
Why it hid: `dscale` (the gate/scale gradient) *is* offset by `row_idx` and is a per-row reduction, so it stays correct — the corruption is confined to the full-width `dA`/`dB` halves of `grad_fc1_output`.
## Affected code
`python/triton_dist/kernels/nvidia/swiglu.py`:
- `_swiglu_backward_kernel` — input pointers advanced by `row_idx`, outputs not.
- `_swiglu_backward_kernel_persistent` — same, in both the prologue and the per-row loop tail.
## Fix (6 lines)
```diff
@@ _swiglu_backward_kernel @@
dC_ptr += row_idx * dC_row_stride
A_ptr += row_idx * A_row_stride
B_ptr += row_idx * B_row_stride
+ dA_ptr += row_idx * dA_row_stride
+ dB_ptr += row_idx * dB_row_stride
@@ _swiglu_backward_kernel_persistent (prologue) @@
dC_ptr += row_start * dC_row_stride
A_ptr += row_start * A_row_stride
B_ptr += row_start * B_row_stride
+ dA_ptr += row_start * dA_row_stride
+ dB_ptr += row_start * dB_row_stride
@@ _swiglu_backward_kernel_persistent (loop tail) @@
dC_ptr += dC_row_stride
A_ptr += A_row_stride
B_ptr += B_row_stride
+ dA_ptr += dA_row_stride
+ dB_ptr += dB_row_stride
```
## Minimal reproduction (single GPU, no distributed)
```python
import torch, torch.nn.functional as F
from triton_dist.kernels.nvidia.swiglu import swiglu_forward, swiglu_backward
torch.manual_seed(0)
M, dim, dt, dev = 1024, 1536, torch.bfloat16, "cuda"
AB = torch.randn(M, dim*2, device=dev, dtype=dt) * 0.5
dC = torch.randn(M, dim, device=dev, dtype=dt) * 0.5
scale = (torch.rand(M, device=dev).float() * 0.1 + 0.01)
A = AB[:, :dim].float(); B = AB[:, dim:].float(); dCf = dC.float(); sc = scale.unsqueeze(1)
sig = torch.sigmoid(A); si = A * sig
def re(a, b): return ((a.float()-b.float()).norm() / max(b.float().norm().item(), 1e-9)).item()
# forward
sw_out, ctx = swiglu_forward(AB.clone(), scale=scale.clone())
print("forward relerr =", re(sw_out, F.silu(A) * B * sc)) # ~1e-3 (fwd is fine)
# backward
gfc1, ggate = swiglu_backward(dC.clone(), AB.clone(), scale=scale.clone(), ctx=ctx)
dA = dCf*(si*(1-sig)+sig)*B*sc; dB = dCf*si*sc; ref = torch.cat([dA, dB], dim=1)
dscale = (si*B*dCf).sum(dim=1)
print("grad_fc1_output relerr =", re(gfc1, ref)) # ~1.0 BEFORE fix -> ~1.7e-3 AFTER
print("grad_gate(dscale) relerr =", re(ggate, dscale)) # ~1e-7 (correct both before and after -> masks the bug)
```
`grad_fc1_output` relerr goes from **≈1.0 (fully wrong)** to **≈1.7e-3** with the fix; `grad_gate` is correct either way. After the fix a real fused EP-MoE 50B training run converges normally (loss-matches the DeepEP reference; grad-norm no longer diverges).
## Note
The kernel's existing correctness test hard-codes forward-only precision checking, so the backward path was never validated against a reference — hence this slipped through. Happy to expand the repro if useful. (Filed as an issue rather than a PR per our workflow; the fix above is a straightforward drop-in.)
Contributor guide
Research direction
Start in python/triton_dist/kernels/nvidia/swiglu.py by reading _swiglu_backward_kernel and _swiglu_backward_kernel_persistent, including the prologue and loop tail. Run the minimal single-GPU reproduction and compare grad_fc1_output with the reference; update the existing correctness test to cover backward gradients. Done means each row's dA/dB output matches the reference while dscale remains correct.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100