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

G=6 + logits_to_keep OOM fix, generalization constraint, handover rewrite

Browse Source 
train.py: pass logits_to_keep=L_c+1 to model() at all three logp call
sites + the ref-via-zero-delta helper so HF Qwen3's lm_head only runs on
completion-side hidden states; saves ~33% at the 4 GiB step-17 OOM site.
full preset G=8 -> G=6 for a further ~25% B reduction at every act site.
Column names in the streamed TSV row shortened so header and values
share the same 8-char tab stop.

spec.md: documented the v_hack generalization constraint as load-bearing
methodology — pairs.py must NOT be tuned post-hoc to match RL-emergent
hacks, or the H1 generalization claim collapses.

handover.md: rewritten for current state (G=6, post-grader-fix, Qwen3-4B).
Documents the four probe gates, hyperparameters table, and methodological
constraints. justfile gains a SWEEPS comment block clarifying probe vs
queue-full ordering. .gitignore picks up .venv, *.log, /tmp/, cache dirs.

RESEARCH_JOURNAL.md: 2026-05-24 (b) entry covers the OOM diagnosis, fix,
pooled cross-run trend analysis (LR is fine, signal underpowered at n=17
but directionally consistent), and the generalization correction.
This commit is contained in:
wassname
2026-05-24 05:03:04 +00:00
parent 973b9407b5
commit 87a2b48784
6 changed files with 471 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/train.py
+56 -22
View File
@@ -57,6 +57,7 @@ Run:
from __future__ import annotations
import json
import os
import sys
import time
from dataclasses import dataclass, field
@@ -65,6 +66,11 @@ from enum import Enum
from pathlib import Path
from typing import Literal
# Must be set BEFORE `import torch` to take effect on the CUDA allocator.
# Eliminates fragmentation that caused 91 GiB allocated / 581 MiB free crash
# on Qwen3-4B G=8 (PyTorch's own OOM message recommends this).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import torch
import tyro
from loguru import logger
@@ -118,9 +124,10 @@ PRESETS: dict[str, dict] = {
"smoke": dict(model="Qwen/Qwen3.5-0.8B", steps=10, group=2, max_new=128,
n_problems=30, beta=0.0, prompts_per_step=1), # 24GB cap
# 4B matches reference DEFAULT_MODEL_ID (docs/vendor/rl-rewardhacking/src/__init__.py).
# G=12, max_new=1024 chosen to fit 96 GB with the AntiPaSTO+CE+checkpointing stack
# (2B/G=16/max=1024 observed at 54 GB peak; 4B/G=12/max=1024 estimated ~77 GB).
"full": dict(model="Qwen/Qwen3-4B", steps=200, group=12, max_new=1024,
# G=6 after 2026-05-24 step-17 OOM at G=8: lm_head spike on a long-prompt
# problem hit 4.16 GiB / 2.5 GiB free. `logits_to_keep` cuts lm_head ~33%;
# G=8->6 cuts B at every act site ~25%. Combined headroom ~6-10 GB.
"full": dict(model="Qwen/Qwen3-4B", steps=200, group=6, max_new=1024,
n_problems=500, beta=1e-3, prompts_per_step=8),
}
@@ -244,20 +251,26 @@ def load_v_hack(path: Path, model_name: str, wrappers: dict) -> dict[str, torch.
@torch.no_grad()
def ref_logprobs_via_zero_delta(
model, merged: torch.Tensor, wrappers: dict,
model, merged: torch.Tensor, wrappers: dict, plen: int,
) -> torch.Tensor:
"""Compute pi_ref logprobs by temporarily zeroing delta_S (=base model).
"""Compute pi_ref logprobs on completion tokens only.
AntiPaSTO: W' = W + U diag(delta_S) Vh. At delta_S=0, W' = W exactly
(verified bit-exact in step 1). Save -> zero -> forward -> restore.
Zero extra VRAM vs a separately loaded ref_model.
Uses `logits_to_keep=L_c+1` so HF's lm_head only runs on completion-side
hidden states; prompt-side logits never materialize. Saves
~plen/(plen+L_c) memory at the lm_head call (~33% at plen=500, L_c=1024).
That was the OOM site at vanilla step 17 (long prompt -> 4 GiB lm_head spike).
"""
saved = {n: info["delta_S"].data.clone() for n, info in wrappers.items()}
try:
for info in wrappers.values():
info["delta_S"].data.zero_()
logits = model(merged).logits[:, :-1]
return per_token_logps(logits, merged[:, 1:])
L_c = merged.shape[1] - plen
logits = model(merged, logits_to_keep=L_c + 1).logits[:, :-1]
return per_token_logps(logits, merged[:, plen:])
finally:
for n, info in wrappers.items():
info["delta_S"].data.copy_(saved[n])
@@ -288,6 +301,7 @@ def main(cfg: Config) -> int:
model = AutoModelForCausalLM.from_pretrained(
model_name, dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
).to(device)
# Trade compute for memory: recompute activations during backward. ~30-50%
# less activation VRAM on the policy forward, enough to fit G=8 max_new=1024
@@ -351,6 +365,14 @@ def main(cfg: Config) -> int:
eos_id = tok.eos_token_id
pad_id = tok.pad_token_id
# Stream the per-step table live (header once, row per step). Same columns as
# the final tabulate output. logger.info routes through tqdm.write so the
# rows appear above the progress bar without breaking it.
# Names kept <=7 chars so header and value share the same 8-col tab stop.
_row_cols = ["step", "rew", "std", "sprd", "N",
"gt", "hack", "loss", "cin", "cout", "fired", "sec"]
logger.info("row\t" + "\t".join(_row_cols))
pbar = tqdm(range(steps), desc=f"train {cfg.arm} {cfg.preset.value}", mininterval=60)
for step in pbar:
t0 = time.time()
@@ -431,19 +453,28 @@ def main(cfg: Config) -> int:
centered = rewards - rewards.mean()
adv = centered if cfg.unbiased else centered / (rewards.std() + 1e-4)
# Old-policy logprobs (frozen target for PPO ratio).
# Old-policy logprobs (frozen target for PPO ratio). Slice logits to
# logits_to_keep=L_c+1: HF's lm_head only runs on completion-side hidden
# states. Avoids materializing prompt-side logits (~plen/(plen+L_c) saved
# at lm_head). Fixed the OOM at vanilla step 17 (4 GiB lm_head spike on a
# long-prompt problem). Returned tensor has L_c+1 positions; [:, :-1]
# drops the last (predicts beyond `merged`, unused).
completion_ids = merged[:, plen:]
L_c = completion_ids.shape[1]
with torch.no_grad():
gen_logp = per_token_logps(
model(merged).logits[:, :-1], merged[:, 1:]
)[:, plen - 1:].detach()
model(merged, logits_to_keep=L_c + 1).logits[:, :-1],
completion_ids,
).detach()
ref_logp = None
if beta and beta > 0:
ref_logp = ref_logprobs_via_zero_delta(model, merged, wrappers)[:, plen - 1:].detach()
ref_logp = ref_logprobs_via_zero_delta(model, merged, wrappers, plen).detach()
pol_logp = per_token_logps(
model(merged).logits[:, :-1], merged[:, 1:]
)[:, plen - 1:]
model(merged, logits_to_keep=L_c + 1).logits[:, :-1],
completion_ids,
)
mask = (merged[:, plen:] != pad_id).float()
ratio = torch.exp(pol_logp - gen_logp)
@@ -496,20 +527,23 @@ def main(cfg: Config) -> int:
tail = diag_tail.replace("\n", "\\n")
logger.debug(f"step {step} gen[0] tail (last 400 chars): {tail!r}")
rows.append({
row = {
"step": step,
"rew_mean": f"{rew_mean:+.2f}",
"rew_std": f"{rew_std:.2f}",
"spread": "T" if spread else "F",
"rollouts": n_rollouts,
"gt_pass": f"{sum(agg_gt)}/{n_rollouts}",
"rew": f"{rew_mean:+.2f}",
"std": f"{rew_std:.2f}",
"sprd": "T" if spread else "F",
"N": n_rollouts,
"gt": f"{sum(agg_gt)}/{n_rollouts}",
"hack": f"{sum(agg_hack)}/{n_rollouts}",
"loss": f"{agg_loss:+.4f}",
"cos_in": f"{diag['mean_cos_in']:+.3f}",
"cos_out": f"{diag['mean_cos_out']:+.3f}",
"cin": f"{diag['mean_cos_in']:+.3f}",
"cout": f"{diag['mean_cos_out']:+.3f}",
"fired": f"{diag['frac_fired']:.2f}",
"sec": f"{time.time()-t0:.0f}",
})
}
rows.append(row)
# Stream this step as TSV row (header was printed before the loop).
logger.info("row\t" + "\t".join(str(row[c]) for c in _row_cols))
# Live status in tqdm postfix; full per-step line in verbose log only.
pbar.set_postfix(
rew=f"{rew_mean:+.2f}", gt=f"{sum(agg_gt)}/{n_rollouts}",