fixed initialization

modeling_molmo.py

@@ -1759,9 +1759,14 @@ class Molmo(nn.Module):
         self.__num_fwd_flops: Optional[int] = None
 
         self.total_vocab_size = config.vocab_size + config.additional_vocab_size + config.skill_vocab_size
-        torch.nn.init.xavier_uniform_(self.transformer.skill_ff_out.weight)
-        if self.transformer.skill_ff_out.bias is not None:
-            torch.nn.init.zeros_(self.transformer.skill_ff_out.bias)
+
+    def init_weights(self):
+        if hasattr(self.transformer, "skill_ff_out"):
+            nn.init.xavier_uniform_(self.transformer.skill_ff_out.weight)
+            if self.transformer.skill_ff_out.bias is not None:
+                nn.init.zeros_(self.transformer.skill_ff_out.bias)
+        if hasattr(self.transformer.wte, "skill_embedding"):
+            nn.init.xavier_uniform_(self.transformer.wte.skill_embedding)
 
     def reset_parameters(self):
         if self.vision_backbone is not None:
@@ -1773,6 +1778,9 @@ class Molmo(nn.Module):
         if hasattr(self.transformer.wte, "new_embedding"):
             nn.init.normal_(self.transformer.wte.new_embedding, std=self.config.new_embedding_init_range)
 
+        if hasattr(self.transformer.wte, "skill_embedding"):
+            nn.init.xavier_uniform_(self.transformer.wte.skill_embedding)
+
         if hasattr(self.transformer, "wpe"):
             nn.init.normal_(self.transformer.wpe, mean=0.0, std=1.0)
 
@@ -1780,6 +1788,9 @@ class Molmo(nn.Module):
 
         if hasattr(self.transformer, "ff_out"):
             nn.init.normal_(self.transformer.ff_out, mean=0.0, std=0.02)
+        
+        if hasattr(self.transformer, "skill_ff_out"):
+            nn.init.normal_(self.transformer.skill_ff_out, mean=0.0, std=0.02)
 
         if self.config.block_group_size == 1:
             for block in self.transformer.blocks: