Closes#43, #44, #47.
#43 Multi-trainer wiring (sft/dpo/pretrain):
- New utils/v028_features.apply_v028_speed_memory(model, tcfg, base_model,
console) — single shared helper for use_cut_ce, quantization_aware="fp8",
kernel_auto_compose. Each feature degrades silently to a yellow advisory
if the underlying lib is missing; never crashes training kick-off.
- Helpers supports_v028_features(task) and warn_unsupported_features(tcfg, task)
drive both the schema validator and runtime advisories.
- soup_cli/trainer/dpo.py and trainer/pretrain.py now call the helper after
model load (post-LoRA, post-QAT) — same hook point as SFT.
- soup_cli/config/schema.py validator
_validate_v028_speed_memory_sft_only renamed
_validate_v028_speed_memory_supported_tasks; allowlist now {sft, dpo,
pretrain}. GRPO/KTO/ORPO/SimPO/IPO/PPO/RewardModel/Embedding still error
out at config-load with a precise multi-trainer message.
#44 Selective gradient-checkpoint hooks:
- New utils/gradient_ckpt.install_selective_hooks(model, granularity)
iterates ``model.named_modules()`` looking for transformer-block-shaped
names (numeric suffix on layer path), wraps each module's ``forward``
with torch.utils.checkpoint.checkpoint based on tier:
- selective: only attention sub-modules
- medium: every second transformer block
- full: every transformer block
- Returns hook count so callers can fall back to HF native checkpointing
when zero blocks were found.
#47 CrossDocCollator:
- New soup_cli/data/collators.CrossDocCollator wraps any base data
collator and injects a block-diagonal causal ``cross_doc_attn_mask``
built from per-example ``doc_lengths``. Preferred over TRL's
``packing_strategy="attention_free"`` flag (best-effort across TRL
versions). Degrades gracefully when doc_lengths is missing or shapes
don't match — base attention_mask preserved, no crash.
Tests: +16 in tests/test_part_c.py covering apply_v028_speed_memory
(no-features, cut_ce graceful failure), supports/warn helpers extension,
schema gate (dpo + pretrain accept, kto still rejects), selective hook
installation across full/medium/selective with fake transformer-shaped
models, CrossDocCollator passthrough + strip + injection. One existing
test in test_training_speed.py updated: dpo+use_cut_ce now accepted.
Known limitations:
- 7 trainers (GRPO/KTO/ORPO/SimPO/IPO/PPO/RewardModel/Embedding) still
reject v0.28.0 flags at config-load. Each is a 5-line addition once
schema validation is satisfied; tracked as a v0.33.x follow-up.
- install_selective_hooks doesn't undo earlier hooks — caller must be
re-init aware. Not an issue for the typical "construct wrapper, train,
exit" flow but worth noting.
- CrossDocCollator emits ``cross_doc_attn_mask`` (not ``attention_mask``)
to avoid clobbering the base collator's contract; downstream consumers
must read the new key explicitly. The plan calls for "preferred over
TRL's packing_strategy" which we satisfy via opt-in collation, not
silent override.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
55d1b9312c