fix cpu omp num threads set (#2755)

* fix cpu omp num threads set * fix OMP_NUM_THREADS * consider no-cpu usage * fix style
huggingface · May 10, 2024 · 0f0a856 · 0f0a856
1 parent 0292728
commit 0f0a856
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
@@ -1034,8 +1034,8 @@ def _validate_launch_command(args):
         defaults is not None and defaults.compute_environment != ComputeEnvironment.AMAZON_SAGEMAKER
     )
     if is_aws_env_disabled and args.num_cpu_threads_per_process is None:
-        args.num_cpu_threads_per_process = 1
-        if args.use_cpu and args.num_processes >= 1:
+        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
             )

diff --git a/src/accelerate/state.py b/src/accelerate/state.py
@@ -236,7 +236,10 @@ def __init__(self, cpu: bool = False, **kwargs):
                 kwargs["rank"] = dist_information.rank
                 kwargs["world_size"] = dist_information.world_size
 
-                if self.distributed_type == DistributedType.MULTI_CPU and get_int_from_env(["OMP_NUM_THREADS"], 0) > 0:
+                if (
+                    self.distributed_type == DistributedType.MULTI_CPU
+                    and get_int_from_env(["OMP_NUM_THREADS"], 0) == 0
+                ):
                     import psutil
 
                     num_cpu_threads_per_process = int(