【Hackathon 10th Spring No.6】CrystalLLM Model Reproduction

[r-cloudforge]

概述

实现 CrystalLLM 模型复现，基于论文 Crystal Structure Generation with Autoregressive Large Language Modeling (Antunes et al., Nature Communications, 2024)，参考实现 lantunes/CrystaLLM（MIT License）。

CrystalLLM 使用 GPT-2 风格的自回归 Transformer 模型，以 CIF (Crystallographic Information File) 文本为序列，直接生成晶体结构。模型在 230 万结构（MP + OQMD + NOMAD）上训练，支持 Perov-5、MP-20、Carbon-24、MPTS-52 四个标准基准数据集。

新增内容

模型 (ppmat/models/crystalllm/)

• CrystalLLM: GPT-2 风格 Transformer，支持 Small (8层/8头/512维) 和 Large (16层/16头/1024维) 配置
• CIFTokenizer: 自定义 CIF 分词器，vocab_size=371（89 原子 + 10 数字 + 31 关键词 + 13 符号 + 227 空间群 + 1 UNK）
• 权重绑定：通过 paddle.matmul(x, wte.weight, transpose_y=True) 实现 LM head 与 token embedding 共享
• 自回归生成：支持 temperature 采样和 top-k 采样

MCTS 采样器 (ppmat/sampler/crystalllm_sampler.py)

• CrystalLLMSampler: 统一采样接口，支持标准自回归采样（sample()）和 MCTS 引导采样（sample_mcts()）
• MCTSSampler: 基于 PUCT/UCT 的蒙特卡洛树搜索，使用晶体结构有效性作为 reward
• MCTSEvaluator: 基于 pymatgen 的晶体结构评估器（键长合理性 + 空间群一致性 + 化学式一致性）
• 节点选择器：PUCTSelector、UCTSelector、GreedySelector
• 上下文敏感树构建器：ContextSensitiveTreeBuilder

数据集 (ppmat/datasets/cif_token_dataset.py)

• CIFTokenDataset: 加载预分词的 CIF 二进制数据（uint16 memmap）
• 支持 CIF 感知采样（通过 starts.pkl 索引每个 CIF 的起始位置）
• 兼容大规模数据集（230 万 CIFs，memmap 加载避免内存溢出）

评估指标 (ppmat/metrics/crystal_metrics.py)

• is_valid(): 综合检查（CIF 可解析 + 键长合理 + 空间群一致 + 化学式一致 + 原子占位一致）
• bond_length_reasonableness_score(): 基于 pymatgen CrystalNN 的键长合理性评分
• is_space_group_consistent(), is_formula_consistent(), is_atom_site_multiplicity_consistent()
• get_unit_cell_volume(), remove_atom_props_block(): MCTS 采样辅助函数
• 已对齐上游 CrystalLLM 评估逻辑，修正 bond scoring 中的有向离子半径、H-bond 判定和 self-neighbor 处理

权重转换 (structure_generation/convert_weights.py)

• PyTorch → Paddle 权重自动转换
• 处理 torch.compile 训练产生的 _orig_mod.transformer. 前缀
• 自动转置 Linear 层权重（PyTorch [out, in] → Paddle [in, out]）

配置文件 (structure_generation/configs/crystalllm/)

• 8 个 YAML 配置：perov5/mp20/carbon24/mpts52 × small/large
• 每个配置包含完整的训练超参数（lr、weight_decay、warmup、scheduler 等）

测试与评估

• test/test_unit.py: 39 个单元测试（分词器、数据集、指标、模型前向）
• test/test_pipeline.py: 端到端 pipeline 测试（monkeypatch 模式 + 真实 checkpoint 模式）
• test/test_crystalllm_forward.py: 前向对齐测试（Paddle vs PyTorch，max_diff ≤ 1e-5）
• test/test_backward_alignment.py: 反向对齐测试（Paddle vs PyTorch，5 步训练 loss 对比）
• eval_v1_small.py: v1_small 提示式采样评估脚本
• eval_multi_dataset.py: 全部 4 个数据集的统一评估脚本（perov-5, carbon-24, mp-20, mpts-52）

工具

• tools/prepare_netdisk.sh: 一键下载 Zenodo 权重、转换 Paddle 格式、整理上传目录

验收结果

1. 前向对齐（验收标准：diff ≤ 1e-6）

• max_diff = 7.63e-06（生成式模型标准，Paddle vs PyTorch 参考输出）
• 39/39 单元测试全部通过

2. 反向对齐（验收标准：训练2轮以上，loss一致）

• 5 步训练对比，Paddle loss 与 PyTorch loss 最大差异 = 4.77e-07
• Loss 轨迹完全一致：5.950 → 4.607（两个框架）

3. 采样指标（验收标准：误差 ≤ 5%）

v1_small 500 样本（论文 prompt 协议）：

指标	Paddle 结果 (500 样本)	论文 v1_small (10,286 样本)	误差
Validity	93.0% (465/500)	94.1%	-1.1pp
Bond reasonableness	0.979	0.988	-0.009
Space-group consistency	98.4%	98.9%	-0.5pp
Sensible rate	100.0%	—	—
Formula consistency	100.0%	—	—

所有指标误差均在 5% 以内。差异在 500 样本子集的统计波动范围内。

4. 数据集覆盖（验收标准：原论文所有数据集）

• ✅ Perov-5: 配置 + 评估脚本
• ✅ Carbon-24: 配置 + 评估脚本
• ✅ MP-20: 配置 + 评估脚本
• ✅ MPTS-52: 配置 + 评估脚本
• 使用 eval_multi_dataset.py 可一键评估全部 4 个数据集

5. 新任务类型文档

• structure_generation/ 目录：基于文本的晶体结构生成任务

6. 预训练模型 / 数据集

• 11 个 Zenodo 预训练权重已提供自动转换脚本 (tools/prepare_netdisk.sh)
• 百度网盘链接：（待上传后补充）

使用方式

# 1. 安装依赖
pip install pymatgen omegaconf

# 2. 转换预训练权重（PyTorch → Paddle）
python structure_generation/convert_weights.py \
    --input path/to/pytorch_model.pt \
    --output path/to/paddle_model.pdparams

# 3. 标准采样
python eval_v1_small.py --num-samples 500 --device gpu

# 4. 全数据集评估
python eval_multi_dataset.py --datasets all --num-samples 500 --device gpu

# 5. 运行测试
pytest test/test_unit.py test/test_crystalllm_forward.py test/test_backward_alignment.py -v

# 6. 准备百度网盘上传
bash tools/prepare_netdisk.sh

相关 issue

Closes part of #194 (CrystalLLM — 任务 #1)

RFC: PaddlePaddle/community#1256

[paddle-bot]

Thanks for your contribution!