Recipes

default_adam_optimizer_with_cosine_annealing_recipe()

Default optimizer scheduler config for Geneformer. See OptimizerSchedulerConfig for defaults.

Source code in bionemo/geneformer/run/recipes.py

def default_adam_optimizer_with_cosine_annealing_recipe() -> OptimizerSchedulerConfig:
    """Default optimizer scheduler config for Geneformer. See OptimizerSchedulerConfig for defaults."""
    return OptimizerSchedulerConfig()

default_trainer_config_recipe()

Default trainer config for Geneformer.

Source code in bionemo/geneformer/run/recipes.py

def default_trainer_config_recipe() -> TrainingConfig:
    """Default trainer config for Geneformer."""
    return TrainingConfig(max_steps=55000, limit_val_batches=2, val_check_interval=100)

experiment_config_recipe()

Default experiment config for Geneformer. Used in testing.

Source code in bionemo/geneformer/run/recipes.py

def experiment_config_recipe() -> ExperimentConfig:
    """Default experiment config for Geneformer. Used in testing."""
    return ExperimentConfig(
        save_every_n_steps=100,
        result_dir="./results",
        experiment_name="default_experiment",
        restore_from_checkpoint_path=None,
        save_last_checkpoint=True,
        metric_to_monitor_for_checkpoints="reduced_train_loss",
        save_top_k=2,
        create_tensorboard_logger=False,
    )

finetune_test_recipe(args)

Recipe for finetuning a regression head on the masked tokens.

Source code in bionemo/geneformer/run/recipes.py

def finetune_test_recipe(args) -> MainConfig[ExposedFineTuneSeqLenBioBertConfig, GeneformerPretrainingDataConfig]:
    """Recipe for finetuning a regression head on the masked tokens."""
    data_path = args.data_path
    result_dir = args.result_dir

    parallel_config = ParallelConfig(
        tensor_model_parallel_size=1, pipeline_model_parallel_size=1, num_devices=1, accumulate_grad_batches=2
    )
    training_config = TrainingConfig(
        max_steps=10, limit_val_batches=2, val_check_interval=2, precision="bf16-mixed", accelerator="gpu"
    )
    data_config = GeneformerPretrainingDataConfig(
        seq_length=128,
        micro_batch_size=2,
        num_dataset_workers=0,
        data_dir=data_path,
    )
    experiment_config = ExperimentConfig(
        save_every_n_steps=training_config.val_check_interval,
        result_dir=result_dir,
        experiment_name="test-experiment",
        restore_from_checkpoint_path=None,
        save_last_checkpoint=True,
        metric_to_monitor_for_checkpoints="reduced_train_loss",
        save_top_k=2,
        create_tensorboard_logger=False,
    )

    optim_config = OptimizerSchedulerConfig(lr_scheduler="cosine")
    geneformer_config = geneformer_10m_finetune_config(
        seq_length=data_config.seq_length, initial_ckpt_path=args.initial_ckpt_path
    )

    return MainConfig(
        data_config=data_config,
        parallel_config=parallel_config,
        training_config=training_config,
        bionemo_model_config=geneformer_config,
        optim_config=optim_config,
        experiment_config=experiment_config,
    )

geneformer_106m_experiment_config(result_dir)

Experiment config for Geneformer 106m.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_106m_experiment_config(result_dir) -> ExperimentConfig:
    """Experiment config for Geneformer 106m."""
    return ExperimentConfig(
        save_every_n_steps=100,
        result_dir=result_dir,
        experiment_name="geneformer-106m",
        restore_from_checkpoint_path=None,
    )

geneformer_106m_model_config(seq_length=2048, precision='bf16-mixed', nemo1_init_path=None, initial_ckpt_path=None, biobert_spec_option=BiobertSpecOption.bert_layer_with_transformer_engine_spec)

