Views
No views yet

pip install transformers1# pip install transformers
2from transformers import AutoModelForCausalLM, AutoTokenizer
3checkpoint = "HuggingFaceTB/SmolLM-360M"
4device = "cuda" # for GPU usage or "cpu" for CPU usage
5tokenizer = AutoTokenizer.from_pretrained(checkpoint)
6# for multiple GPUs install accelerate and do `model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")`
7model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
8inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to(device)
9outputs = model.generate(inputs)
10print(tokenizer.decode(outputs[0]))
11
12* _Using `torch.bfloat16`_
13```python
14# pip install accelerate
15import torch
16from transformers import AutoTokenizer, AutoModelForCausalLM
17checkpoint = "HuggingFaceTB/SmolLM-360M"
18tokenizer = AutoTokenizer.from_pretrained(checkpoint)
19# for fp16 use `torch_dtype=torch.float16` instead
20model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.bfloat16)
21inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to("cuda")
22outputs = model.generate(inputs)
23print(tokenizer.decode(outputs[0]))1>>> print(f"Memory footprint: {model.get_memory_footprint() / 1e6:.2f} MB")
2Memory footprint: 723.56 MBbitsandbytes1# pip install bitsandbytes accelerate
2from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
3# to use 4bit use `load_in_4bit=True` instead
4quantization_config = BitsAndBytesConfig(load_in_8bit=True)
5checkpoint = "HuggingFaceTB/SmolLM-360M"
6tokenizer = AutoTokenizer.from_pretrained(checkpoint)
7model = AutoModelForCausalLM.from_pretrained(checkpoint, quantization_config=quantization_config)
8inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to("cuda")
9outputs = model.generate(inputs)
10print(tokenizer.decode(outputs[0]))1>>> print(f"Memory footprint: {model.get_memory_footprint() / 1e6:.2f} MB")
2# load_in_8bit
3Memory footprint: 409.07 MB
4# load_in_4bit
5>>> print(f"Memory footprint: {model.get_memory_footprint() / 1e6:.2f} MB")
6Memory footprint: 251.79 MB1@misc{allal2024SmolLM,
2 title={SmolLM - blazingly fast and remarkably powerful},
3 author={Loubna Ben Allal and Anton Lozhkov and Elie Bakouch and Leandro von Werra and Thomas Wolf},
4 year={2024},
5}