Merge pull request #4 from huggingface/moe-128

128 experts

Author: ArthurZucker

src/transformers/modeling_utils.py

@@ -5883,10 +5883,10 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict):
             generic_name = re.sub(r"\.\d+\.", ".*.", param_name)
             param_byte_count //= torch.distributed.get_world_size() if tp_plan_regex.search(generic_name) else 1
 
-        total_byte_count[device] += param_byte_count
+        parameter_count[device] += param_byte_count
 
     # This will kick off the caching allocator to avoid having to Malloc afterwards
-    for device, byte_count in total_byte_count.items():
+    for device, byte_count in parameter_count.items():
         if device.type == "cuda":
             index = device.index if device.index is not None else torch.cuda.current_device()
             device_memory = torch.cuda.mem_get_info(index)[0]

src/transformers/models/llama4/configuration_llama4.py

@@ -142,6 +142,7 @@ class Llama4TextConfig(PretrainedConfig):
         "layers.*.feed_forward.experts.gate_up_proj": "local_packed_rowwise",  # row because not linear
         "layers.*.feed_forward.experts.down_proj": "local_colwise",  # col because not linear
         "layers.*.feed_forward.experts": "local",
+        "layers.*.feed_forward.down_proj": "local_rowwise",
         "layers.*.feed_forward": "gather",
     }
 
