RL logprob parity scripts: MaxText trainer vs vLLM sampler (Qwen3.5 35B/397B)

tools/rl_logprob_parity/README.md

@@ -0,0 +1,47 @@
+# MaxText trainer vs vLLM sampler logprob parity (Qwen3.5-35B-A3B / 397B-A17B)
+
+Scripts behind the "Compare logprobs between MaxText trainer and vLLM sampler" tables (TIS band
+[0.999, 1.002] statistics). They are experiment scripts, not tests: paths, worktrees and the output
+directory (`/mnt/disks/persist/pr4925_repro`) are hard-coded for the host they ran on (see "Environment").
+`run_all.sh` holds the exact invocation for every table cell, in dependency order.
+
+## Setup measured
+* Trainer: MaxText Qwen3.5 (nnx), fsdp=8, bf16 (fp8 rows: qwix `fp8_full` dynamic quantization).
+* Sampler, native: vLLM + tpu-inference torchax path (`MODEL_IMPL_TYPE=vllm`), tp=8 with
+  `attn_dp_size=4` (attention DP-4 x TP-2), `--enable-expert-parallel` (EP-8), bf16 KV cache,
+  chunked prefill, block 256. FP8 sampler = `Qwen/Qwen3.5-35B-A3B-FP8`.
+* Sampler, MaxText-in-vLLM: `maxtext_vllm_adapter` (`MODEL_IMPL_TYPE=flax_nnx`), attention DP-4 x TP-2, MoE TP-8.
+* Data: 8 x 512 real tokens (MaxText `docs/*.md`, Qwen3.5 tokenizer), produced by `run_35b_real.py`
+  into `real35b_L3.npz` (tokens + layer-3 hidden state); every other script reads that file.
+
+## Table rows -> scripts
+| rows | sampler side | trainer side |
+|---|---|---|
+| 35B prompt tokens, layers 4/8, bf16/bf16 | `torchax_prefix.py` (also captures per-layer expert ids for replay) | `maxtext_prefix.py` |
+| 35B prompt tokens, 40 (full), bf16/bf16 and bf16/fp8 (+ fp8 layers 4/8) | `torchax_prefix_gen.py` (`NL`, `MODEL`) | `maxtext_prefix_gen.py` (`NL`, `MODES`, `TX_NPZ`, `SUFFIX`) |
+| 35B fp8/fp8 (qwix trainer), 40 (full) | same FP8-sampler npz | `maxtext_prefix_gen.py` with `FULLCALL=1 EXTRA='{"quantization": "fp8_full", "use_qwix_quantization": true}'` |
+| 35B output tokens (decode path) | `vllm_generate.py` (rollouts, `enable_return_routed_experts`) | `maxtext_score.py` (`GEN_ONLY_REPLAY=1` = replay row); `compare_output.py` prints the row |
+| 35B MaxText-in-vLLM, 40 (full) | `adapter_prompt_logprobs.py` (real engine `prompt_logprobs`) | `maxtext_full_logprobs.py` |
+| 397B layers 4/8, bf16/fp8 | `torchax_prefix_397b.py` | `maxtext_prefix_397b.py` (`MODE=own|replay`, 8-layer partial restore) |
+| tables | `build_tables.py` renders the markdown tables from the npz outputs | |
+
+Row semantics: layers 4/8 = logit lens (model's own final norm + lm_head on the hidden state after
+layer N, both sides); 40 = true logprobs. "Expert replay Y" = the sampler's per-layer top-k expert ids
+fed to the trainer's `RoutedMoE.get_topk` (class patch in `maxtext_prefix*.py` / `maxtext_score.py`),
+weights recomputed from the trainer's fp32 gate logits. Output-token replay covers decode positions
+only: the engine's `routed_experts` rows for prefill positions are zero-filled under attn_dp +
+chunked prefill, and `maxtext_score.py` maps those rows (and all prefill rows) to the trainer's own routing.
+
+## Running
+`run_all.sh` — steps 0-8; each step needs the whole TPU. Wrappers:
+* `run_maxtext.sh <script.py> [args]` — MaxText-side environment (conda + tpu-inference/vLLM/MaxText trees on `PYTHONPATH`).
+* `run_vllm.sh <script.py> [args]` — vLLM/tpu-inference environment with the production serving env vars
+  (`USE_MOE_EP_KERNEL=0`, `NEW_MODEL_DESIGN=1`, `VLLM_ENABLE_V1_MULTIPROCESSING=0`, ...); `ADAPTER=1` switches to
+  the MaxText-in-vLLM adapter (`MODEL_IMPL_TYPE=flax_nnx`, MaxText tree on `PYTHONPATH`).
+* `wait_tpu.sh <log> <wrapper> <script.py> [args]` — waits for `/dev/vfio/*` to be free, runs, retries if the TPU was taken.
+
+## Environment
+A conda env plus `PYTHONPATH` pointing at a tpu-inference worktree (Qwen3.5 canonical weight mapping), a vLLM
+worktree and this MaxText tree (`src/`); `HF_HOME` with the HF checkpoints; MaxText-format checkpoints under
+`/mnt/disks/persist`. The vLLM engine is started in-process (`vllm.LLM`) so the harness can call the loaded,
+sharded model directly.

