4x Tesla P40 训练 Llama-3.3-70B 大模型技术方案

✅ 解决方案：纯 FP32 训练管线 (Pure FP32 Pipeline)

这是唯一稳健的方案。虽然 FP32 显存占用比 FP16 大一倍（主要在激活值和梯度），但由于我们已经使用了 4-bit 权重，剩下的空间足够 batch_size=1 的 FP32 训练。

1. 禁用 AMP: fp16=False, bf16=False。
2. 强制 FP32 权重: 将原本可能默认为 FP16 的非量化层 (lm_head, LayerNorm, LoRA adapters) 强制转换为 float32。
3. BitsAndBytes 配置: 显式指定 bnb_4bit_compute_dtype=torch.float32。

4. 完整实施教程

步骤 1: 环境安装

conda create -n llama_p40 python=3.10 -y
conda activate llama_p40
# 安装 PyTorch (CUDA 11.8)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# 安装核心依赖
pip install --upgrade transformers peft bitsandbytes trl accelerate unsloth

步骤 2: 训练脚本配置 (关键代码)

创建一个 python 脚本 (e.g., train_p40.py)，关键部分如下：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

# 1. 配置 4-bit 量化 (关键：使用 FP32 计算)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4", # 推荐 nf4 精度更高
    bnb_4bit_compute_dtype=torch.float32, # [关键] P40 必须用 FP32 计算
    bnb_4bit_use_double_quant=True,
)

# 2. 加载模型 (关键：分片 + 强制 float32)
model = AutoModelForCausalLM.from_pretrained(
    "/path/to/Llama-3.3-70B-Instruct",
    quantization_config=bnb_config,
    device_map="auto", # 自动分片到多卡
    torch_dtype=torch.float32, # [关键] 强制模型加载为 float32
    low_cpu_mem_usage=True,
)

# 3. 后处理：手动将所有非量化层转为 FP32
# BitsAndBytes 会锁定量化层，我们只需要转换剩下的
for name, module in model.named_modules():
    if "norm" in name.lower() or "lm_head" in name.lower():
        module.to(torch.float32)

# 4. 配置 LoRA
model = prepare_model_for_kbit_training(model)
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj","k_proj","v_proj","o_proj"],
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# 5. 再次确保 LoRA 层也是 FP32
for name, module in model.named_modules():
    if "lora_" in name:
        module.to(torch.float32)

# 6. 训练参数 (关键：禁用 AMP)
sft_config = SFTConfig(
    output_dir="./output",
    per_device_train_batch_size=1, # 显存有限，BS=1
    gradient_accumulation_steps=8, # 累积梯度弥补 BS
    fp16=False, # [关键] 必须关闭
    bf16=False, # [关键] 必须关闭
    optim="paged_adamw_8bit", # 节省优化器显存
    max_length=2048, # 根据显存调整
)

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    # ... 其他数据集配置
)
trainer.train()

步骤 3: 运行训练

不要使用 accelerate launch 或 torchrun，因为我们使用的是 device_map="auto" (模型并行)，而不是 DDP (数据并行)。直接用 Python 运行：

CUDA_VISIBLE_DEVICES=0,1,2,3 python train_p40.py

5. 常见问题 (Troubleshooting)

• Q: 为什么不使用 DeepSpeed?
  • A: DeepSpeed 依赖较新的 CUDA 特性，在 CUDA 11.8 + P40 上编译极易失败。且 ZeRO-3 在高延迟的 PCIe 3.0 上效率不佳。device_map 是更简单的替代方案。
• Q: 报错 RuntimeError: "_amp_foreach_non_finite_check_and_unscale_cuda" not implemented for 'BFloat16'
  • A: 你开启了 AMP (fp16=True)。即使你没有显式使用 BF16，某些内部组件也可能触发它。请确保 fp16=False 且所有层都是 float32。
• Q: 显存还是不够怎么办？
  • A:
    1. 减小 max_length (e.g. 2048 -> 1024).
    2. 减小 per_device_train_batch_size (e.g. 1).
    3. 启用 gradient_checkpointing=True (Unsloth/TRL 默认支持).