examples/quantization/quantize_mixed_precision_moe.py (5)

48-65: Fix potential IndexError when some amax files are missing

The current implementation will cause an IndexError if some (but not all) amax files are missing. When files are missing, state_dict_list will have fewer elements than world_size, but the loop at line 60 still iterates over the full range.

Apply this fix to handle partial file loading correctly:

 def load_and_preprocess_state_dict(modelopt_state_root, world_size=8):
-    state_dict_list = []
+    state_dict_list = {}
     # load amax from state dict
     for rank in range(world_size):
         amax_file = f"{modelopt_state_root}/amax_dict_rank{rank}-mp{world_size}.pt"
         if os.path.exists(amax_file):
-            state_dict_list.append(torch.load(amax_file, map_location="cuda:0"))
+            state_dict_list[rank] = torch.load(amax_file, map_location="cuda:0")
         else:
             print(f"WARNING: amax file not found: {amax_file}")
 
     if not state_dict_list:
         print("ERROR: No amax files loaded!")
         return {}
 
     # calculate the max across all TP ranks
-    merged_state_dict = state_dict_list[0]
-    for rank in range(world_size):
-        for key, amax in state_dict_list[rank].items():
+    merged_state_dict = next(iter(state_dict_list.values())).copy()
+    for rank, state_dict in state_dict_list.items():
+        for key, amax in state_dict.items():
             if key in merged_state_dict.items():
                 amax = torch.max(amax.to(0), merged_state_dict[key].to(0))
             merged_state_dict[key] = amax.to(0)

---

242-323: Refactor to eliminate code duplication for activation scales processing

The activation scales extraction logic is duplicated between rank-0 (lines 242-249) and non-rank-0 (lines 315-322) branches. This violates the DRY principle.

Extract the common logic before the rank check:

+    # Process activation scales for all ranks
+    if os.path.isdir(args.act_scales):
+        # Extract activation scales
+        renamed_state_dict = load_and_preprocess_state_dict(
+            modelopt_state_root=args.act_scales, world_size=8)
+        scales = get_scales_from_amax(start_layer=start_layer,
+                                      end_layer=end_layer,
+                                      renamed_state_dict=renamed_state_dict)
+        new_safetensors.update(scales)
+
     if args.rank == 0:
-        if os.path.isdir(args.act_scales):
-            # Extract activation scales
-            renamed_state_dict = load_and_preprocess_state_dict(
-                modelopt_state_root=args.act_scales, world_size=8)
-            scales = get_scales_from_amax(start_layer=start_layer,
-                                          end_layer=end_layer,
-                                          renamed_state_dict=renamed_state_dict)
-            new_safetensors.update(scales)
-        else:
+        if not os.path.isdir(args.act_scales):
             input_scales = safe_open(args.act_scales, "pt")
             for k in input_scales.keys():
                 new_safetensors.update({k: input_scales.get_tensor(k)})
                 new_json['weight_map'][k] = args.act_scales.split("/")[-1]

         # ... rest of rank-0 specific logic ...
     else:
-        if os.path.isdir(args.act_scales):
-            # Extract activation scales
-            renamed_state_dict = load_and_preprocess_state_dict(
-                modelopt_state_root=args.act_scales, world_size=8)
-            scales = get_scales_from_amax(start_layer=start_layer,
-                                          end_layer=end_layer,
-                                          renamed_state_dict=renamed_state_dict)
-            new_safetensors.update(scales)
-
         file_name = get_file_name(start_layer)

---

60-64: Fix incorrect dictionary key check

Line 62 incorrectly uses .items() when checking if a key exists in the dictionary.

     for rank in range(world_size):
         for key, amax in state_dict_list[rank].items():
-            if key in merged_state_dict.items():
+            if key in merged_state_dict:
                 amax = torch.max(amax.to(0), merged_state_dict[key].to(0))
             merged_state_dict[key] = amax.to(0)

---

66-78: Remove duplicate keys in mapping dictionary

The mapping dictionary has duplicate keys for "ffn.shared_experts" (lines 70-72), which means only the last value will be retained.

     mapping = {
         "ffn.shared_experts.w1": "mlp.shared_experts.gate_proj",
         "ffn.shared_experts.w2": "mlp.shared_experts.down_proj",
         "ffn.shared_experts.w3": "mlp.shared_experts.up_proj",
-        "ffn.shared_experts": "mlp.shared_experts",
-        "ffn.shared_experts": "mlp.shared_experts",
-        "ffn.shared_experts": "mlp.shared_experts",
         "ffn.w1": "mlp.gate_proj",
         "ffn.w2": "mlp.down_proj",
         "ffn.w3": "mlp.up_proj",
         "head": "lm_head",
         "attn": "self_attn",
     }

---

288-288: Use num_layer variable instead of hardcoded value

Line 288 uses a hardcoded value of 61 for the layer count, which should use the num_layer variable for consistency and maintainability.

-        for l in range(61):
+        for l in range(num_layer):

examples/quantization/quantize_mixed_precision_moe.py (2)

269-272: Consider preserving the directory structure more explicitly

The current implementation copies the entire act_scales directory into output_dir. This might create an unexpected nested structure if act_scales has a complex path.

Consider copying to a specific subdirectory or just the contents:

 if os.path.isdir(args.act_scales):
-    shutil.copytree(args.act_scales, output_dir, dirs_exist_ok=True)
+    # Option 1: Copy to a specific subdirectory
+    act_scales_dest = os.path.join(output_dir, "act_scales")
+    shutil.copytree(args.act_scales, act_scales_dest, dirs_exist_ok=True)
+    # Option 2: Copy only the contents (if you want flat structure)
+    # for item in os.listdir(args.act_scales):
+    #     src = os.path.join(args.act_scales, item)
+    #     dst = os.path.join(output_dir, item)
+    #     if os.path.isdir(src):
+    #         shutil.copytree(src, dst, dirs_exist_ok=True)
+    #     else:
+    #         shutil.copy2(src, dst)
 else:
     shutil.copy(args.act_scales, output_dir)

---

52-56: Consider using Python's logging module instead of print statements

For better logging control and production readiness, consider using Python's logging module.

Add at the top of the file:

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

Then replace print statements:

-            print(f"WARNING: amax file not found: {amax_file}")
+            logger.warning(f"amax file not found: {amax_file}")

     if not state_dict_list:
-        print("ERROR: No amax files loaded!")
+        logger.error("No amax files loaded!")
         return {}

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro