init

modeling_flexbert.py

@@ -1537,17 +1537,26 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
         self._init_weights(reset_params=False)
 
     def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
+        # Handle the XOR condition
         assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
-        if module:
-            self._init_module_weights(module)
+        
+        if module is not None:
+            # Add basic initialization for common module types
+            if isinstance(module, (nn.Linear, nn.Embedding)):
+                module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+                if isinstance(module, nn.Linear) and module.bias is not None:
+                    module.bias.data.zero_()
+            elif isinstance(module, nn.LayerNorm):
+                module.bias.data.zero_()
+                module.weight.data.fill_(1.0)
         else:
             assert isinstance(reset_params, bool)
             self.bert._init_weights(reset_params=reset_params)
             self.lm_head._init_weights(reset_params=reset_params)
-
+            
             if not self.config.tie_word_embeddings:
                 init_weights(self.config, self.decoder, self.config.hidden_size, type_of_module=ModuleType.final_out)
-
+                
     @classmethod
     def from_composer(
         cls,