| # Recipe: distill Beam into your own model |
|
|
| This is the path that produced Dory-Beam-4B, written as steps you can reuse with another student model. Every |
| step names the script that does it. Costs are what we actually paid on Modal; Beam API calls used credits from |
| Reflection. Paths are relative to this `code/` folder. |
|
|
| ## 1. Build a prompt pool with checkable answers |
|
|
| `scripts/build_pool.py` draws math problems with closed-form answers from DeepMath-103K (difficulty ≥ 5) and |
| multiple-choice science questions from OpenScienceReasoning-2. It drops any prompt that shares a 13-gram with the |
| benchmarks you will report. Checkable answers matter more than anything else here: every later step filters on |
| them. |
|
|
| ## 2. Collect Beam's reasoning at `medium` |
|
|
| Two ways to call Beam, both used for this model: |
| - **Reflection API:** `scripts/generate.py` (driven nightly by `scripts/run_main.sh`). It resumes where it left |
| off and stops at the daily output-token budget. |
| - **Self-hosted open weights:** `modal_app/beam_selfhost.py` runs SGLang with NVFP4 Beam on 2× B300. |
| - It produces about 8–12k output tokens/s for about $17/h, roughly **$0.40–0.60 per million output tokens**. |
| - It takes any `reasoning_effort` (several at once with `"low,high"`). |
| - Resumed runs skip finished prompts. |
|
|
| Settings: temperature 0.7, top-p 0.99, max 16k tokens. Beam at medium writes about 1.8k tokens per answer on |
| these problems. |
|
|
| `scripts/verify.py` grades each answer: math-verify for math, a letter match for multiple choice. |
| `scripts/build_dataset.py` then keeps the **shortest correct** trace per prompt. Shortest vs random correct made no |
| difference in our runs. 62k prompts gave 52,660 training examples. |
|
|
| Cost: about $140 of self-hosted generation plus API credits. |
|
|
| ## 3. Collect the other effort levels (optional) |
|
|
| The same `beam_selfhost.py` answered 20k of the training prompts at `low`, `high` and `xhigh`. |
| `scripts/build_effort_dataset.py` then tags each example with the student's own effort syntax, here Dory's |
| `{reasoning effort: ...}`. It keeps half the medium examples untagged, so the default stays medium-like. |
|
|
| Cost: about $70. |
|
|
| Beam's levels are not evenly spaced on these problems: low runs 0.65× the length of medium, high 1.33×, and xhigh |
| only 1.08× of high. Your distilled levels will inherit that spacing. |
|
|
| ## 4. Fine-tune |
|
|
| The trainer is `scripts/train_dory.py`, launched by `modal_app/train_dory.py`: |
| - full-parameter SFT; |
| - lr 1e-5, packed 16k sequences with document masks, and a loss only on the assistant turn. |
|
|
| Dory-specific parts: Dory has no public training code, so `dory/modeling.py` is a port checked op-for-op against |
| vLLM (`scripts/check_dory.py`). Its nGPT weights are re-normalized after each step, and the loop count is sampled |
| per step so every loop setting keeps working. |
|
|
| For another student, use your usual SFT stack with the same data. Two epochs on medium took 2 h 10 min on 4× B200 |
| (about $67). One more epoch with effort tags took 40 min on 8× B200 (about $40). |
|
|
| ## 5. Blend with the base model |
|
|
| The fine-tune alone copied Beam's brevity but not its reliability: 2–3.5× shorter answers, 7–10 points lower on |
| AIME. Interpolating base and fine-tuned weights and then re-normalizing fixed that without more training |
| (`scripts/merge_dory.py`; drop the re-normalization for a non-nGPT model): |
| - 50/50 recovered most of the accuracy and kept the effort levels; |
| - 30/70 (30% fine-tune) tied the base model everywhere, at 1.8× fewer tokens. |
|
|
| Sweep the ratio on a small eval before committing. |
|
|
| ## 6. Evaluate honestly |
|
|
| - **Use a generous token limit.** We use 64k. A 32k limit cut off a tenth of base Dory's AIME answers and made |
| the distilled model look 3–7 points better than it was. `modal_app/eval_dory.py` can resume cut-off answers |
| (`continue_truncated`). |
| - **Pair everything.** Same problems and sample indices for every model, with bootstrap CIs over problems |
| (`scripts/audit_eval.py`). Treat "not significant" as a tie, not as a win. |
| - **Compare under a budget fairly.** Don't score answers cut off at a hard limit as wrong: close the thinking at |
| the budget and let the model answer (`budget_force`, `scripts/spawn_budget.py`, `scripts/summarize_budget.py`). |
| The hard cutoff made base Dory look far worse than it was. |
| - **Check what you didn't train on.** We checked IFEval (it held up) and coding (`modal_app/code_eval.py`), which |
| dropped 2–5 points. 7,000 test-verified Beam coding traces did not win it back. |
|
|
| All our evaluation, including experiments, came to about $180. |
|
|
| ## Pitfalls we hit |
|
|
| - **Self-hosted Beam (NVFP4 or FP8) matches the API's accuracy but writes about 35% longer reasoning.** Use one |
| teacher for all effort levels, or the levels mix two styles. |
| - **SGLang fell over on our eval run with several samples per request at a 30k limit.** Single-sample requests |
| were stable, so `beam_selfhost.py` can split samples into separate requests. Keep the prompt plus `max_tokens` |
| inside the context length; one long GPQA prompt overflowed a 32k context. |
|
|
| The generation and training steps above came to about $320 of Modal, plus the evaluation you choose to run. The |
| project's full bill was about $630; the rest went to ablations, extra evals and the experiments listed as not |
| helping on the model card. |
|
|
