feat: integrate OPUS data selection into training pipeline

[NSR9]

Summary

• Integrates OPUS (arXiv:2602.05400) dynamic per-step data selection into the LLM training pipeline
• Each training step scores N candidate samples and trains on the best rho*N, guided by AdamW preconditioner geometry
• Implemented as a composable middleware (OpusDataSelector) with minimal changes to existing train.py and main.py
• Fully toggleable — disabled by default (opus.enabled: false), zero overhead when off

Architecture

DataLoader (2x batch) → OpusDataSelector.select_batch() → selected batch (1x) → train_epoch (unchanged)

• Scoring pass runs on unwrapped model_engine.module (bypasses DeepSpeed ZeRO gradient sync)
• Uses standard CE (not FusedLinearCE) for scoring to support ghost hook gradient capture
• Preconditioner uses DeepSpeed's official safe_get_full_optimizer_state() API
• RandomInDistributionProxyProvider auto-resets on epoch boundaries

New files (llm/src/llm/opus/)

File	Description
config.py	OpusConfig dataclass with from_dict() factory
ghost.py	GhostCollector — forward/backward hooks for activation+gradient capture
countsketch.py	CountSketchProjector — dimensionality reduction for gradient sketches
preconditioner.py	AdamWPreconditionerView — reads optimizer state for scoring geometry
selector.py	OpusSelector — scoring + Boltzmann sampling selection
proxy.py	Proxy data providers (random in-distribution, benchmark)
data_selector.py	OpusDataSelector — composable middleware encapsulating the full pipeline
distributed.py	Thin wrappers around torch.distributed

Modified files

• llm/main.py — Parse opus: config, create candidate+proxy dataloaders, init OpusDataSelector
• llm/src/llm/train.py — Add opus_selector param, call select_batch() before forward, refresh preconditioner after step
• llm/configs/config.yaml — Add opus: section (disabled by default)

Test plan

• 16 unit tests covering config, countsketch, ghost, preconditioner, selector, data_selector, public API
• All existing logger tests pass (no regressions)
• Import smoke test passes
• GPU smoke test: 10 steps with OPUS enabled on single GPU
• A/B comparison: 100 steps OPUS vs baseline loss curves

🤖 Generated with Claude Code