MatterSim 昇腾NPU训练性能优化

本项目是 MatterSim 模型在昇腾910C NPU环境上的性能调优工程，使用1 Die进行性能调优测试，包含优化后的训练代码及相关测试脚本。

1. 优化效果

经过多层面优化，训练性能显著提升：最初版本使用 torch/torch_npu版本2.6，单epoch性能是0.48s，升级到 torch/torch_npu 2.9 版本，单epoch性能0.39s。然后通过算子融合、AMP混合精度训练、激活函数融合、GateMLP重构等多维度优化，单epoch性能优化到0.28s，Loss精度从0.0255 降低到 0.0115，MAE(e)与MAE(f)收敛正常。

优化后训练日志样例：

2026-05-13 00:58:46.095 | INFO     | mattersim.forcefield.potential:train_model:282 - Epoch: 994 / 1000
2026-05-13 00:58:46.380 | INFO     | mattersim.forcefield.potential:train_one_epoch:751 - train: Loss: 0.0114, MAE(e): 0.1161, MAE(f): 1.0383, MAE(s): 0.0000, Time: 0.28s, lr: 0.00000125

2026-05-13 00:58:46.382 | INFO     | mattersim.forcefield.potential:train_model:282 - Epoch: 995 / 1000
2026-05-13 00:58:46.662 | INFO     | mattersim.forcefield.potential:train_one_epoch:751 - train: Loss: 0.0115, MAE(e): 0.1178, MAE(f): 1.0378, MAE(s): 0.0000, Time: 0.28s, lr: 0.00000125

2026-05-13 00:58:46.665 | INFO     | mattersim.forcefield.potential:train_model:282 - Epoch: 996 / 1000
2026-05-13 00:58:46.948 | INFO     | mattersim.forcefield.potential:train_one_epoch:751 - train: Loss: 0.0115, MAE(e): 0.1171, MAE(f): 1.0378, MAE(s): 0.0000, Time: 0.28s, lr: 0.00000125

2026-05-13 00:58:46.950 | INFO     | mattersim.forcefield.potential:train_model:282 - Epoch: 997 / 1000
2026-05-13 00:58:47.232 | INFO     | mattersim.forcefield.potential:train_one_epoch:751 - train: Loss: 0.0115, MAE(e): 0.1175, MAE(f): 1.0378, MAE(s): 0.0000, Time: 0.28s, lr: 0.00000125

2026-05-13 00:58:47.234 | INFO     | mattersim.forcefield.potential:train_model:282 - Epoch: 998 / 1000
2026-05-13 00:58:47.517 | INFO     | mattersim.forcefield.potential:train_one_epoch:751 - train: Loss: 0.0115, MAE(e): 0.1183, MAE(f): 1.0375, MAE(s): 0.0000, Time: 0.28s, lr: 0.00000125

2026-05-13 00:58:47.520 | INFO     | mattersim.forcefield.potential:train_model:282 - Epoch: 999 / 1000
2026-05-13 00:58:47.803 | INFO     | mattersim.forcefield.potential:train_one_epoch:751 - train: Loss: 0.0114, MAE(e): 0.1167, MAE(f): 1.0371, MAE(s): 0.0000, Time: 0.28s, lr: 0.00000125

2. 目录结构

├── mattersim/                           # 已包含优化代码的完整源码（已打patch）
│   ├── src/mattersim/                   # 核心源码
│   │   ├── training/                    # 训练相关代码
│   │   │   └── finetune_mattersim.py    # 微调训练入口（已优化）
│   │   ├── forcefield/                  # 力场相关
│   │   │   └── potential.py             # 模型训练逻辑（已优化）
│   │   ├── networks/                    # 网络结构
│   │   │   ├── m3gnet.py                # M3GNet核心网络（已优化）
│   │   │   ├── layers.py                # 网络层定义（已优化）
│   │   │   ├── message_passing.py       # 消息传递（已优化）
│   │   │   └── angle_encoding.py        # 角度编码（已优化）
│   │   └── datasets/                    # 数据集处理
│   │       └── build.py                 # DataLoader构建（已优化）
│   ├── pretrained_models/               # 预训练模型
│   └── tests/data/                     # 测试数据
├── patch/                              # 优化补丁文件
│   └── mattersim-028s.patch            # 性能优化补丁
├── test/                               # 测试脚本
│   ├── finetune.sh                     # 训练测试脚本
│   └── run_profiling.sh                # NPU性能分析脚本
├── MatterSim_NPU_train_optimize.md     # 详细优化说明文档
└── README.md                           # 本文档

