Config-driven model knobs: decoupled attention head dim, adapter pre-norm, weights-only warm-start

[amazloumi]

Summary

• model.head_dim_override (default 0 → dim // n_heads) decouples the attention head width from the residual width, so n_heads * head_dim need not equal dim. dim % n_heads == 0 is now required only when the head width is inferred; a negative override is rejected at config time.
• head_dim stays a derived property rather than a resolved field. load_config overlays TOML and CLI values onto an already-constructed default config and carries unspecified fields forward, so a stored value would go stale as soon as dim or n_heads came from TOML.
• CrossAttentionBlock forwards head_dim to CrossAttention, but only when the override is set. Cross-attention has its own head count (vlm.cross_attention_n_heads), so under the default coupling it keeps its existing dim // n_heads width.
• adapter.pre_norm applies a norm-registry-selected norm to the vision features before mlp_2layer's first projection, exposed as ln_q; "" (default) builds no module. This is what keeps a model-specific adapter type out of core.
• checkpoint.exclude_from_loading is now honored when warm-starting from checkpoint.load_path and ignored on a resume, which always restores full state. It was previously declared but never read, and its docs described FQN prefixes rather than the DCP state keys CheckpointManager.load accepts.
• Overlaps other in-flight examples/-realignment work in kempnerforge/model/adapter.py, kempnerforge/config/adapter.py and tests/unit/test_adapter.py (ragged attentional_pool), plus CHANGELOG.md; whichever lands second will conflict in those files.

Testing

• uv run ruff check kempnerforge/ tests/ scripts/ passes

• uv run ruff format --check kempnerforge/ tests/ scripts/ passes — 166 files

• uv run pyright kempnerforge/ passes (0 errors)

• uv run pytest tests/unit/ -v --timeout=120 — 1785 passed, 3 skipped (+58 new)

• uv run pytest examples/vlm/tests -q — 33 passed; uv run pytest examples/vlm/eval/tests/unit -q — 104 passed

• uv run pytest tests/integration -q (4 GPUs) — 77 passed

• 4-GPU tests/distributed/ in full — 99 passed, 2 skipped, on each of 4 ranks (test_checkpoint.py + test_resilience.py alone: 13 passed per rank)

• Defaults unchanged: dim=256, n_heads=4, n_kv_heads=2, 4 GPUs, seed 42, pre-tokenized data, 40 steps. Two runs of clean origin/main (separate checkout, own venv) are bit-identical to each other, and this branch is bit-identical to both at all 40 logged steps — step 1 10.4375, step 40 9.375. Losses read at full precision from TensorBoard event files, not the 4-dp console line. The comparator refuses instead of reporting equality when a run is missing, empty, truncated, step-shifted or non-finite, or when the control pair itself disagrees; all 9 refusal/power cases self-tested in the same job.

• Decoupled head_dim, 4-GPU FSDP2: head_dim_override=96 with n_heads=4, dim=256 (4 * 96 = 384 != 256), 19,925,248 params — loss 10.4375 → 10.0625 over 25 steps, grad_norm 0.961 → 1.938.

• Decoupled head_dim, 2 nodes × 4 GPUs via scripts/slurm/multinode.sh, FSDP2 over dp_shard=8 — loss 10.4375 → 9.5625 over 30 steps, grad_norm 0.664 → 1.617, 1.03M tok/s, no wedge or rank divergence. A second 2-node run under tp=2 × dp_shard=4 also completes 30 steps cleanly.

• TP is unaffected by this change. Clean origin/main with a coupled head_dim is already flat under tp=2, at both model sizes, so the flat loss is pre-existing and not this PR's. 4 GPUs, 25 steps, seed 42:

  tree	head_dim	mesh	step 1	step 25	grad_norm 1 → 25
  dim=256, n_heads=4 (~20M)
  origin/main	64 coupled	dp_shard=4	10.4375	9.8125	0.859 → 1.844
  this branch	64 coupled	dp_shard=4	10.4375	9.8125	0.859 → 1.844
  origin/main (×2 runs)	64 coupled	tp=2 × dp_shard=2	10.3750	10.3750	0.006 → 0.015
  this branch	64 coupled	tp=2 × dp_shard=2	10.3750	10.3750	0.006 → 0.015
  this branch	96 decoupled	tp=2 × dp_shard=2	10.3750	10.3750	0.005 → 0.015
  dim=1024, n_heads=8 (~150M)
  origin/main	128 coupled	dp_shard=4	10.5625	10.1250	5.625 → 6.375
  this branch	192 decoupled	dp_shard=4	10.5625	10.4375	5.312 → 4.062
  origin/main	128 coupled	tp=2 × dp_shard=2	10.3750	10.3750	0.011 → 0.013
  this branch	192 decoupled	tp=2 × dp_shard=2	10.3750	10.3750	0.011 → 0.013

  Under tp=2 with the coupled config, branch and origin/main are bit-identical across all 25 steps (comparator PASS, same-code control pair bit-identical). The decoupled head dim trains normally under FSDP2 at both sizes.

• TP shards from n_heads * head_dim, not dim: apply_tensor_parallel builds a plan of bare ColwiseParallel() / RowwiseParallel() and contains no head-width arithmetic, so widths come from each nn.Linear's own out_features. Verified under gloo with dim=256, n_heads=4, n_kv_heads=2, head_dim=96, tp=2: global q_proj (384, 256) → per-rank (192, 256), k_proj (192, 256) → (96, 256), o_proj (256, 384) → (256, 192). A dim // tp shard would be (128, 256). Forward+backward gives finite non-zero grads.

• SIGTERM preemption drill, twice. (a) Cooperative path, dim=1024/12 layers, 4 GPUs, SIGTERM to the worker ranks at step 300 of 900: Shutdown requested at step 300 — saving emergency checkpoint, agent exits 0, latest -> step_300, step=300, batches_yielded=300; resume logs Resumed from step 300, 2,457,600 tokens seen (= 300 × 4 × 4 × 512, exact) and Resumed DataLoader: epoch=0, skip_batches=300, then runs to 900. Two identical runs of this config are not bit-reproducible (same-code control pair max|d| 3.125e-02), so the comparator refused an exact verdict; against that noise floor the resumed 900-step trajectory differs by max|d| 3.125e-02 — ratio 1.000, i.e. no more than two identical runs differ from each other. (b) Hard-kill path, dim=256, SIGTERM to the agent and all 4 workers at step 31 of 40: step_31 was left with zero shards and auto-resume correctly fell back to the newest complete checkpoint step_30 (skip_batches=30, 245,760 tokens, exact); that config is bit-reproducible, and the resumed 40-step trajectory is bit-identical to an uninterrupted run.

• Head-geometry consumer sweep: every consumer already reads config.head_dim, so all follow the override — Attention (q/k/v/o widths, q_norm/k_norm, SDPA scale), KVCache, MoTAttention, CrossAttention, MoMa's Attention, RoPE precompute in Transformer and pipeline_parallel, generate's cache sizing, the dense and MoE FLOP models, num_params_estimate. No site assumes n_heads * head_dim == dim: every output reshape uses -1 and o_proj maps back to dim. Tested per-arch — with the override set, every Attention/MoTAttention/CrossAttention module in a joint_decoder, cross_attention, mot and moma build reports the overridden width.

• Cross-attention regression guard: a coupled model (dim=64, n_heads=8) with cross_attention_n_heads=2 keeps head_dim=32 and q_proj (64, 64), verified equal to clean origin/main. Forwarding the model's dim // n_heads unconditionally would make it head_dim=8, q_proj (16, 64).

• TOML overlay guard: a TOML setting dim=256, n_heads=4 with no override resolves head_dim=64 on both this branch and clean origin/main. A head_dim resolved into a field in __post_init__ reads 128 here, carried from the 4096/32 default — silently wrong for configs/train/debug.toml, debug_moe.toml and hf_wikitext.toml.

• Validation cases: negative override, indivisible dim with no override (and accepted with one), unknown adapter.pre_norm — all rejected at config time, before distributed init, with the options named.

• Warm-start/resume asymmetry, both directions: exclude_from_loading=["optimizer"] reaches load(exclude_keys=...) on a load_path warm start, and is passed as None on a resume — including when a stale load_path is also set. An empty list passes None.

• pre_norm default path: no ln_q module, no ln_q state-dict key, and the projections' RNG draw identical to a build without the argument. With it set, forward equals proj2(act(proj1(ln_q(x)))) exactly, and reset_parameters restores weight=1/bias=0 after a meta-device to_empty.

• uv run pytest tests/smoke/ --smoke (4 GPUs) — 16 passed, 3 skipped, no failures; identical counts on clean origin/main. Covers dense/moe × fsdp/tp2_fsdp, plus compile, FP8, grad-accum, full AC, sigmoid router, pp2 and inline eval.

• Re-run with a pre-tokenized --data-path to unskip the last three: both TestAutoResume tests pass on this branch — RNG restore, StatefulDataLoader replay with monotonic batches_yielded, train_state.pt ownership gate, scheduler continuity, i.e. the path restore_checkpoint sits on.

• Open, under investigation: in that same --data-path run, TestRealConfigs::test_fp8_7b_config failed on this branch (18 passed, 1 failed) and passed on clean origin/main (19 passed), same 4 GPUs, same job. The config resolves byte-identically on both trees (head_dim=128, ffn=14336, 8,030,261,248 params) and the run is text-only with no resume, so none of the four knobs is on its path — but that is an argument, not evidence. Re-running the test alone and interleaved across both trees to separate a real regression from suite-position / GPU-state effects. Not claiming this one is clean until that says so.

• uv run pytest tests/e2e --e2e (4 GPUs) — 5 failed / 26 passed on this branch vs 4 failed / 27 passed on clean origin/main in the same session. The 4 shared failures (test_checkpoint_save_and_resume, test_moe_checkpoint_resume, test_pp_checkpoint_save_and_resume, test_sigterm_triggers_emergency_checkpoint) are the ones Several tests/e2e/ training tests fail on main (stale checkpoint assertions + PP/resume subprocess failures) #133 already records as pre-existing on main. The one-test delta is test_hf_dataset_single_gpu, which Several tests/e2e/ training tests fail on main (stale checkpoint assertions + PP/resume subprocess failures) #133 records as network-flaky; re-run alone on the branch it passes 3/3, so the branch and main fail the same 4 pre-existing tests.

• Re-checked against the design plan's core-admission test at completion: three general model knobs entering by general config fields (pre_norm also through the norm registry) plus one correctness fix to an inert config field, defaults unchanged, useful to any decoupled-head-geometry model rather than one experiment.

• If training loop / parallelism / optimizers changed: parallelism application order (TP → EP → FP8 → AC → FSDP2) is untouched — kempnerforge/distributed/ has no diff.

Closes #197

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/config/adapter.py	100.00% <100.00%> (ø)
kempnerforge/config/checkpoint.py	100.00% <ø> (ø)
kempnerforge/config/model.py	100.00% <100.00%> (ø)
kempnerforge/model/adapter.py	97.43% <100.00%> (+0.15%)	⬆️
kempnerforge/model/attention.py	87.87% <ø> (ø)
kempnerforge/model/cross_attention.py	100.00% <100.00%> (ø)
kempnerforge/model/transformer.py	94.65% <100.00%> (+0.02%)	⬆️
kempnerforge/training/entry.py	94.77% <100.00%> (+0.03%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.