Incompatible devices for jitted computation

[versae]

Incompatible devices for jitted computation

[rjpower]

Ah interesting, I hit this same error when trying to use llama2-lora. I found if I disabled the "best-effort sharding" in hf_checkpoints I could get training to work:

levanter/src/levanter/compat/hf_checkpoints.py

Line 1072 in 9adba01

if jax.device_count() > 1:

[dlwh]

ok i think can be fixed with a long-standing thing I've been meaning to fix in Haliax (but isn't quite compatible with some of the CPU-only testing I do)

[dlwh]

can you try #622, which is a simpler fix than the haliax thing

[dlwh]

Well that's a remarkably unhelpful error message. I have a strong suspicion it's because of #588 and it's now using an incompatible mesh. Can you try reverting that commit and seeing if that fixes? If it fixes, it's a pretty easy patch

[versae]

Indeed. I just checked out 2bb1252 and it trains nicely, but I had to reduce parallelism to 8 to make it fit into the v4-64, strange.