[jinyouzhi]

Summary

Adds training/deepspeed_finetune_demo/autotp_rollout_e2e.py and run_autotp_rollout_e2e.sh: AutoTP + ZeRO training → rollout engines at a different TP size, three ways:

• ucp: DeepSpeed checkpoint → Universal checkpoint → load.
• gather: gather whole parameters, slice the rollout shard.
• affine: autotp_storage_affine_maps (ZeRO's own storage) → plan_transfer → execute_transfers; no full parameter is materialized.

Rollout engines start from NaN weights; every path is checked bit for bit against whole parameters rebuilt with plain collectives. Plans are built once and reused for repeated syncs. --replica_policy (lowest / local / balanced) chooses which ZeRO-0/1/2 replica each target block reads from.

Companion to #1020 (AutoEP); this one covers dense AutoTP weights.

Depends on deepspeedai/DeepSpeed#8809 (executor) and deepspeedai/DeepSpeed#8885 (ZeRO / AutoTP storage maps).

Results

Qwen3-8B, 8× NVIDIA L20, train TP 2 → rollout TP 4, --paths gather,affine, 3 training steps + 3 repeated syncs, all bitwise equal:

	ZeRO-2 gather	ZeRO-2 affine lowest	ZeRO-2 affine balanced	ZeRO-3 gather	ZeRO-3 affine
maps + plan (once)	-	0.63 + 0.57 s	0.70 + 0.70 s	-	0.82 + 0.60 s
first sync	1.01 s	1.16 s	0.48 s	1.71 s	1.18 s
repeated syncs	0.78-0.82 s	1.05-1.35 s	0.49-0.55 s	1.22-1.28 s	0.45-0.74 s
device memory added	3.48 GB	0.30 GB	0.16 GB	3.48 GB	0.66 GB
sent / local copy	-	35.2 / 5.0 GB	15.1 / 25.1 GB	-	34.0 / 6.3 GB
executor rounds	-	96	41	-	32

Each map has 3,192 pieces on both sides.

[jinyouzhi]

Verification: AffineIR resharding across training world size, ZeRO stage, AutoTP size and AutoEP size

I ran standalone experiments (two scripts, below) to check that the zero_affine storage maps plus plan_transfer / execute_transfers can move a live training state to an engine with a different layout mid-training, without changing the result. Each phase trains 20 steps, for 60 in total. At each switch a new engine is built and every shard it owns (bf16 params, fp32 master, Adam exp_avg / exp_avg_sq) is pre-filled with NaN. AffineIR then moves the state into it, Adam step is copied, and training continues.

Oracle. Every AffineIR run (affine) is repeated with a reference path (gather). The reference reads each full tensor from the source with safe_get_full_*, plus a plain all-gather over the TP or EP group in sections 3 and 4. It writes the tensor back with DeepSpeed's own setters: safe_set_full_* for ZeRO-1/2, safe_set_local_* for ZeRO-3. It shares no code with AffineIR, so a bitwise match between the two runs means the move is exact.

Checks at every switch

• bf16 param, fp32 master and Adam states, read through the oracle, are bitwise equal before and after the switch, and no NaN is left.
• Logits of a fixed batch from the old engine and from the new engine are compared. They are bitwise equal unless the new layout changes the arithmetic itself (section 3, and one case in section 4).

Checks at the end: the whole 60-step loss curve and the final logits are compared bitwise between the affine and gather runs.

1. Changing the training world size is consistent

run	process world	phases (stage x DP)	every switch: state / logits bitwise	affine vs gather: losses / final logits bitwise
4-2-4	4	Z3x4 → Z3x2 → Z3x4	✅ / ✅	✅ / ✅
4-8-4	8	Z3x4 → Z3x8 → Z3x4	✅ / ✅	✅ / ✅

Analysis:

• Shrink and grow both covered. The DP size shrinks (4→2) and grows (4→8), and in both directions the ZeRO-3 partition count and the padded partition size change.
• Within a switch: the full state and the logits are unchanged across the switch.
• Over the whole run: training after the switch is bitwise identical to the run that resharded through the gather oracle.
• Baselines without a switch:
  • Z3 DP4 for 60 steps gives the same losses bitwise in a 4-process and an 8-process world, and when it "reshards" 4→4→4 through AffineIR.
  • Against these baselines, the resized runs differ by at most ~1e-3 relative loss.
  • That difference is expected: gradients are reduced over a different number of ranks. The oracle comparison above shows it does not come from the move.
• Negative control: moving only parameters and leaving Adam states at zero breaks the state check and diverges by up to 0.39 relative loss. The checks detect an incomplete move.

2. Changing the ZeRO stage is consistent

run	process world	phases (stage x DP)	every switch: state / logits bitwise	affine vs gather: losses / final logits bitwise
Z3↔Z2	4	Z3x4 → Z2x4 → Z3x4	✅ / ✅	✅ / ✅
Z3↔Z1	4	Z3x4 → Z1x4 → Z3x4	✅ / ✅	✅ / ✅
stage + DP	4	Z3x4 → Z2x2 → Z1x4	✅ / ✅	✅ / ✅
stage + DP	8	Z2x8 → Z3x4 → Z1x8	✅ / ✅	✅ / ✅

Analysis:

• The layouts really differ.
  • ZeRO-1/2 replicate the bf16 parameter on every DP rank and split the fp32 master and Adam states into one contiguous range per rank of the flattened group. Those ranges cross parameter boundaries, and some ranks hold an empty range.
  • ZeRO-3 partitions every parameter evenly, with padding.
  • The maps on the two sides therefore have different holders, piece boundaries and replication. Any switch between stages has to re-tile the state, not just copy it.
• Results. Both directions are covered (3→1/2→3), and stage and DP size can change in the same switch. In every case the state and logits are unchanged across the switch, and the loss curve matches the gather oracle bitwise.
• The difference from a ZeRO-3-only run (≤1.2e-3 relative loss) is the expected stage-dependent reduction order.

3. Changing the AutoTP size is consistent

A second script (tpep.py, below) uses the same protocol on a 3-layer MLP under AutoTP: l1 column-parallel, l2 row-parallel, l3 replicated. A phase is Z<stage>TP<tp>, with DP = world / TP.

• AffineIR side: autotp_storage_affine_maps(model, tp_group=...) for the bf16 params and autotp_storage_affine_maps(model, optimizer_states=True, ...) for the master and Adam states. These are the shipped composed maps (TP layout × ZeRO storage).
• Oracle side: the ZeRO-full tensor of each rank, all-gathered over the TP group and concatenated along the split dim the oracle knows independently. Writing back uses DeepSpeed's setters on the destination's own TP slice.

run	world	phases	every switch: state bitwise / logits max abs diff	affine vs gather: losses / final logits bitwise	vs a run that keeps the first phase's config: max rel loss diff
Z1	4	Z1TP2 → Z1TP4 → Z1TP2	✅ / 4.9e-4	✅ / ✅	9.6e-4
Z2	4	Z2TP2 → Z2TP4 → Z2TP2	✅ / 4.9e-4	✅ / ✅	9.6e-4
Z3	4	Z3TP2 → Z3TP4 → Z3TP2	✅ / 2.0e-3	✅ / ✅	1.4e-3
stage + TP	4	Z1TP2 → Z3TP4 → Z2TP2	✅ / 4.9e-4	✅ / ✅	1.4e-3
TP up to 8	8	Z2TP2 → Z2TP8 → Z2TP4	✅ / 2.0e-3	✅ / ✅	1.6e-3
stage + TP	8	Z3TP4 → Z1TP8 → Z2TP2	✅ / 2.0e-3	✅ / ✅	1.3e-3

Analysis:

• What changes at a switch.
  • The TP size changes the slice each rank holds: rows of l1, columns of l2.
  • DP = world / TP changes with it, so the ZeRO partitioning of each slice changes as well.
  • The TP size and the ZeRO stage can change in the same switch.
• The state moves exactly. At every switch the full bf16 params, master and Adam states are bitwise equal and no NaN is left. The whole loss curve and the final logits match the gather oracle bitwise.
• Logits are not bitwise across TP sizes, and should not be.
  • The row-parallel l2 all-reduces partial sums, and the number of partial sums is the TP size. A different TP size adds the same terms in a different order.
  • The reference runs keep one configuration and still reshard at every switch (e.g. Z1TP2 → Z1TP2). Their logits are bitwise equal at every switch.
• The loss difference comes from the TP size, not from the switch. Two runs without any switch that differ only in TP size (TP2 vs TP4, same stage) already differ by up to 3.2e-3. The switched runs stay within 1.6e-3 of their reference.

4. Changing the AutoEP size is consistent

Same script and protocol on the AutoEP MockMoETransformer from tests/unit/v1/moe/autoep_test_utils.py: 2 layers, 8 experts, top-2, preset_model: mixtral, use_grouped_mm: false. A phase is Z<stage>EP<ep>; DP is the whole world, and the expert-data-parallel group is world / EP.

• AffineIR side, parameters: expert_storage_affine_maps() for expert params, zero_param_storage_affine_map for the rest.
• AffineIR side, expert optimizer states. The shipped API stops short of these: autotp_storage_affine_maps(optimizer_states=True) raises NotImplementedError when AutoEP is present. The script composes them itself, the same way the parameter maps are built:
  • take nest_affine_maps(layer.expert_affine_maps()[w], {ep_rank: ZeRO optimizer map of that rank's packed tensor});
  • as a check, build the parameter maps the same way. They equal the shipped expert_storage_affine_maps() exactly in every run.
• Oracle side: the packed expert tensor of each EP rank, all-gathered over the EP group and concatenated on the expert dim.

run	world	phases	every switch: state / logits bitwise	affine vs gather: losses / final logits bitwise	vs a run that keeps the first phase's config: max rel loss diff
Z2	4	Z2EP2 → Z2EP4 → Z2EP2	✅ / ✅	✅ / ✅	1.7e-2
Z3	4	Z3EP2 → Z3EP4 → Z3EP2	✅ / ✅	✅ / ✅	1.3e-2
stage + EP	4	Z2EP2 → Z3EP4 → Z1EP2	✅ / ❌ (see below)	✅ / ✅	1.4e-2
EP up to 8	8	Z2EP2 → Z2EP8 → Z2EP4	✅ / ✅	✅ / ✅	1.1e-2
stage + EP	8	Z3EP4 → Z1EP8 → Z2EP2	✅ / ✅	✅ / ✅	9.5e-3

Analysis:

• What changes at a switch.
  • The EP size changes which experts each rank holds: 4 → 2 → 1 experts per rank at EP 2/4/8. Expert params are packed per rank as [local_experts, ...].
  • The expert-data-parallel group changes from world/2 to world/4 or world/8 ranks, so the ZeRO partitioning of the packed tensors changes too. Non-expert params stay data-parallel over the whole world.
• The state moves exactly. At every switch the full state is bitwise equal and no NaN is left. All loss curves and final logits match the gather oracle bitwise. Unlike TP, the expert computation has no cross-rank reduction, so in 4 of the 5 runs the logits are bitwise equal across the switch as well.
• The ❌ in the stage + EP run is a GEMM alignment effect, not a resharding error.
  • Its state checks pass. The gather oracle shows exactly the same logits difference.
  • The first divergence is a single element (of 8192) in layer 0's MoE output. Routing in layer 1 then amplifies it.
  • Tracing it to the expert: with world 4 and EP 4, the expert-data-parallel group has a single rank. ZeRO-3 then places the packed expert weights at an odd element offset (10329) in its storage, i.e. only 2-byte aligned.
  • On one GPU, _run_experts_for_loop with the same weights and inputs reproduces the post-switch output bitwise when the weights start at an odd element offset (1, 3 or 10329). With freshly allocated, aligned weights, the mismatched row comes out as before the switch. The GEMM result depends on the weight address alignment.
  • The same effect shows when only the EP size changes and the stage stays ZeRO-3: a diagnostic run Z2EP2 → Z3EP2 → Z3EP4 shows it at the Z3EP2 → Z3EP4 switch. It is data-dependent, since a rounding difference must land on a bf16 boundary to show. The Z3 run in the table crosses the same switch at step 20 with bitwise-equal logits.
• The loss difference comes from the EP size, not from the switch. Two runs without any switch that differ only in EP size (EP2 vs EP4, same stage) already differ by up to 1.8e-2. That is larger than for TP because discrete top-k routing amplifies rounding differences in the reduced expert gradients. The switched runs stay within 1.7e-2 of their reference.
• Two DeepSpeed issues found on the way:
  • ZeRO-1/2 + AutoEP leaves expert params without optimizer states. ZeRO-1/2 links optimizer states right after stepping the first param group, so the expert groups get none. Both safe_get_full_optimizer_state and AffineIR's zero_optimizer_storage_shard then fail on expert params. This is the known [BUG] ZeRO stage 1/2 marks optimizer-state fragments linked before later groups initialize DeepSpeed#8911 (fixes in #8956 and #8643); the script re-links after each step as a workaround.
  • The F.linear replacement leaks across engines. Building a ZeRO-3 engine from a plain module replaces torch.nn.functional.linear process-wide, as documented. On 3-D inputs the replacement gives different bf16 results from the stock op; I checked this standalone. From then on every engine in the process, ZeRO-1/2 included, computes its linear layers differently, so the script sets memory_efficient_linear: false.

AffineIR is the path that moved the data

These hold for both scripts unless noted.

1. Byte accounting. At every switch, the bytes execute_transfers wrote (local copies plus received) equal the bytes the target maps hold. Examples:
   • Z3x4→Z2x2: wrote 33,161,584 B, the maps hold 33,161,584 B.
   • Z2TP2→Z2TP8: 29,274,224 B on both sides.
   • Z3EP4→Z1EP8: 6,661,872 B on both sides.
2. No other path available. During resharding, every full-tensor entry point (safe_get/set_*, tensor_fragment.{get,set}_full_hp_param, ZeRO-3 {get,set}_full_hp_param / set_local_hp_param) is patched to raise. The runs still finish, bitwise equal to the oracle.
3. NaN pre-fill. Any element the plan missed would stay NaN; none did.
4. Sabotage (z3x.py only). Swapping the pieces of ranks 0 and 1 in one source map makes planning fail with InvalidMapError: source: rank 1 does not hold all of Box(...) before any data moves.

Scope and limitations

• Same process world. "World size" here is the number of ranks the engine trains on, inside a fixed process world (4 or 8). execute_transfers runs on the whole process group, so adding or removing processes was not tested.

• Two monkeypatches in z3x.py. A ZeRO-3 engine on a subgroup with idle ranks needed them:

  • the world-group dist.barrier() in stage-3 init is scoped to the DP group;
  • the world-group collectives in zero.utils.assert_lst_len_same_as_other_ranks and get_lst_from_rank0 are disabled.

  These are DeepSpeed limitations, not AffineIR ones. tpep.py needs neither, since every phase uses all ranks.

• Adam step is copied separately, since it is not part of the storage maps.

• Expert optimizer-state maps are composed by the script (section 4); the shipped API covers expert parameters only.

• One parallelism kind per run. TP and EP are tested separately, each combined with DP and ZeRO; AutoEP+AutoTP folding is not.

• Not covered: offload, fp16.

• Note on the oracle. stage3.set_full_hp_param slices the full value by the padded partition size. On a padded last partition it would read out of range, so the oracle uses safe_set_local_*. I found this from the code and did not reproduce it separately.

Environment and reproduction

• Hardware: 8× NVIDIA H20, torch 2.9.0a0+50eac811a6.nv25.09, Python 3.12.
• DeepSpeed: b62afa821, an ancestor of the current master.
  • The AffineIR, AutoEP and ZeRO-1/2 files are unchanged since that commit.
  • stage3.py changed only for fp32-weight master reuse (#8842) and AutoEP+AutoTP folding offload (#8829). Neither path is exercised here.
• Commands:
  • torchrun --nproc_per_node=<4|8> z3x.py <name> <phases> 20 <affine|gather|noopt>. A phase is <stage>x<dp>; the engine of each phase runs on global ranks [0, dp).
  • torchrun --nproc_per_node=<4|8> tpep.py <name> <phases> 20 <affine|gather>. A phase is <stage>t<tp> or <stage>e<ep>, e.g. 1t2-3t4-2t2. It imports MockMoETransformer from the DeepSpeed source tree's tests/.

z3x.py (sections 1 and 2)

# ZeRO resharding with AffineIR across ZeRO stages and DP sizes, checked against a gather/scatter oracle.
# usage: torchrun --nproc_per_node=W z3x.py <name> <phases, e.g. 3x4-2x2-3x4> <steps per phase> [affine|gather|noopt]
# A phase is <stage>x<dp>; a bare <dp> means ZeRO-3. The engine of a phase lives on global ranks [0, dp).
import json
import math
import os
import sys

import torch

import deepspeed
import deepspeed.comm as dist
import deepspeed.runtime.zero.utils  # noqa: F401  (patched through sys.modules below)
from deepspeed.accelerator import get_accelerator
from deepspeed.checkpoint.affine_transfer import plan_transfer
from deepspeed.checkpoint.affine_transfer_executor import TransferEndpoints, TransferRequest, execute_transfers
from deepspeed.checkpoint.zero_affine import (zero_optimizer_storage_affine_maps, zero_optimizer_storage_shard,
                                              zero_param_storage_affine_map, zero_param_storage_shard)
from deepspeed.utils import (safe_get_full_fp32_param, safe_get_full_optimizer_state, safe_set_full_fp32_param,
                             safe_set_full_optimizer_state, safe_set_local_fp32_param, safe_set_local_optimizer_state)


def parse(token):
    stage, dp = token.split('x') if 'x' in token else ('3', token)
    return int(stage), int(dp)


NAME, PHASES, STEPS = sys.argv[1], [parse(t) for t in sys.argv[2].split('-')], int(sys.argv[3])
MODE = sys.argv[4] if len(sys.argv) > 4 else 'affine'
MOVE_OPT = MODE != 'noopt'
SABOTAGE = 'SABOTAGE' in os.environ  # swap ranks 0 and 1 of l1.weight's source maps; the checks must fail
GLOBAL_BATCH, DIN, HID = 16, 1000, 1023  # odd sizes so ZeRO partitions are padded / unaligned
STATE_KEYS = ('exp_avg', 'exp_avg_sq')

deepspeed.init_distributed()
rank, world = dist.get_rank(), dist.get_world_size()
if 'DUMP_AFTER' in os.environ:
    import faulthandler
    import signal
    faulthandler.register(signal.SIGUSR1, all_threads=False)
get_accelerator().set_device(int(os.environ.get('LOCAL_RANK', rank)))
dev = get_accelerator().current_device_name()
singletons = [dist.new_group([r]) for r in range(world)]
dp_groups = {dp: dist.new_group(list(range(dp))) for dp in sorted({dp for _, dp in PHASES})}
if min(dp for _, dp in PHASES) < world:
    # ZeRO-3's first-trace sanity check broadcasts on the world group, which idle ranks never join.
    zero_utils = sys.modules['deepspeed.runtime.zero.utils']
    zero_utils.assert_lst_len_same_as_other_ranks = lambda lst: None
    zero_utils.get_lst_from_rank0 = lambda lst: list(lst)


class MPU:

    def __init__(self, dp):
        self.dp, self.group = dp, dp_groups[dp]

    def get_data_parallel_group(self):
        return self.group

    def get_data_parallel_world_size(self):
        return self.dp

    def get_data_parallel_rank(self):
        return dist.get_rank(group=self.group)

    def get_data_parallel_group_ranks(self):
        return list(range(self.dp))

    def get_model_parallel_group(self):
        return singletons[rank]

    def get_model_parallel_world_size(self):
        return 1

    def get_model_parallel_rank(self):
        return 0

    def get_model_parallel_src_rank(self):
        return rank


class Net(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(DIN, HID)
        self.l2 = torch.nn.Linear(HID, HID)
        self.l3 = torch.nn.Linear(HID, 1)

    def forward(self, x, y=None):
        out = self.l3(torch.nn.functional.gelu(self.l2(torch.nn.functional.gelu(self.l1(x)))))
        return out if y is None else torch.nn.functional.mse_loss(out.float(), y)


teacher = torch.randn(DIN, generator=torch.Generator().manual_seed(7)) / math.sqrt(DIN)
eval_x = torch.randn(GLOBAL_BATCH, DIN, generator=torch.Generator().manual_seed(4242)).to(dev, torch.bfloat16)


def batch(step, dp, dp_rank):
    g = torch.Generator().manual_seed(1000 + step)
    x = torch.randn(GLOBAL_BATCH, DIN, generator=g)
    y = torch.sin(3 * x @ teacher).unsqueeze(1)
    micro = GLOBAL_BATCH // dp
    s = slice(dp_rank * micro, (dp_rank + 1) * micro)
    return x[s].to(dev, torch.bfloat16), y[s].to(dev)


class Phase:
    """The engine of one phase, or None on ranks outside its DP group."""

    def __init__(self, stage, dp, engine):
        self.stage, self.dp, self.engine = stage, dp, engine

    def params(self):
        return dict(self.engine.module.named_parameters()) if self.engine else {}

    def release_gathered(self):
        if self.engine and self.stage == 3:
            self.engine.empty_partition_cache()


def build(stage, dp):
    cfg = {
        'train_micro_batch_size_per_gpu': GLOBAL_BATCH // dp,
        'gradient_accumulation_steps': 1,
        'optimizer': {
            'type': 'Adam',
            'params': {
                'lr': 1e-3
            }
        },
        'bf16': {
            'enabled': True
        },
        'zero_optimization': {
            'stage': stage
        },
    }
    torch.manual_seed(42)
    # Engine init barriers on the world group; scope those to the engine's own group.
    world_barrier = dist.barrier
    if dp < world:
        dist.barrier = lambda group=None, **k: world_barrier(group=group or dp_groups[dp], **k)
    if stage == 3:
        with deepspeed.zero.Init(data_parallel_group=dp_groups[dp], config_dict_or_path=cfg):
            model = Net()
    else:
        model = Net()
    engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=cfg, mpu=MPU(dp))
    dist.barrier = world_barrier
    if rank == 0:
        print(f'built ZeRO-{stage} DP{dp}: {type(engine.optimizer).__name__}', flush=True)
    return Phase(stage, dp, engine)


def train(phase, first, last, losses):
    dp_rank = dist.get_rank(group=dp_groups[phase.dp])
    for step in range(first, last):
        loss = phase.engine(*batch(step, phase.dp, dp_rank))
        phase.engine.backward(loss)
        phase.engine.step()
        mean = loss.detach().float().clone()
        dist.all_reduce(mean, group=dp_groups[phase.dp])
        losses.append((mean / phase.dp).item())


def logits(phase):
    """Outputs of a fixed batch, computed by every rank of the engine; rank 0's copy is returned."""
    if phase.engine is None:
        return None
    with torch.no_grad():
        out = phase.engine(eval_x).float().cpu()
    phase.release_gathered()
    return out


def bcast(obj):
    holder = [obj]
    dist.broadcast_object_list(holder, src=0)
    return holder[0]


def adam_step(phase):
    steps = {s['step'] for s in phase.engine.optimizer.optimizer.state.values()} if phase.engine else None
    (step, ) = bcast(steps)
    return step


def set_adam_step(phase, step):
    if phase.engine:
        for state in phase.engine.optimizer.optimizer.state.values():
            state['step'] = step


def fresh(stage, dp):
    """A new engine whose every shard is NaN, with Adam states created by a no-op step."""
    if rank >= dp:
        dist.barrier()
        return Phase(stage, dp, None)
    phase = build(stage, dp)
    engine = phase.engine
    loss = engine(*batch(0, dp, rank)) * 0
    engine.backward(loss)
    engine.step()
    phase.release_gathered()
    with torch.no_grad():
        for p in engine.module.parameters():
            zero_param_storage_shard(p).fill_(float('nan'))
            for key in (None, ) + STATE_KEYS:
                zero_optimizer_storage_shard(p, key).fill_(float('nan'))
    dist.barrier()
    return phase


def storage_maps(phase):
    """name -> (param map, optimizer map), built by the engine's ranks and shared with every rank."""
    out = None
    if phase.engine:
        params = phase.params()
        names = sorted(params)
        opt_maps = zero_optimizer_storage_affine_maps([params[n] for n in names])
        out = {}
        for name, opt_map in zip(names, opt_maps):
            param_map = zero_param_storage_affine_map(params[name], replica_group=dp_groups[phase.dp])
            out[name] = (param_map, opt_map)
    return bcast(out)


class Forbidden:
    """Makes every full-tensor gather/scatter entry point raise while AffineIR moves the data."""

    def __enter__(self):
        # deepspeed.utils re-exports a class named tensor_fragment, which shadows the module attribute.
        tf = sys.modules['deepspeed.utils.tensor_fragment']
        from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 as Z3
        self.saved = []
        targets = [(globals(), n) for n in list(globals()) if n.startswith('safe_')]
        targets += [(vars(tf), n) for n in ('get_full_hp_param', 'set_full_hp_param')]
        targets += [(Z3, n) for n in ('get_full_hp_param', 'set_full_hp_param', 'set_local_hp_param')]
        for owner, name in targets:
            get = owner.get if isinstance(owner, dict) else lambda n, o=owner: getattr(o, n)
            self.saved.append((owner, name, get(name)))

            def boom(*a, _n=name, **k):
                raise RuntimeError(f'{_n} called during AffineIR resharding')

            if isinstance(owner, dict):
                owner[name] = boom
            else:
                setattr(owner, name, boom)
        return self

    def __exit__(self, *exc):
        for owner, name, value in self.saved:
            if isinstance(owner, dict):
                owner[name] = value
            else:
                setattr(owner, name, value)


def sabotage(param_map):
    import copy
    bad = copy.deepcopy(param_map)
    bad.pieces_by_rank[0], bad.pieces_by_rank[1] = bad.pieces_by_rank[1], bad.pieces_by_rank[0]
    return bad


def held_bytes(param_map, dtype):
    elements = sum(piece.numel for pieces in param_map.pieces_by_rank.values() for piece in pieces)
    return elements * torch.tensor([], dtype=dtype).element_size()


def reshard_affine(src, dst):
    with Forbidden():
        _reshard_affine(src, dst)


def _reshard_affine(src, dst):
    """Move every parameter shard, fp32 master and Adam state through AffineIR plans."""
    src_maps, dst_maps = storage_maps(src), storage_maps(dst)
    assert src_maps.keys() == dst_maps.keys()
    if SABOTAGE:
        src_maps['l1.weight'] = tuple(sabotage(m) for m in src_maps['l1.weight'])
    moved = torch.zeros(1, dtype=torch.int64, device=dev)
    expected = 0
    src.release_gathered()
    dst.release_gathered()
    src_params, dst_params = src.params(), dst.params()
    identity = {r: r for r in range(world)}
    kinds = [('param', torch.bfloat16)] + [(k, torch.float32) for k in ((None, ) + STATE_KEYS if MOVE_OPT else (None, ))]
    for kind, dtype in kinds:
        which = 0 if kind == 'param' else 1

        def shard(p):
            return zero_param_storage_shard(p) if kind == 'param' else zero_optimizer_storage_shard(p, kind)

        requests = []
        for name in sorted(src_maps):
            plan = plan_transfer(dst_maps[name][which], src_maps[name][which])
            src_buf = {rank: shard(src_params[name])} if src.engine else {}
            dst_buf = {rank: shard(dst_params[name])} if dst.engine else {}
            requests.append(TransferRequest(f'{name}:{kind}', plan, src_buf, dst_buf))
            expected += held_bytes(dst_maps[name][which], dtype)
        stats = execute_transfers(requests, TransferEndpoints(identity, identity), dtype)
        moved += stats.local_copy_bytes + stats.received_bytes
    dist.all_reduce(moved)
    if rank == 0:
        report.append(f'  AffineIR wrote {moved.item()} bytes into target shards, target maps hold {expected}')
    if not MOVE_OPT:
        with torch.no_grad():
            for p in dst.params().values():
                for key in STATE_KEYS:
                    zero_optimizer_storage_shard(p, key).zero_()
    set_adam_step(dst, adam_step(src) if MOVE_OPT else 0)


def write_full(phase, p, full, key):
    """Oracle write of a full fp32 value through DeepSpeed's own setters, plus the bf16 copy for the master."""
    if phase.stage == 3:
        # set_full_hp_param slices by the padded partition size, so scatter padded local partitions instead.
        part = p.ds_tensor.ds_numel
        dp_rank = dist.get_rank(group=dp_groups[phase.dp])
        padded = torch.zeros(part * phase.dp, dtype=torch.float32, device=dev)
        padded[:full.numel()] = full
        local = padded.narrow(0, part * dp_rank, part)
        if key is None:
            safe_set_local_fp32_param(p, local)
            p.ds_tensor.data.view(-1).copy_(local.to(p.ds_tensor.dtype))
        else:
            safe_set_local_optimizer_state(p, local, key)
    else:
        value = full.view(p.shape)
        if key is None:
            safe_set_full_fp32_param(p, value)
            p.data.copy_(value.to(p.dtype))
        else:
            safe_set_full_optimizer_state(p, value, key)


def reshard_gather(src, dst):
    """Oracle: gather each full tensor on the source, broadcast it, and scatter it into the target."""
    shapes = bcast({n: tuple(p.ds_shape if hasattr(p, 'ds_shape') else p.shape)
                    for n, p in src.params().items()} if src.engine else None)
    src.release_gathered()
    dst.release_gathered()
    src_params, dst_params = src.params(), dst.params()
    for name in sorted(shapes):
        numel = math.prod(shapes[name])
        for key in (None, ) + STATE_KEYS:
            full = torch.empty(numel, dtype=torch.float32, device=dev)
            if src.engine:
                p = src_params[name]
                got = safe_get_full_fp32_param(p) if key is None else safe_get_full_optimizer_state(p, key)
                if rank == 0:
                    full.copy_(got.reshape(-1))
            dist.broadcast(full, src=0)
            if dst.engine:
                with torch.no_grad():
                    write_full(dst, dst_params[name], full, key)
    set_adam_step(dst, adam_step(src))


def snapshot(phase):
    """Full bf16 params, fp32 master and Adam states on rank 0 of the engine's group."""
    if phase.engine is None:
        return None
    out = {}
    for name, p in phase.engine.module.named_parameters():
        if phase.stage == 3:
            with deepspeed.zero.GatheredParameters(p):
                lp = p.data.clone()
        else:
            lp = p.data.clone()
        out[name] = [lp] + [safe_get_full_fp32_param(p)] + [safe_get_full_optimizer_state(p, k) for k in STATE_KEYS]
    phase.release_gathered()
    return {n: [t.cpu() for t in ts] for n, ts in out.items()} if rank == 0 else None


def label(stage, dp):
    return f'Z{stage}DP{dp}'


losses, report, saved = [], [], {'end': [], 'start': []}
stage, dp = PHASES[0]
current = build(stage, dp) if rank < dp else Phase(stage, dp, None)
for index, (next_stage, next_dp) in enumerate(PHASES):
    if index:
        before = snapshot(current)
        end_logits = logits(current)
        target = fresh(next_stage, next_dp)
        (reshard_gather if MODE == 'gather' else reshard_affine)(current, target)
        after = snapshot(target)
        start_logits = logits(target)
        if rank == 0:
            exact = all(torch.equal(a, b) for n in before for a, b in zip(before[n], after[n]))
            nans = any(torch.isnan(t).any().item() for ts in after.values() for t in ts)
            same_logits = torch.equal(end_logits, start_logits)
            report.append(f'switch {label(stage, dp)}->{label(next_stage, next_dp)} at step {index * STEPS}: '
                          f'bf16 param+master+adam bitwise equal {exact}, NaN left {nans}, '
                          f'logits bitwise equal {same_logits}')
            saved['end'].append(end_logits)
            saved['start'].append(start_logits)
        if current.engine:
            current.engine.destroy()
        current, stage, dp = target, next_stage, next_dp
    if current.engine:
        train(current, index * STEPS, (index + 1) * STEPS, losses)
    dist.barrier()

final_logits = logits(current)
if rank == 0:
    saved['final'] = final_logits
    json.dump(losses, open(f'/tmp/z3x_{NAME}.json', 'w'))
    torch.save(saved, f'/tmp/z3x_{NAME}_logits.pt')
    print(f'### {NAME}: world {world}, phases {sys.argv[2]} x {STEPS} steps, mode {MODE}')
    for line in report:
        print(line)
    print('losses', ' '.join(f'{x:.5f}' for x in losses))
dist.barrier()

tpep.py (sections 3 and 4)

# Resharding a live training state across AutoTP / AutoEP sizes and ZeRO stages with AffineIR, checked
# against a gather/scatter oracle.
# usage: torchrun --nproc_per_node=W tpep.py <name> <phases> <steps per phase> <affine|gather>
# A phase is <stage>t<tp> (AutoTP, DP = W / tp) or <stage>e<ep> (AutoEP, DP = W); a run uses one kind.
import json
import math
import os
import re
import sys

import torch

import deepspeed
import deepspeed.comm as dist
from deepspeed.accelerator import get_accelerator
from deepspeed.checkpoint.affine import ParamAffineMap, nest_affine_maps
from deepspeed.checkpoint.affine_transfer import plan_transfer
from deepspeed.checkpoint.affine_transfer_executor import TransferEndpoints, TransferRequest, execute_transfers
from deepspeed.checkpoint.zero_affine import (autotp_storage_affine_maps, zero_optimizer_storage_affine_maps,
                                              zero_optimizer_storage_shard, zero_param_storage_affine_map,
                                              zero_param_storage_shard)
from deepspeed.utils import (groups, safe_get_full_fp32_param, safe_get_full_optimizer_state, safe_set_full_fp32_param,
                             safe_set_full_optimizer_state, safe_set_local_fp32_param, safe_set_local_optimizer_state)

sys.path.insert(0, os.path.join(os.path.dirname(deepspeed.__file__), '..', 'tests'))
from unit.v1.moe.autoep_test_utils import MockMoETransformer  # noqa: E402


def parse(token):
    stage, kind, size = re.fullmatch(r'(\d)([te])(\d+)', token).groups()
    return int(stage), kind, int(size)


NAME, PHASES, STEPS = sys.argv[1], [parse(t) for t in sys.argv[2].split('-')], int(sys.argv[3])
MODE = sys.argv[4]
KIND = PHASES[0][1]
assert all(kind == KIND for _, kind, _ in PHASES)
STATE_KEYS = ('exp_avg', 'exp_avg_sq')
KEYS = ('lp', None) + STATE_KEYS  # bf16 param, fp32 master, Adam states
GLOBAL_BATCH = 16
DIN, HID = 1001, 1024  # AutoTP model; DIN is odd so ZeRO partitions of TP shards are padded
NUM_EXPERTS, SEQ, HIDDEN, VOCAB = 8, 8, 64, 100  # AutoEP model
TP_DIM = {'l1.weight': 0, 'l1.bias': 0, 'l2.weight': 1}  # the oracle's own knowledge of the TP split
GROUP_STATE = ('_TENSOR_MODEL_PARALLEL_GROUP', '_MODEL_PARALLEL_GROUP', '_DATA_PARALLEL_GROUP',
               '_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE', '_MPU_TENSOR_MODEL_PARALLEL_RANK', 'mesh_device', 'mpu')

deepspeed.init_distributed()
rank, world = dist.get_rank(), dist.get_world_size()
get_accelerator().set_device(int(os.environ.get('LOCAL_RANK', rank)))
dev = get_accelerator().current_device_name()


class Net(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(DIN, HID)
        self.l2 = torch.nn.Linear(HID, HID)
        self.l3 = torch.nn.Linear(HID, 1)

    def forward(self, x):
        return self.l3(torch.nn.functional.gelu(self.l2(torch.nn.functional.gelu(self.l1(x)))))


def make_model():
    torch.manual_seed(42)
    if KIND == 't':
        return Net()
    return MockMoETransformer(num_experts=NUM_EXPERTS, hidden_size=HIDDEN, intermediate_size=128)


def inputs(seed):
    g = torch.Generator().manual_seed(seed)
    if KIND == 't':
        teacher = torch.randn(DIN, generator=torch.Generator().manual_seed(7)) / math.sqrt(DIN)
        x = torch.randn(GLOBAL_BATCH, DIN, generator=g)
        return x, torch.sin(3 * x @ teacher).unsqueeze(1)
    x = torch.randn(GLOBAL_BATCH, SEQ, HIDDEN, generator=g)
    return x, 0.1 * torch.randn(GLOBAL_BATCH, SEQ, VOCAB, generator=g)


eval_x = inputs(4242)[0].to(dev, torch.bfloat16)


def config(stage, size, dp):
    cfg = {
        'train_micro_batch_size_per_gpu': GLOBAL_BATCH // dp,
        'gradient_accumulation_steps': 1,
        'optimizer': {
            'type': 'Adam',
            'params': {
                'lr': 1e-3
            }
        },
        'bf16': {
            'enabled': True
        },
        'zero_optimization': {
            'stage': stage,
            # A ZeRO-3 engine built from a plain module swaps torch.nn.functional.linear for its own
            # process-wide, for every later engine too. On 3-D inputs that version rounds the bias add
            # separately, so the same weights would give different bf16 logits depending on whether a
            # ZeRO-3 engine was built yet. Keep the stock linear so engines are comparable bit for bit.
            'memory_efficient_linear': False
        },
    }
    if KIND == 't':
        specs = [{
            'patterns': [r'.*l1\.weight$'],
            'partition_type': 'column'
        }, {
            'patterns': [r'.*l2\.weight$'],
            'partition_type': 'row'
        }, {
            'patterns': [r'.*l3\.weight$'],
            'partition_type': 'skip'
        }]
        cfg['tensor_parallel'] = {
            'autotp_size': size,
            'partition_config': {
                'use_default_specs': False,
                'layer_specs': specs
            }
        }
    else:
        cfg['expert_parallel'] = {
            'enabled': True,
            'autoep_size': size,
            'preset_model': 'mixtral',
            'use_grouped_mm': False
        }
    return cfg


class Phase:

    def __init__(self, stage, size):
        from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
        self.stage, self.size = stage, size
        self.dp = world // size if KIND == 't' else world
        self.dp_rank = rank // size if KIND == 't' else rank
        if KIND == 't':
            # DeepSpeed keeps one TP mesh per process; drop the previous engine's so this one builds its own.
            for group in (groups._TENSOR_MODEL_PARALLEL_GROUP, groups._DATA_PARALLEL_GROUP):
                if group is not None:
                    groups._release_process_group_references(group)
        # No model_parameters: AutoEP replaces the expert modules inside initialize, so a list taken
        # beforehand would hand the optimizer the discarded expert parameters.
        self.engine, _, _, _ = deepspeed.initialize(model=make_model(), config=config(stage, size, self.dp))
        self.group_state = {name: getattr(groups, name) for name in GROUP_STATE}
        self.tp_group = groups.get_tensor_model_parallel_group() if KIND == 't' else None
        self.ep_layers = {
            name: module
            for name, module in self.engine.module.named_modules() if isinstance(module, AutoEPMoELayer)
        }
        if rank == 0:
            print(f'built {label(stage, size)}: {type(self.engine.optimizer).__name__}', flush=True)

    def activate(self):
        for name, value in self.group_state.items():
            setattr(groups, name, value)

    def params(self):
        return dict(self.engine.module.named_parameters())

    def release_gathered(self):
        if self.stage == 3:
            self.engine.empty_partition_cache()

    def layout(self, name):
        """(dim, group, size) of the model-parallel split of ``name``, or None when every rank holds it whole."""
        if KIND == 't':
            return (TP_DIM[name], self.tp_group, self.size) if name in TP_DIM and self.size > 1 else None
        if '.experts.' in name and self.size > 1:
            layer = self.ep_layers[name.split('.experts.')[0]]
            return 0, layer.ep_group, self.size
        return None


def link_all_optimizer_states(phase):
    """Work around deepspeedai/DeepSpeed#8911 (fixes pending in #8956 and #8643).

    ZeRO-1/2 links optimizer states to parameters right after stepping its first param group, so
    every later group -- AutoEP's expert groups -- is left with none, and both the oracle's
    safe_get_full_optimizer_state and AffineIR's zero_optimizer_storage_shard fail on them.
    """
    optimizer = phase.engine.optimizer
    if phase.stage < 3:
        optimizer._hp_optimizer_states_linked = False
        optimizer._lazy_init_hp_params_optimizer_state()


def label(stage, size):
    return f'Z{stage}{"TP" if KIND == "t" else "EP"}{size}'


def train(phase, first, last, losses):
    phase.activate()
    micro = GLOBAL_BATCH // phase.dp
    rows = slice(phase.dp_rank * micro, (phase.dp_rank + 1) * micro)
    for step in range(first, last):
        x, y = inputs(1000 + step)
        out = phase.engine(x[rows].to(dev, torch.bfloat16))
        loss = torch.nn.functional.mse_loss(out.float(), y[rows].to(dev))
        phase.engine.backward(loss)
        phase.engine.step()
        link_all_optimizer_states(phase)
        mean = loss.detach().float().clone()
        dist.all_reduce(mean)
        losses.append((mean / world).item())


def logits(phase):
    phase.activate()
    with torch.no_grad():
        out = phase.engine(eval_x).float().cpu()
    phase.release_gathered()
    return out


def adam_step(phase):
    steps = {s['step'] for s in phase.engine.optimizer.optimizer.state.values()}
    gathered = [None] * world
    dist.all_gather_object(gathered, steps)
    (step, ) = set().union(*gathered)
    return step


def set_adam_step(phase, step):
    for state in phase.engine.optimizer.optimizer.state.values():
        state['step'] = step


def fresh(stage, size):
    """A new engine whose every shard is NaN, with Adam states created by a no-op step."""
    phase = Phase(stage, size)
    micro = GLOBAL_BATCH // phase.dp
    x, y = inputs(0)
    out = phase.engine(x[:micro].to(dev, torch.bfloat16))
    phase.engine.backward(torch.nn.functional.mse_loss(out.float(), y[:micro].to(dev)) * 0)
    phase.engine.step()
    link_all_optimizer_states(phase)
    phase.release_gathered()
    with torch.no_grad():
        for p in phase.engine.module.parameters():
            zero_param_storage_shard(p).fill_(float('nan'))
            for key in (None, ) + STATE_KEYS:
                zero_optimizer_storage_shard(p, key).fill_(float('nan'))
    dist.barrier()
    return phase


# ---------------------------------------------------------------------------------------------------
# Oracle: full tensors through DeepSpeed's own getters/setters and plain collectives, no AffineIR.
# ---------------------------------------------------------------------------------------------------


def local_full(phase, p, key):
    """This rank's model-parallel shard of ``key``, whole across ZeRO."""
    if key == 'lp':
        if phase.stage == 3:
            with deepspeed.zero.GatheredParameters([p]):
                return p.data.clone()
        return p.data.clone()
    return safe_get_full_fp32_param(p) if key is None else safe_get_full_optimizer_state(p, key)


def full_tensor(phase, name, p, key):
    shard = local_full(phase, p, key).contiguous()
    layout = phase.layout(name)
    if layout is None:
        return shard
    dim, group, size = layout
    parts = [torch.empty_like(shard) for _ in range(size)]
    dist.all_gather(parts, shard, group=group)
    return torch.cat(parts, dim)


def write_local(phase, name, p, full, key):
    """Scatter a full tensor into this rank's storage of ``key``."""
    layout = phase.layout(name)
    if layout is not None:
        dim, group, size = layout
        full = full.chunk(size, dim)[dist.get_rank(group=group)]
    if phase.stage == 3:
        # set_full_hp_param slices by the padded partition size, so scatter padded local partitions instead.
        part = p.ds_tensor.ds_numel
        dp_size = dist.get_world_size(group=p.ds_process_group)
        dp_rank = dist.get_rank(group=p.ds_process_group)
        padded = torch.zeros(part * dp_size, dtype=torch.float32, device=dev)
        padded[:full.numel()] = full.reshape(-1)
        local = padded.narrow(0, part * dp_rank, part)
        if key is None:
            safe_set_local_fp32_param(p, local)
            p.ds_tensor.data.view(-1).copy_(local.to(p.ds_tensor.dtype))
        else:
            safe_set_local_optimizer_state(p, local, key)
    else:
        value = full.reshape(p.shape)
        if key is None:
            safe_set_full_fp32_param(p, value)
            p.data.copy_(value.to(p.dtype))
        else:
            safe_set_full_optimizer_state(p, value, key)


def reshard_gather(src, dst):
    src.release_gathered()
    dst.release_gathered()
    src_params, dst_params = src.params(), dst.params()
    for name in sorted(src_params):
        for key in (None, ) + STATE_KEYS:
            full = full_tensor(src, name, src_params[name], key)
            with torch.no_grad():
                write_local(dst, name, dst_params[name], full, key)
    set_adam_step(dst, adam_step(src))


def snapshot(phase):
    phase.release_gathered()
    out = {name: [full_tensor(phase, name, p, key).cpu() for key in KEYS] for name, p in sorted(phase.params().items())}
    phase.release_gathered()
    return out


# ---------------------------------------------------------------------------------------------------
# AffineIR
# ---------------------------------------------------------------------------------------------------


def ep_storage_maps(phase):
    """Param and optimizer maps of an AutoEP model.

    The shipped API describes expert *parameter* storage only. Expert optimizer storage is composed
    here the same way: the EP placement layout nested with ZeRO's optimizer storage of each EP rank's
    packed tensor. The parameter maps composed this way must equal the shipped ones.
    """
    params = phase.params()
    names = sorted(params)
    opt_list = zero_optimizer_storage_affine_maps([params[n] for n in names])
    param_maps, opt_maps, local_inner, layouts = {}, {}, {}, {}
    for name, opt_map in zip(names, opt_list):
        if '.experts.' not in name:
            param_maps[name] = zero_param_storage_affine_map(params[name])
            opt_maps[name] = opt_map
            continue
        layer_name, weight = name.split('.experts.')
        layer = phase.ep_layers[layer_name]
        param_maps[name] = layer.expert_storage_affine_maps()[weight]
        layouts[name] = layer.expert_affine_maps()[weight]
        local_inner[name] = (zero_param_storage_affine_map(params[name]).to_dict(), opt_map.to_dict())
    ep_group = next(iter(phase.ep_layers.values())).ep_group
    gathered = [None] * world
    dist.all_gather_object(gathered, (dist.get_rank(group=ep_group), local_inner))
    by_ep_rank = dict(gathered)
    composed_ok = True
    for name, layout in layouts.items():
        param_inner = {r: ParamAffineMap.from_dict(maps[name][0]) for r, maps in by_ep_rank.items()}
        opt_inner = {r: ParamAffineMap.from_dict(maps[name][1]) for r, maps in by_ep_rank.items()}
        composed = nest_affine_maps(layout, param_inner)
        composed_ok &= composed.to_dict() == param_maps[name].to_dict()
        opt_maps[name] = nest_affine_maps(layout, opt_inner)
    return param_maps, opt_maps, composed_ok


def storage_maps(phase):
    phase.activate()
    phase.release_gathered()
    if KIND == 't':
        model = phase.engine.module
        return (autotp_storage_affine_maps(model, tp_group=phase.tp_group),
                autotp_storage_affine_maps(model, optimizer_states=True, tp_group=phase.tp_group), True)
    return ep_storage_maps(phase)


class Forbidden:
    """Makes every full-tensor gather/scatter entry point raise while AffineIR moves the data."""

    def __enter__(self):
        tf = sys.modules['deepspeed.utils.tensor_fragment']
        from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 as Z3
        self.saved = []
        targets = [(globals(), n) for n in list(globals()) if n.startswith('safe_')]
        targets += [(vars(tf), n) for n in ('get_full_hp_param', 'set_full_hp_param')]
        targets += [(Z3, n) for n in ('get_full_hp_param', 'set_full_hp_param', 'set_local_hp_param')]
        for owner, name in targets:
            original = owner[name] if isinstance(owner, dict) else getattr(owner, name)
            self.saved.append((owner, name, original))

            def boom(*a, _n=name, **k):
                raise RuntimeError(f'{_n} called during AffineIR resharding')

            if isinstance(owner, dict):
                owner[name] = boom
            else:
                setattr(owner, name, boom)
        return self

    def __exit__(self, *exc):
        for owner, name, value in self.saved:
            if isinstance(owner, dict):
                owner[name] = value
            else:
                setattr(owner, name, value)


def held_bytes(param_map, dtype):
    elements = sum(piece.numel for pieces in param_map.pieces_by_rank.values() for piece in pieces)
    return elements * torch.tensor([], dtype=dtype).element_size()


def reshard_affine(src, dst):
    src_param_maps, src_opt_maps, src_ok = storage_maps(src)
    dst_param_maps, dst_opt_maps, dst_ok = storage_maps(dst)
    if rank == 0 and KIND == 'e':
        report.append(f'  composed EP param maps equal the shipped expert_storage_affine_maps: {src_ok and dst_ok}')
    src_params, dst_params = src.params(), dst.params()
    identity = {r: r for r in range(world)}
    moved = torch.zeros(1, dtype=torch.int64, device=dev)
    expected = 0
    with Forbidden():
        for key in KEYS:
            dtype = torch.bfloat16 if key == 'lp' else torch.float32
            src_maps, dst_maps = (src_param_maps, dst_param_maps) if key == 'lp' else (src_opt_maps, dst_opt_maps)

            def shard(p):
                return zero_param_storage_shard(p) if key == 'lp' else zero_optimizer_storage_shard(p, key)

            requests = []
            for name in sorted(src_params):
                plan = plan_transfer(dst_maps[name], src_maps[name])
                requests.append(
                    TransferRequest(f'{name}:{key}', plan, {rank: shard(src_params[name])},
                                    {rank: shard(dst_params[name])}))
                expected += held_bytes(dst_maps[name], dtype)
            stats = execute_transfers(requests, TransferEndpoints(identity, identity), dtype)
            moved += stats.local_copy_bytes + stats.received_bytes
    dist.all_reduce(moved)
    if rank == 0:
        report.append(f'  AffineIR wrote {moved.item()} bytes into target shards, target maps hold {expected}')
    set_adam_step(dst, adam_step(src))


losses, report, saved = [], [], {'end': [], 'start': []}
stage, _, size = PHASES[0]
current = Phase(stage, size)
for index, (next_stage, _, next_size) in enumerate(PHASES):
    if index:
        current.activate()
        before = snapshot(current)
        end_logits = logits(current)
        target = fresh(next_stage, next_size)
        (reshard_gather if MODE == 'gather' else reshard_affine)(current, target)
        target.activate()
        after = snapshot(target)
        start_logits = logits(target)
        if rank == 0:
            exact = all(torch.equal(a, b) for n in before for a, b in zip(before[n], after[n]))
            nans = any(torch.isnan(t).any().item() for ts in after.values() for t in ts)
            diff = (end_logits - start_logits).abs().max().item()
            report.append(f'switch {label(stage, size)}->{label(next_stage, next_size)} at step {index * STEPS}: '
                          f'bf16 param+master+adam bitwise equal {exact}, NaN left {nans}, logits bitwise equal '
                          f'{torch.equal(end_logits, start_logits)} (max abs diff {diff:.2e})')
            saved['end'].append(end_logits)
            saved['start'].append(start_logits)
        # The old engine is kept alive: destroying it could release process groups the new one shares.
        current, stage, size = target, next_stage, next_size
    train(current, index * STEPS, (index + 1) * STEPS, losses)
    dist.barrier()

final_logits = logits(current)
if rank == 0:
    saved['final'] = final_logits
    json.dump(losses, open(f'/tmp/tpep_{NAME}.json', 'w'))
    torch.save(saved, f'/tmp/tpep_{NAME}_logits.pt')
    print(f'### {NAME}: world {world}, phases {sys.argv[2]} x {STEPS} steps, mode {MODE}')
    for line in report:
        print(line)
    print('losses', ' '.join(f'{x:.5f}' for x in losses))
dist.barrier()

AI assistance was used to write the experiment script and this summary.

[jinyouzhi]

Re-run of the resharding verification on current master, plus aliased fp32 masters (#8842)

I re-ran every experiment from the verification above in one recorded environment on current DeepSpeed master. This ties every number to that environment and covers the ZeRO-3 changes merged since b62afa821: #8842 (fp32-weight master reuse), #8892 (fp32 master restore) and #8829. I also added a section 5 for the case #8842 introduces, where the ZeRO-3 master is the weight partition itself. z3x.py and tpep.py are unchanged from the comment above.

Environment

• 8× NVIDIA H20, torch 2.13.0+cu129, NCCL 2.29.7, Python 3.12.3.
• DeepSpeed master a4c989b76.
• 55 runs for sections 1–4 and 6 for section 5, each phase 20 steps (60 in total). Protocol and checks as above: NaN pre-filled targets, gather oracle that shares no code with AffineIR, full-tensor entry points patched to raise during the AffineIR move.

Result: at every switch of every run, the bf16 params, masters and Adam states are bitwise equal and no NaN is left. Every affine run matches its gather run bitwise in the whole loss curve, the logits at each switch and the final logits.

1. Training world size (z3x.py)

run	world	phases	every switch: state / logits bitwise	affine vs gather: losses / final logits bitwise	vs Z3x4 without switch: max rel loss diff
4-2-4	4	Z3x4 → Z3x2 → Z3x4	✅ / ✅	✅ / ✅	7.2e-4
4-8-4	8	Z3x4 → Z3x8 → Z3x4	✅ / ✅	✅ / ✅	1.0e-3

• Z3x4 for 60 steps gives bitwise-equal losses in a 4-process and an 8-process world, and when it reshards 4→4→4 through AffineIR.
• Negative control: leaving Adam states at zero fails the state check at both switches and diverges by up to 0.39 relative loss.
• Sabotage: swapping the pieces of ranks 0 and 1 in one source map fails planning with InvalidMapError: source: rank 1 does not hold all of Box((255750,) ...) before any data moves.

2. ZeRO stage (z3x.py)

run	world	phases	every switch: state / logits bitwise	affine vs gather: losses / final logits bitwise	vs a run that keeps the first phase's config: max rel loss diff
Z3↔Z2	4	Z3x4 → Z2x4 → Z3x4	✅ / ✅	✅ / ✅	1.2e-3
Z3↔Z1	4	Z3x4 → Z1x4 → Z3x4	✅ / ✅	✅ / ✅	1.0e-3
stage + DP	4	Z3x4 → Z2x2 → Z1x4	✅ / ✅	✅ / ✅	7.2e-4
stage + DP	8	Z2x8 → Z3x4 → Z1x8	✅ / ✅	✅ / ✅	9.8e-4

• Correction to the comment above: it compared every run to a ZeRO-3-only run. That does not work for the run that starts at Z2x8. ZeRO-1/2 and ZeRO-3 (zero.Init) draw the initial weights differently, so that run differs from Z3x4 already at step 0, before any switch (0.40 relative). The column above therefore compares each run to a reference with the same first phase (Z3x4, or Z2x8 → Z2x8 → Z2x8 for the last row).

3. AutoTP size (tpep.py)

run	world	phases	every switch: state bitwise / logits max abs diff	affine vs gather: losses / final logits bitwise	vs a run that keeps the first phase's config: max rel loss diff
Z1	4	Z1TP2 → Z1TP4 → Z1TP2	✅ / 4.9e-4	✅ / ✅	9.6e-4
Z2	4	Z2TP2 → Z2TP4 → Z2TP2	✅ / 4.9e-4	✅ / ✅	9.6e-4
Z3	4	Z3TP2 → Z3TP4 → Z3TP2	✅ / 2.0e-3	✅ / ✅	1.4e-3
stage + TP	4	Z1TP2 → Z3TP4 → Z2TP2	✅ / 4.9e-4	✅ / ✅	1.4e-3
TP up to 8	8	Z2TP2 → Z2TP8 → Z2TP4	✅ / 2.0e-3	✅ / ✅	1.9e-3
stage + TP	8	Z3TP4 → Z1TP8 → Z2TP2	✅ / 2.0e-3	✅ / ✅	1.3e-3

• The reference runs (e.g. Z1TP2 → Z1TP2 → Z1TP2) reshard at every switch and keep bitwise-equal logits.
• Runs without a switch that differ only in TP size (TP2 vs TP4) differ by up to 3.2e-3, more than any switched run differs from its reference.

4. AutoEP size (tpep.py)

run	world	phases	every switch: state / logits bitwise	affine vs gather: losses / final logits bitwise	vs a run that keeps the first phase's config: max rel loss diff
Z2	4	Z2EP2 → Z2EP4 → Z2EP2	✅ / ✅	✅ / ✅	1.1e-2
Z3	4	Z3EP2 → Z3EP4 → Z3EP2	✅ / ❌ at → Z3EP4	✅ / ✅	1.1e-2
stage + EP	4	Z2EP2 → Z3EP4 → Z1EP2	✅ / ❌ at → Z3EP4	✅ / ✅	1.5e-2
EP up to 8	8	Z2EP2 → Z2EP8 → Z2EP4	✅ / ✅	✅ / ✅	9.0e-3
stage + EP	8	Z3EP4 → Z1EP8 → Z2EP2	✅ / ✅	✅ / ✅	1.1e-2

• The composed expert parameter maps equal the shipped expert_storage_affine_maps() in all 24 checks.
• Both ❌ are at the switch into Z3EP4 with world 4. That is the case traced in the comment above: the expert-data-parallel group has one rank and ZeRO-3 places the packed expert weights at an odd element offset, which changes the GEMM result.
  • The gather oracle gives bitwise the same logits at those switches, and the state is bitwise equal.
  • It is data-dependent: this time the Z3 run also hits it, while the diagnostic run Z2EP2 → Z3EP2 → Z3EP4 keeps bitwise logits at both switches. The earlier run showed the opposite. I did not re-trace the first divergent element here.
• Runs without a switch that differ only in EP size (EP2 vs EP4) differ by up to 1.8e-2.
• [BUG] ZeRO stage 1/2 marks optimizer-state fragments linked before later groups initialize DeepSpeed#8911 is still open, so tpep.py keeps its re-link workaround.

5. ZeRO-3 with aliased masters (#8842, z3_alias.py)

Since #8842, when the weight partition already has the master dtype (fp32 weights, or bf16 with bf16_master_weights_and_grads) and the optimizer is not offloaded, ZeRO-3 uses the weight partition as the master. zero_optimizer_storage_shard(p) and zero_param_storage_shard(p) are then views of the same memory.

z3_alias.py is z3x.py's ZeRO-3 DP protocol with three precision modes. It checks the alias by data_ptr and moves each storage kind in its own dtype. The oracle reads full weights through GatheredParameters and full masters and Adam states through safe_get_full_*.

run	precision	phases	masters aliased (src / dst)	weights + master in one batch	every switch: weights + master + Adam bitwise, no NaN
ref	bf16 weights, fp32 master	Z3x4 → Z3x2 → Z3x4	0/6 / 0/6	n/a (different dtypes)	✅
bf16m	bf16 weights, bf16_master_weights_and_grads + bf16_optimizer_states	Z3x4 → Z3x2 → Z3x4	6/6 / 6/6	rejected	✅
bf16m, master skipped	same	Z3x4 → Z3x2 → Z3x4	6/6 / 6/6	–	✅
fp32	fp32 weights and master	Z3x4 → Z3x2 → Z3x4	6/6 / 6/6	rejected	✅
fp32, master skipped	same	Z3x4 → Z3x2 → Z3x4	6/6 / 6/6	–	✅
fp32	same	Z3x4 → Z3x4 → Z3x4	6/6 / 6/6	–	✅

• The maps need no change. With aliasing, the master's storage map is the weight's map. Moving the master separately, or skipping it because it is the same memory, both give a bitwise-exact state.
• Both in one batch is refused, not silently doubled. Sending the weights and the master to the same target in one execute_transfers call is rejected on every rank with two targets overlap within the transfer batch, before any data moves. A caller that wants one batch must send one of them, e.g. deduplicating by data_ptr.
• dtypes follow the config. Without offload, bf16_master_weights_and_grads requires bf16_optimizer_states, so the master and both Adam states are bf16.
• Not covered: offload_states / reload_states, which re-point the aliased master at the moved weight buffer, and ZeRO-1/2, where #8842 does not apply.

Reproduction

• torchrun --nproc_per_node=<4|8> z3x.py <name> <phases> 20 <affine|gather|noopt> and tpep.py <name> <phases> 20 <affine|gather>, as in the comment above.
• torchrun --nproc_per_node=4 z3_alias.py <name> <dp sizes, e.g. 4-2-4> 20 <ref|bf16m|fp32> [dedup] [onebatch].

z3_alias.py (section 5)

# ZeRO-3 resharding across DP sizes with AffineIR when the master copy may alias the weight partition (#8842).
# usage: torchrun --nproc_per_node=4 z3_alias.py <name> <dp sizes, e.g. 4-2-4> <steps per phase> <ref|bf16m|fp32>
#        [dedup] [onebatch]
#   ref:   bf16 weights, fp32 master (no alias)
#   bf16m: bf16 weights, bf16_master_weights_and_grads + bf16_optimizer_states (master aliases the weights)
#   fp32:  fp32 weights and master (master aliases the weights)
#   dedup: do not move the master when it aliases the weight partition
#   onebatch: first try weights and master in one batch, which must be rejected when they alias
import json
import math
import sys

import torch

import deepspeed
import deepspeed.comm as dist
from deepspeed.accelerator import get_accelerator
from deepspeed.checkpoint.affine_transfer import plan_transfer
from deepspeed.checkpoint.affine_transfer_executor import TransferEndpoints, TransferRequest, execute_transfers
from deepspeed.checkpoint.zero_affine import (zero_optimizer_storage_shard, zero_param_storage_shard,
                                              zero_storage_affine_map)
from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_optimizer_state

NAME, PHASES, STEPS = sys.argv[1], [int(x) for x in sys.argv[2].split('-')], int(sys.argv[3])
MODE = sys.argv[4]
assert MODE in ('ref', 'bf16m', 'fp32')
DEDUP = 'dedup' in sys.argv[5:]
ONEBATCH = 'onebatch' in sys.argv[5:]
MOVE_OPT = True
IN_DTYPE = torch.float32 if MODE == 'fp32' else torch.bfloat16
GLOBAL_BATCH, DIN, HID = 16, 1000, 1023  # odd sizes so ZeRO-3 partitions are padded
STATE_KEYS = ('exp_avg', 'exp_avg_sq')

deepspeed.init_distributed()
rank, world = dist.get_rank(), dist.get_world_size()
if 'DUMP_AFTER' in __import__('os').environ:
    import faulthandler
    import signal
    faulthandler.register(signal.SIGUSR1, all_threads=False)
get_accelerator().set_device(rank)
dev = get_accelerator().current_device_name()
singletons = [dist.new_group([r]) for r in range(world)]
dp_groups = {dp: dist.new_group(list(range(dp))) for dp in sorted(set(PHASES))}


class MPU:

    def __init__(self, dp):
        self.dp, self.group = dp, dp_groups[dp]

    def get_data_parallel_group(self):
        return self.group

    def get_data_parallel_world_size(self):
        return self.dp

    def get_data_parallel_rank(self):
        return dist.get_rank(group=self.group)

    def get_data_parallel_group_ranks(self):
        return list(range(self.dp))

    def get_model_parallel_group(self):
        return singletons[rank]

    def get_model_parallel_world_size(self):
        return 1

    def get_model_parallel_rank(self):
        return 0

    def get_model_parallel_src_rank(self):
        return rank


class Net(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(DIN, HID)
        self.l2 = torch.nn.Linear(HID, HID)
        self.l3 = torch.nn.Linear(HID, 1)

    def forward(self, x, y):
        h = torch.nn.functional.gelu(self.l2(torch.nn.functional.gelu(self.l1(x))))
        return torch.nn.functional.mse_loss(self.l3(h).float(), y)


teacher = torch.randn(DIN, generator=torch.Generator().manual_seed(7)) / math.sqrt(DIN)


def batch(step, dp, dp_rank):
    g = torch.Generator().manual_seed(1000 + step)
    x = torch.randn(GLOBAL_BATCH, DIN, generator=g)
    y = torch.sin(3 * x @ teacher).unsqueeze(1)
    micro = GLOBAL_BATCH // dp
    s = slice(dp_rank * micro, (dp_rank + 1) * micro)
    return x[s].to(dev, IN_DTYPE), y[s].to(dev)


def build(dp):
    import os
    import traceback
    if os.environ.get('TRACE_WORLD') and dp < world:
        import torch.distributed as tdist
        for fn in ('barrier', 'all_reduce', 'broadcast', 'all_gather', 'all_gather_into_tensor', 'all_gather_object',
                   'broadcast_object_list', 'reduce_scatter_tensor'):
            orig = getattr(tdist, fn)

            def wrapped(*a, _orig=orig, _fn=fn, **k):
                g = k.get('group')
                if g is None or g is tdist.group.WORLD or tdist.get_world_size(g) == world:
                    print(f'[rank{rank}] WORLD {_fn}\n' + ''.join(traceback.format_stack(limit=8)[:-1]), flush=True)
                return _orig(*a, **k)

            setattr(tdist, fn, wrapped)
    cfg = {
        'train_micro_batch_size_per_gpu': GLOBAL_BATCH // dp,
        'gradient_accumulation_steps': 1,
        'optimizer': {
            'type': 'Adam',
            'params': {
                'lr': 1e-3
            }
        },
        **({} if MODE == 'fp32' else {
            'bf16': {
                'enabled': True,
                **({
                    'bf16_master_weights_and_grads': True,
                    'bf16_optimizer_states': True
                } if MODE == 'bf16m' else {})
            }
        }),
        'zero_optimization': {
            'stage': 3
        },
    }
    torch.manual_seed(42)
    # Stage3ZeroOptimizer.__init__ barriers on the world group; scope those to the engine's own group.
    world_barrier = dist.barrier
    if dp < world:
        dist.barrier = lambda group=None, **k: world_barrier(group=group or dp_groups[dp], **k)
    if min(PHASES) < world:
        # The first-trace sanity check broadcasts on the world group, which idle ranks never join.
        zero_utils = sys.modules['deepspeed.runtime.zero.utils']
        zero_utils.assert_lst_len_same_as_other_ranks = lambda lst: None
        zero_utils.get_lst_from_rank0 = lambda lst: list(lst)
    with deepspeed.zero.Init(data_parallel_group=dp_groups[dp], config_dict_or_path=cfg):
        model = Net()
    engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=cfg, mpu=MPU(dp))
    dist.barrier = world_barrier
    return engine


def train(engine, dp, first, last, losses):
    dp_rank = dist.get_rank(group=dp_groups[dp])
    for step in range(first, last):
        loss = engine(*batch(step, dp, dp_rank))
        engine.backward(loss)
        engine.step()
        mean = loss.detach().float().clone()
        dist.all_reduce(mean, group=dp_groups[dp])
        losses.append((mean / dp).item())


def meta(engine):
    """name -> (numel, ZeRO-3 partition numel), the same on every rank."""
    out = {name: (p.ds_numel, p.ds_tensor.ds_numel) for name, p in engine.module.named_parameters()} if engine else None
    holder = [out]
    dist.broadcast_object_list(holder, src=0)
    return holder[0]


def storage_info(engine):
    """Dtype of every storage kind and how many parameters have a master that aliases the weight partition."""
    out = None
    if engine:
        params = list(engine.module.parameters())
        p = params[0]
        dtypes = {'param': zero_param_storage_shard(p).dtype}
        dtypes.update({k: zero_optimizer_storage_shard(p, k).dtype for k in (None, ) + STATE_KEYS})
        aliased = sum(
            zero_optimizer_storage_shard(q).data_ptr() == zero_param_storage_shard(q).data_ptr() for q in params)
        out = (dtypes, aliased, len(params))
    holder = [out]
    dist.broadcast_object_list(holder, src=0)
    return holder[0]


def adam_step(engine):
    steps = {s['step'] for s in engine.optimizer.optimizer.state.values()} if engine else None
    holder = [steps]
    dist.broadcast_object_list(holder, src=0)
    (step, ) = holder[0]
    return step


def fresh(dp):
    """A new engine whose every shard is NaN, with Adam states created by a no-op step."""
    if rank >= dp:
        dist.barrier()
        return None
    engine = build(dp)
    print(f'[rank{rank}] fresh: built', flush=True)
    loss = engine(*batch(0, dp, rank)) * 0
    print(f'[rank{rank}] fresh: forward', flush=True)
    engine.backward(loss)
    print(f'[rank{rank}] fresh: backward', flush=True)
    engine.step()
    print(f'[rank{rank}] fresh: step', flush=True)
    engine.empty_partition_cache()
    with torch.no_grad():
        for p in engine.module.parameters():
            zero_param_storage_shard(p).fill_(float('nan'))
            for key in (None, ) + STATE_KEYS:
                zero_optimizer_storage_shard(p, key).fill_(float('nan'))
    dist.barrier()
    return engine


def reshard(src, src_dp, dst, dst_dp):
    """Move every parameter shard, fp32 master and Adam state from src (DP src_dp) to dst (DP dst_dp)."""
    src_meta, dst_meta = meta(src), meta(dst)
    assert src_meta.keys() == dst_meta.keys()
    for engine in (src, dst):
        if engine:
            engine.empty_partition_cache()
    src_params = dict(src.module.named_parameters()) if src else {}
    dst_params = dict(dst.module.named_parameters()) if dst else {}
    identity = {r: r for r in range(world)}
    src_dtypes, src_aliased, total = storage_info(src)
    dst_dtypes, dst_aliased, _ = storage_info(dst)
    assert src_dtypes == dst_dtypes
    notes = [f'aliased masters src {src_aliased}/{total} dst {dst_aliased}/{total}, dtypes {src_dtypes}']
    aliased = src_aliased == total and dst_aliased == total
    kinds = [(k, src_dtypes[k]) for k in ('param', None) + STATE_KEYS]

    def requests_for(kind):

        def shard(p):
            return zero_param_storage_shard(p) if kind == 'param' else zero_optimizer_storage_shard(p, kind)

        requests = []
        for name in sorted(src_meta):
            numel, src_part = src_meta[name]
            _, dst_part = dst_meta[name]
            plan = plan_transfer(zero_storage_affine_map(numel, list(range(dst_dp)), dst_part),
                                 zero_storage_affine_map(numel, list(range(src_dp)), src_part))
            src_buf = {rank: shard(src_params[name])} if src else {}
            dst_buf = {rank: shard(dst_params[name])} if dst else {}
            requests.append(TransferRequest(f'{name}:{kind}', plan, src_buf, dst_buf))
        return requests

    if ONEBATCH and src_dtypes['param'] == src_dtypes[None]:
        try:
            execute_transfers(requests_for('param') + requests_for(None), TransferEndpoints(identity, identity),
                              src_dtypes['param'])
            notes.append('weights+master in one batch: accepted')
        except Exception as caught:
            notes.append(f'weights+master in one batch: {type(caught).__name__}: {str(caught)[:160]}')
    for kind, dtype in kinds:
        if kind is None and DEDUP and aliased:
            notes.append('master skipped (aliases the weights)')
            continue
        execute_transfers(requests_for(kind), TransferEndpoints(identity, identity), dtype)
    if not MOVE_OPT:
        with torch.no_grad():
            for p in dst.module.parameters() if dst else []:
                for key in STATE_KEYS:
                    zero_optimizer_storage_shard(p, key).zero_()
    step = adam_step(src) if MOVE_OPT else 0
    if dst:
        for state in dst.optimizer.optimizer.state.values():
            state['step'] = step
    return notes


def snapshot(engine):
    """Full weights, master and Adam states on rank 0 of the engine's group."""
    if engine is None:
        return None
    engine.empty_partition_cache()
    out = {}
    for name, p in engine.module.named_parameters():
        with deepspeed.zero.GatheredParameters([p]):
            weight = p.data.clone()
        out[name] = [weight, safe_get_full_fp32_param(p)] + [safe_get_full_optimizer_state(p, k) for k in STATE_KEYS]
    return {n: [t.cpu() for t in ts] for n, ts in out.items()} if rank == 0 else None


losses, report = [], []
dp = PHASES[0]
engine = build(dp) if rank < dp else None
for index, next_dp in enumerate(PHASES):
    if index:
        before = snapshot(engine)
        target = fresh(next_dp)
        notes = reshard(engine, dp, target, next_dp)
        after = snapshot(target)
        if rank == 0:
            exact = all(torch.equal(a, b) for n in before for a, b in zip(before[n], after[n]))
            nans = any(torch.isnan(t).any().item() for ts in after.values() for t in ts)
            report.append(f'switch DP{dp}->DP{next_dp} at step {index * STEPS}: '
                          f'weights+master+adam bitwise equal {exact}, NaN left {nans}')
            report.extend('  ' + note for note in notes)
        if engine:
            engine.destroy()
        engine, dp = target, next_dp
    if engine:
        train(engine, dp, index * STEPS, (index + 1) * STEPS, losses)
    dist.barrier()

if rank == 0:
    json.dump(losses, open(f'/tmp/z3a_{NAME}.json', 'w'))
    print(f'### {NAME}: mode {MODE}, phases {PHASES} x {STEPS} steps, dedup {DEDUP}, onebatch {ONEBATCH}')
    for line in report:
        print(line)
    print('losses', ' '.join(f'{x:.5f}' for x in losses))
dist.barrier()