starcoder2

发布时间:2025-05-15 00:58

StarCoder 2

<p align="center"><a href="https://huggingface.co/bigcode">[ 模型与数据集]</a> | <a href="https://arxiv.org/abs/2402.19173">[论文]</a></p>

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参数更改为包含你想训练的代码/文本的列名。

评估

要评估StarCoder2及其衍生模型,你可以使用BigCode-Evaluation-Harness来评估代码LLM。你还可以查看BigCode排行榜。

网址:starcoder2 https://m.mxgxt.com/news/view/1196718

相关内容

人人都是程序员,英伟达联合推出 StarCoder2 模型
NVIDIA 与合作伙伴共同发布面向企业应用全新开放获取大数据模型

随便看看