MatterGen 模型NPU 迁移适配完整指南

本项目完成 MatterGen 模型从 GPU 到华为昇腾 NPU 的迁移，包括训练和推理功能。

** 数据集和预训练权重**

数据集: dataset/mp_20/ (22MB) 仓库目录下已提供

预训练权重: ckpt/mp_20_base/ (440MB) 文件较大，需要按照如下指导自行下载

环境说明

硬件环境：昇腾910C 1卡；
软件环境：
    Docker镜像: quay.io/ascend/vllm-ascend:v0.18.0rc1-a3-openeuler
    CANN: 配套版本 cann-8.5.1
    PyTorch: torch_npu（镜像内置 2.9.0.post1+gitdc51c2d）

验证状态

训练验证

项目	结果	详情
训练速度	正常	14.7 it/s
预训练权重微调	成功	loss收敛正常

推理验证

项目	结果	详情
材料生成	成功	生成4个新材料结构
推理时间	正常	~62秒（batch_size=4）
输出格式	完整	CIF文件（3D晶体结构）

快速开始（完整流程）

步骤 1: 下载源代码

git clone https://gitee.com/mirrors_microsoft/mattergen.git
cd mattergen
git checkout a245cf2  # 切换到指定版本便于打补丁

步骤 2: 应用 NPU 迁移补丁

所有迁移修改已整合为单一补丁文件 patch/mattergen-npu-migration.patch（183行，修改8个文件）。

应用补丁：

cd mattergen

# 检查补丁是否可以应用
git apply --check ../patch/mattergen-npu-migration.patch

# 应用补丁
git apply ../patch/mattergen-npu-migration.patch

# 验证补丁应用成功（应显示8个修改文件）
git status

补丁包含的修改：

文件	修改说明	行数
`scripts/finetune.py`	transfer_to_npu + Lightning patches	+18
`scripts/run.py`	transfer_to_npu	+4
`common/utils/globals.py`	NPU 设备检测 + PyG CPU fallback	+10
`conf/lightning_module/default.yaml`	移除 `verbose` 参数（PyTorch 2.9兼容）	-1
`diffusion/run.py`	NPU 种子设置	+7
`diffusion/sampling/predictors_correctors.py`	设备不匹配修复	+2
`common/diffusion/predictors_correctors.py`	设备不匹配修复	+2
`common/gemnet/layers/basis_utils.py`	NumPy 2.x 兼容性修复	+5

步骤 3: 安装依赖

pip install -e . -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
pip install torch_scatter --no-build-isolation
pip install torch_sparse --no-build-isolation
pip install torch_cluster --no-build-isolation

步骤 4: 准备数据集

数据集已同步发布在 dataset/mp_20/ 目录下（22MB），直接拷贝即可：

