Split the Megatron generation refit mixin out of megatron_worker.py
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
`nemo_rl/models/generation/megatron/megatron_worker.py` now holds two unrelated mixins and a set of refit-only module-level symbols. After #3739 it is 1978 lines, up from 1284 — the largest file in the package by a wide margin (`megatron_generation.py` 669, `utils.py` 268, `config.py` 208).
Current layout:
| Lines | Symbol | Concern |
|---|---|---|
| [113](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/megatron/megatron_worker.py#L113), [127](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/megatron/megatron_worker.py#L127) | `_inference_optimized_transformer_layer_spec`, `_configure_inference_optimized_layer_spec` | model build |
| [159](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/megatron/megatron_worker.py#L159) | `_resolve_mxfp8_refit_backend` | refit |
| [164-188](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/megatron/megatron_worker.py#L164-L188) | `_MegatronRefitTask` | refit |
| [189-200](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/megatron/megatron_worker.py#L189-L200) | `_MegatronBulkRefitPiece` | refit |
| [201-1163](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/megatron/megatron_worker.py#L201) | `MegatronGenerationMixin` | generation |
| [1164-1978](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/megatron/megatron_worker.py#L1164) | `MegatronGenerationRefitMixin` | refit (~815 lines, 41% of the file) |
The three refit dataclasses/helpers sit between the model-build helpers and a class that does not use them, so reading either concern means skipping over the other.
The sibling backend already has the split: `nemo_rl/models/generation/dynamo/` keeps [`dynamo_worker.py`](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/dynamo/dynamo_worker.py) (304 lines) separate from [`refit.py`](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/generation/dynamo/refit.py) (271 lines), which holds `DynamoRefitChannel` and its endpoint dataclass. The same package already carries a `megatron/utils.py`, so a second module here needs no new convention.
Proposed: move `MegatronGenerationRefitMixin` plus `_resolve_mxfp8_refit_backend`, `_MegatronRefitTask` and `_MegatronBulkRefitPiece` into `nemo_rl/models/generation/megatron/refit.py`. The move is mechanical — the mixin's only external references are the import and base-class list in `megatron_policy_worker.py` ([:67](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/policy/workers/megatron_policy_worker.py#L67), [:416](https://github.com/NVIDIA-NeMo/RL/blob/b11d237e3021504186fd34f3ef39826b5e6e3b54/nemo_rl/models/policy/workers/megatron_policy_worker.py#L416)), and the three underscore-prefixed symbols are file-private.
Deliberately not done in #3739: that PR is already large, and a pure relocation would make its diff harder to review. Filing so the move is not lost.
Contributor guide
Research direction
Start with nemo_rl/models/generation/megatron/megatron_worker.py and locate MegatronGenerationRefitMixin plus the three refit-only symbols listed in the issue. Read the import and base-class references in nemo_rl/models/policy/workers/megatron_policy_worker.py, then move the refit code to nemo_rl/models/generation/megatron/refit.py and update those references. Done means the refit symbols live in refit.py and megatron_worker.py retains the generation and model-build concerns.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 2/5
- Estimated time
- Half a day
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100