Profiling
About
The d9d.internals.profiling package wraps the PyTorch profiler for distributed jobs.
Internal API
If you use the standard d9d training loop, you do not need to call this package. d9d profiles the job based on its configuration. This page is for users who extend d9d.
Problems It Solves
Profiling a large distributed job has three problems:
- File naming: Thousands of ranks that write to the same file name cause race conditions.
- Storage space: Raw Chrome trace JSON files can grow to several GiB quickly.
- Synchronization: All ranks must profile the same steps without manual work.
The Profiler class solves them. It names each trace file after the DeviceMesh coordinates of the rank. It compresses each trace into a .tar.gz archive right after export. It runs the same periodic schedule (wait, warmup, active) on all ranks.
API Reference
d9d.internals.profiling
Distributed-aware wrapper around the PyTorch profiler.
Profiler
Periodic job profiler based on torch.profiler.
The profiling cycle repeats every period_steps steps. Each trace is exported, compressed and named after
the coordinates of the rank in the device mesh.
__init__(save_dir, period_steps, warmup_steps, active_steps, record_shapes, with_stack, dist_context)
Constructs the Profiler object.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
save_dir
|
Path
|
The directory to save the traces to. |
required |
period_steps
|
int
|
The length of one profiling cycle in steps (wait + warmup + active). |
required |
warmup_steps
|
int
|
The number of steps the profiler runs without recording, to warm up. |
required |
active_steps
|
int
|
The number of steps to record. |
required |
record_shapes
|
bool
|
Whether to record the input shapes of operators. |
required |
with_stack
|
bool
|
Whether to record the Python call stacks of operators. |
required |
dist_context
|
DistributedContext
|
The distributed context. |
required |
open(start_step)
Profiles the code in the block.
The profiler records CPU and CUDA activity on all threads. Each finished trace is saved to
save_dir/step_<step>/ as a .tar.gz archive.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
start_step
|
int
|
The current global step. The profiler counts steps from it. |
required |
Yields:
| Type | Description |
|---|---|
|
The configured |