Kaggle双T4 GPU上Accelerate微调ChatGLM-6B LoRA报错求助

我在Kaggle的2*T4 GPU环境下，用transformers、peft工具结合LoRA方法微调ChatGLM-6B模型。模型结构如下：

传统AutoModel.from_pretrained加载方式需要先把15GB的模型加载到CPU，但Kaggle的CPU内存只有13GB，无法完成加载。于是改用Accelerate的load_checkpoint_and_dispatch加载模型：

from transformers import AutoTokenizer, AutoModel, AutoConfig
from accelerate import load_checkpoint_and_dispatch, init_empty_weights
from huggingface_hub import snapshot_download

FilePath = snapshot_download(repo_id='THUDM/chatglm-6b')

config = AutoConfig.from_pretrained(FilePath, load_in_8bit=True, trust_remote_code=True)
with init_empty_weights():
    model = AutoModel.from_config(config, trust_remote_code=True).half()
model = load_checkpoint_and_dispatch(
    model, FilePath, device_map='auto', no_split_module_classes=["GLMBlock"]
)

这个方法成功把模型分配到CPU和GPU中：

之后通过peft添加LoRA适配器：

from peft import get_peft_model, LoraConfig, TaskType

peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=False, r=32, lora_alpha=32, lora_dropout=0.1, bias='none',
    target_modules=['query_key_value',],
)
model = get_peft_model(model, peft_config)

此时模型可以正常生成输出：

outputs = model(**tokenizer(['Hello world!'], return_tensors='pt').to(model.device))

但用accelerator.prepare包装训练数据加载器、验证数据加载器、模型和优化器后：

accelerator = Accelerator()

train_dataloader, val_dataloader, model, optimizer = \
        accelerator.prepare(train_dataloader, val_dataloader, model, optimizer)

执行训练代码时出现RuntimeError：

train_loss = []
epoch_correct_num, epoch_total_num = 0, 0
model.train()
for batch in tqdm(train_dl):
        labels = batch['labels']
        outputs = model(**batch)
        loss, logits = outputs.loss, outputs.logits
        optim.zero_grad()
#         loss.backward()
        accelerator.backward(loss)
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 2.0)
        optim.step()
        scheduler.step()

错误截图如下：

求解决这个问题的方法。

解决方法 #

1. 避免重复分配设备 #

用load_checkpoint_and_dispatch加载模型时已经通过device_map='auto'完成了设备分配，再用accelerator.prepare处理模型会触发重复分配逻辑，导致设备冲突。仅用accelerator.prepare处理数据加载器、优化器和调度器：

accelerator = Accelerator()
train_dataloader, val_dataloader, optimizer, scheduler = accelerator.prepare(
    train_dataloader, val_dataloader, optimizer, scheduler
)
# 模型不传入prepare，保持原有设备分配状态

2. 替换梯度裁剪方法 #

直接使用torch.nn.utils.clip_grad_norm_会遍历模型所有参数（包括CPU上的冻结主模型参数），引发跨设备操作错误。改用Accelerate提供的适配方法：

# 替换原有的torch.nn.utils.clip_grad_norm_
accelerator.clip_grad_norm_(model.parameters(), 2.0)

3. 限定优化器仅更新可训练参数 #

确保优化器只针对LoRA适配器的可训练参数更新，避免尝试修改冻结的主模型参数：

# 显式筛选可训练参数
optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=2e-4)

4. 调整设备映射策略（可选） #

若CPU内存仍吃紧，可改用balanced设备映射模式，更均匀地将模型层分配到GPU上，减少CPU参与：

model = load_checkpoint_and_dispatch(
    model, FilePath, device_map='balanced', no_split_module_classes=["GLMBlock"]
)

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