From 9362efc52a983830845a2048e1ff562fbbdf4cd5 Mon Sep 17 00:00:00 2001 From: pengdurice Date: Thu, 15 Oct 2026 19:32:17 +0000 Subject: [PATCH] AutoEP + AutoTP folding: Universal Checkpointing under ZeRO-3 Signed-off-by: pengdurice --- deepspeed/checkpoint/ds_to_universal.py | 35 ++- deepspeed/runtime/engine.py | 17 +- docs/_pages/config-json.md | 2 +- docs/code-docs/source/autoep.rst | 5 +- .../test_autoep_autotp_sharded_attention.py | 203 ++++++++++++++++-- 5 files changed, 220 insertions(+), 42 deletions(-) diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index 58bde9d78e68..1540716f63b9 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/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) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index f5066ce0cad9..916808419665 100644 --- a/deepspeed/runtime/engine.py +++ b/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 diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index a658724ee577..f473d9f1037d 100644 --- a/docs/_pages/config-json.md +++ b/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 diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index a20891eb904c..c09554855138 100644 --- a/docs/code-docs/source/autoep.rst +++ b/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. diff --git a/tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py b/tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py index 1cc88cf6ef26..d798fa2c5dc8 100644 --- a/tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py +++ b/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):