Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve Profiling Docs #1298

Merged
merged 1 commit into from
Oct 1, 2024
Merged

Improve Profiling Docs #1298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -679,7 +679,9 @@ We support profiling with Nsight Systems, the PyTorch Profiler, and PyTorch Memo

## Nsight Systems Profiling

To use the Nsight Systems profiling, set config options `profile`, `profile_step_start`, and `profile_step_stop`. Launch training with:
To use the Nsight Systems profiling, set config options `profile`, `profile_step_start`, and `profile_step_stop` (see [here](https://github.com/EleutherAI/gpt-neox/blob/main/configs/neox_arguments.md) for argument usage, and [here](https://github.com/EleutherAI/gpt-neox/blob/main/configs/prof.yml) for a sample config).

To populate nsys metrics, launch training with:

```
nsys profile -s none -t nvtx,cuda -o <path/to/profiling/output> --force-overwrite true \
Expand All @@ -689,22 +691,22 @@ $TRAIN_PATH/train.py --conf_dir configs <config files>

The generated output file can then by viewed with the Nsight Systems GUI:

![Alt text](images/nsight_profiling.png)
![nsight-prof](images/nsight_profiling.png)

## PyTorch Profiling

To use the built-in PyTorch profiler, set config options `profile`, `profile_step_start`, and `profile_step_stop`.
To use the built-in PyTorch profiler, set config options `profile`, `profile_step_start`, and `profile_step_stop` (see [here](https://github.com/EleutherAI/gpt-neox/blob/main/configs/neox_arguments.md) for argument usage, and [here](https://github.com/EleutherAI/gpt-neox/blob/main/configs/prof.yml) for a sample config).

The PyTorch profiler will save traces to your `tensorboard` log directory. You can view these traces within
TensorBoard by following the steps [here](https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html).

![Alt text](images/pytorch_profiling.png)
![torch-prof](images/pytorch_profiling.png)

## PyTorch Memory Profiling

To use PyTorch Memory Profiling, set config options `memory_profiling` and `memory_profiling_path`.
To use PyTorch Memory Profiling, set config options `memory_profiling` and `memory_profiling_path` (see [here](https://github.com/EleutherAI/gpt-neox/blob/main/configs/neox_arguments.md) for argument usage, and [here](https://github.com/EleutherAI/gpt-neox/blob/main/configs/prof.yml) for a sample config).

![Alt text](images/memory_profiling.png)
![mem-prof](images/memory_profiling.png)

View the generated profile with the [memory_viz.py](https://github.com/pytorch/pytorch/blob/main/torch/cuda/_memory_viz.py) script. Run with:

Expand Down
17 changes: 17 additions & 0 deletions configs/prof.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
# Sample profiling config
{
# Turns on nsys and pytorch profiling
"profile": true,

# pytorch profiler options
"profile_step_start": 10,
"profile_step_stop": 12,

# pytorch memory profiler options
"memory_profiling": true,
"memory_profiling_path": tensorboard,


# All trace files (pytorch, nsys, tensorboard, etc) will be written here
"tensorboard_dir": "tensorboard",
}
4 changes: 2 additions & 2 deletions megatron/neox_arguments/neox_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -732,8 +732,8 @@ class NeoXArgsLogging(NeoXArgsTemplate):

profile: bool = False
"""
Enable nsys profiling. When using this option,
nsys options should be specified in commandline.
Enable nsys and pytorch profiling. When using this option with nsys,
nsys options should be directly specified in commandline.
An example nsys commandline is
```
nsys profile -s none -t nvtx,cuda -o <path/to/output_file>
Expand Down
Loading