AutoEP + AutoTP folding: Universal Checkpointing under ZeRO-3

deepspeed/checkpoint/ds_to_universal.py

@@ -605,6 +605,10 @@ def _merge_tp_slice_files(args, uc_info, tp_degree, slice_shapes, temp_dir, excl
     # if a pattern was not used by any of the workers, then it was not used at all -> assert/alert
     sets = [set(lst) for lst in unmatched_patterns_lists]
     unmatched_patterns = list(set.intersection(*sets))
+    # A pattern for a parameter merged elsewhere (AutoEP experts) has nothing to match here.
+    unmatched_patterns = [
+        pattern for pattern in unmatched_patterns if not any(re.match(pattern, name) for name in exclude_param_names)
+    ]
     if args.strict:
         assert not unmatched_patterns, f'Unused patterns={unmatched_patterns} while merging tp slices'
     elif unmatched_patterns:
@@ -630,13 +634,6 @@ def _shape_numel(shape):
     return math.prod(shape)
 
 
-def _zero3_rank_from_file(path):
-    match = re.search(r'(?:bf16_)?zero_pp_rank_([0-9]+)_mp_rank_', os.path.basename(path))
-    if match is None:
-        raise ValueError(f"Cannot parse ZeRO rank from checkpoint file name: {path}")
-    return int(match.group(1))
-
-
 def _zero3_tp_dp_ranks_from_file(path):
     # A ZeRO-3 checkpoint file is named zero_pp_rank_{dp_rank}_mp_rank_{tp_rank}_*.
     # Under AutoTP, the mp_rank enumerates tensor-parallel ranks and the pp_rank
@@ -709,7 +706,7 @@ def _autoep_expert_param_names_by_rank(model_files):
     expert_param_names = set()
     metadata_by_rank = {}
     for model_file in model_files:
-        rank = _zero3_rank_from_file(model_file)
+        rank = _zero3_tp_dp_ranks_from_file(model_file)
         model_state = torch.load(model_file, map_location=torch.device('cpu'), weights_only=False)
         autoep_metadata = _get_autoep_metadata(model_state)
         if autoep_metadata is not None:
@@ -720,9 +717,11 @@ def _autoep_expert_param_names_by_rank(model_files):
 
 
 def _rank_map_from_files(files, description):
+    # Keyed by (tp_rank, dp_rank): under AutoEP + AutoTP folding the files of different
+    # tensor-parallel ranks share a data-parallel rank.
     rank_map = {}
     for path in files:
-        rank = _zero3_rank_from_file(path)
+        rank = _zero3_tp_dp_ranks_from_file(path)
         if rank in rank_map:
             raise RuntimeError(f"Duplicate ZeRO rank {rank} in {description} files: "
                                f"{rank_map[rank]} and {path}")
@@ -836,6 +835,7 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files
         for rank, model_state in model_states_by_rank.items() if _get_autoep_metadata(model_state) is not None
     }
     _validate_autoep_expert_shapes(model_states_by_rank, metadata_by_rank)
+    dp_degree = len({dp_rank for _, dp_rank in model_states_by_rank})
 
     expert_fragments = {}
     num_experts_by_param = {}
@@ -867,8 +867,8 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files
                 'exp_avg_sq': optimizer_sub_state.get('exp_avg_sq'),
             }
             partition_metadata = partition_groups[sub_group_id] if sub_group_id < len(partition_groups) else {}
-            partition_count = partition_metadata.get('partition_count', len(model_states_by_rank))
-            partition_rank = partition_metadata.get('partition_rank', rank)
+            partition_count = partition_metadata.get('partition_count', dp_degree)
+            partition_rank = partition_metadata.get('partition_rank', rank[1])
 
             offset = 0
             for param_name, shape in sub_group_shape.items():
@@ -1286,12 +1286,6 @@ def main(args):
         if not has_zero3_partitioned_autoep:
             autoep_expert_param_names = set()
         else:
