[EZ] Fix set_torch_num_threads in multi-node. (#2164)

EugenHotaj · Felipe Mello · commit 8600c49ce565 · 2024-12-20T08:26:12.000-08:00
diff --git a/torchtune/training/_distributed.py b/torchtune/training/_distributed.py
@@ -111,7 +111,7 @@ def set_torch_num_threads() -> None:
     things like CPU affinity is set.
     """
     num_threads = os.cpu_count() // (
-        torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1
+        torch.cuda.device_count() if torch.cuda.is_available() else 1
     )
     torch.set_num_threads(num_threads)
     _log.info(f"Set intra op parallelism no. of threads to {num_threads}")

Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,7 @@ def set_torch_num_threads() -> None:`
`111`	`111`	`things like CPU affinity is set.`
`112`	`112`	`"""`
`113`	`113`	`num_threads = os.cpu_count() // (`
`114`		`- torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1`
	`114`	`+ torch.cuda.device_count() if torch.cuda.is_available() else 1`
`115`	`115`	`)`
`116`	`116`	`torch.set_num_threads(num_threads)`
`117`	`117`	`_log.info(f"Set intra op parallelism no. of threads to {num_threads}")`