ESM-2 预训练

ESM-2 的预训练检查点可在 8M、650M 和 3B 模型大小下使用。这些模型由 bionemo-framework 团队训练，旨在重现 Lin 等人，《科学》(2023) 的原始训练结果，使用了更新的 UniProt 数据并利用了 bionemo 训练基础设施。完整的预训练数据和训练/测试拆分可用。

模型收敛

在 NVIDIA 验证集上评估的验证困惑度。

ESM-2 Pre-training Convergence

模型大小	50 万次更新时的困惑度
8M	10.26
650M	7.14
3B	6.42

预训练配方

8M650M3B

esm2_8m_ckpt_path = load("esm2/8m:2.0")

训练脚本

训练参数	值
GPU 数量	32
GPU 类型	A100
批量大小（每个设备）	64

train_esm2 \
  --create-tensorboard-logger \
  --resume-if-exists \
  --wandb-project=<wandb-project-name> \
  --save-top-k=10 \
  --train-cluster-path=/data/train_clusters.parquet \  # (1)!
  --train-database-path=/data/train.db \
  --valid-cluster-path=/data/valid_clusters.parquet \
  --valid-database-path=/data/validation.db \
  --num-steps=500_000 \
  --metric-to-monitor-for-checkpoints=val_loss \
  --micro-batch-size=64 \
  --num-nodes=4 \
  --num-gpus=8 \
  --val-check-interval=10000 \
  --limit-val-batches=1.0 \
  --result-dir=/results/esm2_pretrain_8m \
  --experiment-name=esm2_pretrain_8m \
  --num-layers=6 \
  --hidden-size=320 \
  --num-attention-heads=20 \
  --ffn-hidden-size=1280;

此处的路径必须挂载到 bionemo-framework docker 镜像中。

esm2_650m_ckpt_path = load("esm2/nv_650m:2.1")

训练脚本

训练参数	值
GPU 数量	64
GPU 类型	H100
批量大小（每个设备）	32

train_esm2 \
  --create-tensorboard-logger \
  --resume-if-exists \
  --wandb-project=<wandb-project-name> \
  --save-top-k=10 \
  --train-cluster-path=/data/train_clusters.parquet \  # (1)!
  --train-database-path=/data/train.db \
  --valid-cluster-path=/data/valid_clusters.parquet \
  --valid-database-path=/data/validation.db \
  --num-steps=500_000 \
  --metric-to-monitor-for-checkpoints=val_loss \
  --micro-batch-size=32 \
  --num-nodes=8 \
  --num-gpus=8 \
  --val-check-interval=10000 \
  --limit-val-batches=1.0 \
  --result-dir=/results/esm2_pretrain_650m \
  --experiment-name=esm2_pretrain_650m \
  --min-seq-length=1024 \
  --max-seq-length=1024 \
  --num-layers=33 \
  --hidden-size=1280 \
  --num-attention-heads=20 \
  --ffn-hidden-size=5120;

此处的路径必须挂载到 bionemo-framework docker 镜像中。

esm2_3b_ckpt_path = load("esm2/nv_3b:2.1")

训练脚本

训练参数	值
GPU 数量	128
GPU 类型	H100
批量大小（每个设备）	16
预热步数	20,000

train_esm2 \
  --create-tensorboard-logger \
  --resume-if-exists \
  --wandb-project=<wandb-project-name> \
  --save-top-k=10 \
  --train-cluster-path=/data/train_clusters.parquet \  # (2)!
  --train-database-path=/data/train.db \
  --valid-cluster-path=/data/valid_clusters.parquet \
  --valid-database-path=/data/validation.db \
  --num-steps=500_000 \
  --warmup-steps=20_000 \  # (1)!
  --metric-to-monitor-for-checkpoints=val_loss \
  --micro-batch-size=16 \
  --num-nodes=16 \
  --num-gpus=8 \
  --val-check-interval=2500 \
  --limit-val-batches=1.0 \
  --result-dir=/results/esm2_pretrain_3b \
  --experiment-name=esm2_pretrain_3b \
  --min-seq-length=1024 \
  --max-seq-length=1024 \
  --num-layers=36 \
  --hidden-size=2560 \
  --num-attention-heads=40 \
  --ffn-hidden-size=10240;

对于 ESM-2 3B 的已发布训练配方，我们不得不将预热步数增加 10 倍，这可能是使用 fp16 精度训练的。这给了我们一个总体上相似的初始曲线，但避免了在约 2,000 步时的收敛问题。
此处的路径必须挂载到 bionemo-framework docker 镜像中。