import os # Must pop before PyTorch imports os.environ.pop('NCCL_MIN_NCHANNELS', None) os.environ.pop('NCCL_MAX_NCHANNELS', None) os.environ.pop('NCCL_NVLS_ENABLE', None) import argparse import math import random import torch import torch.distributed as dist import deep_ep from deep_ep.utils.envs import init_dist, dist_print, get_rdma_gbs from deep_ep.utils.testing import bench_kineto def generate_stress_ops(rank_idx: int, num_ranks: int, num_sends: int, shape: tuple, num_max_inflight_tensors: int): assert num_max_inflight_tensors > 0 send_times = {(s, d): [] for s in range(num_ranks) for d in range(num_ranks) if s != d} recv_times = {(s, d): [] for s in range(num_ranks) for d in range(num_ranks) if s != d} for _ in range(num_sends): src_rank_idx = random.randint(0, num_ranks - 1) dst_rank_idx = (src_rank_idx + (1 if random.randint(0, 1) else -1)) % num_ranks st = random.randint(0, 10 ** 8) rt = st + random.randint(1, 3 * 10 ** 6) send_times[(src_rank_idx, dst_rank_idx)].append(st) recv_times[(src_rank_idx, dst_rank_idx)].append(rt) ops = [] for (src_rank_idx, dst_rank_idx) in send_times: n = len(send_times[(src_rank_idx, dst_rank_idx)]) sorted_send = sorted(send_times[(src_rank_idx, dst_rank_idx)]) sorted_recv = sorted(recv_times[(src_rank_idx, dst_rank_idx)]) for i in range(n): # Keep each send after its previous `num_max_inflight_tensors` recv to avoid deadlock if i >= num_max_inflight_tensors: sorted_send[i] = max( sorted_send[i], sorted_recv[i - num_max_inflight_tensors] + 1) sorted_recv[i] = max(sorted_recv[i], sorted_send[i] + 1) tensor = torch.randn(shape, dtype=torch.bfloat16, device='cuda') if src_rank_idx == rank_idx: ops.append(('send', sorted_send[i], dst_rank_idx, i, tensor)) if dst_rank_idx == rank_idx: ops.append(('recv', sorted_recv[i], src_rank_idx, i, tensor)) ops.sort(key=lambda x: (x[1], x[3])) return ops # noinspection PyShadowingNames @torch.inference_mode() def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) shape = (args.num_tokens, args.hidden) num_max_tensor_bytes = math.prod(shape) * 2 num_max_inflight_tensors = args.num_max_inflight_tensors buffer = deep_ep.PPBuffer( group, num_max_tensor_bytes=num_max_tensor_bytes, num_max_inflight_tensors=num_max_inflight_tensors, explicitly_destroy=True) # Print configs assert num_ranks > 1 dist_print(f'Config:\n' f' > Ranks: {num_ranks}\n' f' > Shape: {shape}\n' f' > Max inflight tensors: {num_max_inflight_tensors}\n', once_in_node=True) # Run stress tests dist_print('Running stress tests:', once_in_node=True) for seed in range(args.num_stress_iterations): dist_print(f' > Testing with {seed=} ...', once_in_node=True) torch.manual_seed(42 + seed) random.seed(42 + seed) ops = generate_stress_ops(rank_idx, num_ranks, args.num_sends, shape, num_max_inflight_tensors) prev = 0 for j, (op, timestamp, peer, _, tensor) in enumerate(ops): if op == 'send': buffer.send(tensor, peer) else: result = torch.empty_like(tensor) buffer.recv(result, peer) assert torch.equal(result, tensor), \ f'Rank {rank_idx}: mismatch at op {j}' if timestamp > prev: torch.cuda._sleep(int((timestamp - prev) / 10 ** 8 * args.num_sleep_cycles)) prev = timestamp dist_print(' > All stress tests passed', once_in_node=True) dist_print(once_in_node=True) # Profiling dist_print('Profiling PP send and recv:', once_in_node=True) num_approx_rdma_cycles = int(num_max_tensor_bytes * 2 / get_rdma_gbs() * 1.5) def get_trace_path(prefix: str): return (None if not args.dump_profile_traces else f'{args.dump_profile_traces}/{prefix}_rank{rank_idx}.json') for hide_rdma_latency in (True, False): for num_concurrent in range(1, min(3, num_max_inflight_tensors) + 1): send_tensors = [torch.randn(shape, dtype=torch.bfloat16, device='cuda') for _ in range(num_concurrent)] recv_tensors = [torch.empty(shape, dtype=torch.bfloat16, device='cuda') for _ in range(num_concurrent)] def loop(_hide_rdma_latency=hide_rdma_latency): torch.zeros((131072, 32768), dtype=torch.int, device='cuda') for t in send_tensors: buffer.send(t, (rank_idx + 1) % num_ranks) if _hide_rdma_latency: torch.cuda._sleep(num_approx_rdma_cycles * num_concurrent) for t in recv_tensors: buffer.recv(t, (rank_idx - 1) % num_ranks) send_t, recv_t = bench_kineto( loop, kernel_names=('send_impl', 'recv_impl'), barrier_comm_profiling=True, trace_path=get_trace_path(f'pp_{num_concurrent}_{hide_rdma_latency}')) dist_print( f' > PP: {rank_idx:3}/{num_ranks:3} | ' f'hide={int(hide_rdma_latency)}, concurrent={num_concurrent} | ' f'send: {send_t * 1e6:.3f} us, ' f'{2 * num_max_tensor_bytes / send_t / 1e9:.3f} GB/s | ' f'recv: {recv_t * 1e6:.3f} us, ' f'{(2 if hide_rdma_latency else 1) * num_max_tensor_bytes / recv_t / 1e9:.3f} GB/s') dist_print(once_in_node=True) # Destroy the runtime and communication group buffer.destroy() dist.destroy_process_group() if __name__ == '__main__': parser = argparse.ArgumentParser(description='Test PP send/recv kernels') parser.add_argument('--num-processes', type=int, default=4) parser.add_argument('--num-tokens', type=int, default=4096) parser.add_argument('--hidden', type=int, default=7168) parser.add_argument('--num-max-inflight-tensors', type=int, default=4) parser.add_argument('--num-stress-iterations', type=int, default=4) parser.add_argument('--num-sends', type=int, default=128) parser.add_argument('--num-sleep-cycles', type=int, default=10 ** 7) parser.add_argument('--dump-profile-traces', type=str, default='') args = parser.parse_args() if args.dump_profile_traces: os.makedirs(args.dump_profile_traces, exist_ok=True) torch.multiprocessing.spawn(test, args=(args.num_processes, args), nprocs=args.num_processes)