GR00T N1.7在昇腾910B上的训练部署简记

发表于 2026-06-25  266 次阅读


文章目录

本文记录将 NVIDIA GR00T N1.7(3B VLA 模型)迁移到华为 Ascend 910B 进行分布式训练的全过程。从基础镜像选型、环境适配、mx_driving 补丁、代码 patch、调优排错到最终跑通 DeepSpeed ZeRO-2 8 卡训练,逐一展开。

一、环境概览

组件版本
硬件8 × Ascend 910B (aarch64)
基础镜像drivingsdk:26.0.0-cann9.0.0-910b-ubuntu22.04
Python3.10 (conda 内置环境 torch2.7.1)
PyTorch2.7.1 + torch_npu 2.7.1.post2
CANN9.0.0
transformers4.57.3
diffusers0.35.1
DeepSpeed0.18.4 (ZeRO-2)
数据集LIBERO-10 (379 episodes, 95K steps)
步时~8s/step (global_batch_size=640)

二、基础镜像选型

2.1 镜像来源

华为云 SWR 提供了官方预置镜像,集成了完整的 DrivingSDK 套件:

swr.cn-south-1.myhuaweicloud.com/ascendhub/drivingsdk:26.0.0-cann9.0.0-910b-ubuntu22.04

这个镜像的定位是"开箱即用"——内置 CANN 9.0.0 驱动栈、conda 环境管理器、以及预装的 torch2.7.1 conda 环境(含 torch_npu 2.7.1.post2)。

2.2 为什么选这个镜像

  • 标准 PyTorch op 的 NPU 后端支持(全量)
  • torch_npu.optim.NpuFusedAdamW 等融合优化器
  • DeepSpeed ZeRO 系列的 hccl 通信后端
  • Diffusers / Transformers 的 NPU patcher 接口

2.3 Dockerfile 参考

虽然基础镜像自带大部分环境,但 GR00T N1.7 需要额外的 ffmpeg 视频解码库和若干 Python 包。Dockerfile 扩展层如下:

FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/drivingsdk:26.0.0-cann9.0.0-910b-ubuntu22.04

