vision: Molmo2-Stage2.py smoke recipe OOMs on the step-0 pre-train checkpoint with HF init
- 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
Assessment
This issue has not been assessed yet.