Geneformer 106m model config settings.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_106m_model_config(
    seq_length: int = 2048,
    precision: PrecisionTypes = "bf16-mixed",
    nemo1_init_path: Optional[str] = None,
    initial_ckpt_path: Optional[str] = None,
    biobert_spec_option: BiobertSpecOption = BiobertSpecOption.bert_layer_with_transformer_engine_spec,
) -> ExposedGeneformerPretrainConfig:
    """Geneformer 106m model config settings."""
    geneformer_config = ExposedGeneformerPretrainConfig(
        num_layers=12,
        hidden_size=768,
        ffn_hidden_size=3072,
        num_attention_heads=12,
        seq_length=seq_length,
        fp32_residual_connection=False,
        hidden_dropout=0.02,
        init_method_std=0.02,
        kv_channels=None,
        apply_query_key_layer_scaling=False,
        make_vocab_size_divisible_by=128,
        masked_softmax_fusion=True,
        fp16_lm_cross_entropy=False,
        params_dtype=precision,
        pipeline_dtype=precision,
        autocast_dtype=precision,
        gradient_accumulation_fusion=False,
        layernorm_zero_centered_gamma=False,
        layernorm_epsilon=1.0e-12,
        activation_func="gelu",
        qk_layernorm=False,
        apply_residual_connection_post_layernorm=False,
        bias_activation_fusion=True,
        bias_dropout_fusion=True,
        get_attention_mask_from_fusion=True,
        attention_dropout=0.1,
        share_embeddings_and_output_weights=True,
        enable_autocast=False,
        biobert_spec_option=biobert_spec_option,
        nemo1_ckpt_path=nemo1_init_path,
        initial_ckpt_path=initial_ckpt_path,
    )
    return geneformer_config

geneformer_106m_parallel_config()

Base parallel config for Geneformer.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_106m_parallel_config() -> ParallelConfig:
    """Base parallel config for Geneformer."""
    return ParallelConfig(
        tensor_model_parallel_size=1,
        pipeline_model_parallel_size=1,
        accumulate_grad_batches=1,
        ddp="megatron",
        num_devices=8,
        num_nodes=1,
    )

geneformer_106m_pretrain_recipe(args)

Recipe for pretraining the 106m model. Uses 8 GPUs for data parallelism.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_106m_pretrain_recipe(
    args,
) -> MainConfig[ExposedGeneformerPretrainConfig, GeneformerPretrainingDataConfig]:
    """Recipe for pretraining the 106m model. Uses 8 GPUs for data parallelism."""
    data_config: GeneformerPretrainingDataConfig = geneformer_data_recipe(data_dir=args.data_path)
    parallel_config = geneformer_106m_parallel_config()
    training_config = geneformer_base_training_config()
    bionemo_model_config = geneformer_106m_model_config(initial_ckpt_path=args.initial_ckpt_path)
    optim_config = geneformer_base_optimizer_scheduler_config()
    experiment_config = geneformer_106m_experiment_config(result_dir=args.result_dir)
    wandb_config = geneformer_106m_wandb_config()
    main_config = MainConfig[ExposedGeneformerPretrainConfig, GeneformerPretrainingDataConfig](
        data_config=data_config,
        parallel_config=parallel_config,
        training_config=training_config,
        bionemo_model_config=bionemo_model_config,
        optim_config=optim_config,
        experiment_config=experiment_config,
        wandb_config=wandb_config,
    )
    return main_config

geneformer_106m_wandb_config()

Wandb config for Geneformer 106m.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_106m_wandb_config() -> WandbConfig:
    """Wandb config for Geneformer 106m."""
    wandb_config = WandbConfig(
        entity="geneformer-106m_pretraining",
        project="geneformer-106m_pretraining",
        group="geneformer-106m",
        tags=["geneformer-106m"],
        offline=True,
        anonymous=True,
        id="1",
        log_model=False,
    )
    return wandb_config

geneformer_10m_experiment_config(result_dir)

Experiment config for Geneformer 10m.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_10m_experiment_config(result_dir) -> ExperimentConfig:
    """Experiment config for Geneformer 10m."""
    return ExperimentConfig(
        save_every_n_steps=100,
        result_dir=result_dir,
        experiment_name="geneformer-10m",
        restore_from_checkpoint_path=None,
    )

geneformer_10m_finetune_config(seq_length=2048, precision='bf16-mixed', nemo1_init_path=None, initial_ckpt_path=None, biobert_spec_option=BiobertSpecOption.bert_layer_with_transformer_engine_spec)