CMD ["bash"]
```

> 实际部署中未使用完整 Dockerfile,而是在容器内逐层验证和安装,更灵活地排错。Dockerfile 适合固化环境后一键重建。

三、容器环境验证与增强

镜像 pull 下来后,第一步不是直接跑训练,而是验证环境和补齐缺失组件

3.1 硬件可用性检查

npu-smi info

检查 8 卡是否全部可见、驱动版本是否匹配 CANN 9.0.0。关键输出:

NPU-ID Name                  Health  Power(Tot) Temp(C)
0      910B                   OK      65.0W      42
1      910B                   OK      65.0W      41
...(8 卡均 OK)

3.2 conda 环境确认

镜像自带 conda 环境和预激活的 torch2.7.1

conda activate torch2.7.1
python -c "import torch; print(torch.__version__)"
python -c "import torch_npu; print(torch_npu.__version__)"
# 输出: 2.7.1, 2.7.1.post2

确认 NPU 设备数和通信能力:

python -c "import torch; import torch_npu; print('NPU count:', torch.npu.device_count())"
python -c "
import torch
import torch.distributed as dist
dist.init_process_group(backend='hccl', init_method='env://', world_size=8, rank=0)
print('hccl OK')
dist.destroy_process_group()
"

3.3 系统依赖安装

GR00T 的数据 pipeline 使用 torchvision 解码视频,依赖系统级 ffmpeg 库:

apt update && apt install -y \
    libavcodec-dev libavformat-dev libavutil-dev \
    libswscale-dev libswresample-dev libavfilter-dev \
    libavdevice-dev pkg-config ffmpeg

3.4 Python 依赖安装

GR00T N1.7 的 pyproject.toml 声明了很多 NVIDIA 专有依赖(flash-attn、triton、tensorrt),需要在 NPU 上裁剪掉。实际安装的命令:

pip install --upgrade pip

这里有个重要的注意事项:不要直接 pip install -e .,因为 Isaac-GR00T 的 pyproject.toml 会尝试安装 flash-attn(CUDA 专属),在昇腾上会直接报错。需要先 patch pyproject.toml 再安装。

3.5 DrivingSDK 获取

# 克隆 DrivingSDK 仓库(含 mx_driving 源码)
git clone https://gitcode.com/Ascend/DrivingSDK.git /data/DrivingSDK

DrivingSDK 中包含了 GR00T-N1.7 示例目录和 mx_driving 的 patcher 源码,后面 mx_driving 补丁需要从这里的源码拷贝文件。

四、mx_driving 补丁(最关键的一步)

4.1 问题定位

pip 安装的 mx_driving v1.0.20260421 是旧版 wheel,存在两个关键缺陷:

  1. API 缺失:缺少 GR00T N1.7 需要的 TransformersNPUDiffusersNPU patcher 类
  2. Patcher API 过旧:使用旧版链式调用,不支持新版 Patcher().add().apply() 的 fluent 接口

直接 import 时报错:

from mx_driving.patcher import Patcher, TransformersNPU, DiffusersNPU
# ImportError: cannot import name 'TransformersNPU'

4.2 修复方案

从 DrivingSDK 源码中拷贝新版 patcher 文件,合并新旧 API 后重建 wheel。

修复文件清单(从 /data/DrivingSDK/mx_driving/patcher/ 源码拷贝):

文件来源说明
`transformers_patch.py`DrivingSDK 源码新增 — TransformersNPU patcher
`diffusers_patch.py`DrivingSDK 源码新增 — DiffusersNPU patcher
`patch.py`DrivingSDK 源码新版补丁管理框架
`version.py`DrivingSDK 源码版本追踪
`reporting.py`DrivingSDK 源码补丁应用报告
`patcher_logger.py`DrivingSDK 源码日志组件
`patcher.py`**重写**新版 fluent API:`Patcher().add().apply()`
`__init__.py`**重写**桥接新旧导出,加 `default_patcher_builder = None`

4.3 关键细节:_C.so 扩展

重建 wheel 直接覆盖安装会丢掉原有的 C 扩展 _C.cpython-310-aarch64-linux-gnu.so。这是 mx_driving 的底层算子绑定,没有它会报:

ModuleNotFoundError: No module named 'mx_driving._C'

解决:从同版本 CANN 环境的已有安装(如 bevfusion 的 conda 环境)中拷贝 _C.so 到 wheel 包内,或安装 wheel 后手动恢复:

# 从同级 conda 环境借 _C
cp /path/to/other/env/lib/python3.10/site-packages/mx_driving/_C*.so \
   /opt/conda/envs/torch2.7.1/lib/python3.10/site-packages/mx_driving/

4.4 最终 wheel 包

mx_driving-1.0.20260421.1-py3-none-any.whl

打包含 8 个 patcher 文件 + 旧版 _C.so,安装后 Patcher(TransformersNPU, DiffusersNPU).add(...).apply() 和旧版 API 同时可用。

五、GR00T 代码适配(5 个文件 patch)

核心适配工作集中在 5 个文件,覆盖分布式后端、tensor 处理、模型结构和依赖声明。

5.1 分布式后端——nccl → nccl + hccl(3 处)

原代码仅检查 backend == "nccl",昇腾 NPU 的分布式通信后端为 hccl。需要在 3 个文件中添加 hccl 分支:

文件函数修改内容
`dist_utils.py``_collective_device()`加 `else backend == "hccl"` → return npu device
`utils.py`device 选择加 `backend in ("nccl", "hccl")`
`sharded_mixture_dataset.py``_get_default_pg_tensor_device()`加 hccl → npu device

5.2 pin_memory——trainer.py

pin_memory 是 CUDA 专有特性,NPU 上会报:

RuntimeError: cannot pin 'npuFloatType'

修复:硬编码 pin_memory=False

pin_memory_device = "npu"
pin_memory = False  # NPU 不支持 pin_memory,硬编码关掉

5.3 get_device_capability 容错——dit.py

_is_spark_sm121() 在每次前向传播时都调用 torch.cuda.get_device_capability(),NPU 上返回 None 导致解包异常:

try:
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", ".*get_device_capability.*")
        major, minor = torch.cuda.get_device_capability()
        return (major, minor) == (12, 1)
except (TypeError, RuntimeError):
    return False

5.4 Qwen3-VL 骨干——qwen3_backbone.py

GR00T N1.7 使用 Qwen3-VL 作为视觉-语言骨干。原代码在初始化时尝试 import flash_attn,NPU 上没有这个库。加 torch_npu 优先检测,确保在无 flash_attn 时仍设置正确的 attention 实现:

# 优先检测 torch_npu
try:
    import torch_npu
    _has_npu = True
except ImportError:
    _has_npu = False

if _has_npu:
# NPU: 使用 transformers 自带的 flash_attention_2 实现
attn_implementation = "flash_attention_2"
else:
# CUDA: 原有逻辑
...
```

