Refit method families on GenerationInterface have grown per-transport instead of behind one contract
- 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
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