allenai / allenai/OLMo-core

vision: Molmo2-Stage2.py smoke recipe OOMs on the step-0 pre-train checkpoint with HF init

Open
#846 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`

The documented 1-GPU smoke recipe in `src/scripts/train/Molmo2-Stage2.py` (with `--trainer.load_path=null`, i.e. HF init from `allenai/Molmo2-4B`) OOMs before step 1, inside `CheckpointerCallback.pre_train`.

**Root cause:** `CheckpointerCallback.pre_train_checkpoint=None` means "save at step 0 unless a checkpoint was loaded". HF init loads no trainer checkpoint, so the step-0 save fires, and torch DCP's `_get_optim_state_dict → _init_optim_state` calls `optim.step()` — force-allocating the full fp32 Adam states (~34 GiB for the 4B model) on top of the fp32 master + bf16 compute copies. On one 80 GB H100 that OOMs:

```
OutOfMemoryError: CUDA out of memory. Tried to allocate 96.00 MiB ... 77.93 GiB is allocated by PyTorch
File ".../train/callbacks/checkpointer.py", line 239, in pre_train
...
File "torch/distributed/checkpoint/state_dict.py", line 648, in _init_optim_state
optim.step(closure=None)
```

The default stage-1-checkpoint init path never hits this (`trainer.checkpoint_loaded` suppresses the save), which is probably why the smoke recipe was never observed failing.

**Suggested fixes:** document `--trainer.callbacks.checkpointer.pre_train_checkpoint=false` in the smoke recipe, and/or avoid materializing optimizer state for a step-0 save (it is empty by construction).

Found while porting the stage-2 recipe into allenai/open-instruct (open-instruct#1856); we work around it by exposing `pre_train_checkpoint=False` for HF-init runs.

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.