Imported from asckaya/opsd (AGENTS.md). Install upstream with npx skills add asckaya/opsd. Copyright stays with the author.

OPSD Implementation Notes

Source-of-truth files (read in this order):

1. ALGO.md — Algorithm spec (Part 1) + paper REF excerpts (Part 2).
2. slime_plugins/opsd/README.md — User-facing argument table, project conventions, opt-in extensions; must stay in sync with slime/utils/arguments.py.
3. CLAUDE.md / AGENTS.md — Top-level project instructions (one is a symlink to the other), including health-monitoring tables.

ALGO.md Part 1 ⊇ paper: paper is the project algorithm at N=1 with τ_k = y* (ground truth). When in doubt, prefer Part 1 semantics; prefer Part 2 numerics for paper-faithful hyperparameters.

Implementation guidelines

Algorithmic contract

• Default loss = α · L_distill = α · (1/T) Σ KL(q_mix^t ‖ p^t) (forward, full-vocab). No GRPO PG signal mixed in unless --opsd-mix-with-policy-loss is explicitly set.
• Teacher = frozen initial-policy snapshot (--opsd-freeze-teacher=True default). Teacher forwards run against the snapshot, not the live student weights.
• Rollout produces K+1 samples per prompt: 1 random student y, the other K go to the privileged candidate pool. The student is by construction never one of the τ_k seen by any teacher.
• Quality scoring uses all three ALGO.md §1.1-step-3 terms: length penalty, format penalty, Conf. Conf is rank-normalized across a sample's candidates by default (--opsd-quality-conf-norm rank); change with explicit knob if you want raw log-prob behavior.
• Diversity selection = k-center greedy on the configured distance metric. Default token_jsd (ALGO.md §1.1-step-4 recommended); unigram_jsd is opt-in and changes selection ordering relative to teacher forwards.
• Mixture-weight softmax is on raw distributions (no temperature). ALGO.md §1.1-step-6 has no temperature term; the --opsd-temperature knob is a project extension defaulting to 1.0 (identity).
• Per-position KL clamp τ=0.05 (--opsd-jsd-token-clip) applied AFTER sum-over-vocab, preserving non-negativity. The legacy per-(position, vocab-entry) clip --opsd-pointwise-kl-clip is one-sided and can drive per-token KL negative — keep it default-off.

Numerical / framework invariants

• KL must be non-negative on every step. Monitor train/opsd_kl and train/opsd_kl_clamped; they should stay equal. If they diverge, investigate clip config.
• train/opsd_w_entropy ∈ [0, 1]. Health band: 0.3–0.8. If it pins at 0 or 1, mixture is degenerate — check [opsd_dbg/mixture] logs for which logit term saturates (Δ / h / g).
• train/lr-pg_0/1 reads param_group["lr"] directly; should equal --lr after the first scheduler step. Do NOT revert to opt_param_scheduler.get_lr(...).
• Pipeline-parallel size must be 1 for OPSD: plugin.py:before_train_step_hook calls model(...) directly, bypassing pipeline scheduling. TP / DP / CP are fine.
• --opsd-freeze-teacher=True requires --enable-weights-backuper (validated in arguments.py).

Where each invariant lives

When changing OPSD behavior

1. Read ALGO.md Part 1 first. Verify your change is consistent with the algorithmic contract; if it isn't, the change must land as an opt-in knob, not a default change.
2. Update ALGO.md if the algorithmic contract changes (Part 1 for spec, Part 2 for paper REF).
3. Update slime_plugins/opsd/README.md if any user-facing argument, project convention, or opt-in extension is added / removed / re-defaulted.
4. Update slime/utils/arguments.py with matching defaults and help text. Keep the argparse help short; longer rationale goes into the README's project-conventions or opt-in tables.
5. Run the diagnostic loop: launch the 1.7B script for a few train steps, check [opsd_dbg/mixture] and [opsd_dbg/sample] stdout + train/opsd_* TB scalars match the health bands below.
6. Sweep the docs in the same commit. If a rename, a removed flag, or an algorithm change landed, this file, slime_plugins/opsd/README.md, ALGO.md, BUG.md, and any affected docstrings must reflect the new state in the same commit — not a follow-up.

Soft rules

• Prefer method-consistent semantics over convenience shortcuts when the two conflict.
• Default values must be paper-aligned where paper specifies them; opt-in extensions never change defaults.
• The mixture-weight w_entropy derivation in mixture_weights is sensitive to Python operator precedence — keep the explicit parens around the negation ((-(w * log_w).sum(0)).clamp(min=0)). Do NOT collapse to -(w * log_w).sum(0).clamp(min=0) — it parses as -((w*log_w).sum(0).clamp(min=0)), which clamps the negative sum to 0 then negates to -0.
• Keep [opsd_dbg/mixture] and [opsd_dbg/sample] rank-0 diagnostic prints until the health-band monitoring is verified stable in production. Once removed, document the removal in this file.

Coding style

• Prefer concise code. If a refactor adds layers (helper classes, wrapper functions, indirection) without a concrete reason, the original was right. Three straight lines beat three abstractions.
• Do not use getattr / hasattr as a silent fallback for values that should always be present. Missing fields, mistyped flags, or unset configs are bugs — raise on them, don't paper over with a default.
  • When the assumption is internal and a violation would be a programmer error, an assert is enough.
  • When the assumption is on user-supplied data (CLI args, configs, batch contents), raise an exception with a message that names the field.
  • For args registered in slime/utils/arguments.py with a default, read them directly (args.opsd_xxx) — the registered default already covers "not set". The only acceptable getattr(obj, "x", default) is when x is a genuinely optional knob and the default is the documented behavior.

Quick reference: paper Table 6 vs current scripts

Health monitoring (TB / stdout quick lookup)