3. 环境要求

3.1 硬件环境

NPU: 昇腾910C，1 卡，该代码工程同样支持DDP多卡训练；
内存: 建议32GB以上

3.2 软件环境

Docker镜像: quay.io/ascend/vllm-ascend:v0.18.0rc1-a3-openeuler
CANN: 配套版本 cann-8.5.1
PyTorch: torch_npu（镜像内置 2.9.0.post1+gitdc51c2d）

4. 快速开始

4.1 创建容器

docker run -itd \
  --net=host \
  --shm-size 32g \
  --privileged \
  --name mattersim-npu \
  --device=/dev/davinci6 \
  --device=/dev/davinci7 \
  --device=/dev/davinci_manager \
  --device=/dev/devmm_svm \
  --device=/dev/hisi_hdc \
  --device=/dev/davinci_mini_manage \
  -v /home/:/home/ \
  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver:ro \
  -v /usr/local/Ascend/firmware:/usr/local/Ascend/firmware:ro \
  -v /etc/ascend_install.info:/etc/ascend_install.info:ro \
  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi:ro \
  -v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi:ro \
  -e ASCEND_VISIBLE_DEVICES=6,7 \
  -e LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/driver:/usr/local/lib:$LD_LIBRARY_PATH \
  quay.io/ascend/vllm-ascend:v0.18.0rc1-a3-openeuler \
  /bin/bash -c "while true; do sleep 3600; done"

注意: 请根据实际NPU卡号修改 --device=/dev/davinci* 和 ASCEND_VISIBLE_DEVICES 参数。

4.2 进入容器

docker exec -it mattersim-npu /bin/bash

4.3 方式一：使用已优化的代码（推荐）

直接使用本工程 mattersim/ 目录下的代码（已包含所有优化）：

cd mattersim
pip install -e . -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com

4.4 方式二：从原始代码打补丁

# 下载原始代码
git clone https://github.com/microsoft/mattersim.git && cd mattersim
git checkout 5b1ee33b615faae41abd56581336827b2a1c49d3

# 应用优化补丁
git apply ../patch/mattersim-028s.patch

# 安装
pip install -e . -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com

4.5 运行训练测试

cd test
bash finetune.sh

5. 性能优化要点

5.1 优化措施总览

序号	优化类别	优化措施	影响文件
1	训练框架优化	自定义NPUHuberLoss损失函数	finetune_mattersim.py
2	训练框架优化	算子融合（Operator Fusion）	finetune_mattersim.py
3	训练框架优化	自动混合精度（AMP）	potential.py
4	算子级优化	Polynomial cutoff函数优化（Horner方法）	message_passing.py
5	算子级优化	向量化角度编码计算	angle_encoding.py
6	算子级优化	激活函数融合优化	layers.py
7	算子级优化	GatedMLP结构重构	layers.py
8	算子级优化	npu_repeat_interleave_indices替代函数	m3gnet.py, scaling.py, message_passing.py
9	算子级优化	预计算index_map减少重复调用	m3gnet.py
10	内存优化	合并gather操作减少内存访问	message_passing.py
11	数据加载优化	Dataloader预加载与持久化worker	build.py

5.2 训练框架优化

5.2.1 自定义NPUHuberLoss损失函数

PyTorch原生HuberLoss在NPU上可能存在算子兼容性问题，使用NPU友好的基本算子组合实现：

class NPUHuberLoss(torch.nn.Module):
    def __init__(self, delta: float = 1.0, reduction: str = "mean"):
        super().__init__()
        self.delta = delta
        self.reduction = reduction

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        diff = pred - target
        abs_diff = torch.abs(diff)
        quadratic = torch.clamp(abs_diff, max=self.delta)
        linear = abs_diff - quadratic
        loss = 0.5 * quadratic * quadratic + self.delta * linear
        if self.reduction == "mean":
            return loss.mean()
        elif self.reduction == "sum":
            return loss.sum()
        else:
            return loss

