[hwchen2017]

Problem

Since #8680 landed, modal-torch-latest / DeepSpeedAI CI has failed for every merge-queue entry that runs this test on Modal's L40S runners (#8794, #8701, #8670). Example: run 38009529440.

FAILED tests/unit/v1/autotp/test_autotp_training.py::TestTiedRowParallelOutputHeadTraining::test_training_matches_unsharded_reference[35-3-False-0]
FAILED tests/unit/v1/autotp/test_autotp_training.py::TestTiedRowParallelOutputHeadTraining::test_training_matches_unsharded_reference[35-3-True-0]
Mismatched elements: 1 / 670 (0.1%)
Greatest absolute difference: 1.1173542588949203e-06 at index (1, 1, 45) (up to 1e-06 allowed)

#8680's own merge-queue run and the next nightly run both used the AWS fallback, where the test passes, so the failure wasn't caught there.

Root cause

The row-parallel head (LinearAllreduceWithReplicatedInput) slices the hidden vector, computes a partial matmul on each rank, and all-reduces the partial logits. The reference is one unsharded fp32 matmul, and the two summation orders round differently. The tied head uses the embedding's N(0,1) weights, so logits reach about 14, where one fp32 ulp is about 1e-6. Logits that are near zero through cancellation can therefore differ by more than atol=1e-6. Even the fp32 reference differs from a float64 ground truth by up to 2.2e-6.

Fix

Use torch.testing's default fp32 tolerance (atol=1e-5, rtol=1.3e-6) for every fp32 assert_close in test_autotp_training.py. That covers 17 checks, including both row-parallel logits checks, and matches the 19 fp32 checks in the file that already use the default. The half-precision tolerances in run_tp_layer_fwd_bwd are unchanged.

The default is enough for what these tests check:

• The sharding bugs they target (wrong shard offset, missing all-reduce, dropped tied-gradient contribution, tied weights drifting apart) move values by 0.1 to 10.
• A simulated precision regression (all-reduce inputs rounded to bf16) still fails every check it failed under the old tolerances. The one exception is the SGD-update check, which can't detect it under either tolerance.
• The custom values were tighter than the default only for values below about 1, and looser for larger ones such as the losses.

Evidence (2x H200)

• In a sweep of 40 init seeds, every parametrization of the tied test has seeds that fail the old logits check (worst ratio to the allowed tolerance: 1.19). Under the default, the worst case uses under 30% of the limit.
• Changing only the seed to 8199 reproduces the CI failure: [35-3-False-0] and [35-3-True-0] fail with 1 / 670 mismatched and 1.25e-06 (up to 1e-06 allowed).
• Every other changed check stays under 3% of the default limit, recorded per call across the 4 affected test classes.

Testing

Hardware: 2x NVIDIA H200, torch 2.12.0+cu130, Python 3.12.

• pytest tests/unit/v1/autotp/test_autotp_training.py -k "TestRowParallelOutputHeadTraining or TestTiedRowParallelOutputHeadTraining": 8 passed.
• pytest tests/unit/v1/autotp/test_autotp_training.py -m sequential -k "TestTiedEmbeddingAliasesAndFallback or TestTiedVocabParallelLMHead": 30 passed.
• Seed-8199 reproduction, which fails before this change: passes with the default logits tolerance.
• pre-commit run --files tests/unit/v1/autotp/test_autotp_training.py: passed.

No L40S is available locally. This PR's merge-queue run selects this test file, so it will check the fix on Modal's L40S runners unless the run falls back to AWS.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-10T03:47:49.376484Z	5858a4f	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5858a4fc0a

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[delock]

Let's run AWS and modal test on this PR manually to ensure the new tolerence work on both platform. @hwchen2017

https://github.com/deepspeedai/DeepSpeed/actions/runs/38029376284
https://github.com/deepspeedai/DeepSpeed/actions/runs/38034637564