[Achyuthan-S]

The AutoTP TP2 → TP4 test I promised on deepspeedai#9045, for your branch at 48a3575.

What it adds (test_an_autotp_trainer_syncs_into_a_tp4_copy in TestAffineWeightSync):

• An AutoTP TP2 × DP2 trainer on four ranks, under ZeRO-1, ZeRO-2, ZeRO-3, and ZeRO-3 with the model built under zero.Init. In the last case the norm, which AutoTP leaves alone, keeps a partition over the whole world.
• autotp_storage_affine_maps(engine.module) as the source maps, into one TP4 copy over the same ranks (_layouts(4), so the up and down shards are uneven: 5, 5, 4, 4).
• Each rank's shard is compared with layout.extract of the full weights from engine._consolidated_16bit_state_dict(), before and after a training step. That consolidation gathers over ZeRO and TP on the engine's own path and shares no code with the maps.
• Every rank builds the model from the same seed, as TestAutoTPStorageMaps does. DeepSpeed broadcasts the initial weights only within each data-parallel group, so without the seed the two TP ranks hold different down.bias and the replicas disagree. AutoTP's dataloader check also needs the same batch on both TP ranks.
• The comparison is gathered before it is asserted, so a wrong shard on one rank fails every rank instead of leaving the others waiting in the next collective, which is how it hung before I seeded it.
• The module docstring and affine-transfer.rst no longer call the AutoTP source untested.

Validation on macOS arm64 CPU, Python 3.9.6, torch 2.8.0, gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4, from tests/):

• The new test: 4 passed.
• The whole file: 23 passed, 1 skipped (CUDA graphs) and 1 failed. The failure is test_shards_given_at_each_sync_follow_offload_and_reload, which fails the same way on unmodified 48a3575 here: the fused_adam op needs ninja to build, and this machine doesn't have it.
• Swapping the two TP ranks' storage inside autotp_storage_affine_maps makes all four cases fail on up.weight.
• yapf and flake8 pass on the changed files.
• Not run on GPU yet; you offered to run it with the PR's other tests on 4 GPUs.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.