[0z5a]

Add an opt-in vocabulary-parallel KL route to OPSD. Teacher logits move directly from the actual local vocabulary shard to a host cache, and student training calls the KL operator from DeepSpeed #8791. Eager greedy rollout gathers only the final position's logits so token selection remains global.

The default dense route retains full-vocabulary heads on both models. This also fixes compatibility with the dependency's automatic sharding of tied HuggingFace heads. Cache copies complete before use, cache storage is freed on failure, and the existing causal shift, token mean, temperature-squared scaling and Engine GAS scaling are preserved.

Refs the vocabulary-parallel proposal #8790. This is the OPSD application integration; Core operators remain in deepspeedai/DeepSpeed#8791 and are not copied into this PR. deepspeedai/DeepSpeed#8791 is now merged into DeepSpeed master. The original matrix below was tested on Core cdc335793bd14dbb01a340c19d6dbcca9f86dc19. The update linked below was tested on DeepSpeed master with deepspeedai/DeepSpeed#8838 (Triton kernels for the same ops). The tested Examples base is 0c6dd5ed1ec0f0967ae0564534ecfa6f157c6746.

Incremental review

Incremental review report and validation · diff from the reviewed head.

Update: skip mask-free chunks on the vocabulary-parallel route, with new results · diff from the previous head · Update: copy teacher-cache chunks from pinned memory · diff. At the current head, with DeepSpeed#8838, C is 1.22× / 1.17× the throughput of B at B2/ISL1024 / B4/ISL2048 on the replay fixtures, with both routes reading the teacher cache through the new pinned copies, and uses 3.9 GiB/rank less peak memory at B4/ISL2048. The matrix below belongs to the original head.

The correction rejects mismatched benchmark labels/configs and prepares the output directory before model initialization. Updated-source checks: 63 CPU tests passed, 1 skipped, plus fresh-directory online smokes for all four implementation labels, with 20 applied optimizer updates and 0 skipped updates on 2 H100 GPUs. Both ranks changed FP32 optimizer masters and kept the teacher frozen in every case. Independent archive/report readback passed and owned workers and locks closed. The linked report records the configs and runtime; the performance matrix below belongs to the original submitted head.

Supported route

• Forward/reverse KL; student ZeRO-0/2; eager single-sample greedy generation. Teacher DP uses the existing ZeRO-3 route.
• Actual ordered TP members, vocabulary shard ranges, physical widths and tokenizer token-ID semantics are validated. TP peers must use identical tokens and masks; separate DP replicas may use different samples and lengths. TP=1 remains local with DP>1.
• JSD, stochastic/graph/shared-prefill rollout, teacher CPU offload, PP/SP and compiled or accelerated HybridEngine training are rejected. Frozen Qwen3 teachers have a known Core replicated-parameter hook initialization limitation; this PR does not patch Core.

Validation

The complete finite pretrained workload used pinned Qwen2.5-0.5B-Instruct student (7ae557604adf67be50417f59c2c2f167def9a775) and Qwen2.5-1.5B-Instruct teacher (989aa7980e4cf806f80c7fef2b1adb7bc71aa306). Each batch ran on two H100 GPUs, TP2/DP1, BF16, student ZeRO-0, GAS2, Torch 2.12.1+cu130 / CUDA 13.0, Transformers 5.15.1 and driver 595.91.07.

All 56 cases passed, with 1,900 actual optimizer updates and zero skipped updates: 955 online and 945 replay. This includes two five-update smoke cases and three independent repetitions of A/B/C online and replay workloads at each input shape. Each measured run has five warmup plus 30 measured updates. FP32 optimizer-master changes are checked separately from BF16 model hashes; Engine.global_steps alone is not accepted as proof of an applied update.

A uses unchanged upstream OPSD source with explicit full-logits initializer configuration. B uses the same sharded heads as C, gathers full training/cache logits, and applies dense KL with local-slice gather backward. C uses this vocabulary-parallel route. B/C is the comparison that isolates this integration from head-layout differences. All paths use the same public math SDPA backend and start from the same initial state. Replay uses the same real rollout fixture for each shape and excludes generation; online includes new generation, teacher/cache and synchronized forward/backward/GAS optimizer update.

