Composable per-term losses

[snimu]

What

Introduces a composable per-term loss framework. A run defines a list of named loss terms under a single shared losses key; each env selects which terms apply via enabled_losses; the trainer sums the enabled terms over one shared forward/backward. Echo ships as a preset term — {core: SFT cross-entropy, filters: [role], weight: per-token alpha} — rather than bespoke code.

The default losses = [rl] reproduces today's DPPO+KL training bit-for-bit. It too is just a preset that expands to the full loss config.

Changes vs main

• Removed [trainer.loss] (DefaultLossConfig / CustomLossConfig). Its DPPO+KL knobs (kl_tau, dppo_mask_low/high, adv_tau) now live on a [[losses]] type = "rl" term; a custom loss is a type = "custom" term.
• Added the unified, shared losses list (term types rl / echo / custom), the echo overlay capability (new — not present on main), and per-env enabled_losses + loss_overrides.
• Unchanged: the DPPO+KL math (relocated verbatim into the rl term), the opd / sft cores (still dispatched by training_mode), and — with default losses — every training run's numerics.

How it works

A loss term = { core, filters, weight }:

• Cores run trainer-side (they need logprobs): rl (DPPO+KL), echo (weighted masked NLL), custom (your import path). The opd (teacher-KL) and sft (masked NLL) cores are dispatched by training_mode and are not loss-list terms.
• Filters + weights run orchestrator-side (they need role attribution / GRPO advantage). Echo's per-role token selection and per-token alpha are resolved here and shipped to the trainer.
• Every enabled term differentiates the same shared forward → the trainer sums them into one backward.

Key behaviors:

• Shared config. losses is set once at the top level and propagated to both trainer.losses and orchestrator.losses (propagate_shared_fields). Don't set it under [trainer]/[orchestrator] separately — a validator rejects a mismatch.
• Per-env selection. enabled_losses (default None = all terms) picks which terms apply to an env; loss_overrides deep-merges per-env params into a named term (currently echo-only, e.g. a different alpha).
• Per-env primary disable. If an env's enabled_losses omits the rl/custom primary, that env's completion mask is zeroed (no RL gradient) while echo still applies — i.e. an echo-only env. The zero_advantage filter is gated so echo-only envs aren't dropped when their advantage is 0.
• Echo semantics. Echo is a per-role CE overlay on context tokens (system / user / tool / assistant). It is computed on the rollout's temperature-scaled logprobs (not a true T=1 NLL) — scale alpha to compensate. alpha = 0 keeps a token supervised with zero gradient; negative alpha suppresses tokens (anti-echo). Echo masks/weights are kept separate from loss_mask/advantages, so RL and echo may cover overlapping tokens and their gradients simply sum.
• Prompt-role echo needs prompt_attribution. Supervising system/user/tool tokens requires a renderer that emits prompt attribution; under MITO (renderer = None, which sft mode forces) those tokens silently no-op, so config validation warns. Assistant-role echo has no such requirement.
• Multi-run. Each run's orchestrator.losses is validated to equal the trainer's via a MultiRunManager config hook, so all runs served by one trainer agree on the term list.

Config surface

# top-level (shared) — propagated to trainer + orchestrator
[[losses]]
type = "rl"                  # the DPPO+KL primary; default term
kl_tau = 1e-3                # dppo_mask_low/high, adv_tau also live here

[[losses]]
type = "echo"
name = "echo"                # unique; referenced by enabled_losses
assistant = { alpha = 0.5 }  # any of system/user/assistant/tool (>=1 required)
tool = { alpha = 0.3, tool_names = ["calc"] }   # restrict to named tools

# per-env selection (orchestrator)
[[orchestrator.train.env]]
id = "math-env"
enabled_losses = ["rl", "echo"]                          # None (default) = all terms
loss_overrides = { echo = { system = { alpha = 0.01 } } } # deep-merged per env

Term fields — rl: kl_tau, dppo_mask_low, dppo_mask_high, adv_tau. echo: name, system/user/assistant/tool (each { alpha, [tool_names] }, ≥1 required), optional filter = { import_path, kwargs }. custom: name, import_path, kwargs.

Validation

Unique term names · ≤1 primary (rl/custom) term · enabled_losses ⊆ defined terms · ≤1 echo term per env · loss_overrides keys must be defined echo terms · training_mode = "rl" requires an rl/custom term · trainer.losses == orchestrator.losses · MITO + prompt-role echo warning.

