Add an AutoTP equivalence check for uneven TP size - #1008
Conversation
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. Compare the loss curves instead, which separates a correct partition from a plausible-looking wrong one. 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 tests the uneven split of Qwen3's 16 attention / 8 KV heads (6/6/4 per rank), where a shard boundary is most likely to be computed wrongly -- an even split divides every dimension exactly and hides off-by-one bugs. A tolerance alone would not settle whether the 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 and calibrates what the uneven split should cost. Over 500 steps the two land on top of each other: means within a few percent, and the same worst step on both backends (467 on GPU, 379 on CPU), with the even control the further of the two from the baseline at that step. The spike belongs to the training trajectory, not to how the heads were divided. Everything that could make the runs diverge for another reason is pinned: same seed and batch order across ranks, fp32, dropout asserted zero, world_size required to equal autotp_size so spare ranks cannot silently add data parallelism, and a fixed thread cap on CPU. 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, since both runs start from the same weights there and training dynamics cannot yet have amplified anything. Verified on 3 and 4 GPUs with NCCL and on 3 and 4 CPU ranks with gloo, 500 steps each. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: youzhiji <youzhiji@5090d-8.sh.intel.com>
|
test deepspeed #8185 c02df7ff84799c94f8decf50c2511811cfcf01df Run on 4x RTX 5090 |
Very cool work, can you elaborate your testing environment, e.g. how many GPUs did you use to test out AutoTP? |
Thank you for your reviewing. That was completed on 4x RTX 5090. |
2 participants
Motivation
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.pytrains one model for N steps and records every step's loss. Running itat 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's2048 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
0.00e+007.69e-052.52e-03(step 467)0.00e+007.58e-052.72e-03(step 467)6.82e-082.75e-051.15e-03(step 379)6.82e-082.79e-052.39e-03(step 379)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:
world_sizemust equalautotp_size. Spare ranks would silently become a data-parallel dimension, averaging gradients over more samples and changing what is being compared.OMP_NUM_THREADSidentically at every width. Without it each rank sizes its thread pool for the whole machine and several ranks oversubscribe it: 27s per step instead of 1.5s.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.pychecks every step rather than the final loss, becausereassociation error jitters while a sharding bug compounds. It checks the first
step far more tightly (
--forward-rtol, default 1e-6): both runs start from thesame 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
Testing
reproducible bit-for-bit across repeated runs.
compare_loss.py: accepts reassociation-scale noise, rejects acompounding 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.