INNER CODE UNIT · Python

print_rank_0

HIT-SCIR/Chinese-Mixtral-8x7B · train.py:177

def print_rank_0(*args, **kwargs):
    if torch.distributed.is_initialized():
        if torch.distributed.get_rank() == 0:
            print(*args, **kwargs)
    else:
        print(*args, **kwargs)


class EvaluateFirstStepCallback(TrainerCallback):
    # Make Trainer evaluate before first training step.
    # `--logging_first_step` will not work, maybe this is a bug.
    # We use this callback to deal with this.
    # See https://discuss.huggingface.co/t/how-to-evaluate-before-first-training-step/18838/7
    def on_step_begin(self, args, state, control, **kwargs):
        if state.global_step == 0:
            control.should_evaluate = True

View source record →