Emit affine maps from AutoTP layers

deepspeed/checkpoint/affine_ir_spec.md

@@ -399,8 +399,30 @@ whose name matched a metadata key could never be confused for one.
   and for debugging a checkpoint that will not load.
 - **`locations` is per piece, not per parameter.** This is the field that expresses
   `bigcodetype`, whose kv block is held identically by every rank (§5).
-
-### 6.4 Size
+- **A reader that prefers the map must still account for the categories.** Additivity means
+  both are in the file, so a reader taking the map never consults the category branches —
+  which leaves their patterns looking unused, and a strict conversion rejects them. They are
+  superseded, not unused, and the reader has to say so. This is not cosmetic: it is the
+  difference between a checkpoint that converts and one that aborts.
+
+### 6.4 Where a map is built
+
+Not where the file is assembled. Collecting model-level metadata sees only the conversion
+schema, which does not carry per-rank extents — those are resolved while a layer is built
+and are not recoverable later from a shape alone. So a map is derived at the point a layer
+records its metadata, and collection gathers what the layers produced.
+
+Two consequences worth stating, because both are easy to get wrong:
+
+- A parameter the tensor-parallel machinery never touches is still describable — one piece
+  held by every rank — but it reaches no layer, so nothing derives its map. It needs the
+  degree of the group, which only the partitioned layers know. Producing a map for it means
+  taking that from elsewhere in the model rather than skipping it: a parameter with no map
+  falls back to its name category, which is what the IR exists to replace.
+- A layout the machinery refuses to describe produces no map, and conversion falls back to
+  the categories. That is the intended escape, not a gap.
+
+### 6.5 Size
 
 Measured on a fused-QKV parameter at TP=8: 96 pieces, 6836 bytes of compact JSON —
 about **71 bytes per piece**. Extrapolated by tensor count:

deepspeed/checkpoint/ds_to_universal.py

@@ -346,7 +346,22 @@ def get_matched_affine_map(name_):
         assert len(matched_) <= 1, f'Got more than one matching affine map patterns={matched_} for {name_}'
         return ParamAffineMap.from_dict(affine_params[matched_[0]]) if matched_ else None
 
+    def consume_superseded_patterns(name_):
+        """Mark the category patterns for this parameter as used.
+
+        A checkpoint carrying an affine map also carries the category patterns, so an older
+        converter can still read it. This converter prefers the map and never consults those
+        branches, which would otherwise leave their patterns looking unused and fail the
+        strict check. They are superseded here, not unused.
+        """
+        for patterns_ in (replicated_parameters, parameters_to_average, parameters_with_row_parallelism,
+                          vocabulary_parameters, parameters_with_2_sub_params_cat_dim_0):
+            get_matched_pattern(patterns_, name_)
+        get_matched_sub_params_pattern(name_)
+
     matched_affine_map = get_matched_affine_map(name)
+    if matched_affine_map is not None:
+        consume_superseded_patterns(name)
     matched_sub_params_shape, matched_sub_params_pattern = get_matched_sub_params_pattern(name)
 
     step_merged = _merge_zero_shards(slice_base_path, "step", tp_degree, per_tp_shapes)

deepspeed/module_inject/layers.py

@@ -39,6 +39,41 @@ def _normalize_uc_shape(value):
     return tuple(value) if value is not None else None
 
 
