[graph_trainer] Fix H100 CI: pre-install spmd_types (torch nightly hard-dep)

[SherlockNoMad]

Problem

The GraphTrainer 8 GPU H100 Integration Tests workflow has been failing on main for a while. The current (and recurring most-recently) failure is during dependency install:

Collecting spmd-types==0.2.0 (from torch)
  WARNING: Retrying ... ConnectionResetError(104, 'Connection reset by peer'): /packages/.../spmd_types-0.2.0-py3-none-any.whl.metadata
ERROR: Could not install packages due to an OSError: HTTPSConnectionPool(host='files.pythonhosted.org', port=443): Max retries exceeded

Root cause

The nightly torch wheel declares spmd_types==0.2.0 as a hard dependency. The Meta-internal H100 runner can reach download.pytorch.org but not the public PyPI file host (files.pythonhosted.org). The torch install step deliberately clears PIP_EXTRA_INDEX_URL (the reachable in-cluster cache) "so the default cpu index can't supply a +cpu torch" — so when pip resolves torch's spmd_types dependency it falls back to public PyPI, which the runner cannot reach, and the install fails on every run.

(The docker image bakes spmd_types==0.2.1 from .ci/docker/requirements.txt, which does not satisfy torch's exact ==0.2.0 pin, so pip tries to fetch 0.2.0 regardless.)

Fix

Pre-install spmd_types from the in-cluster cache alongside torch's other pure-python deps (the workflow already does this for filelock, sympy, etc. for exactly this reason). Once it is already satisfied, the subsequent torch install does not try to fetch it from public PyPI. The pin matches the version torch nightly requires (0.2.0).

- python -m pip install filelock typing-extensions "setuptools<82" sympy networkx jinja2 fsspec numpy
+ python -m pip install filelock typing-extensions "setuptools<82" sympy networkx jinja2 fsspec numpy spmd_types==0.2.0

Notes

• The shared .github/scripts/run_8xgpu_integration_tests.sh (main H100 workflow) and the RL H100 workflow have the same latent issue and likely want the same one-line fix; left out here to keep this PR scoped to the graph_trainer H100 CI.
• This unblocks the install so the test suite can actually run again; any test-level failures surfaced afterward (e.g. nightly hash drift) will be addressed on top.

Test

Triggered via ciflow/h100.8 on this PR.

[tianyu-l]

IIUC, right now graph trainer doesn't require spmd_types package, can we skip for now?

[SherlockNoMad]

It's not torchtitan/graph_trainer pulling it — spmd_types==0.2.0 is a hard dependency of the torch nightly wheel itself (Collecting spmd-types==0.2.0 (from torch)), so pip install torch resolves it regardless of what graph_trainer needs. The H100 runner can reach download.pytorch.org but not public PyPI (files.pythonhosted.org), and we clear PIP_EXTRA_INDEX_URL for the torch install, so the fetch dies with "Connection reset". There's no clean way to skip a single torch hard-dep via pip, so I pre-install it from the reachable in-cluster cache — the same line already does this for torch's other pure-python deps (filelock/sympy/…). Confirmed working in the latest run (it resolved from the cache, no connection reset).

The image bakes spmd_types==0.2.1 (from .ci/docker/requirements.txt) which doesn't satisfy torch's exact ==0.2.0, which is why pip still tries to re-fetch. If you'd rather keep the workflow untouched, the alternative is aligning the .ci/docker pin to 0.2.0 — happy to switch to that. Either way it can't be fully skipped while the nightly hard-deps it.

[aditvenk]

@pianpwk -- is it expected for spmd_types to be a hard dependency of torch?

[pianpwk]

I think this happens for CI wheels? I don't think running generic pip install torch gets it. never mind, it gets installed with nightlies

It sounds like we could update pytorch's dependency to 0.2.1? I think it should stay pinned there for a while (that version should be sufficient for titan + spmd types), while we figure this out

pytorch/pytorch#186803

[aditvenk]

thanks, but is it expected for spmd_types to be a hard dependency for the nightly? i thought it is an optional dependency

[pianpwk]

Let me figure this out.. I think many things are true right now:

• pip install torch getting spmd_types is desirable for spmd_types' discoverability.
• technically pytorch will run without spmd_types installed (e.g. building from source shouldn't install it), though with the nightly situation this is rarer. So technically it's not a hard requirement.
• But titan CI via torch nightly, needs it to workout, so in some sense it is a requirement?

[aditvenk]

My expectation was that if someone wanted the spmd_types-parts of torch to work, they would have to install it directly or via torchannex (for which spmd_types can be a hard dependency). It's possible I am mistaken...

[fegin]

TorchTitan has a hard dependency on spmd_types now because the core components, e.g., sharding.py import spmd_types. So yes, graph_trainer depends on spmd_types.

[aditvenk]

Yes I expect torchtitan will have a hard dependency on spmd_types. I was surprised by torch's

[tianyu-l]

Is numerics change expected?

[SherlockNoMad]

Yes — expected upstream cu130-nightly drift, not a torchtitan change. In the same CI run the trace path still bitwise-matches eager (test_*_aot_fx_trace*_vs_eager pass), so eager and traced drifted together; only the hardcoded eager self-deterministic baselines needed refreshing. Same recurring nightly drift re-baselined before in #3482 / #3484. (DSv3's loss moved too, hence the loss update on the DSv3 classes.)