INNER CODE UNIT · Python

is_last_rank

openhackathons-org/gpubootcamp · archived/ai/Megatron/English/Python/jupyter_notebook/Megatron-LM/megatron/__init__.py:46

def is_last_rank():
    return torch.distributed.get_rank() == (
        torch.distributed.get_world_size() - 1)

def print_rank_last(message):
    """If distributed is initialized, print only on last rank."""
    if torch.distributed.is_initialized():
        if is_last_rank():
            print(message, flush=True)
    else:
        print(message, flush=True)

View source record →