[Experiment] Build Blackwell train-inference mismatch baseline for VIME post-training

[inaniloquentee]

Background

This issue tracks a Blackwell/B-series post-training experiment in RL-Align/vime. The first step is to intentionally construct a measurable train-inference mismatch scenario on VIME and use it as the baseline. Later we will plug in Blackwell-specific operators/kernels and compare against the same baseline configuration.

VIME already exposes train/rollout logprob consistency through train/train_rollout_logprob_abs_diff, which measures the mean absolute difference between Megatron recomputed logprobs and rollout-engine logprobs on the same sampled tokens.

Useful prior references from upstream VIME context:

• [Analysis] train_rollout_logprob_abs_diff: vLLM (vime) vs SGLang (slime)  vllm-project/vime#64 analyzed train_rollout_logprob_abs_diff between vime/vLLM and slime/SGLang.
• RFC: Sync upstream slime PRs into vime (tracking) vllm-project/vime#107 includes GB200 post-training experiment notes for Qwen3-30B-A3B in colocate and non-colocate setups.

Goal

Create a reproducible B-card post-training baseline where the rollout path and training path have a controlled, measurable mismatch.

The baseline should answer:

1. Can RL-Align/vime post-training run stably on B200/GB200 under this setup?
2. Can we observe a non-trivial train/train_rollout_logprob_abs_diff signal?
3. Is the mismatch stable enough to serve as an A/B baseline for future Blackwell operator integration?
4. After integrating B-card operators, can we compare correctness and performance under the same workload?

Baseline Design

Use Megatron as the training backend and vLLM as the rollout backend on Blackwell GPUs.

To intentionally create a train-inference mismatch, start with one or more controlled differences between the training and rollout paths:

• BF16 Megatron training plus FP8 vLLM rollout/inference.
• Rollout-side logprobs from vLLM compared against Megatron recomputed logprobs.
• Sampling filters enabled on the rollout side, for example top_p < 1.0 or top_k > 0.
• Optional follow-up: compare processed_logprobs against other vLLM logprob modes if available.

Recommended first baseline:

Hardware: B200 / GB200
Model: Qwen3-30B-A3B
Training backend: Megatron
Rollout backend: vLLM
Training precision: BF16
Rollout/inference precision: FP8 if available
Algorithm: GRPO
Primary metric: train/train_rollout_logprob_abs_diff

This should produce a measurable mismatch while still being close to a realistic production post-training setup.

Hardware Matrix

Stage	Hardware	Mode	Purpose
P0	1 B200/GB200 node	colocate	Smoke test and config validation
P1	2 GB200 nodes, 8 GPUs total	colocate	Stable baseline
P2	4 GB200 nodes, 16 GPUs total	non-colocate	Production-like train/rollout split
P3	same as P1/P2	candidate B-card operator path	A/B comparison

Suggested Baseline Config

Model and data:

--hf-checkpoint /path/to/Qwen3-30B-A3B-FP8
--ref-load /path/to/Qwen3-30B-A3B_torch_dist
--prompt-data /path/to/dapo-math-17k
--apply-chat-template
--input-key prompt
--label-key label
--rm-type deepscaler

Training and rollout:

--train-backend megatron
--advantage-estimator grpo
--num-rollout 100
--rollout-batch-size 32
--n-samples-per-prompt 8
--global-batch-size 256
--rollout-max-response-len 8192
--rollout-temperature 1.0
--rollout-top-p 0.95
--rollout-top-k 40
--balance-data

Megatron parallelism:

--tensor-model-parallel-size 4
--sequence-parallel
--pipeline-model-parallel-size 1
--context-parallel-size 2
--expert-model-parallel-size 8
--expert-tensor-parallel-size 1
--use-dynamic-batch-size
--max-tokens-per-gpu 8192
--recompute-granularity full
--recompute-method uniform
--recompute-num-layers 1
--attention-dropout 0.0
--hidden-dropout 0.0
--attention-softmax-in-fp32
--attention-backend flash

vLLM rollout:

--rollout-num-gpus 8
--rollout-num-gpus-per-engine 2
--vllm-gpu-memory-utilization 0.8
--vllm-router-policy consistent_hash
--vllm-enable-expert-parallel
--vllm-all2all-backend deepep_high_throughput

