nemo_rl.weight_sync.megatron_weight_synchronizer#
Module Contents#
Classes#
Weight synchronization for the Megatron generation backend, both colocation modes. |
API#
- class nemo_rl.weight_sync.megatron_weight_synchronizer.MegatronWeightSynchronizer(
- policy: Any,
- generation: Any,
- *,
- colocated: bool,
- train_cluster: Optional[Any] = None,
- inference_cluster: Optional[Any] = None,
- refit_timeout_s: Optional[float] = None,
Bases:
nemo_rl.weight_sync.interfaces.WeightSynchronizerWeight synchronization for the Megatron generation backend, both colocation modes.
Colocated is the degenerate path: generation either aliases the training weights outright (reshardless) or re-partitions them into the worker’s dedicated inference model inside
prepare_for_generation(when the configured inference layout/impl differs) — a genuine parallelism-changing transfer, but one the worker performs internally on wake. Sync therefore reduces to dropping training-only buffers and re-entering inference mode.Non-colocated generation keeps the Megatron engine lifecycle here and delegates the transfer to native MCore refit, packed collective, or M2N.
Initialization
- init_communicator() None#
Wire the cross-group refit collective (non-colocated only).
Colocated generation shares the training worker group, so there is nothing to wire.
- sync_weights(
- *,
- timer: Optional[nemo_rl.utils.timer.Timer] = None,
- kv_scales: Optional[dict[str, float]] = None,
- property is_stale: bool#
- shutdown() None#
Release any resources owned by the delegated transport.