将 nanochat 移植到华为昇腾 NPU 实践指南

项目地址：github.com/cyzlmh/nanochat-ascend

前言

nanochat 是 Andrej Karpathy 开源的一个轻量级 LLM 训练框架，以其简洁的代码结构和清晰的训练流程著称。随着国产 AI 芯片的快速发展，越来越多的开发者希望在华为昇腾 NPU 上进行大模型训练。本文将分享 nanochat 从 CUDA 生态迁移到昇腾 CANN 生态的完整过程，包括适配要点、实验结果对比以及待改进的方向。

一、适配内容与技术要点

1.1 核心适配概览

本次适配的核心目标是在保持 nanochat 原有代码结构的前提下，使其能够在昇腾 NPU 上完成完整的训练流程。主要涉及以下几个层面的修改：

1.2 设备自动检测与分布式初始化

在 nanochat/common.py 中，设备检测逻辑进行了扩展，将 NPU 纳入自动检测流程：

def autodetect_device_type():
    if torch.cuda.is_available():
        device_type = "cuda"
    elif torch.backends.mps.is_available():
        device_type = "mps"
    elif hasattr(torch, "npu") and torch.npu.is_available():
        device_type = "npu"  # 新增 NPU 检测
    else:
        device_type = "cpu"
    return device_type

分布式训练方面，昇腾使用 HCCL（Huawei Collective Communication Library）替代 NCCL：

elif is_ddp_requested and device_type == "npu":
    device = torch.device("npu", ddp_local_rank)
    torch.npu.set_device(device)
    dist.init_process_group(backend="hccl", device_id=device)  # 使用 HCCL
    dist.barrier()

需要注意的地方：

• HCCL 的初始化参数与 NCCL 基本一致，但需要确保 torch_npu 正确安装
• 多卡训练时，RANK、LOCAL_RANK、WORLD_SIZE 等环境变量由 torchrun 自动设置

1.3 BF16 混合精度路径适配

昇腾 910B 对 BF16 有良好的硬件支持。训练脚本和引擎代码中将 NPU 与 CUDA 同等对待：

# scripts/base_train.py 和 scripts/chat_sft.py
if device_type in ("cuda", "npu"):  # NPU 纳入 BF16 路径
    ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16)

关键注意点：

• 昇腾的 BF16 算子覆盖率已经相当完善，但在某些边界情况下可能需要回退到 FP32
• 建议在训练初期监控 nan/inf 的出现情况

1.4 优化器融合算子兼容性修复

这是适配过程中遇到的主要技术难点之一。在原始 nanochat 中，fused AdamW/Muon 优化器将超参数保持为 CPU 上的 0-D 张量，但在昇腾 NPU 上，某些融合算子要求标量张量与目标张量位于同一设备。

解决方案是在融合步骤中将标量移动到目标设备：

# nanochat/optim.py
def fused_step(self, params, grads, exp_avgs, exp_avg_sqs, ...):
    # 将 CPU 标量移动到 NPU 设备
    lr = self.lr.to(device) if isinstance(self.lr, torch.Tensor) else self.lr
    beta1 = self.beta1.to(device) if isinstance(self.beta1, torch.Tensor) else self.beta1
    # ... 后续计算

1.5 MFU 计算与峰值算力映射

为了准确计算 Model FLOPs Utilization (MFU)，get_peak_flops() 函数中添加了昇腾 910B 的 BF16 峰值算力数据：

_PEAK_FLOPS_TABLE = (
    # Ascend NPU
    (["ascend", "910b"], 313e12),  # 约 313 TFLOPS (BF16)
    # NVIDIA GPUs...
)

注意：313 TFLOPS 是基于昇腾 910B 理论峰值估算的近似值，实际 MFU 计算结果会受此估计值影响。

二、实验结果与原版对比

2.1 验证环境

2.2 预训练阶段 (nanochat-train)

2.3 SFT 阶段 (nanochat-sft)

