add codonfm 5b arch params

Signed-off-by: Bruno Alvisio <balvisio@nvidia.com>

Author: balvisio

bionemo-recipes/recipes/codonfm_ptl_te/codonfm_ckpt_te_conversion.py

@@ -26,14 +26,39 @@
 
 import argparse
 import logging
+import os
 
 import torch
+from safetensors.torch import save_file as safetensors_save_file
 
 from src.utils.load_checkpoint import load_checkpoint
 
 
 logger = logging.getLogger(__name__)
 
+ALLOWED_HYPERPARAMETER_KEYS = (
+    "vocab_size",
+    "hidden_size",
+    "num_hidden_layers",
+    "num_attention_heads",
+    "intermediate_size",
+    "hidden_act",
+    "hidden_dropout_prob",
+    "attention_probs_dropout_prob",
+    "initializer_range",
+    "layer_norm_eps",
+    "pad_token_id",
+    "position_embedding_type",
+    "classifier_dropout",
+    "rotary_theta",
+    "ignore_index",
+    "loss_type",
+    "lora",
+    "lora_alpha",
+    "lora_r",
+    "lora_dropout",
+)
+
 # PYTorch -> TE keymap
 PYTORCH_TO_TE_KEYMAP = {
     "model.layers.*.pre_attn_layer_norm.weight": "model.layers.*.self_attention.layernorm_qkv.layer_norm_weight",
@@ -300,6 +325,11 @@ def convert_state_dict(src: dict, keymap: dict):
     return dst_state_dict
 
 
+def filter_hyper_parameters(hyper_parameters: dict) -> dict:
+    """Keep only conversion-compatible hyperparameter keys."""
+    return {key: value for key, value in hyper_parameters.items() if key in ALLOWED_HYPERPARAMETER_KEYS}
+
+
 def main():
     """Main function."""
     logging.basicConfig(level=logging.INFO)
@@ -325,6 +355,7 @@ def main():
     # Load source checkpoint (automatically detects format)
     logger.info(f"Loading checkpoint from {args.src}")
     src_checkpoint = load_checkpoint(args.src, map_location="cpu")
+    src_checkpoint["hyper_parameters"] = filter_hyper_parameters(src_checkpoint["hyper_parameters"])
 
     # Perform conversion based on direction
     if args.direction == "pytorch2te":
@@ -341,11 +372,19 @@ def main():
         dst_state_dict = split_qkv(converted_state_dict, src_checkpoint["hyper_parameters"])
 
     # Prepare final checkpoint
-    dst_checkpoint = {"state_dict": dst_state_dict, "hyper_parameters": src_checkpoint["hyper_parameters"]}
+    dst_checkpoint = {
+        "state_dict": dst_state_dict,
+        "hyper_parameters": src_checkpoint["hyper_parameters"],
+    }
 
     # Save the converted checkpoint in pickled format
     torch.save(dst_checkpoint, args.dst)
-    logger.info(f"Successfully converted checkpoint from {args.src} to {args.dst}")
+    logger.info(f"Successfully converted checkpoint saved to {args.dst}")
+
+    # Save the state_dict in safetensors format alongside the .ckpt file
+    safetensors_path = os.path.splitext(args.dst)[0] + ".safetensors"
+    safetensors_save_file(dst_state_dict, safetensors_path)
+    logger.info(f"Successfully saved safetensors checkpoint to {safetensors_path}")
 
 
 if __name__ == "__main__":

bionemo-recipes/recipes/codonfm_ptl_te/data_scripts/check_codon_frequency.py

@@ -23,7 +23,7 @@
 from tqdm import tqdm
 
 
-sys.path.append("/workspace/codon_fm")
+sys.path.append("/workspace/codonfm")
 from src.tokenizer import Tokenizer
 
 

bionemo-recipes/recipes/codonfm_ptl_te/notebooks/4-EnCodon-Downstream-Task-riboNN.ipynb

@@ -108,6 +108,12 @@
     ")\n",
     "download_checkpoint(\n",
     "    repo_id=\"nvidia/NV-CodonFM-Encodon-TE-1B-v1\", local_dir=\"/data/checkpoints/NV-CodonFM-Encodon-TE-1B-v1\"\n",
+    ")\n",
+    "download_checkpoint(\n",
+    "    repo_id=\"nvidia/NV-CodonFM-Encodon-TE-5B-v1\", local_dir=\"/data/checkpoints/NV-CodonFM-Encodon-TE-5B-v1\"\n",
+    ")\n",
+    "download_checkpoint(\n",
+    "    repo_id=\"nvidia/NV-CodonFM-Encodon-TE-Cdwt-5B-v1\", local_dir=\"/data/checkpoints/NV-CodonFM-Encodon-TE-Cdwt-5B-v1\"\n",
     ")"
    ]
   },
@@ -123,6 +129,8 @@
     "    \"/data/checkpoints/NV-CodonFM-Encodon-TE-80M-v1\",\n",
     "    \"/data/checkpoints/NV-CodonFM-Encodon-TE-600M-v1\",\n",
     "    \"/data/checkpoints/NV-CodonFM-Encodon-TE-Cdwt-1B-v1\",\n",
+    "    \"/data/checkpoints/NV-CodonFM-Encodon-TE-5B-v1\",\n",
+    "    \"/data/checkpoints/NV-CodonFM-Encodon-TE-Cdwt-5B-v1\",\n",
     "]\n",
     "\n",
     "checkpoint_path = checkpoint_paths[0]\n",

bionemo-recipes/recipes/codonfm_ptl_te/src/config.py

@@ -251,6 +251,12 @@ def get_logger_config(args: Any) -> fdl.Config:
         "num_attention_heads": 16,
         "num_hidden_layers": 18,
     },
+    "encodon_5b": {
+        "hidden_size": 4096,
+        "intermediate_size": 16384,
+        "num_attention_heads": 32,
+        "num_hidden_layers": 24,
+    },
     "encodon_10b": {
         "hidden_size": 5120,
         "intermediate_size": 20480,

bionemo-recipes/recipes/codonfm_ptl_te/src/runner.py

@@ -113,6 +113,7 @@ def get_parser():  # noqa: D103
             "encodon_80m",
             "encodon_600m",
             "encodon_1b",
+            "encodon_5b",
             "encodon_10b",
         ],
     )