Skip to content

Add AutoTP=2 smoke test config and script - #1004

Merged
tohtana merged 5 commits into
masterfrom
gma/opsd-autotp
Aug 4, 2026
Merged

Add AutoTP=2 smoke test config and script#1004
tohtana merged 5 commits into
masterfrom
gma/opsd-autotp

Conversation

@delock

@delockdelock commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Verified on 2x RTX 5090 with Qwen2.5-0.5B-Instruct:

  • Student and teacher both use AutoTP=2
  • HybridEngineRollout.generate() produces identical output on both ranks
  • OPSD training loop runs end-to-end with temperature=0 (greedy)
  • No DistributedSampler in pure-TP mode (all ranks see same data)
  • Teacher uses ZeRO-0 when AutoTP is enabled (AutoTP+ZeRO-3 not supported)

Requires DeepSpeed >= commit 53a2ac4 (AutoTP KV cache consistency fix).

@delock
Add AutoTP=2 smoke test config and script
ac68cb7
Verified on 2x RTX 5090 with Qwen2.5-0.5B-Instruct:
- Student and teacher both use AutoTP=2
- HybridEngineRollout.generate() produces identical output on both ranks
- OPSD training loop runs end-to-end with temperature=0 (greedy)
- No DistributedSampler in pure-TP mode (all ranks see same data)
- Teacher uses ZeRO-0 when AutoTP is enabled (AutoTP+ZeRO-3 not supported)
Requires DeepSpeed >= commit 53a2ac4 (AutoTP KV cache consistency fix).
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
@delock
delock requested a review from tjruwase as a code ownerJuly 21, 2026 14:42
@PKUWZP
PKUWZP self-requested a review July 21, 2026 14:43
@delock

delock commented Jul 22, 2026

Copy link
Copy Markdown
ContributorAuthor

Note DeepSpeed AutoTP for now does not have column parallel in lm_head, we expect deepspeedai/DeepSpeed#8146 and its follow up (inplace softmax) would significantly reduce memory overhead in case of long sequence OPSD.

@delock
Enable AutoTP + ZeRO-3 for teacher and student; fix teacher ZeRO-3 gate
6e1791d
- teacher.py: drop the stale guard that forced ZeRO-0 whenever AutoTP was
enabled. AutoTP + ZeRO-3 is now supported by recent DeepSpeed, so the two
can be combined. Also gate ZeRO-3 on a real data-parallel dimension
(world_size > autotp_size) instead of world_size > 1, so pure-TP runs
(autotp_size == world_size, no offload) no longer pay needless ZeRO-3
per-forward gather overhead.
- configs/smoke_ds_zero3_offload_autotp2.json: student ZeRO-3 + param/optimizer
offload + AutoTP=2 DeepSpeed config.
- configs/smoke_hybrid_autotp_zero3.json: OPSD smoke with teacher and student
both on AutoTP=2 + ZeRO-3 (teacher offload).
- test_teacher_autotp_zero3.py / test_student_autotp_zero3.py: isolated
AutoTP + ZeRO-3 probes (env-tunable OFFLOAD/AUTOTP).
Verified on 4x RTX 4080 SUPER with Qwen2.5-0.5B/1.5B-Instruct:
- Teacher AutoTP=2 + ZeRO-3 (offload and 2x2) runs on upstream DeepSpeed master.
- Student AutoTP=2 + ZeRO-3 (offload and 2x2) and the full OPSD smoke
(teacher + student) require DeepSpeed PR #8168 (AutoTP + ZeRO-3 checkpoint
consolidation).
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
@delock

delock commented Jul 28, 2026

Copy link
Copy Markdown
ContributorAuthor

Now teacher support autotp + zero3, including autotp + zero3 offload.

tohtana
tohtana requested changes Jul 29, 2026
View reviewed changes

@tohtanatohtana 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.

Hi @delock,
This is a great improvement of the OPSD example! I left a few comments. I think the sampler issue should be addressed before we merge.

Comment threadtraining/opsd/main.py Outdated
Comment threadtraining/opsd/test_student_autotp_zero3.py Outdated
Comment threadtraining/opsd/test_teacher_autotp_zero3.py Outdated
delock added 2 commits July 29, 2026 22:35
@delock
Address review: fix DistributedSampler rank and use HF model IDs
8bcc48b
main.py: pass num_replicas=dp_size and rank=global_rank//tp_size to
DistributedSampler so data shards only along the real DP dimension.
Without this, the default sampler used the full world_size as replicas,
splitting data across TP-group ranks and deadlocking the all-reduce.
Verified end-to-end on CPU (4 procs, TP=2/DP=2) with Qwen2.5-0.5B-Instruct
on DeepSpeed master: full OPSD loop (rollout+teacher+student) runs without
deadlock.
test_*_autotp_zero3.py: replace local model paths with HF model IDs so the
defaults work out of the box.
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
@delock
Guard autotp_size()==0 in pure-DP mode
6cc7b76
DeepSpeed uses 0 as the sentinel for 'AutoTP not configured'
(tensor_parallel_config.autotp_size defaults to 0, and engine.autotp_size()
returns it verbatim). Without a tensor_parallel section in the DS config
(pure-DP runs), dist_world_size() // tp_size hit ZeroDivisionError.
Normalize 0 -> 1 so the DP size math matches DeepSpeed's own convention
(autotp_size not in (0, 1) means real TP). Verified pure-DP (world=2,
TP=1, DP=2) now runs end-to-end on CPU.
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@gmail.com>
@delock

delock commented Jul 29, 2026

Copy link
Copy Markdown
ContributorAuthor

@tohtana thanks for the catches! I have updated the code and verified it on TP=2xDP=2.

@delock
delock requested a review from tohtanaJuly 29, 2026 15:21
@delock

delock commented Aug 3, 2026

Copy link
Copy Markdown
ContributorAuthor

Hi @tohtana now with latest autotp+zero3 support on DeepSpeed this PR will be ready to merge, thanks!

tohtana
tohtana approved these changes Aug 4, 2026
View reviewed changes

@tohtanatohtana 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.

Thank you @delock! Let's merge this.

@tohtana
tohtana merged commit c230e2d into masterAug 4, 2026
2 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tohtanatohtanatohtana approved these changes
@tjruwasetjruwaseAwaiting requested review from tjruwasetjruwase is a code owner
@PKUWZPPKUWZPAwaiting requested review from PKUWZP

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@delock@tohtana