Repository: pytorch
Issue: [RFC] Cross-Process Performance Analysis: Straggler Detection
Labels: high priority,triage review,oncall: distributed,feature,module: c10d,module: ddp

### 🚀 The feature, motivation and pitch

## Motivation: Limitation of Existing Profiling Approach
To conduct PyTorch distributed training performance analysis, currently a recommended way is profiling by using [PyTorch profiler](https://pytorch.org/docs/stable/profiler.html) or [Nsight](https://docs-dev.nersc.gov/cgpu/software/profiling/). As a result, every profiled process can output a trace file, and users usually only need to take one or two representative files to analyze (mostly for the case of data parallelism).

However, one limitation of such approach is no easy support for cross-process performance analysis. A prominent example is **straggler detection**, because it's hard to identify straggling by only looking into a small number of individual trace files, and if the training runs at a relatively large scale, it's also impractical for users to eyeball a large number of trace files.


## Impact
Note that as we scale out trainings to more and more machines nowadays, straggler issues can become increasingly common, especially given that the most well-adopted PyTorch distributed training implementation, DDP is highly sensitive to stragglers -- the allreduce-based gradient synchronization has to wait for the slowest process at every iteration.

I believe that straggler issues can occur in various cases as DDP training is scaled out, and detecting and mitigating stragglers can have a high impact for improving large-scale training performance. A few cases:

1. Embedding table lookup can be expensive with large embeddings, especially these embedding tables may need to be partitioned across machines and hence incur all-to-all communication overheads. Such workloads may have a large variance across processes in the forward pass and hence cause stragglers. This can be found in training large-scale recommendation systems. (This one can be quite relevant to Ads companies like Meta: @zhaojuanmao, @rohan-varma)
2. If a massive amount of training data is used for training a relatively small model, the I/O cost can become a bottleneck. In this case, stragglers can occur if the data size of the input examples/records have a very large variance. Reading the input data from (potentially multi-homed) cloud storage may make this even worse. (This can be common when there is a strict low-latency requirement, and the performance of open-source data loader is not good enough.)
3. If a customized data loader has a very computationally expensive `transform` phase, and a high variance in the computation amount of on-the-fly data transformation may also contribute to straggling.

On the other hand, PyTorch users may not be even aware of such pattern ([example](https://github.com/pytorch/pytorch/issues/59433)).


## Detailed Requests
Can PyTorch provide an easy way for straggler detection? Particularly:

1. It should be easy to install such a detection component, and there should be no intrusive code change.
2. Straggling can be defined programmatically, and a metric can be reported to measure the extent of straggling.
3. Report whether the stragglers are the same across different iterations. If so, the straggling is more likely to be host/machine-specific, or it may imply the lack of sufficient data shuffling, etc.

With such feature, it will be much easier for industry powerful users to run fleet-wide straggler detections, and PyTorch users can be more alert to stragglers.


## Possible Implementations

1. We can develop a new type of [debugging DDP communication hook](https://pytorch.org/docs/stable/ddp_comm_hooks.html#debugging-communication-hooks). In the hook, before allreduce is kicked off, first measure the start time in the hook state, and then summarize different start times across processes at the same iteration. However, in the hook, it's impossible to also measure the total runtime of an iteration. To evaluate the extent of straggling, may need to bring some side inputs (e.g., regarding the total runtime of an iteration) when parsing a hook state that records allreduce starts times.
2. We can build an analysis layer on top of existing DDP logging (`_get_ddp_logging_data()`), which has already collected all the raw data needed for this purpose. CC: @zhaojuanmao, @rohan-varma 

### Alternatives

As mentioned above, currently there isn't an easy way to detect stragglers. After all, it's impractical to eyeball many trace files generated by profilers for the case of large-scale training.

### Additional context

This can be considered as a DDP feature request. To facilitate the triage:
1) Labels: **feature**, **module:ddp**, **oncall:distributed**.
2) CC PyTorch Distributed DDP developers: @mrshenli @pritamdamania87 @zhaojuanmao @rohan-varma 

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501