Emit affine maps from AutoTP layers

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
+    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,7 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
     vocabulary_patterns = []
     parameter_with_sub_params = []
     unsupported_parameter_patterns = {}
+    affine_maps = {}
     original_vocab_size = None
 
     # Tied parameters are reachable under several module attributes, but the optimizer -- and
@@ -658,6 +712,10 @@ def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]
                 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 +761,13 @@ 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 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
 
 

tests/unit/checkpoint/test_autotp_uc_checkpoint.py

@@ -1679,32 +1679,16 @@ class affine_resume_checkpoint(DistributedFixture):
     world_size = 2
 
     def run(self, tmpdir, affine_layout):
-        from deepspeed.checkpoint.affine import (AFFINE_MAP_FORMAT_VERSION, contiguous_split_map, replicated_map)
-        from deepspeed.checkpoint.constants import AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION
+        from deepspeed.module_inject.layers import collect_autotp_universal_checkpoint_info
 
         engine = _affine_resume_engine(self.world_size)
         for step in range(4):
             _affine_resume_step(engine, self.world_size, step)
 
-        # Phase 1 does not yet emit AutoTP affine metadata. Supply the known fixture
-        # layout through the checkpoint API; all weights and moments come from training.
-        # LinearAllreduce adds this bias AFTER reduction, so its real scale is one.
-        maps = {
-            r"^fc1\.weight$": contiguous_split_map((16, 16), [8, 8], 0),
-            r"^fc1\.bias$": contiguous_split_map((16, ), [8, 8], 0),
-            r"^fc2\.weight$": contiguous_split_map((16, 16), [8, 8], 1),
-            r"^fc2\.bias$": replicated_map((16, ), self.world_size),
-        }
-        uc_info = {
-            UNIVERSAL_CHECKPOINT_VERSION_KEY: UNIVERSAL_CHECKPOINT_VERSION_VALUE,
-            AFFINE_MAP: {
-                AFFINE_MAP_VERSION: AFFINE_MAP_FORMAT_VERSION,
-                AFFINE_MAP_PARAMS: {
-                    name: layout.to_dict()
-                    for name, layout in maps.items()
-                },
-            },
-        }
+        # The affine layout comes from the producer rather than from the test, so this
+        # exercises the metadata a real job would write. All weights and moments come
+        # from training. LinearAllreduce adds fc2's bias AFTER reduction, so its scale is one.
+        uc_info = collect_autotp_universal_checkpoint_info(engine.module)
         engine.save_checkpoint(tmpdir,
                                tag="affine_resume",
                                client_state={UNIVERSAL_CHECKPOINT_INFO: uc_info} if affine_layout else {})
@@ -1745,3 +1729,37 @@ def test_resume_matches_uninterrupted_training(self, affine_resume_checkpoint, t
                                            rtol=2e-5,
                                            msg=lambda message: f"snapshot {index}, {name}: {message}")
         engine.destroy()
+
+
+class TestAffineMapProducer(DistributedTest):
+    """The producer must emit the layout that `affine_resume_checkpoint` supplies by hand.
+
+    That fixture's maps are not a guess: a full train -> save -> convert -> resume cycle
+    reproduces uninterrupted training through them. Requiring the producer to match them
+    exactly is what makes emitted metadata trustworthy without re-running the whole cycle.
+    """
+
+    world_size = 2
+
+    def test_producer_matches_the_verified_fixture_layout(self):
+        from deepspeed.checkpoint.affine import contiguous_split_map, replicated_map
+        from deepspeed.checkpoint.constants import AFFINE_MAP, AFFINE_MAP_PARAMS
+        from deepspeed.module_inject.layers import collect_autotp_universal_checkpoint_info
+
+        engine = _affine_resume_engine(self.world_size)
+        emitted = collect_autotp_universal_checkpoint_info(engine.module)
+        maps = emitted.get(AFFINE_MAP, {}).get(AFFINE_MAP_PARAMS, {})
+
+        expected = {
+            r"^fc1\.weight$": contiguous_split_map((16, 16), [8, 8], 0).to_dict(),
+            r"^fc1\.bias$": contiguous_split_map((16, ), [8, 8], 0).to_dict(),
+            r"^fc2\.weight$": contiguous_split_map((16, 16), [8, 8], 1).to_dict(),
+            r"^fc2\.bias$": replicated_map((16, ), self.world_size).to_dict(),
+        }
+
+        assert set(maps) == set(expected), (f"producer emitted maps for {sorted(maps)}, expected exactly "
+                                            f"{sorted(expected)}")
+        for pattern, want in expected.items():
+            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()