[AI算法] LLM训练-构建transformers custom model
文章目录
- 1. 继承与实现基础结构
- 2. 支持 DeepSpeed 和 Accelerate 的注意事项
- a. 模型输出格式
- b. 设备管理
- c. 分布式训练兼容性
- d. DeepSpeed 特定优化
- 3. 训练脚本集成建议
- 4. 测试与调试建议
在使用 Hugging Face 的 transformers 库时，若要自定义一个继承自 PreTrainedModel 的模型，并确保其在训练过程中支持 DeepSpeed 或 Accelerate 等加速框架，需要注意以下关键点：
1. 继承与实现基础结构
继承 PreTrainedModel
from transformers import PreTrainedModel, PretrainedConfigclass MyCustomModel(PreTrainedModel):config_class = MyCustomConfig # 自定义配置类base_model_prefix = "my_model" # 模型前缀名def __init__(self, config):super().__init__(config)# 初始化模型结构
实现必要的方法
forward()：必须正确返回 loss（用于训练）和输出。
save_pretrained() / from_pretrained()：确保模型可保存和加载。
2. 支持 DeepSpeed 和 Accelerate 的注意事项
a. 模型输出格式
返回的输出应为 Seq2SeqLMOutput 或 CausalLMOutputWithPast 等标准输出类型，包含 loss, logits 等字段。
例如：
from transformers.modeling_outputs import CausalLMOutputWithPastdef forward(...):...return CausalLMOutputWithPast(loss=loss,logits=logits,past_key_values=past_key_values,hidden_states=hidden_states,attentions=attentions,)
b. 设备管理
不要在模型内部硬编码 .to(device)，让 Accelerate 或 DeepSpeed 控制设备放置。
使用 accelerator.prepare(model, optimizer, dataloader) 来自动处理设备分配。
c. 分布式训练兼容性
避免使用不支持分布式训练的操作（如某些自定义 gather/scatter 操作）。
使用 PyTorch 原生支持的并行方式（如 nn.parallel.DistributedDataParallel）。
d. DeepSpeed 特定优化
若使用 DeepSpeed ZeRO，请避免在模型中使用 torch.nn.DataParallel。
使用 deepspeed.initialize() 替代常规优化器初始化。
在 deepspeed 配置文件中指定 train_batch_size、gradient_accumulation_steps 等参数。
3. 训练脚本集成建议
- 使用 Accelerate
from accelerate import Accelerator
accelerator = Accelerator()
model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)for batch in train_dataloader:outputs = model(**batch)loss = outputs.lossaccelerator.backward(loss)optimizer.step()optimizer.zero_grad()
- 使用 DeepSpeed
安装 DeepSpeed 并使用其启动脚本：deepspeed --num_gpus=4 train.py --deepspeed --deepspeed_config ds_config.json
示例 ds_config.json：
json{"train_batch_size": 32,"gradient_accumulation_steps": 1,"optimizer": {"type": "AdamW","params": {"lr": 3e-5}},"zero_optimization": {"stage": 2}}
4. 测试与调试建议
- 使用 transformers.Trainer 进行快速验证是否能正常训练。
- 启用 fp16 或 bf16 加速训练时，确保模型计算图支持混合精度。
- 使用 torch.compile() 可进一步提升性能（PyTorch 2.0+）。