From bad6c3ec4269fc20fe70ab0bf882ac659b7cbd92 Mon Sep 17 00:00:00 2001 From: Jari Kolehmainen Date: Wed, 16 Sep 2026 18:36:55 +0000 Subject: [PATCH] Keep Triton FLCE C=1 by default and use C=8 for GPT-OSS (#1472) ## 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 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../ops/fused_linear_cross_entropy.py | 35 ++++++++------- src/liger_kernel/ops/fused_linear_kl_div.py | 4 +- src/liger_kernel/transformers/functional.py | 2 + .../transformers/model/gpt_oss.py | 4 ++ test/ops/test_fused_linear_cross_entropy.py | 43 ++++++------------- 5 files changed, 38 insertions(+), 50 deletions(-) diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 86e55a8b..0dd376db 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -21,21 +21,10 @@ from liger_kernel.utils import infer_device # The optimal maximum block size depends on your hardware, your kernel, and your dtype MAX_FUSED_SIZE = 2048 if infer_device() == "npu" else 65536 // 2 -# Memory-budget constant for the token-chunk geometry below. The historical -# formula (inc_factor = cdiv(V, H), i.e. budget C=1) sizes chunks to a memory -# FLOOR of ~BT x H transient logits, which at LLM vocab shapes forces many tiny -# chunks and a launch-bound wall. At llama_3 shapes (V=128256, H=4096) C=1 gives -# inc_factor=32 -> chunk_size=BT/32 (256 rows @ BT=8192) = 32 loop iterations. -# Widening the transient budget to C x BT x H collapses the loop. Measured -# Blackwell B200 (full fwd+bwd, bf16): V=128256/H=4096 BT=8192 123.4ms (C=1) -> -# 24.7ms (C=16) = 5.0x; BT=4096 119.3 -> 15.1 = 7.9x; BT=1024 116.9 -> 9.07 = -# 12.9x; small V=32000 BT=8192 11.0 -> 5.4 = 2.0x (flat for C>=8). C=16 is the -# knee (gains flatten beyond it at BT>=4096). Numerics are partition-invariant -# per row: loss is bit-identical across C; input-grad differs only at the ~1e-5 -# GEMM reduction-order level; weight-grad differs only in addmm chunk-accumulation -# order (same class as the fused_linear_jsd chunk-memory idiom). Same constant on -# all devices. -_CHUNK_MEM_CONST = 16 +# Public callers default to the historical minimum transient-logits budget: +# approximately BT x H elements regardless of vocabulary size. Integrations +# with stricter reduction-order requirements may explicitly request more. +_CHUNK_MEM_CONST = 1 _TORCH_VERSION = Version(torch.__version__.split("+")[0]) _ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0") @@ -80,6 +69,7 @@ def fused_linear_cross_entropy_forward( token_grad_output=None, compute_gradients=None, weight_requires_grad=None, + chunk_mem_const: int = _CHUNK_MEM_CONST, ): assert isinstance(return_z_loss, bool), f"return_z_loss must be True or False. Got: {return_z_loss}" assert isinstance(return_token_accuracy, bool), ( @@ -88,6 +78,8 @@ def fused_linear_cross_entropy_forward( assert isinstance(return_predicted_tokens, bool), ( f"return_predicted_tokens must be True or False. Got: {return_predicted_tokens}" ) + if isinstance(chunk_mem_const, bool) or not isinstance(chunk_mem_const, int) or chunk_mem_const < 1: + raise ValueError(f"chunk_mem_const must be a positive integer. Got: {chunk_mem_const!r}") device = _input.device # ``compute_gradients`` / ``weight_requires_grad`` let backward re-enter this # function with *detached* saved tensors (whose ``.requires_grad`` is False) and @@ -107,8 +99,8 @@ def fused_linear_cross_entropy_forward( BT, H = _input.shape V = weight.shape[0] - # widen the transient logits budget to C x BT x H (C=1 is a memory floor, not a perf target) - inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H) + # Bound transient logits to C x BT x H. + inc_factor = triton.cdiv(V, chunk_mem_const * H) chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # (BT + inc_factor - 1) // inc_factor chunk_size = min(chunk_size, BT) # a single chunk covers BT when the budget allows; never exceed BT num_chunks = triton.cdiv(BT, chunk_size) # (BT + chunk_size - 1) // chunk_size @@ -494,6 +486,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): return_predicted_tokens: bool = False, ce_impl=None, ce_mode=None, + chunk_mem_const: int = _CHUNK_MEM_CONST, ): """ Fusing the last linear layer with cross-entropy loss @@ -519,6 +512,8 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): Default: False. return_token_accuracy (bool): When `return_token_accuracy` is `True`, computes and returns per-token accuracy without materializing logits. Default: `False` return_predicted_tokens (bool): When `return_predicted_tokens` is `True`, returns per-token predicted class indices (argmax) without materializing logits. Default: `False` + chunk_mem_const (int): transient-logits budget multiplier. Defaults to + `1`; larger values trade memory for fewer chunk reductions. """ # With reduction="none" the loss is per-token, so backward receives a @@ -550,6 +545,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): ce_impl=ce_impl, ce_mode=ce_mode, compute_gradients=False if ctx.defer_grads else None, + chunk_mem_const=chunk_mem_const, ) ) @@ -568,6 +564,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): # reduction="none" path re-enters forward from backward to recompute grads. ce_impl=ce_impl, ce_mode=ce_mode, + chunk_mem_const=chunk_mem_const, ) ctx.weight_requires_grad = weight.requires_grad else: @@ -612,7 +609,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_backward( grad_output, grad_input, grad_weight, grad_bias ) - return ( + gradients = ( grad_input, grad_weight, None, @@ -630,4 +627,6 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): None, # return_predicted_tokens None, # ce_impl None, # ce_mode + None, # chunk_mem_const ) + return gradients[: len(ctx.needs_input_grad)] diff --git a/src/liger_kernel/ops/fused_linear_kl_div.py b/src/liger_kernel/ops/fused_linear_kl_div.py index ca9c5e25..df7b174f 100644 --- a/src/liger_kernel/ops/fused_linear_kl_div.py +++ b/src/liger_kernel/ops/fused_linear_kl_div.py @@ -17,8 +17,8 @@ MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2 # Memory-budget constant for the token-chunk geometry below. The historical # formula (inc_factor = cdiv(V, H), i.e. budget C=1) sizes chunks to a memory # FLOOR of ~BT x H transient logits, which at LLM vocab shapes forces many tiny -# chunks and a launch-bound wall. Same constant and rationale as -# fused_linear_cross_entropy.py; see the measured numbers in its comment. +# chunks and a launch-bound wall. Fused KL retains its independently tuned +# C=16 policy; fused linear cross entropy defaults to the C=1 floor. _CHUNK_MEM_CONST = 16 _TORCH_VERSION = Version(torch.__version__.split("+")[0]) _ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0") diff --git a/src/liger_kernel/transformers/functional.py b/src/liger_kernel/transformers/functional.py index b95976d7..f089f24f 100644 --- a/src/liger_kernel/transformers/functional.py +++ b/src/liger_kernel/transformers/functional.py @@ -96,6 +96,7 @@ def liger_fused_linear_cross_entropy( use_token_scaling: bool = False, return_token_accuracy: bool = False, return_predicted_tokens: bool = False, + chunk_mem_const: int = 1, ): apply_args = ( input, @@ -116,6 +117,7 @@ def liger_fused_linear_cross_entropy( ) if getattr(LigerFusedLinearCrossEntropyFunction, "supports_inner_impl_dispatch", False): apply_args += (None, None) + apply_args += (chunk_mem_const,) loss, z_loss, token_accuracy, predicted_tokens = LigerFusedLinearCrossEntropyFunction.apply(*apply_args) if not return_z_loss and not return_token_accuracy and not return_predicted_tokens: diff --git a/src/liger_kernel/transformers/model/gpt_oss.py b/src/liger_kernel/transformers/model/gpt_oss.py index 8787fde6..df0bab19 100644 --- a/src/liger_kernel/transformers/model/gpt_oss.py +++ b/src/liger_kernel/transformers/model/gpt_oss.py @@ -11,6 +11,8 @@ from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result from liger_kernel.transformers.model.output_classes import LigerMoeCausalLMOutputWithPast +_GPT_OSS_FLCE_CHUNK_MEM_CONST = 8 + def lce_forward( self, @@ -138,6 +140,7 @@ def lce_forward( output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) + flce_chunk_mem_const = kwargs.pop("chunk_mem_const", _GPT_OSS_FLCE_CHUNK_MEM_CONST) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( @@ -175,6 +178,7 @@ def lce_forward( labels=labels, shift_labels=shift_labels, hidden_size=self.config.hidden_size, + chunk_mem_const=flce_chunk_mem_const, **kwargs, ) loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result) diff --git a/test/ops/test_fused_linear_cross_entropy.py b/test/ops/test_fused_linear_cross_entropy.py index 5ff8a73b..10b72970 100644 --- a/test/ops/test_fused_linear_cross_entropy.py +++ b/test/ops/test_fused_linear_cross_entropy.py @@ -273,12 +273,8 @@ def test_fused_linear_cross_entropy_propagates_backend_to_inner_ce(monkeypatch, # fixture, and the accumulation path itself requires a CUDA SM80+ device. # # The production token-chunk loop sizes chunks from the ``_CHUNK_MEM_CONST`` -# memory budget (``inc_factor = cdiv(V, C * H)``). At the tiny shapes used here -# the default ``C=16`` collapses to a single chunk, which would never exercise -# the per-chunk grad_weight accumulation these tests target. The class-local -# autouse ``_force_multichunk`` fixture pins that budget knob to ``C=1`` *only -# for these migrated tests* so the real production loop splits these shapes into -# multiple chunks -- the cross-backend tests above keep the production ``C=16``. +# memory budget (``inc_factor = cdiv(V, C * H)``). The public default is C=1, +# so the small shapes below naturally exercise the real multi-chunk path. # =========================================================================== # --------------------------------------------------------------------------- @@ -324,16 +320,20 @@ requires_nvidia_cuda = pytest.mark.skipif( ) -def _production_chunk_size(BT: int, V: int, H: int) -> int: - """The chunk size the production loop will pick for these dims, using the - *current* module-level ``_CHUNK_MEM_CONST`` (pinned to 1 by the class-local - autouse fixture). Mirrors the arithmetic in - ``fused_linear_cross_entropy_forward``.""" - inc_factor = triton.cdiv(V, flce_ops._CHUNK_MEM_CONST * H) +def _production_chunk_size(BT: int, V: int, H: int, chunk_mem_const: int | None = None) -> int: + """Mirror the production token-chunk geometry for the given dimensions.""" + chunk_mem_const = flce_ops._CHUNK_MEM_CONST if chunk_mem_const is None else chunk_mem_const + inc_factor = triton.cdiv(V, chunk_mem_const * H) chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) return min(chunk_size, BT) +def test_triton_flce_chunk_memory_budget(): + assert flce_ops._CHUNK_MEM_CONST == 1 + assert _production_chunk_size(BT=8192, V=128256, H=4096) == 256 + assert _production_chunk_size(BT=8192, V=128256, H=4096, chunk_mem_const=8) == 2048 + + # --------------------------------------------------------------------------- # Dispatch recorder # --------------------------------------------------------------------------- @@ -597,24 +597,7 @@ def _mk_inputs( class TestFusedLinearCrossEntropyAddmm: - """Regression suite for chunk buffers and the ``grad_weight`` addmm path. - - Grouped in a class so the ``_force_multichunk`` autouse fixture (which pins - ``_CHUNK_MEM_CONST=1``) is scoped to *only* these migrated tests and never - perturbs the production ``C=16`` used by the cross-backend tests above. - """ - - @pytest.fixture(autouse=True) - def _force_multichunk(self, monkeypatch): - """Pin the existing ``_CHUNK_MEM_CONST`` memory budget to 1 for each test. - - This is the *production* chunk-sizing knob (``inc_factor = cdiv(V, C*H)``), - not a test-only chunk API. At ``C=1`` the small shapes below split into - several token chunks, so the genuine per-chunk grad_weight accumulation - loop runs more than once -- which is exactly what these regression tests - cover. Being class-local, it does not affect any other test in the module. - """ - monkeypatch.setattr(flce_ops, "_CHUNK_MEM_CONST", 1) + """Regression suite for chunk buffers and dW accumulation.""" # ----------------------------------------------------------------------- # 1. Dispatch: the low-precision default path uses addmm(out=) and no param upcast