NVIDIA-NeMo / NVIDIA-NeMo/RL

Megatron Bridge TrainingConfig support

Open
#1,319 2 comments 0 reactions 0 assignees View on GitHub
community-request enhancement external t-mcore training waiting-on-maintainers x-inflection
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Is your feature request related to a problem? Please describe.**

[Here](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/policy/megatron_policy_worker.py#L644) only `micro_batch_size`, `global_batch_size`, and `train_iters` are supported but other params may be needed too.

i.e. in async GRPO the `decrease_batch_size_if_needed` may be useful since some of GPUs don't participate in train loop

**Describe the solution you'd like**

That's a design decision with many options:

- leave it as is - I assume current logic is there to support train params for both FSDP and Megatron and it may not be a big deal to support all megatron bridge params
- add every TrainingConfig param as `train_`
- make `TrainingConfig` a part of `policy` or `policy.megatron_cfg` - I think it's the cleanest solution but deprecates previous configs

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.