RL logprob parity scripts: MaxText trainer vs vLLM sampler (Qwen3.5 35B/397B) - #5047
RL logprob parity scripts: MaxText trainer vs vLLM sampler (Qwen3.5 35B/397B)#5047wenxindongwork wants to merge 2 commits into
Conversation
…ables) Experiment scripts that produced the MaxText-trainer vs vLLM-sampler logprob comparison tables (tpu-inference native and MaxText-in-vLLM samplers; bf16/fp8; expert routing replay; 4/8/full layers; prompt and sampled tokens), with the env wrappers and a README mapping each table row to its script.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a suite of experimental scripts and utilities under tools/rl_logprob_parity/ to analyze and compare log-probability parity between the MaxText trainer and the vLLM sampler for Qwen3.5 models. The feedback highlights several critical issues and areas for improvement: a referenced script (run_397b_layer.py) is missing from the repository, potential runtime crashes and NaN propagation could occur under certain edge cases in maxtext_score.py and compare_replay.py, helper class D is duplicated in build_tables.py, and several shell scripts contain fragile, non-portable hardcoded absolute paths and process IDs.
| export OUT_DIR=/mnt/disks/persist/pr4925_repro | ||
| mkdir -p $OUT_DIR | ||
| cd $MAXTEXT_WT | ||
| exec python /home/wenxindong_google_com/.claude/jobs/7ea88918/tmp/run_397b_layer.py "$@" |
There was a problem hiding this comment.
The script run_397b_layer.py is executed here, but it is completely missing from the repository and this pull request. Anyone attempting to run run.sh will encounter a 'file not found' error. Please ensure that run_397b_layer.py is added to the repository or update the script to point to the correct file.
| rows = routed.shape[1] | ||
| for b in range(B): | ||
| n = min(rows, L); ridx[:, b, :n, :] = routed[b, :n, :, :K].transpose(1, 0, 2) | ||
| if os.environ.get("GEN_ONLY_REPLAY") == "1": | ||
| ridx[:, :, :S, :] = -1 # own routing on all prefill rows 0..S-1 (row S-1 predicts the first decode token and is a prefill row) | ||
| zero_rows = np.all(routed[:, :, :, :K] == 0, axis=(2, 3)) # engine zero-fills rows it did not capture (prefill under attn_dp+chunked prefill) | ||
| for b in range(B): | ||
| n = min(rows, L); ridx[:, b, :n, :][:, zero_rows[b, :n], :] = -1 |
There was a problem hiding this comment.
If the routed experts array is empty or has mismatched shapes (which can happen if any r is None or shapes mismatch in vllm_generate.py), routed is saved as an empty 1D array. In that case, accessing routed.shape[1] will raise an IndexError, and slicing routed[:, :, :, :K] will crash. Wrapping this block in a dimension check prevents a runtime crash.
| rows = routed.shape[1] | |
| for b in range(B): | |
| n = min(rows, L); ridx[:, b, :n, :] = routed[b, :n, :, :K].transpose(1, 0, 2) | |
| if os.environ.get("GEN_ONLY_REPLAY") == "1": | |
| ridx[:, :, :S, :] = -1 # own routing on all prefill rows 0..S-1 (row S-1 predicts the first decode token and is a prefill row) | |
| zero_rows = np.all(routed[:, :, :, :K] == 0, axis=(2, 3)) # engine zero-fills rows it did not capture (prefill under attn_dp+chunked prefill) | |
| for b in range(B): | |
| n = min(rows, L); ridx[:, b, :n, :][:, zero_rows[b, :n], :] = -1 | |
| if routed.ndim >= 4: | |
| rows = routed.shape[1] | |
| for b in range(B): | |
| n = min(rows, L); ridx[:, b, :n, :] = routed[b, :n, :, :K].transpose(1, 0, 2) | |
| if os.environ.get("GEN_ONLY_REPLAY") == "1": | |
| ridx[:, :, :S, :] = -1 # own routing on all prefill rows 0..S-1 (row S-1 predicts the first decode token and is a prefill row) | |
| zero_rows = np.all(routed[:, :, :, :K] == 0, axis=(2, 3)) # engine zero-fills rows it did not capture (prefill under attn_dp+chunked prefill) | |
| for b in range(B): | |
| n = min(rows, L); ridx[:, b, :n, :][:, zero_rows[b, :n], :] = -1 | |
| else: | |
| log("Warning: routed experts array is empty or invalid; skipping replay index population.") |
| Pt = mt[f"{mode}_lens{k}_logprobs_2seq"]; Pi = tx[f"lens{k}_logprobs_2seq"] | ||
| kl = np.sum(np.exp(Pt) * (Pt - Pi), -1) # KL(train || sampler) per token, first 2 sequences |
There was a problem hiding this comment.
If any log-probability in Pt is -inf (which is common for masked tokens or padding), np.exp(Pt) * (Pt - Pi) will produce NaN because 0 * -inf is undefined in NumPy. This will propagate NaN to the mean and max KL divergence values. Using np.where to safely ignore -inf values prevents this issue.
| Pt = mt[f"{mode}_lens{k}_logprobs_2seq"]; Pi = tx[f"lens{k}_logprobs_2seq"] | |
| kl = np.sum(np.exp(Pt) * (Pt - Pi), -1) # KL(train || sampler) per token, first 2 sequences | |
| Pt = mt[f"{mode}_lens{k}_logprobs_2seq"]; Pi = tx[f"lens{k}_logprobs_2seq"] | |
| kl = np.sum(np.where(Pt > -1e9, np.exp(Pt) * (Pt - Pi), 0.0), -1) # KL(train || sampler) per token, first 2 sequences |
| 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 = [] |
There was a problem hiding this comment.
The helper class D is defined locally multiple times throughout this file. Defining it once at the module level is much cleaner and improves maintainability.
| rows = [] | |
| class D(dict): | |
| files = property(lambda self: list(self.keys())) | |
| rows = [] |
| class D(dict): | ||
| files = property(lambda self: list(self.keys())) | ||
| m = D(m) |
There was a problem hiding this comment.
Remove the duplicate local definition of class D since it is now defined at the module level.
m = D(m)| 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,)) |
There was a problem hiding this comment.
Remove the duplicate local definition of class D since it is now defined at the module level.
lens_rows("35B", "tpu-inference native", "bf16 / fp8", D(m), tx40f, D(m), depths=(40,))| for attempt in $(seq 1 12); do | ||
| while [ -n "$(sudo lsof -t /dev/vfio/* 2>/dev/null)" ]; do sleep 10; done | ||
| sleep 5; echo "attempt $attempt start $(date)" > /mnt/disks/persist/pr4925_repro/adapter.log | ||
| bash /home/wenxindong_google_com/.claude/jobs/7ea88918/tmp/run_adapter.sh >> /mnt/disks/persist/pr4925_repro/adapter.log 2>&1 |
There was a problem hiding this comment.
Using hardcoded absolute paths to a specific user's temporary directory (/home/wenxindong_google_com/...) makes the script non-portable and fragile. Using $(dirname "$0") allows the script to dynamically locate the companion script relative to its own location, making it instantly portable.
| bash /home/wenxindong_google_com/.claude/jobs/7ea88918/tmp/run_adapter.sh >> /mnt/disks/persist/pr4925_repro/adapter.log 2>&1 | |
| bash "$(dirname "$0")/run_adapter.sh" >> /mnt/disks/persist/pr4925_repro/adapter.log 2>&1 |
| @@ -0,0 +1,8 @@ | |||
| #!/bin/bash | |||
| while kill -0 3487041 2>/dev/null; do sleep 20; done; sleep 10 | |||
| L=/mnt/disks/persist/pr4925_repro; T=/home/wenxindong_google_com/.claude/jobs/7ea88918/tmp | |||
There was a problem hiding this comment.
Similar to the other scripts, using a hardcoded absolute path for T makes the script non-portable. Defining T dynamically using $(dirname "$0") ensures it works seamlessly on any host.
| L=/mnt/disks/persist/pr4925_repro; T=/home/wenxindong_google_com/.claude/jobs/7ea88918/tmp | |
| L=/mnt/disks/persist/pr4925_repro; T="$(dirname "$0")" |
| @@ -0,0 +1,6 @@ | |||
| #!/bin/bash | |||
| while kill -0 1665431 2>/dev/null; do sleep 15; done | |||
There was a problem hiding this comment.
Waiting on a hardcoded process ID (1665431) is extremely risky for reproducibility. If this PID happens to be active and reused by an unrelated process on the host machine, the script will hang indefinitely. Consider passing the PID as an argument or removing the hardcoded PID wait entirely.
Drop the real-engine native row script, replay diagnostics and the older chain scripts (not used for the reported numbers). Replace the per-script env and TPU-retry wrappers with run_maxtext.sh / run_vllm.sh / wait_tpu.sh, and record the exact invocation for every table cell in run_all.sh. Guard build_tables.py against the one npz whose producer is not part of this set.
1 participant
Description
Experiment scripts that produced the MaxText-trainer vs vLLM-sampler logprob comparison tables (TIS band [0.999, 1.002] statistics), added under
tools/rl_logprob_parity/.Setup measured:
fp8_fulldynamic quantization.MODEL_IMPL_TYPE=vllm, attention DP-4 x TP-2, EP-8, bf16 KV, chunked prefill) and MaxText-in-vLLM (maxtext_vllm_adapter, DP-4 x TP-2, MoE TP-8).RoutedMoE.get_topk, weights recomputed from the trainer's fp32 gate logits), depth (4 / 8 via logit lens, full model), prompt tokens vs sampled tokens (decode path).Contents: the scripts for each table row (sampler side, trainer side, table builder), the env wrappers / TPU-retry loops,
chain_matrix.shwith the exact invocation for every cell, and a README mapping rows to scripts and giving the run order.These are experiment scripts, not tests: output directory, checkpoint and worktree paths are hard-coded for the host they ran on (documented in the README). Kept byte-identical to what produced the reported numbers.
Tests
Not applicable (no library code changed). Scripts were run on 8 x v7x to produce the tables;
build_tables.pyrenders them from the saved npz outputs.Checklist
Before submitting this PR, please make sure (put X in square brackets):