| .. _profiler: | |
| ======== | |
| Profiler | |
| ======== | |
| Last updated: 2025-11-04 | |
| VeOmni offers a profiler function for users to trace training: | |
| .. code-block:: python | |
| 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) | |