Input shape	Mode	Path	Response token/s mean [min, max], 3 runs	GPU allocated peak GiB/rank	Teacher CPU cache MiB/rank
B1, ISL38-56	online	A: upstream dense	41.35 [40.14, 42.72]	5.287	33.906
B1, ISL38-56	online	B: gathered reference	40.72 [40.03, 41.68]	4.307	33.906
B1, ISL38-56	online	C: vocabulary parallel	41.26 [40.35, 42.06]	4.291	16.953
B1, ISL38-56	replay	A: upstream dense	543.82 [536.57, 553.01]	5.287	33.906
B1, ISL38-56	replay	B: gathered reference	534.07 [529.79, 537.84]	4.307	33.906
B1, ISL38-56	replay	C: vocabulary parallel	528.32 [525.21, 534.30]	4.291	16.953
B2, ISL1024	online	A: upstream dense	3.64 [3.46, 3.89]	10.343	630.594
B2, ISL1024	online	B: gathered reference	6.02 [5.91, 6.08]	9.628	630.594
B2, ISL1024	online	C: vocabulary parallel	3.51 [3.46, 3.58]	9.204	315.297
B2, ISL1024	replay	A: upstream dense	26.93 [26.08, 27.64]	10.330	630.594
B2, ISL1024	replay	B: gathered reference	26.22 [25.91, 26.55]	9.611	630.594
B2, ISL1024	replay	C: vocabulary parallel	21.72 [19.89, 22.65]	9.204	315.297
B4, ISL2048	online	A: upstream dense	100.02 [98.97, 101.14]	30.614	2448.188
B4, ISL2048	online	B: gathered reference	96.94 [96.20, 97.40]	29.828	2448.188
B4, ISL2048	online	C: vocabulary parallel	86.59 [85.10, 88.95]	29.998	1224.094
B4, ISL2048	replay	A: upstream dense	309.00 [307.61, 310.27]	30.549	2448.188
B4, ISL2048	replay	B: gathered reference	293.84 [293.10, 294.21]	29.826	2448.188
B4, ISL2048	replay	C: vocabulary parallel	205.56 [199.69, 216.27]	29.998	1224.094

Short inputs are 16 real math/instruction prompts, 38-56 tokens. The 1024/2048-token inputs are constructed code/document capacity requests. OSL64 is a limit; long-input responses can end earlier, and throughput counts actual response tokens once per TP group. C halves the teacher-cache payload relative to B; this does not imply halving complete GPU peak. The table reports complete measured step peaks and variability. These finite small-model runs do not establish distillation convergence, held-out quality, 14B/32B or TP8/8K qualification.

At the original head, with #8791's eager kernels, there is no demonstrated throughput speedup. On identical replay fixtures, C has about 17% lower throughput than B at B2/ISL1024 and 30% lower throughput at B4/ISL2048. At the largest online shape, allocated peak is slightly higher for C (29.998 vs 29.828 GiB/rank); its lower reserved pool is not evidence of a lower live allocated peak. The short-input throughput ranges overlap.

Online responses can diverge after training: at B2/ISL1024 the 90 measured updates generated 1260/2064/1224 effective tokens for A/B/C. Do not attribute the online token/s difference solely to execution cost. All three replay paths at that shape measured the same 1224 tokens, and all B4/ISL2048 groups measured 46080 tokens.

Default SDPA on the tested long-input early-EOS workload produced finite logits and KL logit gradients but NaN model gradients and a skipped optimizer update. The matched math-backend diagnostic had finite gradients and changed FP32 masters at the same learning rate. All reported A/B/C runs therefore use --sdpa-backend math; production attention-backend defaults are unchanged. The original negative outputs are preserved.

Additional checks at the original submitted head passed: 58 CPU tests with one explicit skip; real CUDA forward/reverse KL against an independent FP64 dense oracle across FP32/BF16, non-divisible vocabularies, temperatures and chunk sizes; uniform rejection of bad metadata, token semantics and TP inputs; real TP/DP Engine integration with ZeRO-2 student, ZeRO-3 teacher, empty-mask DP samples, GAS and new-process checkpoint resume. Full raw archives were independently read and checksummed, and all owned GPU workers and locks closed.

Examples' unchanged pre-commit configuration skips OPSD paths. Separate scoped Flake8, AST and git diff --check checks passed. The original performance matrix was executed on the Python source published in f74f9906b8237b29a67ce253e594ad3e435c4240. Validation of the subsequent benchmark preflight/output-path correction is recorded in the incremental review above; the full performance matrix has not been rerun for that correction.

