Add an AutoTP equivalence check for uneven TP size - #1008

Motivation

Verify DeepSpeed#8185

The existing AutoTP smoke tests check that a sharded run exits cleanly. That is a
weak signal: a shard cut on the wrong boundary also exits cleanly, it just trains
a different model. This adds a check that compares what the runs actually compute.

What it does

train.py trains one model for N steps and records every step's loss. Running it
at AutoTP=1 and AutoTP=3 and diffing the curves exercises the uneven split of
Qwen3's 16 attention / 8 KV heads (6/6/4 per rank), which is where a shard
boundary is most likely to be computed wrongly. Even splits divide every dimension
exactly and hide off-by-one bugs; an uneven split also forces the head count,
rather than the raw element count, to drive the partition. Splitting q_proj's
2048 output features as 683/683/682 would land mid-head and corrupt the attention
reshape.

Why there is a control run

A tolerance alone cannot settle whether an uneven split is accurate, only
whether it fits under a number chosen after the fact. So the runs also include
AutoTP=4 as a control. It divides evenly, so its gap to the baseline is pure
floating-point reassociation, which calibrates what the uneven split ought to cost.

Results — 500 steps, Qwen3-0.6B, fp32

The uneven split and the even control land on top of each other. Their means agree
to within a few percent, and on both backends the worst step is the same step
(467 on GPU, 379 on CPU), with the even control the further of the two from
the baseline there. The spike belongs to the training trajectory at that point,
not to how the heads were divided.

So the conclusion is not merely "AutoTP=3 stayed under the tolerance". It is that
splitting 16 heads unevenly across 3 ranks costs nothing in accuracy beyond what
an evenly-divisible tensor-parallel run already costs.

The GPU runs match exactly at step 0 while the CPU runs differ by 6.82e-08:
gloo reduces in a different order than a single rank does, so the sharded forward
is not bit-identical there. That is why the forward check uses a small tolerance
rather than demanding equality.

What is pinned, and why

Anything that could make the runs diverge for a reason other than sharding is fixed:

The training data is random tokens, so the loss itself is meaningless. What matters
is only that differently-sharded runs agree on it.

How the comparison reads

compare_loss.py checks every step rather than the final loss, because
reassociation error jitters while a sharding bug compounds. It checks the first
step far more tightly (--forward-rtol, default 1e-6): both runs start from the
same weights and no optimizer step has happened yet, so a gap there is a wrong
forward, not accumulated drift, and training dynamics cannot be blamed for it.

Usage

cd training/autotp_equivalence
bash run_gpu.sh 500 0,1,2,3   # NCCL
bash run_cpu.sh 500           # gloo
python -m pytest tests/ -v

Testing

• 500 steps at AutoTP=1/3/4 on GPUs (NCCL) and on CPU ranks (gloo); results above,
  reproducible bit-for-bit across repeated runs.
• 12 unit tests for compare_loss.py: accepts reassociation-scale noise, rejects a
  compounding gap, catches a wrong forward at the first step without mistaking later
  drift for one, reports the worst step rather than the last, and never hides the
  worst step when sampling a long run.