-            # The AutoEP files are matched by data-parallel rank alone, so the files of different
-            # tensor-parallel ranks under AutoEP + AutoTP folding cannot be told apart.
-            _, autoep_tp_degree, _ = _build_zero3_rank_grid(model_files)
-            if autoep_tp_degree > 1:
-                raise NotImplementedError("Universal checkpoint conversion of an AutoEP + AutoTP folded ZeRO-3 "
-                                          "checkpoint is not supported yet. Load it with the same topology.")
             _validate_zero3_model_optim_rank_sets(model_files, optim_files)
         param_shapes = _parse_model_states_stage3(model_files)
         # Recover the tensor-parallel / data-parallel grid from the file names. With
@@ -1337,7 +1331,12 @@ def main(args):
                 dict((k, v) for sub in grid_tp for k, v in sub.items()) for grid_tp in param_shapes_grid
             ]
             slice_shapes = {name: [param_shapes_by_tp[tp][name] for tp in range(tp_degree)] for name in param_keys}
-            _merge_tp_slice_files(args, uc_info, tp_degree, slice_shapes, temp_dir)
+            _merge_tp_slice_files(args,
+                                  uc_info,
+                                  tp_degree,
+                                  slice_shapes,
+                                  temp_dir,
+                                  exclude_param_names=autoep_expert_param_names)
         else:
             # Plain ZeRO-3 (no tensor parallelism): DP-only merge, unchanged behavior.
             _merge_zero3_slice_files(args, param_keys, dp_degree, temp_dir)

deepspeed/runtime/engine.py