2.4 下游任务评估结果

2.5 与 CUDA 版本的对比分析

性能层面：

• 当前 MFU (~20%) 低于 H100 上良好调优的 CUDA 基线（通常可达 40-50%）
• 这主要是因为本分支优先保证正确性和可移植性，而非极致性能调优
• FlashAttention 加速在昇腾上暂未启用，attention 计算回退到 PyTorch SDPA

稳定性层面：

• 4 卡分布式训练运行稳定，未出现 OOM 或通信错误
• BF16 训练全程无 nan/inf 问题

三、待改进的地方

3.1 性能优化方向

1. FlashAttention 支持

   • 当前 attention 使用 PyTorch SDPA 回退路径
   • 昇腾有自研的 FlashAttention 算子，需要集成测试

2. torch.compile 启用

   • 目前 NPU 路径禁用了 torch.compile（出于兼容性考虑）
   • 需要验证昇腾对 inductor/triton 算子的支持情况

3. 算子级调优

   • 部分算子的 NPU 实现效率仍有提升空间
   • 数据流水线（dataloader）可以进一步优化

3.2 功能完善

1. RL 阶段验证
   • 目前仅完成 smoke 测试，需要更全面的稳定性验证
   • GRPO/PPO 等算法的长时间运行测试
2. 多节点训练
   • 当前验证集中在单节点多卡场景
   • 多节点 HCCL 通信需要进一步测试
3. 成本估算精度
   • 当前使用 3 CNY/小时/卡的粗略估计

四、其他改动：SwanLab 与实验追踪

4.1 为什么选择 SwanLab？

原版 nanochat 使用 WandB 进行实验追踪，但考虑到：

1. 国内访问 WandB 的稳定性问题
2. SwanLab 对国产硬件生态的更好支持
3. 成本因素（SwanLab 提供免费额度）

因此在本分支中增加 SwanLab 作为首选实验追踪工具。

4.2 实现方式

新增的 nanochat/experiment_logger.py 模块提供统一的日志接口：

# 使用方式
python scripts/base_train.py --logger=swanlab  # 使用 SwanLab
python scripts/base_train.py --logger=wandb    # 使用 WandB
python scripts/base_train.py --logger=none     # 禁用日志

也可以通过环境变量设置默认后端：

export NANOCHAT_LOGGER=swanlab
export SWANLAB_API_KEY=<your_key>
export SWANLAB_WORKSPACE=<your_workspace>

4.3 SwanLab 实验快照

为了方便社区复现，dev/export_swanlab_snapshot.py 实现了实验数据导出功能，可以将 SwanLab 的指标历史导出为静态图表和 CSV 文件，存放在 docs/data/swanlab/ 目录下。

五、快速开始

如果有一台昇腾 NPU 机器，可以尝试以下步骤：

# 1. 克隆仓库
git clone https://github.com/cyzlmh/nanochat-ascend.git
cd nanochat-ascend

# 2. 配置环境（确保 CANN 和 torch_npu 已安装）
export SWANLAB_API_KEY=<your_key>

# 3. 运行端到端训练脚本
export NNPU=4
export DEPTH=26
export DEVICE_BATCH_SIZE=16
bash runs/run_ascend.sh

详细的适配说明请参考 docs/ascend_npu.md。

六、总结

本次适配工作验证了 nanochat 在昇腾 NPU 上的可行性，核心训练流程（预训练、SFT、评估）已能稳定运行。虽然当前 MFU 与 CUDA 版本存在差距，但这主要是出于稳定性优先的设计选择，而非硬件限制。

对于希望在国产 AI 芯片上进行 LLM 训练的开发者，这个项目可以作为一个起点。后续将持续优化性能、完善功能，并跟进昇腾软件栈的更新。

欢迎社区贡献和反馈！

参考链接

• 本项目：github.com/cyzlmh/nanochat-ascend
• 上游项目：github.com/karpathy/nanochat
• 昇腾文档：www.hiascend.com