mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-07-21 21:22:05 +08:00
281 lines
10 KiB
Python
281 lines
10 KiB
Python
"""
|
||
LLM 编码器验证脚本（统一版）
|
||
|
||
对四种 LLM encoder 中的任意一种做四项检查：
|
||
TEST 1: encoder 独立输出（shape / NaN / 梯度 / 缓存 / batch 一致性）
|
||
TEST 2: 与 LNPModel 集成（任务头维度 / biodist 和为1 / 反向传播 / fusion token 数）
|
||
TEST 3: 可复现性（同种子相同输出，不同种子不同输出）
|
||
TEST 4: 过拟合 sanity（小数据上 loss 下降、R² > 0）
|
||
|
||
用法:
|
||
python verify_llm_encoder.py \
|
||
--model_path <llm_model_path> \
|
||
--llm_type <biot5|molt5|moleculestm|selfiested> \
|
||
--lnp_repo_path <仓库根目录>
|
||
|
||
不同 encoder 的可训练主体属性名不同，本脚本会按 llm_type 自动选择：
|
||
biot5 -> encoder._llm
|
||
molt5 -> encoder._encoder
|
||
moleculestm -> encoder._gnn
|
||
selfiested -> encoder._model
|
||
"""
|
||
|
||
import argparse
|
||
import sys
|
||
import random
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
|
||
# llm_type -> (encoder 模块, 类名, 冻结主体属性名)
|
||
LLM_REGISTRY = {
|
||
"biot5": ("lnp_ml.modeling.encoders.llm_encoder", "LLMEncoder", "_llm"),
|
||
"molt5": ("lnp_ml.modeling.encoders.molt5_encoder", "MolT5Encoder", "_encoder"),
|
||
"moleculestm": ("lnp_ml.modeling.encoders.moleculestm_encoder", "MoleculeSTMEncoder", "_gnn"),
|
||
"selfiested": ("lnp_ml.modeling.encoders.selfiested_encoder", "SELFIESTEDEncoder", "_model"),
|
||
}
|
||
|
||
|
||
def set_seed(seed: int = 42) -> None:
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
torch.cuda.manual_seed_all(seed)
|
||
print(f"[Seed] All seeds fixed to {seed}")
|
||
|
||
|
||
def build_encoder(llm_type, model_path, d_model, device):
|
||
"""按 llm_type 构造 encoder，处理 model_path / checkpoint_path 参数名差异。"""
|
||
import importlib
|
||
module_name, class_name, _ = LLM_REGISTRY[llm_type]
|
||
module = importlib.import_module(module_name)
|
||
EncoderCls = getattr(module, class_name)
|
||
|
||
if llm_type == "moleculestm":
|
||
return EncoderCls(checkpoint_path=model_path, d_model=d_model, device=device)
|
||
else:
|
||
return EncoderCls(model_path=model_path, d_model=d_model, device=device)
|
||
|
||
|
||
def get_frozen_body(encoder, llm_type):
|
||
"""取出该 encoder 的冻结主体（用于检查参数是否冻结）。"""
|
||
_, _, attr = LLM_REGISTRY[llm_type]
|
||
return getattr(encoder, attr)
|
||
|
||
|
||
SAMPLE_SMILES = [
|
||
"CCCCCCCCCCCCCCCC(=O)OCC(COC(=O)CCCCCCCCCCCCCCC)OC(=O)CCCCCCC",
|
||
"CCN(CC)CCOC(=O)C1=CC=CC=C1",
|
||
"CC(C)(C)OC(=O)N1CCC(CC1)C(=O)O",
|
||
"O=C(O)c1ccccc1O",
|
||
]
|
||
|
||
|
||
def make_tabular(batch_size, device):
|
||
return {
|
||
"comp": torch.randn(batch_size, 5, device=device),
|
||
"phys": torch.randn(batch_size, 12, device=device),
|
||
"help": torch.randn(batch_size, 4, device=device),
|
||
"exp": torch.randn(batch_size, 32, device=device),
|
||
}
|
||
|
||
|
||
def test_1_standalone(llm_type, model_path, d_model, device):
|
||
print("\n" + "=" * 60)
|
||
print("TEST 1: encoder standalone")
|
||
print("=" * 60)
|
||
|
||
print("\n[1a] Testing output shape...")
|
||
encoder = build_encoder(llm_type, model_path, d_model, device)
|
||
out = encoder(SAMPLE_SMILES)
|
||
assert out.shape == (len(SAMPLE_SMILES), d_model), f"bad shape {out.shape}"
|
||
print(f" Output shape: {out.shape} OK")
|
||
|
||
print("[1b] Checking for NaN/Inf...")
|
||
assert not torch.isnan(out).any() and not torch.isinf(out).any()
|
||
print(f" No NaN/Inf OK")
|
||
print(f" Output stats: mean={out.mean():.4f}, std={out.std():.4f}, "
|
||
f"min={out.min():.4f}, max={out.max():.4f}")
|
||
|
||
print("[1c] Checking gradients...")
|
||
body = get_frozen_body(encoder, llm_type)
|
||
n_trainable_body = sum(p.numel() for p in body.parameters() if p.requires_grad)
|
||
assert n_trainable_body == 0, f"frozen body has {n_trainable_body} trainable params"
|
||
print(f" Frozen body ({0} trainable params) OK")
|
||
n_proj = sum(p.numel() for p in encoder.projection.parameters() if p.requires_grad)
|
||
assert n_proj > 0
|
||
print(f" Projection trainable ({n_proj:,} params) OK")
|
||
loss = encoder(SAMPLE_SMILES).sum()
|
||
loss.backward()
|
||
grad_ok = any(p.grad is not None for p in encoder.projection.parameters())
|
||
assert grad_ok
|
||
print(f" Gradient flows to projection OK")
|
||
|
||
print("[1e] Single vs batch consistency...")
|
||
encoder.eval()
|
||
with torch.no_grad():
|
||
batch_out = encoder(SAMPLE_SMILES)
|
||
single_out = torch.cat([encoder([s]) for s in SAMPLE_SMILES], dim=0)
|
||
assert torch.allclose(batch_out, single_out, atol=1e-4)
|
||
print(f" Batch == single OK")
|
||
print("\nTEST 1 PASSED")
|
||
|
||
|
||
def test_2_integration(llm_type, model_path, d_model, device):
|
||
print("\n" + "=" * 60)
|
||
print("TEST 2: Integration with LNPModel")
|
||
print("=" * 60)
|
||
|
||
from lnp_ml.modeling.models import LNPModelWithoutMPNN
|
||
model = LNPModelWithoutMPNN(
|
||
d_model=d_model, use_llm=True, llm_type=llm_type,
|
||
llm_model_path=model_path, llm_device=device,
|
||
).to(device)
|
||
|
||
print("\n[2a] Testing full forward pass shapes...")
|
||
tab = make_tabular(len(SAMPLE_SMILES), device)
|
||
outputs = model(SAMPLE_SMILES, tab)
|
||
expected = {"size": 1, "pdi": 4, "ee": 3, "delivery": 1, "biodist": 7, "toxic": 2}
|
||
for k, dim in expected.items():
|
||
assert outputs[k].shape == (len(SAMPLE_SMILES), dim), f"{k}: {outputs[k].shape}"
|
||
print(f" outputs['{k}']: {tuple(outputs[k].shape)} OK")
|
||
|
||
print("[2b] Checking biodist sums to 1...")
|
||
sums = outputs["biodist"].sum(dim=-1)
|
||
assert torch.allclose(sums, torch.ones_like(sums), atol=1e-4)
|
||
print(f" biodist sums: {[round(s.item(),4) for s in sums]} OK")
|
||
|
||
print("[2c] Testing backward pass...")
|
||
loss = sum(o.sum() for o in outputs.values())
|
||
loss.backward()
|
||
proj_grad = any(p.grad is not None for p in model.llm_encoder.projection.parameters())
|
||
assert proj_grad
|
||
print(f" llm_encoder.projection gets gradients OK")
|
||
body = get_frozen_body(model.llm_encoder, llm_type)
|
||
body_grad = any(p.grad is not None and p.grad.abs().sum() > 0 for p in body.parameters())
|
||
assert not body_grad
|
||
print(f" LLM body has no gradients (frozen) OK")
|
||
|
||
print("[2d] Checking fusion token count...")
|
||
# 无 MPNN 时主干 7 个 token + 1 个 V5 = 8
|
||
assert model.n_fusion_tokens == 8, f"n_fusion_tokens={model.n_fusion_tokens}"
|
||
print(f" fusion.n_tokens={model.n_fusion_tokens} OK")
|
||
print("\nTEST 2 PASSED")
|
||
|
||
|
||
def test_3_reproducibility(llm_type, model_path, d_model, device):
|
||
print("\n" + "=" * 60)
|
||
print("TEST 3: Reproducibility")
|
||
print("=" * 60)
|
||
|
||
print("\n[3a] Running twice with same seed=42...")
|
||
set_seed(42)
|
||
enc1 = build_encoder(llm_type, model_path, d_model, device)
|
||
enc1.eval()
|
||
with torch.no_grad():
|
||
o1 = enc1(SAMPLE_SMILES)
|
||
set_seed(42)
|
||
enc2 = build_encoder(llm_type, model_path, d_model, device)
|
||
enc2.eval()
|
||
with torch.no_grad():
|
||
o2 = enc2(SAMPLE_SMILES)
|
||
assert torch.allclose(o1, o2, atol=1e-5)
|
||
print(f" Same seed -> identical outputs OK")
|
||
|
||
print("[3b] Running with different seeds...")
|
||
set_seed(123)
|
||
enc3 = build_encoder(llm_type, model_path, d_model, device)
|
||
enc3.eval()
|
||
with torch.no_grad():
|
||
o3 = enc3(SAMPLE_SMILES)
|
||
assert not torch.allclose(o1, o3, atol=1e-5)
|
||
print(f" Different seeds -> different outputs OK")
|
||
print("\nTEST 3 PASSED")
|
||
|
||
|
||
def test_4_overfit(llm_type, model_path, d_model, device):
|
||
print("\n" + "=" * 60)
|
||
print("TEST 4: Overfit sanity (gradient + metric check)")
|
||
print("=" * 60)
|
||
set_seed(42)
|
||
|
||
from lnp_ml.modeling.models import LNPModelWithoutMPNN
|
||
model = LNPModelWithoutMPNN(
|
||
d_model=d_model, use_llm=True, llm_type=llm_type,
|
||
llm_model_path=model_path, llm_device=device,
|
||
).to(device)
|
||
|
||
n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
print(f"\n Trainable params: {n_trainable:,}")
|
||
|
||
smiles = SAMPLE_SMILES * 2 # 8 samples
|
||
tab = make_tabular(len(smiles), device)
|
||
target = torch.randn(len(smiles), 1, device=device)
|
||
|
||
opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=1e-3)
|
||
lossfn = nn.MSELoss()
|
||
|
||
print(f" Training 50 steps on {len(smiles)} samples...")
|
||
first_loss = None
|
||
model.train()
|
||
for step in range(1, 51):
|
||
opt.zero_grad()
|
||
pred = model.forward_delivery(smiles, tab)
|
||
loss = lossfn(pred, target)
|
||
loss.backward()
|
||
opt.step()
|
||
if first_loss is None:
|
||
first_loss = loss.item()
|
||
if step % 10 == 0:
|
||
print(f" step {step:3d}: loss={loss.item():.4f}")
|
||
|
||
last_loss = loss.item()
|
||
assert last_loss < first_loss
|
||
print(f" Loss decreased: {first_loss:.4f} -> {last_loss:.4f} OK")
|
||
|
||
model.eval()
|
||
with torch.no_grad():
|
||
pred = model.forward_delivery(smiles, tab)
|
||
ss_res = ((pred - target) ** 2).sum().item()
|
||
ss_tot = ((target - target.mean()) ** 2).sum().item()
|
||
r2 = 1 - ss_res / (ss_tot + 1e-9)
|
||
rmse = (ss_res / len(smiles)) ** 0.5
|
||
print(f"\n After 50 steps on {len(smiles)} samples (overfit test):")
|
||
print(f" R2 = {r2:.4f} (should be > 0 after overfitting)")
|
||
print(f" RMSE = {rmse:.4f}")
|
||
assert r2 > 0
|
||
print(f" R2 > 0 OK")
|
||
print("\nTEST 4 PASSED")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--model_path", required=True, help="LLM 权重路径或 HuggingFace ID")
|
||
parser.add_argument("--llm_type", required=True,
|
||
choices=["biot5", "molt5", "moleculestm", "selfiested"])
|
||
parser.add_argument("--lnp_repo_path", default=".", help="仓库根目录")
|
||
parser.add_argument("--d_model", type=int, default=256)
|
||
args = parser.parse_args()
|
||
|
||
if args.lnp_repo_path not in sys.path:
|
||
sys.path.insert(0, args.lnp_repo_path)
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
print(f"Device: {device}")
|
||
set_seed(42)
|
||
|
||
test_1_standalone(args.llm_type, args.model_path, args.d_model, device)
|
||
test_2_integration(args.llm_type, args.model_path, args.d_model, device)
|
||
test_3_reproducibility(args.llm_type, args.model_path, args.d_model, device)
|
||
test_4_overfit(args.llm_type, args.model_path, args.d_model, device)
|
||
|
||
print("\n" + "=" * 60)
|
||
print("ALL TESTS PASSED")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|