多GPU环境下Llama-2-7b-chat-hf微调显存不足问题求助

我使用搭载6张3070 Ti（总计48GB显存）、32GB内存及150GB SSD的Ubuntu 20.04桌面设备，尝试用从Kaggle的diagnoise-me数据集中截取的小数据集，对meta-llama/Llama-2-7b-chat-hf进行微调以做概念验证：

• 最初用DataParallel搭配8GB内存运行脚本时，因内存不足失败；升级至32GB内存后，仍报CUDA显存不足错误：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate
64.00 MiB. GPU 0 has a total capacty of 7.58 GiB of which 17.88 MiB is free. Including non-PyTorch memory, this process has 7.53 GiB memory
in use. Of the allocated memory 7.37 GiB is allocated by PyTorch, and
5.71 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid
fragmentation. See documentation for Memory Management and
PYTORCH_CUDA_ALLOC_CONF

• 用glances观察到仅第一张GPU在工作，任务未分配至所有GPU；改用DistributedDataParallel后，32GB内存仍出现内存不足。

想请教：

• 推荐使用ZeRO、DP、DDP、PP、TP或它们的组合？还是Accelerate、Optimum（NVIDIA或其他）？亦或是TrainingArguments？
• 希望得到方向指引及可在我设备上测试的运行代码示例。

环境复现步骤 #

mkdir test && cd test
python3 -m venv venv
source venv/bin/activate
pip install torch transformers
# 创建代码文件
# 下载数据集到当前目录

接近可用的DP实现代码 #

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
import json

# 从JSON文件加载数据集
with open('small_en_medical_dialog.json', 'r') as file:
    dataset = json.load(file)

# 格式化数据
formatted_data = []
for item in dataset:
    input_text = "Description: " + item["Description"] + " Patient: " + item["Patient"]
    target_text = item["Doctor"]
    formatted_data.append({"input": input_text, "target": target_text})

