PyTorch NCCL Flight Recorder

gjit_JNPR · August 1, 2024, 6:46pm

The recently published LLaMa-3 paper (https://arxiv.org/pdf/2407.21783) mentions the use of a PyTorch built-in feature called NCCL Flight Recorder as shown below:. On further exploring PyTorch, I did not find any such feature that can record NCCL events.

Question:

Is this feature open sourced?
If yes, is there any documentation or summary of it?

gnadathur · August 20, 2024, 8:42pm

Thanks for reaching out.
Flight recorder will be released as prototype API in 2.5 release with documentation.
cc: @c-p-i-o

c-p-i-o · August 20, 2024, 8:48pm

Hey Jit! Yes, flight recorder is open sourced and available in the PyTorch branch. We just haven’t broadly made any announcements as we’ve been testing and changing some of the underlying implementation.

We’re writing up some documentation next few days on how to enable it and make use of the generated data to detect stuck jobs. Please stand by. As @gnadathur mentioned, FR will be marked Prototype in upcoming 2.5 release.

gjit_JNPR · August 20, 2024, 8:58pm

Hi Chiraj and Gokul,

Thanks for your responses! Looking forward to being able to use the NCCL Flight Recorder.

Jay_you_Dosen · August 30, 2024, 12:41pm

Could you please tell me which branch the flight recorder is in? thanks

c-p-i-o · August 30, 2024, 1:42pm

Flight recorder collection portion has been available in 2.4 and in the main branch.
The tool to analyze FR traces is under /tools/flight_recorder/fr_trace.py in the main branch.

c-p-i-o · August 30, 2024, 1:43pm

Also, I’m finishing a tutorial here. Almost done.

github.com

pytorch/tutorials/blob/cpio/flight_recorder/prototype_source/flight_recorder_tutorial.rst

(prototype) Flight Recorder for Debugging
=========================================
**Author**: `Chirag Pandya <https://github.com/c-p-i-o>`, `Junjie Wang <https://github.com/fduwjj>`

This tutorial introduces a new tool for debugging stuck jobs during distributed training.

Background and Motivation
--------------------------
An AI distributed training job refers to the process of training a machine learning model using multiple devices, such
as GPUs or CPUs, connected in a network. This approach allows for faster and more efficient training of large models
that require significant computational resources.
An engineer’s goal is to complete an AI training job as fast as possible and make continuous improvements such that
subsequent training can be done faster. A trained usable model is the final desired outcome.
One of the biggest impediment to completing training is the concept of a "stuck job".

A distributed AI training job is considered "stuck" when it stops making meaningful progress for an extended period of
time.

A job can get stuck for various reasons:
    - Data Starvation: This happens when the training job is not receiving data at the expected rate. This could be due to

This file has been truncated. show original

Jay_you_Dosen · September 26, 2024, 6:40am

Thanks. It’s very helpful.

jinqinn · September 27, 2024, 2:04am

Thanks. its very helpful.

c-p-i-o · October 10, 2024, 8:36pm

Hey @jinqinn, the official tutorial is published here. Let us know if you’ve had any success with Flight Recorder.

tmm1 · January 21, 2025, 5:51am

The tutorial says I can run torchfrtrace directly when using nightly, but it didn’t work for me.

❯ uv run torchfrtrace -h
Traceback (most recent call last):
  File ".venv/bin/torchfrtrace", line 5, in <module>
    from tools.flight_recorder.fr_trace import main
ModuleNotFoundError: No module named 'tools'

danielz · January 21, 2025, 6:09pm

Yes, I have the same issue. This is probably because we are using a non-nightly build.

If you install the PyTorch nightly build or build from scratch with USE_DISTRIBUTED=1, you can directly use the following command directly:

torchfrtrace <dump dir containing trace files> [-o <output file>]

EDIT: Actually, I still have the same error even in the nightly build docker.

c-p-i-o · January 21, 2025, 11:20pm

Oops. It seems related to PYTHONPATH not correctly set for the module.

$pwd
/home/cpio/local/pytorch
$export PYTHONPATH=$(pwd):$PYTHONPATH
$torchfrtrace
Traceback (most recent call last):
  File "/home/cpio/local/b/pytorch-env/bin/torchfrtrace", line 33, in <module>
    sys.exit(load_entry_point('torch', 'console_scripts', 'torchfrtrace')())
  File "/home/cpio/local/pytorch/tools/flight_recorder/fr_trace.py", line 44, in main
    assert args.trace_dir, "Trace directory trace_dir is required"
AssertionError: Trace directory trace_dir is required

Let me see if I can resolve this.

danielz · January 21, 2025, 11:31pm

Yes, I had to point PYTHONPATH to pytorch/tools

c-p-i-o · January 22, 2025, 5:35pm

Thanks for the quick test! I’ll figure out the packaging issue so that you don’t have to set PYTHONPATH.
I quickly tried adding an empty __init__.py file in the tools directory but that didn’t seem to help.