Fix device map in limited mem (#1190)

Author: tastelikefeet

.dev_scripts/dockerci.sh

@@ -3,8 +3,8 @@ MODELSCOPE_CACHE_DIR_IN_CONTAINER=/modelscope_cache
 CODE_DIR=$PWD
 CODE_DIR_IN_CONTAINER=/swift
 echo "$USER"
-gpus='0,1 2,3 4,5 6,7'
-cpu_sets='0-15 16-31 32-47 48-63'
+gpus='0,1 2,3'
+cpu_sets='0-15 16-31'
 cpu_sets_arr=($cpu_sets)
 is_get_file_lock=false
 CI_COMMAND=${CI_COMMAND:-bash .dev_scripts/ci_container_test.sh python tests/run.py --parallel 2 --run_config tests/run_config.yaml}

swift/llm/infer.py

@@ -140,7 +140,7 @@ def prepare_model_template(args: InferArguments,
     else:
         logger.info(f'device_count: {torch.cuda.device_count()}')
         if device_map is None:
-            device_map = 'auto'
+            device_map = 'auto' if torch.cuda.device_count() > 1 else 'cuda:0'
     if device_map == 'auto':
         model_kwargs['low_cpu_mem_usage'] = True
     model_kwargs['device_map'] = device_map

swift/llm/rlhf.py

@@ -49,6 +49,8 @@ def llm_rlhf(args: RLHFArguments) -> Dict[str, Any]:
         model_kwargs = {'low_cpu_mem_usage': True}
         if is_dist() and not is_ddp_plus_mp():
             model_kwargs['device_map'] = {'': local_rank}
+        elif torch.cuda.device_count() == 1:
+            model_kwargs['device_map'] = 'cuda:0'
         else:
             model_kwargs['device_map'] = 'auto'
 

swift/llm/sft.py

@@ -56,6 +56,8 @@ def llm_sft(args: SftArguments) -> Dict[str, Union[str, Any]]:
         model_kwargs = {'low_cpu_mem_usage': True}
         if is_dist() and not is_ddp_plus_mp():
             model_kwargs['device_map'] = {'': local_rank}
+        elif torch.cuda.device_count() == 1:
+            model_kwargs['device_map'] = 'cuda:0'
         elif not use_torchacc():
             model_kwargs['device_map'] = 'auto'