Function print_rank_0

megatron/__init__.py:17–23

If distributed is initialized print only on rank 0.

(*message)

def print_rank_0(*message):
    """If distributed is initialized print only on rank 0."""
    if torch.distributed.is_initialized():
        if torch.distributed.get_rank() == 0:
            print(*message, flush=True)
    else:
        print(*message, flush=True)


from .initialize import initialize_megatron