def create_set_process_group(k: int):
    """Creates a new distributed group using sets of k GPUs.

    For example, if you have 16 GPUs and input k=4, the resulting process groups
    will have ranks:
        process group 0 ranks: [ 0,  1,  2,  3]
        process group 1 ranks: [ 4,  5,  6,  7]
        process group 2 ranks: [ 8,  9, 10, 11]
        process group 3 ranks: [12, 13, 14, 15]

    Args:
        k (int): Number of GPUs to use in set size.

    Returns:
        A handle of distributed group that can be given to collective calls.
    """
    world_size = distributed.get_world_size()
    if world_size % k != 0:
        raise RuntimeError(f'{world_size=} must be divisible by {k=}.')
    start = distributed.get_rank() // k * k
    ranks = tuple(range(start, start + k))
    return create_process_group_ranks(ranks)
Creates a new distributed group using sets of k GPUs.

    For example, if you have 16 GPUs and input k=4, the resulting process groups
    will have ranks:
        process group 0 ranks: [ 0,  1,  2,  3]
        process group 1 ranks: [ 4,  5,  6,  7]
        process group 2 ranks: [ 8,  9, 10, 11]
        process group 3 ranks: [12, 13, 14, 15]

    Args:
        k (int): Number of GPUs to use in set size.

    Returns:
        A handle of distributed group that can be given to collective calls.
    
llmfoundry/models/utils/config_moe_args.py