本文基于CANN开源社区的多个仓库进行应用案例讲解

CANN组织地址：https://atomgit.com/cann

hccl仓库地址：https://atomgit.com/cann/hccl

前言

大语言模型（LLM）如GPT、LLaMA动辄数十亿甚至千亿参数，单卡训练根本不可能。如何在昇腾NPU上高效训练这些大模型？

本文将展示如何使用CANN的各个组件，构建完整的大模型训练方案。

大模型训练的挑战

1. 内存挑战

以GPT-3 175B为例：

参数量：175B
FP32存储：175B × 4 bytes = 700GB
加上梯度：700GB × 2 = 1.4TB
加上优化器状态（Adam）：1.4TB × 2 = 2.8TB

单卡显存：32GB
需要的卡数：2.8TB / 32GB ≈ 90张卡！

2. 计算挑战

训练一个token的FLOPs：
FLOPs ≈ 6 × 参数量
     = 6 × 175B
     = 1050 GFLOPs

训练1万亿token需要：
总FLOPs = 1050G × 1T = 1.05 × 10^24 FLOPs

单卡算力：300 TFLOPs
需要的时间：1.05 × 10^24 / 300T ≈ 3500万秒 ≈ 405天

3. 通信挑战

每个训练step需要同步的梯度：700GB
如果带宽是100GB/s，通信时间：7秒
如果计算时间是1秒，通信占比：7/(7+1) = 87.5%

CANN的解决方案

1. 数据并行（Data Parallel）

使用HCCL进行梯度同步：

import torch
import torch.nn as nn
import torch.distributed as dist
import torch_npu

