StarCoder 2

[🤗 模型与数据集] | [论文]

StarCoder2是一系列代码生成模型（3B、7B和15B），在来自The Stack v2的600多种编程语言以及一些自然语言文本（如维基百科、Arxiv和GitHub问题）上进行训练。这些模型使用分组查询注意力机制，上下文窗口为16,384个标记，滑动窗口注意力为4,096个标记。3B和7B模型在3万亿以上的标记上训练，而15B模型在4万亿以上的标记上训练。更多详情请查看论文。

快速开始

StarCoder2模型主要用于代码补全，它们不是指令模型，像"编写一个计算平方根的函数"这样的命令效果不佳。

安装

首先，我们需要安装requirements.txt中列出的所有库

pip install -r requirements.txt
# 导出你的HF令牌，可以在这里找到：https://huggingface.co/settings/account
export HF_TOKEN=xxx

模型使用和内存占用

以下是一些加载模型和生成代码的示例，以及最大模型StarCoder2-15B的内存占用情况。确保你已从源代码安装了transformers（如果你使用了requirements.txt，应该已经安装了）

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

在CPU/GPU/多GPU上运行模型

使用全精度

# pip install git+https://github.com/huggingface/transformers.git # TODO: 合并PR到主分支
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "bigcode/starcoder2-15b"
device = "cuda" # 使用GPU，或使用"cpu"进行CPU运算

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
# 使用多个GPU，可以这样做：`model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")`
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]))

使用torch.bfloat16

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

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

# 对于fp16，使用`torch_dtype=torch.float16`
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]))

>>> print(f"内存占用：{model.get_memory_footprint() / 1e6:.2f} MB")
内存占用：32251.33 MB

通过`bitsandbytes`进行量化

使用8位精度（int8）

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

# 使用4位，改用`load_in_4bit=True`
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]))

>>> print(f"内存占用：{model.get_memory_footprint() / 1e6:.2f} MB")
# load_in_8bit
内存占用：16900.18 MB
# load_in_4bit
>>> print(f"内存占用：{model.get_memory_footprint() / 1e6:.2f} MB")
内存占用：9224.60 MB

你也可以使用pipeline进行生成：

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():") )

文本生成推理：

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

更多详情，请参见此处。

微调

这里我们展示如何微调StarCoder2模型。更多微调资源，你可以查看StarCoder的GitHub仓库和SantaCoder-Finetuning。

设置

安装pytorch参见文档，例如以下命令适用于cuda 12.1：

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

安装要求（这将从源代码安装transformers以支持StarCoder2架构）：

pip install -r requirements.txt

在运行任何脚本之前，确保你已登录wandb和HuggingFace Hub以推送检查点：

wandb login
huggingface-cli login

现在一切就绪，你可以克隆仓库并进入相应目录。

训练

为了高效且低成本地微调，我们使用PEFT库进行低秩适应（LoRA）训练，并使用bitsandbytes进行4位量化。我们还使用了来自TRL的SFTTrainer。

在这个例子中，我们将在the-stack-smol的Rust子集上微调StarCoder2-3b。这只是为了说明目的；对于更大更干净的Rust代码数据集，你可以使用The Stack dedup。

启动训练：

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)"

如果你想在其他文本数据集上进行微调，需要将dataset_text_field参数更改为包含你想训练的代码/文本的列名。