Skip to content

Commit

Permalink
Fix max_memory example on README (#944)
Browse files Browse the repository at this point in the history
* Fix `max_memory` example on README

- The new `max_memory` syntax expects a dictionary
- This change also accounts for multiple devices

* Fix model name in `from_pretrained` on README
  • Loading branch information
MilesCranmer authored Jan 25, 2024
1 parent f1c7574 commit 94c7f2c
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,11 @@ model = AutoModelForCausalLM.from_pretrained(
'decapoda-research/llama-7b-hf',
device_map='auto',
load_in_8bit=True,
max_memory=f'{int(torch.cuda.mem_get_info()[0]/1024**3)-2}GB')
max_memory={
i: f'{int(torch.cuda.mem_get_info(i)[0]/1024**3)-2}GB'
for i in range(torch.cuda.device_count())
}
)
```

A more detailed example, can be found in [examples/int8_inference_huggingface.py](examples/int8_inference_huggingface.py).
Expand Down

0 comments on commit 94c7f2c

Please sign in to comment.