linkedin / linkedin/Liger-Kernel
[RFC] Testing Methodology for FP32 Convergence Tests on Fused SwiGLU/dSwiGLU Kernels
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
## Background
`LigerMLP` (#1357) co-optimizes the full **forward and backward dataflow** of the SwiGLU MLP, fusing the gate/up GEMM + SwiGLU into a single Triton kernel (rather than just the element-wise SiLU-and-gating stage), in order to reduce the memory I/O of intermediate tensors and speed up training.
Current validation status:
- Unit tests (`test/transformers/test_mlp.py`, judged by cosine similarity) all pass.
- bf16 convergence tests pass overall, except for `mini_gpt_oss`, `mini_qwen3_moe`, and `mini_pixtral`: the first two also fail on the unmodified branch and are unrelated to this change; `mini_pixtral` fails due to a rounding-path difference (~3.18e-2 relative error), not a formula bug.
- The one systematic blocker is the **fp32 convergence tests**: across the three files under `test/convergence/fp32/` (`test_mini_models.py`, `test_mini_models_with_logits.py`, `test_mini_models_multimodal.py`), about 25 cases — including `mini_llama3`, `mini_llava`, `mini_qwen2`, `mini_qwen2_vl`, `mini_qwen2_5_vl`, `mini_qwen3`, `mini_mistral`, `mini_ministral`, `mini_granite3`, `mini_exaone4`, and `mini_mllama` — all fail on Loss / Top-k logprobs mismatches.
## Root Cause
By default, `tl.dot` computes fp32 matmuls via the **TF32** Tensor Core path, which is less precise than true IEEE FP32. The reference implementation (`nn.Linear`) uses full FP32. This creates a systematic numerical gap on fp32 inputs, which is the root cause of the widespread convergence test failures — not a bug in the kernel logic.
## **Approaches I tried but didn't work**
**Forcing `input_precision="ieee"` in `tl.dot` to get true FP32**: triggers a shared memory overflow (reports needing 114712 bytes vs. a hardware limit of 101376 bytes). The reason is that TF32 compiles to Tensor Core MMA instructions with a compact data layout, while IEEE cannot use Tensor Cores and falls back to a software FMA path that needs more shared memory. Even the smallest tile configuration in the current autotune search space overflows, so there is no fallback option.
## Candidate Directions (open for discussion)
1. **Relax fp32 tolerances**: simple to implement, but lacks a principled basis for "how much error is acceptable" — treats the symptom, not the cause.
2. **Force IEEE precision + redesign a finer-grained autotune search space**: fixes the precision issue but comes with a clear performance cost, and it's unverified whether it can fully avoid the shared memory limit.
3. **Change the testing methodology itself**
## Question for Maintainers
For operators like this one — inherently precision-sensitive due to their reliance on Tensor Cores — does the community/maintainers already have a preferred **testing methodology**?
## References
- Implementation PR: #1357
- Full failure logs and attempted approaches: https://github.com/linkedin/Liger-Kernel/pull/1357#issuecomment-5293320296
- Original RFC for this kernel: #1347
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading test/convergence/fp32/test_mini_models.py, test_mini_models_with_logits.py, and test_mini_models_multimodal.py, then review implementation PR #1357 and the linked failure logs. Compare the FP32 failures with the existing unit and bf16 convergence tests. Done requires a maintainer-approved testing methodology and any agreed changes to the FP32 test approach.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, testing-qa
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100