Add mixtral support to new CLI

Author: qihqi

jetstream_pt/cli.py

@@ -11,6 +11,7 @@
 from jetstream_pt import fetch_models
 from jetstream_pt import environment, engine, quantize_model, torchjax
 from jetstream_pt import config
+from transformers import AutoTokenizer
 
 
 FLAGS = flags.FLAGS
@@ -25,13 +26,13 @@
 
 def shard_weights(env, weights, weight_shardings):
   """Shard weights according to weight_shardings"""
-  for k, v in weight_shardings.items():
-    print("SHARDING", k, v)
   sharded = {}
   for key, val in weights.items():
     sharding = env.sharding_by_axis(weight_shardings.get(key, -1))
     with jax.default_device(jax.devices("cpu")[0]):
       arr = torch_xla2.tensor.t2j(val)
+    
+    print("SHARDING", key, sharding)
     arr = jax.device_put(arr, sharding)
     sharded[key] = torchjax.to_torch(arr)
   return sharded
@@ -48,7 +49,9 @@ def create_engine(devices):
       FLAGS.max_output_length,
       quant_config.enable_weight_quantization,
   )
+  tokenizer = AutoTokenizer.from_pretrained(FLAGS.model_id)
   env = environment.JetEngineEnvironment(env_data)
+  env.hf_tokenizer = tokenizer
   model = fetch_models.instantiate_model_from_repo_id(FLAGS.model_id, env)
 
   weight_shardings = model.get_sharding_annotations()

jetstream_pt/engine.py

@@ -36,6 +36,7 @@
 from jetstream_pt import cache_manager
 from jetstream_pt import quantize
 from jetstream_pt import torchjax
+from jetstream_pt.hf_tokenizer import HFTokenizerAdapter
 from jetstream_pt.environment import JetEngineEnvironment, JetEngineEnvironmentData, QuantizationConfig
 from jetstream_pt.third_party.llama import model_exportable as llama_model, model_args
 from jetstream_pt.third_party.gemma import config as gemma_config, model as gemma_model
@@ -705,6 +706,8 @@ def get_tokenizer(self) -> tokenizer_pb2.TokenizerParameters:
   def build_tokenizer(
       self, metadata: tokenizer_pb2.TokenizerParameters  # pylint: disable=all
   ) -> tokenizer_api.Tokenizer:
+    if self.env.hf_tokenizer is not None:
+      return HFTokenizerAdapter(self.env.hf_tokenizer)
     if "llama-3" in self.env.model_type:
       return token_utils.TikToken(metadata)
 

jetstream_pt/environment.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import dataclasses
-from typing import Tuple
+from typing import Tuple, Any
 
 import jax
 import jax.numpy as jnp
@@ -141,6 +141,10 @@ def __init__(self, data: JetEngineEnvironmentData):
     self.testing_seed = self._data.testing_seed
     self.ring_buffer = self._data.ring_buffer
 
+    # If not None, then use this tokenizer without 
+    # trying to create new ones.
+    self.hf_tokenizer = None
+
     if not self.ring_buffer:
       self.lazy_cache_update = True
       self.ragged_mha = True

jetstream_pt/fetch_models.py

@@ -12,6 +12,7 @@
     QuantizationConfig,
 )
 from jetstream_pt.third_party.llama import model_exportable as llama_model
+from jetstream_pt.third_party.mixtral import model as mixtral_model
 
 FLAGS = flags.FLAGS
 
@@ -38,12 +39,15 @@ class ModelInfo:
   num_layers: int
   num_heads: int
   head_dim: int
+  n_reps: int # repeatition for GQA
 
 
-_llama2_7 = ModelInfo(llama_model.Transformer, 32, 32, 128)
-_llama2_13 = ModelInfo(llama_model.Transformer, 40, 40, 128)
-_llama2_70 = ModelInfo(llama_model.Transformer, 80, 8, 128)
-_llama3_8 = ModelInfo(llama_model.Transformer, 32, 8, 128)
+_llama2_7 = ModelInfo(llama_model.Transformer, 32, 32, 128, 1)
+_llama2_13 = ModelInfo(llama_model.Transformer, 40, 40, 128, 1)
+_llama2_70 = ModelInfo(llama_model.Transformer, 80, 8, 128, 4)
+_llama3_8 = ModelInfo(llama_model.Transformer, 32, 8, 128, 4)
+
+_mixtral_87 = ModelInfo(mixtral_model.Transformer, 32, 8, 128, 4)
 
 
 model_id_to_class = {
@@ -57,8 +61,8 @@ class ModelInfo:
     "google/gemma-2b-it": None,
     "google/gemma-7b": None,
     "google/gemma-7b-it": None,
-    "mistralai/Mixtral-8x7B-v0.1": None,
-    "mistralai/Mixtral-8x7B-Instruct-v0.1": None,
+    "mistralai/Mixtral-8x7B-v0.1": _mixtral_87,
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": _mixtral_87,
 }
 
 