Geneformer 10m finetuning config settings.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_10m_finetune_config(
    seq_length: int = 2048,
    precision: PrecisionTypes = "bf16-mixed",
    nemo1_init_path: Optional[str] = None,
    initial_ckpt_path: Optional[str] = None,
    biobert_spec_option=BiobertSpecOption.bert_layer_with_transformer_engine_spec,
) -> ExposedFineTuneSeqLenBioBertConfig:
    """Geneformer 10m finetuning config settings."""
    geneformer_config = ExposedFineTuneSeqLenBioBertConfig(
        num_layers=6,
        hidden_size=256,
        ffn_hidden_size=512,
        num_attention_heads=4,
        seq_length=seq_length,
        fp32_residual_connection=False,
        hidden_dropout=0.02,
        init_method_std=0.02,
        kv_channels=None,
        apply_query_key_layer_scaling=False,
        make_vocab_size_divisible_by=128,
        masked_softmax_fusion=True,
        fp16_lm_cross_entropy=False,
        params_dtype=precision,
        pipeline_dtype=precision,
        autocast_dtype=precision,
        gradient_accumulation_fusion=False,
        layernorm_zero_centered_gamma=False,
        layernorm_epsilon=1.0e-12,
        activation_func="gelu",
        qk_layernorm=False,
        apply_residual_connection_post_layernorm=False,
        bias_activation_fusion=True,
        bias_dropout_fusion=True,
        get_attention_mask_from_fusion=True,
        attention_dropout=0.1,
        share_embeddings_and_output_weights=True,
        enable_autocast=False,
        biobert_spec_option=biobert_spec_option,
        nemo1_ckpt_path=nemo1_init_path,
        initial_ckpt_path=initial_ckpt_path,
    )
    return geneformer_config

geneformer_10m_finetune_recipe(args)

Recipe for finetuning the 10m model on a token regression head. Used as an example and for testing.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_10m_finetune_recipe(
    args,
) -> MainConfig[ExposedFineTuneSeqLenBioBertConfig, GeneformerPretrainingDataConfig]:
    """Recipe for finetuning the 10m model on a token regression head. Used as an example and for testing."""
    data_config: GeneformerPretrainingDataConfig = geneformer_data_recipe(data_dir=args.data_path)
    parallel_config = simple_parallel_recipe()
    training_config = default_trainer_config_recipe()
    bionemo_model_config = geneformer_finetuning_regression_head_recipe(initial_ckpt_path=args.initial_ckpt_path)
    optim_config = default_adam_optimizer_with_cosine_annealing_recipe()
    experiment_config = experiment_config_recipe()
    wandb_config = WandbConfig(
        project="bionemo2-demo",
        entity="nvidia",
        offline=True,
        tags=[],
        group="dev",
        id="dev",
        log_model=False,
        anonymous=True,
    )
    main_config = MainConfig[ExposedFineTuneSeqLenBioBertConfig, GeneformerPretrainingDataConfig](
        data_config=data_config,
        parallel_config=parallel_config,
        training_config=training_config,
        bionemo_model_config=bionemo_model_config,
        optim_config=optim_config,
        experiment_config=experiment_config,
        wandb_config=wandb_config,
    )
    return main_config

geneformer_10m_model_config(seq_length=2048, precision='bf16-mixed', nemo1_init_path=None, initial_ckpt_path=None, biobert_spec_option=BiobertSpecOption.bert_layer_with_transformer_engine_spec)

Geneformer 10m model config settings.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_10m_model_config(
    seq_length: int = 2048,
    precision: PrecisionTypes = "bf16-mixed",
    nemo1_init_path: Optional[str] = None,
    initial_ckpt_path: Optional[str] = None,
    biobert_spec_option: BiobertSpecOption = BiobertSpecOption.bert_layer_with_transformer_engine_spec,
) -> ExposedGeneformerPretrainConfig:
    """Geneformer 10m model config settings."""
    geneformer_config = ExposedGeneformerPretrainConfig(
        num_layers=6,
        hidden_size=256,
        ffn_hidden_size=512,
        num_attention_heads=4,
        seq_length=seq_length,
        fp32_residual_connection=False,
        hidden_dropout=0.02,
        init_method_std=0.02,
        kv_channels=None,
        apply_query_key_layer_scaling=False,
        make_vocab_size_divisible_by=128,
        masked_softmax_fusion=True,
        fp16_lm_cross_entropy=False,
        params_dtype=precision,
        pipeline_dtype=precision,
        autocast_dtype=precision,
        gradient_accumulation_fusion=False,
        layernorm_zero_centered_gamma=False,
        layernorm_epsilon=1.0e-12,
        activation_func="gelu",
        qk_layernorm=False,
        apply_residual_connection_post_layernorm=False,
        bias_activation_fusion=True,
        bias_dropout_fusion=True,
        get_attention_mask_from_fusion=True,
        attention_dropout=0.1,
        share_embeddings_and_output_weights=True,
        enable_autocast=False,
        biobert_spec_option=biobert_spec_option,
        nemo1_ckpt_path=nemo1_init_path,
        initial_ckpt_path=initial_ckpt_path,
    )
    return geneformer_config

geneformer_10m_pretrain_recipe(args)

