LLM 微调实战：使用 Code-Llama 训练自定义代码数据集

模型初始化

加载预训练的 Code-Llama 模型及分词器。为了节省显存，这里演示了使用 8bit 量化加载的方式。

from datetime import datetime
import os
import sys
import torch

from peft import (
    LoraConfig,
    get_peft_model,
    get_peft_model_state_dict,
    prepare_model_for_int8_training,
)
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForSeq2Seq
from datasets import load_dataset

# 加载数据集
train_dataset = load_dataset('json', data_files='train_data.jsonl', split='train')
eval_dataset = load_dataset('json', data_files='val_data.jsonl', split='train')

# 指定基座模型
base_model = 'CodeLlama-7b-Instruct-hf'

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    load_in_8bit=True,
    torch_dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True
)

tokenizer = AutoTokenizer.from_pretrained(base_model)

微调前效果测试

在开始训练前，先对基座模型进行测试，观察其在未微调状态下的代码生成表现。

tokenizer.pad_token = tokenizer.eos_token
prompt = """You are programming coder.

Now answer the question:

{}"""
prompts = [prompt.format(train_dataset[i]['question']) for i in [1, 20, 32, 45, 67]]

model_input = tokenizer(prompts, return_tensors="pt", padding=True).to("cuda")

model.eval()
with torch.no_grad():
    outputs = model.generate(**model_input, max_new_tokens=300)
    outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

print(outputs)

微调配置

定义 LoRA 配置参数。LoRA 通过低秩分解更新部分权重，显著减少训练参数量。我们针对注意力机制的投影层（q_proj, k_proj, v_proj, o_proj）应用 LoRA。

tokenizer.add_eos_token = True
tokenizer.pad_token_id = 0
tokenizer.padding_side = "left"

def tokenize(prompt):
    result = tokenizer(
        prompt,
        truncation=True,
        max_length=512,
        padding=False,
        return_tensors=None,
    )
    # 自监督学习：标签与输入一致
    result["labels"] = result["input_ids"].copy()
    return result

def generate_and_tokenize_prompt(data_point):
    full_prompt = f"""You are a powerful programming model. Your job is to answer questions about a database. You are given a question.

You must output the code that answers the question.

### Input:
{data_point["question"]}

### Response:
{data_point["answer"]}
"""
    return tokenize(full_prompt)

tokenized_train_dataset = train_dataset.map(generate_and_tokenize_prompt)
tokenized_val_dataset = eval_dataset.map(generate_and_tokenize_prompt)

model.train() 
model = prepare_model_for_int8_training(model)

config = LoraConfig(
    r=16,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)

if torch.cuda.device_count() > 1:
    model.is_parallelizable = True
    model.model_parallel = True

训练参数设置

配置训练超参数，包括批次大小、学习率、步数等。使用梯度累积来模拟更大的批次。

batch_size = 128
per_device_train_batch_size = 32
gradient_accumulation_steps = batch_size // per_device_train_batch_size
output_dir = "code-llama-ft"

training_args = TrainingArguments(
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    warmup_steps=100,
    max_steps=400,
    learning_rate=3e-4,
    fp16=True,
    logging_steps=10,
    optim="adamw_torch",
    evaluation_strategy="steps",
    save_strategy="steps",
    eval_steps=20,
    save_steps=20,
    output_dir=output_dir,
    load_best_model_at_end=False,
    group_by_length=True,
    report_to="none",
    run_name=f"codellama-{datetime.now().strftime('%Y-%m-%d-%H-%M')}",
)

trainer = Trainer(
    model=model,
    train_dataset=tokenized_train_dataset,
    eval_dataset=tokenized_val_dataset,
    args=training_args,
    data_collator=DataCollatorForSeq2Seq(
        tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
    ),
)

开始训练

启动训练过程。注意在较新版本的 PyTorch 中可启用编译加速。

model.config.use_cache = False

old_state_dict = model.state_dict
model.state_dict = (lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())).__get__(
    model, type(model)
)
if torch.__version__ >= "2" and sys.platform != "win32":
    print("compiling the model")
    model = torch.compile(model)

trainer.train()

推理测试

训练完成后，加载微调后的适配器（Adapter）进行测试。这允许我们在不修改基座模型权重的情况下应用新知识。

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base_model = 'CodeLlama-7b-Instruct-hf'
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    load_in_8bit=True,
    torch_dtype=torch.float16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(base_model)

output_dir = "code-llama-ft"
model = PeftModel.from_pretrained(model, output_dir)

eval_prompt = """You are a powerful programming model. Your job is to answer questions about a database. You are given a question.

You must output the code that answers the question.

### Input:
Write a function in Java that takes an array and returns the sum of the numbers in the array, or 0 if the array is empty. Except the number 13 is very unlucky, so it does not count any 13, or any number that immediately follows a 13.

### Response:
"""

model_input = tokenizer(eval_prompt, return_tensors="pt").to("cuda")

model.eval()
with torch.no_grad():
    outputs = model.generate(**model_input, max_new_tokens=100)[0]
print(tokenizer.decode(outputs, skip_special_tokens=True))

总结与最佳实践

通过上述流程，我们成功实现了对 Code-Llama 模型的定制化微调。在实际应用中，建议关注以下几点：

1. 数据质量：微调效果高度依赖数据的质量。确保问答对准确且符合目标场景的编码风格。
2. 超参数调整：根据显存情况调整 r 值和学习率。过大的学习率可能导致灾难性遗忘。
3. 评估指标：除了人工查看生成结果，可结合 BLEU、ROUGE 或代码执行通过率等指标进行评估。
4. 资源优化：LoRA 仅训练少量参数，使得单卡微调成为可能。如需更大规模训练，可考虑多卡分布式部署。

此方案为垂直领域代码大模型的落地提供了轻量级且高效的解决方案。