linkedin / linkedin/Liger-Kernel

Fused MoE kernels crash on Qwen3.5-397B (int32 offset overflow)

Open
#1,296 1 comment 2 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

### 🐛 Describe the bug

The fused MoE Triton kernels in `liger_kernel/ops/fused_moe_kernels.py` compute the per-expert weight offset as `expert_idx * stride_w_E` in **int32**. When a model's expert weight tensor exceeds 2^31 elements, this offset overflows to a negative value, producing an out-of-bounds pointer and a CUDA `illegal memory access`.

This only manifests on large MoE models. Smaller ones stay under the limit and work fine, which makes the bug easy to miss.

- **liger-kernel version:** 0.8.0
- **Model that triggers it:** Qwen3.5-397B-A17B (`num_experts=512`, `hidden_size=4096`, `moe_intermediate_size=1024`)
- **Setup:** LoRA SFT, 2 nodes × 8×H200, DeepSpeed ZeRO-3, gradient checkpointing, `use_liger_kernel=True`

Forward pass dies inside `_fused_up_proj_swiglu_kernel`:

```
File ".../liger_kernel/ops/fused_moe.py", line 203, in forward
_fused_up_proj_swiglu_kernel[...](...)
...
torch.AcceleratorError: CUDA error: an illegal memory access was encountered
```

Same root cause also fires in the backward `_moe_bwd_dW2_kernel` / `_moe_bwd_dW1_kernel`.

In each kernel the expert offset is built like this (forward up-proj shown):

```python
gate_up_proj_ptr + expert_idx * stride_w_E + n_idx[:, None] * stride_w_N + k_idx[None, :] * stride_w_K
```

`expert_idx` is loaded from an int32 tensor (`tile_expert`) — or derived via `pid0 // N_M_TILES` — and `stride_w_E` fits in int32, so Triton evaluates `expert_idx * stride_w_E` in int32.

For `gate_up_proj` of shape `(E, 2*I, H)`:

```
stride_w_E = 2 * I * H = 2 * 1024 * 4096 = 8,388,608
max offset = (E-1) * stride_w_E = 511 * 8,388,608 = 4,286,578,688
```

`4,286,578,688 > 2^31-1 (2,147,483,647)` → overflow → negative offset → OOB access once the routed expert index gets high enough (~256+).

Quick check across a few models (max expert offset into `gate_up_proj` vs int32 max):

| Model | E | hidden | moe_inter | max expert offset | overflows int32? |
|---|---|---|---|---|---|
| Qwen3.5-35B-A3B | 256 | 2048 | 512 | 534,773,760 | no |
| Qwen3.5-122B-A10B | 256 | 3072 | 1024 | 1,604,321,280 | no |
| **Qwen3.5-397B-A17B** | **512** | **4096** | **1024** | **4,286,578,688** | **yes** |

### Reproduce

Multi-node LoRA SFT of **Qwen3.5-397B-A17B** via TRL `SFTTrainer` (DeepSpeed ZeRO-3, gradient checkpointing, `use_liger_kernel=True`, bf16). It crashes on the first training step inside the fused MoE kernel. The same recipe works fine on the smaller Qwen3.5-35B / 122B MoE checkpoints — only the 397B one triggers it.

### Versions

Environment Report:
-------------------
Operating System: Linux-6.8.0-71-generic-x86_64-with-glibc2.39
Python version: 3.13.12
Liger Kernel version: 0.8.0
PyTorch version: 2.12.1+cu130
CUDA version: 13.0
HIP(ROCm) version: Not available
Triton version: 3.7.1
Transformers version: 5.13.0
XPU version: XPU Not Available

Contributor guide

Open the contributing guide

Research direction

Start in liger_kernel/ops/fused_moe_kernels.py and trace the offset expressions used by _fused_up_proj_swiglu_kernel, _moe_bwd_dW2_kernel, and _moe_bwd_dW1_kernel. Reproduce with the Qwen3.5-397B-A17B configuration, then verify that expert offsets no longer overflow and that the first training step completes without a CUDA illegal memory access.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.