class GPTModel(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList([
            TransformerLayer(hidden_size) for _ in range(num_layers)
        ])
        self.lm_head = nn.Linear(hidden_size, vocab_size)
  
    def forward(self, input_ids):
        x = self.embedding(input_ids)
        for layer in self.layers:
            x = layer(x)
        logits = self.lm_head(x)
        return logits

def train_data_parallel():
    # 初始化分布式环境
    dist.init_process_group(backend='hccl')
    rank = dist.get_rank()
    world_size = dist.get_world_size()
  
    # 设置设备
    torch.npu.set_device(rank)
  
    # 创建模型
    model = GPTModel(
        vocab_size=50000,
        hidden_size=4096,
        num_layers=32
    ).npu()
  
    # 包装为分布式模型（自动使用HCCL同步梯度）
    model = nn.parallel.DistributedDataParallel(
        model,
        device_ids=[rank],
        bucket_cap_mb=25,  # 梯度bucket大小
        gradient_as_bucket_view=True  # 减少内存拷贝
    )
  
    # 优化器
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
  
    # 训练循环
    for epoch in range(10):
        for batch in dataloader:
            input_ids = batch['input_ids'].npu()
            labels = batch['labels'].npu()
          
            # 前向传播
            logits = model(input_ids)
            loss = nn.CrossEntropyLoss()(
                logits.view(-1, logits.size(-1)),
                labels.view(-1)
            )
          
            # 反向传播（DDP自动调用HCCL AllReduce同步梯度）
            optimizer.zero_grad()
            loss.backward()
          
            # 梯度裁剪
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
          
            # 更新参数
            optimizer.step()
          
            if rank == 0 and step % 100 == 0:
                print(f"Epoch {epoch}, Step {step}, Loss: {loss.item()}")

if __name__ == '__main__':
    train_data_parallel()

2. 模型并行（Model Parallel）

将模型切分到多张卡：

class GPTModelParallel(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers, num_devices=4):
        super().__init__()
        self.num_devices = num_devices
        layers_per_device = num_layers // num_devices
      
        # Embedding在设备0
        self.embedding = nn.Embedding(vocab_size, hidden_size).npu(0)
      
        # Transformer层分布在多个设备
        self.layer_groups = nn.ModuleList()
        for device_id in range(num_devices):
            start_layer = device_id * layers_per_device
            end_layer = start_layer + layers_per_device
          
            layers = nn.ModuleList([
                TransformerLayer(hidden_size)
                for _ in range(layers_per_device)
            ]).npu(device_id)
          
            self.layer_groups.append(layers)
      
        # LM Head在最后一个设备
        self.lm_head = nn.Linear(hidden_size, vocab_size).npu(num_devices - 1)
  
    def forward(self, input_ids):
        # Embedding
        x = self.embedding(input_ids)
      
        # 逐设备处理
        for device_id, layers in enumerate(self.layer_groups):
            # 移动到当前设备
            x = x.npu(device_id)
          
            # 在当前设备上执行
            for layer in layers:
                x = layer(x)
      
        # LM Head
        logits = self.lm_head(x)
        return logits

# 使用
model = GPTModelParallel(
    vocab_size=50000,
    hidden_size=4096,
    num_layers=32,
    num_devices=4
)

3. 流水线并行（Pipeline Parallel）

将batch切分成多个micro-batch：

class PipelineParallel:
    def __init__(self, model, num_microbatches=4):
        self.model = model
        self.num_microbatches = num_microbatches
  
    def forward(self, input_ids, labels):
        batch_size = input_ids.size(0)
        microbatch_size = batch_size // self.num_microbatches
      
        losses = []
      
        # 将batch切分成多个micro-batch
        for i in range(self.num_microbatches):
            start = i * microbatch_size
            end = start + microbatch_size
          
            micro_input = input_ids[start:end]
            micro_labels = labels[start:end]
          
            # 前向传播
            logits = self.model(micro_input)
            loss = nn.CrossEntropyLoss()(
                logits.view(-1, logits.size(-1)),
                micro_labels.view(-1)
            )
          
            # 反向传播（梯度累积）
            loss = loss / self.num_microbatches
            loss.backward()
          
            losses.append(loss.item())
      
        return sum(losses)

# 使用
pipeline = PipelineParallel(model, num_microbatches=4)
loss = pipeline.forward(input_ids, labels)
optimizer.step()
optimizer.zero_grad()

4. 混合并行（3D Parallel）

结合数据并行、模型并行、流水线并行：

def setup_3d_parallel(
    model,
    data_parallel_size=2,
    model_parallel_size=4,
    pipeline_parallel_size=2
):
    """
    3D并行配置
  
    总卡数 = data_parallel_size × model_parallel_size × pipeline_parallel_size
           = 2 × 4 × 2 = 16张卡
    """
    rank = dist.get_rank()
    world_size = dist.get_world_size()
  
    # 计算当前rank在3D空间中的位置
    pipeline_rank = rank // (data_parallel_size * model_parallel_size)
    model_rank = (rank // data_parallel_size) % model_parallel_size
    data_rank = rank % data_parallel_size
  
    print(f"Rank {rank}: DP={data_rank}, MP={model_rank}, PP={pipeline_rank}")
  
    # 创建通信组
    # 数据并行组（同一个模型副本）
    dp_group = dist.new_group([
        i for i in range(world_size)
        if i // data_parallel_size == rank // data_parallel_size
    ])
  
    # 模型并行组（同一个pipeline stage）
    mp_group = dist.new_group([
        i for i in range(world_size)
        if (i // data_parallel_size) % model_parallel_size == model_rank
    ])
  
    # 流水线并行组
    pp_group = dist.new_group([
        i for i in range(world_size)
        if i % (data_parallel_size * model_parallel_size) == 
           rank % (data_parallel_size * model_parallel_size)
    ])
  
    return dp_group, mp_group, pp_group

ZeRO优化

使用ZeRO减少内存占用：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

def train_with_zero():
    # ZeRO Stage 1: 分片优化器状态
    # ZeRO Stage 2: 分片优化器状态 + 梯度
    # ZeRO Stage 3: 分片优化器状态 + 梯度 + 参数
  
    model = GPTModel(
        vocab_size=50000,
        hidden_size=4096,
        num_layers=32
    )
  
    # 使用FSDP（类似ZeRO Stage 3）
    model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,  # 完全分片
        cpu_offload=False,  # 不offload到CPU
        mixed_precision=True  # 混合精度
    ).npu()
  
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
  
    # 训练
    for batch in dataloader:
        input_ids = batch['input_ids'].npu()
        labels = batch['labels'].npu()
      
        logits = model(input_ids)
        loss = nn.CrossEntropyLoss()(
            logits.view(-1, logits.size(-1)),
            labels.view(-1)
        )
      
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

混合精度训练

使用FP16/BF16加速训练：

from torch.cuda.amp import autocast, GradScaler

def train_with_mixed_precision():
    model = GPTModel(...).npu()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    scaler = GradScaler()
  
    for batch in dataloader:
        input_ids = batch['input_ids'].npu()
        labels = batch['labels'].npu()
      
        # 使用autocast自动混合精度
        with autocast():
            logits = model(input_ids)
            loss = nn.CrossEntropyLoss()(
                logits.view(-1, logits.size(-1)),
                labels.view(-1)
            )
      
        # 缩放loss并反向传播
        optimizer.zero_grad()
        scaler.scale(loss).backward()
      
        # 梯度裁剪
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
      
        # 更新参数
        scaler.step(optimizer)
        scaler.update()

梯度累积

模拟更大的batch size：

def train_with_gradient_accumulation(accumulation_steps=4):
    model = GPTModel(...).npu()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
  
    for epoch in range(10):
        for i, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].npu()
            labels = batch['labels'].npu()
          
            # 前向传播
            logits = model(input_ids)
            loss = nn.CrossEntropyLoss()(
                logits.view(-1, logits.size(-1)),
                labels.view(-1)
            )
          
            # 缩放loss
            loss = loss / accumulation_steps
          
            # 反向传播（梯度累积）
            loss.backward()
          
            # 每accumulation_steps步更新一次
            if (i + 1) % accumulation_steps == 0:
                # 梯度裁剪
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
              
                # 更新参数
                optimizer.step()
                optimizer.zero_grad()

完整训练脚本

import torch
import torch.nn as nn
import torch.distributed as dist
import torch_npu
from torch.utils.data import DataLoader, DistributedSampler
from transformers import GPT2Tokenizer

def main():
    # 1. 初始化分布式环境
    dist.init_process_group(backend='hccl')
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    torch.npu.set_device(rank)
  
    # 2. 加载数据
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
    dataset = load_dataset('wikitext', 'wikitext-103-v1')
  
    # 使用DistributedSampler确保数据不重复
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )
  
    dataloader = DataLoader(
        dataset,
        batch_size=8,
        sampler=sampler,
        num_workers=4,
        pin_memory=True
    )
  
    # 3. 创建模型
    model = GPTModel(
        vocab_size=50257,
        hidden_size=4096,
        num_layers=32,
        num_heads=32
    ).npu()
  
    # 4. 分布式数据并行
    model = nn.parallel.DistributedDataParallel(
        model,
        device_ids=[rank],
        bucket_cap_mb=25
    )
  
    # 5. 优化器和学习率调度
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=1e-4,
        betas=(0.9, 0.95),
        weight_decay=0.1
    )
  
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=100000
    )
  
    # 6. 混合精度
    scaler = torch.cuda.amp.GradScaler()
  
    # 7. 训练循环
    global_step = 0
    for epoch in range(10):
        sampler.set_epoch(epoch)  # 确保每个epoch的shuffle不同
      
        for batch in dataloader:
            input_ids = batch['input_ids'].npu()
            labels = batch['labels'].npu()
          
            # 前向传播
            with torch.cuda.amp.autocast():
                logits = model(input_ids)
                loss = nn.CrossEntropyLoss()(
                    logits.view(-1, logits.size(-1)),
                    labels.view(-1)
                )
          
            # 反向传播
            optimizer.zero_grad()
            scaler.scale(loss).backward()
          
            # 梯度裁剪
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
          
            # 更新参数
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
          
            global_step += 1
          
            # 日志
            if rank == 0 and global_step % 100 == 0:
                print(f"Step {global_step}, Loss: {loss.item():.4f}, "
                      f"LR: {scheduler.get_last_lr()[0]:.6f}")
          
            # 保存检查点
            if rank == 0 and global_step % 10000 == 0:
                torch.save({
                    'step': global_step,
                    'model_state_dict': model.module.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict(),
                }, f'checkpoint_{global_step}.pt')
  
    # 8. 清理
    dist.destroy_process_group()

