Add encodon 10b configuration

Author: balvisio

bionemo-recipes/models/codonfm/modeling_codonfm_te.py

@@ -156,6 +156,12 @@ def __init__(
         "num_attention_heads": 32,
         "num_hidden_layers": 24,
     },
+    "encodon_10b": {
+        "hidden_size": 5120,
+        "intermediate_size": 20480,
+        "num_attention_heads": 40,
+        "num_hidden_layers": 34,
+    },
 }
 
 

bionemo-recipes/recipes/codonfm_native_te/hydra_config/encodon_10b.yaml

@@ -0,0 +1,33 @@
+defaults:
+  - defaults
+  - _self_
+
+# Training config
+model_preset: encodon_10b
+num_train_steps: 500_000
+
+precision: bf16-mixed
+
+use_sequence_packing: true
+dataset:
+  data_path: ???
+  micro_batch_size: 4
+  num_workers: 1
+  max_seq_length: 2048
+
+# WandB config
+wandb_init_args:
+  name: "codonfm_native_te_10b"
+
+# Learning rate scheduler config
+lr_scheduler_kwargs:
+  num_warmup_steps: 2_000
+  num_training_steps: 500_000
+
+checkpoint:
+  ckpt_dir: ???
+  resume_from_checkpoint: true
+  save_every_n_steps: 1_000
+
+logger:
+  frequency: 100

bionemo-recipes/recipes/codonfm_native_te/modeling_codonfm_te.py

@@ -162,6 +162,12 @@ def __init__(
         "num_attention_heads": 32,
         "num_hidden_layers": 24,
     },
+    "encodon_10b": {
+        "hidden_size": 5120,
+        "intermediate_size": 20480,
+        "num_attention_heads": 40,
+        "num_hidden_layers": 34,
+    },
 }