Skip to content

Run multi-rank CPU unit tests in CI via LOCAL_SIZE - #8381

Open
delock wants to merge 6 commits into
deepspeedai:masterfrom
delock:ci/cpu-multi-rank-local-size
Open

Run multi-rank CPU unit tests in CI via LOCAL_SIZE#8381
delock wants to merge 6 commits into
deepspeedai:masterfrom
delock:ci/cpu-multi-rank-local-size

Conversation

@delock

@delock delock commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

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-latest runner (single-socket ubuntu-24.04), every unit test that needs more than one rank is skipped:

SKIPPED tests/unit/common.py:275: Skipping test because not enough GPUs are available: 2 required, 1 available

Root cause: CPU_Accelerator.device_count() reports the number of NUMA nodes (accelerator/cpu_accelerator.py), which is 1 on the single-socket runner, and DistributedExec._launch_procs() in tests/unit/common.py gates process count on device_count(). Multi-rank CPU tests do not actually need one device per rank — ranks are ordinary processes communicating over gloo.

Change

Set LOCAL_SIZE=4 for the unit-tests job. device_count() reads LOCAL_SIZE first, so the launch gate now admits world_size<=4 tests. This is the same signal the DeepSpeed launcher sets for spawned processes; the unit-test harness re-sets LOCAL_SIZE per 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_SIZE has 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.

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>
@delock
delock requested a review from loadams as a code owner September 1, 2026 01:21

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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'

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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'

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

@delock

delock commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator Author

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>
@delock

delock commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator Author

Experiment results: enabling multi-rank CPU tests via LOCAL_SIZE=4

Baseline comparison: same-commit master run (all multi-rank tests skipped) vs this branch. Zero regressions: every failure is a test that was previously skipped.

What now runs

~1800 previously-skipped multi-rank tests execute on the CPU runner (gloo). Final complete numbers from the last run (split invocations):

  • unit/v1 half (complete): 916 passed / 131 failed / 206 skipped
  • sequential (complete): 54 passed / 9 failed
  • non-v1 half: killed by the 150m guard (still too slow on 4 vCPU), 59 failures listed inline

Failure inventory & root causes (triaged)

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 pin device_ids) — dropped this class from 114 failures to 7.
  • reduce_boolean_flags / autotp tests use current_device_name() instead of current_device() (a LOCAL_RANK string on CPU, not a device).
  • fp16-config tests get a capability skipif (not get_accelerator().is_fp16_supported()) — GH ubuntu-24.04 runners are hardware-heterogeneous w.r.t. AVX512-FP16, so hardcoded fp16 was a runner lottery.
  • fork_rng probes the device module for a per-device RNG instead of matching accelerator names.

CI infrastructure findings (need maintainer decisions)

  1. --maxfail=100 in PYTEST_OPTS + no timeout in DistributedExec._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.)
  2. 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.
  3. 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant