From d6dbf590b5bf40ff4b0f6a9d1ce32c0351d5accd Mon Sep 17 00:00:00 2001 From: Iris Zhang Date: Fri, 17 Jul 2026 15:25:54 -0700 Subject: [PATCH] examples CI: shrink CPU DDP arm batch to fix gloo timeout 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 --- .github/workflows/examples.yaml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/examples.yaml b/.github/workflows/examples.yaml index 5a0fe19..8256157 100644 --- a/.github/workflows/examples.yaml +++ b/.github/workflows/examples.yaml @@ -48,7 +48,11 @@ jobs: - name: Run DDP example on CPU. run: | source .venv/bin/activate - CUDA_VISIBLE_DEVICES="" torchrun --standalone --nnodes=1 --nproc_per_node=2 -m distributed_shampoo.examples.cifar10_example parallelism=ddp optimizer=shampoo optimizer.precondition_frequency=15 '~optimizer.grafting_config' '+optimizer.grafting_config={_target_:distributed_shampoo.AdamPreconditionerConfig,beta2:0.999,epsilon:1e-8}' epochs=1 local_batch_size=1024 backend=gloo + # Keep the per-rank batch small here: two gloo ranks each at + # local_batch_size=1024 exhaust memory on the 4-core CPU runner, so one + # rank stalls/OOMs and the peer hangs until the 30-min gloo recv timeout + # in backward. 128 exercises the same DDP + Shampoo path within budget. + CUDA_VISIBLE_DEVICES="" torchrun --standalone --nnodes=1 --nproc_per_node=2 -m distributed_shampoo.examples.cifar10_example parallelism=ddp optimizer=shampoo optimizer.precondition_frequency=15 '~optimizer.grafting_config' '+optimizer.grafting_config={_target_:distributed_shampoo.AdamPreconditionerConfig,beta2:0.999,epsilon:1e-8}' epochs=1 local_batch_size=128 backend=gloo - name: Run DDP example on GPU. if: steps.gpu_check.outputs.has_gpu == 'true' run: |