NVIDIA-NeMo / NVIDIA-NeMo/RL

Expose hardcoded Megatron infrastructure parameters to user config

Open
#2,229 1 comment 0 reactions 0 assignees View on GitHub
community-request
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

NeMo RL hardcodes several Megatron infrastructure parameters (I/O, timeouts, diagnostics) that the underlying framework already supports as configurable. In order to maximize performance on varied storage backends (e.g. NFS), we'd like these exposed in user config.

## Parameters to expose

All in `nemo_rl/models/megatron/setup.py`:

| Parameter | Line | Current | Issue |
|-----------|------|---------|-------|
| `async_save` | 904 | `False` | Checkpoints block training; async plumbing exists but is disabled |
| `init_process_group` timeout | 188 | 600s (default) | Too short for 120B weight conversion; env vars don't override it |
| `check_for_nan_in_grad` | 608 | `True` | Unnecessary overhead in production |
| `logging_level` | 600 | `0` | Not overridable |
| `fully_parallel_save/load`, `load_rng` | 904 | Hardcoded | No user override for debugging or reproducibility |

## Proposed config

```yaml
checkpointing:
async_save: true
fully_parallel_save: true
fully_parallel_load: true
load_rng: false

megatron_cfg:
distributed_timeout_minutes: 120
check_for_nan_in_grad: false
logging_level: 0
```

## Environment

- NeMo RL v0.5.0rc0, Megatron-LM (`3rdparty/Megatron-LM`, branch `yifu/superv3`)
- 6x DGX B200, Kubernetes 1.31, KubeRay 1.6.0, NFS storage
- Model: Nemotron 3 Super 120B, GRPO with DAPO-Math-17k

## Workarounds

- `async_save`: patch line 904 (`async_save=True`)
- Timeout: patch line 188 (`timeout=timedelta(minutes=120)`)
- NaN check: patch line 608 (`check_for_nan_in_grad=False`)

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.