【MIIT program】support matterchat by learncat163 · Pull Request #305 · PaddlePaddle/PaddleMaterials

learncat163 · 2026-07-07T05:59:15Z

MatterChat diff (PyTorch and PaddlePaddle)

MatterChat 由三个模块串联组成: CHGNet (晶体编码器) → Q-Former (跨模态桥接) → Mistral-7B (大语言模型)。本文档分别验证每个模块从 PyTorch 迁移到 PaddlePaddle 后的数值对齐情况。

1. CHGNet Diff

CHGNet 负责将晶体结构 (CIF) 编码为逐原子 embedding。验证用 GaN.cif (4 原子) 作为输入, 对比 Paddle 输出与 PyTorch 参考的 [N_atom, 64] material embedding。

结果

指标	值	阈值	状态
max_diff	8.05e-07	1e-4	PASS
mean_diff	1.82e-07	1e-4	PASS

代码样例

PyTorch 参考数据生成

import numpy as np
import torch
from pymatgen.core import Structure
from Model.chgnet_lib.model.model_embedding import CHGNet  # MatterChat 原始 PT 代码

device = torch.device("cuda")
chgnet = CHGNet.load().to(device).float().eval()

struct = Structure.from_file("fix_inputs/GaN.cif")
with torch.no_grad():
    atom_feas = chgnet.predict_structure_embedding(struct)
if isinstance(atom_feas, (list, tuple)):
    atom_feas = atom_feas[0]

# 保存为参考数据
ref = atom_feas.detach().cpu().float().numpy()
np.save("fix_outputs/raw_stage_01_material_embed.npy", ref)
print(f"shape={ref.shape}")  # -> (4, 64)

PaddlePaddle 对比

import numpy as np
import paddle
from pymatgen.core import Structure
from ppmat.models.matterchat.chgnet.model.model_embedding import CHGNet

# 加载 PyTorch 参考数据
ref_embed = np.load("fix_outputs/raw_stage_01_material_embed.npy").astype(np.float64)

# 构建 CHGNet 并加载权重 (CPU)
paddle.set_device("cpu")
chgnet = CHGNet()
chgnet_weights = load_sharded_weights(prefix="material_encoder.")
set_state_dict_filtered(chgnet, chgnet_weights, prefix="material_encoder.")
chgnet.eval()

# GPU 推理
paddle.set_device("gpu")
gpu_chgnet = CHGNet()
gpu_chgnet.set_state_dict({k: v.cuda() for k, v in chgnet.state_dict().items()})
gpu_chgnet.eval()

struct = Structure.from_file("fix_inputs/GaN.cif")
with paddle.no_grad():
    atom_feas = gpu_chgnet.predict_structure_embedding(struct)

# 对比
diff = np.abs(atom_feas.numpy().astype(np.float64) - ref_embed)
print(f"max_diff={diff.max():.2e}  mean_diff={diff.mean():.2e}")
# -> max_diff=8.05e-07  mean_diff=1.82e-07  (PASS, thr=1e-4)

2. Q-Former Diff

Q-Former 是 BLIP-2 风格的 BERT 变体, 用 32 个 query token 对 CHGNet 输出做 cross-attention, 产生 [1, 32, 768] 的查询表示。验证 Q-Former 前向输出与基线一致。

结果

指标	值	阈值	状态
max_diff	0.00e+00	1e-4	PASS
mean_diff	0.00e+00	1e-4	PASS

Q-Former 输出与基线完全一致 (diff=0)。

代码样例

PyTorch 参考数据生成

import numpy as np
import torch
from Model.material_Q_former_base import BertConfig, BertLMHeadModel  # MatterChat 原始 PT 代码

device = torch.device("cuda")

# 加载预训练 Q-Former 权重
full_ckpt = torch.load("model_weight/model_weights.pkl", map_location="cpu", weights_only=False)
full_state = full_ckpt["state_dict"]