@@ -107,6 +111,7 @@ def construct_env_data_from_model_id(
       else input_length + output_length
   )
 
+  model_info = model_id_to_class.get(repo_id)
   env_data = JetEngineEnvironmentData(
       tokenizer_path=tokenizer_path,
       checkpoint_path=checkpoint_path,
@@ -119,8 +124,8 @@ def construct_env_data_from_model_id(
       bf16_enable=True,
       sharding_config_path="",
       shard_on_batch=shard_on_batch,
+      n_reps=model_info.n_reps,
   )
-  model_info = model_id_to_class.get(repo_id)
   env_data.cache_shape = (
       batch_size,
       model_info.num_heads,

jetstream_pt/hf_tokenizer.py

@@ -0,0 +1,54 @@
+from jetstream.engine import tokenizer_api
+
+
+class HFTokenizerAdapter(tokenizer_api.Tokenizer):
+
+  def __init__(self, tokenizer):
+    self.tokenizer = tokenizer
+
+
+  def encode(
+      self, s: str, **kwargs
+  ): 
+    """Tokenize a string.
+    Args:
+        s: String to tokenize.
+        **kwargs: Additional keyword arguments.
+    Returns:
+        tokens: Tokenized into integers.
+        true_length: Actual length of the non-padded sequence
+          if padding is used.
+    """
+    return self(s)
+
+  def decode(self, token_ids: list[int], **kwargs) -> str:
+    """Processess input token ids to generate a string.
+    Args:
+      token_ids: List of token ids.
+      **kwargs: Additional keyword arguments.
+    Returns:
+      str: String generated from the token ids.
+    """
+    return self.decode(token_ids)
+
+  @property
+  def pad_id(self) -> int:
+    """ID of the pad token."""
+    return (self.tokenizer.pad_token_id 
+            if self.tokenizer.pad_token_id else 0)
+
+
+  @property
+  def eos_id(self) -> int:
+    """ID of EOS token."""
+    return self.tokenizer.eos_token_id
+
+  @property
+  def bos_id(self) -> int:
+    """ID of BOS token."""
+    return self.tokenizer.bos_token_id
+
+  @property
+  def stop_tokens(self) -> set[int]:
+    """ID of the stop token."""
+    return {self.eos_id, self.pad_id}

jetstream_pt/third_party/llama/model_exportable.py

@@ -75,6 +75,11 @@ def __init__(
     self.annotate_sharding("w1.weight", 0)
     self.annotate_sharding("w2.weight", 1)
     self.annotate_sharding("w3.weight", 0)
+    if LinearLayer != torch.nn.Linear:
+      self.annotate_sharding("w1.weight_scaler", 0)
+      self.annotate_sharding("w2.weight_scaler", 0)
+      self.annotate_sharding("w3.weight_scaler", 0)
+
 
   def forward(self, x):
     result = self.w2(F.silu(self.w1(x)) * self.w3(x))

jetstream_pt/third_party/mixtral/model.py

@@ -22,11 +22,12 @@
 from torch.nn import functional as F
 from .config import ModelArgs, find_multiple
 from jetstream_pt.layers import Attention, get_quantized_linear_layer, get_quantized_embedding_layer
+from jetstream_pt.model_base import ModuleBase
 
 import jax
 
 
-class Transformer(nn.Module):
+class Transformer(ModuleBase):
 
   def __init__(self, config: ModelArgs, env) -> None:
     super().__init__()
@@ -37,6 +38,7 @@ def __init__(self, config: ModelArgs, env) -> None:
     self.tok_embeddings = Embedding(
         config.vocab_size, config.dim, device=config.device
     )
+
     self.layers = nn.ModuleList(
         TransformerBlock(config, env, layer_id)
         for layer_id in range(config.n_layer)
@@ -47,6 +49,14 @@ def __init__(self, config: ModelArgs, env) -> None:
         config.dim, config.vocab_size, bias=False, device=config.device
     )
 
+    self.hf_name("norm", "model.norm")
+    self.hf_name("layers", "model.layers")
+    self.hf_name('output', 'lm_head')
+    self.hf_name('tok_embeddings', 'model.embed_tokens')
+
+    self.annotate_sharding("tok_embeddings.weight", 1)
+    self.annotate_sharding("output.weight", 0)
+
     self.max_batch_size = -1
     self.max_seq_length = -1
 
@@ -140,8 +150,20 @@ def get_weight_sharding_type():
         "output.weight": "ColumnParallelLinear",
     }
 
+  @classmethod
+  def from_hf_model_id(cls, model_id, env):
+    name = {
+      "mistralai/Mixtral-8x7B-v0.1": "Mixtral-8x7B-v0.1",
+      "mistralai/Mixtral-8x7B-Instruct-v0.1": "Mixtral-8x7B-v0.1",
+    }.get(model_id)
+    assert name
+    args = ModelArgs.from_name(name)
+    args.device = 'meta'
+    model = cls(args, env)
+    return model
 
-class TransformerBlock(nn.Module):
+
+class TransformerBlock(ModuleBase):
 
   def __init__(self, config: ModelArgs, env, layer_id) -> None:
     super().__init__()
@@ -154,10 +176,37 @@ def __init__(self, config: ModelArgs, env, layer_id) -> None:
         device=config.device,
         layer_id=layer_id,
     )
+    self.hf_name("attention", "self_attn")
+    self.attention.hf_name("wq", "q_proj")
+    self.attention.hf_name("wk", "k_proj")
+    self.attention.hf_name("wv", "v_proj")
+    self.attention.hf_name("wo", "o_proj")
+
+    self.attention.annotate_sharding("wq", 0)
+    self.attention.annotate_sharding("wk", 0)
+    self.attention.annotate_sharding("wv", 0)
+    self.attention.annotate_sharding("wo", 1)
+
     self.block_sparse_moe = MOEFeedForward(config, config.device, env)
     self.ffn_norm = RMSNorm(config.dim, config.norm_eps)
     self.attention_norm = RMSNorm(config.dim, config.norm_eps)
 
+    self.hf_name("attention_norm", "input_layernorm")
+    self.hf_name("ffn_norm", "post_attention_layernorm")
+    self._register_load_state_dict_pre_hook(self.load_hook)
+
+  def load_hook(self, state_dict, prefix, *args):
+    if prefix + "block_sparse_moe.experts" in state_dict:
+      w1s, w2s, w3s = [], [], []
+      for i in range(8):
+        exp_prefix = f"{prefix}block_sparse_moe.experts.{i}."
+        w1s.append(state_dict.pop(exp_prefix + ".w1"))
+        w2s.append(state_dict.pop(exp_prefix + ".w2"))
+        w3s.append(state_dict.pop(exp_prefix + ".w3"))
+      state_dict[prefix + "block_sparse_moe.cond_ffn.w1"] = torch.cat(w1s)
+      state_dict[prefix + "block_sparse_moe.cond_ffn.w2"] = torch.cat(w2s)
+      state_dict[prefix + "block_sparse_moe.cond_ffn.w3"] = torch.cat(w3s)
+
   def forward(
       self,
       x: Tensor,
@@ -189,7 +238,7 @@ def forward(
     return out
 
 
-class Int8ConditionalFeedForward(nn.Module):
+class Int8ConditionalFeedForward(ModuleBase):
 
   def __init__(self, config):
     super().__init__()
@@ -215,12 +264,20 @@ def __init__(self, config):
     self.register_buffer("w2", w2)
     self.register_buffer("w3", w3)
 
+    self.annotate_sharding("w1", 1)
+    self.annotate_sharding("w2", 2)
+    self.annotate_sharding("w3", 1)
+
     w1_scaler = torch.empty(config.num_experts, config.intermediate_size)
     w2_scaler = torch.empty(config.num_experts, config.dim)
     w3_scaler = torch.empty(config.num_experts, config.intermediate_size)
+
     self.register_buffer("w1_scaler", w1_scaler)
     self.register_buffer("w2_scaler", w2_scaler)
     self.register_buffer("w3_scaler", w3_scaler)
+    self.annotate_sharding("w1_scaler", 1)
+    self.annotate_sharding("w2_scaler", -1)
+    self.annotate_sharding("w3_scaler", 1)
 
   def forward(self, x: Tensor, expert_indices: Tensor) -> Tensor:
     seq_len = x.shape[0]
@@ -266,7 +323,7 @@ def forward_for_long_seq_len(self, x, expert_indices):
       return expert_outs[seq_indexes, expert_indices]
 
 
-class ConditionalFeedForward(nn.Module):
+class ConditionalFeedForward(ModuleBase):
 
   def __init__(self, config):
     super().__init__()
@@ -280,6 +337,9 @@ def __init__(self, config):
     self.w3 = nn.Parameter(
         torch.empty(config.num_experts, config.intermediate_size, config.dim)
     )
+    self.annotate_sharding("w1", 1)
+    self.annotate_sharding("w2", 2)
+    self.annotate_sharding("w3", 1)
 
   def forward(self, x: Tensor, expert_indices: Tensor) -> Tensor:
     seq_len = x.shape[0]
@@ -318,7 +378,7 @@ def forward_for_long_seq_len(self, x, expert_indices):
       return expert_outs[seq_indexes, expert_indices]
 
 
-class MOEFeedForward(nn.Module):
+class MOEFeedForward(ModuleBase):
 
   def __init__(self, config, device, env) -> None:
     super().__init__()
@@ -352,7 +412,7 @@ def forward(self, x: Tensor) -> Tensor:
     return expert_outs
 
 
-class RMSNorm(nn.Module):
+class RMSNorm(ModuleBase):
 
   def __init__(self, dim: int, eps: float = 1e-5):
     super().__init__()