[training] Adding NUMA support for pytorch (pytorch#150597)

Summary:
X-link: pytorch/FBGEMM#3926

X-link: facebookresearch/FBGEMM#1014


Add entry point and environment variable to control NUMA binding / assigment for distributed pytorch runs (training).

Test Plan: build and run tests for modified libraries locally

Differential Revision: D72321369

Author: efiks

Diff for: torch/_utils_internal.py

@@ -274,3 +274,8 @@ def record_chromium_event_internal(
     event: dict[str, Any],
 ):
     return None
+
+def maybe_enable_numa_binding(
+    device_type: Optional[str] = None, device_id: Optional[int] = None
+):
+    return None

Diff for: torch/distributed/distributed_c10d.py

@@ -44,7 +44,10 @@
     Store,
     Work,
 )
-from torch._utils_internal import set_pytorch_distributed_envs_from_justknobs
+from torch._utils_internal import (
+    maybe_enable_numa_binding,
+    set_pytorch_distributed_envs_from_justknobs,
+)
 from torch.monitor import _WaitCounter
 from torch.overrides import handle_torch_function, has_torch_function
 from torch.utils._typing_utils import not_none
@@ -1669,6 +1672,14 @@ def init_process_group(
         # default map below.
         backend = Backend.default_device_backend_map.get(device_id.type)
 
+
+    # If device_id is provide try to bind current process to the
+    # NUMA node attached to the device
+    if device_id is not None:
+        maybe_enable_numa_binding(device_type=device_id.type, device_id=device_id.index)
+    else:
+        maybe_enable_numa_binding()
+
     # If we still cannot figure it out, e.g.
     # >>> init_process_group()
     # we set it to `undefined` and rely on lazy init.