[Feature] Move metrics composition out of model_type into experiment/family metrics bundles

[dhruvdcoder]

Is your feature request related to a problem? Please describe.

model_type currently owns too much:

1. Train/generate recipe — loss, predictor
2. All metric wiring — reported_metrics, diagnostic_metrics (while experiments separately own post_hoc_evaluator)

Because task scoring differs (STAR wants prediction EM; TinyGSM/IWSLT want post-hoc only), we clone model types (mlm_seq2seq vs mlm_tinygsm_seq2seq vs mlm_iwslt_seq2seq) that share the same loss/predictor and differ mostly in which metrics are listed. That creates a growing task × model config matrix for packaging, not for new generation modes.

We want a full cleanup: metrics composition becomes a first-class, experiment-owned surface (composable packages under each model family), and model_type shrinks to loss + predictor only. Generation modes stay coarse (unconditional / seq2seq / …).

Which component does this affect?

• Core framework (src/xlm/)
• Maintained Model family (xlm-models/)
• Task / dataset
• CLI / commands
• Documentation
• Other

Describe the solution you'd like

Ownership

Concern	Owner
loss, predictor	model_type (model + mode only)
Atomic metric / post-hoc YAML	Core: configs/lightning_train/metrics/, post_hoc_evaluator/
Run-level composition of reported + diagnostic + post-hoc	xlm-models/<family>/configs/metrics_bundle/ (or inline on experiments)
Stable log fields (generated_text, refs)	predictor to_dict contract

Core must not host task×model bundles. src/xlm/configs/lightning_train/ keeps only reusable atomics. Family-qualified scoring bundles (with model-specific update_fns) live next to that family's experiments.

Config shape

# experiment (under xlm-models/<family>/configs/experiment/)
defaults:
  - override /model_type: mlm_seq2seq          # loss + predictor only
  - /metrics_bundle: tinygsm                   # resolved from this family's config path

A bundle is a small @package _global_ config that fills reported_metrics / diagnostic_metrics / post_hoc_evaluator by composing existing atomics under metrics/ and post_hoc_evaluator/.

Reusable pieces (examples):

• Family metrics_bundle/lm_train_basic.yaml — accumulated_loss on train/val/test lm
• Family metrics_bundle/star_seq2seq.yaml — prediction EM + token accuracy with that family's update_fns
• Family metrics_bundle/tinygsm.yaml — lm loss only + Gsm8kCodeEval post-hoc
• Family metrics_bundle/iwslt14.yaml — lm loss only + iwslt14_mt post-hoc
• Optional ILM metrics_bundle/diagnostics_stopping.yaml — length loss / token CE

Not one mega-YAML for the whole repo — one composition story so opening the experiment (or its bundle) shows the full scoring picture.

Migration

1. Strip metrics from all model_type YAMLs (keep loss/predictor/Harness).
2. Add xlm-models/<family>/configs/metrics_bundle/ packages for existing experiments; point experiments at them.
3. Delete task forks that only existed for metrics (mlm_tinygsm_seq2seq, *_iwslt_seq2seq, flexmdm_tinygsm_seq2seq, …) once unused.
4. Update docs/guide/metrics.md, eval.md, adding-a-model / adding-a-task.
5. Experiments must compose a bundle explicitly (or a documented default) so forgetting metrics fails loudly or pulls only lm_train_basic.

Worked example: TinyGSM × MLM (before → after)

Today

tinygsm_mlm uses a forked model type so prediction EM is absent; post-hoc lives on the experiment:

# xlm-models/mlm/configs/experiment/tinygsm_mlm.yaml (today)
defaults:
  - override /datamodule: tinygsm_mlm
  - override /model_type: mlm_tinygsm_seq2seq   # fork of mlm_seq2seq minus EM
  - override /model: rotary_transformer_mlm
# ...
post_hoc_evaluator:
  _target_: xlm.tasks.composite_eval.CompositePostHocEvaluator
  evaluators:
    prediction:
      _target_: xlm.tasks.tinygsm.Gsm8kCodeEval

# xlm-models/mlm/configs/model_type/mlm_tinygsm_seq2seq.yaml (today)
# Same loss/predictor as mlm_seq2seq, but only accumulated_loss — no prediction EM.
defaults:
  - /metrics@reported_metrics.train.lm.accumulated_loss: accumulated_loss
  - /metrics@reported_metrics.val.lm.accumulated_loss: accumulated_loss
  - /metrics@reported_metrics.test.lm.accumulated_loss: accumulated_loss
loss:
  _target_: mlm.loss_mlm.MLMLoss
predictor:
  _target_: mlm.predictor_mlm.MLMPredictor
  # ...
reported_metrics:
  train/val/test.lm.accumulated_loss: ...

STAR uses shared mlm_seq2seq with EM baked into that model_type:

# star_easy_mlm.yaml (today)
defaults:
  - override /model_type: mlm_seq2seq   # includes prediction EM + token_accuracy

After

One seq2seq model_type for MLM; TinyGSM and STAR differ only by family-local metrics bundles.

# xlm-models/mlm/configs/model_type/mlm_seq2seq.yaml (after)
# @package _global_
lightning_module:
  _target_: xlm.harness.Harness
loss:
  _target_: mlm.loss_mlm.MLMLoss
predictor:
  _target_: mlm.predictor_mlm.MLMPredictor
  tokenizer: ${lightning_module:tokenizer}
  noise_schedule: ${lightning_module:noise_schedule}
  max_steps: ${block_size}
  max_new_tokens: ${block_size}
  top_k: 2
  top_p: null
