Megatron Bridge TrainingConfig support
- 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
Assessment
This issue has not been assessed yet.