Rebrand: Ministral DLM → Nemotron Labs Diffusion

config.json

@@ -5,14 +5,14 @@
   "adaptive_mask_rate": false,
   "ar_loss_weight": 1.0,
   "architectures": [
-    "MinistralDiffEncoderModel"
+    "NemotronLabsDiffusionEncoderModel"
   ],
   "attention_bias": false,
   "attention_dropout": 0.0,
   "attn_implementation": null,
   "auto_map": {
-    "AutoConfig": "configuration_ministral_dlm.MinistralDLMConfig",
-    "AutoModel": "modeling_ministral_dlm.MinistralDiffEncoderModel"
+    "AutoConfig": "configuration_nemotron_labs_diffusion.NemotronLabsDiffusionConfig",
+    "AutoModel": "modeling_nemotron_labs_diffusion.NemotronLabsDiffusionEncoderModel"
   },
   "block_size": 32,
   "bos_token_id": 1,
@@ -34,7 +34,7 @@
   "mask_token_id": 100,
   "max_position_embeddings": 262144,
   "mlp_bias": false,
-  "model_type": "ministral_dlm",
+  "model_type": "nemotron_labs_diffusion",
   "multi_sampling": null,
   "num_ar_layers": 0,
   "num_attention_heads": 32,

configuration_ministral_dlm.py → configuration_nemotron_labs_diffusion.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Ministral DLM model configuration"""
+"""Nemotron Labs Diffusion model configuration"""
 
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_rope_utils import rope_config_validation
@@ -22,7 +22,7 @@ from transformers.utils import logging
 logger = logging.get_logger(__name__)
 
 
-class MinistralDLMConfig(PretrainedConfig):
+class NemotronLabsDiffusionConfig(PretrainedConfig):
     r"""
     This is the configuration class to store the configuration of a [`Ministral3Model`] for diffusion language models.
     It is used to instantiate a Ministral model according to the specified arguments, defining the model architecture.
@@ -114,7 +114,7 @@ class MinistralDLMConfig(PretrainedConfig):
             Adaptive permutation ratio for global.
     """
 
-    model_type = "ministral_dlm"
+    model_type = "nemotron_labs_diffusion"
     keys_to_ignore_at_inference = ["past_key_values"]
 
     # Default tensor parallel plan for base model `Ministral`
@@ -243,5 +243,5 @@ class MinistralDLMConfig(PretrainedConfig):
         )
 
 
-__all__ = ["MinistralDLMConfig"]
+__all__ = ["NemotronLabsDiffusionConfig"]
 

modeling_ministral.py

@@ -25,7 +25,7 @@ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
 # from transformers.utils.generic import maybe_autocast
-from .configuration_ministral_dlm import MinistralDLMConfig
+from .configuration_nemotron_labs_diffusion import NemotronLabsDiffusionConfig
 
 #ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask_older_torch
 
@@ -110,7 +110,7 @@ def _get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_positi
 class Ministral3Attention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
-    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+    def __init__(self, config: NemotronLabsDiffusionConfig, layer_idx: int):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
@@ -234,7 +234,7 @@ class Ministral3RMSNorm(nn.Module):
 
 
 class Ministral3DecoderLayer(GradientCheckpointingLayer):
-    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+    def __init__(self, config: NemotronLabsDiffusionConfig, layer_idx: int):
         super().__init__()
         self.hidden_size = config.hidden_size
 
@@ -284,7 +284,7 @@ class Ministral3DecoderLayer(GradientCheckpointingLayer):
 
 @auto_docstring
 class Ministral3PreTrainedModel(PreTrainedModel):
-    config: MinistralDLMConfig
+    config: NemotronLabsDiffusionConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["Ministral3DecoderLayer"]
@@ -304,7 +304,7 @@ class Ministral3PreTrainedModel(PreTrainedModel):
 class Ministral3RotaryEmbedding(nn.Module):
     inv_freq: torch.Tensor  # fix linting for `register_buffer`
 
-    def __init__(self, config: MinistralDLMConfig, device=None):
+    def __init__(self, config: NemotronLabsDiffusionConfig, device=None):
         super().__init__()
         self.max_seq_len_cached = config.max_position_embeddings
         self.original_max_seq_len = config.max_position_embeddings
@@ -323,7 +323,7 @@ class Ministral3RotaryEmbedding(nn.Module):
 
     @staticmethod
     def compute_default_rope_parameters(
-        config: Optional[MinistralDLMConfig] = None,
+        config: Optional[NemotronLabsDiffusionConfig] = None,
         device: Optional["torch.device"] = None,
         seq_len: Optional[int] = None,
     ) -> tuple["torch.Tensor", float]:
@@ -370,7 +370,7 @@ class Ministral3RotaryEmbedding(nn.Module):
 
 @auto_docstring
 class Ministral3Model(Ministral3PreTrainedModel):
-    def __init__(self, config: MinistralDLMConfig):
+    def __init__(self, config: NemotronLabsDiffusionConfig):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size

modeling_ministral_dlm.py → modeling_nemotron_labs_diffusion.py

@@ -29,11 +29,11 @@ import math
 
 from .chat_utils import generate_with_prefix_cache_block_diff
 from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale
-from .configuration_ministral_dlm import MinistralDLMConfig
+from .configuration_nemotron_labs_diffusion import NemotronLabsDiffusionConfig
 
 
 @dataclass
-class MinistralDiffOutputWithPast(ModelOutput):
+class NemotronLabsDiffusionOutputWithPast(ModelOutput):
     loss: torch.FloatTensor | None = None
     logits: torch.FloatTensor | None = None
     causal_logits: torch.FloatTensor | None = None
@@ -87,7 +87,7 @@ def _extract_draft_kv_cache(past_key_values: DynamicCache, clean_len: int, block
 
 
 # with reference to https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb
-class MinistralFlexAttention(Ministral3Attention):
+class NemotronLabsDiffusionFlexAttention(Ministral3Attention):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
  
@@ -434,14 +434,14 @@ def gumbel_topk(log_w: torch.Tensor, k: int) -> torch.Tensor:
     return mask
 
 
-class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
+class NemotronLabsDiffusionEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
     """
     A single model with:
       - a bidirectional encoder + diffusion‐LM head over A
       - a causal decoder + LM head over B, conditioned on F_A
     """
 
-    def __init__(self, config: MinistralDLMConfig):
+    def __init__(self, config: NemotronLabsDiffusionConfig):
         super().__init__(config)
 
         self.mask_token_id = config.mask_token_id
@@ -450,7 +450,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
         diffusion_config.diffusion_lm = True
 
         if config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
-            diffusion_config.attn_class = MinistralFlexAttention
+            diffusion_config.attn_class = NemotronLabsDiffusionFlexAttention
         elif config.dlm_paradigm in ['bidirectional', 'autoregressive']:
             diffusion_config.attn_class = Ministral3Attention
 
@@ -867,7 +867,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
                     else: 
                         loss = (loss, num_mask_tokens)
 
-        return MinistralDiffOutputWithPast(
+        return NemotronLabsDiffusionOutputWithPast(
             loss=loss if not is_teacher else logits,
             logits=logits,
             causal_logits=causal_logits,
@@ -1109,4 +1109,4 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
         return x[:, : -(block_length * 2)], nfe
 
 
-__all__ = ["MinistralDiffEncoderModel", "MinistralFlexAttention"]
+__all__ = ["NemotronLabsDiffusionEncoderModel", "NemotronLabsDiffusionFlexAttention"]