mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
bad6c3ec42
## Summary - keep the public Triton fused linear cross entropy chunk-memory default at `C=1` - add an explicit per-call `chunk_mem_const` override while preserving backward compatibility for existing autograd callers - use `C=8` only in the GPT-OSS patched training path, where C=1 causes a deterministic BF16 convergence mismatch - keep every other model and direct/public FLCE caller at C=1 - retain the reusable-logits, direct-dX, and direct-dW optimizations from #1468 - lock the default and override chunk geometry in regression coverage ## Why the override is scoped The exact mini GPT-OSS convergence case fails with C=1 and passes with C=4, C=8, and C=16. However, setting C=8 globally causes the BF16 Qwen3.5-MoE convergence case to fail. The model-scoped override preserves both trajectories: - public/default FLCE: C=1 - GPT-OSS patched loss: C=8 - fused linear KL: unchanged independently tuned policy ## H200 memory and latency BF16, `BT=8192`, `H=4096`, `V=128256`: | Policy | Chunk size | Chunks | Time | Incremental peak allocation | |---|---:|---:|---:|---:| | Public default C=1 | 256 | 32 | 53.039 ms | 1128.66 MiB | | GPT-OSS override C=8 | 2048 | 4 | 37.748 ms | 1568.04 MiB | | Previous global C=16 | 4096 | 2 | 35.074 ms | 2068.04 MiB | The default therefore retains the minimum memory budget, while GPT-OSS uses the smallest requested release override with stable convergence. ## Testing done On H200 with PyTorch 2.9.1.6+cu129 and Triton 3.6.0: - **144 passed**: FLCE ops suite - **144 passed**: FLCE ops suite with `LIGER_KERNEL_STRICT=1` - **141 passed**: transformer-facing FLCE suite - **71 passed, 3 dependency-gated skips**: FP32 + BF16 text-model convergence suites - exact BF16 mini GPT-OSS convergence passes with the scoped C=8 override - BF16 mini Qwen3.5-MoE convergence passes because its default remains C=1 - Ruff lint/format and `git diff --check` pass The broader six-file convergence attempt also exposed four PaliGemma multimodal failures in the ad hoc Transformers 5.15 environment. Those configurations explicitly set `fused_linear_cross_entropy=False`, so they are unrelated to this change. No owned B300 pod was available for an additional Blackwell run. --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]>