[WS2] TP-aware deterministic logprob for cross-config alignment (Qwen3-8B TP=2 CP=2 BF16)

[KJLdefeated]

Background

From roadmap #83 WS2 cross-config alignment task split. Target: Qwen3-8B, TP=2, BF16. The goal is to make rollout-side old-logp and training-side recomputed logp follow same numerical contract under TP and dlop is within the #108 contract.

The selected-token logprob is logit_target - logsumexp_vocab(logits), and under vocab-parallel TP the logsumexp over the full vocabulary is a cross-rank reduction.

Goal

Implement and validate a deterministic TP-aware selected-token logprob path for the WS2 cross-config stack.

• Model: Qwen3-8B dense
• Dtype: BF16 logits with FP32 accumulation for max / sumexp / exp / log; downcast at final write
• Parallelism: TP=2 (vocab-parallel)
• Vocab contract: Qwen3-8B vocab (~151k) sharded across TP=2
• TP is the numerical merge axis: deterministic cross-vocab-shard (max, sumexp) LSE merge + masked single-owner selected-logit gather
• Comparison target: vocab-logprob LSE drift first, then selected-token dlogp on active response tokens

Non-goals

• CP would not get into this op, since logp only shards the logits in vocab dimensions.
• PP consistency
• FP8 rollout logits / FP8 logprob path / quantized LM head (separate representation-drift issue)
• MoE / MLA routing effects on the LM head
• ROCm / domestic accelerator backend
• Real vLLM vs real Megatron drift-to-zero, which belongs to WS3
• Tokenizer / mask / padding correctness ownership (logical-input alignment is a validated precondition here, not fixed here)
• Performance tuning before the deterministic contract and validation harness are landed

Planned PRs:

• PR 1 — Logprob TP contract + dispatch metadata @ryankert01 @KJLdefeated [WS2][PR1][Logp] Add TP-aware logprob contract and dispatch metadata #259
  Define ShardingSpec / ReductionSpec for the logprob op, role tags (train/infer), and required metadata: vocab_start/vocab_end per rank, padded-vs-real vocab size, tp_rank/tp_world, active-token mask, ignore index. Explicitly declare CP as a non-merge axis (logp: Shard(sequence) → Shard(sequence), no cross-CP arithmetic).
  Acceptance: contract documented; missing vocab-shard bounds or reduction metadata fails loudly; no silent fallback.

• PR 2 — Single-GPU logprob harness registration @hihaluemen [WS2][PR2][Logp] feat: add TP=1 logprob comparison harness #262
  Register logprob in the cross-implementation harness against the existing WS1 deterministic logp. No communication yet.
  Acceptance: TP=1 path is bitwise equal to the current det-logp (it's the same code path, so this is a regression guard, not new math); harness reports LSE drift and active-token-only dlogp at max/p95/p99.

• PR 3 — Deterministic vocab-parallel TP logprob reference @KJLdefeated [WS2][PR3][Logp] Add deterministic vocab-parallel TP logprob reference #265
  per shard: m_l = max(local_logits), s_l = Σ exp(local - m_l) in FP32
  gather (m_l, s_l) partials; NCCL is transport, not the numerical reducer
  merge on every rank in fixed global vocab-shard index order: M = max_l m_l, S = Σ_l s_l · exp(m_l − M), LSE = M + log S
  selected logit: masked single-owner gather (exactly one rank holds y_t, so the sum has one non-zero term and is order-independent by construction) downcast at final write only
  Acceptance: TP=1 bitwise-equals PR2 reference; TP=2 bitwise-equals TP=1 (this should be achievable, not just within [WS1] Ground-truth harness + numerical contract for batch-invariant ops #108 tolerance — the merge is order-fixed and every rank runs identical ops).

• PR 4 — Cross-config integration + distributed drift report @hihaluemen [WS2][PR4][Logp] feat: add distributed TP/CP logprob drift report #289
  Wire into the WS2 planner/runtime for the scoped target and emit artifacts in the [WS2] Tolerance contract alignment & WS2 drift report format #116 format. Sweep TP=1/2/4 × CP=1/2 × BF16 vs FP32 reference.
  Acceptance: backend selectable explicitly; provenance records TP/CP topology, dtype, reduction spec, LSE merge order; unsupported materialization fails rather than silently switching; reports include LSE drift and active-token dlogp per-rank, with exact launch commands.

• PR5: Deterministic GRPO loss on the TP-aware logprob path @KJLdefeated [WS2][PR5][Logp] Add deterministic DP-aware GRPO loss on the TP-aware logprob path #291
  Bind the GRPO objective — ratio, clip, KL, advantage normalization — to a declared reduction contract so both the loss and its gradient are reproducible across TP/CP/DP layouts.

  • Scope: LossReductionSpec declaring the normalizer (token_mean / seq_mean / seq_sum_token_mean), whether n_non_ignore is global or microbatch-local, and the summation order across tokens → sequences → CP ranks → DP ranks. Cross-CP loss reduction becomes a fixed-order gather-and-merge, not an NCCL SUM.
  • Acceptance: the KL=0 identity — when logp == old_logp the ratio is exactly 1.0, pg_loss == -A·logp bitwise, and the gradient matches the analytic form; bitwise-stable across repeat runs; layout invariance (TP/CP/DP) within [WS1] Ground-truth harness + numerical contract for batch-invariant ops #108 tolerance.
  • Explicit non-goal: fusing the loss. That's a follow-up once the reference exists.

• PR6: TP-aware deterministic fused logprob kernel (Optional)
  Promote the WS1 fused online-softmax logp kernel to vocab-parallel TP, emitting rank-local partials that satisfy the PR2 merge contract, with fail-closed capability gating and no silent fallback.

  • Scope: kernel emits (max, sumexp, selected_logit) in FP32 per token, matching PR2's partial layout exactly; merge goes through PR2's fixed-order path rather than a collective reducer. Capability gate rejects row-parallel heads, undeclared padding, and uneven shards.
  • Acceptance: two independent oracles, asserted separately — fused-vs-PR2-reference (bitwise on the merge, tolerance on the rank-local kernel) and PR2-reference-vs-FP32-oracle. Never fused-vs-fused.
  • Worth knowing: TE's parallel_cross_entropy already implements this exact shape (online_softmax_kernel → all_gather_into_tensor of (m, d, X_y) → rank-ordered merge loop). It's usable as a third-party oracle for PR6, which gets you out of the self-mirroring problem for free.