[AutoTP] Support frozen AutoTP models and manage AutoTP mode per model

deepspeed/module_inject/auto_tp.py

@@ -221,17 +221,20 @@ def __init__(self,
                  vocab_parallel_lm_head=False,
                  model_config=None,
                  tp_grain_size: int = 1,
-                 training_mode: bool = False):
+                 training_mode: Optional[bool] = None):
         self.module = module
         self.all_reduce_linears = all_reduce_linears
         self.prefix = prefix
         self.state_dict = state_dict
 
         self.mp_size = None
         self.mp_group = None
+        # None keeps the legacy behavior of following the process-wide AutoTP mode.
+        if training_mode is None:
+            training_mode = is_autotp_training_mode()
         # Per-model TP metadata threaded through every layer / helper so each AutoTP instance
-        # shards by its own model's kv-head / grain values.
-        self.tp_meta = AutoTPMeta.from_model_config(model_config, tp_grain_size)
+        # shards by its own model's kv-head / grain values and keeps its own training mode.
+        self.tp_meta = AutoTPMeta.from_model_config(model_config, tp_grain_size, training=training_mode)
         self.linear_layer_setting = linear_layer_setting
         self.orig_layer_impl = orig_layer_impl
         self.linear_policies = None
@@ -394,7 +397,7 @@ def tp_parser(model):
 
     def set_tensor_parallel_config(self, mp_size, mp_group):
 
-        if is_autotp_training_mode():
+        if self.training_mode:
             self.mp_group = groups.get_tensor_model_parallel_group()
             self.mp_size = groups.get_tensor_model_parallel_world_size()
             return
@@ -937,9 +940,12 @@ def register_replicated_grad_hooks(self, model):
         first optimizer step. The hook runs when the gradient is computed, before accumulation,
         so it holds under gradient accumulation and under the DeepCompile pass, whose compiled
         backward still accumulates leaf gradients through the autograd engine.
+
+        Returns the ``(name, param)`` pairs skipped because they were frozen, so the engine can
+        reject them if they are unfrozen later and would train without the reduction.
         """
         if self.partition_config is None or self.mp_group is None or self.mp_size <= 1:
-            return
+            return []
 
         def make_grad_allreduce_hook(group):
 
@@ -952,10 +958,16 @@ def hook(grad):
 
         model_type = self._get_model_type()
         registered = []
+        frozen = []
         for param_name, param in model.named_parameters():
             spec = self.partition_config.find_matching_spec(param_name, model_type)
             if spec is None or not spec.grad_allreduce:
                 continue
+            # A frozen parameter, such as one of a distillation teacher, has no gradient to reduce,
+            # and PyTorch refuses hooks on tensors that do not require gradients.
+            if not param.requires_grad:
+                frozen.append((param_name, param))
+                continue
             if getattr(param, "_ds_grad_allreduce_registered", False):
                 continue
             param.register_hook(make_grad_allreduce_hook(self.mp_group))
@@ -966,6 +978,7 @@ def hook(grad):
                 f"AutoTP: registered tensor-parallel grad all-reduce for {len(registered)} replicated "
                 f"parameters, e.g. {registered[0]!r}",
                 ranks=[0])
+        return frozen
 
     def update_mp_params(self, child, name=None):
         if getattr(child, "replaced", False) == True:

deepspeed/module_inject/layers.py

@@ -188,10 +188,12 @@ def set_autotp_mode(training=False):
         DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE
 
 
-def add_bias(input, bias):
+def add_bias(input, bias, training=None):
     if bias is None:
         return input
-    if is_autotp_training_mode():
+    if training is None:
+        training = is_autotp_training_mode()
+    if training:
         # Training mode - avoid inplace to ensure correct autograd
         input = input + bias
         return input
@@ -565,8 +567,12 @@ def _should_materialize_tp_partition(self):
         return self.mp_group is not None
 
     def is_training_mode(self):
-        global DEEPSPEED_AUTOTP_MODE
-        return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING
+        # A model keeps the mode it was sharded in, so a frozen inference teacher and a training
+        # student can live in one process without sharing the process-wide AutoTP mode. Layers
+        # built without an explicit mode keep following set_autotp_mode.
+        if self.tp_meta.training is None:
+            return is_autotp_training_mode()
+        return self.tp_meta.training
 
     def _freeze_partition_sizes(self, total_size):
         """Resolve the tensor parallel split of this layer once, while the layer is built.
@@ -896,7 +902,7 @@ def forward(self, input):
         if not self.defer_collectives_to_compiler:
             output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
         if self.bias is not None:
-            output = add_bias(output, self.bias)
+            output = add_bias(output, self.bias, self.is_training_mode())
         return output
 
     @torch.no_grad()
@@ -992,7 +998,7 @@ def forward(self, input):
                 input = ColumnParallel.apply(self.mp_group, input)
             output = torch.matmul(input, self.weight.transpose(-1, -2))
             if self.bias is not None:
-                output = add_bias(output, self.bias)
+                output = add_bias(output, self.bias, self.is_training_mode())
         else:
             output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias)
 
@@ -1737,7 +1743,7 @@ def forward(self, input):
         if self.mp_group is not None:
             dist.inference_all_reduce(output, group=self.mp_group)
         if self.bias is not None:
-            output = add_bias(output, self.bias)
+            output = add_bias(output, self.bias, self.is_training_mode())
         return output
 
 
@@ -2165,7 +2171,7 @@ def forward(self, input):
             input = ColumnParallel.apply(self.mp_group, input)
         output = torch.matmul(input, self.weight.transpose(-1, -2))
         if self.bias is not None:
-            output = add_bias(output, self.bias)
+            output = add_bias(output, self.bias, self.is_training_mode())
         return output
 
     @torch.no_grad()
@@ -2294,7 +2300,7 @@ def forward(self, input):
         if not self.defer_collectives_to_compiler:
             output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
         if self.bias is not None:
-            output = add_bias(output, self.bias)
+            output = add_bias(output, self.bias, self.is_training_mode())
         return output
 
     @torch.no_grad()

deepspeed/module_inject/replace_module.py

@@ -15,7 +15,6 @@
 from .replace_policy import replace_policies, generic_policies
 from .auto_tp import AutoTP, ReplaceWithTensorSlicing, Loading
 from .layers import TensorParallelOcShardConv2d, TensorParallelIcShardConv2d
-from deepspeed.module_inject.layers import is_autotp_training_mode
 from deepspeed import comm as dist
 
 from .load_checkpoint import load_model_with_checkpoint
@@ -330,7 +329,7 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
 
     def replace_fn(child, _policy, layer_id=0, prefix="", state_dict=None):
         # copy relevant state from child -> new module
-        if not is_autotp_training_mode() and config.replace_with_kernel_inject:
+        if not training_mode and config.replace_with_kernel_inject:
             new_module = replace_with_policy(child,
                                              _policy,
                                              config.triangular_masking,
@@ -342,7 +341,7 @@ def replace_fn(child, _policy, layer_id=0, prefix="", state_dict=None):
         return new_module
 
     def set_lm_head(module):
-        if is_autotp_training_mode():
+        if training_mode:
             # we need to handle autoTP training mode separately.
             return
 
@@ -486,7 +485,7 @@ def conv2d_parallel_shard_weights(model, rank, world_size):
         set_lm_head(replaced_module)
         print(f"checkpoint loading time at rank {rank}: {time.time()-start_time} sec")
 
-    if not is_autotp_training_mode() and config.save_mp_checkpoint_path is not None:
+    if not training_mode and config.save_mp_checkpoint_path is not None:
         from collections import OrderedDict
         import json
         num_partitions = 8

deepspeed/module_inject/tp_shard.py

@@ -58,16 +58,19 @@ class AutoTPMeta:
     num_attention_heads: Optional[int] = None
     n_embd: Optional[int] = None
     tp_grain_size: int = 1
+    # Whether the model was sharded for training. None means unknown (a layer built outside
+    # AutoTP), in which case the layer falls back to the process-wide AutoTP mode.
+    training: Optional[bool] = None
 
     @classmethod
-    def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta":
+    def from_model_config(cls, model_config, tp_grain_size: int = 1, training: Optional[bool] = None) -> "AutoTPMeta":
         """The single source of truth for reading kv-head / hidden / attention-head counts.
 
         ``model_config`` may be ``None`` for callers that build a synthetic AutoTP with no real
         model (some unit tests); they get an empty meta and the resulting even-grain split.
         """
         if model_config is None:
-            return cls(tp_grain_size=tp_grain_size)
+            return cls(tp_grain_size=tp_grain_size, training=training)
         # Multimodal configs (e.g. vision-language models) keep the head counts only under
         # text_config while the outer config carries the modality heads, so descend to the
         # text config when it exists.
@@ -85,7 +88,8 @@ def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta"
         return cls(num_kv_heads=num_kv_heads,
                    num_attention_heads=num_attention_heads,
                    n_embd=n_embd,
-                   tp_grain_size=tp_grain_size)
+                   tp_grain_size=tp_grain_size,
+                   training=training)
 
 
 def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None):

deepspeed/runtime/engine.py

@@ -1056,7 +1056,7 @@ def _apply_autotp_partitioning(self, model, tp_config):
 
         def finalize_autotp(autotp=None, attach_uc_metadata=False, require_vocab_parallel_lm_head=False):
             if autotp is not None:
-                autotp.register_replicated_grad_hooks(model)
+                self._autotp_frozen_replicated_params = autotp.register_replicated_grad_hooks(model)
 
             from deepspeed.module_inject.layers import VocabParallelLinear
             vocab_parallel_heads = [module for module in model.modules() if isinstance(module, VocabParallelLinear)]
@@ -3119,10 +3119,22 @@ def _reduce_autoep_folding_tp_replicated_gradients(self):
                 continue
             reduce_autoep_folding_gradient(folding_spec, param, param.grad, tp_group=tp_group, param_name=param_name)
 
+    def _check_autotp_frozen_replicated_params(self):
+        # A replicated parameter frozen at initialize got no tensor-parallel gradient all-reduce, so
+        # training it now would let the ranks drift apart silently instead of failing.
+        unfrozen = [
+            name for name, param in getattr(self, "_autotp_frozen_replicated_params", ()) if param.requires_grad
+        ]
+        if unfrozen:
+            raise RuntimeError(f"AutoTP replicated parameters {unfrozen} were frozen at deepspeed.initialize and "
+                               "have no tensor-parallel gradient all-reduce. Set requires_grad before "
+                               "deepspeed.initialize instead of unfreezing them afterwards.")
+
     def _backward_prologue(self):
         if is_functorch_transforming():
             return
         self._start_timers(self.engine_timers.backward_timers)
+        self._check_autotp_frozen_replicated_params()
 
         # When necessary internal APIs are not available, we disable direct calls to tensor.backward()
         # and limit to engine.backward(loss) only.