tools/rl_logprob_parity/adapter_prompt_logprobs.py

@@ -0,0 +1,45 @@
+"""MaxText-in-vLLM (OOT adapter) 35B sampler: attn_dp=4 x tp=2 (MoE runs TP-8: EP not reachable on this path), bf16,
+real engine prompt_logprobs on the real token sequences."""
+import tpu_raiden.frameworks.jax._tpu_raiden_jax  # noqa: F401
+import tpu_inference  # noqa: F401
+import os, sys, time, numpy as np
+sys.path.insert(0, "/home/wenxindong_google_com/work/maxtext/.claude/worktrees/pr4925/src/maxtext/integration/vllm")
+import maxtext_vllm_adapter; maxtext_vllm_adapter.register()
+from vllm import LLM, SamplingParams, TokensPrompt
+OUT = "/mnt/disks/persist/pr4925_repro"; MODEL = "Qwen/Qwen3.5-35B-A3B"; ATTN_DP = int(os.environ.get("ATTN_DP", "4"))
+CKPT = "gs://maxtext-model-checkpoints/qwen3.5-35b-a3b/unscanned/0/items"
+t0 = time.time(); log = lambda *a: print(f"[{time.time()-t0:5.0f}s]", *a, flush=True)
+tokens = np.load(f"{OUT}/real35b_L3.npz")["tokens"]; B, S = tokens.shape
+mt_cfg = {"model_name": "qwen3.5-35b-a3b", "load_parameters_path": CKPT, "weight_dtype": "bfloat16", "dtype": "bfloat16",
+          "attention": "vllm_rpa", "allow_split_physical_axes": True, "scan_layers": False, "enable_nnx": True, "pure_nnx": True,
+          "pure_nnx_decoder": True, "prefuse_moe_weights": True, "enable_dp_attention": ATTN_DP > 1, "log_config": False,
+          "enable_checkpointing": True, "async_checkpointing": False, "checkpoint_storage_use_ocdbt": True,
+          "checkpoint_storage_use_zarr3": True, "convert_checkpoint_if_possible": False,
+          "float32_logits": True, "float32_gate_logits": True, "float32_weight_sum": True}
+sharding = {"sharding_strategy": {"enable_dp_attention": True, "attn_dp_size": ATTN_DP}} if ATTN_DP > 1 else None
+llm = LLM(model=MODEL, dtype="bfloat16", tensor_parallel_size=8, enable_expert_parallel=True,
+          hf_overrides={"architectures": ["MaxTextForCausalLM"]},
+          additional_config={"maxtext_config": mt_cfg, **({"sharding": sharding} if sharding else {})},
+          max_model_len=4096, max_num_seqs=16, max_num_batched_tokens=2048, block_size=256, enable_chunked_prefill=True,
+          enable_prefix_caching=False, gpu_memory_utilization=0.5, language_model_only=True,
+          limit_mm_per_prompt={"image": 0, "video": 0}, disable_log_stats=True, kv_cache_dtype="bfloat16")
+log("adapter engine up")
+# text-only run: the MaxText adapter has no M-RoPE hook (get_mrope_input_positions_fn is None) -> use plain 1-D positions
+runner = llm.llm_engine.model_executor.driver_worker.model_runner
+log(f"uses_mrope={runner.uses_mrope}; mrope fn={runner.get_mrope_input_positions_fn}")
+for obj in (runner, getattr(runner, "persistent_batch_manager", None), getattr(runner, "input_batch", None)):
+  if obj is not None and hasattr(obj, "uses_mrope"):
+    obj.uses_mrope = False
+def _text_mrope(prompt_token_ids, mm_features):
+  pos = np.arange(len(prompt_token_ids), dtype=np.int64); return np.stack([pos, pos, pos]), 0
+runner.get_mrope_input_positions_fn = _text_mrope
+sp = SamplingParams(max_tokens=1, temperature=0.0, prompt_logprobs=1)
+outs = llm.generate([TokensPrompt(prompt_token_ids=[int(t) for t in row]) for row in tokens], sp)
+lp = np.full((B, S), np.nan, np.float32); top1 = np.full((B, S), -1, np.int32)
+for b, o in enumerate(outs):
+  for i, d in enumerate(o.prompt_logprobs):
+    if d is None: continue
+    tid = int(tokens[b, i]); lp[b, i] = d[tid].logprob if tid in d else np.nan
+    top1[b, i] = max(d.items(), key=lambda kv: kv[1].logprob)[0]
+np.savez(f"{OUT}/adapter35b_prompt_logprobs_dp{ATTN_DP}.npz", logp=lp, top1=top1)
+log(f"saved; mean logp={np.nanmean(lp[:, 1:]):.3f}; top-1 acc vs actual={np.mean(top1[:, 1:] == tokens[:, 1:]):.3f}")

