Gemlite generate.py fix #2372
Gemlite generate.py fix #2372
Conversation
Signed-off-by: mobicham <[email protected]>
Signed-off-by: mobicham <[email protected]>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2372
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 862e0cb with merge base 6243040 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| def _load_model(checkpoint_path, device, precision): | ||
| checkpoint = torch.load( | ||
| str(checkpoint_path), mmap=True, weights_only=True, map_location=device | ||
| str(checkpoint_path), mmap=True, weights_only=True, map_location="cpu" |
There was a problem hiding this comment.
actually is the map_location change correct?
There was a problem hiding this comment.
Yes, without this, you can't load a Llama3-8B fp16 on a 24GB gpu, you just get OOM.
This is the part I mentioned:
Also, the current code gives OOM on smaller gpus. By putting the weights on cpu first, we can avoid this issue.
There was a problem hiding this comment.
oh OK, but do we need to change it somewhere or it has to run in cpu
I think if user request to run on some device but it fails, we should probably not silently change the device, but instead ask user to use a different device instead
There was a problem hiding this comment.
No it still runs on gpu, there's a model.to(device): https://github.com/mobicham/ao/blob/862e0cbd90f8cc5f992ae66e779022912fa4d93a/torchao/_models/llama/generate.py#L245-L256
The issue is that, loading the weights via map_location='cuda' + doing model.to(device) for some reason uses more vram leading to oom. Maybe there's a cleaner way of doing it.
There was a problem hiding this comment.
oh OK, makes sense, this should be OK, some quantized format does not support loading in CPU but for full/half precision model it should be fine
* fix get_plain() with FMA mode * update * fix in_features/out_feature meta-data mismatch * update gemlite slice test * add packing_bitwidth support * add packing_bitwidth support and cleanup * update default gemlite layout * cleanup * fix symmetric use-case and relax _same_meta_data * _copy() meta data * fix (4,) in autoquant * Add dynamic mode in gemlite layout * mode explanation Signed-off-by: mobicham <[email protected]> * use weights_only instead of static * generate fix Signed-off-by: mobicham <[email protected]> * remove set_packing_bitwidth --------- Signed-off-by: mobicham <[email protected]>
Fixes
generate.pybenchmarking with gemlite.Also, the current code gives OOM on smaller gpus. By putting the weights on cpu first, we can avoid this issue.