diff --git a/deep_ep/utils/envs.py b/deep_ep/utils/envs.py index f6e34d98..2da65697 100644 --- a/deep_ep/utils/envs.py +++ b/deep_ep/utils/envs.py @@ -162,7 +162,9 @@ def check_nvlink_connections(group: dist.ProcessGroup) -> None: # noinspection PyTypeChecker devices = os.environ.get('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7').strip(',').split(',') physical_device_idx = int(devices[torch.cuda.current_device()]) - physical_device_indices = [0, ] * group.size() + physical_device_indices = [ + 0, + ] * group.size() dist.all_gather_object(physical_device_indices, physical_device_idx, group) # Check whether they are all connected via NVLink @@ -202,8 +204,7 @@ def get_nvlink_gbs(factor: float = 0.9) -> float: """ # noinspection PyBroadException try: - result = subprocess.run(['nvidia-smi', 'nvlink', '-s'], - capture_output=True, text=True, check=True) + result = subprocess.run(['nvidia-smi', 'nvlink', '-s'], capture_output=True, text=True, check=True) output = result.stdout pattern = r'GPU \d+:.*?(?=^GPU \d+:|^$)' match = re.search(pattern, output, re.MULTILINE | re.DOTALL) @@ -242,16 +243,58 @@ def check_fast_rdma_atomic_support(nic_name: str = _DEFAULT_NIC_NAME) -> bool: return False +def _get_local_gpu_count() -> int: + """ + Get the number of GPUs/ranks sharing local RDMA NICs. + """ + + try: + result = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True) + count = len(re.findall(r'^GPU \d+:', result.stdout, re.MULTILINE)) if result.returncode == 0 else 0 + if count > 0: + return count + except Exception: + pass + + if torch.cuda.is_available(): + count = torch.cuda.device_count() + if count > 0: + return count + + return 1 + + +def _get_rdma_nic_prefix(nic_name: str) -> str: + """ + Get the RDMA NIC prefix used to count peer NIC devices. + """ + match = re.match(r'(.+_)\d+$', nic_name) + return match.group(1) if match else nic_name + + +def _get_active_rdma_nic_count(ibstat_output: str, nic_name: str) -> int: + """ + Count active RDMA NICs from `ibstat` output using the selected NIC prefix. + """ + nic_prefix = _get_rdma_nic_prefix(nic_name) + nic_pattern = rf'{re.escape(nic_prefix)}\d+' if nic_prefix != nic_name else re.escape(nic_name) + ca_blocks = re.findall(r"^CA '[^']+'.*?(?=^CA '|\Z)", ibstat_output, re.MULTILINE | re.DOTALL) + count = sum( + 1 for block in ca_blocks + if re.match(rf"^CA '{nic_pattern}'", block) and re.search(r'\bState:\s*Active\b', block) and re.search(r'\bRate:\s*\d+', block)) + return max(count, 1) + + @functools.lru_cache() def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float: """ - Get the RDMA bandwidth in GB/s, cached. + Get the per-GPU RDMA bandwidth in GB/s, cached. Arguments: nic_name: the NIC device name. Returns: - gbs: the RDMA bandwidth in GB/s (0 if detection fails). + gbs: the RDMA bandwidth in GB/s per local GPU/rank (0 if detection fails). """ # noinspection PyBroadException try: @@ -262,7 +305,18 @@ def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float: match = re.search(pattern, output, re.DOTALL) assert match rate = int(match.group(1)) - return rate / 8 + nic_count = _get_active_rdma_nic_count(output, nic_name) + # Try to detect 802.3ad LACP bonding slaves and multiply by slave count + result = subprocess.run(f'cat /sys/class/infiniband/{nic_name}/device/net/*/upper_*/bonding/slaves', + shell=True, + capture_output=True, + text=True) + slave_count = len(result.stdout.strip().split()) + if slave_count >= 2: + nic_count *= slave_count + gpu_count = _get_local_gpu_count() + + return rate * nic_count / gpu_count / 8 except Exception as e: print(f'Failed to get RDMA connection speed: {e}') return 0 diff --git a/setup.py b/setup.py index 81a1e266..78822d06 100644 --- a/setup.py +++ b/setup.py @@ -58,16 +58,17 @@ def get_package_version(): status_output = subprocess.check_output(status_cmd).decode('ascii').strip() if status_output: print(f'Warning: Git working directory is not clean. Uncommitted changes:\n{status_output}') - assert False, 'Git working directory is not clean' + raise AssertionError('Git working directory is not clean') cmd = ['git', 'rev-parse', '--short', 'HEAD'] revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip() - except: + except Exception: revision = '+local' return f'{public_version}{revision}' class CustomBuildPy(build_py): + def run(self): # Make clusters' cache setting default into `envs.py` self.generate_default_envs() @@ -98,9 +99,7 @@ def generate_default_envs(self): cxx_flags = ['-O3', '-Wno-deprecated-declarations', '-Wno-unused-variable', '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes'] nvcc_flags = ['-O3', '-Xcompiler', '-O3', '--extended-lambda', '--diag-suppress=128,2417'] sources = ['csrc/python_api.cpp', 'csrc/kernels/legacy/layout.cu', 'csrc/kernels/legacy/intranode.cu'] - include_dirs = [f'{current_dir}/deep_ep/include', - f'{current_dir}/third-party/fmt/include', - '/usr/local/cuda/include/cccl'] + include_dirs = [f'{current_dir}/deep_ep/include', f'{current_dir}/third-party/fmt/include', '/usr/local/cuda/include/cccl'] library_dirs = [] nvcc_dlink = [] extra_link_args = ['-lcuda'] @@ -136,7 +135,7 @@ def generate_default_envs(self): nvcc_flags.append('-DDISABLE_SM90_FEATURES') # Disable internode and low-latency kernels - assert False, 'Not implemented' + raise AssertionError('Not implemented') else: # Prefer H800 series os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '9.0') @@ -189,30 +188,29 @@ def generate_default_envs(self): if name in os.environ: persistent_envs.append((name, os.environ[name])) if len(persistent_envs) > 0: - print(f' > Persistent envs:') + print(' > Persistent envs:') for k, v in persistent_envs: print(f' > {k}: {v}') print() - setuptools.setup( - name='deep_ep', - version=get_package_version(), - packages=setuptools.find_packages(include=['deep_ep', 'deep_ep.*']), - package_data={ - 'deep_ep': [ - 'include/deep_ep/**/*', - ] - }, - ext_modules=[ - CUDAExtension(name='deep_ep._C', - include_dirs=include_dirs, - library_dirs=library_dirs, - sources=sources, - extra_compile_args=extra_compile_args, - extra_link_args=extra_link_args) - ], - cmdclass={ - 'build_ext': BuildExtension, - 'build_py': CustomBuildPy - } - ) + setuptools.setup(name='deep_ep', + version=get_package_version(), + packages=setuptools.find_packages(include=['deep_ep', 'deep_ep.*']), + package_data={'deep_ep': [ + 'include/deep_ep/**/*', + ]}, + install_requires=[ + "deepxtrace>=0.1.0", + ], + ext_modules=[ + CUDAExtension(name='deep_ep._C', + include_dirs=include_dirs, + library_dirs=library_dirs, + sources=sources, + extra_compile_args=extra_compile_args, + extra_link_args=extra_link_args) + ], + cmdclass={ + 'build_ext': BuildExtension, + 'build_py': CustomBuildPy + })