fix: align RotaryEmbedding with Qwen2Moe pattern for transformers compat

Replace custom rope_type handling with the standard Qwen2Moe pattern: use compute_default_rope_parameters for "default" type, only look up ROPE_INIT_FUNCTIONS for non-default types. Also adds partial_rotary_factor support.

modeling_llada2_moe.py

@@ -92,32 +92,41 @@ ALL_LAYERNORM_LAYERS.append(LLaDA2MoeRMSNorm)
 
 
 class LLaDA2MoeRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for register_buffer
+
     def __init__(self, config: LLaDA2MoeConfig, device=None):
         super().__init__()
-        # BC: "rope_type" was originally "type"
-        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
-            self.rope_type = config.rope_scaling.get(
-                "rope_type", config.rope_scaling.get("type")
-            )
-        else:
-            self.rope_type = "linear"
-        # BC: "default" was removed from ROPE_INIT_FUNCTIONS in newer transformers
-        if self.rope_type == "default":
-            self.rope_type = "linear"
-        # Ensure rope_scaling has a factor for linear rope (defaults to no scaling)
-        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
-            config.rope_scaling.setdefault("factor", 1.0)
-        if hasattr(config, "rope_parameters") and config.rope_parameters is not None:
-            config.rope_parameters.setdefault("factor", 1.0)
         self.max_seq_len_cached = config.max_position_embeddings
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
 
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: LLaDA2MoeConfig = None,
+        device=None,
+        seq_len: int = None,
+    ):
+        base = config.rope_parameters["rope_theta"]
+        partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        dim = int(head_dim * partial_rotary_factor)
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
 
     @torch.no_grad()
     @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)