StarCoder2-15B: A 15B Parameter Model for Programming Languages with Grouped Query Attention and Fill-in-the-Middle Training

Technical Overview of StarCoder2-15B

StarCoder2-15B is a massive language model specifically trained on over 600 programming languages, making it a powerful tool for developers and researchers. This 15 billion parameter model leverages several advanced techniques to enhance its performance and efficiency:

• Grouped Query Attention (GQA): GQA reduces the computational load by grouping queries together, which is particularly useful for large models like StarCoder2-15B.
• Context Window of 16,384 Tokens: The model can process a long context window, allowing it to understand and generate more coherent code sequences.
• Sliding Window Attention of 4,096 Tokens: This technique helps the model maintain focus on relevant parts of the input without losing context.
• Fill-in-the-Middle (FiM) Objective: Instead of predicting the next token, FiM trains the model to predict tokens in the middle of a sequence, which can improve its ability to handle complex code structures.

Training Details

The training process for StarCoder2-15B was extensive and resource-intensive:

• Data Source: The model was trained on The Stack v2, a dataset of 600+ programming languages, with opt-out requests excluded.
• Token Count: Trained on over 4 trillion tokens, ensuring a diverse and comprehensive understanding of various coding styles and practices.
• Hardware: The training was conducted using the NVIDIA NeMo™ Framework on the NVIDIA Eos Supercomputer, which is built with NVIDIA DGX H100 systems.

Intended Use

StarCoder2-15B is designed to be a powerful tool for code generation and understanding, but it's important to note that it is not an instruction model. Commands like "Write a function that computes the square root" are less effective compared to providing context and partial code snippets. The model excels when given specific code contexts and tasks.

Getting Started with StarCoder2-15B

To use StarCoder2-15B, you need to install the transformers library from source:

pip install git+https://github.com/huggingface/transformers.git

Running the Model on CPU/GPU/Multi-GPU

Here are some examples of how to run the model in different environments:

• Full Precision (CPU/GPU):

from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "bigcode/starcoder2-15b"
device = "cuda"  # Use 'cuda' for GPU or 'cpu' for CPU

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

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

• Using torch.bfloat16 (Multi-GPU):

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "bigcode/starcoder2-15b"
device = "cuda"  # Use 'cuda' for GPU or 'cpu' for CPU

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, torch_dtype=torch.bfloat16, device_map="auto")

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

Limitations

While StarCoder2-15B is a powerful tool, it has some limitations:

• Non-Instructional: As mentioned, the model is not designed to follow explicit