mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 14:23:20 +08:00
feat: 支持固定超参训练并部署 4 专家稀疏 MoE 模型
This commit is contained in:
Michelle0574
parent
c0203ebe76
commit
aeddf8e1ee
@ -491,6 +491,8 @@ def main(
|
||||
n_trials: int = 20,
|
||||
epochs_per_trial: int = 30,
|
||||
patience: int = 10,
|
||||
fixed_params_json: Optional[Path] = None,
|
||||
fixed_epoch_mean: Optional[int] = None,
|
||||
# 训练参数
|
||||
batch_size: int = 32,
|
||||
# 最终训练参数
|
||||
@ -613,75 +615,86 @@ def main(
|
||||
# 预热 RDKit 缓存（在整个训练流程中共享）
|
||||
rdkit_cache = warmup_rdkit_cache(full_dataset.smiles)
|
||||
|
||||
# 运行 Optuna 调参
|
||||
logger.info(f"\nRunning {n_folds}-fold Optuna with {n_trials} trials...")
|
||||
study_path = output_dir / "optuna_study.sqlite3"
|
||||
|
||||
best_params, epoch_mean, study = run_optuna_cv(
|
||||
full_dataset=full_dataset,
|
||||
strata=strata,
|
||||
device=device,
|
||||
n_trials=n_trials,
|
||||
epochs_per_trial=epochs_per_trial,
|
||||
patience=patience,
|
||||
batch_size=batch_size,
|
||||
n_folds=n_folds,
|
||||
use_mpnn=use_mpnn,
|
||||
chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
|
||||
unimol_cache=(unimol_cache if use_unimol else None),
|
||||
seed=seed,
|
||||
study_path=study_path,
|
||||
pretrain_state_dict=pretrain_state_dict,
|
||||
pretrain_config=pretrain_config,
|
||||
load_delivery_head=load_delivery_head,
|
||||
rdkit_cache=rdkit_cache,
|
||||
use_moe=use_moe,
|
||||
moe_n_experts=moe_n_experts,
|
||||
moe_top_k=moe_top_k,
|
||||
moe_expert_hidden_mult=moe_expert_hidden_mult,
|
||||
moe_jitter_noise=moe_jitter_noise,
|
||||
set_transformer_block=set_transformer_block,
|
||||
llm_kwargs=llm_kwargs,
|
||||
use_retrieval=use_retrieval,
|
||||
freeze_backbone_epochs=freeze_backbone_epochs,
|
||||
)
|
||||
|
||||
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...")
|
||||
study_path = output_dir / "optuna_study.sqlite3"
|
||||
|
||||
best_params, epoch_mean, study = run_optuna_cv(
|
||||
full_dataset=full_dataset,
|
||||
strata=strata,
|
||||
device=device,
|
||||
n_trials=n_trials,
|
||||
epochs_per_trial=epochs_per_trial,
|
||||
patience=patience,
|
||||
batch_size=batch_size,
|
||||
n_folds=n_folds,
|
||||
use_mpnn=use_mpnn,
|
||||
chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
|
||||
unimol_cache=(unimol_cache if use_unimol else None),
|
||||
seed=seed,
|
||||
study_path=study_path,
|
||||
pretrain_state_dict=pretrain_state_dict,
|
||||
pretrain_config=pretrain_config,
|
||||
load_delivery_head=load_delivery_head,
|
||||
rdkit_cache=rdkit_cache,
|
||||
use_moe=use_moe,
|
||||
moe_n_experts=moe_n_experts,
|
||||
moe_top_k=moe_top_k,
|
||||
moe_expert_hidden_mult=moe_expert_hidden_mult,
|
||||
moe_jitter_noise=moe_jitter_noise,
|
||||
set_transformer_block=set_transformer_block,
|
||||
llm_kwargs=llm_kwargs,
|
||||
use_retrieval=use_retrieval,
|
||||
freeze_backbone_epochs=freeze_backbone_epochs,
|
||||
)
|
||||
|
||||
# 保存最佳参数
|
||||
with open(output_dir / "best_params.json", "w") as f:
|
||||
json.dump(best_params, f, indent=2)
|
||||
|
||||
|
||||
with open(output_dir / "epoch_mean.json", "w") as f:
|
||||
json.dump({"epoch_mean": epoch_mean}, f)
|
||||
|
||||
|
||||
# 保存 Optuna 试验历史
|
||||
trials_history = []
|
||||
for trial in study.trials:
|
||||
trials_history.append({
|
||||
"number": trial.number,
|
||||
"value": trial.value,
|
||||
"params": trial.params,
|
||||
"user_attrs": trial.user_attrs,
|
||||
"state": str(trial.state),
|
||||
})
|
||||
|
||||
with open(output_dir / "optuna_trials.json", "w") as f:
|
||||
json.dump(trials_history, f, indent=2)
|
||||
|
||||
if study is not None:
|
||||
trials_history = []
|
||||
for trial in study.trials:
|
||||
trials_history.append({
|
||||
"number": trial.number,
|
||||
"value": trial.value,
|
||||
"params": trial.params,
|
||||
"user_attrs": trial.user_attrs,
|
||||
"state": str(trial.state),
|
||||
})
|
||||
|
||||
with open(output_dir / "optuna_trials.json", "w") as f:
|
||||
json.dump(trials_history, f, indent=2)
|
||||
|
||||
# 全量数据训练
|
||||
logger.info(f"\n{'='*60}")
|
||||
logger.info("FINAL TRAINING ON FULL DATA")
|
||||
logger.info(f"{'='*60}")
|
||||
logger.info(f"Using best params with epochs={epoch_mean}")
|
||||
logger.info(f"SWA: {use_swa}")
|
||||
|
||||
|
||||
# 创建全量 DataLoader
|
||||
full_loader = DataLoader(
|
||||
full_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn
|
||||
)
|
||||
|
||||
|
||||
# 计算类权重
|
||||
class_weights = compute_class_weights_from_loader(full_loader)
|
||||
|
||||
|
||||
# 保存类权重信息
|
||||
class_weights_info = {
|
||||
"pdi": class_weights.pdi.tolist() if class_weights.pdi is not None else None,
|
||||
@ -690,9 +703,7 @@ def main(
|
||||
}
|
||||
with open(output_dir / "class_weights.json", "w") as f:
|
||||
json.dump(class_weights_info, f, indent=2)
|
||||
|
||||
# 架构超参一次性解析：Optuna 采样值优先于 CLI 默认值。
|
||||
# 解析结果必须同时喂给 create_model 和 config，否则 load_model 回读时结构对不上。
|
||||
|
||||
arch = {
|
||||
"moe_n_experts": best_params.get("moe_n_experts", moe_n_experts),
|
||||
"moe_top_k": best_params.get("moe_top_k", moe_top_k),
|
||||
|
||||
@ -3,7 +3,7 @@
|
||||
"lr": 0.000980484886752915,
|
||||
"weight_decay": 0.00014367509113739864,
|
||||
"backbone_lr_ratio": 0.13093310334961422,
|
||||
"moe_n_experts": 2,
|
||||
"moe_n_experts": 4,
|
||||
"moe_top_k": 2,
|
||||
"moe_expert_hidden_mult": 2,
|
||||
"llm_lora_r": 8,
|
||||
|
||||
@ -1,126 +1,126 @@
|
||||
{
|
||||
"train": [
|
||||
{
|
||||
"loss": 5.662022154286222,
|
||||
"loss_size": 0.9464401806581695,
|
||||
"loss_pdi": 0.6576592483610477,
|
||||
"loss_ee": 1.0441054870497506,
|
||||
"loss_delivery": 1.0691580114499577,
|
||||
"loss_biodist": 0.7718977461446006,
|
||||
"loss_toxic": 0.47160032813279135,
|
||||
"loss_moe_lb": 1.9999999887538407,
|
||||
"loss_aux_moe": 1.096376850709038,
|
||||
"loss_aux_llm": 1.2355621509113401
|
||||
"loss": 5.694288951046062,
|
||||
"loss_size": 1.015234741259296,
|
||||
"loss_pdi": 0.6592788083373375,
|
||||
"loss_ee": 1.027098409409793,
|
||||
"loss_delivery": 1.0114885227016683,
|
||||
"loss_biodist": 0.7630386647750746,
|
||||
"loss_toxic": 0.5493527664890829,
|
||||
"loss_moe_lb": 2.045894391131851,
|
||||
"loss_aux_moe": 1.0363417644545716,
|
||||
"loss_aux_llm": 1.1829107824080396
|
||||
},
|
||||
{
|
||||
"loss": 4.549718969273117,
|
||||
"loss_size": 0.9265210998227011,
|
||||
"loss_pdi": 0.6320651516599475,
|
||||
"loss_ee": 0.9564013177493833,
|
||||
"loss_delivery": 0.8463795804682205,
|
||||
"loss_biodist": 0.48631338606465535,
|
||||
"loss_toxic": 0.25031149114991696,
|
||||
"loss_moe_lb": 1.9999999910030726,
|
||||
"loss_aux_moe": 0.8730144033928946,
|
||||
"loss_aux_llm": 1.0696029401612732
|
||||
"loss": 4.587394336484513,
|
||||
"loss_size": 0.9309988603839334,
|
||||
"loss_pdi": 0.6420835014784111,
|
||||
"loss_ee": 0.9480358497151788,
|
||||
"loss_delivery": 0.809673407519201,
|
||||
"loss_biodist": 0.5112692536610477,
|
||||
"loss_toxic": 0.30279909285171974,
|
||||
"loss_moe_lb": 2.0554285611746446,
|
||||
"loss_aux_moe": 0.9779460283242306,
|
||||
"loss_aux_llm": 0.9288384830754883
|
||||
},
|
||||
{
|
||||
"loss": 4.183341939494295,
|
||||
"loss_size": 0.8995583306927726,
|
||||
"loss_pdi": 0.6350391439671786,
|
||||
"loss_ee": 0.9038635166186206,
|
||||
"loss_delivery": 0.8323076782080362,
|
||||
"loss_biodist": 0.44391101247297143,
|
||||
"loss_toxic": 0.1671132598740031,
|
||||
"loss_moe_lb": 1.9999999865046088,
|
||||
"loss_aux_moe": 0.8456886384003567,
|
||||
"loss_aux_llm": 1.0173823231796049
|
||||
"loss": 4.160386994199933,
|
||||
"loss_size": 0.9387357568122306,
|
||||
"loss_pdi": 0.5914091456611201,
|
||||
"loss_ee": 0.9026568987459507,
|
||||
"loss_delivery": 0.84944020264413,
|
||||
"loss_biodist": 0.3876211567307418,
|
||||
"loss_toxic": 0.18268741205563102,
|
||||
"loss_moe_lb": 2.1000781711542382,
|
||||
"loss_aux_moe": 0.9373761130461715,
|
||||
"loss_aux_llm": 0.9303386138295228
|
||||
},
|
||||
{
|
||||
"loss": 3.9456988685535936,
|
||||
"loss_size": 0.925582969104344,
|
||||
"loss_pdi": 0.586557297211773,
|
||||
"loss_ee": 0.867810919599713,
|
||||
"loss_delivery": 0.7656616399521535,
|
||||
"loss_biodist": 0.4162127935099152,
|
||||
"loss_toxic": 0.2137195658123226,
|
||||
"loss_moe_lb": 1.9999999887538407,
|
||||
"loss_aux_moe": 0.7506198146784643,
|
||||
"loss_aux_llm": 0.9585366686981804
|
||||
"loss": 3.809376889804624,
|
||||
"loss_size": 0.888136670766574,
|
||||
"loss_pdi": 0.5791616782827197,
|
||||
"loss_ee": 0.8553679860987753,
|
||||
"loss_delivery": 0.8118948299648627,
|
||||
"loss_biodist": 0.30093744889182866,
|
||||
"loss_toxic": 0.15547859341890183,
|
||||
"loss_moe_lb": 2.102432388179707,
|
||||
"loss_aux_moe": 1.1255526001442153,
|
||||
"loss_aux_llm": 0.9655606655298539
|
||||
},
|
||||
{
|
||||
"loss": 3.4442428957741216,
|
||||
"loss_size": 0.8259714532573268,
|
||||
"loss_pdi": 0.5488103347004585,
|
||||
"loss_ee": 0.8086997242468708,
|
||||
"loss_delivery": 0.7543281304808158,
|
||||
"loss_biodist": 0.3298270061330975,
|
||||
"loss_toxic": 0.09599270570566351,
|
||||
"loss_moe_lb": 1.9999999955015362,
|
||||
"loss_aux_moe": 0.8137425728282839,
|
||||
"loss_aux_llm": 1.044384686873769
|
||||
"loss": 3.1629226522625618,
|
||||
"loss_size": 0.7107668774870207,
|
||||
"loss_pdi": 0.5377992931401955,
|
||||
"loss_ee": 0.7858308424364846,
|
||||
"loss_delivery": 0.7225946192190332,
|
||||
"loss_biodist": 0.23169103110173964,
|
||||
"loss_toxic": 0.10101612598441963,
|
||||
"loss_moe_lb": 2.1294895545491634,
|
||||
"loss_aux_moe": 0.8120556049189478,
|
||||
"loss_aux_llm": 0.8958881971127582
|
||||
},
|
||||
{
|
||||
"loss": 3.108477974837681,
|
||||
"loss_size": 0.7671496489981435,
|
||||
"loss_pdi": 0.5366503028374798,
|
||||
"loss_ee": 0.7869719379353073,
|
||||
"loss_delivery": 0.685029683248052,
|
||||
"loss_biodist": 0.2828179093183212,
|
||||
"loss_toxic": 0.08221106674908749,
|
||||
"loss_moe_lb": 1.9999999910030726,
|
||||
"loss_aux_moe": 0.7168521201413758,
|
||||
"loss_aux_llm": 0.9446638939234445
|
||||
"loss": 2.988117564399287,
|
||||
"loss_size": 0.7456601148225227,
|
||||
"loss_pdi": 0.5200111323370123,
|
||||
"loss_ee": 0.7696504339856921,
|
||||
"loss_delivery": 0.702117268874679,
|
||||
"loss_biodist": 0.19998157122787438,
|
||||
"loss_toxic": 0.08327638518992143,
|
||||
"loss_moe_lb": 2.1596880116552675,
|
||||
"loss_aux_moe": 0.763107772424536,
|
||||
"loss_aux_llm": 0.8947458216604197
|
||||
},
|
||||
{
|
||||
"loss": 2.983976681277437,
|
||||
"loss_size": 0.7495474740159962,
|
||||
"loss_pdi": 0.5227343659355955,
|
||||
"loss_ee": 0.7467741780685928,
|
||||
"loss_delivery": 0.6811002301368511,
|
||||
"loss_biodist": 0.24798926155803339,
|
||||
"loss_toxic": 0.11302385037876929,
|
||||
"loss_moe_lb": 1.9999999955015362
|
||||
"loss": 2.7413068245042047,
|
||||
"loss_size": 0.6423791460692883,
|
||||
"loss_pdi": 0.5119175066081982,
|
||||
"loss_ee": 0.7589052460103665,
|
||||
"loss_delivery": 0.6498366547924168,
|
||||
"loss_biodist": 0.1679054581612911,
|
||||
"loss_toxic": 0.09372632653026332,
|
||||
"loss_moe_lb": 2.148664339533392
|
||||
},
|
||||
{
|
||||
"loss": 3.014890724757932,
|
||||
"loss_size": 0.7314784642098084,
|
||||
"loss_pdi": 0.48637263083233023,
|
||||
"loss_ee": 0.769600696158859,
|
||||
"loss_delivery": 0.8273954027912246,
|
||||
"loss_biodist": 0.21140338072799286,
|
||||
"loss_toxic": 0.08676587497249338,
|
||||
"loss_moe_lb": 1.9999999955015362
|
||||
"loss": 2.7408997472727075,
|
||||
"loss_size": 0.6384303912801562,
|
||||
"loss_pdi": 0.49059138151834597,
|
||||
"loss_ee": 0.7721629592607606,
|
||||
"loss_delivery": 0.7107182373556326,
|
||||
"loss_biodist": 0.15433137111787526,
|
||||
"loss_toxic": 0.0748163205291417,
|
||||
"loss_moe_lb": 2.154493790752483
|
||||
},
|
||||
{
|
||||
"loss": 2.774192355713754,
|
||||
"loss_size": 0.6206410539881239,
|
||||
"loss_pdi": 0.4826473706173447,
|
||||
"loss_ee": 0.7343562110415045,
|
||||
"loss_delivery": 0.6946425725758638,
|
||||
"loss_biodist": 0.19993741290186937,
|
||||
"loss_toxic": 0.1272664367724298,
|
||||
"loss_moe_lb": 1.9999999932523043
|
||||
"loss": 2.610751993251297,
|
||||
"loss_size": 0.6624566971693399,
|
||||
"loss_pdi": 0.4767126595636584,
|
||||
"loss_ee": 0.7076389159796372,
|
||||
"loss_delivery": 0.636365042331646,
|
||||
"loss_biodist": 0.15519743205382014,
|
||||
"loss_toxic": 0.06698143800538262,
|
||||
"loss_moe_lb": 2.1544579132547916
|
||||
},
|
||||
{
|
||||
"loss": 2.683339330385316,
|
||||
"loss_size": 0.5581827477885867,
|
||||
"loss_pdi": 0.4963476042140205,
|
||||
"loss_ee": 0.7145765797709519,
|
||||
"loss_delivery": 0.6466561276817097,
|
||||
"loss_biodist": 0.1851363978436533,
|
||||
"loss_toxic": 0.1616253986507698,
|
||||
"loss_moe_lb": 1.9999999932523043
|
||||
"loss": 2.581533803130096,
|
||||
"loss_size": 0.6987560360499148,
|
||||
"loss_pdi": 0.4562466161431007,
|
||||
"loss_ee": 0.7065791735109294,
|
||||
"loss_delivery": 0.6151787998541346,
|
||||
"loss_biodist": 0.1382373322223155,
|
||||
"loss_toxic": 0.06480917761778845,
|
||||
"loss_moe_lb": 2.137455944745046
|
||||
},
|
||||
{
|
||||
"loss": 2.5718357652988075,
|
||||
"loss_size": 0.6112437141391466,
|
||||
"loss_pdi": 0.4928069238392812,
|
||||
"loss_ee": 0.7109674007262824,
|
||||
"loss_delivery": 0.6155119884829476,
|
||||
"loss_biodist": 0.18839857628885306,
|
||||
"loss_toxic": 0.05047845465344166,
|
||||
"loss_moe_lb": 1.9999999842553768
|
||||
"loss": 2.3695818658144967,
|
||||
"loss_size": 0.5367685263308714,
|
||||
"loss_pdi": 0.47306264318385216,
|
||||
"loss_ee": 0.6626323349070999,
|
||||
"loss_delivery": 0.5844551013817765,
|
||||
"loss_biodist": 0.13707212931564394,
|
||||
"loss_toxic": 0.06755097333027735,
|
||||
"loss_moe_lb": 2.1183437378901355
|
||||
}
|
||||
],
|
||||
"val": []
|
||||
|
||||
Binary file not shown.
@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:b96d3867be53a3e23abdf291baeec4a9d610da58dc790220a345e7012187aaa0
|
||||
size 53008646
|
||||
oid sha256:5b1c826c11338f313eea643f839fbc27c4c1e0f0f19d58d672d405cc731caa2a
|
||||
size 58261734
|
||||
|
||||
16
scripts/hparams_sparse_moe.json
Normal file
16
scripts/hparams_sparse_moe.json
Normal 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"
|
||||
}
|
||||
@ -11,6 +11,8 @@ N_TRIALS=${N_TRIALS:-20}
|
||||
EPOCHS=${EPOCHS:-20} # 与 nested CV 的 EPOCHS 保持一致
|
||||
PATIENCE=${PATIENCE:-5} # 与 nested CV 的 PATIENCE 保持一致
|
||||
N_FOLDS=${N_FOLDS:-3}
|
||||
FIXED_PARAMS=${FIXED_PARAMS:-} # 非空则跳过 Optuna，直接用这份超参
|
||||
FIXED_EPOCH_MEAN=${FIXED_EPOCH_MEAN:-11}
|
||||
BATCH=${BATCH:-8}
|
||||
REG_BYPASS=${REG_BYPASS:-off}
|
||||
FREEZE=${FREEZE:-3}
|
||||
@ -63,6 +65,12 @@ if [ -n "${PRETRAIN}" ] && [ -f "${PRETRAIN}" ]; then
|
||||
elif [ -n "${PRETRAIN}" ]; then
|
||||
echo "[$(date '+%F %T')] 警告：${PRETRAIN} 不存在，跳过预训练初始化" >>"${LOG}"
|
||||
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
|
||||
wait_free
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user