模型训练

Pydantic 配置

BioNeMo 2 为模型提供了两个入口点，分别使用 argparse 和 pydantic。两者都在下面的“模型”部分中进行了文档记录。基于 Pydantic 的配置旨在接受一个配置 yaml 文件作为输入，以及特定于上下文的参数（例如，我们是否应该从现有检查点恢复？）。这些 YAML 配置通过 Pydantic 验证器，在本例中称为“MainConfig”。此配置由几个其他 Pydantic 模型组成，有关详细信息，请参见类定义。为了使用各种标准模型的合理默认值预填充配置，我们提供了“recipes”。这些是实例化配置对象然后将其序列化为 YAML 配置文件的简单方法。从该文件，您可以直接提交它，或者修改各种参数以满足您的用例。例如，Weights and biases、设备、精度和数据集选项都非常有用。然后，您将提交此配置以进行训练。

当使用 pip 安装 esm2 或 geneformer 时，这两个工作流程被打包为可执行文件。这些命令将显示为

bionemo-geneformer-recipe
bionemo-esm2-recipe
bionemo-geneformer-train
bionemo-esm2-train

ESM-2

运行

首先，我们有一个实用程序函数，用于下载完整/测试数据和模型检查点，称为 download_bionemo_data，我们下面的示例目前正在使用它。如果对象尚未在您的本地系统上，这将下载该对象，然后无论哪种方式都返回路径。例如，如果您连续运行两次，您应该期望第二次运行它几乎立即返回路径。

注意：NVIDIA 员工应使用 pbss 而不是 ngc 作为数据源。

export MY_DATA_SOURCE="ngc"

或者对于拥有新数据等的 NVIDIA 内部员工

export MY_DATA_SOURCE="pbss"

# The fastest transformer engine environment variables in testing were the following two
TEST_DATA_DIR=$(download_bionemo_data esm2/testdata_esm2_pretrain:2.0 --source $MY_DATA_SOURCE); \
ESM2_650M_CKPT=$(download_bionemo_data esm2/650m:2.0 --source $MY_DATA_SOURCE); \

train_esm2     \
    --train-cluster-path ${TEST_DATA_DIR}/2024_03_sanity/train_clusters_sanity.parquet     \
    --train-database-path ${TEST_DATA_DIR}/2024_03_sanity/train_sanity.db     \
    --valid-cluster-path ${TEST_DATA_DIR}/2024_03_sanity/valid_clusters.parquet     \
    --valid-database-path ${TEST_DATA_DIR}/2024_03_sanity/validation.db     \
    --result-dir ./results     \
    --experiment-name test_experiment     \
    --num-gpus 1  \
    --num-nodes 1 \
    --val-check-interval 10 \
    --num-dataset-workers 1 \
    --num-steps 10 \
    --max-seq-length 1024 \
    --limit-val-batches 2 \
    --micro-batch-size 2 \
    --restore-from-checkpoint-path ${ESM2_650M_CKPT}

使用 Pydantic 配置运行

或者，我们为执行相同的工作流程提供了经过验证和序列化的配置文件入口点。这些可以使用 bionemo-esm2-recipe 入口点生成。Recipes 适用于 8m、650m 和 3b ESM2 模型。您可以通过设置 --recipe 参数来选择要使用的预设配置。然后，输出是一个序列化的配置文件，可以用于相关的 bionemo-esm2-train 命令。

# The fastest transformer engine environment variables in testing were the following two
TEST_DATA_DIR=$(download_bionemo_data esm2/testdata_esm2_pretrain:2.0 --source $MY_DATA_SOURCE); \
bionemo-esm2-recipe \
--train-cluster-path ${TEST_DATA_DIR}/2024_03_sanity/train_clusters_sanity.parquet     \
--train-database-path ${TEST_DATA_DIR}/2024_03_sanity/train_sanity.db     \
--valid-cluster-path ${TEST_DATA_DIR}/2024_03_sanity/valid_clusters.parquet     \
--valid-database-path ${TEST_DATA_DIR}/2024_03_sanity/validation.db     \
--result-dir ./results     \
--dest my_config.yaml\
--recipe esm2_8m_recipe

⚠️ 重要提示： 检查并编辑输出的 my_config.yaml 的内容，根据您的需要进行调整

注意：要从现有检查点继续训练，只需将路径 --initial-ckpt-path 传递给 recipe 命令。这将使用正确的字段填充 YAML，以确保从现有检查点初始化预训练。

要使用传递的配置提交训练作业，首先使用您选择的任何其他执行参数更新 yaml 文件：设备数量、worker 数量、步数等。其次，调用我们的训练入口点。为此，我们需要三件事

• 配置文件，上一步生成的 YAML
• 模型配置类型，在本例中为预训练配置。这将根据预训练所需的参数验证配置 YAML 中的参数。或者，可以在此处指定使用自定义任务头进行微调之类的事情。这允许混合/匹配数据模块和各种任务。
• 数据配置类型，这指定如何解析、验证和准备 DataModule。这可能会根据任务而变化，例如，预训练 ESM2 使用面向蛋白质簇的采样方法。在推理或微调预训练模型的情况下，简单的 fasta 文件可能就足够了。DataConfig 类型和 DataModule 类型之间存在一对一的关系。

⚠️ 警告： 此设置不配置 Weights and Biases。编辑您的 config YAML 并使用您的 WandB 详细信息填充它。

bionemo-esm2-train \
--data-config-cls bionemo.esm2.run.config_models.ESM2DataConfig \
--model-config-cls bionemo.esm2.run.config_models.ExposedESM2PretrainConfig \
--config my_config.yaml

注意：data-config-cls 和 model-config-cls 都具有对应于 ESM2DataConfig 和 ExposedESM2PretrainingConfig 的默认值

