当前位置：首页 > news >正文

Accelerate 与 torchrun 分布式训练LLM对比

news 2025/9/29 7:48:54

Accelerate 与 torchrun 分布式训练LLM对比

1. 概述

在训练大型语言模型（LLM）时，分布式训练是必不可少的技术。本项目中实现了两种不同的分布式训练方式：

使用torchrun的传统DDP（DistributedDataParallel）方式
使用Accelerate库的现代化分布式训练方式

2. 启动方式对比

torchrun 方式

# 启动命令示例
torchrun --nproc_per_node 2 train_pretrain.py

Accelerate 方式

# 启动命令示例
python train_pretrain_accelerate.py --use_accelerate

3. 代码层面差异分析

3.1 初始化分布式环境

torchrun方式（train_pretrain.py）：

def init_distributed_mode():if not ddp: returnglobal ddp_local_rank, DEVICEdist.init_process_group(backend="nccl")ddp_rank = int(os.environ["RANK"])ddp_local_rank = int(os.environ["LOCAL_RANK"])ddp_world_size = int(os.environ["WORLD_SIZE"])DEVICE = f"cuda:{ddp_local_rank}"torch.cuda.set_device(DEVICE)

Accelerate方式（train_pretrain_accelerate.py）：

# Accelerate会自动处理分布式初始化，无需手动编写初始化代码
if args.use_accelerate and ACCELERATE_AVAILABLE:# 创建Accelerator实例accelerator = Accelerator(mixed_precision="bf16" if args.dtype == "bfloat16" else "no",gradient_accumulation_steps=args.accumulation_steps)# 准备模型、优化器和数据加载器model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader)

3.2 模型和数据加载器处理

torchrun方式：

# 手动将模型包装为DistributedDataParallel
if ddp:model._ddp_params_and_buffers_to_ignore = {"pos_cis"}model = DistributedDataParallel(model, device_ids=[ddp_local_rank])# 手动处理数据采样器
train_sampler = DistributedSampler(train_ds) if ddp else None
train_loader = DataLoader(train_ds,batch_size=args.batch_size,pin_memory=True,drop_last=False,shuffle=False,num_workers=args.num_workers,sampler=train_sampler
)

Accelerate方式：

# Accelerate自动处理分布式数据并行
train_loader = DataLoader(train_ds,batch_size=args.batch_size,pin_memory=(not (args.use_accelerate and ACCELERATE_AVAILABLE)),drop_last=False,shuffle=(train_sampler is None),num_workers=args.num_workers,sampler=train_sampler
)# 由Accelerator统一准备模型、优化器和数据加载器
model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader
)

3.3 损失计算和反向传播

torchrun方式：

with ctx:res = model(X)loss = loss_fct(res.logits.view(-1, res.logits.size(-1)), Y.view(-1)).view(Y.size())loss = (loss * loss_mask).sum() / loss_mask.sum()loss += res.aux_lossloss = loss / args.accumulation_stepsscaler.scale(loss).backward()if (step + 1) % args.accumulation_steps == 0:scaler.unscale_(optimizer)torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)scaler.step(optimizer)scaler.update()optimizer.zero_grad(set_to_none=True)

Accelerate方式：

with ctx:res = model(X)loss = loss_fct(res.logits.view(-1, res.logits.size(-1)), Y.view(-1)).view(Y.size())loss = (loss * loss_mask).sum() / loss_mask.sum()if hasattr(res, 'aux_loss'):loss += res.aux_loss# Accelerate处理梯度累积和缩放
accelerator.backward(loss)if (step + 1) % args.accumulation_steps == 0:# 梯度裁剪if accelerator.sync_gradients:accelerator.clip_grad_norm_(model.parameters(), args.grad_clip)optimizer.step()optimizer.zero_grad()

3.4 模型保存

torchrun方式：

if isinstance(model, torch.nn.parallel.DistributedDataParallel):state_dict = model.module.state_dict()
else:state_dict = model.state_dict()state_dict = {k: v.half() for k, v in state_dict.items()}  # 半精度保存
torch.save(state_dict, ckp)

Accelerate方式：

# 使用Accelerate保存模型
unwrapped_model = accelerator.unwrap_model(model)
state_dict = unwrapped_model.state_dict()
state_dict = {k: v.half() for k, v in state_dict.items()}
torch.save(state_dict, ckp)

4. 主要区别总结

特性	torchrun + DDP	Accelerate
分布式初始化	手动调用`dist.init_process_group`	Accelerator自动处理
模型包装	手动使用`DistributedDataParallel`	`accelerator.prepare()`自动处理
数据加载器	手动使用`DistributedSampler`	`accelerator.prepare()`自动处理
梯度累积	手动实现逻辑	内置支持，通过`gradient_accumulation_steps`参数
混合精度	手动使用`GradScaler`	通过`mixed_precision`参数配置
梯度裁剪	手动调用`clip_grad_norm_`	使用`accelerator.clip_grad_norm_`
模型保存	手动处理DDP模块	使用`accelerator.unwrap_model()`
学习率调度	手动管理	可与accelerator.prepare()一起使用
多平台支持	仅支持PyTorch DDP	支持DDP、DeepSpeed、FSDP等多种后端