!pip install bitsandbytes
We can then load models in 8-bit quantization by simply adding a load_in_8bit=True flag to from_pretrained.
python
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True, pad_token_id=0)
Now, let's run our example again and measure the memory usage.
thon
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result