qformer_config = BertConfig.from_pretrained("bert-base-uncased")
qformer_config.encoder_width = 64
qformer_config.add_cross_attention = True
qformer_config.cross_attention_freq = 2
qformer_config.query_length = 32
qformer = BertLMHeadModel.from_pretrained("bert-base-uncased", config=qformer_config)

qformer_state = {}
for k, v in full_state.items():
    if k.startswith("model.Qformer.") or k.startswith("model.query_tokens"):
        qformer_state[k.replace("model.", "", 1)] = v
qformer.load_state_dict(qformer_state, strict=False)
qformer = qformer.to(device).float().eval()

query_tokens = full_state["model.query_tokens"].to(device).float()

# 输入: 上一步 CHGNet 输出的 material embedding
material_embed = torch.from_numpy(
    np.load("fix_outputs/raw_stage_01_material_embed.npy")
).to(device).float()
material_att = torch.ones((1, material_embed.shape[0]), dtype=torch.long, device=device)

with torch.no_grad():
    query_output = qformer.bert(
        query_embeds=query_tokens,
        encoder_hidden_states=material_embed.unsqueeze(0),
        encoder_attention_mask=material_att,
        return_dict=True,
    )
qf_out = query_output.last_hidden_state[:, :query_tokens.shape[1], :]

# 保存为参考数据
ref = qf_out.detach().cpu().float().numpy()
np.save("fix_outputs/raw_stage_02_qformer_out.npy", ref)
print(f"shape={ref.shape}")  # -> (1, 32, 768)

PaddlePaddle 对比

import numpy as np
import paddle
from ppmat.models.matterchat.q_former.q_former_base import BertConfig, BertLMHeadModel

# 加载基线数据
material_embed = np.load("fix_outputs/raw_stage_01_material_embed.npy")
qt_param = np.load("fix_outputs/paddle_stage_02_query_tokens.npy")

# 构建 Q-Former (CPU)
paddle.set_device("cpu")
config = BertConfig()
config.encoder_width = 64
config.add_cross_attention = True
config.cross_attention_freq = 2
config.query_length = 32

qformer = BertLMHeadModel(config)
qformer.cls = None
qformer.bert.embeddings.word_embeddings = None
qformer.bert.embeddings.position_embeddings = None
for layer in qformer.bert.encoder.layer:
    layer.output = None
    layer.intermediate = None

qformer_weights = load_sharded_weights(prefix="Qformer.")
set_state_dict_filtered(qformer, qformer_weights, prefix="Qformer.")
qformer.eval()

# GPU 推理
paddle.set_device("gpu")
gpu_qformer = BertLMHeadModel(config)
gpu_qformer.set_state_dict({k: v.cuda() for k, v in qformer.state_dict().items()})
gpu_qformer.eval()

qt_gpu = paddle.to_tensor(qt_param).cast(paddle.float32).cuda()
emb_gpu = paddle.to_tensor(material_embed).cuda().unsqueeze(0)
att_gpu = paddle.ones([1, emb_gpu.shape[1]], dtype=paddle.int64).cuda()

with paddle.no_grad():
    query_output = gpu_qformer.bert(
        query_embeds=qt_gpu,
        encoder_hidden_states=emb_gpu,
        encoder_attention_mask=att_gpu,
        return_dict=True,
    )
qf_out = query_output.last_hidden_state[:, :qt_gpu.shape[1], :]

print(f"max_diff={np.abs(qf_out.numpy().astype(np.float64) - ref).max():.2e}")
# -> max_diff=0.00e+00  (PASS, thr=1e-4)

3. Mistral LLM Diff

Mistral-7B 是 32 层 decoder-only Transformer, 含 RoPE、GQA、SwiGLU。由于模型规模大 (7.3B 参数), float16 跨框架累积误差不可避免, 因此阈值放宽至 1e-3。验证逐层 hidden state、最终 hidden_last 及 Top-20 token 匹配。

结果汇总

测试项	通过	阈值	max_diff	状态
32 层 hidden state	32/32	1e-3	2.97e-04 (layer_31)	PASS
hidden_last	1/1	1e-3	1.12e-04	PASS
Top-20 token	20/20	—	—	PASS

