nemo_rl.weight_sync.megatron_weight_synchronizer

Module Contents

Classes

MegatronWeightSynchronizer

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,

)

Bases: nemo_rl.weight_sync.interfaces.WeightSynchronizer

Weight 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 adds the cross-group collective: the training and inference workers are disjoint actor groups that rendezvous in mcore’s reshard-capable weight swap, with the engine suspended around the transfer. That wiring (a joint refit process group over the configured copy-service backend) is established once in init_communicator.

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,

) → Optional[dict[str, float]]

property is_stale: bool

shutdown() → None

Nothing to tear down; the collective lives in the worker groups.