DataConfigCls 和 ModelConfigCls 也可以引用用户本地定义的类型。只要 python 知道如何导入指定的路径，就可以配置它们。例如，您可能有一个自定义 Dataset/DataModule，您想将其与现有 recipe 混合使用。在这种情况下，您可以定义一个 DataConfig 对象，其中泛型指定为您的 DataModule 类型，然后将配置类型传递给训练 recipe。

Geneformer

运行

与 ESM-2 类似，您可以通过我们的实用程序函数下载数据集和检查点。

TEST_DATA_DIR=$(download_bionemo_data single_cell/testdata-20241203 --source $MY_DATA_SOURCE); \
GENEFORMER_10M_CKPT=$(download_bionemo_data geneformer/10M_240530:2.0 --source $MY_DATA_SOURCE); \
train_geneformer     \
    --data-dir ${TEST_DATA_DIR}/cellxgene_2023-12-15_small_processed_scdl    \
    --result-dir ./results     \
    --restore-from-checkpoint-path ${GENEFORMER_10M_CKPT} \
    --experiment-name test_experiment     \
    --num-gpus 1  \
    --num-nodes 1 \
    --val-check-interval 10 \
    --num-dataset-workers 0 \
    --num-steps 55 \
    --seq-length 128 \
    --limit-val-batches 2 \
    --micro-batch-size 2

要进行微调，您需要指定模型和损失的不同组合。将上一步输出的配置文件路径作为 --restore-from-checkpoint-path 传递，并将 --training-model-config-class 更改为新创建的 model-config-class。

虽然目前没有 CLI 选项可以热插拔不同的数据模块和处理函数，但您可以复制 sub-projects/bionemo-geneformer/geneformer/scripts/train_geneformer.py 并修改初始化的 DataModule 类。

简单的微调示例（注意：请将 --restore-from-checkpoint-path 更改为上次训练运行输出的检查点目录路径）

TEST_DATA_DIR=$(download_bionemo_data single_cell/testdata-20241203 --source $MY_DATA_SOURCE); \
train_geneformer     \
    --data-dir ${TEST_DATA_DIR}/cellxgene_2023-12-15_small_processed_scdl    \
    --result-dir ./results     \
    --experiment-name test_finettune_experiment     \
    --num-gpus 1  \
    --num-nodes 1 \
    --val-check-interval 10 \
    --num-dataset-workers 0 \
    --num-steps 55 \
    --seq-length 128 \
    --limit-val-batches 2 \
    --micro-batch-size 2 \
    --training-model-config-class FineTuneSeqLenBioBertConfig \
    --restore-from-checkpoint-path results/test_experiment/dev/checkpoints/test_experiment--val_loss=4.3506-epoch=1-last

使用 Pydantic 配置运行

或者，我们为执行相同的工作流程提供了经过验证和序列化的配置文件入口点。Recipes 适用于 10m 和 106m geneformer 模型。此外，我们还提供了一个微调示例 recipe，其中目标是对 token ID 进行“回归”而不是传统的掩码语言模型方法。在实践中，您可能需要实现自己的 DataModule、DataConfig 和 Finetuning 模型。您可以使用相同的总体方法，但要针对您的任务进行自定义。

TEST_DATA_DIR=$(download_bionemo_data single_cell/testdata-20241203 --source $MY_DATA_SOURCE); \
bionemo-geneformer-recipe \
    --recipe 10m-pretrain \
    --dest my_config.json \
    --data-path ${TEST_DATA_DIR}/cellxgene_2023-12-15_small_processed_scdl \
    --result-dir ./results

⚠️ 重要提示： 检查并编辑输出的 my_config.yaml 的内容，根据您的需要进行调整

注意：要从现有检查点预训练，只需将路径 --initial-ckpt-path 传递给 recipe 命令。这将使用正确的字段填充 YAML，以确保从现有检查点初始化预训练。

要使用传递的配置提交训练作业，首先使用您选择的任何其他执行参数更新 yaml 文件：设备数量、worker 数量、步数等。其次，调用我们的训练入口点。为此，我们需要三件事

• 配置文件，上一步生成的 YAML
• 模型配置类型，在本例中为预训练配置。这将根据预训练所需的参数验证配置 YAML 中的参数。或者，可以在此处指定使用自定义任务头进行微调之类的事情。这允许混合/匹配数据模块和各种任务。
• 数据配置类型，这指定如何解析、验证和准备 DataModule。这可能会根据任务而变化，例如，在微调时，您可能希望使用包含 PERTURB-seq 的自定义 Dataset/DataModule。在这种情况下，默认的预训练 DataConfig 和 DataModule 将不足。有关其他示例用例，请参见 ESM2。

⚠️ 警告： 此设置不配置 Weights and Biases。编辑您的 config YAML 并使用您的 WandB 详细信息填充它。

bionemo-geneformer-train \
--data-config-cls bionemo.geneformer.run.config_models.GeneformerPretrainingDataConfig \
--model-config-cls bionemo.geneformer.run.config_models.ExposedGeneformerPretrainConfig \
--config my_config.yaml

注意：data-config-cls 和 model-config-cls 都具有对应于 GeneformerPretrainingDataConfig 和 ExposedGeneformerPretrainConfig 的默认值

DataConfigCls 和 ModelConfigCls 也可以引用用户本地定义的类型。只要 python 知道如何导入指定的路径，就可以配置它们。例如，您可能有一个自定义 Dataset/DataModule，您想将其与现有 recipe 混合使用。在这种情况下，您可以定义一个 DataConfig 对象，其中泛型指定为您的 DataModule 类型，然后将配置类型传递给训练 recipe。