allenai / allenai/OLMo-core

vision: compile_vision/compile_connector backward hits an inductor stride AssertionError when crop counts vary across batches

Aperta
#848 3 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
1.5k
Fork
315
Merge medio
1g 9h
PR unite (30g)
11

Descrizione

**Branch/rev:** `vision` @ `c54d228`

With `compile_vision=true`/`compile_connector=true` (the `MultimodalTransformerTrainModuleConfig` defaults), training on a mixture where per-batch crop counts swing hard — e.g. a text-dominant mixture (90% zero-crop text, 10% image VQA) alternating between dummy-crop and image-heavy packs — fails in the compiled backward:

```
AssertionError: expected size 8==8, stride 839808==840704 at dim=0; expected size 16==16, stride 52488==52544 at dim=1
File "torch/_inductor/output_code.py", line 638, in __call__
... (during loss.backward(), AOT-autograd saved-tensor stride guard)
```

The forward compiles under one saved-activation stride layout and the backward is invoked with another after the crop count changes. Not reproduced on mixtures with steadier crop counts (a 10-step run on the 3-source debug mixture passes with vision compiled).

**Workaround:** run vision/connector eager (`compile_vision=false compile_connector=false`); only the LM needs compilation (for flex-attention memory).

**Repro sketch:** stage-2 debug mixture with the tulu4 (text-only) source weighted to 0.9, 2×H100, seq 16384, `torch 2.10` — fails within the first steps.

Found while porting the stage-2 recipe into allenai/open-instruct (open-instruct#1856).

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.