Profiler

Last updated: 2025-11-04

VeOmni offers a profiler function for users to trace training:

from veomni.utils import helper

# Before train loop, create your profiler
if args.train.global_rank == 0:
    if args.train.enable_profiling:
        profiler = helper.create_profiler(
            start_step=args.train.profile_start_step,
            end_step=args.train.profile_end_step,
            trace_dir=args.train.profile_trace_dir,
            record_shapes=args.train.profile_record_shapes,
            profile_memory=args.train.profile_profile_memory,
            with_stack=args.train.profile_with_stack,
        )
        profiler.start()

for epoch in range(args.train.num_train_epochs):
    data_iterator = iter(train_dataloader)
    for _ in range(args.train.train_steps):
        ...  # Your training code
        profiler.step()
        if global_step == args.train.profile_end_step:
            profiler.stop()
            # upload file to merlin
            helper.upload_trace(args.train.wandb_project, args.train.wandb_name, args.train.profile_trace_dir)