tools/rl_logprob_parity/build_tables.py

@@ -0,0 +1,126 @@
+"""Assemble the TIS-band tables (35B, 397B) from all result npz files present."""
+import numpy as np, os, glob
+L = "/mnt/disks/persist/pr4925_repro"; S = 512
+tok = np.load(f"{L}/real35b_L3.npz")["tokens"]
+def band(lt, li):
+  m = np.isfinite(lt) & np.isfinite(li); lt = lt[m]; li = li[m]; d = lt - li; r = np.exp(d)
+  return dict(n=len(d), inb=np.mean((r >= 0.999) & (r <= 1.002)), p1=np.mean((r >= 0.99) & (r <= 1.01)), p5=np.mean((r >= 0.95) & (r <= 1.05)),
+              med=np.median(np.abs(d)), p99=np.percentile(np.abs(d), 99), mx=np.abs(d).max())
+def fmt(b):
+  return f"{b['inb']:6.2%} | {b['p1']:5.1%} | {b['p5']:5.1%} | {b['med']:.4f} / {b['p99']:.3f} / {b['mx']:.2f}"
+def hidden(a, b):
+  D = a.shape[-1]; a = a.reshape(-1, D); b = b.reshape(-1, D)
+  c = np.sum(a * b, -1) / (np.linalg.norm(a, axis=-1) * np.linalg.norm(b, axis=-1) + 1e-30)
+  return f"— | — | — | — | layer-contribution error: relL2 {np.linalg.norm(a-b)/np.linalg.norm(a):.1%}, tokens with contribution cos<0.999: {np.mean(c<0.999):.1%} (no logprob at depth 1);"
+def load(p):
+  return np.load(p) if os.path.exists(p) else None
+rows = []
+def row(model, sampler, dtype, replay, depth, text, note=""):
+  rows.append((model, sampler, dtype, replay, depth, text, note))
+HDR = "| sampler | trainer/sampler dtype | expert replay | layers | in [0.999,1.002] | in ±1% | in ±5% | \\|Δlogp\\| median / p99 / max | note |\n|---|---|---|---|---|---|---|---|---|"
+
+# ---------------- 35B ----------------
+mt8 = load(f"{L}/maxtext_prefix.npz"); tx8 = load(f"{L}/torchax_prefix_dp4tp2_ep1.npz")
+mt40 = load(f"{L}/maxtext_prefix_nl40_bf16.npz"); tx40 = load(f"{L}/torchax_prefix_dp4tp2_ep1_nl40.npz")
+tx8f = load(f"{L}/torchax_prefix_dp4tp2_ep1_nl8_fp8.npz"); mt8f = load(f"{L}/maxtext_prefix_nl8_fp8replay.npz")
+tx40f = load(f"{L}/torchax_prefix_dp4tp2_ep1_nl40_fp8.npz"); mt40f = load(f"{L}/maxtext_prefix_nl40_fp8replay.npz")
+mtfull = load(f"{L}/maxtext35b_full_logprobs.npz"); vfull = load(f"{L}/vllm35b_prompt_logprobs_dp4.npz"); afull = load(f"{L}/adapter35b_prompt_logprobs_dp4.npz")
+mt8q = load(f"{L}/maxtext_prefix_nl40_fp8trainer_full.npz")
+z = np.load(f"{L}/real35b_L3.npz")
+def lens_rows(model, sampler, dtype, mt, tx, mt_own, prefix_own="own", prefix_rep="replay", depths=(4, 8), extra_full=None):
+  for rep, pre, m in (("N", prefix_own, mt_own), ("Y", prefix_rep, mt)):
+    if m is None or tx is None: continue
+    for k in depths:
+      if f"{pre}_lens{k}_logp_actual" in m.files and f"lens{k}_logp_actual" in tx.files:
+        note = "full model (true logprobs)" if k >= 40 else "logit lens (probe)"
+        row(model, sampler, dtype, rep, str(k) if k < 40 else "40 (full)", fmt(band(m[f"{pre}_lens{k}_logp_actual"], tx[f"lens{k}_logp_actual"])), note)
+# native bf16
+if mt8 is not None and tx8 is not None:
+  row("35B", "tpu-inference native", "bf16 / bf16", "N", "1", hidden(z["train"] - z["h_in"], tx8["y3_on_maxtext_h"] - z["h_in"]), "layer 3 alone, identical input")
+  row("35B", "tpu-inference native", "bf16 / bf16", "Y", "1", hidden(mt8["replay_y3"] - z["h_in"], tx8["y3_on_maxtext_h"] - z["h_in"]), "layer 3 alone, identical input")
+lens_rows("35B", "tpu-inference native", "bf16 / bf16", mt8, tx8, mt8)
+if mt40 is not None and tx40 is not None:
+  lens_rows("35B", "tpu-inference native", "bf16 / bf16", mt40, tx40, mt40, depths=(40,))
+if mtfull is not None and vfull is not None:
+  row("35B", "tpu-inference native (real engine, chunked prefill)", "bf16 / bf16", "N", "40 (full)", fmt(band(mtfull["logp"][:, :-1].reshape(-1), vfull["logp"][:, 1:].reshape(-1))), "true logprobs via prompt_logprobs")
+if mt40 is not None and vfull is not None:
+  eng = vfull["logp"][:, 1:].reshape(-1)
+  row("35B", "tpu-inference native (real engine, chunked prefill)", "bf16 / bf16", "Y", "40 (full)", fmt(band(mt40["replay_lens40_logp_actual"].reshape(8, S)[:, :-1].reshape(-1), eng)),
+      "trainer replays the harness-captured torchax routing (engine's own prefill routing not exportable under attn_dp+chunked prefill); trainer own-routing in the same loop: 29.26%")
+# real engine with a truncated N-layer model (production scheduler/chunked prefill), vs trainer lens at the same depth
+for k in (4, 8):
+  e = load(f"{L}/vllm35b_prompt_logprobs_dp4_nl{k}.npz")
+  if e is not None and mt8 is not None:
+    eng = e["logp"][:, 1:].reshape(-1)
+    for rep, pre in (("N", "own"), ("Y", "replay")):
+      tr = mt8[f"{pre}_lens{k}_logp_actual"].reshape(8, S)[:, :-1].reshape(-1)
+      row("35B", "tpu-inference native (real engine, N-layer model)", "bf16 / bf16", rep, str(k), fmt(band(tr, eng)),
+          "engine truncated to N layers via config; trainer = logit lens" + (" (replay routing from harness capture)" if rep == "Y" else ""))
+# native fp8 sampler, bf16 trainer
+if tx8f is not None:
+  own = mt8  # trainer own-routing outputs do not depend on the sampler
+  m = {}
+  if own is not None:
+    for k in own.files: m[k] = own[k]
+  if mt8f is not None:
+    for k in mt8f.files: m[k] = mt8f[k]
+  class D(dict):
+    files = property(lambda self: list(self.keys()))
+  m = D(m)
+  row("35B", "tpu-inference native", "bf16 / fp8", "N", "1", hidden(z["train"] - z["h_in"], tx8f["y3_on_maxtext_h"] - z["h_in"]), "layer 3 alone")
+  if "replay_y3" in m: row("35B", "tpu-inference native", "bf16 / fp8", "Y", "1", hidden(m["replay_y3"] - z["h_in"], tx8f["y3_on_maxtext_h"] - z["h_in"]), "layer 3 alone")
+  lens_rows("35B", "tpu-inference native", "bf16 / fp8", m, tx8f, m)
+if tx40f is not None and mt40f is not None:
+  m = {}
+  if mt40 is not None:
+    for k in mt40.files: m[k] = mt40[k]
+  for k in mt40f.files: m[k] = mt40f[k]
+  class D(dict):
+    files = property(lambda self: list(self.keys()))
+  lens_rows("35B", "tpu-inference native", "bf16 / fp8", D(m), tx40f, D(m), depths=(40,))
+# fp8 (qwix) trainer + fp8 sampler
+if mt8q is not None and tx40f is not None:
+  lens_rows("35B", "tpu-inference native", "fp8 (MaxText qwix fp8_full, dynamic) / fp8", mt8q, tx40f, mt8q, depths=(40,))
+# MaxText-in-vLLM
+if mtfull is not None and afull is not None:
+  row("35B", "MaxText in vLLM (attn_dp=4 x tp=2; MoE TP-8, EP not reachable)", "bf16 / bf16", "N", "40 (full)", fmt(band(mtfull["logp"][:, :-1].reshape(-1), afull["logp"][:, 1:].reshape(-1))), "true logprobs via prompt_logprobs")
+if afull is not None and vfull is not None:
+  row("35B", "calibration: MaxText-in-vLLM engine vs native engine", "bf16 / bf16", "–", "40 (full)", fmt(band(afull["logp"][:, 1:].reshape(-1), vfull["logp"][:, 1:].reshape(-1))), "two samplers vs each other")
+mi = load(f"{L}/real35b_L3_dp4tp2.npz")
+if mi is not None:
+  row("35B", "MaxText in vLLM (layer harness, attn_dp=4 x tp=2, MoE TP-8)", "bf16 / bf16", "N", "1", hidden(z["train"] - z["h_in"], mi["infer"] - z["h_in"]), "layer 3 alone")
+
+# ---------------- 397B ----------------
+mo = load(f"{L}/maxtext397_prefix_own.npz"); mr = load(f"{L}/maxtext397_prefix_replay.npz"); t3 = load(f"{L}/torchax397_prefix_dp4tp2_ep1.npz")
+if mo is not None and t3 is not None:
+  h = mo["own_h_after3"]
+  row("397B", "tpu-inference native", "bf16 / fp8", "N", "1", hidden(mo["own_y3"] - h, t3["y3_on_maxtext_h"] - h), "layer 3 alone (sampler FP8: bf16 397B does not fit)")
+  if mr is not None: row("397B", "tpu-inference native", "bf16 / fp8", "Y", "1", hidden(mr["replay_y3"] - h, t3["y3_on_maxtext_h"] - h), "layer 3 alone")
+  m = {}
+  for k in mo.files: m[k] = mo[k]
+  if mr is not None:
+    for k in mr.files: m[k] = mr[k]
+  class D(dict):
+    files = property(lambda self: list(self.keys()))
+  lens_rows("397B", "tpu-inference native", "bf16 / fp8", D(m), t3, D(m))
+qo = load(f"{L}/maxtext397_prefix_own_fp8trainer.npz"); qr = load(f"{L}/maxtext397_prefix_replay_fp8trainer.npz")
+if qo is not None and t3 is not None:
+  m = {}
+  for k in qo.files: m[k] = qo[k]
+  if qr is not None:
+    for k in qr.files: m[k] = qr[k]
+  class D(dict):
+    files = property(lambda self: list(self.keys()))
+  lens_rows("397B", "tpu-inference native", "fp8 (MaxText qwix fp8_full, dynamic) / fp8", D(m), t3, D(m), depths=(8,))
+
+out = []
+for model in ("35B", "397B"):
+  out.append(f"\n### {model}\n{HDR}")
+  for r in rows:
+    if r[0] == model: out.append(f"| {r[1]} | {r[2]} | {r[3]} | {r[4]} | {r[5]} | {r[6]} |")
+  if model == "397B":
+    out.append("| tpu-inference native | bf16 / bf16 | – | – | n/a | n/a | n/a | n/a | bf16 397B (794 GB) exceeds 8x v7x HBM (758 GB) |")
+    out.append("| any | any | any | 60 (full) | n/a | n/a | n/a | n/a | bf16 trainer cannot hold 60 layers on this host; prefix of 8 only |")
+    out.append("| MaxText in vLLM | any | – | – | n/a | n/a | n/a | n/a | adapter cannot load 397B (bf16 only, 794 GB) nor FP8 checkpoints |")
+  out.append("| MaxText in vLLM | bf16 / fp8, fp8 / fp8 | – | – | n/a | n/a | n/a | n/a | adapter has no FP8 checkpoint path |" if model == "35B" else "")
+print("\n".join(o for o in out if o is not None))