逐层误差

层	max_diff	mean_diff	状态
layer_00	5.14e-07	7.92e-09	PASS
layer_01	1.16e-04	1.06e-07	PASS
layer_05	1.22e-04	1.70e-07	PASS
layer_10	1.22e-04	2.25e-07	PASS
layer_15	1.22e-04	3.15e-07	PASS
layer_20	1.22e-04	5.78e-07	PASS
layer_25	1.22e-04	8.60e-07	PASS
layer_30	1.06e-04	1.51e-06	PASS
layer_31	2.97e-04	1.99e-06	PASS
hidden_last	1.12e-04	1.78e-05	PASS

代码样例

PyTorch 参考数据生成

import numpy as np
import torch
from transformers import LlamaTokenizer
from transformers.models.mistral.modeling_mistral import MistralForCausalLM
from transformers import MistralConfig as HFMistralConfig

device = torch.device("cuda")
MODEL_WEIGHT_DIR = "model_weight/Mistral-7B-Instruct-v0.3"

# Tokenize
tokenizer = LlamaTokenizer.from_pretrained(MODEL_WEIGHT_DIR, use_fast=False)
tokenizer.add_special_tokens({"pad_token": "[PAD]", "bos_token": "<s>",
                              "eos_token": "</s>", "unk_token": "<unk>"})
prompt = "[INST] What is the chemical formula of this material? [/INST]"
tokens = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=64)
input_ids = tokens["input_ids"].to(device)

# 加载 Mistral-7B
hf_config = HFMistralConfig.from_pretrained(MODEL_WEIGHT_DIR)
hf_config._attn_implementation = "eager"
llm = MistralForCausalLM.from_pretrained(MODEL_WEIGHT_DIR, config=hf_config,
                                         torch_dtype=torch.float32).to(device).eval()
llm.resize_token_embeddings(len(tokenizer))

# input embeddings
with torch.no_grad():
    inputs_embeds = llm.model.embed_tokens(input_ids)
np.save("fix_outputs/raw_stage_05_input_embeds.npy",
        inputs_embeds.detach().cpu().float().numpy())

# 构建因果掩码
seq_len = input_ids.shape[1]
min_dtype = torch.finfo(inputs_embeds.dtype).min
causal_mask = torch.full((seq_len, seq_len), min_dtype, dtype=inputs_embeds.dtype, device=device)
causal_mask = torch.triu(causal_mask, diagonal=1)[None, None, :, :].expand(1, 1, -1, -1)

# 逐层前向, 保存每层 hidden state 作为参考
hidden_states = inputs_embeds
position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
for li in range(32):
    with torch.no_grad():
        hidden_states = llm.model.layers[li](
            hidden_states, attention_mask=causal_mask,
            position_ids=position_ids, use_cache=True,
        )[0]
    np.save(f"fix_outputs/raw_stage_06_hidden_layer{li:02d}.npy",
            hidden_states.detach().cpu().float().numpy())

# 最终 norm + hidden_last
with torch.no_grad():
    hidden_final = llm.model.norm(hidden_states)
np.save("fix_outputs/raw_stage_06_hidden_last.npy",
        hidden_final[0, -1].detach().cpu().float().numpy())

PaddlePaddle 对比

import json
import numpy as np
import paddle
from ppmat.models.matterchat.mistral.configuration_mistral import MistralConfig
from ppmat.models.matterchat.mistral.modeling_mistral import (
    MistralForCausalLM,
    MistralDecoderLayer,
    MistralRMSNorm,
)

# 加载 PyTorch 参考数据 (逐层 hidden state)
refs = {}
for li in range(32):
    refs[f"stage_06_hidden_layer{li:02d}"] = np.load(
        f"fix_outputs/raw_stage_06_hidden_layer{li:02d}.npy"
    ).astype(np.float64)
refs["stage_05_input_embeds"] = np.load(
    "fix_outputs/raw_stage_05_input_embeds.npy"
).astype(np.float64)

