Skip to content

Add OPSD (On-Policy Distillation) training example#1002

Merged
delock merged 25 commits into
masterfrom
gma/opsd
Jul 15, 2026
Merged

Add OPSD (On-Policy Distillation) training example#1002
delock merged 25 commits into
masterfrom
gma/opsd

Conversation

@delock

@delock delock commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

This PR moved example code from @PKUWZP 's OPSD PR in DeepSpeed (deepspeedai/DeepSpeed#8027).

Note the basic OPSD infrastructures (OPSD trainer, rollout engines) are still in the original PR and need to be merged seperately. This PR will work when the original PR merged.

@delock delock requested a review from tjruwase as a code owner June 24, 2026 09:30
Comment thread training/opsd/configs/opsd_vllm_disjoint.json Outdated
Comment thread training/opsd/configs/opsd_vllm_disjoint.json Outdated
Comment thread training/opsd/configs/opsd_hybrid_engine.json
@delock

delock commented Jul 1, 2026

Copy link
Copy Markdown
Contributor Author

Hi @tjruwase I have changed according to your first and second comments. For the third comment, I left an explaination. Thanks!

delock added 2 commits July 2, 2026 10:34
Entry point, configs, data, and tests for on-policy distillation
using DeepSpeed's hybrid engine rollout and vLLM backend.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…m_dtype to engine_dtype

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Comment thread training/opsd/README.md Outdated
delock added 6 commits July 3, 2026 17:51
…DeepSpeed

- Delete vLLM configs, scripts (opsd_vllm_disjoint.json, smoke_vllm.json, train_opsd_vllm.sh)
- Add trainer.py, config.py, losses.py, utils.py (moved from DeepSpeed)
- Add benchmarks/ (5 hybrid engine benchmarks moved from DeepSpeed)
- Update main.py imports (trainer, config now local)
- Update test imports (losses, utils now local)
- Rewrite README (remove all vLLM sections)

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…ants from JSON

- Subclass DeepSpeed's RolloutConfig to add temperature/top_p/etc
- Remove weight_sync_interval from JSON configs (vLLM remnant)

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…per fix

- Remove 14B/multi-GPU benchmarks (bench_14b_rollout, bench_autotp_gc,
  bench_hybrid_tp, bench_hybrid_tp_opt)
- Fix bench_decode_1p1r: wrap model for HybridEngineRollout
- Add --graph-capture CLI flag

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…nfigs

temperature=0 causes logits/0 = inf → NaN loss. The correct default
for knowledge distillation is temperature=1.0 (standard softmax).

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Comment thread training/opsd/teacher.py
delock added 2 commits July 12, 2026 22:45
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
@PKUWZP PKUWZP self-requested a review July 12, 2026 16:30

@PKUWZP PKUWZP left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall this PR looks good. My only suggestion is that we remove all "Microsoft Corporation Copyright" from all files in this PR. After that we should be able to merge it.

Comment thread training/opsd/benchmarks/bench_decode_1p1r.py Outdated
Comment thread training/opsd/scripts/train_opsd_hybrid.sh Outdated
Comment thread training/opsd/tests/test_losses.py Outdated
Comment thread training/opsd/tests/test_teacher_caching.py Outdated
Comment thread training/opsd/config.py Outdated
@delock

delock commented Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

@tjruwase @PKUWZP thanks for the comments. I'll fix the copyright lines.
For teacher configuration. Let me do some study to see if deepspeed can always apply to teacher logits computation.

delock added 7 commits July 13, 2026 11:31
… at TP=1 / under ZeRO-3

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…xample

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…es review)

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
delock added 8 commits July 13, 2026 15:51
…loading, else stage 0

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…lease

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…fig knob

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
…plicitly

Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
@delock

delock commented Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

@tjruwase @PKUWZP thanks for the comments. I'll fix the copyright lines. For teacher configuration. Let me do some study to see if deepspeed can always apply to teacher logits computation.

Now deepspeed is always used for teacher no matter offload is true or not. And all Microsoft file head had been fixed.

I plan to enable tensor parallel in a sperate PR, which is useful in OPSD. This add some complexity so better put in a seperate PR.

@delock delock requested a review from PKUWZP July 13, 2026 15:04
@delock delock merged commit 79cdcf4 into master Jul 15, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants