From d8459bd946384f49afe84816d5967c3d31a01dab Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 17 Jul 2026 13:54:05 -0400 Subject: [PATCH 1/4] Auto-select the monolithic head-loss implementation (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Losses stay configured flat: the head groups combinable losses that share a softmax into one fused kernel and picks the backend from a single head-level `loss_implementation` knob (auto/compiled/triton/per_loss, default `auto`). `auto` uses triton when a group is triton-eligible and triton is available, else the compiled path; `per_loss` keeps the unfused per-loss behavior. Grouping is by effective logits scale (one softmax serves one scale); non-combinable losses (e.g. DPO) stay standalone. The nested `type: monolithic` loss is removed as a user-facing concept — `MonolithicLossConfig` is now synthesized internally by `LanguageModelHeadConfig.get_effective_losses`. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/config.py | 52 +++++++++++- fast_llm/layers/language_model/head.py | 5 +- fast_llm/layers/language_model/loss/config.py | 19 ++++- tests/layers/test_lm_head.py | 83 ++++++++++++++----- 4 files changed, 131 insertions(+), 28 deletions(-) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index 298e691b9..8a2f1b115 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -8,7 +8,14 @@ from fast_llm.layers.common.normalization.config import NormalizationConfig from fast_llm.layers.common.peft.config import PeftConfig from fast_llm.layers.decoder.config import DecoderBlockConfig -from fast_llm.layers.language_model.loss.config import LanguageModelLossConfig, LanguageModelLossKwargs +from fast_llm.layers.language_model.loss.config import ( + CombinableLossConfig, + LanguageModelLabelEntropyLossConfig, + LanguageModelLossConfig, + LanguageModelLossKwargs, + LossImplementation, + MonolithicLossConfig, +) from fast_llm.utils import Assert if typing.TYPE_CHECKING: @@ -108,6 +115,13 @@ class LanguageModelHeadConfig(BlockConfig): "If not specified, a cross-entropy loss with respect to the targets will be used.", hint=FieldHint.core, ) + loss_implementation: LossImplementation = Field( + default=LossImplementation.auto, + desc="How to realize the losses. `auto`/`compiled`/`triton` fuse the combinable losses into a single" + " shared-softmax kernel (`auto` picks triton when available and eligible, else compiled); `per_loss`" + " runs each loss on its own softmax.", + hint=FieldHint.expert, + ) # TODO: Cleanup output_weight: ParameterConfig = Field( desc="Configuration for the LM output layer (weight). Ignored for tied embeddings", @@ -179,6 +193,42 @@ def get_layer( def _validate(self) -> None: super()._validate() assert LM_HEAD_LOSS_NAME not in self.losses + # Surface fusion/grouping errors (e.g. an ineligible `triton` set) at config time, incl. `--validate`. + self.get_effective_losses() + + def get_effective_losses(self) -> dict[str, LanguageModelLossConfig]: + # The top-level losses the head builds. Combinable losses are fused into a shared-softmax + # `MonolithicLoss` unless `loss_implementation` is `per_loss`; a single softmax serves one effective + # scale, so they are grouped by `logits_scale_factor` (the common head scale applies to all). Each + # group takes the slot of its first member; non-combinable losses (e.g. DPO) stay standalone. + losses = self.losses or {"cross_entropy": LanguageModelLabelEntropyLossConfig()} + if self.loss_implementation == LossImplementation.per_loss: + return dict(losses) + use_triton = { + LossImplementation.auto: None, + LossImplementation.compiled: False, + LossImplementation.triton: True, + }[self.loss_implementation] + scale_groups: dict[float, dict[str, LanguageModelLossConfig]] = {} + slots: list[float | tuple[str, LanguageModelLossConfig]] = [] + for name, loss in losses.items(): + if isinstance(loss, CombinableLossConfig): + if loss.logits_scale_factor not in scale_groups: + scale_groups[loss.logits_scale_factor] = {} + slots.append(loss.logits_scale_factor) + scale_groups[loss.logits_scale_factor][name] = loss + else: + slots.append((name, loss)) + named = len(scale_groups) > 1 + effective, group_index = {}, 0 + for slot in slots: + if isinstance(slot, tuple): + effective[slot[0]] = slot[1] + else: + name = f"monolithic_{group_index}" if named else "monolithic" + effective[name] = MonolithicLossConfig(losses=scale_groups[slot], use_triton=use_triton) + group_index += 1 + return effective def get_reference_models(self) -> set[str]: return {reference_model for loss in self.losses.values() for reference_model in loss.get_reference_models()} diff --git a/fast_llm/layers/language_model/head.py b/fast_llm/layers/language_model/head.py index 22c750082..30038de0f 100644 --- a/fast_llm/layers/language_model/head.py +++ b/fast_llm/layers/language_model/head.py @@ -20,7 +20,6 @@ LanguageModelHeadConfig, LanguageModelKwargs, ) -from fast_llm.layers.language_model.loss.config import LanguageModelLabelEntropyLossConfig from fast_llm.layers.language_model.loss.loss import LanguageModelLoss from fast_llm.tensor import TensorMeta from fast_llm.utils import Assert, safe_merge_dicts @@ -93,9 +92,7 @@ def __init__( lr_scale=self._lr_scale, peft=self._peft, ) - loss_configs = ( - self._config.losses if self._config.losses else {"cross_entropy": LanguageModelLabelEntropyLossConfig()} - ) + loss_configs = self._config.get_effective_losses() loss_coefficient = ( 1.0 if self._config.prediction_loss_coefficient is None diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 81af67743..a3151283e 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -222,6 +222,18 @@ class PolicyMetricsLevel(enum.StrEnum): auto = "auto" +class LossImplementation(enum.StrEnum): + # Fuse combinable losses over one shared softmax, using triton when the group is triton-eligible and + # triton is available, else the compiled path. + auto = "auto" + # Fuse combinable losses, forcing the `torch.compile` path. + compiled = "compiled" + # Fuse combinable losses, forcing triton (errors if a group has no triton kernel). + triton = "triton" + # No fusion: each loss runs its own softmax, honoring its own `use_triton`. + per_loss = "per_loss" + + @config_class() class LanguageModelPolicyGradientLossConfig(LanguageModelLossConfig): """Shared base for policy-gradient losses (GRPO, GSPO).""" @@ -275,9 +287,12 @@ def loss_class(self) -> "type[LanguageModelGSPOLoss]": return LanguageModelGSPOLoss -@config_class(dynamic_type={LanguageModelLossConfig: "monolithic"}) +@config_class() class MonolithicLossConfig(LanguageModelLossConfig): - """A composite loss that runs one vocabulary softmax and shares it across its combinable child losses.""" + """A composite loss that runs one vocabulary softmax and shares it across its combinable child losses. + + Not user-selectable: the head synthesizes it internally from a flat loss set (see + `LanguageModelHeadConfig.get_effective_losses`), so it is not registered as a dynamic `type`.""" _abstract: typing.ClassVar[bool] = False diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index a48436fe1..39e15288b 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -9,9 +9,9 @@ from fast_llm.functional.triton import triton_available from fast_llm.layers.attention.config import AttentionKwargs from fast_llm.layers.block.config import BlockKwargs -from fast_llm.layers.language_model.config import LM_HEAD_LOSS_NAME, LanguageModelKwargs +from fast_llm.layers.language_model.config import LM_HEAD_LOSS_NAME, LanguageModelHeadConfig, LanguageModelKwargs from fast_llm.layers.language_model.head import LanguageModelHead -from fast_llm.layers.language_model.loss.config import LanguageModelLossKwargs +from fast_llm.layers.language_model.loss.config import LanguageModelLossKwargs, MonolithicLossConfig from fast_llm.models.gpt.config import GPTModelConfig from fast_llm.utils import Assert from tests.layers.test_lm_losses import ( @@ -118,24 +118,16 @@ def get_config(self) -> GPTModelConfig: losses["gspo_loss"] = {"type": "gspo", "metrics": self.gspo_metrics or "none"} if isinstance(self.gspo_loss, float): losses["gspo_loss"]["weight"] = self.gspo_loss - if self.loss_implementation in ("fused", "fused_triton") and losses: - # Wrap the combinable losses in a single `monolithic` loss that shares one softmax; keep the - # child keys so the registered metric names match the per-loss configuration. `fused` pins the - # compiled path and `fused_triton` the triton path, so both are exercised in every environment. - combinable = { - name: loss - for name, loss in losses.items() - if loss["type"] in ("label", "distillation", "z_loss", "grpo", "gspo") - } - if combinable: - losses = {name: loss for name, loss in losses.items() if name not in combinable} - losses["monolithic"] = { - "type": "monolithic", - "losses": combinable, - "use_triton": self.loss_implementation == "fused_triton", - } if losses: head_config["losses"] = losses + # The head auto-groups the combinable losses into one shared-softmax kernel; the test's implementation + # label maps onto the real selector (`fused` forces the compiled backend, `fused_triton` the triton one). + head_config["loss_implementation"] = { + "per_loss": "per_loss", + "fused": "compiled", + "fused_triton": "triton", + "auto": "auto", + }[self.loss_implementation] return GPTModelConfig.from_dict( { @@ -490,9 +482,9 @@ def _add_configs(base_name: str, **kwargs): _add_configs("label_and_distillation_loss_zero_weight", label_loss=True, distillation_loss=0.0) _add_configs("distillation_loss_temperature", distillation_loss=True, distillation_temperature=2.0) -# Monolithic loss type: the combinable losses are wrapped in a single `monolithic` loss that shares one -# softmax pass; the head treats it as an ordinary loss. These configs must match their per-loss equivalents -# above (validated against the same independent reference). +# Fused paths: the head auto-groups the combinable losses into one shared-softmax kernel (`fused` forces the +# compiled backend, `fused_triton` the triton one). These configs must match their per-loss equivalents above +# (validated against the same independent reference). _add_configs("fused", loss_implementation="fused") _add_configs("fused_bfloat16", loss_implementation="fused", compute_dtype=DataType.bfloat16) _add_configs("fused_logit_scaling", loss_implementation="fused", logits_scale_factor=5.0) @@ -628,6 +620,12 @@ def _add_configs(base_name: str, **kwargs): ) ) +# `auto` backend selection: triton when available and eligible, else compiled. Exercised over the +# interpreter-safe distribution kernel (distillation, alone and with z-loss); label-family `auto` resolves to +# the explicit compiled/triton cases above. +_add_configs("auto_distillation_loss", loss_implementation="auto", distillation_loss=True) +_add_configs("auto_distillation_and_z_loss", loss_implementation="auto", distillation_loss=True, z_loss=0.5) + @pytest.mark.slow @pytest.mark.parametrize( @@ -729,3 +727,46 @@ def test_lm_head(test_config: LMHeadTestConfig): head.final_norm.weight.grad_buffer, ref_normalization_weight_grad, threshold, min_threshold ) Assert.rms_close_relative(logit_weight.grad_buffer, ref_logit_weight_grad, threshold, min_threshold) + + +def _head_config(losses: dict, loss_implementation: str | None = None) -> LanguageModelHeadConfig: + config = {"normalization": {"type": "rms_norm"}, "losses": losses} + if loss_implementation is not None: + config["loss_implementation"] = loss_implementation + return LanguageModelHeadConfig.from_dict(config) + + +def test_get_effective_losses(): + # `auto` (default): combinable losses sharing one scale fuse into a single monolithic group, backend unset. + effective = _head_config({"ce": {"type": "label"}, "z": {"type": "z_loss"}}).get_effective_losses() + Assert.eq(list(effective), ["monolithic"]) + Assert.custom(isinstance, effective["monolithic"], MonolithicLossConfig) + Assert.eq(list(effective["monolithic"].losses), ["ce", "z"]) + Assert.eq(effective["monolithic"].use_triton, None) + + # `per_loss`: unchanged flat set, no grouping. + effective = _head_config({"ce": {"type": "label"}, "z": {"type": "z_loss"}}, "per_loss").get_effective_losses() + Assert.eq(list(effective), ["ce", "z"]) + + # Distinct effective scales can't share one softmax, so they land in separate groups. + effective = _head_config( + {"ce": {"type": "label"}, "z": {"type": "z_loss", "logits_scale_factor": 2.0}} + ).get_effective_losses() + Assert.eq(list(effective), ["monolithic_0", "monolithic_1"]) + + # Non-combinable losses (DPO) stay standalone alongside a fused group. + effective = _head_config( + {"ce": {"type": "label"}, "dpo": {"type": "dpo", "reference_model": "ref"}} + ).get_effective_losses() + Assert.eq(set(effective), {"monolithic", "dpo"}) + Assert.custom(isinstance, effective["monolithic"], MonolithicLossConfig) + + # `compiled` / `triton` map to the explicit backend on every group. + Assert.eq( + _head_config({"ce": {"type": "label"}}, "compiled").get_effective_losses()["monolithic"].use_triton, False + ) + Assert.eq(_head_config({"ce": {"type": "label"}}, "triton").get_effective_losses()["monolithic"].use_triton, True) + + # `triton` on a set with no shared triton kernel is rejected at config time. + with pytest.raises(ValueError): + _head_config({"ce": {"type": "label"}, "d": {"type": "distillation", "reference_model": "t"}}, "triton") From d017fc28c6000e1fc607ed45865aa0916114fd17 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 17 Jul 2026 14:18:35 -0400 Subject: [PATCH 2/4] Review fixes: reserve the `monolithic` loss-name prefix; minor style (#507) Guard against a user loss whose name collides with a synthesized fused-group key (which would silently drop an entry), mirroring the `lm_head_loss` reserved name. Split a compound initializer and trim a redundant comment. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/config.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index 8a2f1b115..2fd61a06b 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -193,7 +193,9 @@ def get_layer( def _validate(self) -> None: super()._validate() assert LM_HEAD_LOSS_NAME not in self.losses - # Surface fusion/grouping errors (e.g. an ineligible `triton` set) at config time, incl. `--validate`. + # `get_effective_losses` synthesizes fused-group names with a `monolithic` prefix; keep it reserved. + assert not any(name.startswith("monolithic") for name in self.losses) + # Surface fusion/grouping errors (e.g. an ineligible `triton` set) at config time. self.get_effective_losses() def get_effective_losses(self) -> dict[str, LanguageModelLossConfig]: @@ -220,7 +222,8 @@ def get_effective_losses(self) -> dict[str, LanguageModelLossConfig]: else: slots.append((name, loss)) named = len(scale_groups) > 1 - effective, group_index = {}, 0 + effective = {} + group_index = 0 for slot in slots: if isinstance(slot, tuple): effective[slot[0]] = slot[1] From 767e10d9faab8cc0d74d4fc360cf49ccb9a16103 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 20 Jul 2026 16:52:19 -0400 Subject: [PATCH 3/4] Disable IS-ratio clipping in the GRPO/GSPO model tests (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The distributed-consistency model tests compare gradients across parallel layouts. GRPO/GSPO clip the importance-sampling ratio at `1 ± epsilon`, and that clip bound is a discontinuity: tiny floating-point differences between layouts (single-GPU vs pipeline/tensor-parallel reduction order) flip a token or segment across the boundary, changing its gradient contribution discontinuously and intermittently blowing past the comparison tolerance — the same failure mode as MoE top-k routing. Disable the clip (large epsilons) for the `llama_grpo`/`llama_gspo` model configs so the comparison is smooth. The clip itself stays covered by the single-GPU `test_lm_losses`/`test_lm_head` tests, which don't compare across layouts. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/utils/model_configs.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/utils/model_configs.py b/tests/utils/model_configs.py index 63b02e66a..b61a7a43f 100644 --- a/tests/utils/model_configs.py +++ b/tests/utils/model_configs.py @@ -710,7 +710,14 @@ def update_and_add_testing_config( "llama_grpo", # Metrics default to `auto` (→ `basic` when pipeline_parallel == 1); pin `none` so this loss-mechanics # config doesn't register the metric family (metrics are covered by the subtests in `test_lm_losses`). - updates={("model", "base_model", "head", "losses"): {"grpo": {"type": "grpo", "metrics": "none"}}}, + # Disable IS-ratio clipping (huge epsilons): the clip bound is a discontinuity that tiny floating-point + # differences between parallel layouts flip, making the distributed-vs-single gradient comparison flaky + # (same failure mode as MoE top-k routing). The clip itself is covered by `test_lm_losses`. + updates={ + ("model", "base_model", "head", "losses"): { + "grpo": {"type": "grpo", "metrics": "none", "epsilon_low": 1e6, "epsilon_high": 1e6} + } + }, groups={ ModelTestingGroup.basic: ModelTestingGroupAction.normal, ModelTestingGroup.checkpoint: ModelTestingGroupAction.not_implemented, @@ -727,8 +734,13 @@ def update_and_add_testing_config( "llama", "llama_gspo", # Metrics default to `auto`; pin `none` for the same reason as `llama_grpo` (keep this config focused - # on loss mechanics; metrics are covered by `test_lm_losses`). - updates={("model", "base_model", "head", "losses"): {"gspo": {"type": "gspo", "metrics": "none"}}}, + # on loss mechanics; metrics are covered by `test_lm_losses`). Clipping disabled (huge epsilons) for the + # same reason as `llama_grpo`: the clip discontinuity makes the distributed-vs-single comparison flaky. + updates={ + ("model", "base_model", "head", "losses"): { + "gspo": {"type": "gspo", "metrics": "none", "epsilon_low": 1e6, "epsilon_high": 1e6} + } + }, # `ms*` (micro_batch_splits>1) and `ce*` (cross_entropy_splits>1) both produce # multiple kernel calls per micro-batch; GSPO's per-document geometric mean can't be # reconstructed from per-fragment `exp(mean)` values, so we skip these variants. From 3220bae61866039782c19d5f5e4086c2ccb6d492 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 20 Jul 2026 17:19:54 -0400 Subject: [PATCH 4/4] Review fixes: narrow reserved-name guard, rename local, trim enum comments (#507) - Reserve only the synthesized fused-group names (monolithic / monolithic_) instead of the whole `monolithic` prefix, so flat loss names like `monolithic_kd` are allowed. - Rename `named` -> `multiple_groups` and annotate the `effective` local. - Drop redundant per-member comments on `LossImplementation`, keeping only the non-obvious `auto` fallback note. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/config.py | 13 ++++++++----- fast_llm/layers/language_model/loss/config.py | 6 +----- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index 2fd61a06b..9ed80515b 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -193,8 +193,11 @@ def get_layer( def _validate(self) -> None: super()._validate() assert LM_HEAD_LOSS_NAME not in self.losses - # `get_effective_losses` synthesizes fused-group names with a `monolithic` prefix; keep it reserved. - assert not any(name.startswith("monolithic") for name in self.losses) + # `get_effective_losses` synthesizes fused-group names `monolithic` / `monolithic_`; keep them reserved. + assert not any( + name == "monolithic" or (name.startswith("monolithic_") and name.removeprefix("monolithic_").isdigit()) + for name in self.losses + ) # Surface fusion/grouping errors (e.g. an ineligible `triton` set) at config time. self.get_effective_losses() @@ -221,14 +224,14 @@ def get_effective_losses(self) -> dict[str, LanguageModelLossConfig]: scale_groups[loss.logits_scale_factor][name] = loss else: slots.append((name, loss)) - named = len(scale_groups) > 1 - effective = {} + multiple_groups = len(scale_groups) > 1 + effective: dict[str, LanguageModelLossConfig] = {} group_index = 0 for slot in slots: if isinstance(slot, tuple): effective[slot[0]] = slot[1] else: - name = f"monolithic_{group_index}" if named else "monolithic" + name = f"monolithic_{group_index}" if multiple_groups else "monolithic" effective[name] = MonolithicLossConfig(losses=scale_groups[slot], use_triton=use_triton) group_index += 1 return effective diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index a3151283e..27b73c715 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -223,14 +223,10 @@ class PolicyMetricsLevel(enum.StrEnum): class LossImplementation(enum.StrEnum): - # Fuse combinable losses over one shared softmax, using triton when the group is triton-eligible and - # triton is available, else the compiled path. + # `auto` picks triton when the group is triton-eligible and triton is available, else the compiled path. auto = "auto" - # Fuse combinable losses, forcing the `torch.compile` path. compiled = "compiled" - # Fuse combinable losses, forcing triton (errors if a group has no triton kernel). triton = "triton" - # No fusion: each loss runs its own softmax, honoring its own `use_triton`. per_loss = "per_loss"