5.2.2 算子融合（Operator Fusion）

昇腾CANN支持算子融合，通过设置 jit_compile=False 启用图模式优化：

if args.enable_op_fusion and NPU_AVAILABLE:
    torch_npu.npu.set_compile_mode(jit_compile=False)

5.2.3 自动混合精度（AMP）

使用FP16进行计算，FP32进行累加，减少显存占用并提升计算吞吐量：

scaler = torch_npu.npu.amp.GradScaler()
with torch_npu.npu.amp.autocast():
    result = self.forward(...)
    loss_, e_mae, f_mae, s_mae = self.loss_calc(...)
scaler.scale(loss_).backward()
scaler.step(self.optimizer)
scaler.update()

5.3 算子级优化

5.3.1 Polynomial Cutoff函数优化（Horner方法）

将多项式计算从3次pow调用减少到1次：

# 原始: 1 - 6*r^5 + 15*r^4 - 10*r^3  (3次pow)
# 优化: 1 + r^3 * (-10 + r * (15 - 6*r))  (1次pow)
def polynomial(r: torch.Tensor, cutoff: float) -> torch.Tensor:
    ratio = r / cutoff
    r3 = ratio ** 3
    result = 1.0 + r3 * (-10.0 + ratio * (15.0 - 6.0 * ratio))
    return torch.clamp(result, min=0.0)

5.3.2 向量化角度编码计算

将循环逐个计算改为向量化批量计算，减少小算子kernel launch次数。

5.3.3 激活函数融合优化

使用PyTorch原生 F.silu() 替代手动组合 x * sigmoid(x)，利用CANN预优化的融合算子。

5.3.4 GatedMLP结构重构

共享层计算次数减半，避免重复计算。

5.3.5 npu_repeat_interleave_indices替代函数

def npu_repeat_interleave_indices(repeats: torch.Tensor) -> torch.Tensor:
    cumsum = torch.cumsum(repeats.view(-1), dim=0)
    total = int(cumsum[-1].item())
    positions = torch.arange(total, dtype=repeats.dtype, device=repeats.device)
    return torch.bucketize(positions, cumsum, right=True)

5.4 数据加载优化

DataLoader_pyg(
    ...,
    prefetch_factor=2 if num_workers > 0 else None,
    persistent_workers=True if num_workers > 0 else False,
)

6. 性能分析（Profiling）

如需进行性能分析，可使用提供的profiling脚本：

cd test
bash run_profiling.sh

或手动启用：

python finetune_mattersim.py \
    --load_model_path ./pretrained_models/mattersim-v1.0.0-1M.pth \
    --train_data_path ./tests/data/high_level_water.xyz \
    --enable_profiling \
    --profiler_output_dir ./profiler \
    --use_amp

生成的trace文件可通过TensorBoard可视化分析。

7. 修改文件清单

文件路径	修改类型	主要修改内容
`src/mattersim/training/finetune_mattersim.py`	新增+修改	NPUHuberLoss类、算子融合参数
`src/mattersim/forcefield/potential.py`	新增+修改	AMP支持、Profiling集成
`src/mattersim/networks/message_passing.py`	修改	Polynomial优化、gather合并、index_map传递
`src/mattersim/networks/angle_encoding.py`	重构	向量化RBF计算
`src/mattersim/networks/layers.py`	重构	激活函数融合、GatedMLP重构
`src/mattersim/networks/m3gnet.py`	修改	npu_repeat_interleave_indices、预计算index_map
`src/mattersim/networks/scaling.py`	修改	npu_repeat_interleave_indices
`src/mattersim/datasets/build.py`	修改	prefetch_factor、persistent_workers

8. 详细文档

更多优化细节请参阅：MatterSim_NPU_train_optimize.md

9. 注意事项

性能分析（Profiling） 会引入额外开销，仅用于调试阶段
AMP混合精度 需确保NPU驱动和CANN版本支持
环境变量 TASK_QUEUE_ENABLE=2 和 CPU_AFFINITY_CONF=1 有助于提升性能

10. 参考资料