Auto-select the monolithic head-loss implementation (#507)

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,45 @@ 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)
+        # 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]:
+        # 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()}

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

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
 

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")