Exclude fake-backed axes from get_all_one_dimensional_meshes

[adityasingh2400]

Fixes #2447.

get_all_one_dimensional_meshes() filters _single_axis_meshes on ndim == 1 and size() > 1. That filter is not the same predicate the rest of ParallelDims uses to decide whether an axis is live, and the gap is visible on every dense run.

Root cause: build_mesh() unflattens the sparse mesh as ("pp", "dp_replicate", "efsdp", "ep") and passes backend_override[name] = "fake" for any axis where _mesh_exist(name, degree) is False. For efsdp, _mesh_exist returns self.ep > 1, so with EP disabled the axis gets a fake process group. Its size, however, is dp_shard * cp * tp / ep, which for an ordinary dense job is greater than 1. The size() > 1 filter therefore lets a fake-backed axis through. get_optional_mesh() does consult _mesh_exist, so today get_optional_mesh("efsdp") returns None while get_all_one_dimensional_meshes() returns an efsdp entry for the same ParallelDims.

Two consequences. The reported one is that the startup line prints efsdp in "Successfully created meshes with active dimensions" for dense models. The one that matters more is that set_pg_timeouts() in torchtitan/distributed/utils.py builds its group list straight from this accessor, so it calls mesh.get_group() on the fake axis and applies a timeout to a process group that cannot carry a collective.

Repro on 2 CPU ranks with gloo, dp_replicate=1, dp_shard=2, cp=1, tp=1, pp=1, ep=1, before the change:

ep_enabled: False
one-d meshes: ['batch', 'loss', 'efsdp', 'fsdp']
  batch: size=2 backend=gloo
  loss: size=2 backend=gloo
  efsdp: size=2 backend=fake
  fsdp: size=2 backend=gloo
get_optional_mesh('efsdp'): None

After:

ep_enabled: False
one-d meshes: ['batch', 'loss', 'fsdp']
  batch: size=2 backend=gloo
  loss: size=2 backend=gloo
  fsdp: size=2 backend=gloo
get_optional_mesh('efsdp'): None

Fix: add self._mesh_exist(k, v.size()) to the filter, which is exactly the predicate get_optional_mesh already applies, so the two accessors agree. The change is subtractive only for axes that build_mesh gave a fake backend. fsdp and, under full_dtensor or spmd_types, dp_shard are the axes _mesh_exist keeps alive at size 1, and they are still filtered out by size() > 1 exactly as before. For every other axis _mesh_exist reduces to degree > 1, so the only behavior change is dropping efsdp when EP is off.

Tests: added TestOneDimensionalMeshesSkipFakeAxes with two 8-rank cases. With ep=1 the efsdp axis has size 8 yet must not be reported, and every reported axis must have a non-fake backend. With ep=2 the efsdp and ep axes must still be reported. Also updated the existing test_world_size_8_mesh_operations assertion that encoded the old behavior, and refreshed the docstring example and notes.

Test evidence: I have no GPU and torchtitan main needs a torch nightly newer than 2.11.0 to import the test module (torch.distributed.fsdp.DataParallelMeshDims), so I ran the two new cases as a standalone script over 8 real gloo ranks on CPU with the same assertions:

case ep=1 PASS, axes: ['batch', 'fsdp', 'loss', 'tp']
case ep=2 PASS, axes: ['batch', 'efsdp', 'ep', 'fsdp', 'loss', 'tp']
ALL PASS

black is clean on both changed files.

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.