NVIDIA-NeMo / NVIDIA-NeMo/RL

Feature request: add straggler detection to DTensor training path

Open
#1,992 6 comments 0 reactions 0 assignees View on GitHub
community-request t-pytdensor waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.