5.5 pyproject.toml 依赖裁剪

原声明的 NVIDIA 专属依赖在昇腾上不可用,需删除:

# 删除以下行
# flash-attn>=2.6.0,
# triton>=3.0.0,
# tensorrt>=10.0.0,
# uv>=0.5.0,

DeepSpeed 和 NCCL 相关依赖保留——DeepSpeed 通过 hccl 后端兼容 NPU。

5.6 最终 patch 文件

合并上述改动为一个 patch 文件:

gr00t_n1d7.patch  — 10 个文件,+296/-101 行

在 Isaac-GR00T 仓库上应用:

git clone https://github.com/NVIDIA/Isaac-GR00T.git
cd Isaac-GR00T
git checkout e5749287857afd97b78f1147166137de29746392
cp /path/to/gr00t_n1d7.patch ./
git apply gr00t_n1d7.patch
pip install -e .

六、训练脚本与参数调优

6.1 训练脚本

两个训练脚本:

train_8p.sh——完整训练(20000 步):

#!/bin/bash
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export HCCL_CONNECT_TIMEOUT=600
export LD_LIBRARY_PATH=/opt/conda/envs/torch2.7.1/lib/python3.10/site-packages/torch/lib:\
/opt/conda/envs/torch2.7.1/lib/python3.10/site-packages/torch_npu/lib:$LD_LIBRARY_PATH

