feat: 支持固定超参训练并部署 4 专家稀疏 MoE 模型

This commit is contained in:
Michelle0574 2026-08-15 12:22:21 +00:00
parent c0203ebe76
commit aeddf8e1ee
7 changed files with 192 additions and 157 deletions

View File

@ -491,6 +491,8 @@ def main(
n_trials: int = 20, n_trials: int = 20,
epochs_per_trial: int = 30, epochs_per_trial: int = 30,
patience: int = 10, patience: int = 10,
fixed_params_json: Optional[Path] = None,
fixed_epoch_mean: Optional[int] = None,
# 训练参数 # 训练参数
batch_size: int = 32, batch_size: int = 32,
# 最终训练参数 # 最终训练参数
@ -613,7 +615,17 @@ def main(
# 预热 RDKit 缓存(在整个训练流程中共享) # 预热 RDKit 缓存(在整个训练流程中共享)
rdkit_cache = warmup_rdkit_cache(full_dataset.smiles) rdkit_cache = warmup_rdkit_cache(full_dataset.smiles)
# 运行 Optuna 调参 if fixed_params_json is not None:
if fixed_epoch_mean is None:
raise typer.BadParameter("--fixed-params-json 必须配合 --fixed-epoch-mean 使用")
logger.info(f"Skipping Optuna, loading fixed params from {fixed_params_json}")
with open(fixed_params_json) as f:
best_params = json.load(f)
epoch_mean = fixed_epoch_mean
study = None
logger.info(f"Fixed params: {best_params}")
logger.info(f"Fixed epoch_mean: {epoch_mean}")
else:
logger.info(f"\nRunning {n_folds}-fold Optuna with {n_trials} trials...") logger.info(f"\nRunning {n_folds}-fold Optuna with {n_trials} trials...")
study_path = output_dir / "optuna_study.sqlite3" study_path = output_dir / "optuna_study.sqlite3"
@ -654,6 +666,7 @@ def main(
json.dump({"epoch_mean": epoch_mean}, f) json.dump({"epoch_mean": epoch_mean}, f)
# 保存 Optuna 试验历史 # 保存 Optuna 试验历史
if study is not None:
trials_history = [] trials_history = []
for trial in study.trials: for trial in study.trials:
trials_history.append({ trials_history.append({
@ -691,8 +704,6 @@ def main(
with open(output_dir / "class_weights.json", "w") as f: with open(output_dir / "class_weights.json", "w") as f:
json.dump(class_weights_info, f, indent=2) json.dump(class_weights_info, f, indent=2)
# 架构超参一次性解析Optuna 采样值优先于 CLI 默认值。
# 解析结果必须同时喂给 create_model 和 config否则 load_model 回读时结构对不上。
arch = { arch = {
"moe_n_experts": best_params.get("moe_n_experts", moe_n_experts), "moe_n_experts": best_params.get("moe_n_experts", moe_n_experts),
"moe_top_k": best_params.get("moe_top_k", moe_top_k), "moe_top_k": best_params.get("moe_top_k", moe_top_k),

View File

@ -3,7 +3,7 @@
"lr": 0.000980484886752915, "lr": 0.000980484886752915,
"weight_decay": 0.00014367509113739864, "weight_decay": 0.00014367509113739864,
"backbone_lr_ratio": 0.13093310334961422, "backbone_lr_ratio": 0.13093310334961422,
"moe_n_experts": 2, "moe_n_experts": 4,
"moe_top_k": 2, "moe_top_k": 2,
"moe_expert_hidden_mult": 2, "moe_expert_hidden_mult": 2,
"llm_lora_r": 8, "llm_lora_r": 8,

View File

@ -1,126 +1,126 @@
{ {
"train": [ "train": [
{ {
"loss": 5.662022154286222, "loss": 5.694288951046062,
"loss_size": 0.9464401806581695, "loss_size": 1.015234741259296,
"loss_pdi": 0.6576592483610477, "loss_pdi": 0.6592788083373375,
"loss_ee": 1.0441054870497506, "loss_ee": 1.027098409409793,
"loss_delivery": 1.0691580114499577, "loss_delivery": 1.0114885227016683,
"loss_biodist": 0.7718977461446006, "loss_biodist": 0.7630386647750746,
"loss_toxic": 0.47160032813279135, "loss_toxic": 0.5493527664890829,
"loss_moe_lb": 1.9999999887538407, "loss_moe_lb": 2.045894391131851,
"loss_aux_moe": 1.096376850709038, "loss_aux_moe": 1.0363417644545716,
"loss_aux_llm": 1.2355621509113401 "loss_aux_llm": 1.1829107824080396
}, },
{ {
"loss": 4.549718969273117, "loss": 4.587394336484513,
"loss_size": 0.9265210998227011, "loss_size": 0.9309988603839334,
"loss_pdi": 0.6320651516599475, "loss_pdi": 0.6420835014784111,
"loss_ee": 0.9564013177493833, "loss_ee": 0.9480358497151788,
"loss_delivery": 0.8463795804682205, "loss_delivery": 0.809673407519201,
"loss_biodist": 0.48631338606465535, "loss_biodist": 0.5112692536610477,
"loss_toxic": 0.25031149114991696, "loss_toxic": 0.30279909285171974,
"loss_moe_lb": 1.9999999910030726, "loss_moe_lb": 2.0554285611746446,
"loss_aux_moe": 0.8730144033928946, "loss_aux_moe": 0.9779460283242306,
"loss_aux_llm": 1.0696029401612732 "loss_aux_llm": 0.9288384830754883
}, },
{ {
"loss": 4.183341939494295, "loss": 4.160386994199933,
"loss_size": 0.8995583306927726, "loss_size": 0.9387357568122306,
"loss_pdi": 0.6350391439671786, "loss_pdi": 0.5914091456611201,
"loss_ee": 0.9038635166186206, "loss_ee": 0.9026568987459507,
"loss_delivery": 0.8323076782080362, "loss_delivery": 0.84944020264413,
"loss_biodist": 0.44391101247297143, "loss_biodist": 0.3876211567307418,
"loss_toxic": 0.1671132598740031, "loss_toxic": 0.18268741205563102,
"loss_moe_lb": 1.9999999865046088, "loss_moe_lb": 2.1000781711542382,
"loss_aux_moe": 0.8456886384003567, "loss_aux_moe": 0.9373761130461715,
"loss_aux_llm": 1.0173823231796049 "loss_aux_llm": 0.9303386138295228
}, },
{ {
"loss": 3.9456988685535936, "loss": 3.809376889804624,
"loss_size": 0.925582969104344, "loss_size": 0.888136670766574,
"loss_pdi": 0.586557297211773, "loss_pdi": 0.5791616782827197,
"loss_ee": 0.867810919599713, "loss_ee": 0.8553679860987753,
"loss_delivery": 0.7656616399521535, "loss_delivery": 0.8118948299648627,
"loss_biodist": 0.4162127935099152, "loss_biodist": 0.30093744889182866,
"loss_toxic": 0.2137195658123226, "loss_toxic": 0.15547859341890183,
"loss_moe_lb": 1.9999999887538407, "loss_moe_lb": 2.102432388179707,
"loss_aux_moe": 0.7506198146784643, "loss_aux_moe": 1.1255526001442153,
"loss_aux_llm": 0.9585366686981804 "loss_aux_llm": 0.9655606655298539
}, },
{ {
"loss": 3.4442428957741216, "loss": 3.1629226522625618,
"loss_size": 0.8259714532573268, "loss_size": 0.7107668774870207,
"loss_pdi": 0.5488103347004585, "loss_pdi": 0.5377992931401955,
"loss_ee": 0.8086997242468708, "loss_ee": 0.7858308424364846,
"loss_delivery": 0.7543281304808158, "loss_delivery": 0.7225946192190332,
"loss_biodist": 0.3298270061330975, "loss_biodist": 0.23169103110173964,
"loss_toxic": 0.09599270570566351, "loss_toxic": 0.10101612598441963,
"loss_moe_lb": 1.9999999955015362, "loss_moe_lb": 2.1294895545491634,
"loss_aux_moe": 0.8137425728282839, "loss_aux_moe": 0.8120556049189478,
"loss_aux_llm": 1.044384686873769 "loss_aux_llm": 0.8958881971127582
}, },
{ {
"loss": 3.108477974837681, "loss": 2.988117564399287,
"loss_size": 0.7671496489981435, "loss_size": 0.7456601148225227,
"loss_pdi": 0.5366503028374798, "loss_pdi": 0.5200111323370123,
"loss_ee": 0.7869719379353073, "loss_ee": 0.7696504339856921,
"loss_delivery": 0.685029683248052, "loss_delivery": 0.702117268874679,
"loss_biodist": 0.2828179093183212, "loss_biodist": 0.19998157122787438,
"loss_toxic": 0.08221106674908749, "loss_toxic": 0.08327638518992143,
"loss_moe_lb": 1.9999999910030726, "loss_moe_lb": 2.1596880116552675,
"loss_aux_moe": 0.7168521201413758, "loss_aux_moe": 0.763107772424536,
"loss_aux_llm": 0.9446638939234445 "loss_aux_llm": 0.8947458216604197
}, },
{ {
"loss": 2.983976681277437, "loss": 2.7413068245042047,
"loss_size": 0.7495474740159962, "loss_size": 0.6423791460692883,
"loss_pdi": 0.5227343659355955, "loss_pdi": 0.5119175066081982,
"loss_ee": 0.7467741780685928, "loss_ee": 0.7589052460103665,
"loss_delivery": 0.6811002301368511, "loss_delivery": 0.6498366547924168,
"loss_biodist": 0.24798926155803339, "loss_biodist": 0.1679054581612911,
"loss_toxic": 0.11302385037876929, "loss_toxic": 0.09372632653026332,
"loss_moe_lb": 1.9999999955015362 "loss_moe_lb": 2.148664339533392
}, },
{ {
"loss": 3.014890724757932, "loss": 2.7408997472727075,
"loss_size": 0.7314784642098084, "loss_size": 0.6384303912801562,
"loss_pdi": 0.48637263083233023, "loss_pdi": 0.49059138151834597,
"loss_ee": 0.769600696158859, "loss_ee": 0.7721629592607606,
"loss_delivery": 0.8273954027912246, "loss_delivery": 0.7107182373556326,
"loss_biodist": 0.21140338072799286, "loss_biodist": 0.15433137111787526,
"loss_toxic": 0.08676587497249338, "loss_toxic": 0.0748163205291417,
"loss_moe_lb": 1.9999999955015362 "loss_moe_lb": 2.154493790752483
}, },
{ {
"loss": 2.774192355713754, "loss": 2.610751993251297,
"loss_size": 0.6206410539881239, "loss_size": 0.6624566971693399,
"loss_pdi": 0.4826473706173447, "loss_pdi": 0.4767126595636584,
"loss_ee": 0.7343562110415045, "loss_ee": 0.7076389159796372,
"loss_delivery": 0.6946425725758638, "loss_delivery": 0.636365042331646,
"loss_biodist": 0.19993741290186937, "loss_biodist": 0.15519743205382014,
"loss_toxic": 0.1272664367724298, "loss_toxic": 0.06698143800538262,
"loss_moe_lb": 1.9999999932523043 "loss_moe_lb": 2.1544579132547916
}, },
{ {
"loss": 2.683339330385316, "loss": 2.581533803130096,
"loss_size": 0.5581827477885867, "loss_size": 0.6987560360499148,
"loss_pdi": 0.4963476042140205, "loss_pdi": 0.4562466161431007,
"loss_ee": 0.7145765797709519, "loss_ee": 0.7065791735109294,
"loss_delivery": 0.6466561276817097, "loss_delivery": 0.6151787998541346,
"loss_biodist": 0.1851363978436533, "loss_biodist": 0.1382373322223155,
"loss_toxic": 0.1616253986507698, "loss_toxic": 0.06480917761778845,
"loss_moe_lb": 1.9999999932523043 "loss_moe_lb": 2.137455944745046
}, },
{ {
"loss": 2.5718357652988075, "loss": 2.3695818658144967,
"loss_size": 0.6112437141391466, "loss_size": 0.5367685263308714,
"loss_pdi": 0.4928069238392812, "loss_pdi": 0.47306264318385216,
"loss_ee": 0.7109674007262824, "loss_ee": 0.6626323349070999,
"loss_delivery": 0.6155119884829476, "loss_delivery": 0.5844551013817765,
"loss_biodist": 0.18839857628885306, "loss_biodist": 0.13707212931564394,
"loss_toxic": 0.05047845465344166, "loss_toxic": 0.06755097333027735,
"loss_moe_lb": 1.9999999842553768 "loss_moe_lb": 2.1183437378901355
} }
], ],
"val": [] "val": []

Binary file not shown.

View File

@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1 version https://git-lfs.github.com/spec/v1
oid sha256:b96d3867be53a3e23abdf291baeec4a9d610da58dc790220a345e7012187aaa0 oid sha256:5b1c826c11338f313eea643f839fbc27c4c1e0f0f19d58d672d405cc731caa2a
size 53008646 size 58261734

View File

@ -0,0 +1,16 @@
{
"dropout": 0.19323529834582587,
"lr": 0.000980484886752915,
"weight_decay": 0.00014367509113739864,
"backbone_lr_ratio": 0.13093310334961422,
"moe_n_experts": 4,
"moe_top_k": 2,
"moe_expert_hidden_mult": 2,
"llm_lora_r": 8,
"d_model": 256,
"num_heads": 8,
"n_attn_layers": 4,
"fusion_strategy": "attention",
"head_hidden_dim": 128,
"set_transformer_block": "sab"
}

View File

@ -11,6 +11,8 @@ N_TRIALS=${N_TRIALS:-20}
EPOCHS=${EPOCHS:-20} # 与 nested CV 的 EPOCHS 保持一致 EPOCHS=${EPOCHS:-20} # 与 nested CV 的 EPOCHS 保持一致
PATIENCE=${PATIENCE:-5} # 与 nested CV 的 PATIENCE 保持一致 PATIENCE=${PATIENCE:-5} # 与 nested CV 的 PATIENCE 保持一致
N_FOLDS=${N_FOLDS:-3} N_FOLDS=${N_FOLDS:-3}
FIXED_PARAMS=${FIXED_PARAMS:-} # 非空则跳过 Optuna直接用这份超参
FIXED_EPOCH_MEAN=${FIXED_EPOCH_MEAN:-11}
BATCH=${BATCH:-8} BATCH=${BATCH:-8}
REG_BYPASS=${REG_BYPASS:-off} REG_BYPASS=${REG_BYPASS:-off}
FREEZE=${FREEZE:-3} FREEZE=${FREEZE:-3}
@ -63,6 +65,12 @@ if [ -n "${PRETRAIN}" ] && [ -f "${PRETRAIN}" ]; then
elif [ -n "${PRETRAIN}" ]; then elif [ -n "${PRETRAIN}" ]; then
echo "[$(date '+%F %T')] 警告:${PRETRAIN} 不存在,跳过预训练初始化" >>"${LOG}" echo "[$(date '+%F %T')] 警告:${PRETRAIN} 不存在,跳过预训练初始化" >>"${LOG}"
fi fi
if [ -n "${FIXED_PARAMS}" ]; then
if [ ! -f "${FIXED_PARAMS}" ]; then
echo "[$(date '+%F %T')] 错误:${FIXED_PARAMS} 不存在" >>"${LOG}"; exit 1
fi
EXTRA+=(--fixed-params-json "${FIXED_PARAMS}" --fixed-epoch-mean "${FIXED_EPOCH_MEAN}")
fi
for attempt in $(seq 1 "${MAX_RETRY}"); do for attempt in $(seq 1 "${MAX_RETRY}"); do
wait_free wait_free