/tests/elastic/test_agrs.py

import argparse
import math
import random
import torch
import torch.distributed as dist
import numpy as np

import deep_ep
from deep_ep.utils.envs import init_dist, dist_print


def all_gather_ref(shape: tuple, rank_idx: int, num_ranks: int, round_idx: int = 0):
    ref_list = []
    for i in range(num_ranks):
        torch.manual_seed(42 + round_idx * 43 + i)
        ref_list.append(torch.randn(shape, dtype=torch.bfloat16, device='cuda'))
    return ref_list[rank_idx], torch.stack(ref_list, dim=0)


def generate_stress_ops(
    num_ops: int,
    num_max_inflight_agrs: int,
    shape: tuple,
    rank_idx: int,
    num_ranks: int,
) -> tuple[list[tuple], tuple[torch.Tensor], tuple[torch.Tensor]]:
    tensors, refs = zip(*(all_gather_ref(shape, rank_idx, num_ranks, round_idx=i) for i in range(num_ops)), strict=True)
    unprocessed = random.sample(range(num_ops), num_ops)
    inflight, ops = [], [('create_session', (-1,))]
    limit = num_max_inflight_agrs
    while unprocessed or inflight:
        max_g = min(len(unprocessed), limit)
        choices = []
        if max_g > 0:
            choices.append('ag')
        if inflight:
            choices.append('fetch')
        else:
            choices.append('destroy')
        op = random.choice(choices)
        if op == 'ag':
            b = tuple(unprocessed[-random.randint(1, max_g):])
            limit -= len(b)
            del unprocessed[-len(b):]
            inflight.append(b)
            ops.append(('ag', b))
        elif op == 'fetch':
            ops.append(('fetch', inflight.pop(random.randrange(len(inflight)))))
        else:
            ops.extend([('destroy_session', (-1,)), ('create_session', (-1,))])
            limit = num_max_inflight_agrs

    ops.append(('destroy_session', (-1,)))
    return ops, tensors, refs


def do_all_gather(buffer: deep_ep.ElasticBuffer,
                  is_inplace: bool, is_batched: bool,
                  tensors: tuple[torch.Tensor, ...],
                  start_event: torch.cuda.Event | None = None):
    # Copy into buffer if inplace
    if is_inplace:
        ag_tensors = buffer.agrs_get_inplace_tensor(tuple(t.shape for t in tensors), torch.bfloat16)
        for x, y in zip(ag_tensors, tensors, strict=True):
            x.copy_(y)
    else:
        ag_tensors = tensors

    # Record event
    if start_event is not None:
        torch.zeros(int(256e6 // 4), dtype=torch.int, device='cuda')  # flush L2 cache
        start_event.record()

    # Do all-gather
    if is_batched:
        *out_tensors, handle = buffer.all_gather(ag_tensors)
        return out_tensors, [handle]
    else:
        out_tensors, handles = [], []
        for t in ag_tensors:
            out_tensor, handle = buffer.all_gather(t)
            out_tensors.append(out_tensor)
            handles.append(handle)
        return out_tensors, handles


# noinspection PyTypeChecker,PyCallingNonCallable,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)

    # Print configs
    shape = (32, 64, 2048)
    num_max_inflight_agrs = args.num_max_inflight_agrs
    num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_num_max_session_bytes(
        group,
        [shape for _ in range(num_max_inflight_agrs)],
        torch.bfloat16
    )
    num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_buffer_size_hint(
        group, num_max_session_bytes)
    dist_print(f'Config:\n'
               f' > Ranks: {num_ranks}\n'
               f' > Shape: {shape}\n'
               f' > Max inflight AGRS: {num_max_inflight_agrs}\n',
               once_in_node=True)

    # Create buffer
    buffer = deep_ep.ElasticBuffer(group, explicitly_destroy=True, num_bytes=num_max_session_bytes)
    buffer.agrs_set_config(num_max_session_bytes, num_max_inflight_agrs)

    # Run stress tests
    dist_print('Running stress tests:', once_in_node=True)
    for seed in range(args.num_stress_iterations):
        random.seed(42 + seed)
        num_ops = 128
        ops, tensors, refs = generate_stress_ops(num_ops, num_max_inflight_agrs, shape, rank_idx, num_ranks)
        results = [None] * num_ops
        handles = dict()
        torch.cuda.synchronize()
        for op, indices in ops:
            if op == 'create_session':
                buffer.create_agrs_session()
            elif op == 'destroy_session':
                buffer.destroy_agrs_session()
            elif op == 'ag':
                is_inplace, is_batched = random.random() < 0.5, random.random() < 0.8
                handles[indices] = do_all_gather(buffer, is_inplace, is_batched, tuple(tensors[i] for i in indices))
            elif op == 'fetch':
                out_tensors, wait_handles = handles[indices]
                for h in wait_handles:
                    h()
                for out, idx in zip(out_tensors, indices, strict=True):
                    results[idx] = out.clone()

        for i in range(num_ops):
            assert results[i] is not None and torch.equal(results[i], refs[i]), \
                f'Rank {rank_idx}: stress mismatch at seed={seed}, op={i}'
        dist_print(f' > Seed {seed} passed ({num_ops} ops)', once_in_node=True)
    dist_print(once_in_node=True)

    # Destroy the buffer
    dist_print(f'Profiling all-gather:', once_in_node=True)
    buffer.destroy()

    # Profiling
    num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_num_max_session_bytes(
        group,
        [(2 ** 26,) for _ in range(num_max_inflight_agrs)],
        torch.bfloat16
    )
    num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_buffer_size_hint(
        group, num_max_session_bytes)
    buffer = deep_ep.ElasticBuffer(group, explicitly_destroy=True, num_bytes=num_max_session_bytes)
    buffer.agrs_set_config(num_max_session_bytes, num_max_inflight_agrs)
    for num_bytes in (2 ** p for p in range(20, 27)):
        # Create tensors
        shape = (num_bytes // 2, )
        tensors = tuple(torch.randn(shape, dtype=torch.bfloat16, device='cuda') for _ in range(num_max_inflight_agrs))

        # Tests
        for is_inplace in (False, True):
            for is_batched in (False, True):
                num_tests = 50
                start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
                end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
                torch.cuda.synchronize()

                for i in range(num_tests):
                    with buffer.agrs_new_session():
                        _, wait_handles = do_all_gather(buffer, is_inplace, is_batched, tensors, start_event=start_events[i])
                        for h in wait_handles:
                            h()
                    end_events[i].record()
                torch.cuda.synchronize()

                times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events, strict=True)])[1:]
                avg_t = np.average(times)
                unit = ('MB', 1e6) if num_bytes >= 1e6 else ('KB', 1e3)
                bandwidth_info = f', {num_bytes * num_ranks * num_max_inflight_agrs / avg_t / 1e9:.3f} GB/s' if num_ranks > 1 else ''

                dist_print(
                    f' > Rank: {rank_idx:3}/{num_ranks:3} | '
                    f'{num_ranks} x {(num_bytes / unit[1]):.0f} {unit[0]} | '
                    f'avg: {avg_t / num_max_inflight_agrs * 1e6:.3f} us'
                    f'{bandwidth_info}'
                    f' (inplace={int(is_inplace)}, batched={int(is_batched)})')
    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 all_gather kernels')
    parser.add_argument('--num-processes', type=int, default=8)
    parser.add_argument('--num-max-inflight-agrs', type=int, default=4)
    parser.add_argument('--num-stress-iterations', type=int, default=4)
    args = parser.parse_args()

    torch.multiprocessing.spawn(test, args=(args.num_processes, args), nprocs=args.num_processes)