accelerate launch \
--num_processes 8 \
--num_machines 1 \
--machine_rank 0 \
--main_process_port 29500 \
--mixed_precision bf16 \
gr00t/experiment/launch_finetune.py \
--dataset libero \
--config_path configs/train_config.yaml \
--output_dir results/Checkpoints/gr00t_n1.7_8p \
--num_train_steps 20000
```

关键参数说明:

参数原因
-----------------
`dataloader_num_workers`0NPU 不支持 fork 子进程创建 DataLoader
`embodiment_tag`LIBERO_PANDAN1.7 tag 体系变了,预注册标签自带 modality config
`mixed_precision`bf16NPU 原生支持 bf16,无需 fp16 转换
`optimizer`adamw_torchNPU 兼容标准 PyTorch AdamW,无需自定义
`tf32`FalseNPU 的 tf32 实现精度与 CUDA 不同,显式关闭

6.2 DeepSpeed ZeRO-2 配置

{
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "reduce_scatter": true,
        "contiguous_gradients": true,
        "overlap_comm": true
    },
    "bf16": {
        "enabled": true
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 1e-4,
            "weight_decay": 0.01
        }
    },
    "communication_data_type": "bf16",
    "gradient_accumulation_steps": 8
}

hccl 通信后端适配:DeepSpeed 在 NPU 上会自动检测 hccl 后端,关键环境变量确保通信正常:

export HCCL_CONNECT_TIMEOUT=600       # 8 卡建连超时
export HCCL_ALGO="ring"               # 通信算法
export HCCL_BUFFER_SIZE=200           # 通信缓冲区

七、排错记录

#错误根因修复
1`ImportError: TransformersNPU`mx_driving pip 包版本过旧从源码补 8 个 patcher 文件,打 wheel
2`Two accelerators: npu and npu`mx_driving._C 和 torch_npu 双重注册`transfer_to_npu()` 在 import torch_npu 前生效
3`Unknown tag: fourier_gr1`N1.7 tag 体系变了改用 `LIBERO_PANDA`
4`Cosmos-Reason2-2B 403`Gated HF 模型需授权申请授权后下载到本地
5`No backend type for cpu tensor`hccl 不认 CPU tensor3 处 `nccl` → `nccl + hccl` 适配
6`Cannot re-init NPU in forked subprocess`DataLoader fork 限制`num_workers=0`
7`cannot pin 'npuFloatType'`pin_memory 是 CUDA 特有trainer.py 硬编码 False
8`get_device_capability returns None`NPU API 未实现该接口`catch_warnings` + `try-except` 容错
9`ModuleNotFoundError: _C`wheel 覆盖安装丢 C 扩展从同级环境拷贝 `_C.so` 恢复

排错要点详解

错误 2(双重注册)是最隐晦的坑。mx_driving 的 __init__.py 在 import 时会自动调用 transfer_to_npu(),但如果此时 torch_npu 尚未 import,注册顺序会错乱。解决顺序

# 正确的 import 顺序
import torch_npu  # 必须先初始化 NPU backend
from mx_driving.patcher import Patcher, TransformersNPU, DiffusersNPU

错误 4(HF Gated Model):GR00T N1.7 依赖 Cosmos-Reason2-2B(NVIDIA 发布的 gated 模型),昇腾机器上没有 HF token 授权。需要在 HF 官网申请访问权限,然后:

huggingface-cli login --token YOUR_HF_TOKEN
# 或将模型下载到本地后挂载

错误 5(CPU tensor):hccl 通信要求所有 tensor 在 NPU 上。部分算子(如 dataset collate 中的 pad)会返回 CPU tensor,需要显式 .to("npu")。这 3 处 "nccl+hccl" 的修改本质上都是确保通信 tensor 正确落在 NPU device 上。

八、训练结果

8.1 Loss 曲线

8 卡 910B,DeepSpeed ZeRO-2,LIBERO-10 数据集(94 shards,95K steps),每卡约 9GB HBM 占用。

Loss 下降曲线健康,无震荡:

Step     Loss
0        1.28
500      1.18
1000     0.90
2000     0.53
5000     0.28
10000    0.16
20000    0.12

Grad Norm 稳定在 0.3~2.0 区间。

8.2 性能数据

指标
步时~8s/step
Global batch size640 (8 卡 × 16 per_device × 8 grad_accum)
每卡显存~9GB HBM
分布式DeepSpeed ZeRO-2 (hccl)

8.3 关键检查点验证

  • MX-DRIVING PATCHER 正确应用了以下 NPU 适配:
  • Qwen3RMSNorm(NPU 版 LayerNorm 替换)
  • RoPE(NPU 版旋转位置编码)
  • Attention scale(NPU bf16 scale)
  • 总共 9 个 patcher 生效

九、输出物

文件说明
`gr00t_n1d7.patch`适配补丁(10 文件,+296/-101 行)
`mx_driving-1.0.20260421.1-py3-none-any.whl`mx_driving 补丁 wheel
`train_8p.sh`完整训练脚本(20000 步)
`train_performance_8p.sh`性能测试脚本(1000 步 + FPS 计算)
`deploy.sh`一键部署脚本
`pyproject_ascend.toml`适配后依赖声明(去 NVIDIA 专有包)

以上文件均在 ascend_patch/ 目录下。

---

整个适配过程去掉镜像 pull 和环境准备,核心 patch 工作约 2-3 天。其中 mx_driving patcher 的 wheel 打包和 _C.so 恢复占了最大头,属于典型的"非官方环境适配"的成本。代码层面的改动不大(+296/-101 行),但每个改动都是踩坑后才确定的。昇腾生态整体可用性尚可,但在非主流模型(非 LLaMA/GLM)的适配中,mx_driving patcher 的版本兼容性是需要重点关注的风险点。


scanz个人博客