For colocate mode:

--colocate
--vllm-gpu-memory-utilization 0.5

Metrics To Track

Primary correctness and mismatch metrics:

• train/train_rollout_logprob_abs_diff
• rollout/log_probs
• rollout/rollout_log_probs
• train/ppo_kl
• train/pg_clipfrac
• train/loss
• rollout/raw_reward

Performance metrics:

• total step time
• rollout generation time
• Megatron logprob recomputation time
• weight sync time
• training time
• tokens/sec
• peak GPU memory
• vLLM KV cache utilization, if available

Suggested summary table:

Run	Hardware	Mode	Precision	Operator Path	Final Reward	logprob_abs_diff Mean	Step Time Mean	Notes
baseline	GB200	colocate	BF16 train + FP8 rollout	default	TBD	TBD	TBD	mismatch baseline
candidate	GB200	colocate	BF16 train + FP8 rollout	B-card operator	TBD	TBD	TBD	compare against baseline
baseline	GB200	non-colocate	BF16 train + FP8 rollout	default	TBD	TBD	TBD	production-like split
candidate	GB200	non-colocate	BF16 train + FP8 rollout	B-card operator	TBD	TBD	TBD	compare against baseline

Acceptance Criteria

The baseline is valid if:

• Training completes at least 30 rollout steps without NCCL, HTTP, CUDA illegal access, or OOM failures.
• rollout/raw_reward does not collapse.
• train/train_rollout_logprob_abs_diff is present in logs.
• The mismatch is non-trivial and measurable, ideally consistently above the near-aligned range observed in prior runs.
• Metrics are stable enough to compare against a later candidate run.
• Commit hash, Docker image, hardware type, run mode, exact config, and wandb/log links are recorded.

Follow-up Candidate Work

After the baseline is established, integrate Blackwell/B-card specific operators and rerun the same config.

The candidate comparison should check:

1. Does the B-card operator reduce train_rollout_logprob_abs_diff?
2. Does it preserve or improve rollout/raw_reward?
3. Does it improve rollout throughput or logprob recomputation time?
4. Does it introduce numerical drift under BF16/FP8 mixed paths?
5. Does it behave consistently between colocate and non-colocate modes?

Notes and Risks

• If mismatch is too small, increase the intentional rollout/training difference by enabling stronger sampling filters or using FP8 rollout.
• If mismatch is too large and reward collapses, reduce sampling filters before changing model parallelism.
• top_p < 1 or top_k > 0 can make rollout logprobs differ from Megatron replay logprobs because rollout-side processed logprobs may include sampling filter semantics.
• For GB200 multi-node runs, record NCCL, MNNVL, CUDA, PyTorch, vLLM, and Megatron environment details.
• Keep the first baseline simple. The immediate goal is not best reward, but a clean and measurable mismatch baseline for future B-card operator A/B comparison.

Control-Variable Drift and Operator Attribution Protocol

The intended experiment is not just to observe a train-rollout logprob gap. The goal is to use the gap as a controlled signal to locate which numerical operator or operator contract contributes most to train-inference drift, then develop Blackwell/B-card Megatron or vLLM CUDA kernels that improve operator-level train-inference consistency.

This is a reasonable direction, with one guardrail: do not assume the first observed drift is caused by a CUDA operator. First rule out semantic and framework-level mismatches such as weight-sync timing, token shift, batch layout, TP/CP slicing, top-p/top-k processed-logprob semantics, temperature handling, or fallback path changes. Kernel work should start only after the drift remains reproducible under controlled semantics and points to a specific numerical boundary.

Mismatch Knob Ladder

Run the experiments as a control-variable ladder. Start from a near-aligned baseline, then enable one mismatch knob at a time. Each step should compare against the previous step, not just against an absolute threshold.

Suggested ladder:

1. Near-aligned sanity run: same model state, same dtype where possible, top_p=1.0, top-k disabled, temperature fixed, minimal filters.
2. Enable FP8 rollout only: BF16 Megatron train/replay vs FP8 vLLM rollout, with sampling filters still disabled.
3. Enable sampling filter semantics only: keep dtype aligned, then test top_p < 1.0 and top_k > 0 separately.
4. Combine FP8 rollout plus sampling filters.
5. Add colocate/non-colocate and multi-node differences only after the smaller cases are understood.
6. Add RL-Kernel or candidate Blackwell operator paths only after the baseline and attribution signals are stable.

