allenai / allenai/OLMo-core

vision: FlexAttention in eager mode (compile_model=false) OOMs at seq 16384 — Molmo2-Stage2.py smoke recipe

Open
#847 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.