diff --git a/.gitmodules b/.gitmodules
new file mode 100644
index 0000000..dc547d3
--- /dev/null
+++ b/.gitmodules
@@ -0,0 +1,4 @@
+[submodule "slime"]
+ path = slime
+ url = git@github.com:HJSang/slime.git
+ branch = crisp
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 0000000..fa67a7b
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,92 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## What this is
+
+CRISP / OPSD (On-Policy Self-Distillation) trains reasoning LLMs to think more concisely by distilling their own concise behavior back into themselves. There is **one teacher and one student that are the same base model**, differing only by prompt:
+
+- **Student** generates a rollout from the *question-only* prompt (`sft_prompt`).
+- **Teacher** (frozen ref model) re-scores those *same response tokens* under a longer *conciseness* prompt (`sd_prompt`, e.g. "Solve concisely").
+- Training minimizes per-token **reverse-KL** (default) or **JSD** between student and teacher logits on the response positions, over **ALL** rollouts (no correctness filtering — verification is metrics-only).
+
+No ground-truth answers, token budgets, or difficulty estimators are used in the loss.
+
+## Commands
+
+All paths below are relative to `workspace/`.
+
+```bash
+# --- Environment (assumes a cluster image with verl/torch/sglang preinstalled) ---
+bash scripts/sft/setup_sft.sh # adds the two extras (math-verify, tensordict) + prints versions
+# Fresh install instead: pip install -r requirements.txt (pins verl@deec5d02, torch 2.9.1, sglang 0.5.9)
+
+# --- Unit tests (CPU, fast) ---
+pytest src/self_distill_hybrid/test_opsd_jsd.py # JSD/reverse-KL/entropy loss correctness + chunk invariance
+pytest src/self_distill_hybrid/test_opsd_jsd.py::TestJSDLossEquivalence::test_loss_values_match # single test
+
+# --- Data pipeline (run in order) ---
+cd src/data
+python process_eval_data.py --data_dir ../../data --output_dir ../../data/processed # DAPO train/val split + MATH-500/AIME parquets
+python prepare_length_prune_data.py batch \ # builds 4 teacher-strength variants (concise/20/50/80pct)
+ --input-parquet ../../data/DAPO-Math-17k-dedup/distinct-prompts-with-rewards.parquet \
+ --output-root ../../data
+# `single` subcommand builds one variant: --teacher-style {concise,percent_reduce} --percent-reduce N
+
+# --- Training (8x H100/H200 80GB). See README "Quick Start" for the full env-var block ---
+MODEL_PATH=/path/to/Qwen3-8B \
+SD_PROMPTS_PATH=./workspace/data/length_prune_concise/self_distill_prompts.parquet \
+OPSD_LOSS_TYPE=reverse_kl TEACHER_UPDATE_FREQ=50 ... \
+bash workspace/scripts/sft/train_opsd.sh
+
+# --- Checkpoint export (FSDP shards -> HF format) ---
+bash scripts/sft/merge_checkpoints.sh
+```
+
+`train_opsd.sh` is the single entry point: it runs `process_eval_data.py` on-cluster, auto-detects val parquets, then launches `python -m self_distill_hybrid.main_opsd` with a long list of Hydra overrides. Tune behavior via the env vars it reads (documented in the script header and the README "Key Hyperparameters" table), not by editing the python.
+
+## Architecture
+
+Built on a **pinned fork of [VERL](https://github.com/volcengine/verl)** (`deec5d02`, between v0.7.0 and v0.7.1) using its **HybridEngine**: sglang for generation and FSDP for training are *colocated* on the same GPUs, with weights synced between them each step. The repo adds OPSD-specific subclasses rather than patching verl in-tree; the custom math scorer is loaded via verl's `custom_reward_function.path`, not an overlay.
+
+The training stack (`workspace/src/self_distill_hybrid/`) layers on verl:
+
+- **`main_opsd.py`** — Hydra/Ray entry point. Maps the worker to `Role.ActorRolloutRef` (not `ActorRollout`) **specifically so the ref model is materialized** — that ref model *is* the frozen teacher. Builds the dataset, an optional generation-based val dataset + reward manager, then runs the trainer.
+- **`opsd_trainer.py` (`OPSDTrainer`)** — the per-step loop in `fit()`:
+ 1. **Generate**: swaps `raw_prompt` (initially the teacher `sd_prompt`) to `sft_prompt` (question-only) and generates student rollouts via sglang; then `sleep_replicas()` to free rollout GPU memory for the backward pass.
+ 2. **Verify**: `verify_batch` scores correctness for metrics only — results never filter the batch.
+ 3. **Train**: `build_opsd_batch` + dispatch to `update_opsd`.
+ 4. **Sync**: `CheckpointEngineManager.update_weights()` pushes fresh student weights into sglang so the next generation is on-policy. Skipping this would break the on-policy assumption.
+- **`opsd_worker.py` (`OPSDWorker`)** — subclasses verl's `AsyncActorRolloutRefWorker`. Adds two registered methods:
+ - `update_opsd`: two forward passes per micro-batch — teacher (`ref_module_fsdp`, no-grad) and student (`actor_module_fsdp`, with-grad) — then the divergence loss on response logits. Has padded and unpadded (flash-attn varlen) logit paths, and `_liger` loss variants (logsumexp mixture + progressive teacher-chunk freeing) for lower peak memory.
+ - `update_teacher`: optional hard-copy of student shards into the ref model every `TEACHER_UPDATE_FREQ` steps (progressive compression). Copies FSDP shards **in lock-step without `summon_full_params`** — full materialization OOMs Qwen3-14B on 80GB.
+- **`sd_dataset.py` / `sd_verifier.py`** — dataset loading + the batch builders and the dual-path math verifier.
+
+### The load-bearing invariant
+
+Teacher and student sequences share the **same response tokens** but have **different (different-length) prompts**. The reverse-KL/JSD loss aligns teacher vs. student response logits **by position**, so per-sample response-token counts must match exactly between the two sides. A one-sided truncation would silently misalign every subsequent sample in the flattened batch. This is enforced in two places, and **must stay enforced** when touching the data path:
+
+1. `SelfDistillDataset` (`sd_dataset.py`) drops rows whose teacher prompt exceeds `data.max_prompt_length`.
+2. `build_opsd_batch` / `_tokenize_sequence` (`sd_verifier.py`) **refuse** (return `None` → drop the pair) any sequence over `opsd.sft_max_length` rather than truncating.
+
+The `assert teacher_logits.shape[0] == student_logits.shape[0]` in `opsd_worker._opsd_training_step` is a defensive cross-check, not the primary guarantee. If a large fraction of samples is being silently dropped, **raise `SFT_MAX_LENGTH`** — don't relax the drop.
+
+### Data columns (parquet → batch)
+
+`prepare_length_prune_data.py` emits per row:
+- **`sft_prompt`** = student prompt = original DAPO-Math question, unchanged → used for **generation** and as the **student** logit prompt.
+- **`sd_prompt`** = teacher prompt = question + conciseness instruction (from `config/prompts.json`) → used only as the **teacher** logit prompt.
+- `ground_truth` (verification), `question` (logging), `teacher_solution` (empty for length-pruning).
+
+Prompt templates live in `workspace/config/prompts.json` (`length_prune_teacher`, `length_prune_teacher_percent_reduce`, `opsd_qwen3_student`, etc.). Qwen3 thinking mode is pinned via `data.apply_chat_template_kwargs.enable_thinking: true` in the config.
+
+### Verification / scoring
+
+Both the in-trainer verifier (`sd_verifier.verify_response`) and the val reward fn (`src/rewards/dual_path_math_verify.py`) use **dual-path math_verify**: (1) regex-extract `Answer: X` → wrap in `\boxed{}` → sympy symbolic-equivalence; (2) fallback math_verify over the full response to catch in-prose `\boxed{...}` from Qwen3 thinking mode. Correct iff *either* path matches. The reward fn dispatches only `math*`/`aime*` data_sources to this scorer and delegates everything else to verl's `default_compute_score`.
+
+## Gotchas
+
+- **Do NOT set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments`** — sglang's `torch_memory_saver` refuses to initialize under that allocator and crashes rollout init (see the comment block atop `train_opsd.sh`). Address reverse-KL/JSD OOMs by lowering the `chunk_size` in the loss functions in `opsd_worker.py` instead, or halving `MICRO_BATCH_SIZE`.
+- **DP padding uses `dp_world = total_gpus // ulysses_sp`, not `total_gpus`** (TP is rollout-only and doesn't collapse the actor's DP axis). The OPSD batch is padded to a multiple of `dp_world` before dispatch; getting this wrong yields `AssertionError("only support equal chunk")`.
+- `execution-configs/` holds the per-ablation hyperparameter sets (Qwen3-8B/14B × teacher-update-freq `tu1`..`tu100` × compression strength). These are the canonical reproductions — prefer them over hand-rolled env vars.
+- `*.pyc` under `src/.../__pycache__/` for cpython-312 **and** 314 are checked in; ignore them.
diff --git a/METHOD.md b/METHOD.md
new file mode 100644
index 0000000..645d1c8
--- /dev/null
+++ b/METHOD.md
@@ -0,0 +1,211 @@
+# CRISP — Method & Algorithm (with code-vs-paper verification)
+
+This document describes the CRISP / OPSD reasoning-compression algorithm **as implemented in `workspace/`** and verifies it against the paper *CRISP: Compressed Reasoning via Iterative Self-Policy Distillation* (`crisp_compressed_reasoning_via_iterative_self_policy_distillation.pdf`).
+
+**Verdict up front:** the implementation faithfully matches the paper. The default loss, the on-policy training loop, the teacher-refresh mechanism, the prompt templates, and the training hyperparameters all correspond exactly. The code additionally implements a few options the paper does not center on (JSD loss, memory-efficient "liger" loss variants, a word-limit teacher template, and an OPSD-with-reference-solution teacher) — these are supersets, not contradictions. Details and the exact correspondence are below.
+
+---
+
+## 1. The idea in one paragraph
+
+A reasoning model already knows how to be concise; it just needs permission. CRISP takes **one model** and conditions it two ways via the prompt:
+
+- **Student** `π_θ(· | x)` — the original math prompt `x`, no special instruction.
+- **Teacher** `π_θ̃(· | x, c)` — the same problem prefixed with a **conciseness instruction** `c` ("Solve concisely, be direct…").
+
+Training generates **student** rollouts and minimizes the **per-token reverse KL** between the student and the (stop-gradient) teacher distribution on the student's own tokens. No ground-truth answers, no token budgets, no reward model, no difficulty estimator. The conciseness signal emerges from the KL objective and adapts to problem difficulty automatically.
+
+---
+
+## 2. Problem formulation (paper §3.1)
+
+A reasoning model `π_θ` maps input `x` to output `y = (r, a)` — a reasoning trace `r` inside `…` followed by an answer `a`. Goal: learn `θ*` that produces **shorter traces while maintaining accuracy**. The student gets the original DAPO-17K prompt `x`; the teacher gets the same prompt prefixed with conciseness instruction `c`.
+
+---
+
+## 3. Training objective (paper §3.2, Eq. 1)
+
+CRISP minimizes per-token **reverse KL** between student and a **stop-gradient** teacher, on **student-generated** rollouts:
+
+```
+L(θ) = E_{x~D, y~π_θ(·|x)} [ Σ_t D_KL( π_θ(· | x, y_ Note: the YAML default `opsd.teacher_update_freq: 0` is the frozen baseline; the README launch examples and execution-configs set `M=50` for the paper's main results.
+
+---
+
+## 5. Training algorithm (paper Algorithm 1)
+
+```
+Input: model π_θ, dataset D = {x_i}, conciseness instruction c,
+ learning rate η, teacher update interval M
+Output: compressed model π_θ*
+
+Initialize teacher: θ̃ ← θ_0
+for each training step k = 1, 2, …:
+ if k mod M == 0:
+ θ̃ ← θ # periodic refresh
+ sample batch {x_1, …, x_B} ~ D
+ for each x_i in batch:
+ y_i ~ π_θ(· | x_i) # student rollout (on-policy)
+ for each token position t = 1 … |y_i|:
+ q_t ← π_θ (· | x_i, y_{i, Solve the following math problem step by step. The last line of your response should be of the form Answer: $Answer … Remember to put your answer on its own line after "Answer:".
+
+**Teacher** `π_θ̃(·|x,c)` — conciseness instruction `c` (`prompts.json → length_prune_teacher`):
+
+> **Solve the following math problem concisely and correctly. Be direct — avoid unnecessary elaboration, redundant steps, or restating the problem. Focus only on the key reasoning steps needed to reach the answer.** The last line of your response should be of the form Answer: $Answer … Remember to put your answer on its own line after "Answer:".
+
+**Soft-budget teacher** `π_θ̃(·|x,c_p)` (ablation, `prompts.json → length_prune_teacher_percent_reduce`):
+
+> Solve the following math problem correctly **using {p}% fewer tokens than you normally would.** Be more concise — cut unnecessary elaboration, redundant steps, and verbose explanations while preserving correctness. …
+
+The ablation (paper Table 4: qualitative "be concise" beats explicit `p ∈ {20,50,80}%` targets) is reproduced by the `length_prune_concise` vs `length_prune_{20,50,80}pct` data variants and the matching `execution-configs/*-{20,50,80}pct.json`. ✅
+
+---
+
+## 7. The load-bearing implementation invariant
+
+Teacher and student sequences share the **same response tokens** but have **different-length prompts** (teacher prompt is longer). The reverse-KL loss aligns teacher vs. student response logits **by position**, so per-sample response-token counts must match exactly. The code enforces this by **dropping** (never truncating) any over-length pair:
+
+- `sd_dataset.SelfDistillDataset` drops rows whose teacher prompt exceeds `data.max_prompt_length`.
+- `sd_verifier.build_opsd_batch` / `_tokenize_sequence` return `None` (drop the pair) if either side exceeds `opsd.sft_max_length`.
+- `opsd_worker._opsd_training_step` asserts `teacher_logits.shape[0] == student_logits.shape[0]` as a defensive cross-check.
+
+This is an engineering detail not in the paper's math, but it is exactly what makes the per-token KL of Eq. 1 well-defined across the two differently-prompted sequences.
+
+---
+
+## 8. Training hyperparameters (paper §5.1 & Appendix F) — match ✅
+
+| Setting | Paper | Code (README launch / execution-configs) |
+|---|---|---|
+| Models | Qwen3-8B, Qwen3-14B | `MODEL_PATH=…/Qwen3-8B` / `…/Qwen3-14B` |
+| Data | ~13,600 DAPO-Math-17k, **no GT in loss** | `DAPO-Math-17k-dedup`, train on all rollouts |
+| Epochs | 1 (~100 steps to converge; step 100 = default ckpt) | `TOTAL_EPOCHS=1` |
+| Learning rate | 1e-6 | `LEARNING_RATE=1e-6` |
+| Batch size | 32 | `TRAIN_BATCH_SIZE=32` |
+| Teacher update `M` | 50 | `TEACHER_UPDATE_FREQ=50` |
+| Rollout | single (n=1), temperature 1.0, max 8,192 tokens | `rollout.n=1`, `SD_TEMPERATURE=1.0`, `SD_MAX_TOKENS=8192` |
+| Eval token budgets | 8,192 and 30,000 | `VAL_MAX_TOKENS=30000` (and 8K variant) |
+| Eval sampling | mean@8 | `val_kwargs.n=8` |
+| Hardware | 1 node × 8 H200 | `N_GPUS=8`, `nnodes=1` |
+| Framework | verl HybridEngine + sglang | verl `deec5d02`, `rollout.name=sglang` |
+| Parallelism | FSDP + Ulysses-SP deg 4 (train), TP deg 2 (infer) | `ULYSSES_SP_SIZE=4`, `TP_SIZE=2` |
+| Precision | bfloat16, gradient checkpointing, CPU offload | `autocast(bfloat16)`, `enable_gradient_checkpointing=true`, `param_offload`/`optimizer_offload=true` |
+
+Benchmarks (MATH-500, AIME 2024, AIME 2025) and the dual-path math grading are implemented in `process_eval_data.py` and `rewards/dual_path_math_verify.py` (mirrors veRL's `math_dapo` grading, with an added `Answer:`-line extraction path). ✅
+
+---
+
+## 9. Emergent properties (paper §5.3–5.5) — observable in code
+
+These are not enforced by the loss; they emerge. The code logs the signals needed to reproduce the paper's findings:
+
+- **Difficulty-adaptive compression** (~1.6× more on easy vs. hard): emerges from the KL objective; the trainer logs per-`data_source` response-token counts (`val/{ds}/avg_response_tokens`).
+- **Entropy preservation** (Finding 3, the central contrast with RL length penalties): `opsd_worker` computes and logs `opsd/student_entropy`, `opsd/teacher_entropy`, `opsd/entropy_diff` every step.
+- **Training-time accuracy rises with no correctness reward** (Fig. 2): `sd/accuracy` is logged (metrics-only verification).
+
+---
+
+## 10. Where the code goes beyond the paper (supersets, not conflicts)
+
+| Code capability | Status vs. paper |
+|---|---|
+| `OPSD_LOSS_TYPE=jsd` (+ `OPSD_BETA`) — symmetric Jensen-Shannon divergence | Extra option. The paper's CRISP method is **reverse KL**; it discusses forward KL only as a negative comparison (Appendix I). JSD is an available alternative, not the paper's method. |
+| `_compute_*_liger` loss variants (logsumexp mixture + progressive teacher-chunk freeing) | Memory-engineering only; numerically equivalent (verified in `test_opsd_jsd.py`). Toggled by `USE_LIGER`. |
+| `prompts.json → length_prune_teacher_word_limit` | A word-cap teacher variant; paper only reports the percentage soft-budget. |
+| `prompts.json → opsd_qwen3_teacher` (includes a **Reference Solution**) and `self_distill*` templates | The **OPSD-with-ground-truth** teacher family (Zhao et al. 2026 baseline / classic self-distillation), where the teacher gets privileged info. **CRISP uses the conciseness teacher only** (`length_prune_*`) and is explicitly the no-GT variant. The repo is historically named "OPSD" but the paper's method is the `length_prune_concise` config. |
+| DP-padding to `total_gpus // ulysses_sp`, `expandable_segments` ban, chunked losses | Pure infra correctness/memory; orthogonal to the algorithm. |
+
+---
+
+## 11. Summary
+
+| Claim in paper | Verified in code? |
+|---|---|
+| One model, two prompt conditionings (student `x`, teacher `x,c`) | ✅ `sft_prompt` vs `sd_prompt` |
+| Per-token **reverse** KL `D_KL(student‖teacher)` | ✅ `_compute_reverse_kl_loss` |
+| Stop-gradient teacher (no grad through teacher) | ✅ `torch.no_grad()` on `ref_module_fsdp` |
+| On-policy (train on student's own rollouts) | ✅ generation from `sft_prompt`, then train on those tokens |
+| No ground-truth answers / no reward in loss | ✅ verification is metrics-only, never filters |
+| Periodic teacher refresh `θ̃ ← θ` every `M=50` | ✅ `update_teacher` @ `TEACHER_UPDATE_FREQ` |
+| Loss normalized by `|y|` | ✅ `kl_sum / n_tokens` |
+| Exact student/teacher/soft-budget prompts | ✅ byte-identical in `prompts.json` |
+| Hyperparameters (lr, batch, rollout, hardware, parallelism) | ✅ README / execution-configs |
+
+**The `workspace/` implementation is consistent with the paper's described method and Algorithm 1.** The only things to keep in mind: the algorithm's "default" in the paper is *reverse KL + M=50*, whereas the bare YAML default is *reverse KL + M=0 (frozen teacher)*; use the README launch block or `execution-configs/` to reproduce the paper's numbers.
diff --git a/SLIME_DESIGN.md b/SLIME_DESIGN.md
new file mode 100644
index 0000000..2b16ac1
--- /dev/null
+++ b/SLIME_DESIGN.md
@@ -0,0 +1,164 @@
+# slime Design Notes & CRISP Integration Plan
+
+Notes on the [THUDM/slime](https://github.com/THUDM/slime) framework (submodule at `slime/`, pinned `ee72ab5`, v0.3.0-25), focused on the three subsystems that matter for porting CRISP: **rollout generation**, the **training engine**, and **weight update** — plus a concrete plan for running CRISP on slime.
+
+**One-liner:** Megatron-Core for training + SGLang for inference, orchestrated by Ray. Nearly everything is pluggable via `--*-function-path` hooks (rollout fn, reward fn, data source, advantage fn, loss hooks).
+
+---
+
+## 0. Top-level control flow
+
+Two drivers at the repo root:
+
+```
+train.py — synchronous loop; supports colocate (train+rollout time-slice the same GPUs)
+train_async.py — disaggregated only; one-step pipelining (rollout k+1 generates while step k trains)
+```
+
+Per `rollout_id`, `train.py` does:
+
+```python
+rollout_data_ref = rollout_manager.generate(rollout_id) # sglang rollouts → per-DP-rank ray Boxes
+actor_model.async_train(rollout_id, rollout_data_ref) # Megatron fwd/bwd
+actor_model.update_weights() # push new weights into sglang engines
+# (+ optional offload/onload dance in colocate mode, periodic eval/save)
+```
+
+`train_async.py` overlaps generation and training, syncing weights every `--update-weights-interval` rollouts (always after the in-flight generation finishes — never mid-generation).
+
+GPU layout (`slime/ray/placement_group.py`): one Ray placement group; bundles sorted by node-IP + GPU id. **Colocate** = rollout engines share the actor's GPUs (offset 0); **disaggregated** = rollout GPUs appended after actor GPUs.
+
+---
+
+## 1. Rollout generation
+
+### Actors and topology
+
+- **`RolloutManager`** (`slime/ray/rollout.py`, CPU-only Ray actor) owns:
+ - the **`DataSource`** (`slime/rollout/data_source.py`) — prompt dataset with epoch/offset/shuffle state, checkpointable; `RolloutDataSourceWithBuffer` adds a buffer for partial/aborted rollouts;
+ - the **rollout function** (`--rollout-function-path`, default `slime/rollout/sglang_rollout.py::generate_rollout`) and eval function;
+ - the **SGLang fleet**.
+- **Fleet topology:** `RolloutServer` (one per *model* — multi-model via `--sglang-config` YAML, each with its own router) → `ServerGroup`s (homogeneous engines; supports prefill/decode **PD-disaggregation**, encoder groups, placeholders) → **`SGLangEngine`** Ray actors (`slime/backends/sglang_utils/sglang_engine.py`).
+- Each `SGLangEngine` is a thin wrapper that **spawns an sglang HTTP server subprocess** and registers it with an `sglang_router`. All generation goes over HTTP through the router (load balancing / consistent-hash sessions). This is server-based, unlike verl's in-process engines.
+- Health monitors restart dead engines; `RolloutServer.recover()` re-creates them and flags `num_new_engines` so the next weight update reconnects + re-pushes weights.
+
+### Generation flow (`sglang_rollout.py`)
+
+1. `RolloutManager.generate(rollout_id)` → rollout fn with the data source.
+2. Asyncio fan-out of **groups** (`n_samples_per_prompt` samples per prompt; semaphore = `sglang_server_concurrency × num_engines`).
+3. Each sample POSTs `input_ids` to the router `/generate` with `return_logprob=True`. Response **tokens and logprobs are appended directly to the `Sample`** (token-in/token-out — no retokenization, exact tokens for training).
+4. Reward via `rm_hub` (per-sample or whole-group `--group-rm`).
+5. **Over-sampling + dynamic-filter loop**: keeps submitting `over_sampling_batch_size` prompts and filtering groups (e.g. drop zero-std GRPO groups) until `rollout_batch_size` good groups are collected; then **aborts** stragglers. With `--partial-rollout`, aborted partial generations return to the buffer and resume next rollout (their off-policy prefix can be loss-masked).
+6. `Sample` (`slime/utils/types.py`) is the universal currency: `tokens`, `response_length`, `loss_mask`, `reward`, `rollout_log_probs`, `teacher_log_probs`, `metadata`, status.
+
+### Conversion to train data (in the manager, not the trainer)
+
+`_convert_samples_to_train_data`:
+- reward post-processing **including GRPO group normalization** (mean-center ± std-normalize per group);
+- loss masks; `rollout_mask_sums` (per-rollout mask totals, so token-weighted loss stays correct when first-fit packing splits a rollout across micro-batches);
+- optional passthroughs: `rollout_log_probs` (for off-policy correction), `teacher_log_probs` (OPD), routed experts (R3: rollout routing replay for MoE).
+
+`_split_train_data_by_dp` → `build_dp_schedule`: sequence-length-balanced first-fit packing into micro-batches, one `ray.put` **Box per DP rank**.
+
+---
+
+## 2. Training engine
+
+- **`RayTrainGroup`** (`slime/ray/actor_group.py`) spawns one **`MegatronTrainRayActor`** (`slime/backends/megatron_utils/actor.py`) per GPU; `torch.distributed` + full Megatron-Core parallelism (TP / PP / EP / CP / VPP, distributed optimizer).
+
+### The key design trick: weight tags (`TensorBackuper`)
+
+There is **one set of live Megatron GPU buffers**. CPU-backed snapshots live under tags:
+
+| tag | role |
+|---|---|
+| `actor` | the trainable policy (re-backed-up after every train step) |
+| `ref` | frozen reference for KL (`--ref-load`); refreshable via `--ref-update-interval` |
+| `teacher` | OPD teacher (`--opd-teacher-load`, megatron mode) |
+| `old_actor` / `rollout_actor` | behavior policy queue for off-policy correction (`--keep-old-actor`) |
+
+`_switch_model(tag)` restores a tag into the live model. Ref/teacher forward passes are **weight swaps, not second resident models** — unlike verl's `ActorRolloutRef` (separately materialized FSDP ref model). Memory cost is CPU RAM + swap time, not GPU.
+
+### `train_actor` sequence (per rollout)
+
+1. swap→`ref`, `forward_only` → `ref_log_probs` (if KL enabled)
+2. swap→`teacher`, `forward_only` → `teacher_log_probs` (OPD megatron mode)
+3. swap→`actor` (or `old_actor`), `forward_only` → `log_probs` (old/behavior log-probs)
+4. `compute_advantages_and_returns` (`loss.py`) — estimator: grpo / gspo / ppo / reinforce++ / custom
+5. `train()` (`model.py`) — Megatron pipeline-scheduled fwd/bwd over micro-batches; `loss_type` ∈ {policy_loss, sft_loss, value_loss, custom_loss}
+6. re-backup `actor`; optionally refresh `ref` every `--ref-update-interval` rollouts
+
+Colocate offload: `torch_memory_saver` pause/resume + destroy/reload of NCCL process groups (`sleep`/`wake_up`).
+
+---
+
+## 3. Weight update (train → rollout)
+
+Common protocol (all transports): `weight_version += 1` → `pause_generation` + `flush_cache` on engines → stream weights in **bounded buckets** (`--update-weight-buffer-size`) → `continue_generation`. Per bucket: TP all-gather → EP all-gather (experts) → **Megatron→HF conversion on the fly** (`megatron_to_hf/` per-arch converters; optional fp8/int4 quantization) → send.
+
+| Class | Mode | Transport |
+|---|---|---|
+| `UpdateWeightFromTensor` | colocate | flatten bucket → CUDA-IPC serialize → gloo-gather to the engine-owning rank → `engine.update_weights_from_tensor` (zero-copy, same GPUs) |
+| `UpdateWeightFromDistributed` | disaggregated | dedicated NCCL group per PP-source rank (`slime-pp_{k}` = train rank + all engine GPUs); broadcast; a Ray `Lock` serializes buckets to avoid NCCL deadlock |
+| `UpdateWeightFromDistributedDelta` | disaggregated | only changed params, delta-encoded — for fast frequent sync |
+| `UpdateWeightFromDisk` | fallback | save HF checkpoint; engines `update_weights_from_disk` |
+
+`--check-weight-update-equal` snapshots engine weights and bit-compares after the first push. CI also verifies `weight_version` propagated to engines.
+
+---
+
+## 4. Native on-policy distillation (OPD) in slime
+
+`examples/on_policy_distillation/` + `slime/rollout/on_policy_distillation.py` + `loss.py::apply_opd_kl_to_advantages`.
+
+- **Two teacher modes** (`--opd-type`):
+ - **`sglang`** — teacher on an external SGLang server. During rollout, a "reward func" POSTs the student's `sample.tokens` with `max_new_tokens=0, return_logprob=True, logprob_start_len=0` — a **prefill-only scoring pass** — and `post_process_rewards` trims the returned logprobs to the response span → `sample.teacher_log_probs`. Scalar reward = 0.
+ - **`megatron`** — teacher checkpoint loaded as the `teacher` weight tag; `teacher_log_probs` computed by a training-side forward pass. Requires same architecture as the policy.
+- **Loss formulation** — *not* a separate loss; an additive advantage penalty, orthogonal to the estimator:
+
+ ```
+ adv_t -= opd_kl_coef · (log π_student(y_t) − log π_teacher(y_t))
+ ```
+
+ i.e. the Tinker-cookbook **sampled-token reverse-KL** pushed through the policy-gradient machinery (reward 0 for pure distillation).
+
+### Crucial difference vs. our verl CRISP implementation
+
+| | CRISP on verl (`workspace/`) | slime OPD |
+|---|---|---|
+| KL estimator | **full-vocabulary** per-token `KL(q_t‖p_t)` from teacher+student logits (exact, dense gradient) | **sampled-token** `log q(y_t) − log p(y_t)` via advantage (unbiased policy-gradient estimate, higher variance) |
+| Teacher | same model, **conciseness prompt**, periodic hard-copy from student | separate (usually larger) model, static |
+| Teacher pass | training-side forward with swapped *prompt* (`sd_prompt`) | sglang prefill-only scoring of `sample.tokens` (same prompt) or megatron forward |
+| GT answers | never in loss | reward = 0 in pure distillation (same property) |
+
+---
+
+## 5. CRISP-on-slime integration plan
+
+Goal: reproduce CRISP (teacher = same model + conciseness instruction, reverse KL on student rollouts, teacher refresh every M steps) on slime's infra.
+
+### What's free
+
+- On-policy rollouts + per-step weight sync: core loop.
+- Teacher log-probs on student tokens + reverse-KL training signal: `--use-opd --opd-kl-coef`.
+- No-GT training: OPD already uses reward 0.
+- DAPO-math-17k data, Qwen3-8B recipes: `examples/on_policy_distillation/run-qwen3-8B-opd.sh` is the template.
+
+### What needs building
+
+1. **Conciseness-prompt teacher scoring (the core change).** A custom CRISP reward func (variant of `on_policy_distillation.reward_func`): instead of scoring `sample.tokens` verbatim, build
+ `teacher_input_ids = tokenize(conciseness_prompt(question)) + response_tokens`
+ and request logprobs with `logprob_start_len = len(teacher_prompt_ids)`; store the response-span logprobs as `sample.teacher_log_probs`. This is exactly our `sd_prompt`/`sft_prompt` pair, expressed as a prefill-only scoring call. Alignment invariant from `METHOD.md` carries over: same response tokens after different prompts → positions align by construction (no padding/truncation games needed since sglang scores exact token ids).
+2. **Teacher refresh every M steps (progressive compression).**
+ - *Megatron mode*: mirror `--ref-update-interval` — add `--opd-teacher-update-interval` that re-backups `actor → teacher` (one `weights_backuper.backup("teacher")` call in `train_actor`). Cleanest.
+ - *sglang mode*: point the teacher scoring at **the same updatable engines** (teacher = current policy + conciseness prompt). That is M=1 — which the paper showed is catastrophically unstable — so for M>1 either keep a second non-updatable model entry in `--sglang-config` refreshed manually, or use megatron mode.
+3. **Choice of KL estimator.**
+ - *Option A (use slime as-is)*: sampled-token reverse KL via `apply_opd_kl_to_advantages`. Cheap; different gradient variance than the paper; needs empirical validation that compression/accuracy match Table 2.
+ - *Option B (faithful)*: add a `full_kl` distillation loss type — teacher forward keeps full logits (or top-k) and computes exact per-token reverse KL like `opsd_worker._compute_reverse_kl_loss`. Fully specified in `workspace/slime_crisp/FULL_KL_PLAN.md`: sglang `top_logprobs_num` on the existing scoring call + bucketed reverse KL as a `custom_loss` (`--custom-loss-function-path`), ~5-file plumbing patch to the slime fork.
+ - Recommendation: start with A (a few hundred lines total, mostly config), benchmark on MATH-500 against the verl numbers, fall back to B only if the sampled estimator can't reproduce the compression–accuracy trade-off.
+4. **Config**: student prompt = original DAPO prompt (already the OPD example's data path); temperature 1.0; `n_samples_per_prompt=1` (CRISP uses single rollouts; note GRPO normalization is degenerate at n=1 — use reward 0 + pure OPD penalty, estimator effectively REINFORCE with zero advantage base); batch 32; lr 1e-6; `rollout_max_response_len=8192`.
+5. **Eval**: MATH-500/AIME via slime's `--eval-datasets` + a `rm_hub` scorer mirroring `dual_path_math_verify` (slime ships `math_dapo_utils`/`math_utils` which match the paper's footnote-1 grading).
+
+### Suggested first milestone
+
+Qwen3-8B, megatron-mode teacher initialized from the same checkpoint, custom conciseness-prompt scoring func, `--opd-teacher-update-interval 50`, sampled-KL estimator (Option A), 100 steps on DAPO-17k → compare MATH-500 acc/len against `workspace/` Table-2 numbers (86.6% / 1,921 tok for 8B).
diff --git a/crisp_compressed_reasoning_via_iterative_self_policy_distillation.pdf b/crisp_compressed_reasoning_via_iterative_self_policy_distillation.pdf
index 0bf4216..75ec398 100644
Binary files a/crisp_compressed_reasoning_via_iterative_self_policy_distillation.pdf and b/crisp_compressed_reasoning_via_iterative_self_policy_distillation.pdf differ
diff --git a/slime b/slime
new file mode 160000
index 0000000..e0ef620
--- /dev/null
+++ b/slime
@@ -0,0 +1 @@
+Subproject commit e0ef620588a118b3bbddd6d1504d3d3fa7b2891a
diff --git a/workspace/slime_crisp/FULL_KL_PLAN.md b/workspace/slime_crisp/FULL_KL_PLAN.md
new file mode 100644
index 0000000..a323acb
--- /dev/null
+++ b/workspace/slime_crisp/FULL_KL_PLAN.md
@@ -0,0 +1,190 @@
+# Plan: full-vocab reverse KL `KL(q_t‖p_t)` on slime (milestone 2)
+
+Closes the one fundamental estimator gap between this port and the paper/verl implementation
+(see README "What differs"): replace the sampled-token advantage penalty with a per-token
+**distribution-level** reverse KL, computed from full student logits and (top-K-truncated)
+teacher log-probs — recovering verl's `_compute_reverse_kl_loss` objective as K→V.
+
+Status: **implemented** (loss: `crisp_full_kl_loss.py`; rollout top-K: `crisp_opd.py`;
+plumbing: slime submodule branch `crisp`, commit `e0ef620`; 23 unit tests passing).
+Not yet run on GPUs — the A/B in §6 needs both arms on the cluster.
+
+**Correction found during testing** (encoded in `test_full_kl_loss.py`): approximation
+fidelity is governed by **`q_tail`** (the *student's* mass outside teacher top-K), not by
+teacher coverage alone — reverse KL is an expectation under `q`, so a diffuse student hides
+large log-ratios in the tail bucket (the bucketed KL stays a lower bound). In CRISP's
+self-distillation regime q≈p so q_tail is tiny, and the tail term's gradient actively shrinks
+it; treat rising `q_tail` (logged per step) as the signal to raise K.
+
+---
+
+## 1. Why and what
+
+| | milestone 1 (sampled) | milestone 2 (this plan) | verl reference |
+|---|---|---|---|
+| objective | `adv_t = −β(log q(y_t) − log p(y_t))` through PG loss | `(1/N) Σ_t KL_K(q_t‖p_t)` direct loss | `(1/N) Σ_t KL(q_t‖p_t)` full vocab |
+| gradient | unbiased, **one-sample** estimate per token | dense over teacher top-K + tail bucket | dense over full vocab |
+| teacher info needed | `log p(y_t)` (1 float/token) | top-K `(id, log p)` pairs/token | full teacher logits |
+
+The training-side student logits are already full-vocab (the training forward); the only thing
+the sampled estimator throws away is the **teacher's** distribution. We recover it via sglang's
+`top_logprobs_num` on the existing prefill-only scoring call — no second training-side model,
+no change to the teacher-refresh machinery.
+
+### The truncated objective: bucketed reverse KL
+
+For each response position `t`, with teacher top-K ids `x_1..x_K` (from sglang), teacher
+log-probs `log p_k`, and student probabilities `q_k = q_t(x_k)` (exact, from training logits):
+
+```
+q_tail = 1 − Σ_k q_k p_tail = 1 − Σ_k p_k
+KL_K(q_t‖p_t) = Σ_k q_k (log q_k − log p_k) + q_tail (log q_tail − log p_tail)
+```
+
+Properties: a true KL on the coarsened (K+1)-bucket space — non-negative, zero iff the
+coarsened distributions match, **exactly the verl objective at K=V**. Reasoning teachers are
+low-entropy (paper Fig. 8: ~0.3 nats), so top-K coverage `Σ_k p_k` should be ≥99.9% at K=256;
+we log it every step to quantify the truncation (§5). Differentiable in the student logits
+everywhere (including the tail term, which pushes mass off non-teacher tokens).
+
+---
+
+## 2. Architecture decision
+
+**Chosen: B1 — sglang teacher top-K + custom Megatron loss.**
+
+Rejected alternatives:
+- **B2: Megatron-mode teacher** (`teacher` weight tag) — the teacher must see a *different,
+ longer prompt* than the student, so the teacher forward needs a second token sequence per
+ sample plumbed through the data pipeline; and since slime's weight tags time-share one set of
+ GPU buffers, teacher logits would have to be *stored* between the teacher phase and the train
+ phase — full logits are `T×V≈8192×150k×2B ≈ 2.4 GB/sample`, so storage forces top-K anyway.
+ Same approximation, far more plumbing. Only advantage: bit-identical teacher precision.
+- **B3: resident second model** (verl-style) — contradicts slime's core memory design; a deep
+ fork we don't want to maintain.
+
+B1 keeps everything from milestone 1 (teacher server, conciseness prompt, `/update_weights_from_disk`
+refresh, alignment check) and adds one payload field + one custom loss.
+
+---
+
+## 3. Implementation steps
+
+### Step 1 — rollout side: fetch teacher top-K (ours, `crisp_opd.py`)
+
+In `reward_func`, when `crisp_kl_mode == "full"` (new knob in `crisp_config.yaml`):
+
+- add `"top_logprobs_num": K` (`crisp_teacher_topk`, default 256) to the scoring payload;
+- parse `meta_info["input_top_logprobs"]` (per-position list of `[logprob, token_id, ...]`),
+ trim to the response span exactly like `input_token_logprobs`;
+- store `sample.teacher_top_ids` (T×K ints) and `sample.teacher_top_logprobs` (T×K floats);
+- keep storing sampled `teacher_log_probs` too (cheap; enables logging both estimators).
+
+⚠️ Verify the sglang field name/shape on the deployed version first (one curl); it has shifted
+across sglang releases.
+
+### Step 2 — plumbing: carry T×K fields to the loss (small slime fork patch)
+
+The pipeline whitelists keys at four places; new fields need ~5 mechanical edits. Fork the
+submodule to `HJSang/slime` branch `crisp` (candidate upstream PR: "generic top-K teacher
+log-probs for OPD"):
+
+1. `slime/utils/types.py` — `Sample.teacher_top_ids`, `Sample.teacher_top_logprobs`.
+2. `slime/ray/rollout.py::_convert_samples_to_train_data` — pass-through (mirror the existing
+ `teacher_log_probs` lines).
+3. `slime/ray/rollout.py::_split_train_data_by_dp` — add both keys to the per-rank key list.
+4. `slime/backends/megatron_utils/actor.py::_get_rollout_data` — move to GPU; **assert CP==1
+ for now** (the `slice_log_prob_with_cp` helper is 1-D/per-token; our milestone runs
+ `context-parallel-size 1`, so defer 2-D CP slicing).
+5. `slime/backends/megatron_utils/model.py` — add both keys to the `get_batch([...])` list in
+ the train `forward_step`.
+
+### Step 3 — the loss (ours, new `crisp_full_kl_loss.py`)
+
+Wire: `--loss-type custom_loss --custom-loss-function-path slime_crisp.crisp_full_kl_loss.full_kl_loss_function`.
+slime's contract (`loss.py::loss_function`): `func(args, batch, logits, sum_of_sample_mean) → (loss, metrics)`,
+with `logits` = full student `[1, T, V]` (vocab-parallel under TP).
+
+```
+full_kl_loss_function(args, batch, logits, sum_of_sample_mean):
+ 1. slice response-position logits # reuse get_log_probs_and_entropy's slicing path
+ 2. per token-chunk (chunk_size≈128, fp32 upcast — mirrors verl's chunked loss):
+ lse = vocab_parallel_logsumexp(logits_chunk) # max + sum-exp all-reduce over TP
+ log_q_K = vocab_parallel_gather(logits_chunk, top_ids) − lse # local mask+gather, all-reduce(SUM)
+ q_K = exp(log_q_K); q_tail = clamp(1 − Σ q_K, min=eps)
+ p_tail = clamp(1 − Σ exp(teacher_top_logprobs), min=eps)
+ kl_chunk = Σ_k q_K (log_q_K − log p_K) + q_tail (log q_tail − log p_tail)
+ 3. loss = sum_of_sample_mean(kl_per_token) # --calculate-per-token-loss ⇒ global token mean = verl parity
+ 4. zero-token guard: loss += 0 * logits.sum() # same trick as sft_loss_function
+ 5. metrics: kl, teacher_topk_coverage (Σ p_K), q_tail, student_entropy
+ (reuse compute_entropy_from_logits under no_grad), sampled-KL for comparison
+```
+
+New primitive `vocab_parallel_topk_log_probs(logits, ids_TK, tp_group)` (~40 lines): can't reuse
+`compute_log_probs` directly (it gathers 1 id/position); implement K-id gather as
+local-shard-range mask + `torch.gather` + `all_reduce(SUM)`, sharing one logsumexp per position.
+Degrade to plain torch when TP world size == 1 → CPU-unit-testable.
+
+### Step 4 — config & mode switching
+
+- `crisp_config.yaml`: `crisp_kl_mode: full|sampled`, `crisp_teacher_topk: 256`.
+- Launch script (full mode): drop `--use-opd`/`--opd-*`, add `--loss-type custom_loss
+ --custom-loss-function-path ...` and `--disable-compute-advantages-and-returns` — skips the
+ whole advantage pipeline *and the extra old-log-prob forward pass* (the custom loss needs only
+ the training forward; per-step cost roughly matches verl's teacher-fwd + student-fwd/bwd,
+ with the teacher fwd amortized into rollout-time scoring).
+- Everything else (data, prompts, refresh, reward zeroing) unchanged from milestone 1.
+
+---
+
+## 4. Tests (extend `test_crisp_opd.py` + new `test_full_kl_loss.py`)
+
+1. **Math**: bucketed `KL_K` vs exact `KL` on a toy vocab (V=50): equal at K=V; lower bound and
+ monotone non-decreasing in K; zero iff distributions equal; gradient matches
+ `torch.autograd.gradcheck` on the K=V case against a direct full-KL implementation
+ (this is the verl-equivalence proof at small scale).
+2. **Primitive**: `vocab_parallel_topk_log_probs` TP=1 path vs plain `log_softmax + gather`.
+3. **Parsing**: `input_top_logprobs` trimming + id-echo alignment check (extend the existing
+ alignment test to the top-K field).
+4. **Plumbing smoke** (GPU, 1 node): `--num-rollout 2`, assert `teacher_top_*` reach the loss
+ with the right shapes and `teacher_topk_coverage > 0.99`.
+
+---
+
+## 5. Observability
+
+Per-step logs (all cheap, from the loss metrics): `loss` (= bucketed KL), `kl_sampled` (same
+batch — directly measures estimator variance/bias), `teacher_topk_coverage`, `q_tail`, plus
+the existing `raw_reward` accuracy curve. **`q_tail` is the fidelity diagnostic** (see status
+note above); if it rises above ~1e-2, raise K before trusting the run.
+
+---
+
+## 6. Validation: the A/B that motivates all of this
+
+Same data, prompts, hyperparameters, M=50, 100 steps, Qwen3-8B:
+
+| arm | loss | expectation |
+|---|---|---|
+| A (milestone 1) | sampled-token KL via OPD penalty | ? — the open question |
+| B (this plan) | bucketed full KL, K=256 | ≈ verl reference |
+| verl reference | `workspace/` Table-2 run | MATH-500 86.6% / 1,921 tok |
+
+Acceptance for B: MATH-500 accuracy within ~1 pt and response length within ~10% of the verl
+run at step 100. Ablate K ∈ {64, 256, 1024} only if coverage logging shows it matters.
+If A ≈ B, the sampled estimator is vindicated (cheaper, zero fork); if A ≪ B, milestone 2
+becomes the default recipe.
+
+---
+
+## 7. Risks & costs
+
+- **sglang `top_logprobs_num` overhead**: K floats+ids per scored token serialized over HTTP
+ (≈ T×256×8B ≈ 8 MB per 4k-token response). Measure scoring throughput; drop to K=64 or gzip
+ if the teacher GPU becomes the bottleneck.
+- **API drift**: `input_top_logprobs` name/shape varies across sglang versions — verify first.
+- **CP > 1 unsupported** (asserted) until 2-D CP slicing is added; irrelevant to the 8B recipe.
+- **Numerics**: `q_tail`/`p_tail` need `clamp(min=1e-8)`; upcast chunks to fp32 like verl.
+- **Fork maintenance**: ~5-file mechanical patch; pin to our fork branch and propose upstream.
+
+Estimated effort: Step 1 ~0.5 d, Step 2 ~0.5 d, Step 3+4 ~1.5 d, A/B runs = cluster time.
diff --git a/workspace/slime_crisp/README.md b/workspace/slime_crisp/README.md
new file mode 100644
index 0000000..fec42b3
--- /dev/null
+++ b/workspace/slime_crisp/README.md
@@ -0,0 +1,73 @@
+# CRISP on slime
+
+Runs CRISP — teacher = the **same model** conditioned on a conciseness prompt, reverse-KL
+distillation on the student's own rollouts, teacher refresh every M steps — on
+[slime](../../slime). Two KL estimators, switched via `CRISP_KL_MODE`:
+
+- **`sampled`** (milestone 1): sampled-token reverse KL through slime's OPD advantage penalty
+ (`--use-opd --opd-type sglang`). Stock slime.
+- **`full`** (milestone 2, [`FULL_KL_PLAN.md`](FULL_KL_PLAN.md)): bucketed full-vocab reverse KL
+ over the teacher's top-K (`crisp_full_kl_loss.py`, wired as a Megatron `custom_loss`).
+ Equals the verl objective exactly at K=V (gradient-verified in tests). Requires the patched
+ slime submodule (branch `crisp`: plumbs `teacher_top_ids/logprobs`, ~30 lines).
+
+Design rationale: [`SLIME_DESIGN.md`](../../SLIME_DESIGN.md) §5; algorithm: [`METHOD.md`](../../METHOD.md).
+
+## What differs from the verl implementation (`workspace/src/`)
+
+| | verl (`self_distill_hybrid/`) | this (slime) |
+|---|---|---|
+| KL estimator | full-vocabulary per-token `KL(q_t‖p_t)` from logits | **sampled-token** `log q(y_t) − log p(y_t)` as an advantage penalty (slime OPD) |
+| Teacher pass | training-side forward, swapped prompt | **prefill-only scoring call** to a teacher sglang server |
+| Teacher refresh | `update_teacher` FSDP shard copy | actor saves HF dump (`--save-hf`) → teacher server `/update_weights_from_disk` |
+| Response tokens in training | decoded text **re-tokenized** (+EOS appended) | **exact rollout token ids** (no re-tokenization; cleaner) |
+| In-training accuracy metric | dual-path `math_verify` (sympy) | `math_dapo` minerva (`Answer:` extraction) — what the paper's benchmarks used (footnote 1); curves may differ a few points from verl logs |
+| Teacher logit precision | trainer bf16 forward | sglang inference kernels (numerically close, not bit-equal) |
+
+Verified identical (audited, with regression tests): teacher prompt strings (byte-equal to
+`config/prompts.json`), question/GT extraction (0 mismatches on 2,000 real rows), train split
+(row-for-row equal to verl's seed-42 80% split, 13,918 rows), teacher-refresh window semantics
+(steps 1–50 use θ₀, 51–100 use θ₅₀ in both), sampling params, optimizer settings, no
+correctness filtering. Whether the sampled-token estimator reproduces the paper's
+compression–accuracy trade-off is exactly what this milestone tests.
+
+## Files
+
+- `crisp_opd.py` — the three slime hooks:
+ - `reward_func` (`--custom-rm-path`): builds `teacher_prompt(question) + response_tokens`,
+ requests a prefill-only forward from the teacher server, stores per-token log-probs on
+ `sample.teacher_log_probs` (with a token-id alignment check), returns math correctness
+ (metrics only, logged as the training-accuracy curve).
+ - `post_process_rewards` (`--custom-reward-post-process-path`): zeroes training rewards —
+ pure distillation; the signal is slime's `adv_t -= opd_kl_coef·(logπ_s(y_t) − logπ_t(y_t))`.
+ - `generate_rollout` (`--rollout-function-path`): default rollout + teacher refresh
+ (θ̃←θ) every `crisp_teacher_update_interval` rollouts via `/update_weights_from_disk`.
+- `crisp_full_kl_loss.py` — milestone-2 loss: `vocab_parallel_topk_log_probs` (TP-aware,
+ autograd-correct), `bucketed_reverse_kl` (top-K + tail bucket), and the slime `custom_loss`
+ entry point `full_kl_loss_function` (logs `teacher_topk_coverage`, `q_tail`, and the
+ same-batch `kl_sampled` estimator-comparison diagnostic).
+- `crisp_config.yaml` / `crisp_config_full_kl.yaml` — per-mode knobs, merged onto args via
+ `--custom-config-path`.
+- `prepare_crisp_slime_data.py` — DAPO parquet → slime jsonl
+ (`prompt` = original DAPO content, `label` = GT, `metadata.question` = bare question).
+- `run-qwen3-8b-crisp.sh` — 8-GPU launch: actor 4 | rollout 3 | teacher 1; paper recipe
+ (batch 32, lr 1e-6, temp 1.0, 8192-token rollouts, n=1, 100 steps, M=50).
+ `CRISP_KL_MODE=full bash run-qwen3-8b-crisp.sh` for milestone 2.
+- `test_crisp_opd.py`, `test_full_kl_loss.py` — 23 unit tests (no GPU/slime runtime needed).
+
+## Invariants to keep
+
+1. **Same response token ids** under both prompts — enforced by the token-id echo check in
+ `reward_func` (the slime analogue of verl's drop-don't-truncate rule, see `METHOD.md` §7).
+2. `--save-interval` == `crisp_teacher_update_interval`, `--save-hf` == `crisp_teacher_hf_path`
+ (fixed path, overwritten each save).
+3. Synchronous `train.py` only — `train_async.py` would refresh the teacher mid-pipeline.
+4. Full-KL mode: CP must be 1 and `qkv_format` must be `thd` (asserted); the slime submodule
+ must be on the `crisp` branch (the plumbing patch).
+
+## Success criterion (vs paper Table 2, Qwen3-8B @30K budget)
+
+MATH-500: base 77.7% / 4,661 tok → CRISP 86.6% / 1,921 tok (58.8% reduction).
+Run both estimators (`CRISP_KL_MODE=sampled` vs `full`) and compare against the verl
+reference; the `kl_sampled` diagnostic logged by the full-KL loss measures the estimator gap
+on identical batches. See [`FULL_KL_PLAN.md`](FULL_KL_PLAN.md) §6 for acceptance criteria.
diff --git a/workspace/slime_crisp/__init__.py b/workspace/slime_crisp/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/workspace/slime_crisp/crisp_config.yaml b/workspace/slime_crisp/crisp_config.yaml
new file mode 100644
index 0000000..dd1e9b7
--- /dev/null
+++ b/workspace/slime_crisp/crisp_config.yaml
@@ -0,0 +1,24 @@
+# CRISP-specific knobs, merged onto slime args via --custom-config-path.
+# (slime's parser setattr's every key here onto the args Namespace.)
+# This is the milestone-1 (sampled-token KL) config; the full-vocab KL
+# variant lives in crisp_config_full_kl.yaml.
+
+crisp_kl_mode: sampled
+
+# Teacher refresh interval M (paper default: 50; 0 = frozen teacher).
+# MUST equal --save-interval in the launch script, and requires synchronous
+# train.py (train_async.py would refresh mid-pipeline).
+crisp_teacher_update_interval: 50
+
+# Where the actor dumps HF weights for the teacher to reload.
+# MUST equal --save-hf in the launch script. Fixed path (no {rollout_id}) so
+# each save overwrites the previous dump.
+crisp_teacher_hf_path: /root/crisp_teacher_hf
+
+# Optional: teacher /generate URL; defaults to --rm-url when unset.
+# crisp_teacher_url: http://127.0.0.1:13141/generate
+
+# Optional teacher-prompt overrides (defaults = paper Figure-3 conciseness
+# instruction, identical to workspace/config/prompts.json length_prune_teacher).
+# crisp_teacher_prompt_prefix: "..."
+# crisp_teacher_prompt_suffix: "..."
diff --git a/workspace/slime_crisp/crisp_config_full_kl.yaml b/workspace/slime_crisp/crisp_config_full_kl.yaml
new file mode 100644
index 0000000..acde7ea
--- /dev/null
+++ b/workspace/slime_crisp/crisp_config_full_kl.yaml
@@ -0,0 +1,23 @@
+# CRISP full-KL mode (milestone 2, FULL_KL_PLAN.md) — merged onto slime args
+# via --custom-config-path. Same knobs as crisp_config.yaml plus the KL mode.
+# Pair with the full-KL CLI flags set by run-qwen3-8b-crisp.sh under
+# CRISP_KL_MODE=full (--loss-type custom_loss etc.).
+
+crisp_kl_mode: full
+crisp_teacher_topk: 256 # bucketed-KL K; coverage logged per step
+crisp_kl_chunk_size: 128 # token chunk for the loss (memory bound)
+
+# Teacher refresh interval M (paper default: 50; 0 = frozen teacher).
+# MUST equal --save-interval; synchronous train.py only.
+crisp_teacher_update_interval: 50
+
+# Where the actor dumps HF weights for the teacher to reload.
+# MUST equal --save-hf. Fixed path so each save overwrites the previous dump.
+crisp_teacher_hf_path: /root/crisp_teacher_hf
+
+# Optional: teacher /generate URL; defaults to --rm-url when unset.
+# crisp_teacher_url: http://127.0.0.1:13141/generate
+
+# Optional teacher-prompt overrides (defaults = paper Figure-3 instruction).
+# crisp_teacher_prompt_prefix: "..."
+# crisp_teacher_prompt_suffix: "..."
diff --git a/workspace/slime_crisp/crisp_full_kl_loss.py b/workspace/slime_crisp/crisp_full_kl_loss.py
new file mode 100644
index 0000000..c5114e6
--- /dev/null
+++ b/workspace/slime_crisp/crisp_full_kl_loss.py
@@ -0,0 +1,281 @@
+"""Full-vocab (top-K truncated) reverse KL distillation loss for slime.
+
+Implements CRISP milestone 2 (FULL_KL_PLAN.md): a distribution-level per-token
+reverse KL between the student and the conciseness-conditioned teacher,
+replacing the sampled-token OPD advantage penalty.
+
+For each response position t, with teacher top-K ids x_1..x_K and log-probs
+log p_k (fetched by crisp_opd.reward_func via sglang ``top_logprobs_num``) and
+exact student probabilities q_k from the training logits:
+
+ q_tail = 1 - sum_k q_k p_tail = 1 - sum_k p_k
+ KL_K(q_t || p_t) = sum_k q_k (log q_k - log p_k)
+ + q_tail (log q_tail - log p_tail)
+
+This is the reverse KL of the coarsened (K+1)-bucket distributions: >= 0, zero
+iff the coarsened distributions match, a lower bound of the exact reverse KL,
+and equal to it (verl's ``opsd_worker._compute_reverse_kl_loss`` objective) at
+K=V.
+
+Approximation fidelity is governed by ``q_tail`` — the STUDENT's probability
+mass outside the teacher's top-K — not by teacher coverage alone (reverse KL
+is an expectation under q, so a diffuse student hides large log-ratios in the
+tail bucket). In CRISP's self-distillation regime q ~ p and q_tail is tiny;
+the tail term's gradient also actively pushes student mass back onto the
+teacher's support. Both ``q_tail`` and ``teacher_topk_coverage`` are logged
+per step — treat rising q_tail as the signal to raise K.
+
+Wiring (full-KL mode, see run-qwen3-8b-crisp.sh):
+
+ --loss-type custom_loss
+ --custom-loss-function-path slime_crisp.crisp_full_kl_loss.full_kl_loss_function
+ --disable-compute-advantages-and-returns
+ --recompute-loss-function # frees the per-chunk softmax graph
+ --calculate-per-token-loss # global token mean == verl normalization
+
+Supports TP (vocab-parallel logits) via two small autograd functions modeled
+on slime's ``_VocabParallelEntropy``; degrades to plain ops when no process
+group is active, so the math is unit-testable on CPU. CP must be 1 (asserted
+upstream in the plumbing patch) and qkv_format must be "thd".
+"""
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_warned_temperature = False
+
+
+def _tp_world_size(tp_group) -> int:
+ import torch.distributed as dist
+
+ if tp_group is None or not dist.is_available() or not dist.is_initialized():
+ return 1
+ return dist.get_world_size(group=tp_group)
+
+
+class _AllReduceSumKeepGrad(torch.autograd.Function):
+ """all_reduce(SUM) whose backward passes the (TP-replicated) grad through.
+
+ Correct because the downstream loss is identical on every TP rank, so the
+ upstream gradient is already the full dL/d(sum); each rank then applies its
+ local Jacobian (here: identity onto its own summand).
+ """
+
+ @staticmethod
+ def forward(ctx, x: torch.Tensor, tp_group) -> torch.Tensor:
+ import torch.distributed as dist
+
+ x = x.clone()
+ dist.all_reduce(x, op=dist.ReduceOp.SUM, group=tp_group)
+ return x
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor):
+ return grad_output, None
+
+
+class _VocabParallelLogSumExp(torch.autograd.Function):
+ """logsumexp over the (TP-sharded) vocab dim. [N, V_local] -> [N, 1]."""
+
+ @staticmethod
+ def forward(ctx, logits: torch.Tensor, tp_group) -> torch.Tensor:
+ import torch.distributed as dist
+
+ logits_max = logits.max(dim=-1, keepdim=True).values
+ dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group)
+ sum_exp = (logits - logits_max).exp().sum(dim=-1, keepdim=True)
+ dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=tp_group)
+ lse = logits_max + sum_exp.log()
+ ctx.save_for_backward(logits, lse)
+ return lse
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor):
+ logits, lse = ctx.saved_tensors
+ # d lse / d logits_local = softmax_local; recomputed here so the
+ # forward never stores an [N, V] probability tensor.
+ return grad_output * (logits - lse).exp(), None
+
+
+def vocab_parallel_topk_log_probs(
+ logits: torch.Tensor, ids: torch.Tensor, tp_group=None
+) -> torch.Tensor:
+ """Student log-probs at arbitrary token ids under vocab parallelism.
+
+ Args:
+ logits: [N, V_local] (full V when tp world size is 1). Differentiable.
+ ids: [N, K] global token ids.
+ tp_group: tensor-parallel process group (None -> single shard).
+
+ Returns:
+ [N, K] log q(ids), differentiable in ``logits``.
+ """
+ tp_size = _tp_world_size(tp_group)
+ v_local = logits.size(-1)
+
+ if tp_size == 1:
+ gathered = logits.gather(-1, ids)
+ lse = logits.logsumexp(dim=-1, keepdim=True)
+ return gathered - lse
+
+ import torch.distributed as dist
+
+ shard_start = dist.get_rank(group=tp_group) * v_local
+ in_shard = (ids >= shard_start) & (ids < shard_start + v_local)
+ local_ids = (ids - shard_start).clamp(0, v_local - 1)
+ # Out-of-shard gathers are masked to 0 and contributed by the owning rank
+ # via the SUM all-reduce; their grads are killed by the same mask.
+ gathered = logits.gather(-1, local_ids) * in_shard
+ gathered = _AllReduceSumKeepGrad.apply(gathered, tp_group)
+ lse = _VocabParallelLogSumExp.apply(logits, tp_group)
+ return gathered - lse
+
+
+def bucketed_reverse_kl(
+ log_q_topk: torch.Tensor,
+ teacher_top_logprobs: torch.Tensor,
+ eps: float = 1e-8,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Per-token reverse KL over teacher-top-K buckets + one tail bucket.
+
+ Args:
+ log_q_topk: [N, K] student log-probs at the teacher's top-K ids (grad).
+ teacher_top_logprobs: [N, K] teacher log-probs at the same ids (no grad).
+ eps: tail-mass clamp for numerical safety.
+
+ Returns:
+ (kl [N], teacher_coverage [N], q_tail [N]) — kl is differentiable;
+ the diagnostics are detached.
+ """
+ q_k = log_q_topk.exp()
+ p_coverage = teacher_top_logprobs.exp().sum(dim=-1)
+ q_tail = (1.0 - q_k.sum(dim=-1)).clamp(min=eps)
+ p_tail = (1.0 - p_coverage).clamp(min=eps)
+
+ kl_topk = (q_k * (log_q_topk - teacher_top_logprobs)).sum(dim=-1)
+ kl_tail = q_tail * (q_tail.log() - p_tail.log())
+ kl = kl_topk + kl_tail
+ return kl, p_coverage.detach(), q_tail.detach()
+
+
+def full_kl_loss_function(args, batch, logits: torch.Tensor, sum_of_sample_mean):
+ """slime ``custom_loss`` entry point: mean per-token KL_K(q_t || p_t).
+
+ Contract (slime loss.py::loss_function): receives the full vocab-parallel
+ student logits [1, T, V] for the packed micro-batch and returns
+ ``(loss, metrics)``; reduction/rescaling is handled by the caller through
+ ``sum_of_sample_mean`` and ``--calculate-per-token-loss``.
+ """
+ from megatron.core import mpu
+
+ assert args.qkv_format == "thd", (
+ f"full_kl_loss_function supports qkv_format='thd' only, got {args.qkv_format!r}"
+ )
+ teacher_top_ids = batch.get("teacher_top_ids")
+ teacher_top_logprobs = batch.get("teacher_top_logprobs")
+ assert teacher_top_ids is not None and teacher_top_logprobs is not None, (
+ "batch is missing teacher_top_ids/teacher_top_logprobs. Set crisp_kl_mode: full "
+ "in the custom config so crisp_opd.reward_func fetches the teacher top-K, and "
+ "run with the patched slime (crisp branch) that plumbs these fields."
+ )
+
+ assert logits.size(0) == 1, f"{logits.shape}"
+ raw_logits = logits # unscaled [1, T, V], for the kl_sampled diagnostic
+ logits = logits.squeeze(0)
+
+ # Mirror get_log_probs_and_entropy: train-time log-probs are defined at the
+ # rollout temperature. The teacher's sglang input logprobs are raw
+ # (temperature-free), so the objective only matches the paper's at 1.0.
+ rollout_temperature = getattr(args, "rollout_temperature", 1.0)
+ if rollout_temperature != 1.0:
+ global _warned_temperature
+ if not _warned_temperature:
+ logger.warning(
+ "full_kl_loss_function: rollout_temperature=%s scales student logits but "
+ "teacher top-K log-probs from sglang are raw; the KL mixes temperatures.",
+ rollout_temperature,
+ )
+ _warned_temperature = True
+ logits = logits / rollout_temperature
+
+ tp_group = mpu.get_tensor_model_parallel_group()
+ chunk_size = getattr(args, "crisp_kl_chunk_size", None) or 128
+
+ kl_list = []
+ coverage_sum = torch.tensor(0.0, device=logits.device)
+ q_tail_sum = torch.tensor(0.0, device=logits.device)
+ n_kl_tokens = 0
+
+ offset = 0
+ for total_length, response_length, top_ids, top_lps in zip(
+ batch["total_lengths"], batch["response_lengths"], teacher_top_ids, teacher_top_logprobs, strict=True
+ ):
+ end = offset + total_length
+ start = end - response_length
+ offset = end
+ if response_length == 0:
+ kl_list.append(logits.new_zeros((0,)))
+ continue
+
+ assert top_ids.size(0) == response_length, (
+ f"teacher_top_ids rows ({top_ids.size(0)}) != response_length ({response_length}); "
+ "teacher scoring must cover exactly the response span (see crisp_opd.reward_func)."
+ )
+ # Position i of the response is predicted by logits at sequence
+ # position (start - 1 + i): same shift as _extract_per_sample.
+ resp_logits = logits[start - 1 : end - 1]
+
+ sample_kl = []
+ for c in range(0, response_length, chunk_size):
+ d = min(c + chunk_size, response_length)
+ log_q_k = vocab_parallel_topk_log_probs(
+ resp_logits[c:d].float(), top_ids[c:d], tp_group
+ )
+ kl, coverage, q_tail = bucketed_reverse_kl(log_q_k, top_lps[c:d])
+ sample_kl.append(kl)
+ coverage_sum += coverage.sum()
+ q_tail_sum += q_tail.sum()
+ n_kl_tokens += d - c
+ kl_list.append(torch.cat(sample_kl, dim=0))
+
+ kl_per_token = torch.cat(kl_list, dim=0)
+ loss = sum_of_sample_mean(kl_per_token)
+ # Keep the graph connected when every sample in the micro-batch is empty
+ # (same guard as sft_loss_function).
+ if kl_per_token.numel() == 0:
+ loss = loss + 0 * logits.sum()
+
+ metrics = {
+ "loss": loss.clone().detach(),
+ "teacher_topk_coverage": coverage_sum / max(1, n_kl_tokens),
+ "q_tail": q_tail_sum / max(1, n_kl_tokens),
+ }
+
+ # Estimator-comparison diagnostic: the milestone-1 sampled-token reverse KL
+ # (log q(y_t) - log p(y_t)) on the same batch, when sampled teacher
+ # log-probs were collected alongside the top-K.
+ sampled_teacher = batch.get("teacher_log_probs")
+ if sampled_teacher:
+ from slime.backends.megatron_utils.loss import get_log_probs_and_entropy
+
+ with torch.no_grad():
+ # raw_logits: get_log_probs_and_entropy applies the rollout
+ # temperature itself, so passing the scaled logits would double-scale.
+ _, lp = get_log_probs_and_entropy(
+ raw_logits,
+ args=args,
+ unconcat_tokens=batch["unconcat_tokens"],
+ total_lengths=batch["total_lengths"],
+ response_lengths=batch["response_lengths"],
+ with_entropy=False,
+ max_seq_lens=batch.get("max_seq_lens", None),
+ )
+ sampled_kl = torch.cat(
+ [q - t for q, t in zip(lp["log_probs"], sampled_teacher, strict=True)], dim=0
+ )
+ metrics["kl_sampled"] = sum_of_sample_mean(sampled_kl).detach()
+
+ return loss, metrics
diff --git a/workspace/slime_crisp/crisp_opd.py b/workspace/slime_crisp/crisp_opd.py
new file mode 100644
index 0000000..28cc6b5
--- /dev/null
+++ b/workspace/slime_crisp/crisp_opd.py
@@ -0,0 +1,254 @@
+"""CRISP on slime: conciseness-prompted teacher scoring for on-policy distillation.
+
+Implements the CRISP method (teacher = same model conditioned on a conciseness
+instruction; per-token reverse KL on the student's own rollouts) on top of
+slime's sglang-mode OPD machinery (``--use-opd --opd-type sglang``).
+
+Three entry points, wired via slime CLI args:
+
+ --custom-rm-path slime_crisp.crisp_opd.reward_func
+ --custom-reward-post-process-path slime_crisp.crisp_opd.post_process_rewards
+ --rollout-function-path slime_crisp.crisp_opd.generate_rollout
+
+``reward_func`` scores each student rollout under the conciseness-conditioned
+teacher: it builds ``teacher_prompt(question) + response_tokens`` and requests
+a prefill-only forward (``max_new_tokens=0, return_logprob=True``) from the
+teacher sglang server (``--rm-url``). The returned per-token log-probs are
+stored on ``sample.teacher_log_probs``; slime's ``apply_opd_kl_to_advantages``
+then turns them into the sampled-token reverse-KL penalty
+``adv_t -= opd_kl_coef * (log pi_student(y_t) - log pi_teacher(y_t))``.
+
+The scalar return value of ``reward_func`` is math correctness (metrics only,
+mirroring CRISP's metrics-only verification); ``post_process_rewards`` zeroes
+it for training so the learning signal is pure distillation.
+
+``generate_rollout`` wraps slime's default rollout and implements CRISP's
+periodic teacher refresh (theta_tilde <- theta every M rollouts) by telling the
+teacher server to reload weights from the actor's latest HF dump
+(``/update_weights_from_disk``). Requires ``--save-hf`` pointing at a fixed
+path and ``--save-interval`` == the refresh interval.
+
+CRISP-specific knobs (set via ``--custom-config-path`` YAML; all land on args):
+
+ crisp_teacher_update_interval: 50 # M; 0 = frozen teacher
+ crisp_teacher_hf_path: /path/to/hf_dump # must equal --save-hf
+ crisp_teacher_url: http://ip:port/generate # optional; defaults to --rm-url
+ crisp_teacher_prompt_prefix / _suffix # optional template overrides
+ crisp_kl_mode: sampled | full # full also fetches teacher top-K
+ crisp_teacher_topk: 256 # K for crisp_kl_mode: full
+
+Alignment invariant (see METHOD.md): teacher and student share the exact same
+response token ids after different prompts, so per-position log-probs align by
+construction. The token ids echoed back by the sglang server are checked
+against the student's response tokens to enforce this.
+"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Must match workspace/config/prompts.json -> "length_prune_teacher"
+# (the paper's Figure-3 conciseness instruction).
+DEFAULT_TEACHER_PREFIX = (
+ "Solve the following math problem concisely and correctly. Be direct — avoid "
+ "unnecessary elaboration, redundant steps, or restating the problem. Focus only "
+ "on the key reasoning steps needed to reach the answer.\n\n"
+ "The last line of your response should be of the form Answer: $Answer (without "
+ "quotes) where $Answer is the answer to the problem.\n\n"
+)
+DEFAULT_TEACHER_SUFFIX = '\n\nRemember to put your answer on its own line after "Answer:".'
+
+_TOKENIZER = None
+
+
+def _get_tokenizer(args):
+ """Lazily load and cache the tokenizer (shared with the student model)."""
+ global _TOKENIZER
+ if _TOKENIZER is None:
+ from slime.utils.processing_utils import load_tokenizer
+
+ _TOKENIZER = load_tokenizer(args.hf_checkpoint, trust_remote_code=True)
+ return _TOKENIZER
+
+
+async def _post(url, payload):
+ """Indirection over slime's retrying HTTP helper (patchable in tests)."""
+ from slime.utils.http_utils import post
+
+ return await post(url, payload)
+
+
+def _compute_correct(response: str, label: str) -> float:
+ """Metrics-only math correctness via slime's math_dapo scorer (paper footnote 1)."""
+ from slime.rollout.rm_hub.math_dapo_utils import compute_score
+
+ return 1.0 if compute_score(response, label) > 0 else 0.0
+
+
+def build_teacher_prompt_ids(args, question: str) -> list[int]:
+ """Tokenize the conciseness-conditioned teacher prompt for ``question``."""
+ prefix = getattr(args, "crisp_teacher_prompt_prefix", None) or DEFAULT_TEACHER_PREFIX
+ suffix = getattr(args, "crisp_teacher_prompt_suffix", None) or DEFAULT_TEACHER_SUFFIX
+ messages = [{"role": "user", "content": prefix + question + suffix}]
+ tokenizer = _get_tokenizer(args)
+ return tokenizer.apply_chat_template(
+ messages,
+ tokenize=True,
+ add_generation_prompt=True,
+ **(getattr(args, "apply_chat_template_kwargs", None) or {}),
+ )
+
+
+async def reward_func(args, sample, **kwargs):
+ """Score one student rollout under the conciseness-conditioned teacher.
+
+ Sets ``sample.teacher_log_probs`` (len == ``sample.response_length``) and
+ returns math correctness in {0.0, 1.0} — metrics only, zeroed for training
+ by ``post_process_rewards``.
+ """
+ if sample.response_length == 0:
+ sample.teacher_log_probs = []
+ if getattr(args, "crisp_kl_mode", "sampled") == "full":
+ sample.teacher_top_ids = []
+ sample.teacher_top_logprobs = []
+ return 0.0
+
+ question = (sample.metadata or {}).get("question")
+ assert question, (
+ "CRISP reward_func requires sample.metadata['question'] (the bare problem "
+ "text). Prepare data with slime_crisp/prepare_crisp_slime_data.py."
+ )
+
+ teacher_prompt_ids = build_teacher_prompt_ids(args, question)
+ response_tokens = list(sample.tokens[-sample.response_length :])
+
+ url = getattr(args, "crisp_teacher_url", None) or args.rm_url
+ assert url, "CRISP requires --rm-url (or crisp_teacher_url) pointing at the teacher /generate endpoint"
+
+ # Prefill-only scoring pass: teacher prompt + the student's exact response
+ # token ids. logprob_start_len=0 mirrors slime's OPD example; we trim to
+ # the response span below.
+ payload = {
+ "input_ids": teacher_prompt_ids + response_tokens,
+ "sampling_params": {
+ "temperature": 0.0,
+ "max_new_tokens": 0,
+ "skip_special_tokens": False,
+ },
+ "return_logprob": True,
+ "logprob_start_len": 0,
+ }
+ # Full-KL mode (milestone 2): also fetch the teacher's top-K distribution
+ # per position, consumed by crisp_full_kl_loss.full_kl_loss_function.
+ full_kl = getattr(args, "crisp_kl_mode", "sampled") == "full"
+ if full_kl:
+ payload["top_logprobs_num"] = int(getattr(args, "crisp_teacher_topk", 256))
+
+ output = await _post(url, payload)
+
+ # input_token_logprobs entries are [logprob, token_id, ...]; the first
+ # entry of the sequence has logprob None, but it can never fall inside the
+ # response span because the teacher prompt is non-empty.
+ entries = output["meta_info"]["input_token_logprobs"][-sample.response_length :]
+ echoed_ids = [e[1] for e in entries]
+ assert echoed_ids == response_tokens, (
+ "CRISP alignment violation: token ids echoed by the teacher server do not "
+ "match the student's response tokens. Teacher and student must share the "
+ f"exact response token ids (got {len(echoed_ids)} vs {len(response_tokens)} "
+ "tokens; first mismatch at index "
+ f"{next((i for i, (a, b) in enumerate(zip(echoed_ids, response_tokens)) if a != b), -1)})."
+ )
+ teacher_log_probs = [e[0] for e in entries]
+ assert all(lp is not None for lp in teacher_log_probs), (
+ "Teacher server returned None log-probs inside the response span; "
+ "check logprob_start_len handling."
+ )
+ sample.teacher_log_probs = teacher_log_probs
+
+ if full_kl:
+ # input_top_logprobs: per input position, a list of [logprob, token_id, ...]
+ # entries (length top_logprobs_num). Trim to the response span like
+ # input_token_logprobs.
+ top_entries = output["meta_info"]["input_top_logprobs"][-sample.response_length :]
+ k = payload["top_logprobs_num"]
+ assert all(pos is not None and len(pos) == k for pos in top_entries), (
+ "input_top_logprobs malformed inside the response span (None or wrong K); "
+ "check the deployed sglang version's top_logprobs_num support."
+ )
+ sample.teacher_top_logprobs = [[e[0] for e in pos] for pos in top_entries]
+ sample.teacher_top_ids = [[e[1] for e in pos] for pos in top_entries]
+
+ if sample.label is None:
+ return 0.0
+ return _compute_correct(sample.response, sample.label)
+
+
+def post_process_rewards(args, samples, **kwargs):
+ """Zero task rewards for pure distillation; keep correctness for logging.
+
+ Returns ``(raw_rewards, train_rewards)``: raw correctness goes into the
+ logged ``raw_reward`` (training-accuracy curve, paper Fig. 2); training
+ rewards are all 0.0 so advantages come solely from the OPD KL penalty.
+ """
+ missing = [
+ i for i, s in enumerate(samples) if s.teacher_log_probs is None and s.response_length > 0
+ ]
+ assert not missing, (
+ f"{len(missing)} samples are missing teacher_log_probs (indices {missing[:5]}...). "
+ "Was reward_func wired via --custom-rm-path?"
+ )
+ if getattr(args, "crisp_kl_mode", "sampled") == "full":
+ missing_topk = [
+ i for i, s in enumerate(samples) if s.teacher_top_ids is None and s.response_length > 0
+ ]
+ assert not missing_topk, (
+ f"crisp_kl_mode=full but {len(missing_topk)} samples are missing teacher_top_ids "
+ f"(indices {missing_topk[:5]}...). reward_func must run with the same custom config."
+ )
+ raw_rewards = [float(s.reward) if s.reward is not None else 0.0 for s in samples]
+ return raw_rewards, [0.0] * len(samples)
+
+
+def _refresh_teacher_weights(args, rollout_id):
+ """POST /update_weights_from_disk to the teacher server (theta_tilde <- theta)."""
+ import os
+
+ import requests
+
+ hf_path = getattr(args, "crisp_teacher_hf_path", None)
+ assert hf_path, (
+ "crisp_teacher_update_interval is set but crisp_teacher_hf_path is not. "
+ "Set both in the --custom-config-path YAML (and point --save-hf at the same path)."
+ )
+ if not os.path.exists(os.path.join(hf_path, "config.json")):
+ raise FileNotFoundError(
+ f"Teacher refresh at rollout {rollout_id}: no HF dump at {hf_path}. "
+ "Ensure --save-hf points there and --save-interval == crisp_teacher_update_interval."
+ )
+
+ url = getattr(args, "crisp_teacher_url", None) or args.rm_url
+ base = url.rsplit("/", 1)[0]
+ logger.info(f"CRISP teacher refresh at rollout {rollout_id}: loading {hf_path} into {base}")
+ resp = requests.post(f"{base}/update_weights_from_disk", json={"model_path": hf_path}, timeout=1200)
+ resp.raise_for_status()
+ body = resp.json()
+ if not body.get("success", True):
+ raise RuntimeError(f"Teacher refresh failed: {body}")
+ logger.info(f"CRISP teacher refresh complete at rollout {rollout_id}")
+
+
+def generate_rollout(args, rollout_id, data_source, evaluation=False):
+ """slime's default sglang rollout + CRISP periodic teacher refresh.
+
+ Refreshes the teacher *before* generating rollout ``M, 2M, ...`` — by then
+ the actor has saved its HF dump at rollout ``M-1`` (slime saves when
+ ``(rollout_id + 1) % save_interval == 0``), so the teacher becomes the
+ student after exactly M training steps: progressive compression.
+ """
+ interval = getattr(args, "crisp_teacher_update_interval", 0) or 0
+ if not evaluation and interval > 0 and rollout_id > 0 and rollout_id % interval == 0:
+ _refresh_teacher_weights(args, rollout_id)
+
+ from slime.rollout.sglang_rollout import generate_rollout as default_generate_rollout
+
+ return default_generate_rollout(args, rollout_id, data_source, evaluation=evaluation)
diff --git a/workspace/slime_crisp/prepare_crisp_slime_data.py b/workspace/slime_crisp/prepare_crisp_slime_data.py
new file mode 100644
index 0000000..3cccc6a
--- /dev/null
+++ b/workspace/slime_crisp/prepare_crisp_slime_data.py
@@ -0,0 +1,125 @@
+"""Convert DAPO-Math-17k-dedup parquet into slime jsonl for CRISP training.
+
+Each output row:
+ {
+ "prompt": "", # student prompt (unchanged)
+ "label": "", # metrics-only verification
+ "metadata": {"question": ""} # used to build the teacher prompt
+ }
+
+slime wraps ``prompt`` into a user message and applies the chat template
+(``--apply-chat-template --input-key prompt --label-key label``); the CRISP
+reward_func rebuilds the teacher prompt from ``metadata.question`` at rollout
+time. Extraction logic mirrors workspace/src/data/prepare_length_prune_data.py
+so the student/teacher prompt pair is identical to the verl pipeline.
+
+Split parity: same fixed-seed pandas shuffle + first-80% train split as
+workspace/src/data/prepare_length_prune_data.py::sample_and_split, so the
+training set is row-for-row identical to the verl pipeline's.
+
+Usage:
+ python prepare_crisp_slime_data.py \
+ --input-parquet ../data/DAPO-Math-17k-dedup/distinct-prompts-with-rewards.parquet \
+ --output ../data/crisp_slime/dapo_math_crisp_train.jsonl
+"""
+
+import argparse
+import json
+import os
+
+import pandas as pd
+
+DEFAULT_SEED = 42
+
+# The instruction header DAPO-Math prepends to every question; stripping it
+# yields the bare question for teacher-prompt construction. Must match
+# workspace/src/data/prepare_length_prune_data.py.
+DAPO_PREFIX = (
+ "Solve the following math problem step by step. "
+ "The last line of your response should be of the form Answer: "
+ "$Answer (without quotes) where $Answer is the answer to the problem.\n\n"
+)
+
+
+def extract_user_content(prompt) -> str:
+ """Return the original DAPO user-message content (the student prompt)."""
+ if isinstance(prompt, str):
+ messages = json.loads(prompt)
+ else:
+ messages = list(prompt)
+ return messages[0]["content"]
+
+
+def extract_question(content: str) -> str:
+ """Strip the DAPO instruction header to get the bare question text."""
+ if content.startswith(DAPO_PREFIX):
+ return content[len(DAPO_PREFIX) :].strip()
+ return content.strip()
+
+
+def extract_ground_truth(reward_model) -> str:
+ if isinstance(reward_model, str):
+ reward_model = json.loads(reward_model)
+ if isinstance(reward_model, dict):
+ return str(reward_model.get("ground_truth", ""))
+ return ""
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--input-parquet", required=True)
+ parser.add_argument("--output", required=True, help="Output train jsonl path")
+ parser.add_argument(
+ "--val-output", default=None, help="Optional val jsonl path (the held-out 1-train_frac)"
+ )
+ parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
+ parser.add_argument(
+ "--train-frac",
+ type=float,
+ default=0.8,
+ help="Train fraction (default 0.8, matching the verl pipeline's split)",
+ )
+ parser.add_argument(
+ "--max-rows", type=int, default=None, help="Optional cap on TRAIN rows (smoke tests)"
+ )
+ args = parser.parse_args()
+
+ df = pd.read_parquet(args.input_parquet)
+ print(f"Loaded {len(df)} rows from {args.input_parquet}")
+
+ # Same fixed-seed shuffle + first-80% split as the verl pipeline
+ # (prepare_length_prune_data.sample_and_split) -> identical train set.
+ df = df.sample(frac=1.0, random_state=args.seed).reset_index(drop=True)
+ n_train = int(len(df) * args.train_frac)
+ train_df, val_df = df.iloc[:n_train], df.iloc[n_train:]
+ if args.max_rows:
+ train_df = train_df.iloc[: args.max_rows]
+
+ def write_jsonl(split_df, path):
+ os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+ n_written = n_skipped = 0
+ with open(path, "w", encoding="utf-8") as f:
+ for _, row in split_df.iterrows():
+ content = extract_user_content(row["prompt"])
+ question = extract_question(content)
+ gt = extract_ground_truth(row["reward_model"])
+ if not question or not gt:
+ n_skipped += 1
+ continue
+ f.write(
+ json.dumps(
+ {"prompt": content, "label": gt, "metadata": {"question": question}},
+ ensure_ascii=False,
+ )
+ + "\n"
+ )
+ n_written += 1
+ print(f"Wrote {n_written} rows to {path} ({n_skipped} skipped: missing question/GT)")
+
+ write_jsonl(train_df, args.output)
+ if args.val_output:
+ write_jsonl(val_df, args.val_output)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/workspace/slime_crisp/run-qwen3-8b-crisp.sh b/workspace/slime_crisp/run-qwen3-8b-crisp.sh
new file mode 100644
index 0000000..42fa5c6
--- /dev/null
+++ b/workspace/slime_crisp/run-qwen3-8b-crisp.sh
@@ -0,0 +1,196 @@
+#!/bin/bash
+# =============================================================================
+# CRISP on slime — Qwen3-8B, milestone 1
+#
+# Teacher = the SAME Qwen3-8B conditioned on a conciseness prompt, served on a
+# dedicated GPU; refreshed from the student every M=50 rollouts (progressive
+# compression). Student trains via slime's sglang-mode OPD (sampled-token
+# reverse KL as an advantage penalty), reward = 0 (pure distillation).
+#
+# Paper recipe (Table 2 / §5.1): batch 32, lr 1e-6, temp 1.0, top-p 1.0,
+# 8192-token rollouts, single rollout per prompt, ~100 steps, M=50.
+#
+# GPU layout (8 GPUs): actor 4 | rollout 3 | teacher 1 (GPU 7).
+# Uses synchronous train.py — REQUIRED for correct teacher-refresh ordering.
+#
+# KL estimator (CRISP_KL_MODE):
+# sampled (default) — milestone 1: sampled-token reverse KL via slime's OPD
+# advantage penalty (--use-opd).
+# full — milestone 2 (FULL_KL_PLAN.md): bucketed full-vocab
+# reverse KL over teacher top-K as a custom Megatron
+# loss. Requires the patched slime (crisp branch).
+#
+# Prereqs:
+# 1. HF checkpoint at $HF_CKPT, converted Megatron ckpt at $MCORE_CKPT
+# (tools/convert_hf_to_torch_dist.py, see slime OPD example README)
+# 2. Data: python prepare_crisp_slime_data.py --input-parquet ... --output $PROMPT_DATA
+# 3. This repo's workspace/ dir on PYTHONPATH (for slime_crisp.*)
+# =============================================================================
+
+set -ex
+
+# ---- Paths (override via env) ----
+HF_CKPT=${HF_CKPT:-/root/Qwen3-8B}
+MCORE_CKPT=${MCORE_CKPT:-/root/Qwen3-8B_torch_dist}
+SAVE_DIR=${SAVE_DIR:-/root/Qwen3-8B_crisp_slime}
+TEACHER_HF_PATH=${TEACHER_HF_PATH:-/root/crisp_teacher_hf} # must match crisp_config.yaml
+PROMPT_DATA=${PROMPT_DATA:-/root/data/dapo_math_crisp.jsonl}
+SLIME_DIR=${SLIME_DIR:-/root/slime}
+MEGATRON_DIR=${MEGATRON_DIR:-/root/Megatron-LM}
+CRISP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # .../workspace/slime_crisp
+WORKSPACE_DIR="$(dirname "${CRISP_DIR}")" # .../workspace
+
+# ---- Teacher server (same base checkpoint, dedicated GPU) ----
+TEACHER_IP=${TEACHER_IP:-127.0.0.1}
+TEACHER_PORT=${TEACHER_PORT:-13141}
+LOG_FILE="/tmp/sglang_teacher_$(date +%s).log"
+
+CUDA_VISIBLE_DEVICES=7 python3 -m sglang.launch_server \
+ --model-path ${HF_CKPT} \
+ --host 0.0.0.0 \
+ --port ${TEACHER_PORT} \
+ --tp 1 \
+ --chunked-prefill-size 4096 \
+ --mem-fraction-static 0.7 \
+ > "${LOG_FILE}" 2>&1 &
+
+echo "Waiting for teacher server..."
+until curl -sf http://${TEACHER_IP}:${TEACHER_PORT}/health_generate > /dev/null; do
+ tail -n 5 "${LOG_FILE}" || true
+ sleep 5
+done
+echo "Teacher up at ${TEACHER_IP}:${TEACHER_PORT}"
+
+export PYTHONUNBUFFERED=1
+source "${SLIME_DIR}/scripts/models/qwen3-8B.sh"
+
+CKPT_ARGS=(
+ --hf-checkpoint ${HF_CKPT}
+ --load ${MCORE_CKPT}
+ --save ${SAVE_DIR}
+ --save-interval 50 # == crisp_teacher_update_interval (M)
+ --save-hf ${TEACHER_HF_PATH} # fixed path: each save overwrites; teacher reloads it
+)
+
+ROLLOUT_ARGS=(
+ --prompt-data ${PROMPT_DATA}
+ --input-key prompt
+ --label-key label
+ --apply-chat-template
+ # Pin Qwen3 thinking mode (mirrors verl opsd_trainer.yaml; guards against
+ # tokenizer default drift). Also forwarded to the teacher prompt in crisp_opd.
+ --apply-chat-template-kwargs '{"enable_thinking": true}'
+ --rollout-shuffle
+ --num-rollout 100 # paper: step 100 is the sweet spot
+ --rollout-batch-size 32
+ --n-samples-per-prompt 1 # CRISP: single rollout per prompt
+ --rollout-max-response-len 8192
+ --rollout-temperature 1.0
+ --rollout-top-p 1.0
+ --global-batch-size 32 # 1 optimizer step per rollout
+)
+
+CRISP_KL_MODE=${CRISP_KL_MODE:-sampled}
+
+if [ "${CRISP_KL_MODE}" = "full" ]; then
+ CRISP_CONFIG=${CRISP_DIR}/crisp_config_full_kl.yaml
+ DISTILL_ARGS=(
+ # Distribution-level loss replaces the PG/OPD pipeline entirely.
+ --loss-type custom_loss
+ --custom-loss-function-path slime_crisp.crisp_full_kl_loss.full_kl_loss_function
+ # Skip advantages AND the old-log-prob forward pass — the custom loss
+ # needs only the training forward.
+ --disable-compute-advantages-and-returns
+ # Recompute the loss fn in backward: frees the per-chunk softmax graph,
+ # bounding loss memory to the [T, V] logits already resident.
+ --recompute-loss-function
+ --advantage-estimator grpo # unused (advantages disabled); schema only
+ --entropy-coef 0.00
+ --calculate-per-token-loss # global token mean == verl normalization
+ )
+else
+ CRISP_CONFIG=${CRISP_DIR}/crisp_config.yaml
+ DISTILL_ARGS=(
+ --advantage-estimator grpo # degenerate at n=1 + reward 0: pure OPD penalty
+ --use-opd
+ --opd-type sglang
+ --opd-kl-coef 1.0
+ --entropy-coef 0.00
+ # Global token-mean loss reduction: closest match to verl OPSD's
+ # kl_sum / n_tokens normalization (default slime reduction is per-rollout mean).
+ --calculate-per-token-loss
+ )
+fi
+
+CRISP_ARGS=(
+ --custom-rm-path slime_crisp.crisp_opd.reward_func
+ --custom-reward-post-process-path slime_crisp.crisp_opd.post_process_rewards
+ --rollout-function-path slime_crisp.crisp_opd.generate_rollout
+ --rm-url http://${TEACHER_IP}:${TEACHER_PORT}/generate
+ --custom-config-path ${CRISP_CONFIG}
+)
+
+PERF_ARGS=(
+ --tensor-model-parallel-size 2
+ --sequence-parallel
+ --pipeline-model-parallel-size 1
+ --context-parallel-size 1
+ --recompute-granularity full
+ --recompute-method uniform
+ --recompute-num-layers 1
+ --use-dynamic-batch-size
+ --max-tokens-per-gpu 16384
+)
+
+OPTIMIZER_ARGS=(
+ --optimizer adam
+ --lr 1e-6
+ --lr-decay-style constant
+ # Match verl's AdamW defaults (train_opsd.sh leaves these unset).
+ --weight-decay 0.01
+ --adam-beta1 0.9
+ --adam-beta2 0.999
+)
+
+SGLANG_ARGS=(
+ --rollout-num-gpus-per-engine 1
+ --sglang-mem-fraction-static 0.75
+)
+
+MISC_ARGS=(
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --accumulate-allreduce-grads-in-fp32
+ --attention-softmax-in-fp32
+ --attention-backend flash
+)
+
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 7 --disable-usage-stats \
+ --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+ --runtime-env-json="{
+ \"env_vars\": {
+ \"PYTHONPATH\": \"${MEGATRON_DIR}:${WORKSPACE_DIR}\",
+ \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\"
+ }
+ }" \
+ -- python3 ${SLIME_DIR}/train.py \
+ --actor-num-nodes 1 \
+ --actor-num-gpus-per-node 4 \
+ --rollout-num-gpus 3 \
+ ${MODEL_ARGS[@]} \
+ ${CKPT_ARGS[@]} \
+ ${ROLLOUT_ARGS[@]} \
+ ${OPTIMIZER_ARGS[@]} \
+ ${DISTILL_ARGS[@]} \
+ ${PERF_ARGS[@]} \
+ ${SGLANG_ARGS[@]} \
+ ${MISC_ARGS[@]} \
+ ${CRISP_ARGS[@]}
+
+# ---- cleanup ----
+pkill -9 sglang || true
+sleep 3
+ray stop --force || true
diff --git a/workspace/slime_crisp/test_crisp_opd.py b/workspace/slime_crisp/test_crisp_opd.py
new file mode 100644
index 0000000..477af77
--- /dev/null
+++ b/workspace/slime_crisp/test_crisp_opd.py
@@ -0,0 +1,310 @@
+"""Unit tests for the CRISP teacher-scoring logic (no GPU / no slime runtime).
+
+Heavy imports in crisp_opd are lazy, so these tests run with only the module's
+pure-python logic: payload construction, response-span trimming, alignment
+checks, and reward post-processing.
+
+Run: pytest workspace/slime_crisp/test_crisp_opd.py
+"""
+
+import asyncio
+from argparse import Namespace
+from dataclasses import dataclass, field
+
+import pytest
+
+import crisp_opd
+
+
+# ---------------------------------------------------------------------------
+# Fakes
+# ---------------------------------------------------------------------------
+
+
+class FakeTokenizer:
+ """apply_chat_template -> deterministic ids: [100, len(content)]."""
+
+ def apply_chat_template(self, messages, tokenize, add_generation_prompt, **kwargs):
+ assert tokenize and add_generation_prompt
+ assert messages[0]["role"] == "user"
+ self.last_content = messages[0]["content"]
+ self.last_kwargs = kwargs
+ return [100, len(messages[0]["content"]) % 1000]
+
+
+@dataclass
+class FakeSample:
+ tokens: list = field(default_factory=list)
+ response_length: int = 0
+ response: str = ""
+ label: str | None = None
+ reward: float | None = None
+ metadata: dict = field(default_factory=dict)
+ teacher_log_probs: list | None = None
+ teacher_top_ids: list | None = None
+ teacher_top_logprobs: list | None = None
+
+
+def make_args(**overrides):
+ base = dict(
+ hf_checkpoint="fake",
+ rm_url="http://teacher/generate",
+ apply_chat_template_kwargs={"enable_thinking": True},
+ )
+ base.update(overrides)
+ return Namespace(**base)
+
+
+def fake_post_factory(captured, prompt_len, response_tokens, logprob_value=-0.5, corrupt_index=None):
+ """Return an async fake for crisp_opd._post echoing sglang's logprob format."""
+
+ async def fake_post(url, payload):
+ captured["url"] = url
+ captured["payload"] = payload
+ ids = list(response_tokens)
+ if corrupt_index is not None:
+ ids[corrupt_index] += 1
+ entries = [[None, payload["input_ids"][0]]]
+ for tok in payload["input_ids"][1 : prompt_len]:
+ entries.append([-1.0, tok]) # prompt span (discarded by trimming)
+ for tok in ids:
+ entries.append([logprob_value, tok]) # response span
+ meta_info = {"input_token_logprobs": entries}
+ if "top_logprobs_num" in payload:
+ k = payload["top_logprobs_num"]
+ # per position: K entries of [logprob, token_id]; None on position 0
+ meta_info["input_top_logprobs"] = [None] + [
+ [[-0.1 * (j + 1), 1000 + j] for j in range(k)]
+ for _ in range(len(payload["input_ids"]) - 1)
+ ]
+ return {"meta_info": meta_info}
+
+ return fake_post
+
+
+@pytest.fixture(autouse=True)
+def fake_tokenizer(monkeypatch):
+ tok = FakeTokenizer()
+ monkeypatch.setattr(crisp_opd, "_TOKENIZER", tok)
+ yield tok
+ crisp_opd._TOKENIZER = None
+
+
+# ---------------------------------------------------------------------------
+# reward_func
+# ---------------------------------------------------------------------------
+
+
+def test_reward_func_scores_response_span(monkeypatch, fake_tokenizer):
+ response_tokens = [11, 12, 13]
+ sample = FakeSample(
+ tokens=[1, 2, 3] + response_tokens, # student prompt ids + response ids
+ response_length=3,
+ response="... Answer: 42",
+ label="42",
+ metadata={"question": "What is 6*7?"},
+ )
+ captured = {}
+ monkeypatch.setattr(crisp_opd, "_post", fake_post_factory(captured, prompt_len=2, response_tokens=response_tokens))
+ monkeypatch.setattr(crisp_opd, "_compute_correct", lambda response, label: 1.0)
+
+ reward = asyncio.run(crisp_opd.reward_func(make_args(), sample))
+
+ # teacher logprobs trimmed to exactly the response span
+ assert sample.teacher_log_probs == [-0.5, -0.5, -0.5]
+ assert reward == 1.0
+ # payload: teacher prompt ids (FakeTokenizer -> 2 ids) + student response ids
+ assert captured["payload"]["input_ids"][-3:] == response_tokens
+ assert len(captured["payload"]["input_ids"]) == 2 + 3
+ assert captured["payload"]["sampling_params"]["max_new_tokens"] == 0
+ assert captured["payload"]["return_logprob"] is True
+ assert captured["url"] == "http://teacher/generate"
+ # teacher prompt uses the conciseness instruction around the bare question
+ assert fake_tokenizer.last_content.startswith(crisp_opd.DEFAULT_TEACHER_PREFIX)
+ assert "What is 6*7?" in fake_tokenizer.last_content
+ assert fake_tokenizer.last_content.endswith(crisp_opd.DEFAULT_TEACHER_SUFFIX)
+ # chat-template kwargs (e.g. enable_thinking) forwarded to the teacher side
+ assert fake_tokenizer.last_kwargs == {"enable_thinking": True}
+
+
+def test_reward_func_alignment_violation_raises(monkeypatch):
+ response_tokens = [11, 12, 13]
+ sample = FakeSample(
+ tokens=[1, 2] + response_tokens,
+ response_length=3,
+ metadata={"question": "q"},
+ label="1",
+ )
+ monkeypatch.setattr(
+ crisp_opd,
+ "_post",
+ fake_post_factory({}, prompt_len=2, response_tokens=response_tokens, corrupt_index=1),
+ )
+
+ with pytest.raises(AssertionError, match="alignment violation"):
+ asyncio.run(crisp_opd.reward_func(make_args(), sample))
+
+
+def test_reward_func_requires_question_metadata(monkeypatch):
+ sample = FakeSample(tokens=[1, 11], response_length=1, metadata={})
+ with pytest.raises(AssertionError, match="question"):
+ asyncio.run(crisp_opd.reward_func(make_args(), sample))
+
+
+def test_reward_func_empty_response(monkeypatch):
+ sample = FakeSample(tokens=[1, 2], response_length=0, metadata={"question": "q"})
+ reward = asyncio.run(crisp_opd.reward_func(make_args(), sample))
+ assert reward == 0.0
+ assert sample.teacher_log_probs == []
+
+
+def test_reward_func_prompt_override(monkeypatch, fake_tokenizer):
+ response_tokens = [11]
+ sample = FakeSample(
+ tokens=[1, 2] + response_tokens, response_length=1, metadata={"question": "q"}, label=None
+ )
+ monkeypatch.setattr(crisp_opd, "_post", fake_post_factory({}, prompt_len=2, response_tokens=response_tokens))
+ args = make_args(crisp_teacher_prompt_prefix="BE BRIEF: ", crisp_teacher_prompt_suffix=" END")
+
+ reward = asyncio.run(crisp_opd.reward_func(args, sample))
+
+ assert fake_tokenizer.last_content == "BE BRIEF: q END"
+ assert reward == 0.0 # label None -> metrics skipped
+
+
+def test_reward_func_full_kl_mode_fetches_topk(monkeypatch):
+ response_tokens = [11, 12, 13]
+ sample = FakeSample(
+ tokens=[1, 2] + response_tokens,
+ response_length=3,
+ metadata={"question": "q"},
+ label=None,
+ )
+ captured = {}
+ monkeypatch.setattr(crisp_opd, "_post", fake_post_factory(captured, prompt_len=2, response_tokens=response_tokens))
+ args = make_args(crisp_kl_mode="full", crisp_teacher_topk=4)
+
+ asyncio.run(crisp_opd.reward_func(args, sample))
+
+ assert captured["payload"]["top_logprobs_num"] == 4
+ # top-K trimmed to exactly the response span, [R, K]
+ assert len(sample.teacher_top_ids) == 3
+ assert len(sample.teacher_top_ids[0]) == 4
+ assert sample.teacher_top_ids[0] == [1000, 1001, 1002, 1003]
+ assert sample.teacher_top_logprobs[0] == [-0.1, -0.2, -0.30000000000000004, -0.4]
+ # sampled log-probs still collected alongside (for the kl_sampled diagnostic)
+ assert sample.teacher_log_probs == [-0.5, -0.5, -0.5]
+
+
+def test_reward_func_sampled_mode_skips_topk(monkeypatch):
+ response_tokens = [11]
+ sample = FakeSample(tokens=[1, 2] + response_tokens, response_length=1, metadata={"question": "q"})
+ captured = {}
+ monkeypatch.setattr(crisp_opd, "_post", fake_post_factory(captured, prompt_len=2, response_tokens=response_tokens))
+
+ asyncio.run(crisp_opd.reward_func(make_args(), sample))
+
+ assert "top_logprobs_num" not in captured["payload"]
+ assert sample.teacher_top_ids is None
+
+
+# ---------------------------------------------------------------------------
+# post_process_rewards
+# ---------------------------------------------------------------------------
+
+
+def test_post_process_zeroes_training_rewards():
+ samples = [
+ FakeSample(reward=1.0, response_length=3, teacher_log_probs=[-0.1] * 3),
+ FakeSample(reward=0.0, response_length=2, teacher_log_probs=[-0.2] * 2),
+ FakeSample(reward=None, response_length=0, teacher_log_probs=[]),
+ ]
+ raw, train = crisp_opd.post_process_rewards(make_args(), samples)
+ assert raw == [1.0, 0.0, 0.0] # correctness preserved for logging
+ assert train == [0.0, 0.0, 0.0] # pure distillation
+
+
+def test_post_process_detects_missing_teacher_logprobs():
+ samples = [FakeSample(reward=1.0, response_length=3, teacher_log_probs=None)]
+ with pytest.raises(AssertionError, match="teacher_log_probs"):
+ crisp_opd.post_process_rewards(make_args(), samples)
+
+
+def test_post_process_full_mode_requires_topk():
+ samples = [
+ FakeSample(reward=1.0, response_length=3, teacher_log_probs=[-0.1] * 3, teacher_top_ids=None)
+ ]
+ with pytest.raises(AssertionError, match="teacher_top_ids"):
+ crisp_opd.post_process_rewards(make_args(crisp_kl_mode="full"), samples)
+
+
+# ---------------------------------------------------------------------------
+# teacher refresh gating
+# ---------------------------------------------------------------------------
+
+
+def test_refresh_gating(monkeypatch):
+ calls = []
+ monkeypatch.setattr(crisp_opd, "_refresh_teacher_weights", lambda args, rid: calls.append(rid))
+ inner_calls = []
+
+ def fake_default(args, rollout_id, data_source, evaluation=False):
+ inner_calls.append(rollout_id)
+ return "rollout-output"
+
+ import sys
+ import types
+
+ fake_mod = types.ModuleType("slime.rollout.sglang_rollout")
+ fake_mod.generate_rollout = fake_default
+ fake_pkg_rollout = types.ModuleType("slime.rollout")
+ fake_pkg = types.ModuleType("slime")
+ monkeypatch.setitem(sys.modules, "slime", fake_pkg)
+ monkeypatch.setitem(sys.modules, "slime.rollout", fake_pkg_rollout)
+ monkeypatch.setitem(sys.modules, "slime.rollout.sglang_rollout", fake_mod)
+
+ args = make_args(crisp_teacher_update_interval=50)
+ for rid in [0, 1, 49, 50, 99, 100]:
+ assert crisp_opd.generate_rollout(args, rid, data_source=None) == "rollout-output"
+ assert calls == [50, 100] # not at 0; only at multiples of M
+ assert inner_calls == [0, 1, 49, 50, 99, 100]
+
+ # eval never refreshes; interval 0 disables
+ crisp_opd.generate_rollout(args, 50, data_source=None, evaluation=True)
+ crisp_opd.generate_rollout(make_args(crisp_teacher_update_interval=0), 50, data_source=None)
+ assert calls == [50, 100]
+
+
+# ---------------------------------------------------------------------------
+# Parity with the verl pipeline (workspace/src, workspace/config)
+# ---------------------------------------------------------------------------
+
+
+def test_teacher_prompt_matches_prompts_json():
+ """The conciseness instruction must stay byte-identical to the verl pipeline's."""
+ import json
+ import pathlib
+
+ prompts_path = pathlib.Path(__file__).resolve().parents[1] / "config" / "prompts.json"
+ cfg = json.loads(prompts_path.read_text())["length_prune_teacher"]
+ assert crisp_opd.DEFAULT_TEACHER_PREFIX == cfg["prefix"]
+ assert crisp_opd.DEFAULT_TEACHER_SUFFIX == cfg["suffix"]
+
+
+def test_dapo_header_matches_verl_prep():
+ """Question extraction must strip the same DAPO header as the verl pipeline."""
+ import pathlib
+ import sys
+
+ verl_data_dir = str(pathlib.Path(__file__).resolve().parents[1] / "src" / "data")
+ sys.path.insert(0, verl_data_dir)
+ try:
+ import prepare_length_prune_data as verl_prep
+
+ import prepare_crisp_slime_data as slime_prep
+
+ dapo_prompt = [{"role": "user", "content": slime_prep.DAPO_PREFIX + "What is 2+2?"}]
+ assert verl_prep.extract_question_from_dapo(dapo_prompt) == "What is 2+2?"
+ assert slime_prep.extract_question(dapo_prompt[0]["content"]) == "What is 2+2?"
+ finally:
+ sys.path.remove(verl_data_dir)
diff --git a/workspace/slime_crisp/test_full_kl_loss.py b/workspace/slime_crisp/test_full_kl_loss.py
new file mode 100644
index 0000000..76eba10
--- /dev/null
+++ b/workspace/slime_crisp/test_full_kl_loss.py
@@ -0,0 +1,171 @@
+"""Unit tests for the bucketed full-vocab reverse KL (milestone 2).
+
+Pure-torch tests of the loss math and the TP=1 path of the vocab-parallel
+primitives — no GPU, no megatron, no distributed init.
+
+Run: pytest workspace/slime_crisp/test_full_kl_loss.py
+"""
+
+import pytest
+import torch
+
+from crisp_full_kl_loss import bucketed_reverse_kl, vocab_parallel_topk_log_probs
+
+V = 13 # toy vocab
+
+
+def make_dists(n_tokens=5, seed=0, vocab=V):
+ g = torch.Generator().manual_seed(seed)
+ student_logits = torch.randn(n_tokens, vocab, generator=g, dtype=torch.float64)
+ teacher_logits = torch.randn(n_tokens, vocab, generator=g, dtype=torch.float64)
+ return student_logits, teacher_logits
+
+
+def exact_reverse_kl(student_logits, teacher_logits):
+ """Direct full-vocab KL(q||p) — the verl objective (opsd_worker)."""
+ log_q = student_logits.log_softmax(-1)
+ log_p = teacher_logits.log_softmax(-1)
+ return (log_q.exp() * (log_q - log_p)).sum(-1)
+
+
+def bucketed_from_logits(student_logits, teacher_logits, k):
+ """Run the production pipeline: teacher top-K -> gather -> bucketed KL."""
+ log_p = teacher_logits.log_softmax(-1)
+ top_lps, top_ids = log_p.topk(k, dim=-1)
+ log_q_k = vocab_parallel_topk_log_probs(student_logits, top_ids, tp_group=None)
+ return bucketed_reverse_kl(log_q_k, top_lps)
+
+
+# ---------------------------------------------------------------------------
+# vocab_parallel_topk_log_probs (TP=1 path)
+# ---------------------------------------------------------------------------
+
+
+def test_topk_log_probs_matches_log_softmax_gather():
+ student_logits, _ = make_dists()
+ ids = torch.randint(0, V, (5, 4))
+ out = vocab_parallel_topk_log_probs(student_logits, ids, tp_group=None)
+ expected = student_logits.log_softmax(-1).gather(-1, ids)
+ torch.testing.assert_close(out, expected)
+
+
+def test_topk_log_probs_gradient():
+ student_logits, _ = make_dists()
+ leaf = student_logits.clone().requires_grad_(True)
+ ids = torch.randint(0, V, (5, 4))
+
+ def fn(x):
+ return vocab_parallel_topk_log_probs(x, ids, tp_group=None)
+
+ assert torch.autograd.gradcheck(fn, (leaf,), raise_exception=True)
+
+
+# ---------------------------------------------------------------------------
+# bucketed_reverse_kl: math properties
+# ---------------------------------------------------------------------------
+
+
+def test_equals_exact_kl_at_k_equals_v():
+ student_logits, teacher_logits = make_dists()
+ kl, coverage, q_tail = bucketed_from_logits(student_logits, teacher_logits, k=V)
+ exact = exact_reverse_kl(student_logits, teacher_logits)
+ # At K=V the tail buckets carry ~0 mass (clamped at eps), so the bucketed
+ # KL reduces to the exact full-vocab reverse KL.
+ torch.testing.assert_close(kl, exact, atol=1e-6, rtol=1e-6)
+ torch.testing.assert_close(coverage, torch.ones_like(coverage), atol=1e-6, rtol=0)
+
+
+def test_gradient_matches_exact_kl_at_k_equals_v():
+ """The verl-equivalence proof: at K=V, gradients match the direct loss."""
+ student_logits, teacher_logits = make_dists()
+
+ leaf_bucketed = student_logits.clone().requires_grad_(True)
+ kl, _, _ = bucketed_from_logits(leaf_bucketed, teacher_logits, k=V)
+ kl.mean().backward()
+
+ leaf_exact = student_logits.clone().requires_grad_(True)
+ exact_reverse_kl(leaf_exact, teacher_logits).mean().backward()
+
+ torch.testing.assert_close(leaf_bucketed.grad, leaf_exact.grad, atol=1e-6, rtol=1e-5)
+
+
+def test_nonnegative_and_zero_iff_equal():
+ student_logits, teacher_logits = make_dists(n_tokens=20, seed=3)
+ kl, _, _ = bucketed_from_logits(student_logits, teacher_logits, k=6)
+ assert (kl >= -1e-9).all()
+
+ # identical distributions -> zero at any K
+ kl_same, _, _ = bucketed_from_logits(student_logits, student_logits, k=6)
+ torch.testing.assert_close(kl_same, torch.zeros_like(kl_same), atol=1e-9, rtol=0)
+
+
+def test_monotone_nondecreasing_in_k():
+ """Coarsening loses information: KL_K is non-decreasing in K (up to eps)."""
+ student_logits, teacher_logits = make_dists(n_tokens=10, seed=7)
+ means = []
+ for k in (2, 4, 8, V):
+ kl, _, _ = bucketed_from_logits(student_logits, teacher_logits, k=k)
+ means.append(kl.mean().item())
+ for lo, hi in zip(means, means[1:], strict=False):
+ assert hi >= lo - 1e-7, f"KL_K not monotone: {means}"
+
+
+def test_lower_bounds_exact_kl():
+ student_logits, teacher_logits = make_dists(n_tokens=10, seed=11)
+ exact = exact_reverse_kl(student_logits, teacher_logits)
+ for k in (2, 4, 8):
+ kl, _, _ = bucketed_from_logits(student_logits, teacher_logits, k=k)
+ assert (kl <= exact + 1e-7).all()
+
+
+def test_self_distillation_regime_small_k_is_near_exact():
+ """Student ~ teacher on a shared peaked support (the CRISP regime).
+
+ The bucketed-KL gap is governed by q_tail — the STUDENT's mass outside the
+ teacher top-K — not by teacher coverage alone (reverse KL is an expectation
+ under q, so a diffuse student hides arbitrarily large per-token log-ratios
+ in the tail bucket). In CRISP, student and teacher are the same base model
+ (training loss ~1e-2), so q_tail is tiny and small K is near exact. The
+ loss logs q_tail per step as the fidelity diagnostic.
+ """
+ g = torch.Generator().manual_seed(5)
+ teacher_logits = torch.randn(8, V, generator=g, dtype=torch.float64) * 12.0 # peaked
+ student_logits = teacher_logits + 0.3 * torch.randn(8, V, generator=g, dtype=torch.float64)
+
+ kl, coverage, q_tail = bucketed_from_logits(student_logits, teacher_logits, k=4)
+ exact = exact_reverse_kl(student_logits, teacher_logits)
+
+ assert coverage.min() > 0.99
+ assert q_tail.max() < 0.01
+ torch.testing.assert_close(kl, exact, atol=5e-3, rtol=0.05)
+
+
+def test_diffuse_student_gap_is_bounded_by_tail():
+ """Counterpart: a diffuse student under-counts via the tail bucket, but the
+ bucketed KL stays a lower bound and q_tail flags the regime."""
+ g = torch.Generator().manual_seed(5)
+ student_logits = torch.randn(8, V, generator=g, dtype=torch.float64) # diffuse
+ teacher_logits = torch.randn(8, V, generator=g, dtype=torch.float64) * 12.0 # peaked
+
+ kl, coverage, q_tail = bucketed_from_logits(student_logits, teacher_logits, k=4)
+ exact = exact_reverse_kl(student_logits, teacher_logits)
+
+ assert coverage.min() > 0.99 # teacher coverage alone does NOT imply small gap...
+ assert (kl <= exact + 1e-7).all() # ...but the bucketed KL never overshoots
+ assert q_tail.max() > 0.3 # and q_tail exposes the diffuse-student regime
+
+
+def test_tail_term_pushes_mass_toward_teacher_support():
+ """Gradient sanity: student mass outside teacher top-K must be penalized."""
+ teacher_logits = torch.full((1, V), -10.0, dtype=torch.float64)
+ teacher_logits[0, :2] = 5.0 # teacher concentrated on tokens {0, 1}
+ student_logits = torch.zeros(1, V, dtype=torch.float64).requires_grad_(True) # uniform student
+
+ kl, _, q_tail = bucketed_from_logits(student_logits, teacher_logits, k=2)
+ kl.sum().backward()
+
+ # Increasing logits of teacher-supported tokens decreases the KL...
+ assert (student_logits.grad[0, :2] < 0).all()
+ # ...and increasing tail-token logits increases it.
+ assert (student_logits.grad[0, 2:] > 0).all()
+ assert q_tail.item() > 0.5 # uniform student leaves most mass in the tail