Reproduce the public five-update pretrained smoke from training/opsd:

PYTHONPATH="$PWD" torchrun --standalone --nproc_per_node=2 \
  benchmarks/bench_vocab_parallel_opsd.py \
  --config configs/opsd_vocab_parallel_smoke.json \
  --mode online --implementation vocab_parallel --sdpa-backend math \
  --fixture artifacts/smoke-fixture --warmup-updates 0 --measure-updates 5 \
  --output-jsonl artifacts/smoke.jsonl
python -m pytest tests/ -v

AI assistance was used for implementation, test preparation and evidence analysis.

[0z5a]

Incremental review and corrections

Reviewed the original PR head f74f9906b8237b29a67ce253e594ad3e435c4240.

Two benchmark defects were identified and corrected:

1. --implementation dense_default could be paired with a vocabulary-parallel config and label the executed vocabulary-parallel loss as dense. The benchmark now rejects implementation/backend mismatches before model initialization. upstream_dense and dense_default require the dense backend (including its omitted default); gathered_reference and vocab_parallel require the vocabulary-parallel backend used to build their sharded heads. An explicitly sharded dense-default head is also rejected.
2. The documented smoke command failed on its first measured update if artifacts/ did not exist. The benchmark now creates the result parent directory and checks append access before model initialization, preserving existing output contents.

Corrected in 94fe46b. Review only this increment.

Validation of the corrected source:

• python -m pytest tests/ -q: 63 passed, 1 skipped, including backend mismatches for all four implementation labels, explicit sharded dense heads, fresh nested output directories, preservation of existing results, and an invalid output target.
• Fresh-directory online smokes passed for all four labels: upstream_dense, dense_default, gathered_reference and vocab_parallel. Each completed 5 actual optimizer updates, for 20 updates total, 0 skipped, with changed FP32 optimizer masters and a frozen teacher on both ranks. Each run created artifacts/smoke.jsonl and both rank receipts without pre-creating artifacts/.
• Hardware/runtime: 2 NVIDIA H100 GPUs, TP2/DP1, BF16, ZeRO-0 student, GAS2, Torch 2.12.1+cu130 / CUDA 13.0, Transformers 5.15.1, driver 595.91.07; pinned Qwen2.5-0.5B / 1.5B student/teacher revisions and the math SDPA backend documented in the PR description.
• Scoped Flake8, Python AST and git diff --check passed. The repository's unchanged pre-commit configuration skips OPSD paths. DCO passed for the correction commit.

The smokes used the documented five-update command with each implementation and its corresponding backend/head config. upstream_dense omitted the candidate-only loss_backend field to preserve the unchanged upstream config schema; the benchmark accepts its dense default. The earlier private setup failure with that field present is preserved and excluded from the pass counts. Complete raw outputs were archived; an independent reader checked archive/member SHA256 values, bound all result rows and eight rank receipts to the archived bytes, checked actual updates and frozen teachers, and verified worker closure and GPU-lock release. Executed source hashes match the correction commit.

The original 56-case / 1,900-update performance table remains evidence from the original submitted source. Those runs used correctly paired configs and prepared output directories; these findings do not invalidate their labels or measurements. The incremental smokes validate the corrected entry paths and output creation; it is not a rerun of the full performance matrix and supports no new speedup or convergence claim.

AI assistance was used for this review, correction and validation.

[0z5a]

Update: skip mask-free chunks on the vocabulary-parallel route

DeepSpeed#8791 has merged, and DeepSpeed#8838 adds Triton kernels for the same ops. The new commit e9fd43c (diff from the previous head) changes streamed_distillation_loss in two ways.

• Skip chunks without response tokens. Previously the vocabulary-parallel route computed every sequence chunk, so that a per-rank skip could never leave the op's collectives unmatched. With a 2048-token prompt and at most 64 response tokens, that meant 9 chunks were fetched from the host cache and computed per microstep where 1–2 are needed.
  • Now the TP group agrees on which chunks have response tokens through one MAX all-reduce of a per-chunk flag. Chunks without them are skipped on both routes.
  • The dense route still skips the same chunks, now with one host sync per call instead of one per chunk.
  • The test that pinned the old behaviour now checks that a mask-free chunk is neither fetched nor computed, and that backward stays safe.
