diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index fdfdee9b1c..8a147f9c0d 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -38,6 +38,13 @@ def setup_pytorch_extension( # Source files sources = all_files_in_dir(Path(csrc_source_files), name_extension="cpp") + build_native_cp_transport = bool(int(os.getenv("NVTE_WITH_NCCL_DEVICE_CP", "0"))) + if build_native_cp_transport: + sources.extend( + path + for path in all_files_in_dir(Path(csrc_source_files), name_extension="cu") + if path.name == "cp_native_transport.cu" + ) # Header files include_dirs = get_cuda_include_dirs() @@ -87,16 +94,37 @@ def setup_pytorch_extension( libraries.append("nvshmem_host") cxx_flags.append("-DNVTE_ENABLE_NVSHMEM") + extra_compile_args = {"cxx": cxx_flags} + if build_native_cp_transport: + nvcc_flags = ["-O3", "-std=c++17"] + nccl_home = os.getenv("NCCL_HOME") + nccl_include_dir = os.getenv("NVTE_NCCL_INCLUDE_DIR") + nccl_library_dir = os.getenv("NVTE_NCCL_LIBRARY_DIR") + if nccl_home: + nccl_home = Path(nccl_home) + nccl_include_dir = nccl_include_dir or str(nccl_home / "include") + nccl_library_dir = nccl_library_dir or str(nccl_home / "lib") + if nccl_include_dir: + include_dirs.append(Path(nccl_include_dir)) + if nccl_library_dir: + library_dirs.append(Path(nccl_library_dir)) + libraries.append("nccl") + cxx_flags.append("-DNVTE_WITH_NCCL_DEVICE_CP") + nvcc_flags.append("-DNVTE_WITH_NCCL_DEVICE_CP") + extra_compile_args["nvcc"] = nvcc_flags + # Construct PyTorch CUDA extension sources = [str(path) for path in sources] include_dirs = [str(path) for path in include_dirs] - from torch.utils.cpp_extension import CppExtension + from torch.utils.cpp_extension import CppExtension, CUDAExtension + + extension_cls = CUDAExtension if build_native_cp_transport else CppExtension - return CppExtension( + return extension_cls( name="transformer_engine_torch", sources=[str(src) for src in sources], include_dirs=[str(inc) for inc in include_dirs], - extra_compile_args={"cxx": cxx_flags}, + extra_compile_args=extra_compile_args, libraries=[str(lib) for lib in libraries], library_dirs=[str(lib_dir) for lib_dir in library_dirs], ) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 0f36a8816d..cc099808ea 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -5,15 +5,38 @@ import os import sys import logging +import copy from contextlib import nullcontext import torch import torch.distributed as dist from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( get_cu_seqlens_on_cp_rank, ) +from transformer_engine.pytorch.attention.native_cp_transport import ( + destroy_native_cp_transport, + initialize_native_cp_transport, + set_native_cp_parent_group, +) from transformer_engine.pytorch.attention.dot_product_attention.utils import combine_and_quantize import transformer_engine_torch as tex from test_attention_with_cp import model_configs_flash_attn, model_configs_fused_attn + + +class _LogicalCPGroup: + """Minimal topology descriptor used by native-transport tests.""" + + def __init__(self, ranks, rank): + self.ranks = tuple(ranks) + self.cp_size = len(self.ranks) + self.cp_rank = self.ranks.index(rank) + + def size(self): + return self.cp_size + + def rank(self): + return self.cp_rank + + from transformer_engine.pytorch import ( autocast, DotProductAttention, @@ -180,6 +203,9 @@ def run_dpa_with_cp( scaling_mode="delayed", f16_O="False", is_training="True", + native_cp_transport="False", + logical_cp_ring="False", + max_seqlen=None, log_level=logging.WARNING, ): """Test DotProductAttention module with context parallelism""" @@ -202,6 +228,14 @@ def run_dpa_with_cp( if kernel_backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" config = model_configs_fused_attn[model] + config = copy.deepcopy(config) + native_cp_transport = native_cp_transport == "True" + logical_cp_ring = logical_cp_ring == "True" + if logical_cp_ring and not native_cp_transport: + raise ValueError("logical_cp_ring requires native_cp_transport=True") + if max_seqlen is not None: + config.max_seqlen_q = int(max_seqlen) + config.max_seqlen_kv = int(max_seqlen) assert config.attn_mask_type in [ "causal", "no_mask", @@ -229,6 +263,14 @@ def run_dpa_with_cp( cp_comm_ranks = range(world_size) assert rank in cp_comm_ranks cp_comm_group = dist.new_group(cp_comm_ranks, backend="nccl") + cp_group = cp_comm_group + cp_rank = rank + cp_global_ranks = tuple(cp_comm_ranks) + if logical_cp_ring: + offsets = [0, *range(1, world_size, 2), *reversed(range(2, world_size, 2))] + cp_global_ranks = tuple(cp_comm_ranks[offset] for offset in offsets) + cp_group = _LogicalCPGroup(cp_global_ranks, rank) + cp_rank = cp_group.rank() if cp_comm_type == "a2a+p2p": assert world_size % 2 == 0, ( "{cp_comm_type=} requires world_size % 2 = 0 as it assumes the a2a level has cp_size" @@ -390,17 +432,17 @@ def run_dpa_with_cp( ) for x in [q_, k_, v_, dout_] ] - seq_idx = torch.tensor([rank, 2 * world_size - rank - 1], device=q_.device) + seq_idx = torch.tensor([cp_rank, 2 * world_size - cp_rank - 1], device=q_.device) q_, k_, v_, dout_ = [x.index_select(seq_dim, seq_idx) for x in [q_, k_, v_, dout_]] q_, k_, v_, dout_ = [ x.view(*x.shape[:seq_dim], -1, *x.shape[(seq_dim + 2) :]) for x in [q_, k_, v_, dout_] ] elif qkv_format == "thd": seq_idx_q = tex.thd_get_partitioned_indices( - cu_seqlens_q_padded, q_.shape[0], world_size, rank + cu_seqlens_q_padded, q_.shape[0], world_size, cp_rank ) seq_idx_kv = tex.thd_get_partitioned_indices( - cu_seqlens_kv_padded, k_.shape[0], world_size, rank + cu_seqlens_kv_padded, k_.shape[0], world_size, cp_rank ) q_, dout_ = [x.index_select(0, seq_idx_q) for x in [q_, dout_]] k_, v_ = [x.index_select(0, seq_idx_kv) for x in [k_, v_]] @@ -438,10 +480,20 @@ def run_dpa_with_cp( bias_ = bias_.index_select(seq_q_dim, bias_seq_idx) bias_ = bias_.view(*shape_before_seq, -1, seq_kv_size) bias_.requires_grad = True + + if native_cp_transport: + kv_bytes = (k_.numel() + v_.numel()) * k_.element_size() + pair_bytes = 2 * kv_bytes + initialize_native_cp_transport( + cp_comm_group, + ((pair_bytes + 255) // 256) * 256 + pair_bytes, + ) + if logical_cp_ring: + set_native_cp_parent_group(cp_group, cp_comm_group) # set up environment core_attn.set_context_parallel_group( - cp_comm_sub_groups if cp_comm_type == "a2a+p2p" else cp_comm_group, - cp_comm_ranks, + cp_comm_sub_groups if cp_comm_type == "a2a+p2p" else cp_group, + cp_global_ranks, torch.cuda.Stream(), cp_comm_type, ) @@ -562,7 +614,7 @@ def run_dpa_with_cp( dq_, dk_, dv_, out_ = [dq_, dk_, dv_, out_] cu_seqlens_q_padded = cu_seqlens_q_padded // world_size cu_seqlens_q = get_cu_seqlens_on_cp_rank( - cu_seqlens_q, cu_seqlens_q_padded, world_size, rank, True, True + cu_seqlens_q, cu_seqlens_q_padded, world_size, cp_rank, True, True ) cu_pads_q = cu_seqlens_q_padded - cu_seqlens_q num_pads_q = cu_pads_q[1:] - cu_pads_q[:-1] @@ -582,7 +634,7 @@ def run_dpa_with_cp( ) cu_seqlens_kv_padded = cu_seqlens_kv_padded // world_size cu_seqlens_kv = get_cu_seqlens_on_cp_rank( - cu_seqlens_kv, cu_seqlens_kv_padded, world_size, rank, True, True + cu_seqlens_kv, cu_seqlens_kv_padded, world_size, cp_rank, True, True ) cu_pads_kv = cu_seqlens_kv_padded - cu_seqlens_kv num_pads_kv = cu_pads_kv[1:] - cu_pads_kv[:-1] @@ -733,6 +785,9 @@ def run_dpa_with_cp( ) logging.info(f"[Rank {rank}] CP vs no-CP: {names[i]} matches") + if native_cp_transport: + destroy_native_cp_transport(cp_comm_group) + # destroy distribution group dist.destroy_process_group() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index a6a8b0b26a..e271b1ad0b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -48,7 +48,10 @@ META_QKV, ) from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager -from transformer_engine.pytorch.distributed import get_distributed_world_size +from transformer_engine.pytorch.distributed import ( + get_distributed_world_size, + is_logical_process_group, +) from transformer_engine.pytorch.jit import no_torch_dynamo from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( attn_forward_func_with_cp, @@ -757,7 +760,7 @@ def forward( ), f"FlashAttention does not support qkv_layout = {qkv_layout}!" cp_size = 1 - if isinstance(cp_group, dist_group_type): + if isinstance(cp_group, dist_group_type) or is_logical_process_group(cp_group): cp_size = get_distributed_world_size(cp_group) elif isinstance(cp_group, list): for group in cp_group: @@ -1828,7 +1831,7 @@ def forward( ), f"FusedAttention does not support qkv_layout = {qkv_layout}!" cp_size = 1 - if isinstance(cp_group, dist_group_type): + if isinstance(cp_group, dist_group_type) or is_logical_process_group(cp_group): cp_size = get_distributed_world_size(cp_group) elif isinstance(cp_group, list): for group in cp_group: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 030b1d9cdc..eb924635a3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -25,6 +25,7 @@ from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.jit import jit_fuser from transformer_engine.pytorch.graph import is_graph_capturing +from transformer_engine.pytorch.attention.native_cp_transport import get_native_cp_transport from transformer_engine.pytorch.constants import ( dist_group_type, TE_DType, @@ -33,6 +34,7 @@ get_distributed_world_size, get_distributed_rank, gather_along_first_dim, + is_logical_process_group, reduce_scatter_along_first_dim, ) @@ -60,10 +62,24 @@ def flash_attn_p2p_communicate( - rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm + rank, + send_tensor, + send_dst, + recv_tensor, + recv_src, + cp_group, + batch_p2p_comm, + native_channel=0, ): """Point-to-point communications of KV and dKV in Attention with context parallelism""" send_recv_ops = [] + native_transport = get_native_cp_transport(cp_group) + if native_transport is not None: + return [ + native_transport.send_recv( + send_tensor, send_dst, recv_tensor, recv_src, channel=native_channel + ) + ] if batch_p2p_comm: if rank % 2 == 0: @@ -1557,11 +1573,22 @@ def forward( # synchronize fwd results correction across steps fwd_results_correction_done = torch.cuda.Event() - p2p_comm_buffers = [None for _ in range(cp_size)] k_shape = k.shape k_numel = k.numel() v_shape = v.shape - p2p_comm_buffers[0] = torch.cat((k.view(-1), v.view(-1)), dim=-1) + p2p_shape = (k_numel + v.numel(),) + native_transport = get_native_cp_transport(cp_group) + if native_transport is not None and fp8: + raise RuntimeError("Native CP transport does not support FP8 attention yet") + if native_transport is None: + p2p_comm_buffers = [None for _ in range(cp_size)] + p2p_comm_buffers[0] = torch.cat((k.view(-1), v.view(-1)), dim=-1) + else: + native_pair = native_transport.attention_buffer_pair(p2p_shape, k.dtype) + arena_buffers = (native_pair[0][0], native_pair[0][1], native_pair[1][0]) + p2p_comm_buffers = [arena_buffers[i % 3] for i in range(cp_size)] + p2p_comm_buffers[0][:k_numel].copy_(k.view(-1)) + p2p_comm_buffers[0][k_numel:].copy_(v.view(-1)) send_recv_reqs = [[], []] # P2P communication and compute: each rank has cp_size steps @@ -1576,7 +1603,8 @@ def forward( req.wait() if i < (cp_size - 1): - p2p_comm_buffers[i + 1] = torch.empty_like(p2p_comm_buffers[i]) + if native_transport is None: + p2p_comm_buffers[i + 1] = torch.empty_like(p2p_comm_buffers[i]) send_recv_reqs[i % 2] = flash_attn_p2p_communicate( rank, p2p_comm_buffers[i], @@ -1585,6 +1613,7 @@ def forward( recv_src, cp_group, batch_p2p_comm, + native_channel=0, ) kv_inputs[i % 2] = p2p_comm_buffers[i] @@ -1953,6 +1982,8 @@ def forward( kv_fp8 = None kv = p2p_comm_buffers[-1] + if native_transport is not None: + kv = kv.clone() if fp8: q_fp8, kv_fp8 = [ Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) @@ -2003,6 +2034,7 @@ def forward( ctx.cp_size_a2a = cp_size_a2a ctx.rank_a2a = rank_a2a ctx.cp_group = cp_group + ctx.native_cp_transport = native_transport ctx.cp_global_ranks = cp_global_ranks ctx.cp_stream = cp_stream ctx.dropout_p = dropout_p @@ -2238,10 +2270,15 @@ def backward(ctx, dout, *_args): if isinstance(dout, QuantizedTensorStorage): dout = dout.dequantize(dtype=bwd_nominal_dtype) dq_buffer = torch.empty_like(q) - p2p_comm_buffers = [ - torch.empty((2, *kv.shape), dtype=kv.dtype, device=kv.device), - torch.empty((2, *kv.shape), dtype=kv.dtype, device=kv.device), - ] + if ctx.native_cp_transport is None: + p2p_comm_buffers = [ + torch.empty((2, *kv.shape), dtype=kv.dtype, device=kv.device), + torch.empty((2, *kv.shape), dtype=kv.dtype, device=kv.device), + ] + else: + p2p_comm_buffers = list( + ctx.native_cp_transport.attention_buffer_pair(kv.shape, kv.dtype) + ) p2p_comm_buffers[0][0].copy_(kv) if ctx.use_fused_attention: bwd_output_te_dtype = TE_DType[bwd_nominal_dtype] @@ -2341,7 +2378,14 @@ def backward(ctx, dout, *_args): send_tensor = send_tensor[1] recv_tensor = recv_tensor[1] send_recv_reqs = flash_attn_p2p_communicate( - rank, send_tensor, send_dst, recv_tensor, recv_src, ctx.cp_group, batch_p2p_comm + rank, + send_tensor, + send_dst, + recv_tensor, + recv_src, + ctx.cp_group, + batch_p2p_comm, + native_channel=1, ) kv = p2p_comm_buffers[i % 2][0] @@ -2729,6 +2773,13 @@ def backward(ctx, dout, *_args): # [b, h, sq, 2*cp, sk//(2*cp)] -> [b, h, sq, sk] attn_dbias = attn_dbias.view(*attn_dbias.shape[:-2], -1) + if ctx.native_cp_transport is not None and not ctx.fp8: + dkv_out = torch.empty((dk.numel() + dv.numel(),), dtype=dk.dtype, device=dk.device) + dkv_out[: dk.numel()].copy_(dk.reshape(-1)) + dkv_out[dk.numel() :].copy_(dv.reshape(-1)) + dk = dkv_out[: dk.numel()].view_as(dk) + dv = dkv_out[dk.numel() :].view_as(dv) + nvtx_range_pop(f"{nvtx_label}") return ( @@ -3081,7 +3132,7 @@ def backward(ctx, dout, *_args): rank = get_distributed_rank(ctx.cp_group) (*saved_tensors,) = ctx.saved_tensors - (q, k, v, cu_seqlens_q, cu_seqlens_q_padded) = saved_tensors[:5] + q, k, v, cu_seqlens_q, cu_seqlens_q_padded = saved_tensors[:5] cu_seqlens_kv_per_step = saved_tensors[5:7] out_per_step = saved_tensors[7:9] softmax_lse_per_step = saved_tensors[9:11] @@ -4041,9 +4092,19 @@ def attn_forward_func_with_cp( cp_group = cp_group[0] cp_comm_type = "a2a" else: - assert isinstance( - cp_group, dist_group_type - ), f"cp_group must be {dist_group_type} type for {cp_comm_type=}!" + assert isinstance(cp_group, dist_group_type) or is_logical_process_group( + cp_group + ), f"cp_group must be a ProcessGroup or logical CP descriptor for {cp_comm_type=}!" + + if is_logical_process_group(cp_group): + if cp_comm_type != "p2p": + raise RuntimeError("Logical CP groups only support cp_comm_type='p2p'.") + if fp8: + raise RuntimeError("Native CP transport does not support FP8 attention yet.") + if softmax_type != "vanilla" or return_max_logit: + raise RuntimeError( + "Native CP transport currently supports vanilla softmax without max-logit output." + ) assert qkv_format in [ "bshd", @@ -4286,9 +4347,9 @@ def get_batch_on_this_cp_rank( raise ValueError(f"Unsupported qvk_format: {qvk_format}!") if qvk_format == "thd": # Get context parallel size and rank - cp_size = torch.distributed.get_world_size(group=cp_group) + cp_size = get_distributed_world_size(cp_group) if cp_size > 1: - cp_rank = torch.distributed.get_rank(group=cp_group) + cp_rank = get_distributed_rank(cp_group) # Calculate the chunk sizes for each sequence total_slices_of_any_sequence = 2 * cp_size diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 2dc42be18a..e6e7271efd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -40,6 +40,7 @@ ) from transformer_engine.pytorch.distributed import ( get_distributed_world_size, + is_logical_process_group, checkpoint, set_all_rng_states, CudaRNGStatesTracker, @@ -1236,7 +1237,9 @@ def forward( # adjust max_seqlen and cu_seqlens for CP cp_size = 1 - if isinstance(self.cp_group, dist_group_type): + if isinstance(self.cp_group, dist_group_type) or is_logical_process_group( + self.cp_group + ): cp_size = get_distributed_world_size(self.cp_group) elif isinstance(self.cp_group, list): for group in self.cp_group: diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index d95d327c78..f8672ab381 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -26,6 +26,7 @@ from transformer_engine.pytorch.distributed import ( get_distributed_world_size, get_distributed_rank, + is_logical_process_group, ) from transformer_engine.pytorch.attention.dot_product_attention import DotProductAttention @@ -609,7 +610,7 @@ def set_context_parallel_group( across each CP sub-group (e.g., via NVLink), then exchanging KV with p2p between sub-groups (e.g., via IBLink). """ - if isinstance(cp_group, dist_group_type): + if isinstance(cp_group, dist_group_type) or is_logical_process_group(cp_group): self.cp_size = get_distributed_world_size(cp_group) self.cp_rank = get_distributed_rank(cp_group) elif isinstance(cp_group, list): diff --git a/transformer_engine/pytorch/attention/native_cp_transport.py b/transformer_engine/pytorch/attention/native_cp_transport.py new file mode 100644 index 0000000000..038aaaa078 --- /dev/null +++ b/transformer_engine/pytorch/attention/native_cp_transport.py @@ -0,0 +1,174 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""NCCL Device API transport for context-parallel rings.""" + +import math +import os +import weakref +from typing import Iterable, Optional + +import torch + +import transformer_engine_torch as tex + +_group_transports = weakref.WeakKeyDictionary() + + +class _Work: + """Stream dependency compatible with ProcessGroup work handles.""" + + def __init__(self, handle: int, channel: int) -> None: + self.handle = handle + self.channel = channel + + def wait(self) -> bool: + """Wait until the native operation has completed.""" + if self.handle: + tex.cp_native_transport_wait(self.handle, self.channel) + self.handle = 0 + return True + + +class NativeCPTransport: + """One symmetric arena attached to a borrowed parent NCCL communicator.""" + + def __init__(self, parent_group, payload_bytes: int) -> None: + if not hasattr(tex, "cp_native_transport_create"): + raise RuntimeError("Transformer Engine was not built with native CP transport") + # This transport uses one shared context and VA-based signals. NCCL's + # larger GIN defaults only reserve unused device state. + os.environ.setdefault("NCCL_GIN_NCONTEXTS", "1") + os.environ.setdefault("NCCL_GIN_SIGNAL_POOL_SIZE", "64") + os.environ.setdefault("NCCL_GIN_COUNTER_POOL_SIZE", "64") + torch.distributed.barrier(group=parent_group, device_ids=[torch.cuda.current_device()]) + backend = parent_group._get_backend(torch.device("cuda")) + if not hasattr(backend, "_comm_ptr"): + raise RuntimeError("ProcessGroupNCCL does not expose _comm_ptr()") + + self._parent = weakref.ref(parent_group) + self._parent_rank = { + global_rank: rank + for rank, global_rank in enumerate( + torch.distributed.get_process_group_ranks(parent_group) + ) + } + self.handle, self.arena = tex.cp_native_transport_create( + int(backend._comm_ptr()), int(payload_bytes) + ) + self.handle = int(self.handle) + + @property + def payload_bytes(self) -> int: + """Return the usable size of the symmetric arena in bytes.""" + return 0 if self.arena is None else self.arena.numel() + + def _view(self, offset: int, shape: Iterable[int], dtype: torch.dtype) -> torch.Tensor: + shape = tuple(int(dim) for dim in shape) + size = math.prod(shape) * torch.empty((), dtype=dtype).element_size() + if offset + size > self.payload_bytes: + raise RuntimeError( + f"Native CP arena needs {offset + size} bytes, has {self.payload_bytes}" + ) + return self.arena.narrow(0, offset, size).view(dtype).view(shape) + + def attention_buffer_pair(self, shape: Iterable[int], dtype: torch.dtype): + """Return two contiguous ``[KV, dKV]`` work buffers.""" + shape = tuple(shape) + pair_bytes = 2 * math.prod(shape) * torch.empty((), dtype=dtype).element_size() + return self._view(0, (2, *shape), dtype), self._view( + (pair_bytes + 255) // 256 * 256, (2, *shape), dtype + ) + + def send_recv( + self, + send_tensor: torch.Tensor, + send_global_rank: int, + recv_tensor: torch.Tensor, + recv_global_rank: int, + channel: int = 0, + ) -> _Work: + """Launch one native send and receive through the parent communicator.""" + try: + send_peer = self._parent_rank[int(send_global_rank)] + recv_peer = self._parent_rank[int(recv_global_rank)] + except KeyError as error: + raise ValueError(f"Peer {error.args[0]} is outside the parent group") from error + channel = tex.cp_native_transport_send_recv( + self.handle, send_tensor, recv_tensor, send_peer, recv_peer, int(channel) + ) + return _Work(self.handle, int(channel)) + + def all_reduce(self, tensor: torch.Tensor, group, channel: int = 2) -> torch.Tensor: + """Ring sum over a dynamic-CP subgroup.""" + ranks = ( + group.ranks + if hasattr(group, "ranks") + else torch.distributed.get_process_group_ranks(group) + ) + result = tensor.contiguous().clone() + if len(ranks) == 1: + return result + + size = tensor.nbytes + send = self._view(0, tensor.shape, tensor.dtype) + recv = self._view((size + 255) // 256 * 256, tensor.shape, tensor.dtype) + send.copy_(tensor) + rank = group.rank() + dst, src = ranks[(rank + 1) % len(ranks)], ranks[(rank - 1) % len(ranks)] + for _ in range(len(ranks) - 1): + self.send_recv(send, dst, recv, src, channel).wait() + result.add_(recv) + send, recv = recv, send + return result + + def destroy(self) -> None: + """Collectively release the native transport and its symmetric arena.""" + if not self.handle: + return + parent = self._parent() + if parent is None: + raise RuntimeError("Parent ProcessGroup was released before its native transport") + torch.distributed.barrier(group=parent, device_ids=[torch.cuda.current_device()]) + tex.cp_native_transport_destroy(self.handle) + self.handle, self.arena = 0, None + torch.distributed.barrier(group=parent, device_ids=[torch.cuda.current_device()]) + + +def initialize_native_cp_transport(parent_group, payload_bytes: int) -> NativeCPTransport: + """Collectively initialize one transport per parent ProcessGroup.""" + transport = get_native_cp_transport(parent_group) + if transport is None: + transport = NativeCPTransport(parent_group, payload_bytes) + _group_transports[parent_group] = transport + elif transport.payload_bytes < payload_bytes: + raise RuntimeError( + f"Existing native CP arena has {transport.payload_bytes} bytes; " + f"requested {payload_bytes}" + ) + return transport + + +def set_native_cp_parent_group(cp_group, parent_group) -> None: + """Route a dynamic-CP subgroup through its parent's native transport.""" + transport = get_native_cp_transport(parent_group) + if transport is None: + raise RuntimeError("Native CP parent transport is not initialized") + _group_transports[cp_group] = transport + + +def get_native_cp_transport(group) -> Optional[NativeCPTransport]: + """Return the live native transport mapped to ``group``, if any.""" + transport = _group_transports.get(group) + return transport if transport is not None and transport.handle else None + + +def destroy_native_cp_transport(parent_group) -> None: + """Destroy the transport mapped to ``parent_group`` and remove its aliases.""" + transport = _group_transports.get(parent_group) + if transport is not None: + transport.destroy() + for group, mapped in list(_group_transports.items()): + if mapped is transport: + del _group_transports[group] diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index e4bc744e7e..bdce80f5a6 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -22,6 +22,16 @@ class CommOverlapP2P; namespace transformer_engine::pytorch { +#ifdef NVTE_WITH_NCCL_DEVICE_CP +std::tuple cp_native_transport_create(int64_t nccl_comm_ptr, + int64_t payload_bytes); +void cp_native_transport_destroy(int64_t handle); +int64_t cp_native_transport_send_recv(int64_t handle, at::Tensor send_tensor, + at::Tensor recv_tensor, int64_t send_peer, int64_t recv_peer, + int64_t channel); +void cp_native_transport_wait(int64_t handle, int64_t channel); +#endif + /*************************************************************************************************** * Router fusion **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/cp_native_transport.cu b/transformer_engine/pytorch/csrc/extensions/cp_native_transport.cu new file mode 100644 index 0000000000..b3cff5ef2e --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/cp_native_transport.cu @@ -0,0 +1,408 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../extensions.h" + +#ifdef NVTE_WITH_NCCL_DEVICE_CP + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace transformer_engine::pytorch { +namespace { + +#define NVTE_CP_NCCL_CHECK(call) \ + do { \ + const ncclResult_t result_ = (call); \ + TORCH_CHECK(result_ == ncclSuccess, #call, " failed: ", ncclGetErrorString(result_)); \ + } while (0) + +#define NVTE_CP_CUDA_CHECK(call) \ + do { \ + const cudaError_t result_ = (call); \ + TORCH_CHECK(result_ == cudaSuccess, #call, " failed: ", cudaGetErrorString(result_)); \ + } while (0) + +constexpr size_t kArenaAlignment = 256; +constexpr int kThreads = 256; +constexpr int kMaxCopyBlocks = 16; +constexpr int kNumChannels = 3; // forward, backward, aux-loss +constexpr int kGinContexts = 1; +constexpr int kGinQueueDepth = 8; +using Counter = unsigned long long; // NOLINT(runtime/int) + +size_t align_up(size_t value) { + return (value + kArenaAlignment - 1) / kArenaAlignment * kArenaAlignment; +} + +struct NativeCPTransport { + ncclComm_t comm = nullptr; + ncclDevComm dev_comm{}; + ncclWindow_t window = nullptr; + void *allocation = nullptr; + void *payload = nullptr; + size_t payload_bytes = 0; + size_t payload_offset = 0; + size_t signal_offset = 0; + size_t signal_shadow_offset = 0; + size_t ready_offset = 0; + int device = -1; + int rank = -1; + int nranks = 0; + bool dev_comm_created = false; + bool window_registered = false; + std::array streams{}; + std::array ready_events{}; + std::array done_events{}; + std::array outstanding{}; + std::vector ready_expected; + + ~NativeCPTransport() noexcept { + int previous_device = -1; + bool restore_device = false; + if (device >= 0) { + if (cudaGetDevice(&previous_device) == cudaSuccess) { + restore_device = previous_device != device; + } + if (previous_device != device) cudaSetDevice(device); + } + for (cudaStream_t stream : streams) { + if (stream != nullptr) cudaStreamSynchronize(stream); + } + for (cudaEvent_t event : ready_events) { + if (event != nullptr) cudaEventDestroy(event); + } + for (cudaEvent_t event : done_events) { + if (event != nullptr) cudaEventDestroy(event); + } + for (cudaStream_t stream : streams) { + if (stream != nullptr) cudaStreamDestroy(stream); + } + if (dev_comm_created) ncclDevCommDestroy(comm, &dev_comm); + if (window_registered) ncclCommWindowDeregister(comm, window); + if (allocation != nullptr) ncclMemFree(allocation); + if (restore_device) cudaSetDevice(previous_device); + } +}; + +NativeCPTransport *unwrap(int64_t handle) { + TORCH_CHECK(handle != 0, "Native CP transport handle is null"); + return reinterpret_cast(handle); +} + +__device__ __forceinline__ Counter system_load(Counter *ptr) { return atomicAdd_system(ptr, 0ULL); } + +__device__ __forceinline__ void copy_to_lsa_peer(void *dst_void, const void *src_void, + size_t bytes) { + auto *dst = static_cast(dst_void); + const auto *src = static_cast(src_void); + const uintptr_t packed = + reinterpret_cast(dst) | reinterpret_cast(src) | bytes; + if ((packed & (alignof(uint4) - 1)) == 0) { + auto *dst4 = reinterpret_cast(dst); + const auto *src4 = reinterpret_cast(src); + const size_t count4 = bytes / sizeof(uint4); + for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < count4; + index += blockDim.x * gridDim.x) { + dst4[index] = src4[index]; + } + return; + } + for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < bytes; + index += blockDim.x * gridDim.x) { + dst[index] = src[index]; + } +} + +__global__ void cp_native_prepare_kernel(ncclDevComm dev_comm, ncclWindow_t window, int rank, + int recv_peer, unsigned int channel, size_t ready_offset) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + const ncclTeam world = ncclTeamWorld(dev_comm); + const ncclTeam lsa = ncclTeamLsa(dev_comm); + const bool recv_is_lsa = ncclTeamRankIsMember(lsa, world, recv_peer); + const size_t ready_index = (static_cast(rank) * kNumChannels + channel) * sizeof(Counter); + + // A direct store has no implicit ncclRecv rendezvous. Advertise that this + // rank has finished with its receive buffer before the peer may overwrite it. + if (recv_is_lsa) { + const int recv_lsa_rank = ncclTeamRankToTeam(lsa, world, recv_peer); + auto *remote_ready = static_cast( + ncclGetLsaPointer(window, ready_offset + ready_index, recv_lsa_rank)); + __threadfence_system(); + atomicAdd_system(remote_ready, 1ULL); + } else { + ncclGin gin{dev_comm, 0}; + gin.signal(world, recv_peer, ncclGin_VASignalInc{window, ready_offset + ready_index}, + ncclCoopThread{}); + gin.flush(ncclCoopThread{}); + } +} + +__global__ void cp_native_send_recv_kernel(ncclDevComm dev_comm, ncclWindow_t window, + size_t send_offset, size_t recv_offset, size_t bytes, + int send_peer, int recv_peer, int rank, + unsigned int channel, Counter ready_expected, + size_t signal_offset, size_t signal_shadow_offset, + size_t ready_offset) { + const ncclTeam world = ncclTeamWorld(dev_comm); + const ncclTeam lsa = ncclTeamLsa(dev_comm); + const bool send_is_lsa = ncclTeamRankIsMember(lsa, world, send_peer); + const bool recv_is_lsa = ncclTeamRankIsMember(lsa, world, recv_peer); + + if (send_is_lsa) { + const size_t ready_index = + (static_cast(send_peer) * kNumChannels + channel) * sizeof(Counter); + if (threadIdx.x == 0) { + auto *local_ready = + static_cast(ncclGetLocalPointer(window, ready_offset + ready_index)); + while (system_load(local_ready) < ready_expected) { +#if __CUDA_ARCH__ >= 700 + __nanosleep(64); +#endif + } + } + __syncthreads(); + const int send_lsa_rank = ncclTeamRankToTeam(lsa, world, send_peer); + void *send_local = ncclGetLocalPointer(window, send_offset); + void *recv_remote = ncclGetLsaPointer(window, recv_offset, send_lsa_rank); + copy_to_lsa_peer(recv_remote, send_local, bytes); + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence_system(); + auto *remote_signal = static_cast(ncclGetLsaPointer( + window, + signal_offset + (static_cast(rank) * kNumChannels + channel) * sizeof(Counter), + send_lsa_rank)); + atomicAdd_system(remote_signal, 1ULL); + } + __syncthreads(); + } else if (blockIdx.x == 0) { + ncclGin gin{dev_comm, 0}; + const size_t ready_index = + (static_cast(send_peer) * kNumChannels + channel) * sizeof(Counter); + gin.waitSignal(ncclCoopCta{}, window, ready_offset + ready_index, ready_expected); + + const size_t completion_index = + (static_cast(rank) * kNumChannels + channel) * sizeof(Counter); + gin.put(world, send_peer, window, recv_offset, window, send_offset, bytes, + ncclGin_VASignalInc{window, signal_offset + completion_index}, ncclGin_None{}, + ncclCoopCta{}); + gin.flush(ncclCoopCta{}); + } + + if (recv_is_lsa) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + const size_t completion_index = + (static_cast(recv_peer) * kNumChannels + channel) * sizeof(Counter); + auto *local_signal = + static_cast(ncclGetLocalPointer(window, signal_offset + completion_index)); + auto *local_shadow = static_cast( + ncclGetLocalPointer(window, signal_shadow_offset + completion_index)); + const Counter expected = *local_shadow + gridDim.x; + while (system_load(local_signal) < expected) { +#if __CUDA_ARCH__ >= 700 + __nanosleep(64); +#endif + } + *local_shadow = expected; + __threadfence_system(); + } + __syncthreads(); + } else if (blockIdx.x == 0) { + ncclGin gin{dev_comm, 0}; + const size_t completion_index = + (static_cast(recv_peer) * kNumChannels + channel) * sizeof(Counter); + auto *local_shadow = static_cast( + ncclGetLocalPointer(window, signal_shadow_offset + completion_index)); + const Counter expected = *local_shadow + 1ULL; + gin.waitSignal(ncclCoopCta{}, window, signal_offset + completion_index, expected); + if (threadIdx.x == 0) { + *local_shadow = expected; + __threadfence_system(); + } + __syncthreads(); + } +} + +void validate_tensor(const NativeCPTransport &transport, const at::Tensor &tensor, + const char *name) { + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.get_device() == transport.device, name, " is on CUDA device ", + tensor.get_device(), " but transport uses ", transport.device); + const auto begin = reinterpret_cast(transport.payload); + const auto end = begin + transport.payload_bytes; + const auto tensor_begin = reinterpret_cast(tensor.data_ptr()); + const auto tensor_end = tensor_begin + tensor.nbytes(); + TORCH_CHECK(tensor_begin >= begin && tensor_end <= end, name, + " must be a view of the native CP transport arena"); +} + +} // namespace + +std::tuple cp_native_transport_create(int64_t nccl_comm_ptr, + int64_t payload_bytes) { + TORCH_CHECK(nccl_comm_ptr != 0, "nccl_comm_ptr must not be null"); + TORCH_CHECK(payload_bytes > 0, "payload_bytes must be positive"); + TORCH_CHECK(payload_bytes <= std::numeric_limits::max() - 2 * kArenaAlignment, + "payload_bytes is too large"); + + std::unique_ptr transport(new NativeCPTransport()); + transport->comm = reinterpret_cast(nccl_comm_ptr); + transport->device = c10::cuda::current_device(); + transport->payload_bytes = static_cast(payload_bytes); + + int runtime_version = 0; + NVTE_CP_NCCL_CHECK(ncclGetVersion(&runtime_version)); +#if NCCL_VERSION_CODE < NCCL_VERSION(2, 28, 7) + TORCH_CHECK(false, "Native CP transport requires NCCL 2.28.7 or newer"); +#endif + TORCH_CHECK(runtime_version == NCCL_VERSION_CODE, + "NCCL Device API GIN requires matching compile/runtime versions; ", "compiled with ", + NCCL_VERSION_CODE, ", loaded ", runtime_version); + + ncclCommProperties_t properties = NCCL_COMM_PROPERTIES_INITIALIZER; + NVTE_CP_NCCL_CHECK(ncclCommQueryProperties(transport->comm, &properties)); + TORCH_CHECK(properties.deviceApiSupport, + "The parent NCCL communicator does not support NCCL Device API"); + transport->rank = properties.rank; + transport->nranks = properties.nRanks; + const int lsa_size = ncclTeamLsa(transport->comm).nRanks; + const bool needs_gin = transport->nranks != lsa_size; + TORCH_CHECK(transport->nranks == lsa_size || properties.ginType != NCCL_GIN_TYPE_NONE, + "The parent communicator spans multiple LSA domains but GIN is unavailable"); + + const size_t peer_channel_bytes = + static_cast(transport->nranks) * kNumChannels * sizeof(Counter); + transport->signal_offset = 0; + transport->signal_shadow_offset = align_up(transport->signal_offset + peer_channel_bytes); + transport->ready_offset = align_up(transport->signal_shadow_offset + peer_channel_bytes); + transport->payload_offset = align_up(transport->ready_offset + peer_channel_bytes); + TORCH_CHECK( + transport->payload_bytes <= std::numeric_limits::max() - transport->payload_offset, + "Native CP transport allocation size overflow"); + const size_t allocation_bytes = transport->payload_offset + transport->payload_bytes; + + NVTE_CP_NCCL_CHECK(ncclMemAlloc(&transport->allocation, allocation_bytes)); + NVTE_CP_CUDA_CHECK(cudaMemset(transport->allocation, 0, transport->payload_offset)); + NVTE_CP_NCCL_CHECK(ncclCommWindowRegister(transport->comm, transport->allocation, + allocation_bytes, &transport->window, + NCCL_WIN_DEFAULT)); + transport->window_registered = true; + + ncclDevCommRequirements_t requirements = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER; + if (needs_gin) { +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 7) + requirements.ginContextCount = kGinContexts; + requirements.ginQueueDepth = kGinQueueDepth; + requirements.ginConnectionType = NCCL_GIN_CONNECTION_FULL; +#else + TORCH_CHECK(false, "Multi-node native CP transport requires NCCL 2.29.7 or newer"); +#endif + } + NVTE_CP_NCCL_CHECK(ncclDevCommCreate(transport->comm, &requirements, &transport->dev_comm)); + transport->dev_comm_created = true; + + int least_priority = 0; + int greatest_priority = 0; + NVTE_CP_CUDA_CHECK(cudaDeviceGetStreamPriorityRange(&least_priority, &greatest_priority)); + transport->ready_expected.resize(transport->nranks * kNumChannels, 0); + for (int channel = 0; channel < kNumChannels; ++channel) { + NVTE_CP_CUDA_CHECK(cudaStreamCreateWithPriority(&transport->streams[channel], + cudaStreamNonBlocking, greatest_priority)); + NVTE_CP_CUDA_CHECK( + cudaEventCreateWithFlags(&transport->ready_events[channel], cudaEventDisableTiming)); + NVTE_CP_CUDA_CHECK( + cudaEventCreateWithFlags(&transport->done_events[channel], cudaEventDisableTiming)); + } + + transport->payload = static_cast(transport->allocation) + transport->payload_offset; + auto options = + at::TensorOptions().device(at::Device(at::kCUDA, transport->device)).dtype(at::kByte); + at::Tensor arena = at::from_blob(transport->payload, {payload_bytes}, [](void *) {}, options); + const auto handle = reinterpret_cast(transport.release()); + return {handle, arena}; +} + +void cp_native_transport_destroy(int64_t handle) { delete unwrap(handle); } + +int64_t cp_native_transport_send_recv(int64_t handle, at::Tensor send_tensor, + at::Tensor recv_tensor, int64_t send_peer, int64_t recv_peer, + int64_t channel) { + NativeCPTransport *transport = unwrap(handle); + at::cuda::CUDAGuard device_guard(at::Device(at::kCUDA, transport->device)); + TORCH_CHECK(channel >= 0 && channel < kNumChannels, + "channel is outside the transport channel range"); + TORCH_CHECK(send_peer >= 0 && send_peer < transport->nranks, + "send_peer is outside the parent communicator"); + TORCH_CHECK(recv_peer >= 0 && recv_peer < transport->nranks, + "recv_peer is outside the parent communicator"); + TORCH_CHECK(send_peer != transport->rank || recv_peer != transport->rank, + "Native CP send/recv is unnecessary when both peers are self"); + validate_tensor(*transport, send_tensor, "send_tensor"); + validate_tensor(*transport, recv_tensor, "recv_tensor"); + TORCH_CHECK(send_tensor.nbytes() == recv_tensor.nbytes(), + "send_tensor and recv_tensor must have the same byte size"); + TORCH_CHECK(!transport->outstanding[channel], "Native CP transport channel ", channel, + " is reused before its previous work is waited"); + + const auto base = reinterpret_cast(transport->allocation); + const size_t send_offset = reinterpret_cast(send_tensor.data_ptr()) - base; + const size_t recv_offset = reinterpret_cast(recv_tensor.data_ptr()) - base; + const size_t bytes = send_tensor.nbytes(); + TORCH_CHECK(bytes > 0, "Native CP transport does not accept an empty payload"); + const size_t bytes_per_copy_block = kThreads * sizeof(uint4); + const int copy_blocks = static_cast( + std::min(kMaxCopyBlocks, (bytes + bytes_per_copy_block - 1) / bytes_per_copy_block)); + + const cudaStream_t caller_stream = at::cuda::getCurrentCUDAStream().stream(); + const int channel_index = static_cast(channel); + Counter &ready_expected = + transport->ready_expected[static_cast(send_peer) * kNumChannels + channel_index]; + ++ready_expected; + NVTE_CP_CUDA_CHECK(cudaEventRecord(transport->ready_events[channel_index], caller_stream)); + NVTE_CP_CUDA_CHECK(cudaStreamWaitEvent(transport->streams[channel_index], + transport->ready_events[channel_index], 0)); + cp_native_prepare_kernel<<<1, 1, 0, transport->streams[channel_index]>>>( + transport->dev_comm, transport->window, transport->rank, static_cast(recv_peer), + static_cast(channel), transport->ready_offset); + NVTE_CP_CUDA_CHECK(cudaGetLastError()); + cp_native_send_recv_kernel<<streams[channel_index]>>>( + transport->dev_comm, transport->window, send_offset, recv_offset, bytes, + static_cast(send_peer), static_cast(recv_peer), transport->rank, + static_cast(channel), ready_expected, transport->signal_offset, + transport->signal_shadow_offset, transport->ready_offset); + NVTE_CP_CUDA_CHECK(cudaGetLastError()); + NVTE_CP_CUDA_CHECK( + cudaEventRecord(transport->done_events[channel_index], transport->streams[channel_index])); + transport->outstanding[channel_index] = true; + return channel; +} + +void cp_native_transport_wait(int64_t handle, int64_t channel) { + NativeCPTransport *transport = unwrap(handle); + at::cuda::CUDAGuard device_guard(at::Device(at::kCUDA, transport->device)); + TORCH_CHECK(channel >= 0 && channel < kNumChannels, + "channel is outside the transport channel range"); + const int channel_index = static_cast(channel); + TORCH_CHECK(transport->outstanding[channel_index], "Native CP transport channel ", channel, + " has no outstanding work"); + const cudaStream_t caller_stream = at::cuda::getCurrentCUDAStream().stream(); + NVTE_CP_CUDA_CHECK(cudaStreamWaitEvent(caller_stream, transport->done_events[channel_index], 0)); + transport->outstanding[channel_index] = false; +} + +} // namespace transformer_engine::pytorch + +#endif // NVTE_WITH_NCCL_DEVICE_CP diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 18da5d0e9f..4d39cbe499 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -135,6 +135,18 @@ void init_extension() { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m) +#ifdef NVTE_WITH_NCCL_DEVICE_CP + m.def("cp_native_transport_create", &transformer_engine::pytorch::cp_native_transport_create, + py::arg("nccl_comm_ptr"), py::arg("payload_bytes")); + m.def("cp_native_transport_destroy", &transformer_engine::pytorch::cp_native_transport_destroy, + py::arg("handle")); + m.def("cp_native_transport_send_recv", + &transformer_engine::pytorch::cp_native_transport_send_recv, py::arg("handle"), + py::arg("send_tensor"), py::arg("recv_tensor"), py::arg("send_peer"), py::arg("recv_peer"), + py::arg("channel") = 0); + m.def("cp_native_transport_wait", &transformer_engine::pytorch::cp_native_transport_wait, + py::arg("handle"), py::arg("channel") = 0); +#endif m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("output") = py::none(), py::arg("noop") = py::none()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index b80e58fe20..74ada0b94a 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -166,6 +166,8 @@ def set_tensor_model_parallel_attributes( @lru_cache def get_distributed_world_size(group: Optional[dist_group_type] = None) -> int: """Return world size for the distributed group.""" + if is_logical_process_group(group): + return int(group.cp_size) if not torch.distributed.is_initialized(): return 1 return torch.distributed.get_world_size(group=group) @@ -174,6 +176,8 @@ def get_distributed_world_size(group: Optional[dist_group_type] = None) -> int: @lru_cache def get_distributed_rank(group: Optional[dist_group_type] = None) -> int: """Return my rank for the distributed group.""" + if is_logical_process_group(group): + return int(group.cp_rank) if not torch.distributed.is_initialized(): raise RuntimeError( "torch.distributed is not initialized. Call torch.distributed.init_process_group() " @@ -182,6 +186,16 @@ def get_distributed_rank(group: Optional[dist_group_type] = None) -> int: return torch.distributed.get_rank(group=group) +def is_logical_process_group(group: Any) -> bool: + """Return whether ``group`` is a topology-only CP descriptor.""" + return ( + group is not None + and hasattr(group, "ranks") + and hasattr(group, "cp_size") + and hasattr(group, "cp_rank") + ) + + def initialize_affine_weight_gpu( weight: torch.Tensor, init_method: Callable,