`nemo_rl.models.megatron.router_replay`#

Module Contents#

`router_replay_enabled`
`configure_vllm_for_router_replay`	Apply vLLM settings required for Megatron router replay correctness.
`validate_router_replay_config`
`_iter_model_modules`
`_unwrap_model_config`
`_global_moe_layer_numbers`
`_router_replay_instances_for_model`
`_local_layer_numbers_for_model`
`_normalize_routed_experts_for_mcore`
`_payload_indices_for_moe_layers`
`_router_replay_validation_enabled`
`_validate_replay_tensor`
`_install_missing_route_fallback_patch`
`_get_tensor_model_parallel_world_size`
`_get_tensor_model_parallel_rank`
`_split_for_sequence_parallel`
`build_router_replay_tensors`	Build MCore RouterReplay tensors in model-local router order.
`build_router_replay_assignments`	Pair model-owned MCore RouterReplay instances with their replay tensors.
`set_router_replay_forward`
`set_router_replay_backward`
`clear_router_replay`
`clear_global_router_replay_instances`

`_ROUTER_REPLAY_VALIDATE_ENV`
`_MISSING_ROUTE_SENTINEL`
`_MISSING_ROUTE_FALLBACK_PATCH_ATTR`

nemo_rl.models.megatron.router_replay._ROUTER_REPLAY_VALIDATE_ENV#: ‘NRL_ROUTER_REPLAY_VALIDATE’

nemo_rl.models.megatron.router_replay._MISSING_ROUTE_FALLBACK_PATCH_ATTR#: ‘_nrl_missing_route_fallback_patch’

nemo_rl.models.megatron.router_replay.router_replay_enabled( config: nemo_rl.models.policy.PolicyConfig, ) → bool#

nemo_rl.models.megatron.router_replay.configure_vllm_for_router_replay( config: nemo_rl.models.policy.PolicyConfig, ) → None#: Apply vLLM settings required for Megatron router replay correctness.

nemo_rl.models.megatron.router_replay.validate_router_replay_config( config: nemo_rl.models.policy.PolicyConfig, ) → None#

nemo_rl.models.megatron.router_replay._iter_model_modules( model: Any, ) → collections.abc.Iterable[Any]#

nemo_rl.models.megatron.router_replay._unwrap_model_config(model: Any) → Optional[Any]#

nemo_rl.models.megatron.router_replay._global_moe_layer_numbers(model_config: Any) → list[int]#

nemo_rl.models.megatron.router_replay._router_replay_instances_for_model( model: Any, ) → list[tuple[Any, int]]#

nemo_rl.models.megatron.router_replay._local_layer_numbers_for_model(model: Any) → set[int]#

nemo_rl.models.megatron.router_replay._normalize_routed_experts_for_mcore( routed_experts: torch.Tensor, ) → torch.Tensor#

nemo_rl.models.megatron.router_replay._payload_indices_for_moe_layers( *, global_moe_layers: list[int], num_payload_layers: int, total_num_layers: int, ) → dict[int, int]#

nemo_rl.models.megatron.router_replay._router_replay_validation_enabled() → bool#

nemo_rl.models.megatron.router_replay._validate_replay_tensor( replay_tensor: torch.Tensor, model_config: Any, *, layer_number: int, payload_idx: int, ) → None#

nemo_rl.models.megatron.router_replay._install_missing_route_fallback_patch() → None#

nemo_rl.models.megatron.router_replay._get_tensor_model_parallel_world_size() → int#

nemo_rl.models.megatron.router_replay._get_tensor_model_parallel_rank() → int#

nemo_rl.models.megatron.router_replay._split_for_sequence_parallel( model_config: Any, routed_experts: torch.Tensor, ) → torch.Tensor#

nemo_rl.models.megatron.router_replay.build_router_replay_tensors( model: Any, routed_experts: torch.Tensor, ) → list[torch.Tensor]#: Build MCore RouterReplay tensors in model-local router order.

nemo_rl.models.megatron.router_replay.build_router_replay_assignments( model: Any, routed_experts: torch.Tensor, ) → list[tuple[Any, torch.Tensor]]#: Pair model-owned MCore RouterReplay instances with their replay tensors.

nemo_rl.models.megatron.router_replay.set_router_replay_forward( model: Any, routed_experts: torch.Tensor, ) → None#

nemo_rl.models.megatron.router_replay.set_router_replay_backward(model: Any) → None#

nemo_rl.models.megatron.router_replay.clear_router_replay(model: Optional[Any] = None) → None#

nemo_rl.models.megatron.router_replay.clear_global_router_replay_instances() → None#