Repository navigation
Conversation
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Incremental review and correctionsReviewed the original PR head Two benchmark defects were identified and corrected:
Corrected in 94fe46b. Review only this increment. Validation of the corrected source:
The smokes used the documented five-update command with each implementation and its corresponding backend/head config. 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. |
1 participant
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. The tested Core dependency is
cdc335793bd14dbb01a340c19d6dbcca9f86dc19. This PR should merge after deepspeedai/DeepSpeed#8791 makes that API available in DeepSpeed master; that PR is currently approved and unmerged. The tested Examples base is0c6dd5ed1ec0f0967ae0564534ecfa6f157c6746.Incremental review
Incremental review report and validation · diff from the reviewed 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
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_stepsalone 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.
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.
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 --checkchecks passed. The original performance matrix was executed on the Python source published inf74f9906b8237b29a67ce253e594ad3e435c4240. 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/ -vAI assistance was used for implementation, test preparation and evidence analysis.