add device list to clear memory

auto_round/compressors/base.py

@@ -1459,12 +1459,12 @@ def _quantize_via_rtn_blockwise(self, all_to_quantized_module_names: list[str])
                 if is_fp8_model(self.model):
                     convert_fp8_model_to_16b_model(block, dtype=self.amp_dtype, device=self.device)
 
-                if is_auto_device_mapping(self.device_map):
+                if is_auto_device_mapping(self.device_map) and len(self.device_list) > 1:
                     set_auto_device_map_for_block_with_tuning(
                         block, self.device_map, input_ids, self.low_gpu_mem_usage, self.batch_size, self.device
                     )
                 # Dispatch model if needed
-                if len(self.device_list) > 0:
+                if len(self.device_list) > 1:
                     from accelerate.hooks import AlignDevicesHook, add_hook_to_module
 
                     for _, m in block.named_modules():
@@ -2498,7 +2498,7 @@ def _quantize_block(
                     set_module(block, n, new_layer)
         # card_0_in_high_risk indicates that card_0 memory is already in high usage (90%) w/o any weights
         # loss_device is used to calculate loss on the second device if available and card_0_in_high_risk
-        if self.device_map == "auto" or ((isinstance(self.device_map, str) and "," in self.device_map)):
+        if is_auto_device_mapping(self.device_map) and len(self.device_list) > 1:
             card_0_in_high_risk, loss_device = set_auto_device_map_for_block_with_tuning(
                 block, self.device_map, input_ids, self.low_gpu_mem_usage, self.batch_size, device
             )
@@ -2699,7 +2699,7 @@ def _quantize_block(
         )
         logger.info(dump_info)
         if self.low_gpu_mem_usage:
-            clear_memory()  # clear cached memory during training
+            clear_memory(self.device_list)  # clear cached memory during training
         if len(unquantized_layer_names) != 0:
             logger.info(f"{unquantized_layer_names} have not been quantized")
         with torch.no_grad():
@@ -2721,15 +2721,15 @@ def _quantize_block(
             if len(self.device_list) > 1:
                 accelerate.hooks.remove_hook_from_submodules(block)
             mv_module_from_gpu(block)
-            clear_memory(input_ids)
+            clear_memory(input_ids, self.device_list)
 
             return q_outputs, output
 
         else:
             if len(self.device_list) > 1:
                 accelerate.hooks.remove_hook_from_submodules(block)
             mv_module_from_gpu(block)
-            clear_memory(input_ids)
+            clear_memory(input_ids, self.device_list)
             return None, output
 
     def _split_inputs(self, inputs: dict) -> tuple[torch.Tensor, dict]: