[LiRunGuo]

Status: draft, part of RFC #9026. The RFC has not been accepted by maintainers yet. This PR stays a draft until it is, so that the design can be discussed on code.

• Per CONTRIBUTING, a tutorial and a DeepSpeedExamples PR will follow once the design is agreed.
• This is the first step of the bottom-up plan agreed with @Achyuthan-S: declare a sync from the storage maps plus a target map of the inference copy, and plan it with plan_transfer.
• [checkpoint] Batch replica selection for affine transfer plans #8936's batch policy and other transports come later.

What it is. A building block for rollout engines that keep their own tensor-parallel copy of a ZeRO model. No rollout engine is wired to it in this PR. It is validated on synthetic modules in the unit tests and on Qwen3 checkpoints in the benchmark below. Wiring it into an engine (for example as the sync of #9023's AutoTPRollout) is a follow-up.

What it adds. A new module, deepspeed/runtime/rollout/affine_weight_sync.py, and a section in docs/code-docs/source/affine-transfer.rst. No existing source file changes.

• AffineWeightSync(source_maps, target_maps, dtype, group=None, budget=StagingBudget(), limits=None):
  • plans every parameter once with plan_transfer, without communication;
  • sync(source_buffers, target_buffers) then runs the copies of all parameters as one execute_transfers batch ([checkpoint] Execute affine transfer plans with bounded peer staging #8809), so no rank gathers a whole parameter;
  • a rank that lacks one of its shards is reported through the executor's preflight_error, so every rank refuses the batch before any data moves, rather than one rank raising while the others wait.
• Replica choice. A block that several ranks hold is read from the target's own process when that process holds it. Otherwise it is read from the holder that sends the fewest elements so far; ties go to the lowest rank.
  • As in [checkpoint] Batch replica selection for affine transfer plans #8936's local_first_least_loaded, local copies do not count as load, and single-holder regions do.
  • One difference remains: plan_transfer consults a selector only for regions with several holders, so the single-holder load of a parameter is counted once that parameter is planned.
  • I'll switch to [checkpoint] Batch replica selection for affine transfer plans #8936's plan_transfers with its policy once it lands.
• zero_param_shaped_storage_map(param): zero_param_storage_affine_map nested into the parameter's own shape, as autotp_storage_affine_maps does for AutoTP models. A tensor-parallel layout describes a parameter in its own shape, and both sides of a plan must describe the same tensor.
• replicated_layout_map(layout, replicas): one replica's tensor-parallel layout, keyed by TP rank (an AutoTP layout map, or contiguous_split_map), held by every replica and keyed by global rank. It is built with nest_affine_maps and zero_storage_affine_map, as test_zero_affine.py's _one_holder_each builds one holder per layout shard.
• WeightArena: a rank's shards as views of one buffer that is never reallocated.
  • Each view starts at a 256-byte-aligned address; the buffer is started at an aligned element of a slightly larger allocation.
  • The executor writes targets in place, so a CUDA graph captured on the arena reads the new weights after every sync.

Why not a duplicate

• feat(rollout): add AutoTP rollout engine with in-memory weight sync #9023 (AutoTPRollout) refreshes its inference copy by gathering each full fp32 master weight with safe_get_full_fp32_param and cutting out the local shard. That moves every parameter to every rank. This PR moves each element only to the ranks that need it.
• [checkpoint] Execute affine transfer plans with bounded peer staging #8809 (executor), [AffineIR] Affine storage maps for live ZeRO parameters and optimizer states #8885 (storage maps) and [checkpoint] Batch replica selection for affine transfer plans #8936 (batch replica policy) are the pieces this builds on. It changes none of them.
• The relay, prepared transfers and pipelining ([RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026 §5, @0z5a) are not touched.

What it does not do (yet)

• Holders of one parameter in two layouts, such as the trainer's ZeRO-3 partitions plus a surviving TP shard: plan_transfer refuses replicas with different boundaries.
• Partial targets, such as refilling one lost TP rank.
• Relay across nodes. That waits for relay_map ([RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026 §1).
• Fused targets (q/k/v into one qkv tensor), which work by passing views of the fused tensor as the targets of the separate weights, are not tested here.

Tests (tests/unit/runtime/rollout/test_affine_weight_sync.py)

• Live engines, 4 ranks: ZeRO stages 1, 2 and 3 into two TP=2 replicas, one TP=4 replica (uneven shards), and TP=1 replicas on half of the ranks.
  • Every rank's shard equals layout.extract of the gathered parameter after two optimizer steps.
  • The arena keeps its address.
  • Under ZeRO-1/2 no byte crosses processes.
• Load:
  • reads of a parameter held by three ranks spread evenly over them;
  • with a single-holder parameter and a replicated one, local copies cost nothing and single-holder load counts, so each sender sends one parameter.
• Missing shard: when one rank lacks a shard, every rank raises AffineTransferExecutionError and no target changes.
• CUDA graphs: a graph captured before a sync reads the synced weights.
• CPU-only checks:
  • arena views of eight arenas alive at once are 256-byte aligned by address, disjoint and zeroed;
  • replicated_layout_map gives each replica its layout shard and refuses replicas that do not fit.

Results: 4× RTX A6000 (driver 595.58.03), torch 2.13.0+cu130, CUDA 13.0, Python 3.12. Commands run from tests/:

pytest unit/runtime/rollout/test_affine_weight_sync.py                     # 18 passed
CUDA_VISIBLE_DEVICES= DS_ACCELERATOR=cpu LOCAL_SIZE=4 \
  pytest unit/runtime/rollout/test_affine_weight_sync.py                   # gloo, 4 ranks: 17 passed, 1 skipped (the CUDA-graph test)
pytest unit/runtime/rollout                                                # 95 passed, 1 skipped
pytest unit/checkpoint/test_affine_shard_map.py unit/checkpoint/test_affine_transfer_plan.py \
  unit/checkpoint/test_affine_transfer_executor.py \
  unit/checkpoint/test_affine_transfer_executor_external_group.py \
  unit/checkpoint/test_zero_affine.py                                      # 285 passed, 1 skipped

pre-commit passes on the changed files.

Measurements

• Setup: 4× RTX 4090 on one node (PCIe, no NVLink; driver 595.58.03), bf16, ZeRO-3 over the 4 ranks (stage3_param_persistence_threshold: 0, no optimizer). Each number is the median of 5 syncs after one warm-up.
• Baselines:
  • "gather" is GatheredParameters on each parameter, then a copy of this rank's slice into an arena. That is feat(rollout): add AutoTP rollout engine with in-memory weight sync #9023's approach, but on the bf16 parameter rather than the fp32 master, so it moves half the bytes it would there.
  • "affine" is this PR, at a staging budget of 8 MiB (the executor's default) and of 64 MiB.
• Correctness: after the timed syncs, the two arenas were compared element by element, and every weight was identical.
• Peak memory: "peak" is the largest max_memory_allocated increase on any rank during a sync.

Model	Inference layout	gather	affine, 8 MiB	affine, 64 MiB	Planning (once)
Qwen3-1.7B	TP2 × 2 replicas	0.193 s, +602 MiB	0.146 s, +44 MiB	0.143 s, +380 MiB	0.11 s
	TP4	0.195 s, +602 MiB	0.042 s, +45 MiB	0.042 s, +336 MiB	0.11 s
	TP1 × 4 replicas	0.191 s, +602 MiB	0.209 s, +42 MiB	0.192 s, +380 MiB	0.12 s
	TP2 on 2 of the 4 ranks	0.195 s, +602 MiB	0.105 s, +29 MiB	0.104 s, +253 MiB	0.09 s
Qwen3-8B	TP4	0.778 s, +1220 MiB	0.107 s, +36 MiB	0.096 s, +372 MiB	0.14 s

• Where the gain comes from: it grows as each rank needs a smaller part of the model.
  • With TP1 every rank needs the whole model, and both methods move about the same bytes (9.6 GiB for Qwen3-1.7B).
  • Gather's peak is the largest parameter gathered whole: the embedding, 594 MiB for 1.7B and 1187 MiB for 8B.
  • The affine sync's peak is its staging, bounded by the budget.
• Not measured yet: NVLink and InfiniBand. The two-node comparison is the next step of [RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026 (open question 5).

To reproduce, run the script below with torchrun --nproc_per_node 4 pr_weight_sync_bench.py --model <Qwen3-1.7B or Qwen3-8B> --layouts 0-3:2 0-3:4 0-3:1 2-3:2 --out <dir>. It needs transformers. Layouts are <ranks>:<tp>, so 0-3:2 is two TP=2 replicas on ranks 0–3; Qwen3-8B was run with 0-3:4 only, since two copies at TP1 or TP2 do not fit in 24 GB.

pr_weight_sync_bench.py

#!/usr/bin/env python
"""Draft PR for RFC #9026 (2026-11-08): refreshing a tensor-parallel inference copy of a ZeRO-3 model.

Times ``AffineWeightSync`` (one plan per parameter from the storage maps, every sync one batch of the #8809 executor)
against gathering each parameter whole and copying this rank's slice, which is what #9023's AutoTPRollout does (there
on the fp32 master, here on the bf16 parameter, so the gather moves half the bytes it would there). Inference layouts
are given as <ranks>:<tp> (e.g. 0-3:2 = two replicas of TP=2 on ranks 0-3); every rank holds a ZeRO-3 partition.
The two copies are compared element for element after the timed syncs. Imports DeepSpeed from PYTHONPATH (the PR
tree), run with torchrun:
  torchrun --nproc_per_node 4 bench/native/pr_weight_sync_bench.py --model $MODELS/Qwen3-1.7B \
    --layouts 0-3:2 0-3:4 0-3:1 --out <dir>
"""

import argparse
import datetime
import json
import os
import statistics
import time
from pathlib import Path

import torch

import deepspeed
import deepspeed.comm as dist
from deepspeed.checkpoint.affine import contiguous_split_map, replicated_map
from deepspeed.checkpoint.affine_transfer_executor import StagingBudget
from deepspeed.checkpoint.zero_affine import zero_param_storage_shard
from deepspeed.runtime.rollout.affine_weight_sync import (AffineWeightSync, WeightArena, replicated_layout_map,
                                                          zero_param_shaped_storage_map)

# Megatron-style tensor parallelism for Llama / Qwen-like models: these weights are split along their rows (output
# features, vocabulary) and these along their columns (input features); everything else (norms) is replicated.
ROWS = ("q_proj.weight", "k_proj.weight", "v_proj.weight", "gate_proj.weight", "up_proj.weight",
        "embed_tokens.weight", "lm_head.weight")
COLUMNS = ("o_proj.weight", "down_proj.weight")


def even(size, parts):
    base, extra = divmod(size, parts)
    return [base + (part < extra) for part in range(parts)]


def layout_of(name, shape, tp):
    """(layout map, split axis or None)"""
    if name.endswith(ROWS):
        return contiguous_split_map(shape, even(shape[0], tp), 0), 0
    if name.endswith(COLUMNS):
        return contiguous_split_map(shape, even(shape[1], tp), 1), 1
    return replicated_map(shape, tp), None


def parse(spec):
    ranks, tp = spec.split(":")
    lo, hi = (int(x) for x in ranks.split("-"))
    ranks, tp = list(range(lo, hi + 1)), int(tp)
    assert len(ranks) % tp == 0, spec
    return [ranks[i:i + tp] for i in range(0, len(ranks), tp)]


def now():
    torch.cuda.synchronize()
    dist.barrier()
    return time.perf_counter()


def everyone(value):
    out = [None] * dist.get_world_size()
    dist.all_gather_object(out, value)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--layouts", nargs="+", required=True)
    ap.add_argument("--repeats", type=int, default=5)
    ap.add_argument("--budget-mib", type=int, nargs="+", default=[8, 64], help="StagingBudget per peer, MiB")
    ap.add_argument("--out", required=True)
    args = ap.parse_args()

    deepspeed.init_distributed(dist_backend="nccl", timeout=datetime.timedelta(minutes=30))
    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
    rank, world = dist.get_rank(), dist.get_world_size()
    device = torch.device("cuda", torch.cuda.current_device())
    from transformers import AutoModelForCausalLM
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16)
    config = {
        "train_micro_batch_size_per_gpu": 1,
        "bf16": {
            "enabled": True
        },
        "zero_optimization": {
            "stage": 3,
            "stage3_param_persistence_threshold": 0
        },
    }
    engine, *_ = deepspeed.initialize(model=model, config=config)
    engine.empty_partition_cache()
    params = dict(engine.module.named_parameters())
    shapes = {name: tuple(param.ds_shape) for name, param in params.items()}
    numel = sum(torch.Size(shape).numel() for shape in shapes.values())
    source_maps = {name: zero_param_shaped_storage_map(param) for name, param in params.items()}
    sources = {name: zero_param_storage_shard(param) for name, param in params.items()}
    res = {
        "model": Path(args.model).name,
        "gpu": torch.cuda.get_device_name(),
        "world": world,
        "nodes": len(set(everyone(os.uname().nodename))),
        "torch": torch.__version__,
        "params": len(params),
        "bytes": 2 * numel,
        "layouts": {},
    }
    if rank == 0:
        print(f"[bench] {res['model']}: {len(params)} parameters, {2 * numel / 2**30:.2f} GiB bf16, {world} ranks on "
              f"{res['nodes']} node(s), {res['gpu']}")

    for spec in args.layouts:
        replicas = parse(spec)
        tp = len(replicas[0])
        tp_rank = next((ranks.index(rank) for ranks in replicas if rank in ranks), None)
        layouts = {name: layout_of(name, shape, tp) for name, shape in shapes.items()}
        target_maps = {name: replicated_layout_map(layouts[name][0], replicas) for name in params}
        arenas = {}
        if tp_rank is not None:
            local = {name: layouts[name][0].shard_shapes[tp_rank] for name in params}
            arenas = {method: WeightArena(local, torch.bfloat16, device) for method in ("affine", "gather")}
        out = {"replicas": replicas, "affine": {}}

        def gather():
            for name, param in params.items():
                with deepspeed.zero.GatheredParameters([param]):
                    if tp_rank is None:
                        continue
                    layout, axis = layouts[name]
                    full = param.data
                    if axis is not None:
                        start = sum(even(full.shape[axis], tp)[:tp_rank])
                        full = full.narrow(axis, start, layout.shard_shapes[tp_rank][axis])
                    arenas["gather"][name].copy_(full)

        def timed(run):
            times, peaks, stats = [], [], None
            for repeat in range(args.repeats + 1):
                torch.cuda.reset_peak_memory_stats()
                base = torch.cuda.memory_allocated()
                t = now()
                stats = run()
                seconds = now() - t
                peaks.append(torch.cuda.max_memory_allocated() - base)
                if repeat:  # the first call warms up
                    times.append(seconds)
            return {
                "median_s": statistics.median(times),
                "times_s": times,
                "peak_extra_mib": max(everyone(max(peaks))) / 2**20,
            }, stats

        out["gather"], _ = timed(gather)
        for mib in args.budget_mib:
            t = now()
            sync = AffineWeightSync(source_maps,
                                    target_maps,
                                    torch.bfloat16,
                                    budget=StagingBudget(bytes_per_peer=mib << 20))
            plan_s = now() - t
            targets = dict(arenas["affine"].items()) if arenas else {}
            result, stats = timed(lambda: sync.sync(sources, targets))
            result["plan_s"] = plan_s
            result["segments"] = sum(len(plan.segments) for plan in sync.plans.values())
            for key in ("rounds", "chunks"):
                result[key] = max(everyone(getattr(stats, key)))
            for key in ("sent_bytes", "received_bytes", "local_copy_bytes"):
                result[key] = sum(everyone(getattr(stats, key)))
            out["affine"][f"{mib}MiB"] = result
            del sync

        bad = [name for name in params if arenas and not torch.equal(arenas["affine"][name], arenas["gather"][name])]
        out["mismatched"] = sum(len(b) for b in everyone(bad))
        res["layouts"][spec] = out
        if rank == 0:
            best = min(out["affine"].values(), key=lambda r: r["median_s"])
            print(f"[bench] {spec} (TP{tp} x{len(replicas)}): gather {out['gather']['median_s']:.3f} s "
                  f"(+{out['gather']['peak_extra_mib']:.0f} MiB); affine " + ", ".join(
                      f"{k} {r['median_s']:.3f} s (+{r['peak_extra_mib']:.0f} MiB, plan {r['plan_s']:.2f} s, "
                      f"{r['rounds']} rounds, {r['sent_bytes'] / 2**30:.2f} GiB sent)" for k, r in out["affine"].items()) +
                  f"; best/gather {best['median_s'] / out['gather']['median_s']:.2f}; mismatched weights "
                  f"{out['mismatched']}",
                  flush=True)
        del arenas
        torch.cuda.empty_cache()

    if rank == 0:
        Path(args.out).mkdir(parents=True, exist_ok=True)
        (Path(args.out) / "result.json").write_text(json.dumps(res, indent=1))
    dist.barrier()


if __name__ == "__main__":
    main()

Compatibility

• Validated:
  • ZeRO-1/2/3, bf16, on GPU (unit tests);
  • ZeRO-3, bf16, with Qwen3-1.7B and Qwen3-8B (benchmark);
  • ZeRO-1/2/3, bf16, on CPU (gloo, 4 ranks).
• Not validated:
  • offload (ZeRO-3 parameter offload keeps the partition on the host, while the executor needs device buffers);
  • quantized and NVMe storage (refused by the storage maps);
  • an AutoTP trainer (autotp_storage_affine_maps as the source maps);
  • multi-node.

AI assistance: AI assistance was used. The code, tests and this description were drafted with an AI coding agent (Claude Code). The submitter reviewed every changed line and ran the tests above.

Maintenance contact: @LiRunGuo. cc @Achyuthan-S @0z5a @jinyouzhi

[nathon-lee]

Hi， @LiRunGuo
Thanks for the detailed write-up — reading this was a great way to understand how the affine building blocks (plan_transfer + execute_transfers) compose into a real weight-sync path. The shard-to-shard approach and the in-place WeightArena for CUDA-graph validity are really neat.

A couple of questions, mostly to learn:

In _LocalFirstLeastLoaded, the inline load increment in call is later overwritten by settle(), so it only balances multiple multi-holder regions within a single parameter. Is that the intended scope, and did you consider balancing across parameters during planning too?
For the group path, the maps are keyed by global rank while the executor can use group.rank() — is the expectation that callers always key their source/target maps by global rank even for a custom weight-update group? A one-line note in the AffineWeightSync docstring might help future readers.
Looking forward to the follow-ups once #8936 lands. Thanks again!

[Achyuthan-S]

Thanks @LiRunGuo, this is a clean first step, and it follows the RFC closely: the sync is declared entirely in maps, planned by plan_transfer and run by #8809, without changing any of them.

I ran it on CPU (gloo, 4 ranks), merged with current master: 17 passed and 1 skipped (the CUDA-graph test), and flake8 is clean. Turning off the local-first branch of the selector makes 3 of the tests fail, so the load tests do pin the policy.

A few thoughts, none blocking for a draft:

1. Why the selector can't change the result. This is Theorem 4 of the correctness note: once both maps pass plan_transfer's checks and the source replicas agree (P2), every plan writes the same bits, whichever replica it reads. Under ZeRO-1/2 the replicas are the data-parallel copies of each bf16 parameter, which agree after the step's all-gather, and under ZeRO-3 each element has one holder. A sentence in the module docstring would tell reviewers that the selector only affects cost.
2. AutoTP trainers. The module docstring and the docs name autotp_storage_affine_maps as the source for an AutoTP model, while Compatibility lists an AutoTP trainer as not validated. A test (an AutoTP TP2 trainer into a TP4 copy, say) or a note that it's untested would keep the two consistent. Happy to help with that test.
3. Where the two map helpers live. zero_param_shaped_storage_map and replicated_layout_map don't depend on rollout. Would they fit in zero_affine.py, next to zero_param_storage_affine_map and autotp_storage_affine_maps, so the checkpoint side can reuse them? Either way works for me.
4. Prepared transfers. The plans and the arena's addresses don't change between syncs, so once the staged prepare/run from §5 lands, the sync could keep the prepared handle and skip the per-call validation and schedule build.

On "What it does not do": the two-layout constructor from #9026 would let a surviving TP shard and the trainer's partitions go through this same path when the split is along rows. I can open it as a small separate PR that this one can use, if that suits you and @0z5a.

The RTX 4090 table is a nice result: 7–8× faster than gathering for Qwen3-8B at TP4, with the peak set by the staging budget rather than the largest parameter.

[LiRunGuo]

Thanks @nathon-lee and @Achyuthan-S. I pushed two commits for the points below: 70831cb (docstrings, docs and one test, no behavior change), and 48a3575, which moves the planning onto #8936 now that it has merged (after merging current master).

@nathon-lee

1. Selector scope. The load did carry across parameters: settle(before, plan) set it to before (what the earlier parameters' plans send) plus what this plan sends, so the increment in __call__ only balanced the regions of the parameter being planned. That workaround is gone now. With [checkpoint] Batch replica selection for affine transfer plans #8936 merged, the sync plans every parameter with one plan_transfers call, and its policy sees every request of the sync at once. The default is [checkpoint] Batch replica selection for affine transfer plans #8936's local_first_least_loaded, which is what _LocalFirstLeastLoaded imitated, so the selector class is deleted.
   • test_reads_of_replicated_parameters_spread_over_their_holders still pins the balance. Each of its six parameters is a single region with three holders, so there is nothing to balance inside a parameter, and the reads still split two parameters per holder.
2. Keys with a custom group. Yes: the maps stay keyed by global rank whatever the group, and the sync translates them to group ranks for the executor. The docstring of group and the docs now say so, and a new test runs the sync over a group that leaves rank 0 out, so group ranks and global ranks differ.

@Achyuthan-S

Thanks for running it, and for the selector check.

1. The selector changes only the cost. Added to the module docstring and the docs: once both maps pass the planner's checks, every choice writes the same values, provided the replicas agree (P2). Under ZeRO-1/2 the replicas are the data-parallel copies, equal after the step's all-gather; under ZeRO-3 each element has one holder.
2. AutoTP. The docstring and docs now say that autotp_storage_affine_maps should serve as the source for an AutoTP trainer but is not tested yet, which matches Compatibility. I'd gladly take your help with the AutoTP TP2 → TP4 test; it would be a good second case for the map-only interface.
3. Where the helpers live. I kept them in the rollout module for now, so this PR still changes no existing source file. If the maintainers prefer zero_affine.py, next to zero_param_storage_affine_map and autotp_storage_affine_maps, I'll move them in this PR.
4. Prepared transfers. Agreed. The plans and the arena's addresses are fixed after construction, so the sync can keep the prepared handle and skip the per-call validation and schedule build. I'll do it once the staged prepare/run lands.

#8936. AffineWeightSync now takes an optional replica_policy, passed to plan_transfers together with the location→process maps and the element size, so a policy sees local copies and bytes. A cost-aware policy from the policy layer can then be dropped in without touching the sync. The same test runs with first_holder and checks that all reads then come from rank 0, so the argument reaches the planner.

The two-layout constructor. Yes, a separate PR suits me, row case first, as @0z5a agreed on #9026; thanks for opening #9052. The sync declares its sources only through maps, so it can take a surviving TP shard plus the trainer's partitions as soon as #9052 lands, and I'll add that case to this PR's tests then.

Tests on 48a3575 (on current master, dfa7d34), 4× RTX 3090:

• the new tests: 21 passed on GPU (NCCL); 20 passed and 1 skipped on CPU (gloo, 4 ranks), where the skip is the CUDA graph test;
• tests/unit/runtime/rollout: 116 passed, 1 skipped, 1 failed. Both are master's tests that need the transformer_inference extension, which doesn't build here (CUDA 13): the failure is TestZero3SynchronizedEarlyStop, raising from that build. This PR touches neither the test nor the extension;
• the affine tests (test_affine_shard_map, test_affine_transfer_plan, test_affine_transfer_executor, test_affine_transfer_executor_external_group, test_zero_affine, test_affine_replica_policy): 350 passed, 1 skipped (the external-group test is CPU only).

yapf and flake8 (the versions pinned in pre-commit) pass on the changed files.

[Achyuthan-S]

Thanks @LiRunGuo, that addresses everything. I read 48a3575 merged with current master: the sync now plans with one plan_transfers call that carries the location→process maps and the element size, so a policy sees local copies and bytes, and the first_holder case of the spread test shows the argument reaches the planner. The new offload/reload test also covers the stale-storage case from #9026: the sync resolves the shards at every call, so a reload_states() between syncs is picked up.

• replica_policy is where the cost-aware policy will plug in ([RFC] Batch replica selection for affine transfer plans #8935, [RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026 §3), and this sync will be its first user.
• I'll write the AutoTP TP2 → TP4 test: an AutoTP TP2 trainer with autotp_storage_affine_maps as the source maps, into a TP4 copy, checked against layout.extract like the ZeRO cases. I'll send it as a PR to your branch.
• Once Describe a parameter held in several layouts as one affine map #9052 lands, a surviving TP shard plus the trainer's partitions can go through the same sync for row splits.

[LiRunGuo]

Thanks @Achyuthan-S.

• AutoTP TP2 → TP4 test. Yes please. The branch is rollout-affine-weight-sync on my fork (head 48a3575); the commits need a sign-off for the DCO check. I'll run it with the PR's other tests on 4 GPUs and move AutoTP in the description from "should work, untested" to tested.
• Describe a parameter held in several layouts as one affine map #9052. The two-layout case (a surviving TP shard plus the trainer's partitions) goes into this PR's tests when it lands, as said above.
• Prepared transfers. [checkpoint] Expose the affine transfer schedule as schedule_transfers #9027 is approved, so the prepared handle comes next. The sync resolves its shards at every call today, which is how it follows a reload_states(). With a handle kept across syncs it would rely on the handle check you and @jinyouzhi describe on [RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026 to notice the reload.