vision: FlexAttention in eager mode (compile_model=false) OOMs at seq 16384 — Molmo2-Stage2.py smoke recipe
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 315
- Avg merge
- 1d 9h
- Merged PRs (30d)
- 11
Description
**Branch/rev:** `vision` @ `c54d228`
`Molmo2-Stage2.py`'s documented smoke recipe passes `--train_module.compile_model=false` while `USE_FLEX_ATTN=True` sets `OLMO2_FLEX_ATTN=1`. With the flex backend active but uncompiled, attention runs eager flex, which materializes enough intermediates at seq 16384 to OOM two 80 GB H100s (~77 GiB/rank; the identical run passes with `compile_model=true`, peaking ~10 GiB lower):
```
[olmo_core.nn.attention:471] Using attention backend 'flex'
OutOfMemoryError: CUDA out of memory ... 76.68 GiB is allocated by PyTorch
File "olmo_core/nn/transformer/block.py", line 161, in forward
...
```
**Suggested fixes:** the smoke recipe should keep compile on (matching production), or the flex backend could warn/fall back to SDPA when the model is not compiled.
Found while porting the stage-2 recipe into allenai/open-instruct (open-instruct#1856); together with the pre-train-checkpoint OOM (#846) this makes the documented 1-GPU smoke path unrunnable as written.
Contributor guide
Assessment
This issue has not been assessed yet.