AutoTP: send CPU copies in the first-batch consistency check

[pengdurice]

AutoTP: send CPU copies in the first-batch consistency check

Summary

With AutoTP, the engine checks the first batch: tensor-parallel rank 0 sends its inputs to the other ranks of its TP
group with broadcast_object_list (check_dataloader_inputs_same_across_ranks in engine.py), and each peer compares
them with its own. The inputs are accelerator tensors, so they are pickled with their device. Each peer unpickles
them onto rank 0's device and opens a context there, which stays for the life of the process.

On MiniMax-M3 with AutoTP 4 on 4 nodes, every TP rank 0 GPU showed 3 other processes holding 524 MiB each:
1.5 GiB of that GPU lost, and those GPUs were the ones that ran out of memory.

The fix sends CPU copies. The peers' comparison already moves both sides to its own device
(compare_tensors_in_structures), so the check itself is unchanged.

Used in large-scale training

The same change (the same line) was in the DeepSpeed study trees of these MiniMax-M3 runs: all 60 layers (426B) built
from the config (random weights), 4 nodes x 8 H200, AutoTP 4 x AutoEP 8, ZeRO-3, bf16, together with the ZeRO-3
master-copy fix (a separate PR).

This change	Setup	Result
no	32,768 tokens per GPU, no CPU offload	out of memory in step 2 only on the TP rank 0 GPUs (ranks 0, 4, ..., 28), 1.50 GiB short; each held 3 x 524 MiB of other processes' contexts
yes	the same	out of memory in step 2 on many ranks, not only the TP rank 0 GPUs: the model does not fit at this length without offload
yes	CPU optimizer offload, 8,192 / 32,768 / 65,536 tokens per GPU	trained: 30.92 / 59.16 / 106.32 s per step (mean of steps 2-4), 43.5 / 63.4 / 59.1 GiB peak

Duplicate check

Searched open pull requests on 2026-10-15 ("AutoTP", "dataloader check", "broadcast_object_list", "CUDA context"): none
changes this check.

The change

deepspeed/runtime/engine.py, the source rank's object becomes
torch.utils._pytree.tree_map(lambda x: x.cpu() if torch.is_tensor(x) else x, args) (5 lines added, 1 removed,
including a 2-line comment).

Tests

TestTpDataloaderCheckSendsCpuCopies (2 ranks, AutoTP 2, tests/unit/v1/autotp/test_autotp_training.py; every
weight left unsharded with partition_type: "skip", as in TestTpDataloaderCorrectness): wraps
deepspeed.comm.broadcast_object_list after deepspeed.initialize, runs the first forward, and requires every tensor
the source rank sends to be on the CPU; the existing TestTpDataloaderCorrectness checks that inconsistent batches
are still caught.

Hardware: 1 node of 8 H200 (driver 580.126.09), torch 2.8.0+cu126, pytest 9.1.1.

pytest unit/v1/autotp/test_autotp_training.py -k TpDataloader
# this branch:                         3 passed
# master 5f266e26d + only the new test: 2 passed, 1 failed (TestTpDataloaderCheckSendsCpuCopies:
#                                       AssertionError: ['cuda', 'cuda'], the source rank sent GPU tensors)

pre-commit run --files passes on both changed files.

Compatibility

Cell	Validated
AutoTP 2, ZeRO-0, bf16 (test)	yes
AutoTP 4 + AutoEP 8, ZeRO-3, 4 nodes	yes, with the same change in a study tree (see above)
Accelerators other than CUDA	the copy uses .cpu(); not run

AI assistance

This change, its test and this description were prepared with AI assistance (Claude).