INNER CODE UNIT · Python

maybe_make_profiler

Guitaricet/relora · torchrun_main.py:322

def maybe_make_profiler(args):
    if not args.profile: return None
    global_rank = dist.get_rank()
    profiler_logging_dir = os.path.join(f"profiler_logs/{args.run_name}")
    prof = torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
        on_trace_ready=torch.profiler.tensorboard_trace_handler(profiler_logging_dir, worker_name=f"rank{global_rank}"),
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
    )
    print(f"Rank {global_rank} profiling results will be saved to {profiler_logging_dir}")
    prof.start()
    return prof


def main(args):
    # seed all

View source record →