Recipe for pretraining the 10m model.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_10m_pretrain_recipe(
    args,
) -> MainConfig[ExposedGeneformerPretrainConfig, GeneformerPretrainingDataConfig]:
    """Recipe for pretraining the 10m model."""
    data_config: GeneformerPretrainingDataConfig = geneformer_data_recipe(data_dir=args.data_path)
    parallel_config = simple_parallel_recipe()
    training_config = geneformer_base_training_config()
    bionemo_model_config = geneformer_10m_model_config(initial_ckpt_path=args.initial_ckpt_path)
    optim_config = geneformer_base_optimizer_scheduler_config()
    experiment_config = geneformer_10m_experiment_config(result_dir=args.result_dir)
    wandb_config = geneformer_10m_wandb_config()
    main_config = MainConfig[ExposedGeneformerPretrainConfig, GeneformerPretrainingDataConfig](
        data_config=data_config,
        parallel_config=parallel_config,
        training_config=training_config,
        bionemo_model_config=bionemo_model_config,
        optim_config=optim_config,
        experiment_config=experiment_config,
        wandb_config=wandb_config,
    )
    return main_config

geneformer_10m_wandb_config()

Wandb config for Geneformer 10m.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_10m_wandb_config() -> WandbConfig:
    """Wandb config for Geneformer 10m."""
    wandb_config = WandbConfig(
        entity="geneformer-10m_pretraining",
        project="geneformer-10m_pretraining",
        group="geneformer-10m",
        tags=["geneformer-10m"],
        offline=True,
        anonymous=True,
        id="1",
        log_model=False,
    )
    return wandb_config

geneformer_base_optimizer_scheduler_config()

Base optimizer scheduler config for Geneformer.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_base_optimizer_scheduler_config() -> OptimizerSchedulerConfig:
    """Base optimizer scheduler config for Geneformer."""
    return OptimizerSchedulerConfig(lr=1e-3, lr_scheduler="cosine")  # Matches bionemo1

geneformer_base_parallel_config()

Base parallel config for Geneformer.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_base_parallel_config() -> ParallelConfig:
    """Base parallel config for Geneformer."""
    return ParallelConfig(
        tensor_model_parallel_size=1,
        pipeline_model_parallel_size=1,
        accumulate_grad_batches=1,
        ddp="megatron",
        num_devices=1,
        num_nodes=1,
    )

geneformer_base_training_config()

Base training config for Geneformer.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_base_training_config() -> TrainingConfig:
    """Base training config for Geneformer."""
    return TrainingConfig(
        max_steps=400000, limit_val_batches=8, val_check_interval=100, precision="bf16-mixed"
    )  # matches bionemo1

geneformer_data_recipe(data_dir)

Recipe that produces the base geneformer small data configuration.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_data_recipe(data_dir) -> GeneformerPretrainingDataConfig:
    """Recipe that produces the base geneformer small data configuration."""
    return GeneformerPretrainingDataConfig(data_dir=data_dir)

geneformer_finetuning_regression_head_recipe(precision='bf16-mixed', nemo1_init_path=None, initial_ckpt_path=None, initial_ckpt_skip_keys_with_these_prefixes=None)

Recipe for finetuning a regression head on the masked tokens.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_finetuning_regression_head_recipe(
    precision: PrecisionTypes = "bf16-mixed",
    nemo1_init_path: Optional[str] = None,
    initial_ckpt_path: Optional[str] = None,
    initial_ckpt_skip_keys_with_these_prefixes: Optional[List[str]] = None,
) -> ExposedFineTuneSeqLenBioBertConfig:
    """Recipe for finetuning a regression head on the masked tokens."""
    partial_finetuning_config = partial(
        ExposedFineTuneSeqLenBioBertConfig,
        params_dtype=precision,
        pipeline_dtype=precision,
        autocast_dtype=precision,
        nemo1_ckpt_path=nemo1_init_path,
        initial_ckpt_path=initial_ckpt_path,
        biobert_spec_option=BiobertSpecOption.bert_layer_with_transformer_engine_spec,
    )
    if initial_ckpt_skip_keys_with_these_prefixes:
        finetuning_config = partial_finetuning_config(
            initial_ckpt_skip_keys_with_these_prefixes=initial_ckpt_skip_keys_with_these_prefixes
        )
    else:
        # Use the sensible default when None is passed
        finetuning_config = partial_finetuning_config()
    return finetuning_config

geneformer_tiny_config(seq_length=2048, precision='bf16-mixed', nemo1_init_path=None, initial_ckpt_path=None, biobert_spec_option=BiobertSpecOption.bert_layer_with_transformer_engine_spec)