if __name__ == '__main__':
    main()

性能优化技巧

1. 数据加载优化

# 使用多进程加载数据
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=8,  # 多进程
    pin_memory=True,  # 固定内存
    prefetch_factor=2  # 预取
)

2. 通信优化

# 使用梯度压缩
model = nn.parallel.DistributedDataParallel(
    model,
    device_ids=[rank],
    bucket_cap_mb=25,
    gradient_as_bucket_view=True,
    broadcast_buffers=False  # 不广播buffer
)

3. 内存优化

# 使用梯度检查点
from torch.utils.checkpoint import checkpoint

class TransformerLayer(nn.Module):
    def forward(self, x):
        # 使用checkpoint减少内存
        return checkpoint(self._forward, x)
  
    def _forward(self, x):
        # 实际的前向传播
        return self.attention(x) + self.ffn(x)

实际案例：训练LLaMA-7B

# 配置
model_config = {
    'vocab_size': 32000,
    'hidden_size': 4096,
    'num_layers': 32,
    'num_heads': 32,
    'intermediate_size': 11008
}

# 硬件配置
num_gpus = 8
batch_size_per_gpu = 4
gradient_accumulation_steps = 4
effective_batch_size = num_gpus * batch_size_per_gpu * gradient_accumulation_steps
# = 8 × 4 × 4 = 128

