Feature request: add straggler detection to DTensor training path
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
The Megatron training path has full straggler detection support via `megatron.core.utils.StragglerDetector`, integrated into both the forward pass and data processing:
- **`nemo_rl/models/megatron/train.py`**: `StragglerDetector` is passed to `forward_step()`, `forward_step_with_preprocess()`, and `megatron_forward_backward()`, wrapping the model forward pass with timing instrumentation.
- **`nemo_rl/models/megatron/data.py`**: `StragglerDetector` is passed to `get_processed_microbatch_iterator()` and `process_microbatch()`, instrumenting the data loading/processing path.
The DTensor training path (`nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py`) has **zero** straggler detection support. The `train()` method calls `automodel_forward_backward()` without any timing instrumentation, making it impossible to identify slow workers in multi-node DTensor training.
## Motivation
In multi-node distributed training (especially on cloud instances with heterogeneous networking), straggler workers are a common source of throughput degradation. Without straggler detection, users have no visibility into which workers are falling behind or why.
## Proposed approach
We have validated that runtime-patching `DTensorPolicyWorkerV2` to use `nvidia_resiliency_ext.straggler.Detector` works correctly:
1. Create a class-level singleton `Detector` instance in `DTensorPolicyWorkerV2.__init__()`
2. Wrap `automodel_forward_backward()` calls with `detection_section("forward_backward")`
3. Generate straggler reports every N steps (configurable)
The `nvidia_resiliency_ext.straggler.Detector` API is straightforward:
```python
from nvidia_resiliency_ext.straggler import Detector, Section
# Initialize (once per worker)
detector = Detector(scores_to_compute=["relative_perf_scores"], gather_on_rank0=True)
# Wrap training sections
with detector.detection_section("forward_backward"):
mb_results = automodel_forward_backward(...)
# Report every N steps
if step % report_interval == 0:
report = detector.generate_report()
if report is not None: # Only rank 0 gets the report
log_straggler_report(report)
```
This approach:
- Does not require any changes to `automodel_forward_backward()` itself
- Is consistent with how the Megatron path instruments its forward pass
- Uses the same NVIDIA resiliency extension that NeMo RL already depends on (via Megatron)
- Can be gated behind a config flag (e.g., `dtensor_cfg.enable_straggler_detection`)
## Suggested config additions
```yaml
dtensor_cfg:
enable_straggler_detection: false # default off
straggler_report_interval: 10 # report every N steps
```
## References
- Megatron straggler detection: `megatron.core.utils.StragglerDetector`
- NVIDIA Resiliency Extension: `nvidia_resiliency_ext.straggler.Detector`
- DTensor worker: `nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py`
Contributor guide
Assessment
This issue has not been assessed yet.