fix: add missing use_deterministic_attn parameter to MoonViT3dEncoder (#22)

- fix: add missing use_deterministic_attn parameter to MoonViT3dEncoder (f13e4d071155abcae05f85cc3c1b5fa56ea1b9fc)
- also fix weight init issue (bed55d493c07900df87ed51274581025cbf8315a)


Co-authored-by: Ace Eldeib <ace-coreweave@users.noreply.huggingface.co>

modeling_deepseek.py

@@ -1244,7 +1244,7 @@ class DeepseekV3PreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
+        if isinstance(module, nn.Linear) and hasattr(module, "weight"):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
                 module.bias.data.zero_()

modeling_kimi_k25.py

@@ -562,7 +562,8 @@ class MoonViT3dEncoder(nn.Module):
                  hidden_dim: int,
                  num_layers: int,
                  block_cfg: dict,
-                 video_attn_type: str = 'spatial_temporal') -> None:
+                 video_attn_type: str = 'spatial_temporal',
+                 use_deterministic_attn: bool = False) -> None:
         super().__init__()
 
         assert video_attn_type == 'spatial_temporal', f'video_attn_type must be "spatial_temporal", got {video_attn_type}'
@@ -572,7 +573,7 @@ class MoonViT3dEncoder(nn.Module):
         self.blocks = nn.ModuleList([
             MoonViTEncoderLayer(
                 **block_cfg,
-                use_deterministic_attn=self.use_deterministic_attn)
+                use_deterministic_attn=use_deterministic_attn)
             for _ in range(num_layers)
         ])
         self.final_layernorm = nn.LayerNorm(hidden_dim)
@@ -637,6 +638,20 @@ class MoonViT3dPretrainedModel(PreTrainedModel):
     _no_split_modules = ['PackingTransformer']
     _supports_flash_attn_2 = True
     _supports_sdpa = True
+    
+    def _init_weights(self, module):
+        # Default PreTrainedModel._init_weights treats nn.Linear subclasses (e.g.
+        # compressed_tensors.CompressedLinear) as Linear but those layers may have
+        # no `weight` after compression; skip in that case.
+        std = getattr(self.config, "initializer_range", 0.02)
+        if isinstance(module, (nn.Linear, nn.Conv2d)) and hasattr(module, "weight"):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
 
     def __init__(self, config, *inputs, **kwargs):
         super().__init__(config, *inputs, **kwargs)
@@ -785,7 +800,7 @@ class KimiK25PreTrainedModel(PreTrainedModel):
         if hasattr(module, "class_embedding"):
             module.class_embedding.data.normal_(mean=0.0, std=std)
 
-        if isinstance(module, (nn.Linear, nn.Conv2d)):
+        if isinstance(module, (nn.Linear, nn.Conv2d)) and hasattr(module, "weight"):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
                 module.bias.data.zero_()