Colab加载合并权重的微调Llama2模型时GPU未启用问题排查

我已完成Llama2模型的微调，重新加载基础模型并合并了LoRA权重，随后保存该合并后的模型以消除对基础模型的依赖。合并保存的操作代码如下：

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map=device_map,
)
model = PeftModel.from_pretrained(base_model, new_model)
model = model.merge_and_unload()
model.save_pretrained(...path/to/model)

现在尝试加载path/to/model路径下的模型，代码如下：

model_config = transformers.AutoConfig.from_pretrained(
    model_id,
    use_auth_token=hf_auth
)

model = transformers.AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    config=model_config,
    device_map='auto',
    offload_folder="offload",
    torch_dtype=float16,
    use_auth_token=hf_auth,
    offload_state_dict = True,
)
model.eval()

在Colab中运行该模型时，发现完全未使用GPU，仅使用CPU导致运行崩溃，想了解GPU未被使用的原因。

• device_map='auto'未正确触发GPU分配：Colab环境中transformers版本兼容问题或模型元数据缺失，可能导致自动设备映射失效。可以手动指定GPU设备，将device_map='auto'替换为device_map='cuda:0'，或者加载后手动迁移模型到GPU：

  model = model.to('cuda')
  

• torch_dtype参数格式错误：代码中使用的float16是Python原生类型，应改为torch.float16，否则模型可能默认采用CPU兼容的精度，无法适配GPU加速。
• offload_state_dict=True强制CPU卸载：该参数会优先将模型权重卸载到CPU，若GPU显存充足，建议移除该参数，避免强制走CPU路径。
• 合并模型时的设备残留标记：合并LoRA权重时，若基础模型存在部分权重在CPU上，合并后的模型可能保留CPU设备标记。可在保存前将模型迁移到CPU再保存：

  model = model.to('cpu')
  model.save_pretrained(...path/to/model)
  

• Colab GPU未启用或初始化失败：先检查Colab是否已开启GPU（菜单栏「代码执行程序」>「更改运行时类型」>选择GPU），再通过torch.cuda.is_available()验证GPU是否可用。

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