MACE的LoRA微调实操

LoRA在部分线性层上增加低秩增量，冻结原有参数，主要更新较小的一组适配参数。本篇从一个PBE小模型出发，使用PBEsol标签做40轮微调，演示数据、配置、训练、测试的操作顺序。

先准备数据与环境#

本组数据来自已经完成的Si₂、Si₈ PBE与PBEsol计算，共48个几何：32训练、8验证、8独立测试。同一几何在两个泛函下可以分别有标签，但训练、验证、测试的几何ID互不重叠。基模型是此前在PBE数据上训练得到的小模型，不是通用材料基础模型。

下载两种微调输入、实际结果、小模型、数据和绘图脚本：

python3 -m zipfile -e mace-lora-replay-wave962-*.zip .
cd mace-lora-replay
bash read.sh
bash plot.sh

默认读取附带结果，不启动新训练。文件data/中的extxyz保存结构、REF_energy、REF_forces和应力等标签；不能把模型自身的预测重新标成DFT真值。一个完整的双原子训练帧如下：

2
Lattice="0.0 2.77297309 2.77297309 2.77297309 0.0 2.77297309 2.77297309 2.77297309 0.0" Properties=species:S:1:pos:R:3:REF_forces:R:3 REF_energy=-11.23714237 REF_stress="_JSON [[0.02735552809284883, 0.001948741536648442, 0.00818160826329791], [0.001948741536648442, 0.027485975352711876, -0.007389524067632682], [0.00818160826329791, -0.007389524067632682, 0.026808148691218253]]" REF_virials="_JSON [[-1.1665733850509004, -0.08310386124812395, -0.3489037540960845], [-0.08310386124812395, -1.1721362936152147, 0.3151254136361152], [-0.3489037540960845, 0.3151254136361152, -1.1432304527083057]]" geometry_id=Si2-00 split=train xc=PBEsol config_type=Si2 pbc="T T T"
Si       0.01560990       2.74746179       2.74898155      -0.54675093       0.60544867       0.13674030
Si       1.35844832       1.40949393       1.36837537       0.54675093      -0.60544867      -0.13674030

MACE版本0.3.16，Python3.11、PyTorch2.8.0/CUDA12.6、e3nn0.4.4；训练使用一张V10016G。小例采用float64。更换模型规模、材料或训练分布后，应重新检查误差。

在自己的解压目录准备新训练#

module load mace/0.3.16
python setup.py
python prepare.py

setup.py针对本篇单元素Si模型生成一个个人目录内的训练程序副本，保留单元素原子能量数组的维度；它不改变公共MACE安装。脚本会核对匹配的0.3.16源码，版本不同应采用对应方法。

已经有private或work-new时，请保留原结果并使用新的解压目录。下面两种方法使用同一个基模型与同一组数据，修改自己的任务时，先调整configs/内的YAML，再准备目录。

查看LoRA的完整配置#

cat configs/lora.yml

name: lora
energy_key: REF_energy
forces_key: REF_forces
stress_key: REF_stress
virials_key: REF_virials
E0s: foundation
model: MACE
batch_size: 4
valid_batch_size: 4
max_num_epochs: 40
default_dtype: float64
seed: 29
loss: stress
energy_weight: 1.0
forces_weight: 100.0
stress_weight: 1.0
compute_stress: true
ema: false
ema_decay: 0.999
swa: false
num_workers: 0
device: cuda
log_dir: logs
model_dir: models
checkpoints_dir: checkpoints
results_dir: results
lr: 0.001
weight_decay: 0.0
config_type_weights: '{''Si2'':1.0,''Si8'':1.0}'
foundation_model: ../../models/base-PBE.model
multiheads_finetuning: false
train_file: ../../data/PBEsol-train.xyz
valid_file: ../../data/PBEsol-valid.xyz
foundation_head: Default
lora: true
lora_rank: 2
lora_alpha: 2.0

关键设置是lora: true、lora_rank: 2与lora_alpha: 2.0。foundation_model指向此前训练的PBE权重，目标训练和验证文件为PBEsol；这里关闭多头微调，得到一个目标模型。

E0s: foundation沿用基模型的原子能量参考。本例LoRA会学习目标数据相对基模型的变化，但这不意味着不同材料、不同泛函之间都可不检查能量参考直接复用。

提交训练，再读取新模型#

cd work-new/lora
sbatch submit.sh

任务结束后：

ls models
tail -n 12 train.out

其中models/lora.model为保存的模型，所附实例采用验证集选定的权重；不要把“共训练40轮”理解成必定使用最后一轮权重。

如需对新模型输出预测，在资料包根目录执行：

cd ../..
bash evaluate.sh lora

这个命令再提交一个GPU读取任务，使用mace_eval_configs读取独立测试集，输出work-new/evaluate-lora-Default.xyz。新模型包含CUDA相关序列化对象，示例读取安排在申请到的GPU上，不假定在无GPU环境仅改map_location就能加载。

实际更新了多少参数#

原模型有60,760个参数；本次优化器实际更新4,262个LoRA参数，约为原模型参数量的7.0%。训练日志中的某些“参数数目”行包含冻结参数，不能只拿该行总数判断实际更新量。保存时，适配增量合并到普通模型权重，使用模型时不需要再单独拼接LoRA层。

参数较少不代表运行时间或显存一定按同比例下降；邻居图、力求导与原模型前向计算仍然存在。

看独立测试数据#

这些是8个未参与训练或选模的测试几何上的结果。力还做了能量中心差分检查。资料包CSV保留逐帧能量和力RMSE摘要，默认绘图读取这些摘要，不会自动重新运行模型。

这批数据仅在Si近平衡小范围内。低测试误差不代表模型已能描述熔化、缺陷、断键或陌生元素；准备生产势函数时，还需补充所需构型和独立验证。

多头回放的区别与操作见MACE多头回放微调。官方训练概念可参阅MACE微调说明。