allenai / allenai/OLMo-core

FSDP2 post-backward reduce-scatter hangs on GDN hybrid models (Olmo-Hybrid-7B), 4/4 runs

Aperta
#829 2 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
1.5k
Fork
315
Merge medio
1g 9h
PR unite (30g)
11

Descrizione

## Summary

SFT of `allenai/Olmo-Hybrid-7B` (GDN + attention hybrid) on a single 8xH100 node
reliably goes silent mid-training. The stalled operation is the FSDP2 post-backward
gradient reduce-scatter. All 8 ranks enqueue the same collective and none of them
completes, so the job holds its GPUs until the process-group timeout fires.

Reproduced 4/4 times across 4 different nodes, at unpredictable steps.

## Symptom

Training logs stop mid-run with no traceback and no NCCL error. With OLMo-core's
default arrangement the process group is built with a long timeout, so nothing aborts
the job — one occurrence held 8 GPUs for 2h40m before a human noticed. Setting the
timeout to 30 minutes produces this:

```
[Rank 4] Watchdog caught collective operation timeout:
WorkNCCL(SeqNum=27732, OpType=_REDUCE_SCATTER_BASE,
NumelIn=215585536, NumelOut=26948192, Timeout(ms)=1800000)
ran for 1800003 milliseconds before timing out.

[Rank 4] failure detected by watchdog at work sequence id: 27732
PG status: last enqueued work: 27753, last completed work: 27731
```

Every rank reports the same `SeqNum` with identical shapes, and every rank has enqueued
past it (`last enqueued 27753` vs `last completed 27731`). So this is not rank divergence
or a shape mismatch: the collective is posted everywhere and never executes. The CPU runs
~20 collectives ahead while the device is stalled, which is why the last logged step is a
few steps beyond where execution actually froze.

With `TORCH_NCCL_TRACE_BUFFER_SIZE` set, the stack is identical on all ranks:

```
#0 reduce_scatter_tensor torch/distributed/distributed_c10d.py:4591
#1 wrapper torch/distributed/c10d_logger.py:83
#2 __call__ torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py:125
#3 foreach_reduce torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py:543
#4 decorate_context torch/utils/_contextlib.py:124
#5 post_backward torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py:566
#6 backward torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py:900
#7 apply torch/autograd/function.py:317
```

## Reproductions

| run | hung at step | node | checkpointing | fla |
|---|---|---|---|---|
| 1 | 580 | jupiter-cs-aus-137 | async | 0.5.0 |
| 2 | 350 | jupiter-cs-aus-123 | sync | 0.5.0 |
| 3 | 65 | jupiter-cs-aus-149 | sync | 0.5.0 |
| 4 | 205 | jupiter-cs-aus-219 | async | 0.5.2 |

## What we ruled out

- **Checkpointing.** Identical with `save_async` true and false, and run 3 hung at step 65
with the first checkpoint scheduled at 172, i.e. before any save had occurred.
- **A bad node.** Four distinct nodes.
- **fla version.** Same signature on 0.5.0 and 0.5.2 (the latter being the version
OLMo-core mainline moved to in #798).
- **Checkpoint deletion / storage.** Runs used `keep_last_n_checkpoints -1`, so no deletion.

Not ruled out, and the obvious remaining suspects: `torch.compile`, the GDN `fla` kernels
themselves, and their interaction with selected-modules activation checkpointing. We could
not test with compile disabled because the model does not fit at seq 32768 without it
(OOM at ~67 GiB allocated with 12.25 GiB more requested, even after eliminating
fragmentation with `expandable_segments`).

## Configuration

- `allenai/Olmo-Hybrid-7B`, 1 node x 8 H100, seq len 32768
- global batch 1,048,576 tokens (per-device 1, grad accum 4)
- FSDP2, `torch.compile` enabled, activation checkpointing over selected modules
(`blocks.*.attention` etc. with `determinism_check="none"`), `flash_2`
- OLMo-core pinned at `fa6c5014c9f6e9ee789da2d9c20d5126fee8df0d` (2026-06-22)
- torch 2.10, driven from allenai/open-instruct's `olmo_core_finetune.py`

Happy to re-run with extra instrumentation if that would help — the reproduction is
cheap to trigger (it has never survived past step 580) and we have FlightRecorder wired
up already.

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.