tools/rl_logprob_parity/compare_output.py

@@ -0,0 +1,15 @@
+import numpy as np, sys
+L="/mnt/disks/persist/pr4925_repro"; TAG=sys.argv[1] if len(sys.argv)>1 else "bf16"; SUF=sys.argv[2] if len(sys.argv)>2 else ""
+ro=np.load(f"{L}/rollout35b_{TAG}.npz"); mt=np.load(f"{L}/maxtext35b_score_{TAG}{SUF}.npz"); S=512; G=ro["gen_ids"].shape[1]
+def band(lt, li, name):
+  m=np.isfinite(lt)&np.isfinite(li); d=lt[m]-li[m]; r=np.exp(d)
+  print(f"  {name:<40} n={m.sum()} in[0.999,1.002]={np.mean((r>=0.999)&(r<=1.002)):6.2%} ±1%={np.mean((r>=0.99)&(r<=1.01)):5.1%} ±5%={np.mean((r>=0.95)&(r<=1.05)):5.1%} [0.8,1.25]={np.mean((r>=0.8)&(r<=1.25)):5.1%} | |dlogp| med={np.median(np.abs(d)):.4f} mean={np.mean(np.abs(d)):.4f} p99={np.percentile(np.abs(d),99):.3f} max={np.abs(d).max():.2f} | mean ratio={r.mean():.4f}")
+gen_lp = ro["gen_logp"].reshape(-1)                       # sampler decode-path logprob of sampled token j (input position S-1+j predicts it)
+for mode in ("own","replay"):
+  tr = mt[f"{mode}_logp"][:, S-1:S-1+G].reshape(-1)
+  band(gen_lp, tr, f"OUTPUT tokens (decode path), {mode} routing")
+plp = ro["prompt_logp"][:, 1:].reshape(-1)
+for mode in ("own","replay"):
+  tr = mt[f"{mode}_logp"][:, :S-1].reshape(-1)
+  band(plp, tr, f"PROMPT tokens (prefill path), {mode} routing")
+print("  sampled-token stats: sampler mean logp", np.nanmean(gen_lp), "; trainer(own) mean", np.nanmean(mt["own_logp"][:, S-1:S-1+G]))