nemo_automodel.recipes.decision.train_clef

Post-train Cloudflare Clef (Qwen3.5 backbone + joint schema head) on SystemOne records.

Clef reads a state and a schema of typed questions and returns one logit per allowed option in a single forward pass: the Qwen3.5 backbone produces final hidden states and a small transformer “joint schema head” routes evidence from the state to every question and scores the options jointly.

This recipe follows the post-training recipe Cloudflare describe for Clef - LoRA on the backbone plus a fully trained head - on top of NeMo AutoModel’s FSDP2 + PEFT stack:

The backbone is built with NeMoAutoModelForImageTextToText (FSDP2, LoRA, activation checkpointing).
The joint schema head (~128M parameters) is replicated on every rank in fp32; its gradients are averaged across the data-parallel group after backward, like the projection head in distill_bi_encoder.
The head consumes the final hidden states returned by the backbone’s public output_hidden_states interface.
The head reads LM-head rows for option tokens. The LM head is frozen and sharded, so the few rows it needs are gathered once into a LexicalTable.
mode: eval scores the model on validation/test (the before numbers); mode: train runs SFT (loss.mode: sft) or RL (loss.mode: rl). For RL the reference log-probabilities of the starting policy are pre-computed in one inference pass, so no second model copy is needed.
Each immutable stage uses the native PEFT checkpoint and DCP joint-head state. init_from initializes a new stage’s weights (SFT -> RL); optimizers and schedules start afresh. Publisher export is a separate operation.

Run with::

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

Module Contents

Classes

Name	Description
`TrainClefDecisionRecipe`	LoRA + joint-schema-head post-training for Clef decision models (see module docstring).

Functions

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

Data

MODES

logger

API

class nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe(
    cfg: nemo_automodel.recipes._typed_config.RecipeConfig | nemo_automodel.components.config.loader.ConfigNode
)

Bases: BaseRecipe

LoRA + joint-schema-head post-training for Clef decision models (see module docstring).

cfg

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._dp_group() -> torch.distributed.ProcessGroup | None

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._forward(
    batch: list[tuple[dict[str, typing.Any], nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord]]
) -> list[list[torch.Tensor]]

Run backbone + head on a batch.

Parameters:

batch

list[tuple[dict[str, Any], EncodedRecord]]

Record identities and publisher-ordered encoded question/option spans.

Returns: list[list[torch.Tensor]]

For each record, a list with one tensor of shape [options_i] (fp32 logits) per question, in Clef order.

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._load_stage(
    path: pathlib.Path
) -> None

Initialize model/head weights, leaving optimizer and schedule state fresh.

Parameters:

path

Path

Published native CLEF checkpoint or its LOWEST_VAL/LATEST pointer.

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._report(
    split: str
) -> None

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._require_stage_checkpoints() -> None

Keep training, validation selection, and final reporting on the same saved stage.

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._save_stage(
    path: pathlib.Path,
    selection_score: float | None = None
) -> None

Publish immutable weights through native checkpoint/PEFT components.

Parameters:

path

Path

Fresh checkpoint directory, shared by all data-parallel ranks.

selection_score

float | NoneDefaults to None

Higher-is-better validation score for the native LOWEST_VAL pointer.

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._score(
    items: list[tuple[dict[str, typing.Any], nemo_automodel.components.models.clef.joint_schema_head.EncodedRecord]],
    batch_size: int
) -> dict[str, list[list[float]]]

Per-record log-probabilities {record_id: [[logp per option] per question]}, gathered on all ranks.

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._shard(
    items: list,
    shuffle: bool,
    epoch: int = 0
) -> list

Equal-length per-rank shard (FSDP needs every rank to run the same number of forwards).

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._sync_head_grads() -> None

Average the replicated head’s gradients over the data-parallel group.

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._train() -> None

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe._val_score(
    metric: str
) -> float | None

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe.evaluate(
    split: str
) -> dict[str, typing.Any]

Score split and return the metrics summary (identical on every rank).

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe.run_train_validation_loop() -> None

Evaluate (mode: eval) or train (mode: train), then report validation/test metrics.

nemo_automodel.recipes.decision.train_clef.TrainClefDecisionRecipe.setup() -> None

Build the distributed backbone, the head, the encoded datasets and the optimizers.

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

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

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

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