fix

Signed-off-by: Alexandros Koumparoulis <[email protected]>

Author: akoumpa

nemo_automodel/_transformers/auto_model.py

@@ -48,12 +48,16 @@
 
 
 @contextmanager
-def local_torch_dtype(dtype: torch.dtype, model_class_name: str | None = None):
+def local_torch_dtype(
+    dtype: torch.dtype, model_class_name: str | None = None, default_dtype: torch.dtype = torch.bfloat16
+):
     """
     Locally change the torch default dtype to `dtype`, and restore the old one upon exiting the context.
     If `model_class_name` is provided, it's used to provide a more helpful error message if `dtype` is not valid.
     """
     # Just a more helping error before we set `torch.set_default_dtype` later on which would crash in this case
+    if isinstance(dtype, str):
+        dtype = default_dtype
     if not dtype.is_floating_point:
         if model_class_name is not None:
             error_message = (
@@ -62,7 +66,6 @@ def local_torch_dtype(dtype: torch.dtype, model_class_name: str | None = None):
         else:
             error_message = f"Cannot set `{dtype}` as torch's default as it's not a floating-point dtype"
         raise ValueError(error_message)
-
     original_dtype = torch.get_default_dtype()
     try:
         torch.set_default_dtype(dtype)
@@ -426,9 +429,10 @@ def _retry(**override):
             _download_model_weights(hf_config, pretrained_model_name_or_path)
             logger.info(f"Using custom model implementation for {architectures[0]}")
             kwargs.pop("trust_remote_code", None)
-            # TODO: restore weights after initialization.
-            with local_torch_dtype(torch_dtype, ModelRegistry.model_arch_name_to_cls[architectures[0]].__name__):
-                return ModelRegistry.model_arch_name_to_cls[architectures[0]](hf_config)
+            # TODO(@akoumpa): restore weights after initialization.
+            model_cls = ModelRegistry.model_arch_name_to_cls[architectures[0]]
+            with local_torch_dtype(torch_dtype, model_cls.__name__):
+                return model_cls(hf_config)
 
         # 3. fallback to parent class
         model = None

tests/unit_tests/_transformers/test_auto_model.py

@@ -103,6 +103,7 @@ def test_from_pretrained_uses_registry_when_available(self):
             # Prepare a fake custom model class and return value
             custom_model_instance = Mock()
             custom_cls = Mock(return_value=custom_model_instance)
+            custom_cls.__name__ = "MockMockMock"
             mock_registry.model_arch_name_to_cls = {"CustomArch": custom_cls}
 
             returned = NeMoAutoModelForCausalLM.from_pretrained("dummy/path")
@@ -130,6 +131,7 @@ def test_from_config_uses_registry_when_available(self):
             # Registry provides a custom class
             custom_model_instance = Mock()
             custom_cls = Mock(return_value=custom_model_instance)
+            custom_cls.__name__ = "MockMockMock"
             mock_registry.model_arch_name_to_cls = {"CustomArch": custom_cls}
 
             returned = NeMoAutoModelForCausalLM.from_config(cfg)
@@ -160,6 +162,7 @@ def test_from_pretrained_registry_downloads_checkpoint_files_rank0(self):
             # Prepare a fake custom model class and return value
             custom_model_instance = Mock()
             custom_cls = Mock(return_value=custom_model_instance)
+            custom_cls.__name__ = "MockMockMock"
             mock_registry.model_arch_name_to_cls = {"CustomArch": custom_cls}
 
             returned = NeMoAutoModelForCausalLM.from_pretrained("dummy/repo-id")
@@ -194,6 +197,7 @@ def test_from_pretrained_registry_downloads_when_dist_uninitialized(self):
             # Prepare a fake custom model class and return value
             custom_model_instance = Mock()
             custom_cls = Mock(return_value=custom_model_instance)
+            custom_cls.__name__ = "MockMockMock"
             mock_registry.model_arch_name_to_cls = {"CustomArch": custom_cls}
 
             returned = NeMoAutoModelForCausalLM.from_pretrained("dummy/repo-id")
@@ -240,7 +244,8 @@ def test_from_config_with_string_calls_autoconfig(self):
             # Verify AutoConfig.from_pretrained was called with the string
             mock_autoconfig.assert_called_once_with(
                 "hf-internal-testing/tiny-random-gpt2",
-                trust_remote_code=False
+                trust_remote_code=False,
+                attn_implementation="flash_attention_2",
             )
             # Verify the model was returned
             assert model is mock_model
@@ -539,6 +544,7 @@ def test_packed_sequence_and_cp_overrides_from_pretrained(
             else:
                 custom_model_instance = Mock()
                 custom_cls = Mock(return_value=custom_model_instance)
+                custom_cls.__name__ = "MockMockMock"
                 mock_registry.model_arch_name_to_cls = {"CustomArch": custom_cls}
 
             mock_hf_loader.return_value = MagicMock(config={})