[RFC] Performance Profiling Tools

🚀 The feature

1. Store usage statistics in Prefetcher

   • By tracking statistics within Prefetcher, we can reasonably determine whether upstream processes or downstream processes are faster. For example, the emptiness of the buffer queue may imply consumers are faster than producers. Users can insert this into various points in the pipeline to examine various behaviors. A common pattern we expect is to examine whether the pipeline is IO bound or compute bound.
   • #1141

2. DataLoader2 main process

   • torch profilers (e.g. torch.profiler.profile) currently work with DataLoader2, however, it only tracks functions and DataPipes that are executed within the main process. Nonetheless, we should validate that the information is helpful if most of the computations take place within the main process (e.g. using InProcessReadingService or dispatching process.
   • After 1 is completed, we can add APIs to DataLoader2 to fetch the relevant statistics from Prefetcher's buffer, such as the one that exists at the end of the main loop. It should allow users to examine whether the model is consuming faster than the preparation of samples.
   • PR pending
   • Tutorial pending

3. DataLoader2 worker process profiling

   • Two main options under considerations are:
     1. Attaching the profiler to worker process in order to get worker level metrics/trace. This will allow us to use existing profilers without re-implementing their features.
     2. MultiprocessingReadingService can provide methods to retrieve and aggregate metrics from certain DataPipes (mainly Prefetcher)

4. Integration with other tools (e.g. tracers)

   • We will likely want main and worker processes' to be visible within tracers (e.g. useful when integrated with TorchTNT).

Motivation, pitch

This set of tools and features aim to answer the questions:

1. Is my model training bottlenecked by data loading?
2. If so, which part of the pipeline? IO? Compute?

Alternatives

No response

Additional context

Comments and suggestions are welcomed.

Start by reading the Prefetcher and DataLoader2 entry points, including InProcessReadingService and MultiprocessingReadingService, then review how torch.profiler.profile currently observes the main process. The RFC is complete only when the profiling scope, worker-process metrics, aggregation APIs, and tracer integration have concrete designs and validation plans.