NVIDIA-NeMo / NVIDIA-NeMo/RL

Refit method families on GenerationInterface have grown per-transport instead of behind one contract

Open
#4,006 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

`WeightSynchronizer` (`nemo_rl/weight_sync/interfaces.py:48`) is a clean 4-method contract — `init_communicator` / `sync_weights` / `reconcile_communicator` / `shutdown`. `GenerationInterface` underneath it is not: every transport has added its own parallel family of methods, and backends now differ in which family they implement and in whether the method is declared at all.

Current state on `GenerationInterface` (`nemo_rl/models/generation/interfaces.py`):

| Transport | Setup | Plan | Transfer |
|---|---|---|---|
| packed collective (`refit_transport: null`) | `init_collective` :449 | `prepare_refit_info` :501 | `update_weights_from_collective` :509 |
| `nccl_reshard` | `init_nccl_reshard_comm_group` — *not declared* | `prepare_nccl_reshard_refit_info` :535 | `nccl_reshard_refit` :539 |
| `mcore` native | `init_collective` :449 — `uses_native_refit` branch | — *(none)* | `update_weights_from_collective` :509 — same branch |
| colocated IPC | — | — | `update_weights_via_ipc_zmq` :505 |

Three further methods that both non-colocated synchronizers call unconditionally are duck-typed, not on the interface: `set_refit_membership`, `rebuild_collective`, `rebuild_nccl_reshard_comm_group`. With one implementation (vLLM) that read as a vLLM detail; there are now two (`MegatronGeneration` added them in #3739), and a third backend joining either transport fails with a runtime `AttributeError` rather than at import.

Consequences today:

- Adding a transport means adding a method triple to a shared interface that most backends must stub out.
- Which methods a backend must implement is discoverable only by reading each synchronizer.
- `MegatronGeneration.init_collective` serves `null` and `mcore` on two different branches while `nccl_reshard` reaches the class through a different method entirely (`rebuild_collective`), so the same lifecycle stage has three entry points.

Rough direction (not prescriptive): collapse to one `setup` / `plan` / `transfer` triple parameterized by transport, or move the per-transport methods off `GenerationInterface` onto a transport-owned object the synchronizer holds. Either way, declare `set_refit_membership` / `rebuild_*` on whatever contract survives.

No action needed in #3739 — that PR follows the existing shape correctly; this is the accumulated cost of the shape itself.

Contributor guide

Open the contributing guide

Research direction

Start with nemo_rl/models/generation/interfaces.py and nemo_rl/weight_sync/interfaces.py, then trace the synchronizers and the transport-specific entry points listed in the issue. Define the surviving contract and transport ownership boundaries before changing callers; done means transport lifecycles share one discoverable contract and set_refit_membership plus the rebuild methods are declared wherever required.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend-api-design
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.