Skip to content

examples CI: shrink CPU DDP arm batch to fix gloo timeout#269

Open
wz337 wants to merge 1 commit into
facebookresearch:mainfrom
wz337:export-D112602955
Open

examples CI: shrink CPU DDP arm batch to fix gloo timeout#269
wz337 wants to merge 1 commit into
facebookresearch:mainfrom
wz337:export-D112602955

Conversation

@wz337

@wz337 wz337 commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Summary:
The examples GitHub workflow's "Run DDP example on CPU" step launched two gloo
ranks each at local_batch_size=1024 on the 4-core 4-core-ubuntu-gpu-t4 runner.
The combined memory of two 1024-batch processes plus Shampoo preconditioner state
exceeds the runner, so one rank stalls/OOMs and the peer hangs until the 30-min
gloo recv timeout in backward:

RuntimeError: [gloo/transport/tcp/unbound_buffer.cc:78]
Timed out waiting 1800000ms for recv operation to complete
-> ChildFailedError: distributed_shampoo.examples.cifar10_example FAILED

This is a pre-existing latent flake: the example DDP code (cifar10_example.py,
examples/utils.py, examples/parallelism.py, ddp.yaml) is byte-identical to prior
syncs; it only surfaced once a run got past the CIFAR download. The single-process
CPU steps at batch 1024 pass -- it is the 2x memory of the DDP arm that tips over.

Fix: reduce the CPU DDP step to local_batch_size=128 (the config default), which
exercises the same DDP + gloo + Shampoo path at ~8x lower activation memory. The
GPU DDP step and the single-process steps are unchanged.

Differential Revision: D112602955

Summary:
The examples GitHub workflow's "Run DDP example on CPU" step launched two gloo
ranks each at local_batch_size=1024 on the 4-core `4-core-ubuntu-gpu-t4` runner.
The combined memory of two 1024-batch processes plus Shampoo preconditioner state
exceeds the runner, so one rank stalls/OOMs and the peer hangs until the 30-min
gloo recv timeout in backward:

  RuntimeError: [gloo/transport/tcp/unbound_buffer.cc:78]
    Timed out waiting 1800000ms for recv operation to complete
  -> ChildFailedError: distributed_shampoo.examples.cifar10_example FAILED

This is a pre-existing latent flake: the example DDP code (cifar10_example.py,
examples/utils.py, examples/parallelism.py, ddp.yaml) is byte-identical to prior
syncs; it only surfaced once a run got past the CIFAR download. The single-process
CPU steps at batch 1024 pass -- it is the 2x memory of the DDP arm that tips over.

Fix: reduce the CPU DDP step to local_batch_size=128 (the config default), which
exercises the same DDP + gloo + Shampoo path at ~8x lower activation memory. The
GPU DDP step and the single-process steps are unchanged.

Differential Revision: D112602955
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 17, 2026
@meta-codesync

meta-codesync Bot commented Jul 17, 2026

Copy link
Copy Markdown

@wz337 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D112602955.

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

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant