unit tests

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

Author: adil-a

nemo_automodel/_diffusers/auto_diffusion_pipeline.py

@@ -14,47 +14,23 @@
 
 import logging
 import os
-import sys
-import types
 from typing import Any, Dict, Iterable, Optional, Tuple
 
 import torch
 import torch.nn as nn
 
-# diffusers is an optional dependency. Some CI environments may have it missing
-# or installed with incompatible transitive deps. Import defensively so that
-# helper functions in this module (and unit tests) can run without diffusers.
-from nemo_automodel.shared.import_utils import safe_import
-
-try:  # pragma: no cover - exercised indirectly via unit tests
-    ok, diffusers = safe_import("diffusers")
-except Exception:
-    # diffusers can fail with non-ImportError exceptions (e.g. missing optional deps).
-    ok, diffusers = False, None
-
-if ok and hasattr(diffusers, "DiffusionPipeline"):
-    DiffusionPipeline = diffusers.DiffusionPipeline
-else:  # pragma: no cover
-    # Provide a minimal stub module/class so tests can patch
-    # `diffusers.DiffusionPipeline.from_pretrained` even when diffusers fails to import.
-    diffusers_stub = sys.modules.get("diffusers")
-    if diffusers_stub is None:
-        diffusers_stub = types.ModuleType("diffusers")
-        sys.modules["diffusers"] = diffusers_stub
-
-    class DiffusionPipeline:  # type: ignore[no-redef]
-        @classmethod
-        def from_pretrained(cls, *args, **kwargs):
-            raise RuntimeError(
-                "diffusers is required for NeMoAutoDiffusionPipeline.from_pretrained. "
-                "Install a compatible diffusers + deps stack to use this feature."
-            )
-
-    setattr(diffusers_stub, "DiffusionPipeline", DiffusionPipeline)
-
 from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager
 from nemo_automodel.shared.utils import dtype_from_str
 
+# diffusers is an optional dependency
+try:
+    from diffusers import DiffusionPipeline
+
+    DIFFUSERS_AVAILABLE = True
+except Exception:
+    DIFFUSERS_AVAILABLE = False
+    DiffusionPipeline = object
+
 logger = logging.getLogger(__name__)
 
 