For every knob, record the delta in:

• train/train_rollout_logprob_abs_diff
• dlogp mean/max/p50/p90/p99, if audit diagnostics are enabled
• worst-token examples and whether they correlate with filtered tokens, long contexts, vocab shards, or specific experts
• backend/operator decisions and fallback reasons
• batch layout and rollout consistency metadata

Per-Knob Workflow

For each knob:

1. Record the metric delta from the previous run.
2. Classify the drift as one of:
   • expected numerical drift, for example FP8 quantization error;
   • semantic drift, for example rollout processed logprobs include top-p/top-k filtering semantics that Megatron replay does not reproduce;
   • framework drift, for example weight-sync timing, token alignment, batch layout, TP/CP slicing, or fallback path mismatch;
   • candidate operator drift, where semantics are aligned but the numerical boundary still diverges.
3. If drift increases, split the path into smaller probes before changing kernels.
4. Only promote a candidate operator hypothesis when the same samples, tokens, model state, and metadata point to a stable numerical boundary.

Operator Attribution Probes

When a knob introduces non-trivial drift, add probes in this order:

1. Verify sample and token identity: same prompt, response tokens, loss mask, shifted target IDs, and response slice.
2. Verify model state: same checkpoint/weight version and no train-step update between rollout logprob capture and Megatron recomputation.
3. Compare pre-softmax logits for the selected positions when feasible.
4. Compare logsumexp and selected-token logprob separately.
5. Compare vocab-parallel reduce behavior across TP ranks.
6. Compare top-p/top-k retained token sets and whether train replay uses the same processed-logprob contract.
7. Compare the native Megatron path, vLLM path, RL-Kernel path, and candidate Blackwell path under the same fixed tokens.

The likely operator boundaries to inspect first are:

• LM-head GEMM / output projection;
• vocab-parallel logsumexp or softmax reduction;
• selected-logprob gather and masking;
• fused linear_logp style kernels;
• FP8 quantization/dequantization and scaling policy in the rollout path;
• sampling-filter processed-logprob kernels when top_p or top_k is enabled.

Blackwell/B-Card Operator A/B Gate

Candidate Blackwell operator work should use two levels of validation:

1. Operator-level parity/microbench:
   • fixed hidden states, LM-head weights, target IDs, masks, TP metadata, and dtype policy;
   • compare Megatron native, vLLM native, RL-Kernel/native fallback, and candidate Blackwell kernels;
   • report selected logprob max/mean error, p99 error, backward parity when applicable, throughput, and memory.
2. End-to-end post-training A/B:
   • same workload, same seeds, same prompts, same rollout settings, same hardware, same commit/docker image;
   • baseline default operator path vs candidate B-card operator path;
   • compare logprob drift, reward stability, step time, rollout time, training time, and memory.

Do not count an operator as improving train-inference consistency unless it reduces drift under fixed semantics without hiding a framework fallback or changing the sampled-token contract.

Dependency on Phase 2 Observability

Small P0/P1 mismatch-baseline runs can start before all Phase 2 work is merged, because they validate configs and produce initial drift signals. However, operator attribution and B-card kernel conclusions should wait for the Phase 2 observability stack:

• execution decision records to show native/optimized/fallback paths;
• consistency metadata and batch fingerprints to rule out layout mismatches;
• dlogp audit diagnostics to identify token-level drift distribution;
• linear_logp backend metadata, fallback reason, timing, and optional memory probes.

This keeps framework/semantic inconsistency from being misclassified as a CUDA operator problem.

Additional Acceptance Criteria

The experiment should be considered ready for Blackwell operator development only if:

• the near-aligned sanity run has low and stable logprob drift;
• each mismatch knob has an isolated metric delta;
• at least one drift source is localized to a numerical operator boundary rather than sample alignment, weight sync, or processed-logprob semantics;
• the candidate operator can be tested with fixed tensors before the full training loop;
• the final A/B uses the same workload and metadata contract for baseline and candidate runs.