From 501b508abc9d9b3c65ee767e773b06989f53f761 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Sat, 14 Nov 2026 19:06:29 -0500 Subject: [PATCH] Test the weight sync from an AutoTP trainer into a TP4 copy An AutoTP TP2 x DP2 trainer under ZeRO-1, ZeRO-2 and ZeRO-3, and under ZeRO-3 with the model built in zero.Init, syncs into one TP4 copy over the same four ranks, with autotp_storage_affine_maps as the source maps. Each shard is checked against the engine's own 16-bit consolidation, before and after a training step. The module docstring and the docs no longer call the AutoTP source untested. Signed-off-by: Achyuthan Sivasankar --- .../runtime/rollout/affine_weight_sync.py | 4 +- docs/code-docs/source/affine-transfer.rst | 4 +- .../rollout/test_affine_weight_sync.py | 74 ++++++++++++++++++- 3 files changed, 77 insertions(+), 5 deletions(-) diff --git a/deepspeed/runtime/rollout/affine_weight_sync.py b/deepspeed/runtime/rollout/affine_weight_sync.py index 57ed82985fbc..4c1daa3e3a89 100644 --- a/deepspeed/runtime/rollout/affine_weight_sync.py +++ b/deepspeed/runtime/rollout/affine_weight_sync.py @@ -7,8 +7,8 @@ A rollout engine that generates with its own tensor-parallel copy of the weights refreshes that copy after every optimizer step. Gathering each parameter whole and cutting out the local shard sends every parameter to every rank. Instead, each parameter is described twice with affine maps (:mod:`deepspeed.checkpoint.affine`): how the trainer -stores it (:func:`zero_param_shaped_storage_map`; ``autotp_storage_affine_maps`` should serve for an AutoTP model, -but that is not tested yet) and which block each inference rank needs (:func:`replicated_layout_map`). +stores it (:func:`zero_param_shaped_storage_map`, or ``autotp_storage_affine_maps`` for an AutoTP model) and which +block each inference rank needs (:func:`replicated_layout_map`). :func:`~deepspeed.checkpoint.affine_transfer.plan_transfers` composes the two into copies from trainer shards straight into inference shards, once per layout, and every sync runs the copies of all parameters as one :func:`~deepspeed.checkpoint.affine_transfer_executor.execute_transfers` batch. diff --git a/docs/code-docs/source/affine-transfer.rst b/docs/code-docs/source/affine-transfer.rst index d44fc65ae715..1a60a053920e 100644 --- a/docs/code-docs/source/affine-transfer.rst +++ b/docs/code-docs/source/affine-transfer.rst @@ -96,8 +96,8 @@ model refreshes that copy after every optimizer step. ``deepspeed.runtime.rollout.affine_weight_sync`` does it shard to shard: no rank gathers a whole parameter. Both sides are maps keyed by global rank, also when the copies run over a custom process group. The trainer side comes from -``zero_param_shaped_storage_map`` (``autotp_storage_affine_maps`` should serve -for an AutoTP model, but that is not tested yet). The inference side is a +``zero_param_shaped_storage_map``, or from ``autotp_storage_affine_maps`` for +an AutoTP model. The inference side is a tensor-parallel layout held by every replica, from ``replicated_layout_map``:: diff --git a/tests/unit/runtime/rollout/test_affine_weight_sync.py b/tests/unit/runtime/rollout/test_affine_weight_sync.py index dcfffbbb80bf..13b3ac0ca28e 100644 --- a/tests/unit/runtime/rollout/test_affine_weight_sync.py +++ b/tests/unit/runtime/rollout/test_affine_weight_sync.py @@ -13,7 +13,7 @@ from deepspeed.checkpoint.affine import contiguous_split_map, replicated_map from deepspeed.checkpoint.affine_replica_policy import first_holder from deepspeed.checkpoint.affine_transfer_executor import AffineTransferExecutionError -from deepspeed.checkpoint.zero_affine import zero_param_storage_shard +from deepspeed.checkpoint.zero_affine import autotp_storage_affine_maps, zero_param_storage_shard from deepspeed.runtime.rollout.affine_weight_sync import (AffineWeightSync, WeightArena, replicated_layout_map, zero_param_shaped_storage_map) from unit.common import DistributedTest @@ -147,6 +147,78 @@ def test_sync_writes_every_replica_in_place(self, zero_stage, replicas): assert not torch.equal(second_values, first_values) engine.destroy() + # An AutoTP TP2 x DP2 trainer and one TP4 copy over the same four ranks, whose shards are uneven. + @pytest.mark.parametrize('zero_stage, zero_init', [(1, False), (2, False), (3, False), (3, True)]) + def test_an_autotp_trainer_syncs_into_a_tp4_copy(self, zero_stage, zero_init): + dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32 + config = { + 'train_micro_batch_size_per_gpu': 1, + 'optimizer': { + 'type': 'Adam', + 'params': { + 'lr': 1e-2, + 'torch_adam': True + } + }, + 'zero_optimization': { + 'stage': zero_stage + }, + 'tensor_parallel': { + 'autotp_size': 2, + 'partition_config': { + 'use_default_specs': + False, + 'layer_specs': [ + { + 'patterns': [r'.*up\.weight$'], + 'partition_type': 'column' + }, + { + 'patterns': [r'.*down\.weight$'], + 'partition_type': 'row' + }, + ], + }, + }, + } + if dtype == torch.bfloat16: + config['bf16'] = {'enabled': True} + # Every rank builds the model from the same seed, as loading a checkpoint would give: the TP ranks must hold + # equal copies of the weights AutoTP leaves whole, and AutoTP checks that they see the same batch. + torch.manual_seed(0) + if zero_init: + # zero.Init partitions every parameter over the whole world, and the norm, which AutoTP leaves alone, + # keeps that partition. + with deepspeed.zero.Init(config_dict_or_path=config): + model = MLPBlock() + else: + model = MLPBlock() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + params = dict(engine.module.named_parameters()) + layouts = _layouts(4) + rank = dist.get_rank() + + target_maps = {name: replicated_layout_map(layouts[name], [[0, 1, 2, 3]]) for name in params} + sync = AffineWeightSync(autotp_storage_affine_maps(engine.module), target_maps, dtype) + arena = WeightArena({name: layouts[name].shard_shapes[rank] for name in params}, dtype, engine.device) + for step in range(2): + if zero_stage == 3: + engine.empty_partition_cache() + sync.sync({name: zero_param_storage_shard(param) for name, param in params.items()}, dict(arena.items())) + # The full weights from the engine's own 16-bit consolidation, which gathers over ZeRO and TP and shares no + # code with the affine maps. + full = _all(engine._consolidated_16bit_state_dict())[0] + wrong = [] + for name in params: + if not torch.equal(arena[name].cpu(), layouts[name].extract(full[name], rank)): + wrong.append(name) + # Every rank fails together, rather than one rank failing while the others wait in the next collective. + assert _all(wrong) == [[]] * dist.get_world_size() + for x, y in random_dataloader(engine, 1, HIDDEN, engine.device, dtype): + engine.backward(engine(x, y)) + engine.step() + engine.destroy() + @pytest.mark.parametrize('policy', ['default', 'first_holder']) def test_reads_of_replicated_parameters_spread_over_their_holders(self, policy): # Ranks 0-2 each hold all six parameters whole, and rank 3 needs them all. Each parameter is a single region,