# no reported_metrics / diagnostic_metrics / metrics defaults here
tags:
  model_type: mlm_seq2seq

# xlm-models/mlm/configs/metrics_bundle/lm_train_basic.yaml
# @package _global_
defaults:
  - /metrics@reported_metrics.train.lm.accumulated_loss: accumulated_loss
  - /metrics@reported_metrics.val.lm.accumulated_loss: accumulated_loss
  - /metrics@reported_metrics.test.lm.accumulated_loss: accumulated_loss
reported_metrics:
  train:
    lm:
      accumulated_loss:
        prefix: train/lm
        update_fn: mlm.metrics_mlm.mean_metric_update_fn
  val:
    lm:
      accumulated_loss:
        prefix: val/lm
        update_fn: mlm.metrics_mlm.mean_metric_update_fn
  test:
    lm:
      accumulated_loss:
        prefix: test/lm
        update_fn: mlm.metrics_mlm.mean_metric_update_fn

# xlm-models/mlm/configs/metrics_bundle/tinygsm.yaml
# @package _global_
defaults:
  - lm_train_basic
  # no prediction EM / token_accuracy
post_hoc_evaluator:
  _target_: xlm.tasks.composite_eval.CompositePostHocEvaluator
  evaluators:
    prediction:
      _target_: xlm.tasks.tinygsm.Gsm8kCodeEval

# xlm-models/mlm/configs/metrics_bundle/star_seq2seq.yaml
# @package _global_
defaults:
  - lm_train_basic
  - /metrics@reported_metrics.val.prediction.exact_match: seq2seq_exact_match
  - /metrics@reported_metrics.test.prediction.exact_match: seq2seq_exact_match
  - /metrics@reported_metrics.val.prediction.token_accuracy: seq2seq_token_accuracy
  - /metrics@reported_metrics.test.prediction.token_accuracy: seq2seq_token_accuracy
reported_metrics:
  val:
    prediction:
      exact_match:
        prefix: val/prediction
        update_fn: mlm.metrics_mlm.seq2seq_exact_match_update_fn
      token_accuracy:
        prefix: val/prediction
        update_fn: mlm.metrics_mlm.seq2seq_token_accuracy_update_fn
  test:
    prediction:
      exact_match:
        prefix: test/prediction
        update_fn: mlm.metrics_mlm.seq2seq_exact_match_update_fn
      token_accuracy:
        prefix: test/prediction
        update_fn: mlm.metrics_mlm.seq2seq_token_accuracy_update_fn

# xlm-models/mlm/configs/experiment/tinygsm_mlm.yaml (after)
defaults:
  - override /datamodule: tinygsm_mlm
  - override /noise_schedule: dummy
  - override /model_type: mlm_seq2seq          # shared with STAR
  - override /model: rotary_transformer_mlm
  - /metrics_bundle: tinygsm                  # scoring story lives here (mlm config path)
# log_predictions unchanged (generated_text, answer, …)
# inline post_hoc_evaluator removed — owned by the bundle

# star_easy_mlm.yaml (after) — same model_type, different bundle
defaults:
  - override /model_type: mlm_seq2seq
  - /metrics_bundle: star_seq2seq

Delete after migration: mlm_tinygsm_seq2seq.yaml (and analogous forks for FlexMDM/ILM/IWSLT).

Same pattern for IWSLT × MLM: bundle = lm_train_basic + post_hoc_evaluator → iwslt14_mt; still model_type: mlm_seq2seq. FlexMDM gets xlm-models/flexmdm/configs/metrics_bundle/tinygsm.yaml with flexmdm.metrics_flexmdm.* update_fns — still no task×model matrix in core.

Optional: skip a metrics_bundle Hydra group and compose atomics directly in each experiment's defaults:; bundles are only for reuse (e.g. STAR easy/medium/hard).

Describe alternatives you've considered

1. Partial move — only prediction reported metrics + post_hoc leave model_type; keep accumulated_loss/diagnostics on model_type. Smaller diff, but scoring story remains split.
2. Keep cloning model_types per task — status quo; does not scale.
3. Host task×model bundles under src/xlm/configs/lightning_train/metrics_bundle/ — rejected: core must not own family-qualified task scoring wiring.
4. Single monolithic metrics YAML per experiment (no reusable family bundles) — clear per run, poor reuse across STAR easy/medium/hard.

Additional context

• Triggering cases: TinyGSM / IWSLT model_type forks vs STAR *_seq2seq.
• Post-hoc is already experiment-scoped for TinyGSM/IWSLT; step metrics should join that story under family config paths.
• Step update_fns stay model-specific inside family bundles (xlm-models/mlm/..., flexmdm/..., ilm/...), not inside model_type and not in core.
• Related: Harness diagnostic vs reported intent; docs/guide/metrics.md.

Acceptance criteria

• For MLM: star_easy_mlm and tinygsm_mlm share model_type: mlm_seq2seq and differ only by metrics_bundle (under xlm-models/mlm/configs/).
• mlm_tinygsm_seq2seq (and parallel forks) removed or deprecated.
• model_type YAMLs contain no reported_metrics / diagnostic_metrics / metrics defaults.
• No task×model metrics bundles under src/xlm/configs/lightning_train/.
• Docs describe: model_type = train/generate; family metrics_bundle / experiment = all metrics + post-hoc; core = atomics only.
• Migration checklist for remaining families (ILM, FlexMDM, ARLM, …).