# 训练配置
max_steps = 100000
learning_rate = 3e-4
warmup_steps = 2000

# 预计训练时间
# 假设每step 2秒
total_time = max_steps * 2 / 3600  # 小时
print(f"预计训练时间: {total_time:.1f}小时")

总结

使用CANN训练大模型的关键技术：

  • 数据并行：使用HCCL同步梯度
  • 模型并行：切分模型到多卡
  • 流水线并行：切分batch提升效率
  • ZeRO优化：减少内存占用
  • 混合精度：加速训练
  • 梯度累积：模拟大batch

通过合理组合这些技术，可以在昇腾NPU上高效训练大语言模型。

相关链接

hccl仓库地址：https://atomgit.com/cann/hccl

CANN组织地址：https://atomgit.com/cann

runtime仓库地址：https://atomgit.com/cann/runtime

这篇是我在昇腾NPU上训练大模型的经验总结，如有错误欢迎指正。

# 人工智能
Logo
DeepSeek技术社区

欢迎加入DeepSeek 技术社区。在这里，你可以找到志同道合的朋友，共同探索AI技术的奥秘。

更多推荐

初识智能体

在过去两年中，几乎所有人都体验过大语言模型（如 ChatGPT、Claude、DeepSeek）的强大能力：写周报、润色邮件、解答编程疑难。被动响应：你问一句，它答一句。如果你的 Prompt 没有写清楚约束条件，它无法主动追问或修正目标。知行脱节：大模型能用流畅的文字告诉你“如何部署一个高可用 K8s 集群”，但它无法替你登录服务器敲下命令、查看报错日志并自动回滚。信息孤岛：大模型的知识止步于预

avatar DeepSeek技术社区

DeepSeek总结的时间旅行者的主键

来源：https://www.pgedge.com/blog/the-time-traveler-s-primary-key。

avatar DeepSeek技术社区
cover

Windows 下部署 New API 中转站全过程实录（含 6 个真实踩坑）

avatar DeepSeek技术社区