Huggingface与PyTorch Lightning兼容性咨询：含PEFT LORA及FSDP场景

二者完全兼容，以下是实操中的关键要点和代码示例：

1. 模型与PEFT初始化 #

先加载HuggingFace预训练模型，再用PEFT封装LoRA模块：

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

# 加载基础模型
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 配置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"
)

# 封装为PEFT模型
peft_model = get_peft_model(model, lora_config)
peft_model.print_trainable_parameters()  # 确认仅LoRA层可训练

2. FSDP分布式配置 #

初始化Fabric时指定FSDP策略，并配置正确的自动包装规则，确保Transformer主层被分片，LoRA层保留在合适的设备上：

from lightning.fabric import Fabric
import torch

# 初始化Fabric，启用FSDP
fabric = Fabric(
    strategy="fsdp",
    accelerator="cuda",
    devices=4,
    precision="bf16-mixed"
)

# 配置FSDP自动包装策略（针对Transformer模型）
from lightning.fabric.strategies.fsdp import FSDPStrategy
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

fabric.strategy.config.auto_wrap_policy = transformer_auto_wrap_policy(
    transformer_layer_cls={LlamaDecoderLayer}
)

# 准备模型、优化器、数据加载器
optimizer = torch.optim.AdamW(peft_model.parameters(), lr=2e-4)
peft_model, optimizer = fabric.setup(peft_model, optimizer)
train_dataloader = fabric.setup_dataloader(train_dataloader)

3. 训练循环与模型保存 #

训练过程遵循Fabric的分布式规范，保存时仅导出LoRA适配器（无需保存完整基础模型）：

fabric.launch()
fabric.seed_everything(42)

for epoch in range(3):
    peft_model.train()
    for batch in train_dataloader:
        with fabric.no_backward_sync(peft_model, enabled=(not fabric.is_last_batch)):
            outputs = peft_model(**batch)
            loss = outputs.loss
            fabric.backward(loss)
        optimizer.step()
        optimizer.zero_grad()
    fabric.print(f"Epoch {epoch+1} loss: {loss.item()}")

# 保存LoRA适配器
fabric.save("lora_adapter", {"model": peft_model})
# 或者直接用PEFT的原生方法
peft_model.save_pretrained("lora_adapter")

常见注意事项 #

• 版本兼容：确保PEFT≥0.4.0、PyTorch Lightning Fabric≥2.0.0、Transformers≥4.28.0，避免API不兼容问题
• 分片策略：无需对LoRA层单独分片，FSDP会自动处理主模型的分片，LoRA层参数量小，全量保留不影响效率
• 设备分配：不要手动指定模型设备，交给Fabric的setup方法自动分配，避免分布式环境下的设备冲突

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