前言

torchtitan 是 Meta 开源的 LLM 训练框架，torchtitan-npu 是它的昇腾NPU适配版本。不需要学新的训练框架，用 PyTorch FSDP 的方式就能在昇腾NPU上训练大模型。

---

torchtitan-npu 和 torchtitan 的关系

torchtitan-npu 不是 fork —— 是在 torchtitan 上层加了昇腾NPU 适配层。

原始 torchtitan：
PyTorch → FSDP → cudatensor → CUDA Kernel

torchtitan-npu：
PyTorch → FSDP → npu_tensor → torch_npu → CANN Kernel

改动集中在：

1. Device 抽象层（把 cuda 换成 npu）
2. 通信后端（把 NCCL 换成 HCCL）
3. 算子适配（把 CUDA 特有算子换成 CANN 的算子）

---

环境搭建

依赖版本（严格匹配）

组件	版本	说明
CANN	8.0+	必须 ≥ 8.0（torchtitan 用了 CANN 8.0 的新 API）
PyTorch	2.1.0+	2.0.x 有 bug
torch_npu	对应 PyTorch 版本	用 pip install torch-npu==2.1.0
torchtitan-npu	main 分支	持续更新

安装步骤

# 1. 安装 CANN 8.0（已装跳过）
# 见 CANN 安装指南

# 2. 安装 PyTorch + torch_npu
pip install torch==2.1.0 torchvision==0.16.0
pip install torch-npu==2.1.0  # 昇腾NPU 的 PyTorch 后端

# 3. 克隆 torchtitan-npu
git clone https://atomgit.com/cann/torchtitan-npu.git
cd torchtitan-npu

# 4. 安装依赖
pip install -r requirements.txt

---

配置 LLaMA 模型训练

torchtitan-npu 的配置文件在 configs/ 目录下，按模型规模分：

配置文件	模型规模	推荐卡数
llama_7b.toml	7B	8 卡
llama_13b.toml	13B	16 卡
llama_70b.toml	70B	64 卡

配置文件结构

[model]
name = "llama"
vocab_size = 32000
d_model = 4096
n_heads = 32
n_layers = 32

[optimizer]
name = "AdamW"
lr = 3e-4
betas = [0.9, 0.95]
weight_decay = 0.1

[distributed]
backend = "hccl"          # ← 用 HCCL，不用 NCCL
fsdp = true                # ← 开 FSDP（数据并行）
tp_size = 1                # ← Tensor Parallel Size（单机不用）

[checkpoint]
save_interval = 1000        # 每 1000 step 存一次
keep_last_n = 3             # 只保留最近 3 个 checkpoint

---

FSDP + HCCL：数据并行的配置

FSDP（Fully Sharded Data Parallel）是 PyTorch 原生支持的数据并行策略，把模型参数、梯度、优化器状态分片到各卡。

HCCL 后端配置

# train.py（torchtitan-npu 的入口）
import torch
import torch.distributed as dist

# 初始化 HCCL 后端（替代 NCCL）
dist.init_process_group(
    backend="hccl",        # ← 关键：用 HCCL
    init_method="env://"   # 环境变量里读 RANK/MASTER_ADDR
)

# FSDP 配置
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = LLaMAModel(...)
model = FSDP(model, cpu_offload=None)  # 不 offload 到 CPU

# 优化器（AdamW）
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

HCCL 的环境变量

# 多卡训练的启动命令
export MASTER_ADDR="localhost"
export MASTER_PORT=29500
export RANK=0              # 当前卡的 rank
export WORLD_SIZE=8        # 总卡数

# 启动训练
torchrun --nproc_per_node=8 train.py --config configs/llama_7b.toml

---

Tensor Parallel：模型并行的配置

7B 模型单卡放得下，70B 放不下——需要 Tensor Parallel（TP）。

TP 的原理

不加 TP：每卡存完整的权重（70B × 2byte = 140GB，放不下）

加 TP（tp_size=8）：
每卡只存 1/8 的权重（140GB / 8 = 17.5GB，放得下）
前向/反向时，用 HCCL AllReduce 同步中间结果

torchtitan-npu 的 TP 配置

[distributed]
tp_size = 8               # ← Tensor Parallel Size = 8
pp_size = 1               # ← Pipeline Parallel 不用（70B 用 TP+FSDP 够了）

注意：TP 的通信量和 seq_len 成正比。长 context（>4K）时，TP 的通信开销大——此时建议用 TP+PP 混用（3D 并行）。

---

训练监控和 checkpoint 管理

监控指标

torchtitan-npu 用 torch_npu.profiler 打性能数据：

from torch_npu.npu.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.NPU],
    record_shapes=True,
    profile_memory=True
) as prof:
    for step, batch in enumerate(dataloader):
        loss = model(batch)
        loss.backward()
        optimizer.step()

        if step % 100 == 0:
            prof.step()  # ← 打一个 step 的 profile

# 导出 trace（用 Chrome 的 chrome://tracing 看）
prof.export_chrome_trace("train_trace.json")

checkpoint 的保存/加载

# 保存（FSDP 模型要用 state_dict）
from torch.distributed.fsdp import FullStateDictConfig, StateDictType

# 配置：把分片的参数合并成完整参数再存
save_policy = FullStateDictConfig(offload_to_cpu=True)
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
    torch.save(model.state_dict(), "checkpoint_1000.pt")

# 加载（继续训练）
checkpoint = torch.load("checkpoint_1000.pt")
model.load_state_dict(checkpoint)

---

常见错误和排错方法

错误 1：HCCL 初始化失败

现象：RuntimeError: HCCL initialization failed

原因：

• MASTER_ADDR / MASTER_PORT 没设对
• 或者防火墙拦了 HCCL 的通信端口

排错：

# 检查 HCCL 是否能通
hccn_tools list

# 输出示例（8 卡都通的才正常）：
Device 0: 192.168.1.10:29500 ✓
Device 1: 192.168.1.11:29500 ✓
...

错误 2：FSDP 的 OOM

现象：RuntimeError: NPU out of memory

原因：FSDP 的 cpu_offload 没开，中间激活值占爆了。

解决：

model = FSDP(
    model,
    cpu_offload=CPUOffload(offload_params=True)  # ← 参数 offload 到 CPU
)

错误 3：checkpoint 加载失败

现象：KeyError: 'layers.0.attention.wq.weight'

原因：保存时是 FSDP 分片格式，加载时用了非 FSDP 格式。

解决：保存和加载都用 FULL_STATE_DICT 格式（见上方代码）。

---

用 torchtitan-npu 训 LLaMA-2 7B，先跑通 8 卡 FP16——能跑通再调优化（混合精度、梯度累积、学习率 schedule）。

配置文件在 configs/llama_7b.toml，改 vocab_size 和 d_model 要跟模型定义一致。

torchtitan-npu 的代码在 AtomGit 的 cann/torchtitan-npu 仓库：

https://atomgit.com/cann/torchtitan-npu

configs/ 目录下是所有模型的配置文件。examples/ 目录下是训练启动脚本。