From 204a6b7b4485cd366e73cb512f58414e6b2e1955 Mon Sep 17 00:00:00 2001 From: Yigong Qin Date: Tue, 14 Jul 2026 17:54:46 -0700 Subject: [PATCH 1/5] chore: pin Megatron-Bridge to YigongQin fork (zero-kl) Signed-off-by: YigongQin --- 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 554c7b9324..7767a134a2 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 554c7b9324225aa863eee52e8b8fdde7abced2b1 +Subproject commit 7767a134a278b0adf4e2d378ed158b87bd743e15 From db09bc263703356791480a04c6cdd808512cee0c Mon Sep 17 00:00:00 2001 From: Yigong Qin Date: Tue, 14 Jul 2026 18:54:38 -0700 Subject: [PATCH 2/5] knobs wiring for kl study Signed-off-by: YigongQin --- .../generation/megatron/megatron_worker.py | 32 ++++- nemo_rl/models/megatron/moe_routing_record.py | 96 ++++++++++++++ nemo_rl/models/megatron/router_replay.py | 119 +++++++++++++++++- nemo_rl/models/megatron/setup.py | 66 +++++++++- nemo_rl/models/policy/__init__.py | 10 ++ nemo_rl/models/policy/lm_policy.py | 21 +++- .../policy/workers/megatron_policy_worker.py | 49 +++++++- pyproject.toml | 6 +- 8 files changed, 386 insertions(+), 13 deletions(-) create mode 100644 nemo_rl/models/megatron/moe_routing_record.py diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 0a0721dd99..9f9bd89e86 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -94,6 +94,14 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: ) from megatron.core.utils import get_attr_wrapped_model + from nemo_rl.models.megatron.router_replay import ( + assert_megatron_inference_router_replay_ready, + rebuild_global_router_replay_registry, + reset_moe_routing_metadata_buffer, + router_replay_enabled, + sync_model_config_for_router_replay, + ) + pg_collection = get_attr_wrapped_model(self.model, "pg_collection") buffer_size_gb = mcore_generation_config["buffer_size_gb"] @@ -107,7 +115,10 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: # The value may be overwritten by `recompute_kv_cache_after_weight_updates`. kv_cache_management_mode = mcore_generation_config["kv_cache_management_mode"] - needs_static_kv_pointers = kv_cache_management_mode != "persist" + static_kv_memory_pointers = mcore_generation_config.get( + "static_kv_memory_pointers", + kv_cache_management_mode != "persist", + ) materialize_only_last_token_logits = mcore_generation_config[ "materialize_only_last_token_logits" @@ -117,6 +128,14 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: mamba_inference_state_config = MambaInferenceStateConfig.from_model(self.model) is_hybrid_model = mamba_inference_state_config is not None + if is_hybrid_model and self.cfg.get("megatron_cfg", {}).get( + "zero_train_gen_mismatch" + ): + # Match the train scan's fp32 boundary-state precision so gen SSM cache + # doesn't diverge from train (gen defaults to bf16 otherwise). + mcore_generation_config.setdefault( + "mamba_inference_ssm_states_dtype", "float32" + ) if is_hybrid_model: if ( mcore_generation_config.get("mamba_inference_ssm_states_dtype") @@ -139,6 +158,10 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: if logging_step_interval is None: logging_step_interval = 0 + if router_replay_enabled(self.cfg): + rebuild_global_router_replay_registry(self.model) + sync_model_config_for_router_replay(self.model, self.cfg) + # flashinfer's fused-RoPE kernel only dispatches fp16/bf16 q/k. use_flashinfer_fused_rope = self.model.config.params_dtype in ( torch.float16, @@ -152,7 +175,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: max_tokens=max_tokens, max_sequence_length=mcore_generation_config["max_model_len"], kv_cache_management_mode=KVCacheManagementMode(kv_cache_management_mode), - static_kv_memory_pointers=needs_static_kv_pointers, + static_kv_memory_pointers=static_kv_memory_pointers, use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, use_flashinfer_fused_rope=use_flashinfer_fused_rope, sampling_backend="flashinfer", @@ -182,6 +205,11 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: self.inference_context = DynamicInferenceContext( self.model.config, inference_config ) + if router_replay_enabled(self.cfg): + reset_moe_routing_metadata_buffer(self.inference_context) + assert_megatron_inference_router_replay_ready( + self.model, self.inference_context, self.cfg + ) self.inference_wrapped_model = GPTInferenceWrapper( self.model, self.inference_context ) diff --git a/nemo_rl/models/megatron/moe_routing_record.py b/nemo_rl/models/megatron/moe_routing_record.py new file mode 100644 index 0000000000..cc0dd4abd7 --- /dev/null +++ b/nemo_rl/models/megatron/moe_routing_record.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Record MoE routing during Megatron-inference generation as R3 ``routed_experts``. + +Upstream router replay (R3, ``nemo_rl.models.megatron.router_replay``) provides the +*replay* side for the Megatron training/logprob forward, but only records routing on +the vLLM generation path. The colocated Megatron-inference generation path therefore +has no recorder. This module bridges that gap: it converts the dynamic-inference +engine's per-sample ``routing_indices`` into the ``[batch, seq, layers, topk]`` +``routed_experts`` tensor that ``router_replay.set_router_replay_forward`` / +``build_router_replay_assignments`` consume (see ``_normalize_routed_experts_for_mcore``, +which accepts ``[B, S, L, K]``). +""" + +from typing import Optional + +import torch + + +def coerce_routing_to_3d(routing: torch.Tensor) -> torch.Tensor: + """Normalize per-sample routing to ``[num_tokens, num_layers, topk]``.""" + if not isinstance(routing, torch.Tensor): + raise TypeError(f"routing_indices must be a torch.Tensor, got {type(routing)}") + if routing.ndim == 3: + return routing + raise ValueError( + f"routing_indices must be 3D [tokens, layers, topk], got shape {tuple(routing.shape)}" + ) + + +def align_routing_rows_to_token_count( + routing: torch.Tensor, num_tokens: int +) -> torch.Tensor: + """Pad/trim routing rows to ``num_tokens`` so they replay on a full-sequence forward. + + Dynamic inference accumulates ``[num_tokens, layers, topk]`` (sometimes one fewer + row than the padded sequence length). Repeat the last recorded row when an extra + step is required; trim when there are more. + """ + if routing.ndim != 3: + raise ValueError(f"Expected 3D routing tensor, got {tuple(routing.shape)}") + num_rows = routing.shape[0] + if num_rows == num_tokens: + return routing + if num_rows > num_tokens: + return routing[:num_tokens].contiguous() + if num_rows == 0: + raise ValueError("Cannot align empty routing indices to a non-empty sequence") + pad_rows = num_tokens - num_rows + last = routing[-1:].expand(pad_rows, -1, -1) + return torch.cat([routing, last], dim=0) + + +def build_routed_experts_batch( + routing_per_sample: list[Optional[torch.Tensor]], + seq_lengths: torch.Tensor, + seq_dim: int, +) -> Optional[torch.Tensor]: + """Build the R3 ``routed_experts`` tensor ``[batch, seq_dim, layers, topk]``. + + Each sample's routing is aligned to its (unpadded) sequence length and right-padded + with zeros to ``seq_dim`` (the generation ``output_ids`` sequence length), so that + ``routed_experts[i][:seq_lengths[i]]`` is the per-token routing and the tail is pad. + Returns ``None`` when no sample carries routing (e.g. dense models / replay disabled). + """ + if not any(r is not None for r in routing_per_sample): + return None + if any(r is None for r in routing_per_sample): + raise ValueError( + "routing_indices must be present for every sample when router replay is enabled" + ) + aligned = [ + align_routing_rows_to_token_count( + coerce_routing_to_3d(r), int(seq_lengths[i].item()) + ) + for i, r in enumerate(routing_per_sample) + ] + num_layers, topk = aligned[0].shape[1], aligned[0].shape[2] + out = torch.zeros( + len(aligned), seq_dim, num_layers, topk, dtype=aligned[0].dtype + ) + for i, r in enumerate(aligned): + n = min(r.shape[0], seq_dim) + out[i, :n] = r[:n] + return out diff --git a/nemo_rl/models/megatron/router_replay.py b/nemo_rl/models/megatron/router_replay.py index d81763b33d..e46086ea22 100644 --- a/nemo_rl/models/megatron/router_replay.py +++ b/nemo_rl/models/megatron/router_replay.py @@ -47,6 +47,58 @@ def configure_vllm_for_router_replay(config: PolicyConfig) -> None: vllm_kwargs["enable_return_routed_experts"] = True +def sync_model_config_for_router_replay(model: Any, config: PolicyConfig) -> None: + """Ensure the built mcore model config enables inference routing recording.""" + if not router_replay_enabled(config): + return + model_config = _unwrap_model_config(model) + if model_config is None: + return + model_config.moe_enable_routing_replay = True + # Fused TE top-k bypasses RouterReplay.RECORD; force the replay-capable path. + model_config.moe_router_fusion = False + + +def assert_megatron_inference_router_replay_ready( + model: Any, inference_context: Any, config: PolicyConfig +) -> None: + """Fail fast when router replay is on but the inference engine cannot record routing.""" + if not router_replay_enabled(config): + return + + from megatron.core.transformer.moe.router_replay import RouterReplay + + model_config = getattr(model, "config", None) + if model_config is None or not getattr( + model_config, "moe_enable_routing_replay", False + ): + raise RuntimeError( + "policy.router_replay.enabled=true requires model.config." + "moe_enable_routing_replay=True before Megatron inference starts." + ) + + if getattr(model_config, "num_moe_experts", None) in (None, 0): + raise RuntimeError( + "policy.router_replay.enabled=true requires a MoE model with " + "num_moe_experts > 0." + ) + + if not RouterReplay.global_router_replay_instances: + raise RuntimeError( + "policy.router_replay.enabled=true but no RouterReplay instances " + "were registered on the model. Ensure setup_model_config applied " + "moe_enable_routing_replay before the model was built." + ) + + if getattr(inference_context, "moe_routing_metadata", None) is None: + raise RuntimeError( + "policy.router_replay.enabled=true but the Megatron dynamic " + "inference context has no moe_routing_metadata. The inference " + "engine was likely initialized before moe_enable_routing_replay " + "was enabled; restart the job after enabling router replay." + ) + + def validate_router_replay_config(config: PolicyConfig) -> None: if not router_replay_enabled(config): return @@ -54,8 +106,11 @@ def validate_router_replay_config(config: PolicyConfig) -> None: generation = config.get("generation") or {} megatron_cfg = config.get("megatron_cfg") or {} - if generation.get("backend") != "vllm": - raise ValueError("router_replay.enabled requires vLLM generation.") + generation_backend = generation.get("backend") + if generation_backend not in ("vllm", "megatron"): + raise ValueError( + "router_replay.enabled requires vLLM or Megatron generation." + ) if not megatron_cfg.get("enabled", False): raise ValueError("router_replay.enabled requires the Megatron policy backend.") @@ -99,8 +154,37 @@ def _unwrap_model_config(model: Any) -> Optional[Any]: return None +def _hybrid_moe_layer_numbers(hybrid_layer_pattern: str, num_layers: int) -> list[int]: + # Deferred import: megatron.core is only available inside the Megatron worker venv. + from megatron.core.models.hybrid.hybrid_layer_allocation import ( + Symbols, + parse_hybrid_pattern, + ) + + main_pattern = parse_hybrid_pattern(hybrid_layer_pattern).main_pattern or "" + # '|' marks a pipeline-stage boundary, not a layer; every other symbol occupies one slot. + layer_symbols = [symbol for symbol in main_pattern if symbol != Symbols.PIPE] + if len(layer_symbols) != num_layers: + raise ValueError( + f"hybrid_layer_pattern main segment has {len(layer_symbols)} layers " + f"but num_layers={num_layers} (pattern={hybrid_layer_pattern!r})" + ) + return [ + layer_idx + 1 + for layer_idx, symbol in enumerate(layer_symbols) + if symbol == Symbols.MOE + ] + + def _global_moe_layer_numbers(model_config: Any) -> list[int]: num_layers = int(getattr(model_config, "num_layers")) + + # Hybrid Mamba/attention models place MoE layers via 'E' symbols of hybrid_layer_pattern + # and leave moe_layer_freq at its default of 1, which would wrongly claim every layer is MoE. + hybrid_layer_pattern = getattr(model_config, "hybrid_layer_pattern", None) + if hybrid_layer_pattern: + return _hybrid_moe_layer_numbers(hybrid_layer_pattern, num_layers) + moe_layer_freq = getattr(model_config, "moe_layer_freq", 1) if isinstance(moe_layer_freq, int): @@ -527,3 +611,34 @@ def clear_global_router_replay_instances() -> None: from megatron.core.transformer.moe.router_replay import RouterReplay RouterReplay.clear_global_router_replay_instances() + + +def rebuild_global_router_replay_registry(model: Any) -> None: + """Re-register policy RouterReplay instances after a temporary model build. + + Reference-model setup builds a throwaway Megatron model whose RouterReplay + objects register globally, then clears the global list. Megatron inference + recording uses that global list, so we must restore the live policy model's + instances before generation starts. + """ + from megatron.core.transformer.moe.router_replay import RouterReplay + from megatron.core.utils import unwrap_model + + instances = _router_replay_instances_for_model(unwrap_model(model)) + if not instances: + return + + RouterReplay.clear_global_router_replay_instances() + # Preserve module-walk order to match RouterReplay() instantiation order. + RouterReplay.global_router_replay_instances.extend( + replay_instance for replay_instance, _ in instances + ) + + +def reset_moe_routing_metadata_buffer(inference_context: Any) -> None: + """Drop cached MoE routing CUDA-graph buffers so they resize to the live registry.""" + metadata = getattr(inference_context, "moe_routing_metadata", None) + if metadata is None: + return + metadata.routing_indices_buffer = None + metadata.num_moe_layers = None diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index e8d7f81450..8e1ddeae28 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -140,6 +140,7 @@ def _safe_post_init(self, **kwargs): ) from nemo_rl.models.megatron.router_replay import ( clear_global_router_replay_instances, + rebuild_global_router_replay_registry, router_replay_enabled, validate_router_replay_config, ) @@ -741,9 +742,17 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.moe_permute_fusion = config["megatron_cfg"]["moe_permute_fusion"] + if "use_mamba_mem_eff_path" in config["megatron_cfg"]: + model_cfg.use_mamba_mem_eff_path = config["megatron_cfg"]["use_mamba_mem_eff_path"] + if "moe_grouped_gemm" in config["megatron_cfg"]: model_cfg.moe_grouped_gemm = config["megatron_cfg"]["moe_grouped_gemm"] + + # Enable the dynamic-inference engine's MoE routing recorder when router replay (R3) + # is on, so the Megatron-inference generate path produces routing_indices to replay. model_cfg.moe_enable_routing_replay = router_replay_enabled(config) + if router_replay_enabled(config): + model_cfg.moe_router_fusion = False def _apply_mtp_config(model_cfg: Any, config: PolicyConfig) -> None: @@ -867,6 +876,9 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: ): model_cfg.use_te_rng_tracker = True + if config["megatron_cfg"].get("batch_invariant_mode"): + model_cfg.batch_invariant_mode = True + # FP8 configuration fp8_cfg = config["megatron_cfg"].get("fp8_cfg", None) if fp8_cfg is not None and fp8_cfg.get("enabled", False): @@ -1129,6 +1141,52 @@ def draft_pre_wrap_hook(model: list[MegatronModule]) -> list[MegatronModule]: return draft_pre_wrap_hook +def _enable_batch_invariant_kernels_if_requested(config: PolicyConfig) -> None: + """Enable global batch-invariant kernel patches (cuBLAS workspace shrink, FA num_splits=1, TE GEMM pin). + + Bridge's initialize_megatron() does not call megatron.training.initialize, so we + mirror the legacy ``if args.batch_invariant_mode: enable_batch_invariant_mode()`` + hook here, calling it before initialize_megatron so the patches are active from the + start of the worker lifetime. + """ + if not config["megatron_cfg"].get("batch_invariant_mode"): + return + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + enable_batch_invariant_mode, + ) + + print("Enabling batch invariant mode globally", flush=True) + enable_batch_invariant_mode() + + +def _apply_zero_train_gen_mismatch(config: PolicyConfig) -> None: + """Propagate zero_train_gen_mismatch flag to its constituent sub-knobs. + + When True, forces batch_invariant_mode=True, use_mamba_mem_eff_path=False, + and defaults env vars for batch-invariant TE/cuBLAS/MoE/Mamba kernels + if not already set by the environment. + Router replay and moe_grouped_gemm must be configured explicitly. + """ + if not config.get("megatron_cfg", {}).get("zero_train_gen_mismatch"): + return + import os + + mc = config["megatron_cfg"] + mc["batch_invariant_mode"] = True + mc.setdefault("use_mamba_mem_eff_path", False) + # Default to cuBLAS workspace shrink (te_gemm_cublas_pinned) so the flag + # delivers batch-invariant TE GEMM without requiring the env var to be set. + os.environ.setdefault("NRL_BI_KERNELS", "te_gemm_cublas_pinned") + # Starve PyTorch's own cuBLAS workspace so non-TE aten::mm/addmm paths also + # pick workspace-free (splitK=1, reduction=NONE) algorithms. + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":0:0") + os.environ.setdefault("CUBLASLT_WORKSPACE_SIZE", "0") + # Force fixed-order MoE unpermute+index_put to eliminate scatter_add nondeterminism. + os.environ.setdefault("NRL_FIXED_ORDER_MOE_COMBINE", "1") + # Pin Triton autotune config for Mamba SSM/conv kernels across train and gen. + os.environ.setdefault("MAMBA_DETERMINISTIC", "1") + + _BRIDGE_SIGNAL_HANDLER_PATCHED = False @@ -1170,6 +1228,8 @@ def setup_model_and_optimizer( ): state = GlobalState() _patch_bridge_signal_handler_for_worker_threads() + _apply_zero_train_gen_mismatch(policy_cfg) + _enable_batch_invariant_kernels_if_requested(policy_cfg) state.cfg = megatron_cfg # TODO: Freeze state.cfg @@ -1475,6 +1535,7 @@ def setup_reference_model_state( megatron_cfg: ConfigContainer, pretrained_path: str, pre_load_checkpoint_hook: Optional[Callable] = None, + policy_model: Optional[Any] = None, ) -> dict: """Setup the reference model for inference and return its state dict.""" # Create reference checkpoint config @@ -1604,7 +1665,10 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: else: print("Reference model not loaded") finally: - clear_global_router_replay_instances() + if policy_model is not None: + rebuild_global_router_replay_registry(policy_model) + else: + clear_global_router_replay_instances() return reference_state_dict diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 55f511bb56..773d20a027 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -354,6 +354,16 @@ class MegatronConfig(TypedDict): clear_memory_caches_before_refit: NotRequired[bool] # FP8 quantization settings for the Megatron training backend. fp8_cfg: NotRequired[Fp8Config] + # Use batch-invariant kernels (cuBLAS workspace shrink, FA num_splits=1, TE GEMM pin) + # for deterministic execution regardless of batch size. Required for zero-KL. + batch_invariant_mode: NotRequired[bool] + # Disable the memory-efficient Mamba SSM path so training and generation kernels match. + # Required for zero-KL on hybrid Mamba models. + use_mamba_mem_eff_path: NotRequired[bool] + # Master switch: when True, forces batch_invariant_mode=True and + # use_mamba_mem_eff_path=False to eliminate train/gen KL mismatch sources. + # Router replay and moe_grouped_gemm must be configured explicitly. + zero_train_gen_mismatch: NotRequired[bool] class DraftConfigDisabled(TypedDict): diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 64bd8d4788..0cbdcae248 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -814,10 +814,23 @@ def generate( output_is_replicated=["tensor_parallel", "pipeline_parallel"], common_kwargs={"greedy": greedy}, ) - result = BatchedDataDict.from_batches( - self.worker_group.get_all_worker_results(futures), - pad_value_dict={"output_ids": self.cfg["generation"]["_pad_token_id"]}, - ) + assert self.cfg["generation"] is not None, "Generation config is not set" + worker_batches = self.worker_group.get_all_worker_results(futures) + if self.cfg["generation"]["backend"] == "megatron": + # Coordinator-based Megatron inference: only the DP=0 submitter returns + # tensors; other ranks return empty shells. Take the submitter's result + # directly — from_batches would also flatten any ndim>3 output tensors + # (e.g. routed_experts [B, S, L, K] for router replay). + result: BatchedDataDict[GenerationOutputSpec] = next( + wb + for wb in worker_batches + if wb.get("output_ids") is not None and wb["output_ids"].numel() > 0 + ) + else: + result = BatchedDataDict.from_batches( + worker_batches, + pad_value_dict={"output_ids": self.cfg["generation"]["_pad_token_id"]}, + ) required_keys = [ "output_ids", diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 477f1c65ee..a3e3f7ab95 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -50,7 +50,11 @@ from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.named_sharding import NamedSharding -from nemo_rl.models.generation.interfaces import GenerationDatumSpec +from nemo_rl.models.generation.interfaces import ( + GenerationDatumSpec, + GenerationOutputSpec, +) +from nemo_rl.models.megatron.moe_routing_record import build_routed_experts_batch from nemo_rl.models.generation.megatron.megatron_worker import ( MegatronGenerationMixin, MegatronGenerationRefitMixin, @@ -400,6 +404,7 @@ def __init__( pre_load_checkpoint_hook=getattr( self, "_pre_load_checkpoint_hook", None ), + policy_model=self.model, ) self.model = self.move_model(self.model, "cuda") log_gpu_memory_diagnostics( @@ -1568,6 +1573,48 @@ def use_reference_model(self): if self.should_disable_forward_pre_hook: self.enable_forward_pre_hook() + def _parse_result_to_batched_data_dict( + self, + data: BatchedDataDict[GenerationDatumSpec], + result: list, + ) -> BatchedDataDict[GenerationOutputSpec]: + """Pack inference results, additionally recording MoE routing for R3 replay. + + Extends the mixin's packing with the Megatron-inference router-replay recorder. + The dynamic-inference engine exposes per-sample ``routing_indices`` (numpy + ``[tokens, layers, topk]``); we convert it into the ``routed_experts`` + ``[batch, seq, layers, topk]`` tensor that the train/logprob forward replays + (same key/shape the vLLM path already produces). No-op for dense models or when + router replay is disabled / the engine recorded no routing. + """ + out = super()._parse_result_to_batched_data_dict(data, result) + if router_replay_enabled(self.cfg): + routing_per_sample = [ + getattr(result[i], "routing_indices", None) + for i in range(len(result)) + ] + routed_experts = build_routed_experts_batch( + [ + torch.from_numpy(r) if r is not None else None + for r in routing_per_sample + ], + out["unpadded_sequence_lengths"], + out["output_ids"].shape[1], + ) + if routed_experts is None: + missing = sum(r is None for r in routing_per_sample) + raise RuntimeError( + "policy.router_replay.enabled=true requires Megatron " + "inference to return routing_indices for every generated " + f"sample, but {missing}/{len(routing_per_sample)} samples " + "were missing that field. This usually means " + "moe_enable_routing_replay was not enabled on the model " + "config before the inference engine was initialized." + ) + # Set after super()'s from_batches: it flattens 4D tensors into 1D rows. + out["routed_experts"] = routed_experts + return out + @wrap_with_nvtx_name("megatron_policy_worker/get_topk_logits") def get_topk_logits( self, diff --git a/pyproject.toml b/pyproject.toml index 07796fe0c4..b0408acbcb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -340,7 +340,7 @@ link-mode = "copy" # (all in opencv_python_headless.libs/libav*.so shipped by 4.x wheels, which bundle FFmpeg 5.1.6). # The timm override is needed because sglang requires timm==1.0.16. override-dependencies = [ - "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.15", + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.17.0", "nvidia-cublas==13.5.1.27; sys_platform != 'darwin'", "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", "nvidia-cudnn-frontend==1.23.0", @@ -493,12 +493,12 @@ requires-dist = ["torch", "packaging", "ninja"] [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.15.0+42b8400" +version = "2.17.0" requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"] [[tool.uv.dependency-metadata]] name = "transformer-engine-torch" -version = "2.15.0+42b8400" +version = "2.17.0" requires-dist = ["torch", "transformer-engine"] [[tool.uv.dependency-metadata]] From 8db6804400cd820c2203784a7bee49a57df908e8 Mon Sep 17 00:00:00 2001 From: Yigong Qin Date: Tue, 21 Jul 2026 21:30:52 -0700 Subject: [PATCH 3/5] add fa4 to mcore venv Signed-off-by: YigongQin --- nemo_rl/models/megatron/setup.py | 9 ++++++--- pyproject.toml | 8 +++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 8e1ddeae28..292009c287 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -283,6 +283,8 @@ def validate_and_set_config( f"policy.hf_config_overrides.rope_scaling." ) + _apply_zero_train_gen_mismatch(config) + megatron_cfg, model_cfg = setup_model_config( config, rank, @@ -1163,9 +1165,9 @@ def _apply_zero_train_gen_mismatch(config: PolicyConfig) -> None: """Propagate zero_train_gen_mismatch flag to its constituent sub-knobs. When True, forces batch_invariant_mode=True, use_mamba_mem_eff_path=False, - and defaults env vars for batch-invariant TE/cuBLAS/MoE/Mamba kernels - if not already set by the environment. - Router replay and moe_grouped_gemm must be configured explicitly. + attention_backend=flash (FA4 via TE), and defaults env vars for + batch-invariant TE/cuBLAS/MoE/Mamba kernels if not already set by the + environment. Router replay and moe_grouped_gemm must be configured explicitly. """ if not config.get("megatron_cfg", {}).get("zero_train_gen_mismatch"): return @@ -1174,6 +1176,7 @@ def _apply_zero_train_gen_mismatch(config: PolicyConfig) -> None: mc = config["megatron_cfg"] mc["batch_invariant_mode"] = True mc.setdefault("use_mamba_mem_eff_path", False) + mc.setdefault("attention_backend", "flash") # Default to cuBLAS workspace shrink (te_gemm_cublas_pinned) so the flag # delivers batch-invariant TE GEMM without requiring the env var to be set. os.environ.setdefault("NRL_BI_KERNELS", "te_gemm_cublas_pinned") diff --git a/pyproject.toml b/pyproject.toml index b0408acbcb..c165f9fbda 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -182,6 +182,8 @@ mcore = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'", "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')", + # TE v2.15 Megatron attention_backend=flash requires flash_attn.cute (FA4). + "flash-attn-4==4.0.0b13 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", "mamba-ssm", "causal-conv1d", "fast-hadamard-transform", @@ -340,7 +342,7 @@ link-mode = "copy" # (all in opencv_python_headless.libs/libav*.so shipped by 4.x wheels, which bundle FFmpeg 5.1.6). # The timm override is needed because sglang requires timm==1.0.16. override-dependencies = [ - "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.17.0", + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.15", "nvidia-cublas==13.5.1.27; sys_platform != 'darwin'", "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", "nvidia-cudnn-frontend==1.23.0", @@ -493,12 +495,12 @@ requires-dist = ["torch", "packaging", "ninja"] [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.17.0" +version = "2.15.0+42b8400" requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"] [[tool.uv.dependency-metadata]] name = "transformer-engine-torch" -version = "2.17.0" +version = "2.15.0+42b8400" requires-dist = ["torch", "transformer-engine"] [[tool.uv.dependency-metadata]] From 5c1819756988e60bb5e1bebb89514123a7933880 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Wed, 22 Jul 2026 16:28:47 -0700 Subject: [PATCH 4/5] move patch to nemorl Signed-off-by: YigongQin --- .../Megatron-Bridge-workspace/Megatron-Bridge | 2 +- nemo_rl/models/megatron/setup.py | 18 +++-- nemo_rl/models/policy/workers/patches.py | 76 +++++++++++++++++++ 3 files changed, 89 insertions(+), 7 deletions(-) diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 7767a134a2..554c7b9324 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 7767a134a278b0adf4e2d378ed158b87bd743e15 +Subproject commit 554c7b9324225aa863eee52e8b8fdde7abced2b1 diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 292009c287..31843679fc 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1150,7 +1150,12 @@ def _enable_batch_invariant_kernels_if_requested(config: PolicyConfig) -> None: mirror the legacy ``if args.batch_invariant_mode: enable_batch_invariant_mode()`` hook here, calling it before initialize_megatron so the patches are active from the start of the worker lifetime. + + ``zero_train_gen_mismatch`` applies TE cuBLAS workspace shrink via + ``patches.apply_te_gemm_cublas_pinned_patch`` instead of Megatron BIK. """ + if config.get("megatron_cfg", {}).get("zero_train_gen_mismatch"): + return if not config["megatron_cfg"].get("batch_invariant_mode"): return from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( @@ -1165,21 +1170,22 @@ def _apply_zero_train_gen_mismatch(config: PolicyConfig) -> None: """Propagate zero_train_gen_mismatch flag to its constituent sub-knobs. When True, forces batch_invariant_mode=True, use_mamba_mem_eff_path=False, - attention_backend=flash (FA4 via TE), and defaults env vars for - batch-invariant TE/cuBLAS/MoE/Mamba kernels if not already set by the - environment. Router replay and moe_grouped_gemm must be configured explicitly. + attention_backend=flash (FA4 via TE), applies TE cuBLAS workspace shrink via + patches.py, and defaults env vars for cuBLAS/MoE/Mamba determinism if not + already set by the environment. Router replay and moe_grouped_gemm must be + configured explicitly. """ if not config.get("megatron_cfg", {}).get("zero_train_gen_mismatch"): return import os + from nemo_rl.models.policy.workers.patches import apply_te_gemm_cublas_pinned_patch + mc = config["megatron_cfg"] mc["batch_invariant_mode"] = True mc.setdefault("use_mamba_mem_eff_path", False) mc.setdefault("attention_backend", "flash") - # Default to cuBLAS workspace shrink (te_gemm_cublas_pinned) so the flag - # delivers batch-invariant TE GEMM without requiring the env var to be set. - os.environ.setdefault("NRL_BI_KERNELS", "te_gemm_cublas_pinned") + apply_te_gemm_cublas_pinned_patch() # Starve PyTorch's own cuBLAS workspace so non-TE aten::mm/addmm paths also # pick workspace-free (splitK=1, reduction=NONE) algorithms. os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":0:0") diff --git a/nemo_rl/models/policy/workers/patches.py b/nemo_rl/models/policy/workers/patches.py index 5a0d5b0ab8..a7ce9fe9c6 100644 --- a/nemo_rl/models/policy/workers/patches.py +++ b/nemo_rl/models/policy/workers/patches.py @@ -12,8 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib import os from importlib.util import find_spec +from typing import Callable, Optional + +# Original TE cuBLAS workspace sizer, saved by apply_te_gemm_cublas_pinned_patch(). +_TE_CUBLAS_WS_SIZE_FN_ORIG: Optional[Callable[[], int]] = None + +# Minimum workspace that satisfies TE's NVFP4 alpha-scratch guard in cublaslt_gemm.cu. +_TE_CUBLAS_WS_PINNED_BYTES: int = 4 def _get_transformer_engine_file(relative_path: str) -> str: @@ -104,3 +112,71 @@ def apply_transformer_engine_patch(): except Exception as e: print(f"Error checking/patching transformer_engine: {e}") + + +def apply_te_gemm_cublas_pinned_patch( + target_bytes: int = _TE_CUBLAS_WS_PINNED_BYTES, +) -> None: + """Shrink TE's cuBLAS workspace so cuBLASLt picks workspace-free algorithms. + + Mirrors megatron.core.transformer.custom_layers.batch_invariant_kernels. + ``_shrink_te_cublas_workspace_for_invariance``. Intended for zero-KL / + ``zero_train_gen_mismatch`` only — call from ``_apply_zero_train_gen_mismatch`` + in setup.py, not from generic batch-invariant mode. + """ + global _TE_CUBLAS_WS_SIZE_FN_ORIG + if _TE_CUBLAS_WS_SIZE_FN_ORIG is not None: + return + try: + te_gemm_mod = importlib.import_module( + "transformer_engine.pytorch.cpp_extensions.gemm" + ) + except ImportError: + print( + "te_gemm_cublas_pinned: transformer_engine.pytorch.cpp_extensions.gemm " + "is not importable; skipping workspace shrink." + ) + return + if not hasattr(te_gemm_mod, "get_cublas_workspace_size_bytes"): + print( + "te_gemm_cublas_pinned: TE gemm module has no get_cublas_workspace_size_bytes " + "(TE version mismatch?); skipping workspace shrink." + ) + return + + _TE_CUBLAS_WS_SIZE_FN_ORIG = te_gemm_mod.get_cublas_workspace_size_bytes + te_gemm_mod.get_cublas_workspace_size_bytes = lambda: int(target_bytes) + ws_fn = getattr(te_gemm_mod, "get_cublas_workspace", None) + if ws_fn is not None and hasattr(ws_fn, "cache_clear"): + try: + ws_fn.cache_clear() + except Exception: # pylint: disable=broad-except + pass + print( + f"[zero_train_gen_mismatch] shrunk TE cuBLAS workspace to {target_bytes} bytes " + "(te_gemm_cublas_pinned via patches.py). " + "Set CUBLASLT_LOG_LEVEL=5 to verify cuBLASLt picks a stable algo across batch sizes." + ) + + +def restore_te_gemm_cublas_pinned_patch() -> None: + """Restore TE's original cuBLAS workspace sizer (for tests).""" + global _TE_CUBLAS_WS_SIZE_FN_ORIG + if _TE_CUBLAS_WS_SIZE_FN_ORIG is None: + return + try: + te_gemm_mod = importlib.import_module( + "transformer_engine.pytorch.cpp_extensions.gemm" + ) + except ImportError: + _TE_CUBLAS_WS_SIZE_FN_ORIG = None + return + if hasattr(te_gemm_mod, "get_cublas_workspace_size_bytes"): + te_gemm_mod.get_cublas_workspace_size_bytes = _TE_CUBLAS_WS_SIZE_FN_ORIG + ws_fn = getattr(te_gemm_mod, "get_cublas_workspace", None) + if ws_fn is not None and hasattr(ws_fn, "cache_clear"): + try: + ws_fn.cache_clear() + except Exception: # pylint: disable=broad-except + pass + _TE_CUBLAS_WS_SIZE_FN_ORIG = None From 34c844a49e4fbd3002d2f849f029aaee8cde6886 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Thu, 23 Jul 2026 15:29:03 -0700 Subject: [PATCH 5/5] setting up zero-train-mismatch for 1.5B and 30B Signed-off-by: YigongQin --- nemo_rl/models/generation/megatron/config.py | 3 + nemo_rl/models/megatron/setup.py | 26 +- .../policy/workers/moe_determinism_patches.py | 334 ++++++++++++++++++ .../.env.template | 8 + .../README.md | 40 +++ .../run_qwen1.5b_zero_kl_precision.sh | 254 +++++++++++++ .../run_qwen30ba3b_zero_kl_precision.sh | 295 ++++++++++++++++ .../models/megatron/test_megatron_setup.py | 31 ++ .../policy/test_moe_determinism_patches.py | 95 +++++ tests/unit/models/policy/test_patches.py | 57 +++ 10 files changed, 1134 insertions(+), 9 deletions(-) create mode 100644 nemo_rl/models/policy/workers/moe_determinism_patches.py create mode 100644 research/megatron-inference-true-on-policy/.env.template create mode 100644 research/megatron-inference-true-on-policy/README.md create mode 100755 research/megatron-inference-true-on-policy/run_qwen1.5b_zero_kl_precision.sh create mode 100755 research/megatron-inference-true-on-policy/run_qwen30ba3b_zero_kl_precision.sh create mode 100644 tests/unit/models/policy/test_moe_determinism_patches.py diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 7cfb88cc1f..8ce94699c3 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -37,6 +37,9 @@ class MCoreGenerationSpecificArgs(TypedDict): num_cuda_graphs: int use_cuda_graphs_for_non_decode_steps: bool cuda_graph_impl: str + # Layer spec used by Megatron generation. + # Options are "transformer_engine" and "inference_optimized". + transformer_impl: NotRequired[Literal["transformer_engine", "inference_optimized"]] # Inference CUDA-graph scope. Options: # - 'none': inference runs in eager mode (no CUDA graphs). # - 'layer': graphs are owned at the per-layer boundary (TransformerLayer / MambaLayer). diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 31843679fc..9edea94a23 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1170,30 +1170,38 @@ def _apply_zero_train_gen_mismatch(config: PolicyConfig) -> None: """Propagate zero_train_gen_mismatch flag to its constituent sub-knobs. When True, forces batch_invariant_mode=True, use_mamba_mem_eff_path=False, - attention_backend=flash (FA4 via TE), applies TE cuBLAS workspace shrink via - patches.py, and defaults env vars for cuBLAS/MoE/Mamba determinism if not - already set by the environment. Router replay and moe_grouped_gemm must be - configured explicitly. + attention_backend=flash (FA4 via TE), and the Transformer Engine generation + layer spec so generation and scoring share the patched MoE unpermute path. + Also applies TE cuBLAS workspace shrink via patches.py, MoE fixed-order + unpermute via moe_determinism_patches.py, and defaults env vars for + cuBLAS/MoE determinism if not already set by the environment. Router + replay and moe_grouped_gemm must be configured explicitly. """ if not config.get("megatron_cfg", {}).get("zero_train_gen_mismatch"): return import os + from nemo_rl.models.policy.workers.moe_determinism_patches import ( + apply_moe_unpermute_determinism_patch, + ) from nemo_rl.models.policy.workers.patches import apply_te_gemm_cublas_pinned_patch mc = config["megatron_cfg"] mc["batch_invariant_mode"] = True - mc.setdefault("use_mamba_mem_eff_path", False) mc.setdefault("attention_backend", "flash") + generation = config.get("generation") + if generation is not None and "mcore_generation_config" in generation: + generation["mcore_generation_config"]["transformer_impl"] = ( + "transformer_engine" + ) apply_te_gemm_cublas_pinned_patch() + apply_moe_unpermute_determinism_patch() # Starve PyTorch's own cuBLAS workspace so non-TE aten::mm/addmm paths also # pick workspace-free (splitK=1, reduction=NONE) algorithms. os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":0:0") os.environ.setdefault("CUBLASLT_WORKSPACE_SIZE", "0") - # Force fixed-order MoE unpermute+index_put to eliminate scatter_add nondeterminism. - os.environ.setdefault("NRL_FIXED_ORDER_MOE_COMBINE", "1") - # Pin Triton autotune config for Mamba SSM/conv kernels across train and gen. - os.environ.setdefault("MAMBA_DETERMINISTIC", "1") + # Force fixed-order MoE unpermute to eliminate scatter_add nondeterminism. + _BRIDGE_SIGNAL_HANDLER_PATCHED = False diff --git a/nemo_rl/models/policy/workers/moe_determinism_patches.py b/nemo_rl/models/policy/workers/moe_determinism_patches.py new file mode 100644 index 0000000000..ad645a7e4b --- /dev/null +++ b/nemo_rl/models/policy/workers/moe_determinism_patches.py @@ -0,0 +1,334 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + + +from __future__ import annotations + +import logging +from typing import Any, Callable, Dict, Optional, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine + from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, + ) + +_NRL_LOGGER = logging.getLogger(__name__) + +# One-shot guard so the unpermute-combine-path diagnostic surfaces which combine +# branch actually executes (fused vs fixed-order vs scatter_add) without flooding. +_NRL_UNPERMUTE_PATH_SEEN: set[str] = set() + +_UNPERMUTE_ORIG: Optional[Callable[..., torch.Tensor]] = None +_TOKEN_DISPATCHER_UNPERMUTE_ORIG: Optional[Callable[..., torch.Tensor]] = None +_DYNAMIC_STEP_BOOKKEEPING_ORIG: Optional[Callable[..., Dict[str, Any]]] = None +_ASYNC_BOOKKEEP_ORIG: Optional[Callable[..., Any]] = None + +_MOE_UNPERMUTE_PATCHED = False +_ROUTER_REPLAY_INFERENCE_PATCHED = False + + +def _nrl_log_unpermute_path(path: str) -> None: + if path not in _NRL_UNPERMUTE_PATH_SEEN: + _NRL_UNPERMUTE_PATH_SEEN.add(path) + _NRL_LOGGER.warning("[moe-combine] unpermute executed via '%s'", path) + + +def _unpermute_fixed_order_combine( + permuted_tokens: torch.Tensor, + sorted_indices: torch.Tensor, + restore_shape: torch.Size, +) -> torch.Tensor: + """Sum expert outputs per token in stable (permute) order via [T, max_slots, H].sum(1). + + Avoids atomic ``scatter_add_`` / ``index_add_``. Uses the same accumulation dtype as + ``scatter_add_`` (``permuted_tokens.dtype`` after gating weights). + """ + num_tokens, hidden = restore_shape + num_permuted = permuted_tokens.size(0) + if num_permuted == 0: + return torch.zeros( + restore_shape, dtype=permuted_tokens.dtype, device=permuted_tokens.device + ) + + sort_perm = torch.argsort(sorted_indices, stable=True) + dest = sorted_indices[sort_perm] + vals = permuted_tokens[sort_perm] + + seq = torch.arange(num_permuted, device=permuted_tokens.device, dtype=torch.long) + if num_permuted > 1: + change = dest.new_ones(num_permuted, dtype=torch.bool) + change[1:] = dest[1:] != dest[:-1] + else: + change = dest.new_ones(1, dtype=torch.bool) + group_id = change.long().cumsum(0) - 1 + num_groups = int(group_id[-1].item()) + 1 + group_sizes = torch.bincount(group_id, minlength=num_groups) + starts = torch.zeros(num_groups, dtype=torch.long, device=permuted_tokens.device) + if num_groups > 1: + starts[1:] = group_sizes.cumsum(0)[:-1] + slot = seq - starts[group_id] + max_slots = int(group_sizes.max().item()) + + contrib = torch.zeros( + num_tokens, max_slots, hidden, dtype=permuted_tokens.dtype, device=permuted_tokens.device + ) + contrib[dest, slot] = vals + return contrib.sum(dim=1) + + +def _patched_unpermute( + permuted_tokens: torch.Tensor, + sorted_indices: torch.Tensor, + restore_shape: torch.Size, + probs: Optional[torch.Tensor] = None, + routing_map: Optional[torch.Tensor] = None, + fused: bool = False, + drop_and_pad: bool = False, + pad_offsets: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """``moe_utils.unpermute`` with fixed-order combine.""" + import megatron.core.transformer.moe.moe_utils as moe_utils + + if fused: + if not moe_utils.HAVE_TE or moe_utils.fused_unpermute is None: + raise ValueError("fused_unpermute is not available. Please install TE >= 2.1.0.") + _nrl_log_unpermute_path("fused_unpermute") + extra_kwargs = {} + if moe_utils.is_te_min_version("2.12.0"): + extra_kwargs["pad_offsets"] = pad_offsets + return moe_utils.fused_unpermute( + permuted_tokens, + sorted_indices, + merging_probs=probs, + restore_shape=restore_shape, + **extra_kwargs, + ) + + input_dtype = permuted_tokens.dtype + + if probs is not None: + assert routing_map is not None, "Mask must be provided to permute the probs." + if drop_and_pad: + num_experts = routing_map.size(1) + num_permuted_tokens = sorted_indices.size(0) + capacity = num_permuted_tokens // num_experts + num_unpermuted_tokens = probs.size(0) + + probs_T_1D = probs.T.contiguous().view(-1) + indices_dim0 = torch.arange(num_experts, device=routing_map.device).unsqueeze(-1) + indices_dim1 = sorted_indices.view(num_experts, capacity) + indices_1D = (indices_dim0 * num_unpermuted_tokens + indices_dim1).view(-1) + permuted_probs = probs_T_1D.index_select(0, indices_1D) + else: + permuted_probs = probs.T.contiguous().masked_select(routing_map.T.contiguous()) + permuted_tokens = permuted_tokens * permuted_probs.unsqueeze(-1) + + _nrl_log_unpermute_path("fixed_order_combine") + output_tokens = _unpermute_fixed_order_combine( + permuted_tokens, sorted_indices, restore_shape + ) + return output_tokens.to(dtype=input_dtype) + + +def _nrl_dynamic_step_context_bookkeeping(self: "TextGenerationController") -> Dict[str, Any]: + """Early MoE routing reconstruction before KV blocks are released (a6e829b).""" + from torch.cuda.nvtx import range_pop, range_push + + context = self.inference_wrapped_model.inference_context + active_request_count = context.total_request_count - context.paused_request_count + active_request_slice = slice(context.paused_request_count, context.total_request_count) + + range_push("transfer_samples_to_cpu") + sampled_tokens_cpu, sampled_mtp_tokens_cpu = self._transfer_samples_to_cpu( + active_request_count + ) + range_pop() + + range_push("active_request_mask") + active_request_ids = context.request_ids[active_request_slice].long() + active_sequence_lengths = context.get_active_sequence_lengths() + active_sequence_lengths += 1 + max_sequence_lengths = context.get_max_sequence_lengths() + + active_request_mask = ( + sampled_tokens_cpu + != context.active_request_metadata["termination_id"][:active_request_count] + ).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte() + + if self._get_stop_word_finished_ids_callback is not None: + request_ids_list = active_request_ids.tolist() + stop_word_finished_ids = self._get_stop_word_finished_ids_callback(request_ids_list) + if stop_word_finished_ids: + for idx, request_id in enumerate(request_ids_list): + if request_id in stop_word_finished_ids: + active_request_mask[idx] = 0 + + finished_idxs = ( + torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count + ) + finished_request_ids = context.request_ids[finished_idxs] + + finished_routing_indices: Dict[int, Any] = {} + if context.moe_enable_routing_replay and finished_idxs.numel() > 0: + for fidx in finished_idxs.tolist(): + req_id = int(context.request_ids[fidx].item()) + blocks = context.request_to_kv_block_ids[fidx] + valid = blocks[blocks >= 0].tolist() + if not valid: + continue + total_tokens = int( + active_sequence_lengths[fidx - context.paused_request_count].item() + ) + routing = context.kv_block_allocator.reconstruct_routing_from_blocks( + valid, total_tokens - 1 + ) + if routing is not None: + finished_routing_indices[req_id] = routing + + new_sample_copy = sampled_tokens_cpu.clone() + range_pop() + + range_push("update_requests") + update_result = context.update_requests( + active_request_mask, new_sample_copy, sampled_mtp_tokens_cpu + ) + range_pop() + + return { + "active_request_ids": active_request_ids, + "finished_request_ids": finished_request_ids, + "sample": sampled_tokens_cpu, + "finished_routing_indices": finished_routing_indices, + **(update_result or {}), + } + + +async def _nrl_async_bookkeep( + self: "DynamicInferenceEngine", + step_result: Optional[Dict[str, Any]], + context_state: Dict[str, Any], + step_time: float, +): + """Apply pre-reconstructed routing before upstream post_process (a6e829b).""" + if step_result is not None: + finished_routing_indices = step_result.get("finished_routing_indices") + if finished_routing_indices: + for request_id, routing in finished_routing_indices.items(): + if request_id in self.requests: + self.get_request(request_id).routing_indices = routing + step_result = dict(step_result) + step_result.pop("finished_routing_block_ids", None) + assert _ASYNC_BOOKKEEP_ORIG is not None + return await _ASYNC_BOOKKEEP_ORIG(self, step_result, context_state, step_time) + + +def apply_moe_unpermute_determinism_patch() -> None: + """Patch MoE unpermute and the token dispatcher's cached import.""" + global _UNPERMUTE_ORIG, _TOKEN_DISPATCHER_UNPERMUTE_ORIG, _MOE_UNPERMUTE_PATCHED + if _MOE_UNPERMUTE_PATCHED: + return + try: + import megatron.core.transformer.moe.moe_utils as moe_utils + import megatron.core.transformer.moe.token_dispatcher as token_dispatcher + except ImportError: + print( + "moe_determinism_patches: Megatron MoE modules are not importable; " + "skipping unpermute patch." + ) + return + + _UNPERMUTE_ORIG = moe_utils.unpermute + _TOKEN_DISPATCHER_UNPERMUTE_ORIG = token_dispatcher.unpermute + moe_utils.unpermute = _patched_unpermute + # token_dispatcher imports unpermute by value, so changing only the source + # module leaves its already-bound call site on the original implementation. + token_dispatcher.unpermute = _patched_unpermute + _MOE_UNPERMUTE_PATCHED = True + print( + "[moe_determinism_patches] patched moe_utils.unpermute and " + "token_dispatcher.unpermute with fixed-order combine." + ) + + +def apply_router_replay_inference_patches() -> None: + """Patch dynamic inference for early router-replay reconstruction (a6e829b).""" + global _DYNAMIC_STEP_BOOKKEEPING_ORIG, _ASYNC_BOOKKEEP_ORIG, _ROUTER_REPLAY_INFERENCE_PATCHED + if _ROUTER_REPLAY_INFERENCE_PATCHED: + return + try: + from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine + from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, + ) + except ImportError: + print( + "moe_determinism_patches: Megatron inference modules are not importable; " + "skipping router-replay inference patches." + ) + return + + _DYNAMIC_STEP_BOOKKEEPING_ORIG = TextGenerationController._dynamic_step_context_bookkeeping + TextGenerationController._dynamic_step_context_bookkeeping = ( + _nrl_dynamic_step_context_bookkeeping + ) + + _ASYNC_BOOKKEEP_ORIG = DynamicInferenceEngine.async_bookkeep + DynamicInferenceEngine.async_bookkeep = _nrl_async_bookkeep + + _ROUTER_REPLAY_INFERENCE_PATCHED = True + print( + "[moe_determinism_patches] patched TextGenerationController._dynamic_step_context_bookkeeping " + "and DynamicInferenceEngine.async_bookkeep for early router-replay routing reconstruction." + ) + + +def apply_moe_determinism_patches(*, router_replay: bool = False) -> None: + """Apply all requested MoE determinism / router-replay runtime patches.""" + apply_moe_unpermute_determinism_patch() + if router_replay: + apply_router_replay_inference_patches() + + +def restore_moe_determinism_patches() -> None: + """Restore Megatron entry points patched by this module (for tests).""" + global _MOE_UNPERMUTE_PATCHED, _ROUTER_REPLAY_INFERENCE_PATCHED + + if _MOE_UNPERMUTE_PATCHED and _UNPERMUTE_ORIG is not None: + import megatron.core.transformer.moe.moe_utils as moe_utils + import megatron.core.transformer.moe.token_dispatcher as token_dispatcher + + moe_utils.unpermute = _UNPERMUTE_ORIG + if _TOKEN_DISPATCHER_UNPERMUTE_ORIG is not None: + token_dispatcher.unpermute = _TOKEN_DISPATCHER_UNPERMUTE_ORIG + _MOE_UNPERMUTE_PATCHED = False + + if _ROUTER_REPLAY_INFERENCE_PATCHED: + from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine + from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, + ) + + if _DYNAMIC_STEP_BOOKKEEPING_ORIG is not None: + TextGenerationController._dynamic_step_context_bookkeeping = ( + _DYNAMIC_STEP_BOOKKEEPING_ORIG + ) + if _ASYNC_BOOKKEEP_ORIG is not None: + DynamicInferenceEngine.async_bookkeep = _ASYNC_BOOKKEEP_ORIG + _ROUTER_REPLAY_INFERENCE_PATCHED = False + + _NRL_UNPERMUTE_PATH_SEEN.clear() diff --git a/research/megatron-inference-true-on-policy/.env.template b/research/megatron-inference-true-on-policy/.env.template new file mode 100644 index 0000000000..a7bd18b401 --- /dev/null +++ b/research/megatron-inference-true-on-policy/.env.template @@ -0,0 +1,8 @@ +# Copy to .env and fill in your values. +RL_DIR=/path/to/RL +CONTAINER_IMAGE=/path/to/nemo_rl_v0.7.0.sqsh +HF_TOKEN=hf_your_token_here +HF_HOME=/path/to/huggingface/cache +WANDB_API_KEY=your_wandb_api_key_here +WANDB_ENTITY=your_wandb_entity_here +WANDB_PROJECT=your_wandb_project_here diff --git a/research/megatron-inference-true-on-policy/README.md b/research/megatron-inference-true-on-policy/README.md new file mode 100644 index 0000000000..8534466bdb --- /dev/null +++ b/research/megatron-inference-true-on-policy/README.md @@ -0,0 +1,40 @@ +# Megatron-Inference True On-Policy Study + +True on-policy RL training using Megatron-Inference to minimize training-generation mismatch (`gen_kl_error` -> 0). + + +## Setup + +```bash +# 1. Clone the RL repo +git clone -b yigongq/minf-onpolicy git@github.com:YigongQin/RL.git +cd RL && git submodule update --init --recursive + +# 2. Download the container (one-time) +enroot import --output nemo_rl_v0.7.0.sqsh 'docker://nvcr.io#nvidia/nemo-rl:v0.7.0' +# This creates nemo_rl_v0.7.0.sqsh in the current directory. +# Move it to your preferred location and update CONTAINER_IMAGE in the scripts. + +# 3. Create .env with your config +cp .env.template .env +# Edit .env: set RL_DIR, CONTAINER_IMAGE, HF_TOKEN, HF_HOME, WANDB_API_KEY, WANDB_PROJECT +``` + +## Run + +```bash +# Qwen2.5-1.5B (1 node, 8GPU, TP=1) +# ZERO_TRAIN_GEN_MISMATCH is set to true by default +sbatch --export=PRECISION=bf16 run_qwen1.5b_zero_kl_precision.sh +sbatch --export=PRECISION=bf16,ZERO_TRAIN_GEN_MISMATCH=false run_qwen1.5b_zero_kl_precision.sh +sbatch --export=PRECISION=mxfp8 run_qwen1.5b_zero_kl_precision.sh +sbatch --export=PRECISION=mxfp8,ZERO_TRAIN_GEN_MISMATCH=false run_qwen1.5b_zero_kl_precision.sh + + +# Qwen3-30B-A3B (1 node, 8GPU, TP=1) +sbatch --export=PRECISION=bf16 run_qwen30ba3b_zero_kl_precision.sh +sbatch --export=PRECISION=mxfp8 run_qwen30ba3b_zero_kl_precision.sh +sbatch --export=PRECISION=bf16,ZERO_TRAIN_GEN_MISMATCH=false run_qwen30ba3b_zero_kl_precision.sh +sbatch --export=PRECISION=mxfp8,ZERO_TRAIN_GEN_MISMATCH=false run_qwen30ba3b_zero_kl_precision.sh +``` + diff --git a/research/megatron-inference-true-on-policy/run_qwen1.5b_zero_kl_precision.sh b/research/megatron-inference-true-on-policy/run_qwen1.5b_zero_kl_precision.sh new file mode 100755 index 0000000000..380e12ce1a --- /dev/null +++ b/research/megatron-inference-true-on-policy/run_qwen1.5b_zero_kl_precision.sh @@ -0,0 +1,254 @@ +#!/bin/bash +#SBATCH --job-name=qwen1.5b-zero-kl-precision +#SBATCH --account=your_account_here +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gpus-per-node=8 +#SBATCH --time=00:30:00 +#SBATCH --exclusive +#SBATCH --output=logs/qwen1.5b-zero-kl-precision-%j-%x.out +#SBATCH --error=logs/qwen1.5b-zero-kl-precision-%j-%x.err + +# ============================================================================= +# Qwen2.5-1.5B — Megatron colocated inference, zero-KL precision sweep (BF16 vs MXFP8) +# +# Enables zero_train_gen_mismatch which activates batch-invariant kernels (TE GEMM +# workspace pinned, FA4 num_splits=1) and use_mamba_mem_eff_path=False for hybrid models. +# +# Megatron colocated inference (m-inf) with FlashAttention 4 via TE v2.15. +# zero_train_gen_mismatch wires attention_backend=flash automatically. +# +# PRECISION (default bf16). Both use examples/configs/grpo_math_1B_megatron.yaml; +# mxfp8 adds inline policy.megatron_cfg.fp8_cfg overrides (enabled=true, recipe=mxfp8). +# +# Usage (from RL/ directory): +# sbatch --export=PRECISION=bf16 run_qwen1.5b_zero_kl_precision.sh +# sbatch --export=PRECISION=bf16,ZERO_TRAIN_GEN_MISMATCH=false run_qwen1.5b_zero_kl_precision.sh +# sbatch --export=PRECISION=mxfp8 run_qwen1.5b_zero_kl_precision.sh +# sbatch --export=PRECISION=mxfp8,ZERO_TRAIN_GEN_MISMATCH=false run_qwen1.5b_zero_kl_precision.sh +# +# RESTART=true (default false) — wipe CKPT_DIR and start a fresh W&B run. +# ZERO_TRAIN_GEN_MISMATCH=true (default) — enable zero-KL train/gen mismatch patches. +# EXP_TAG — optional suffix on run name (default: today's date). +# MAX_STEPS — default 2000. +# NRL_FORCE_REBUILD_VENVS=true — first build / after pyproject or uv.lock changes. +# NRL_USE_WARM_UV_CACHE=true — read wheels from NRL_WARM_UV_CACHE_DIR (default ${RL_DIR}/uv_cache). +# ============================================================================= + +set -euo pipefail + +SCRIPT_DIR="${SLURM_SUBMIT_DIR:-$(cd "$(dirname "$0")" && pwd)}" + +GPUS_PER_NODE="${SLURM_GPUS_ON_NODE:-${SLURM_GPUS_PER_NODE:-8}}" +NUM_NODES="${SLURM_NNODES:-1}" + +PRECISION="${PRECISION:-bf16}" +MAX_STEPS="${MAX_STEPS:-2000}" +NRL_USE_WARM_UV_CACHE="${NRL_USE_WARM_UV_CACHE:-false}" +ZERO_TRAIN_GEN_MISMATCH="${ZERO_TRAIN_GEN_MISMATCH:-true}" + +if [[ -f "${SCRIPT_DIR}/.env" ]]; then + set -a; source "${SCRIPT_DIR}/.env"; set +a +fi + +export TORCH_CUDA_ARCH_LIST='9.0 10.0' +: "${RL_DIR:?Set RL_DIR in .env}" +: "${CONTAINER_IMAGE:?Set CONTAINER_IMAGE in .env}" +: "${HF_TOKEN:?Set HF_TOKEN in .env}" +: "${HF_HOME:?Set HF_HOME in .env}" +: "${WANDB_API_KEY:?Set WANDB_API_KEY in .env}" +: "${WANDB_ENTITY:?Set WANDB_ENTITY in .env}" +: "${WANDB_PROJECT:?Set WANDB_PROJECT in .env}" + +NRL_WARM_UV_CACHE_DIR="${NRL_WARM_UV_CACHE_DIR:-${RL_DIR}/uv_cache}" + +NRL_RAY_VENVS_MOUNT_HOST="${NRL_RAY_VENVS_MOUNT_HOST:-}" +if [[ -n "${NRL_RAY_VENVS_MOUNT_HOST}" ]]; then + mkdir -p "${NRL_RAY_VENVS_MOUNT_HOST}" + NEMO_RL_VENV_CONTAINER="/opt/ray_venvs" + NRL_RAY_VENVS_MOUNT_SUFFIX=",${NRL_RAY_VENVS_MOUNT_HOST}:/opt/ray_venvs" +else + mkdir -p "${RL_DIR}/venvs" + NEMO_RL_VENV_CONTAINER="/opt/nemo-rl/venvs" + NRL_RAY_VENVS_MOUNT_SUFFIX="" +fi + +if [[ ! "$PRECISION" =~ ^(bf16|mxfp8)$ ]]; then + echo "ERROR: Invalid PRECISION '${PRECISION}'. Must be: bf16, mxfp8" >&2 + exit 1 +fi + +GRPO_CONFIG="examples/configs/grpo_math_1B_megatron.yaml" +MXFP8_OVERRIDES="" +case "$PRECISION" in + bf16) + WANDB_RUN_SUFFIX="bf16" + ;; + mxfp8) + WANDB_RUN_SUFFIX="mxfp8" + MXFP8_OVERRIDES="\ + policy.megatron_cfg.fp8_cfg.enabled=true \ + policy.megatron_cfg.fp8_cfg.fp8_recipe=mxfp8" + ;; +esac + +NRL_NVTE_DEBUG="${NRL_NVTE_DEBUG:-1}" +NRL_NVTE_DEBUG_LEVEL="${NRL_NVTE_DEBUG_LEVEL:-1}" + +_RUN_SUFFIX="${EXP_TAG:-$(date +%Y-%m-%d)}" +WANDB_RUN_NAME="qwen-1.5b-zero-kl-${WANDB_RUN_SUFFIX}-m-inf-${_RUN_SUFFIX}" +CKPT_DIR="${CKPT_DIR:-${RL_DIR}/results/${WANDB_RUN_NAME}}" +LOG_DIR_EXP="${LOG_DIR_EXP:-${RL_DIR}/logs/${WANDB_RUN_NAME}}" +SAVE_PERIOD="${SAVE_PERIOD:-250}" +KEEP_TOP_K="${KEEP_TOP_K:-3}" +RESTART="${RESTART:-false}" +if [[ "${RESTART}" != "true" && "${RESTART}" != "false" ]]; then + echo "ERROR: RESTART must be true or false (got '${RESTART}')." >&2 + exit 1 +fi +if [[ "${ZERO_TRAIN_GEN_MISMATCH}" != "true" && "${ZERO_TRAIN_GEN_MISMATCH}" != "false" ]]; then + echo "ERROR: ZERO_TRAIN_GEN_MISMATCH must be true or false (got '${ZERO_TRAIN_GEN_MISMATCH}')." >&2 + exit 1 +fi + +if [[ "${RESTART}" == "true" ]]; then + if [[ -z "${CKPT_DIR}" || "${CKPT_DIR}" == "/" ]]; then + echo "ERROR: refusing to RESTART-wipe unsafe CKPT_DIR='${CKPT_DIR}'." >&2 + exit 1 + fi + if [[ -d "${CKPT_DIR}" ]]; then + echo "RESTART=true: wiping existing checkpoints in ${CKPT_DIR}" >&2 + rm -rf "${CKPT_DIR}" + fi +fi + +mkdir -p "${CKPT_DIR}" "${LOG_DIR_EXP}" + +WANDB_RUN_ID_PIN="${CKPT_DIR}/.wandb_run_id" +if [[ -z "${WANDB_RUN_ID:-}" ]]; then + _pinned_id="" + if [[ -f "${WANDB_RUN_ID_PIN}" ]]; then + _pinned_id="$(head -n1 "${WANDB_RUN_ID_PIN}" 2>/dev/null | tr -d '[:space:]')" + fi + if [[ -n "${_pinned_id}" ]]; then + WANDB_RUN_ID="${_pinned_id}" + elif [[ "${RESTART}" == "true" ]]; then + WANDB_RUN_ID="$(printf '%s-%s-%s' "${WANDB_PROJECT}" "${WANDB_RUN_NAME}" "${SLURM_JOB_ID:-$(date +%s)}" | md5sum | cut -c1-32)" + else + WANDB_RUN_ID="$(printf '%s-%s' "${WANDB_PROJECT}" "${WANDB_RUN_NAME}" | md5sum | cut -c1-32)" + fi +fi +printf '%s\n' "${WANDB_RUN_ID}" > "${WANDB_RUN_ID_PIN}" + +if [[ "${RESTART}" == "true" ]]; then + WANDB_RESUME="${WANDB_RESUME:-never}" +else + WANDB_RESUME="${WANDB_RESUME:-allow}" +fi + +echo "==============================================" +echo "Qwen2.5-1.5B zero-KL precision study" +echo " Mode: m-inf (FA4 via TE v2.15)" +echo " Precision: ${PRECISION}" +echo " Config: ${GRPO_CONFIG}" +echo " zero_train_gen_mismatch: ${ZERO_TRAIN_GEN_MISMATCH}" +echo " Max steps: ${MAX_STEPS}" +echo " Nodes: ${NUM_NODES}" +echo " GPUs/node: ${GPUS_PER_NODE}" +echo " Total GPUs: $((NUM_NODES * GPUS_PER_NODE))" +echo " RL dir: ${RL_DIR}" +echo " Container: ${CONTAINER_IMAGE}" +echo " Job ID: ${SLURM_JOB_ID:-interactive}" +echo " Nodes: ${SLURM_NODELIST:-N/A}" +echo " Run name: ${WANDB_RUN_NAME}" +echo " Ckpt dir: ${CKPT_DIR} (save_period=${SAVE_PERIOD}, keep_top_k=${KEEP_TOP_K}, restart=${RESTART})" +echo " Log dir: ${LOG_DIR_EXP}" +echo " W&B run id: ${WANDB_RUN_ID} (resume=${WANDB_RESUME})" +echo " Time: $(date)" +echo "==============================================" + +mkdir -p logs + +GRPO_ARGS="\ + --config ${GRPO_CONFIG} \ + grpo.max_num_steps=${MAX_STEPS} \ + cluster.num_nodes=${NUM_NODES} cluster.gpus_per_node=${GPUS_PER_NODE} \ + grpo.val_at_start=true grpo.val_at_end=true grpo.val_period=10 \ + checkpointing.enabled=true \ + checkpointing.checkpoint_dir=${CKPT_DIR} \ + checkpointing.save_period=${SAVE_PERIOD} \ + checkpointing.keep_top_k=${KEEP_TOP_K} \ + logger.log_dir=${LOG_DIR_EXP} \ + logger.wandb_enabled=true \ + logger.wandb.project=${WANDB_PROJECT}" + +if [[ "${NRL_FORCE_REBUILD_VENVS:-false}" == "true" ]]; then + BASE_CMD="NRL_FORCE_REBUILD_VENVS=${NRL_FORCE_REBUILD_VENVS:-true} uv run --extra mcore examples/run_grpo.py ${GRPO_ARGS}" +else + BASE_CMD="uv run examples/run_grpo.py ${GRPO_ARGS}" +fi + +MINF_FLAGS="\ + policy.generation.backend=megatron \ + +policy.megatron_cfg.zero_train_gen_mismatch=${ZERO_TRAIN_GEN_MISMATCH} \ + policy.generation.mcore_generation_config.enable_chunked_prefill=false \ + policy.max_total_sequence_length=1024" + +CMD="${BASE_CMD} ${MINF_FLAGS} \ + logger.wandb.name=${WANDB_RUN_NAME}" +if [[ -n "${MXFP8_OVERRIDES}" ]]; then + CMD="${CMD} ${MXFP8_OVERRIDES}" +fi + +if [[ $# -gt 0 ]]; then + CMD="${CMD} $*" +fi + +cd "${RL_DIR}" + +export CONTAINER="${CONTAINER_IMAGE}" +export MOUNTS="/lustre:/lustre,${RL_DIR}:/opt/nemo-rl${NRL_RAY_VENVS_MOUNT_SUFFIX}" +if [[ -d /scratch ]]; then + export MOUNTS="${MOUNTS},/scratch:/scratch" +fi +export GPUS_PER_NODE + +if [[ "${NRL_USE_WARM_UV_CACHE}" == "true" ]]; then + mkdir -p "${NRL_WARM_UV_CACHE_DIR}" + NRL_WARM_UV_CACHE_EXPORT="export UV_CACHE_DIR=${NRL_WARM_UV_CACHE_DIR} && " +else + NRL_WARM_UV_CACHE_EXPORT="" +fi + +export COMMAND="\ + ${NRL_WARM_UV_CACHE_EXPORT}\ + export NEMO_RL_VENV_DIR=${NEMO_RL_VENV_CONTAINER} && \ + export NVTE_DEBUG=${NRL_NVTE_DEBUG} && \ + export NVTE_DEBUG_LEVEL=${NRL_NVTE_DEBUG_LEVEL} && \ + export PYTHONUNBUFFERED=1 && \ + export UV_HTTP_TIMEOUT=900 && \ + export HF_HOME=${HF_HOME} && \ + export TORCH_CUDA_ARCH_LIST='${TORCH_CUDA_ARCH_LIST}' && \ + export HF_TOKEN=${HF_TOKEN} && \ + export WANDB_API_KEY=${WANDB_API_KEY} && \ + export WANDB_ENTITY=${WANDB_ENTITY} && \ + export WANDB_RUN_ID=${WANDB_RUN_ID} && \ + export WANDB_RESUME=${WANDB_RESUME} && \ + export CUDA_DEVICE_MAX_CONNECTIONS=1 && \ + cd /opt/nemo-rl && \ + ${CMD}" + +echo "" +echo "COMMAND:" +echo "${CMD}" +echo "" +echo "==============================================" + +export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_CONTAINER}" + +source ray.sub + +echo "" +echo "==============================================" +echo "Job completed at: $(date)" +echo "==============================================" diff --git a/research/megatron-inference-true-on-policy/run_qwen30ba3b_zero_kl_precision.sh b/research/megatron-inference-true-on-policy/run_qwen30ba3b_zero_kl_precision.sh new file mode 100755 index 0000000000..0841c3b88f --- /dev/null +++ b/research/megatron-inference-true-on-policy/run_qwen30ba3b_zero_kl_precision.sh @@ -0,0 +1,295 @@ +#!/bin/bash +#SBATCH --job-name=qwen30ba3b-zero-kl-precision +#SBATCH --account=your_account_here +#SBATCH --partition=batch +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gpus-per-node=8 +#SBATCH --time=00:30:00 +#SBATCH --exclusive +#SBATCH --output=logs/qwen30ba3b-zero-kl-precision-%j-%x.out +#SBATCH --error=logs/qwen30ba3b-zero-kl-precision-%j-%x.err + +# ============================================================================= +# Qwen3-30B-A3B MoE — Megatron colocated inference, zero-KL precision sweep (BF16 vs MXFP8) +# +# Enables zero_train_gen_mismatch which activates batch-invariant kernels (TE GEMM +# workspace pinned, FA4 num_splits=1). Router replay is enabled separately via +# policy.router_replay.enabled=true (MoE routing must be set explicitly). +# +# MODE (default m-inf-fa4): +# m-inf-fa4 — FlashAttention 4 (+policy.megatron_cfg.attention_backend=flash) +# m-inf — TE unfused attention +# +# PRECISION (default bf16): +# bf16 — examples/configs/grpo_math_qwen30ba3b_megatron.yaml +# mxfp8 — examples/configs/grpo_math_qwen30ba3b_megatron_mxfp8.yaml +# +# Usage (from RL/ directory): +# sbatch --export=PRECISION=bf16 run_qwen30ba3b_zero_kl_precision.sh +# sbatch --export=PRECISION=mxfp8 run_qwen30ba3b_zero_kl_precision.sh +# sbatch --export=PRECISION=bf16,ZERO_TRAIN_GEN_MISMATCH=false run_qwen30ba3b_zero_kl_precision.sh +# sbatch --export=PRECISION=mxfp8,ZERO_TRAIN_GEN_MISMATCH=false run_qwen30ba3b_zero_kl_precision.sh +# +# RESTART=true (default false) — wipe CKPT_DIR and start a fresh W&B run. +# EXP_TAG — optional suffix on run name (default: today's date). +# MAX_STEPS — default 2000. +# NRL_FORCE_REBUILD_VENVS=true — first build / after pyproject or uv.lock changes. +# NRL_USE_WARM_UV_CACHE=true — read wheels from NRL_WARM_UV_CACHE_DIR (default ${RL_DIR}/uv_cache). +# ============================================================================= + +set -euo pipefail + +SCRIPT_DIR="${SLURM_SUBMIT_DIR:-$(cd "$(dirname "$0")" && pwd)}" + +GPUS_PER_NODE="${SLURM_GPUS_ON_NODE:-${SLURM_GPUS_PER_NODE:-8}}" +NUM_NODES="${SLURM_NNODES:-1}" + +MODE="${MODE:-m-inf-fa4}" +PRECISION="${PRECISION:-bf16}" +MAX_STEPS="${MAX_STEPS:-2000}" +NRL_USE_WARM_UV_CACHE="${NRL_USE_WARM_UV_CACHE:-false}" + +if [[ -f "${SCRIPT_DIR}/.env" ]]; then + set -a; source "${SCRIPT_DIR}/.env"; set +a +fi + +export TORCH_CUDA_ARCH_LIST='9.0 10.0' +: "${RL_DIR:?Set RL_DIR in .env}" +: "${CONTAINER_IMAGE:?Set CONTAINER_IMAGE in .env}" +: "${HF_TOKEN:?Set HF_TOKEN in .env}" +: "${HF_HOME:?Set HF_HOME in .env}" +: "${WANDB_API_KEY:?Set WANDB_API_KEY in .env}" +: "${WANDB_ENTITY:?Set WANDB_ENTITY in .env}" +: "${WANDB_PROJECT:?Set WANDB_PROJECT in .env}" + +NRL_WARM_UV_CACHE_DIR="${NRL_WARM_UV_CACHE_DIR:-${RL_DIR}/uv_cache}" + +NRL_RAY_VENVS_MOUNT_HOST="${NRL_RAY_VENVS_MOUNT_HOST:-}" +if [[ -n "${NRL_RAY_VENVS_MOUNT_HOST}" ]]; then + mkdir -p "${NRL_RAY_VENVS_MOUNT_HOST}" + NEMO_RL_VENV_CONTAINER="/opt/ray_venvs" + NRL_RAY_VENVS_MOUNT_SUFFIX=",${NRL_RAY_VENVS_MOUNT_HOST}:/opt/ray_venvs" +else + mkdir -p "${RL_DIR}/venvs" + NEMO_RL_VENV_CONTAINER="/opt/nemo-rl/venvs" + NRL_RAY_VENVS_MOUNT_SUFFIX="" +fi + +if [[ "$MODE" != "m-inf" && "$MODE" != "m-inf-fa4" ]]; then + echo "ERROR: MODE must be m-inf or m-inf-fa4 (got '${MODE}')." >&2 + exit 1 +fi + +case "$MODE" in + m-inf-fa4) + ATTENTION_FLAGS="+policy.megatron_cfg.attention_backend=flash" + TE_ATTN_PREFIX='unset NVTE_FUSED_ATTN NVTE_FLASH_ATTN NVTE_UNFUSED_ATTN && export NVTE_FLASH_ATTN=1 && export NVTE_FUSED_ATTN=0 && export NVTE_UNFUSED_ATTN=0 && ' + ;; + m-inf) + ATTENTION_FLAGS="+policy.megatron_cfg.attention_backend=unfused" + TE_ATTN_PREFIX='unset NVTE_FUSED_ATTN NVTE_FLASH_ATTN NVTE_UNFUSED_ATTN && export NVTE_FLASH_ATTN=0 && export NVTE_FUSED_ATTN=0 && export NVTE_UNFUSED_ATTN=1 && ' + ;; +esac + +if [[ ! "$PRECISION" =~ ^(bf16|mxfp8)$ ]]; then + echo "ERROR: Invalid PRECISION '${PRECISION}'. Must be: bf16, mxfp8" >&2 + exit 1 +fi + +GRPO_CONFIG="examples/configs/grpo_math_qwen30ba3b_megatron.yaml" +MXFP8_OVERRIDES="" +case "$PRECISION" in + bf16) + WANDB_RUN_SUFFIX="bf16" + ;; + mxfp8) + WANDB_RUN_SUFFIX="mxfp8" + MXFP8_OVERRIDES="\ + policy.megatron_cfg.fp8_cfg.enabled=true \ + policy.megatron_cfg.fp8_cfg.fp8_recipe=mxfp8 \ + policy.megatron_cfg.optimizer.use_precision_aware_optimizer=false" + ;; +esac + +NRL_NVTE_DEBUG="${NRL_NVTE_DEBUG:-1}" +NRL_NVTE_DEBUG_LEVEL="${NRL_NVTE_DEBUG_LEVEL:-1}" + +_RUN_SUFFIX="${EXP_TAG:-$(date +%Y-%m-%d)}" +WANDB_RUN_NAME="qwen30ba3b-zero-kl-${WANDB_RUN_SUFFIX}-${MODE}-${_RUN_SUFFIX}" +CKPT_DIR="${CKPT_DIR:-${RL_DIR}/results/${WANDB_RUN_NAME}}" +LOG_DIR_EXP="${LOG_DIR_EXP:-${RL_DIR}/logs/${WANDB_RUN_NAME}}" +SAVE_PERIOD="${SAVE_PERIOD:-10}" +KEEP_TOP_K="${KEEP_TOP_K:-3}" +RESTART="${RESTART:-false}" +if [[ "${RESTART}" != "true" && "${RESTART}" != "false" ]]; then + echo "ERROR: RESTART must be true or false (got '${RESTART}')." >&2 + exit 1 +fi + +if [[ "${RESTART}" == "true" ]]; then + if [[ -z "${CKPT_DIR}" || "${CKPT_DIR}" == "/" ]]; then + echo "ERROR: refusing to RESTART-wipe unsafe CKPT_DIR='${CKPT_DIR}'." >&2 + exit 1 + fi + if [[ -d "${CKPT_DIR}" ]]; then + echo "RESTART=true: wiping existing checkpoints in ${CKPT_DIR}" >&2 + rm -rf "${CKPT_DIR}" + fi +fi + +mkdir -p "${CKPT_DIR}" "${LOG_DIR_EXP}" + +WANDB_RUN_ID_PIN="${CKPT_DIR}/.wandb_run_id" +if [[ -z "${WANDB_RUN_ID:-}" ]]; then + _pinned_id="" + if [[ -f "${WANDB_RUN_ID_PIN}" ]]; then + _pinned_id="$(head -n1 "${WANDB_RUN_ID_PIN}" 2>/dev/null | tr -d '[:space:]')" + fi + if [[ -n "${_pinned_id}" ]]; then + WANDB_RUN_ID="${_pinned_id}" + elif [[ "${RESTART}" == "true" ]]; then + WANDB_RUN_ID="$(printf '%s-%s-%s' "${WANDB_PROJECT}" "${WANDB_RUN_NAME}" "${SLURM_JOB_ID:-$(date +%s)}" | md5sum | cut -c1-32)" + else + WANDB_RUN_ID="$(printf '%s-%s' "${WANDB_PROJECT}" "${WANDB_RUN_NAME}" | md5sum | cut -c1-32)" + fi +fi +printf '%s\n' "${WANDB_RUN_ID}" > "${WANDB_RUN_ID_PIN}" + +if [[ "${RESTART}" == "true" ]]; then + WANDB_RESUME="${WANDB_RESUME:-never}" +else + WANDB_RESUME="${WANDB_RESUME:-allow}" +fi + +echo "==============================================" +echo "Qwen3-30B-A3B zero-KL precision study" +echo " Mode: ${MODE}" +echo " Precision: ${PRECISION}" +echo " Config: ${GRPO_CONFIG}" +echo " zero_train_gen_mismatch: true" +echo " Router replay: true" +echo " Max steps: ${MAX_STEPS}" +echo " Nodes: ${NUM_NODES}" +echo " GPUs/node: ${GPUS_PER_NODE}" +echo " Total GPUs: $((NUM_NODES * GPUS_PER_NODE))" +echo " RL dir: ${RL_DIR}" +echo " Container: ${CONTAINER_IMAGE}" +echo " Job ID: ${SLURM_JOB_ID:-interactive}" +echo " Nodes: ${SLURM_NODELIST:-N/A}" +echo " Run name: ${WANDB_RUN_NAME}" +echo " Ckpt dir: ${CKPT_DIR} (save_period=${SAVE_PERIOD}, keep_top_k=${KEEP_TOP_K}, restart=${RESTART})" +echo " Log dir: ${LOG_DIR_EXP}" +echo " W&B run id: ${WANDB_RUN_ID} (resume=${WANDB_RESUME})" +echo " Time: $(date)" +echo "==============================================" + +mkdir -p logs + +STUDY_OVERRIDES="\ + grpo.num_prompts_per_step=4 \ + grpo.num_generations_per_prompt=16 \ + policy.train_global_batch_size=64 \ + policy.megatron_cfg.moe_token_dispatcher_type=alltoall \ + policy.megatron_cfg.moe_router_dtype=fp32 \ + policy.max_total_sequence_length=128" + +GRPO_ARGS="\ + --config ${GRPO_CONFIG} \ + grpo.max_num_steps=${MAX_STEPS} \ + cluster.num_nodes=${NUM_NODES} cluster.gpus_per_node=${GPUS_PER_NODE} \ + ${STUDY_OVERRIDES} \ + grpo.val_at_start=false grpo.val_at_end=true grpo.val_period=10 \ + checkpointing.enabled=true \ + checkpointing.checkpoint_dir=${CKPT_DIR} \ + checkpointing.save_period=${SAVE_PERIOD} \ + checkpointing.keep_top_k=${KEEP_TOP_K} \ + logger.log_dir=${LOG_DIR_EXP} \ + logger.wandb_enabled=true \ + logger.wandb.project=${WANDB_PROJECT}" + +if [[ "${NRL_FORCE_REBUILD_VENVS:-false}" == "true" ]]; then + BASE_CMD="NRL_FORCE_REBUILD_VENVS=${NRL_FORCE_REBUILD_VENVS:-true} uv run --extra mcore examples/run_grpo.py ${GRPO_ARGS}" +else + BASE_CMD="uv run examples/run_grpo.py ${GRPO_ARGS}" +fi + +# Colocated m-inf: generation shares training megatron_cfg (incl. TP). Do not set +# mcore_generation_config.tensor_model_parallel_size — that key is absent from +# grpo_math_1B.yaml's mcore_generation_config and Hydra struct-rejects it. +MINF_FLAGS="\ + policy.generation.backend=megatron \ + +policy.megatron_cfg.zero_train_gen_mismatch=true \ + policy.router_replay.enabled=true \ + policy.megatron_cfg.tensor_model_parallel_size=1 \ + +policy.megatron_cfg.cuda_graph_impl=none \ + +policy.megatron_cfg.moe_pad_experts_for_cuda_graph_inference=false \ + policy.megatron_cfg.moe_permute_fusion=false \ + policy.generation.mcore_generation_config.kv_cache_management_mode=recompute \ + +policy.generation.mcore_generation_config.static_kv_memory_pointers=false \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=False \ + policy.generation.mcore_generation_config.cuda_graph_impl=none \ + policy.generation.mcore_generation_config.num_cuda_graphs=0 \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope=none \ + policy.generation.mcore_generation_config.buffer_size_gb=16" + +CMD="${BASE_CMD} ${MINF_FLAGS} \ + logger.wandb.name=${WANDB_RUN_NAME} \ + ${ATTENTION_FLAGS}" +if [[ -n "${MXFP8_OVERRIDES}" ]]; then + CMD="${CMD} ${MXFP8_OVERRIDES}" +fi + +if [[ $# -gt 0 ]]; then + CMD="${CMD} $*" +fi + +cd "${RL_DIR}" + +export CONTAINER="${CONTAINER_IMAGE}" +export MOUNTS="/lustre:/lustre,${RL_DIR}:/opt/nemo-rl${NRL_RAY_VENVS_MOUNT_SUFFIX}" +if [[ -d /scratch ]]; then + export MOUNTS="${MOUNTS},/scratch:/scratch" +fi +export GPUS_PER_NODE + +if [[ "${NRL_USE_WARM_UV_CACHE}" == "true" ]]; then + mkdir -p "${NRL_WARM_UV_CACHE_DIR}" + NRL_WARM_UV_CACHE_EXPORT="export UV_CACHE_DIR=${NRL_WARM_UV_CACHE_DIR} && " +else + NRL_WARM_UV_CACHE_EXPORT="" +fi + +export COMMAND="${TE_ATTN_PREFIX}\ + ${NRL_WARM_UV_CACHE_EXPORT}\ + export NEMO_RL_VENV_DIR=${NEMO_RL_VENV_CONTAINER} && \ + export NVTE_DEBUG=${NRL_NVTE_DEBUG} && \ + export NVTE_DEBUG_LEVEL=${NRL_NVTE_DEBUG_LEVEL} && \ + export PYTHONUNBUFFERED=1 && \ + export UV_HTTP_TIMEOUT=900 && \ + export HF_HOME=${HF_HOME} && \ + export TORCH_CUDA_ARCH_LIST='${TORCH_CUDA_ARCH_LIST}' && \ + export HF_TOKEN=${HF_TOKEN} && \ + export WANDB_API_KEY=${WANDB_API_KEY} && \ + export WANDB_ENTITY=${WANDB_ENTITY} && \ + export WANDB_RUN_ID=${WANDB_RUN_ID} && \ + export WANDB_RESUME=${WANDB_RESUME} && \ + export CUDA_DEVICE_MAX_CONNECTIONS=1 && \ + export NRL_INSTALL_FA3=${NRL_INSTALL_FA3:-0} && \ + cd /opt/nemo-rl && \ + ${CMD}" + +echo "" +echo "COMMAND:" +echo "${CMD}" +echo "" +echo "==============================================" + +export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_CONTAINER}" + +source ray.sub + +echo "" +echo "==============================================" +echo "Job completed at: $(date)" +echo "==============================================" diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 1f004a18c5..6d5fecb2fe 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -2481,3 +2481,34 @@ def attach_fresh_draft(): restored_chunk.draft_model.weight, owner_chunk.draft_model.weight, ) + + +@pytest.mark.mcore +def test_zero_train_gen_mismatch_forces_te_generation_spec(): + """Zero train/gen mismatch must keep generation on the train-side TE MoE path.""" + from nemo_rl.models.megatron.setup import _apply_zero_train_gen_mismatch + + config = { + "megatron_cfg": {"zero_train_gen_mismatch": True}, + "generation": { + "mcore_generation_config": { + "transformer_impl": "inference_optimized", + } + }, + } + + with ( + patch( + "nemo_rl.models.policy.workers.moe_determinism_patches." + "apply_moe_unpermute_determinism_patch" + ), + patch( + "nemo_rl.models.policy.workers.patches.apply_te_gemm_cublas_pinned_patch" + ), + ): + _apply_zero_train_gen_mismatch(config) + + assert ( + config["generation"]["mcore_generation_config"]["transformer_impl"] + == "transformer_engine" + ) diff --git a/tests/unit/models/policy/test_moe_determinism_patches.py b/tests/unit/models/policy/test_moe_determinism_patches.py new file mode 100644 index 0000000000..3129823f90 --- /dev/null +++ b/tests/unit/models/policy/test_moe_determinism_patches.py @@ -0,0 +1,95 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from nemo_rl.models.policy.workers.moe_determinism_patches import ( + _nrl_dynamic_step_context_bookkeeping, + _patched_unpermute, + _unpermute_fixed_order_combine, + apply_moe_unpermute_determinism_patch, + apply_router_replay_inference_patches, + restore_moe_determinism_patches, +) + + +class TestUnpermuteFixedOrderCombine: + def test_sums_per_token(self): + permuted = torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]) + sorted_indices = torch.tensor([0, 0, 1]) + out = _unpermute_fixed_order_combine(permuted, sorted_indices, torch.Size([2, 2])) + assert torch.allclose(out, torch.tensor([[3.0, 0.0], [3.0, 0.0]])) + +class TestApplyMoeDeterminismPatches: + def setup_method(self): + restore_moe_determinism_patches() + + def teardown_method(self): + restore_moe_determinism_patches() + + def test_unpermute_patch_is_idempotent(self): + fake_mod = MagicMock() + fake_mod.HAVE_TE = False + fake_mod.fused_unpermute = None + fake_mod.is_te_min_version = lambda _v: False + fake_mod.unpermute = MagicMock(return_value="orig") + fake_dispatcher = MagicMock() + fake_dispatcher.unpermute = MagicMock(return_value="dispatcher_orig") + + with patch.dict( + sys.modules, + { + "megatron.core.transformer.moe.moe_utils": fake_mod, + "megatron.core.transformer.moe.token_dispatcher": fake_dispatcher, + }, + ): + apply_moe_unpermute_determinism_patch() + apply_moe_unpermute_determinism_patch() + assert fake_mod.unpermute is _patched_unpermute + assert fake_dispatcher.unpermute is _patched_unpermute + restore_moe_determinism_patches() + assert fake_mod.unpermute() == "orig" + assert fake_dispatcher.unpermute() == "dispatcher_orig" + + def test_patched_unpermute_uses_fixed_order(self): + permuted = torch.tensor([[1.0], [2.0]]) + sorted_indices = torch.tensor([0, 0]) + out = _patched_unpermute(permuted, sorted_indices, torch.Size([1, 1])) + assert torch.allclose(out, torch.tensor([[3.0]])) + + def test_router_replay_inference_patch_replaces_methods(self): + pytest.importorskip("megatron") + from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine + from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, + ) + + orig_bookkeeping = TextGenerationController._dynamic_step_context_bookkeeping + orig_async_bookkeep = DynamicInferenceEngine.async_bookkeep + try: + apply_router_replay_inference_patches() + assert TextGenerationController._dynamic_step_context_bookkeeping is ( + _nrl_dynamic_step_context_bookkeeping + ) + assert DynamicInferenceEngine.async_bookkeep is not orig_async_bookkeep + finally: + restore_moe_determinism_patches() + assert ( + TextGenerationController._dynamic_step_context_bookkeeping is orig_bookkeeping + ) + assert DynamicInferenceEngine.async_bookkeep is orig_async_bookkeep diff --git a/tests/unit/models/policy/test_patches.py b/tests/unit/models/policy/test_patches.py index e8cacbcd4a..52d63a60c0 100644 --- a/tests/unit/models/policy/test_patches.py +++ b/tests/unit/models/policy/test_patches.py @@ -21,7 +21,9 @@ from nemo_rl.models.policy.workers.patches import ( _get_transformer_engine_file, + apply_te_gemm_cublas_pinned_patch, apply_transformer_engine_patch, + restore_te_gemm_cublas_pinned_patch, ) @@ -445,3 +447,58 @@ def permutation_kernel(x): assert captured.out.count("Successfully patched") == 1 finally: os.unlink(tmp_path) + + +class TestApplyTeGemmCublasPinnedPatch: + """Tests for apply_te_gemm_cublas_pinned_patch.""" + + def setup_method(self): + restore_te_gemm_cublas_pinned_patch() + + def teardown_method(self): + restore_te_gemm_cublas_pinned_patch() + + def test_shrinks_workspace_and_is_idempotent(self, capsys): + orig_fn = MagicMock(return_value=4096) + mock_ws_fn = MagicMock() + mock_ws_fn.cache_clear = MagicMock() + mock_gemm_mod = MagicMock() + mock_gemm_mod.get_cublas_workspace_size_bytes = orig_fn + mock_gemm_mod.get_cublas_workspace = mock_ws_fn + + with patch( + "nemo_rl.models.policy.workers.patches.importlib.import_module", + return_value=mock_gemm_mod, + ): + apply_te_gemm_cublas_pinned_patch() + apply_te_gemm_cublas_pinned_patch() + + assert mock_gemm_mod.get_cublas_workspace_size_bytes() == 4 + mock_ws_fn.cache_clear.assert_called_once() + captured = capsys.readouterr() + assert captured.out.count("[zero_train_gen_mismatch] shrunk TE cuBLAS workspace") == 1 + + def test_skips_when_te_gemm_module_missing(self, capsys): + with patch( + "nemo_rl.models.policy.workers.patches.importlib.import_module", + side_effect=ImportError("no te"), + ): + apply_te_gemm_cublas_pinned_patch() + + captured = capsys.readouterr() + assert "is not importable" in captured.out + + def test_restore_puts_back_original(self): + orig_fn = MagicMock(return_value=4096) + mock_gemm_mod = MagicMock() + mock_gemm_mod.get_cublas_workspace_size_bytes = orig_fn + mock_gemm_mod.get_cublas_workspace = MagicMock() + + with patch( + "nemo_rl.models.policy.workers.patches.importlib.import_module", + return_value=mock_gemm_mod, + ): + apply_te_gemm_cublas_pinned_patch() + restore_te_gemm_cublas_pinned_patch() + + assert mock_gemm_mod.get_cublas_workspace_size_bytes is orig_fn