nemo_automodel.recipes.decision.train_clm_heads

Post-train Contrastive Language Model (CLM) decision heads on SystemOne records.

CLM keeps its Qwen3-8B encoder frozen (the released heads, the clm-serve action cache and the vLLM pooling server all depend on the exact encoder embedding space), so post-training touches only the two projection heads and the logit scale. The recipe therefore runs in two phases:

Embed every unique state / option text once with NeMoAutoModelBiEncoder (last-token pooling, L2 normalisation and causal attention, matching the clm-serve input contract). Texts are sharded across all ranks and cached on disk, so re-running SFT / RL / eval never re-encodes.
Train or evaluate the heads on rank 0 (they have ~20M parameters):
- mode: eval scores heads.init_checkpoint on validation/test (the before numbers).
- mode: train with loss.mode: sft fits a custom typed-decision SFT objective: label-smoothed CE plus Brier and RPS terms. This does not implement the publisher’s soft preference-target fine-tuning objective.
- mode: train with loss.mode: rl maximises the per-option reward tables (exact expected reward or GRPO) with a KL penalty to the heads the stage started from.

The trained heads are written in the contrastive-lm checkpoint format, so CLM_CKPT=<out>/heads.pt clm-serve serves them unchanged.

Run with::

automodel examples/decision/clm_soc_triage.yaml —nproc-per-node 8 —mode train —loss.mode sft

Module Contents

Classes

Name	Description
`TrainCLMHeadsRecipe`	Embed-once, train-heads recipe for CLM decision models (see module docstring).

Functions

Name	Description
`main`	CLI entry point: `python -m nemo_automodel.recipes.decision.train_clm_heads --config <yaml>`.

Data

MODES

logger

API

class nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe(
    cfg: nemo_automodel.components.config.loader.ConfigNode
)

Bases: BaseRecipe

Embed-once, train-heads recipe for CLM decision models (see module docstring).

cfg

= RecipeConfig(cfg)

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe._cache_path(
    texts: list[str]
) -> pathlib.Path

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe._embed_all(
    texts: list[str]
) -> nemo_automodel.components.models.clm.heads.EmbeddingTable

Embed texts with the frozen encoder, sharded over all ranks, and cache the result.

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe._embedding_config() -> nemo_automodel.components.models.clm.heads.CLMEmbeddingCacheConfig

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe._forward(
    heads: nemo_automodel.components.models.clm.heads.CLMHeadPair,
    records: list[dict[str, typing.Any]]
) -> tuple[list[list[nemo_automodel.components.loss.decision_batch.QuestionLogits]], list[list[nemo_automodel.components.datasets.decision.records.QuestionTarget]]]

Score records while retaining actual question and option output order.

Parameters:

heads

CLMHeadPair

CLM state/action projection heads.

records

list[dict[str, Any]]

Records in input batch order.

Returns: list[list[QuestionLogits]]

Identified logits and targets grouped by record. Each logit values

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe._raise_cache_failure(
    failure: Exception | None,
    phase: str
) -> None

Propagate cache filesystem failures before any rank enters the next phase.

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe._report(
    heads: nemo_automodel.components.models.clm.heads.CLMHeadPair,
    split: str
) -> dict

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe._train(
    heads: nemo_automodel.components.models.clm.heads.CLMHeadPair
) -> nemo_automodel.components.models.clm.heads.CLMHeadPair

Train heads in place; return a copy of the best heads by validation train.eval_metric.

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe.evaluate(
    heads: nemo_automodel.components.models.clm.heads.CLMHeadPair,
    records: list[dict[str, typing.Any]],
    batch_size: int = 256
) -> tuple[dict, list[dict]]

Return (summary, per-question predictions) for records.

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe.run_train_validation_loop() -> None

Evaluate (eval) or train (sft/rl) the heads, then report before/after metrics.

nemo_automodel.recipes.decision.train_clm_heads.TrainCLMHeadsRecipe.setup() -> None

Initialise distributed state, load data, and make sure every text has a cached embedding.

nemo_automodel.recipes.decision.train_clm_heads.main(
    config_path: str | None = None
) -> None

CLI entry point: python -m nemo_automodel.recipes.decision.train_clm_heads --config <yaml>.

nemo_automodel.recipes.decision.train_clm_heads.MODES = ('eval', 'train')

nemo_automodel.recipes.decision.train_clm_heads.logger = logging.getLogger(__name__)