fixed formatting

Author: parambole

MaxText/convert_qwen3_moe.py

@@ -26,11 +26,9 @@
 import gc
 import os
 import pathlib
-import logging
 
 import numpy as np
 import torch
-import jax
 from safetensors import safe_open
 from tqdm import tqdm
 
@@ -59,23 +57,27 @@ def hf_to_maxtext_mapping(layer_idx: int, num_experts: int) -> dict:
       "lm_head.weight": "decoder.logits_dense.kernel",
   }
   # Layer-specific mappings for a pure MoE/scanned model
-  mapping.update(
-      {
-          f"model.layers.{layer_idx}.input_layernorm.weight": f"decoder.layers.{layer_idx}.pre_self_attention_layer_norm.scale",
-          f"model.layers.{layer_idx}.post_attention_layernorm.weight": f"decoder.layers.{layer_idx}.post_self_attention_layer_norm.scale",
-          f"model.layers.{layer_idx}.self_attn.q_proj.weight": f"decoder.layers.{layer_idx}.self_attention.query.kernel",
-          f"model.layers.{layer_idx}.self_attn.k_proj.weight": f"decoder.layers.{layer_idx}.self_attention.key.kernel",
-          f"model.layers.{layer_idx}.self_attn.v_proj.weight": f"decoder.layers.{layer_idx}.self_attention.value.kernel",
-          f"model.layers.{layer_idx}.self_attn.o_proj.weight": f"decoder.layers.{layer_idx}.self_attention.out.kernel",
-          f"model.layers.{layer_idx}.self_attn.q_norm.weight": f"decoder.layers.{layer_idx}.self_attention.query_norm.scale",
-          f"model.layers.{layer_idx}.self_attn.k_norm.weight": f"decoder.layers.{layer_idx}.self_attention.key_norm.scale",
-          f"model.layers.{layer_idx}.mlp.gate.weight": f"decoder.layers.{layer_idx}.moe_block.gate.kernel",
-      }
-  )
+  mapping.update({
+      f"model.layers.{layer_idx}.input_layernorm.weight": (
+          f"decoder.layers.{layer_idx}.pre_self_attention_layer_norm.scale"
+      ),
+      f"model.layers.{layer_idx}.post_attention_layernorm.weight": (
+          f"decoder.layers.{layer_idx}.post_self_attention_layer_norm.scale"
+      ),
+      f"model.layers.{layer_idx}.self_attn.q_proj.weight": f"decoder.layers.{layer_idx}.self_attention.query.kernel",
+      f"model.layers.{layer_idx}.self_attn.k_proj.weight": f"decoder.layers.{layer_idx}.self_attention.key.kernel",
+      f"model.layers.{layer_idx}.self_attn.v_proj.weight": f"decoder.layers.{layer_idx}.self_attention.value.kernel",
+      f"model.layers.{layer_idx}.self_attn.o_proj.weight": f"decoder.layers.{layer_idx}.self_attention.out.kernel",
+      f"model.layers.{layer_idx}.self_attn.q_norm.weight": f"decoder.layers.{layer_idx}.self_attention.query_norm.scale",
+      f"model.layers.{layer_idx}.self_attn.k_norm.weight": f"decoder.layers.{layer_idx}.self_attention.key_norm.scale",
+      f"model.layers.{layer_idx}.mlp.gate.weight": f"decoder.layers.{layer_idx}.moe_block.gate.kernel",
+  })
 
   # MoE expert mappings
   for i in range(num_experts):
-    mapping[f"model.layers.{layer_idx}.mlp.experts.{i}.gate_proj.weight"] = f"decoder.layers.{layer_idx}.moe_block.{i}.wi_0"
+    mapping[f"model.layers.{layer_idx}.mlp.experts.{i}.gate_proj.weight"] = (
+        f"decoder.layers.{layer_idx}.moe_block.{i}.wi_0"
+    )
     mapping[f"model.layers.{layer_idx}.mlp.experts.{i}.up_proj.weight"] = f"decoder.layers.{layer_idx}.moe_block.{i}.wi_1"
     mapping[f"model.layers.{layer_idx}.mlp.experts.{i}.down_proj.weight"] = f"decoder.layers.{layer_idx}.moe_block.{i}.wo"
 
@@ -163,6 +165,7 @@ def convert_hf_to_maxtext(base_model_path: str, model_params: dict) -> dict:
   moe["wo"] = np.zeros((num_experts, num_layers, moe_intermediate_size, hidden_size), dtype=np.float16)
 
   # Loop through layers and populate the stacked arrays
+  # pylint: disable=unsupported-assignment-operation
   for l in tqdm(range(num_layers), desc="Stacking layer weights"):
     ln["pre_self_attention_layer_norm"]["scale"][l, :] = (
         chkpt_vars[f"decoder.layers.{l}.pre_self_attention_layer_norm.scale"].to(torch.float16).numpy()
@@ -268,5 +271,5 @@ def main(args):
   parser.add_argument("--use-ocdbt", type=str2bool, default=True, help="Use OCDBT format for saving.")
   parser.add_argument("--use-zarr3", type=str2bool, default=True, help="Use Zarr3 format for saving.")
 
-  args = parser.parse_args()
-  main(args)
+  parsed_args = parser.parse_args()
+  main(parsed_args)

MaxText/layers/decoders.py

@@ -29,7 +29,6 @@
 
 from MaxText.common_types import DecoderBlockType, Config, MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE
 from MaxText import max_logging
-from MaxText import max_utils
 from MaxText.inference import page_manager
 from MaxText.layers import linears
 from MaxText.layers import quantizations
@@ -444,14 +443,7 @@ def scan_decoder_layers(self, cfg, decoder_layer, length, metadata_axis_name, me
         length=length,
         metadata_params={nn.PARTITION_NAME: metadata_axis_name},
     )
-    return scan_fn(
-        config=cfg,
-        mesh=mesh,
-        name=metadata_axis_name,
-        quant=self.quant,
-        model_mode=model_mode,
-        **kwargs
-    )
+    return scan_fn(config=cfg, mesh=mesh, name=metadata_axis_name, quant=self.quant, model_mode=model_mode, **kwargs)
 
   def get_pipeline_stage_module(self, decoder_blocks):
     """get pipeline stage module"""