fix bug

Signed-off-by: He, Xin3 <xin3.he@intel.com>

Author: xin3he

auto_round/compressors/utils.py

@@ -205,9 +205,10 @@ def collect_best_params(block, low_gpu_mem_usage: bool = False):
         if hasattr(m, "orig_layer"):
             params[n] = {}
             for key in m.params.keys():
-                params[n][key] = copy.deepcopy(m.params[key].data)
                 if low_gpu_mem_usage:
-                    params[n][key] = params[n][key].cpu()
+                    params[n][key] = m.params[key].data.cpu()
+                else:
+                    params[n][key] = copy.deepcopy(m.params[key].data)
     return params
 
 

auto_round/utils/device.py

@@ -855,10 +855,13 @@ def estimate_tuning_block_mem(
             else:
                 output_memory_gb = 0.0
 
+            if has_moe:
+                pparent_module = get_module(block, layer_name.rsplit(".", 2)[0]) if "." in layer_name else block
+                is_moe_expert = "expert" in layer_name.lower() and isinstance(pparent_module, torch.nn.ModuleList)
+            else:
+                is_moe_expert = False
+
             # memory * 2, because it contains grad tensor.
-            # Check if this is a MoE expert layer by layer name (e.g., "mlp.experts.0.gate_proj")
-            parent_module = get_module(block, layer_name.rsplit(".", 1)[0]) if "." in layer_name else block
-            is_moe_expert = "expert" in layer_name.lower() and isinstance(parent_module, torch.nn.ModuleList)
             layer_memory_dict[layer_name] = {
                 "param_memory": param_memory_gb * 2,
                 "output_memory": output_memory_gb * 2,