StarCoder 2 Documentation

To install StarCoder 2, install the dependencies listed in requirements.txt and export your Hugging Face token to access the models.

pip install -r requirements.txt
export HF_TOKEN=xxx

pip install -r requirements.txt
# export your HF token, found here: https://huggingface.co/settings/account
export HF_TOKEN=xxx

To prepare for fine-tuning, install PyTorch (e.g., for CUDA 12.1) and the project requirements. You must also be logged into wandb and the Hugging Face Hub to push checkpoints.

# Example for CUDA 12.1
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

# Install requirements (installs transformers from source)
pip install -r requirements.txt

# Login to services
wandb login
huggingface-cli login

You can run StarCoder 2 using Docker with the Text-Generation-Inference (TGI) container. Ensure you provide your Hugging Face token via the HUGGING_FACE_HUB_TOKEN environment variable.

docker run -p 8080:80 -v $PWD/data:/data -e HUGGING_FACE_HUB_TOKEN=<YOUR BIGCODE ENABLED TOKEN> -d  ghcr.io/huggingface/text-generation-inference:latest --model-id bigcode/starcoder2-15b --max-total-tokens 8192

To reduce memory usage, you can load the model using torch.bfloat16 or torch.float16. Using torch.bfloat16 requires the accelerate library.

pip install accelerate

For float16, use torch_dtype=torch.float16 instead of bfloat16.

# pip install accelerate
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

checkpoint = "bigcode/starcoder2-15b"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# for fp16 use `torch_dtype=torch.float16` instead
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.bfloat16)

inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to("cuda")
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

You can use bitsandbytes to load the model in 8-bit or 4-bit precision to significantly reduce the memory footprint.

To use 4-bit, set load_in_4bit=True in the BitsAndBytesConfig instead of load_in_8bit=True.

# pip install bitsandbytes accelerate
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# to use 4bit use `load_in_4bit=True` instead
quantization_config = BitsAndBytesConfig(load_in_8bit=True)

checkpoint = "bigcode/starcoder2-15b_16k"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained("bigcode/starcoder2-15b_16k", quantization_config=quantization_config)

inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to("cuda")
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

Fine-tuning is performed efficiently using PEFT (LoRA), bitsandbytes (4-bit quantization), and SFTTrainer from the TRL library. Use accelerate launch to start the training script.

If fine-tuning on custom datasets, ensure you update the --dataset_text_field argument to match the column name containing your code/text.

accelerate launch finetune.py \
        --model_id "bigcode/starcoder2-3b" \
        --dataset_name "bigcode/the-stack-smol" \
        --subset "data/rust" \
        --dataset_text_field "content" \
        --split "train" \
        --max_seq_length 1024 \
        --max_steps 10000 \
        --micro_batch_size 1 \
        --gradient_accumulation_steps 8 \
        --learning_rate 2e-5 \
        --warmup_steps 20 \
        --num_proc "$(nproc)"

The pipeline API provides a high-level way to perform text generation with StarCoder 2.

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
checkpoint = "bigcode/starcoder2-15b"

model = AutoModelForCausalLM.from_pretrained(checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
print( pipe("def hello():") )