Skip to content

Commit

Permalink
fix default value for rank size in cpu threads_per_process assignment…
Browse files Browse the repository at this point in the history
… logic (#3009)

* fix default value for rank size

* fix style

* apply int in case ratio is decimal

* style quality fix
  • Loading branch information
rbrugaro authored Aug 15, 2024
1 parent a452327 commit 00785cd
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/accelerate/commands/launch.py
Original file line number Diff line number Diff line change
Expand Up @@ -1126,7 +1126,8 @@ def _validate_launch_command(args):
args.num_cpu_threads_per_process = get_int_from_env(["OMP_NUM_THREADS"], 1)
if args.use_cpu and args.num_processes >= 1 and get_int_from_env(["OMP_NUM_THREADS"], 0) == 0:
local_size = get_int_from_env(
["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"],
max(int(args.num_processes / args.num_machines), 1),
)
threads_per_process = int(psutil.cpu_count(logical=False) / local_size)
if threads_per_process > 1:
Expand Down

0 comments on commit 00785cd

Please sign in to comment.