Nemorl vs Megatron: fwd/bwd slowdown (1.1–1.3x) and additional overhead from Nemorl runtime components
- 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:
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
Assessment
This issue has not been assessed yet.