Function tb_wandb_log

megatron/logging.py:360–369

(
    key, value, iteration_no, use_wandb, tensorboard_writer=None, all_ranks=False
)

def tb_wandb_log(
    key, value, iteration_no, use_wandb, tensorboard_writer=None, all_ranks=False
):
    # logs to both tb and wandb (if present) from the zeroth rank
    do_log = torch.distributed.get_rank() == 0 or all_ranks
    if do_log and value is not None:
        if tensorboard_writer:
            tensorboard_writer.add_scalar(key, value, iteration_no)
        if use_wandb:
            wandb.log({key: value}, step=iteration_no)