profile_model() — pytorch Function Reference

Architecture documentation for the profile_model() function in distributed.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration called by 1

Entity Profile

DynamoBenchmarks→ Configuration→ profile_model() — pytorch Function Reference

Dependency Diagram

graph TD
  6bdabe92_d399_b13a_52e9_438095848d71["profile_model()"]
  24e4ab44_f8ed_b835_f57a_8a13ebaabe74["run_model()"]
  24e4ab44_f8ed_b835_f57a_8a13ebaabe74 -->|calls| 6bdabe92_d399_b13a_52e9_438095848d71
  style 6bdabe92_d399_b13a_52e9_438095848d71 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/distributed.py lines 35–44

def profile_model(args, model, inputs, rank):
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
        for i in range(args.repeat):
            with record_function("Forward"):
                outputs = model(*inputs)
                loss = reduce_to_scalar_loss(outputs)
            with record_function("Backward"):
                loss.backward()
    if rank == 0:
        prof.export_chrome_trace(args.trace_file)