Check in nemo change and updated packages

sub-packages/bionemo-geneformer/src/bionemo/geneformer/api.py

@@ -50,8 +50,11 @@ class GeneformerNeMo1LightningModuleConnector(GenericBioBertNeMo1LightningModule
     def tokenizer(self):
         nemo1_settings = self.get_nemo1_config()
         fmt_vocab, vocab_tar_path = nemo1_settings["tokenizer"]["vocab_file"].split(":")
-        fmt_medians, medians_tar_path = nemo1_settings["data"]["medians_file"].split(":")
-        assert fmt_vocab == fmt_medians and fmt_vocab == "nemo"
+        assert fmt_vocab == "nemo"
+        # TODO add another function to pull out the medians file from a nemo1 checkpoint, if the user wants it.
+        #  It's not needed for checkpoint conversion though.
+        # fmt_medians, medians_tar_path = nemo1_settings["data"]["medians_file"].split(":")
+        # assert fmt_vocab == fmt_medians and fmt_vocab == "nemo"
         nemo1_path = str(self)
         with tarfile.open(nemo1_path, "r") as old_ckpt:
             vocab_gene_ens_dict = json.loads(old_ckpt.extractfile(f"./{vocab_tar_path}").readlines()[0])

sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py

@@ -116,7 +116,7 @@
 ]
 
 MODEL_PRECISION: str = "bf16-mixed"
-USE_TE: bool = False  # TODO use this for high level decisions around whether we're ready to switch to TE
+USE_TE: bool = True  # TODO use this for high level decisions around whether we're ready to switch to TE
 
 
 @pytest.fixture()
@@ -140,7 +140,7 @@ def geneformer_config():
         hidden_dropout=0.02,
         init_method_std=0.02,
         kv_channels=None,
-        apply_query_key_layer_scaling=True,
+        apply_query_key_layer_scaling=False,
         make_vocab_size_divisible_by=128,
         masked_softmax_fusion=True,  # TODO(@jstjohn) check this
         fp16_lm_cross_entropy=False,
@@ -453,16 +453,23 @@ def test_nemo1_checkpoint_conversion(
     tmpdir: Path, geneformer_config: GeneformerConfig, cells: List[List[str]], seed: int = 42
 ):
     with megatron_parallel_state_utils.distributed_model_parallel_state(32):
-        converter = GeneformerNeMo1LightningModuleConnector(nemo1_checkpoint_path)
+        converter = GeneformerNeMo1LightningModuleConnector(nemo1_release_checkpoint_path)
         assert isinstance(converter.tokenizer, GeneTokenizer)
         diffs = compare_dataclasses(converter.config, geneformer_config)
         skip_fields = {"nemo1_ckpt_path", "return_only_hidden_states", "init_method", "output_layer_init_method"}
         filt = [d for d in diffs if d["field"] not in skip_fields]
         assert filt == []
         out_config = tmpdir / "out_config"
         converter.apply(out_config)  # currently crashes in here during self.nemo_save(out_path, trainer)
-        assert io.is_distributed_ckpt(out_config / "checkpoint")
-        assert False  # TODO test weights are the right ones.
+        assert io.is_distributed_ckpt(out_config / "weights")
+        geneformer_config_logit = deepcopy(geneformer_config)
+        # Set up the model to return logits and switch to the released 10M checkpoint
+        geneformer_config_logit.set_hparam("return_only_hidden_states", False)  # return logits
+        geneformer_config_logit.set_hparam("initial_ckpt_path", str(out_config))  # release checkpoint is important
+
+        mean_loss = _get_loss_from_model(geneformer_config_logit, seed)
+        target: float = 2.368649959564209
+        assert mean_loss < target or mean_loss == pytest.approx(target, abs=1e-2, rel=None)
 
 
 @pytest.mark.skipif(USE_TE, reason="This per-layer test does not yet support TE mapping.")

sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/connector.py

@@ -66,6 +66,7 @@ def apply(self, output_path: Path) -> Path:
             old_weights = torch.load(ckpt_file)
         target = self.init()
         trainer = self.nemo_setup(target)
+        target.trainer = trainer
         self.convert_state(old_weights, target)
         self.nemo_save(output_path, trainer)
 
@@ -87,6 +88,7 @@ def is_te_mapping(model: BioBertLightningModule) -> bool:
     def convert_state(self, source: Dict[str, torch.Tensor], target: BioBertLightningModule) -> BioBertLightningModule:
         """Convert the input state_dict keys from nemo1 biobert to nemo2 biobert."""
         te_mapping = self.is_te_mapping(target)  # check for TE layers.
+        target.module.cpu()
         new_state_dict_from_old = {}
         for k, v in source.items():
             new_key = nemo1_to_nemo2_biobert_key_mapping(k, new_model_prefix="", te_mapping=te_mapping)
@@ -124,6 +126,7 @@ def config(self) -> BioBertGenericConfig[MegatronBioBertModelT]:
             "fp32_residual_connection": False,
             "bias_activation_fusion": True,
             "bias_dropout_fusion": True,
+            "apply_query_key_layer_scaling": False,
             "share_embeddings_and_output_weights": True,
             "fp16": autocast_dtype == torch.float16,
             "bf16": autocast_dtype == torch.bfloat16,

sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/model.py

@@ -406,7 +406,7 @@ class BioBertGenericConfig(
     num_attention_heads: int = 8
     num_layers: int = 6
     init_method_std: float = 0.02
-    biobert_spec_option: BiobertSpecOption = BiobertSpecOption.bert_layer_local_spec
+    biobert_spec_option: BiobertSpecOption = BiobertSpecOption.bert_layer_with_transformer_engine_spec
 
     # TODO: Move this to better places?
     get_attention_mask_from_fusion: bool = False