@@ -4602,13 +4602,16 @@ def _load_checkpoint(self,
         folding_spec = getattr(self, "_autoep_folding_spec", None)
         folded_autoep_tp = folding_spec is not None and folding_spec.tp_size > 1
         ep_group_name = f"ep_size_{folding_spec.ep_size}" if folded_autoep_tp else None
-        DeepSpeedEngine._validate_autoep_folding_checkpoint_metadata(
-            checkpoint,
-            folding_spec=folding_spec,
-            family="dense",
-            zero_partition_group="dense_dp",
-            zero_partition_count=folding_spec.dp_size if folded_autoep_tp else None,
-            tp_rank=groups.get_tensor_model_parallel_rank() if folded_autoep_tp else None)
+        # A universal checkpoint is re-sliced from zero/ for this runtime's layout, so the layout the
+        # source run recorded in its model states need not match.
+        if not self.load_universal_checkpoint():
+            DeepSpeedEngine._validate_autoep_folding_checkpoint_metadata(
+                checkpoint,
+                folding_spec=folding_spec,
+                family="dense",
+                zero_partition_group="dense_dp",
+                zero_partition_count=folding_spec.dp_size if folded_autoep_tp else None,
+                tp_rank=groups.get_tensor_model_parallel_rank() if folded_autoep_tp else None)
 
         fetch_z3_params = False
         z3_params_to_fetch = None

docs/_pages/config-json.md

@@ -1144,7 +1144,7 @@ Use a built-in preset but override specific naming/weight fields for a fine-tune
 - `autoep_size` must divide `num_experts` for all detected MoE layers
 - AutoEP combines with AutoTP (`tensor_parallel.autotp_size > 1`) by folding them onto the same ranks, at ZeRO stages 0 to 3, with `expert_tensor_parallel_size` 1 and without ZeRO offload or DeepCompile
 - AutoEP with ZeRO Stage 3 is supported only without sequence parallelism, hpZeRO secondary tensor groups, non-1 `expert_tensor_parallel_size`, or quantized gradients
-- ZeRO Stage 3 saves AutoEP checkpoints partition-natively and supports same-topology save/load, module-only loads, optimizer-state-skipping loads, and universal checkpoint conversion (not yet with AutoTP folding). Universal loads can resume at a different data-parallel world size, a different `autoep_size`, or both (when the target `autoep_size` divides the expert count), including weights-only/module-only loads from the converted `fp32.pt` parameter files
+- ZeRO Stage 3 saves AutoEP checkpoints partition-natively and supports same-topology save/load, module-only loads, optimizer-state-skipping loads, and universal checkpoint conversion, also with AutoTP folding. Universal loads can resume at a different data-parallel world size, a different `autoep_size`, a different `tensor_parallel.autotp_size` (folded or not), or any combination (when the target `autoep_size` divides the expert count), including weights-only/module-only loads from the converted `fp32.pt` parameter files
 
 ### Python cyclic garbage collection
 

docs/code-docs/source/autoep.rst

@@ -64,10 +64,11 @@ including models built under ``deepspeed.zero.Init``. Stage 3 AutoEP checkpoints
 partition-natively in the ``zero_pp_rank_*`` shard files and support
 same-topology load, module-only loads (``load_module_only``),
 optimizer-state-free loads (``load_optimizer_states=False``), and Universal
-Checkpoint conversion, except with AutoTP folding. Optimizer-including
+Checkpoint conversion, also with AutoTP folding. Optimizer-including
 Universal Checkpoint loads can
 resume with a different data-parallel world size, a different ``autoep_size``,
-or both, when the target ``autoep_size`` divides the model's expert count.
+a different ``tensor_parallel.autotp_size`` (folded or not), or any
+combination, when the target ``autoep_size`` divides the model's expert count.
 Weights-only/module-only Universal Checkpoint loads use the converted
 ``fp32.pt`` parameter files and support the same data-parallel and
 ``autoep_size`` topology changes.

tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py

@@ -31,8 +31,8 @@
 from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
 from deepspeed.module_inject.auto_tp import AutoTP
 from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, TensorParallel_Layer, VocabParallelLinear
-from deepspeed.utils import groups, safe_get_full_fp32_param, safe_get_full_grad
-from unit.common import DistributedTest
+from deepspeed.utils import groups, safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
+from unit.common import DistributedFixture, DistributedTest
 
 transformers = pytest.importorskip("transformers")
 
@@ -345,8 +345,7 @@ class TestZero3FoldedCheckpoint(DistributedTest):
 
     Train one step, save, load into an engine built the same way from a different random
     initialization, then train one more step on both: the losses and every parameter must agree.
-    SGD momentum makes that step depend on the restored optimizer state. Universal Checkpoint
-    conversion of such a checkpoint is not supported yet and must say so.
+    SGD momentum makes that step depend on the restored optimizer state.
     """
     world_size = 4
 
@@ -393,18 +392,194 @@ def build_engine(seed):
             if not torch.allclose(got, want, rtol=1e-6, atol=1e-7):
                 differing.append(name)
 
-        conversion_error = None
-        args = SimpleNamespace(input_folder=os.path.join(save_dir, "step1"),
-                               output_folder=os.path.join(save_dir, f"universal_rank{dist.get_rank()}"))
-        try:
-            ds_to_universal.main(args)
-        except NotImplementedError as error:
-            conversion_error = str(error)
-
         torch.testing.assert_close(resumed_losses, losses, rtol=1e-6, atol=1e-6)
         assert not differing, f"{len(differing)} parameters differ after resuming: {differing}"
-        assert conversion_error is not None, "Universal conversion accepted a folded ZeRO-3 checkpoint"
-        assert "folded ZeRO-3 checkpoint is not supported" in conversion_error, conversion_error
+
+
+UNIVERSAL_SOURCE = (2, 2)  # TP, EP of the run that saves, on 4 ranks
+UNIVERSAL_STATES = ("fp32", "exp_avg", "exp_avg_sq")
+
+
+def _universal_config(tp_size, ep_size, precision):
+    config = _folded_config(tp_size, ep_size, zero_stage=3)
+    # Universal Checkpointing converts Adam's states.
+    config["optimizer"] = {"type": "AdamW", "params": {"lr": 1e-3, "weight_decay": 0.01}}
+    if tp_size == 1:
+        del config["tensor_parallel"]
+    if precision == "bf16":
+        config["bf16"] = {"enabled": True}
+    return config
+
+
+def _build_universal_engine(tp_size, ep_size, precision, seed, load_universal=False):
+    config = _universal_config(tp_size, ep_size, precision)
+    if load_universal:
+        config["checkpoint"] = {"load_universal": True}
+    torch.manual_seed(seed)
+    model = transformers.Glm4MoeForCausalLM(_tiny_glm4_moe_config()).float()
+    engine, _, _, _ = deepspeed.initialize(model=model, config=config)
+    return engine
+
+
+def _logical_state(engine, key, tp_size):
+    """Every parameter, or one of its Adam states, at full size under names every layout shares."""
+
+    def full(param):
+        if key == "fp32":
+            return safe_get_full_fp32_param(param)
+        return safe_get_full_optimizer_state(param, key)
+
+    values = {}
+    moe_layers = []
+    for module_name, module in engine.module.named_modules():
+        prefix = f"{module_name}." if module_name else ""
+        if isinstance(module, AutoEPMoELayer):
+            moe_layers.append((prefix, module))
+        for name, param in module.named_parameters(recurse=False):
+            value = full(param)
+            assert value is not None, f"{prefix}{name} has no {key}"
+            if tp_size > 1 and isinstance(module, LinearAllreduce) and name == "weight":
+                value = _gather_tp_shards(value, dim=1)
+            elif tp_size > 1 and isinstance(module, LinearLayer):
+                value = _gather_tp_shards(value, dim=0)
+            values[f"{prefix}{name}"] = value
+    # Each EP rank holds a block of experts; the full tensor stacks them in EP-rank order.
+    for prefix, module in moe_layers:
+        for wname in ("w1", "w2", "w3"):
+            local = values[f"{prefix}experts.{wname}"]
+            shards = [torch.empty_like(local) for _ in range(dist.get_world_size(group=module.ep_group))]
+            dist.all_gather(shards, local.contiguous(), group=module.ep_group)
+            values[f"{prefix}experts.{wname}"] = torch.cat(shards, dim=0)
+    return {name: value.detach().float().cpu() for name, value in values.items()}
+
+
+def _universal_mismatches(got, want, rtol, atol):
+    assert got.keys() == want.keys(), sorted(set(got) ^ set(want))
+    mismatches = []
+    for name, expected in want.items():
+        actual = got[name]
+        if actual.shape != expected.shape:
+            mismatches.append(f"{name}: shape {tuple(actual.shape)}, expected {tuple(expected.shape)}")
+        elif not torch.allclose(actual, expected, rtol=rtol, atol=atol):
+            relative = (actual - expected).abs().max().item() / max(expected.abs().max().item(), 1e-12)
+            mismatches.append(f"{name}: largest relative difference {relative:.3e}")
+    return mismatches
+
+
+class _FoldedUniversalSource(DistributedFixture):
+    """Train a folded TP 2 x EP 2 run one step, save it, convert it, and record what a load must restore.
+
+    The record holds every parameter and Adam state at full size after the step, then the losses of
+    the next accumulation window and the parameters after the next step.
+    """
+    world_size = 4
+
+    def run(self, universal_precision, tmpdir):
+        tp_size, ep_size = UNIVERSAL_SOURCE
+        engine = _build_universal_engine(tp_size, ep_size, universal_precision, seed=1234)
+        logical_rank, logical_world_size = dist.get_rank() // tp_size, dist.get_world_size() // tp_size
+        _run_to_boundary(engine, logical_rank, logical_world_size)
+        engine.step()
+        save_dir = str(tmpdir)
+        engine.save_checkpoint(save_dir, tag="step1")
+
+        dist.barrier()
+        error = None
+        if dist.get_rank() == 0:
+            try:
+                ds_to_universal.main(
+                    SimpleNamespace(input_folder=os.path.join(save_dir, "step1"),
+                                    output_folder=os.path.join(save_dir, "step1_universal"),
+                                    num_extract_workers=1,
+                                    num_merge_workers=1,
+                                    keep_temp_folder=False,
+                                    strict=True,
+                                    inject_missing_state=False))
+            except Exception as exc:  # report on every rank instead of leaving them in the next collective
+                error = f"{type(exc).__name__}: {exc}"
+        errors = [None] * dist.get_world_size()
+        dist.all_gather_object(errors, error)
+        assert errors[0] is None, f"universal conversion failed: {errors[0]}"
+
+        states = {key: _logical_state(engine, key, tp_size) for key in UNIVERSAL_STATES}
+        losses = _run_to_boundary(engine, logical_rank, logical_world_size)
+        engine.step()
+        params_after = _logical_state(engine, "fp32", tp_size)
+        all_losses = [None] * dist.get_world_size()
+        dist.all_gather_object(all_losses, losses)
+        if dist.get_rank() == 0:
+            torch.save({
+                "states": states,
+                "next_losses": all_losses,
+                "next_params": params_after
+            }, os.path.join(save_dir, "source_record.pt"))
+        dist.barrier()
+
+
+@pytest.fixture(params=["fp32", "bf16"])
+def universal_precision(request):
+    return request.param
+
+
+@pytest.fixture
+def folded_universal_source(request):
+    _FoldedUniversalSource()(request)
+
+
+UNIVERSAL_TARGETS = [
+    (4, 2, 2),  # the same layout
+    (4, 2, 4),  # EP 4: each EP group spans both TP groups
+    (4, 4, 2),  # TP 4: one data-parallel group
+    (8, 4, 2),  # more ranks with the same 2 data-parallel groups
+    (2, 1, 2),  # AutoEP without AutoTP, 2 data-parallel groups
+]
+
+
+@pytest.mark.parametrize("world_size, tp_size, ep_size", UNIVERSAL_TARGETS)
+class TestZero3FoldedUniversalCheckpoint(DistributedTest):
+    """A folded ZeRO-3 checkpoint converted to the universal format loads into another layout.
+
+    Every parameter and Adam state must come back exactly. Layouts with the source's 2 data-parallel
+    groups read the same batches, so they must also repeat its next step.
+    """
+
+    def test_universal_load_restores_and_resumes(self, folded_universal_source, universal_precision, tmpdir,
+                                                 world_size, tp_size, ep_size):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("AutoEP folding runs on an accelerator")
+        save_dir = str(tmpdir)
+        record = torch.load(os.path.join(save_dir, "source_record.pt"), weights_only=False)
+        engine = _build_universal_engine(tp_size, ep_size, universal_precision, seed=999, load_universal=True)
+        engine.load_checkpoint(save_dir, tag="step1_universal")
+
+        # Every rank finishes its collectives before any assertion, so a failure cannot leave others waiting.
+        problems = []
+        for key in UNIVERSAL_STATES:
+            loaded = _logical_state(engine, key, tp_size)
+            problems += [f"{key}: {m}" for m in _universal_mismatches(loaded, record["states"][key], 1e-6, 1e-8)]
+
+        source_tp, source_ep = UNIVERSAL_SOURCE
+        logical_world_size = world_size // tp_size
+        same_batches = logical_world_size == 4 // source_tp
+        same_layout = (world_size, tp_size, ep_size) == (4, source_tp, source_ep)
+        # In bf16 another TP split rounds the forward differently, so there only the same layout repeats the step.
+        if same_batches and (universal_precision == "fp32" or same_layout):
+            logical_rank = dist.get_rank() // tp_size
+            losses = _run_to_boundary(engine, logical_rank, logical_world_size)
+            engine.step()
+            want_losses = record["next_losses"][logical_rank * source_tp]
+            if max(abs(got - want) for got, want in zip(losses, want_losses)) > 1e-5:
+                problems.append(f"next losses {losses}, expected {want_losses}")
+            # The attention key bias has a zero true gradient, so its gradient is rounding noise that
+            # another layout adds up in another order, and AdamW turns that noise into steps of about lr.
+            after = {
+                name: value
+                for name, value in _logical_state(engine, "fp32", tp_size).items() if not name.endswith("k_proj.bias")
+            }
+            want = {name: record["next_params"][name] for name in after}
+            problems += [f"after the next step: {m}" for m in _universal_mismatches(after, want, 1e-4, 1e-6)]
+
+        assert not problems, f"{len(problems)} problems:\n" + "\n".join(problems[:40])
 
 
 class _LlamaWithoutTPPlan(transformers.LlamaForCausalLM):