把 R2.1 scalar-d512 sidecar 蒸餾回 OLMoE-1B-7B base
1 · 研究目標與假設
把已凍結的 R2.1 sidecar 能力「吸收」回 OLMoE base 的 MoE 模組,蒸餾完成後可拆掉 sidecar,
得到一個純 base 模型 olmoe_0.1,在 R2.1 heldout 上接近原本「base + sidecar 同時掛」的成績,
同時 R1(原能力)不退化。
核心數學身分(per-block, in-block self-consistent target):
target(h) = original_ref_block(h) + side_update(h) side_update(h) = sidecar_on_block(h) − base_now_block(h) ← base 自己抵銷 loss = MSE( live_base_block(h), target(h) )
訓練分兩階段 sweep: P1 freeze router、unfreeze experts; P2 freeze (剛訓好的) experts、unfreeze router,並加 KL router-anchor 正則化把 router 拉回 pristine。
2 · Sidecar 結構與 forward
Sidecar 是一個並聯(side-attached)模組,逐層包住 OLMoE 的 MoE block。 forward 是嚴格 加法殘差,base MoE 的輸出原封不動地保留:
side_hidden = down_proj( SiLU( gate_proj(h) ) * up_proj(h) ) # SwiGLU bottleneck need_gate = σ( gate_w · h ) # per-token scalar, sigmoid quality_gate = σ( quality_w · side_hidden ) # optional, per-token scalar gate = need_gate * quality_gate side_update = α · gate · side_hidden # α = 0.2 output = base_moe(h) + side_update # ← base 不變,只加
實作位置:scripts/train_gate_c_sidecar_smoke.py:249–396(class GatedSidecarMoeBlock)。
attachment:decoder_layer.mlp = wrapper(line 460),wrapper 內含 凍結的 base_moe。
2.1 名稱解碼:scalar_d512_g1_alpha02
| Token | 含義 | 對應參數 |
|---|---|---|
scalar | gate 是 per-token scalar(R1),不是 per-channel | sidecar_gate_groups=1 |
d512 | SwiGLU 內層 bottleneck 寬度 | intermediate_size=512 |
g1 | gate group 數 = 1(=scalar) | 同上 |
alpha02 | 整條 side path 的縮放係數 | sidecar_scale=0.2 |
2.2 參數量(per layer)
| 模組 | trainable params | 狀態 |
|---|---|---|
| base MoE block(2048→1024,64 experts,top-8) | ~50 M | frozen |
| sidecar up_proj 2048→512 | ~1.05 M | trainable |
| sidecar gate_proj 2048→512 | ~1.05 M | trainable |
| sidecar down_proj 512→2048 | ~1.05 M | trainable |
| need_gate(2048→1) | ~2 K | trainable |
| quality_gate(512→1) | ~0.5 K | trainable |
| 合計 sidecar / layer | ~3.1 M | trainable |
| 16 層總計 | ~50 M | ≈ base MoE 一層 |
base_moe(h) 在 sidecar 掛載與否時計算完全一樣(同一個凍結 module、同一輸入)。
要拆掉 sidecar 只需 decoder_layer.mlp = wrapper.base_moe,還原成 stock OLMoE 行為,無任何 weight 變更。
這個可分離性是「蒸餾回 base」這個 well-posed problem 的前提。
3 · 為什麼選 sidecar 而不是 LoRA
LoRA 是業界默認的 PEFT,我們也認真考慮過。但在 MoE + 持續學習這個場景下,sidecar 的形狀更合適。 以下是按照本研究的實際需求(可分離、可蒸餾、可逆、不破壞 base 能力)逐條比較:
| 面向 | LoRA | Sidecar(本研究) |
|---|---|---|
| 對 base 權重的影響 | 修改:W' = W + BA。inference 通常 merge,merge 後 base 不再可分離 |
不修改:out = base(h) + α·g·side(h)。base 權重逐字節相同 |
| 計算路徑 | 串接在原 matmul 上,inference 時必須 merge 或多算一條 low-rank path | 純並聯,base 計算路徑零干擾。g≈0 時等價於完全沒掛 |
| MoE 適配 | 64 個 expert 每個都要 LoRA → 參數膨脹;只 LoRA 共享部分(router/attention)又摸不到 expert 容量 | 整個 MoE block 外掛一條 SwiGLU side-path,跟 expert 數量解耦 |
| 能力 ON/OFF | merge 後是 binary;未 merge 可選但要載入 adapter | 每 token 連續調節:need_gate(σ)學「要不要用」,quality_gate(σ)學「能不能信任輸出」 |
| 原能力衰退風險 | merge 後改寫 base,catastrophic forgetting 不可逆 | base 不動,gate 可學成「無關 token 完全跳過」(g→0,加法項=0) |
| 可組合性(多任務) | 多 LoRA stack 會互相干擾 weight space | 多 sidecar 可同時掛,各有獨立 gate,互不干擾 base path |
| 蒸餾回 base 的可行性 | LoRA 本來就是設計成 merge,無「蒸餾回 base」這個概念 | 因為 base path 仍存在,side_update = sidecar_on − base_now 可在訓練 base 時用 base-cancellation 識別取得乾淨 target |
這份比較不是「LoRA 在所有任務都比較差」——它在 deploy 簡單、生態成熟、merge 後零推論開銷上仍佔優勢。 但對「凍結 base、累加能力、隨時可拆、可蒸餾回 base」這個具體需求集,sidecar 的形狀就是答案。
4 · 過去 sidecar 變體實驗(seed6k 系列)
在現在做蒸餾之前,我們訓練了 3 個 sidecar,用同樣的結構(scalar_d512_g1_alpha02),
只變動訓練資料集。目的:確定 sidecar 形狀對不同 task family 都能學起來,
再選最強的一支做蒸餾實驗。
| Sidecar 變體 | 訓練資料 | d | Gate | α | Heldout pass@1 |
|---|---|---|---|---|---|
sidecar_seed6k_r1ds_… |
seed6k_r1_ds(5.3k) | 512 | scalar (need) | 0.2 | 0.636(89/140) |
sidecar_seed6k_r2_1_… ← 蒸餾用 |
seed6k_r2_1(5.3k) | 512 | scalar (need) | 0.2 | 0.500(70/140) |
sidecar_seed6k_r2_… |
seed6k_r2_balanced(5.3k) | 512 | scalar (need) | 0.2 | 0.414(58/140) |
4.1 沒做過的 sweep(誠實揭露)
- Dimension 掃描:沒有 d=256 / d=1024 對照。512 是直接賭一把選的。
- Gate 形狀掃描:沒有 per-channel gate(g>1)、MLP gate 對照。
- α 掃描:沒有 α=0.1 / 0.4 對照。0.2 是試出來「能訓收斂、不會 overshoot」的點。
這意味著 d512 / scalar / α=0.2 在「特定資料量(5.3k)+ 特定 task family」上能訓起來, 但不能宣稱是最佳設計點。若 R2.2 階段成績不如預期,這三個 sweep 是首要重新檢視的設計變數。
5 · 已完成驗證(全部落地到 main)
5.1 同-harness 基線校準
過去 doc 引用的 70/114(R2.1/R1)是另一個「嚴格 per-record budget」harness 出的數字。
我們的 gate 全部改用同一支 router_bias_probe.py + 同一組 flag:
--system-prompt "return only one complete fenced Python code block"
--max-new-tokens 512 (flat)
--generation-batch-size 8
--bias-json nobias.json # {"bias_by_layer": [[0.0]*64 for _ in range(16)]}
| 條件 | R2.1 / 140 | R1 / 140 |
|---|---|---|
| Pristine OLMoE base | 12 | 3 |
| Base + R2.1 sidecar attached | 61 | 117 |
5.2 根因:bf16 expert update round-to-zero
第一批多 task 跑出「loss 看似正在下降但 generation 行為毫無變化」的怪異結果。隔離出單層 50-step 診斷:
bf16 lr=1e-5: loss[0] = 2.899e-08 → loss[49] = 2.899e-08 (Δ = 0%) fp32 同條件 : loss drop 89%–100%
原因:expert weight scale ≈ 0.05,bf16 在這個量級的 ULP ≈ 4e-4;而 AdamW 在 eps-dominated regime 下
per-step weight delta ≈ lr · grad / (sqrt(v) + eps) ≈ 1e-8,被 bf16 截斷為零。
模型 dtype 改 fp32 後立即見效:
| 條件 | Task 1 R2.1 | Task 1 R1 | Generation 行為 |
|---|---|---|---|
| bf16(舊) | 15 | 5 | 未吸收,輸出近 base |
| fp32(1b1371d) | 30 | 44 | 已吸收 sidecar 圍欄碼塊風格 |
5.3 工具/介面落地
- 主腳本
scripts/distill_sidecar_into_moe.py(compute_block_distill_loss/phase_total_loss/router_anchor_loss/selection_shift_rate/set_phase_trainable/GateThresholds/gate_decision) - 33 個單元測試(
tests/test_distill_sidecar_into_moe.py)含 10 個 gate 測試, 其中對「實際觀測到的 task-1 (r2=15,r1=5)」釘住CONTINUE不再誤判為 ROLLBACK - 診斷工具
scripts/_diag_distill_signal.py/scripts/_diag_precision_sweep.py(可重複的精度對照) - Runtime 用的是 elaborate sidecar trainer(
gate_mode/quality_gate_init_bias/group_gate_mode/group_gate_beta/gate_groups/sidecar_down_init),這點曾誤判為 committed scalar 版本,於 8289bc0 修正
5.4 已 commit 的 7 個 fix(main HEAD d834445)
| Commit | 內容 |
|---|---|
| d834445 | fix: base-relative gate(rollback 對齊 pristine base,不對齊 sidecar-on) |
| 1b1371d | fix(distill): 訓練改 fp32(根因:bf16 expert updates round to zero) |
| ba4645f | fix: load_jsonl 需要 Path 而非 str |
| 8289bc0 | fix: attach_sidecars 用 elaborate runtime signature |
| ae8201e | fix: run_quick140_eval 加 --generation-batch-size 8 |
| f5465b6 | fix: gate 門檻校準到同-harness 基線 |
| 65703e0 | fix: run_quick140_eval 補上 --system-prompt |
6 · 產生的 Impact
6.1 證實了 bf16 在「small-Δ optimizer regime」會靜默失效
這個現象在文獻中以「mixed precision must keep fp32 master weights」一帶而過,但我們把它落到具體的 ULP × per-step Δ 計算,並用 4-arm 精度掃描(bf16/fp32 × lr1e-5/1e-3)做出可被別人復現的證據。 對任何想直接用 bf16 做 LoRA / 蒸餾 / RL 微調 small-update 工作的人,這份診斷有警示價值。
6.2 推翻了直覺式的「sidecar-on relative」門檻
蒸餾後的純 base 是新模型,不是 sidecar+base 的混合體。 拿它去跟「sidecar+base 在 R1 上的 117」比較會結構性把所有結果判成回歸,即便它已經明顯優於 pristine base 的 3。 這在使用者直接質疑(「117/140 是 sidecar+base 對吧?那拆掉後應該跟 pristine base 比啊?」)後才被抓出。 現在 gate 的語意是:ROLLBACK 對 pristine base 的退化、DONE 對 sidecar-on 的吸收比例 ≥ 0.9。
6.3 建立了可重複的同-harness 校準 SOP
所有後續實驗(包含未來其他 sidecar 任務)都能用同一條校準腳本對齊到 12/3 / 61/117, 不必再追問「上次那個數字是哪支 harness 出的」。
6.4 觀察到的多-task 不穩定(open finding)
fp32 修好後,5-task 跑出震盪而非單調上升:
| Task | R2.1 | R1 | Selection-shift | Decision |
|---|---|---|---|---|
| 1 | 30 | 44 | 0.27 | continue |
| 2 | 19 | 40 | 0.33 | continue |
| 3 | 35 | 61 | 0.35 | continue |
| 4 | 33 | 66 | 0.35 | continue |
| 5 | 22 | 54 | 0.35 | stop_no_progress |
最佳 R2.1=35(task 3)離目標 ≥55 還差不少,且工具側 prune_checkpoints({0, last_good_id})
在非單調軌跡下會把 task 3 的 best weights 默默刪掉——這是接下來必修的 tooling bug。
7 · 後續研究方向
7.1 立即:debug 不穩定(系統性,one variable at a time)
- H1 — lr 過高:把
lr 1e-3 → 1e-4,3-task 探針。 預判:若 R2.1 變單調且 selection_shift 維持 ≤0.27,即確認;然後跑滿 max-tasks 直到 DONE。 原 background jobbw9pg8tav被中斷,需重跑 - H2 — steps-per-layer 過多:50 → 10 或 20。假說:單層過度擬合該 batch,逐層累積偏移。
- H3 — router-anchor 太弱:1e-3 → 1e-2。假說:router 漂移(selection_shift 單調上升)失控。
- H4 — 蒸餾順序:目前 front→back(因為 base-cancellation identity),改回 back→front(原 doc 設計)做對照。
- H5 — 目標函數:per-block feature MSE 換成 end-to-end KL/CE 到 sidecar-on logits。 理由:sidecar 本來就用 KL/CE 訓出來的,per-block MSE 可能 asymptote 低於 sidecar-on。
7.2 必修工具:prune_checkpoints 保留 best
scripts/distill_sidecar_into_moe.py::main() 在 gate loop 中追蹤
best_task_id = argmax accepted r2,prune set 改成
{0, best_task_id, last_good_id},並把 best_task_id 寫進
task_metrics.jsonl。TDD 擴充 test_prune_checkpoints_keeps_only_requested 至 3-keep 情境。
任何下一次 multi-task 之前必須先做完。
7.3 中期:目標達成的成功標準
R2.1 吸收
140 上 ≥ 55(= round(0.9 × 61))
R1 不退化
140 上 ≥ 3(pristine base 基線)且實務上希望 ≥ 30 以維持一般 instruct 行為
7.4 交付精度
訓練必須 fp32(累積 sub-ULP updates),但 olmoe_0.1 最終可 cast 回 bf16(~14 GB)交付——
推論精度損失可忽略(run_quick140_eval 已經 bf16 評估,所以 gate 通過的數字本身就是「bf16 可交付」的數字)。
7.5 開放問題
- 蒸餾上限:per-block MSE 在 sidecar 是 KL-trained 的情況下,理論能不能達到 sidecar-on 同等表現?還是會被 asymptotic gap 卡住?H5 就是在測這個。
- R1 退化路徑:目前 task 3-4 看到 R1 短暫超過 pristine 很多(44 → 66), 不知是真的有遷移效應(sidecar 對 Python 能力的一般化)還是 generation 偏置造成的虛假提升。 需要在 R1 上加更嚴格的子集 audit。
- 規模外推:這條 pipeline 之後要套到 R2.2(進階 ds-code-bench 方向)的 sidecar 上, 需要驗證同一套精度與 gate 設定能不能搬。
8 · 復現入口
環境
- GPU: RTX PRO 6000 Blackwell Max-Q, 96 GB
- Host RAM 61 GB · Disk 554 GB free
- fp32 checkpoint ≈ 28 GB / task · multi-task peak ≈ 84 GB(pristine + best + current)
關鍵 artifacts
R2.1 sidecar:
results/gate_i_seed6k_r2_1/sidecar_seed6k_r2_1_scalar_d512_g1_alpha02_v0/
sidecar_state.pt
training_manifest.json
R2.1 data:
/home/weiciao/projects/olmoe-data/datagen/frozen/ds_data_v1_seed6k_r2_1_20260518/
sft_train.jsonl
heldout_quick140.jsonl
heldout_quick140_execution_accepted.jsonl
R1 data (regression check):
/home/weiciao/projects/olmoe-data/datagen/frozen/ds_data_v1_seed6k_r1_ds_20260517/
heldout_quick140.jsonl
heldout_quick140_execution_accepted.jsonl
Calibration outputs:
results/distill_sidecar_r2_1_into_moe_v0/calibration/
nobias.json RESULT.txt BASE_R1_RESULT.txt precision_sweep.log
多-task 跑法
cd /home/weiciao/projects/olmoe bash scripts/_run_distill_multi8.sh # fp32, lr 1e-3, 50 steps/layer, anchor 1e-3 # 或下一步的 H1 probe: bash scripts/_run_distill_lr1e4_t3.sh # 同上但 lr 1e-4, 3 tasks