@@ -150,6 +151,7 @@ def __init__(
         vocab_size=202048,
         hidden_size=5120,
         intermediate_size=8192,
+        intermediate_size_mlp=16384,
         num_hidden_layers=48,
         num_attention_heads=40,
         num_key_value_heads=8,
@@ -167,10 +169,12 @@ def __init__(
         attention_dropout=0.0,
         num_experts_per_tok=1,
         num_local_experts=16,
+        interleave_moe_layer_step=1,
+        use_qk_norm=True,
         output_router_logits=False,
         router_aux_loss_coef=0.001,
         router_jitter_noise=0.0,
-        rope_scaling="llama3",
+        rope_scaling=None,
         **kwargs,
     ):
         super().__init__(
@@ -184,6 +188,7 @@ def __init__(
         self.max_position_embeddings = max_position_embeddings
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
+        self.intermediate_size_mlp = intermediate_size_mlp
         self.num_hidden_layers = num_hidden_layers
         self.num_attention_heads = num_attention_heads
         self.rope_scaling = rope_scaling
@@ -201,9 +206,11 @@ def __init__(
         self.rope_theta = rope_theta
         self.attention_dropout = attention_dropout
         self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads
+        self.use_qk_norm = use_qk_norm
 
         self.num_experts_per_tok = num_experts_per_tok
         self.num_local_experts = num_local_experts
+        self.interleave_moe_layer_step = interleave_moe_layer_step
         self.output_router_logits = output_router_logits
         self.router_aux_loss_coef = router_aux_loss_coef
         self.router_jitter_noise = router_jitter_noise

src/transformers/models/llama4/convert_llama4_weights_to_hf.py

@@ -55,6 +55,14 @@
     # Unused keys in load hooks (explicitly removed)
     r'layers.(\d+).attention.wqkv._extra_state':             None,
     r'layers.(\d+).attention.wo._extra_state':               None,
+
+    # MLP layer variant
+    r"layers.(\d+).feed_forward.w1.weight":                  r"language_model.model.layers.\1.feed_forward.gate_proj.weight",               # might need to be fused for efficiency?
+    r"layers.(\d+).feed_forward.w3.weight":                  r"language_model.model.layers.\1.feed_forward.up_proj.weight",                 # might need to be fused for efficiency?
+    # r"layers.(\d+).feed_forward.mlp.fc1_weight":             r"language_model.model.layers.\1.feed_forward.gate_up_proj.weight",
+    r"layers.(\d+).feed_forward.mlp.fc2_weight":             r"language_model.model.layers.\1.feed_forward.down_proj.weight",
+    r"layers.(\d+).feed_forward.mlp.layer_norm.weight":      r"language_model.model.layers.\1.post_attention_layernorm.weight",
+
     # Vision encoder mapping
     r"vision_embeddings.vision_encoder.conv1._linear":                                            r"vision_model.patch_embedding.linear",
     r'vision_embeddings.vision_adapter.mlp.c_fc':                                                 r"vision_model.vision_adapter.mlp.fc1",
@@ -142,6 +150,9 @@ def get_concat_dim(key):
         "experts.gate_proj",
         "experts.up_proj",
         "expert.down_proj",
+        # "feed_forward.up_proj",
+        # "feed_forward.gate_proj",
+        "feed_forward.down_proj",
         "global_gate_stats",
         # vision dim1 sharded stuff
         "mlp.fc2.weight", # covers all rowparallels across vis
@@ -166,6 +177,20 @@ def safe_load(filename):
     return shard
 
 
+# Unpack mlp projections - possibly to be removed when they are fused
+def preprocess_keys(state_dict):
+    new_state_dict = dict()
+    for key, value in state_dict.items():
+        if "mlp.fc1_weight" in key:
+            prefix = key.split("mlp.fc1_weight")[0]
+            w1, w3 = value.chunk(2, dim=0)
+            new_state_dict[prefix + "w1.weight"] = w1
+            new_state_dict[prefix + "w3.weight"] = w3
+        else:
+            new_state_dict[key] = value
+    return new_state_dict
+
+
 def write_model(
     model_path,
     input_base_path,
@@ -194,14 +219,17 @@ def write_model(
     rms_norm_eps = params["norm_eps"]
     rope_theta = params["rope_theta"]
 
-    # some constans from original code
-    rope_scaling = {
-        "rope_type": "llama3",
-        "factor": 8.0,
-        "low_freq_factor": 1.0,
-        "high_freq_factor": 4.0,
-        "original_max_position_embeddings": 8192,
-    }
+    config_kwargs = {}
+    if params["use_scaled_rope"]:
+        # some constans from original code
+        rope_scaling = {
+            "rope_type": "llama3",
+            "factor": 8.0,
+            "low_freq_factor": 1.0,
+            "high_freq_factor": 4.0,
+            "original_max_position_embeddings": 8192,
+        }
+        config_kwargs.update(dict(rope_scaling=rope_scaling))
 
     # compute additional params for weight conversion
     num_heads_per_shard = num_heads // num_shards
@@ -211,9 +239,10 @@ def write_model(
     num_key_value_heads = params["n_kv_heads"]  # for GQA / MQA
 
     num_experts = params["moe_args"]["num_experts"]
+    interleave_moe_layer_step = params["moe_args"].get("interleave_moe_layer_step", 1)
 
     bos_token_id = 200000
-    eos_token_id = [200001, 200002, 200003] if instruct else 200001
+    eos_token_id = [200001, 200002, 200003, 200008] if instruct else 200001
     pad_token_id = 200008
 
     text_config = Llama4TextConfig(
@@ -224,13 +253,16 @@ def write_model(
         rope_theta=rope_theta,
         num_hidden_layers=num_layers,
         intermediate_size=8192,
-        rope_scaling=rope_scaling,
+        intermediate_size_mlp=16384,
         num_local_experts=num_experts,
+        interleave_moe_layer_step=interleave_moe_layer_step,
+        use_qk_norm=params["use_qk_norm"],
         bos_token_id=bos_token_id,
         eos_token_id=eos_token_id,
         pad_token_id=pad_token_id,
         tie_word_embeddings=False,  # Constant set to False
         torch_dtype=torch_dtype,
+        **config_kwargs,
     )
     # default vision config frmo params
 
@@ -273,6 +305,7 @@ def write_model(
                 safe_load(os.path.join(input_base_path, f"consolidated.{i:02d}.pth"))
                 for i in tqdm(range(num_shards), desc="Loading shards", unit="shard")
             ]
+        loaded = [preprocess_keys(d) for d in loaded]
 
         all_keys_raw = list(loaded[0].keys())
         repeated_keys = []
@@ -354,7 +387,7 @@ def write_model(
                 if gate_key == new_key:
                     state_dict[new_key] = torch.cat(current_parameter, dim=concat_dim)
                 elif new_key == up_key:
-                    if "shared" in new_key:
+                    if "experts" not in new_key:
                         gate_proj = state_dict.pop(gate_key)
                         up_proj = torch.cat(current_parameter, dim=concat_dim)
                         state_dict[gate_key] = gate_proj
@@ -365,11 +398,11 @@ def write_model(
                     else:
                         gate_proj = state_dict.pop(gate_key)
                         gate_proj = [
-                            gate_proj.reshape(16, -1, 8, 1024)[:, :, k, :].reshape(16, -1, 1024) for k in range(8)
+                            gate_proj.reshape(num_experts, -1, 8, 1024)[:, :, k, :].reshape(num_experts, -1, 1024) for k in range(8)
                         ]
                         gate_proj = torch.cat(gate_proj, dim=-1)
 
-                        up_proj = [k.reshape(16, -1, 8, 1024).reshape(16, -1, 1024) for k in current_parameter]
+                        up_proj = [k.reshape(num_experts, -1, 8, 1024).reshape(num_experts, -1, 1024) for k in current_parameter]
                         up_proj = torch.cat(up_proj, dim=-1)
 
                         gate_up_proj = torch.cat((gate_proj, up_proj), dim=-1)
@@ -432,10 +465,7 @@ def write_model(
         print("Loading the checkpoint in a Llama4 model.")
         state_dict.pop("")
         model.load_state_dict(state_dict, strict=True, assign=True)
-        print("Model reloaded successfully. Checking logits...")
-        # ipdb.set_trace()
-        # zero_out = model.forward(inputs_embeds=torch.zeros((1,743, 4096)))
-        # ipdb.set_trace()
+        print("Model reloaded successfully.")
         print("Saving the model.")
         model.save_pretrained(model_path, safe_serialization=safe_serialization)
         del state_dict, model
@@ -448,8 +478,7 @@ def write_model(
         model = Llama4ForConditionalGeneration.from_pretrained(
             model_path, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="eager"
         )
-        # ipdb.set_trace()
-        model.eval()
+
         model.generation_config.top_p = 0.9
         model.generation_config.temperature = 0.6
         print("Model reloaded successfully.")
@@ -458,7 +487,7 @@ def write_model(
 
         tokenizer = AutoTokenizer.from_pretrained(model_path)
         inputs = tokenizer(["Roses are red,"], return_tensors="pt").to(model.device)
-        out = model.generate(**inputs, max_new_tokens=10)
+        out = model.generate(**inputs, max_new_tokens=4)
         print(tokenizer.batch_decode(out))
     # generation config
     if instruct:

src/transformers/models/llama4/modeling_llama4.py

@@ -96,18 +96,23 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 # Phi3MLP
 class Llama4TextMLP(nn.Module):
-    def __init__(self, config):
+    def __init__(self, config, intermediate_size=None):
         super().__init__()
+
+        if intermediate_size is None:
+            intermediate_size = config.intermediate_size
+
         self.config = config
-        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
         self.activation_fn = ACT2FN[config.hidden_act]
 
     def forward(self, x):
         down_proj = self.activation_fn(self.gate_proj(x)) * self.up_proj(x)
         return self.down_proj(down_proj)
 
+
 class Llama4TextL2Norm(torch.nn.Module):
     def __init__(self, dim: int=None, eps: float = 1e-6):
         super().__init__()
@@ -193,7 +198,7 @@ class Llama4TextRotaryEmbedding(nn.Module):
     def __init__(self, config: Llama4TextConfig, device=None):
         super().__init__()
         # BC: "rope_type" was originally "type"
-        self.rope_type = "llama3"
+        self.rope_type = "llama3" if config.rope_scaling is not None else "default"
 
         self.max_seq_len_cached = config.max_position_embeddings
         self.original_max_seq_len = config.max_position_embeddings
@@ -319,7 +324,8 @@ def __init__(self, config, layer_idx):
         self.o_proj = nn.Linear(
             config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
         )
-        self.qk_norm = Llama4TextL2Norm()
+        if self.config.use_qk_norm:
+            self.qk_norm = Llama4TextL2Norm()
 
     def forward(
         self,
@@ -341,16 +347,7 @@ def forward(
             query_states, key_states, position_embeddings.to(query_states.device)
         )
 
-        # because L2 is computed on the shards, we need to find an appropriate reshape
-        # here, to make sure in TP but also non TP settings. Logits diverge otherwise
-        if query_states.shape[-1] == self.num_attention_heads * self.head_dim:
-            query_states = self.qk_norm(
-                query_states.view(input_shape[0], input_shape[1], self.pretraining_tp, -1)
-            ).reshape(hidden_shape)
-            key_states = self.qk_norm(
-                key_states.view(input_shape[0], input_shape[1], self.pretraining_tp, -1)
-            ).reshape((*input_shape, self.pretraining_tp, -1))
-        else:
+        if self.config.use_qk_norm:
             query_states = self.qk_norm(query_states)
             key_states = self.qk_norm(key_states)
 
@@ -394,7 +391,11 @@ def __init__(self, config, layer_idx):
         self.hidden_size = config.hidden_size
 
         self.self_attn = Llama4TextAttention(config, layer_idx)
-        self.feed_forward = Llama4TextMoe(config)
+        self.is_moe_layer = (layer_idx + 1) % config.interleave_moe_layer_step == 0
+        if self.is_moe_layer:
+            self.feed_forward = Llama4TextMoe(config)
+        else:
+            self.feed_forward = Llama4TextMLP(config, intermediate_size=config.intermediate_size_mlp)
 
         self.input_layernorm = Llama4TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = Llama4TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -457,7 +458,11 @@ def forward(
         residual = hidden_states
 
         hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states, router_logits = self.feed_forward(hidden_states)
+        hidden_states = self.feed_forward(hidden_states)
+        if self.is_moe_layer:
+            hidden_states, router_logits = hidden_states
+        else:
+            router_logits = None
         hidden_states = residual + hidden_states.view(residual.shape)
 
         outputs = (hidden_states,)