diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index 298e691b9..2fd61a06b 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,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()} 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")