Instructions to use google/gemma-2-2b-it with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use google/gemma-2-2b-it with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="google/gemma-2-2b-it")
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=40)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Inference
Local Apps Settings

vLLM

How to use google/gemma-2-2b-it with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "google/gemma-2-2b-it"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "google/gemma-2-2b-it",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/google/gemma-2-2b-it

SGLang

How to use google/gemma-2-2b-it with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "google/gemma-2-2b-it" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "google/gemma-2-2b-it",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "google/gemma-2-2b-it" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "google/gemma-2-2b-it",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use google/gemma-2-2b-it with Docker Model Runner:
```
docker model run hf.co/google/gemma-2-2b-it
```

GPU training makes loss=nan

#37

by hidonbush - opened Sep 5, 2024

Discussion

hidonbush

Sep 5, 2024

When I was testing the model, I encountered the problem that if I use GPU, the loss becomes nan. If it's on CPU, everything is OK.

from transformers import AutoTokenizer,AutoModelForSequenceClassification
from peft import get_peft_model, LoraConfig, TaskType, PeftModel
import torch

peft_config = LoraConfig(
task_type=TaskType.SEQ_CLS,
inference_mode=False,
r=8,
lora_alpha=32,
lora_dropout=0.1,
target_modules=['down_proj','o_proj','k_proj','q_proj','gate_proj','up_proj','v_proj']
)

temp=AutoModelForSequenceClassification.from_pretrained("gemma2b",num_labels=2,torch_dtype = torch.bfloat16,device_map='auto')
model = get_peft_model(temp, peft_config)
tokenizer = AutoTokenizer.from_pretrained("gemma2b")
testdata=tokenizer('I like it too'+'',return_tensors='pt',padding='max_length',max_length=10).to('cuda')

print(model(**testdata,labels=torch.tensor(1).to('cuda')))

transformers is the latest
PyTorch:2.4
CUDA:12.1
Python:3.10

GopiUppari

Google org Sep 9, 2024

Hi @hidonbush ,

I executed the provided code in Google Colab with the runtime type set to 'T4 GPU' using the following library versions, and I did not encounter a NaN loss when using the GPU. For more details, please refer to the attached Gist notebook : https://colab.research.google.com/gist/Gopi-Uppari/e285595fac24fc55ae09a688b8e9d9b9/gpu-training-makes-loss-nan.ipynb#scrollTo=pzy7UtnkTuwx

torch: 2.4.0+cu121
Python: 3.10.12
peft: 0.12.0
transformers: 4.44.2

Please let us know if the issue still persists.

Thank you.

hidonbush

Sep 10, 2024

Hi @GopiUppari ,
I also tested the code on T4 and V100, it worked.
But if the GPU is Ampere or higher like 3090 or 4090, RTX6000, the problem occurs.
Here are the packages in the virtual conda environment:

accelerate 0.34.2
aiohappyeyeballs 2.4.0
aiohttp 3.10.5
aiosignal 1.3.1
async-timeout 4.0.3
attrs 24.2.0
certifi 2024.8.30
charset-normalizer 3.3.2
filelock 3.13.1
frozenlist 1.4.1
fsspec 2024.2.0
huggingface-hub 0.24.6
idna 3.8
Jinja2 3.1.3
lightning 2.4.0
lightning-utilities 0.11.7
MarkupSafe 2.1.5
mpmath 1.3.0
multidict 6.0.5
networkx 3.2.1
numpy 1.26.3
nvidia-cublas-cu12 12.4.2.65
nvidia-cuda-cupti-cu12 12.4.99
nvidia-cuda-nvrtc-cu12 12.4.99
nvidia-cuda-runtime-cu12 12.4.99
nvidia-cudnn-cu12 9.1.0.70
nvidia-cufft-cu12 11.2.0.44
nvidia-curand-cu12 10.3.5.119
nvidia-cusolver-cu12 11.6.0.99
nvidia-cusparse-cu12 12.3.0.142
nvidia-nccl-cu12 2.20.5
nvidia-nvjitlink-cu12 12.4.99
nvidia-nvtx-cu12 12.4.99
packaging 24.1
peft 0.12.0
pillow 10.2.0
pip 24.2
psutil 6.0.0
pytorch-lightning 2.4.0
PyYAML 6.0.2
regex 2024.7.24
requests 2.32.3
safetensors 0.4.5
setuptools 72.1.0
sympy 1.12
tokenizers 0.19.1
torch 2.4.1+cu124
torchaudio 2.4.1+cu124
torchmetrics 1.4.1
torchvision 0.19.1+cu124
tqdm 4.66.5
transformers 4.44.2
triton 3.0.0
typing_extensions 4.9.0
urllib3 2.2.2
wheel 0.43.0
yarl 1.11.0

nbroad

Sep 10, 2024

Your labels tensor is the wrong shape. This is what it should look like to not get a nan loss

print(model(**testdata,labels=torch.tensor([1]).to('cuda')))

# SequenceClassifierOutputWithPast(loss=tensor(0.6836, device='cuda:0', dtype=torch.bfloat16, grad_fn=<NllLossBackward0>), logits=tensor([[0.8047, 0.8203]], device='cuda:0', dtype=torch.bfloat16, grad_fn=<IndexBackward0>), past_key_values=None, hidden_states=None, attentions=None)

hidonbush

Sep 11, 2024

•

edited Sep 11, 2024

Okay guys, there is the solution:
https://github.com/huggingface/transformers/issues/32390
There might be some problem with the default attention implementation in the precision of bf16

pauliussasnauskas

Sep 18, 2024

Okay guys, there is the solution:
https://github.com/huggingface/transformers/issues/32390
There might be some problem with the default attention implementation in the precision of bf16

This fixed it for me. It appears there is some bug with the sdpa attention implementation in bf16. The solution for now is to change to a different attention implementation. I was loading the model via:

model = AutoModelForSequenceClassification.from_pretrained(model_config.model_name_or_path, num_labels=1, trust_remote_code=model_config.trust_remote_code, **model_kwargs)

When I set model_kwargs.attn_implementation = "flash_attention_2" it stopped being nans

Upload images, audio, and videos by dragging in the text input, pasting, or clicking here.

Tap or paste here to upload images

· Sign up or log in to comment