[jinyouzhi]

Motivation

Teacher-student training compares the models' next-token distributions, not
just the probability of a single ground-truth token. DeepSpeed already offers
vocabulary-parallel forward/reverse KL, but does not support the interior
0 < beta < 1 generalized Jensen-Shannon divergence used by GKD-style
distillation. Applying a local softmax independently on each vocabulary shard
does not compute the intended full-vocabulary objective.

TRL currently has two relevant implementations on its main branch:
experimental.GKDTrainer computes generalized JSD from materialized logits;
DistillationTrainer uses completion filtering and tiled LM-head projection
to compute KL/JSD without materializing full logits. Those implementations do
not call DeepSpeed's vocabulary-parallel divergence ops. In particular,
support for DeepSpeed/ZeRO or FSDP parameter sharding is not itself support
for a vocabulary-parallel softmax and divergence.

This PR supplies the missing Core logits-based primitive, as a building block
for teacher-student workloads and a possible subsequent TRL integration. It
does not claim to integrate TRL or replace its memory-efficient tiled loss.
Refs #8790.

Changes

• Add vocab_parallel_jsd(student_logits, teacher_logits, beta=0.5, tp_group=None), returning FP32 per-token divergence replicated on TP ranks.
• Add fused Triton paired softmax statistics, divergence and backward kernels.
  Reuse global normalization and reduce the two KL-to-mixture expectations
  together. Recompute probabilities in registers during backward rather than
  retaining full FP32 probability tensors; keep bounded eager chunks as fallback.
• Enable JSD on an active Intel Triton backend without changing XPU's global
  Triton capability flag or existing-op dispatch.
• Support gradients to both inputs, detached teachers, uneven vocabulary
  shards, empty local shards, and non-contiguous causal-shift views.
• Preserve the existing Core convention: beta=0 selects forward KL and beta=1
  selects reverse KL. These are explicit endpoint branches, not continuous
  limits of the weighted interior JSD.
• Leave temperature, optional T-squared scaling, masking and normalization to
  the caller, following the existing vocabulary-parallel per-token ops.

Not a duplicate

• Add vocabulary-parallel logprobs, entropy and KL ops #8791 / Add Triton kernels for vocab-parallel logprobs, entropy and KL #8838 implement logprobs, entropy and forward/reverse KL, not
  interior-beta generalized JSD. The new API reuses their KL endpoints.
• Sum the OPSD JSD loss in float32 for bf16 and fp16 logits #8947 fixes counting and accumulation in the existing dense rollout loss.
  This PR does not edit runtime/rollout/opsd.py or its tests.
• DeepSpeedExamples#1019 integrates vocabulary-parallel KL into OPSD. This PR
  does not change Examples, cache handling, rollout or trainer integration.
• [VocabTP] vocab_parallel_ce support fused LM-head projection + loss for the HF forward #8945 connects fused-linear CE to HF forward. This PR is independent of it:
  no LM-head projection changes or loss-function hooks.

Scope and limitations

This is a new opt-in logits-based API; existing ops and training defaults are
unchanged. Unlike the existing single-pass KL forward, JSD needs global
normalization before evaluating the nonlinear mixture: statistics and
divergence are two input passes. It still consumes materialized local logits.
Top-k restriction, pointwise
clipping, differing student/teacher shard layouts, fused-linear divergence,
DeepCompile support and automatic TRL integration are not included.

Every TP rank must supply identically ordered student/teacher vocabulary
shards and backpropagate the same replicated per-token loss.

Validation

Hardware: up to 4 Intel Arc Pro B70 GPUs, XCCL, driver 1.15.39122+11,
torch 2.14.0.dev20260804+xpu, Intel Triton 3.8.0. Final validation: 67 CPU
regressions passed, 11 fused XPU cases passed, and 4 distributed cases passed
at each of TP1/2/4. Existing entropy/KL Triton-only tests still skip on XPU
because their accelerator-wide flag remains disabled; JSD uses its own checked
Intel backend dispatch. The final TP cases cover FP32/BF16 interior JSD versus
an FP64 oracle, empty shards/disjoint support, and head SGD update parity.

On NVIDIA L20, driver 580.173.02, torch 2.14.1+cu130 (CUDA 13.0),
Triton 3.8.0, NCCL: the full test file passed, 113 passed with no skips.
This includes all JSD TP1/2/4 cases and FP32/FP16/BF16 single-device kernels.
The existing ds conda environment was used without dependency changes;
tested source hashes match the local implementation and tests.