+def _derive_affine_map(*, tp_world_size, logical_shape, partition_dim, partition_sizes, sub_param_shard_widths,
+                       replicated, unsupported_reason):
+    """Describe this parameter's layout geometrically, from what the layer already knows.
+
+    The layer is the only place the per-rank extents exist: `_freeze_partition_sizes` resolves
+    them while the layer is built, and they are not recoverable later from a shape alone. So a
+    map is derived here rather than where the model-level metadata is collected.
+
+    Returns None when the layout is not describable yet, in which case conversion falls back to
+    the pattern categories.
+    """
+    from deepspeed.checkpoint.affine import replicated_map, contiguous_split_map, sub_param_map
+
+    if unsupported_reason or not logical_shape or not tp_world_size:
+        return None
+
+    if replicated:
+        return replicated_map(logical_shape, tp_world_size)
+
+    if partition_dim is None:
+        return None
+
+    if sub_param_shard_widths:
+        widths = [list(w) for w in sub_param_shard_widths]
+        return sub_param_map(shape=logical_shape,
+                             sub_dim_sizes=[sum(w) for w in widths],
+                             shard_widths=widths,
+                             partition_dim=partition_dim)
+
+    if partition_sizes:
+        return contiguous_split_map(logical_shape, list(partition_sizes), partition_dim)
+
+    return None
+
+
 def _build_param_uc_conversion_meta(*,
                                     partition_type,
                                     partition_dim=None,
@@ -47,13 +82,14 @@ def _build_param_uc_conversion_meta(*,
                                     original_shape=None,
                                     is_bias=False,
                                     replicated=False,
+                                    affine_map=None,
                                     unsupported_reason=None):
     """Build the conversion-facing subset of parameter UC metadata.
 
     This is the only schema that should flow into model-level
     `UNIVERSAL_CHECKPOINT_INFO` via `collect_autotp_universal_checkpoint_info()`.
     """
-    return {
+    meta = {
         'partition_type': partition_type,
         'partition_dim': partition_dim,
         'sub_param_shape': _normalize_uc_shape(sub_param_shape),
@@ -63,6 +99,11 @@ def _build_param_uc_conversion_meta(*,
         'replicated': replicated,
         'unsupported_reason': unsupported_reason,
     }
+    if affine_map is not None:
+        # Only present for a layout that can be described, so the schema an existing
+        # layer publishes is unchanged. Stored as plain scalars like every other field.
+        meta['affine_map'] = affine_map.to_dict()
+    return meta
 
 
 def _build_param_uc_restore_meta(*,
@@ -78,6 +119,7 @@ def _build_param_uc_restore_meta(*,
                                  original_shape=None,
                                  is_bias=False,
                                  replicated=False,
+                                 affine_map=None,
                                  unsupported_reason=None):
     """Build the restore-facing parameter UC metadata.
 
@@ -119,6 +161,7 @@ def _build_param_uc_restore_meta(*,
                                         original_shape=original_shape,
                                         is_bias=is_bias,
                                         replicated=replicated,
+                                        affine_map=affine_map,
                                         unsupported_reason=unsupported_reason),
     }
 
@@ -462,6 +505,13 @@ def _set_param_uc_meta(self,
                            unsupported_reason=None):
         if param is None:
             return
+        affine_map = _derive_affine_map(tp_world_size=getattr(self, 'tp_world_size', None),
+                                        logical_shape=logical_shape or original_shape,
+                                        partition_dim=partition_dim,
+                                        partition_sizes=partition_sizes,
+                                        sub_param_shard_widths=sub_param_shard_widths,
+                                        replicated=replicated,
+                                        unsupported_reason=unsupported_reason)
         setattr(
             param, DS_AUTOTP_UC_META,
             _build_param_uc_restore_meta(partition_type=partition_type,
@@ -476,6 +526,7 @@ def _set_param_uc_meta(self,
                                          original_shape=original_shape,
                                          is_bias=is_bias,
                                          replicated=replicated,
+                                         affine_map=affine_map,
                                          unsupported_reason=unsupported_reason))
 
     def _mark_uc_metadata(self):
@@ -614,7 +665,9 @@ 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.constants import (AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS, ORIGINAL_VOCAB_SIZE,
+    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,
                                                 SUB_PARAM_SHARD_WIDTHS, TP_REPLICATED_PARAMETER_PATTERNS,
                                                 UNIVERSAL_CHECKPOINT_VERSION_KEY, UNIVERSAL_CHECKPOINT_VERSION_VALUE,
@@ -626,6 +679,9 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
     vocabulary_patterns = []
     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
@@ -637,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
@@ -650,14 +708,23 @@ 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')
             if unsupported_reason:
                 unsupported_parameter_patterns[pattern] = unsupported_reason
                 continue
 
+            affine_map = conversion_meta.get('affine_map')
+            if affine_map is not None:
+                affine_maps[pattern] = affine_map
+
             if conversion_meta.get('replicated'):
                 replicated_patterns.append(pattern)
 
@@ -703,6 +770,17 @@ 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.
+        uc_info[AFFINE_MAP] = {
+            AFFINE_MAP_VERSION: AFFINE_MAP_FORMAT_VERSION,
+            AFFINE_MAP_PARAMS: affine_maps,
+        }
     return uc_info