NVIDIA-NeMo / NVIDIA-NeMo/RL

Nemorl vs Megatron: fwd/bwd slowdown (1.1–1.3x) and additional overhead from Nemorl runtime components

Open
#2,129 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

### Summary

When comparing Nemorl and Megatron under similar workloads, we observe:

- A **moderate slowdown (1.1–1.3x)** in the pure compute portion (fwd + bwd) for Nemorl vs Megatron.
- Additional, more noticeable overhead from Nemorl-side runtime components (e.g., `ray submit`, `data_processing`, `model_and_config_setup`), which pushes total step time further beyond Megatron.

This issue documents the current findings and requests investigation/optimizations on the Nemorl side.

### Context and Methodology

Comparison data and details are here:

Image

Image

Image

Key points about methodology:

1. **Different sequence packing**

- Nemorl and Megatron use different sequence packing strategies.
- Because of this, we **cannot strictly guarantee identical compute per step** if we only compare wall-clock step times.

2. **Token-based normalization**

- To make the comparison as fair as possible, we recorded the **number of computed tokens per step** for both systems.
- We then **selected steps whose token counts are as close as possible** and used those for the detailed comparison.
- The selection procedure and concrete examples are described in the doc linked above.

This approach is intended to approximate “equal compute” across the two systems before comparing timings.

---

### Observations

1. **Pure compute (fwd + bwd)**

- When isolating the fwd + bwd portion, Nemorl is **~1.1–1.3x slower** than Megatron for comparable token counts per step.
- This suggests that the **core compute slowdown itself is relatively modest**.

2. **Non-compute overhead in Nemorl**

- The more significant gap in end-to-end step time appears to come from Nemorl’s runtime / orchestration overhead, including but not limited to:
- `ray submit`
- `data_processing`
- `model_and_config_setup`
- These components do not exist (or are much lighter) in the Megatron baseline, so they contribute to the total slowdown beyond the 1.1–1.3x pure compute factor.

3. **Overall effect**

- From a user perspective, the **total step latency** difference is larger than what the 1.1–1.3x fwd/bwd numbers alone would suggest.
- To make Nemorl competitive for our workloads, it would be helpful to reduce or hide as much of this non-compute overhead as possible.

---

### Requested Actions / Questions

1. **Confirm and reproduce**

- Can the Nemorl team reproduce the ~1.1–1.3x fwd/bwd slowdown vs Megatron using a similar token-matched methodology?
- Are there known reasons in the current implementation (e.g., kernels, communication patterns) that explain this level of compute overhead?

2. **Analyze non-compute overhead**

- Break down the contributions from:
- `ray submit`
- `data_processing`
- `model_and_config_setup`
- any other Nemorl-specific orchestration components
- Identify which parts are:
- intrinsic to the design, vs
- implementation details that can be optimized (e.g., caching, reusing models/configs, reducing per-step Ray overhead).

3. **Potential optimizations**

Some possible directions (to be validated by maintainers):

- **Megatron-LM may open torch.compile**
- We will find a longer step time with smaller compute tokens in MLM.

4. **Documentation / guidance**

- It would be helpful to have guidance in the Nemorl docs on:
- How to configure Nemorl to minimize non-compute overhead for long-running training runs.
- Recommended best practices (e.g., how to structure jobs to avoid repeated setup costs).

---

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.