Run multi-rank CPU unit tests in CI via LOCAL_SIZE - #8381
Conversation
cpu-torch-latest runs on a single-socket runner where CPU_Accelerator.device_count() reports 1 NUMA node, so the per-device gate in tests/unit/common.py skips every test that needs more than one rank. CPU ranks are plain processes over gloo and need no per-rank hardware, so advertise 4 local devices via LOCAL_SIZE, the env var device_count() reads first. The test harness re-sets LOCAL_SIZE per worker, so this value only affects the launch gate. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: af3841378c
ℹ️ 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".
| # multi-rank test. CPU ranks are plain processes over gloo, so advertise 4 | ||
| # local devices to let world_size<=4 tests run. The test harness re-sets | ||
| # LOCAL_SIZE per worker, so this value only affects the launch gate. | ||
| LOCAL_SIZE: '4' |
There was a problem hiding this comment.
Keep CUDA-only distributed tests out of the CPU run
In the cpu-torch-latest job, advertising four devices admits every distributed test with world_size <= 4, not only CPU-safe tests. For example, tests/unit/ulysses_alst/test_ulysses_sp_hf.py:240-263 defines an unguarded two-rank test that creates tensors on cuda:<rank>; because CPU_Accelerator.is_available() returns true, the harness does not skip it, and the all-unit pytest invocation at workflow line 283 will fail on the CPU-only PyTorch installation. Scope this override to an explicitly CPU-compatible subset or add CPU capability checks before enabling the previously skipped tests.
Useful? React with 👍 / 👎.
| # multi-rank test. CPU ranks are plain processes over gloo, so advertise 4 | ||
| # local devices to let world_size<=4 tests run. The test harness re-sets | ||
| # LOCAL_SIZE per worker, so this value only affects the launch gate. | ||
| LOCAL_SIZE: '4' |
There was a problem hiding this comment.
Add the required sign-off trailer
This is a non-merge commit, but its commit message has no Signed-off-by trailer. Add the author sign-off so the commit satisfies the repository's commit and CI requirements.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
|
This is an on going work. Note I'll rebase when revealed bugs are fixed. The goal is fix all issues exposed by this CI, then turn on multi-rank test in CPU workflow. |
The now-executing multi-rank CPU tests exposed three test-infrastructure
bugs, all GPU assumptions that were never exercised before:
- create_ddp_model passed device_ids=[rank] to DDP; on CPU the module
lives on the shared 'cpu' device and torch requires device_ids=None
there (106 failures).
- reduce_boolean_flags built a device from current_device(), which is a
LOCAL_RANK string on CPU ('Invalid device string'), and fed a 0-dim
tensor to all_gather_into_tensor, which gloo rejects.
- train_cifar passed devices=['cpu'] to fork_rng, but torch.cpu has no
get_rng_state; fork_rng already saves the CPU/global RNG.
Left for later batches: onebit optimizer NoneType (runtime), ulysses
gloo allgather/numerics (runtime), autotp CUDA-only assert, autocast
tests (hardcoded nccl backend).
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Review follow-up: don't special-case the accelerator implementation. - create_ddp_model now checks the model's torch device type, mirroring the exact precondition torch's own DDP constructor enforces, instead of matching the DeepSpeed accelerator name. - train_cifar probes the device module for a per-device RNG instead of matching the accelerator name; backends without get_rng_state (cpu) need no fork_rng device entries because the global RNG is always saved. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…fp16 skips Follow-up triage of the LOCAL_SIZE run: - test_zero_user_backward had four more DDP reference sites pinning device_ids=[rank]; route them all through one wrap_ddp_reference helper that only pins device ids for indexed devices. - test_autotp_training passed get_accelerator().current_device() (a LOCAL_RANK string on CPU) to .to(device=...); use current_device_name() everywhere. - Tests with hardcoded fp16 configs (multiple models, MoE, MoE-TP) fail with 'Type fp16 is not supported' on runners whose CPU lacks AVX512-FP16; skip them on accelerators without fp16 support instead of depending on runner hardware lottery. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
With LOCAL_SIZE=4 the suite now runs to ~63% and then all xdist workers go silent for hours until the 6h job limit cancels the run - the pool worker cleanup hang that DS_DISABLE_REUSE_DIST_ENV was added for. Fresh pools per test cost some wall time but let the run finish and print the failure summary. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Both full-suite attempts wedge before printing a summary: once the 100th failure trips PYTEST_OPTS' --maxfail, pytest-xdist's interrupt path stalls forever in mp pool teardown (no timeout guards _close_pool), and the 6h job limit cancels the run. Run the suite as two fresh-worker halves, override maxfail so all failures are listed, and cap each half with timeout so the sequential tail always runs. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Experiment results: enabling multi-rank CPU tests via
|
| Count | Test | Root cause |
|---|---|---|
| 92 | v1/zero/test_offload_states.py |
Asserts memory_allocated() drops after offload — a VRAM concept. CPU_Accelerator.memory_allocated() returns RSS, which never shrinks on free. Test/accelerator semantic gap; needs discussion (gate on device semantics or a better CPU metric). |
| 24 | v1/zero/test_zero_autocast.py |
Hardcoded dist_backend='nccl' + a CUDA/NCCL-oriented bf16 gate. Needs accelerator-derived backend + capability-based skip. |
| 25 | ulysses SP tests | Numerical mismatches on CPU (real correctness questions, need deep dive). |
| 16 | onebit optimizer tests | NoneType.size at onebit/adam.py:108 — runtime bug on CPU path. |
| ~20 | pipe / zeropp / moe-checkpoint / coalesce | Smaller follow-ups, same GPU-assumption patterns. |
Test-side fixes already in this branch (validated by CI)
- All 5 DDP reference sites route through
wrap_ddp_reference()(CPU modules must not pindevice_ids) — dropped this class from 114 failures to 7. reduce_boolean_flags/ autotp tests usecurrent_device_name()instead ofcurrent_device()(a LOCAL_RANK string on CPU, not a device).- fp16-config tests get a capability skipif (
not get_accelerator().is_fp16_supported()) — GHubuntu-24.04runners are hardware-heterogeneous w.r.t. AVX512-FP16, so hardcoded fp16 was a runner lottery. fork_rngprobes the device module for a per-device RNG instead of matching accelerator names.
CI infrastructure findings (need maintainer decisions)
--maxfail=100inPYTEST_OPTS+ no timeout inDistributedExec._close_pool= permanent wedge: when the 100th failure trips the interrupt, teardown (pool.starmap(_dist_destroy),close/join) can block forever on wedged gloo peers; the run then dies at the 6h job limit without printing any summary. This is what cancelled the first three attempts. (This branch works around it by splitting the suite and raising maxfail.)- Suite size vs 4 vCPU: the full multi-rank suite does not fit one 6h job. Options: split invocations (done here), lower
-n, or a dedicated larger runner. CPU_Accelerator.device_count()'s NUMA-node semantics remain the root cause of the original gap; the long-term fix could be a test-harness-level exemption for CPU (ranks are processes, not devices).
Happy to split the commits into separate PRs (workflow change / mechanical test fixes / triage follow-ups) per maintainer preference — the failure inventory above is intended as the working list.
This is an on going work. Note I'll rebase when revealed bugs are fixed. The goal is fix all issues exposed by this CI, then turn on multi-rank test in CPU workflow.
Problem
On the
cpu-torch-latestrunner (single-socketubuntu-24.04), every unit test that needs more than one rank is skipped:Root cause:
CPU_Accelerator.device_count()reports the number of NUMA nodes (accelerator/cpu_accelerator.py), which is 1 on the single-socket runner, andDistributedExec._launch_procs()intests/unit/common.pygates process count ondevice_count(). Multi-rank CPU tests do not actually need one device per rank — ranks are ordinary processes communicating over gloo.Change
Set
LOCAL_SIZE=4for theunit-testsjob.device_count()readsLOCAL_SIZEfirst, so the launch gate now admitsworld_size<=4tests. This is the same signal the DeepSpeed launcher sets for spawned processes; the unit-test harness re-setsLOCAL_SIZEper worker in_dist_run(), so the CI-level value only affects the gate and cannot leak into test bodies.Evidence this is safe
TestDistIsendIrecv(tests/unit/comm/test_dist.py, world_size=2) and the autotp universal-checkpoint test (tests/unit/checkpoint/test_autotp_uc_checkpoint.py, world_size=4) already bypass the per-device gate for CPU and run green in this very CI job.LOCAL_SIZEhas a single reader in the codebase (CPU_Accelerator.device_count()); the launcher only writes it.Expected impact
~207 world_size=2 and ~85 world_size=4 tests that are currently skipped will now execute on CPU CI (world_size>=8 stays skipped). Some of them may have latent failures — this PR intentionally surfaces them so they can be triaged.