wassname/evil_MoE
1
0
Fork 0
mirror of https://github.com/wassname/evil_MoE.git synced 2026-08-20 21:11:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: mix=0 no-teacher ablation path (pure on-policy, pool kept for v_grad+partition)

Browse Source 
train.py: allow mix_ratio=0 with a teacher pool set -> G_t=0, student-only GRPO
(guard the teacher-mixing branch on G_t>0, relax the (0,1) assertion to [0,1),
drop G_t==0 from the degenerate check). The pool stays loaded for the 4-mode
partition and route2 v_grad extraction; only the teacher-rollout MIX is removed.
Smoke (mix=0 + normal mix=0.5 + vanilla) all green.

Also: fill A4 long-run figure (fig:longrun) in main.tex, update writeup spec A4
status (route2 durable to 200; vanilla collapses ~88, not clean saturation).

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-06-02 23:26:26 +00:00
co-authored by Claudypoo
parent e00292860f
commit 62e510ff57
3 changed files with 42 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/train.py
+13 -6
View File
@@ -488,14 +488,19 @@ def main(cfg: Config) -> int:
G_s = group
G_t = 0
if cfg.teacher_pool_dir is not None:
if not (0.0 < cfg.mix_ratio < 1.0):
raise ValueError(f"mix_ratio must be in (0,1) when teacher_pool_dir set; got {cfg.mix_ratio}")
# mix=0 is the NO-TEACHER ablation: pure on-policy GRPO (G_t=0, no teacher
# rollouts injected) while the pool is still loaded for the 4-mode partition
# and route2 v_grad extraction. Using the pairs for v_grad is allowed under
# the no-cheat invariant; mixing teacher rollouts into training is the thing
# mix=0 removes. mix in [0,1).
if not (0.0 <= cfg.mix_ratio < 1.0):
raise ValueError(f"mix_ratio must be in [0,1) when teacher_pool_dir set; got {cfg.mix_ratio}")
G_t = round(group * cfg.mix_ratio)
G_s = group - G_t
if G_s == 0 or G_t == 0:
if G_s == 0:
raise ValueError(
f"degenerate split: G={group} mix_ratio={cfg.mix_ratio} -> G_s={G_s}, G_t={G_t}. "
f"Pick mix_ratio so both halves are non-empty, or drop --teacher-pool-dir."
f"degenerate split: G={group} mix_ratio={cfg.mix_ratio} -> G_s={G_s}. "
f"Pick mix_ratio < 1 so the student half is non-empty."
)
for path in sorted(cfg.teacher_pool_dir.glob("prompt_*.jsonl.gz")):
# path.stem on 'prompt_0004.jsonl.gz' is 'prompt_0004.jsonl' (only one
@@ -862,8 +867,10 @@ def main(cfg: Config) -> int:
model.config.use_cache = True
_tg = time.perf_counter()
teacher_sample: list[dict] | None = None
if teacher_pool:
if teacher_pool and G_t > 0:
# Mixed-pool: G_s live student + G_t cached teacher rollouts.
# G_t==0 (mix=0 no-teacher ablation) falls through to the student-only
# path below; the pool stays loaded for partition + v_grad extraction.
# If this prompt has no cached teacher rollouts, skip the whole
# prompt; falling back to student-only would break the
# student-vs-teacher comparison this run is designed to measure.