Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

[RFC] Vocabulary-parallel per-token ops for RL and on-policy distillation

Open
  • Is something wrong?
#8,790 4 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

@jinyouzhi is already working on this.

Since Oct 10, 2026.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
15/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
distributed-systems, machine-learning

Research direction

This is a design RFC, not a ready task. Start with the existing vocab-parallel cross-entropy from #8173 under deepspeed.sequence and the OPSD example it was built for. Settle the open questions (temperature, masking, module placement) before writing the proposed vocab_parallel_logprobs, vocab_parallel_entropy and vocab_parallel_kl_div ops. Done means agreed API, a full-vocab oracle test, and docs for PR-1.

Written by the indexing model from the issue text.

Description

Motivation

#8173 added vocabulary-parallel cross-entropy for AutoTP. The original motivation was sharding the OPSD teacher across TP ranks, which serves the On-Policy Distillation Trainer item on the roadmap (#8104). Cross-entropy alone does not cover RL and distillation losses. They need per-token quantities computed over the full vocabulary:

Method Per-token quantity needed
GRPO / PPO log-prob of the sampled token, plus entropy
Policy-gradient OPD (Thinking Machines, Qwen3, verl) sampled-token log-probs of student and teacher (k1/k3 KL estimators)
Teacher top-k distillation (verl forward_kl_topk) student log-probs at the teacher's top-k ids
Dense KL distillation (DeepSpeedExamples OPSD, TRL GKD) full-vocab forward or reverse KL
Token selection (TIP, arXiv 2604.14084) student entropy and teacher–student KL

With AutoTP's vocab-parallel LM head, each rank holds only [B, T, V/tp] logits. Today the only way to compute these quantities is to all-gather the full [B, T, V] logits. For example, the DeepSpeedExamples OPSD trainer gathers the teacher's full logits and caches them on CPU. On long sequences with 128K–256K vocabularies, these logits dominate peak memory. Megatron users get sharded versions of these ops from verl (vocab_parallel_log_probs_from_logits, vocab_parallel_entropy, _VocabParallelKLDivergence). AutoTP has nothing equivalent.

Proposal

Add three small ops in deepspeed.sequence.vocab_parallel_ops that work directly on the local shard:

vocab_parallel_logprobs(logits, index, tp_group, vocab_start_index, vocab_end_index, ignore_index=-100)  # index: [...] or [..., k]
vocab_parallel_entropy(logits, tp_group)
vocab_parallel_kl_div(student_logits, teacher_logits, tp_group, reverse=False)  # KL(t||s), or KL(s||t) if reverse
  • Each op returns an fp32 per-token tensor, replicated across the TP group. Masking, reduction, and temperature are left to the caller, which is where existing trainers already handle them.
  • The ops compose into the losses above. For example, teacher top-k forward KL is two lines on top of vocab_parallel_logprobs.
  • Communication: one MAX and two SUM all-reduces of [N]-sized per-token statistics. The KL op packs student and teacher statistics into the same collectives. The logits are never gathered.
  • Backward follows the Megatron convention: every rank backprops the same loss, so gradients are closed-form, local, and need no communication.
  • Memory: each op saves only its inputs and per-token statistics, and processes rows in bounded chunks.

Measured on 1× Intel Arc Pro B70, T=8192, V_local=16032 (128K vocab / TP8), bf16. Peak is fwd+bwd, excluding inputs:

Op Naive PyTorch Proposed
logprobs 1506 MiB 443 MiB
entropy 3012 MiB 586 MiB
reverse KL 3012 MiB 698 MiB
forward KL 2510 MiB 634 MiB

Plan

# Item Status Notes
PR-1 The three ops, with tests against a full-vocab oracle and docs ✅ #8791 Perf follow-up #8838: fused Triton kernels (eager fallback), ~20× faster fwd+bwd, no extra fp32 copies, no API change
PR-2 Wire the ops into the OPSD example, so the sharded teacher caches V/tp logits instead of gathering full [B, T, V] logits 🔄 deepspeedai/DeepSpeedExamples#1019 (@0z5a) Covers forward and reverse KL. The teacher cache shrinks by tp. Benchmarks are in the PR.
3 Optional: a fused-linear or chunked variant that avoids materializing the logits shard ⚠️ CE only, #8792 Whether KL needs it depends on PR-2's numbers.
Possible next steps (not yet owned)
Item Why
Teacher top-k cache The other caching option in the original plan. It needs a vocab-parallel top-k. The student side already works through vocab_parallel_logprobs(index=[..., k]), and the cache does not depend on the teacher's TP degree.
Generalized JSD (0 < β < 1) The Core generalized_jsd_loss supports it. It composes from the same softmax statistics plus one [N] reduction.
Vocab-parallel token selection for rollout (argmax, top-k sampling) Avoids gathering the last-position logits during rollout and enables stochastic rollout with a sharded head. Related to #8197.
vocab_parallel_lm_head for init_inference Lets an inference-initialized teacher produce sharded logits.

Open questions

  • JSD / generalized JSD (TRL GKD): is there demand? It composes from the same statistics, and a pair of KLs approximates it.
  • Temperature and T² scaling: leave both to the caller, or add a temperature argument?
  • Reduction and mask: should the ops stay per-token, or also offer a masked-mean option?
  • Student-support top-k with a tail bucket (TRL SDPO): this needs a vocab-parallel top-k. Is it worth adding?
  • Teacher and student with different TP degrees or sharding: out of scope for now.
  • Module placement: deepspeed.sequence sits next to the existing vocab-parallel CE, or should this go under a future deepspeed.runtime.rlhf?

cc @delock @PKUWZP

Dominant language
Python
Stars
43.3k
Forks
5.1k
Avg merge
4d 20h
Merged PRs (30d)
69

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepspeedai/DeepSpeed

All issues in deepspeedai/DeepSpeed

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.