fixing convert

Signed-off-by: Peter St. John <[email protected]>

Author: pstjohn

bionemo-recipes/models/amplify/src/amplify/state.py

@@ -210,12 +210,12 @@ def scale_weights(ctx):
         target.to(cast_dtype)
         logger.info(f"Casting model to {cast_dtype} complete.")
     else:
-        assert target_orig_dtypes == extract_dtypes(target.named_parameters()), (
-            f"dtype mismatch between source and target state dicts. "
-            f"Left side is { {k: v for k, v in target_orig_dtypes.items() if v != torch.bfloat16} }, "
-            f"Right side is "
-            f"{ {k: v for k, v in extract_dtypes(target.named_parameters()).items() if v != torch.bfloat16} }"
-        )
+        target_new_dtypes = extract_dtypes(target.named_parameters())
+        for key in target_orig_dtypes.keys():
+            if key in target_new_dtypes:  # For tied weights, these parameters may disappear.
+                assert target_orig_dtypes[key] == target_new_dtypes[key], (
+                    f"dtype mismatch for key {key}: {target_orig_dtypes[key]} vs {target_new_dtypes[key]}"
+                )
 
     return target
 

bionemo-recipes/models/esm2/src/esm/convert.py

@@ -71,11 +71,8 @@ def convert_esm_hf_to_te(model_hf: nn.Module, **config_kwargs) -> nn.Module:
             _pad_decoder_weights,
             _pad_bias,
         ],
-        state_dict_ignored_entries=["lm_head.decoder.weight"],
     )
 
-    output_model.tie_weights()
-
     return output_model
 
 

bionemo-recipes/models/esm2/src/esm/state.py

@@ -210,12 +210,12 @@ def scale_weights(ctx):
         target.to(cast_dtype)
         logger.info(f"Casting model to {cast_dtype} complete.")
     else:
-        assert target_orig_dtypes == extract_dtypes(target.named_parameters()), (
-            f"dtype mismatch between source and target state dicts. "
-            f"Left side is { {k: v for k, v in target_orig_dtypes.items() if v != torch.bfloat16} }, "
-            f"Right side is "
-            f"{ {k: v for k, v in extract_dtypes(target.named_parameters()).items() if v != torch.bfloat16} }"
-        )
+        target_new_dtypes = extract_dtypes(target.named_parameters())
+        for key in target_orig_dtypes.keys():
+            if key in target_new_dtypes:  # For tied weights, these parameters may disappear.
+                assert target_orig_dtypes[key] == target_new_dtypes[key], (
+                    f"dtype mismatch for key {key}: {target_orig_dtypes[key]} vs {target_new_dtypes[key]}"
+                )
 
     return target
 

bionemo-recipes/models/esm2/tests/test_distributed_fp8.py

@@ -215,6 +215,7 @@ def is_main_process(self) -> bool:
                 dict_2 = pickle.loads(state_2.detach().numpy(force=True).tobytes())
                 recipe_1 = dict_1.pop("recipe")
                 recipe_2 = dict_2.pop("recipe")
+                breakpoint()
                 torch.testing.assert_close(dict_1, dict_2)
                 assert recipe_1 == recipe_2
 

bionemo-recipes/models/esm2/tests/test_fp8.py

@@ -20,12 +20,16 @@
 from torch.distributed.checkpoint.state_dict import get_model_state_dict
 from transformer_engine.common import recipe as recipe_module
 from transformer_engine.pytorch import fp8
-from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensor
 
 from esm.collator import MLMDataCollatorWithFlattening
 from esm.modeling_esm_te import NVEsmConfig, NVEsmForMaskedLM
 
 
+try:
+    from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensor
+except ImportError:  # TE nightly uses a new import path for QuantizedTensor
+    from transformer_engine.pytorch.quantized_tensor import QuantizedTensor
+
 ALL_RECIPES = [
     recipe_module.DelayedScaling(),
     recipe_module.Float8CurrentScaling(),

bionemo-recipes/models/llama3/state.py

@@ -210,12 +210,12 @@ def scale_weights(ctx):
         target.to(cast_dtype)
         logger.info(f"Casting model to {cast_dtype} complete.")
     else:
-        assert target_orig_dtypes == extract_dtypes(target.named_parameters()), (
-            f"dtype mismatch between source and target state dicts. "
-            f"Left side is { {k: v for k, v in target_orig_dtypes.items() if v != torch.bfloat16} }, "
-            f"Right side is "
-            f"{ {k: v for k, v in extract_dtypes(target.named_parameters()).items() if v != torch.bfloat16} }"
-        )
+        target_new_dtypes = extract_dtypes(target.named_parameters())
+        for key in target_orig_dtypes.keys():
+            if key in target_new_dtypes:  # For tied weights, these parameters may disappear.
+                assert target_orig_dtypes[key] == target_new_dtypes[key], (
+                    f"dtype mismatch for key {key}: {target_orig_dtypes[key]} vs {target_new_dtypes[key]}"
+                )
 
     return target