The adept/fuyu-8b
model is broken for device_map=auto
(e.g. shard layers across gpu).
Edit: Started writing this and was going to upload to make patching model on other stuff easier but apparently they fixed it and didn't close the gh issue:
huggingface/transformers#29880
Just pushing incase other things need to be fixed and will push model.
File paths for the env install:
/data/graham/code/hf-fuyu/.venv/lib/python3.11/site-packages/transformers/models/persimmon/modeling_persimmon.py
- this is probably where the OOM will occur, e.g. line ~338 in which does the upcast:
attn_weights = nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query_states.dtype)
which you might want to change dtype to float16 (debateable wether this impacts results)
- this is probably where the OOM will occur, e.g. line ~338 in which does the upcast:
/data/graham/code/hf-fuyu/.venv/lib/python3.11/site-packages/transformers/models/fuyu/modeling_fuyu.py
- this file contains the
gather_continuous_embeddings
and similar that is what likely is broken in the base model
- this file contains the
use python >=3.11 and then just pdm sync
Very confused. Seems like sometimes the continuous_embeddings[batch_idx]
are not on the correct device on the base model and other times they are. Not clear if its due to the size of the inputs or the way infer_auto_device_map
shards across devices.