# 定义自定义Dataset类
class DoctorPatientDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.data = data
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        input_encoding = self.tokenizer(
            item['input'],
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        target_encoding = self.tokenizer(
            item['target'],
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        return {
            'input_ids': input_encoding['input_ids'].flatten(),
            'attention_mask': input_encoding['attention_mask'].flatten(),
            'labels': target_encoding['input_ids'].flatten()
        }

# 初始化模型和分词器
model_name = "meta-llama/Llama-2-7b-chat-hf"

tokenizer = AutoTokenizer.from_pretrained(model_name)
# 如果未设置padding token则设置
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(model_name)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model = torch.nn.DataParallel(model)

# 训练设置
optimizer = AdamW(model.parameters(), lr=5e-5)
num_epochs = 3

# 创建数据集和DataLoader
train_dataset = DoctorPatientDataset(formatted_data, tokenizer)
train_dataloader = DataLoader(train_dataset, batch_size=3, shuffle=True)

# 训练循环
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch in train_dataloader:
        optimizer.zero_grad()
        outputs = model(input_ids=batch['input_ids'],
                        attention_mask=batch['attention_mask'],
                        labels=batch['labels'])
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(train_dataloader)}")

# 保存微调后的模型
model.module.save_pretrained("my_finetuned_model")

# 推理
model.eval()
query = "Hello, how are you?"
input_ids = tokenizer.encode(query, return_tensors='pt')
with torch.no_grad():
    output = model(input_ids=input_ids)
    response_ids = output.logits.argmax(-1)
    response = tokenizer.decode(response_ids, skip_special_tokens=True)

print(response)

方案选择建议 #

1. 优先用DDP+ZeRO Stage 2：Llama-2-7b单卡全精度加载需~28GB，远超单卡8GB显存。ZeRO Stage 2可拆分优化器状态与梯度到多卡，结合DDP分布式训练，能充分利用所有GPU显存，解决单卡负载过高问题。
2. 用Hugging Face Accelerate简化配置：Accelerate封装了DDP、ZeRO等复杂逻辑，无需手动编写分布式启动代码，适合快速验证。
3. 弃用DataParallel：DP设计缺陷导致仅单卡存完整模型，其他卡仅做计算，显存利用率极低，这就是你之前只有第一张GPU工作的核心原因。
4. TrainingArguments+Trainer：来自Transformers库，内置梯度累积、FP16等显存优化选项，适合快速上手；若需自定义训练循环，优先选Accelerate。

可测试的代码示例（Accelerate+ZeRO Stage 2+LoRA） #

第一步：安装额外依赖 #

在激活的虚拟环境中执行：

pip install accelerate peft bitsandbytes

• bitsandbytes：支持量化训练，进一步降低显存占用；
• peft：实现LoRA微调，仅训练部分参数，显存占用大幅降低，适合小样本验证；
• accelerate：处理分布式训练与ZeRO优化逻辑。

第二步：配置Accelerate #

运行以下命令生成配置文件，按提示选择：

• 选择Yes启用分布式训练；
• 选择DDP作为分布式框架；
• 开启ZeRO优化并选择Stage 2；
• 其余选项默认即可。

accelerate config

第三步：训练代码（保存为train.py） #

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from torch.utils.data import Dataset
import json
import torch
from peft import LoraConfig, get_peft_model

# 加载并格式化数据集（适配Llama-2-Chat对话格式）
with open('small_en_medical_dialog.json', 'r') as file:
    dataset = json.load(file)

formatted_data = []
for item in dataset:
    text = f"[INST] Description: {item['Description']} Patient: {item['Patient']} [/INST] {item['Doctor']}"
    formatted_data.append({"text": text})

# 自定义数据集类
class MedicalDialogDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.data = data
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        text = self.data[idx]['text']
        encoding = self.tokenizer(
            text,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        # 标签与input_ids一致，将padding部分设为-100避免计算损失
        encoding['labels'] = encoding['input_ids'].clone()
        padding_mask = encoding['attention_mask'] == 0
        encoding['labels'][padding_mask] = -100
        return {k: v.flatten() for k, v in encoding.items()}

# 初始化分词器与模型
model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 启用FP16加载模型，自动分配到多卡
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 配置LoRA微调（仅训练部分参数）
lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 查看可训练参数占比

# 训练参数设置
training_args = TrainingArguments(
    output_dir="./llama2-medical-finetune",
    per_device_train_batch_size=2,  # 根据单卡显存调整，2适配8GB显存
    gradient_accumulation_steps=2,  # 梯度累积等效增大batch size
    learning_rate=2e-5,
    num_train_epochs=3,
    fp16=True,  # 启用FP16训练节省显存
    logging_steps=10,
    save_strategy="epoch",
    remove_unused_columns=False,
)

# 创建数据集与Trainer
train_dataset = MedicalDialogDataset(formatted_data, tokenizer)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

# 启动训练
trainer.train()

# 保存模型
model.save_pretrained("./llama2-medical-lora")
tokenizer.save_pretrained("./llama2-medical-lora")

# 推理示例
model.eval()
query = "[INST] Description: Patient has a fever and cough. Patient: I've had a fever for 2 days and can't stop coughing. [/INST]"
inputs = tokenizer(query, return_tensors="pt").to(model.device)
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        temperature=0.7,
        top_p=0.9
    )
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response.split("[/INST]")[-1].strip())

第四步：启动训练 #

执行以下命令启动分布式训练：

accelerate launch train.py

额外优化建议 #

• 梯度检查点：加载模型时添加gradient_checkpointing=True，进一步降低显存占用，但会增加训练时间；
• 4-bit量化：加载模型时设置load_in_4bit=True，单卡显存占用可降至~4GB；
• 调整batch size：若仍爆显存，降低per_device_train_batch_size，同时增大gradient_accumulation_steps保持有效batch size不变。

内容的提问来源于stack exchange，提问作者chriscrutt