Geneformer tiny model config settings, used in testing.

Source code in bionemo/geneformer/run/recipes.py

def geneformer_tiny_config(
    seq_length: int = 2048,
    precision: PrecisionTypes = "bf16-mixed",
    nemo1_init_path: Optional[str] = None,
    initial_ckpt_path: Optional[str] = None,
    biobert_spec_option: BiobertSpecOption = BiobertSpecOption.bert_layer_with_transformer_engine_spec,
) -> ExposedGeneformerPretrainConfig:
    """Geneformer tiny model config settings, used in testing."""
    geneformer_config = ExposedGeneformerPretrainConfig(
        num_layers=2,
        hidden_size=32,
        ffn_hidden_size=4 * 32,
        num_attention_heads=2,
        seq_length=seq_length,
        fp32_residual_connection=False,
        hidden_dropout=0.02,
        init_method_std=0.02,
        kv_channels=None,
        apply_query_key_layer_scaling=False,
        make_vocab_size_divisible_by=128,
        masked_softmax_fusion=True,
        fp16_lm_cross_entropy=False,
        params_dtype=precision,
        pipeline_dtype=precision,
        autocast_dtype=precision,
        gradient_accumulation_fusion=False,
        layernorm_zero_centered_gamma=False,
        layernorm_epsilon=1.0e-12,
        activation_func="gelu",
        qk_layernorm=False,
        apply_residual_connection_post_layernorm=False,
        bias_activation_fusion=True,
        bias_dropout_fusion=True,
        get_attention_mask_from_fusion=True,
        attention_dropout=0.1,
        share_embeddings_and_output_weights=True,
        enable_autocast=False,
        biobert_spec_option=biobert_spec_option,
        nemo1_ckpt_path=nemo1_init_path,
        initial_ckpt_path=initial_ckpt_path,
    )
    return geneformer_config

pretrain_tiny_test_recipe(args)

Recipe for pretraining a tiny model. Used in testing.

Source code in bionemo/geneformer/run/recipes.py

def pretrain_tiny_test_recipe(args) -> MainConfig[ExposedGeneformerPretrainConfig, GeneformerPretrainingDataConfig]:
    """Recipe for pretraining a tiny model. Used in testing."""
    data_path = args.data_path
    result_dir = args.result_dir

    parallel_config = ParallelConfig(
        tensor_model_parallel_size=1, pipeline_model_parallel_size=1, num_devices=1, accumulate_grad_batches=2
    )
    training_config = TrainingConfig(
        max_steps=10, limit_val_batches=2, val_check_interval=2, precision="bf16-mixed", accelerator="gpu"
    )
    data_config = GeneformerPretrainingDataConfig(
        seq_length=128,
        micro_batch_size=2,
        num_dataset_workers=0,
        data_dir=data_path,
    )
    experiment_config = ExperimentConfig(
        save_every_n_steps=training_config.val_check_interval,
        result_dir=result_dir,
        experiment_name="test-experiment",
        restore_from_checkpoint_path=None,
        save_last_checkpoint=True,
        metric_to_monitor_for_checkpoints="reduced_train_loss",
        save_top_k=2,
        create_tensorboard_logger=False,
    )

    optim_config = OptimizerSchedulerConfig(lr_scheduler="cosine")
    geneformer_config = geneformer_tiny_config(
        seq_length=data_config.seq_length, initial_ckpt_path=args.initial_ckpt_path
    )

    return MainConfig(
        data_config=data_config,
        parallel_config=parallel_config,
        training_config=training_config,
        bionemo_model_config=geneformer_config,
        optim_config=optim_config,
        experiment_config=experiment_config,
    )

simple_parallel_recipe(tensor_model_parallel_size=1, pipeline_model_parallel_size=1, num_devices=1, accumulate_grad_batches=1)

Simple parallel config for Geneformer, only used in testing.

Source code in bionemo/geneformer/run/recipes.py

def simple_parallel_recipe(
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
    num_devices: int = 1,
    accumulate_grad_batches: int = 1,
) -> ParallelConfig:
    """Simple parallel config for Geneformer, only used in testing."""
    assert (
        num_devices >= tensor_model_parallel_size * pipeline_model_parallel_size
    ), "devices must be divisible by tensor_model_parallel_size * pipeline_model_parallel_size"
    return ParallelConfig(
        tensor_model_parallel_size=tensor_model_parallel_size,
        pipeline_model_parallel_size=pipeline_model_parallel_size,
        accumulate_grad_batches=accumulate_grad_batches,
        num_devices=num_devices,
    )