Emit a map for parameters AutoTP leaves untouched

An untouched parameter is identical on every rank, so it is describable as one
replicated piece -- but building the map needs the tp degree, which only the
partitioned layers carry. Take it from one of those and fill in the rest after
the walk.

Add the coverage invariant: every parameter placed by a name category must also
carry a map, or conversion still depends on the category. Only the layouts
AutoTP refuses to describe are exempt.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Author: Achyuthan-S

deepspeed/module_inject/layers.py

@@ -665,7 +665,7 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
     restore-time per-parameter details such as `sub_param_sizes` or
     `target_partition_shape`, which stay on the parameter metadata object.
     """
-    from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION
+    from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, replicated_map
     from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION,
                                                 AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS, ORIGINAL_VOCAB_SIZE,
                                                 PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_SUB_PARAMS,
@@ -680,6 +680,8 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
     parameter_with_sub_params = []
     unsupported_parameter_patterns = {}
     affine_maps = {}
+    untouched_shapes = {}
+    tp_world_size = None
     original_vocab_size = None
 
     # Tied parameters are reachable under several module attributes, but the optimizer -- and
@@ -691,6 +693,8 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
         marker = getattr(module, "_mark_uc_metadata", None)
         if marker is not None:
             marker()
+        if tp_world_size is None:
+            tp_world_size = getattr(module, 'tp_world_size', None)
 
         for param_name, param in module.named_parameters(recurse=False):
             full_name = f"{module_name}.{param_name}" if module_name else param_name
@@ -704,7 +708,12 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
                 # ranks. Classify it as TP-replicated; otherwise it falls through to
                 # the converter's default dim-0 concat and is wrongly expanded (e.g.
                 # LayerNorm/RMSNorm weights [H] -> [H * tp_degree]).
+                #
+                # Such a parameter is describable -- one piece held by every rank -- but the
+                # map needs the tp degree, which only the partitioned layers carry. Record
+                # the shape and build the map once the loop has seen one of them.
                 replicated_patterns.append(pattern)
+                untouched_shapes[pattern] = tuple(param.shape)
                 continue
 
             unsupported_reason = conversion_meta.get('unsupported_reason')
@@ -761,6 +770,10 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
         uc_info[SUB_PARAM_SHARD_WIDTHS] = sub_param_shard_widths
     if original_vocab_size is not None:
         uc_info[ORIGINAL_VOCAB_SIZE] = original_vocab_size
+    if tp_world_size:
+        for pattern, shape in untouched_shapes.items():
+            affine_maps[pattern] = replicated_map(shape, tp_world_size).to_dict()
+
     if affine_maps:
         # Published alongside the pattern lists rather than instead of them, so a converter
         # that predates the map simply does not see the key and takes the categories.

tests/unit/checkpoint/test_autotp_uc_checkpoint.py

@@ -1763,3 +1763,87 @@ def test_producer_matches_the_verified_fixture_layout(self):
             assert maps[pattern] == want, (f"emitted map for {pattern} differs from the layout the resume "
                                            f"fixture verifies:\n  emitted  {maps[pattern]}\n  expected {want}")
         engine.destroy()
+
+
+class AffineCoverageModel(torch.nn.Module):
+    """Covers each conversion category: column split, row split, replicated, untouched."""
+
+    def __init__(self, hidden_dim=16, vocab_size=26):
+        super().__init__()
+        self.embed = torch.nn.Embedding(vocab_size, hidden_dim)  # AutoTP leaves this alone
+        self.norm = torch.nn.LayerNorm(hidden_dim)  # untouched, and not 2-D
+        self.fc1 = torch.nn.Linear(hidden_dim, hidden_dim)  # column
+        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)  # row
+        self.lm_head = torch.nn.Linear(hidden_dim, vocab_size)  # vocabulary
+
+    def forward(self, x):
+        h = self.norm(self.embed(x))
+        return self.lm_head(self.fc2(self.fc1(h))).sum()
+
+
+class TestAffineMapCoverage(DistributedTest):
+    """Every parameter the converter can place must carry an affine map.
+
+    A parameter with no map falls back to its name category, which is the behaviour the IR
+    exists to replace. The only parameters allowed to have no map are the ones AutoTP itself
+    refuses to describe, which conversion rejects anyway.
+    """
+
+    world_size = 2
+
+    def test_every_convertible_parameter_has_a_map(self):
+        from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS,
+                                                    AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS,
+                                                    PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_SUB_PARAMS,
+                                                    TP_REPLICATED_PARAMETER_PATTERNS, VOCABULARY_PARAMETER_PATTERNS)
+        from deepspeed.module_inject.layers import collect_autotp_universal_checkpoint_info
+
+        config = {
+            "train_micro_batch_size_per_gpu": 1,
+            "zero_allow_untested_optimizer": True,
+            "zero_optimization": {
+                "stage": 1
+            },
+            "tensor_parallel": {
+                "autotp_size": self.world_size,
+                "partition_config": {
+                    "use_default_specs":
+                    False,
+                    "layer_specs": [
+                        {
+                            "patterns": [r".*fc1\.weight$"],
+                            "partition_type": "column"
+                        },
+                        {
+                            "patterns": [r".*fc2\.weight$"],
+                            "partition_type": "row"
+                        },
+                        {
+                            "patterns": [r".*lm_head\.weight$"],
+                            "partition_type": "column",
+                            "gather_output": True
+                        },
+                    ],
+                },
+            },
+        }
+        model = AffineCoverageModel()
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+        engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config)
+        info = collect_autotp_universal_checkpoint_info(engine.module)
+
+        mapped = set(info.get(AFFINE_MAP, {}).get(AFFINE_MAP_PARAMS, {}))
+        unsupported = set(info.get(AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS, {}))
+
+        categorised = set()
+        for key in (TP_REPLICATED_PARAMETER_PATTERNS, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
+                    VOCABULARY_PARAMETER_PATTERNS):
+            categorised.update(info.get(key, []))
+        for entry in info.get(PARAMETER_WITH_SUB_PARAMS, []):
+            categorised.update(entry["patterns"])
+
+        assert categorised, "model exercised no conversion category, so this proves nothing"
+        missing = categorised - mapped - unsupported
+        assert not missing, (f"these parameters are placed by a name category but carry no affine map, "
+                             f"so conversion still depends on the category: {sorted(missing)}")
+        engine.destroy()