[AutoTP] Support frozen AutoTP models and manage AutoTP mode per model by jinyouzhi · Pull Request #8851 · deepspeedai/DeepSpeed

jinyouzhi · 2026-10-17T16:19:23Z

Summary

Two bugs show up when a frozen teacher runs next to an AutoTP student in one process, as in on-policy distillation (deepspeedai/DeepSpeedExamples#1019, #8850).

1. Frozen replicated parameters break deepspeed.initialize.
register_replicated_grad_hooks registers a gradient all-reduce hook on every replicated (non-TP) parameter. PyTorch refuses hooks on tensors with requires_grad=False, so a fully or partially frozen model fails during initialize.

Frozen params are now skipped. They get no gradient, so their TP replicas cannot diverge.
To keep that guarantee, the engine remembers which replicated params were frozen at initialize. If any of them becomes trainable later, the next backward raises a clear error instead of letting TP replicas drift apart silently.

2. The AutoTP training mode leaks into init_inference.
deepspeed.initialize with AutoTP sets a process-global training-mode flag. A teacher passed to deepspeed.init_inference afterwards still sees training mode. As a result:

it is sharded with the student's TP group instead of its own tp_size;
its LM head is not handled as for inference;
its parameters are left trainable.

The fix makes the mode per model instead of process-wide:

AutoTPMeta, which AutoTP already passes to every layer it builds, gains a training field. AutoTP fills it from its training_mode argument, which deepspeed.initialize sets to True and init_inference to False.
Each TensorParallel_Layer takes its mode from tp_meta, so the forward path (RowParallel, add_bias) follows how that model was sharded. The student keeps training-mode behavior after the teacher is built, and vice versa.
AutoTP.set_tensor_parallel_config and replace_transformer_layer (kernel injection, set_lm_head, MP checkpoint saving) use the explicit training_mode instead of the global flag.
Backward compatibility: AutoTP(training_mode=None) (the new default) and layers built without tp_meta still follow set_autotp_mode, so existing direct callers behave as before.

Related work

I searched open and closed PRs and issues for register_replicated_grad_hooks, AUTOTP_MODE and frozen AutoTP models, and found no existing fix. Reported while integrating #1019.

Sibling paths checked:

Every is_autotp_training_mode() reader in deepspeed/: auto_tp.py (set_tensor_parallel_config) and replace_module.py (kernel injection, set_lm_head, MP checkpoint saving) now use the explicit mode; the remaining reads are the fallbacks for callers that pass no mode.
All AutoTP(...) call sites in deepspeed/ (engine, TpTrainingManager, InferenceEngine) already pass training_mode explicitly. No code in deepspeed/ builds TP layers outside AutoTP.

Tests

New tests in tests/unit/v1/autotp/test_tp_plan_e2e.py. Each fails without the corresponding fix and passes with it.

test_frozen_qwen3_initializes_and_runs: 2 ranks; the frozen model's outputs match the unsharded reference.
test_unfreezing_replicated_param_after_initialize_raises: 2 ranks.
TestInferenceTeacherAfterTrainingStudent: 4 ranks. A student with TP=2 is trained, then a teacher is built with init_inference(tp_size=4) and checked to be sharded 4 ways and frozen.

Results at head 68647c93e:

Tests	Result
`tests/unit/v1/autotp/test_tp_plan_e2e.py`	8 passed
`tests/unit/module_inject` + the other `tests/unit/v1/autotp` files	171 passed, 37 skipped
`tests/unit/v1/autotp/test_autotp_training.py` (skipped on XPU by `skip_on_device`; run locally with that skip removed)	42 passed, 2 skipped (bf16 overflow skip)
`tests/unit/v1/autotp/test_autotp_training.py` with `DS_ACCELERATOR=cpu`	11 passed, 33 skipped
`test_autoep_autotp_sharded_attention.py`, `test_autoep_autotp_runtime.py`, `test_per_head_muon_accelerator.py`, `test_checkpoint_sharding.py` (`-k "not True"`)	44 passed, 4 skipped, 3 failed

The 3 failures (TestTiedOutputHead::test_row_parallel_tied_head_is_refused[False] and two test_per_head_on_a_shard_is_the_shard_of_per_head[2-*]) fail the same way on master a4811fbf9. The zero_init=True cases of test_autoep_autotp_sharded_attention.py were deselected because test_grads_and_norm_match_reference[3-True-4-2-2] hangs on master as well.

Hardware and software:

8× Intel Arc Pro B70 (2/4/8 ranks depending on the test)
Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11)
torch 2.14.0.dev20260804+xpu, transformers 5.17.0

Compatibility matrix

The changes are mode bookkeeping and hook registration; no numerics change on the existing training path.

	Validated	Not validated
