allenai / allenai/molmo2

Vision CP crashes when a packed sample cannot populate every CP rank

Open
#26 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
726
Forks
48
PR merge metrics
No merged PRs in 30d

Description

### Reproduction

- Molmo2-4B SFT checkpoint on 8 A100s with CP=8
- `model.apply_cp_to_vision_backbone=true`
- Debug mixture, dynamic packing, 36,864-token context, device batch size 1

### Actual behavior

On the first training step, ranks 5-7 receive image boundaries with identical
start and end indices. The vision backbone passes the resulting zero-image tensor
to SDPA, which fails with:

```text
RuntimeError: batch size must be positive
```

The remaining ranks wait in NCCL until the watchdog terminates the job.
**`compute_image_shard_boundaries()` tries to move whole image/frame boundary units
between ranks, but it cannot make every shard nonempty when there are fewer usable
units than CP ranks.**

Could you kindly provide some help and guidance on this? Thank you! @chrisc36 @sangho-vision @zixianm

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.