@@ -123,8 +99,13 @@ def from_pretrained(
         torch_dtype: Any = "auto",
         move_to_device: bool = True,
         **kwargs,
-    ) -> DiffusionPipeline:
-        pipe: DiffusionPipeline = DiffusionPipeline.from_pretrained(
+    ) -> "DiffusionPipeline":
+        if not DIFFUSERS_AVAILABLE:
+            raise RuntimeError(
+                "diffusers is required for NeMoAutoDiffusionPipeline.from_pretrained. "
+                "Install diffusers with a compatible version."
+            )
+        pipe = DiffusionPipeline.from_pretrained(
             pretrained_model_name_or_path,
             *model_args,
             torch_dtype=torch_dtype,

nemo_automodel/components/models/mistral3/model.py

@@ -144,7 +144,7 @@ def __init__(
             bos_token_id=bos_token_id,
             eos_token_id=eos_token_id,
             tie_word_embeddings=tie_word_embeddings,
-            ignore_keys_at_rope_validation=["llama_4_scaling_beta"],
+            ignore_keys_at_rope_validation={"llama_4_scaling_beta"},
             **kwargs,
         )
 
@@ -195,7 +195,10 @@ def __init__(self, config: Ministral3Config, device=None):
 
         self.config = config
 
-        self.rope_type = self.config.rope_parameters["type"]
+        # Support both transformers v4 ("type") and v5 ("rope_type") key names
+        self.rope_type = self.config.rope_parameters.get("rope_type") or self.config.rope_parameters.get(
+            "type", "default"
+        )
         rope_init_fn = self.compute_default_rope_parameters
         if self.rope_type != "default":
             rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

nemo_automodel/components/models/nemotron_parse/model.py

@@ -248,8 +248,6 @@ def forward(
         attention_mask: Optional[torch.Tensor] = None,
         encoder_hidden_states: Optional[torch.FloatTensor] = None,
         encoder_attention_mask: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        cross_attn_head_mask: Optional[torch.Tensor] = None,
         past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         use_cache: Optional[bool] = None,
@@ -282,7 +280,7 @@ def forward(
 
         if self.config._attn_implementation == "flash_attention_2":
             attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
-        elif self.config._attn_implementation == "sdpa" and not output_attentions and cross_attn_head_mask is None:
+        elif self.config._attn_implementation == "sdpa" and not output_attentions:
             attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
                 attention_mask, input_shape, inputs_embeds, past_key_values_length
             )
@@ -294,7 +292,7 @@ def forward(
         if encoder_hidden_states is not None and encoder_attention_mask is not None:
             if self.config._attn_implementation == "flash_attention_2":
                 encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None
-            elif self.config._attn_implementation == "sdpa" and cross_attn_head_mask is None and not output_attentions:
+            elif self.config._attn_implementation == "sdpa" and not output_attentions:
                 encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
                     encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
                 )
@@ -310,12 +308,6 @@ def forward(
         all_self_attns = () if output_attentions else None
         all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
 
-        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
-            if attn_mask is not None and attn_mask.size()[0] != len(self.layers):
-                raise ValueError(
-                    f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for {attn_mask.size()[0]}."
-                )
-
         for idx, decoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
@@ -331,23 +323,17 @@ def forward(
                     attention_mask,
                     encoder_hidden_states,
                     encoder_attention_mask,
-                    head_mask[idx] if head_mask is not None else None,
-                    cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
-                    None,
+                    None,  # past_key_values
                     output_attentions,
-                    False,
+                    False,  # use_cache
                 )
             else:
                 layer_outputs = decoder_layer(
                     hidden_states,
                     attention_mask=attention_mask,
                     encoder_hidden_states=encoder_hidden_states,
                     encoder_attention_mask=encoder_attention_mask,
-                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
-                    cross_attn_layer_head_mask=(
-                        cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
-                    ),
-                    past_key_value=None,
+                    past_key_values=None,
                     output_attentions=output_attentions,
                     use_cache=False,
                 )

nemo_automodel/components/utils/model_utils.py

@@ -324,9 +324,11 @@ def register_empty_parameter(module, name, param):
                 # (e.g., TransformerEngine sets tensor_model_parallel on weights)
                 if param_cls is nn.Parameter:
                     kwargs = {"requires_grad": param.requires_grad}
+                    is_hf_initialized = None
                 else:
                     kwargs = module._parameters[name].__dict__.copy()
                     kwargs["requires_grad"] = param.requires_grad
+                    is_hf_initialized = kwargs.pop("_is_hf_initialized", None)
             module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
             if is_hf_initialized is not None:
                 setattr(module._parameters[name], "_is_hf_initialized", is_hf_initialized)

tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py

@@ -19,6 +19,14 @@
 import pytest
 import torch
 
+# Check if diffusers can be imported properly (may fail due to peft/transformers incompatibility)
+try:
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _choose_device
+    DIFFUSERS_AVAILABLE = True
+except Exception:
+    DIFFUSERS_AVAILABLE = False
+
+pytestmark = pytest.mark.skipif(not DIFFUSERS_AVAILABLE, reason="diffusers not available or incompatible with current transformers version")
 
 MODULE_PATH = "nemo_automodel._diffusers.auto_diffusion_pipeline"
 
@@ -128,16 +136,20 @@ def test_from_pretrained_basic_flow_moves_modules_and_returns_pipeline(caplog):
     m1, m2 = DummyModule(), DummyModule()
     dummy_pipe = DummyPipeline({"unet": m1, "text_encoder": m2})
 
+    mock_diffusion_pipeline = MagicMock()
+    mock_diffusion_pipeline.from_pretrained.return_value = dummy_pipe
+
     with (
-        patch("diffusers.DiffusionPipeline.from_pretrained", return_value=dummy_pipe) as mock_hf_from,
+        patch(f"{MODULE_PATH}.DIFFUSERS_AVAILABLE", True),
+        patch(f"{MODULE_PATH}.DiffusionPipeline", mock_diffusion_pipeline),
         patch.object(torch.nn.Module, "to") as mock_to,
         patch(f"{MODULE_PATH}.torch.cuda.is_available", return_value=False),
     ):
         caplog.set_level(logging.WARNING)
         out = NeMoAutoDiffusionPipeline.from_pretrained("dummy")
 
     assert out is dummy_pipe
-    assert mock_hf_from.call_count == 1
+    assert mock_diffusion_pipeline.from_pretrained.call_count == 1
     # Both modules should be moved to device once
     assert mock_to.call_count == 2
 
@@ -146,8 +158,12 @@ def test_from_pretrained_skips_move_when_flag_false():
     from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
 
     dummy_pipe = DummyPipeline({"unet": DummyModule()})
+    mock_diffusion_pipeline = MagicMock()
+    mock_diffusion_pipeline.from_pretrained.return_value = dummy_pipe
+
     with (
-        patch("diffusers.DiffusionPipeline.from_pretrained", return_value=dummy_pipe),
+        patch(f"{MODULE_PATH}.DIFFUSERS_AVAILABLE", True),
+        patch(f"{MODULE_PATH}.DiffusionPipeline", mock_diffusion_pipeline),
         patch.object(torch.nn.Module, "to") as mock_to,
     ):
         out = NeMoAutoDiffusionPipeline.from_pretrained("dummy", move_to_device=False)
@@ -172,8 +188,12 @@ def test_from_pretrained_parallel_scheme_applies_managers_and_sets_attrs():
 
     parallel_scheme = {"unet": mgr_unet, "text_encoder": mgr_text}
 
+    mock_diffusion_pipeline = MagicMock()
+    mock_diffusion_pipeline.from_pretrained.return_value = dummy_pipe
+
     with (
-        patch("diffusers.DiffusionPipeline.from_pretrained", return_value=dummy_pipe),
+        patch(f"{MODULE_PATH}.DIFFUSERS_AVAILABLE", True),
+        patch(f"{MODULE_PATH}.DiffusionPipeline", mock_diffusion_pipeline),
         patch(f"{MODULE_PATH}.torch.distributed.is_initialized", return_value=True),
     ):
         out = NeMoAutoDiffusionPipeline.from_pretrained("dummy", parallel_scheme=parallel_scheme, move_to_device=False)
@@ -196,8 +216,12 @@ def test_from_pretrained_parallel_scheme_logs_and_continues_on_errors(caplog):
     mgr = Mock()
     mgr.parallelize.side_effect = RuntimeError("boom")
 
+    mock_diffusion_pipeline = MagicMock()
+    mock_diffusion_pipeline.from_pretrained.return_value = dummy_pipe
+
     with (
-        patch("diffusers.DiffusionPipeline.from_pretrained", return_value=dummy_pipe),
+        patch(f"{MODULE_PATH}.DIFFUSERS_AVAILABLE", True),
+        patch(f"{MODULE_PATH}.DiffusionPipeline", mock_diffusion_pipeline),
         patch(f"{MODULE_PATH}.torch.distributed.is_initialized", return_value=True),
         caplog.at_level(logging.WARNING),
     ):