Tests

tests/unit/train/rl/test_loss_terms.py (term registry, multi-term compute_loss, echo core, default==rl-only) · tests/unit/orchestrator/test_echo.py (role attribution, per-token alpha, filter) · tests/unit/orchestrator/test_batch.py (echo plumbing) · tests/unit/orchestrator/test_filters.py (echo-only zero-advantage gating) · losses propagation + validators in tests/unit/test_configs.py.

Deferred (follow-ups)

term_weights dict (multiple independent-core terms; would lift the ≤1-echo-per-env limit) · a custom weight fn · per-env trainer-side core kwargs (per-env kl_tau). Custom cores (type = "custom") and custom echo filters (filter.import_path) already ship. The IPO ablation TOMLs under the configs/private submodule are migrated separately (owner-managed).

---

Example configs

Each block shows only the loss-relevant keys; [model], [orchestrator], [trainer], [inference] are as usual (see configs/gsm8k/rl.toml for a minimal full config).

1 — Default RL (no change from main). With no losses key, the default is [rl] → DPPO+KL exactly as before. To tune the knobs (the former [trainer.loss]):

[[losses]]
type = "rl"
kl_tau = 5e-4
dppo_mask_low = 0.2
dppo_mask_high = 0.2
adv_tau = 1.0

2 — RL + assistant echo. Overlay CE on the model's own assistant tokens (no renderer requirement). Works with any renderer:

[[losses]]
type = "rl"

[[losses]]
type = "echo"
assistant = { alpha = 0.5 }

3 — RL + prompt-role echo. Supervise system/user/tool context tokens. Requires a renderer that emits prompt_attribution (i.e. not MITO):

[[losses]]
type = "rl"

[[losses]]
type = "echo"
name = "echo"
system = { alpha = 1.0 }
user   = { alpha = 0.5 }
tool   = { alpha = 0.3, tool_names = ["python", "search"] }  # only these tools

[orchestrator.renderer]
name = "auto"   # a prompt_attribution-capable renderer (not renderer = None)

4 — Multi-env: echo on some envs, not others. A shared term list; each env opts in. The math env gets RL-only; the chat env adds echo with a per-env alpha override:

[[losses]]
type = "rl"

[[losses]]
type = "echo"
name = "echo"
assistant = { alpha = 0.5 }

[[orchestrator.train.env]]
id = "math-env"
name = "math"
enabled_losses = ["rl"]            # RL only here

[[orchestrator.train.env]]
id = "chat-env"
name = "chat"
enabled_losses = ["rl", "echo"]                                   # RL + echo
loss_overrides = { echo = { assistant = { alpha = 0.1 } } }       # softer echo here

5 — Echo-only env (primary disabled). Omit the primary from enabled_losses: the env's completion mask is zeroed (no RL gradient) and only echo trains. Useful for a distillation/imitation env mixed alongside RL envs:

[[losses]]
type = "rl"

[[losses]]
type = "echo"
name = "echo"
assistant = { alpha = 1.0 }

[[orchestrator.train.env]]
id = "rl-env"
name = "math"
enabled_losses = ["rl"]

[[orchestrator.train.env]]
id = "imitate-env"
name = "demos"
enabled_losses = ["echo"]   # echo-only: no RL signal, not dropped on zero advantage

6 — Anti-echo (negative alpha). Suppress a role's tokens instead of reinforcing them:

[[losses]]
type = "rl"

[[losses]]
type = "echo"
name = "suppress-tools"
tool = { alpha = -0.2 }   # push probability mass away from tool-message tokens

7 — Echo with a custom filter. Narrow role-selected tokens further with an imported callable:

[[losses]]
type = "rl"

[[losses]]
type = "echo"
name = "echo"
user = { alpha = 0.5 }
filter = { import_path = "my_pkg.filters.drop_boilerplate", kwargs = { min_len = 8 } }

8 — Custom loss core (replaces the RL primary). Point at your own core (def core(inputs, **kwargs) -> LossOutputs). At most one primary (rl/custom) term per run:

[[losses]]
type = "custom"
name = "my-loss"
import_path = "my_pkg.losses.my_core"
kwargs = { beta = 0.1 }