add tp-support for mixtral

src/transformers/modeling_utils.py

@@ -2475,7 +2475,8 @@ def save_pretrained(
             if safe_serialization:
                 # At some point we will need to deal better with save_function (used for TPU and other distributed
                 # joyfulness), but for now this enough.
-                safe_save_file(shard, os.path.join(save_directory, shard_file), metadata={"format": "pt"})
+                safe_save_file(shard, os.path.join(save_directory, 
+                                f"{shard_file.split('.')[0]}-tp_{(torch.distributed.get_rank() if torch.distributed.is_initialized() else 0):0>2d}.safetensors"), metadata={"format": "pt"})
             else:
                 save_function(shard, os.path.join(save_directory, shard_file))
 

src/transformers/models/mixtral/configuration_mixtral.py

@@ -135,6 +135,7 @@ def __init__(
         output_router_logits=False,
         router_aux_loss_coef=0.001,
         moe_layer_frequency=1,
+        tp_size=1,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -162,6 +163,7 @@ def __init__(
         self.output_router_logits = output_router_logits
         self.router_aux_loss_coef = router_aux_loss_coef
         self.moe_layer_frequency = moe_layer_frequency
+        self.tp_size = tp_size
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,

src/transformers/models/mixtral/modeling_mixtral.py

@@ -270,10 +270,10 @@ def __init__(self, config: MixtralConfig, layer_idx: Optional[int] = None):
                 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                 f" and `num_heads`: {self.num_heads})."
             )
-        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
-        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads // config.tp_size * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads // config.tp_size * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads // config.tp_size * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads // config.tp_size * self.head_dim, self.hidden_size, bias=False)
 
         self.rotary_emb = MixtralRotaryEmbedding(
             self.head_dim,
@@ -762,9 +762,9 @@ def __init__(self, config: MixtralConfig):
         self.ffn_dim = config.intermediate_size
         self.hidden_dim = config.hidden_size
 
-        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
-        self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
-        self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
+        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim // config.tp_size, bias=False)
+        self.w2 = nn.Linear(self.ffn_dim // config.tp_size, self.hidden_dim, bias=False)
+        self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim // config.tp_size, bias=False)
 
         self.act_fn = ACT2FN[config.hidden_act]