• No fp32 copies before the op. The shards are passed to the op in their own dtype, since it computes in fp32 internally. The fp32 copies doubled what the kernels read and what autograd keeps per chunk. A temperature other than 1 still divides in fp32 first.

Results. 2× H100, TP2/DP1, BF16, same workload, replay fixtures and protocol as the original matrix: 5 warm-up and 30 measured updates, 3 runs per arm. All arms ran on one node in one session. Response tokens/s, mean [min, max]:

Input shape	B: gathered reference	C: this PR as submitted	C: with this update
B2, ISL1024	32.45 [32.32, 32.60]	32.26 [32.13, 32.42]	41.54 [41.38, 41.63]
B4, ISL2048	344.54 [343.76, 345.51]	307.23 [305.80, 309.34]	415.38 [415.16, 415.60]
B4, ISL2048 peak allocated GiB/rank	29.826	29.997	25.938

• Core version. These runs use DeepSpeed master (with #8791) plus #8838.
• Earlier slowdown. The 17%/30% slowdowns reported above used #8791's eager kernels. With #8838 but without this update, C is 0.99× / 0.89× of B.
• Numerics, skipping. Per-update losses, the post-training fixed-response loss and the final FP32 optimizer masters are bitwise identical with and without the skipping, across all 6 runs.
• Numerics, bf16. Passing bf16 instead of fp32 copies is not bitwise neutral with the #8838 kernels: per-update losses drift by up to about 1e-2 over 35 updates, comparable to the drift between the eager and Triton ops. The fixed-response loss before training is identical; the kernels' reductions differ by about 1 ulp between the two input dtypes, and that reaches the gradients. With the single-pass kernels proposed on DeepSpeed#8838, bf16 and fp32 inputs gave bitwise identical losses and optimizer masters where I compared them: a 2-update B4/ISL2048 run, and the op alone on a 1024 × 75968 shard.
• Path check. Every C run counted its KL op and Triton kernel calls, to confirm the Triton path ran.
• Online runs. Online smokes with the update passed both shapes. Online throughput is not compared, since generated lengths diverge after training.

Checks.

• CPU: 63 passed, 1 skipped.
• GPU: the 2× H100 integration runs above.

Not validated. The multi-rank skip agreement is exercised by these TP2 runs. I have not re-run the ZeRO-3 teacher / DP integration probe on this commit.

Next. The largest remaining entry in the profile is the pageable host-to-device copy of teacher-cache chunks. My guess is that the non-contiguous chunk slices of the pinned buffer cause the pageable path; that is not yet verified, and it is a follow-up.

AI assistance was used for implementation, test preparation and evidence analysis.

[0z5a]

Update: copy teacher-cache chunks from pinned memory

Follow-up to the "Next" item in my previous comment. The new commit cea3bf7 (diff from e9fd43c) changes only TeacherLogitCache.chunk_to_device. It also adds a CUDA test.

Cause. The teacher cache is a pinned [B, T, V] host buffer. Each loss chunk was copied with cpu_logits[:, start:end].to(device). For B > 1 that slice is strided, and PyTorch cannot DMA a strided host tensor directly. It first makes the slice contiguous in newly allocated pageable host memory, then copies that temporary synchronously. Every chunk therefore paid a host copy plus a pageable transfer, even though the buffer is pinned.

Change. Each row's slice cpu_logits[b, start:end] is contiguous within the pinned buffer.

• The chunk is now copied row by row into a contiguous device tensor. The row copies are non-blocking when the buffer is pinned.
• The pinned host allocator records the copies. A cache freed while copies are still queued keeps its block until they finish, so free() and the cache lifetime are unchanged.
• A dtype override is applied on the device after the copy. The stored-dtype chunk then briefly coexists with the converted one on the device; before, the conversion used a host temporary. All runs below fetch in the stored dtype. Host destinations keep the previous path.
• No staging buffer or layout change is needed.

Microbenchmark. One H100, bf16, chunks fetched from the replay shapes' caches. Each call is timed between two device synchronizations:

Chunk	Before	After	Contiguous pinned slice, same bytes
[4, 256, 75968] (155.6 MB) from [4, 2112, 75968]	25.2 ms (6.2 GB/s)	2.82 ms (55.2 GB/s)	2.81 ms
[4, 256, 151936] (311.2 MB) from [4, 2112, 151936]	52.6 ms (5.9 GB/s)	5.61 ms (55.4 GB/s)	5.60 ms

• Where the old time goes. The profiler reports the old copy as Memcpy HtoD (Pageable -> Device). In the first row, 10.1 ms of the 25.2 ms is the host-side .contiguous() and 15.0 ms is the pageable transfer.
• New path. The new copies are Pinned -> Device, one per row.
• Staging buffer. A reusable pinned staging buffer would take 9.5 ms for the same chunk, most of it the host copy.
• Lifetime check. In a test that frees the cache while its copies are queued behind a busy stream:
  • the allocator does not hand the block out early;
  • the copied values stay intact;
  • the next allocation of the same size reuses the cached block, so no new pinned buffer is allocated.

Results. 2× H100, TP2/DP1, BF16. Same workload, replay fixtures and protocol as before: 5 warm-up and 30 measured updates, 3 runs per arm, all on one node in one session. B (gathered reference) reads the same teacher cache, so it gets faster too. Its cache is twice as wide, so it gains more. Response tokens/s, mean [min, max]:

Input shape	B, e9fd43c	B, this update	C, e9fd43c	C, this update
B2, ISL1024	32.61 [32.43, 32.75]	37.32 [37.21, 37.40]	40.71 [40.63, 40.79]	45.43 [45.27, 45.53]
B4, ISL2048	345.93 [345.43, 346.25]	372.23 [371.95, 372.65]	413.28 [412.97, 413.77]	435.94 [435.69, 436.42]

• Speedup from this update. B2/ISL1024: B 1.144×, C 1.116×. B4/ISL2048: B 1.076×, C 1.055×.
• C relative to B. On the same commit, C is 1.22× / 1.17× of B. With e9fd43c's cache for both it was 1.25× / 1.19×. The ratio shrinks because B's chunks are twice as large.
• Attribution. Each update fetches 4 chunks: two microsteps, each with one 256-wide and one 63-wide chunk.
  • The microbenchmark predicts savings per update at B4/ISL2048 of 54 ms for C and 114 ms for B. Measured: 64 ms and 105 ms.
  • In a 3-update profile of C at B4/ISL2048, the device time of the fetch copies drops from 136.8 ms (16 pageable copies) to 28.0 ms (64 pinned row copies).
  • The Pageable -> Device time still left in a whole-run profile comes from outside the cache fetches. That profile also spans model loading and setup.
• Numerics. Across all 3 runs × 2 shapes × both B and C, the following are bitwise identical to e9fd43c: per-update losses, the fixed-response loss before and after training, the final parameters and the FP32 optimizer masters. The fetch pattern (count and bytes) is identical too. Peak allocated memory is unchanged.
• Online. An online B4/ISL2048 smoke of C gave identical losses and optimizer masters. Throughput was +1.2%, since rollout dominates.
• Pinned buffers. The number of pinned host allocations per run is the same as with e9fd43c, so the freed cache is still reused by the next step.
• Core version. DeepSpeed master plus #8838 at 97d4f66, which includes the single-pass forward. The e9fd43c columns were re-measured on this core and node, so they differ slightly from my previous table.

Checks.

• CPU: 63 passed, 2 skipped. The new CUDA test is skipped on CPU.
• With CUDA visible: 65 passed. The new test compares device chunks to host slices bit for bit, covering:
  • 1 and 3 rows;
  • pinned and pageable caches;
  • the last partial chunk;
  • both dtype overrides.
• Device integration harnesses. On 4 GPUs: the trainer harness at TP2 × DP2 and at TP1 × DP4, with a ZeRO-3 teacher and a ZeRO-2 student. On 2 GPUs: the numeric harness. Both pass, and their outputs are identical to e9fd43c. This also covers the ZeRO-3 teacher / DP integration run that my previous comment listed as not re-run.

Not validated.

• Non-CUDA accelerators. The row copies are generic, but non_blocking relies on the accelerator's pinned allocator recording the copy.
• The trainer harness uses one row per rank, so B > 1 is covered by the replay runs and the CUDA test.

AI assistance was used for implementation, test preparation and evidence analysis.