config fixes

configuration_rnd.py

@@ -7,6 +7,7 @@ extending Qwen3MoeConfig with RND1-specific parameters.
 
 from typing import Optional
 from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
+from transformers import AutoConfig
 
 
 class RND1Config(Qwen3MoeConfig):
@@ -34,10 +35,19 @@ class RND1Config(Qwen3MoeConfig):
         **kwargs,
     ):
         # Force non-causal and no caching for RND1
-        kwargs['use_cache'] = False
-        kwargs['is_causal'] = False
+        kwargs["use_cache"] = False
+        kwargs["is_causal"] = False
         super().__init__(**kwargs)
 
+        # `head_dim` needs to be 128 for Qwen3MoE
+        # need to ensure that the config has this attr if directly passing config to RND1LM at instantiation
+        if not hasattr(self, "head_dim"):
+            self.head_dim = 128
+
+        # Note that in transformers 4.57.0 there is an error in the config
+        # num_hidden_layers is defaulted to 24
+        self.num_hidden_layers = 48
+
         # RND1-specific parameters
         self.moe_backend = moe_backend
         self.num_diffusion_steps = num_diffusion_steps
@@ -55,9 +65,12 @@ class RND1Config(Qwen3MoeConfig):
         the correct custom classes are automatically resolved.
         """
         data = super().to_dict()
-        data.setdefault("auto_map", {
-            "AutoConfig": "configuration_rnd.RND1Config",
-            "AutoModel": "modeling_rnd.RND1Model",
-            "AutoModelForMaskedLM": "modeling_rnd.RND1LM",
-        })
-        return data
+        data.setdefault(
+            "auto_map",
+            {
+                "AutoConfig": "configuration_rnd.RND1Config",
+                "AutoModel": "modeling_rnd.RND1Model",
+                "AutoModelForMaskedLM": "modeling_rnd.RND1LM",
+            },
+        )
+        return data