DeepSpeed

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

[AutoTP] Support frozen AutoTP models and manage AutoTP mode per model (#8851) ## 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. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>