NVIDIA-NeMo / NVIDIA-NeMo/RL

Split the Megatron generation refit mixin out of megatron_worker.py

Open Beginner friendly
#4,016 0 comments 0 reactions 0 assignees View on GitHub
Feature
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.