linkedin / linkedin/Liger-Kernel

[RFC] Testing Methodology for FP32 Convergence Tests on Fused SwiGLU/dSwiGLU Kernels

Open
#1,383 0 comments 1 reaction 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.