Fix _init_weights and RotaryEmbedding for transformers v5.x compatibility

modeling_bailing_moe_v2.py

@@ -201,9 +201,21 @@ class BailingMoeV2RotaryEmbedding(nn.Module):
         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)
+        if self.rope_type == "default":
+            # ROPE_INIT_FUNCTIONS does not contain "default"; compute inv_freq
+            # directly. Use explicit float32 dtype to prevent bf16 overflow
+            # when the model is instantiated under torch_dtype=bfloat16 (e.g.
+            # via from_pretrained). With rope_theta=600000, base**x overflows
+            # bf16 (max ~65504) for most frequency bands.
+            base = float(getattr(config, "rope_theta", 10000.0))
+            dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+            inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+            if device is not None:
+                inv_freq = inv_freq.to(device)
+            self.attention_scaling = 1.0
+        else:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.original_inv_freq = self.inv_freq
 
@@ -1060,13 +1072,13 @@ class BailingMoeV2PreTrainedModel(PreTrainedModel):
     def _init_weights(self, module):
         std = self.config.initializer_range
         if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
             if module.bias is not None:
-                module.bias.data.zero_()
+                torch.nn.init.zeros_(module.bias)
         elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
             if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
+                torch.nn.init.zeros_(module.weight[module.padding_idx])
 
 
 BAILINGMOEV2_INPUTS_DOCSTRING = r"""