From 6b3bdd18348e9755c00b321f836f1ece40cfe1e8 Mon Sep 17 00:00:00 2001 From: Evgeny Fiksman Date: Sat, 5 Apr 2025 07:58:25 -0700 Subject: [PATCH] Adding NUMA support for pytorch (#3926) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/1014 X-link: https://github.com/pytorch/pytorch/pull/150597 Add entry point and environment variable to control NUMA binding / assigment for distributed pytorch runs (training). Differential Revision: D72321369 --- fbgemm_gpu/test/sparse/common.py | 2 +- fbgemm_gpu/test/test_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/fbgemm_gpu/test/sparse/common.py b/fbgemm_gpu/test/sparse/common.py index 8abddca752..02729f87a7 100644 --- a/fbgemm_gpu/test/sparse/common.py +++ b/fbgemm_gpu/test/sparse/common.py @@ -16,7 +16,7 @@ import fbgemm_gpu import torch from hypothesis import HealthCheck, settings -from torch._utils_internal import get_file_path_2 # @manual=//caffe2:utils_internal +from torch._utils_internal import get_file_path_2 # @manual=//caffe2:_utils_internal from torch.testing._internal.optests import generate_opcheck_tests # pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. diff --git a/fbgemm_gpu/test/test_utils.py b/fbgemm_gpu/test/test_utils.py index 522c822e57..b7e7a8e306 100644 --- a/fbgemm_gpu/test/test_utils.py +++ b/fbgemm_gpu/test/test_utils.py @@ -141,7 +141,7 @@ def decorator(test_class: unittest.TestCase) -> unittest.TestCase: if not has_optests(): return test_class import torch.testing._internal.optests as optests - from torch._utils_internal import ( # @manual=//caffe2:utils_internal + from torch._utils_internal import ( # @manual=//caffe2:_utils_internal get_file_path_2, )