feat(megatron): enable batch-invariant logprob parity

examples/configs/grpo_math_1B.yaml

@@ -188,6 +188,14 @@ policy:
     # parallelism, sequence packing, or top-k/top-p training-time filtering.
     use_fused_linear_logprobs: false
     fused_linear_logprobs_chunk_size: 256
+    # Opt in to bitwise-identical Megatron generation and policy logprobs.
+    # Requires Megatron generation, BF16, matching training/generation TP, CP=1,
+    # attention_backend=flash, and flash_attention_version=3/4. Sampling supports
+    # the usual temperature, top-k, and top-p settings.
+    batch_invariant_mode: false
+    batch_invariant_backend: "te_native"  # Options: "te_native", "triton", "deepgemm".
+    batch_invariant_collective: "ordered"  # Options: "ordered", "multimem" (NVLS).
+    flash_attention_version: null  # Set to 3 or 4 when enabling the mode.
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false

nemo_rl/algorithms/loss/utils.py

@@ -51,8 +51,10 @@ def prepare_loss_input(
         logits: Logits from the model.
         data: Microbatch data. Will be updated if sampling_params is not None.
         loss_fn: Loss function.
-        vocab_parallel_rank: Vocab parallel rank.
-        vocab_parallel_group: Vocab parallel group.
+        vocab_parallel_rank: Vocab parallel rank. Leave unset with
+            vocab_parallel_group when logits contain the full vocabulary.
+        vocab_parallel_group: Vocab parallel group. Leave unset with
+            vocab_parallel_rank when logits contain the full vocabulary.
         context_parallel_group: Context parallel group.
         sampling_params: Sampling parameters.
         d2t: Draft to target token mapping.
@@ -304,11 +306,12 @@ def prepare_packed_loss_input(
             f"got {loss_fn.input_type}. Use SequencePackingLossWrapper with "
             f"prepare_loss_input for other types."
         )
-    assert vocab_parallel_group is not None, (
-        "prepare_packed_loss_input requires vocab_parallel_group (Megatron TP)."
+    assert (vocab_parallel_group is None) == (vocab_parallel_rank is None), (
+        "vocab_parallel_rank and vocab_parallel_group must either both be provided "
+        "or both be None for full-vocabulary logits."
     )
-    assert vocab_parallel_rank is not None, (
-        "vocab_parallel_rank must be provided with vocab_parallel_group."
+    resolved_vocab_parallel_rank = (
+        0 if vocab_parallel_rank is None else vocab_parallel_rank
     )
 
     input_ids = data["input_ids"]
@@ -345,8 +348,8 @@ def prepare_packed_loss_input(
         packed_rolled_targets,
         cu_seqlens_q_padded,
         unpacked_seqlen,
-        vocab_start_index=vocab_parallel_rank * logits.shape[-1],
-        vocab_end_index=(vocab_parallel_rank + 1) * logits.shape[-1],
+        vocab_start_index=resolved_vocab_parallel_rank * logits.shape[-1],
+        vocab_end_index=(resolved_vocab_parallel_rank + 1) * logits.shape[-1],
         group=vocab_parallel_group,
         inference_only=False,
         cp_group=context_parallel_group,
@@ -371,8 +374,9 @@ def prepare_packed_loss_input(
                     packed_rolled_targets,
                     cu_seqlens_q_padded,
                     unpacked_seqlen,
-                    vocab_start_index=vocab_parallel_rank * logits.shape[-1],
-                    vocab_end_index=(vocab_parallel_rank + 1) * logits.shape[-1],
+                    vocab_start_index=resolved_vocab_parallel_rank * logits.shape[-1],
+                    vocab_end_index=(resolved_vocab_parallel_rank + 1)
+                    * logits.shape[-1],
                     group=vocab_parallel_group,
                     inference_only=False,
                     cp_group=context_parallel_group,

nemo_rl/distributed/model_utils.py

@@ -1049,7 +1049,7 @@ def from_parallel_logits_to_logprobs_packed_sequences(
     unpacked_seqlen: int,
     vocab_start_index: int,
     vocab_end_index: int,
-    group: torch.distributed.ProcessGroup,
+    group: Optional[torch.distributed.ProcessGroup],
     inference_only: bool = False,
     cp_group: Optional[torch.distributed.ProcessGroup] = None,
     chunk_size: Optional[int] = None,
@@ -1070,7 +1070,8 @@ def from_parallel_logits_to_logprobs_packed_sequences(
         unpacked_seqlen (int): The length of the unpacked sequence tensor.
         vocab_start_index (int): Starting vocabulary index for this worker's partition.
         vocab_end_index (int): Ending vocabulary index for this worker's partition.
-        group (torch.distributed.ProcessGroup): Process group for distributed communication.
+        group (torch.distributed.ProcessGroup, optional): Tensor-parallel process
+            group. If None, logits already contain the full vocabulary.
         inference_only (bool, optional): If True, tensors won't be saved for backward pass. Defaults to False.
         cp_group (torch.distributed.ProcessGroup, optional): Context parallelism process group. Defaults to None.
         chunk_size (int, optional): Sequence dimension chunk size for computing the log probabilities.
@@ -1110,8 +1111,23 @@ def from_parallel_logits_to_logprobs_packed_sequences(
         target = rolled_targets.unsqueeze(0)
         vocab_parallel_logits = vocab_parallel_logits.unsqueeze(0)
 
+    # Batch-invariant Megatron scoring gathers the full vocabulary before using
+    # the same canonical FP32 log_softmax as generation. Keeping this branch in
+    # the packed helper lets packed and unpacked callers share that contract.
+    if group is None:
+        logits = vocab_parallel_logits.to(torch.float32)
+        logits, _ = apply_top_k_top_p(
+            logits,
+            top_k=sampling_params.top_k if sampling_params is not None else None,
+            top_p=sampling_params.top_p if sampling_params is not None else 1.0,
+        )
+        probs = (
+            torch.nn.functional.log_softmax(logits, dim=-1)
+            .gather(dim=-1, index=target.unsqueeze(-1))
+            .squeeze(-1)
+        )
     # Apply distributed log probability computation
-    if need_top_k_or_top_p_filtering(sampling_params):
+    elif need_top_k_or_top_p_filtering(sampling_params):
         if chunk_size is not None:
             probs: torch.Tensor = ChunkedDistributedLogprobWithSampling.apply(  # type: ignore
                 vocab_parallel_logits,

nemo_rl/models/generation/megatron/megatron_worker.py

@@ -291,7 +291,14 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None:
             ),
             logging_step_interval=logging_step_interval,
             num_speculative_tokens=num_speculative_tokens,
-            logprobs_mode="processed_logprobs",
+            # Sampling parameters control token selection, but batch-invariant
+            # generation reports raw model logprobs. Policy scoring mirrors this
+            # contract so parity is independent of temperature/top-k/top-p.
+            logprobs_mode=(
+                "raw_logprobs"
+                if self.cfg["megatron_cfg"].get("batch_invariant_mode")
+                else "processed_logprobs"
+            ),
             max_requests=max_requests,
         )
 

nemo_rl/models/megatron/batch_invariant.py

@@ -0,0 +1,31 @@
+# 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.
+
+import math
+
+
+def batch_invariant_token_multiple(configured_multiple: int, tp_size: int) -> int:
+    """Return a token multiple compatible with batch-invariant MCore inference."""
+    if configured_multiple < 1:
+        raise ValueError("configured_multiple must be positive.")
+    if tp_size < 1:
+        raise ValueError("tp_size must be positive.")
+
+    # Import lazily so non-Megatron policy drivers do not import Megatron-Core.
+    # TOKEN_ROUNDER is MCore's single source of truth for eager and graphed
+    # batch-invariant inference token alignment.
+    from megatron.core.inference.batch_dimensions_utils import TOKEN_ROUNDER
+
+    inference_multiple = ((TOKEN_ROUNDER + tp_size - 1) // tp_size) * tp_size
+    return math.lcm(configured_multiple, inference_multiple)

nemo_rl/models/megatron/data.py

@@ -30,6 +30,7 @@
 from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank
+from nemo_rl.models.megatron.batch_invariant import batch_invariant_token_multiple
 from nemo_rl.models.megatron.common import _round_up_to_multiple
 from nemo_rl.utils.r3_trace import (
     r3_trace_verify_forward_enabled,
@@ -213,6 +214,17 @@ def get_microbatch_iterator(
             cfg["make_sequence_length_divisible_by"],
             pack_seq_dim_size,
         )
+        if cfg["megatron_cfg"].get("batch_invariant_mode"):
+            # Packed scoring only needs its total token dimension aligned; padding
+            # every constituent sequence to the inference bucket size wastes work.
+            pad_packed_seq_to_multiple_of = batch_invariant_token_multiple(
+                pad_packed_seq_to_multiple_of,
+                cfg["megatron_cfg"]["tensor_model_parallel_size"],
+            )
+            if pad_full_seq_to is not None:
+                pad_full_seq_to = _round_up_to_multiple(
+                    pad_full_seq_to, pad_packed_seq_to_multiple_of
+                )
         micro_batch_size = 1
     else:
         raw_iterator = data.make_microbatch_iterator(mbs)

nemo_rl/models/megatron/setup.py

@@ -228,6 +228,7 @@ def _sync_distrib_opt(distrib_opt):
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.models.generation.megatron.config import (
     dedicated_inference_megatron_cfg,
+    merged_inference_megatron_cfg,
 )
 from nemo_rl.models.megatron.community_import import (
     import_model_from_hf_name,
@@ -259,6 +260,109 @@ def _sync_distrib_opt(distrib_opt):
 TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase)
 
 
+def enable_batch_invariant_mode(config: PolicyConfig) -> None:
+    """Enable Megatron-Core batch-invariant kernels before CUDA initialization.
+
+    The mode is deliberately limited to the topology for which Megatron
+    generation and policy scoring can execute the same arithmetic.
+    Megatron-Core performs the remaining model-specific validation when the
+    provider is finalized. Sampling parameters do not affect this validation:
+    generation samples from the processed distribution while generation and
+    policy scoring both report raw model logprobs.
+
+    Args:
+        config: Policy configuration for this Megatron worker.
+
+    Raises:
+        ValueError: If batch-invariant mode is requested with an unsupported
+            NeMo-RL topology, generation backend, or precision.
+        AssertionError: If the installed Transformer Engine cannot pin the
+            requested FlashAttention version.
+    """
+    megatron_cfg = config["megatron_cfg"]
+    if not megatron_cfg.get("batch_invariant_mode"):
+        return
+
+    required_fields = (
+        "batch_invariant_backend",
+        "batch_invariant_collective",
+        "flash_attention_version",
+    )
+    missing_fields = [field for field in required_fields if field not in megatron_cfg]
+    if missing_fields:
+        raise ValueError(
+            "batch_invariant_mode=True requires policy.megatron_cfg fields: "
+            f"{', '.join(missing_fields)}."
+        )
+
+    if config["precision"] != "bfloat16":
+        raise ValueError(
+            "batch_invariant_mode=True requires policy.precision='bfloat16'."
+        )
+    if megatron_cfg["context_parallel_size"] != 1:
+        raise ValueError(
+            "batch_invariant_mode=True currently requires training context "
+            "parallel size 1."
+        )
+    if megatron_cfg.get("use_fused_linear_logprobs"):
+        raise ValueError(
+            "batch_invariant_mode=True is incompatible with "
+            "use_fused_linear_logprobs=True because generation parity requires "
+            "the shared float-log_softmax-gather path."
+        )
+    if megatron_cfg.get("attention_backend") != "flash":
+        raise ValueError(
+            "batch_invariant_mode=True requires "
+            "policy.megatron_cfg.attention_backend='flash'."
+        )
+    if megatron_cfg["flash_attention_version"] not in (3, 4):
+        raise ValueError(
+            "batch_invariant_mode=True requires "
+            "policy.megatron_cfg.flash_attention_version to be 3 or 4."
+        )
+
+    generation_cfg = config.get("generation")
+    if generation_cfg is None or generation_cfg["backend"] != "megatron":
+        raise ValueError(
+            "batch_invariant_mode=True requires policy.generation.backend='megatron'."
+        )
+    inference_cfg = merged_inference_megatron_cfg(config)
+    matching_fields = (
+        "tensor_model_parallel_size",
+        "context_parallel_size",
+        "batch_invariant_mode",
+        "batch_invariant_backend",
+        "batch_invariant_collective",
+        "attention_backend",
+        "flash_attention_version",
+    )
+    mismatched_fields = [
+        field
+        for field in matching_fields
+        if inference_cfg[field] != megatron_cfg[field]
+    ]
+    if mismatched_fields:
+        raise ValueError(
+            "Training and generation must use the same Megatron settings for "
+            f"batch invariance: {', '.join(mismatched_fields)}."
+        )
+
+    collective = megatron_cfg["batch_invariant_collective"]
+
+    # Keep optional Megatron GPU kernels out of imports when the mode is disabled.
+    from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+        assert_te_supports_batch_invariant_attention,
+    )
+    from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+        enable_batch_invariant_mode as enable_mcore_batch_invariant_mode,
+    )
+
+    assert_te_supports_batch_invariant_attention()
+    enable_mcore_batch_invariant_mode(
+        backend=megatron_cfg["batch_invariant_backend"], collective=collective
+    )
+
+
 def destroy_parallel_state():
     """Safely destroy parallel state and reset async call tracking.
 
@@ -350,12 +454,15 @@ def validate_and_set_config(
         generation_cfg = config["generation"]
         # set generation colocated
         is_generation_colocated = generation_cfg["colocated"]["enabled"]
-        # set sampling params
-        sampling_params = TrainingSamplingParams(
-            top_k=generation_cfg["top_k"],
-            top_p=generation_cfg["top_p"],
-            temperature=generation_cfg["temperature"],
-        )
+        # Batch-invariant Megatron inference returns raw model logprobs even
+        # when token sampling uses temperature, top-k, or top-p. Match
+        # Megatron-RL by recomputing raw training logprobs as well.
+        if not config["megatron_cfg"].get("batch_invariant_mode"):
+            sampling_params = TrainingSamplingParams(
+                top_k=generation_cfg["top_k"],
+                top_p=generation_cfg["top_p"],
+                temperature=generation_cfg["temperature"],
+            )
 
     # Explicitly set NCCL_CUMEM_ENABLE to 1 to avoid the P2P initialization error for PyNCCLCommunicator.
     # See https://github.com/NVIDIA-NeMo/RL/issues/564 for more details.
@@ -1140,6 +1247,21 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None:
                 f"Available backends are: {list(AttnBackend.__members__.keys())}"
             )
 
+    if "batch_invariant_mode" in config["megatron_cfg"]:
+        model_cfg.batch_invariant_mode = config["megatron_cfg"]["batch_invariant_mode"]
+    if "batch_invariant_backend" in config["megatron_cfg"]:
+        model_cfg.batch_invariant_backend = config["megatron_cfg"][
+            "batch_invariant_backend"
+        ]
+    if "batch_invariant_collective" in config["megatron_cfg"]:
+        model_cfg.batch_invariant_collective = config["megatron_cfg"][
+            "batch_invariant_collective"
+        ]
+    if "flash_attention_version" in config["megatron_cfg"]:
+        model_cfg.flash_attention_version = config["megatron_cfg"][
+            "flash_attention_version"
+        ]
+
     # These overrides need to be applied before the workers spawn.
     if "transformer_impl" in config["megatron_cfg"]:
         model_cfg.transformer_impl = config["megatron_cfg"]["transformer_impl"]