update config & modeling

config.json

@@ -28,6 +28,6 @@
     "num_experts": 8,
     "output_router_logits": false,
     "router_aux_loss_coef": 0.001,
-    "torch_dtype": "float16",
+    "torch_dtype": "bfloat16",
     "transformers_version": "4.35.0"
 }

modeling_grok1.py

@@ -7,14 +7,16 @@ from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 
 try:
-    from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+    from transformers.modeling_attn_mask_utils import \
+        _prepare_4d_causal_attention_mask
 
     HAS_MASK_UTILS = True
 except ImportError:
     HAS_MASK_UTILS = False
 
 from .configuration_grok1 import Grok1Config
-from .modeling_grok1_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from .modeling_grok1_outputs import (MoeCausalLMOutputWithPast,
+                                     MoeModelOutputWithPast)
 
 logger = logging.get_logger(__name__)
 
@@ -549,7 +551,7 @@ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int]
 
 
 class Grok1Model(Grok1PretrainedModel):
-    def __init__(self, config: Grok1Config) -> None:
+    def __init__(self, config: Grok1Config, **kwargs) -> None:
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
@@ -787,7 +789,7 @@ class Grok1Model(Grok1PretrainedModel):
 class Grok1ModelForCausalLM(Grok1PretrainedModel):
     _tied_weights_keys = ["lm_head.weight"]
 
-    def __init__(self, config: Grok1Config):
+    def __init__(self, config: Grok1Config, **kwargs):
         super().__init__(config)
         self.model = Grok1Model(config)
         self.vocab_size = config.vocab_size