# 从项目根目录拷贝数据集
mkdir -p mattergen/datasets/cache/mp_20
cp -r dataset/mp_20/* mattergen/datasets/cache/mp_20/

# 验证数据集
ls -lh mattergen/datasets/cache/mp_20/
# 应显示: train/ val/ test/ 三个目录

步骤 5: 准备预训练权重

预训练权重文件较大，无法上传代码仓，需要按照如下指导下载：

#  cd 到 mattergen目录下 执行如下命令，可以触发程序自动下载缺失的权重文件
export HF_ENDPOINT=https://hf-mirror.com
# 调用推理测试命令（并非真实运行） → 自动触发下载缺失文件
mattergen-generate results/test --pretrained-name=mattergen_base --batch_size=4 --num_batches=1

# 验证权重
ls -lh mattergen/microsoft/mattergen/checkpoints/mp_20_base/checkpoints/
# 应显示: last.ckpt (约440MB)

步骤 6: 配置 NPU

# 查看NPU状态
npu-smi info

# 选择空闲的卡（HBM-Usage较低），  finetune_mattergen.sh 脚本中进行配置
export ASCEND_RT_VISIBLE_DEVICES=6

步骤 7: 启动训练

cd test
bash finetune_mattergen.sh

# 日志自动生成: finetune_mattergen_npu6_<时间戳>.log
# 训练完成标志: "Trainer.fit stopped: max_epochs=2 reached."

步骤 8: 推理生成材料结构

cd test
bash infer-test.sh

# 日志自动生成: infer-test_npu5_<时间戳>.log
# 输出目录: inference_results_<时间戳>/

推理输出：

inference_results_<时间戳>/
├── cif_files/                     # CIF文件（解压后）
│   ├── gen_0.cif                  # 第1个材料结构
│   ├── gen_1.cif                  # 第2个材料结构
│   ├── gen_2.cif                  # 第3个材料结构
│   └── gen_3.cif                  # 第4个材料结构
├── generated_crystals_cif.zip     # CIF压缩包
└── generated_crystals.extxyz      # 扩展XYZ格式

生成的材料示例：

IrZnErIr - 铱锌铒化合物
Mn5 - 锰化合物
Ir3 - 铱化合物
Sr1 - 锶化合物

步骤 9: 查看3D结构（使用 MatterViz 网站）

推理完成后，生成的 CIF 文件可通过 MatterViz 网站进行3D可视化。

在线可视化方法：

访问 MatterViz 网站: https://matterviz.janosh.dev/structure
拖动上传 CIF 文件（位于 inference_results_<时间戳>/cif_files/）
网站自动渲染3D晶体结构，支持交互式旋转、缩放、查看

示例图片：

项目结构

docker-gen/
├── dataset/
│   └── mp_20/                      # 数据集（22MB，已同步发布）
│       ├── train/                  # 训练集
│       ├── val/                    # 验证集
│       └── test/                   # 测试集
│       └── README.md               # 数据集说明
├── patch/
│   └── mattergen-npu-migration.patch   # 统一迁移补丁（183行，8文件）
├── test/                           # 测试脚本和结果
│   ├── finetune_mattergen.sh       # 训练脚本
│   ├── infer-test.sh               # 推理脚本
├── README.md                   # 本文档（完整流程）
└── MIGRATION_GUIDE.md          # 迁移技术细节

训练结果

Epoch	Loss	速度
0	0.37	14.7 it/s

推理性能

推理时间: ~62秒
生成数量: 4个材料结构
推理速度: ~14.7 it/s
输出格式: CIF（晶体结构文件）

核心迁移技术点

1. transfer_to_npu 机制（关键）

使用 torch_npu.contrib.transfer_to_npu 自动转换所有 CUDA 操作到 NPU：

import torch_npu
from torch_npu.contrib import transfer_to_npu

原理：替换 torch.cuda.* API 为 torch.npu.*，让 Lightning 检测到 "cuda" 后自动转换。

2. Lightning 兼容性 Patches

# Patch torch.cuda.get_device_capability
def patched_get_device_capability(device=None): return (8, 0)
patched_get_device_capability.__module__ = 'torch_npu_monkey'
torch.cuda.get_device_capability = patched_get_device_capability

# Patch torch._C._cuda_clearCublasWorkspaces
if not hasattr(torch._C, '_cuda_clearCublasWorkspaces'):
    torch._C._cuda_clearCublasWorkspaces = lambda: None

3. PyG CPU Fallback（globals.py）

torch_scatter 不支持 NPU，图构建必须在 CPU 执行：

@lru_cache
def get_pyg_device() -> torch.device:
    return torch.device("cpu")  # torch_scatter only works on CPU

4. 设备不匹配修复

修复 step_size 设备不匹配问题：

step_size[grad_norm == 0, :] = self.max_step_size if isinstance(self.max_step_size, float) else self.max_step_size.to(step_size.device)

5. NumPy 2.x 兼容性

修复 np.math.factorial 在 NumPy 2.x 中不存在的问题：

import math  # 替代 np.math
math.factorial(...)  # 替代 np.math.factorial(...)

6. PyTorch 2.9 兼容性

修复 ReduceLROnPlateau 不支持 verbose 参数：

scheduler:
  _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
  factor: 0.6
  patience: 100
  min_lr: 1e-6
  # verbose: true  # 已移除，PyTorch 2.9不支持
  _partial_: true

常见问题

Q1: 训练速度慢（0.28 it/s）

原因：缺少 transfer_to_npu 导入
解决：应用补丁，确保包含 from torch_npu.contrib import transfer_to_npu

Q2: 推理报错 "Expected NPU tensor"

原因：设备不匹配
解决：补丁已修复 predictors_correctors.py 设备不匹配问题

Q3: NumPy 报错 "module 'numpy' has no attribute 'math'"

原因：NumPy 2.x 移除了 np.math
解决：补丁已修复 basis_utils.py，使用标准库 math.factorial

Q4: PyTorch 报错 "ReduceLROnPlateau got unexpected keyword 'verbose'"

原因：PyTorch 2.9 移除了 verbose 参数
解决：补丁已修复 default.yaml，移除 verbose: true

Q5: 如何查看生成的材料结构？

访问 MatterViz 网站：https://matterviz.janosh.dev/structure

迁移总结

训练成功: 14.7 it/s
推理成功: 生成4个新材料结构
transfer_to_npu: 自动转换 CUDA→NPU
统一补丁: 单一补丁文件，修改8个源码文件（183行）
环境兼容: PyTorch 2.9 + NumPy 2.x + torch_npu 2.9

适用版本: MatterGen commit a245cf2
迁移平台: 华为昇腾 NPU
补丁文件: patch/mattergen-npu-migration.patch (183行，8文件)
验证状态: 训练推理