ZeRO	New tests: stage 0. Existing AutoTP suite: stages 0/1/2.	Frozen model with ZeRO-3 + AutoTP
Precision	New tests: fp32. Existing suite: fp16/bf16.
Offload		Not exercised
Parallelism	TP=2 and TP=4; student TP=2 with teacher TP=4	DP>1 together with a frozen teacher
Accelerator	XPU	CUDA, HPU

AI assistance

This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.

cc @0z5a I try to fix the bug you described in Vocab-parallel pipeline with OPSD.
@delock And address AutoTP_MODE issue, I think OPSD will reach the issue soon.

Distillation runs a trainable student and a frozen teacher in one process. - register_replicated_grad_hooks registered a gradient hook on every replicated parameter, which PyTorch refuses for parameters that do not require gradients, so initializing a frozen model with AutoTP failed. Frozen parameters have no gradient to reduce and are now skipped. - deepspeed.initialize switches the process-wide AutoTP mode to training and never resets it, so a teacher passed to init_inference afterwards was sharded with the student's TP group and left trainable. init_inference now builds its engine in inference mode and restores the caller's mode, and each AutoTP layer records the mode it was built in instead of reading the global. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Replicated params frozen at deepspeed.initialize get no gradient all-reduce hook, so making them trainable later would let TP replicas drift silently. Raise a clear error in the backward prologue instead. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Layers and AutoTP read the training mode from AutoTPMeta, which AutoTP fills from its explicit training_mode argument. A frozen inference teacher and a training student therefore keep their own mode without init_inference having to save and restore the process-wide flag. Callers that build AutoTP or layers without a mode keep following set_autotp_mode. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Layers built without tp_meta captured the global AutoTP mode at construction, so a later set_autotp_mode(training=True) no longer reached them and the in-place inference bias add broke autograd. Resolve the mode at call time when tp_meta does not set it. Also let the 4-rank teacher/student test run on CPU. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

chatgpt-codex-connector · 2026-10-21T17:06:36Z

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-21T17:14:52.855659Z	`63bb7f5`	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 63bb7f58cb

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-10-21T17:14:50Z

+        if training_mode is None:
+            training_mode = is_autotp_training_mode()


Preserve the dynamic mode for default AutoTP instances

When a direct AutoTP caller leaves training_mode=None and later switches modes with set_autotp_mode (for example, constructing under inference and enabling training before replacement or forward), this eagerly collapses None to the current global value and stores that boolean in every layer's tp_meta. Unlike the pre-change behavior and TensorParallel_Layer.is_training_mode()'s new fallback, those layers can no longer observe the later mode change; after gradients are re-enabled, an inference-built LinearAllreduce still uses the in-place bias/inference collective path and can fail backward. Keep None in the per-layer metadata while resolving a separate effective mode only for construction-time routing.

Useful? React with 👍 / 👎.

Fixed in 11a3cb3: tp_meta now keeps training=None, and None is resolved to the current mode only for construction-time routing. test_layer_without_explicit_mode_follows_set_autotp_mode now also covers a layer built with a default AutoTP's tp_meta (fails before the fix).

A direct AutoTP caller that leaves training_mode=None and switches modes later with set_autotp_mode expects its layers to follow the switch, as they did before. Resolving None to the current mode before building tp_meta pinned every layer to the mode at construction. Keep None in tp_meta and resolve it only for construction-time routing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

jinyouzhi and others added 2 commits October 17, 2026 15:39

jinyouzhi mentioned this pull request Oct 17, 2026

[REQUEST] Add an OPSD training batch builder for teacher-prefix assembly and response-logit alignment #8850

Open

pengdurice mentioned this pull request Oct 19, 2026

AutoTP: sequence sharding for training #8860

Draft

jinyouzhi mentioned this pull request Oct 20, 2026

AutoTP: shard an untied input embedding along the vocabulary (`vocab_… #8859

Draft

jinyouzhi marked this pull request as ready for review October 21, 2026 17:06

jinyouzhi requested review from hwchen2017, loadams, tjruwase and tohtana as code owners October 21, 2026 17:06

chatgpt-codex-connector Bot reviewed Oct 21, 2026

View reviewed changes

Navigation Menu

Repository navigation

[AutoTP] Support frozen AutoTP models and manage AutoTP mode per model - #8851

[AutoTP] Support frozen AutoTP models and manage AutoTP mode per model#8851
jinyouzhi wants to merge 5 commits into
deepspeedai:masterfrom
jinyouzhi:autotp-frozen-replicated-hook

jinyouzhi commented Oct 17, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot commented Oct 21, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Oct 21, 2026

Uh oh!

jinyouzhi Oct 23, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		if training_mode is None:
		training_mode = is_autotp_training_mode()

Conversation

jinyouzhi commented Oct 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Related work

Tests

Compatibility matrix

AI assistance

Uh oh!

chatgpt-codex-connector Bot commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

jinyouzhi Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

jinyouzhi commented Oct 17, 2026 •

edited

Loading

chatgpt-codex-connector Bot commented Oct 21, 2026 •

edited

Loading