emb = refs["stage_05_input_embeds"].astype(np.float32)
seq_len = emb.shape[1]

# 构建因果掩码 + position ids
min_dt = float(paddle.finfo(paddle.float32).min)
causal_mask = paddle.triu(
    paddle.full([seq_len, seq_len], min_dt, dtype="float32"), diagonal=1
)[None, None].expand([1, 1, -1, -1])
pos_ids = paddle.arange(seq_len).unsqueeze(0)
cache_pos = paddle.arange(seq_len)

# CPU 加载完整 Mistral-7B, 提取每层权重
paddle.set_device("cpu")
llm = MistralForCausalLM(MistralConfig(vocab_size=32769))
load_sharded_weights(llm, weight_dir=WEIGHT_DIR)
llm.eval()

layer_sds = [
    {k: v for k, v in llm.model.layers[i].state_dict().items()}
    for i in range(32)
]

# 逐层 GPU 推理 + 对比
paddle.set_device("gpu")
hidden = paddle.to_tensor(emb).cuda()
cm = causal_mask.cuda()
pos_ids = pos_ids.cuda()
cache_pos = cache_pos.cuda()

for li in range(32):
    gpu_layer = MistralDecoderLayer(llm.config, li)
    gpu_layer.set_state_dict({k: v.cuda() for k, v in layer_sds[li].items()})
    gpu_layer.eval()

    with paddle.no_grad():
        hidden = gpu_layer(
            hidden, attention_mask=cm,
            position_ids=pos_ids, use_cache=False,
            cache_position=cache_pos,
        )[0]

    ref = refs[f"stage_06_hidden_layer{li:02d}"]
    diff = np.abs(hidden.cpu().numpy().astype(np.float64) - ref)
    status = "PASS" if diff.max() < 1e-3 else "FAIL"
    print(f"layer_{li:02d}: max={diff.max():.2e} mean={diff.mean():.2e} [{status}]")
# -> layer_00: max=5.14e-07 mean=7.92e-09 [PASS]
# -> layer_01: max=1.16e-04 mean=1.06e-07 [PASS]
# -> ...
# -> layer_31: max=2.97e-04 mean=1.99e-06 [PASS]

learncat163 · 2026-07-07T06:12:14Z

MatterChat 推理对话示例

1. 单次推理

1.1 一行命令推理

from ppmat.models import build_model_from_name
from pymatgen.io.cif import CifParser

model, _ = build_model_from_name('matterchat_full')
model.to('gpu').eval()

struct = CifParser('Si.cif').get_structures()[0]
print(model.chat(struct, 'what is the chemical formula of this material?'))
# -> The chemical formula of this material is Si.

1.2 指定 config + 本地权重

from omegaconf import OmegaConf
from pymatgen.io.cif import CifParser
from ppmat.models import build_model
from ppmat.utils import save_load

config = OmegaConf.to_container(OmegaConf.load('structure_generation/configs/matterchat/matterchat_full.yaml'), resolve=True)
model = build_model(config['Model'])
save_load.load_pretrain(model, './matterchat_full/')
model.to('gpu').eval()

struct = CifParser('Si.cif').get_structures()[0]
print(model.chat(struct, 'what is the chemical formula of this material?'))

2. 多轮对话示例

对同一晶体连续提问多个问题:

from ppmat.models import build_model_from_name
from pymatgen.io.cif import CifParser

model, _ = build_model_from_name("matterchat_full")
model.to("gpu").eval()

struct = CifParser("GaN.cif").get_structures()[0]

# 4 个标准问题
prompts = [
    "what is the chemical formula of this material?",
    "what is the space group of this material?",
    "Is this material stable or not?",
    "What is the bandgap of this material?",
]

for prompt in prompts:
    answer = model.chat(struct, prompt, max_new_tokens=64)
    print(f"Q: {prompt}")
    print(f"A: {answer}")
    print()

输出:

Q: what is the chemical formula of this material?
A: The chemical formula of this material is GaN.

Q: what is the space group of this material?
A: The space group of this material is P6_3mc.

