fix (#1383)

tastelikefeet · web-flow · commit 3e024e91d479 · 2024-07-13T13:46:35.000+08:00
diff --git a/swift/llm/rlhf.py b/swift/llm/rlhf.py
@@ -66,7 +66,7 @@ def llm_rlhf(args: RLHFArguments) -> Dict[str, Any]:
         assert len(args.device_max_memory) == n_gpu / local_world_size
         model_kwargs['max_memory'] = {
             i: mem
-            for i, mem in zip(list(range(local_rank, n_gpu, local_world_size)), args.device_max_memory)
+            for i, mem in zip(list(range(max(local_rank, 0), n_gpu, local_world_size)), args.device_max_memory)
         }
 
     # quantization
diff --git a/swift/llm/sft.py b/swift/llm/sft.py
@@ -65,7 +65,7 @@ def llm_sft(args: SftArguments) -> Dict[str, Union[str, Any]]:
         assert len(args.device_max_memory) == n_gpu / local_world_size
         model_kwargs['max_memory'] = {
             i: mem
-            for i, mem in zip(list(range(local_rank, n_gpu, local_world_size)), args.device_max_memory)
+            for i, mem in zip(list(range(max(local_rank, 0), n_gpu, local_world_size)), args.device_max_memory)
         }
 
     if args.quant_method == 'hqq':

Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ def llm_rlhf(args: RLHFArguments) -> Dict[str, Any]:`
`66`	`66`	`assert len(args.device_max_memory) == n_gpu / local_world_size`
`67`	`67`	`model_kwargs['max_memory'] = {`
`68`	`68`	`i: mem`
`69`		`- for i, mem in zip(list(range(local_rank, n_gpu, local_world_size)), args.device_max_memory)`
	`69`	`+ for i, mem in zip(list(range(max(local_rank, 0), n_gpu, local_world_size)), args.device_max_memory)`
`70`	`70`	`}`
`71`	`71`
`72`	`72`	`# quantization`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ def llm_sft(args: SftArguments) -> Dict[str, Union[str, Any]]:`
`65`	`65`	`assert len(args.device_max_memory) == n_gpu / local_world_size`
`66`	`66`	`model_kwargs['max_memory'] = {`
`67`	`67`	`i: mem`
`68`		`- for i, mem in zip(list(range(local_rank, n_gpu, local_world_size)), args.device_max_memory)`
	`68`	`+ for i, mem in zip(list(range(max(local_rank, 0), n_gpu, local_world_size)), args.device_max_memory)`
`69`	`69`	`}`
`70`	`70`
`71`	`71`	`if args.quant_method == 'hqq':`