PYTHONPATH=tests:. OMP_NUM_THREADS=4 pytest \
  tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py -q \
  -k 'not TestVocabParallel and not triton'
PYTHONPATH=tests:. OMP_NUM_THREADS=4 pytest \
  tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py -q -k jsd_triton
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 OMP_NUM_THREADS=4 PYTHONPATH=tests:. \
  pytest --forked tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py -q \
  -k 'TestVocabParallelJSDTP and (0.2 or empty or optimizer)'
pre-commit run --files deepspeed/ops/triton_ops/vocab_parallel_softmax.py \
  deepspeed/sequence/vocab_parallel_ops.py \
  tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py \
  docs/_tutorials/autotp-training.md
git diff --check
# L20 full-file run:
CUDA_VISIBLE_DEVICES=0,1,2,3 PYTHONPATH=tests:. OMP_NUM_THREADS=4 \
  pytest --forked tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py \
  -q --color=no

Scoped pre-commit hooks and diff checks pass. Numerical checks use an
independent FP64 dense PyTorch reference, nonuniform vocabulary partitions,
FP32/BF16, both-input gradients and frozen inputs. Additional cases cover
disjoint support, undefined masked rows, empty local shards, empty token
batches, causal-shift views, independent large logit shifts and row chunking.
Three SGD updates of a sharded head match the dense reference at TP1/2/4.
This head-only training test does not exercise a DeepSpeed Engine or ZeRO.

Exploratory logits-only forward+backward A/B on B70, N=1024,
V=151936, BF16, beta=0.3, frozen teacher:

TP	Implementation	Peak allocated GiB/rank	Time ms, median
2	Gather both logits, dense PyTorch JSD	6.671	112.17
2	Vocabulary-parallel eager JSD	1.057	58.38
2	Vocabulary-parallel Triton JSD	0.435	3.08
2	Existing Triton reverse KL	0.435	3.29
4	Vocabulary-parallel eager JSD	0.841	31.14
4	Vocabulary-parallel Triton JSD	0.217	3.41
4	Existing Triton reverse KL	0.217	3.16

Five warmups and ten measured trials, synchronized wall time; per-trial MAX
over TP ranks, independently for time and memory. Resident input allocations
are included in peak. Student gradients are included. All arms start with sharded logits, and the dense arm includes
their gather; it is not TRL's existing chunked-projection implementation.
Losses and local student gradients match before timing. Clocks were not
pinned. This benchmark excludes head projection, backbone, optimizer and
rollout; it demonstrates only an op-level comparison, not full-training or
TRL speedup. Reproduction script and raw JSON/logs are retained locally and
should be attached to the PR when submitted. KL's Triton kernels are enabled
only in the XPU benchmark for comparison; their production dispatch is unchanged.

L20, same workload and protocol:

TP	Implementation	Peak allocated GiB/rank	Time ms, median
2	Gather both logits, dense PyTorch JSD	6.671	98.02
2	Vocabulary-parallel eager JSD	1.057	34.43
2	Vocabulary-parallel Triton JSD	0.435	1.95
2	Existing Triton reverse KL	0.435	2.21
4	Gather both logits, dense PyTorch JSD	6.526	103.41
4	Vocabulary-parallel eager JSD	0.841	16.78
4	Vocabulary-parallel Triton JSD	0.217	1.57
4	Existing Triton reverse KL	0.217	1.82

L20 clocks were also not pinned; these remain exploratory op-level timings.

All 15 CUDA sm89 offline kernel compilations pass (FP32/FP16/BF16 and both/
student-only/teacher-only backward variants), in addition to the actual L20
runtime validation above.

Not validated: ZeRO stages, offload, PP/SP,
DeepCompile, real-model convergence or TRL integration. No end-to-end speedup
or convergence claim is made.

AI assistance

Implementation, test preparation and this description were prepared with
GitHub Copilot CLI. Human review and independently run tests are required
before submission.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-30T15:07:30.008179Z	1036bdc	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[delock]

Why we need this behavior change? This seems not related to vocab parallel.

[delock]

Hi @jinyouzhi , thanks for the PR. There seems to be a unrelated hunk in this PR, can you take a look? Thanks!

[jinyouzhi]

@delock Thanks! That hunk belongs to #8959 (fixes #8948); removed it from this PR in c419e94.