Q: Is this material stable or not?
A: This material is not stable.

Q: What is the bandgap of this material?
A: The bandgap of this material is 1.68300.

3. 批量推理 (多个 CIF)

import os
from ppmat.models import build_model_from_name
from pymatgen.io.cif import CifParser

model, _ = build_model_from_name("matterchat_full")
model.to("gpu").eval()

cif_dir = "path/to/cif_files"
prompt = "what is the chemical formula of this material?"

for fname in sorted(os.listdir(cif_dir)):
    if not fname.endswith(".cif"):
        continue
    struct = CifParser(os.path.join(cif_dir, fname)).get_structures()[0]
    answer = model.chat(struct, prompt, max_new_tokens=64)
    print(f"{fname}: {answer}")

paddle-bot · 2026-07-14T12:48:28Z

Thanks for your contribution!

leeleolay

辛苦修改整体的代码规范符合套件风格

leeleolay · 2026-07-16T07:07:14Z

 from ppmat.datasets.oc20_s2ef_dataset import OC20S2EFDataset  # noqa
 from ppmat.datasets.qm9_dataset import QM9Dataset # noqa
 from ppmat.datasets.omol25_dataset import OMol25Dataset
+from ppmat.models.matterchat.trainer import MTDataset  # noqa


使用已有trainer

leeleolay · 2026-07-16T07:07:37Z

使用默认的collator

leeleolay · 2026-07-16T07:08:13Z

套件内已有chgnet，辛苦使用

leeleolay · 2026-07-16T07:08:31Z

使用已有的graph_converter

leeleolay · 2026-07-16T07:09:17Z

ppmatSim已经支持相关的功能，复用已有的

leeleolay · 2026-07-16T07:10:02Z

vasp在这个模型里的作用是？

leeleolay · 2026-07-16T07:10:33Z

不符合已有规范，config不这么处理

support matterchat

1bc4d0b

leeleolay added the MIIT Program label Jul 7, 2026

leeleolay changed the title ~~support matterchat~~ 【MIIT program】support matterchat Jul 7, 2026

paddle-bot Bot added the contributor External developers label Jul 14, 2026

leeleolay reviewed Jul 16, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

【MIIT program】support matterchat - #305

【MIIT program】support matterchat#305
learncat163 wants to merge 1 commit into
PaddlePaddle:developfrom
learncat163:pr-matterchat

learncat163 commented Jul 7, 2026 •

edited

Loading

Uh oh!

learncat163 commented Jul 7, 2026

Uh oh!

paddle-bot Bot commented Jul 14, 2026

Uh oh!

leeleolay left a comment

Uh oh!

leeleolay Jul 16, 2026

Uh oh!

leeleolay Jul 16, 2026

Uh oh!

leeleolay Jul 16, 2026

Uh oh!

leeleolay Jul 16, 2026

Uh oh!

leeleolay Jul 16, 2026

Uh oh!

leeleolay Jul 16, 2026

Uh oh!

leeleolay Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

learncat163 commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

MatterChat diff (PyTorch and PaddlePaddle)

1. CHGNet Diff

结果

代码样例

PyTorch 参考数据生成

PaddlePaddle 对比

2. Q-Former Diff

结果

代码样例

PyTorch 参考数据生成

PaddlePaddle 对比

3. Mistral LLM Diff

结果汇总

逐层误差

代码样例

PyTorch 参考数据生成

PaddlePaddle 对比

Uh oh!

learncat163 commented Jul 7, 2026

MatterChat 推理对话示例

1. 单次推理

1.1 一行命令推理

1.2 指定 config + 本地权重

2. 多轮对话示例

3. 批量推理 (多个 CIF)

Uh oh!

paddle-bot Bot commented Jul 14, 2026

Uh oh!

leeleolay left a comment

Choose a reason for hiding this comment

Uh oh!

leeleolay Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

leeleolay Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

leeleolay Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

leeleolay Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

leeleolay Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

leeleolay Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

leeleolay Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

learncat163 commented Jul 7, 2026 •

edited

Loading