File size: 1,202 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | .. _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)
|