From 897649d0bae3ab28a95dbeee05ae18e1b6fae2da Mon Sep 17 00:00:00 2001 From: "Mikail Khona (NVIDIA)" Date: Wed, 12 Aug 2026 00:31:27 -0400 Subject: [PATCH 001/101] Add mixed-precision (FP32) LM output logits (#6252) Signed-off-by: mkhona Signed-off-by: Deepak Narayanan Signed-off-by: Keshav Santhanam Co-authored-by: Deepak Narayanan Co-authored-by: Keshav Santhanam --- .../detxoify_lm/generate_samples_gpt.py | 1 + gpt_builders.py | 1 + hybrid_builders.py | 1 + .../core/extensions/transformer_engine.py | 5 +- megatron/core/models/gpt/gpt_model.py | 6 + megatron/core/models/hybrid/hybrid_model.py | 6 + megatron/core/tensor_parallel/layers.py | 71 ++++++++-- .../elastification/pretrain_hybrid_flex.py | 1 + megatron/post_training/model_builder.py | 2 + megatron/training/argument_utils.py | 2 + megatron/training/arguments.py | 6 + megatron/training/models/gpt.py | 2 + megatron/training/models/hybrid.py | 4 + .../test_te_lmhead_column_parallel_linear.py | 18 +++ .../unit_tests/tensor_parallel/test_layers.py | 132 +++++++++++++++++- tests/unit_tests/test_argument_utils.py | 23 +++ .../training/models/test_gpt_builder.py | 5 + .../training/models/test_hybrid_builder.py | 6 + 18 files changed, 278 insertions(+), 14 deletions(-) diff --git a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py index 2a2b1d63a21..2db8e62ba40 100644 --- a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py +++ b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py @@ -70,6 +70,7 @@ def model_provider(pre_process=True, post_process=True) -> GPTModel: pre_process=pre_process, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=False, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/gpt_builders.py b/gpt_builders.py index 3512918efe6..00a9b3bd2b5 100644 --- a/gpt_builders.py +++ b/gpt_builders.py @@ -91,6 +91,7 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_ pre_process=pre_process, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=True, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/hybrid_builders.py b/hybrid_builders.py index 7e1c58682ac..4ab668eb8bd 100644 --- a/hybrid_builders.py +++ b/hybrid_builders.py @@ -33,6 +33,7 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, pre_process=pre_process, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=True, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index e806a9fdd09..1f67b1a5b5b 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1749,7 +1749,7 @@ class TELMHeadColumnParallelLinear(TEColumnParallelLinear): ``delay_wgrad_compute`` is forced off to mirror its no-op ``backward_dw``, and ``get/set_extra_state`` match the bf16 LM head's state-dict shim. The LM-head kwargs ``keep_master_weight_for_test``, ``skip_weight_param_allocation``, - ``defer_embedding_wgrad_compute`` buffers, and ``disable_grad_reduce`` are + ``defer_embedding_wgrad_compute`` buffers, ``disable_grad_reduce``, and ``output_dtype`` are accepted to preserve the ``ColumnParallelLinear`` signature but currently raise when set non-default — TE will not support them natively, so they would have to be implemented in this subclass, which has not been done yet. @@ -1776,6 +1776,7 @@ def __init__( tp_comm_buffer_name: Optional[str] = None, disable_grad_reduce: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, + output_dtype: Optional[torch.dtype] = None, ): from megatron.core.fp8_utils import is_mxfp8_output_proj_active @@ -1794,6 +1795,8 @@ def __init__( ) if disable_grad_reduce: raise ValueError("TE output projection does not support disable_grad_reduce.") + if output_dtype is not None: + raise ValueError("TE MXFP8 output projection does not support output_dtype.") te_config = copy.copy(config) # Match ColumnParallelLinear.backward_dw's no-op so the LM head keeps diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 3e292ed5d76..ad5f46c1f97 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -66,6 +66,9 @@ class GPTModel(LanguageModule): Include an output layer (used with pipeline parallelism). Defaults to True. fp16_lm_cross_entropy (bool, optional): Defaults to False. + logit_dtype (torch.dtype, optional): + Dtype for the output-layer GEMM result. Defaults to None, which uses + the hidden-state dtype. parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor parallel ranks. Defaults to True. @@ -100,6 +103,7 @@ def __init__( pre_process: bool = True, post_process: bool = True, fp16_lm_cross_entropy: bool = False, + logit_dtype: Optional[torch.dtype] = None, parallel_output: bool = True, share_embeddings_and_output_weights: bool = False, position_embedding_type: Literal[ @@ -133,6 +137,7 @@ def __init__( self.pre_process = pre_process self.post_process = post_process self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.logit_dtype = logit_dtype self.parallel_output = parallel_output self.share_embeddings_and_output_weights = share_embeddings_and_output_weights self.vp_stage = vp_stage @@ -283,6 +288,7 @@ def __init__( embedding_activation_buffer=self.embedding_activation_buffer, grad_output_buffer=self.grad_output_buffer, tp_group=self.pg_collection.tp, + output_dtype=self.logit_dtype, ) if self.pre_process or self.post_process or self.mtp_process: diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index f0358de57b9..65ab3ba1178 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -3,6 +3,7 @@ import logging from typing import Literal, Optional +import torch from torch import Tensor from megatron.core import tensor_parallel @@ -83,6 +84,8 @@ class HybridModel(LanguageModule, GraphableMegatronModule): post_process (bool, optional): Include an output layer (used with pipeline parallelism). Defaults to True. fp16_lm_cross_entropy (bool, optional): Defaults to False. + logit_dtype (torch.dtype, optional): Dtype for the output-layer GEMM result. + Defaults to None, which uses the hidden-state dtype. parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor parallel ranks. Defaults to True. share_embeddings_and_output_weights (bool, optional): When True, input embeddings and @@ -113,6 +116,7 @@ def __init__( pre_process: bool = True, post_process: bool = True, fp16_lm_cross_entropy: bool = False, + logit_dtype: Optional[torch.dtype] = None, parallel_output: bool = True, share_embeddings_and_output_weights: bool = False, # Mamba with no attention has no need for position embeddings, so none is default @@ -144,6 +148,7 @@ def __init__( self.pre_process = pre_process self.post_process = post_process self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.logit_dtype = logit_dtype self.parallel_output = parallel_output self.share_embeddings_and_output_weights = share_embeddings_and_output_weights self.position_embedding_type = position_embedding_type @@ -323,6 +328,7 @@ def __init__( skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights, tp_group=self.pg_collection.tp, + output_dtype=self.logit_dtype, pg_collection=self.pg_collection, ) diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 0f248bcf399..e842e5d9f39 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -411,21 +411,20 @@ class LinearWithFrozenWeight(torch.autograd.Function): @staticmethod @custom_fwd - def forward(ctx, input, weight, bias, allreduce_dgrad, tp_group): + def forward(ctx, input, weight, bias, allreduce_dgrad, tp_group, output_dtype): """Forward with frozen weight.""" ctx.save_for_backward(weight) ctx.allreduce_dgrad = allreduce_dgrad ctx.tp_group = tp_group - output = torch.matmul(input, weight.t()) - if bias is not None: - output = output + bias - return output + ctx.input_dtype = input.dtype + return _linear_forward(input, weight, bias, output_dtype) @staticmethod @custom_bwd def backward(ctx, grad_output): """Backward with frozen weight.""" (weight,) = ctx.saved_tensors + grad_output = grad_output.to(ctx.input_dtype) if grad_output.dim() > 2: # Work around PyTorch matmul not folding some size-1 leading dims to mm. # Remove this once https://github.com/pytorch/pytorch/issues/186148 is fixed. @@ -439,7 +438,34 @@ def backward(ctx, grad_output): # All-reduce. Note: here async and sync are effectively the same. torch.distributed.all_reduce(grad_input, group=ctx.tp_group) - return grad_input, None, None, None, None + return grad_input, None, None, None, None, None + + +def _linear_forward( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + output_dtype: Optional[torch.dtype], +) -> torch.Tensor: + """Run a linear GEMM with an optional output dtype distinct from its input dtype.""" + if output_dtype is None or output_dtype == input.dtype: + output = torch.matmul(input, weight.t()) + if bias is not None: + output = output + bias + return output + + # Deferred to avoid a circular import: transformer_engine imports tensor-parallel layers. + from megatron.core.extensions.transformer_engine import te_general_gemm + + if te_general_gemm is None: + raise RuntimeError( + "A mixed-precision linear output requires Transformer Engine general_gemm." + ) + + input_shape = input.shape + input_2d = input.reshape(-1, input_shape[-1]) + output = te_general_gemm(weight, input_2d, out_dtype=output_dtype, layout="TN", bias=bias)[0] + return output.reshape(*input_shape[:-1], weight.size(0)) def linear_with_frozen_weight( @@ -453,6 +479,7 @@ def linear_with_frozen_weight( grad_output_buffer: Optional[List[torch.Tensor]] = None, wgrad_deferral_limit: None = None, gtp_remat_size: int = 1, + output_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: """Linear layer execution with weight.requires_grad == False. @@ -493,6 +520,9 @@ def linear_with_frozen_weight( gtp_remat_size (int): GTP shard count. When > 1 the weight is GTP-sharded and must be all-gathered to its full shape before the matmul, mirroring the trainable path. Defaults to 1 (no-op) for the common non-GTP / non-sharded case. + + output_dtype (torch.dtype optional): Optional GEMM output dtype. A dtype different from + the input dtype requires Transformer Engine ``general_gemm``. """ assert grad_output_buffer is None, ( @@ -516,7 +546,7 @@ def linear_with_frozen_weight( if gtp_remat_size > 1: weight = weight.all_gather_and_prefetch(fwd=True) - args = [input, weight, bias, allreduce_dgrad, tp_group] + args = [input, weight, bias, allreduce_dgrad, tp_group, output_dtype] return LinearWithFrozenWeight.apply(*args) @@ -538,6 +568,7 @@ def forward( wgrad_deferral_limit, tp_group, gtp_remat_size, + output_dtype, ): """Forward.""" if gradient_accumulation_fusion and hasattr(weight, "main_grad"): @@ -561,6 +592,7 @@ def forward( ctx.grad_output_buffer = grad_output_buffer ctx.tp_group = tp_group ctx.gtp_remat_size = gtp_remat_size + ctx.input_dtype = input.dtype if sequence_parallel: dim_size = list(input.size()) @@ -572,10 +604,7 @@ def forward( else: total_input = input - output = torch.matmul(total_input, weight.t()) - if bias is not None: - output = output + bias - return output + return _linear_forward(total_input, weight, bias, output_dtype) @staticmethod @custom_bwd @@ -584,6 +613,9 @@ def backward(ctx, grad_output): input, weight = ctx.saved_tensors main_grad = ctx.main_grad use_bias = ctx.use_bias + # TE owns only the forward GEMM here; Megatron retains the backward contract. + # Cast dY to the input dtype to match the legacy FP32-logit-cast backward path. + grad_output = grad_output.to(ctx.input_dtype) # GTP: re-gather weight for dgrad if ctx.gtp_remat_size > 1: @@ -743,12 +775,13 @@ def backward(ctx, grad_output): None, None, None, + None, ) if ctx.allreduce_dgrad: handle.wait() - return grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None + return grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None, None def linear_with_grad_accumulation_and_async_allreduce( @@ -762,6 +795,7 @@ def linear_with_grad_accumulation_and_async_allreduce( wgrad_deferral_limit: Optional[int] = 0, tp_group: Optional[torch.distributed.ProcessGroup] = None, gtp_remat_size: int = 1, + output_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: """Linear layer execution with asynchronous communication and gradient accumulation fusion in backprop. @@ -824,6 +858,9 @@ def linear_with_grad_accumulation_and_async_allreduce( wgrad_deferral_limit (int optional): Limit on the number of micro-batches for which embedding weight gradient GEMM should be deferred. Disable by setting this to 0. Defaults to 0. + + output_dtype (torch.dtype optional): Optional GEMM output dtype. A dtype different from + the input dtype requires Transformer Engine ``general_gemm``. """ tp_group = get_tensor_model_parallel_group_if_none(tp_group) @@ -839,6 +876,7 @@ def linear_with_grad_accumulation_and_async_allreduce( wgrad_deferral_limit, tp_group, gtp_remat_size, + output_dtype, ] if not linear_with_grad_accumulation_and_async_allreduce.warned: @@ -912,6 +950,12 @@ class ColumnParallelLinear(torch.nn.Module): If True, reduction of output gradients across tensor-parallel ranks will be disabled. Defaults to False. This feature is used by Lora Adapter in Nemo to delay and fuse reduction along with other gradients for performance optimization. + output_dtype: + Optional dtype for the GEMM output. When it differs from the input dtype, + Transformer Engine ``general_gemm`` is used. + pg_collection: + Optional process group collection. Used to resolve the generalized tensor + parallel remat group; falls back to the global parallel state when omitted. """ def __init__( @@ -934,6 +978,7 @@ def __init__( disable_grad_reduce: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + output_dtype: Optional[torch.dtype] = None, pg_collection: Optional[ProcessGroupCollection] = None, ): super(ColumnParallelLinear, self).__init__() @@ -951,6 +996,7 @@ def __init__( self.config = config self.disable_grad_reduce = disable_grad_reduce self.tp_group = tp_group + self.output_dtype = output_dtype self.tp_group = get_tensor_model_parallel_group_if_none( self.tp_group, is_expert=self.is_expert @@ -1180,6 +1226,7 @@ def forward( ), tp_group=self.tp_group, gtp_remat_size=self.gtp_remat_size, + output_dtype=self.output_dtype, ) gather_output = self.gather_output diff --git a/megatron/elastification/pretrain_hybrid_flex.py b/megatron/elastification/pretrain_hybrid_flex.py index 13eeca1f7c8..adddab019f7 100644 --- a/megatron/elastification/pretrain_hybrid_flex.py +++ b/megatron/elastification/pretrain_hybrid_flex.py @@ -134,6 +134,7 @@ def model_provider(pre_process=True, post_process=True, vp_stage: Optional[int] hybrid_layer_pattern=args.hybrid_layer_pattern, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=True, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py index cd497907406..39c2be07e02 100644 --- a/megatron/post_training/model_builder.py +++ b/megatron/post_training/model_builder.py @@ -383,6 +383,7 @@ def modelopt_gpt_hybrid_builder( "pre_process": pre_process, "post_process": post_process, "fp16_lm_cross_entropy": args.fp16_lm_cross_entropy, + "logit_dtype": getattr(args, "logit_dtype", None), "parallel_output": True, "share_embeddings_and_output_weights": not args.untie_embeddings_and_output_weights, "position_embedding_type": args.position_embedding_type, @@ -428,6 +429,7 @@ def modelopt_gpt_hybrid_builder( "pre_process": pre_process, "post_process": post_process, "fp16_lm_cross_entropy": args.fp16_lm_cross_entropy, + "logit_dtype": getattr(args, "logit_dtype", None), "parallel_output": True, "share_embeddings_and_output_weights": not args.untie_embeddings_and_output_weights, "position_embedding_type": args.position_embedding_type, diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 70f26c64d56..124083ead83 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -457,6 +457,7 @@ def gpt_config_from_args( kwargs["transformer_layer_spec"] = import_module(args.spec) kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy + kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent kwargs["rotary_base"] = args.rotary_base @@ -507,6 +508,7 @@ def hybrid_config_from_args( kwargs["hybrid_stack_spec"] = import_module(args.spec) kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy + kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["hybrid_layer_pattern"] = args.hybrid_layer_pattern kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e54f8a22252..3497834ef93 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1052,6 +1052,7 @@ def validate_args(args, defaults={}): args.mamba_inference_conv_states_dtype = map_dtype(args.mamba_inference_conv_states_dtype) args.mamba_inference_ssm_states_dtype = map_dtype(args.mamba_inference_ssm_states_dtype) args.mamba_training_ssm_states_dtype = map_dtype(args.mamba_training_ssm_states_dtype) + args.logit_dtype = map_dtype(getattr(args, 'logit_dtype', None)) args.megatron_fsdp_main_params_dtype = map_dtype(args.megatron_fsdp_main_params_dtype) args.megatron_fsdp_main_grads_dtype = map_dtype(args.megatron_fsdp_main_grads_dtype) @@ -2935,6 +2936,11 @@ def _add_mixed_precision_args(parser): group.add_argument('--fp16-lm-cross-entropy', action='store_true', help='Move the cross entropy unreduced loss calculation' 'for lm head to fp16.') + group.add_argument('--output-logit-dtype', type=str, choices=['bf16', 'fp32'], default=None, + dest='logit_dtype', + help='Output dtype for the language-model output-layer GEMM. When the ' + 'requested dtype differs from the input dtype, Transformer Engine ' + 'general_gemm is used. By default, logits use the output-layer input dtype.') group.add_argument('--reuse-grad-buf-for-mxfp8-param-ag', action='store_true', help='If True, reuse the grad buffer for MXFP8 parameter all-gather.') group.add_argument('--mamba-training-ssm-states-dtype', type=str, diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py index 63448b196c0..46dcc9b28f4 100644 --- a/megatron/training/models/gpt.py +++ b/megatron/training/models/gpt.py @@ -169,6 +169,7 @@ class GPTModelConfig(ModelConfig): ### GPT Model initialization ### seq_length: int = 1024 fp16_lm_cross_entropy: bool = False + logit_dtype: torch.dtype | None = None parallel_output: bool = True share_embeddings_and_output_weights: bool = False position_embedding_type: Literal["learned_absolute", "rope", "mrope", "yarn", "none"] = "learned_absolute" @@ -316,6 +317,7 @@ def build_model( vocab_size=padded_vocab_size, max_sequence_length=self._model_config.seq_length, fp16_lm_cross_entropy=self._model_config.fp16_lm_cross_entropy, + logit_dtype=self._model_config.logit_dtype, parallel_output=self._model_config.parallel_output, share_embeddings_and_output_weights=self._model_config.share_embeddings_and_output_weights, position_embedding_type=self._model_config.position_embedding_type, diff --git a/megatron/training/models/hybrid.py b/megatron/training/models/hybrid.py index 99f98920eff..c7d1c91b118 100644 --- a/megatron/training/models/hybrid.py +++ b/megatron/training/models/hybrid.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import Any, Callable, ClassVar, Literal, override +import torch + from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.enums import ModelType from megatron.core.models.hybrid.hybrid_layer_specs import ( @@ -50,6 +52,7 @@ class HybridModelConfig(ModelConfig): builder: ClassVar[str] = "megatron.training.models.hybrid.HybridModelBuilder" transformer: TransformerConfig fp16_lm_cross_entropy: bool = False + logit_dtype: torch.dtype | None = None parallel_output: bool = True share_embeddings_and_output_weights: bool = False hybrid_attention_ratio: float = 0.0 @@ -179,6 +182,7 @@ def build_model( max_sequence_length=self._model_config.seq_length, hybrid_layer_pattern=self._model_config.hybrid_layer_pattern, fp16_lm_cross_entropy=self._model_config.fp16_lm_cross_entropy, + logit_dtype=self._model_config.logit_dtype, parallel_output=self._model_config.parallel_output, share_embeddings_and_output_weights=self._model_config.share_embeddings_and_output_weights, position_embedding_type=self._model_config.position_embedding_type, diff --git a/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py b/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py index dbde4c67df4..7d5d58777af 100644 --- a/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py +++ b/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py @@ -92,6 +92,10 @@ def test_rejects_disable_grad_reduce(self): with pytest.raises(ValueError, match="disable_grad_reduce"): TELMHeadColumnParallelLinear(**self._kwargs(disable_grad_reduce=True)) + def test_rejects_output_dtype(self): + with pytest.raises(ValueError, match="output_dtype"): + TELMHeadColumnParallelLinear(**self._kwargs(output_dtype=torch.float32)) + class TestGPTModelOutputLayerSelection: """Verify GPTModel picks the right output-layer class based on config.""" @@ -118,6 +122,20 @@ def test_default_uses_column_parallel_linear(self): assert isinstance(model.output_layer, tensor_parallel.ColumnParallelLinear) assert not isinstance(model.output_layer, TELMHeadColumnParallelLinear) + @pytest.mark.internal + def test_logit_dtype_is_forwarded_to_output_layer(self): + config = TransformerConfig( + num_layers=2, hidden_size=12, num_attention_heads=4, use_cpu_initialization=True + ) + model = GPTModel( + config=config, + transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(), + vocab_size=100, + max_sequence_length=4, + logit_dtype=torch.float32, + ) + assert model.output_layer.output_dtype == torch.float32 + @pytest.mark.internal @pytest.mark.skipif( not _IS_BLACKWELL, reason="MXFP8 output projection requires Blackwell (SM >= 10)" diff --git a/tests/unit_tests/tensor_parallel/test_layers.py b/tests/unit_tests/tensor_parallel/test_layers.py index dbc27f502c6..cf1e8185017 100644 --- a/tests/unit_tests/tensor_parallel/test_layers.py +++ b/tests/unit_tests/tensor_parallel/test_layers.py @@ -2,11 +2,33 @@ import pytest import torch -from megatron.core.tensor_parallel.layers import linear_with_frozen_weight +from megatron.core.extensions.transformer_engine import te_general_gemm +from megatron.core.tensor_parallel.layers import ( + linear_with_frozen_weight, + linear_with_grad_accumulation_and_async_allreduce, +) from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region from tests.unit_tests.test_utilities import Utils +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_linear_default_output_dtype_preserves_input_dtype(dtype): + Utils.initialize_model_parallel(1, 1) + + try: + input_data = torch.randn(4, 3, 16, device="cuda", dtype=dtype) + weight = torch.randn(32, 16, device="cuda", dtype=dtype) + output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=None + ) + reference = torch.nn.functional.linear(input_data, weight) + + assert output.dtype == input_data.dtype + torch.testing.assert_close(output, reference) + finally: + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("tensor_parallel,allreduce_dgrad", [(1, False), (8, True)]) def test_LinearWithFrozenWeight(tensor_parallel, allreduce_dgrad): Utils.initialize_model_parallel(tensor_parallel, 1) @@ -70,3 +92,111 @@ def test_LinearWithFrozenWeight_3d_input_matches_torch_linear(): assert torch.allclose(input_data.grad, expected_input.grad) Utils.destroy_model_parallel() + + +@pytest.mark.skipif( + te_general_gemm is None, reason="Transformer Engine general_gemm is not available" +) +def test_linear_with_grad_accumulation_supports_fp32_output_and_bf16_backward(): + Utils.initialize_model_parallel(1, 1) + + input_data = torch.randn(4, 3, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(32, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + reference_input = input_data.detach().clone().requires_grad_(True) + reference_weight = weight.detach().clone().requires_grad_(True) + + output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=torch.float32 + ) + output.sum().backward() + + reference_output = torch.nn.functional.linear(reference_input, reference_weight) + reference_output.sum().backward() + fp32_reference_output = torch.nn.functional.linear( + input_data.detach().float(), weight.detach().float() + ) + + assert output.dtype == torch.float32 + assert input_data.grad.dtype == torch.bfloat16 + assert weight.grad.dtype == torch.bfloat16 + assert torch.allclose(output, fp32_reference_output, atol=1e-4, rtol=1e-4) + assert torch.allclose(input_data.grad, reference_input.grad) + assert torch.allclose(weight.grad, reference_weight.grad) + + Utils.destroy_model_parallel() + + +@pytest.mark.skipif( + te_general_gemm is None, reason="Transformer Engine general_gemm is not available" +) +def test_linear_fp32_output_is_bitwise_exact_for_integer_bf16_operands(): + Utils.initialize_model_parallel(1, 1) + + generator = torch.Generator(device="cuda").manual_seed(1234) + input_data = torch.randint( + -8, 9, (4, 3, 512), device="cuda", dtype=torch.int32, generator=generator + ).to(torch.bfloat16) + weight = torch.randint( + -8, 9, (128, 512), device="cuda", dtype=torch.int32, generator=generator + ).to(torch.bfloat16) + + output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=torch.float32 + ) + reference = torch.nn.functional.linear(input_data.float(), weight.float()) + + # K * 8^2 = 32,768, so every possible integer product and partial sum is + # exactly representable in FP32. Compare raw words instead of using a tolerance. + assert output.dtype == torch.float32 + assert torch.equal(output.contiguous().view(torch.int32), reference.view(torch.int32)) + + Utils.destroy_model_parallel() + + +@pytest.mark.skipif( + te_general_gemm is None, reason="Transformer Engine general_gemm is not available" +) +def test_linear_fp32_output_matches_plain_te_general_gemm(): + from transformer_engine.pytorch.cpp_extensions import general_gemm + + try: + from transformer_engine.pytorch.module.base import get_workspace + except ImportError: + get_workspace = None + + Utils.initialize_model_parallel(1, 1) + + input_data = torch.randn(4, 3, 64, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(96, 64, device="cuda", dtype=torch.bfloat16) + wrapped_output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=torch.float32 + ) + + kwargs = { + "out_dtype": torch.float32, + "quantization_params": None, + "gelu": None, + "gelu_in": None, + "accumulate": False, + "layout": "TN", + "out": None, + "bias": None, + "use_split_accumulator": False, + "grad": False, + "ub": None, + "ub_type": None, + "extra_output": None, + "bulk_overlap": False, + } + if get_workspace is not None: + kwargs["workspace"] = get_workspace() + plain_te_output = general_gemm(weight, input_data.reshape(-1, 64), **kwargs)[0] + plain_te_output = plain_te_output.reshape_as(wrapped_output) + + assert wrapped_output.dtype == torch.float32 + assert torch.equal( + wrapped_output.contiguous().view(torch.int32), + plain_te_output.contiguous().view(torch.int32), + ) + + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/test_argument_utils.py b/tests/unit_tests/test_argument_utils.py index 7c0b30d3d56..57a327cddfd 100644 --- a/tests/unit_tests/test_argument_utils.py +++ b/tests/unit_tests/test_argument_utils.py @@ -678,6 +678,29 @@ def test_transformer_callback_fields_are_not_registered_as_cli_args(self): assert not hasattr(args, field_name) +class TestMegatronMixedPrecisionArguments: + """Test language-model logit dtype CLI choices.""" + + @staticmethod + def _parser() -> ArgumentParser: + from megatron.training.arguments import _add_mixed_precision_args + + return _add_mixed_precision_args(ArgumentParser(exit_on_error=False)) + + def test_logit_dtype_defaults_to_input_dtype(self): + args = self._parser().parse_args([]) + assert args.logit_dtype is None + + @pytest.mark.parametrize("dtype", ["bf16", "fp32"]) + def test_logit_dtype_accepts_supported_choices(self, dtype): + args = self._parser().parse_args(["--output-logit-dtype", dtype]) + assert args.logit_dtype == dtype + + def test_logit_dtype_rejects_fp16(self): + with pytest.raises(ArgumentError, match="invalid choice"): + self._parser().parse_args(["--output-logit-dtype", "fp16"]) + + # --------------------------------------------------------------------------- # Tests for pretrain_cfg_container_from_args # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/training/models/test_gpt_builder.py b/tests/unit_tests/training/models/test_gpt_builder.py index 2263525e030..5ae7ab4f418 100644 --- a/tests/unit_tests/training/models/test_gpt_builder.py +++ b/tests/unit_tests/training/models/test_gpt_builder.py @@ -261,6 +261,7 @@ def test_default_values(self): assert config.should_pad_vocab is False assert config.seq_length == 1024 assert config.fp16_lm_cross_entropy is False + assert config.logit_dtype is None assert config.parallel_output is True assert config.share_embeddings_and_output_weights is False assert config.position_embedding_type == "learned_absolute" @@ -279,6 +280,7 @@ def test_custom_initialization(self): transformer=_make_transformer(), seq_length=4096, fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, share_embeddings_and_output_weights=True, position_embedding_type="rope", @@ -287,6 +289,7 @@ def test_custom_initialization(self): ) assert config.seq_length == 4096 assert config.fp16_lm_cross_entropy is True + assert config.logit_dtype == torch.float32 assert config.parallel_output is False assert config.share_embeddings_and_output_weights is True assert config.position_embedding_type == "rope" @@ -678,6 +681,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): vocab_size=32000, seq_length=4096, fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, share_embeddings_and_output_weights=True, position_embedding_type="rope", @@ -699,6 +703,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): assert kw["vocab_size"] == 32000 assert kw["max_sequence_length"] == 4096 assert kw["fp16_lm_cross_entropy"] is True + assert kw["logit_dtype"] == torch.float32 assert kw["parallel_output"] is False assert kw["share_embeddings_and_output_weights"] is True assert kw["position_embedding_type"] == "rope" diff --git a/tests/unit_tests/training/models/test_hybrid_builder.py b/tests/unit_tests/training/models/test_hybrid_builder.py index 9984e224ce3..fa304f7ca34 100644 --- a/tests/unit_tests/training/models/test_hybrid_builder.py +++ b/tests/unit_tests/training/models/test_hybrid_builder.py @@ -3,6 +3,7 @@ from unittest.mock import Mock, call, patch import pytest +import torch from megatron.core.transformer import ModuleSpec from megatron.core.transformer.transformer_config import TransformerConfig @@ -39,6 +40,7 @@ def test_builder_classvar(self): def test_default_values(self): config = HybridModelConfig(transformer=_make_transformer()) assert config.fp16_lm_cross_entropy is False + assert config.logit_dtype is None assert config.parallel_output is True assert config.share_embeddings_and_output_weights is False assert config.hybrid_layer_pattern is None @@ -55,6 +57,7 @@ def test_custom_initialization(self): config = HybridModelConfig( transformer=_make_transformer(), fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, hybrid_attention_ratio=0.25, hybrid_mlp_ratio=0.1, @@ -63,6 +66,7 @@ def test_custom_initialization(self): vocab_size=50000, ) assert config.fp16_lm_cross_entropy is True + assert config.logit_dtype == torch.float32 assert config.parallel_output is False assert config.hybrid_attention_ratio == 0.25 assert config.hybrid_mlp_ratio == 0.1 @@ -329,6 +333,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): seq_length=4096, hybrid_layer_pattern="M-A-", fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, share_embeddings_and_output_weights=True, position_embedding_type="rope", @@ -345,6 +350,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): assert kw["max_sequence_length"] == 4096 assert kw["hybrid_layer_pattern"] == "M-A-" assert kw["fp16_lm_cross_entropy"] is True + assert kw["logit_dtype"] == torch.float32 assert kw["parallel_output"] is False assert kw["share_embeddings_and_output_weights"] is True assert kw["position_embedding_type"] == "rope" From 14346b65a2d0790e451919858f7771078105c5f0 Mon Sep 17 00:00:00 2001 From: Antoni-Joan Solergibert Date: Wed, 12 Aug 2026 07:06:55 +0200 Subject: [PATCH 002/101] Eliminate the ShardedObject all_gather_object on the FullyParallel load path (#5551) Signed-off-by: asolergi-nv --- .../strategies/fully_parallel.py | 73 +++++++++++++++---- megatron/training/checkpointing.py | 7 +- megatron/training/config/training_config.py | 8 ++ .../model_config.yaml | 11 +++ 4 files changed, 82 insertions(+), 17 deletions(-) diff --git a/megatron/core/dist_checkpointing/strategies/fully_parallel.py b/megatron/core/dist_checkpointing/strategies/fully_parallel.py index ca7f28f3cc6..657942cadb0 100644 --- a/megatron/core/dist_checkpointing/strategies/fully_parallel.py +++ b/megatron/core/dist_checkpointing/strategies/fully_parallel.py @@ -178,6 +178,7 @@ def __init__( parallelization_group: Optional[torch.distributed.ProcessGroup] = None, do_cache_distribution: bool = False, exchange_algo: str = 'broadcast', + per_rank_object_load: bool = False, ): self.base_strategy = strategy if parallelization_group is None: @@ -187,6 +188,11 @@ def __init__( self.parallelization_group = parallelization_group self.do_cache_distribution = do_cache_distribution self.exchange_algo = exchange_algo + # When True, every rank loads *all* of its ShardedObjects directly from + # storage instead of loading only its main replicas and exchanging them + # with a WORLD-wide `all_gather_object`. Opt-in; defaults to the legacy + # gather-based exchange. See `load` for correctness rationale. + self.per_rank_object_load = per_rank_object_load self.cached_distribution: Optional[ShardDistribution] = None self.cached_global_metadata: Optional[Metadata] = None @@ -256,15 +262,47 @@ def load( assert ( len(sharded_state_dict) == 0 ), "sharded_state_dict is not empty after deferring tensors and objects" - with debug_time("base_load_ShardedObjects", logger): - # Load sharded objects first - loaded_objects = self.base_strategy.load( - to_load_objects, checkpoint_dir, async_strategy - ) - with debug_time("base_load_ShardedTensors", logger): - # Load sharded tensors separately - loaded_tensors = self.base_strategy.load(to_load_shards, checkpoint_dir, async_strategy) + if self.per_rank_object_load: + # Opt-in: every rank loads *all* of its own objects directly from + # storage, so no inter-rank object exchange is needed. This is correct + # because an object is addressed in the checkpoint by its `unique_key` + # (key + global_offset + global_shape), which excludes `replica_id`; + # each distinct object position is therefore written to disk exactly + # once (by its main replica) and any rank that needs it can read it + # locally. We merge both the main (`to_load_objects`) and non-main + # (`unloaded_objects`) replica maps so this rank loads every object in + # its state dict. This replaces the WORLD-wide `all_gather_object` + # collective with extra (but cheap) local reads of small artifacts + # (RNG states, `_extra_state`, ...). + # + # Objects are loaded together with this rank's tensor shards in a + # single base-strategy `.load()` call: `mcore_to_pyt_state_dict` + # supports a mixed tensor/object state dict, and one call means one + # metadata read and one load plan instead of two. + all_objects_to_load = {**to_load_objects, **unloaded_objects} + with debug_time("base_load_ShardedTensorsAndObjects", logger): + loaded = self.base_strategy.load( + {**to_load_shards, **all_objects_to_load}, checkpoint_dir, async_strategy + ) + # The base strategy returns the loaded values keyed by the same shard + # ids we passed in; split them back into tensors and objects. Tensor + # and object shard ids never collide (a given key is either a tensor + # or an object), so membership in the original maps is an unambiguous + # split. + loaded_tensors = {shard_id: loaded[shard_id] for shard_id in to_load_shards} + loaded_objects = {shard_id: loaded[shard_id] for shard_id in all_objects_to_load} + else: + # Default (legacy): load only this rank's main-replica objects and + # exchange them across ranks below. + with debug_time("base_load_ShardedObjects", logger): + loaded_objects = self.base_strategy.load( + to_load_objects, checkpoint_dir, async_strategy + ) + with debug_time("base_load_ShardedTensors", logger): + loaded_tensors = self.base_strategy.load( + to_load_shards, checkpoint_dir, async_strategy + ) with debug_time("self.exchange_loaded_tensors", logger): @@ -286,14 +324,17 @@ def load( with debug_time("torch.cuda.synchronize", logger): torch.cuda.synchronize() - all_loaded_objects = exchange_loaded_objects_gather_object(loaded_objects) - - if not set(unloaded_objects.keys()).issubset(all_loaded_objects.keys()): - missing_object_shards = set(unloaded_objects.keys()) - all_loaded_objects.keys() - raise CheckpointingException( - f'Missing object shards after fully parallel loading: {missing_object_shards}' - ) - torch.cuda.synchronize() + if self.per_rank_object_load: + # No object exchange: each rank already loaded every object it needs. + all_loaded_objects = loaded_objects + else: + all_loaded_objects = exchange_loaded_objects_gather_object(loaded_objects) + if not set(unloaded_objects.keys()).issubset(all_loaded_objects.keys()): + missing_object_shards = set(unloaded_objects.keys()) - all_loaded_objects.keys() + raise CheckpointingException( + f'Missing object shards after fully parallel loading: {missing_object_shards}' + ) + torch.cuda.synchronize() self.fill_in_deferred_sharded_tensors(sharded_tensors, all_loaded_tensors) self.fill_in_deferred_sharded_objects(sharded_objects, all_loaded_objects) diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index fa859664a46..a81a6c2dfe1 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1748,7 +1748,12 @@ def _load_global_dist_base_checkpoint( ) load_strategy = FullyParallelLoadStrategyWrapper( - load_strategy, process_group, exchange_algo=args.ckpt_fully_parallel_load_exchange_algo + load_strategy, + process_group, + exchange_algo=args.ckpt_fully_parallel_load_exchange_algo, + per_rank_object_load=getattr( + args, 'ckpt_fully_parallel_load_per_rank_objects', False + ), ) if checkpointing_context is not None: checkpointing_context['load_strategy'] = load_strategy diff --git a/megatron/training/config/training_config.py b/megatron/training/config/training_config.py index 1ba743d5962..a379039ff20 100644 --- a/megatron/training/config/training_config.py +++ b/megatron/training/config/training_config.py @@ -548,6 +548,14 @@ class CheckpointConfig: "gather_object": Gather the checkpoint from all ranks in a single operation. """ + ckpt_fully_parallel_load_per_rank_objects: bool = False + """Load ShardedObjects per-rank during fully parallel load of distributed checkpoints. + When True, every rank reads all of its own ShardedObjects (RNG states, + TE `_extra_state`, ...) directly from storage, which removes the WORLD-wide + `all_gather_object` that otherwise exchanges them. Objects are + content-addressable by `unique_key`, so the loaded values are identical. + When False (default), the legacy gather-based object exchange is used.""" + ckpt_fully_parallel_save_process_group: Literal["dp", "ep_dp"] = "dp" """Process group for fully parallel save of distributed checkpoints. "dp"(default): Data parallel process group. diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml index 10d94d27f03..f913c1a62ca 100644 --- a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml +++ b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml @@ -141,6 +141,17 @@ MODEL_ARGS: --ckpt-fully-parallel-load: true --dist-ckpt-strictness: log_all + # Checkpointing optimizations exercised by this load. All three are meant to be + # numerically neutral, so the golden values must stay identical with them on. + # Every rank reads its own ShardedObjects instead of exchanging them with a + # WORLD all_gather_object. + --ckpt-fully-parallel-load-per-rank-objects: true + # Skips the determine_global_metadata all_gather_object that validates sharding. + --no-ckpt-load-validate-sharding-integrity: true + # Redundant TE _extra_state is neither written nor requested on load; only + # delayed-scaling FP8 extra state is ever needed, and this model does not use it. + --ckpt-drop-redundant-extra-state: true + # Validation and functional metrics. eval-interval < train-iters so the # evaluation loop actually runs during the test. --eval-interval: 10 From 15c83d2fcd00e283bb59ff26dce40266a445c615 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 12 Aug 2026 09:26:08 +0000 Subject: [PATCH 003/101] chore: rotate oncall schedule --- .github/oncall_schedule.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/oncall_schedule.json b/.github/oncall_schedule.json index 0fe3ecb86ab..622a830a73e 100644 --- a/.github/oncall_schedule.json +++ b/.github/oncall_schedule.json @@ -1,8 +1,4 @@ [ - { - "user": "janEbert", - "date": "2026-08-05" - }, { "user": "maanug-nv", "date": "2026-08-12" @@ -46,5 +42,9 @@ { "user": "janEbert", "date": "2026-10-21" + }, + { + "user": "maanug-nv", + "date": "2026-10-28" } ] From 81fe7c746f6556e5d8dc24196a1358ef0898dbf7 Mon Sep 17 00:00:00 2001 From: Jiangfei Duan <30710061+JF-D@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:08:25 -0700 Subject: [PATCH 004/101] GTP + PartialCG: Classify latent projections for partial CUDA graphs (#6446) Signed-off-by: Jiangfei Duan --- .../generalized_tensor_parallelism.py | 4 + .../test_gtp_basics.py | 36 +++++ .../test_gtp_partial_cg.py | 144 +++++++++++++----- 3 files changed, 150 insertions(+), 34 deletions(-) diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index de97cdab121..b8d2a8faff2 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -217,6 +217,10 @@ def _classify_param_chain(param_name: str) -> str: if not scope: # CG disabled return U + # MoE latent projections. + if ".mlp.fc1_latent_proj." in n or ".mlp.fc2_latent_proj." in n: + return G if "moe_router" in scope else U + if ".mlp.shared_experts." in n: if _MOE_SHARED_EXPERT_OVERLAP: return U diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py index 297c89578f1..6c594d0be62 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py @@ -502,6 +502,42 @@ def test_graphness_helpers(self): assert gtp_module._chain_is_graphed("GTP_graphed") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="GTPShardedParam requires CUDA") +class TestLatentProjectionChainClassification: + """Classify opted-in latent projections through the public GTP chain API.""" + + FC1 = "decoder.layers.3.mlp.fc1_latent_proj.weight" + FC2 = "decoder.layers.3.mlp.fc2_latent_proj.weight" + + def teardown_method(self, method): + gtp_module.set_cuda_graph_modules(None, cuda_graph_impl="none") + gtp_module.reset_gtp_state() + + def _chains(self, *, cuda_graph_modules=None, cuda_graph_impl="none"): + params = tuple(GTPShardedParam(torch.zeros(1, device="cuda")) for _ in range(2)) + assert all(isinstance(param, GTPShardedParam) for param in params) + + class _Model: + def named_parameters(_self): + return iter(zip((self.FC1, self.FC2), params)) + + gtp_module.classify_gtp_remat_chains( + _Model(), cuda_graph_modules=cuda_graph_modules, cuda_graph_impl=cuda_graph_impl + ) + return tuple(param.chain_id for param in params) + + def test_eager_latent_projections_are_ungraphed(self): + assert self._chains() == (GTPChain.UNGRAPHED.value, GTPChain.UNGRAPHED.value) + + def test_local_router_captures_both_latent_projections(self): + chains = self._chains(cuda_graph_modules={"moe_router"}, cuda_graph_impl="local") + assert chains == (GTPChain.GRAPHED.value, GTPChain.GRAPHED.value) + + def test_unrelated_local_scope_leaves_latent_projections_ungraphed(self): + chains = self._chains(cuda_graph_modules={"mamba", "attn"}, cuda_graph_impl="local") + assert chains == (GTPChain.UNGRAPHED.value, GTPChain.UNGRAPHED.value) + + class TestGroupedDoubleBuffer: """One-block-ahead grouped chains must double-buffer: consecutive MoE layers get distinct gather buffers (else prefetching layer N+1 clobbers layer N's in-use weight). Pure cache-key diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py index addcb9cb657..14641c3a3b4 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py @@ -36,9 +36,10 @@ ) -def _worker_gtp_partial_cg_correctness(rank, world_size, port): - """Compare eager and local attention CUDA graphs with GTP2 x DP2.""" +def _worker_gtp_partial_cg_correctness(rank, world_size, port, partial_cg_modules, opt_in_modules): + """Compare eager and local CUDA graphs with GTP2 x DP2.""" del port + gtp_module._GTP_PARAMS.clear() from megatron.core import parallel_state as ps from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec @@ -56,14 +57,15 @@ def _worker_gtp_partial_cg_correctness(rank, world_size, port): ) from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.transformer_config import TransformerConfig - - hidden = 4096 - num_heads = 32 - ffn_hidden = 16384 - # Four layers force parameters with matching scheduling domains/shapes to reuse the two-slot - # wgrad ring across independently replayed graphs. - num_layers = 4 - sequence_length = 32 + from megatron.core.transformer.transformer_layer import MoETransformerLayer + + latent_projection_case = "moe_latent_proj" in opt_in_modules + hidden = 256 if latent_projection_case else 4096 + num_heads = 8 if latent_projection_case else 32 + ffn_hidden = 512 if latent_projection_case else 16384 + # Use multiple layers to exercise repeated local CUDA-graph execution with GTP parameters. + num_layers = 1 if latent_projection_case else 4 + sequence_length = 16 if latent_projection_case else 32 batch_size = 1 learning_rate = 0.01 steps = 10 @@ -73,6 +75,18 @@ def _worker_gtp_partial_cg_correctness(rank, world_size, port): assert world_size == gtp_degree * dp_degree def make_config(*, partial_cg=False): + moe_options = {} + if latent_projection_case: + moe_options = { + "num_moe_experts": 2, + "moe_router_topk": 1, + "moe_router_pre_softmax": True, + "moe_ffn_hidden_size": ffn_hidden, + "moe_grouped_gemm": True, + "moe_token_dispatcher_type": "allgather", + "moe_aux_loss_coeff": 0.0, + "moe_latent_size": 128, + } return TransformerConfig( num_attention_heads=num_heads, num_layers=num_layers, @@ -83,16 +97,40 @@ def make_config(*, partial_cg=False): hidden_dropout=0.0, attention_dropout=0.0, bias_dropout_fusion=False, + gradient_accumulation_fusion=latent_projection_case, tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_weight_remat_size=gtp_degree, + gtp_remat_opt_in_modules=opt_in_modules, cuda_graph_impl="local" if partial_cg else "none", - cuda_graph_modules=["attn"] if partial_cg else [], + cuda_graph_modules=partial_cg_modules if partial_cg else [], cuda_graph_warmup_steps=2, + **moe_options, ) - def make_attention_stack(config, pg_collection): - spec = copy.deepcopy(get_gpt_layer_with_transformer_engine_spec()) + def make_layer_stack(config, pg_collection): + spec = copy.deepcopy( + get_gpt_layer_with_transformer_engine_spec( + num_experts=2 if latent_projection_case else None, + moe_grouped_gemm=latent_projection_case, + ) + ) + if latent_projection_case: + spec.submodules.input_layernorm = IdentityOp + spec.submodules.self_attention = IdentityOp + spec.submodules.self_attn_bda = IdentityFuncOp + return torch.nn.ModuleList( + [ + MoETransformerLayer( + config, + spec.submodules, + layer_number=1, + pg_collection=pg_collection, + name="decoder.layers.0", + ) + ] + ) + spec.submodules.pre_mlp_layernorm = IdentityOp spec.submodules.mlp = IdentityOp spec.submodules.mlp_bda = IdentityFuncOp @@ -105,6 +143,27 @@ def make_attention_stack(config, pg_collection): ] ) + def get_cudagraph_managers(layers): + if latent_projection_case: + return [ + manager + for layer in layers + for manager in (layer.cudagraph_manager_router, layer.cudagraph_manager_postprocess) + ] + return [layer.cudagraph_manager for layer in layers] + + def get_latent_params(layers): + return [ + param + for name, param in layers.named_parameters() + if "fc1_latent_proj.weight" in name or "fc2_latent_proj.weight" in name + ] + + def make_pg_collection(): + if latent_projection_case: + return ProcessGroupCollection.use_mpu_process_groups() + return ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp", "gtp_remat"]) + def run_step(layers, x): with fp8_autocast(enabled=False): for layer in layers: @@ -162,17 +221,19 @@ def apply_sgd_step(layers, gtp_size): tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=gtp_degree ) model_parallel_cuda_manual_seed(42) - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["tp", "cp", "gtp_remat"] - ) + pg_collection = make_pg_collection() eager_config = make_config() - eager = make_attention_stack(eager_config, pg_collection).cuda() + eager = make_layer_stack(eager_config, pg_collection).cuda() eager_gtp_group = ps.get_gtp_weight_remat_group() eager_dp_group = ps.get_data_parallel_group(with_gtp_remat=False) eager_dp_rank = eager_dp_group.rank() assert eager_gtp_group.size() == gtp_degree assert eager_dp_group.size() == dp_degree assert any(isinstance(param, GTPShardedParam) for param in eager.parameters()) + if latent_projection_case: + eager_latent_params = get_latent_params(eager) + assert len(eager_latent_params) == 2 + assert all(isinstance(param, GTPShardedParam) for param in eager_latent_params) initialize_main_grads(eager) saved_local_weights = {name: param.data.clone() for name, param in eager.named_parameters()} @@ -190,20 +251,20 @@ def apply_sgd_step(layers, gtp_size): apply_sgd_step(eager, eager_gtp_group.size()) del eager, loss, x + torch.cuda.synchronize() ps.destroy_model_parallel() gtp_module.reset_gtp_state() + gtp_module._GTP_PARAMS.clear() - # Optimized path: the same GTP2 x DP2 topology with attention-only local CUDA graphs. + # Optimized path: the same GTP2 x DP2 topology with local CUDA graphs. ps.initialize_model_parallel( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=gtp_degree ) initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) model_parallel_cuda_manual_seed(42) - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["tp", "cp", "gtp_remat"] - ) + pg_collection = make_pg_collection() partial_cg_config = make_config(partial_cg=True) - partial_cg = make_attention_stack(partial_cg_config, pg_collection).cuda() + partial_cg = make_layer_stack(partial_cg_config, pg_collection).cuda() classify_gtp_remat_chains( partial_cg, cuda_graph_modules=partial_cg_config.cuda_graph_modules, @@ -220,7 +281,8 @@ def apply_sgd_step(layers, gtp_size): assert dp_rank == eager_dp_rank gtp_params = [param for param in partial_cg.parameters() if isinstance(param, GTPShardedParam)] assert gtp_params, "GTP not active: no GTPShardedParam found" - assert all(param.chain_id == GTPChain.GRAPHED.value for param in gtp_params) + params_in_scope = get_latent_params(partial_cg) if latent_projection_case else gtp_params + assert all(param.chain_id == GTPChain.GRAPHED.value for param in params_in_scope) for name, param in partial_cg.named_parameters(): param.data.copy_(saved_local_weights[name]) # Production captures after DDP maps every parameter into a main-grad buffer and initializes @@ -241,18 +303,21 @@ def apply_sgd_step(layers, gtp_size): wait_for_gtp_grad_reduction_on_current_stream() eager_grad_norm = global_grad_norm(partial_cg, gtp_group) eager_probe_loss_value = eager_probe_loss.item() + del eager_probe_loss, eager_probe_x + reset_grad_state(partial_cg) create_cudagraphs() assert _CudagraphGlobalRecord.cudagraph_created - runners = [layer.cudagraph_manager.cudagraph_runners[0] for layer in partial_cg] - assert all(runner.gtp_remat for runner in runners) - assert any(runner._gtp_wgrad_ring_slots for runner in runners) + managers = get_cudagraph_managers(partial_cg) + assert all(len(manager.cudagraph_runners) == 1 for manager in managers) + runners = [manager.cudagraph_runners[0] for manager in managers] + assert any(runner.gtp_remat for runner in runners) replay_grad_norms = [] replay_losses = [] for _ in range(3): reset_grad_state(partial_cg) - replay_x = eager_probe_x.detach().clone().requires_grad_() + replay_x = make_replica_input(1234, dp_rank).requires_grad_() replay_loss = run_step(partial_cg, replay_x) replay_loss.backward() wait_for_gtp_grad_reduction_on_current_stream() @@ -280,7 +345,7 @@ def apply_sgd_step(layers, gtp_size): flush=True, ) - del eager_probe_loss, eager_probe_x, replay_loss, replay_x + del replay_loss, replay_x for step in range(steps): reset_grad_state(partial_cg) @@ -295,19 +360,21 @@ def apply_sgd_step(layers, gtp_size): del loss, x finally: torch.cuda.synchronize() - for layer in partial_cg: - for runner in layer.cudagraph_manager.cudagraph_runners: + managers = get_cudagraph_managers(partial_cg) + for manager in managers: + for runner in manager.cudagraph_runners: if runner.fwd_graph is not None: runner.fwd_graph.reset() if runner.bwd_graph is not None: runner.bwd_graph.reset() delete_cuda_graphs() - for layer in partial_cg: - layer.cudagraph_manager.cudagraph_runners.clear() + for manager in managers: + manager.cudagraph_runners.clear() gc.collect() ps.destroy_model_parallel() ps.initialize_model_parallel() gtp_module.reset_gtp_state() + gtp_module._GTP_PARAMS.clear() if rank == 0: for step, (eager_loss, partial_cg_loss) in enumerate(zip(eager_losses, partial_cg_losses)): @@ -324,8 +391,17 @@ def apply_sgd_step(layers, gtp_size): class TestGTPPartialCGCorrectness: - def test_gtp_partial_cg_loss_and_grad_norm_match_eager(self): + @pytest.mark.parametrize( + "partial_cg_modules,opt_in_modules", + [ + pytest.param(["attn"], [], id="attention"), + pytest.param(["moe_router"], ["moe_latent_proj"], id="moe-router-latent-projections"), + ], + ) + def test_gtp_partial_cg_loss_and_grad_norm_match_eager( + self, partial_cg_modules, opt_in_modules + ): """Local-CG loss trajectory and global grad norm must match eager execution.""" if torch.cuda.device_count() < 4: pytest.skip("Requires at least 4 CUDA devices") - _run_distributed(_worker_gtp_partial_cg_correctness, 4) + _run_distributed(_worker_gtp_partial_cg_correctness, 4, partial_cg_modules, opt_in_modules) From 85344907be6a9985410bf528acbedaca77b2a020 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Wed, 12 Aug 2026 07:13:41 -0700 Subject: [PATCH 005/101] Balance LayerWise optimizer shards by Newton-Schulz cost, not parameter size (#6379) Signed-off-by: Deepak Narayanan Co-authored-by: Claude Opus 4.6 --- .../core/optimizer/layer_wise_optimizer.py | 69 +++++++-- .../test_layer_wise_param_layout.py | 134 ++++++++++++++++++ 2 files changed, 194 insertions(+), 9 deletions(-) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index d1d792ca76d..a791d262954 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -213,15 +213,19 @@ def _compute_per_buffer_param_layout( followed by an isolated bucket for that embedding alone. 2. When the chunk's total numel reaches ``bucket_size`` (or all params have been consumed), bin-pack the chunk into ``dp_size`` - shards via greedy LPT — sort by numel descending and assign each - param to the shard with the smallest current load. + shards: sort by estimated Newton-Schulz compute cost descending and + assign each param to the shard with the smallest accumulated compute + load, subject to a per-bucket numel cap that bounds shard-imbalance + padding. Compute loads persist across buckets so expensive + (GTP-sharded) matrices spread over the whole buffer instead of + clustering inside each bucket. 3. Pad each shard to ``max(shard_cursors)`` aligned to :meth:`_shard_divisor`, then emit the bucket. Each bucket therefore spans a contiguous backprop range so that ``overlap_grad_reduce`` can dispatch the bucket's reduce-scatter as soon as the bucket's backward segment finishes — preserving the - original DDP overlap semantics. LPT bin-packing keeps shards close + original DDP overlap semantics. Greedy bin-packing keeps shards close to balanced; for uniform transformer blocks where ``params_per_layer * num_layers`` is a multiple of ``dp_size`` the packing is perfect. @@ -247,6 +251,27 @@ def _compute_per_buffer_param_layout( bucket_id = 0 shard_imbalance_padding_numel = 0 + # Persistent compute loads across buckets so LPT spreads expensive + # (GTP-sharded) params evenly instead of clustering them per bucket. + shard_compute_loads = [0] * dp_size + + def _ns_compute_cost(param): + """Estimate Newton-Schulz compute cost for a parameter. + + Newton-Schulz only runs on matrices, so anything that is not 2D falls + back to its element count. For a 2D param the cost is + ~ max(M,N) * min(M,N)^2, the dominant term in the orthogonalization; + GTP-sharded params reconstruct the full post-AllGather shape first + (GTP always shards along dim 0). + """ + if param.dim() != 2: + return param.data.nelement() + m, n = param.data.shape + if getattr(param, 'is_gtp_weight_remat', False): + m = m * getattr(param, 'gtp_remat_size', 1) + big, small = max(m, n), min(m, n) + return big * small * small + def _emit_bucket( chunk_params: List[torch.nn.Parameter], shared_embedding: bool = False ) -> None: @@ -276,17 +301,33 @@ def _emit_bucket( shard_assignments[shard_id].append((None, numel)) shard_cursors[shard_id] = numel else: - # Greedy LPT: largest first, assign to the least-loaded shard. - # The within-shard order is sorted-by-numel, not backprop — + # Compute-balanced LPT: sort by Newton-Schulz compute cost + # (accounts for full post-AllGather shape under GTP), assign to + # the shard with least accumulated compute load. Compute loads + # persist across buckets; numel cursors reset per bucket. + # A per-bucket numel cap prevents excessive padding. + # The within-shard order is sorted-by-compute-cost, not backprop; # that is fine because all params in the chunk share the same # bucket_id, so DDP's backprop-order iteration still sees # monotonic bucket_ids across the chunk boundary. - for param in sorted(chunk_params, key=lambda p: -p.data.nelement()): + _NUMEL_EPSILON = 0.3 + total_chunk_numel = sum(p.data.nelement() for p in chunk_params) + max_shard_numel = total_chunk_numel / dp_size * (1 + _NUMEL_EPSILON) + for param in sorted(chunk_params, key=lambda p: -_ns_compute_cost(p)): numel = param.data.nelement() - min_shard = min(range(dp_size), key=lambda s: shard_cursors[s]) + candidates = [ + s + for s in range(dp_size) + if pad_param_start(shard_cursors[s]) + numel <= max_shard_numel + ] + if candidates: + min_shard = min(candidates, key=lambda s: shard_compute_loads[s]) + else: + min_shard = min(range(dp_size), key=lambda s: shard_cursors[s]) placement = pad_param_start(shard_cursors[min_shard]) shard_assignments[min_shard].append((param, numel)) shard_cursors[min_shard] = placement + numel + shard_compute_loads[min_shard] += _ns_compute_cost(param) padded_shard_size = pad_to_divisor(max(shard_cursors), shard_divisor) bucket_start_index = buffer_cursor @@ -325,8 +366,18 @@ def _emit_bucket( chunk_params: List[torch.nn.Parameter] = [] chunk_numel = 0 chunk_max_param = 0 - # Mirror _emit_bucket's greedy LPT placement incrementally so we can - # decide, per param, whether it still fits in the current bucket. + # Approximate _emit_bucket's placement so we can decide, per param, + # whether it still fits in the current bucket. This estimates rather + # than mirrors, for two reasons: _emit_bucket sorts the chunk before + # packing it, while this places params in backprop order, and + # _emit_bucket assigns by Newton-Schulz compute cost, while this tracks + # numel. Equal-sized params make both differences vanish, because + # sorting is then a no-op and cost is proportional to numel. Mixed + # sizes send params to different shards under the two orders, so the + # real maximum shard load can exceed the estimated one. _absorbs then + # admits a param that does grow the bucket, leaving a buffer larger + # than closing the bucket early would have produced. The layout stays + # valid either way; see test_mixed_sizes_can_absorb_into_larger_bucket. shard_loads = [0] * dp_size def _absorbs(numel: int) -> bool: diff --git a/tests/unit_tests/distributed/test_layer_wise_param_layout.py b/tests/unit_tests/distributed/test_layer_wise_param_layout.py index a73bf4789a5..2d68a96d7ad 100644 --- a/tests/unit_tests/distributed/test_layer_wise_param_layout.py +++ b/tests/unit_tests/distributed/test_layer_wise_param_layout.py @@ -106,6 +106,15 @@ def _assert_param_within_shard(layout, param, dp_size): class TestSizeMatchingLayout: + """Packing and bucketing rules, exercised with 1-D params. + + Every param here is 1-D, so ``_ns_compute_cost`` falls back to ``nelement()`` + and Newton-Schulz cost equals numel. That keeps these cases focused on the + packing and bucketing rules, but it also means they cannot tell + compute-balanced placement apart from numel-balanced placement: both order + and assign identically when the two metrics agree. ``TestComputeBalancedLayout`` + covers that distinction with 2-D GTP-sharded params, where the metrics diverge. + """ # -- uniform params: all same size, dp_size divides count -- @@ -318,6 +327,32 @@ def test_bucket_has_no_padding_when_params_pack_evenly(self): total_buffer_numel = layout.bucket_indices[-1][1] assert total_buffer_numel == 8 * numel + def test_mixed_sizes_can_absorb_into_larger_bucket(self): + """Absorbing is not always a win: with mixed sizes it can cost space. + + ``_place`` walks params in backprop order while ``_emit_bucket`` sorts the + chunk first, so for mixed sizes the two reach different shard loads. + ``_absorbs`` compares against its own lower estimate and admits the two + 128-element params, after which the sorted packing stacks both onto one + shard and the bucket grows. Closing at the threshold instead would have + emitted 768 + 512 = 1280 elements; absorbing emits 1024 + 384 = 1408. + + Documented rather than fixed: the target case is equal-sized expert + matrices, where sorting is a no-op and the estimate is exact. + """ + dp_size = 2 + # Backprop order is reversed(params), so this list is written back to front. + backprop_order_numels = [192, 192, 256, 128, 128, 192] + params = [_make_param((n,)) for n in reversed(backprop_order_numels)] + cfg = _make_ddp_config() + + layout = _LWO._compute_per_buffer_param_layout(params, 448, dp_size, cfg) + + for param in params: + _assert_param_within_shard(layout, param, dp_size) + assert len(layout.bucket_indices) == 2 + assert layout.bucket_indices[-1][1] == 1408 + # -- bucket alignment -- def test_bucket_dp_divisible(self): @@ -427,3 +462,102 @@ def test_expert_parallel_separate_buffer(self): cfg = _make_ddp_config() layout = _LWO.compute_full_param_layout([dense, expert], None, dp_size, cfg) assert len(layout.layouts) == 2 + + +# --------------------------------------------------------------------------- +# Tests for compute-balanced LPT (Newton-Schulz cost, not numel) +# --------------------------------------------------------------------------- + + +class TestComputeBalancedLayout: + """Placement keys on Newton-Schulz cost, so GTP-sharded params spread out. + + A GTP-sharded param has the numel of its local shard but Newton-Schulz runs on + the full all-gathered matrix, so its cost is far higher than numel suggests. + """ + + def _ns_cost(self, param): + rows, cols = param.data.shape + rows *= getattr(param, 'gtp_remat_size', 1) + big, small = max(rows, cols), min(rows, cols) + return big * small * small + + def _shard_compute_loads(self, layout, params, dp_size): + loads = [0] * dp_size + for param in params: + loads[_get_shard_for_param(layout, param, dp_size)] += self._ns_cost(param) + return loads + + def test_gtp_params_balanced_by_compute_not_numel(self): + """GTP-sharded params dominate cost while having the smallest numel. + + Sorting by numel puts the three cheap-but-large params first and leaves the + expensive GTP ones to fill in, which piles them onto shards that are already + loaded. Sorting by compute cost spreads them instead. + """ + dp_size = 4 + gtp = [ + _make_param((64, 256), is_gtp_weight_remat=True, gtp_remat_size=64) for _ in range(3) + ] + dense = [_make_param((128, 1024)), _make_param((256, 256)), _make_param((256, 256))] + tail = [_make_param((64, 256))] + params = dense + gtp + tail + cfg = _make_ddp_config() + + layout = _LWO._compute_per_buffer_param_layout(params, None, dp_size, cfg) + + # Each GTP param costs 268M against 16.8M for the largest dense param, so no + # two of the three may share a shard. + gtp_shards = [_get_shard_for_param(layout, param, dp_size) for param in gtp] + assert len(set(gtp_shards)) == len(gtp), f"GTP params clustered onto {gtp_shards}" + + loads = self._shard_compute_loads(layout, params, dp_size) + imbalance = max(loads) / (sum(loads) / dp_size) + # Numel-ordered placement gives 3.77x on this input. + assert imbalance < 1.5, f"compute imbalance {imbalance:.2f}x too high: {loads}" + + def test_gtp_params_land_on_different_shards_in_each_bucket(self): + """Expensive params spread across buckets because compute loads do not reset. + + Four buckets, each holding one GTP-sharded matrix plus three dense params of + identical numel. ``shard_cursors`` resets per bucket, so numel gives every + bucket the same starting state and would send all four GTP matrices to shard + 0. ``shard_compute_loads`` carries over, so each bucket sees the previous + ones' cost and picks a different shard. + """ + dp_size = 4 + buckets = 4 + params = [] + for _ in range(buckets): + params.append(_make_param((64, 256), is_gtp_weight_remat=True, gtp_remat_size=64)) + params.extend(_make_param((128, 128)) for _ in range(3)) + gtp_params = [param for param in params if hasattr(param, 'gtp_remat_size')] + cfg = _make_ddp_config() + + # Each group of four params is 4 * 16384 elements, so this cuts one bucket per group. + layout = _LWO._compute_per_buffer_param_layout(params, 4 * 16384, dp_size, cfg) + + assert len(layout.bucket_indices) == buckets + gtp_shards = [_get_shard_for_param(layout, param, dp_size) for param in gtp_params] + assert len(set(gtp_shards)) == buckets, f"GTP params clustered onto {gtp_shards}" + + def test_param_larger_than_cap_is_still_placed(self): + """A param larger than the per-bucket cap still gets placed. + + The cap is ``total_chunk_numel / dp_size * 1.3``. A param above it disqualifies + every shard, so assignment falls back to the previous least-numel rule instead + of wedging. + """ + dp_size = 2 + oversized = _make_param((1024,)) + small_params = [_make_param((64,)) for _ in range(2)] + cfg = _make_ddp_config() + + # cap = (1024 + 64 + 64) / 2 * 1.3 = 748.8, below the oversized param's numel. + layout = _LWO._compute_per_buffer_param_layout( + small_params + [oversized], None, dp_size, cfg + ) + + for param in small_params + [oversized]: + _assert_param_within_shard(layout, param, dp_size) + assert oversized in layout.param_index_map From fad6111eea2d0ecce70597eb8d7795ffb4facfe2 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Wed, 12 Aug 2026 10:46:05 -0700 Subject: [PATCH 006/101] Suppress noisy external-library log lines on non-rank-0 processes (#5590) Signed-off-by: Deepak Narayanan --- megatron/training/initialize.py | 10 ++++++++++ pretrain_gpt.py | 13 +++++++++++++ pretrain_hybrid.py | 13 +++++++++++++ 3 files changed, 36 insertions(+) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 87d6aa65b03..c3eb2cccb2e 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -580,3 +580,13 @@ def setup_logging() -> None: if is_rank0(): logger.info(f'Setting logging level to {logging_level}') logging.getLogger().setLevel(logging_level) + + if not is_rank0(): + for noisy_logger_name in [ + 'GroupedGemmQuantSm100', + 'GroupedGemmDsreluSm100', + 'GroupedGemmSreluSm100', + 'GroupedGemmWgradSm100', + 'absl', + ]: + logging.getLogger(noisy_logger_name).setLevel(logging.ERROR) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 47c1935eb90..8b979d16d4a 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -17,6 +17,19 @@ if rank != 0: warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", category=DeprecationWarning) + + # Some libraries (e.g., CUTLASS DSL) use warnings.catch_warnings() with + # simplefilter("always"), which overrides the filters above. Override + # showwarning as a fallback to suppress warnings that slip through. + _original_showwarning = warnings.showwarning + + def _rank0_only_showwarning(message, category, filename, lineno, file=None, line=None): + if issubclass(category, (UserWarning, FutureWarning, DeprecationWarning)): + return + _original_showwarning(message, category, filename, lineno, file, line) + + warnings.showwarning = _rank0_only_showwarning from functools import lru_cache, partial from typing import Any, List, Optional, Tuple diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 39bc7f30b57..cf665b36887 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -16,6 +16,19 @@ if rank != 0: warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", category=DeprecationWarning) + + # Some libraries (e.g., CUTLASS DSL) use warnings.catch_warnings() with + # simplefilter("always"), which overrides the filters above. Override + # showwarning as a fallback to suppress warnings that slip through. + _original_showwarning = warnings.showwarning + + def _rank0_only_showwarning(message, category, filename, lineno, file=None, line=None): + if issubclass(category, (UserWarning, FutureWarning, DeprecationWarning)): + return + _original_showwarning(message, category, filename, lineno, file, line) + + warnings.showwarning = _rank0_only_showwarning from functools import lru_cache, partial from typing import Any, List, Optional, Tuple From b2f2888e00273d1ae725cf3ec6c92276e2a8e231 Mon Sep 17 00:00:00 2001 From: Ajay Date: Wed, 12 Aug 2026 20:04:10 +0000 Subject: [PATCH 007/101] feat(ci): enhance triage script to assign on-call assignee dynamically (#6469) Signed-off-by: Ajay Balasa --- .gitlab/stages/06.triage.yml | 16 ++- docker/Dockerfile.linting | 2 +- .../python_scripts/resolve_oncall_assignee.py | 117 ++++++++++++++++++ tests/test_utils/test_ci_triage.py | 10 +- 4 files changed, 138 insertions(+), 7 deletions(-) create mode 100644 tests/test_utils/python_scripts/resolve_oncall_assignee.py diff --git a/.gitlab/stages/06.triage.yml b/.gitlab/stages/06.triage.yml index 91f8a4e703c..bac492bc873 100644 --- a/.gitlab/stages/06.triage.yml +++ b/.gitlab/stages/06.triage.yml @@ -49,11 +49,17 @@ triage:linear_write: artifacts: true allow_failure: true script: - - >- - cerno-linear write - --config "${CERNO_CONFIG}" - --plan linear_action_plan.json - --output linear_action_plan_post.json + # The shared GitHub email lookup reads the existing GH_TOKEN CI variable. + - | + ONCALL_ASSIGNEE="$( + python tests/test_utils/python_scripts/resolve_oncall_assignee.py \ + --schedule-file .github/oncall_schedule.json + )" + cerno-linear write \ + --config "${CERNO_CONFIG}" \ + --plan linear_action_plan.json \ + --output linear_action_plan_post.json \ + --assignee "${ONCALL_ASSIGNEE}" artifacts: when: always paths: diff --git a/docker/Dockerfile.linting b/docker/Dockerfile.linting index e9a4de80455..ef3a5813e6a 100644 --- a/docker/Dockerfile.linting +++ b/docker/Dockerfile.linting @@ -24,7 +24,7 @@ RUN --mount=type=secret,id=JET_INDEX_URLS \ # Keep this in the internal-only stage so public CI has no internal service dependency. ARG CI_SERVER_URL -ARG CERNO_COMMIT=5a5fb5360e67f8f09d189871bbc0d768c09c43fa +ARG CERNO_COMMIT=3219afcbfcdbddd1d55e013d94e54a15986e2b1f RUN --mount=type=secret,id=CERNO_TOKEN \ GIT_CONFIG_COUNT=1 \ GIT_CONFIG_KEY_0=http.extraHeader \ diff --git a/tests/test_utils/python_scripts/resolve_oncall_assignee.py b/tests/test_utils/python_scripts/resolve_oncall_assignee.py new file mode 100644 index 00000000000..8cebeb09b7a --- /dev/null +++ b/tests/test_utils/python_scripts/resolve_oncall_assignee.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Resolve the scheduled Megatron-LM on-call user to an NVIDIA email address.""" + +import argparse +import json +import sys +from contextlib import redirect_stdout +from io import StringIO +from pathlib import Path + +_GITHUB_SCRIPTS_DIR = Path(__file__).resolve().parents[3] / ".github" / "scripts" +sys.path.insert(0, str(_GITHUB_SCRIPTS_DIR)) + +from github_slack_utils import get_user_email # noqa: E402 + + +class AssigneeResolutionError(RuntimeError): + """Raised when the current on-call Linear assignee cannot be resolved safely.""" + + +def load_schedule(path: Path) -> list[dict[str, object]]: + """Load and validate the top-level on-call schedule structure.""" + + try: + schedule = json.loads(path.read_text()) + except OSError as exc: + raise AssigneeResolutionError(f"could not read on-call schedule {path}: {exc}") from exc + except json.JSONDecodeError as exc: + raise AssigneeResolutionError(f"on-call schedule {path} is not valid JSON") from exc + + if not isinstance(schedule, list) or not schedule: + raise AssigneeResolutionError(f"on-call schedule {path} must be a non-empty list") + return schedule + + +def current_oncall(schedule: list[dict[str, object]]) -> str: + """Return the current on-call user, matching the existing rotation manager.""" + + if not schedule: + raise AssigneeResolutionError("on-call schedule must be a non-empty list") + entry = schedule[0] + if not isinstance(entry, dict): + raise AssigneeResolutionError("the current on-call schedule entry must be an object") + + username = entry.get("user") + if not isinstance(username, str) or not username.strip(): + raise AssigneeResolutionError("the current on-call schedule entry must have a user") + return username.strip() + + +def resolve_nvidia_email(github_login: str) -> str: + """Resolve a GitHub login with the repository's shared email lookup helper.""" + + # The shared helper reports lookup details on stdout. Suppress those messages + # so command substitution receives only the email printed by ``main`` and CI + # logs do not repeat email addresses from fallback diagnostics. + try: + with redirect_stdout(StringIO()): + email = get_user_email(github_login) + except SystemExit as exc: + raise AssigneeResolutionError( + f"could not look up an email for GitHub user {github_login!r}" + ) from exc + + if not isinstance(email, str): + raise AssigneeResolutionError( + f"GitHub user {github_login!r} did not resolve to an NVIDIA email" + ) + email = email.strip() + local_part, separator, domain = email.rpartition("@") + if ( + separator != "@" + or not local_part + or "@" in local_part + or domain.casefold() != "nvidia.com" + or any(character.isspace() for character in email) + ): + raise AssigneeResolutionError( + f"GitHub user {github_login!r} did not resolve to a valid NVIDIA email" + ) + return email + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse command-line arguments.""" + + parser = argparse.ArgumentParser( + description="Resolve the scheduled Megatron-LM on-call user's NVIDIA email" + ) + parser.add_argument( + "--schedule-file", + type=Path, + default=Path(".github/oncall_schedule.json"), + help="path to the dated GitHub on-call schedule", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + """Print only the current on-call user's NVIDIA email to stdout.""" + + args = parse_args(argv) + try: + username = current_oncall(load_schedule(args.schedule_file)) + email = resolve_nvidia_email(username) + except AssigneeResolutionError as exc: + print(f"Error: {exc}", file=sys.stderr) + return 1 + + print(f"Resolved scheduled on-call GitHub user {username!r}.", file=sys.stderr) + print(email) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_utils/test_ci_triage.py b/tests/test_utils/test_ci_triage.py index 626013c73ef..3c3b3b7ed08 100644 --- a/tests/test_utils/test_ci_triage.py +++ b/tests/test_utils/test_ci_triage.py @@ -129,13 +129,21 @@ def test_cerno_hard_cutover_contract(): assert triage.count("cerno-linear") == 3 assert triage.count("cerno-notify") == 2 assert triage.count('--config "${CERNO_CONFIG}"') == 3 - assert "ARG CERNO_COMMIT=5a5fb5360e67f8f09d189871bbc0d768c09c43fa" in dockerfile assert '"cerno @ git+${CI_SERVER_URL}/dl/nemo/cerno.git@${CERNO_COMMIT}"' in dockerfile assert "id=CERNO_TOKEN" in dockerfile assert "/run/secrets/CERNO_TOKEN" in dockerfile assert "--secret id=CERNO_TOKEN,env=PAT" in build_script +def test_linear_write_assigns_new_issues_to_scheduled_oncall(): + triage = yaml.safe_load(Path(".gitlab/stages/06.triage.yml").read_text()) + script = "\n".join(triage["triage:linear_write"]["script"]) + + assert "resolve_oncall_assignee.py" in script + assert "--schedule-file .github/oncall_schedule.json" in script + assert '--assignee "${ONCALL_ASSIGNEE}"' in script + + def test_notification_rules_use_expected_pipeline_sources(): unit = yaml.safe_load(Path(".gitlab/stages/02.test.yml").read_text()) functional = yaml.safe_load(Path(".gitlab/stages/04.functional-tests.yml").read_text()) From 543628902932815fc3ea29e74e7a526c03ce99d6 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 13 Aug 2026 01:20:34 +0000 Subject: [PATCH 008/101] Skip attention mask materialization in RL training (#5885) Signed-off-by: Teodor-Dumitru Ene --- megatron/rl/rl_utils.py | 7 ++-- megatron/training/utils/common_utils.py | 43 +++++++++++++++++++------ train_rl.py | 4 +++ 3 files changed, 43 insertions(+), 11 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 262c366110f..e43e2f72ced 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -1708,16 +1708,19 @@ def prepare_data_for_update( data_loader = DataLoader(dataset, batch_size=1) logprobs_batch_size = 1 else: - # Always compute standard masks for the original data (we'll need them later) + # Compute the loss mask and position ids for the original data (we'll need them later). + # No dense attention mask: the forward pass masks via PackedSeqParams (see + # get_logprobs), even when sequence packing is disabled. with nvtx_range("rl/get-ltor-masks", time=True): _, original_loss_mask, original_position_ids = get_ltor_masks_and_position_ids( trajs, tokenizer.eod, tokenizer.pad, args.reset_position_ids, - args.reset_attention_mask, + reset_attention_mask=False, eod_mask_loss=False, pad_mask_loss=True, + create_attention_mask=False, ) original_loss_mask[~generation_masks] = 0.0 compute_trajs = trajs diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index 30617ef9b4c..1b20eb5a985 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -377,20 +377,44 @@ def get_ltor_masks_and_position_ids(data, reset_position_ids, reset_attention_mask, eod_mask_loss, - pad_mask_loss): - """Build masks and position id for left to right model.""" + pad_mask_loss, + create_attention_mask=True): + """Build masks and position id for left to right model. + + Args: + data: Token ids, shape [micro_batch_size, seq_length]. + eod_token: End-of-document token id. + pad_token: Padding token id. + reset_position_ids: Restart position ids from 0 after each EOD token. + reset_attention_mask: Additionally mask attention across document boundaries, + turning the shared causal mask into a per-sample block-causal mask. + Requires create_attention_mask, since it modifies the materialized mask. + eod_mask_loss: Zero the loss mask at EOD tokens. + pad_mask_loss: Zero the loss mask at pad tokens. + create_attention_mask: Materialize the dense causal attention mask. + Can be disabled if the attention kernel generates the mask by itself + (e.g. from PackedSeqParams), in which case attention_mask is returned as None. + + Returns: + Tuple of (attention_mask or None, loss_mask, position_ids). + """ + assert create_attention_mask or not reset_attention_mask, \ + "reset_attention_mask requires the attention mask to be created." # Extract batch size and sequence length. micro_batch_size, seq_length = data.size() # Attention mask (lower triangular). - if reset_attention_mask: - att_mask_batch = micro_batch_size + if create_attention_mask: + if reset_attention_mask: + att_mask_batch = micro_batch_size + else: + att_mask_batch = 1 + attention_mask = torch.tril( + torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) + ).view(att_mask_batch, 1, seq_length, seq_length) else: - att_mask_batch = 1 - attention_mask = torch.tril( - torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) - ).view(att_mask_batch, 1, seq_length, seq_length) + attention_mask = None # Loss mask. loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device) @@ -429,7 +453,8 @@ def get_ltor_masks_and_position_ids(data, prev_index = i + 1 # Convert attention mask to binary: - attention_mask = attention_mask < 0.5 + if create_attention_mask: + attention_mask = attention_mask < 0.5 return attention_mask, loss_mask, position_ids diff --git a/train_rl.py b/train_rl.py index acf54680f4a..17eec357b15 100644 --- a/train_rl.py +++ b/train_rl.py @@ -415,6 +415,10 @@ def _model_builder( extra_args_provider=add_inference_args, args_defaults={}, ) + assert not args.reset_attention_mask, ( + "--reset-attention-mask is not supported in RL training: " + "the forward pass masks via PackedSeqParams and never consumes a dense attention mask." + ) if is_hybrid_model(args): model_cfg = hybrid_config_from_args(args) else: From a34880f0adbb8730709675c49898944014771dc7 Mon Sep 17 00:00:00 2001 From: Jiangfei Duan <30710061+JF-D@users.noreply.github.com> Date: Thu, 13 Aug 2026 03:07:12 +0000 Subject: [PATCH 009/101] Avoid unnecessary MoE router host synchronization (#6432) Signed-off-by: Jiangfei Duan --- .../core/transformer/transformer_layer.py | 21 +++-- .../transformer/test_transformer_layer.py | 94 ++++++++++++++++++- 2 files changed, 105 insertions(+), 10 deletions(-) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 5c55f2abe6c..0f7be974418 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1560,8 +1560,6 @@ def transition_cudagraph_scope(self, mode): and "moe" in self.config.recompute_modules and self.config.cuda_graph_impl == "local" ) - if not hasattr(self, '_router_dtoh_event'): - self._router_dtoh_event = torch.cuda.Event() if not hasattr(self, 'cudagraph_manager_router'): self.cudagraph_manager_router = CudaGraphManager( self.config, self, function_name="_forward_mlp_router" @@ -1630,6 +1628,16 @@ def _get_token_dispatcher_attrs(self): return tuple(attr_names), token_dispatcher_attr_outputs + def _synchronize_router_host_outputs(self, attr_outputs): + """Wait for partial-router graph outputs only when they reside on the host.""" + if not any(attr.device.type == "cpu" for attr in attr_outputs): + return + + if not hasattr(self, '_router_dtoh_event'): + self._router_dtoh_event = torch.cuda.Event() + self._router_dtoh_event.record() + self._router_dtoh_event.synchronize() + def _forward_mlp_router(self, hidden_states, padding_mask=None): """ Executes the router phase of the MoE block. @@ -1735,12 +1743,9 @@ def _forward_mlp_partial_cudagraphs( *token_dispatcher_attr_outputs, ) = router_outputs - # After the router graph replays, the captured .copy_() operations that update - # the returned dispatcher tensors via `_maybe_dtoh_and_synchronize` are queued on - # the current stream but may not have completed. Record an event after the router - # graph and wait on it, so we block only until the router's D2H copies complete. - self._router_dtoh_event.record() - self._router_dtoh_event.synchronize() + # CUDA outputs remain ordered by the graph-completion event. Only host outputs need + # a CPU-blocking wait before the eager dispatcher can consume them. + self._synchronize_router_host_outputs(token_dispatcher_attr_outputs) expert_output, mlp_bias = self._forward_mlp_expert_compute( hidden_states, probs, token_dispatcher_attr_outputs diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 93650cf13b0..9f6a20ae453 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -2,28 +2,37 @@ import gc +from unittest.mock import Mock, patch import pytest import torch from megatron.core import parallel_state from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedTensor +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.utils import InferenceMode from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_with_transformer_engine_submodules, ) +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec from megatron.core.tensor_parallel.random import ( HAVE_TE, initialize_rng_tracker, model_parallel_cuda_manual_seed, ) -from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord -from megatron.core.transformer.enums import InferenceCudaGraphScope +from megatron.core.transformer.cuda_graphs import ( + CudaGraphManager, + _CudagraphGlobalRecord, + create_cudagraphs, +) +from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + MoETransformerLayer, TransformerLayer, + TransformerLayerSubmodules, get_transformer_layer_offset, ) from megatron.core.utils import is_te_min_version @@ -378,6 +387,52 @@ def _no_layers_have_manager(block) -> bool: return all(not hasattr(layer, 'cudagraph_manager') for layer in block.layers) +def _make_moe_transformer_layer(*, partial_cudagraph: bool): + config = TransformerConfig( + num_layers=1, + hidden_size=32, + num_attention_heads=4, + ffn_hidden_size=64, + moe_ffn_hidden_size=64, + num_moe_experts=4, + moe_router_topk=2, + moe_router_load_balancing_type="none", + moe_token_dispatcher_type="allgather", + hidden_dropout=0.0, + attention_dropout=0.0, + bias_dropout_fusion=False, + add_bias_linear=False, + use_cpu_initialization=True, + cuda_graph_impl="local" if partial_cudagraph else "none", + cuda_graph_modules=[CudaGraphModule.moe_router] if partial_cudagraph else [], + ) + submodules = TransformerLayerSubmodules( + mlp=get_moe_module_spec(use_te=False, num_experts=4, moe_grouped_gemm=False), + mlp_bda=get_bias_dropout_add, + ) + return MoETransformerLayer(config, submodules) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_moe_router_synchronizes_host_outputs_and_reuses_event(): + layer = object.__new__(MoETransformerLayer) + host_output = Mock() + host_output.device.type = "cpu" + cuda_output = torch.empty(1, device="cuda") + event = Mock() + + with patch.object(torch.cuda, "Event", return_value=event) as event_factory: + layer._synchronize_router_host_outputs((host_output,)) + layer._synchronize_router_host_outputs((cuda_output,)) + assert event.record.call_count == 1 + assert event.synchronize.call_count == 1 + layer._synchronize_router_host_outputs((host_output,)) + + event_factory.assert_called_once_with() + assert event.record.call_count == 2 + assert event.synchronize.call_count == 2 + + @pytest.mark.skipif( not (HAVE_TE and is_te_min_version("1.5.0")), reason="CUDA graph tests require TransformerEngine >= 1.5", @@ -393,6 +448,41 @@ def teardown_method(self, method): _reset_cudagraph_state() gc.collect() + def test_moe_router_partial_cudagraph_forward_matches_eager(self): + eager_layer = _make_moe_transformer_layer(partial_cudagraph=False) + partial_cg_layer = _make_moe_transformer_layer(partial_cudagraph=True) + partial_cg_layer.load_state_dict(eager_layer.state_dict()) + eager_layer.cuda() + partial_cg_layer.cuda() + for param in partial_cg_layer.parameters(): + param.main_grad = torch.zeros_like(param) + + hidden_states = torch.randn(8, 2, 32, device="cuda", requires_grad=True) + eager_output, _ = eager_layer(hidden_states.clone(), attention_mask=None) + eager_output = eager_output.detach().clone() + + # The first forward/backward records the real router and postprocess graph boundaries. + recorded_output, _ = partial_cg_layer(hidden_states.clone(), attention_mask=None) + recorded_output.sum().backward() + create_cudagraphs() + + assert _CudagraphGlobalRecord.cudagraph_created + assert partial_cg_layer.use_partial_cudagraphs + for manager in ( + partial_cg_layer.cudagraph_manager_router, + partial_cg_layer.cudagraph_manager_postprocess, + ): + assert len(manager.cudagraph_runners) == 1 + assert manager.cudagraph_runners[0].fwd_graph is not None + + partial_cg_layer.zero_grad(set_to_none=True) + partial_cg_output, _ = partial_cg_layer(hidden_states.clone(), attention_mask=None) + partial_cg_output = partial_cg_output.detach().clone() + + # All-gather routing metadata stays on CUDA, so replay must not create a host-wait event. + assert not hasattr(partial_cg_layer, '_router_dtoh_event') + torch.testing.assert_close(partial_cg_output, eager_output, rtol=0, atol=0) + def test_empty_scope_transformer_layer_has_per_layer_manager(self): block = _make_cuda_graph_gpt_block( cuda_graph_impl='local', cuda_graph_modules=[], inference_cuda_graph_scope='layer' From 3cd8cf30a20b1eea4ae29ae201d0c021738680d9 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 13 Aug 2026 03:22:27 +0000 Subject: [PATCH 010/101] Refactor GDP dynamic inference onto the shared SSM interface (#6443) Signed-off-by: Deepak Narayanan Signed-off-by: Keshav Santhanam Signed-off-by: Mikail Khona Signed-off-by: Mikail Khona (NVIDIA) Signed-off-by: Kezhi Kong Co-authored-by: Kezhi Kong Co-authored-by: Deepak Narayanan Co-authored-by: Claude Opus 4.6 Co-authored-by: Mikail Khona (NVIDIA) Co-authored-by: Mikail Khona Co-authored-by: root Co-authored-by: Roger Waleffe Co-authored-by: Kezhi Kong --- megatron/core/ssm/gated_delta_product.py | 314 +++++++++-------------- 1 file changed, 122 insertions(+), 192 deletions(-) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 62f86e842d3..1baaea0ea3a 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -17,9 +17,7 @@ from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory from megatron.core.inference.contexts import BaseInferenceContext, DynamicInferenceContext from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( - tensor_get_slice_after, tensor_masked_update, - tensor_merge, ) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gdp_context_parallel import GDPContextParallel @@ -28,6 +26,7 @@ check_fla_sequence_packing_support, get_cu_seqlens, ) +from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer import TransformerConfig from megatron.core.transformer.module import MegatronModule @@ -36,7 +35,7 @@ make_sharded_tensors_for_checkpoint, sharded_state_dict_default, ) -from megatron.core.utils import deprecate_inference_params, is_using_quantization_scales +from megatron.core.utils import deprecate_inference_params try: from causal_conv1d import causal_conv1d_fn, causal_conv1d_update @@ -98,7 +97,7 @@ class GatedDeltaProductMixerSubmodules: out_proj: Union[ModuleSpec, type] = None -class GatedDeltaProductMixer(MegatronModule): +class GatedDeltaProductMixer(SSMDynamicInferenceMixin, MegatronModule): """Gated Delta Product (GDP) sequence mixer for hybrid models. The mixer accepts hidden states with shape ``[sequence, batch, hidden]`` and returns @@ -381,7 +380,12 @@ def forward( conv_state, ssm_state = None, None if inference_context is not None: if inference_context.is_dynamic_batching(): - return self._dynamic_inference(hidden_states, inference_context) + ok, reason = check_fla_sequence_packing_support() + assert ok, reason + assert ( + self.cp.cp_size == 1 + ), "Context parallel is not supported for GDP dynamic inference" + return self.ssm_dynamic_inference(hidden_states, inference_context) assert ( inference_context.is_static_batching() ), "GDP inference must be either static or dynamic batching." @@ -391,6 +395,9 @@ def forward( "Packing is only wired through the training/prefill (chunk) path." ) conv_state, ssm_state = self._get_states_from_cache(inference_context, batch_size) + if inference_context.seqlen_offset > 0: + # The states are updated in place. + return self._static_decode(hidden_states, conv_state, ssm_state) # Build cu_seqlens for the chunked recurrence (FLA) when running with # packed (THD) sequences on the training/prefill path. @@ -437,32 +444,19 @@ def forward( VKQ = VKQ.contiguous() VKQ = rearrange(VKQ, "b l d -> b d l") - # Decode - if inference_context is not None and inference_context.seqlen_offset > 0: - VKQ = causal_conv1d_update( - VKQ, - conv_state, - rearrange(self.conv1d.weight, "d 1 w -> d w"), - self.conv1d.bias, - self.activation, - ) - else: - # Prefill - if conv_state is not None: - # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv - # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. - conv_state.copy_( - F.pad(VKQ, (self.d_conv - VKQ.shape[-1], 0)) - ) # Update state (B D W) - # Train - # causal_conv1d uses seq_idx_packed to reset the convolution boundaries - VKQ = causal_conv1d_fn( - x=VKQ, - weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), - bias=self.cp.get_conv1d_bias(), - activation=self.activation, - seq_idx=seq_idx_packed, - ) + if conv_state is not None: + # Static-batching prefill: seed the conv state from the prompt's tail. + # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv + # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. + conv_state.copy_(F.pad(VKQ, (self.d_conv - VKQ.shape[-1], 0))) # Update state (B D W) + # causal_conv1d uses seq_idx_packed to reset the convolution boundaries + VKQ = causal_conv1d_fn( + x=VKQ, + weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), + bias=self.cp.get_conv1d_bias(), + activation=self.activation, + seq_idx=seq_idx_packed, + ) VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() @@ -506,46 +500,18 @@ def forward( self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=2 ) - # Decode - if inference_context is not None and inference_context.seqlen_offset > 0: - - g_new = g.new_zeros(g.shape[0], g.shape[1], self.num_householder, g.shape[2]) - g_new[:, :, 0] = g - g = rearrange(g_new, '... t n h -> ... (t n) h') - - query_new = query.new_zeros( - query.shape[0], query.shape[1], self.num_householder, query.shape[2], query.shape[3] - ) - query_new[:, :, -1] = query - query = rearrange(query_new, '... t n h d-> ... (t n) h d') - - core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=ssm_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - ) - core_attn_out = rearrange( - core_attn_out, '... (t n) h d -> ... t n h d', n=self.num_householder - )[..., -1, :, :].contiguous() - # Train or Prefill - else: - core_attn_out, last_recurrent_state = chunk_gated_delta_product( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=(ssm_state is not None), - num_householder=self.num_householder, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens_packed, - ) + core_attn_out, last_recurrent_state = chunk_gated_delta_product( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=(ssm_state is not None), + num_householder=self.num_householder, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens_packed, + ) if ssm_state is not None: ssm_state.copy_(last_recurrent_state) @@ -561,112 +527,70 @@ def forward( return out, out_bias + # ================================================================== + # Static / eager inference + # + # ``_static_decode`` implements legacy static-batching decode. It is + # deliberately kept separate from the dynamic inference hooks below so that + # static-batching bookkeeping does not pollute the interface defined by + # ``SSMDynamicInferenceMixin``. Static-batching prefill shares the training + # body in ``forward``, which seeds the conv/SSM state when the caches are + # present. Mirrors ``MambaMixer._static_decode``. + # ================================================================== + def _static_decode( + self, hidden_states: torch.Tensor, conv_state: torch.Tensor, ssm_state: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Single-token static-batching decode step (updates state in place).""" + assert hidden_states.shape[0] == 1, "Only support decoding with 1 token at a time for now" + assert self.cp.cp_size == 1, "Context parallel not supported for GDP inference decode" + + # (1, b, d_model) -> (1, b, proj_dim) + zVKQba, _ = self.in_proj(hidden_states) + + # The decode kernels are batch-first: (1, b, proj_dim) -> (b, 1, proj_dim). + # Static batching has no slot remapping, so batch_indices is None. + y = self.ssm_decode( + zVKQba.transpose(0, 1), conv_state=conv_state, ssm_state=ssm_state, batch_indices=None + ) + + # (b, 1, d_inner) -> (1, b, d_inner), which is what out_proj expects. + return self.out_proj(y.transpose(0, 1)) + # ------------------------------------------------------------------ # Dynamic-batching inference. # - # Mirrors ``MambaMixer._dynamic_inference`` / ``_ssm_decode`` / ``_ssm_prefill`` - # (same ``_ssm_`` naming and the same request-level control flow), but runs - # the Gated Delta Product kernels instead of the Mamba2 scan. The per-request - # recurrent state (short-conv state + matrix-valued SSM state) is read/written - # through the slot-indexed caches owned by ``DynamicInferenceContext``. + # These are the two hooks required by ``SSMDynamicInferenceMixin``; the + # mixin owns the surrounding decode/prefill partitioning and merge. They run + # the Gated Delta Product kernels instead of the Mamba2 scan. The + # per-request recurrent state (short-conv state + matrix-valued SSM state) + # is read/written through the slot-indexed caches owned by + # ``DynamicInferenceContext``. # # MVP scope: this path does not yet support context parallelism (cp_size > 1), # speculative decoding, chunked prefill, Mamba prefix caching, or CUDA-graph # capture. The reshapes mirror the static ``forward`` math with batch/seq # repurposed for the packed dynamic layout. # ------------------------------------------------------------------ - def _dynamic_inference( - self, hidden_states: torch.Tensor, context: DynamicInferenceContext - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Execute one dynamic inference step by separating decode and prefill - requests, running each through the GDP kernels independently, and merging - the results back into packed token order.""" - ok, reason = check_fla_sequence_packing_support() - assert ok, reason - assert self.cp.cp_size == 1, "Context parallel is not supported for GDP dynamic inference" - assert ( - not context.is_chunked_prefill_enabled() - ), "GDP dynamic inference does not support chunked prefill yet." - - # GDP-style layers register as Mamba layers, so the same (conv_state, - # ssm_state) accessor and per-layer slab layout apply. - conv_state, ssm_state = context.mamba_states_cache(self.layer_number - self.pp_layer_offset) - - padded_dims = context.padded_batch_dimensions - token_count = padded_dims.token_count - decode_req_count = padded_dims.decode_req_count - prefill_req_count = padded_dims.prefill_req_count - metadata = context.mamba_metadata - - # Input projection over the full packed batch. - zVKQba, _ = self.in_proj(hidden_states) - - y_decode = None - y_prefill = None - - # --- Decode partition (placed first in the packed batch) --------- - if decode_req_count > 0: - # MVP: exactly one token per decode request (no speculative tokens). - zVKQba_decode = zVKQba[:decode_req_count] if prefill_req_count > 0 else zVKQba - y_decode = self._ssm_decode( - zVKQba_decode.transpose(0, 1), conv_state, ssm_state, metadata.batch_indices_decode - ).transpose(0, 1) - - # --- Prefill partition ------------------------------------------- - if prefill_req_count > 0: - if decode_req_count > 0: - # Mixed batch: gather the prefill tokens out of the packed tensor. - zVKQba_prefill = torch.empty_like(zVKQba) - tensor_get_slice_after( - zVKQba, zVKQba_prefill, metadata.device_decode_prefill, check_bounds=False - ) - else: - zVKQba_prefill = zVKQba - y_prefill = self._ssm_prefill( - zVKQba_prefill, - conv_state=conv_state, - ssm_state=ssm_state, - seq_idx=metadata.seq_idx, - cu_seqlens=metadata.cu_seqlens, - batch_indices=metadata.batch_indices_prefill, - ) - - # --- Merge back into packed token order -------------------------- - if y_decode is not None and y_prefill is not None: - y = torch.empty( - [token_count, 1, y_prefill.shape[-1]], - dtype=y_prefill.dtype, - device=y_prefill.device, - ) - tensor_merge(y_decode, y_prefill, metadata.device_decode_prefill, output_tensor=y) - elif y_decode is not None: - y = y_decode - elif y_prefill is not None: - y = y_prefill - else: - raise RuntimeError("Dynamic inference called with 0 decode and 0 prefill requests") - - # Zero padding positions to avoid corrupting quantization amax calculations. - if is_using_quantization_scales(self.config): - y[context.padding_slice] = 0.0 - - out, out_bias = self.out_proj(y) - return out, out_bias - - def _ssm_decode( + def ssm_decode( self, zVKQba: torch.Tensor, conv_state: torch.Tensor, ssm_state: torch.Tensor, batch_indices: Optional[torch.Tensor] = None, + intermediate_conv_state: Optional[torch.Tensor] = None, + intermediate_ssm_state: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Single-token-per-request decode. ``zVKQba`` is ``[1, decode_req_count, - proj_dim]``; returns ``[1, decode_req_count, d_inner]``. The conv and SSM - states are read/written in place at the slots named by ``batch_indices`` - (``-1`` marks padding slots).""" - seq_len, _, _ = zVKQba.shape + """Single-token-per-request decode. ``zVKQba`` is ``[n, seq_len, + proj_dim]``; returns ``[n, seq_len, d_inner]``. The conv and SSM states + are read/written in place at the slots named by ``batch_indices`` + (``-1`` marks padding slots); ``batch_indices=None`` means static + batching, where the caches are already in request order.""" + _, seq_len, _ = zVKQba.shape assert seq_len == 1, "GDP decode supports one token per request" - zVKQba = zVKQba.squeeze(0) # [n, proj_dim] + assert ( + intermediate_conv_state is None and intermediate_ssm_state is None + ), "GDP decode does not support speculative decoding yet" + zVKQba = zVKQba.squeeze(1) # [n, proj_dim] M = self.num_householder z, VKQ, ba = torch.split( @@ -729,11 +653,14 @@ def _ssm_decode( query_new[:, :, -1] = query query = rearrange(query_new, "n t m h d -> n (t m) h d") - # Gather this step's per-request initial states. ``.clamp`` (NOT in-place) - # returns a new tensor, so ``batch_indices`` keeps its -1 padding sentinels - # for the scatter below; the padding rows' outputs are never scattered back. - gather_idx = batch_indices.clamp(min=0) - initial_state = ssm_state[gather_idx] + if batch_indices is None: + # Static batching: the cache rows are already in request order. + initial_state = ssm_state + else: + # Gather this step's per-request initial states. ``.clamp`` (NOT in-place) + # returns a new tensor, so ``batch_indices`` keeps its -1 padding sentinels + # for the scatter below; the padding rows' outputs are never scattered back. + initial_state = ssm_state[batch_indices.clamp(min=0)] core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule( query, @@ -749,29 +676,37 @@ def _ssm_decode( ..., -1, :, : ].contiguous() # [n, 1, h, d] - # Scatter updated states back into the cache (skips -1 padding slots). - tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) + if batch_indices is None: + ssm_state.copy_(last_recurrent_state) + else: + # Scatter updated states back into the cache (skips -1 padding slots). + tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) - y = rearrange(core_attn_out, "n t h p -> t n (h p)").contiguous() # [1, n, d_inner] + y = rearrange(core_attn_out, "n t h p -> n t (h p)").contiguous() # [n, 1, d_inner] if self.rmsnorm: - z = rearrange(z, "n t h p -> t n (h p)").contiguous() + z = rearrange(z, "n t h p -> n t (h p)").contiguous() y = self.norm(y, z) return y - def _ssm_prefill( + def ssm_prefill( self, zVKQba: torch.Tensor, - conv_state: Optional[torch.Tensor] = None, - ssm_state: Optional[torch.Tensor] = None, - seq_idx: Optional[torch.Tensor] = None, - cu_seqlens: Optional[torch.Tensor] = None, - batch_indices: Optional[torch.Tensor] = None, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + context: DynamicInferenceContext, ) -> torch.Tensor: """Variable-length prefill over all prefill requests in one varlen call. ``zVKQba`` is ``[l, 1, proj_dim]``; returns ``[l, 1, d_inner]``. Fresh requests start from a zero recurrent state (no prefix caching in the MVP); the resulting final conv/SSM states are written back into the caches.""" - is_dynamic_batching = seq_idx is not None + assert ( + not context.is_chunked_prefill_enabled() + ), "GDP dynamic inference does not support chunked prefill yet." + + metadata = context.mamba_metadata + seq_idx = metadata.seq_idx + cu_seqlens = metadata.cu_seqlens + batch_indices = metadata.batch_indices_prefill M = self.num_householder # l b d -> b l d @@ -788,18 +723,14 @@ def _ssm_prefill( dim=-1, ) - if conv_state is not None and is_dynamic_batching: - assert batch_indices is not None - # Capture per-request final conv states (before the conv consumes the - # inputs) and write them into the prefill requests' cache rows. - conv_varlen_states = causal_conv1d_varlen_states( - VKQ.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] - ) - tensor_masked_update(conv_state, batch_indices, conv_varlen_states) - # Maintain channels-last memory layout so causal_conv1d_fn can use seq_idx. - VKQ = VKQ.transpose(1, 2) - else: - VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() + # Capture per-request final conv states (before the conv consumes the + # inputs) and write them into the prefill requests' cache rows. + conv_varlen_states = causal_conv1d_varlen_states( + VKQ.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] + ) + tensor_masked_update(conv_state, batch_indices, conv_varlen_states) + # Maintain channels-last memory layout so causal_conv1d_fn can use seq_idx. + VKQ = VKQ.transpose(1, 2) seqlen = VKQ.size(2) if causal_conv1d_fn is None: @@ -846,15 +777,14 @@ def _ssm_prefill( g=g, beta=beta, initial_state=None, - output_final_state=ssm_state is not None, + output_final_state=True, num_householder=M, use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, ) # Write per-request final SSM states into the cache for subsequent decode. - if ssm_state is not None and is_dynamic_batching: - tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) + tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) y = rearrange(core_attn_out, "b l h p -> l b (h p)").contiguous() if self.rmsnorm: From 3ce4b4ac0f2165032839e86b1f56de8c1ce1657f Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 13 Aug 2026 03:55:23 +0000 Subject: [PATCH 011/101] test(fsdp): mark default overlap check flaky (#6493) Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 8b8cbcfabbc..1886945f3be 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -594,7 +594,20 @@ def test_root_backward_returns_to_resting_memory(distributed_setup): ) -@pytest.mark.parametrize("use_symmetric_memory", [False, True], ids=["default", "symmetric_memory"]) +@pytest.mark.parametrize( + "use_symmetric_memory", + [ + # Both variants' all-gathers are launch-timing sensitive, but default-CTA + # kernels occupy more compute CTAs, making the profiler overlap count less + # stable across ranks (see + # https://github.com/NVIDIA/Megatron-LM/actions/runs/31615942188). + # The symmetric-memory variant uses zero-CTA all-gather kernels, so its overlap + # measurement is more stable and remains enabled. + pytest.param(False, marks=(pytest.mark.flaky, pytest.mark.flaky_in_dev)), + pytest.param(True), + ], + ids=["default", "symmetric_memory"], +) def test_overlaps_communication_and_compute(distributed_setup, use_symmetric_memory): """Forward and backward communication should overlap GEMM compute.""" world_size = distributed_setup.world_size From 301e0d658f23c34e8455d2a877d2ff48e7f81cd5 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Thu, 13 Aug 2026 05:35:43 +0000 Subject: [PATCH 012/101] Paged stashing no longer assumes single layer config (#6419) Signed-off-by: Philip Petrakian --- megatron/core/transformer/moe/paged_stash.py | 29 ++- .../moe/test_paged_stash_runner.py | 206 ++++++++++++++++++ 2 files changed, 227 insertions(+), 8 deletions(-) create mode 100644 tests/unit_tests/transformer/moe/test_paged_stash_runner.py diff --git a/megatron/core/transformer/moe/paged_stash.py b/megatron/core/transformer/moe/paged_stash.py index c38aac11bd2..450aa857d8c 100644 --- a/megatron/core/transformer/moe/paged_stash.py +++ b/megatron/core/transformer/moe/paged_stash.py @@ -986,15 +986,15 @@ def __init__(self, config, copy_main_params, model, optimizer, forward_backward_ # them (NCCL EP only). Both set by check_moe_overflow. self._required_recv_capacity = None self._required_capacity_factor = None - # TransformerConfig objects that must stay in sync for moe_paged_stash: the training - # loop `config` (schedules / paged_stash_reset) plus each VP chunk's GPT root config - # (GPTModel.forward). MoE mlps use the same config reference as that root, so we do - # not track mlp.config separately. + # Config objects that must stay in sync for moe_paged_stash: the training loop config + # (schedules / paged_stash_reset), each model chunk's root config (model forward), and + # every MoE layer config (expert forward). Some models may use a distinct config for + # each layer. seen_cfg_ids = set() self._configs_to_sync_moe_paged_stash = [] def _track_cfg(c): - if c is None: + if c is None or not hasattr(c, 'moe_paged_stash'): return cid = id(c) if cid not in seen_cfg_ids: @@ -1012,6 +1012,16 @@ def _track_cfg(c): model_chunk, "decoder", allow_none=False, return_model_obj=True ) _track_cfg(model_with_decoder.config) + + # Track MoE configs independently from the existing structural discovery below. + # This keeps overflow and retry behavior unchanged for models whose modules share + # the root config while allowing distinct module configs to stay synchronized. + for module in model_with_decoder.modules(): + token_dispatcher = getattr(module, 'token_dispatcher', None) + if token_dispatcher is None or not hasattr(token_dispatcher, 'check_over_budget'): + continue + _track_cfg(getattr(module, 'config', None)) + for layer in model_with_decoder.decoder.layers: transformer_layer = ( layer.mtp_model_layer if isinstance(layer, MultiTokenPredictionLayer) else layer @@ -1039,7 +1049,7 @@ def _track_cfg(c): self.moe_layers.append(mlp) def _set_moe_paged_stash_all(self, value: bool) -> None: - """Set moe_paged_stash on every tracked config (train + per VP chunk root).""" + """Set moe_paged_stash on every tracked training, model, and MoE config.""" for c in self._configs_to_sync_moe_paged_stash: c.moe_paged_stash = value @@ -1234,7 +1244,9 @@ def __call__(self, *args, **kwargs): training = not kwargs['forward_only'] data_iterator = kwargs['data_iterator'] - saved_moe_paged_stash = self.config.moe_paged_stash + saved_moe_paged_stash_values = [ + (config, config.moe_paged_stash) for config in self._configs_to_sync_moe_paged_stash + ] num_tries = 0 while True: num_tries += 1 @@ -1271,7 +1283,8 @@ def __call__(self, *args, **kwargs): for mlp in self.moe_layers: mlp.token_dispatcher.invalidate_ep_bootstrap() nccl_ep_release_context() - self._set_moe_paged_stash_all(saved_moe_paged_stash) + for config, value in saved_moe_paged_stash_values: + config.moe_paged_stash = value break # Overflow or over-budget: prepare_for_rerun clears capacity factor and paged stash. diff --git a/tests/unit_tests/transformer/moe/test_paged_stash_runner.py b/tests/unit_tests/transformer/moe/test_paged_stash_runner.py new file mode 100644 index 00000000000..aee4203cb1b --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_paged_stash_runner.py @@ -0,0 +1,206 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace + +import torch + +from megatron.core.transformer.moe.paged_stash import PagedStashManager, PagedStashRunner + + +def _config(moe_paged_stash): + return SimpleNamespace(moe_paged_stash=moe_paged_stash, moe_expert_rank_capacity_factor=1.5) + + +class _FakeTokenDispatcher: + + def __init__(self, config): + self.config = config + self._comm_manager = SimpleNamespace(moe_expert_rank_capacity_factor=1.5) + self.invalidate_count = 0 + self.reset_count = 0 + + def check_over_budget(self): + return None + + def reset_over_budget(self): + self.reset_count += 1 + + def invalidate_ep_bootstrap(self): + self.invalidate_count += 1 + + +class _FakeMoELayer(torch.nn.Module): + + def __init__(self, config): + super().__init__() + self.config = config + self.token_dispatcher = _FakeTokenDispatcher(config) + + +class _FakeTransformerLayer(torch.nn.Module): + + def __init__(self, mlp): + super().__init__() + self.mlp = mlp + + +class _FakeStack(torch.nn.Module): + + def __init__(self, layer): + super().__init__() + self.layers = torch.nn.ModuleList([layer]) + + +class _FakeMTPPredictionLayer(torch.nn.Module): + + def __init__(self, mtp_model_layer): + super().__init__() + self.mtp_model_layer = mtp_model_layer + + +class _FakeModelChunk(torch.nn.Module): + + def __init__(self, config, decoder_moe, mtp_moe, nested_mtp): + super().__init__() + self.config = config + self.decoder = _FakeStack(_FakeTransformerLayer(decoder_moe)) + mtp_model_layer = _FakeTransformerLayer(mtp_moe) + if nested_mtp: + mtp_model_layer = _FakeStack(mtp_model_layer) + self.mtp = _FakeStack(_FakeMTPPredictionLayer(mtp_model_layer)) + self.mtp_process = True + # Register the decoder MoE through a second path to verify identity deduplication. + self.duplicate_decoder_moe = decoder_moe + self.zero_grad_count = 0 + + def zero_grad_buffer(self): + self.zero_grad_count += 1 + + +def _run_retry( + monkeypatch, training_config, model_config, decoder_config, mtp_config, nested_mtp=False +): + monkeypatch.setattr( + "megatron.core.transformer.multi_token_prediction.MultiTokenPredictionLayer", + _FakeMTPPredictionLayer, + ) + decoder_moe = _FakeMoELayer(decoder_config) + mtp_moe = _FakeMoELayer(mtp_config) + model = _FakeModelChunk(model_config, decoder_moe, mtp_moe, nested_mtp=nested_mtp) + + values_seen_by_forward = [] + + def forward_backward_func(**_): + values_seen_by_forward.append( + ( + training_config.moe_paged_stash, + model_config.moe_paged_stash, + decoder_config.moe_paged_stash, + mtp_config.moe_paged_stash, + ) + ) + return len(values_seen_by_forward) + + release_stash_buffer_calls = [] + fake_stash_manager = SimpleNamespace( + overflow=None, + host_spill=None, + release_stash_buffers=lambda: release_stash_buffer_calls.append(None), + ) + monkeypatch.setattr(PagedStashManager, 'STASH_MGR', fake_stash_manager) + runner = PagedStashRunner( + config=training_config, + copy_main_params=False, + model=[model], + optimizer=None, + forward_backward_func=forward_backward_func, + ) + overflow_results = iter([(1, 0, 0), (0, 0, 0)]) + runner.check_moe_overflow = lambda: next(overflow_results) + + result = runner( + model=[model], data_iterator=None, num_microbatches=1, seq_length=1, forward_only=False + ) + + return SimpleNamespace( + runner=runner, + result=result, + model=model, + decoder_moe=decoder_moe, + mtp_moe=mtp_moe, + values_seen_by_forward=values_seen_by_forward, + release_stash_buffer_calls=release_stash_buffer_calls, + ) + + +def test_retry_preserves_shared_root_config_behavior(monkeypatch): + """Models whose MoE modules share the root config retain their existing behavior.""" + training_config = _config(True) + model_config = _config(True) + + run = _run_retry( + monkeypatch, + training_config=training_config, + model_config=model_config, + decoder_config=model_config, + mtp_config=model_config, + ) + + assert run.runner.moe_layers == [run.decoder_moe, run.mtp_moe] + assert [id(config) for config in run.runner._configs_to_sync_moe_paged_stash] == [ + id(training_config), + id(model_config), + ] + assert run.values_seen_by_forward == [(True, True, True, True), (False, False, False, False)] + assert run.result == 2 + assert run.model.zero_grad_count == 1 + assert len(run.release_stash_buffer_calls) == 1 + assert run.decoder_moe.token_dispatcher.reset_count == 1 + assert run.mtp_moe.token_dispatcher.reset_count == 1 + assert run.decoder_moe.token_dispatcher.invalidate_count == 2 + assert run.mtp_moe.token_dispatcher.invalidate_count == 2 + assert run.decoder_moe.token_dispatcher._comm_manager.moe_expert_rank_capacity_factor == 1.5 + assert run.mtp_moe.token_dispatcher._comm_manager.moe_expert_rank_capacity_factor == 1.5 + assert training_config.moe_paged_stash is True + assert model_config.moe_paged_stash is True + + +def test_retry_disables_and_restores_per_module_configs(monkeypatch): + """Retry must disable direct and nested-MTP MoE configs, then restore each value.""" + training_config = _config(True) + model_config = _config(True) + decoder_moe_config = _config(True) + mtp_moe_config = _config(False) + + run = _run_retry( + monkeypatch, + training_config=training_config, + model_config=model_config, + decoder_config=decoder_moe_config, + mtp_config=mtp_moe_config, + nested_mtp=True, + ) + + assert run.runner.moe_layers == [run.decoder_moe] + assert [id(config) for config in run.runner._configs_to_sync_moe_paged_stash] == [ + id(training_config), + id(model_config), + id(decoder_moe_config), + id(mtp_moe_config), + ] + assert run.values_seen_by_forward == [(True, True, True, False), (False, False, False, False)] + assert run.result == 2 + assert run.model.zero_grad_count == 1 + assert len(run.release_stash_buffer_calls) == 1 + assert run.decoder_moe.token_dispatcher.reset_count == 1 + assert run.mtp_moe.token_dispatcher.reset_count == 0 + assert run.decoder_moe.token_dispatcher.invalidate_count == 2 + assert run.mtp_moe.token_dispatcher.invalidate_count == 0 + assert run.decoder_moe.token_dispatcher._comm_manager.moe_expert_rank_capacity_factor == 1.5 + assert run.mtp_moe.token_dispatcher._comm_manager.moe_expert_rank_capacity_factor == 1.5 + assert ( + training_config.moe_paged_stash, + model_config.moe_paged_stash, + decoder_moe_config.moe_paged_stash, + mtp_moe_config.moe_paged_stash, + ) == (True, True, True, False) From 5714a206e89dcc08f4ac6ed3d8057c6c1d8dcfa2 Mon Sep 17 00:00:00 2001 From: Maanu Grover Date: Thu, 13 Aug 2026 07:04:11 +0000 Subject: [PATCH 013/101] Respect `profile_ranks` on log interval memory snapshots (#6354) Signed-off-by: Maanu Grover --- megatron/training/training.py | 13 +++++++------ megatron/training/utils/utils.py | 2 +- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index ebfa93f6b87..0e9417b8fbe 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -41,6 +41,7 @@ _LEGACY_TRAIN_START_TIME = time.time() # NOTE(asolergi-nv): Legacy timestamp # First-party. +from megatron.core._rank_utils import safe_get_rank from megatron.core import mpu, nccl_allocator, tensor_parallel from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper from megatron.core.distributed import DistributedDataParallel as DDP @@ -2928,12 +2929,12 @@ def training_log( # Dump memory snapshot and print metrics to stdout. if iteration % args.log_interval == 0 or is_first_iteration: - if args.record_memory_history and (is_last_rank() or torch.distributed.get_backend() == 'fake'): - snapshot = torch.cuda.memory._snapshot() - from pickle import dump - - with open(args.memory_snapshot_path, 'wb') as f: - dump(snapshot, f) + should_prof_rank = (args.profile_ranks == [] or safe_get_rank() in args.profile_ranks) # [] is all ranks + if args.record_memory_history and (should_prof_rank or torch.distributed.get_backend() == 'fake'): + rank = safe_get_rank() + base, ext = os.path.splitext(args.memory_snapshot_path) + snapshot_filename = f"{base}_{rank}{ext}" + torch.cuda.memory._dump_snapshot(snapshot_filename) elapsed_time = timers('interval-time').elapsed(barrier=True, reset=should_reset) elapsed_time_per_iteration = elapsed_time / total_iterations diff --git a/megatron/training/utils/utils.py b/megatron/training/utils/utils.py index 6c5b3454798..108a92d585c 100644 --- a/megatron/training/utils/utils.py +++ b/megatron/training/utils/utils.py @@ -44,7 +44,7 @@ def _oom_observer( """Dump a snapshot on OOM so we can inspect what was live at the failure.""" rank = safe_get_rank() base, ext = os.path.splitext(profiling.memory_snapshot_path) - filename = f"{base}_oom_rank-{rank}{ext}" + filename = f"{base}_oom_rank_{rank}{ext}" torch.cuda.memory._dump_snapshot(filename) # logger.info so the message reaches stderr on any profiled rank, not just rank 0. logger.info(f"[OOM] rank {rank} saved memory snapshot to {filename}") From fecb29f99265d5d26907988e1d87f5335df761bc Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 13 Aug 2026 08:43:19 +0000 Subject: [PATCH 014/101] Add unit tests for GDP inference (#6505) Signed-off-by: Deepak Narayanan Signed-off-by: Keshav Santhanam Signed-off-by: Mikail Khona Signed-off-by: Mikail Khona (NVIDIA) Signed-off-by: Kezhi Kong Co-authored-by: Kezhi Kong Co-authored-by: Deepak Narayanan Co-authored-by: Claude Opus 4.6 Co-authored-by: Mikail Khona (NVIDIA) Co-authored-by: Mikail Khona Co-authored-by: root Co-authored-by: Roger Waleffe Co-authored-by: Kezhi Kong --- megatron/core/ssm/gated_delta_product.py | 11 +- .../ssm/test_gdp_dynamic_inference.py | 629 ++++++++++++++++++ 2 files changed, 635 insertions(+), 5 deletions(-) create mode 100644 tests/unit_tests/ssm/test_gdp_dynamic_inference.py diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 1baaea0ea3a..3fa69813680 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -21,6 +21,9 @@ ) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gdp_context_parallel import GDPContextParallel + +# Decode uses the in-repo Triton conv update, which accepts int64 slot indices. +from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update from megatron.core.ssm.packed_seq_helpers import ( build_packed_seq_idx, check_fla_sequence_packing_support, @@ -38,11 +41,10 @@ from megatron.core.utils import deprecate_inference_params try: - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update + from causal_conv1d import causal_conv1d_fn from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states except ImportError: causal_conv1d_fn = None - causal_conv1d_update = None causal_conv1d_varlen_states = None try: @@ -604,9 +606,8 @@ def ssm_decode( dim=-1, ) - # Indexed conv update: reads/writes the per-request conv state rows - # selected by ``batch_indices``, in place. ``self.activation`` must be the - # activation *string* so the kernel enables SiLU (a bool would disable it). + # Indexed conv update into the per-request state rows (``batch_indices`` + # is None for static batching, where the cache is already in order). VKQ = causal_conv1d_update( VKQ, conv_state, diff --git a/tests/unit_tests/ssm/test_gdp_dynamic_inference.py b/tests/unit_tests/ssm/test_gdp_dynamic_inference.py new file mode 100644 index 00000000000..cbb4fbc3821 --- /dev/null +++ b/tests/unit_tests/ssm/test_gdp_dynamic_inference.py @@ -0,0 +1,629 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GDP (Gated Delta Product) inference equivalence tests. + +`GatedDeltaProductMixer` supports two inference paths: static batching +(`StaticInferenceContext`) and dynamic batching (`DynamicInferenceContext`). +These tests assert that both are numerically equivalent to a plain +full-sequence forward, and therefore to each other. + +The reference is a single full-sequence `model.forward` (the +`chunk_gated_delta_product` path). Dynamic prefill runs the same chunk kernel +over a packed var-len layout, and static-batching prefill runs the same chunk +kernel with the recurrent cache seeded; both must reproduce the reference's +last-token logits. + +The single-forward equivalence tests run at TP=1 (they compare raw logits that +are sequence-sharded under sequence-parallel, and the static path does not +support SP). The end-to-end engine tests sweep TP (`_TP_SIZES`) with SP enabled +at TP>1, covering dynamic inference under tensor + sequence parallelism; TP>1 +variants skip when the world has too few GPUs. +""" + +from __future__ import annotations + +import random +import types +from typing import Dict, List, Optional, Sequence, Tuple + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig +from megatron.core.inference.contexts import StaticInferenceContext +from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext +from megatron.core.inference.engines import DynamicInferenceEngine +from megatron.core.inference.inference_request import DynamicInferenceRequest, Status +from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( + GPTInferenceWrapper, +) +from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) +from megatron.core.inference.utils import InferenceMode +from megatron.core.models.hybrid.hybrid_layer_specs import gated_delta_product_inference_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer +from megatron.core.ssm.packed_seq_helpers import check_fla_sequence_packing_support +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.cuda_graphs import delete_cuda_graphs +from megatron.core.utils import is_fa_min_version +from tests.unit_tests.test_utilities import Utils, clear_nvte_env_vars + +try: + import einops # noqa: F401 + import fla # noqa: F401 + import mamba_ssm # noqa: F401 + + HAVE_GDP_DEPS = True +except ImportError: + HAVE_GDP_DEPS = False + +# GDP dynamic inference relies on the same packed-sequence conv1d kernel as the +# training/prefill path (`causal_conv1d_fn(seq_idx=...)`, added in 1.4.0). +_PACKING_OK, _PACKING_REASON = check_fla_sequence_packing_support() + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif(not HAVE_GDP_DEPS, reason="GDP requires fla, mamba_ssm, and einops"), + pytest.mark.skipif(not _PACKING_OK, reason=_PACKING_REASON or "packed-seq support missing"), +] + + +# A short single-chunk prompt is enough to exercise the packed-varlen dynamic +# path; the sizes are kept small to keep the test fast. +_VOCAB_SIZE = 128 +_MAX_SEQ_LEN = 512 +_PROMPT_LEN = 64 + +# bf16 chunk-vs-chunk tolerance. Full-forward and prefill run the same +# `chunk_gated_delta_product` kernel, so they differ only by the packed var-len +# layout and floating-point accumulation order. +_ATOL = 5e-2 +_RTOL = 5e-2 + +# Looser tolerance for the decode-step check: it compares the recurrent decode +# kernel against a full-sequence chunk-kernel recompute (different kernels), so +# it drifts more than the chunk-vs-chunk prefill comparison. Still far tighter +# than the O(1)+ deviations a genuinely broken decode/state-handoff would show. +_DECODE_ATOL = 1e-1 +_DECODE_RTOL = 1e-1 + +# `DynamicInferenceContext` requires at least one attention layer, so the model +# pattern is GDP mixer + attention + MLP. Dynamic batching needs a recent +# flash-attention. +_LAYER_PATTERN = "M*-" +_NUM_LAYERS = len(_LAYER_PATTERN) +requires_dynamic_batching = pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need flash-attn >= 2.7.3 for dynamic batching" +) + +# Tensor-parallel sizes swept by the end-to-end engine tests. TP>1 requires +# sequence-parallel (the inference-optimized linears assert it), which GDP's +# dynamic path supports; static inference does not, so the single-forward +# equivalence tests above stay TP=1. +_TP_SIZES = [1, 2] + + +def _make_config(tp: int = 1) -> TransformerConfig: + """A small but shape-valid GDP config, sharded across `tp` tensor-parallel ranks. + + The in_proj output width (`zVKQba`) is column-parallel, so each rank sees + `proj_dim / tp` channels. The packed-prefill conv slices a channels-last view + out of that and `causal_conv1d_fn` requires its stride (the per-rank width) to + be a multiple of 8. With mamba_num_heads=16 the full width is + `(1+M)*d_inner + (M+1)*ngroups*d_state + (M+1)*nheads = 3*256 + 3*64 + 3*16 + = 1008`, so per-rank widths are 1008 (tp=1) and 504 (tp=2), both aligned. + Production configs satisfy this by having much larger, aligned dimensions. + """ + return TransformerConfig( + num_layers=_NUM_LAYERS, + hidden_size=64, + num_attention_heads=4, + num_query_groups=4, + ffn_hidden_size=128, + normalization="RMSNorm", + bf16=True, + params_dtype=torch.bfloat16, + mamba_num_heads=16, + mamba_head_dim=16, + mamba_num_groups=4, + mamba_state_dim=16, + gdp_num_householder=2, + is_hybrid_model=True, # needed for correct out_proj init + tensor_model_parallel_size=tp, + sequence_parallel=False, + context_parallel_size=1, + ) + + +def _build_model(tp: int = 1) -> HybridModel: + """Build a small GDP hybrid model (mixer + attention + MLP), eval on CUDA.""" + model_parallel_cuda_manual_seed(123) + model = HybridModel( + config=_make_config(tp), + hybrid_stack_spec=gated_delta_product_inference_stack_spec, + vocab_size=_VOCAB_SIZE, + max_sequence_length=_MAX_SEQ_LEN, + hybrid_layer_pattern=_LAYER_PATTERN, + ) + return model.cuda().eval() + + +@requires_dynamic_batching +class TestGDPDynamicInference: + """Static/dynamic GDP inference equivalence against a full-sequence forward. + + These compare raw `model.forward` logits, which are sequence-sharded under + sequence-parallel; combined with the static path not supporting SP, they run + at TP=1 only. TP>1 dynamic inference is covered end-to-end by the engine + tests below, which handle SP through the inference wrapper. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + self.model = _build_model(tp=1) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _input_ids(self) -> torch.Tensor: + """A deterministic single-request prompt: shape [1, _PROMPT_LEN].""" + return torch.arange(_PROMPT_LEN, device="cuda", dtype=torch.long).unsqueeze(0) + + @torch.inference_mode() + def _full_forward_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Reference: plain full-sequence forward -> last-token logits [1, V].""" + # No inference_context: GDP runs the training / chunk_gated_delta_product + # path. This is the ground truth both inference modes must reproduce. + InferenceMode.unset_active() + position_ids = torch.arange(input_ids.shape[1], device="cuda").unsqueeze(0) + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + runtime_gather_output=True, + ) + return logits[:, -1, :].float() + + def _build_dynamic_context(self) -> DynamicInferenceContext: + mamba_config = MambaInferenceStateConfig.from_model(self.model) + assert mamba_config is not None, "GDP hybrid model should expose Mamba inference state" + return DynamicInferenceContext( + model_config=self.model.config, + inference_config=InferenceConfig( + max_sequence_length=_MAX_SEQ_LEN, + buffer_size_gb=1.0, + block_size_tokens=256, + # Materialize all tokens so we can read the prompt's last-token + # logits directly (static batching always uses last-token only). + materialize_only_last_token_logits=False, + mamba_inference_state_config=mamba_config, + num_cuda_graphs=0, + use_cuda_graphs_for_non_decode_steps=False, + max_requests=4, + max_tokens=128, + ), + ) + + @torch.inference_mode() + def _dynamic_prefill_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Dynamic-batching prefill -> last-token logits [1, V].""" + ctx = self._build_dynamic_context() + request = DynamicInferenceRequest( + request_id=0, + prompt_tokens=input_ids.cpu().squeeze(0), + sampling_params=SamplingParams(num_tokens_to_generate=1, termination_id=-1), + ) + ctx.add_request(request) + ctx.initialize_attention_state() + with InferenceMode.active(): + logits = self.model( + input_ids=input_ids, + position_ids=None, + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + # materialize_only_last_token_logits=False -> [1, prompt_len, V]. + return logits[:, -1, :].float() + + @torch.inference_mode() + def _static_prefill_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Static-batching prefill -> last-token logits [1, V].""" + ctx = StaticInferenceContext(max_batch_size=1, max_sequence_length=_MAX_SEQ_LEN) + ctx.sequence_len_offset = 0 + position_ids = torch.arange(input_ids.shape[1], device="cuda").unsqueeze(0) + with InferenceMode.active(): + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + # StaticInferenceContext forces materialize_only_last_token_logits=True, + # so the sequence dimension is already collapsed to the last token. + assert logits.shape[1] == 1 + return logits[:, 0, :].float() + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + def test_constructor(self): + """The GDP stack spec wires a GatedDeltaProductMixer into the mamba layer.""" + assert isinstance(self.model, HybridModel) + mixers = [ + layer.mixer + for layer in self.model.decoder.layers + if hasattr(layer, "mixer") and layer.mixer is not None + ] + assert len(mixers) == 1, f"pattern {_LAYER_PATTERN!r} should yield exactly one mixer layer" + assert isinstance(mixers[0], GatedDeltaProductMixer) + + def test_full_forward_shape(self): + """Sanity check: plain forward returns [batch, seq, vocab].""" + input_ids = self._input_ids() + InferenceMode.unset_active() + position_ids = torch.arange(_PROMPT_LEN, device="cuda").unsqueeze(0) + with torch.inference_mode(): + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + runtime_gather_output=True, + ) + assert logits.shape == (1, _PROMPT_LEN, _VOCAB_SIZE) + + def test_dynamic_prefill_matches_full_forward(self): + """Dynamic-batching prefill reproduces the full-sequence forward.""" + input_ids = self._input_ids() + reference = self._full_forward_last_logits(input_ids) + dynamic = self._dynamic_prefill_last_logits(input_ids) + torch.testing.assert_close(dynamic, reference, atol=_ATOL, rtol=_RTOL) + + def test_static_prefill_matches_full_forward(self): + """Static-batching prefill reproduces the full-sequence forward.""" + input_ids = self._input_ids() + reference = self._full_forward_last_logits(input_ids) + static = self._static_prefill_last_logits(input_ids) + torch.testing.assert_close(static, reference, atol=_ATOL, rtol=_RTOL) + + def test_static_and_dynamic_prefill_agree(self): + """Static and dynamic inference produce equivalent logits. + + This is the central invariant: the two batching strategies must agree. + Anchoring each to the full-sequence forward (above) guarantees this + transitively, but assert it directly as well so a regression in either + path that happens to drift in the same direction is still caught. + """ + input_ids = self._input_ids() + static = self._static_prefill_last_logits(input_ids) + dynamic = self._dynamic_prefill_last_logits(input_ids) + torch.testing.assert_close(static, dynamic, atol=_ATOL, rtol=_RTOL) + + @torch.inference_mode() + def test_decode_step_matches_recompute(self): + """One decode step matches a full-sequence recompute (decode-path check). + + This validates the recurrent decode kernel and the prefill->decode + conv/SSM state handoff independently of any golden snapshot: after + prefilling the prompt, decoding one more token must produce the same + next-token logits as a plain forward over prompt+token. Compared at the + logit level with tolerance, so it is robust to the bf16 numerics that + make exact greedy token equality across the recurrent/chunk kernels + fragile. Uses `StaticInferenceContext`, whose decode calls the same + `ssm_decode` recurrent kernel as the dynamic engine. + """ + prompt = self._input_ids() # [1, P] + + # Ground truth: the next token from the prompt, and the full-recompute + # distribution for the token after it (chunk kernel over prompt+token). + next_token = int(self._full_forward_last_logits(prompt).argmax(dim=-1).item()) + extended = torch.cat( + [prompt, torch.tensor([[next_token]], dtype=torch.int64, device="cuda")], dim=1 + ) + recompute = self._full_forward_last_logits(extended) # [1, V] + + # Incremental path: prefill the prompt, then a single decode step. + ctx = StaticInferenceContext(max_batch_size=1, max_sequence_length=_MAX_SEQ_LEN) + prompt_length = prompt.shape[1] + with InferenceMode.active(): + ctx.sequence_len_offset = 0 + self.model( + input_ids=prompt, + position_ids=torch.arange(prompt_length, device="cuda").unsqueeze(0), + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + ctx.sequence_len_offset = prompt_length + decode_logits = self.model( + input_ids=torch.tensor([[next_token]], dtype=torch.int64, device="cuda"), + position_ids=torch.tensor([[prompt_length]], dtype=torch.int64, device="cuda"), + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + decode_last = decode_logits[:, -1, :].float() + + torch.testing.assert_close(decode_last, recompute, atol=_DECODE_ATOL, rtol=_DECODE_RTOL) + + +# ====================================================================== +# End-to-end engine tests. +# +# The tests above exercise a single forward pass. These drive the full +# `DynamicInferenceEngine` (add requests -> schedule -> prefill -> decode -> +# finish) so that GDP is validated through the same runtime path production +# inference uses: the text-generation controller, the inference-wrapped model, +# the KV/Mamba-state cache, and the request scheduler. +# +# Decoding is greedy (`top_k=1`) so outputs are deterministic. Decode +# correctness (the recurrent kernel `fused_recurrent_gated_delta_rule` plus the +# slot-indexed conv/SSM cache, distinct from the prefill chunk kernel) is +# validated by a byte-for-byte match against committed golden token ids, +# mirroring the Mamba2 `test_dynamic_engine.py::test_simple` style. The golden +# constant is captured from a reference GPU run (see `_GOLDEN_*` below); until it +# is populated the test self-captures and skips with the observed ids. +# ====================================================================== + +# Fixed prompts for the golden-token test. Deterministic (not random) so the +# committed golden ids below are reproducible across machines. Varying lengths +# exercise the scheduler's mixed-length prefill batching. +_GOLDEN_PROMPTS: List[List[int]] = [ + [3, 14, 15, 92, 65, 35, 89, 79], + [2, 71, 82, 81, 8], + [11, 22, 33, 44, 55, 66], + [7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7], +] +_GOLDEN_NUM_TOKENS_TO_GENERATE = 12 + +# Golden generated-token ids per TP size, one list per prompt in `_GOLDEN_PROMPTS`, +# captured from a reference GPU run. TP shards the weights differently, so each TP +# size has its own goldens. Environment-sensitive (FLA / causal_conv1d kernel +# build, GPU arch); re-capture if the kernels or config change. While an entry is +# None, the test self-captures: it prints the observed ids and skips instead of +# failing. Paste them in (see the skip message) to turn it into a hard assertion. +_GOLDEN_GENERATED_TOKENS: Dict[int, Optional[List[List[int]]]] = { + 1: [ + [32, 126, 35, 125, 52, 116, 55, 38, 39, 4, 53, 100], + [88, 10, 105, 95, 105, 44, 2, 100, 127, 59, 23, 18], + [29, 9, 61, 2, 100, 2, 69, 75, 36, 80, 103, 26], + [2, 55, 4, 108, 116, 120, 24, 113, 100, 48, 111, 22], + ], + 2: [ + [12, 51, 63, 22, 2, 40, 45, 30, 55, 10, 31, 29], + [38, 54, 29, 17, 33, 13, 10, 45, 1, 22, 21, 37], + [24, 55, 38, 55, 35, 61, 26, 25, 31, 20, 62, 56], + [53, 26, 31, 61, 16, 19, 42, 41, 49, 18, 53, 26], + ], +} + + +def _make_engine_config(tp: int = 1) -> TransformerConfig: + """GDP config for the engine tests: same shape as `_make_config`, plus the + deterministic inference sampling knobs greedy decoding needs. TP>1 turns on + sequence-parallel, which the inference-optimized linears require.""" + config = _make_config(tp) + config.sequence_parallel = tp > 1 + config.inference_rng_tracker = True + config.inference_sampling_seed = 123 + return config + + +@pytest.mark.internal +@requires_dynamic_batching +@pytest.mark.skipif(not HAVE_GDP_DEPS, reason="GDP requires fla, mamba_ssm, and einops") +@pytest.mark.skipif(not _PACKING_OK, reason=_PACKING_REASON or "packed-seq support missing") +class TestGDPDynamicInferenceEngine: + """End-to-end GDP decoding through `DynamicInferenceEngine`.""" + + SEED = 123 + VOCAB_SIZE = _VOCAB_SIZE + + def teardown_method(self, method): + delete_cuda_graphs() + Utils.destroy_model_parallel() + + # ------------------------------------------------------------------ + # Harness + # ------------------------------------------------------------------ + + def _build_engine( + self, + *, + num_tokens_to_generate: int, + tp: int = 1, + num_requests: Optional[int] = None, + prompt_length: Optional[int] = None, + prompts: Optional[Sequence[Sequence[int]]] = None, + ) -> Tuple[DynamicInferenceEngine, List[DynamicInferenceRequest]]: + """Build a greedy GDP engine plus its requests at TP=`tp`. + + Either pass explicit `prompts` (deterministic token lists) or + `num_requests` + `prompt_length` (random prompts of a fixed length). + Skips if the world is too small for the requested TP size. + """ + if Utils.world_size < tp: + pytest.skip(f"TP={tp} requires at least {tp} GPUs") + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp, pipeline_model_parallel_size=1 + ) + clear_nvte_env_vars() + random.seed(self.SEED) + torch.manual_seed(self.SEED) + model_parallel_cuda_manual_seed( + seed=self.SEED, inference_rng_tracker=True, force_reset_rng=True + ) + + if prompts is not None: + prompt_tensors = [torch.tensor(p, dtype=torch.int64, device="cuda") for p in prompts] + else: + assert num_requests is not None and prompt_length is not None + prompt_tensors = [ + torch.randint( + 0, self.VOCAB_SIZE - 1, (prompt_length,), dtype=torch.int64, device="cuda" + ) + for _ in range(num_requests) + ] + + max_prompt_length = max(int(p.numel()) for p in prompt_tensors) + max_sequence_length = max_prompt_length + num_tokens_to_generate + config = _make_engine_config(tp) + model = HybridModel( + config=config, + hybrid_stack_spec=gated_delta_product_inference_stack_spec, + vocab_size=self.VOCAB_SIZE, + max_sequence_length=max_sequence_length, + parallel_output=True, + hybrid_layer_pattern=_LAYER_PATTERN, + pre_process=parallel_state.is_pipeline_first_stage(), + post_process=parallel_state.is_pipeline_last_stage(), + ).cuda() + for param in model.parameters(): + param.data = param.data.to(config.params_dtype) + model.eval() + + context = DynamicInferenceContext( + model_config=config, + inference_config=InferenceConfig( + max_sequence_length=max_sequence_length, + buffer_size_gb=0.1, + block_size_tokens=256, + materialize_only_last_token_logits=True, + mamba_inference_state_config=MambaInferenceStateConfig.from_model(model), + num_cuda_graphs=None, + use_cuda_graphs_for_non_decode_steps=False, + max_requests=32, + max_tokens=1024, + ), + ) + + wrapped_model = GPTInferenceWrapper(model, context) + wrapped_model.model_is_pipeline_parallel = not ( + parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage() + ) + controller = TextGenerationController( + inference_wrapped_model=wrapped_model, + tokenizer=types.SimpleNamespace( + vocab_size=self.VOCAB_SIZE, detokenize=lambda tokens: "tokenized_prompt" + ), + ) + delete_cuda_graphs() + engine = DynamicInferenceEngine(controller, context) + + requests = [ + DynamicInferenceRequest( + request_id=request_id, + prompt_tokens=prompt_tokens, + sampling_params=SamplingParams( + num_tokens_to_generate=num_tokens_to_generate, + termination_id=-1, # never terminate early -> fixed output length + top_k=1, # greedy -> deterministic + ), + ) + for request_id, prompt_tokens in enumerate(prompt_tensors) + ] + return engine, requests + + @staticmethod + @torch.inference_mode() + def _run_to_completion( + engine: DynamicInferenceEngine, requests: List[DynamicInferenceRequest] + ) -> Dict[int, DynamicInferenceRequest]: + """Add every request, step until the engine drains, return finished requests by id.""" + for request in requests: + engine._add_request(request) + + finished: Dict[int, DynamicInferenceRequest] = {} + # Bound the loop so a scheduling regression fails loudly instead of hanging. + for _ in range(1000): + result = engine.step_modern() + for record in result["finished_request_records"]: + merged = record.merge() + finished[merged.request_id] = merged + if not engine.has_unfinished_requests(): + break + assert not engine.has_unfinished_requests(), "engine did not drain within step budget" + return finished + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_engine_runs_to_completion(self, tp): + """Every request completes and yields exactly the requested token count.""" + num_tokens_to_generate = 8 + engine, requests = self._build_engine( + tp=tp, num_requests=4, prompt_length=8, num_tokens_to_generate=num_tokens_to_generate + ) + finished = self._run_to_completion(engine, requests) + + assert len(finished) == len(requests) + for request in requests: + merged = finished[request.request_id] + assert merged.status == Status.COMPLETED + # termination_id=-1 disables early stop, so the length is exact. + assert len(merged.generated_tokens) == num_tokens_to_generate + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_engine_greedy_matches_golden(self, tp): + """Greedy decode reproduces committed golden token ids (Mamba2-style). + + Deterministic fixed prompts + greedy sampling make the output a stable + fingerprint of the GDP prefill+decode path. Until the TP entry in + `_GOLDEN_GENERATED_TOKENS` is captured from a reference GPU run, the test + prints the observed ids and skips instead of failing. + """ + engine, requests = self._build_engine( + tp=tp, prompts=_GOLDEN_PROMPTS, num_tokens_to_generate=_GOLDEN_NUM_TOKENS_TO_GENERATE + ) + finished = self._run_to_completion(engine, requests) + observed = [finished[r.request_id].generated_tokens for r in requests] + + golden = _GOLDEN_GENERATED_TOKENS.get(tp) + if golden is None: + pytest.skip( + f"golden tokens for TP={tp} not captured yet; paste the following into " + f"_GOLDEN_GENERATED_TOKENS[{tp}]:\n{observed!r}" + ) + + assert observed == golden, ( + f"generated tokens != golden (TP={tp}):\n golden = {golden}\n" + f" observed = {observed}" + ) + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_generate_over_multiple_prompts(self, tp): + """`engine.generate` drives several prompts through to completion at once.""" + engine, requests = self._build_engine( + tp=tp, num_requests=4, prompt_length=8, num_tokens_to_generate=4 + ) + + prompts = [f"prompt{i}" for i in range(len(requests))] + + def mock_tokenize_prompt(tokenizer, prompt, add_BOS=False): + prompt_num = int(prompt[-1]) + return [10 + i for i in range(prompt_num + 2)] + + engine.controller.tokenize_prompt = mock_tokenize_prompt + + finished_records = engine.generate(prompts, requests[0].sampling_params) + finished = [record.merge() for record in finished_records] + + assert len(finished) == len(prompts) + # generate() returns finished requests in request-id order. + assert [r.request_id for r in finished] == sorted(r.request_id for r in finished) + for request in finished: + assert request.status == Status.COMPLETED + assert len(request.generated_tokens) > 0 From 8ba5c60d2aba1c958bdd5c000ff2b241908ad2b2 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Thu, 13 Aug 2026 08:50:41 +0000 Subject: [PATCH 015/101] [fix] fix GTP+DCP ckpt saving/loading for GDP module (#6503) Signed-off-by: Shiqing Fan --- .../core/generalized_tensor_parallel.md | 8 +- megatron/core/ssm/gated_delta_product.py | 81 ++++- .../test_gtp_dcp.py | 278 ++++++++++++++++++ 3 files changed, 363 insertions(+), 4 deletions(-) diff --git a/docs/api-guide/core/generalized_tensor_parallel.md b/docs/api-guide/core/generalized_tensor_parallel.md index 11c01b8a504..eca9e2b7eb8 100644 --- a/docs/api-guide/core/generalized_tensor_parallel.md +++ b/docs/api-guide/core/generalized_tensor_parallel.md @@ -182,7 +182,7 @@ GTP_remat runs under both the standard **Adam** `DistributedOptimizer` and **Muo - **Adam** shards optimizer state over the gtp_remat/egtp_remat-excluded replicate group, like any GTP_remat run (§3.2). - **Muon** keeps matrix params *whole* (Newton–Schulz needs the full 2D weight). A GTP_remat-replicated whole param (e.g. MoE router, latent-proj MLPs by default) then lands on one checkpoint key shared by all GTP_remat peers, so the LayerWise optimizer folds `gtp_rank` into its `replica_id` — exactly one peer writes (the optimizer-state analog of the model-side fold in §3.3). -- **Native-FP8 optimizer-state matching (Muon path).** The save-side dequantize (§3.3) hands DCP a *fresh* BF16 tensor, which breaks the id-based optimizer-param → model-`ShardedTensor` match for every native-FP8 GTP_remat weight. The dequantized copy carries a `_gtp_dequant_src` backlink to the live FP8 param, and `_backfill_gtp_sharded_param_map` reuses the model's **own** entry (backlink first, tagged-name second) — preserving its full offsets (expert axes included) and `replica_id`. Only truly-unmatched params (Mamba `in_proj`, a gathered+split factory) take the per-shard rebuild, which refuses expert-parallel params rather than emit EP-colliding shards. +- **Native-FP8 optimizer-state matching (Muon path).** The save-side dequantize (§3.3) hands DCP a *fresh* BF16 tensor, which breaks the id-based optimizer-param → model-`ShardedTensor` match for every native-FP8 GTP_remat weight. The dequantized copy carries a `_gtp_dequant_src` backlink to the live FP8 param, and `_backfill_gtp_sharded_param_map` reuses the model's **own** entry (backlink first, tagged-name second) — preserving its full offsets (expert axes included) and `replica_id`. Only truly-unmatched params (the SSM `in_proj` weights, gathered+split factories) take the per-shard rebuild, which refuses expert-parallel params rather than emit EP-colliding shards. Neither path adds a GTP_remat-specific checkpoint format or call site. @@ -543,7 +543,9 @@ Because the offsets reconstruct the global shape, the checkpoint is independent **Alignment padding & cross-topology reshard.** When `_gtp_slice_one_param` pads `out_features` to a multiple of `gtp_remat_size · pad_for_alignment`, the saved global describes the *padded* shape, so the helper sets `allow_shape_mismatch=True`. DCP then tolerates a load-side topology whose alignment yields a different padded size — the unpadded data overlaps and the tail pad rows are zeros GTP_remat recomputes. -> Note: Mamba's `in_proj` is a special case: it **all-gathers its GTP_remat shards** back to the logical TP-local size and strips the pad *before* saving, so its global is topology-independent and needs no `allow_shape_mismatch`. +> Note: the SSM `in_proj` weights — Mamba's (`mamba_mixer.py`, split `[z|x|B|C|dt]`) and gated-delta-product's (`gated_delta_product.py`, split householder-major into `z|V*|K*|Q|b*|a`) — are a special case: each **all-gathers its GTP_remat shards** back to the logical TP-local size and strips the pad *before* saving, so its global is topology-independent and needs no `allow_shape_mismatch`. This is required, not just tidier: the split-chunk boundaries do not line up with the GTP_remat slice boundaries, so a raw shard cannot be split at all. The checkpoint therefore matches a non-GTP_remat run byte-for-byte. +> +> On **load**, the split factory's `merge_fn` is wrapped to invert this: it cats the chunks back to the unpadded TP-local width, re-pads with zeros up to `gtp_remat_local_size · gtp_remat_size`, and slices by the GTP_remat rank — mirroring `_gtp_slice_one_param` so the tensor lands in the live shard's layout. `gtp_remat_size == 1` skips both the gather and the pad/slice. **Optimizer state.** The distributed optimizer's master/moment `ShardedObject`s are keyed by `dp_group_idx`. Under GTP_remat/EGTP_remat each peer owns a *different* master shard (the optimizer shards over the gtp_remat/egtp_remat-**excluded** replicate group), so the index is taken from the gtp_remat/egtp_remat-**merged** model-parallel group (`mp_group` for dense, `expt_tp_pp_with_egtp_remat_group` for expert) — giving every peer a distinct key while replicate-group ranks remain true replicas under that key. @@ -739,7 +741,7 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall | `test_gtp_grad_correctness.py` | Gradient + dist-opt + grad-norm numeric parity vs a DP baseline at replicate (DP) > 1. Also the fp32-accumulation reduce-scatter (§2.6): gtp_remat-axis and DDP-axis parity, plus the size-2 bypass. | | `test_gtp_cudagraph_grad.py` | Capture-step grad-norm guard (§1.2): `_backup_grads_before_capture`/`_restore_grads_after_capture` keep a graph capture from clobbering finalized `main_grad` (own params + cross-graph `next_w`, incl. routed-expert `weight_list`). | | `test_gtp_partial_cg.py` | Four-layer partial-CG loss and eager-vs-replay grad-norm parity with two-slot ring reuse across independently replayed graphs (§3.5). | -| `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. | +| `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. Also the SSM `in_proj` gather+split: the gated-delta-product mixer's factory build/merge at MXFP8 alignment, and a full DCP save→load roundtrip of that mixer. | | `test_gtp_muon_dcp.py` | Muon optimizer-state DCP roundtrip (§1.6): `replica_id` fold + native-FP8 backfill matching. | | `test_gtp_recompute_chain.py` | Recompute-chain buffers (§3.1): adjacent nodes never share a gather buffer, dense and grouped, plus dgrad/wgrad parity vs no-recompute. | | `test_gtp_mtp.py` | GTP_remat + MTP shared weights (§3.5), 14 cases over `mtp_use_repeated_layer` × dense/MoE. Both MTP hazards are silent, so each needs its own guard: the async reduce-scatter path is compared numerically against the sync path on an identical model/sharding/batch, and all-gathers issued are tallied against consumes to catch a consume reading a buffer nothing gathered into. | diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 3fa69813680..b85a5a4756c 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -31,14 +31,21 @@ ) from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.tensor_parallel.gtp_api import HAVE_GTP from megatron.core.transformer import TransformerConfig from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.utils import ( + ensure_metadata_has_dp_cp_group, make_sharded_tensors_for_checkpoint, sharded_state_dict_default, ) -from megatron.core.utils import deprecate_inference_params +from megatron.core.utils import deprecate_inference_params, make_tp_sharded_tensor_for_checkpoint + +if HAVE_GTP: + from megatron.core.tensor_parallel.gtp_api import is_gtp_param +else: + is_gtp_param = None try: from causal_conv1d import causal_conv1d_fn @@ -863,6 +870,9 @@ def _get_states_from_cache(self, inference_context, batch_size, *, inference_par def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): """Provide a sharded state dictionary for distributed checkpointing.""" + # Guard for cases metadata is not provided + metadata = ensure_metadata_has_dp_cp_group(metadata) + sharded_state_dict = {} # Parameters self._save_to_state_dict(sharded_state_dict, "", keep_vars=True) @@ -899,6 +909,41 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + (1 + self.num_householder) * self.ngroups_local_tp * self.d_state + self.nheads_local_tp * (1 + self.num_householder) ) + # Under GTP, in_proj.weight is GTP-sliced along axis 0. The [z|V*|K*|Q|b*|a] split + # boundaries don't line up with GTP slice boundaries, so gather the shards back to + # TP-local size (strip the trailing pad rows from the gathered tail) and fall through + # to the same split path the non-GTP run uses — saved ckpt matches a non-GTP run. + in_proj_gtp_remat_size = getattr(self.in_proj.weight, "gtp_remat_size", 1) + in_proj_is_gtp = ( + in_proj_gtp_remat_size > 1 and HAVE_GTP and is_gtp_param(self.in_proj.weight) + ) + if in_proj_is_gtp: + gtp_remat_group = self.in_proj.weight.group + # in_proj.weight was already built at the sharded size by the submodule + # sharded_state_dict above — and, for native-FP8 GTP, dequantized to BF16 there + # (make_tp_sharded_tensor_for_checkpoint). Gather those (BF16) shards back to the + # full TP-local size so the [z|V*|K*|Q|b*|a] split below matches a non-GTP run. + local = sharded_state_dict[f"{prefix}in_proj.weight"].data.contiguous() + gathered = torch.empty( + (local.shape[0] * in_proj_gtp_remat_size,) + local.shape[1:], + dtype=local.dtype, + device=local.device, + ) + torch.distributed.all_gather_into_tensor(gathered, local, group=gtp_remat_group) + if gathered.shape[0] != in_proj_dim: + gathered = gathered[:in_proj_dim].contiguous() + # Gathered weight is replicated across full dp_cp; replica_id needs only the DP slot. + dp_cp_rank = torch.distributed.get_rank(metadata["dp_cp_group"]) + sharded_state_dict[f"{prefix}in_proj.weight"] = make_tp_sharded_tensor_for_checkpoint( + gathered, + f"{prefix}in_proj.weight", + tp_axis=0, + replica_id=(0, 0, dp_cp_rank), + prepend_offsets=sharded_offsets, + tp_group=self.pg_collection.tp, + dp_cp_group=metadata["dp_cp_group"], + ) + assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim, ( in_proj_dim, sharded_state_dict[f"{prefix}in_proj.weight"], @@ -930,6 +975,40 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_state_dict[key], in_proj_split_sections, in_proj_split_names, 0 ) + # GTP load-side inverse of the save-time all-gather (see + # docs/api-guide/core/generalized_tensor_parallel.md §3.3, in_proj note): the checkpoint + # stores the FULL TP-local in_proj.weight (pad stripped) under the per-householder split + # keys, so the default merge_fn cats them back to ``in_proj_dim`` rows with no + # padding. To reload into the live GTP param we must mirror init + # (``_gtp_slice_one_param``): F.pad the merged tensor with zeros up to + # ``gtp_remat_local_size * gtp_remat_size``, then slice by ``gtp_remat_local_rank``. + # gtp_remat_size=1 has no pad/slice. + if in_proj_is_gtp: + factory = sharded_state_dict[f"{prefix}in_proj.weight"] + gtp_remat_local_rank = torch.distributed.get_rank(self.in_proj.weight.group) + gtp_remat_local_size = self.in_proj.weight.data.size(0) + original_merge_fn = factory.merge_fn + + @torch.no_grad() + def _gtp_slice_after_cat( + sub_state_dict, + _orig=original_merge_fn, + _rank=gtp_remat_local_rank, + _size=gtp_remat_local_size, + _gtp_remat_size=in_proj_gtp_remat_size, + ): + full = _orig(sub_state_dict) + aligned_total = _size * _gtp_remat_size + pad_rows = aligned_total - full.shape[0] + if pad_rows > 0: + full = torch.nn.functional.pad(full, (0, 0, 0, pad_rows)) + start = _rank * _size + return full[start : start + _size].contiguous() + + sharded_state_dict[f"{prefix}in_proj.weight"] = replace( + factory, merge_fn=_gtp_slice_after_cat + ) + conv_dim = ( self.d_inner_local_tp * self.num_householder + (1 + self.num_householder) * self.ngroups_local_tp * self.d_state diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py index da42a086a91..af62304c5f8 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py @@ -942,6 +942,272 @@ def _worker_mamba_inproj_optim_param_map(rank, world_size, port): GTPShardedParam._chain_state = {} +# --------------------------------------------------------------------------- +# Gated-delta-product (GDP) in_proj: gather+split under GTP_remat +# +# GDP's ``in_proj.weight`` is GTP-sliced along axis 0 and zero-padded to an alignment multiple, +# while the checkpoint splits it into householder-major chunks whose boundaries do NOT line up with +# the GTP slice boundaries. ``GatedDeltaProductMixer.sharded_state_dict`` therefore all-gathers the +# shards back to the TP-local width and strips the pad before splitting (§3.3), and wraps the +# factory's merge_fn to re-pad + re-slice on load. The three workers below cover that contract. +# --------------------------------------------------------------------------- + +# in_proj width = d_inner(256)*4 + 4*ngroups(2)*d_state(128) + nheads(4)*4 = 2064. With +# pad_for_alignment=32 (what setup_gtp_remat_from_recipe picks for MXFP8) and gtp_remat_size=2 the +# alignment block is 64, so 48 pad rows fire -- the padded-shard case the split path must handle. +_GDP_HIDDEN_SIZE = 256 + + +def _build_gdp_mixer(required_pgs): + """Build a 1-layer GatedDeltaProductMixer. Returns ``(mixer, pg, in_proj_dim)``. + + Callers must have set ``update_gtp_config(pad_for_alignment=32)`` and initialized model + parallel with ``gtp_remat_size=2`` first. + """ + from megatron.core.models.hybrid.hybrid_layer_specs import gdp_stack_spec + from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer + + pg = ProcessGroupCollection.use_mpu_process_groups(required_pgs=required_pgs) + config = TransformerConfig( + num_layers=1, + hidden_size=_GDP_HIDDEN_SIZE, + num_attention_heads=4, + mamba_num_heads=4, + mamba_head_dim=64, + mamba_num_groups=2, + mamba_state_dim=128, + params_dtype=torch.bfloat16, + bf16=True, + ) + mixer = GatedDeltaProductMixer( + config, + gdp_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + config.hidden_size, + layer_number=1, + pg_collection=pg, + ).cuda() + in_proj_dim = ( + mixer.d_inner_local_tp * (1 + mixer.num_householder) + + (1 + mixer.num_householder) * mixer.ngroups_local_tp * mixer.d_state + + mixer.nheads_local_tp * (1 + mixer.num_householder) + ) + in_proj_w = mixer.in_proj.weight + assert isinstance(in_proj_w, GTPShardedParam), "in_proj.weight should be GTP_remat-sharded" + assert in_proj_w.data.size(0) * 2 > in_proj_dim, ( + f"expected GTP alignment padding to fire (got {in_proj_w.data.size(0)} * 2 == " + f"{in_proj_dim}); these tests must cover the strip-pad / re-pad path" + ) + return mixer, pg, in_proj_dim + + +def _gdp_valid_rows(in_proj_w, in_proj_dim): + """Rows of this rank's GTP shard that hold real weights (the rest are alignment pad).""" + local_rows = in_proj_w.data.size(0) + gtp_remat_rank = torch.distributed.get_rank(in_proj_w.group) + return max(0, min(local_rows, in_proj_dim - gtp_remat_rank * local_rows)) + + +def _worker_gdp_inproj_gather_split(rank, world_size, port): + """GatedDeltaProductMixer.sharded_state_dict under GTP_remat. + + Regression for the GDP save crash: the raw GTP shard neither matches ``in_proj_dim`` nor lines + up with the in_proj split boundaries -- the pre-fix code asserted here. Verify the mixer + gathers back to TP-local size, splits into the 6 chunks a non-GTP_remat run would write, and + that the load-side merge_fn re-pads + re-slices back to the live GTP shard. + """ + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + model_parallel_cuda_manual_seed(42) + update_gtp_config(pad_for_alignment=32) # MXFP8 alignment + mixer, _, in_proj_dim = _build_gdp_mixer(['tp', 'cp', 'gtp_remat']) + in_proj_w = mixer.in_proj.weight + + metadata = {'dp_cp_group': ps.get_data_parallel_group(with_context_parallel=True)} + # Pre-fix this raised AssertionError((in_proj_dim, ShardedTensor(...))). + sd = mixer.sharded_state_dict(prefix='mixer.', metadata=metadata) + + factory = sd['mixer.in_proj.weight'] + assert isinstance(factory, ShardedTensorFactory), type(factory) + # Save side: the gathered tensor is the full TP-local width, pad stripped. + assert factory.data.size(0) == in_proj_dim, (factory.data.size(0), in_proj_dim) + + from megatron.core.ssm.gated_delta_product import _get_in_proj_checkpoint_split_layout + + # The chunk names/sizes come from _get_in_proj_checkpoint_split_layout (householder-major: + # z, V0..V(M-1), K0..K(M-1), Q, b0..b(M-1), a). Derive the expectation from that helper so + # this pins "GTP splits exactly like a non-GTP_remat run" rather than a frozen key list. + _, expected_names = _get_in_proj_checkpoint_split_layout( + mixer.d_inner_local_tp, + mixer.ngroups_local_tp * mixer.d_state, + mixer.nheads_local_tp, + mixer.num_householder, + ) + chunks = factory.build_fn(factory.key, factory.data, factory.replica_id, None) + assert [t.key.rsplit('.', 1)[-1] for t in chunks] == expected_names + assert sum(t.data.size(0) for t in chunks) == in_proj_dim + + # Load side: cat the chunks, re-pad, re-slice -> exactly this rank's live GTP shard. + merged = factory.merge_fn([t.data for t in chunks]) + assert tuple(merged.shape) == tuple(in_proj_w.data.shape), ( + tuple(merged.shape), + tuple(in_proj_w.data.shape), + ) + # The pad rows the last GTP rank carries are never written to the ckpt -> back as zeros. + n_valid = _gdp_valid_rows(in_proj_w, in_proj_dim) + torch.testing.assert_close(merged[:n_valid], in_proj_w.data[:n_valid], rtol=0, atol=0) + assert torch.equal(merged[n_valid:], torch.zeros_like(merged[n_valid:])) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + GTPShardedParam._chain_state = {} + + +def _worker_gdp_save_load_roundtrip(rank, world_size, ckpt_base): + """End-to-end DCP save->load of a GatedDeltaProductMixer under GTP_remat. + + Companion to ``_worker_gdp_inproj_gather_split``, which only checks the factory build/merge + functions in isolation. This drives the real ``save``/``load`` so the load-side merge_fn + (re-pad + re-slice back to the live GTP shard) is exercised through DCP, and so a + duplicate-writer replica_id would surface as an 'Invalid access pattern'. + """ + from megatron.core.dist_checkpointing import load, save + from tests.unit_tests.dist_checkpointing import TempNamedDir + + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + model_parallel_cuda_manual_seed(42) + update_gtp_config(pad_for_alignment=32) # MXFP8 alignment + mixer, pg, in_proj_dim = _build_gdp_mixer( + ['tp', 'cp', 'gtp_remat', 'dp_cp', 'dp_cp_gtp_remat'] + ) + in_proj_w = mixer.in_proj.weight + + # ``save_checkpoint_and_time`` threads the gtp_remat-INCLUSIVE group; using the + # gtp_remat-excluding pg.dp_cp here collides replica_ids across gtp_remat peers. + metadata = {'dp_cp_group': pg.dp_cp_gtp_remat} + golden = {k: v.detach().clone() for k, v in mixer.state_dict().items()} + + with TempNamedDir(ckpt_base / 'gdp_gtp_dcp_roundtrip', sync=True) as ckpt_dir: + save(mixer.sharded_state_dict(prefix='mixer.', metadata=metadata), ckpt_dir) + + # Scribble over every param so a no-op load cannot pass. + with torch.no_grad(): + for p in mixer.parameters(): + p.data.fill_(float(rank + 1)) + loaded = load(mixer.sharded_state_dict(prefix='mixer.', metadata=metadata), ckpt_dir) + + # in_proj comes back through the 6-way split + the GTP re-pad/re-slice merge_fn. + merged = loaded['mixer.in_proj.weight'] + assert tuple(merged.shape) == tuple(in_proj_w.data.shape), ( + tuple(merged.shape), + tuple(in_proj_w.data.shape), + ) + n_valid = _gdp_valid_rows(in_proj_w, in_proj_dim) + torch.testing.assert_close( + merged[:n_valid].cpu(), golden['in_proj.weight'][:n_valid].cpu(), rtol=0, atol=0 + ) + + # The rest of the mixer must round-trip too -- a colliding replica_id across gtp_remat + # peers would either fail the load or return another rank's data. + for name in ( + 'A_log', + 'dt_bias', + 'conv1d.weight', + 'norm.weight', + 'out_proj.weight', + 'in_proj.layer_norm_weight', + ): + key = f'mixer.{name}' + assert key in loaded, f"{key} missing from the loaded state dict: {sorted(loaded)}" + torch.testing.assert_close( + loaded[key].cpu(), golden[name].cpu(), rtol=0, atol=0, msg=f"{name} drifted" + ) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + GTPShardedParam._chain_state = {} + + +def _worker_gdp_inproj_optim_param_map(rank, world_size, port): + """GDP ``in_proj`` must survive the optimizer id->ShardedTensor match (Muon path, §1.6). + + Same gap as the Mamba case (``_worker_mamba_inproj_optim_param_map``): the model entry for a + gathered+split ``in_proj`` exposes the *gathered* tensor, so it never id-matches the per-shard + GTP optimizer param and ``get_param_id_to_sharded_param_map`` drops it -> KeyError in + ``Float16OptimizerWithFloat16Params.sharded_state_dict``. Unlike that test, this one drives the + real production backfill (``_backfill_gtp_sharded_param_map``) rather than reproducing its + rebuild, so it also pins that GDP takes the per-shard rebuild branch, not the EP refusal. + """ + from megatron.core.dist_checkpointing.optimizer import ( + get_param_id_to_sharded_param_map, + make_sharded_optimizer_tensor, + ) + from megatron.core.optimizer.optimizer import _backfill_gtp_sharded_param_map + from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + tag_gtp_params_with_names, + ) + + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + model_parallel_cuda_manual_seed(42) + update_gtp_config(pad_for_alignment=32) # MXFP8 alignment + mixer, pg, _ = _build_gdp_mixer(['tp', 'cp', 'gtp_remat', 'dp_cp', 'dp_cp_gtp_remat']) + tag_gtp_params_with_names(mixer) # sets _debug_name, mirrors production setup + in_proj_w = mixer.in_proj.weight + + metadata = {'dp_cp_group': pg.dp_cp_gtp_remat} + model_sd = mixer.sharded_state_dict(prefix='mixer.', metadata=metadata) + + # The gap: the gathered+split factory does not id-match the per-shard optimizer param. + id_map = get_param_id_to_sharded_param_map(model_sd, [in_proj_w]) + assert 0 not in id_map, "expected in_proj to be MISSING from the id map (the KeyError gap)" + + # The production backfill must fill it via the per-shard rebuild. An expert-parallel param + # would raise instead; in_proj is dense, so it must rebuild cleanly. + _backfill_gtp_sharded_param_map(id_map, [[in_proj_w]], model_sd) + assert 0 in id_map, "backfill did not restore in_proj" + entry = id_map[0] + # A plain per-shard ShardedTensor keyed by the tagged name -- NOT the model's gathered+split + # factory (reusing that would hand the optimizer the wrong shape). + assert isinstance(entry, ShardedTensor), type(entry) + assert entry is not model_sd['mixer.in_proj.weight'] + assert entry.key == in_proj_w._debug_name, (entry.key, in_proj_w._debug_name) + assert tuple(entry.local_shape) == tuple(in_proj_w.shape), ( + f"rebuilt local_shape {tuple(entry.local_shape)} != param shape " + f"{tuple(in_proj_w.shape)}" + ) + # The rebuilt entry must describe this rank's GTP slice of the PADDED global (what the + # optimizer shard actually is), not the gathered/pad-stripped width the model entry uses. + gtp_remat_rank = torch.distributed.get_rank(in_proj_w.group) + assert entry.global_offset[0] == gtp_remat_rank * in_proj_w.shape[0], ( + entry.global_offset, + gtp_remat_rank, + ) + assert entry.global_shape[0] == in_proj_w.shape[0] * in_proj_w.gtp_remat_size, ( + entry.global_shape, + in_proj_w.shape, + ) + + # make_sharded_optimizer_tensor must accept it for a same-shape optimizer state tensor. + opt_state = torch.zeros_like(in_proj_w) + osh = make_sharded_optimizer_tensor(entry, opt_state, prefix='optimizer.state.exp_avg') + assert osh is not None + assert tuple(osh.local_shape) == tuple(in_proj_w.shape) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + GTPShardedParam._chain_state = {} + + def _worker_save_load_roundtrip_needs_gtp_inclusive_group(rank, world_size, ckpt_base): """Save->load roundtrip: save and load must use the gtp_remat-INCLUSIVE replica group. @@ -1022,6 +1288,18 @@ def test_mamba_inproj_optim_param_map(self): _require_world_size(4) _worker_mamba_inproj_optim_param_map(dist.get_rank(), 4, None) + def test_gdp_inproj_gather_split(self): + _require_world_size(4) + _worker_gdp_inproj_gather_split(dist.get_rank(), 4, None) + + def test_gdp_save_load_roundtrip(self, tmp_path_dist_ckpt): + _require_world_size(4) + _worker_gdp_save_load_roundtrip(dist.get_rank(), 4, tmp_path_dist_ckpt) + + def test_gdp_inproj_optim_param_map(self): + _require_world_size(4) + _worker_gdp_inproj_optim_param_map(dist.get_rank(), 4, None) + def test_replicated_param_needs_gtp_inclusive_dp_cp(self): _require_world_size(4) _worker_replicated_param_needs_gtp_inclusive_dp_cp(dist.get_rank(), 4, None) From 281075d81f54c2788d37a55dffd331b8e5b9f67a Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 13 Aug 2026 09:05:05 +0000 Subject: [PATCH 016/101] feat(fsdp): support expert parallelism in MFSDP v2 adapter (#6450) Signed-off-by: Jingyue Wu Co-authored-by: svcnvidia-nemo-ci --- .../distributed/fsdp/mcore_fsdp_adapter.py | 47 ++++- .../mfsdp_v2/test_mcore_adapter.py | 191 +++++++++++++++++- 2 files changed, 221 insertions(+), 17 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 284beef606e..0b1c8a9d5f2 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -43,6 +43,7 @@ from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.mamba_layer import MambaLayer +from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import MoETransformerLayer, TransformerLayer from megatron.core.utils import is_te_min_version, log_single_rank @@ -560,7 +561,12 @@ def __init__( dp_group = pg_collection.dp_cp device_type = device.type if device is not None else "cuda" - mesh = DeviceMesh.from_group(dp_group, device_type=device_type, mesh_dim_names=("dp",)) + dp_mesh = DeviceMesh.from_group(dp_group, device_type=device_type, mesh_dim_names=("dp",)) + expert_dp_mesh = None + if config.expert_model_parallel_size > 1: + expert_dp_mesh = DeviceMesh.from_group( + pg_collection.expt_dp, device_type=device_type, mesh_dim_names=("expert_dp",) + ) placements = Placements( dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()] ) @@ -569,6 +575,19 @@ def __init__( # ncclCommWindowRegister: # https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html#window-registration with fully_shard_context(device=device, use_symmetric_memory=ddp_config.nccl_ub): + if expert_dp_mesh is not None: + # Expert parameters are replicated over expert-DP, not the full DP group. + # Their gradients need the EP divisor because the same expert receives + # contributions after dispatch from every EP rank. + for submodule in module.modules(): + if isinstance(submodule, MoELayer): + fully_shard( + submodule.experts, + mesh=expert_dp_mesh, + placements=placements, + mixed_precision_policy=self.mp_policy, + grad_divisor=config.expert_model_parallel_size, + ) for submodule in reversed(list(module.modules())): if submodule is module: # The root is always sharded after selected child units so it is not @@ -577,12 +596,12 @@ def __init__( if any(isinstance(submodule, module_type) for module_type in fsdp_unit_modules): fully_shard( submodule, - mesh=mesh, + mesh=dp_mesh, placements=placements, mixed_precision_policy=self.mp_policy, ) fully_shard( - module, mesh=mesh, placements=placements, mixed_precision_policy=self.mp_policy + module, mesh=dp_mesh, placements=placements, mixed_precision_policy=self.mp_policy ) super().__init__(config=config, module=module) @@ -617,7 +636,6 @@ def _validate_config( "tensor_model_parallel_size", "pipeline_model_parallel_size", "context_parallel_size", - "expert_model_parallel_size", ] if any(getattr(config, parallelism) != 1 for parallelism in unsupported_parallelisms): raise ValueError( @@ -630,7 +648,7 @@ def _validate_config( # The config validates the requested topology, while these checks validate the # materialized topology supplied by the caller's process-group collection. - for group_name in ("tp", "pp", "cp", "ep"): + for group_name in ("tp", "pp", "cp"): group = getattr(pg_collection, group_name, None) if group is not None and group.size() != 1: raise ValueError( @@ -638,10 +656,21 @@ def _validate_config( f"got {group.size()}." ) - if getattr(config, "num_moe_experts", None) is not None or any( - not getattr(parameter, "allreduce", True) for parameter in module.parameters() - ): - raise ValueError("MFSDP v2 does not currently support expert parameters.") + if config.expert_model_parallel_size > 1: + if ( + pg_collection.ep is None + or pg_collection.ep.size() != config.expert_model_parallel_size + ): + actual_ep_size = None if pg_collection.ep is None else pg_collection.ep.size() + raise ValueError( + "MFSDP v2 requires an EP process group matching " + f"expert_model_parallel_size={config.expert_model_parallel_size}, " + f"got {actual_ep_size}." + ) + if pg_collection.expt_dp is None: + raise ValueError("MFSDP v2 with EP requires an explicit expert-DP process group.") + if not any(isinstance(submodule, MoELayer) for submodule in module.modules()): + raise ValueError("MFSDP v2 with EP requires MoE transformer layers.") if ddp_config.data_parallel_sharding_strategy != "optim_grads_params": raise ValueError( "MFSDP v2 requires data_parallel_sharding_strategy='optim_grads_params'." diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py index 3892a354540..597a0a7fe56 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py @@ -2,6 +2,8 @@ """MCore adapter and optimizer integration tests for experimental MFSDP v2.""" +import logging +import os from dataclasses import replace import pytest @@ -12,15 +14,20 @@ from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module import FsdpModule from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer from megatron.core.optimizer.fully_sharded_optimizer import FullyShardedOptimizer from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.transformer_block import TransformerBlock from megatron.core.transformer.transformer_config import TransformerConfig -from megatron.core.transformer.transformer_layer import TransformerLayer +from megatron.core.transformer.transformer_layer import MoETransformerLayer, TransformerLayer from tests.unit_tests.test_utilities import Utils +logger = logging.getLogger(__name__) + def _build_layer(config: TransformerConfig) -> TransformerLayer: return TransformerLayer( @@ -37,13 +44,11 @@ def _build_block(config: TransformerConfig) -> TransformerBlock: ) -class TestMcoreAdapter: +class TestMcoreAdapterDense: """Exercise a dense MCore transformer block over two data-parallel ranks.""" def setup_method(self): Utils.initialize_model_parallel(1, 1) - if torch.distributed.get_world_size() < 2: - pytest.skip("MFSDP v2 MCore integration test requires at least two ranks.") self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() model_parallel_cuda_manual_seed(1234) @@ -72,8 +77,6 @@ def test_wraps_fsdp_unit_modules_before_root(self): megatron_fsdp_version=2, use_distributed_optimizer=False, data_parallel_sharding_strategy="optim_grads_params", - megatron_fsdp_main_params_dtype=torch.float32, - megatron_fsdp_main_grads_dtype=torch.float32, ), module=model, fsdp_unit_modules=[TransformerLayer], @@ -160,8 +163,6 @@ def test_build_train_and_step(self): megatron_fsdp_version=2, use_distributed_optimizer=False, data_parallel_sharding_strategy="optim_grads_params", - megatron_fsdp_main_params_dtype=torch.float32, - megatron_fsdp_main_grads_dtype=torch.bfloat16, ), module=model, pg_collection=self.pg_collection, @@ -232,3 +233,177 @@ def test_build_train_and_step(self): assert torch.isfinite(losses).all() assert torch.isfinite(reference_losses).all() torch.testing.assert_close(losses, reference_losses, rtol=1e-3, atol=0) + + +class TestMcoreAdapterExpertParallel: + """Exercise the MFSDP v2 adapter over an MoE model with EP=2.""" + + def setup_method(self): + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + if self.world_size < 2 or self.world_size % 2: + pytest.skip("MFSDP v2 EP adapter test requires an even world size of at least two.") + Utils.initialize_model_parallel(1, 1, expert_model_parallel_size=2) + self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() + assert self.pg_collection.ep.size() == 2 + assert self.pg_collection.expt_dp.size() == self.world_size // 2 + self.reference_group = torch.distributed.new_group( + [torch.distributed.get_rank()], use_local_synchronization=True + ) + self.reference_pg_collection = ProcessGroupCollection( + tp=self.reference_group, + expt_tp=self.reference_group, + cp=self.reference_group, + pp=self.reference_group, + tp_cp=self.reference_group, + tp_dp_cp=self.reference_group, + ep=self.reference_group, + tp_ep=self.reference_group, + expt_dp=self.reference_group, + dp=self.reference_group, + dp_cp=self.reference_group, + embd=None, + pos_embd=None, + ) + model_parallel_cuda_manual_seed(1234) + + def teardown_method(self): + torch.distributed.destroy_process_group(self.reference_group) + Utils.destroy_model_parallel() + + def test_build_train_and_step(self): + """Shard experts over expert-DP and dense parameters over full DP.""" + # The in-process EP=1 reference needs rank-invariant initialization. GPU expert + # initialization instead uses the globally configured EP=2 rank in its RNG seed. + config = TransformerConfig( + num_layers=2, + hidden_size=64, + num_attention_heads=4, + num_moe_experts=4, + expert_model_parallel_size=2, + moe_layer_freq=[0, 1], + moe_token_dispatcher_type="alltoall", + moe_router_topk=2, + moe_grouped_gemm=True, + moe_ffn_hidden_size=128, + add_bias_linear=False, + use_cpu_initialization=True, + params_dtype=torch.float32, + attention_dropout=0.0, + hidden_dropout=0.0, + gradient_accumulation_fusion=False, + attention_backend=AttnBackend.unfused, + ) + # Pair CPU initialization with an explicit common seed for the reference and EP model. + torch.manual_seed(123) + reference_config = replace(config, expert_model_parallel_size=1) + reference_model = HybridModel( + config=reference_config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=128, + max_sequence_length=8, + hybrid_layer_pattern="*E", + pg_collection=self.reference_pg_collection, + ).cuda() + model = HybridModel( + config=config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=128, + max_sequence_length=8, + hybrid_layer_pattern="*E", + pg_collection=self.pg_collection, + ).cuda() + model.load_state_dict(reference_model.state_dict(), strict=False) + for model_layer, reference_layer in zip( + model.decoder.layers, reference_model.decoder.layers + ): + if not isinstance(model_layer, MoETransformerLayer): + continue + for fc in ("linear_fc1", "linear_fc2"): + model_fc = getattr(model_layer.mlp.experts, fc) + reference_fc = getattr(reference_layer.mlp.experts, fc) + for local, global_ in enumerate(model_layer.mlp.local_expert_indices): + for parameter_name in ("weight", "bias"): + model_parameter = getattr(model_fc, f"{parameter_name}{local}", None) + reference_parameter = getattr( + reference_fc, f"{parameter_name}{global_}", None + ) + if model_parameter is not None: + model_parameter.data.copy_(reference_parameter.data) + reference_model.ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=False) + model = FullyShardedDataParallel( + config=config, + ddp_config=DistributedDataParallelConfig( + use_megatron_fsdp=True, + megatron_fsdp_version=2, + use_distributed_optimizer=False, + data_parallel_sharding_strategy="optim_grads_params", + fsdp_all_gather_in_start_param_sync=False, + ), + module=model, + pg_collection=self.pg_collection, + ) + assert isinstance(model.module, FsdpModule) + assert isinstance(model.module.decoder.layers[1].mlp.experts, FsdpModule) + + optimizer_config = OptimizerConfig( + lr=1.0e-3, weight_decay=0.0, use_distributed_optimizer=False, clip_grad=0.0 + ) + reference_optimizer = get_megatron_optimizer( + optimizer_config, [reference_model], use_gloo_process_groups=False + ) + optimizer = get_megatron_optimizer( + replace(optimizer_config), [model], use_gloo_process_groups=False + ) + assert isinstance(optimizer, FullyShardedOptimizer) + optimizer.reload_model_params() + + local_batch_size = 2 + torch.manual_seed(4321) + input_ids = torch.randint(0, 128, (self.world_size * local_batch_size, 8), device="cuda") + position_ids = torch.arange(8, device="cuda").repeat(self.world_size * local_batch_size, 1) + targets = torch.randn(self.world_size * local_batch_size, 8, 128, device="cuda") + input_slice = slice( + torch.distributed.get_rank() * local_batch_size, + (torch.distributed.get_rank() + 1) * local_batch_size, + ) + reference_losses = [] + for _ in range(5): + reference_optimizer.zero_grad(set_to_none=True) + reference_loss = torch.nn.functional.mse_loss( + reference_model( + input_ids=input_ids, position_ids=position_ids, attention_mask=None + ), + targets, + ) + reference_loss.backward() + reference_success, _, _ = reference_optimizer.step() + assert reference_success + reference_losses.append(reference_loss.detach()) + + losses = [] + for _ in range(5): + optimizer.zero_grad(set_to_none=True) + loss = torch.nn.functional.mse_loss( + model( + input_ids=input_ids[input_slice], + position_ids=position_ids[input_slice], + attention_mask=None, + ), + targets[input_slice], + ) + loss.backward() + success, _, _ = optimizer.step() + assert success + loss = loss.detach() + torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.AVG) + losses.append(loss) + + losses = torch.stack(losses) + reference_losses = torch.stack(reference_losses) + if torch.distributed.get_rank() == 0: + logger.info("MFSDP v2 EP loss curve: %s", losses.tolist()) + logger.info("MFSDP v2 EP reference loss curve: %s", reference_losses.tolist()) + assert torch.isfinite(losses).all() + assert torch.isfinite(reference_losses).all() + assert losses[-1] < losses[0] + torch.testing.assert_close(losses, reference_losses) From 2a75ac12c54ba6a42b34024756820bf819a2e25f Mon Sep 17 00:00:00 2001 From: Deyu Fu Date: Thu, 13 Aug 2026 09:09:29 +0000 Subject: [PATCH 017/101] Complete clamped SwiGLU across expert paths (#5940) Signed-off-by: Deyu Fu --- .../core/extensions/transformer_engine.py | 23 ++- megatron/core/fusions/fused_bias_swiglu.py | 116 ++++++++++-- megatron/core/transformer/mlp.py | 2 + megatron/core/transformer/moe/experts.py | 44 +++-- .../core/transformer/moe/shared_experts.py | 32 +++- .../core/transformer/transformer_config.py | 29 ++- .../unit_tests/fusions/test_swiglu_fusion.py | 161 +++++++++++++++++ .../transformer/moe/test_grouped_mlp.py | 102 ++++++++++- .../transformer/moe/test_shared_experts.py | 165 ++++++++++++++++++ 9 files changed, 631 insertions(+), 43 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 1f67b1a5b5b..8dd339f3cef 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2434,12 +2434,23 @@ def __init__( tp_group_for_te = None if is_te_min_version("2.14.0"): - extra_kwargs["single_grouped_weight"] = getattr( - config, "moe_single_grouped_weight", False - ) - extra_kwargs["single_grouped_bias"] = getattr( - config, "moe_single_grouped_bias", False - ) + # Some TE 2.14 builds predate these keyword arguments. Cache the + # installed constructor signature instead of relying on the version alone. + global _TE_GROUPED_LINEAR_INIT_PARAMS + try: + grouped_linear_init_params = _TE_GROUPED_LINEAR_INIT_PARAMS + except NameError: + grouped_linear_init_params = _TE_GROUPED_LINEAR_INIT_PARAMS = set( + inspect.signature(te.pytorch.GroupedLinear.__init__).parameters + ) + if "single_grouped_weight" in grouped_linear_init_params: + extra_kwargs["single_grouped_weight"] = getattr( + config, "moe_single_grouped_weight", False + ) + if "single_grouped_bias" in grouped_linear_init_params: + extra_kwargs["single_grouped_bias"] = getattr( + config, "moe_single_grouped_bias", False + ) self.te_quant_params: Optional[TEQuantizationParams] = None quant_config = get_quant_config_or_none(name, config.quant_recipe) diff --git a/megatron/core/fusions/fused_bias_swiglu.py b/megatron/core/fusions/fused_bias_swiglu.py index 632470876c9..8d64e3016a9 100644 --- a/megatron/core/fusions/fused_bias_swiglu.py +++ b/megatron/core/fusions/fused_bias_swiglu.py @@ -48,6 +48,31 @@ def weighted_swiglu(y, weights): return res.to(dtype) +@jit_fuser +def clamped_swiglu(y, clamp_value): + """Perform SwiGLU after clamping both halves of the input.""" + dtype = y.dtype + y_1, y_2 = torch.chunk(y.to(torch.float32), 2, -1) + y_1 = y_1.clamp(min=None, max=clamp_value) + y_2 = y_2.clamp(min=-clamp_value, max=clamp_value) + res = F.silu(y_1) * y_2 + return res.to(dtype) + + +@jit_fuser +def bias_clamped_swiglu(y, bias, clamp_value): + """Perform clamped SwiGLU after bias addition.""" + return clamped_swiglu(y + bias, clamp_value) + + +@jit_fuser +def clamped_weighted_swiglu(y, weights, clamp_value): + """Perform token-weighted clamped SwiGLU.""" + dtype = y.dtype + res = clamped_swiglu(y, clamp_value) * weights + return res.to(dtype) + + # gradient of tanh approximation of gelu # gradient of actual gelu is: # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) @@ -97,12 +122,50 @@ def weighted_swiglu_back(g, y, weights): return input_grad.to(input_dtype), weights_grad.to(w_dtype) +@jit_fuser +def clamped_swiglu_back(g, y, clamp_value): + """Compute the input gradient for clamped SwiGLU.""" + dtype = y.dtype + y_1, y_2 = torch.chunk(y.to(torch.float32), 2, -1) + y_1c = y_1.clamp(min=None, max=clamp_value) + y_2c = y_2.clamp(min=-clamp_value, max=clamp_value) + res = torch.cat( + ( + g + * torch.sigmoid(y_1c) + * (1 + y_1c * (1 - torch.sigmoid(y_1c))) + * y_2c + * (y_1 <= clamp_value).to(g.dtype), + g * F.silu(y_1c) * ((y_2 >= -clamp_value) & (y_2 <= clamp_value)).to(g.dtype), + ), + -1, + ) + return res.to(dtype) + + +@jit_fuser +def bias_clamped_swiglu_back(g, y, bias, clamp_value): + """Compute the input gradient for clamped SwiGLU with bias.""" + return clamped_swiglu_back(g, y + bias, clamp_value) + + +@jit_fuser +def clamped_weighted_swiglu_back(g, y, weights, clamp_value): + """Compute input and weight gradients for token-weighted clamped SwiGLU.""" + input_dtype = y.dtype + w_dtype = weights.dtype + input_grad = clamped_swiglu_back(g * weights, y, clamp_value) + weights_grad = clamped_swiglu(y, clamp_value) * g.to(w_dtype) + weights_grad = torch.sum(weights_grad, dim=-1, keepdim=True) + return input_grad.to(input_dtype), weights_grad.to(w_dtype) + + class BiasSwiGLUFunction(torch.autograd.Function): """Custom autograd function for SwiGLU activation with bias support.""" @staticmethod @nvtx_decorator() - def forward(ctx, input, bias, fp8_input_store, cpu_offload_input): + def forward(ctx, input, bias, fp8_input_store, cpu_offload_input, clamp_value): """Forward pass of biased SwiGLU activation. Args: @@ -121,6 +184,9 @@ def forward(ctx, input, bias, fp8_input_store, cpu_offload_input): ctx.save_for_backward(input_for_backward, bias) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = clamp_value + if clamp_value is not None and clamp_value > 0: + return bias_clamped_swiglu(input, bias, clamp_value) return bias_swiglu(input, bias) @staticmethod @@ -140,8 +206,11 @@ def backward(ctx, grad_output): """ input, bias = ctx.saved_tensors input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp = bias_swiglu_back(grad_output, input, bias) - return tmp, tmp, None, None + if ctx.clamp_value is not None and ctx.clamp_value > 0: + tmp = bias_clamped_swiglu_back(grad_output, input, bias, ctx.clamp_value) + else: + tmp = bias_swiglu_back(grad_output, input, bias) + return tmp, tmp, None, None, None class SwiGLUFunction(torch.autograd.Function): @@ -149,7 +218,7 @@ class SwiGLUFunction(torch.autograd.Function): @staticmethod @nvtx_decorator() - def forward(ctx, input, fp8_input_store, cpu_offload_input): + def forward(ctx, input, fp8_input_store, cpu_offload_input, clamp_value): """Forward pass of SwiGLU activation. Args: @@ -166,6 +235,9 @@ def forward(ctx, input, fp8_input_store, cpu_offload_input): ctx.save_for_backward(input_for_backward) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = clamp_value + if clamp_value is not None and clamp_value > 0: + return clamped_swiglu(input, clamp_value) return swiglu(input) @staticmethod @@ -184,29 +256,37 @@ def backward(ctx, grad_output): """ input = ctx.saved_tensors[0] input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp = swiglu_back(grad_output, input) - return tmp, None, None + if ctx.clamp_value is not None and ctx.clamp_value > 0: + tmp = clamped_swiglu_back(grad_output, input, ctx.clamp_value) + else: + tmp = swiglu_back(grad_output, input) + return tmp, None, None, None class WeightedSwiGLUFunction(torch.autograd.Function): @staticmethod - # bias is an optional argument - def forward(ctx, input, weights, fp8_input_store): + def forward(ctx, input, weights, fp8_input_store, clamp_value): input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input ctx.save_for_backward(input_for_backward, weights) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = clamp_value + if clamp_value is not None and clamp_value > 0: + return clamped_weighted_swiglu(input, weights, clamp_value) return weighted_swiglu(input, weights) @staticmethod def backward(ctx, grad_output): input, weights = ctx.saved_tensors input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp, wgrad = weighted_swiglu_back(grad_output, input, weights) - return tmp, wgrad, None + if ctx.clamp_value is not None and ctx.clamp_value > 0: + tmp, wgrad = clamped_weighted_swiglu_back(grad_output, input, weights, ctx.clamp_value) + else: + tmp, wgrad = weighted_swiglu_back(grad_output, input, weights) + return tmp, wgrad, None, None -def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False): +def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False, clamp_value=None): """Implementation of biased SwiGLU that handles different input shapes. This function reshapes the input if necessary, applies the SwiGLU activation @@ -218,6 +298,10 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False uses the bias-free SwiGLU variant. fp8_input_store (bool, optional): Whether to store intermediate values in FP8 format. Defaults to False. + cpu_offload_input (bool, optional): Whether to mark saved activation inputs for CPU + offloading. Defaults to False. + clamp_value (float, optional): Maximum gate value and absolute linear value. When None, + preserve the legacy unclamped SwiGLU behavior. Returns: torch.Tensor: Result of biased SwiGLU activation. @@ -229,14 +313,16 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False assert len(ori_shape) in [2, 3] input = input.view(-1, ori_shape[-1]) if bias is not None: - output = BiasSwiGLUFunction.apply(input, bias, fp8_input_store, cpu_offload_input) + output = BiasSwiGLUFunction.apply( + input, bias, fp8_input_store, cpu_offload_input, clamp_value + ) else: - output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input) + output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input, clamp_value) return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) -def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False): +def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False, clamp_value=None): """ Token-wise-weighted bias swiglu fusion. """ @@ -246,7 +332,7 @@ def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False): if bias is not None: raise NotImplementedError("Bias is not supported for weighted swiglu fusion") else: - output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store) + output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store, clamp_value) return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index cb3b1b0be82..66a69ff3f37 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -281,6 +281,7 @@ def forward( bias_parallel, per_token_scale.unsqueeze(-1), self.config.activation_func_fp8_input_store, + self.config.activation_func_clamp_value, ) elif self.activation_func == quick_gelu and self.config.gated_linear_unit: intermediate_parallel = weighted_bias_quick_geglu_impl( @@ -312,6 +313,7 @@ def forward( self.config.cpu_offloading and self.config.cpu_offloading_activations and HAVE_TE, + self.config.activation_func_clamp_value, ) else: raise ValueError("Only support fusion of gelu and swiglu") diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 59f8deffeca..35817f05e6a 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -377,7 +377,13 @@ def _is_fused_impl_supported(self) -> bool: if not (use_glu_fusion or use_srelu_fusion): return False if self.config.activation_func == F.silu: - pass + if self.config.activation_func_clamp_value is not None: + if not is_te_min_version("2.17.0.dev0"): + return False + try: + from transformer_engine.pytorch.ops import ScaledClampedQGeGLU # noqa: F401 + except ImportError: + return False elif self.config.activation_func == quick_gelu: try: from transformer_engine.pytorch.ops import ScaledClampedQGeGLU # noqa: F401 @@ -479,20 +485,35 @@ def register_grouped_linear_params( ) ops.append(op) - # Activation and post-multiply probs (SwiGLU, clamped quick-GeGLU, or SReLU) + # Activation and post-multiply probs (SwiGLU, clamped GLU, or SReLU). glu_interleave = self.config.moe_mlp_glu_interleave_size activation_recompute_in_mlp = bool(getattr(self, "activation_recompute", False)) if self.config.activation_func == F.silu and self.config.gated_linear_unit: - if ( - "activation_recompute_in_mlp" - in inspect.signature(te.pytorch.ops.ScaledSwiGLU).parameters - ): - op = te.pytorch.ops.ScaledSwiGLU( - glu_interleave_size=glu_interleave, - activation_recompute_in_mlp=activation_recompute_in_mlp, - ) + clamp_value = self.config.activation_func_clamp_value + if clamp_value is not None: + clamped_glu_kwargs = { + "glu_interleave_size": glu_interleave, + "alpha": 1.0, + "limit": clamp_value, + "glu_linear_offset": 0.0, + } + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledClampedQGeGLU).parameters + ): + clamped_glu_kwargs["activation_recompute_in_mlp"] = activation_recompute_in_mlp + op = te.pytorch.ops.ScaledClampedQGeGLU(**clamped_glu_kwargs) else: - op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave) + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledSwiGLU).parameters + ): + op = te.pytorch.ops.ScaledSwiGLU( + glu_interleave_size=glu_interleave, + activation_recompute_in_mlp=activation_recompute_in_mlp, + ) + else: + op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave) elif self.config.activation_func == quick_gelu and self.config.gated_linear_unit: clamp = self.config.activation_func_clamp_value if clamp is not None: @@ -832,6 +853,7 @@ def bias_act_func(intermediate_parallel, bias_parallel, permuted_probs): bias_parallel, permuted_probs, self.config.activation_func_fp8_input_store, + self.config.activation_func_clamp_value, ) elif self.activation_func == quick_gelu and self.config.gated_linear_unit: intermediate_parallel = weighted_bias_quick_geglu_impl( diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 038a162f899..6eef5dee9ff 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/megatron/core/transformer/moe/shared_experts.py @@ -287,6 +287,7 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None): intermediate_parallel, bias_parallel, self.config.activation_func_fp8_input_store, + clamp_value=self.config.activation_func_clamp_value, ) else: raise ValueError("Only support fusion of gelu and swiglu") @@ -296,8 +297,13 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None): if self.config.gated_linear_unit: def glu(x): - x = torch.chunk(x, 2, dim=-1) - return self.config.activation_func(x[0]) * x[1] + x_glu, x_linear = torch.chunk(x, 2, dim=-1) + if (clamp_value := self.config.activation_func_clamp_value) is not None: + x_glu = x_glu.clamp(min=None, max=clamp_value) + x_linear = x_linear.clamp(min=-clamp_value, max=clamp_value) + return self.config.activation_func(x_glu) * ( + x_linear + self.config.glu_linear_offset + ) intermediate_parallel = glu(intermediate_parallel) else: @@ -416,6 +422,15 @@ def _validate_fused_grouped_swiglu(self) -> None: f"fused kernel, but got activation_func={self.config.activation_func}, " f"gated_linear_unit={self.config.gated_linear_unit}." ) + if self.config.activation_func_clamp_value is not None and ( + not is_te_min_version("2.17.0.dev0") + or not hasattr(te.pytorch.ops, "ScaledClampedQGeGLU") + ): + raise RuntimeError( + f"{self.__class__.__name__} requires Transformer Engine >= 2.17.0.dev0 " + "with pytorch.ops.ScaledClampedQGeGLU when " + "activation_func_clamp_value is set." + ) if self.config.moe_shared_expert_glu_interleave_size is None: raise ValueError( f"{self.__class__.__name__} requires " @@ -452,7 +467,7 @@ def _get_fused_grouped_swiglu_recipe(self): return self._fused_grouped_swiglu_recipe def _make_fused_grouped_swiglu_ops(self) -> torch.nn.Module: - """Construct GroupedLinear(num_groups=1) -> ScaledSwiGLU -> GroupedLinear.""" + """Construct the grouped-linear shared-expert MLP operations.""" ops = te.pytorch.ops.Sequential() tp_world_size = get_pg_size(self.tp_group) rng_state_tracker_function = None @@ -475,7 +490,16 @@ def _make_fused_grouped_swiglu_ops(self) -> torch.nn.Module: op._glu_interleave_size = glu_interleave_size ops.append(op) - activation_op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + clamp_value = self.config.activation_func_clamp_value + if clamp_value is None: + activation_op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + else: + activation_op = te.pytorch.ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave_size, + alpha=1.0, + limit=clamp_value, + glu_linear_offset=0.0, + ) # Shared experts are not router-gated. Mark this fused-op instance so # TE can omit the optional forward cuDNN probability tensor without # changing the semantics of routed single-group MLPs. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index aa8af1b5c18..2a35f8b585f 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -227,7 +227,7 @@ class TransformerConfig(ModelParallelConfig): activation_func_clamp_value: Optional[float] = None """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func - is quick_gelu.""" + is quick_gelu or SwiGLU (MoE only).""" num_moe_experts: Optional[int] = None """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None @@ -2227,6 +2227,33 @@ def __post_init__(self): if self.activation_func != F.silu or not self.gated_linear_unit: raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.") + if ( + self.activation_func_clamp_value is not None + and self.activation_func == F.silu + and self.gated_linear_unit + ): + if ( + not math.isfinite(self.activation_func_clamp_value) + or self.activation_func_clamp_value <= 0 + ): + raise ValueError( + "activation_func_clamp_value for SwiGLU must be finite and greater than zero." + ) + if self.num_moe_experts is None: + raise ValueError( + "activation_func_clamp_value for SwiGLU is only supported with MoE." + ) + if self.glu_linear_offset != 0.0: + raise ValueError( + "glu_linear_offset must be zero when activation_func_clamp_value " + "is set for SwiGLU." + ) + if self.use_te_activation_func: + raise ValueError( + "use_te_activation_func must be False " + "when activation_func_clamp_value is not None for SwiGLU" + ) + if self.apply_rope_fusion: if self.multi_latent_attention: warnings.warn( diff --git a/tests/unit_tests/fusions/test_swiglu_fusion.py b/tests/unit_tests/fusions/test_swiglu_fusion.py index c72679cd047..1ba9b3eb891 100644 --- a/tests/unit_tests/fusions/test_swiglu_fusion.py +++ b/tests/unit_tests/fusions/test_swiglu_fusion.py @@ -1,7 +1,51 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import pytest import torch +import torch.nn.functional as F from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl, weighted_bias_swiglu_impl +from megatron.core.transformer.transformer_config import TransformerConfig + + +def _clamped_swiglu_config(**kwargs): + defaults = dict( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + gated_linear_unit=True, + activation_func=F.silu, + activation_func_clamp_value=10.0, + ) + return TransformerConfig(**(defaults | kwargs)) + + +def test_clamped_swiglu_config_accepts_positive_moe_clamp(): + assert _clamped_swiglu_config().activation_func_clamp_value == 10.0 + + +@pytest.mark.parametrize("clamp_value", [0.0, -1.0, float("nan"), float("inf"), float("-inf")]) +def test_clamped_swiglu_config_requires_positive_clamp(clamp_value): + with pytest.raises(ValueError, match="greater than zero"): + _clamped_swiglu_config(activation_func_clamp_value=clamp_value) + + +def test_clamped_swiglu_config_rejects_linear_offset(): + with pytest.raises(ValueError, match="glu_linear_offset must be zero"): + _clamped_swiglu_config(glu_linear_offset=1.0) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"num_moe_experts": None}, "only supported with MoE"), + ({"use_te_activation_func": True}, "use_te_activation_func must be False"), + ], +) +def test_clamped_swiglu_config_rejects_unsupported_paths(kwargs, match): + with pytest.raises(ValueError, match=match): + _clamped_swiglu_config(**kwargs) @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) @@ -39,3 +83,120 @@ def test_weighted_bias_swiglu(input_dtype): assert weights_2.grad.dtype == weights.grad.dtype if input_dtype == torch.float32: assert torch.allclose(weights.grad, weights_2.grad, **tols) + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) +def test_clamped_weighted_bias_swiglu(input_dtype): + clamp_value = 10.0 + + if input_dtype == torch.float32: + tols = dict(rtol=1.0e-6, atol=1.0e-6) + elif input_dtype == torch.bfloat16: + tols = dict(rtol=2.0e-2, atol=1.0e-3) + else: + raise ValueError(f"Invalid input dtype: {input_dtype}") + + x = (torch.randn(16, 64, dtype=input_dtype, device="cuda") * 5.0).requires_grad_(True) + weights = torch.randn(16, 1, dtype=torch.float32, device="cuda", requires_grad=True) + bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda") + + # Reference: clamp and activate in FP32, then restore the input dtype. + y_1, y_2 = torch.chunk(x.to(torch.float32), 2, -1) + y = ( + F.silu(y_1.clamp(min=None, max=clamp_value)) + * y_2.clamp(min=-clamp_value, max=clamp_value) + * weights + ).to(input_dtype) + y.backward(bwd_input) + + x_fused = x.detach().clone().requires_grad_(True) + weights_fused = weights.detach().clone().requires_grad_(True) + y_fused = weighted_bias_swiglu_impl(x_fused, None, weights_fused, clamp_value=clamp_value) + y_fused.backward(bwd_input.detach().clone()) + + assert y_fused.dtype == y.dtype + assert torch.allclose(y, y_fused, **tols) + assert x_fused.grad.dtype == x.grad.dtype + assert torch.allclose(x.grad, x_fused.grad, **tols) + assert weights_fused.grad.dtype == weights.grad.dtype + if input_dtype == torch.float32: + assert torch.allclose(weights.grad, weights_fused.grad, **tols) + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_clamped_bias_swiglu_impl(input_dtype, with_bias): + clamp_value = 10.0 + + if input_dtype == torch.float32: + tols = dict(rtol=1.0e-6, atol=1.0e-6) + elif input_dtype == torch.bfloat16: + tols = dict(rtol=2.0e-2, atol=1.0e-3) + else: + raise ValueError(f"Invalid input dtype: {input_dtype}") + + x = (torch.randn(16, 64, dtype=input_dtype, device="cuda") * 5.0).requires_grad_(True) + bias = ( + torch.randn(64, dtype=input_dtype, device="cuda").requires_grad_(True) + if with_bias + else None + ) + bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda") + + x_fp32 = x.to(torch.float32) + x_effective = x_fp32 + bias.to(torch.float32) if with_bias else x_fp32 + y_1, y_2 = torch.chunk(x_effective, 2, -1) + y = ( + F.silu(y_1.clamp(min=None, max=clamp_value)) * y_2.clamp(min=-clamp_value, max=clamp_value) + ).to(input_dtype) + y.backward(bwd_input) + + x_fused = x.detach().clone().requires_grad_(True) + bias_fused = bias.detach().clone().requires_grad_(True) if with_bias else None + y_fused = bias_swiglu_impl(x_fused, bias_fused, clamp_value=clamp_value) + y_fused.backward(bwd_input.detach().clone()) + + assert y_fused.dtype == y.dtype + assert torch.allclose(y, y_fused, **tols) + assert x_fused.grad.dtype == x.grad.dtype + assert torch.allclose(x.grad, x_fused.grad, **tols) + if with_bias: + assert bias_fused.grad.dtype == bias.grad.dtype + bias_grad_cos = F.cosine_similarity( + bias.grad.flatten().float().unsqueeze(0), bias_fused.grad.flatten().float().unsqueeze(0) + ).item() + assert bias_grad_cos > 0.999, f"bias.grad cosine similarity = {bias_grad_cos:.6f}" + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_bias_swiglu_impl_clamp_none_matches_unclamped(input_dtype, with_bias): + if input_dtype == torch.float32: + tols = dict(rtol=1.0e-6, atol=1.0e-6) + else: + tols = dict(rtol=2.0e-2, atol=1.0e-3) + + x = torch.randn(16, 64, dtype=input_dtype, device="cuda").requires_grad_(True) + bias = ( + torch.randn(64, dtype=input_dtype, device="cuda").requires_grad_(True) + if with_bias + else None + ) + bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda") + + y = bias_swiglu_impl(x, bias) + y.backward(bwd_input) + + x_explicit = x.detach().clone().requires_grad_(True) + bias_explicit = bias.detach().clone().requires_grad_(True) if with_bias else None + y_explicit = bias_swiglu_impl(x_explicit, bias_explicit, clamp_value=None) + y_explicit.backward(bwd_input.detach().clone()) + + assert torch.allclose(y, y_explicit, **tols) + assert torch.allclose(x.grad, x_explicit.grad, **tols) + if with_bias: + bias_grad_cos = F.cosine_similarity( + bias.grad.flatten().float().unsqueeze(0), + bias_explicit.grad.flatten().float().unsqueeze(0), + ).item() + assert bias_grad_cos > 0.999, f"bias.grad cosine similarity = {bias_grad_cos:.6f}" diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index b9e7fa346d2..6b3531b090e 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -38,6 +38,23 @@ def test_op_fuser_transformer_config_args_are_exposed(): assert args.moe_mlp_glu_interleave_size == 16 +def test_clamped_swiglu_allows_te_op_fuser(): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_grouped_gemm=True, + gated_linear_unit=True, + activation_func=F.silu, + activation_func_clamp_value=10.0, + use_transformer_engine_op_fuser=True, + ) + + assert config.activation_func_clamp_value == 10.0 + assert config.use_transformer_engine_op_fuser is True + + def test_remove_glu_interleaving_restores_contiguous_gate_and_linear_halves(): interleaved = torch.tensor([[1, 2, 5, 6, 3, 4, 7, 8], [11, 12, 15, 16, 13, 14, 17, 18]]) expected = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [11, 12, 13, 14, 15, 16, 17, 18]]) @@ -452,11 +469,21 @@ def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): self.activation_recompute_in_mlp = activation_recompute_in_mlp class FakeScaledClampedQGeGLU(torch.nn.Module): - def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False, limit=None): + def __init__( + self, + glu_interleave_size, + *, + activation_recompute_in_mlp=False, + limit=None, + alpha=1.702, + glu_linear_offset=1.0, + ): super().__init__() self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp self.limit = limit + self.alpha = alpha + self.glu_linear_offset = glu_linear_offset class FakeScaledSReLU(torch.nn.Module): def __init__(self, *, activation_recompute_in_mlp=False): @@ -484,8 +511,15 @@ def register_forward_pre_hook(self, hook): ) -def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): - """quick_gelu + clamp value → ScaledClampedQGeGLU(limit=clamp).""" +@pytest.mark.parametrize( + ("activation_func", "expected_alpha", "expected_offset"), + [(quick_gelu, 1.702, 1.0), (F.silu, 1.0, 0.0)], + ids=("quick-geglu", "clamped-swiglu"), +) +def test_make_fused_ops_uses_clamped_qgeglu( + monkeypatch, activation_func, expected_alpha, expected_offset +): + """Clamped quick GeGLU and SwiGLU use the appropriate TE parameters.""" fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -495,10 +529,10 @@ def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): moe_mlp_glu_interleave_size=4, delay_wgrad_compute=False, activation_func_clamp_value=7.0, - activation_func=quick_gelu, + activation_func=activation_func, gated_linear_unit=True, ) - module.activation_func = quick_gelu + module.activation_func = activation_func module.activation_recompute = True common = dict( device="cuda", @@ -520,6 +554,8 @@ def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): assert activation.glu_interleave_size == 4 assert activation.activation_recompute_in_mlp is True assert activation.limit == 7.0 + assert activation.alpha == expected_alpha + assert activation.glu_linear_offset == expected_offset def test_make_fused_ops_uses_scaled_srelu_for_weighted_squared_relu(monkeypatch): @@ -610,12 +646,18 @@ def _install_fake_te_ops_modules( def _make_fused_impl_support_module( - FakeGroupedLinear, *, activation_func, gated_linear_unit, use_fused_weighted_squared_relu=False + FakeGroupedLinear, + *, + activation_func, + gated_linear_unit, + use_fused_weighted_squared_relu=False, + activation_func_clamp_value=None, ): module = TEGroupedMLP.__new__(TEGroupedMLP) torch.nn.Module.__init__(module) module.config = SimpleNamespace( activation_func=activation_func, + activation_func_clamp_value=activation_func_clamp_value, gated_linear_unit=gated_linear_unit, use_fused_weighted_squared_relu=use_fused_weighted_squared_relu, moe_apply_probs_on_input=False, @@ -649,6 +691,27 @@ def test_is_fused_impl_supported_uses_config_activation_for_swiglu(monkeypatch): assert module._is_fused_impl_supported() is True +@pytest.mark.parametrize("activation_func_clamp_value", [None, 10.0], ids=("unclamped", "clamped")) +def test_is_fused_impl_supported_requires_cutedsl_for_swiglu( + monkeypatch, activation_func_clamp_value +): + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + monkeypatch.setattr(experts_module, "HAVE_TE", True) + monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True) + monkeypatch.delenv("NVTE_CUTEDSL_FUSED_GROUPED_MLP", raising=False) + _install_fake_te_ops_modules(monkeypatch, fake_te) + + module = _make_fused_impl_support_module( + FakeGroupedLinear, + activation_func=F.silu, + gated_linear_unit=True, + activation_func_clamp_value=activation_func_clamp_value, + ) + + assert module._is_fused_impl_supported() is False + + def test_is_fused_impl_supported_requires_scaled_fc2_bias(monkeypatch): fake_te, FakeGroupedLinear = _make_fake_te_namespace() @@ -671,6 +734,33 @@ def __init__(self, *args, **kwargs): assert module._is_fused_impl_supported() is False +@pytest.mark.parametrize( + ("te_217_or_later", "include_clamped_qgeglu", "expected"), + [(True, True, True), (False, True, False), (True, False, False)], +) +def test_is_fused_impl_supported_gates_clamped_swiglu( + monkeypatch, te_217_or_later, include_clamped_qgeglu, expected +): + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + monkeypatch.setattr(experts_module, "HAVE_TE", True) + monkeypatch.setattr( + experts_module, "is_te_min_version", lambda version: version == "2.14.0" or te_217_or_later + ) + _install_fake_te_ops_modules( + monkeypatch, fake_te, include_clamped_qgeglu=include_clamped_qgeglu + ) + + module = _make_fused_impl_support_module( + FakeGroupedLinear, + activation_func=F.silu, + gated_linear_unit=True, + activation_func_clamp_value=10.0, + ) + + assert module._is_fused_impl_supported() is expected + + @pytest.mark.parametrize( ("use_fused_weighted_squared_relu", "gated_linear_unit", "expected"), [(True, False, True), (False, False, False), (True, True, False)], diff --git a/tests/unit_tests/transformer/moe/test_shared_experts.py b/tests/unit_tests/transformer/moe/test_shared_experts.py index 798313f188c..c71f665259f 100644 --- a/tests/unit_tests/transformer/moe/test_shared_experts.py +++ b/tests/unit_tests/transformer/moe/test_shared_experts.py @@ -52,6 +52,15 @@ def __init__(self, glu_interleave_size): self.glu_interleave_size = glu_interleave_size +class _FakeTEScaledClampedQGeGLU(torch.nn.Module): + def __init__(self, glu_interleave_size, *, alpha, limit, glu_linear_offset): + super().__init__() + self.glu_interleave_size = glu_interleave_size + self.alpha = alpha + self.limit = limit + self.glu_linear_offset = glu_linear_offset + + class _FakeTESequential(torch.nn.Module): def append(self, module): self.add_module(str(len(self._modules)), module) @@ -90,6 +99,7 @@ def _fake_te_module(linear_cls=_FakeTELinear): ops=SimpleNamespace( GroupedLinear=_FakeTEGroupedLinear, ScaledSwiGLU=_FakeTEScaledSwiGLU, + ScaledClampedQGeGLU=_FakeTEScaledClampedQGeGLU, Sequential=_FakeTESequential, ), fp8_autocast=_FakeFP8Autocast, @@ -123,6 +133,7 @@ def _fake_shared_expert(**config_kwargs): add_bias_linear=False, gated_linear_unit=True, activation_func=F.silu, + activation_func_clamp_value=None, moe_shared_expert_glu_interleave_size=32, delay_wgrad_compute=False, sequence_parallel=False, @@ -185,6 +196,21 @@ def test_validate_fused_grouped_swiglu_requires_te(monkeypatch): shared_expert._validate_fused_grouped_swiglu() +@pytest.mark.parametrize("has_clamped_op", [True, False]) +def test_validate_fused_grouped_swiglu_requires_clamped_te_support(monkeypatch, has_clamped_op): + fake_te = _patch_fake_shared_expert_te(monkeypatch) + if has_clamped_op: + monkeypatch.setattr( + shared_experts_module, "is_te_min_version", lambda version: version != "2.17.0.dev0" + ) + else: + del fake_te.pytorch.ops.ScaledClampedQGeGLU + shared_expert = _fake_shared_expert(activation_func_clamp_value=7.0) + + with pytest.raises(RuntimeError, match="ScaledClampedQGeGLU"): + shared_expert._validate_fused_grouped_swiglu() + + @pytest.mark.parametrize( ("config_kwargs", "bad_linear", "match"), [ @@ -241,6 +267,22 @@ def test_make_fused_grouped_swiglu_ops_builds_grouped_pipeline(monkeypatch): assert fc2_op.weight0 is shared_expert.linear_fc2.weight +def test_make_fused_grouped_swiglu_ops_builds_clamped_activation(monkeypatch): + _patch_fake_shared_expert_te(monkeypatch) + shared_expert = _fake_shared_expert(activation_func_clamp_value=7.0) + + shared_expert._validate_fused_grouped_swiglu() + ops = shared_expert._make_fused_grouped_swiglu_ops() + + activation_op = list(ops.children())[1] + assert isinstance(activation_op, _FakeTEScaledClampedQGeGLU) + assert activation_op.glu_interleave_size == 32 + assert activation_op.alpha == 1.0 + assert activation_op.limit == 7.0 + assert activation_op.glu_linear_offset == 0.0 + assert activation_op._grouped_mlp_unit_activation_scale is True + + def test_fused_grouped_swiglu_ops_replay_linear_pre_forward_hooks(monkeypatch): _patch_fake_shared_expert_te(monkeypatch) shared_expert = _fake_shared_expert() @@ -415,3 +457,126 @@ def test_shared_expert_forward_backward(self, dispatcher_type: str, tp_size, ep_ assert torch.allclose( p_overlap.grad, p_no_overlap.grad ), f"max diff: {torch.max(torch.abs(p_overlap.grad - p_no_overlap.grad))}" + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_shared_expert_glu_linear_offset_overlap_parity(self): + """Nonzero GLU offsets produce identical overlap and synchronous results.""" + Utils.initialize_model_parallel(tensor_model_parallel_size=1, expert_model_parallel_size=1) + shared_expert_kwargs = { + "moe_token_dispatcher_type": "alltoall", + "bias_activation_fusion": False, + "activation_func_clamp_value": None, + "glu_linear_offset": 0.5, + } + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_overlap = self.get_moe_layer( + moe_shared_expert_overlap=True, **shared_expert_kwargs + ).to(dtype=torch.bfloat16) + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_no_overlap = self.get_moe_layer( + moe_shared_expert_overlap=False, **shared_expert_kwargs + ).to(dtype=torch.bfloat16) + moe_layer_no_overlap.load_state_dict(moe_layer_overlap.state_dict()) + + hidden_states = torch.randn( + (32, 2, self.config.hidden_size), + requires_grad=True, + device="cuda", + dtype=torch.bfloat16, + ) + hidden_states_no_overlap = hidden_states.detach().clone().requires_grad_(True) + + shared_expert_overlap = moe_layer_overlap.shared_experts + shared_expert_no_overlap = moe_layer_no_overlap.shared_experts + assert shared_expert_overlap is not None + assert shared_expert_no_overlap is not None + + # Isolate shared-expert parity from nondeterministic routed-token unpermutation. + shared_expert_overlap.pre_forward_comm(hidden_states) + shared_expert_overlap.linear_fc1_forward_and_act() + shared_expert_overlap.linear_fc2_forward() + shared_expert_overlap.post_forward_comm() + output_overlap = shared_expert_overlap.get_output() + output_no_overlap = shared_expert_no_overlap(hidden_states_no_overlap) + torch.testing.assert_close(output_overlap, output_no_overlap) + + output_overlap.mean().backward() + output_no_overlap.mean().backward() + + torch.testing.assert_close(hidden_states.grad, hidden_states_no_overlap.grad) + for p_overlap, p_no_overlap in zip( + shared_expert_overlap.parameters(), shared_expert_no_overlap.parameters() + ): + torch.testing.assert_close(p_overlap.grad, p_no_overlap.grad) + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("bias_activation_fusion", [True, False]) + def test_shared_expert_clamped_swiglu(self, bias_activation_fusion): + """Verify clamped SwiGLU parity for overlapped and synchronous shared experts.""" + Utils.initialize_model_parallel(tensor_model_parallel_size=1, expert_model_parallel_size=1) + clamp_value = 1.0 + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_overlap = self.get_moe_layer( + moe_shared_expert_overlap=True, + moe_token_dispatcher_type="alltoall", + activation_func_clamp_value=clamp_value, + bias_activation_fusion=bias_activation_fusion, + ).to(dtype=torch.bfloat16) + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_no_overlap = self.get_moe_layer( + moe_shared_expert_overlap=False, + moe_token_dispatcher_type="alltoall", + activation_func_clamp_value=clamp_value, + bias_activation_fusion=bias_activation_fusion, + ).to(dtype=torch.bfloat16) + moe_layer_no_overlap.load_state_dict(moe_layer_overlap.state_dict()) + + hidden_states = ( + torch.randn((32, 2, self.config.hidden_size), device="cuda", dtype=torch.bfloat16) * 5.0 + ).requires_grad_(True) + hidden_states_no_overlap = hidden_states.detach().clone().requires_grad_(True) + + output_overlap, _ = moe_layer_overlap(hidden_states) + output_no_overlap, _ = moe_layer_no_overlap(hidden_states_no_overlap) + + cos_out = F.cosine_similarity( + output_overlap.flatten().unsqueeze(0).float(), + output_no_overlap.flatten().unsqueeze(0).float(), + ).item() + assert cos_out > 0.999, ( + f"shared-expert clamp output mismatch (fusion={bias_activation_fusion}): " + f"cos sim = {cos_out:.6f}" + ) + + output_overlap.mean().backward() + output_no_overlap.mean().backward() + + for p_overlap, p_no_overlap in zip( + moe_layer_overlap.parameters(), moe_layer_no_overlap.parameters() + ): + assert torch.allclose(p_overlap.grad, p_no_overlap.grad), ( + f"shared-expert clamp mismatch (fusion={bias_activation_fusion}); " + f"max diff: {torch.max(torch.abs(p_overlap.grad - p_no_overlap.grad))}" + ) + + _set_random_seed(seed_=123, data_parallel_random_init=False) + moe_layer_unclamped = self.get_moe_layer( + moe_shared_expert_overlap=False, + moe_token_dispatcher_type="alltoall", + activation_func_clamp_value=None, + bias_activation_fusion=bias_activation_fusion, + ).to(dtype=torch.bfloat16) + moe_layer_unclamped.load_state_dict(moe_layer_overlap.state_dict()) + + hidden_states_unclamped = hidden_states.detach().clone().requires_grad_(True) + output_unclamped, _ = moe_layer_unclamped(hidden_states_unclamped) + assert not torch.allclose(output_no_overlap, output_unclamped), ( + "Clamping had no observable effect on shared-expert output; " + "activation_func_clamp_value may not be plumbed through." + ) From 87bf52075b05a0155ba454dc8de136ee448d7cd5 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 13 Aug 2026 12:18:27 +0000 Subject: [PATCH 018/101] RL: unify rollout generation into one pipeline (#6495) Signed-off-by: Teodor-Dumitru Ene --- megatron/rl/agent/api.py | 369 +------------ megatron/rl/agent/rollout_pipeline.py | 506 ++++++++++++++++++ megatron/rl/agent/weighted_multi_task.py | 154 ++---- megatron/rl/inference/megatron.py | 3 +- megatron/rl/rl_utils.py | 251 +++++---- megatron/rl/rollout_granularity.py | 11 +- tests/unit_tests/rl/test_rl_utils.py | 87 ++- .../unit_tests/rl/test_rollout_generation.py | 269 +++++++--- 8 files changed, 923 insertions(+), 727 deletions(-) create mode 100644 megatron/rl/agent/rollout_pipeline.py diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index ffa38be06bc..b4ffb35048f 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -1,23 +1,16 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -import asyncio -import time from abc import ABC, abstractmethod -from typing import AsyncIterator, Awaitable, Callable, Generic, NamedTuple, TypeVar +from typing import Awaitable, Callable, Generic, NamedTuple, TypeVar -import numpy as np from pydantic import BaseModel -from megatron.core.inference.utils import asyncio_Queue, asyncio_QueueShutDown -from megatron.core.utils import trace_async_exceptions - from ..__init__ import Request, TypeLookupable from ..inference import ( InferenceInterface, InferenceRequest, InferenceResponse, LLMChatMessage, - ReturnsRaw, ) from ..rollout_granularity import ConsumptionGranularity, SubmissionGranularity @@ -200,359 +193,31 @@ class TokenizedRolloutGenerator(Agent, ABC): async def get_reward_rollouts(self, request: RolloutRequest) -> list[TokenRollout]: ... -class _GranularityConfig(NamedTuple): - submission: SubmissionGranularity - consumption: ConsumptionGranularity - num_groups_per_batch: int - - @classmethod - def from_request(cls, request: GroupedRolloutRequest) -> "_GranularityConfig": - cls._validate(request) - return cls( - submission=request.submission_granularity, - consumption=request.consumption_granularity, - num_groups_per_batch=request.num_groups, - ) - - @property - def prevent_dataset_reorder(self) -> bool: - return self.consumption == "B" - - @staticmethod - def _validate(request: GroupedRolloutRequest) -> None: - assert not ( - request.submission_granularity == "B" and request.consumption_granularity == "G" - ), "Batch submission with group consumption is not supported." - assert not request.filter_groups_with_same_reward, ( - "filter_groups_with_same_reward is not currently supported: dropped groups " - "are not regenerated, so non-streaming callers receive fewer groups than " - "requested and batch-order consumers stall on incomplete batches." - ) - - -class _SubmissionGate: - """Gate capacity is measured in units of the configured submission granularity. - - Each granularity has a single release point: R slots free when inference - completes, so the gate bounds engine concurrency in rollouts. G and B - slots free when the trainer consumes the group/batch, so the gate - enforces the --rl-generation-lag run-ahead cap in groups/batches - respectively. - """ - - def __init__( - self, - *, - capacity: int, - submission: SubmissionGranularity, - ) -> None: - self._sem = asyncio.Semaphore(capacity) - self._submission = submission - self.capacity = capacity - # Observability counters, updated only on the configured submission - # granularity (the only path that touches the semaphore). `held` - # counts slots currently held; `prepare_blocked_seconds` accumulates - # time stage_prepare spent waiting on the semaphore. - self.held = 0 - self.prepare_blocked_seconds = 0.0 - self.acquire_calls = 0 - self.release_calls = 0 - - async def acquire_for(self, granularity: SubmissionGranularity) -> None: - if self._submission == granularity: - start = time.monotonic() - await self._sem.acquire() - self.prepare_blocked_seconds += time.monotonic() - start - self.held += 1 - self.acquire_calls += 1 - - def release_for(self, granularity: SubmissionGranularity) -> None: - if self._submission == granularity: - self._sem.release() - self.held -= 1 - self.release_calls += 1 - - -class _InferWorkItem(NamedTuple): - """One rollout's worth of work flowing from prepare to infer. - - Timestamps are wall-clock monotonic seconds: `prepared_at` is stamped at - construction and `infer_dequeued_at` is filled in via `_replace` when an - infer worker dequeues the item. Zero means "not yet reached". - """ - - group_id: int - rollout_idx: int - batch_id: int - index_in_batch: int - params: GroupRolloutParams - prepared_at: float = 0.0 - infer_dequeued_at: float = 0.0 - - -class _InferredItem(NamedTuple): - """One rollout post-inference, flowing from infer to assemble.""" - - item: _InferWorkItem - episode: EpisodeResult - inferred_at: float = 0.0 - - -class _RolloutPipeline: - """Per-call orchestrator for grouped rollout generation.""" +class EnvAllocation(NamedTuple): + """One env's constant share of every trainer batch.""" - def __init__( - self, - agent: "GroupedRolloutGenerator", - request: GroupedRolloutRequest, - parallel_generation_tasks: int, - ) -> None: - self.agent = agent - self.request = request - self.gran_policy = _GranularityConfig.from_request(request) - self.gate = _SubmissionGate( - capacity=parallel_generation_tasks, - submission=self.gran_policy.submission, - ) - rollouts_per_submission_unit = { - "R": 1, - "G": request.rollouts_per_group, - "B": self.gran_policy.num_groups_per_batch * request.rollouts_per_group, - }[self.gran_policy.submission] - self.num_infer_workers = parallel_generation_tasks * rollouts_per_submission_unit - if not request.streaming: - self.num_infer_workers = min( - self.num_infer_workers, request.num_groups * request.rollouts_per_group - ) - self.infer_queue = asyncio_Queue() - self.assemble_queue = asyncio_Queue() - # Unbounded: flow control is owned entirely by the submission gate. - # Bounding this queue would add a second backpressure that silently - # clamps the run-ahead configured via --rl-generation-lag. - self.output_queue = asyncio_Queue() - # Buffer of pending groups (incomplete groups being filled by - # stage_assemble). Held here so metric collection can report its size. - self._assemble_pending: dict[int, list[_InferredItem]] = {} - # Pending groups waiting for their batch to fill in stage_consume - # (only populated when prevent_dataset_reorder is True). - self._consume_pending: dict[int, list[RolloutGroup]] = {} - # Per-group "output entry" times, keyed by (batch_id, index_in_batch), - # so stage_consume can compute output_queue_dwell when yielding. - self._output_enqueued_at: dict[tuple[int, int], float] = {} - # Observability accumulators. Measured here; snapshot/reset and - # wandb formatting happen in rl_utils during metric logging. - self.infer_queue_dwell: list[float] = [] - self.engine_dwell: list[float] = [] - self.assemble_queue_dwell: list[float] = [] - self.output_queue_dwell: list[float] = [] - self.prepared_count = 0 - self.inferred_count = 0 - self.assembled_count = 0 - self.yielded_count = 0 - - async def stage_prepare(self) -> None: - """Generate gated inference work items.""" - assert ( - self.request.streaming - or self.request.num_groups % self.gran_policy.num_groups_per_batch == 0 - ), "non-streaming requires num_groups to be a multiple of num_groups_per_batch" - group_id = 0 - try: - while self.request.streaming or group_id < self.request.num_groups: - await self.gate.acquire_for("B") - batch_id = group_id // self.gran_policy.num_groups_per_batch - - for index_in_batch in range(self.gran_policy.num_groups_per_batch): - await self.gate.acquire_for("G") - params: GroupRolloutParams = await self.agent.prepare_group_rollout(self.request) - - for rollout_idx in range(self.request.rollouts_per_group): - await self.gate.acquire_for("R") - item = _InferWorkItem( - group_id=group_id, - rollout_idx=rollout_idx, - batch_id=batch_id, - index_in_batch=index_in_batch, - params=params, - prepared_at=time.monotonic(), - ) - await self.infer_queue.put(item) - self.prepared_count += 1 - group_id += 1 - finally: - self.infer_queue.shutdown() - - async def stage_infer(self) -> None: - """Run a persistent pool of inference workers, spawned once per pipeline.""" - workers = [ - asyncio.create_task(self._infer_worker()) for _ in range(self.num_infer_workers) - ] - try: - await asyncio.gather(*workers, return_exceptions=True) - finally: - for worker in workers: - worker.cancel() - self.assemble_queue.shutdown() - - async def _infer_worker(self) -> None: - while True: - try: - item = await self.infer_queue.get() - except asyncio_QueueShutDown: - return - item = item._replace(infer_dequeued_at=time.monotonic()) - if item.prepared_at: - self.infer_queue_dwell.append(item.infer_dequeued_at - item.prepared_at) - await self._infer_one(item) - - @trace_async_exceptions(verbose=True) - async def _infer_one(self, item: _InferWorkItem) -> None: - episode = await item.params.run_episode() - inferred_at = time.monotonic() - self.gate.release_for("R") - if item.infer_dequeued_at: - self.engine_dwell.append(inferred_at - item.infer_dequeued_at) - self.inferred_count += 1 - await self.assemble_queue.put( - _InferredItem(item=item, episode=episode, inferred_at=inferred_at) - ) - - async def stage_assemble(self) -> None: - """Build complete rollout groups from inferred items.""" - pending = self._assemble_pending - try: - while True: - try: - inferred = await self.assemble_queue.get() - except asyncio_QueueShutDown: - break - dequeued_at = time.monotonic() - if inferred.inferred_at: - self.assemble_queue_dwell.append(dequeued_at - inferred.inferred_at) - bucket = pending.setdefault(inferred.item.group_id, []) - bucket.append(inferred) - if len(bucket) < self.request.rollouts_per_group: - continue - completed = pending.pop(inferred.item.group_id) - completed.sort(key=lambda item: item.item.rollout_idx) - rollouts = await asyncio.gather( - *[item.item.params.build_rollout(item.episode) for item in completed] - ) - self.assembled_count += 1 - # NOTE: this filter is currently non-functional dead code: - # _GranularityConfig._validate rejects filter_groups_with_same_reward - # at pipeline construction, so `keep` is always True. Kept for a - # future PR that regenerates dropped groups instead of - # under-delivering to the caller. That PR must also release the - # gate slot on the drop path: G/B slots free on consumption, and - # a dropped group never reaches stage_consume, so its slot (and - # eventually its batch's) would leak permanently. - keep = ( - not self.request.filter_groups_with_same_reward - or np.std([rollout.reward for rollout in rollouts]) > 1e-6 - ) - if keep: - first = completed[0] - output_enqueued_at = time.monotonic() - self._output_enqueued_at[ - (first.item.batch_id, first.item.index_in_batch) - ] = output_enqueued_at - await self.output_queue.put( - RolloutGroup( - rollouts=rollouts, - batch_id=first.item.batch_id, - index_in_batch=first.item.index_in_batch, - ) - ) - finally: - self.output_queue.shutdown() - - def _record_output_dwell(self, group: RolloutGroup) -> None: - """Record how long a group sat in output_queue before being yielded.""" - key = (group.batch_id, group.index_in_batch) - enqueued_at = self._output_enqueued_at.pop(key, 0.0) - if enqueued_at: - self.output_queue_dwell.append(time.monotonic() - enqueued_at) - self.yielded_count += 1 - - async def stage_consume(self) -> AsyncIterator[RolloutGroup]: - if not self.gran_policy.prevent_dataset_reorder: - while True: - try: - group = await self.output_queue.get() - except asyncio_QueueShutDown: - return - self._record_output_dwell(group) - yield group - self.gate.release_for("G") - - next_batch_id = 0 - pending = self._consume_pending - while True: - try: - group = await self.output_queue.get() - except asyncio_QueueShutDown: - return - self._record_output_dwell(group) - pending.setdefault(group.batch_id, []).append(group) - while ( - len(pending.get(next_batch_id, [])) - >= self.gran_policy.num_groups_per_batch - ): - batch = pending.pop(next_batch_id) - batch.sort(key=lambda group: group.index_in_batch) - next_batch_id += 1 - for group in batch: - yield group - self.gate.release_for("G") - self.gate.release_for("B") + agent: "GroupedRolloutGenerator" + env_id: str + num_groups: int class GroupedRolloutGenerator(Agent, ABC): - """An interface to return grouped Rollout objects to support algorithms like GRPO.""" - - parallel_generation_tasks: int = 512 - - def __init__(self, *, parallel_generation_tasks: int | None = None, **kwargs): - super().__init__(**kwargs) - if parallel_generation_tasks is not None: - self.parallel_generation_tasks = parallel_generation_tasks + """Agent contract consumed by RolloutPipeline to generate grouped rollouts (e.g. GRPO).""" @abstractmethod - async def prepare_group_rollout( - self, - request: GroupedRolloutRequest, - ) -> GroupRolloutParams: + async def prepare_group_rollout(self, request: GroupedRolloutRequest) -> GroupRolloutParams: """Return the params for one group's rollouts.""" ... - async def get_grouped_rollouts( - self, request: GroupedRolloutRequest - ) -> AsyncIterator[RolloutGroup]: - assert isinstance( - request.inference_interface, ReturnsRaw - ), "InferenceInterface must support raw_text return to provide rollouts." - pipeline = _RolloutPipeline( - agent=self, - request=request, - parallel_generation_tasks=self.parallel_generation_tasks, - ) - # Expose the live pipeline for observability; rl_utils reads its - # queue sizes, gate state, and timing accumulators during logging. - self._active_pipeline = pipeline - stage_prepare_task = asyncio.create_task(pipeline.stage_prepare()) - infer_task = asyncio.create_task(pipeline.stage_infer()) - assemble_task = asyncio.create_task(pipeline.stage_assemble()) - tasks = (stage_prepare_task, infer_task, assemble_task) - - try: - async for group in pipeline.stage_consume(): - yield group - finally: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - self._active_pipeline = None + def rollout_allocations(self, num_groups: int) -> list[EnvAllocation]: + """Returns each env's per-trainer-batch allocation, in env order.""" + return [ + EnvAllocation( + agent=self, + env_id=getattr(self, "env_id", None) or "rollout", + num_groups=num_groups, + ) + ] class EvaluationAgent(Agent, ABC): diff --git a/megatron/rl/agent/rollout_pipeline.py b/megatron/rl/agent/rollout_pipeline.py new file mode 100644 index 00000000000..35712ce6992 --- /dev/null +++ b/megatron/rl/agent/rollout_pipeline.py @@ -0,0 +1,506 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Caller-owned orchestration of grouped rollout generation over an agent.""" + +import asyncio +import time +from collections import deque +from typing import TYPE_CHECKING, AsyncIterator, NamedTuple + +import numpy as np + +from megatron.core.inference.utils import asyncio_Queue, asyncio_QueueShutDown +from megatron.core.utils import trace_async_exceptions + +from ..inference import ReturnsRaw +from ..rollout_granularity import ( + GRANULARITY_RANK, + ConsumptionGranularity, + SubmissionGranularity, +) +from .api import EpisodeResult, GroupedRolloutRequest, GroupRolloutParams, RolloutGroup + +if TYPE_CHECKING: + from .api import GroupedRolloutGenerator + + +class _GranularityConfig(NamedTuple): + submission: SubmissionGranularity + consumption: ConsumptionGranularity + num_groups_per_batch: int + rollouts_per_group: int + num_groups_per_env: tuple[int, ...] + + @classmethod + def from_request( + cls, request: GroupedRolloutRequest, num_groups_per_env: list[int] + ) -> "_GranularityConfig": + """Build the per-request granularity policy. + + Args: + request: Grouped rollout request carrying the granularity choices. + num_groups_per_env: Groups each env contributes to one batch, in env order. + + Returns: + A validated _GranularityConfig. + """ + cls._validate(request, num_groups_per_env) + return cls( + submission=request.submission_granularity, + consumption=request.consumption_granularity, + num_groups_per_batch=request.num_groups, + rollouts_per_group=request.rollouts_per_group, + num_groups_per_env=tuple(num_groups_per_env), + ) + + def env_of_index(self, index_in_batch: int) -> int: + """Map a batch slot to the env owning it (slots are env-blocked, in env order). + + Args: + index_in_batch: Slot index within one trainer batch. + + Returns: + The env_index owning the slot. + """ + boundary = 0 + for env_index, groups in enumerate(self.num_groups_per_env): + boundary += groups + if index_in_batch < boundary: + return env_index + raise IndexError( + f"index_in_batch {index_in_batch} outside batch of {self.num_groups_per_batch}" + ) + + def units_per_batch(self) -> int: + """Submission units in one batch; gate capacity = depth-in-batches x this. + + Returns: + The number of submission units one trainer batch contains. + """ + return { + "R": self.num_groups_per_batch * self.rollouts_per_group, + "G": self.num_groups_per_batch, + "B": 1, + }[self.submission] + + @staticmethod + def _validate(request: GroupedRolloutRequest, num_groups_per_env: list[int]) -> None: + """Reject invalid granularity, layout, and filter combinations. + + Args: + request: Grouped rollout request to check. + num_groups_per_env: Proposed per-env group layout. + """ + assert ( + GRANULARITY_RANK[request.consumption_granularity] + >= GRANULARITY_RANK[request.submission_granularity] + ), ( + f"Consumption granularity ({request.consumption_granularity}) must be no finer " + f"than submission granularity ({request.submission_granularity})." + ) + assert all( + groups > 0 for groups in num_groups_per_env + ), "Each environment must request at least one group per batch." + assert ( + sum(num_groups_per_env) == request.num_groups + ), "The sum of groups per environment must equal the total number of groups requested." + assert not request.filter_groups_with_same_reward, ( + "filter_groups_with_same_reward is not currently supported: dropped groups " + "are not regenerated, so non-streaming callers receive fewer groups than " + "requested and batch-order consumers stall on incomplete batches." + ) + + +class _SubmissionGate: + """Gate capacity is measured in units of the configured submission granularity. + + Each granularity has a single release point: R slots free when inference + completes, so the gate bounds engine concurrency in rollouts. G and B + slots free when the trainer consumes the group/batch, so the gate + enforces the --rl-generation-lag run-ahead cap in groups/batches + respectively. + """ + + def __init__( + self, + *, + capacity: int, + submission: SubmissionGranularity, + ) -> None: + """Create a gate with `capacity` slots counted at `submission` granularity. + + Args: + capacity: Maximum submission units in flight. + submission: Configured submission granularity. + """ + self._sem = asyncio.Semaphore(capacity) + self._submission = submission + self.capacity = capacity + # Observability counters, updated only on the configured submission + # granularity (the only path that touches the semaphore). `held` + # counts slots currently held; `prepare_blocked_seconds` accumulates + # time stage_prepare spent waiting on the semaphore. + self.held = 0 + self.prepare_blocked_seconds = 0.0 + self.acquire_calls = 0 + self.release_calls = 0 + + async def acquire_for(self, granularity: SubmissionGranularity) -> None: + """Take one slot when crossing a boundary of the configured granularity. + + Args: + granularity: The dispatch boundary being crossed. + """ + if self._submission == granularity: + start = time.monotonic() + await self._sem.acquire() + self.prepare_blocked_seconds += time.monotonic() - start + self.held += 1 + self.acquire_calls += 1 + + def release_for(self, granularity: SubmissionGranularity) -> None: + """Release one slot when work at the given granularity reaches its release point. + + Args: + granularity: The granularity whose release point was just reached. + """ + if self._submission == granularity: + self._sem.release() + self.held -= 1 + self.release_calls += 1 + + +class _InferWorkItem(NamedTuple): + """One rollout's worth of work flowing from prepare to infer. + + Timestamps are wall-clock monotonic seconds: `prepared_at` is stamped at + construction and `infer_dequeued_at` is filled in via `_replace` when an + infer worker dequeues the item. Zero means "not yet reached". + """ + + group_id: int + rollout_idx: int + batch_id: int + index_in_batch: int + params: GroupRolloutParams + env_index: int = 0 + prepared_at: float = 0.0 + infer_dequeued_at: float = 0.0 + + +class _InferredItem(NamedTuple): + """One rollout post-inference, flowing from infer to assemble.""" + + item: _InferWorkItem + episode: EpisodeResult + inferred_at: float = 0.0 + + +class RolloutPipeline: + """Orchestrates grouped rollout generation over an agent, one instance per request. + + Constructed and driven by the caller (e.g. the trainer via run()); + the agent only supplies the env allocations, per-group preparation, and inference calls. + """ + + def __init__( + self, + agent: "GroupedRolloutGenerator", + request: GroupedRolloutRequest, + parallel_generation_tasks: int, + ) -> None: + """Validate the request and size the gate, queues, and worker pool. + + Args: + agent: Agent supplying the env layout, preparation, and inference. + request: Grouped rollout request to serve; one pipeline per request. + parallel_generation_tasks: Submission gate depth in trainer batches. + """ + assert isinstance( + request.inference_interface, ReturnsRaw + ), "InferenceInterface must support raw_text return to provide rollouts." + self.agent = agent + self.request = request + self.allocations = agent.rollout_allocations(request.num_groups) + self.gran_policy = _GranularityConfig.from_request( + request, [allocation.num_groups for allocation in self.allocations] + ) + self.gate = _SubmissionGate( + capacity=parallel_generation_tasks + * self.gran_policy.units_per_batch(), + submission=self.gran_policy.submission, + ) + self.num_infer_workers = ( + parallel_generation_tasks + * self.gran_policy.num_groups_per_batch + * request.rollouts_per_group + ) + if not request.streaming: + self.num_infer_workers = min( + self.num_infer_workers, request.num_groups * request.rollouts_per_group + ) + + # Core queues. + self.infer_queue = asyncio_Queue() + self.assemble_queue = asyncio_Queue() + self.output_queue = asyncio_Queue() + + # Buffers of partial results. + self._assemble_pending: dict[int, list[_InferredItem]] = {} + self._consume_pending: dict[int, list[RolloutGroup]] = {} + self._output_enqueued_at: dict[tuple[int, int], float] = {} + + # Observability accumulators. + self.infer_queue_dwell: list[float] = [] + self.engine_dwell: list[float] = [] + self.assemble_queue_dwell: list[float] = [] + self.output_queue_dwell: list[float] = [] + self.prepared_count = 0 + self.inferred_count = 0 + self.assembled_count = 0 + self.yielded_count = 0 + self.prepared_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) + self.assembled_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) + self.yielded_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) + + async def run(self) -> AsyncIterator[RolloutGroup]: + """Run the pipeline stages; cancels them when the iterator is closed. + + Yields: + RolloutGroup: Groups in consumption-granularity order. + """ + tasks = ( + asyncio.create_task(self.stage_prepare()), + asyncio.create_task(self.stage_infer()), + asyncio.create_task(self.stage_assemble()), + ) + try: + async for group in self.stage_consume(): + yield group + for task in tasks: + task.cancel() + results = await asyncio.gather(*tasks, return_exceptions=True) + failure = next( + ( + result + for result in results + if isinstance(result, BaseException) + and not isinstance(result, asyncio.CancelledError) + ), + None, + ) + expected_end = ( + not self.request.streaming + and self.yielded_count == self.request.num_groups + ) + if failure is not None or not expected_end: + raise RuntimeError( + "RolloutPipeline output stream ended: a pipeline stage died" + + ("" if failure is not None else " (no stage exception was recovered)") + ) from failure + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + async def stage_prepare(self) -> None: + """Generate gated inference work items.""" + assert ( + self.request.streaming + or self.request.num_groups % self.gran_policy.num_groups_per_batch == 0 + ), "non-streaming requires num_groups to be a multiple of num_groups_per_batch" + group_id = 0 + try: + while self.request.streaming or group_id < self.request.num_groups: + await self.gate.acquire_for("B") + batch_id = group_id // self.gran_policy.num_groups_per_batch + + for index_in_batch in range(self.gran_policy.num_groups_per_batch): + env_index = self.gran_policy.env_of_index(index_in_batch) + await self.gate.acquire_for("G") + agent = self.allocations[env_index].agent + params: GroupRolloutParams = await agent.prepare_group_rollout(self.request) + self.prepared_groups_per_env[env_index] += 1 + + for rollout_idx in range(self.request.rollouts_per_group): + await self.gate.acquire_for("R") + item = _InferWorkItem( + group_id=group_id, + rollout_idx=rollout_idx, + batch_id=batch_id, + index_in_batch=index_in_batch, + params=params, + env_index=env_index, + prepared_at=time.monotonic(), + ) + await self.infer_queue.put(item) + self.prepared_count += 1 + group_id += 1 + finally: + self.infer_queue.shutdown() + + async def stage_infer(self) -> None: + """Run a persistent pool of inference workers, spawned once per pipeline.""" + workers = [ + asyncio.create_task(self._infer_worker()) for _ in range(self.num_infer_workers) + ] + try: + await asyncio.gather(*workers, return_exceptions=True) + finally: + for worker in workers: + worker.cancel() + self.assemble_queue.shutdown() + + async def _infer_worker(self) -> None: + while True: + try: + item = await self.infer_queue.get() + except asyncio_QueueShutDown: + return + item = item._replace(infer_dequeued_at=time.monotonic()) + if item.prepared_at: + self.infer_queue_dwell.append(item.infer_dequeued_at - item.prepared_at) + await self._infer_one(item) + + @trace_async_exceptions(verbose=True) + async def _infer_one(self, item: _InferWorkItem) -> None: + """Run one episode for one work item and hand the result to assemble. + + Args: + item: The dequeued work item; its params carry the episode closure. + """ + episode = await item.params.run_episode() + inferred_at = time.monotonic() + self.gate.release_for("R") + if item.infer_dequeued_at: + self.engine_dwell.append(inferred_at - item.infer_dequeued_at) + self.inferred_count += 1 + await self.assemble_queue.put( + _InferredItem(item=item, episode=episode, inferred_at=inferred_at) + ) + + async def stage_assemble(self) -> None: + """Build complete rollout groups from inferred items.""" + pending = self._assemble_pending + try: + while True: + try: + inferred = await self.assemble_queue.get() + except asyncio_QueueShutDown: + break + dequeued_at = time.monotonic() + if inferred.inferred_at: + self.assemble_queue_dwell.append(dequeued_at - inferred.inferred_at) + bucket = pending.setdefault(inferred.item.group_id, []) + bucket.append(inferred) + if len(bucket) < self.request.rollouts_per_group: + continue + completed = pending.pop(inferred.item.group_id) + completed.sort(key=lambda item: item.item.rollout_idx) + rollouts = await asyncio.gather( + *[item.item.params.build_rollout(item.episode) for item in completed] + ) + first = completed[0] + self.assembled_count += 1 + self.assembled_groups_per_env[first.item.env_index] += 1 + # NOTE: this filter is currently non-functional dead code: + # _GranularityConfig._validate rejects filter_groups_with_same_reward + # at pipeline construction, so `keep` is always True. Kept for a + # future PR that regenerates dropped groups instead of + # under-delivering to the caller. That PR must also release the + # gate slot on the drop path: G/B slots free on consumption, + # and a dropped group never reaches stage_consume, so its slot + # (and eventually its batch's) would leak permanently. + keep = ( + not self.request.filter_groups_with_same_reward + or np.std([rollout.reward for rollout in rollouts]) > 1e-6 + ) + if keep: + output_enqueued_at = time.monotonic() + self._output_enqueued_at[ + (first.item.batch_id, first.item.index_in_batch) + ] = output_enqueued_at + await self.output_queue.put( + RolloutGroup( + rollouts=rollouts, + batch_id=first.item.batch_id, + index_in_batch=first.item.index_in_batch, + ) + ) + finally: + self.output_queue.shutdown() + + def _record_output_dwell(self, group: RolloutGroup) -> None: + """Record how long a group sat in output_queue before being yielded.""" + key = (group.batch_id, group.index_in_batch) + enqueued_at = self._output_enqueued_at.pop(key, 0.0) + if enqueued_at: + self.output_queue_dwell.append(time.monotonic() - enqueued_at) + self.yielded_count += 1 + self.yielded_groups_per_env[self.gran_policy.env_of_index(group.index_in_batch)] += 1 + + async def _next_complete_group(self) -> RolloutGroup | None: + """Pop the next group off output_queue and record its dwell.""" + try: + group = await self.output_queue.get() + except asyncio_QueueShutDown: + return None + self._record_output_dwell(group) + return group + + async def stage_consume(self) -> AsyncIterator[RolloutGroup]: + """Deliver groups in the order defined by the consumption granularity.""" + consume = { + "G": self._consume_completion_order, + "B": self._consume_batch_order, + }[self.gran_policy.consumption] + async for group in consume(): + yield group + + async def _consume_completion_order(self) -> AsyncIterator[RolloutGroup]: + """G consumption: deliver groups in completion order, balanced across envs.""" + groups_per_env_per_batch = self.gran_policy.num_groups_per_env + pending_groups_by_env: list[deque[RolloutGroup]] = [ + deque() for _ in groups_per_env_per_batch + ] + delivered_groups_by_env = [0] * len(groups_per_env_per_batch) + while (group := await self._next_complete_group()) is not None: + env_index = self.gran_policy.env_of_index(group.index_in_batch) + pending_groups_by_env[env_index].append(group) + yielded_any = True + while yielded_any: + yielded_any = False + for env, queue in enumerate(pending_groups_by_env): + if queue and delivered_groups_by_env[env] < groups_per_env_per_batch[env]: + yield queue.popleft() + self.gate.release_for("G") + delivered_groups_by_env[env] += 1 + yielded_any = True + if all( + count == quota + for count, quota in zip(delivered_groups_by_env, groups_per_env_per_batch) + ): + delivered_groups_by_env = [0] * len(groups_per_env_per_batch) + # The stream is over; nothing is left to balance against, so drain any pending groups. + for queue in pending_groups_by_env: + while queue: + yield queue.popleft() + self.gate.release_for("G") + + async def _consume_batch_order(self) -> AsyncIterator[RolloutGroup]: + """B consumption: deliver whole batches in dataset order.""" + next_batch_id = 0 + pending = self._consume_pending + while (group := await self._next_complete_group()) is not None: + pending.setdefault(group.batch_id, []).append(group) + while ( + len(pending.get(next_batch_id, [])) + >= self.gran_policy.num_groups_per_batch + ): + batch = pending.pop(next_batch_id) + batch.sort(key=lambda group: group.index_in_batch) + next_batch_id += 1 + for group in batch: + yield group + self.gate.release_for("G") + self.gate.release_for("B") diff --git a/megatron/rl/agent/weighted_multi_task.py b/megatron/rl/agent/weighted_multi_task.py index 2c52784be1c..a11e93fd7a7 100644 --- a/megatron/rl/agent/weighted_multi_task.py +++ b/megatron/rl/agent/weighted_multi_task.py @@ -4,13 +4,12 @@ import logging from typing import Any, Optional, Type -import numpy as np - from .registry import get_agent_class from .api import ( AgentBaseModel, ContrastiveRollout, ContrastiveRolloutGenerator, + EnvAllocation, EvaluationAgent, EvaluationRequest, EvaluationResponse, @@ -70,9 +69,7 @@ def __init__(self, agent_configs: list[AgentConfig]): self.weights.append(config.weight / total_weight) @classmethod - def from_config( - cls, config: list[dict[str, Any]], *, parallel_generation_tasks: int | None = None - ) -> 'WeightedMultiTask': + def from_config(cls, config: list[dict[str, Any]]) -> 'WeightedMultiTask': """Create a WeightedMultiTask from a config list. Args: @@ -89,8 +86,6 @@ def from_config( if not all(k in entry for k in ['agent_type', 'agent_args', 'weight']): raise ValueError(f"Missing required keys in config entry: {entry}") agent_args = entry.get('agent_args', {}) - agent_args['parallel_generation_tasks'] = parallel_generation_tasks - agent_type = get_agent_class(entry['agent_type']) agent_configs.append( AgentConfig( @@ -101,10 +96,7 @@ def from_config( ) ) - instance = cls(agent_configs) - if parallel_generation_tasks is not None: - instance.parallel_generation_tasks = parallel_generation_tasks - return instance + return cls(agent_configs) def _distribute_counts(self, total_count: int, distribute_remainder: bool = True) -> list[int]: """Helper method to distribute counts according to weights. @@ -156,12 +148,52 @@ def _distribute_counts(self, total_count: int, distribute_remainder: bool = True return final_counts + def rollout_allocations(self, num_groups: int) -> list[EnvAllocation]: + """Constant per-batch allocation for each weighted env, in env order.""" + counts = self._distribute_counts(num_groups) + env_ids = [ + getattr(agent, "env_id", None) or f"agent_{idx}" + for idx, agent in enumerate(self.agents) + ] + starved = [ + env_ids[idx] + for idx, count in enumerate(counts) + if count == 0 and self.weights[idx] > 0 + ] + if starved: + raise ValueError( + f"num_groups={num_groups} is too small to give every weighted env a group " + f"per batch (starved envs: {starved}); increase the trainer batch size." + ) + for agent, count in zip(self.agents, counts): + if count > 0 and not isinstance(agent, GroupedRolloutGenerator): + raise TypeError( + f"Agent of type {type(agent)} does not support grouped rollouts" + ) + # Snapshot for metric logging; read back by rl_utils. + self.latest_distribution = { + "env_ids": env_ids, + "agent_groups": list(counts), + "num_groups": num_groups, + } + logger.info( + "WeightedMultiTask layout: num_groups=%d per_agent=%s", + num_groups, + ", ".join(f"{eid}(groups={c})" for eid, c in zip(env_ids, counts)), + ) + return [ + EnvAllocation(agent=agent, env_id=env_id, num_groups=count) + for agent, env_id, count in zip(self.agents, env_ids, counts) + if count > 0 + ] + async def prepare_group_rollout( self, request: GroupedRolloutRequest, ) -> GroupRolloutParams: raise NotImplementedError( - "WeightedMultiTask is a collection of tasks and therefore doesn't implement this method directly. Use get_grouped_rollouts instead to generate grouped rollouts." + "WeightedMultiTask only routes; the pipeline prepares each group via the " + "agent in the matching rollout_allocations entry." ) async def get_rollout_response(self, request, inference_request): @@ -189,104 +221,6 @@ async def get_reward_rollouts(self, request: RolloutRequest) -> list[Rollout]: all_rollouts_lists = await asyncio.gather(*tasks) return [rollout for rollouts in all_rollouts_lists for rollout in rollouts] - async def get_grouped_rollouts(self, request: GroupedRolloutRequest): - """Distribute grouped rollouts across sub-agents according to weights.""" - agent_groups = self._distribute_counts(request.num_groups) - if request.submission_granularity == "B": - # In BATCH mode, pgt counts local batches in flight. agent_groups already - # splits each batch by weight, so copy pgt to every active agent. - agent_pgts = [ - self.parallel_generation_tasks if num_groups > 0 else 0 - for num_groups in agent_groups - ] - else: - # In GROUP/ROLLOUT mode, pgt counts fine-grained work units, so split it by weight. - agent_pgts = self._distribute_counts(self.parallel_generation_tasks) - agent_slots = self._distribute_counts(request.num_groups, distribute_remainder=False) - agent_slots = np.array(agent_slots) / np.gcd.reduce(agent_slots) - - # Snapshot the distribution for observability. Read back by rl_utils - # during per-iteration metric logging. - env_ids = [getattr(a, "env_id", f"agent_{i}") or f"agent_{i}" - for i, a in enumerate(self.agents)] - self.latest_distribution = { - "env_ids": env_ids, - "agent_groups": list(agent_groups), - "agent_pgts": list(agent_pgts), - "agent_slots": agent_slots.tolist(), - "total_pgt": int(sum(agent_pgts)), - "num_groups": request.num_groups, - } - logger.info( - "WeightedMultiTask distribution: sub=%s cons=%s num_groups=%d " - "rollouts_per_group=%d total_pgt=%d per_agent=" - + ", ".join( - f"{eid}(groups={g}, pgt={p}, slots={s:g})" - for eid, g, p, s in zip(env_ids, agent_groups, agent_pgts, agent_slots) - ), - request.submission_granularity, - request.consumption_granularity, - request.num_groups, - request.rollouts_per_group, - int(sum(agent_pgts)), - ) - - # Create tasks for each agent with non-zero groups - generators = [] - for agent, num_groups, pgt in zip( - self.agents, agent_groups, agent_pgts, strict=True - ): - if num_groups > 0: - if not isinstance(agent, GroupedRolloutGenerator): - raise TypeError( - f"Agent of type {type(agent)} does not support grouped rollouts" - ) - agent.parallel_generation_tasks = pgt - agent_request = GroupedRolloutRequest( - num_groups=num_groups, - streaming=request.streaming, - rollouts_per_group=request.rollouts_per_group, - inference_interface=request.inference_interface, - validation=request.validation, - generation_args=request.generation_args, - filter_groups_with_same_reward=request.filter_groups_with_same_reward, - submission_granularity=request.submission_granularity, - consumption_granularity=request.consumption_granularity, - ) - generators.append(agent.get_grouped_rollouts(agent_request)) - else: - generators.append(None) - - while any(generators): - balanced_rollouts = asyncio.Queue() - - async def get_balanced_rollouts_if_remaining(agent_id): - generated_rollouts = 0 - while generated_rollouts < agent_slots[agent_id]: - if generators[agent_id] is None: - return - try: - await balanced_rollouts.put(await anext(generators[agent_id])) - generated_rollouts += 1 - except StopAsyncIteration: - await balanced_rollouts.put(None) - generators[agent_id] = None - return - - tasks = [ - asyncio.create_task(get_balanced_rollouts_if_remaining(agent_id)) - for agent_id in range(len(generators)) - ] - - try: - while balanced_rollouts.qsize() > 0 or not all(task.done() for task in tasks): - rollout = await balanced_rollouts.get() - if rollout is not None: - yield rollout - finally: - for task in tasks: - task.cancel() - async def get_contrastive_rollouts(self, request: RolloutRequest) -> list[ContrastiveRollout]: """Distribute contrastive rollouts across sub-agents according to weights.""" agent_rollouts = self._distribute_counts(request.num_rollouts) diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index b2a6b320a97..1bbb2fb3f25 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -28,7 +28,6 @@ ReturnsRaw, ReturnsTokens, ) -from ..rollout_granularity import get_rl_parallel_generation_tasks from ..server.api import InferenceServer logger = logging.getLogger(__name__) @@ -140,7 +139,7 @@ async def launch(cls, model: GPTModel, **kwargs): concurrency_limit = ( args.grpo_prompts_per_step * args.grpo_group_size - * get_rl_parallel_generation_tasks(args) + * (args.rl_generation_lag + 1) ) custom_limits = httpx.Limits( max_connections=concurrency_limit, diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index e43e2f72ced..14551b685e9 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from functools import partial from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional import numpy as np import torch @@ -71,11 +71,13 @@ RolloutGroup, TokenRollout, ) +from megatron.rl.agent.rollout_pipeline import RolloutPipeline from megatron.rl.agent.weighted_multi_task import WeightedMultiTask +from megatron.rl.inference import ReturnsRaw from megatron.rl.inference.megatron import MegatronLocal from megatron.rl.logging import LOG_DIR as lang_rl_log_dir from megatron.rl.logging import log as lang_rl_log -from megatron.rl.rollout_granularity import get_rl_parallel_generation_tasks +from megatron.rl.rollout_granularity import ConsumptionGranularity, SubmissionGranularity from megatron.rl.sequence_packing_utils import ( compute_packed_inference_logprobs_stats, get_default_packed_seq_params, @@ -537,19 +539,12 @@ def align_unpacked_inference_logprobs( return padded_inference_logprobs -def get_agent(args, parallel_generation_tasks: int | None = None): - """Get an agent based on environment configuration. - - If args.langrl_env_config is provided, uses weighted environment selection. - Otherwise falls back to legacy single environment selection. - """ - with open(args.langrl_env_config, 'r') as f: +def get_agent(env_config_path): + """Build the rollout agent tree from the environment configuration.""" + with open(env_config_path, 'r') as f: config = yaml.safe_load(f) - return WeightedMultiTask.from_config( - config, - parallel_generation_tasks=parallel_generation_tasks, - ) + return WeightedMultiTask.from_config(config) _INFERENCE_INTERFACE = None @@ -569,33 +564,46 @@ def get_inference_interface(args, loop, model): _ROLLOUT_GENERATOR = None -_ROLLOUT_AGENT = None +_ROLLOUT_PIPELINE = None + + +def get_rollout_generator( + inference_interface: ReturnsRaw, + n_prompts: int, + samples_per_group: int, + *, + streaming: bool, + generation_args: dict[str, Any], + filter_groups_with_same_reward: bool, + submission_granularity: SubmissionGranularity, + consumption_granularity: ConsumptionGranularity, + generation_lag: int, + env_config_path: str, +) -> AsyncIterator[RolloutGroup]: + """Return the rollout group iterator for this step. - -def get_rollout_generator(args, inference_interface, n_prompts, samples_per_group): - global _ROLLOUT_GENERATOR, _ROLLOUT_AGENT - if not (streaming := args.rl_partial_rollouts) or _ROLLOUT_GENERATOR is None: - parallel_generation_tasks = get_rl_parallel_generation_tasks(args) - agent = get_agent(args, parallel_generation_tasks=parallel_generation_tasks) + Returns: + The async iterator produced by RolloutPipeline.run(). + """ + global _ROLLOUT_GENERATOR, _ROLLOUT_PIPELINE + if not streaming or _ROLLOUT_GENERATOR is None: request = GroupedRolloutRequest( num_groups=n_prompts, streaming=streaming, rollouts_per_group=samples_per_group, inference_interface=inference_interface, - generation_args={ - 'temperature': args.rl_default_temperature, - 'max_tokens': args.inference_max_seq_length, - 'top_p': args.rl_default_top_p, - 'top_k': args.rl_default_top_k, - }, - filter_groups_with_same_reward=args.grpo_filter_groups_with_same_reward, - submission_granularity=args.rl_submission_granularity, - consumption_granularity=args.rl_consumption_granularity, + generation_args=generation_args, + filter_groups_with_same_reward=filter_groups_with_same_reward, + submission_granularity=submission_granularity, + consumption_granularity=consumption_granularity, ) - # Keep the agent handle so metric logging can read the live rollout - # pipelines (see _collect_rollout_pipeline_metrics). - _ROLLOUT_AGENT = agent - _ROLLOUT_GENERATOR = agent.get_grouped_rollouts(request) + # Keep the pipeline handle so logging can read its queues, gate state, and per-env counters. + _ROLLOUT_PIPELINE = RolloutPipeline( + agent=get_agent(env_config_path), + request=request, + parallel_generation_tasks=generation_lag + 1, + ) + _ROLLOUT_GENERATOR = _ROLLOUT_PIPELINE.run() return _ROLLOUT_GENERATOR @@ -663,7 +671,21 @@ def get_environment_rollouts( with nvtx_range("rl/inference-setup", time=True): # Asyncronously run inference and rollout collection rollout_generator = get_rollout_generator( - args, inference_interface, n_prompts, samples_per_group + inference_interface, + n_prompts, + samples_per_group, + streaming=args.rl_partial_rollouts, + generation_args={ + 'temperature': args.rl_default_temperature, + 'max_tokens': args.inference_max_seq_length, + 'top_p': args.rl_default_top_p, + 'top_k': args.rl_default_top_k, + }, + filter_groups_with_same_reward=args.grpo_filter_groups_with_same_reward, + submission_granularity=args.rl_submission_granularity, + consumption_granularity=args.rl_consumption_granularity, + generation_lag=args.rl_generation_lag, + env_config_path=args.langrl_env_config, ) # NOTE(jbarker): we need to double check this when using PP>1 @@ -1174,91 +1196,96 @@ def _real(grouped): def _collect_rollout_pipeline_metrics() -> dict: - """Snapshot per-pipeline instrumentation into wandb-loggable scalars. + """Snapshot pipeline instrumentation into wandb-loggable scalars. - Walks the live rollout agent (set by get_rollout_generator) and, for each - sub-agent with an active _RolloutPipeline, reads queue sizes, gate state, - per-stage dwell times, and rate counters. Accumulators are reset after - reading; point-in-time values (queue sizes, gate held) are re-read next - call. Keys follow the existing f"{env_id}_{metric}" convention. + Reads the RolloutPipeline held by get_rollout_generator: queue sizes, gate state, + per-stage dwell times, and rate counters, plus per-env group counters for multi-env agents. + + Returns: + Metric name -> value dict; empty when no pipeline exists yet. """ - if _ROLLOUT_AGENT is None: + if _ROLLOUT_PIPELINE is None: return {} - sub_agents = ( - _ROLLOUT_AGENT.agents - if isinstance(_ROLLOUT_AGENT, WeightedMultiTask) - else [_ROLLOUT_AGENT] - ) + pipeline = _ROLLOUT_PIPELINE + dist = getattr(pipeline.agent, "latest_distribution", None) metrics: dict = {} - for sub_agent in sub_agents: - pipeline = getattr(sub_agent, "_active_pipeline", None) - if pipeline is None: - continue - env_id = getattr(sub_agent, "env_id", "") or "rollout" - gate = pipeline.gate - metrics.update({ - # Queue sizes and gate held are point-in-time reads. - f"{env_id}_pipeline_infer_queue_size": pipeline.infer_queue.qsize(), - f"{env_id}_pipeline_assemble_queue_size": pipeline.assemble_queue.qsize(), - f"{env_id}_pipeline_output_queue_size": pipeline.output_queue.qsize(), - f"{env_id}_pipeline_assemble_pending_groups": len(pipeline._assemble_pending), - f"{env_id}_pipeline_consume_pending_groups": len(pipeline._consume_pending), - f"{env_id}_pipeline_gate_capacity": gate.capacity, - f"{env_id}_pipeline_gate_held": gate.held, - f"{env_id}_pipeline_gate_utilization": ( - gate.held / gate.capacity if gate.capacity else 0.0 - ), - # Counters below accumulate since the previous collection. - f"{env_id}_pipeline_gate_prepare_blocked_seconds": gate.prepare_blocked_seconds, - f"{env_id}_pipeline_gate_acquire_calls": gate.acquire_calls, - f"{env_id}_pipeline_gate_release_calls": gate.release_calls, - f"{env_id}_pipeline_prepared_count": pipeline.prepared_count, - f"{env_id}_pipeline_inferred_count": pipeline.inferred_count, - f"{env_id}_pipeline_assembled_count": pipeline.assembled_count, - f"{env_id}_pipeline_yielded_count": pipeline.yielded_count, - }) - for name, samples in ( - ("infer_queue_dwell", pipeline.infer_queue_dwell), - ("engine_dwell", pipeline.engine_dwell), - ("assemble_queue_dwell", pipeline.assemble_queue_dwell), - ("output_queue_dwell", pipeline.output_queue_dwell), - ): - if samples: - arr = np.asarray(samples, dtype=np.float64) - metrics[f"{env_id}_pipeline_mean_{name}_s"] = float(arr.mean()) - metrics[f"{env_id}_pipeline_max_{name}_s"] = float(arr.max()) - metrics[f"{env_id}_pipeline_p50_{name}_s"] = float(np.percentile(arr, 50)) - metrics[f"{env_id}_pipeline_p99_{name}_s"] = float(np.percentile(arr, 99)) - # Reset accumulators; queue sizes and gate held are point-in-time. - pipeline.infer_queue_dwell = [] - pipeline.engine_dwell = [] - pipeline.assemble_queue_dwell = [] - pipeline.output_queue_dwell = [] - pipeline.prepared_count = 0 - pipeline.inferred_count = 0 - pipeline.assembled_count = 0 - pipeline.yielded_count = 0 - gate.prepare_blocked_seconds = 0.0 - gate.acquire_calls = 0 - gate.release_calls = 0 - - # WeightedMultiTask work distribution (agent_slots / agent_pgts). - dist = getattr(_ROLLOUT_AGENT, "latest_distribution", None) + gate = pipeline.gate + metrics.update({ + # Queue sizes and gate held are point-in-time reads. + "rollout_pipeline_infer_queue_size": pipeline.infer_queue.qsize(), + "rollout_pipeline_assemble_queue_size": pipeline.assemble_queue.qsize(), + "rollout_pipeline_output_queue_size": pipeline.output_queue.qsize(), + "rollout_pipeline_assemble_pending_groups": len(pipeline._assemble_pending), + "rollout_pipeline_consume_pending_groups": len(pipeline._consume_pending), + "rollout_pipeline_gate_capacity": gate.capacity, + "rollout_pipeline_gate_held": gate.held, + "rollout_pipeline_gate_utilization": ( + gate.held / gate.capacity if gate.capacity else 0.0 + ), + # Counters below accumulate since the previous collection. + "rollout_pipeline_gate_prepare_blocked_seconds": gate.prepare_blocked_seconds, + "rollout_pipeline_gate_acquire_calls": gate.acquire_calls, + "rollout_pipeline_gate_release_calls": gate.release_calls, + "rollout_pipeline_prepared_count": pipeline.prepared_count, + "rollout_pipeline_inferred_count": pipeline.inferred_count, + "rollout_pipeline_assembled_count": pipeline.assembled_count, + "rollout_pipeline_yielded_count": pipeline.yielded_count, + }) + for name, samples in ( + ("infer_queue_dwell", pipeline.infer_queue_dwell), + ("engine_dwell", pipeline.engine_dwell), + ("assemble_queue_dwell", pipeline.assemble_queue_dwell), + ("output_queue_dwell", pipeline.output_queue_dwell), + ): + if samples: + arr = np.asarray(samples, dtype=np.float64) + metrics[f"rollout_pipeline_mean_{name}_s"] = float(arr.mean()) + metrics[f"rollout_pipeline_max_{name}_s"] = float(arr.max()) + metrics[f"rollout_pipeline_p50_{name}_s"] = float(np.percentile(arr, 50)) + metrics[f"rollout_pipeline_p99_{name}_s"] = float(np.percentile(arr, 99)) + # Per-env group counters, mapped from env_index to env_id via the multi-env layout. + if dist: + active_env_ids = [ + env_id + for env_id, groups in zip(dist["env_ids"], dist["agent_groups"]) + if groups > 0 + ] + for env_index, env_id in enumerate(active_env_ids): + metrics[f"{env_id}_prepared_groups"] = ( + pipeline.prepared_groups_per_env[env_index] + ) + metrics[f"{env_id}_assembled_groups"] = ( + pipeline.assembled_groups_per_env[env_index] + ) + metrics[f"{env_id}_yielded_groups"] = ( + pipeline.yielded_groups_per_env[env_index] + ) + # Reset accumulators; queue sizes and gate held are point-in-time. + pipeline.infer_queue_dwell = [] + pipeline.engine_dwell = [] + pipeline.assemble_queue_dwell = [] + pipeline.output_queue_dwell = [] + pipeline.prepared_count = 0 + pipeline.inferred_count = 0 + pipeline.assembled_count = 0 + pipeline.yielded_count = 0 + pipeline.prepared_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) + pipeline.assembled_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) + pipeline.yielded_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) + gate.prepare_blocked_seconds = 0.0 + gate.acquire_calls = 0 + gate.release_calls = 0 + + # WeightedMultiTask per-batch group distribution. if dist: # An env_id can appear more than once in the config (e.g. an active # entry plus an evaluation-only twin with zero weight). Sum per # env_id so the zero twin does not overwrite the active entry. per_env: dict = {} - for env_id, groups, pgt, slots in zip( - dist["env_ids"], dist["agent_groups"], dist["agent_pgts"], dist["agent_slots"] - ): - g, p, s = per_env.get(env_id, (0, 0, 0.0)) - per_env[env_id] = (g + groups, p + pgt, s + slots) - for env_id, (groups, pgt, slots) in per_env.items(): + for env_id, groups in zip(dist["env_ids"], dist["agent_groups"]): + per_env[env_id] = per_env.get(env_id, 0) + groups + for env_id, groups in per_env.items(): metrics[f"{env_id}_agent_groups"] = groups - metrics[f"{env_id}_agent_pgts"] = pgt - metrics[f"{env_id}_agent_slots"] = slots - metrics["multitask_total_pgt"] = dist["total_pgt"] return metrics @@ -2008,7 +2035,7 @@ def evaluate_and_print_results_rl( rank = torch.distributed.get_rank() if rank == 0: logger.info("Collecting evaluation results...") - agent = get_agent(args) + agent = get_agent(args.langrl_env_config) request = EvaluationRequest( inference_interface=inference_interface, num_prompts=args.rl_prompts_per_eval, @@ -2319,13 +2346,13 @@ def megatron_rl_inference_mode( def rl_inference_interface_shutdown(): global _INFERENCE_INTERFACE global _ROLLOUT_GENERATOR - global _ROLLOUT_AGENT + global _ROLLOUT_PIPELINE if _ROLLOUT_GENERATOR is not None: loop = get_asyncio_loop() loop.run_until_complete(_ROLLOUT_GENERATOR.aclose()) _ROLLOUT_GENERATOR = None - _ROLLOUT_AGENT = None + _ROLLOUT_PIPELINE = None if _INFERENCE_INTERFACE is not None: loop = get_asyncio_loop() diff --git a/megatron/rl/rollout_granularity.py b/megatron/rl/rollout_granularity.py index b9432f0bd4d..aa7ea545c54 100644 --- a/megatron/rl/rollout_granularity.py +++ b/megatron/rl/rollout_granularity.py @@ -7,12 +7,5 @@ SubmissionGranularity = Literal["R", "G", "B"] ConsumptionGranularity = Literal["G", "B"] - -def get_rl_parallel_generation_tasks(args) -> int: - """Return the number of generation slots implied by RL lag and submission granularity.""" - parallel_generation_tasks = args.rl_generation_lag + 1 - if args.rl_submission_granularity != "B": - parallel_generation_tasks *= args.grpo_prompts_per_step - if args.rl_submission_granularity == "R": - parallel_generation_tasks *= args.grpo_group_size - return parallel_generation_tasks +# Coarseness order of the granularity ladder (rollout < group < batch). +GRANULARITY_RANK: dict[str, int] = {"R": 0, "G": 1, "B": 2} diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 72b65689ff5..781d38c378f 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -38,7 +38,6 @@ from megatron.rl import rl_utils from megatron.rl.agent.api import TokenRollout from megatron.rl.inference import ReturnsRaw -from megatron.rl.rollout_granularity import get_rl_parallel_generation_tasks from megatron.rl.sequence_packing_utils import get_default_packed_seq_params from megatron.training.arguments import parse_args, validate_args from megatron.training.global_vars import destroy_global_vars, set_global_variables @@ -214,30 +213,6 @@ def test_rl_granularity_defaults(self): assert args.rl_consumption_granularity == "B" assert args.rl_generation_lag == 0 assert not hasattr(args, "rl_parallel_generation_tasks") - assert get_rl_parallel_generation_tasks(args) == 1 - - @pytest.mark.parametrize( - "submission_granularity, generation_lag, expected_parallel_generation_tasks", - [ - pytest.param("B", 0, 1, id="batch"), - pytest.param("B", 2, 3, id="batch_with_lag"), - pytest.param("G", 0, 8, id="group"), - pytest.param("G", 2, 24, id="group_with_lag"), - pytest.param("R", 0, 32, id="rollout"), - pytest.param("R", 2, 96, id="rollout_with_lag"), - ], - ) - def test_get_rl_parallel_generation_tasks( - self, submission_granularity, generation_lag, expected_parallel_generation_tasks - ): - args = SimpleNamespace( - rl_submission_granularity=submission_granularity, - rl_generation_lag=generation_lag, - grpo_prompts_per_step=8, - grpo_group_size=4, - ) - - assert get_rl_parallel_generation_tasks(args) == expected_parallel_generation_tasks @pytest.mark.parametrize( "rl_partial_rollouts, submission_granularity", @@ -254,49 +229,52 @@ def test_get_rollout_generator_keeps_num_groups_at_trainer_batch_size( """Regression for the removed ``num_groups=1`` streaming override. Previously ``get_rollout_generator`` forced ``num_groups`` to 1 whenever it - streamed with a non-batch submission granularity. For a multi-environment - agent that collapses the per-env group distribution so some environments - receive zero groups (and a degenerate all-zero ``agent_slots``), stalling - ``get_grouped_rollouts``. ``num_groups`` must stay at the trainer batch size - (``n_prompts``) regardless of streaming or submission granularity. + streamed with a non-batch submission granularity, collapsing the per-env + group layout so some environments received zero groups (now a loud + ``ValueError`` from ``rollout_allocations``). ``num_groups`` must stay at + the trainer batch size (``n_prompts``) regardless of streaming or + submission granularity. """ n_prompts = 8 captured = {} rollout_generator = object() + agent = object() - class Agent: - def get_grouped_rollouts(self, request): + class FakePipeline: + def __init__(self, agent, request, parallel_generation_tasks): + captured["agent"] = agent captured["request"] = request - return rollout_generator + captured["parallel_generation_tasks"] = parallel_generation_tasks - def get_agent(_args, parallel_generation_tasks=None): - captured["parallel_generation_tasks"] = parallel_generation_tasks - return Agent() + def run(self): + return rollout_generator monkeypatch.setattr(rl_utils, "_ROLLOUT_GENERATOR", None) - monkeypatch.setattr(rl_utils, "get_agent", get_agent) - - args = SimpleNamespace( - rl_partial_rollouts=rl_partial_rollouts, - rl_submission_granularity=submission_granularity, - rl_consumption_granularity="B", - rl_generation_lag=0, - grpo_prompts_per_step=n_prompts, - grpo_group_size=4, - rl_default_temperature=1.0, - inference_max_seq_length=128, - rl_default_top_p=1.0, - rl_default_top_k=0, - grpo_filter_groups_with_same_reward=False, - ) + monkeypatch.setattr(rl_utils, "_ROLLOUT_PIPELINE", None) + monkeypatch.setattr(rl_utils, "get_agent", lambda _env_config_path: agent) + monkeypatch.setattr(rl_utils, "RolloutPipeline", FakePipeline) + generation_lag = 0 result = rl_utils.get_rollout_generator( - args, inference_interface=ReturnsRaw(), n_prompts=n_prompts, samples_per_group=4 + ReturnsRaw(), + n_prompts, + 4, + streaming=rl_partial_rollouts, + generation_args={'temperature': 1.0, 'max_tokens': 128, 'top_p': 1.0, 'top_k': 0}, + filter_groups_with_same_reward=False, + submission_granularity=submission_granularity, + consumption_granularity="B", + generation_lag=generation_lag, + env_config_path="unused.yaml", ) assert result is rollout_generator + assert captured["agent"] is agent assert captured["request"].num_groups == n_prompts assert captured["request"].streaming == rl_partial_rollouts + # The gate depth handed to the pipeline is lag + 1 trainer batches, + # independent of submission granularity. + assert captured["parallel_generation_tasks"] == generation_lag + 1 assert captured["request"].submission_granularity == submission_granularity @pytest.mark.parametrize( @@ -317,11 +295,6 @@ def get_agent(_args, parallel_generation_tasks=None): "--rl-consumption-granularity R is not currently supported", id="rollout_consumption_unsupported", ), - pytest.param( - {"rl_submission_granularity": "B", "rl_consumption_granularity": "G"}, - "--rl-submission-granularity B with --rl-consumption-granularity G", - id="batch_submit_group_consume_unsupported", - ), ], ) def test_rl_granularity_validation_rejects_unsupported_modes(self, overrides, match): diff --git a/tests/unit_tests/rl/test_rollout_generation.py b/tests/unit_tests/rl/test_rollout_generation.py index b9c26c7600e..ab1d960836c 100644 --- a/tests/unit_tests/rl/test_rollout_generation.py +++ b/tests/unit_tests/rl/test_rollout_generation.py @@ -1,9 +1,9 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import asyncio +from contextlib import aclosing from unittest.mock import MagicMock -import numpy as np import pytest from pydantic import Field, ValidationError @@ -16,9 +16,9 @@ RolloutGenerator, RolloutRequest, TokenRollout, - _SubmissionGate, ) from megatron.rl.agent.reward_only_agent import RewardOnlyAgent +from megatron.rl.agent.rollout_pipeline import RolloutPipeline, _SubmissionGate from megatron.rl.agent.weighted_multi_task import AgentConfig, WeightedMultiTask from megatron.rl.inference import InferenceResponse, LLMChatMessage, ReturnsRaw, ReturnsTokens @@ -60,11 +60,13 @@ def __init__(self, env_id="test", **kwargs): self.env_id = env_id self._call_count = 0 self.prepare_group_rollout_calls = 0 + self.get_rollout_response_calls = 0 async def get_reward_rollouts(self, request): raise NotImplementedError async def get_rollout_response(self, request, inference_request): + self.get_rollout_response_calls += 1 return await request.inference_interface.agenerate(inference_request) async def prepare_group_rollout(self, request): @@ -152,8 +154,10 @@ class TestConsumptionRelease: async def test_group_submission_stalls_until_consumption( self, consumption_granularity, num_groups ): + # Gate capacity in G-submission slots is parallel_generation_tasks + # (a depth in batches) x num_groups (groups per batch). capacity = 4 - gen = MockGenerator(parallel_generation_tasks=capacity) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=num_groups, rollouts_per_group=1, @@ -162,7 +166,8 @@ async def test_group_submission_stalls_until_consumption( submission_granularity="G", consumption_granularity=consumption_granularity, ) - it = gen.get_grouped_rollouts(request) + pipeline = RolloutPipeline(gen, request, parallel_generation_tasks=capacity // num_groups) + it = pipeline.run() try: for pulled in range(1, capacity + 3): # wait_for turns the deadlock failure mode (a slot never freed) @@ -181,7 +186,7 @@ async def test_group_submission_stalls_until_consumption( @pytest.mark.asyncio async def test_batch_submission_releases_once_per_batch(self): - gen = MockGenerator(parallel_generation_tasks=1) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=2, rollouts_per_group=1, @@ -190,12 +195,13 @@ async def test_batch_submission_releases_once_per_batch(self): submission_granularity="B", consumption_granularity="B", ) - it = gen.get_grouped_rollouts(request) + pipeline = RolloutPipeline(gen, request, parallel_generation_tasks=1) + it = pipeline.run() try: await asyncio.wait_for(anext(it), timeout=10) await asyncio.wait_for(anext(it), timeout=10) await _flush() - gate = gen._active_pipeline.gate + gate = pipeline.gate # Batch 0 fully yielded but the consumer hasn't come back yet: its # single batch slot is still held (a per-group release here would # show release_calls == 2 and prepared == 4). @@ -209,6 +215,74 @@ async def test_batch_submission_releases_once_per_batch(self): await it.aclose() +class TestStageFailurePropagation: + """A dead stage must fail run() loudly, never read as a clean end-of-stream. + + A stage that dies runs the queue-shutdown cascade, which reaches + stage_consume exactly like a clean end-of-stream; before run() reaped the + stage tasks, the caller saw StopAsyncIteration and waited forever for + rollouts nobody would ever generate (observed live 2026-07-30: a TypeError + in the first prepare_group_rollout idled a training job to its time limit). + """ + + @pytest.mark.asyncio + async def test_prepare_failure_raises_out_of_run(self): + class BrokenPrepareGenerator(MockGenerator): + async def prepare_group_rollout(self, request): + raise TypeError("agent/pipeline interface mismatch") + + request = GroupedRolloutRequest( + num_groups=2, + rollouts_per_group=2, + inference_interface=MockInferenceInterface(), + submission_granularity="R", + consumption_granularity="G", + ) + pipeline = RolloutPipeline(BrokenPrepareGenerator(), request, parallel_generation_tasks=1) + async with aclosing(pipeline.run()) as it: + with pytest.raises(RuntimeError, match="stage died") as excinfo: + # wait_for turns the pre-fix failure mode (an eternal hang once + # the cascade is mistaken for end-of-stream) into a test failure. + await asyncio.wait_for(anext(it), timeout=10) + assert isinstance(excinfo.value.__cause__, TypeError) + + @pytest.mark.asyncio + async def test_midstream_stage_failure_raises_after_delivered_groups(self): + class BrokenBuildGenerator(MockGenerator): + """First group builds normally; every later group's build_rollout raises.""" + + async def prepare_group_rollout(self, request): + idx = self._call_count + params = await super().prepare_group_rollout(request) + if idx < 1: + return params + + async def broken_build(episode): + raise ValueError("reward model exploded") + + return GroupRolloutParams( + run_episode=params.run_episode, build_rollout=broken_build + ) + + request = GroupedRolloutRequest( + num_groups=2, + rollouts_per_group=2, + inference_interface=MockInferenceInterface(), + submission_granularity="G", + consumption_granularity="G", + ) + pipeline = RolloutPipeline(BrokenBuildGenerator(), request, parallel_generation_tasks=1) + async with aclosing(pipeline.run()) as it: + # Group 0 is healthy and must still be delivered. + group = await asyncio.wait_for(anext(it), timeout=10) + assert len(group.rollouts) == 2 + # Group 1's build_rollout kills stage_assemble; the next pull must + # surface that failure instead of hanging on the drained stream. + with pytest.raises(RuntimeError, match="stage died") as excinfo: + await asyncio.wait_for(anext(it), timeout=10) + assert isinstance(excinfo.value.__cause__, ValueError) + + class TestRewardRollouts: @pytest.mark.asyncio async def test_get_reward_rollouts_matches_per_rollout_composition(self): @@ -246,7 +320,7 @@ def test_grouped_rollout_request_rejects_unknown_granularity(self, field): async def test_filter_groups_with_same_reward_rejected( self, num_groups, submission_granularity, consumption_granularity ): - gen = MockGenerator(parallel_generation_tasks=8) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=num_groups, rollouts_per_group=2, @@ -256,8 +330,7 @@ async def test_filter_groups_with_same_reward_rejected( consumption_granularity=consumption_granularity, ) with pytest.raises(AssertionError, match="filter_groups_with_same_reward"): - async for _ in gen.get_grouped_rollouts(request): - pass + RolloutPipeline(gen, request, parallel_generation_tasks=8) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -307,7 +380,7 @@ async def test_filter_groups_with_same_reward_rejected( ), ], ) - async def test_get_grouped_rollouts( + async def test_grouped_rollout_generation( self, num_slow_calls, streaming, @@ -318,7 +391,7 @@ async def test_get_grouped_rollouts( expected_batch_ids, expected_trajectories, ): - gen = MockGenerator(parallel_generation_tasks=8) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=num_groups, rollouts_per_group=1, @@ -329,7 +402,7 @@ async def test_get_grouped_rollouts( ) groups = [] - async for group in gen.get_grouped_rollouts(request): + async for group in RolloutPipeline(gen, request, parallel_generation_tasks=8).run(): groups.append(group) if request.streaming and len(groups) >= expected_count: break @@ -343,11 +416,14 @@ async def test_get_grouped_rollouts( @pytest.mark.asyncio async def test_rollout_submission_granularity_limits_inference_concurrency(self): - gen = MockGenerator(parallel_generation_tasks=2) + # parallel_generation_tasks is a depth in batches; the R gate admits at + # most depth x (num_groups x rollouts_per_group) rollouts at once. + parallel_generation_tasks = 1 + gen = MockGenerator() inference_interface = MockInferenceInterface(num_slow_calls=100) request = GroupedRolloutRequest( - num_groups=1, - rollouts_per_group=4, + num_groups=2, + rollouts_per_group=2, inference_interface=inference_interface, streaming=True, submission_granularity="R", @@ -355,42 +431,30 @@ async def test_rollout_submission_granularity_limits_inference_concurrency(self) ) groups = [] - async for group in gen.get_grouped_rollouts(request): + pipeline = RolloutPipeline( + gen, request, parallel_generation_tasks=parallel_generation_tasks + ) + async for group in pipeline.run(): groups.append(group) - break + if len(groups) >= 4: + break - assert len(groups) == 1 - assert len(groups[0]) == 4 - assert inference_interface.max_active_requests <= gen.parallel_generation_tasks + assert all(len(group) == 2 for group in groups) + assert inference_interface.max_active_requests <= ( + parallel_generation_tasks * request.num_groups * request.rollouts_per_group + ) @pytest.mark.asyncio @pytest.mark.parametrize( - "submission_granularity, consumption_granularity, expected_parallel_generation_tasks", - [ - pytest.param("B", "B", [4, 4], id="batch_submission"), - pytest.param("G", "G", [3, 1], id="group_submission"), - ], + "submission_granularity, consumption_granularity", + [pytest.param("B", "B", id="batch_batch"), pytest.param("G", "G", id="group_group")], ) - async def test_weighted_multi_task( - self, submission_granularity, consumption_granularity, expected_parallel_generation_tasks - ): + async def test_weighted_multi_task(self, submission_granularity, consumption_granularity): configs = [ AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), ] mt = WeightedMultiTask(configs) - mt.parallel_generation_tasks = 4 - - captured = [] - for agent in mt.agents: - original = agent.get_grouped_rollouts - - async def spy(req, orig=original): - captured.append(req) - async for group in orig(req): - yield group - - agent.get_grouped_rollouts = spy request = GroupedRolloutRequest( num_groups=4, @@ -401,59 +465,91 @@ async def spy(req, orig=original): consumption_granularity=consumption_granularity, ) groups = [] - async for group in mt.get_grouped_rollouts(request): + pipeline = RolloutPipeline(mt, request, parallel_generation_tasks=1) + async for group in pipeline.run(): groups.append(group) assert len(groups) == 4 - # Weights 3:1 → agent "a" produces 3 groups, agent "b" produces 1. + # Weights 3:1 → env "a" owns 3 batch slots, env "b" owns 1; the single + # pipeline routes preparation and generation to the owning sub-agent. env_ids = [g[0].env_id for g in groups] assert sorted(env_ids) == ["a", "a", "a", "b"] - for sub_req in captured: - assert sub_req.num_groups in (1, 3) # distributed proportionally by weight - assert sub_req.streaming == request.streaming - assert sub_req.submission_granularity == request.submission_granularity - assert sub_req.consumption_granularity == request.consumption_granularity - assert [agent.parallel_generation_tasks for agent in mt.agents] == ( - expected_parallel_generation_tasks + assert [agent.prepare_group_rollout_calls for agent in mt.agents] == [3, 1] + assert [agent.get_rollout_response_calls for agent in mt.agents] == [3, 1] + assert mt.latest_distribution["agent_groups"] == [3, 1] + # The pipeline drains fully: every gate slot is released at exhaustion. + assert pipeline.gate.held == 0 + + @pytest.mark.asyncio + async def test_group_consumption_balances_each_batch(self): + """Balanced G: every trainer-batch window holds each env's exact share.""" + configs = [ + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), + ] + mt = WeightedMultiTask(configs) + + request = GroupedRolloutRequest( + num_groups=4, + rollouts_per_group=1, + inference_interface=MockInferenceInterface(num_slow_calls=2), + streaming=True, + submission_granularity="G", + consumption_granularity="G", ) + groups = [] + async for group in RolloutPipeline(mt, request, parallel_generation_tasks=2).run(): + groups.append(group) + if len(groups) >= 12: + break + + for start in range(0, 12, 4): + env_ids = [g[0].env_id for g in groups[start : start + 4]] + assert sorted(env_ids) == ["a", "a", "a", "b"] + @pytest.mark.asyncio @pytest.mark.parametrize( - "num_groups, all_envs_active", - [ - pytest.param(1, False, id="num_groups_1_starves_an_env"), - pytest.param(8, True, id="trainer_batch_size_keeps_all_envs_active"), - ], + "submission_granularity, consumption_granularity", + [pytest.param("B", "G", id="batch_group")], ) - def test_multi_env_distribution_requires_num_groups_above_one( - self, num_groups, all_envs_active + async def test_consumption_finer_than_submission_rejected( + self, submission_granularity, consumption_granularity ): - """Regression for the removed ``num_groups=1`` streaming override. - - With multiple weighted environments, ``num_groups=1`` hands the single - group to one environment and leaves the other with zero groups. It also - collapses ``agent_slots`` (computed without remainder distribution) to all - zeros, so ``np.gcd.reduce`` is 0 and the per-agent slot counts become - ``nan`` -- which stalls ``get_grouped_rollouts``. Keeping ``num_groups`` at - the trainer batch size (> 1) keeps every environment active with a valid, - non-zero slot count. - """ + gen = MockGenerator() + request = GroupedRolloutRequest( + num_groups=2, + rollouts_per_group=1, + inference_interface=MockInferenceInterface(), + submission_granularity=submission_granularity, + consumption_granularity=consumption_granularity, + ) + with pytest.raises(AssertionError, match="no finer"): + RolloutPipeline(gen, request, parallel_generation_tasks=1) + + def test_multi_env_layout_rejects_starving_batch_size(self): + """The layout raises rather than silently starving a weighted env.""" configs = [ AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), ] mt = WeightedMultiTask(configs) - - agent_groups = mt._distribute_counts(num_groups) - agent_slots = mt._distribute_counts(num_groups, distribute_remainder=False) - - assert all(groups > 0 for groups in agent_groups) is all_envs_active - if all_envs_active: - assert min(agent_slots) > 0 - assert np.gcd.reduce(agent_slots) > 0 - else: - assert min(agent_groups) == 0 - assert all(slots == 0 for slots in agent_slots) - assert np.gcd.reduce(agent_slots) == 0 + with pytest.raises(ValueError, match="starved"): + mt.rollout_allocations(1) + assert [a.num_groups for a in mt.rollout_allocations(8)] == [6, 2] + + # Evaluation-only envs take no groups and never count as starved. + mt = WeightedMultiTask( + configs + + [ + AgentConfig( + agent_type=MockGenerator, + agent_args={"env_id": "c"}, + weight=1.0, + evaluation_only=True, + ) + ] + ) + assert [a.num_groups for a in mt.rollout_allocations(8)] == [6, 2] def make_response(epochs, prompt_length, total_len, content="resp", finish_reason="stop"): @@ -589,12 +685,15 @@ async def test_run_episode(self, driver, max_turns, done_at_turn, scripted, expe groups = [] async def _drain(): - async for group in agent.get_grouped_rollouts( - GroupedRolloutRequest( - num_groups=1, rollouts_per_group=1, inference_interface=iface - ) - ): - groups.append(group) + request = GroupedRolloutRequest( + num_groups=1, rollouts_per_group=1, inference_interface=iface + ) + async with aclosing( + RolloutPipeline(agent, request, parallel_generation_tasks=1).run() + ) as iterator: + async for group in iterator: + groups.append(group) + break # Bounded so a wedged pipeline fails fast instead of hanging. await asyncio.wait_for(_drain(), timeout=5.0) From 487bb39488810d6e44a7b3540be42dabae793be0 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 13 Aug 2026 12:24:44 +0000 Subject: [PATCH 019/101] Avoid passing -1 sentinels into inference kernels (#6481) Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Jorge Albericio --- .../inference/contexts/dynamic_context.py | 9 +++++ .../inference/engines/test_dynamic_engine.py | 34 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 0b50f50d6de..4a5a2836e5e 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2474,6 +2474,15 @@ def initialize_attention_state( self._cpu_mha_block_table[:real_bs] = request_to_kv_block_ids_view[:real_bs] if real_bs < padded_bs: self._cpu_mha_block_table[real_bs:padded_bs] = self.kv_block_allocator.dummy_block_idx + # Real rows must avoid having a -1 sentinel in their trailing columns, + # because the kernel treats the -1 sentinel as a real value. + # We cannot avoid writing -1 into `request_to_kv_block_ids`; other logic needs it. + # The only option is to overwrite the -1 with a dummy block index via `masked_fill`. + if real_bs > 0: + _real_rows = self._cpu_mha_block_table[:real_bs] + # masked_fill_ over the pinned int32 view: no index_put/nonzero + # temporaries on the per-step CPU path. + _real_rows.masked_fill_(_real_rows < 0, self.kv_block_allocator.dummy_block_idx) # Max sequence lengths (Python scalars; consumed as kernel launch args). if not self.using_cuda_graph_this_step() and real_bs > 0: diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index af91409ef7b..0336881e22e 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1155,6 +1155,40 @@ def test_cuda_graph_padding_uses_dummy_block(self) -> None: assert (padded_rows != -1).all() assert (padded_rows == context.kv_block_allocator.dummy_block_idx).all() + @pytest.mark.internal + @pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" + ) + @torch.inference_mode() + def test_active_row_block_table_tail_uses_dummy_block(self) -> None: + """A real request's unallocated trailing block-table columns must be + staged as the dummy block, not the -1 sentinel: graphed decode + advertises max_seqlen_k = max_sequence_length, so the kernel's + page-table reach can include the tail of a near-limit request.""" + test_config = DynamicEngineTestConfig( + num_requests=1, + min_prompt_length=8, + max_prompt_length=8, + num_cuda_graphs=1, + context_max_requests=4, + max_sequence_length=512, + ) + env = self._build_test_env(test_config) + context = env.engine.context + + env.engine._add_request(env.requests[0]) + self._run_step(env) # prefill + self._run_step(env) # decode: graphed, one real row + + assert context.using_cuda_graph_this_step() + # The 512-token budget spans two 256-token pages. 8-token prompt allocates only the first. + # So the real row genuinely has an unallocated tail to pin. + block_count = int(context.request_kv_block_counts[0].item()) + staged_row = context._cpu_mha_block_table[0] + assert 0 < block_count < staged_row.numel() + assert (staged_row != -1).all() + assert (staged_row[block_count:] == context.kv_block_allocator.dummy_block_idx).all() + @pytest.mark.internal @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" From 09921c0a9e407154a493bde5ac9e9891a114e728 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 13 Aug 2026 12:25:49 +0000 Subject: [PATCH 020/101] RL: bound per-env wandb keys to the char limit (#6477) Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Jorge Albericio --- megatron/rl/rl_utils.py | 30 +++++++++++++++++++++++++++- tests/unit_tests/rl/test_rl_utils.py | 11 ++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 14551b685e9..19bf761c3ae 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -2,6 +2,7 @@ import copy import gc +import hashlib # Keep this to make the env registered. import itertools @@ -1006,6 +1007,28 @@ def compute_group_stats( return stats +def _bounded_artifact_key(key, limit=100): + """Bound a metric key so the artifact name wandb derives from it stays within wandb's limit. + + wandb rejects artifact names longer than 128 characters with a ValueError + (`NAME_MAXLEN` in wandb/sdk/artifacts/_validators.py) + + A table or plot logged under key K is stored as artifact 'run-{run.id}-{K}' + (13 chars of prefix with the default 8-char run id) + + wandb.plot.* charts log their backing table under '{K}_table' (6 more) + + keys containing characters outside [a-zA-Z0-9_.-] (env ids contain ':') are given + a 7-char CRC suffix by wandb's sanitizer. + + Worst-case key budget: 128 - 13 - 6 - 7 = 102 ~= 100. + """ + if len(key) <= limit: + return key + # 9 = 1 ('_' separator) + 8 (hex chars of the md5 digest). + digest = hashlib.md5(key.encode()).hexdigest()[:8] + return f'{key[: limit - 9]}_{digest}' + def prep_wandb_metrics( wandb_writer: wandb_run.Run, @@ -1376,7 +1399,12 @@ def maybe_log_training_metrics( tokenizer=tokenizer, ) for k, v in env_metrics.items(): - metrics[f"{env_id}_{k}"] = v + # Only table/plot values have a character length limit. + # Scalar keys are fine. + full_key = f"{env_id}_{k}" + if not isinstance(v, (int, float, bool)): + full_key = _bounded_artifact_key(full_key) + metrics[full_key] = v # Per-pipeline instrumentation (queue sizes, gate state, per-stage # timings) and the multi-task work distribution, collected on rank 0 diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 781d38c378f..ef098669da7 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -771,6 +771,17 @@ def test_advantage_calculation(self, num_turns, expected): advs = rl_utils.calculate_grpo_advantages([[-1, 1], [4, 4]], num_turns) torch.testing.assert_close(torch.tensor(advs), torch.tensor(expected), atol=1e-4, rtol=1e-5) + def test_bounded_artifact_key(self): + short = "env_rollout_table" + assert rl_utils._bounded_artifact_key(short) == short + long_key = "nemo_gym:" + "x" * 120 + "_staleness_hist" + bounded = rl_utils._bounded_artifact_key(long_key) + assert len(bounded) == 100 + assert bounded.startswith(long_key[:91]) + # Deterministic, and distinct keys stay distinct after bounding. + assert bounded == rl_utils._bounded_artifact_key(long_key) + assert bounded != rl_utils._bounded_artifact_key(long_key + "y") + @pytest.mark.parametrize( "scenario, expected_turn_lens, expected_traj_lens, expected_num_turns", [ From 211971993e48df60dfed65be369b7fd6ff7832a0 Mon Sep 17 00:00:00 2001 From: Skand Hurkat Date: Thu, 13 Aug 2026 16:14:52 +0000 Subject: [PATCH 021/101] Remove torch cuda empty cache (#5678) Signed-off-by: Skand Hurkat Co-authored-by: Claude Sonnet 4.6 (1M context) --- megatron/training/training.py | 6 ++---- .../nemotron3_super_release_gb200/model_config.yaml | 6 +++++- .../nemotron3_super_release_gb200_sm/model_config.yaml | 6 +++++- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index 0e9417b8fbe..14d753ff6cd 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3166,13 +3166,11 @@ def save_checkpoint_and_time( # Log E2E metrics before save-checkpoint one_logger_utils.track_e2e_metrics() - # Free overlap param-gather buffers and release cached GPU memory so - # that the async checkpoint worker process has enough GPU headroom for - # D2H tensor transfers. + # Free overlap param-gather buffers so that the async checkpoint worker + # process has enough GPU headroom for D2H tensor transfers. for model_chunk in model: if hasattr(model_chunk, 'free_overlap_buffers'): model_chunk.free_overlap_buffers() - torch.cuda.empty_cache() # timer.log() reports the min & max time. We do not need a barrier here. timer_key = 'save-checkpoint-non-persistent' if non_persistent_ckpt else 'save-checkpoint' diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200/model_config.yaml b/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200/model_config.yaml index 321f9d2ca9c..9b925ae44ca 100644 --- a/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200/model_config.yaml +++ b/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200/model_config.yaml @@ -2,7 +2,7 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 32 NCCL_GRAPH_REGISTER: 0 NCCL_NVLS_ENABLE: 0 - PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True,garbage_collection_threshold:0.8" TORCH_NCCL_AVOID_RECORD_STREAMS: 1 TORCH_NCCL_HIGH_PRIORITY: 1 NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: 64 @@ -140,6 +140,10 @@ MODEL_ARGS: --ckpt-fully-parallel-load: true --ckpt-assume-constant-structure: true --dist-ckpt-strictness: log_all + --async-save: true + --async-strategy: nvrx + --async-ckpt-use-cpu-shm: true + --use-persistent-ckpt-worker: true --save-interval: 200 --save-retain-interval: 2000 diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200_sm/model_config.yaml b/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200_sm/model_config.yaml index 430de2aafc4..e0c9a59d88b 100644 --- a/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200_sm/model_config.yaml +++ b/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200_sm/model_config.yaml @@ -2,7 +2,7 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 32 NCCL_GRAPH_REGISTER: 0 NCCL_NVLS_ENABLE: 0 - PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True,garbage_collection_threshold:0.8" TORCH_NCCL_AVOID_RECORD_STREAMS: 1 TORCH_NCCL_HIGH_PRIORITY: 1 NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: 64 @@ -140,6 +140,10 @@ MODEL_ARGS: --ckpt-fully-parallel-load: true --ckpt-assume-constant-structure: true --dist-ckpt-strictness: log_all + --async-save: true + --async-strategy: nvrx + --async-ckpt-use-cpu-shm: true + --use-persistent-ckpt-worker: true --save-interval: 200 --save-retain-interval: 2000 From 87cc590ba03b7490ca2ab27dfc2372cd4d24677d Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Thu, 13 Aug 2026 16:52:28 +0000 Subject: [PATCH 022/101] [fix] GTP+DDP potential data-race: publish DDP params before GTP prefetch reads them (#6388) Signed-off-by: Shiqing Fan Co-authored-by: Jiangfei Duan --- .../core/generalized_tensor_parallel.md | 37 +- .../distributed/distributed_data_parallel.py | 100 +++- .../generalized_tensor_parallelism.py | 20 +- megatron/core/utils.py | 27 +- .../distributed/test_param_readiness.py | 135 ++++++ .../test_gtp_basics.py | 47 ++ .../test_gtp_ddp_param_sync_race.py | 433 ++++++++++++++++++ 7 files changed, 778 insertions(+), 21 deletions(-) create mode 100644 tests/unit_tests/distributed/test_param_readiness.py create mode 100644 tests/unit_tests/generalized_tensor_parallel/test_gtp_ddp_param_sync_race.py diff --git a/docs/api-guide/core/generalized_tensor_parallel.md b/docs/api-guide/core/generalized_tensor_parallel.md index eca9e2b7eb8..379c77f0bde 100644 --- a/docs/api-guide/core/generalized_tensor_parallel.md +++ b/docs/api-guide/core/generalized_tensor_parallel.md @@ -57,6 +57,9 @@ Both `GTP_remat` collectives are prefetched one step ahead, so they overlap the - [wgrad-before-dgrad schedule *(deferred to a follow-up MR)*](#wgrad-before-dgrad-schedule--deferred-to-a-follow-up-mr) - [Recompute-forward prefetch chain *(GTP\_remat + activation recompute)*](#recompute-forward-prefetch-chain--gtp_remat--activation-recompute) - [3.2 DDP buckets with (E)GTP\_remat](#32-ddp-buckets-with-egtp_remat) + - [Ordering invariants](#ordering-invariants) + - [Why this design works](#why-this-design-works) + - [Bucketing and gradient scaling](#bucketing-and-gradient-scaling) - [3.3 Distributed checkpointing (DCP)](#33-distributed-checkpointing-dcp) - [3.4 Prefetch-chain construction and its design assumptions](#34-prefetch-chain-construction-and-its-design-assumptions) - [Grouped-expert chains (one-block-ahead)](#grouped-expert-chains-one-block-ahead) @@ -482,17 +485,40 @@ Under **full-iteration CUDA graphs** the recompute-forward is captured; `wait_as 2. **`is_gtp_weight_remat` / `allreduce` tags** propagated onto the optimizer's master shards — consumed only by the grad-norm dedup filter. 3. **grad-ready hook routing** (`DistributedDataParallel.__init__`) — for a GTP_remat param, DDP registers its backward post-hook via GTP_remat's `register_grad_accum_hook` instead of autograd's `AccumulateGrad`. GTP_remat fires it from `_handle_megatron_grad_accum` **after** the per-param `{wgrad RS → main_grad add}`. This enforces the invariant below; a no-op (plain autograd path) when GTP_remat is inactive. -> **Ordering invariant.** A bucket's DDP gradient reduction (the reduce-scatter / all-to-all + local fp32 accumulation) runs **strictly after every GTP_remat param in that bucket has finished `{GTP_remat wgrad RS → main_grad add}`**. `register_grad_ready` only fires the bucket collective once *all* its params are ready, and for GTP_remat params "ready" is signalled by GTP_remat after the add — never by autograd's `AccumulateGrad`, which (because the wgrad RS is async and its `main_grad` accumulation is deferred to a later backward node) can fire **before** the add and would make the bucket reduce read a stale/empty `main_grad` (notably under `reduce_scatter_with_fp32_accumulation`). +#### Ordering invariants -Everything else — bucketing, the reduce-scatter/all-reduce schedule and its overlap, master-state sharding, grad clipping, the checkpoint format — is unchanged and unaware of GTP_remat. +> **Ordering invariant (gradients).** A bucket's DDP gradient reduction (the reduce-scatter / all-to-all + local fp32 accumulation) runs **strictly after every GTP_remat param in that bucket has finished `{GTP_remat wgrad RS → main_grad add}`**. `register_grad_ready` only fires the bucket collective once *all* its params are ready, and for GTP_remat params "ready" is signalled by GTP_remat after the add — never by autograd's `AccumulateGrad`, which (because the wgrad RS is async and its `main_grad` accumulation is deferred to a later backward node) can fire **before** the add and would make the bucket reduce read a stale/empty `main_grad` (notably under `reduce_scatter_with_fp32_accumulation`). + +**Parameter publication under `--overlap-param-gather`.** DDP publishes a bucket group *lazily*: its parameter all-gather (and post-gather quantize, `_post_param_sync`) is drained from the forward pre-hook of a module owning one of that bucket's parameters. GTP_remat consumes weights *ahead* of that module (§3.4), so when the consumed weight and the prefetch target sit in different bucket groups, the target may not be published yet. GTP_remat therefore asks for it first, at the top of `_all_gather_weight`: + +```text +time --------------------------------------------------------------------------> +compute pre-hook(fc0) -> GEMM(w0) -> [stall] -------> pre-hook(fc1): no-op + | ensure_params_ready(w1) + v +DDP AG +---------------- AG(bucket B1) --->| quantize +GTP ag_strm +-- AG(w1) reads FRESH --> GEMM(w1) +``` + +The request goes through a **backend-agnostic** hook: -**Why this matters:** +- **Contract** (`megatron/core/utils.py`): a backend attaches a zero-argument callable under `PARAM_READY_CALLBACK_ATTR`; a consumer reading `param.data` outside the owning module's pre-hook calls `ensure_params_ready(params)` first. Unmarked params no-op, so the contract is open to FSDP or any other backend. +- **Neither side names the other**: DDP registers one `_BucketParamReadyCallback` per bucket group (weakly held), knowing nothing of GTP_remat; GTP_remat calls `ensure_params_ready`, knowing nothing of DDP. +- **Forward only**: backward re-reads what forward published, and recompute runs *inside* backward, where publishing could gather into the buffer that aliases grads under `--reuse-grad-buf-for-mxfp8-param-ag`. +- **Cost**: publishing early can *start* an undispatched gather and chain-dispatch the next bucket, draining it about one block earlier than the pre-hook would — so that bucket loses some gather/compute overlap. Negligible when the gather was already dispatched early, as `--align-param-gather` and `--overlap-param-gather-with-optimizer-step` do. +- **Not covered — CUDA-graph capture**: no collective may be issued during capture, so the callback no-ops and the captured gather carries no dependency on DDP's. A consumer that captures its reads must publish before launching. + +#### Why this design works + +Everything else — bucketing, the reduce-scatter/all-reduce schedule and its overlap, master-state sharding, grad clipping, the checkpoint format — is unchanged and unaware of GTP_remat. - **Free reuse of a mature stack.** GTP_remat inherits DDP's bucketing + comm/compute overlap, the distributed optimizer's fp32-master + Adam-moment sharding, grad-norm/clip, and the existing checkpoint format — no parallel re-implementation to write or maintain (contrast FSDP, which replaces all of these). - **Orthogonal composability.** Because GTP_remat is a rank-grid sub-axis cut along `out_features` (dim 0, whichever axis TP used), it composes with TP/EP/CP/PP and the DistOpt the same way TP does — no special nesting logic. - **Zero-cost when off.** With GTP_remat disabled the gtp_remat axis is size-1 and the hooks become no-ops, so non-GTP_remat runs hit byte-identical behavior — GTP_remat can be toggled without forking the DDP/optimizer code paths. - **Small, auditable surface.** These three hooks are the whole integration contract, which is what makes the correctness argument below tractable. +#### Bucketing and gradient scaling + DDP groups parameters into **two buffers** by `is_expert_parallel` (MoE tag) — a dense buffer and an expert buffer. GTP_remat/EGTP_remat shards are **merged into** these buffers like ordinary params (no separate GTP_remat/EGTP_remat buckets): they reduce over the replicate group (the default `intra_dp_cp_group` / `intra_expt_dp_group`). The DP collective only covers the replicate axis; the gtp_remat axis is completed separately, and **how both axes are scaled depends on the loss normalization** (`config.calculate_per_token_loss`). In all cases each gtp_remat contribution is summed exactly once: @@ -732,7 +758,7 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall | Test file | What it guards | |-----------|----------------| -| `test_gtp_basics.py` | Core GTP_remat shard/gather, cache ownership, wgrad ring, and DDP bucket alignment. | +| `test_gtp_basics.py` | Core GTP_remat shard/gather, cache ownership, wgrad ring, and DDP bucket alignment. Also pins TE's recompute-phase flag as dtype-agnostic (true under BF16, not just FP8), which the forward-only readiness gate in §3.2 relies on. | | `test_attention_gtp.py` | GTP_remat on attention linears, loss parity vs no-GTP_remat. | | `test_mamba_gtp.py` | GTP_remat on Mamba projection weights. | | `test_tp_gtp.py` | GTP_remat composed with tensor parallelism (`tp_group × gtp_remat_group`). | @@ -746,8 +772,11 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall | `test_gtp_recompute_chain.py` | Recompute-chain buffers (§3.1): adjacent nodes never share a gather buffer, dense and grouped, plus dgrad/wgrad parity vs no-recompute. | | `test_gtp_mtp.py` | GTP_remat + MTP shared weights (§3.5), 14 cases over `mtp_use_repeated_layer` × dense/MoE. Both MTP hazards are silent, so each needs its own guard: the async reduce-scatter path is compared numerically against the sync path on an identical model/sharding/batch, and all-gathers issued are tallied against consumes to catch a consume reading a buffer nothing gathered into. | | `test_gtp_fp8_param_gather.py` | Native-FP8 GTP_remat (§1.3): fp8-vs-BF16 loss parity (TP1/TP2, MoE), post-save-spike guard. | +| `test_gtp_ddp_param_sync_race.py` | Parameter-readiness ordering (§3.2): GTP_remat's ahead-of-consume prefetch must not read a bucket DDP has not published. Structural and numerical (stale-value) guards on the default one-weight-ahead chain, the grouped-expert one-block-ahead chain, and the recompute exclusion. | | `test_gtp_custom_pgs.py` | `pg_collection` plumbing: a custom `gtp_remat` group (permuted ranks, same size) must give the same fwd/bwd results as the MPU groups — catches modules reading `parallel_state` instead of the collection passed to them. | The fp32-accumulation primitive itself is covered outside this suite, by `tests/unit_tests/distributed/test_reduce_scatter_with_fp32_accumulation.py`, which does not require GTP_remat. +The parameter-readiness contract itself (§3.2) is likewise covered outside this suite, by `tests/unit_tests/distributed/test_param_readiness.py` — CPU-only, no GPU or GTP_remat required. It pins the branches the 4-GPU test does not exercise: `align_param_gather`, pre-hooks removed mid-sequence, and a collected DDP or bucket group. + All tests require ≥ 4 GPUs and TransformerEngine >= 2.19; they self-skip when those are unavailable. A green run (skips for unmet hardware/config are acceptable) is the minimum bar for any GTP_remat change. diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py index 81cdbeb3303..763368a9814 100644 --- a/megatron/core/distributed/distributed_data_parallel.py +++ b/megatron/core/distributed/distributed_data_parallel.py @@ -1,6 +1,7 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. import logging +import weakref from contextlib import contextmanager from typing import Optional @@ -11,7 +12,7 @@ from ..process_groups_config import ProcessGroupCollection from ..transformer.cuda_graphs import is_graph_capturing from ..transformer.transformer_config import TransformerConfig -from ..utils import log_single_rank +from ..utils import PARAM_READY_CALLBACK_ATTR, log_single_rank from .data_parallel_base import _BaseDataParallel from .distributed_data_parallel_config import DistributedDataParallelConfig from .param_and_grad_buffer import _ParamAndGradBuffer, group_params_for_buffers, partition_buckets @@ -19,6 +20,64 @@ logger = logging.getLogger(__name__) +class _BucketParamReadyCallback: + """Publishes one bucket group's parameters on demand. + + DDP's side of ``megatron.core.utils.ensure_params_ready``: one instance is stored on every + parameter of the group under ``PARAM_READY_CALLBACK_ATTR``, and calling it makes those + parameters readable. For consumers reading ``param.data`` ahead of the owning module's + pre-hook. + + Publishing may START an undispatched gather, not just wait on one. Holds DDP and the bucket + group weakly; the callback outlives this DDP on re-wrap. No-ops under CUDA-graph capture, so + a consumer that captures its reads must replay this itself. + """ + + def __init__(self, ddp: 'DistributedDataParallel', bucket_group) -> None: + self._ddp = weakref.ref(ddp) + self._bucket_group = weakref.ref(bucket_group) + + def __call__(self) -> None: + bucket_group = self._bucket_group() + if bucket_group is None: + return + + # HOT PATH: already published. True for every microbatch after the first one of an + # iteration, so it precedes the DDP deref and the graph-capture query. + if bucket_group.param_gather_dispatched and bucket_group.param_gather_handle is None: + return + + ddp = self._ddp() + if ddp is None: + # Weakref is dead: the DDP object was garbage-collected while this callback lived on + # (it is stored on the parameters, which outlive the wrapper). Its buffers are gone. + return + + if is_graph_capturing(): + # A captured collective re-runs on EVERY replay: once per microbatch, not once per + # iteration. Consumers that capture their reads must publish before launching. + return + + # No pre-hooks installed means the caller removed them and now drives param sync itself + # (``disable_forward_pre_hook(param_sync=False)``), so the schedule is not ours to touch. + ddp_owns_schedule = bool(ddp.remove_forward_pre_hook_handles) + + if bucket_group.param_gather_handle is not None: + # A gather is in flight over the very buffer we are about to read, so it must be + # waited on. finish_param_sync() does two things: wait for THIS bucket, then start + # the NEXT bucket's gather. The wait is mandatory; starting the next bucket is a + # scheduling decision, so only let it happen while DDP still owns the schedule. + if ddp_owns_schedule: + ddp._finish_param_sync_for_bucket_group(bucket_group) + else: + bucket_group.finish_param_sync(skip_next_bucket_dispatch=True) + elif ddp_owns_schedule: + # No handle and not published (the hot path returned above) => never dispatched, so + # publishing STARTS a gather. Only safe while DDP owns the schedule. + assert not bucket_group.param_gather_dispatched + ddp._finish_param_sync_for_bucket_group(bucket_group) + + class DistributedDataParallel(_BaseDataParallel): """ DDP wrapper which stores grads in contiguous buffers. Also has option of overlapping @@ -334,9 +393,24 @@ def __init__( # Create map from param to bucket group, used in pre_hook. for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]: for bucket_group in bucket_groups: + # One readiness callback per bucket group, stamped on all its params, so a + # consumer reading param.data outside the owning module's pre-hook can publish it + # first. Backend-agnostic: DDP never learns which consumers use it. + ready_callback = ( + _BucketParamReadyCallback(self, bucket_group) + if self.ddp_config.overlap_param_gather + else None + ) for bucket in bucket_group.buckets: for param in bucket.params_list: self.param_to_bucket_group[param] = bucket_group + if ready_callback is not None: + setattr(param, PARAM_READY_CALLBACK_ATTR, ready_callback) + elif hasattr(param, PARAM_READY_CALLBACK_ATTR): + # Re-wrapping a model chunk: a previous DDP may have left a marker + # pointing at ITS bucket group. This DDP owns the parameter now and + # publishes nothing, so the stale callback must go. + delattr(param, PARAM_READY_CALLBACK_ATTR) # Delete references to weight_tensor if they exist since we don't want two parameter copies # if we re-mapped parameters (which happens when we use the distributed optimizer). @@ -455,21 +529,21 @@ def hook(module, *unused): if param not in self.param_to_bucket_group: continue assert param.requires_grad - - # If aligning param all-gather across pipeline stages, all-gather is dispatched - # by start_param_sync calls in core/pipeline_parallelism/schedules.py. - # If overlapping param all-gather with optimizer step, then all-gather has - # already been dispatched in optimizer step. - skip_next_bucket_dispatch = ( - self.ddp_config.align_param_gather - or self.overlap_param_gather_with_optimizer_step - ) - self.param_to_bucket_group[param].finish_param_sync( - skip_next_bucket_dispatch=skip_next_bucket_dispatch - ) + self._finish_param_sync_for_bucket_group(self.param_to_bucket_group[param]) return hook + def _finish_param_sync_for_bucket_group(self, bucket_group): + """Drain one bucket group's param all-gather and run its post-all-gather processing.""" + # If aligning param all-gather across pipeline stages, all-gather is dispatched + # by start_param_sync calls in core/pipeline_parallelism/schedules.py. + # If overlapping param all-gather with optimizer step, then all-gather has + # already been dispatched in optimizer step. + skip_next_bucket_dispatch = ( + self.ddp_config.align_param_gather or self.overlap_param_gather_with_optimizer_step + ) + bucket_group.finish_param_sync(skip_next_bucket_dispatch=skip_next_bucket_dispatch) + def _make_backward_post_hook(self, param: torch.nn.Parameter): """ Creates a backward post-hook to dispatch an all-reduce / reduce-scatter when diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index b8d2a8faff2..ff84be77a38 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -42,7 +42,7 @@ register_capture_comm, register_capture_wgrad_ring_slot, ) -from megatron.core.utils import log_single_rank +from megatron.core.utils import ensure_params_ready, log_single_rank logger = logging.getLogger(__name__) @@ -99,6 +99,10 @@ nvtx_range_pop = nvtx_range_push = round_up_to_nearest_multiple = MagicMock() HAVE_TE = False +# TE's name says fp8, but the flag it returns is set for BF16 activation recompute too (the +# FP8-gated helper is `is_fp8_activation_recompute_enabled`). Alias it so call sites read right. +in_activation_recompute_phase = in_fp8_activation_recompute_phase + class GTPChain(str, Enum): """Prefetch chain identifier for n GTPShardedParam. @@ -1285,10 +1289,20 @@ def _all_gather_weight( weights = self._weights + # 0. Wait for DDP's asynchronous parameter all-gather (and its post-gather quantize), if + # any, to finish on the shards we are about to read. GTP prefetches ahead of the module + # that owns the target weight, so that module's DDP pre-hook may not have run yet and + # the shard can still hold last iteration's values. + # + # Forward only: backward re-reads what forward published, and recompute runs inside + # backward where a dispatch could gather into the grad-aliased buffer. + if fwd and not in_activation_recompute_phase(): + ensure_params_ready(weights) + # 1. Transition state for async gathers. Skip during recompute-forward: it gathers # rowwise (_ag_ticket_fwd) while a bwd-chain prefetch may hold an in-flight columnwise # AG state (_ag_ticket_bwd) on the same weight — clobbering breaks the dgrad consume. - if GTP_CONFIG.check_param_states and not in_fp8_activation_recompute_phase(): + if GTP_CONFIG.check_param_states and not in_activation_recompute_phase(): new_state = GTPWeightState.ASYNC_WAIT if async_op else GTPWeightState.DATA_READY_SYNC for w in weights: w._set_state(new_state) @@ -1595,7 +1609,7 @@ def all_gather_and_prefetch(self, fwd: bool = True, nvtx_label: str = None): # During an activation-recompute forward (runs in backward), route consume + # prefetch through the recompute-forward chain on its own _recompute_* slot # (see __init__) instead of the fwd/bwd chains; lazy-built below. - in_recompute = in_fp8_activation_recompute_phase() + in_recompute = in_activation_recompute_phase() use_recompute_chain = in_recompute and GTP_CONFIG.weight_prefetch # Reaching a forward gather proves the previous backward finished, so the chains it diff --git a/megatron/core/utils.py b/megatron/core/utils.py index c6b83f00d30..940bff83beb 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -24,7 +24,7 @@ from functools import lru_cache, reduce, wraps from importlib.metadata import version from types import TracebackType -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union import numpy import torch @@ -3163,3 +3163,28 @@ def deprecate_inference_params(inference_context, inference_params): ) return inference_params return inference_context + + +#: Attribute a parameter-sharding backend sets on each parameter it publishes asynchronously. +#: See :func:`ensure_params_ready`. +PARAM_READY_CALLBACK_ATTR = "_ensure_param_ready_callback" + + +def ensure_params_ready(params: Iterable[Any]) -> None: + """Make ``params`` readable now, finishing any outstanding backend publication. + + A parameter-sharding backend (DDP with ``overlap_param_gather``, FSDP, ...) publishes values + asynchronously, so only the owning module's forward pre-hook makes ``param.data`` valid. Any + consumer reading it earlier -- ahead of that module, or from another stream -- calls this + first. Backends mark their params with :data:`PARAM_READY_CALLBACK_ATTR`; unmarked params + no-op, so neither side needs to know about the other. + + Callbacks are shared per communication bucket, so each fires once, not once per parameter. + """ + fired = set() + for param in params: + callback = getattr(param, PARAM_READY_CALLBACK_ATTR, None) + if callback is None or id(callback) in fired: + continue + fired.add(id(callback)) + callback() diff --git a/tests/unit_tests/distributed/test_param_readiness.py b/tests/unit_tests/distributed/test_param_readiness.py new file mode 100644 index 00000000000..b66cfdebe00 --- /dev/null +++ b/tests/unit_tests/distributed/test_param_readiness.py @@ -0,0 +1,135 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Parameter readiness: ``ensure_params_ready`` and DDP's ``_BucketParamReadyCallback``. + +The GPU test in ``generalized_tensor_parallel/`` runs one configuration: healthy DDP, pre-hooks +enabled, ``align_param_gather=False``. These cover the branches it never exercises -- +``align_param_gather=True``, pre-hooks removed mid-sequence, a collected DDP or bucket group, +and out-of-order bucket dedup -- which would be contrived to set up on real GPUs. +""" + +from types import SimpleNamespace + +import torch + +from megatron.core.distributed.distributed_data_parallel import ( + DistributedDataParallel, + _BucketParamReadyCallback, +) +from megatron.core.utils import PARAM_READY_CALLBACK_ATTR, ensure_params_ready + + +class _FakeBucketGroup: + def __init__(self): + self.finished = [] + self.param_gather_dispatched = False + self.param_gather_handle = None + + def finish_param_sync(self, skip_next_bucket_dispatch=False): + self.finished.append(skip_next_bucket_dispatch) + self.param_gather_dispatched = True + self.param_gather_handle = None + + +def _fake_ddp(align_param_gather=False, hooks_enabled=True): + ddp = object.__new__(DistributedDataParallel) + ddp.ddp_config = SimpleNamespace(align_param_gather=align_param_gather) + ddp.overlap_param_gather_with_optimizer_step = False + ddp.remove_forward_pre_hook_handles = {object(): object()} if hooks_enabled else {} + return ddp + + +class TestParamReadinessProtocol: + def test_no_marker_is_a_noop(self): + """A parameter no backend owns must be readable without any callback.""" + param = torch.nn.Parameter(torch.zeros(2)) + assert not hasattr(param, PARAM_READY_CALLBACK_ATTR) + ensure_params_ready([param]) # must not raise + + def test_callbacks_are_deduplicated_per_bucket(self): + """Backends share one callback per bucket, so it must fire once per call, in order.""" + calls = [] + shared = lambda: calls.append("shared") # noqa: E731 + other = lambda: calls.append("other") # noqa: E731 + + first, second, third = (torch.nn.Parameter(torch.zeros(1)) for _ in range(3)) + setattr(first, PARAM_READY_CALLBACK_ATTR, shared) + setattr(second, PARAM_READY_CALLBACK_ATTR, shared) + setattr(third, PARAM_READY_CALLBACK_ATTR, other) + plain = torch.nn.Parameter(torch.zeros(1)) + + ensure_params_ready([first, second, plain, first, third]) + + assert calls == ["shared", "other"] + + +class TestBucketParamReadiness: + """Each test must hold a strong reference to the fake DDP. + + The callback stores DDP weakly, so letting ``_fake_ddp()`` stay a temporary would collect it + immediately and make every callback early-return -- turning these into vacuous passes. + """ + + def test_publishes_once_then_is_idempotent(self): + ddp = _fake_ddp() + bucket_group = _FakeBucketGroup() + ready_callback = _BucketParamReadyCallback(ddp, bucket_group) + + ready_callback() + assert bucket_group.finished == [False] + + # Already published this iteration: no second collective. + ready_callback() + assert bucket_group.finished == [False] + + def test_in_flight_gather_is_drained(self): + ddp = _fake_ddp() + bucket_group = _FakeBucketGroup() + bucket_group.param_gather_dispatched = True + bucket_group.param_gather_handle = object() + + _BucketParamReadyCallback(ddp, bucket_group)() + assert bucket_group.finished == [False] + + def test_align_param_gather_skips_next_bucket_dispatch(self): + ddp = _fake_ddp(align_param_gather=True) + bucket_group = _FakeBucketGroup() + _BucketParamReadyCallback(ddp, bucket_group)() + assert bucket_group.finished == [True] + + def test_disabled_forward_hooks_do_not_start_a_gather(self): + """With the pre-hooks removed the caller drives param sync; starting one here could + gather into a buffer it is actively using.""" + ddp = _fake_ddp(hooks_enabled=False) + bucket_group = _FakeBucketGroup() + _BucketParamReadyCallback(ddp, bucket_group)() + assert bucket_group.finished == [] + + def test_in_flight_gather_is_drained_even_with_hooks_disabled(self): + """``disable_forward_pre_hook(param_sync=False)`` leaves hooks removed without a force + sync, so an outstanding all-gather still has to be drained before the buffer is read -- + but without chaining the next bucket's dispatch, which is the caller's schedule now.""" + ddp = _fake_ddp(hooks_enabled=False) + bucket_group = _FakeBucketGroup() + bucket_group.param_gather_dispatched = True + bucket_group.param_gather_handle = object() + + _BucketParamReadyCallback(ddp, bucket_group)() + assert bucket_group.finished == [True] # drained, next dispatch skipped + + def test_dead_bucket_group_is_a_noop(self): + """The bucket group is held weakly so the callback, which lives on parameters that + outlive DDP, cannot retain the bucket-group graph.""" + ddp = _fake_ddp() + ready_callback = _BucketParamReadyCallback(ddp, _FakeBucketGroup()) + assert ready_callback._bucket_group() is None + ready_callback() # must not raise + + def test_dead_ddp_is_a_noop(self): + """The callback holds DDP weakly, so a collected DDP must not resurrect work.""" + bucket_group = _FakeBucketGroup() + ready_callback = _BucketParamReadyCallback(_fake_ddp(), bucket_group) + # No strong reference above, so the referent is already gone. + assert ready_callback._ddp() is None + ready_callback() + assert bucket_group.finished == [] diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py index 6c594d0be62..85e9b942d30 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py @@ -1475,3 +1475,50 @@ def test_rs_input_copies_into_ring(self, monkeypatch): assert rs_input is weights[1]._gtp_graph_wgrad_ring_slot.tensor torch.testing.assert_close(rs_input[:4], wgrad) assert torch.count_nonzero(rs_input[4:]) == 0 + + +class TestActivationRecomputePhaseFlag: + """GTP's forward-only readiness gate rests on TE's recompute flag being dtype-agnostic.""" + + def test_recompute_phase_flag_is_set_under_bf16(self): + """``in_activation_recompute_phase`` must be True for BF16 recompute, not just FP8. + + GTP aliases TE's ``in_fp8_activation_recompute_phase`` under a name without ``fp8`` + because TE assigns the underlying global unconditionally. ``_all_gather_weight`` uses it + to skip publishing during recompute (see + ``test_gtp_ddp_param_sync_race.py::test_recompute_forward_does_not_request_publication``). + If TE ever made it FP8-only, BF16 recompute would start asking DDP to publish inside + backward -- silently, since the failure is a wrong-buffer gather rather than an error. + """ + _requires_multi_gpu(1) + import transformer_engine.pytorch as te + from transformer_engine.pytorch.fp8 import FP8GlobalStateManager + + assert not FP8GlobalStateManager.is_fp8_enabled(), "must run outside fp8_autocast" + + hidden, dtype = 128, torch.bfloat16 + observed = [] + + class _Probe(torch.nn.Module): + def __init__(self): + super().__init__() + self.lin = te.Linear(hidden, hidden, bias=False, params_dtype=dtype, device="cuda") + + def forward(self, x): + observed.append(gtp_module.in_activation_recompute_phase()) + return self.lin(x) + + x = torch.randn(8, hidden, dtype=dtype, device="cuda", requires_grad=True) + out = te.checkpoint( + _Probe(), + x, + distribute_saved_activations=False, + get_rng_state_tracker=None, + tp_group=None, + ) + out.sum().backward() # replays the forward -> second observation + + assert observed == [False, True], ( + f"expected [original fwd=False, recompute fwd=True], got {observed}: " + "in_activation_recompute_phase no longer tracks BF16 activation recompute" + ) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_ddp_param_sync_race.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_ddp_param_sync_race.py new file mode 100644 index 00000000000..74e31bf7e49 --- /dev/null +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_ddp_param_sync_race.py @@ -0,0 +1,433 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""GTP ahead-of-consume weight prefetch vs. DDP's asynchronous parameter all-gather. + +With ``overlap_param_gather=True`` DDP publishes a bucket group only from the forward pre-hook +of a module owning one of its parameters. GTP prefetches ahead of that module, and its gather +input is ``next_w.data`` -- a view into the bucket's ``param_data``. So when the consumed weight +and the prefetch target sit in different bucket groups, GTP reads a buffer whose all-gather is +still in flight or not yet dispatched, and whose post-gather quantize has not run: it gathers +last iteration's weights. + +This is the base prefetch, not the grouped-expert optimization. Default chains look one weight +ahead, grouped-expert chains one MoE block; reading further ahead just widens the window. + +Tests use the real DDP + distributed-optimizer path and none depend on timing: + +- ``..._prefetch_target_bucket_is_not_synced`` -- was the target's bucket published when read? +- ``..._prefetch_reads_stale_weight_values`` -- did the read actually see stale values? +- ``..._grouped_expert_one_block_ahead_prefetch`` -- same, on the widest (grouped) window. +- ``..._recompute_forward_does_not_request_publication`` -- recompute must not ask DDP to publish. + (The TE flag that gate depends on is pinned in ``test_gtp_basics.py``.) +""" + +import contextlib + +import pytest + +from megatron.core.tensor_parallel.gtp_api import HAVE_GTP + +if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.19", allow_module_level=True) + +import torch + +from megatron.core import parallel_state as ps +from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig +from megatron.core.distributed.param_and_grad_buffer import shard_buffer +from megatron.core.tensor_parallel import generalized_tensor_parallelism as gtp_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import PARAM_READY_CALLBACK_ATTR +from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import ( # noqa: F401 + _make_gtp_linear, + _make_gtp_remat_grouped_linear, + _requires_multi_gpu, + _run_distributed, + _torchrun_dist_init, + reset_fp8_state, + reset_gtp_globals, +) + +WORLD = 4 +GTP = 2 +HIDDEN = 128 +SEQ = 8 +BATCH = 2 +DTYPE = torch.bfloat16 + +# Each GTP shard is (HIDDEN / GTP) x HIDDEN = 8192 elements, so any bucket size below that +# forces one bucket group per weight -- which is what puts the prefetch target outside the +# bucket the consuming module's pre-hook drains. +BUCKET_SIZE = 4096 +NUM_EXPERTS = 2 + + +def _make_config(): + return TransformerConfig( + num_layers=1, hidden_size=HIDDEN, num_attention_heads=8, bf16=True, params_dtype=DTYPE + ) + + +def _build_ddp_two_gtp_layers(): + """Two GTP-sharded ``te.Linear`` layers in separate DDP bucket groups. + + Returns ``(ddp_model, [w0, w1])`` where ``w0``/``w1`` are the ``GTPShardedParam``s. + """ + gtp_group = ps.get_gtp_weight_remat_group() + fc0 = _make_gtp_linear(HIDDEN, HIDDEN, gtp_group, dtype=DTYPE) + fc1 = _make_gtp_linear(HIDDEN, HIDDEN, gtp_group, dtype=DTYPE) + + module = torch.nn.Sequential() + module.add_module("fc0", fc0) + module.add_module("fc1", fc1) + + ddp_config = DistributedDataParallelConfig( + use_distributed_optimizer=True, + overlap_param_gather=True, + overlap_grad_reduce=True, + bucket_size=BUCKET_SIZE, + ) + ddp_model = DistributedDataParallel(_make_config(), ddp_config, module) + weights = [fc0.weight, fc1.weight] + # Covers DDP's attach loop against a real DistributedDataParallel. The CPU protocol tests + # build the callback by hand, so without this nothing would catch the registration silently + # not running -- which would make every readiness call a no-op and the fix ineffective. + for w in weights: + assert hasattr(w, PARAM_READY_CALLBACK_ATTR), ( + f"DDP did not attach {PARAM_READY_CALLBACK_ATTR} to a bucketed parameter; the " + "readiness protocol is not wired up" + ) + return ddp_model, weights + + +def _forward(ddp_model, rank): + torch.manual_seed(1234 + rank) + x = torch.randn(SEQ * BATCH, HIDDEN, dtype=DTYPE, device="cuda") + out = x + for layer in ddp_model.module.children(): + out = layer(out) + return out + + +def _param_names(ddp_model): + """``id(param) -> name``; ``GTPShardedParam._debug_name`` is empty on this path.""" + return {id(p): name for name, p in ddp_model.module.named_parameters()} + + +def _record_prefetches(ddp_model, targets): + """Log DDP bucket state at the exact instant GTP reads a prefetched weight's storage. + + The measurement point is ``get_padded_shard`` -- the call that produces the all-gather + input (``generalized_tensor_parallelism.py``, "build gather inputs") -- not the entry to + ``_all_gather_weight``. Anything the gather does to publish the shard beforehand therefore + counts, which keeps the test measuring the read rather than a particular fix's placement. + Scoped to asynchronous forward gathers, i.e. one-block-ahead prefetches. + + Covers the BF16 path only: with ``--fp8-param-gather`` the gather input is ``w.quantized``, + not ``get_padded_shard()``, so this probe would record nothing there. The fix itself is + path-independent (it runs before the input is built either way), but the native-FP8 + configuration is not regression-guarded by this test. + + Returns ``(records, restore)``. + """ + cls = gtp_module.GTPShardedParam + original_gather = cls._all_gather_weight + original_shard = cls.get_padded_shard + records = [] + watched = {id(p) for p in targets} + names = _param_names(ddp_model) + inside_prefetch = [False] + + def patched_gather(self, *args, **kwargs): + async_op = kwargs.get("async_op", args[0] if args else False) + fwd = kwargs.get("fwd", args[1] if len(args) > 1 else True) + previous = inside_prefetch[0] + inside_prefetch[0] = bool(async_op and fwd) + try: + return original_gather(self, *args, **kwargs) + finally: + inside_prefetch[0] = previous + + def patched_shard(self): + if inside_prefetch[0] and id(self) in watched: + bucket_group = ddp_model.param_to_bucket_group.get(self) + records.append( + { + "name": names.get(id(self), ""), + "dispatched": bucket_group.param_gather_dispatched, + "in_flight": bucket_group.param_gather_handle is not None, + # The all-gather input, snapshotted on the caller's stream. + "read_value": self.data.detach().clone(), + } + ) + return original_shard(self) + + cls._all_gather_weight = patched_gather + cls.get_padded_shard = patched_shard + + def restore(): + cls._all_gather_weight = original_gather + cls.get_padded_shard = original_shard + + return records, restore + + +@contextlib.contextmanager +def _gtp_env(): + """MPU at ``gtp_remat_size=GTP`` with padding off; restores global GTP state on exit.""" + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=GTP + ) + orig_pad = gtp_module.GTP_CONFIG.pad_for_alignment + gtp_module.GTP_CONFIG.pad_for_alignment = 0 + try: + yield + finally: + gtp_module.GTP_CONFIG.pad_for_alignment = orig_pad + ps.destroy_model_parallel() + gtp_module.GTPShardedParam._chain_state = {} + gtp_module.get_global_GTP_cache().clear() + ps.initialize_model_parallel() + + +@contextlib.contextmanager +def _patched(module, *names): + """Restore ``module.`` for each name on exit, so tests can monkeypatch freely.""" + saved = {n: getattr(module, n) for n in names} + try: + yield + finally: + for n, v in saved.items(): + setattr(module, n, v) + + +def _assert_published(records, rank, what): + """Every recorded prefetch read must have hit an already-published DDP bucket.""" + assert records, f"no prefetch recorded for {what} -- the chain did not link, test is vacuous" + unsynced = [r for r in records if (not r["dispatched"]) or r["in_flight"]] + detail = "\n".join( + f" {r['name']}: dispatched={r['dispatched']} in_flight={r['in_flight']}" for r in records + ) + assert not unsynced, ( + f"[rank {rank}] {what}: GTP prefetched {len(unsynced)}/{len(records)} weight(s) whose " + f"DDP parameter all-gather had not completed:\n{detail}" + ) + + +def _worker_bucket_not_synced(rank, world_size, _port): + with _gtp_env(): + ddp_model, weights = _build_ddp_two_gtp_layers() + groups = [ddp_model.param_to_bucket_group[w] for w in weights] + assert groups[0] is not groups[1], ( + "test precondition: the two GTP weights must land in different DDP bucket groups; " + f"lower BUCKET_SIZE (currently {BUCKET_SIZE})" + ) + + # The first forward links the chain lazily; prefetch only fires from the second onwards. + ddp_model.start_param_sync(force_sync=True) + _forward(ddp_model, rank) + ddp_model.reset_param_sync_dispatch_state() + + records, restore = _record_prefetches(ddp_model, weights) + try: + _forward(ddp_model, rank) + finally: + restore() + _assert_published(records, rank, "one-weight-ahead prefetch") + + +def _worker_stale_values(rank, world_size, _port): + with _gtp_env(): + ddp_model, weights = _build_ddp_two_gtp_layers() + assert ( + ddp_model.param_to_bucket_group[weights[0]] + is not ddp_model.param_to_bucket_group[weights[1]] + ), "test precondition: the two GTP weights must land in different DDP bucket groups" + + ddp_model.start_param_sync(force_sync=True) + _forward(ddp_model, rank) + ddp_model.reset_param_sync_dispatch_state() + + # Suppress the pre-hook's "dispatch the *next* bucket group" step so the prefetch + # target's bucket is provably un-dispatched rather than merely in flight. This is the + # production `--overlap-param-gather-with-optimizer-step` configuration and makes the + # stale read deterministic instead of timing-dependent. + ddp_model.overlap_param_gather_with_optimizer_step = True + + # STALE = what the buffer holds from the previous iteration. + # FRESH = what this rank owns and what the all-gather is about to publish. + STALE, FRESH = -1.0, float(rank + 1) + for bucket_group in ddp_model.bucket_groups: + for bucket in bucket_group.buckets: + bucket.param_data.fill_(STALE) + local_view = shard_buffer( + bucket.param_data, bucket_group.intra_distributed_optimizer_instance_size + )[bucket_group.intra_distributed_optimizer_instance_rank] + local_view.fill_(FRESH) + + # The whole test rests on param.data being a VIEW into bucket.param_data -- that aliasing + # is what makes the fill above visible to GTP's gather. Assert it: if a future change (or + # the native-FP8 path, which gathers w.quantized instead) breaks the aliasing, the fill + # would silently miss and every read_value would look fresh, turning this into a green + # test that no longer guards anything. + for w in weights: + assert w.data.eq(STALE).any() or w.data.eq(FRESH).all(), ( + f"[rank {rank}] {type(w).__name__}.data does not alias bucket.param_data, so the " + "stale-value fill never reached it -- this test would pass vacuously" + ) + + records, restore = _record_prefetches(ddp_model, weights) + try: + _forward(ddp_model, rank) + finally: + restore() + ddp_model.overlap_param_gather_with_optimizer_step = False + + assert records, "no GTP prefetch was issued -- the chain did not link, test is vacuous" + + stale = [] + for r in records: + n_stale = int((r["read_value"] == STALE).sum().item()) + if n_stale: + stale.append((r["name"], n_stale, r["read_value"].numel())) + detail = "\n".join( + f" {name}: {n}/{total} elements still stale" for name, n, total in stale + ) + assert not stale, ( + f"[rank {rank}] GTP all-gathered weight storage that DDP had not yet refreshed, so " + f"the gathered weight carries the previous iteration's values:\n{detail}" + ) + + +def _worker_recompute_skips_readiness(rank, world_size, _port): + """Activation-recompute must NOT ask DDP to publish parameters. + + Recompute-forward is a forward replayed inside backward, so the shards were already + published by the real forward. Asking again is not merely wasted work: DDP has moved on to + gradient reduction by then, and publishing could dispatch a parameter all-gather mid-backward + into a buffer that aliases the gradient buffer under --reuse-grad-buf-for-mxfp8-param-ag. + """ + patched = _patched(gtp_module, "ensure_params_ready", "in_activation_recompute_phase") + with _gtp_env(), patched: + ddp_model, weights = _build_ddp_two_gtp_layers() + ddp_model.start_param_sync(force_sync=True) + + calls = [] + gtp_module.ensure_params_ready = lambda params: calls.append(tuple(params)) + + # Baseline: a genuine forward gather must consult readiness, otherwise this test could + # pass simply because the call site was deleted. + gtp_module.in_activation_recompute_phase = lambda: False + weights[0]._all_gather_weight(async_op=False, fwd=True) + assert calls, ( + f"[rank {rank}] a real forward gather did not consult param readiness -- the call " + "site is gone, so the recompute assertion below would be vacuous" + ) + + # The actual guarantee: inside recompute, readiness must not be consulted at all. + calls.clear() + gtp_module.in_activation_recompute_phase = lambda: True + weights[0]._all_gather_weight(async_op=False, fwd=True) + assert not calls, ( + f"[rank {rank}] activation-recompute asked DDP to publish parameters " + f"({len(calls)} call(s)); recompute runs inside backward, where the shards are " + "already published and a dispatch could gather into the gradient buffer" + ) + + +class _Experts(torch.nn.Module): + """``.mlp.experts.linear_fc1/2`` -- the names ``_classify_param_chain`` keys on.""" + + def __init__(self, gtp_group): + super().__init__() + self.linear_fc1 = _make_gtp_remat_grouped_linear( + NUM_EXPERTS, HIDDEN, HIDDEN, gtp_group, dtype=DTYPE + ) + self.linear_fc2 = _make_gtp_remat_grouped_linear( + NUM_EXPERTS, HIDDEN, HIDDEN, gtp_group, dtype=DTYPE + ) + + +class _MoEBlock(torch.nn.Module): + def __init__(self, gtp_group): + super().__init__() + self.mlp = torch.nn.Module() + self.mlp.experts = _Experts(gtp_group) + + +def _worker_grouped_one_block_ahead(rank, world_size, _port): + """Grouped-expert chains prefetch one MoE BLOCK ahead -- the widest window. + + ``next_w`` links the same weight role across consecutive blocks, so consuming block 0's fc1 + gathers block 1's fc1 while an entire block of compute still separates it from that module's + DDP pre-hook. Also the only route where ``ensure_params_ready`` receives several parameters + at once (one per expert), exercising the dedup path through the real GTP call. + """ + with _gtp_env(): + gtp_group = ps.get_gtp_weight_remat_group() + model = torch.nn.Module() + model.layers = torch.nn.ModuleList([_MoEBlock(gtp_group) for _ in range(2)]) + model = model.cuda() + + # Stamp chain_id from the parameter NAMES; without this the grouped chains are never + # selected and this would silently retest the default one-weight-ahead chain. + gtp_module.classify_gtp_chains(model) + anchors = [b.mlp.experts.linear_fc1.weight0 for b in model.layers] + chain_ids = {a.chain_id for a in anchors} + assert all("grouped" in c for c in chain_ids), ( + f"expected grouped-expert chains, got {chain_ids} -- the one-block-ahead path is " + "not under test" + ) + + ddp_config = DistributedDataParallelConfig( + use_distributed_optimizer=True, + overlap_param_gather=True, + overlap_grad_reduce=True, + bucket_size=BUCKET_SIZE, + ) + ddp_model = DistributedDataParallel(_make_config(), ddp_config, model) + assert ( + ddp_model.param_to_bucket_group[anchors[0]] + is not ddp_model.param_to_bucket_group[anchors[1]] + ), "test precondition: the two blocks' fc1 weights must land in different bucket groups" + + # First pass links block0.fc1 -> block1.fc1 in the grouped chain; prefetch fires after. + ddp_model.start_param_sync(force_sync=True) + for a in anchors: + a.all_gather_and_prefetch(fwd=True) + ddp_model.reset_param_sync_dispatch_state() + + targets = anchors[1]._weights # every expert weight of block 1's fc1 + records, restore = _record_prefetches(ddp_model, targets) + try: + ddp_model.start_param_sync(force_sync=True) + anchors[0].all_gather_and_prefetch(fwd=True) + finally: + restore() + + _assert_published(records, rank, "one-block-ahead grouped-expert prefetch") + + +class TestGTPDDPParamSyncRace: + """GTP prefetch must not read a parameter before DDP has published it.""" + + def test_prefetch_target_bucket_is_not_synced(self): + """Structural: every prefetch target's DDP bucket group must already be finished.""" + _requires_multi_gpu(WORLD) + _run_distributed(_worker_bucket_not_synced, WORLD) + + def test_prefetch_reads_stale_weight_values(self): + """Numerical: the storage GTP feeds to NCCL must not hold last iteration's values.""" + _requires_multi_gpu(WORLD) + _run_distributed(_worker_stale_values, WORLD) + + def test_recompute_forward_does_not_request_publication(self): + """Scope: only a real forward publishes; recompute runs inside backward.""" + _requires_multi_gpu(WORLD) + _run_distributed(_worker_recompute_skips_readiness, WORLD) + + def test_grouped_expert_one_block_ahead_prefetch(self): + """The widest prefetch window: grouped-expert chains reach a whole MoE block ahead.""" + _requires_multi_gpu(WORLD) + _run_distributed(_worker_grouped_one_block_ahead, WORLD) From d7430d5252f7588f25c5ead0371707e64c3cf32d Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 13 Aug 2026 16:56:18 +0000 Subject: [PATCH 023/101] Cover symmetric memory in CUDA graph test (#5987) Signed-off-by: Jingyue Wu Co-authored-by: Claude Opus 4.8 (1M context) --- .../distributed/mfsdp_v2/test_cuda_graph.py | 44 +++++++++++++++++-- 1 file changed, 40 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py b/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py index 86d764f72f7..73fa812b454 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py @@ -4,7 +4,9 @@ import logging +import pytest import torch +import torch.distributed as dist from torch import nn from torch.distributed.device_mesh import init_device_mesh @@ -17,6 +19,10 @@ logger = logging.getLogger(__name__) +# Sized so NCCL selects its symmetric-memory (ncclSymk*) kernels over ring for the +# sharded Linear collectives; see test_symmetric_memory.py for the rationale. +_HIDDEN = 1024 + class NestedModel(nn.Module): """Model with a root FSDP unit and multiple child FSDP units.""" @@ -38,21 +44,46 @@ def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) -def test_captures_full_iteration(distributed_setup): - """A full training iteration should be CUDA-graphable.""" +# CUDA graph capture without symmetric memory is supported and runs by default. The symmetric +# memory case tracks https://github.com/NVIDIA/Megatron-LM/issues/6154: its MemPool is separate +# from CUDA graph's private pool, so allocation addresses and lifetimes are not yet guaranteed +# across capture and replay. Mark only that case flaky until the required PyTorch allocator and +# buffer-registration support is available. +@pytest.mark.parametrize( + "use_symmetric_memory", + [ + pytest.param(False, id="default"), + pytest.param( + True, id="symmetric_memory", marks=[pytest.mark.flaky, pytest.mark.flaky_in_dev] + ), + ], +) +def test_captures_full_iteration(distributed_setup, use_symmetric_memory): + """A full training iteration should be CUDA-graphable, with and without symmetric memory. + + With ``use_symmetric_memory=True`` the FSDP collective path changes: unshard allocates the + all-gather buffer from a symmetric-memory ``MemPool`` and calls + ``symm_mem.rendezvous()``, and the backward reduce-scatter does the same for the + partial-gradient buffer. Both allocation-from-pool and rendezvous are normally + illegal inside CUDA-graph capture, so this confirms the warmup establishes reusable, + already-rendezvoused buffers that capture and replay cleanly -- reproducing the + non-symmetric-memory loss trajectory exactly. + """ world_size = distributed_setup.world_size device = distributed_setup.device + if use_symmetric_memory and world_size < 2: + pytest.skip("Symmetric memory requires at least 2 ranks.") mesh = init_device_mesh(device.type, (world_size,)) torch.manual_seed(1234) - dim = 8 + dim = _HIDDEN model = NestedModel(dim=dim, num_children=2).to(device) static_input = torch.eye(dim, device=device) static_target = torch.zeros_like(static_input) placements = _flat_placements() - with fully_shard_context(device=device): + with fully_shard_context(device=device, use_symmetric_memory=use_symmetric_memory): for layer in model.layers: fully_shard(layer, mesh=mesh, placements=placements) fully_shard(model, mesh=mesh, placements=placements) @@ -67,6 +98,11 @@ def train_iteration() -> torch.Tensor: optimizer.step() return loss.detach() + if use_symmetric_memory: + # NCCL symmetric-memory window registration can fail when a rendezvous is the + # first collective on a communicator, so establish it before capture. + dist.barrier(group=mesh.get_group(0), device_ids=[device.index]) + capture_stream = torch.cuda.Stream() capture_stream.wait_stream(torch.cuda.current_stream()) From f7b2209770ba98596d5e641842d36caa5a51654c Mon Sep 17 00:00:00 2001 From: nvcsathe Date: Thu, 13 Aug 2026 17:21:56 +0000 Subject: [PATCH 024/101] Fix tool call reasoning boundary (#6411) Signed-off-by: Chaitra Sathe --- .../endpoints/chat_completions.py | 20 +++++++++- .../parsers/deepseek_r1_reasoning_parser.py | 20 ++++++++-- .../text/parsers/qwen3_coder_tool_parser.py | 2 + .../tokenizers/test_text_parsers.py | 37 +++++++++++++++++++ 4 files changed, 73 insertions(+), 6 deletions(-) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index e81b5273e84..b0234e54f70 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -401,14 +401,30 @@ def apply_parsers( message_text, tools, parsers_list, tools_requested, chat_template_kwargs=None ): """Runs CPU-intensive text parsing.""" - meta = {} for parser in parsers_list: if parser not in PARSER_MAPPING: raise ValueError(f"Parser {parser} not found in PARSER_MAPPING") + implicit_reasoning_end_markers = ( + tuple( + marker + for parser_name in parsers_list + for marker in getattr( + PARSER_MAPPING[parser_name], "implicit_reasoning_end_markers", () + ) + ) + if tools_requested + else () + ) + + meta = {} + for parser in parsers_list: prev_text = message_text parsed_text, new_info = PARSER_MAPPING[parser].parse( - message_text, tools=tools, chat_template_kwargs=chat_template_kwargs + message_text, + tools=tools, + chat_template_kwargs=chat_template_kwargs, + implicit_reasoning_end_markers=implicit_reasoning_end_markers, ) if "tool_calls" in new_info: new_info["tool_calls"] = _normalize_tool_calls( diff --git a/megatron/core/tokenizers/text/parsers/deepseek_r1_reasoning_parser.py b/megatron/core/tokenizers/text/parsers/deepseek_r1_reasoning_parser.py index e54088143af..2792dd8e0a5 100644 --- a/megatron/core/tokenizers/text/parsers/deepseek_r1_reasoning_parser.py +++ b/megatron/core/tokenizers/text/parsers/deepseek_r1_reasoning_parser.py @@ -10,8 +10,9 @@ def parse(text: str, **kwargs) -> tuple[str, dict[str, str]]: """Extract reasoning content delimited by `...` tags. Any text before the first `` is discarded. - When no `` follows, the model is still "thinking": all text is reasoning. - Otherwise the text is split at the first ``. + When no `` follows, the model is still "thinking": all text is reasoning, + unless a configured implicit end marker (such as ``) appears. + Otherwise the text is split at the first reasoning boundary. Args: text (str): The text to parse. @@ -24,11 +25,22 @@ def parse(text: str, **kwargs) -> tuple[str, dict[str, str]]: before, think_open, after = text.partition("") remaining = after if think_open else before - if "" not in remaining: + reasoning_end = remaining.find("") + content_start = reasoning_end + len("") + + for marker in kwargs.get("implicit_reasoning_end_markers", ()): + marker_index = remaining.find(marker) + if marker_index >= 0 and (reasoning_end < 0 or marker_index < reasoning_end): + reasoning_end = marker_index + # Preserve the marker for the downstream tool parser. + content_start = marker_index + + if reasoning_end < 0: # No closing tag: treat the remaining text as unterminated reasoning. reasoning_content, content = remaining, "" else: - reasoning_content, _, content = remaining.partition("") + reasoning_content = remaining[:reasoning_end] + content = remaining[content_start:] info = {"reasoning": reasoning_content} if reasoning_content else {} return content, info diff --git a/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py b/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py index 46e0ffa0271..ace63a4b772 100644 --- a/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py +++ b/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py @@ -266,6 +266,8 @@ def extract_tool_calls( class Qwen3CoderToolParser(BaseParser): """Parser for Qwen3 Coder style tool calls.""" + implicit_reasoning_end_markers = ("",) + @staticmethod def parse(text: str, **kwargs) -> tuple[str, dict[str, list[dict]]]: """ diff --git a/tests/unit_tests/tokenizers/test_text_parsers.py b/tests/unit_tests/tokenizers/test_text_parsers.py index 62ed6797149..b5ab1b1dee7 100644 --- a/tests/unit_tests/tokenizers/test_text_parsers.py +++ b/tests/unit_tests/tokenizers/test_text_parsers.py @@ -96,3 +96,40 @@ def test_parser_mapping_registers_nemotron_v3_reasoning(): """Super and Ultra share identical reasoning-extraction logic upstream, so both models are served by a single consolidated parser and registry key.""" assert PARSER_MAPPING["nemotron-v3-reasoning"] is NemotronV3ReasoningParser + + +def test_tool_call_marker_implicitly_ends_reasoning_for_downstream_parser(): + tool_text = ( + "echo hi" "" + ) + model_output = f"I should inspect this first.\n{tool_text}" + tool_parser = PARSER_MAPPING["qwen3-coder-tool"] + + content, reasoning_info = DeepSeekR1ReasoningParser.parse( + model_output, implicit_reasoning_end_markers=tool_parser.implicit_reasoning_end_markers + ) + parsed_content, tool_info = tool_parser.parse( + content, + tools=[ + { + "type": "function", + "function": { + "name": "bash", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}, + }, + } + ], + ) + + assert reasoning_info == {"reasoning": "I should inspect this first.\n"} + assert parsed_content is None + assert tool_info["tool_calls"][0]["function"] == { + "name": "bash", + "arguments": '{"command": "echo hi"}', + } + + +def test_tool_call_marker_does_not_end_reasoning_unless_configured(): + model_output = "reasoningnot enabled" + + assert DeepSeekR1ReasoningParser.parse(model_output) == ("", {"reasoning": model_output}) From 69e282ec2cebaee00603aafec0f9453954d50208 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Thu, 13 Aug 2026 17:34:47 -0400 Subject: [PATCH 025/101] Update release notes page (#6520) Signed-off-by: Philip Petrakian --- docs/get-started/releasenotes.md | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/docs/get-started/releasenotes.md b/docs/get-started/releasenotes.md index b296bdfb254..1572408ea1a 100644 --- a/docs/get-started/releasenotes.md +++ b/docs/get-started/releasenotes.md @@ -9,11 +9,17 @@ # Release Notes +We publish release notes for every [major release](https://github.com/NVIDIA/Megatron-LM/releases). ## Roadmaps Stay up-to-date with the development roadmaps and planned features: -- **[MoE Q3-Q4 2025 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - Comprehensive MoE feature development including DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements. -- **[GPT-OSS Implementation Tracker](https://github.com/NVIDIA/Megatron-LM/issues/1739)** - Advanced features including YaRN RoPE scaling, attention sinks, and custom activation functions. +- **Q3 2026** + - [General Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/5676) +- **Q2 2026** + - [General Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/4815) + - [MoE-focused Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/4815) +- **Q3-Q4 2025** + - [MoE-focused Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729) From 3d0956fe5be5721bcd5c11ba335e5c21e180d3fd Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 13 Aug 2026 19:48:46 +0000 Subject: [PATCH 026/101] test(fsdp): verify optimizer CUDA graphs in v2 (#6498) Signed-off-by: Jingyue Wu --- .../core/optimizer/fully_sharded_optimizer.py | 2 - .../mfsdp_v2/test_mcore_adapter.py | 61 ++++++++++++------- 2 files changed, 38 insertions(+), 25 deletions(-) diff --git a/megatron/core/optimizer/fully_sharded_optimizer.py b/megatron/core/optimizer/fully_sharded_optimizer.py index 9c70c5b6f86..6dcbbad5e9e 100644 --- a/megatron/core/optimizer/fully_sharded_optimizer.py +++ b/megatron/core/optimizer/fully_sharded_optimizer.py @@ -83,8 +83,6 @@ def _validate_config(config: OptimizerConfig, model_chunks: List[MegatronModule] raise ValueError( "MFSDP v2 does not currently support layer-wise distributed optimizer." ) - if config.optimizer_cuda_graph: - raise ValueError("MFSDP v2 does not currently support optimizer CUDA graphs.") def state_dict(self): """Return optimizer state. diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py index 597a0a7fe56..28bb7195cca 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py @@ -18,6 +18,7 @@ from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer from megatron.core.optimizer.fully_sharded_optimizer import FullyShardedOptimizer +from megatron.core.optimizer.optimizer_cuda_graph import OptimizerCudaGraphWrapper from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.enums import AttnBackend @@ -137,7 +138,8 @@ def record_fully_shard_context(*args, **kwargs): assert fully_shard_context_calls == [True] - def test_build_train_and_step(self): + @pytest.mark.parametrize("optimizer_cuda_graph", [False, True], ids=["eager", "cuda_graph"]) + def test_build_train_and_step(self, optimizer_cuda_graph): config = TransformerConfig( num_layers=2, hidden_size=16, @@ -156,16 +158,19 @@ def test_build_train_and_step(self): # DistributedOptimizer path that expects DDP/FSDP buffer metadata. reference_model.ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=False) + ddp_config = DistributedDataParallelConfig( + use_megatron_fsdp=True, + megatron_fsdp_version=2, + use_distributed_optimizer=False, + data_parallel_sharding_strategy="optim_grads_params", + ) + if optimizer_cuda_graph: + # Capturable TE FusedAdam requires the main-gradient and main-weight dtypes + # to match (https://github.com/NVIDIA/TransformerEngine/issues/3358). + ddp_config.megatron_fsdp_main_grads_dtype = ddp_config.megatron_fsdp_main_params_dtype + model = FullyShardedDataParallel( - config=config, - ddp_config=DistributedDataParallelConfig( - use_megatron_fsdp=True, - megatron_fsdp_version=2, - use_distributed_optimizer=False, - data_parallel_sharding_strategy="optim_grads_params", - ), - module=model, - pg_collection=self.pg_collection, + config=config, ddp_config=ddp_config, module=model, pg_collection=self.pg_collection ) reference_optimizer_config = OptimizerConfig( @@ -180,7 +185,9 @@ def test_build_train_and_step(self): reference_optimizer = get_megatron_optimizer( reference_optimizer_config, [reference_model], use_gloo_process_groups=False ) - optimizer_config = replace(reference_optimizer_config) + optimizer_config = replace( + reference_optimizer_config, optimizer_cuda_graph=optimizer_cuda_graph + ) with pytest.raises( ValueError, match="MFSDP v2 currently requires use_distributed_optimizer=False" ): @@ -192,6 +199,8 @@ def test_build_train_and_step(self): optimizer = get_megatron_optimizer(optimizer_config, [model], use_gloo_process_groups=False) assert isinstance(optimizer, FullyShardedOptimizer) optimizer.reload_model_params() + if optimizer_cuda_graph: + optimizer.step = OptimizerCudaGraphWrapper(optimizer.step, cuda_graph_warmup_steps=1) steps = [ [ @@ -215,18 +224,24 @@ def test_build_train_and_step(self): reference_losses.append(torch.stack(microbatch_losses).mean()) losses = [] - for microbatches in steps: - model.zero_grad_buffer() - optimizer.zero_grad(set_to_none=True) - microbatch_losses = [] - for batch in microbatches: - output = model(hidden_states=batch, attention_mask=None) - loss = output.float().square().mean() - (loss / len(microbatches)).backward() - microbatch_losses.append(loss.detach()) - success, _, _ = optimizer.step() - assert success - losses.append(torch.stack(microbatch_losses).mean()) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as prof: + for microbatches in steps: + model.zero_grad_buffer() + optimizer.zero_grad(set_to_none=True) + microbatch_losses = [] + for batch in microbatches: + output = model(hidden_states=batch, attention_mask=None) + loss = output.float().square().mean() + (loss / len(microbatches)).backward() + microbatch_losses.append(loss.detach()) + success, _, _ = optimizer.step() + assert success + losses.append(torch.stack(microbatch_losses).mean()) + + if optimizer_cuda_graph: + # The first step is eager; capture replays once and every subsequent step replays. + graph_launches = sum(event.name == "cudaGraphLaunch" for event in prof.events()) + assert graph_launches == len(steps) - 1 losses = torch.stack(losses) reference_losses = torch.stack(reference_losses) From 0ee415bbf914340364335691bc27f2aa99a5cb56 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Thu, 13 Aug 2026 20:37:20 +0000 Subject: [PATCH 027/101] Update Nemotron 3.5 Lightning golden values (#6445) Signed-off-by: Philip Petrakian --- .../golden_values_dev_dgx_gb200.json | 190 ++++++++++++++++++ .../model_config.yaml | 12 +- .../golden_values_dev_dgx_gb200.json | 190 ------------------ tests/test_utils/recipes/gb200/nemotron.yaml | 50 +---- 4 files changed, 199 insertions(+), 243 deletions(-) create mode 100644 tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200/golden_values_dev_dgx_gb200.json rename tests/functional_tests/test_cases/nemotron/{nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G => nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200}/model_config.yaml (93%) delete mode 100644 tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G/golden_values_dev_dgx_gb200.json diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200/golden_values_dev_dgx_gb200.json new file mode 100644 index 00000000000..8a7f538f6e8 --- /dev/null +++ b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200/golden_values_dev_dgx_gb200.json @@ -0,0 +1,190 @@ +{ + "lm loss": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 11.26482, + "2": 11.25947, + "3": 9.77821, + "4": 10.65636, + "5": 11.47448, + "6": 9.3545, + "7": 9.39757, + "8": 9.06374, + "9": 8.63001, + "10": 8.91632, + "11": 8.50413, + "12": 8.25309, + "13": 8.07036, + "14": 7.94322, + "15": 8.05542, + "16": 7.6059, + "17": 7.6034, + "18": 7.53429, + "19": 7.55707, + "20": 7.3521 + } + }, + "num-zeros": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 13741925.0, + "2": 13135267.0, + "3": 13524615.0, + "4": 561370752.0, + "5": 6051030016.0, + "6": 11747025920.0, + "7": 16715832320.0, + "8": 19022454784.0, + "9": 20591636480.0, + "10": 21548642304.0, + "11": 22057068544.0, + "12": 22505883648.0, + "13": 22916513792.0, + "14": 23244693504.0, + "15": 23021539328.0, + "16": 23019051008.0, + "17": 22831628288.0, + "18": 22545094656.0, + "19": 21749301248.0, + "20": 21136400384.0 + } + }, + "mem-allocated-bytes": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 88825536512.0, + "2": 88825544704.0, + "3": 88825970688.0, + "4": 88826519552.0, + "5": 88824963072.0, + "6": 88826028032.0, + "7": 88827486208.0, + "8": 88828837888.0, + "9": 88830222336.0, + "10": 88829984768.0, + "11": 88828035072.0, + "12": 88827641856.0, + "13": 88828796928.0, + "14": 88829911040.0, + "15": 88829927424.0, + "16": 88830091264.0, + "17": 88830509056.0, + "18": 88828198912.0, + "19": 88828764160.0, + "20": 88828608512.0 + } + }, + "mem-max-allocated-bytes": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 123182538752.0, + "2": 155622637568.0, + "3": 157646094336.0, + "4": 160994197504.0, + "5": 160994197504.0, + "6": 160994197504.0, + "7": 162116304896.0, + "8": 167540277248.0, + "9": 171569627136.0, + "10": 171569627136.0, + "11": 171569627136.0, + "12": 171569627136.0, + "13": 171569627136.0, + "14": 171569627136.0, + "15": 171569627136.0, + "16": 172658491392.0, + "17": 173031964672.0, + "18": 173031964672.0, + "19": 173031964672.0, + "20": 173031964672.0 + } + }, + "mtp_1 loss": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 11.22868, + "2": 11.23172, + "3": 11.05434, + "4": 10.52588, + "5": 10.29408, + "6": 9.88058, + "7": 9.556, + "8": 9.35846, + "9": 9.0526, + "10": 8.92832, + "11": 8.83714, + "12": 8.55226, + "13": 8.41308, + "14": 8.26385, + "15": 8.12458, + "16": 7.73642, + "17": 7.66909, + "18": 7.55402, + "19": 7.55124, + "20": 7.22983 + } + }, + "mtp_2 loss": { + "start_step": 1, + "end_step": 20, + "step_interval": 1, + "values": { + "1": 11.25065, + "2": 11.25147, + "3": 11.10074, + "4": 10.6482, + "5": 10.37554, + "6": 9.9138, + "7": 9.61227, + "8": 9.39007, + "9": 9.04576, + "10": 8.91742, + "11": 8.82202, + "12": 8.51697, + "13": 8.40975, + "14": 8.25986, + "15": 8.10147, + "16": 7.70083, + "17": 7.63184, + "18": 7.53295, + "19": 7.53525, + "20": 7.2173 + } + }, + "iteration-time": { + "start_step": 2, + "end_step": 20, + "step_interval": 1, + "values": { + "2": 112.32996, + "3": 20.94815, + "4": 21.63868, + "5": 22.20665, + "6": 22.81373, + "7": 25.33896, + "8": 24.97021, + "9": 24.01207, + "10": 23.30601, + "11": 23.04278, + "12": 23.13449, + "13": 21.91107, + "14": 22.40248, + "15": 22.62039, + "16": 22.91055, + "17": 22.98766, + "18": 22.93769, + "19": 22.52428, + "20": 22.078 + } + } +} diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G/model_config.yaml b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200/model_config.yaml similarity index 93% rename from tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G/model_config.yaml rename to tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200/model_config.yaml index 28b0943f3d6..402c01bc294 100644 --- a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200/model_config.yaml @@ -93,18 +93,18 @@ MODEL_ARGS: --mtp-use-repeated-layer: true --mtp-loss-scaling-factor: 0.3 - # Data and tokenizer. Square brackets make the functional launcher expand the blend. + # Use the staged corpus shared by the other GB200 functional tests instead of + # the node-local RP2 mount. --seq-length: 8192 --max-position-embeddings: 8192 - --data-path: "[${DATA_BLEND}]" + --data-path: ${DATA_PATH}/text/the_pile/shard00/my-gpt3_00_text_document + --vocab-file: ${DATA_PATH}/text/the_pile/shard00/bpe/vocab.json + --merge-file: ${DATA_PATH}/text/the_pile/shard00/bpe/merges.txt --data-cache-path: ${DATA_CACHE_PATH} - --split: "9999,8,2" + --split: 949,50,1 --dataloader-type: single --num-workers: 8 - --no-mmap-bin-files: true --no-create-attention-mask-in-dataloader: true - --tokenizer-type: ${TOKENIZER_TYPE} - --tokenizer-model: ${TOKENIZER_MODEL_PATH} # BF16 compute and BF16 gradient reduction with FP32 optimizer state. --bf16: true diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G/golden_values_dev_dgx_gb200.json deleted file mode 100644 index 557529b414b..00000000000 --- a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G/golden_values_dev_dgx_gb200.json +++ /dev/null @@ -1,190 +0,0 @@ -{ - "lm loss": { - "start_step": 1, - "end_step": 20, - "step_interval": 1, - "values": { - "1": 10.77293, - "2": 10.7731, - "3": 9.68815, - "4": 11.06906, - "5": 9.77569, - "6": 9.47808, - "7": 9.06072, - "8": 8.8026, - "9": 8.76839, - "10": 8.59208, - "11": 8.38807, - "12": 8.07211, - "13": 8.00935, - "14": 7.92114, - "15": 7.75658, - "16": 7.632, - "17": 7.52375, - "18": 7.44319, - "19": 7.32021, - "20": 7.29368 - } - }, - "num-zeros": { - "start_step": 1, - "end_step": 20, - "step_interval": 1, - "values": { - "1": 8283006.0, - "2": 8215530.0, - "3": 8558295.0, - "4": 1945319168.0, - "5": 10112309248.0, - "6": 15437038592.0, - "7": 19477004288.0, - "8": 21703987200.0, - "9": 23075272704.0, - "10": 23760054272.0, - "11": 23957757952.0, - "12": 24046585856.0, - "13": 23972829184.0, - "14": 23722366976.0, - "15": 23549724672.0, - "16": 23252019200.0, - "17": 22770812928.0, - "18": 22064881664.0, - "19": 21154738176.0, - "20": 20261984256.0 - } - }, - "mem-allocated-bytes": { - "start_step": 1, - "end_step": 20, - "step_interval": 1, - "values": { - "1": 88137433088.0, - "2": 88137441280.0, - "3": 88137605120.0, - "4": 88138850304.0, - "5": 88139759616.0, - "6": 88139350016.0, - "7": 88139849728.0, - "8": 88139620352.0, - "9": 88140791808.0, - "10": 88141807616.0, - "11": 88141611008.0, - "12": 88142356480.0, - "13": 88142135296.0, - "14": 88142127104.0, - "15": 88142004224.0, - "16": 88140652544.0, - "17": 88138448896.0, - "18": 88138235904.0, - "19": 88137572352.0, - "20": 88138219520.0 - } - }, - "mem-max-allocated-bytes": { - "start_step": 1, - "end_step": 20, - "step_interval": 1, - "values": { - "1": 119058399232.0, - "2": 151382818816.0, - "3": 151812358144.0, - "4": 155852455936.0, - "5": 158491279360.0, - "6": 158491279360.0, - "7": 159239569408.0, - "8": 159239569408.0, - "9": 162155266048.0, - "10": 165063098368.0, - "11": 165063098368.0, - "12": 166633291776.0, - "13": 166633291776.0, - "14": 166633291776.0, - "15": 166633291776.0, - "16": 166633291776.0, - "17": 166633291776.0, - "18": 166633291776.0, - "19": 166633291776.0, - "20": 166633291776.0 - } - }, - "mtp_1 loss": { - "start_step": 1, - "end_step": 20, - "step_interval": 1, - "values": { - "1": 10.78576, - "2": 10.78552, - "3": 10.66957, - "4": 10.42839, - "5": 10.18673, - "6": 9.92491, - "7": 9.54198, - "8": 9.20471, - "9": 8.98346, - "10": 8.73057, - "11": 8.52848, - "12": 8.34207, - "13": 8.14723, - "14": 8.02206, - "15": 7.87516, - "16": 7.7349, - "17": 7.55731, - "18": 7.46251, - "19": 7.3314, - "20": 7.23513 - } - }, - "mtp_2 loss": { - "start_step": 1, - "end_step": 20, - "step_interval": 1, - "values": { - "1": 10.81667, - "2": 10.81777, - "3": 10.69916, - "4": 10.46612, - "5": 10.22667, - "6": 9.93677, - "7": 9.51459, - "8": 9.17991, - "9": 9.01314, - "10": 8.68501, - "11": 8.53294, - "12": 8.34325, - "13": 8.12935, - "14": 7.99982, - "15": 7.87422, - "16": 7.73445, - "17": 7.54824, - "18": 7.45019, - "19": 7.3192, - "20": 7.21849 - } - }, - "iteration-time": { - "start_step": 2, - "end_step": 20, - "step_interval": 1, - "values": { - "2": 103.07695, - "3": 26.75034, - "4": 24.96478, - "5": 23.51423, - "6": 24.71646, - "7": 25.16894, - "8": 24.63826, - "9": 23.91178, - "10": 23.75633, - "11": 24.20485, - "12": 23.30083, - "13": 22.73128, - "14": 22.04589, - "15": 22.34995, - "16": 22.52285, - "17": 23.71327, - "18": 22.43595, - "19": 22.85244, - "20": 21.68526 - } - } -} diff --git a/tests/test_utils/recipes/gb200/nemotron.yaml b/tests/test_utils/recipes/gb200/nemotron.yaml index f88e3c4aefa..9aae75eaebf 100644 --- a/tests/test_utils/recipes/gb200/nemotron.yaml +++ b/tests/test_utils/recipes/gb200/nemotron.yaml @@ -47,58 +47,14 @@ spec: cd /opt/megatron-lm export GPUS_PER_NODE={gpus} - DATA_PATH=/mnt/rp2 - DATA_BLEND=( - 0.03846154 "$DATA_PATH/perp_head/head_01_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_02_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_03_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_04_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_05_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_06_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_07_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_08_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_09_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_10_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_11_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_12_text_document" - 0.03846154 "$DATA_PATH/perp_head/head_13_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_01_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_02_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_03_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_04_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_05_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_06_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_07_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_08_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_09_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_10_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_11_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_12_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_13_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_14_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_15_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_16_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_17_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_18_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_19_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_20_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_21_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_22_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_23_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_24_text_document" - 0.02 "$DATA_PATH/perp_middle/middle_25_text_document" - ) ARGUMENTS=( - "DATA_PATH=$DATA_PATH" - "DATA_BLEND=${{DATA_BLEND[*]}}" + "DATA_PATH=/mnt/artifacts" "DATA_CACHE_PATH=/lustre/fsw/coreai_dlalgo_mcore/mcore_ci/data/$RUN_ID/cache/" - "TOKENIZER_TYPE=SentencePieceTokenizer" - "TOKENIZER_MODEL_PATH=$DATA_PATH/tokenizer/tokenizer.model" "OUTPUT_PATH={assets_dir}" "TENSORBOARD_PATH={assets_dir}/tensorboard" "CHECKPOINT_SAVE_PATH={artifacts_dir}/checkpoints" - "CHECKPOINT_LOAD_PATH=$DATA_PATH" + "CHECKPOINT_LOAD_PATH=/mnt/artifacts" "TRAINING_SCRIPT_PATH=pretrain_hybrid.py" "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml" "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json" @@ -110,7 +66,7 @@ spec: bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh "${{ARGUMENTS[@]}}" products: - - test_case: [nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200_2N8G] + - test_case: [nemotron3_5_lightning_nightly_tp1_pp1_cp1_ep8_dgx_gb200] products: - environment: [dev] scope: [L2] From 94939b85d538570f4f2a2c7f7cf41f58253141b0 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee <85179052+lmcafee-nvidia@users.noreply.github.com> Date: Thu, 13 Aug 2026 21:04:03 +0000 Subject: [PATCH 028/101] Handle hidden chunked prefill during async overlap (#6442) Signed-off-by: Lawrence McAfee --- .../text_generation_controller.py | 10 ++- .../test_text_generation_controller.py | 68 +++++++++++++++++++ 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index 0675002bbae..8ff0050b257 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -2189,10 +2189,14 @@ def _build_async_sched_request_state( chunked_prefill_rows = torch.nonzero( active_request_ids == context.chunked_prefill_request_id, as_tuple=True )[0] + # Under memory pressure, the next chunk may not be admitted, so the + # chunked-prefill request can remain hidden for a step while active decode + # requests use overlap. Zero matches is valid; multiple matches are not. assert ( - chunked_prefill_rows.numel() == 1 - ), "The active chunked-prefill request must have exactly one row." - active_request_mask[chunked_prefill_rows[0]] = 1 + chunked_prefill_rows.numel() <= 1 + ), "The chunked-prefill request must have at most one active row." + if chunked_prefill_rows.numel() == 1: + active_request_mask[chunked_prefill_rows[0]] = 1 finished_idxs = ( torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py index edb07b424f6..f6af8a66880 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py +++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py @@ -1767,6 +1767,74 @@ def initialize_survivor_forward(): ) alloc.release_memory_blocks(filler_blocks) + @pytest.mark.internal + def test_async_sched_overlap_handles_hidden_chunked_prefill(self): + """Overlap ignores a continuing chunk outside the active request rows.""" + self.setup_model( + torch.float32, batch_size=3, static=False, block_size_tokens=4, max_requests=4 + ) + controller = self.text_generation_controller + context = controller.inference_wrapped_model.inference_context + context.reset() + + active_slice = slice(0, 3) + context.total_request_count = 3 + context.active_token_count = 3 + context.num_prefill_requests = 1 + context.request_ids[active_slice] = torch.tensor([10, 11, 12], dtype=torch.int32) + context.request_in_prefill_status_tensor[active_slice] = torch.tensor( + [0, 0, 1], dtype=torch.int32 + ) + context.request_query_lengths[active_slice] = 1 + context.request_output_lengths[active_slice] = 16 + context.request_kv_length_offsets[active_slice] = 3 + context.request_last_kv_block_offset[active_slice] = torch.tensor( + [context.block_size_tokens - 1, 0, 0], dtype=torch.int32 + ) + context.request_metadata["termination_id"][active_slice] = 99 + context.chunked_prefill_request_id = 12 + context.build_active_slices(3) + + block_ids = context.kv_block_allocator.allocate_memory_blocks(3) + context.request_to_kv_block_ids[active_slice, 0] = block_ids + context.request_last_kv_block_id[active_slice] = block_ids + context.request_kv_block_counts[active_slice] = 1 + context.token_to_input_ids[active_slice] = torch.tensor([80, 81, 82]) + + alloc = context.kv_block_allocator + alloc.paused_limit = 0 + filler_blocks = alloc.allocate_memory_blocks(alloc.pool_avail) + assert filler_blocks is not None + filler_blocks = filler_blocks.clone() + assert alloc.get_allocatable_count() == 0 + + sampled_tokens = torch.tensor([90, 91, 92], dtype=torch.int64) + request_result = controller._run_async_sched_update_requests( + SimpleNamespace( + sampled_tokens_cpu_view=sampled_tokens, + sampled_mtp_tokens_cpu_view=None, + accepted_tokens_cpu_view=None, + ), + resolved_sequence_lengths=torch.tensor([4, 4, 4]), + ) + + assert request_result.evict_request_ids.tolist() == [10] + assert context.total_request_count == 1 + assert context.request_ids[:2].tolist() == [11, 12] + assert context.get_index_of_chunked_prefill_request(safe=True) == -1 + assert context.get_index_of_chunked_prefill_request(safe=False) == 1 + + active_request_ids, finished_request_ids, active_request_mask = ( + controller._build_async_sched_request_state( + torch.tensor([93]), resolved_sequence_lengths=torch.tensor([5]) + ) + ) + + assert active_request_ids.tolist() == [11] + assert finished_request_ids.numel() == 0 + assert active_request_mask.tolist() == [1] + alloc.release_memory_blocks(filler_blocks) + def test_sample_from_logits(self): self.setup_model(torch.float32) From f9cc5444b4e5a50c74a70a848f5f9da0b6c7ff72 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 13 Aug 2026 21:32:43 +0000 Subject: [PATCH 029/101] Thread socket argument through MInf API (#6513) Signed-off-by: Teodor-Dumitru Ene --- megatron/core/inference/apis/async_llm.py | 1 + megatron/core/inference/apis/llm.py | 1 + megatron/core/inference/apis/serve_config.py | 7 +++++++ tests/unit_tests/inference/high_level_api/test_apis.py | 4 +++- 4 files changed, 12 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/apis/async_llm.py b/megatron/core/inference/apis/async_llm.py index 2c6b20693f3..efbed6c309a 100644 --- a/megatron/core/inference/apis/async_llm.py +++ b/megatron/core/inference/apis/async_llm.py @@ -193,6 +193,7 @@ async def serve(self, serve_config: ServeConfig, *, blocking: bool = True) -> No verbose=serve_config.verbose, num_replicas=serve_config.frontend_replicas, hostname=serve_config.host, + sock=serve_config.sock, ) self._serve_started = True diff --git a/megatron/core/inference/apis/llm.py b/megatron/core/inference/apis/llm.py index 6bbdce570d8..688df260e52 100644 --- a/megatron/core/inference/apis/llm.py +++ b/megatron/core/inference/apis/llm.py @@ -178,6 +178,7 @@ def serve(self, serve_config: ServeConfig, *, blocking: bool = True) -> None: verbose=serve_config.verbose, num_replicas=serve_config.frontend_replicas, hostname=serve_config.host, + sock=serve_config.sock, ) self._serve_started = True diff --git a/megatron/core/inference/apis/serve_config.py b/megatron/core/inference/apis/serve_config.py index aa7c6afe8fd..e92c414cdcf 100644 --- a/megatron/core/inference/apis/serve_config.py +++ b/megatron/core/inference/apis/serve_config.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import socket from dataclasses import dataclass, field @@ -42,3 +43,9 @@ class ServeConfig: The default of 4 matches the existing ``start_text_gen_server`` default of ``num_replicas=4``. """ + + sock: socket.socket | None = None + """Pre-bound listening socket to serve on instead of binding `host:port`. + + Must already be bound to a real port; when set, `host` / `port` are not used for binding. + """ diff --git a/tests/unit_tests/inference/high_level_api/test_apis.py b/tests/unit_tests/inference/high_level_api/test_apis.py index ff6721a4f0d..5391d9f7014 100644 --- a/tests/unit_tests/inference/high_level_api/test_apis.py +++ b/tests/unit_tests/inference/high_level_api/test_apis.py @@ -170,10 +170,12 @@ def test_sync_serve_primary_rank_starts_frontend(self, monkeypatch): monkeypatch.setattr(dist, "get_rank", lambda: 0) monkeypatch.setattr(tgs, "start_text_gen_server", lambda **kw: started.update(kw)) - llm.serve(ServeConfig(port=1234), blocking=False) + sock = MagicMock() + llm.serve(ServeConfig(port=1234, sock=sock), blocking=False) assert llm._serve_started is True assert started["coordinator_addr"] == "tcp://coord:5555" assert started["server_port"] == 1234 + assert started["sock"] is sock class TestNormalizePrompts: From de4a82e5c8c087f33ab1fbd9347b706d4939ff37 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 13 Aug 2026 23:21:53 +0000 Subject: [PATCH 030/101] Clean up SSM ops directory structure (#6444) Signed-off-by: Keshav Santhanam Co-authored-by: Shanmugam Ramasamy --- docs/developer/determinism/op-catalog.md | 2 +- megatron/core/ssm/gated_delta_product.py | 2 +- megatron/core/ssm/mamba_mixer.py | 12 +++++------ megatron/core/ssm/ops/__init__.py | 21 ++++++++++++++++--- megatron/core/ssm/ops/common/__init__.py | 10 +++++++++ .../ops/{ => common}/causal_conv1d_triton.py | 0 .../ops/{ => common}/causal_conv1d_varlen.py | 2 +- .../core/ssm/ops/{ => common}/determinism.py | 0 .../{ => common}/intermediate_extraction.py | 6 +++++- megatron/core/ssm/ops/mamba2/__init__.py | 13 ++++++++++++ .../{ => mamba2}/batch_invariant_decode.py | 8 +++---- .../core/ssm/ops/{ => mamba2}/mamba_ssm.py | 0 megatron/core/ssm/ops/{ => mamba2}/ssd_bmm.py | 2 +- .../ssm/ops/{ => mamba2}/ssd_chunk_scan.py | 2 +- .../ssm/ops/{ => mamba2}/ssd_chunk_state.py | 2 +- .../core/ssm/ops/{ => mamba2}/ssd_combined.py | 0 .../ssm/ops/{ => mamba2}/ssd_state_passing.py | 2 +- .../custom_layers/batch_invariant_kernels.py | 8 +++---- .../common}/test_causal_conv1d_triton.py | 2 +- .../{ => common}/test_causal_conv1d_varlen.py | 2 +- .../test_batch_invariant_decode.py | 4 ++-- .../ssm/ops/{ => mamba2}/test_ssd_bmm.py | 2 +- .../ops/{ => mamba2}/test_ssd_chunk_scan.py | 2 +- .../ops/{ => mamba2}/test_ssd_chunk_state.py | 2 +- .../ssm/ops/{ => mamba2}/test_ssd_combined.py | 7 +++++-- .../{ => mamba2}/test_ssd_state_passing.py | 2 +- .../ssm/ops/{ => mamba2}/test_ssm_kernel.py | 0 27 files changed, 80 insertions(+), 35 deletions(-) create mode 100644 megatron/core/ssm/ops/common/__init__.py rename megatron/core/ssm/ops/{ => common}/causal_conv1d_triton.py (100%) rename megatron/core/ssm/ops/{ => common}/causal_conv1d_varlen.py (99%) rename megatron/core/ssm/ops/{ => common}/determinism.py (100%) rename megatron/core/ssm/ops/{ => common}/intermediate_extraction.py (96%) create mode 100644 megatron/core/ssm/ops/mamba2/__init__.py rename megatron/core/ssm/ops/{ => mamba2}/batch_invariant_decode.py (97%) rename megatron/core/ssm/ops/{ => mamba2}/mamba_ssm.py (100%) rename megatron/core/ssm/ops/{ => mamba2}/ssd_bmm.py (99%) rename megatron/core/ssm/ops/{ => mamba2}/ssd_chunk_scan.py (99%) rename megatron/core/ssm/ops/{ => mamba2}/ssd_chunk_state.py (99%) rename megatron/core/ssm/ops/{ => mamba2}/ssd_combined.py (100%) rename megatron/core/ssm/ops/{ => mamba2}/ssd_state_passing.py (99%) rename tests/unit_tests/ssm/{ => ops/common}/test_causal_conv1d_triton.py (99%) rename tests/unit_tests/ssm/ops/{ => common}/test_causal_conv1d_varlen.py (98%) rename tests/unit_tests/ssm/ops/{ => mamba2}/test_batch_invariant_decode.py (99%) rename tests/unit_tests/ssm/ops/{ => mamba2}/test_ssd_bmm.py (98%) rename tests/unit_tests/ssm/ops/{ => mamba2}/test_ssd_chunk_scan.py (98%) rename tests/unit_tests/ssm/ops/{ => mamba2}/test_ssd_chunk_state.py (99%) rename tests/unit_tests/ssm/ops/{ => mamba2}/test_ssd_combined.py (99%) rename tests/unit_tests/ssm/ops/{ => mamba2}/test_ssd_state_passing.py (97%) rename tests/unit_tests/ssm/ops/{ => mamba2}/test_ssm_kernel.py (100%) diff --git a/docs/developer/determinism/op-catalog.md b/docs/developer/determinism/op-catalog.md index c4381c5b5c6..a1eb27b39fb 100644 --- a/docs/developer/determinism/op-catalog.md +++ b/docs/developer/determinism/op-catalog.md @@ -38,7 +38,7 @@ Selected by `torch.are_deterministic_algorithms_enabled()` or | Vocab-parallel embedding | `megatron/core/tensor_parallel/layers.py` | direct indexing `weight[idx]` (deterministic backward) | `F.embedding` (non-deterministic atomic backward) | | Gated-delta-net kernel | `megatron/core/ssm/gated_delta_net.py` | torch `chunk_gated_delta_rule` | FLA fused kernel | | Gated-delta-net causal conv1d | `megatron/core/ssm/gated_delta_net.py` | `F.conv1d` (plus transposes) | FLA `causal_conv1d` | -| Mamba/SSM Triton ops | `megatron/core/ssm/ops/determinism.py` | one fixed autotune config plus a zero-initialized tiled workspace reduced with an ordered `sum` | timing-based autotune, uninitialized workspace | +| Mamba/SSM Triton ops | `megatron/core/ssm/ops/common/determinism.py` | one fixed autotune config plus a zero-initialized tiled workspace reduced with an ordered `sum` | timing-based autotune, uninitialized workspace | | Transformer Engine attention | `megatron/core/extensions/transformer_engine.py` | requires `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`, under which TE selects only backends that support deterministic execution (including deterministic FlashAttention backward) | TE picks freely, including atomic-accumulation attention backward | | Inference DP scheduling and RL rollout order | `megatron/core/inference/engines/dynamic_engine.py`, `megatron/rl/rl_utils.py` | sort by stable key | completion order | diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index b85a5a4756c..a9d939b6a15 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -23,7 +23,7 @@ from megatron.core.ssm.gdp_context_parallel import GDPContextParallel # Decode uses the in-repo Triton conv update, which accepts int64 slot indices. -from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update +from megatron.core.ssm.ops.common.causal_conv1d_triton import causal_conv1d_update from megatron.core.ssm.packed_seq_helpers import ( build_packed_seq_idx, check_fla_sequence_packing_support, diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 80c8c894c55..bbc435dfb11 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -23,13 +23,13 @@ from megatron.core.inference.utils import InferenceMode from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.ops.batch_invariant_decode import MambaBatchInvariantDecode -from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update -from megatron.core.ssm.ops.intermediate_extraction import ( +from megatron.core.ssm.ops.common.causal_conv1d_triton import causal_conv1d_update +from megatron.core.ssm.ops.common.intermediate_extraction import ( scatter_intermediate_conv, scatter_intermediate_ssm, ) -from megatron.core.ssm.ops.mamba_ssm import selective_state_update +from megatron.core.ssm.ops.mamba2.batch_invariant_decode import MambaBatchInvariantDecode +from megatron.core.ssm.ops.mamba2.mamba_ssm import selective_state_update from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.ssm.utils import _split_tensor_factory from megatron.core.tensor_parallel import get_cuda_rng_tracker @@ -80,7 +80,7 @@ HAVE_MAMBA_SSM = False try: - from megatron.core.ssm.ops.ssd_combined import mamba_chunk_scan_combined_varlen + from megatron.core.ssm.ops.mamba2.ssd_combined import mamba_chunk_scan_combined_varlen HAVE_SSM_OPS_VARLEN = True except ImportError: @@ -837,7 +837,7 @@ def ssm_prefill( conv_bias = self.cp.get_conv1d_bias().to(conv_state_dtype) xBC_pre_conv = xBC if intermediate_conv_out is not None else None - from megatron.core.ssm.ops.causal_conv1d_varlen import causal_conv1d_varlen_fn + from megatron.core.ssm.ops.common.causal_conv1d_varlen import causal_conv1d_varlen_fn xBC_out = causal_conv1d_varlen_fn( x=xBC.squeeze(0).contiguous(), diff --git a/megatron/core/ssm/ops/__init__.py b/megatron/core/ssm/ops/__init__.py index c8a06391843..2a42df6bef9 100644 --- a/megatron/core/ssm/ops/__init__.py +++ b/megatron/core/ssm/ops/__init__.py @@ -1,13 +1,28 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# Triton kernels for Mamba SSM (adapted from vLLM / state-spaces/mamba). + +"""Triton kernels for the linear-attention / SSM mixers. + +Organized by which recurrence a kernel belongs to: + +* `common` -- shared by every variant (causal convolution, slot-indexed state + gather/scatter, autotune determinism). +* `mamba2` -- the Mamba2 state-space duality kernels, adapted from vLLM and + state-spaces/mamba. + +The re-exports below are the historical top-level entry points and are kept so +existing callers do not have to spell out the subpackage. + +This package is internal to Megatron Core. The kernels here are implementation +details of the SSM mixers, not a supported kernel library. +""" try: - from .ssd_combined import mamba_chunk_scan_combined_varlen + from .mamba2.ssd_combined import mamba_chunk_scan_combined_varlen except ImportError: mamba_chunk_scan_combined_varlen = None try: - from .causal_conv1d_varlen import causal_conv1d_varlen_fn + from .common.causal_conv1d_varlen import causal_conv1d_varlen_fn except ImportError: causal_conv1d_varlen_fn = None diff --git a/megatron/core/ssm/ops/common/__init__.py b/megatron/core/ssm/ops/common/__init__.py new file mode 100644 index 00000000000..f968191a055 --- /dev/null +++ b/megatron/core/ssm/ops/common/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Triton kernels shared by every linear-attention / SSM mixer. + +Nothing here is specific to a recurrence. The short causal convolution that +front-runs the recurrence, the slot-indexed state gather/scatter used by prefix +caching, and the autotune-configuration filtering that makes the kernels +deterministic are all common to Mamba2 and its relatives. Recurrence-specific +kernels live in the sibling `mamba2` package. +""" diff --git a/megatron/core/ssm/ops/causal_conv1d_triton.py b/megatron/core/ssm/ops/common/causal_conv1d_triton.py similarity index 100% rename from megatron/core/ssm/ops/causal_conv1d_triton.py rename to megatron/core/ssm/ops/common/causal_conv1d_triton.py diff --git a/megatron/core/ssm/ops/causal_conv1d_varlen.py b/megatron/core/ssm/ops/common/causal_conv1d_varlen.py similarity index 99% rename from megatron/core/ssm/ops/causal_conv1d_varlen.py rename to megatron/core/ssm/ops/common/causal_conv1d_varlen.py index 091a5d2d34b..46e95ed6c76 100644 --- a/megatron/core/ssm/ops/causal_conv1d_varlen.py +++ b/megatron/core/ssm/ops/common/causal_conv1d_varlen.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from megatron.core.ssm.ops.determinism import autotune_configs +from megatron.core.ssm.ops.common.determinism import autotune_configs # Two block-dim regimes: diff --git a/megatron/core/ssm/ops/determinism.py b/megatron/core/ssm/ops/common/determinism.py similarity index 100% rename from megatron/core/ssm/ops/determinism.py rename to megatron/core/ssm/ops/common/determinism.py diff --git a/megatron/core/ssm/ops/intermediate_extraction.py b/megatron/core/ssm/ops/common/intermediate_extraction.py similarity index 96% rename from megatron/core/ssm/ops/intermediate_extraction.py rename to megatron/core/ssm/ops/common/intermediate_extraction.py index ec6552d798e..c44476b6c68 100644 --- a/megatron/core/ssm/ops/intermediate_extraction.py +++ b/megatron/core/ssm/ops/common/intermediate_extraction.py @@ -1,8 +1,12 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Fused gather + conditional-scatter kernels for Mamba intermediate-state +"""Fused gather + conditional-scatter kernels for SSM intermediate-state extraction used by prefix caching. +Nothing here is specific to a recurrence: the kernels only index rows of a state +tensor by slot, so they apply to Mamba2 and any other linear-attention / SSM +mixer that caches per-chunk intermediate states. + These replace the two-step ``states[indices]`` (dense gather) + ``.copy_()`` (scratch write) pattern with a single kernel that: diff --git a/megatron/core/ssm/ops/mamba2/__init__.py b/megatron/core/ssm/ops/mamba2/__init__.py new file mode 100644 index 00000000000..0d835e93112 --- /dev/null +++ b/megatron/core/ssm/ops/mamba2/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Triton kernels implementing the Mamba2 state-space duality (SSD) recurrence. + +The `ssd_*` modules are the chunked scan decomposition -- block matmul, chunk +state, state passing, chunk scan -- plus the `ssd_combined` entry point that +drives them. `mamba_ssm` holds the decode-step kernels and +`batch_invariant_decode` the buffered-replay variant used when +`batch_invariant_mode` is on. + +Kernels that are not specific to this recurrence (the causal convolution, state +gather/scatter, autotune determinism) live in `ops.common`. +""" diff --git a/megatron/core/ssm/ops/batch_invariant_decode.py b/megatron/core/ssm/ops/mamba2/batch_invariant_decode.py similarity index 97% rename from megatron/core/ssm/ops/batch_invariant_decode.py rename to megatron/core/ssm/ops/mamba2/batch_invariant_decode.py index fd532849d19..6ecda1f4c18 100644 --- a/megatron/core/ssm/ops/batch_invariant_decode.py +++ b/megatron/core/ssm/ops/mamba2/batch_invariant_decode.py @@ -7,10 +7,10 @@ import triton import triton.language as tl -from megatron.core.ssm.ops.ssd_bmm import _bmm_chunk_fwd -from megatron.core.ssm.ops.ssd_chunk_scan import _chunk_scan_fwd -from megatron.core.ssm.ops.ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd -from megatron.core.ssm.ops.ssd_state_passing import _state_passing_fwd +from megatron.core.ssm.ops.mamba2.ssd_bmm import _bmm_chunk_fwd +from megatron.core.ssm.ops.mamba2.ssd_chunk_scan import _chunk_scan_fwd +from megatron.core.ssm.ops.mamba2.ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd +from megatron.core.ssm.ops.mamba2.ssd_state_passing import _state_passing_fwd @triton.jit diff --git a/megatron/core/ssm/ops/mamba_ssm.py b/megatron/core/ssm/ops/mamba2/mamba_ssm.py similarity index 100% rename from megatron/core/ssm/ops/mamba_ssm.py rename to megatron/core/ssm/ops/mamba2/mamba_ssm.py diff --git a/megatron/core/ssm/ops/ssd_bmm.py b/megatron/core/ssm/ops/mamba2/ssd_bmm.py similarity index 99% rename from megatron/core/ssm/ops/ssd_bmm.py rename to megatron/core/ssm/ops/mamba2/ssd_bmm.py index 1b3a819ee14..a3a2b0302ba 100644 --- a/megatron/core/ssm/ops/ssd_bmm.py +++ b/megatron/core/ssm/ops/mamba2/ssd_bmm.py @@ -6,7 +6,7 @@ import triton import triton.language as tl -from megatron.core.ssm.ops.determinism import autotune_configs +from megatron.core.ssm.ops.common.determinism import autotune_configs @triton.autotune( diff --git a/megatron/core/ssm/ops/ssd_chunk_scan.py b/megatron/core/ssm/ops/mamba2/ssd_chunk_scan.py similarity index 99% rename from megatron/core/ssm/ops/ssd_chunk_scan.py rename to megatron/core/ssm/ops/mamba2/ssd_chunk_scan.py index 8066965b12b..e8bb4bdcd9b 100644 --- a/megatron/core/ssm/ops/ssd_chunk_scan.py +++ b/megatron/core/ssm/ops/mamba2/ssd_chunk_scan.py @@ -7,7 +7,7 @@ import triton.language as tl from packaging import version -from megatron.core.ssm.ops.determinism import autotune_configs +from megatron.core.ssm.ops.common.determinism import autotune_configs TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0") diff --git a/megatron/core/ssm/ops/ssd_chunk_state.py b/megatron/core/ssm/ops/mamba2/ssd_chunk_state.py similarity index 99% rename from megatron/core/ssm/ops/ssd_chunk_state.py rename to megatron/core/ssm/ops/mamba2/ssd_chunk_state.py index 8a21de1bef7..2e048b23f6a 100644 --- a/megatron/core/ssm/ops/ssd_chunk_state.py +++ b/megatron/core/ssm/ops/mamba2/ssd_chunk_state.py @@ -8,7 +8,7 @@ import triton.language as tl from packaging import version -from megatron.core.ssm.ops.determinism import autotune_configs +from megatron.core.ssm.ops.common.determinism import autotune_configs try: TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0") diff --git a/megatron/core/ssm/ops/ssd_combined.py b/megatron/core/ssm/ops/mamba2/ssd_combined.py similarity index 100% rename from megatron/core/ssm/ops/ssd_combined.py rename to megatron/core/ssm/ops/mamba2/ssd_combined.py diff --git a/megatron/core/ssm/ops/ssd_state_passing.py b/megatron/core/ssm/ops/mamba2/ssd_state_passing.py similarity index 99% rename from megatron/core/ssm/ops/ssd_state_passing.py rename to megatron/core/ssm/ops/mamba2/ssd_state_passing.py index 9c12537d7f1..4591e3a6d63 100644 --- a/megatron/core/ssm/ops/ssd_state_passing.py +++ b/megatron/core/ssm/ops/mamba2/ssd_state_passing.py @@ -7,7 +7,7 @@ import triton import triton.language as tl -from megatron.core.ssm.ops.determinism import autotune_configs +from megatron.core.ssm.ops.common.determinism import autotune_configs @triton.autotune( diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index abb21f36edb..5b8c6e678e6 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -1722,10 +1722,10 @@ def _pin_mamba_autotuners(): kernels = [] try: - from megatron.core.ssm.ops import ssd_bmm as r_bmm - from megatron.core.ssm.ops import ssd_chunk_scan as r_scan - from megatron.core.ssm.ops import ssd_chunk_state as r_state - from megatron.core.ssm.ops import ssd_state_passing as r_pass + from megatron.core.ssm.ops.mamba2 import ssd_bmm as r_bmm + from megatron.core.ssm.ops.mamba2 import ssd_chunk_scan as r_scan + from megatron.core.ssm.ops.mamba2 import ssd_chunk_state as r_state + from megatron.core.ssm.ops.mamba2 import ssd_state_passing as r_pass kernels += [ r_bmm._bmm_chunk_fwd_kernel, diff --git a/tests/unit_tests/ssm/test_causal_conv1d_triton.py b/tests/unit_tests/ssm/ops/common/test_causal_conv1d_triton.py similarity index 99% rename from tests/unit_tests/ssm/test_causal_conv1d_triton.py rename to tests/unit_tests/ssm/ops/common/test_causal_conv1d_triton.py index 624cd8c048b..d44ed3e1bbe 100644 --- a/tests/unit_tests/ssm/test_causal_conv1d_triton.py +++ b/tests/unit_tests/ssm/ops/common/test_causal_conv1d_triton.py @@ -3,7 +3,7 @@ import pytest import torch -from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update +from megatron.core.ssm.ops.common.causal_conv1d_triton import causal_conv1d_update def _requires_cuda(): diff --git a/tests/unit_tests/ssm/ops/test_causal_conv1d_varlen.py b/tests/unit_tests/ssm/ops/common/test_causal_conv1d_varlen.py similarity index 98% rename from tests/unit_tests/ssm/ops/test_causal_conv1d_varlen.py rename to tests/unit_tests/ssm/ops/common/test_causal_conv1d_varlen.py index b09457363bd..7385535080f 100644 --- a/tests/unit_tests/ssm/ops/test_causal_conv1d_varlen.py +++ b/tests/unit_tests/ssm/ops/common/test_causal_conv1d_varlen.py @@ -9,7 +9,7 @@ import pytest import torch -from megatron.core.ssm.ops.causal_conv1d_varlen import causal_conv1d_varlen_fn +from megatron.core.ssm.ops.common.causal_conv1d_varlen import causal_conv1d_varlen_fn try: from causal_conv1d import causal_conv1d_fn diff --git a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py b/tests/unit_tests/ssm/ops/mamba2/test_batch_invariant_decode.py similarity index 99% rename from tests/unit_tests/ssm/ops/test_batch_invariant_decode.py rename to tests/unit_tests/ssm/ops/mamba2/test_batch_invariant_decode.py index 904a6ab9e22..6167947a5d6 100644 --- a/tests/unit_tests/ssm/ops/test_batch_invariant_decode.py +++ b/tests/unit_tests/ssm/ops/mamba2/test_batch_invariant_decode.py @@ -8,11 +8,11 @@ try: from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined - from megatron.core.ssm.ops.batch_invariant_decode import ( + from megatron.core.ssm.ops.mamba2.batch_invariant_decode import ( BatchInvariantDecodeBuffers, batch_invariant_decode_buffered_scan, ) - from megatron.core.ssm.ops.ssd_combined import mamba_chunk_scan_combined_varlen + from megatron.core.ssm.ops.mamba2.ssd_combined import mamba_chunk_scan_combined_varlen HAVE_BATCH_INVARIANT_DECODE = True except ImportError: diff --git a/tests/unit_tests/ssm/ops/test_ssd_bmm.py b/tests/unit_tests/ssm/ops/mamba2/test_ssd_bmm.py similarity index 98% rename from tests/unit_tests/ssm/ops/test_ssd_bmm.py rename to tests/unit_tests/ssm/ops/mamba2/test_ssd_bmm.py index c4c4eef5404..58015f11551 100644 --- a/tests/unit_tests/ssm/ops/test_ssd_bmm.py +++ b/tests/unit_tests/ssm/ops/mamba2/test_ssd_bmm.py @@ -5,7 +5,7 @@ import torch try: - from megatron.core.ssm.ops.ssd_bmm import _bmm_chunk_fwd + from megatron.core.ssm.ops.mamba2.ssd_bmm import _bmm_chunk_fwd HAVE_SSD_OPS = True except (ImportError, Exception): diff --git a/tests/unit_tests/ssm/ops/test_ssd_chunk_scan.py b/tests/unit_tests/ssm/ops/mamba2/test_ssd_chunk_scan.py similarity index 98% rename from tests/unit_tests/ssm/ops/test_ssd_chunk_scan.py rename to tests/unit_tests/ssm/ops/mamba2/test_ssd_chunk_scan.py index dd654743346..969b3240d38 100644 --- a/tests/unit_tests/ssm/ops/test_ssd_chunk_scan.py +++ b/tests/unit_tests/ssm/ops/mamba2/test_ssd_chunk_scan.py @@ -5,7 +5,7 @@ import torch try: - from megatron.core.ssm.ops.ssd_chunk_scan import _chunk_scan_fwd + from megatron.core.ssm.ops.mamba2.ssd_chunk_scan import _chunk_scan_fwd HAVE_SSD_OPS = True except (ImportError, Exception): diff --git a/tests/unit_tests/ssm/ops/test_ssd_chunk_state.py b/tests/unit_tests/ssm/ops/mamba2/test_ssd_chunk_state.py similarity index 99% rename from tests/unit_tests/ssm/ops/test_ssd_chunk_state.py rename to tests/unit_tests/ssm/ops/mamba2/test_ssd_chunk_state.py index b848b6d7c96..c502f4ce7c4 100644 --- a/tests/unit_tests/ssm/ops/test_ssd_chunk_state.py +++ b/tests/unit_tests/ssm/ops/mamba2/test_ssd_chunk_state.py @@ -5,7 +5,7 @@ import torch try: - from megatron.core.ssm.ops.ssd_chunk_state import ( + from megatron.core.ssm.ops.mamba2.ssd_chunk_state import ( _chunk_cumsum_fwd, _chunk_state_fwd, chunk_state_varlen, diff --git a/tests/unit_tests/ssm/ops/test_ssd_combined.py b/tests/unit_tests/ssm/ops/mamba2/test_ssd_combined.py similarity index 99% rename from tests/unit_tests/ssm/ops/test_ssd_combined.py rename to tests/unit_tests/ssm/ops/mamba2/test_ssd_combined.py index 6ca8f466023..ea3661f7768 100644 --- a/tests/unit_tests/ssm/ops/test_ssd_combined.py +++ b/tests/unit_tests/ssm/ops/mamba2/test_ssd_combined.py @@ -5,11 +5,14 @@ import torch try: - from megatron.core.ssm.ops.intermediate_extraction import ( + from megatron.core.ssm.ops.common.intermediate_extraction import ( scatter_intermediate_conv, scatter_intermediate_ssm, ) - from megatron.core.ssm.ops.ssd_combined import is_int_pow_2, mamba_chunk_scan_combined_varlen + from megatron.core.ssm.ops.mamba2.ssd_combined import ( + is_int_pow_2, + mamba_chunk_scan_combined_varlen, + ) HAVE_SSD_OPS = True except (ImportError, Exception): diff --git a/tests/unit_tests/ssm/ops/test_ssd_state_passing.py b/tests/unit_tests/ssm/ops/mamba2/test_ssd_state_passing.py similarity index 97% rename from tests/unit_tests/ssm/ops/test_ssd_state_passing.py rename to tests/unit_tests/ssm/ops/mamba2/test_ssd_state_passing.py index f791d615078..5b28b54e018 100644 --- a/tests/unit_tests/ssm/ops/test_ssd_state_passing.py +++ b/tests/unit_tests/ssm/ops/mamba2/test_ssd_state_passing.py @@ -5,7 +5,7 @@ import torch try: - from megatron.core.ssm.ops.ssd_state_passing import _state_passing_fwd + from megatron.core.ssm.ops.mamba2.ssd_state_passing import _state_passing_fwd HAVE_SSD_OPS = True except (ImportError, Exception): diff --git a/tests/unit_tests/ssm/ops/test_ssm_kernel.py b/tests/unit_tests/ssm/ops/mamba2/test_ssm_kernel.py similarity index 100% rename from tests/unit_tests/ssm/ops/test_ssm_kernel.py rename to tests/unit_tests/ssm/ops/mamba2/test_ssm_kernel.py From d91896769cfbc74e68774486d27b82bc8d78617e Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 13 Aug 2026 23:57:34 +0000 Subject: [PATCH 031/101] Add read/write ledger for bulky inference metadata (#6125) Signed-off-by: Teodor-Dumitru Ene --- .../core/inference/engines/dynamic_engine.py | 15 +- megatron/core/inference/inference_request.py | 33 +++ .../endpoints/chat_completions.py | 11 +- .../endpoints/completions.py | 11 +- megatron/rl/agent/api.py | 7 +- megatron/rl/agent/reward_only_agent.py | 7 +- megatron/rl/inference/api.py | 4 +- megatron/rl/inference/megatron.py | 20 +- megatron/rl/rl_profiling.py | 9 +- megatron/rl/rl_utils.py | 146 +++++---- .../inference/engines/test_dynamic_engine.py | 61 ++++ .../inference/test_inference_request.py | 19 +- tests/unit_tests/rl/test_rl_utils.py | 276 ++++++++++++++---- .../unit_tests/rl/test_rollout_generation.py | 33 +-- 14 files changed, 475 insertions(+), 177 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 63b4176bf37..79a84981f04 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -40,6 +40,7 @@ DynamicInferenceEventType, DynamicInferenceRequest, DynamicInferenceRequestRecord, + FinishedRequestRecord, Status, ) from megatron.core.inference.sampling_params import SamplingParams @@ -356,6 +357,8 @@ def reset(self) -> None: # Generated token count already streamed for each request. self._partial_emit_lengths: Dict[int, int] = {} self._generation_epoch: Optional[int] = None + self.local_metadata_ledger_enabled: bool = False + self.local_metadata_ledger: dict[str, FinishedRequestRecord] = {} # Track requests that should stop due to stop words (detected in post_process_requests) self.stop_word_finished_request_ids: set[int] = set() # Track requests currently being finished due to stop words (to skip extra token) @@ -2335,8 +2338,18 @@ async def async_bookkeep( ] if records_to_send: nvtx_range_push("coordinator_communication") + merged_requests = [r.merge() for r in records_to_send] + if self.local_metadata_ledger_enabled: + # Index every finished request's extra metadata before it is dropped. + for merged in merged_requests: + assert ( + merged.uid not in self.local_metadata_ledger + ), f"finished-request ledger: duplicate uid {merged.uid!r}" + self.local_metadata_ledger[merged.uid] = FinishedRequestRecord.from_request( + merged + ) payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [r.merge().serialize() for r in records_to_send]], + [Headers.ENGINE_REPLY.value, [m.serialize() for m in merged_requests]], use_bin_type=True, ) self.socket_for_receiving_requests.send(payload) diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index a0d2e337c7a..d13859c07b6 100644 --- a/megatron/core/inference/inference_request.py +++ b/megatron/core/inference/inference_request.py @@ -3,6 +3,7 @@ import copy import hashlib import time +import uuid import warnings from dataclasses import asdict, dataclass, field from enum import Enum, auto @@ -366,6 +367,9 @@ class DynamicInferenceRequest(InferenceRequest): """ request_id: int + # `request_id` is per-engine, do not cross DP, and do not persist. + # A uuid is globally unique and can be used to track down individual requests. + uid: str = field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}") prompt: Optional[str] = None prompt_tokens: Optional[torch.Tensor] = None # remaining prompt tokens are used for chunked prefill @@ -748,6 +752,7 @@ def merge_lists(key): # Merged request. request = DynamicInferenceRequest( request_id=self.requests[0].request_id, + uid=self.requests[0].uid, prompt=prompt_text, prompt_tokens=prompt_tokens, prompt_log_probs=self.requests[0].prompt_log_probs, @@ -802,6 +807,34 @@ def deserialize(cls, obj: dict) -> "DynamicInferenceRequestRecord": return request +@dataclass +class FinishedRequestRecord: + """Stores per-request metadata that is not meant to be passed through the RESTful server.""" + + policy_epoch: Optional[list[tuple[int, int]]] + kv_cache_epoch: Optional[list[tuple[int, int]]] + num_evictions: int + + @classmethod + def from_request(cls, request: "DynamicInferenceRequest") -> "FinishedRequestRecord": + """Build the request's non-RESTful metadata from a finished request.""" + # Epoch stamps exist only while the engine has a generation epoch set. + record = cls( + policy_epoch=( + None if request.policy_epoch is None else [tuple(b) for b in request.policy_epoch] + ), + kv_cache_epoch=( + None + if request.kv_cache_epoch is None + else [tuple(b) for b in request.kv_cache_epoch] + ), + num_evictions=sum( + 1 for event in request.events if event.type is DynamicInferenceEventType.EVICT + ), + ) + return record + + @dataclass(kw_only=True) class VLMInferenceRequest(InferenceRequest): """Class for a VLM inference request""" diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index b0234e54f70..23a55d0902d 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -718,8 +718,11 @@ async def chat_completions(): # at submit time (above) and drive both the response shape here and whether the # engine kept the prompt_tokens tensor on the payload. request_idx = 0 + response_uid = None for result_item in batch_results: result = unwrap_serialized_tensors(result_item) + if response_uid is None: + response_uid = result["uid"] text_output = result["generated_text"] # The engine always reports prompt_length (for usage), but drops the @@ -805,12 +808,8 @@ async def chat_completions(): if return_raw_text: prompt_str = tokenizer.detokenize(result["prompt_tokens"]) message["raw_text"] = prompt_str + text_output - # Small RL/debug scalars (a few bytes each); harmless to keep for - # NeMo-RL compatibility. + # Small RL/debug scalars (a few bytes each); harmless to keep for NeMo-RL compatibility. message["generation_log_probs"] = result.get("generated_log_probs", []) - message["policy_epoch"] = result["policy_epoch"] - message["kv_cache_epoch"] = result["kv_cache_epoch"] - message["num_evictions"] = sum(1 for e in result["events"] if e.get("type") == "EVICT") return_log_probs = sampling_params.return_log_probs # Determine finish_reason following vLLM conventions: @@ -859,7 +858,7 @@ async def chat_completions(): prompt_token_count = max(prompt_tokens_counts) if prompt_tokens_counts else 0 cached_token_count = max(cached_tokens_counts) if cached_tokens_counts else 0 response = { - "id": f"chatcmpl-{uuid.uuid4().hex}", + "id": response_uid, "created": int(time.time()), "model": "EMPTY", "object": "chat.completion", diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py index ab217a8e89f..d5795b65c68 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py @@ -3,7 +3,6 @@ import asyncio import logging import time -import uuid from megatron.core.inference.inference_request import unwrap_serialized_tensors from megatron.core.inference.sampling_params import SamplingParams @@ -208,8 +207,11 @@ async def completions(): prompt_tokens_counts = [] request_idx = 0 + response_uid = None for completed_request in batch_results: result = unwrap_serialized_tensors(completed_request) + if response_uid is None: + response_uid = result["uid"] full_text = result["generated_text"] or "" text_output = (prompts_as_strings[request_idx] + full_text) if echo else full_text @@ -296,11 +298,6 @@ async def completions(): "generation_token_ids": result["generated_tokens"], "generation_log_probs": result.get("generated_log_probs", []), } - choice_data["policy_epoch"] = result["policy_epoch"] - choice_data["kv_cache_epoch"] = result["kv_cache_epoch"] - choice_data["num_evictions"] = sum( - 1 for e in result["events"] if e.get("type") == "EVICT" - ) if result["routing_indices"] is not None: choice_data["moe_topk_indices"] = result["routing_indices"] @@ -318,7 +315,7 @@ async def completions(): prompt_token_count = max(prompt_tokens_counts) if prompt_tokens_counts else 0 return jsonify( { - "id": str(uuid.uuid4()), + "id": response_uid, "object": "text_completion", # as per the openAI spec "created": int(time.time()), "model": "EMPTY", diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index b4ffb35048f..310869d142e 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -48,9 +48,6 @@ class Rollout(AgentBaseModel): reward: float = None env_id: str = '' problem_id: str | None = None - policy_epoch: list[list[tuple[int, int]]] - kv_cache_epoch: list[list[tuple[int, int]]] - num_evictions: list[int] class TokenRollout(AgentBaseModel): @@ -62,9 +59,7 @@ class TokenRollout(AgentBaseModel): logprobs: list[list[float]] | None = None env_id: str = '' problem_id: str | None = None - policy_epoch: list[list[tuple[int, int]]] - kv_cache_epoch: list[list[tuple[int, int]]] - num_evictions: list[int] + completion_ids: list[str] = [] Rollouts = list[TokenRollout | Rollout] diff --git a/megatron/rl/agent/reward_only_agent.py b/megatron/rl/agent/reward_only_agent.py index 67d5ca13f57..45e470f8984 100644 --- a/megatron/rl/agent/reward_only_agent.py +++ b/megatron/rl/agent/reward_only_agent.py @@ -202,9 +202,7 @@ async def _rollout_from_episode( ], env_id=self.env_id, problem_id=problem_id, - policy_epoch=[r.policy_epoch for r in responses], - kv_cache_epoch=[r.kv_cache_epoch for r in responses], - num_evictions=[r.num_evictions for r in responses], + completion_ids=[r.completion_id for r in responses], ) else: return Rollout( @@ -212,9 +210,6 @@ async def _rollout_from_episode( reward=reward, env_id=self.env_id, problem_id=problem_id, - policy_epoch=[r.policy_epoch for r in responses], - kv_cache_epoch=[r.kv_cache_epoch for r in responses], - num_evictions=[r.num_evictions for r in responses], ) async def get_rollout_response( diff --git a/megatron/rl/inference/api.py b/megatron/rl/inference/api.py index 7455a8791a8..bdd761cf82b 100644 --- a/megatron/rl/inference/api.py +++ b/megatron/rl/inference/api.py @@ -24,6 +24,4 @@ class InferenceResponse(BaseModel): prompt_length: int | None = None logprobs: list[float] | None = None finish_reason: str - policy_epoch: list[tuple[int, int]] - kv_cache_epoch: list[tuple[int, int]] - num_evictions: int + completion_id: str | None = None diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index 1bbb2fb3f25..6f6610f6289 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -17,6 +17,7 @@ from megatron.core.inference.config import KVCacheManagementMode from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine, EngineState from megatron.core.inference.inference_client import InferenceClient +from megatron.core.inference.inference_request import FinishedRequestRecord from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.utils import log_single_rank from megatron.training.global_vars import get_args, get_tokenizer @@ -83,9 +84,7 @@ async def base_generate(self, request: InferenceRequest) -> InferenceResponse: logprobs=choice.message.generation_log_probs, finish_reason=choice.finish_reason, prompt_length=len(choice.message.prompt_token_ids), - policy_epoch=choice.message.policy_epoch, - kv_cache_epoch=choice.message.kv_cache_epoch, - num_evictions=choice.message.num_evictions, + completion_id=response.id, ) @classmethod @@ -108,6 +107,7 @@ async def launch(cls, model: GPTModel, **kwargs): args.skip_prompt_log_probs = True inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(model=model) + inference_engine.local_metadata_ledger_enabled = True dp_addr = await inference_engine.start_listening_to_data_parallel_coordinator( inference_coordinator_port=41521, launch_inference_coordinator=True, ) @@ -193,6 +193,20 @@ async def suspend(self): self._client.suspend_engines() await self._inference_engine.wait_until(EngineState.SUSPENDED) + def merge_global_request_ledgers(self) -> dict[str, FinishedRequestRecord]: + """Union every engine's local-metadata ledger and clear them.""" + engine = self._inference_engine + local, engine.local_metadata_ledger = engine.local_metadata_ledger, {} + shards = [None] * dist.get_world_size() + dist.all_gather_object(shards, local) + merged: dict[str, FinishedRequestRecord] = {} + for shard in shards: + merged.update(shard) + assert len(merged) == sum(len(shard) for shard in shards), ( + "finished-request ledger: duplicate uids across engine ledgers" + ) + return merged + async def resume(self): if self._inference_engine._state_events[EngineState.RUNNING].is_set(): return diff --git a/megatron/rl/rl_profiling.py b/megatron/rl/rl_profiling.py index 9c0a4916ebd..362b5c1af86 100644 --- a/megatron/rl/rl_profiling.py +++ b/megatron/rl/rl_profiling.py @@ -53,7 +53,8 @@ # Rollout collection breakdown "rl/inference-setup", "rl/collect-rollouts", - "rl/sync-rollouts", + "rl/sync-rollout-state", + "rl/sync-request-ledger", "rl/suspend-engine", # Optimizer offload/restore "rl/offload-optimizer-before-inference", @@ -116,7 +117,7 @@ "rl/rollout-collection": [ "rl/inference-setup", "rl/collect-rollouts", - "rl/sync-rollouts", + "rl/sync-rollout-state", "rl/suspend-engine", "rl/offload-optimizer-before-inference", "rl/restore-optimizer-after-inference", @@ -535,7 +536,7 @@ def get_max(name: str) -> float: phases["training"] = get_max("forward-backward") # Wait/sync time (proxy for load imbalance) - phases["sync_wait"] = get_max("rl/suspend-engine") + get_max("rl/sync-rollouts") + phases["sync_wait"] = get_max("rl/suspend-engine") + get_max("rl/sync-rollout-state") return phases @@ -848,7 +849,7 @@ def analyze_bottlenecks(profile_path: str, top_n: int = 10) -> str: "rl/compute-ref-logprobs", ], "Training": ["forward-backward"], - "Sync/Wait": ["rl/suspend-engine", "rl/sync-rollouts"], + "Sync/Wait": ["rl/suspend-engine", "rl/sync-rollout-state"], } for phase_name, timer_list in phases.items(): diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 19bf761c3ae..bfa15cacd6e 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -28,6 +28,7 @@ from megatron.core import mpu from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.inference.contexts.dynamic_context import HAVE_TORCH_MEMORY_SAVER +from megatron.core.inference.inference_request import FinishedRequestRecord from megatron.core.inference.unified_memory import ( advise_managed_module_parameters_preferred_location, prefetch_managed_module_parameters, @@ -307,6 +308,8 @@ class RLRuntimeState: def __init__(self): self.packing_context = None self.last_collection_iteration = 0 + # Persistent ledger of per-request metadata (e.g. staleness stats). + self.request_ledger = {} self.sequences_this_iteration_on_rank = 0 self.latest_batch_num_sequences = 0 # Derived throughput metrics (set by log_rl_throughput_metrics, read by RLProfiler). @@ -610,7 +613,7 @@ def get_rollout_generator( def get_environment_rollouts( model: LanguageModule, inference_model: LanguageModule, optimizer: MegatronOptimizer, n_prompts: int, samples_per_group: int -): +) -> tuple[GroupedRollouts, dict[str, FinishedRequestRecord]]: """Sample environment rollouts from an LLM. Args: @@ -620,7 +623,7 @@ def get_environment_rollouts( samples_per_group: Amount of trajectories per prompt. Returns: - GroupedRollouts object which is a nested list with each element being a list of rollouts of a group. + (GroupedRollouts, per-request metadata ledger) """ args = get_args() nvtx_range = get_nvtx_range() @@ -716,10 +719,13 @@ def get_environment_rollouts( # Just set up space to collect the rollouts rollouts = [[None for _ in range(samples_per_group)] for _ in range(n_prompts)] - with nvtx_range("rl/sync-rollouts", time=True): + with nvtx_range("rl/sync-rollout-state", time=True): # Wait for Rollouts to be collected # TODO(jbarker): double check why this isn't causing rank 0 memory allocations torch.distributed.broadcast_object_list(rollouts, src=0) + + with nvtx_range("rl/sync-request-ledger", time=True): + request_ledger = inference_interface.merge_global_request_ledgers() logger.debug(f"Got rollouts on rank {rank}") if lang_rl_log_dir and rank == get_pg_rank(inference_pg_collection.tp): @@ -731,7 +737,7 @@ def get_environment_rollouts( ) as f: json.dump([[r.model_dump() for r in group] for group in rollouts], f) - return rollouts + return rollouts, request_ledger def selective_log_softmax(logits, index): @@ -886,7 +892,10 @@ def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[i def compute_group_stats( - rollouts: GroupedRollouts, tokenizer: MegatronTokenizer, seq_len: int, + rollouts: GroupedRollouts, + tokenizer: MegatronTokenizer, + seq_len: int, + request_ledger: dict[str, FinishedRequestRecord], ) -> RolloutStats: """Add group-based rollout stats for logging. @@ -894,6 +903,7 @@ def compute_group_stats( rollouts: Rollouts to generate the stats for. Each inner list is a group (as in GRPO group), i.e. all rollouts are for the same prompt. tokenizer: Tokenizer to tokenize the rollouts in case they are raw strings. seq_len: Maximum sequence length. + request_ledger: Per-turn ledger of additional metadata (e.g. staleness stats). Returns: RolloutStats object containing all the stats. @@ -958,16 +968,19 @@ def compute_group_stats( roll_turn_lens = [len(t) for t in rollout.trajectory] group_turn_lengths.extend(roll_turn_lens) group_traj_lengths.append(sum(roll_turn_lens)) - assert rollout.policy_epoch, "Rollout has no policy_epoch data" - assert rollout.kv_cache_epoch, "Rollout has no kv_cache_epoch data" - group_policy_epoch.append([epoch for turn in rollout.policy_epoch for _, epoch in turn]) - group_kv_epoch.append([epoch for turn in rollout.kv_cache_epoch for _, epoch in turn]) - # completed_epochs is per-turn, so it cannot be masked per-rollout downstream - if rollout.trajectory: - group_completed_epochs.extend( - turn[-1][1] for turn in rollout.policy_epoch - ) - group_num_evictions.append(sum(rollout.num_evictions)) + turn_records = ( + [request_ledger.pop(completion_id) for completion_id in rollout.completion_ids] + if isinstance(rollout, TokenRollout) + else [] + ) + group_policy_epoch.append( + [epoch for record in turn_records for _, epoch in record.policy_epoch] + ) + group_kv_epoch.append( + [epoch for record in turn_records for _, epoch in record.kv_cache_epoch] + ) + group_completed_epochs.extend(record.policy_epoch[-1][1] for record in turn_records) + group_num_evictions.append(sum(record.num_evictions for record in turn_records)) all_policy_epoch.append(group_policy_epoch) all_kv_cache_epoch.append(group_kv_epoch) all_completed_epochs.append(group_completed_epochs) @@ -1085,35 +1098,37 @@ def _real(grouped): [x for x, keep in zip(g, m) if keep] for g, m in zip(grouped, real_mask) ] - # Reward metrics include failures. All other metrics do not. - table_rewards = [r for g in _real(rewards) for r in g] + # Reward metrics include failures. All other metrics do not. Staleness and + # eviction telemetry filters through `joined` instead: placeholders pop no + # ledger records, so they have no epoch rows and never join. traj_lens_real = _real(traj_lens) num_turns_real = _real(num_turns) - policy_epoch_real = _real(policy_epoch) - kv_cache_epoch_real = _real(kv_cache_epoch) group_table = wandb_writer.Table( columns=['group_means', 'group_stds'], data=[[np.mean(g), np.std(g)] for g in rewards], ) + flat_rewards = [r for g in rewards for r in g] + flat_traj_lens = [l for g in traj_lens for l in g] + flat_num_evictions = [e for g in num_evictions for e in g] + flat_policy_epochs = [r for g in policy_epoch for r in g] + flat_kv_epochs = [r for g in kv_cache_epoch for r in g] + joined = [i for i, row in enumerate(flat_policy_epochs) if row] + joined_rewards = [flat_rewards[i] for i in joined] + joined_traj_lens = [flat_traj_lens[i] for i in joined] + joined_num_evictions = [flat_num_evictions[i] for i in joined] + per_rollout_policy_epochs = [flat_policy_epochs[i] for i in joined] + per_rollout_kv_epochs = [flat_kv_epochs[i] for i in joined] # Per-rollout staleness (oldest token) - rollout_policy_staleness = [current_iteration - r[0] for g in policy_epoch_real for r in g] - rollout_kv_staleness = [current_iteration - r[0] for g in kv_cache_epoch_real for r in g] + rollout_policy_staleness = [current_iteration - r[0] for r in per_rollout_policy_epochs] + rollout_kv_staleness = [current_iteration - r[0] for r in per_rollout_kv_epochs] # Per-rollout staleness (newest token) - rollout_policy_last_token_staleness = [ - current_iteration - r[-1] for g in policy_epoch_real for r in g - ] - rollout_kv_last_token_staleness = [ - current_iteration - r[-1] for g in kv_cache_epoch_real for r in g - ] + rollout_policy_last_token_staleness = [current_iteration - r[-1] for r in per_rollout_policy_epochs] + rollout_kv_last_token_staleness = [current_iteration - r[-1] for r in per_rollout_kv_epochs] # Per-token staleness - per_token_policy_staleness = [ - current_iteration - e for g in policy_epoch_real for r in g for e in r - ] - per_token_kv_staleness = [ - current_iteration - e for g in kv_cache_epoch_real for r in g for e in r - ] + per_token_policy_staleness = [current_iteration - e for r in per_rollout_policy_epochs for e in r] + per_token_kv_staleness = [current_iteration - e for r in per_rollout_kv_epochs for e in r] metrics = { 'group_means_hist': wandb_writer.plot.histogram( @@ -1134,7 +1149,28 @@ def _real(grouped): ), 'advantages', 'Advantages' ), - # One row per real rollout. + # Group-level length/turn stats skip groups with all failed rollouts. + 'mean_turn_length': np.mean([np.mean(g) for g in turn_lens if g]), + 'mean_turn_length_std': np.mean([np.std(g) for g in turn_lens if g]), + 'max_turn_length': max(max(g) for g in turn_lens if g), + 'min_turn_length': min(min(g) for g in turn_lens if g), + 'mean_traj_length': np.mean([np.mean(g) for g in traj_lens_real if g]), + 'mean_traj_length_std': np.mean([np.std(g) for g in traj_lens_real if g]), + 'max_traj_length': max(max(g) for g in traj_lens_real if g), + 'min_traj_length': min(min(g) for g in traj_lens_real if g), + 'mean_num_turns': np.mean([np.mean(g) for g in num_turns_real if g]), + 'max_num_turns': max(max(g) for g in num_turns_real if g), + 'min_num_turns': min(min(g) for g in num_turns_real if g), + 'mean_reward': np.mean([np.mean(g) for g in rewards]), + 'mean_advantage': np.mean(advantages), + 'nonzero_groups_ratio': np.count_nonzero(advantages) + / len(advantages), + **failure_metrics, + } + + # Staleness/eviction telemetry: joined rollouts only; skip when nothing joined. + if joined: + metrics |= { 'rollout_table': wandb_writer.Table( columns=[ 'reward', 'traj_length', 'num_evictions', @@ -1142,9 +1178,9 @@ def _real(grouped): 'policy_last_token_staleness', 'kv_last_token_staleness', ], data=list(zip( - table_rewards, - [l for g in traj_lens_real for l in g], - [e for g in _real(num_evictions) for e in g], + joined_rewards, + joined_traj_lens, + joined_num_evictions, rollout_policy_staleness, rollout_kv_staleness, rollout_policy_last_token_staleness, @@ -1156,22 +1192,6 @@ def _real(grouped): columns=['policy_staleness', 'kv_staleness'], data=list(zip(per_token_policy_staleness, per_token_kv_staleness)), ), - # Group-level length/turn stats skip groups with all failed rollouts. - 'mean_turn_length': np.mean([np.mean(g) for g in turn_lens if g]), - 'mean_turn_length_std': np.mean([np.std(g) for g in turn_lens if g]), - 'max_turn_length': max(max(g) for g in turn_lens if g), - 'min_turn_length': min(min(g) for g in turn_lens if g), - 'mean_traj_length': np.mean([np.mean(g) for g in traj_lens_real if g]), - 'mean_traj_length_std': np.mean([np.std(g) for g in traj_lens_real if g]), - 'max_traj_length': max(max(g) for g in traj_lens_real if g), - 'min_traj_length': min(min(g) for g in traj_lens_real if g), - 'mean_num_turns': np.mean([np.mean(g) for g in num_turns_real if g]), - 'max_num_turns': max(max(g) for g in num_turns_real if g), - 'min_num_turns': min(min(g) for g in num_turns_real if g), - 'mean_reward': np.mean([np.mean(g) for g in rewards]), - 'mean_advantage': np.mean(advantages), - 'nonzero_groups_ratio': np.count_nonzero(advantages) - / len(advantages), 'mean_policy_staleness': np.mean(rollout_policy_staleness), 'max_policy_staleness': max(rollout_policy_staleness), 'min_policy_staleness': min(rollout_policy_staleness), @@ -1184,8 +1204,8 @@ def _real(grouped): 'mean_kv_cache_last_token_staleness': np.mean(rollout_kv_last_token_staleness), 'max_kv_cache_last_token_staleness': max(rollout_kv_last_token_staleness), 'min_kv_cache_last_token_staleness': min(rollout_kv_last_token_staleness), - 'total_eviction_count': sum([sum(g) for g in num_evictions]), - 'max_num_evictions': max([max(g) for g in num_evictions]), + 'total_eviction_count': sum(joined_num_evictions), + 'max_num_evictions': max(joined_num_evictions), 'mean_completion_gap': np.mean([current_iteration - s for g in completed_epochs for s in g]), 'per_token_policy_staleness_hist': wandb_writer.plot.histogram( wandb_writer.Table(columns=['staleness'], data=[[s] for s in per_token_policy_staleness]), @@ -1195,8 +1215,9 @@ def _real(grouped): wandb_writer.Table(columns=['staleness'], data=[[s] for s in per_token_kv_staleness]), 'staleness', 'Per-Token KV Cache Staleness' ), - **failure_metrics, - } + } + + if example_group: if tokenizer is None: raise ValueError("If you provide an example group to log, you need to provide a tokenizer too.") @@ -1372,6 +1393,7 @@ def maybe_log_training_metrics( traj_lens=traj_lens, turn_lens=turn_lens, rewards=rewards, num_turns=num_turns, advantages=advantages, policy_epoch=policy_epoch, kv_cache_epoch=kv_cache_epoch, completed_epochs=completed_epochs, num_evictions=num_evictions, current_iteration=current_iteration) + env_stats = lambda cont, idx: [cont[i] for i in idx] group_turn_counts = [sum(nt) for nt in num_turns] @@ -1639,6 +1661,7 @@ def prepare_data_for_update( tokenizer: MegatronTokenizer, sequence_packing: bool, is_correction: bool, + request_ledger: dict[str, FinishedRequestRecord], ) -> tuple[RerunDataIterator, RolloutStats, dict]: """Extract data for the update from raw rollouts. @@ -1649,6 +1672,7 @@ def prepare_data_for_update( tokenizer: Tokenizer to pad/tokenize data. sequence_packing: Use sequence packing if True. is_correction: Prepare data for IS correction if True. + request_ledger: Finished-request records keyed by the completion id. Returns: Tuple of (cycled iterator over dataset batches, group stats, example groups per env). @@ -1668,7 +1692,7 @@ def prepare_data_for_update( with nvtx_range("rl/prepare-data-for-update", time=True): with nvtx_range("rl/compute-group-stats", time=True): - group_stats = compute_group_stats(rollouts, tokenizer, args.seq_length) + group_stats = compute_group_stats(rollouts, tokenizer, args.seq_length, request_ledger) # TODO(vitalyk): why do we need global_advantages here? go inside packing advantages = global_advantages = torch.tensor(group_stats.advantages, dtype=dtype).cuda() @@ -1992,9 +2016,12 @@ def get_grpo_data_iterator( (grpo_iterations * global_batches_per_collection) ): - rollouts = get_environment_rollouts( + rollouts, fresh_ledger = get_environment_rollouts( model, inference_model, optimizer, grpo_prompts_per_step, grpo_group_size ) + # Records persist until their rollout's join pops them (partial rollouts + # deliver across window boundaries); never-joined residents are accepted. + runtime_state.request_ledger.update(fresh_ledger) buffered_rollouts, group_stats, example_groups = prepare_data_for_update( model=model, ref_state_dict=ref_state_dict, @@ -2002,6 +2029,7 @@ def get_grpo_data_iterator( tokenizer=tokenizer, sequence_packing=sequence_packing, is_correction=is_correction, + request_ledger=runtime_state.request_ledger, ) if optimizer_is_on_cpu: nvtx_range = get_nvtx_range() diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 0336881e22e..0806cddf974 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -3318,6 +3318,67 @@ def set_epoch(epoch): assert record[-1].policy_epoch == merged.policy_epoch assert record[-1].kv_cache_epoch is None + @pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" + ) + @torch.inference_mode() + def test_local_metadata_ledger_gated_by_engine_enable(self): + """Only ledger-enabled engines (RL training) index finished requests.""" + PROMPT_LEN = 8 + NUM_TOKENS = 4 + + test_config = DynamicEngineTestConfig( + num_requests=0, + min_prompt_length=PROMPT_LEN, + max_prompt_length=PROMPT_LEN, + num_tokens_to_generate=NUM_TOKENS, + ) + env = self._build_test_env(test_config) + engine = env.engine + engine._generation_epoch = 3 # RL mode: requests are epoch-stamped + # Take the coordinator reply path so finished requests flow through the + # inline ledger indexing in the reply block (socket mocked out). + engine.use_coordinator = True + engine.is_mp_coordinator = True + engine.socket_for_receiving_requests = mock.MagicMock() + + def run_request(request_id): + prompt_tokens = torch.full( + (PROMPT_LEN,), request_id + 1, dtype=torch.int64, device=torch.cuda.current_device() + ) + engine._add_request( + DynamicInferenceRequest( + request_id=request_id, + prompt_tokens=prompt_tokens, + sampling_params=SamplingParams( + num_tokens_to_generate=NUM_TOKENS, termination_id=-1 + ), + ) + ) + finished_records = [] + while engine.has_unfinished_requests(): + result = engine.step_modern() + finished_records.extend(result["finished_request_records"]) + return finished_records + + # Default (plain serving): nothing is indexed, nothing accumulates. + finished_records = run_request(0) + assert len(finished_records) == 1 + assert engine.local_metadata_ledger == {} + + # RL launch (MegatronLocal.launch) enables the ledger: every finished + # request is indexed, no per-request tagging involved. + engine.local_metadata_ledger_enabled = True + finished_records = run_request(1) + assert len(finished_records) == 1 + + # The ledger keys by the request's engine-minted uid — the same string the + # endpoints return as the OpenAI response id. + merged = finished_records[0].merge() + ledger = engine.local_metadata_ledger + assert list(ledger.keys()) == [merged.uid] + assert ledger[merged.uid].policy_epoch == [(0, 3)] + @pytest.mark.internal @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" diff --git a/tests/unit_tests/inference/test_inference_request.py b/tests/unit_tests/inference/test_inference_request.py index 6676ab734e2..2ab50e39798 100644 --- a/tests/unit_tests/inference/test_inference_request.py +++ b/tests/unit_tests/inference/test_inference_request.py @@ -11,6 +11,7 @@ DynamicInferenceEventType, DynamicInferenceRequest, DynamicInferenceRequestRecord, + FinishedRequestRecord, InferenceRequest, compute_block_hashes_batched, deserialize_ndarray, @@ -122,6 +123,8 @@ def test_inference_request_serialize_round_trip_through_msgpack(): dyn_out = DynamicInferenceRequest.deserialize(dyn_data) assert isinstance(dyn_out.routing_indices, np.ndarray) assert dyn_out.routing_indices.tolist() == [[1, 2], [3, 4]] + # The engine-minted uid (the OpenAI response id / ledger key) survives the wire. + assert dyn_out.uid == dyn.uid def test_dynamic_inference_request_post_init_prefix_caching(): @@ -246,6 +249,8 @@ def test_dynamic_inference_request_record_checkpoint_and_merge(): b.generated_text = "bar" a.routing_indices = np.array([[1, 2]]) b.routing_indices = np.array([[3, 4]]) + b.policy_epoch = [(0, 3)] + a.add_event_evict() rec = DynamicInferenceRequestRecord(requests=[a, b]) rec.latency = 4.2 merged = rec.merge() @@ -253,6 +258,14 @@ def test_dynamic_inference_request_record_checkpoint_and_merge(): assert merged.generated_text == "foobar" assert merged.generated_length == 3 and merged.latency == 4.2 assert merged.routing_indices.tolist() == [[1, 2], [3, 4]] + # Every request mints a distinct chatcmpl- uid; merge() keeps the FIRST + # segment's — the id the response and the finished-request ledger key on. + assert a.uid.startswith("chatcmpl-") and a.uid != b.uid + assert merged.uid == a.uid + # FinishedRequestRecord mirrors the merged stamps and counts EVICT events. + finished = FinishedRequestRecord.from_request(merged) + assert finished.policy_epoch == [(0, 3)] and finished.kv_cache_epoch is None + assert finished.num_evictions == 1 # merge() with both generated_text=None propagates None (rather than "None"+"None"). c = DynamicInferenceRequest( @@ -267,7 +280,11 @@ def test_dynamic_inference_request_record_checkpoint_and_merge(): sampling_params=sp, generated_tokens=[12], ) - assert DynamicInferenceRequestRecord(requests=[c, d]).merge().generated_text is None + merged_cd = DynamicInferenceRequestRecord(requests=[c, d]).merge() + assert merged_cd.generated_text is None + # Never-stamped requests record None epochs (non-RL serving). + finished_cd = FinishedRequestRecord.from_request(merged_cd) + assert finished_cd.policy_epoch is None and finished_cd.num_evictions == 0 def test_dynamic_inference_request_serialize_strips_event_add_engine(): diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index ef098669da7..3d7eb2a1a21 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -11,6 +11,7 @@ from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.enums import ModelType +from megatron.core.inference.inference_request import FinishedRequestRecord from megatron.core.models.common.language_module.language_module import LanguageModule from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.gpt.gpt_model import GPTModel @@ -36,7 +37,7 @@ from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope from megatron.core.transformer.module import Float16Module from megatron.rl import rl_utils -from megatron.rl.agent.api import TokenRollout +from megatron.rl.agent.api import Rollout, RolloutGroup, TokenRollout from megatron.rl.inference import ReturnsRaw from megatron.rl.sequence_packing_utils import get_default_packed_seq_params from megatron.training.arguments import parse_args, validate_args @@ -91,9 +92,16 @@ def detokenize(self, tokens): return [str(tok) for tok in tokens] -def make_token_rollout(trajectory, logprobs, generation_mask=None, reward=1.0, problem_id="p"): - """TokenRollout with the per-turn staleness boilerplate derived from the turn count.""" - turns = len(trajectory) +_NEXT_CID = itertools.count() + + +def _mint_cids(num_turns): + return [f"cid-{next(_NEXT_CID)}" for _ in range(num_turns)] + + +def make_token_rollout( + trajectory, logprobs, generation_mask=None, reward=1.0, problem_id="p", completion_ids=None +): return TokenRollout( trajectory=trajectory, reward=reward, @@ -101,12 +109,15 @@ def make_token_rollout(trajectory, logprobs, generation_mask=None, reward=1.0, p logprobs=logprobs, env_id='MEGAENV', problem_id=problem_id, - policy_epoch=[[(0, 0)]] * turns, - kv_cache_epoch=[[(0, 0)]] * turns, - num_evictions=[0] * turns, + completion_ids=_mint_cids(len(trajectory)) if completion_ids is None else completion_ids, ) +def _ledger_for(*rollouts, epoch=0): + """A uniform-epoch finished-request record per turn, keyed by completion id.""" + return {cid: _ledger_record(epoch) for r in rollouts for cid in r.completion_ids} + + class DummyConfigModule(torch.nn.Module): def __init__(self, config): super().__init__() @@ -575,6 +586,7 @@ def test_prepare_data_for_update(self, initialize_model_parallel): # A single-turn rollout whose only turn is short and lacks eod must be rejected: # a single-turn completion has no tool-call boundary to justify stopping early. + # (The eod tripwire fires before the ledger join, so no record is needed.) bad = make_token_rollout( [[1, 2, 3]], [[0.1, 0.2, 0.3]], [[False, True, True]], reward=3.14, problem_id="2" ) @@ -586,6 +598,7 @@ def test_prepare_data_for_update(self, initialize_model_parallel): tokenizer, sequence_packing=False, is_correction=False, + request_ledger={}, ) # Multi-turn rollouts with uneven turn counts: a turn may stop on a tool-call @@ -593,47 +606,68 @@ def test_prepare_data_for_update(self, initialize_model_parallel): # = 5*dp turns, padded up to the next multiple of micro_batch_size*dp (= 2*dp) # -> 6*dp turns. With samples_ratio 1 the calculator is sized by total turns # (6*dp -> 3 microbatches), not by rollout count (the pre-fix bug gave 1). - mt1 = make_token_rollout( - [[1, 2, 3], [1, 2, 3, 4]], - [[-0.1, -0.2], [-0.3, -0.4]], - [[False, True, True], [False, False, True, True]], - problem_id="1", - ) - mt2 = make_token_rollout( - [[1, 2], [1, 2, 3], [1, 2, 3, 4]], - [[-0.1], [-0.2], [-0.3]], - [[False, True], [False, False, True], [False, False, False, True]], - reward=0.0, - problem_id="3", - ) + def mt1(): + return make_token_rollout( + [[1, 2, 3], [1, 2, 3, 4]], + [[-0.1, -0.2], [-0.3, -0.4]], + [[False, True, True], [False, False, True, True]], + problem_id="1", + ) + + def mt2(): + return make_token_rollout( + [[1, 2], [1, 2, 3], [1, 2, 3, 4]], + [[-0.1], [-0.2], [-0.3]], + [[False, True], [False, False, True], [False, False, False, True]], + reward=0.0, + problem_id="3", + ) + + # Every rollout is a distinct engine request with its own completion ids, + # so each dp group gets freshly-built rollouts and the ledger one record + # per turn. + mt_rollouts = [[mt1(), mt2()] for _ in range(dp)] + mt_ledger = _ledger_for(*(r for g in mt_rollouts for r in g)) rl_utils.prepare_data_for_update( [model], {}, - [[mt1, mt2] for _ in range(dp)], + mt_rollouts, tokenizer, sequence_packing=False, is_correction=False, + request_ledger=mt_ledger, ) # 5*dp turns padded to 6*dp; 6*dp / (micro_batch_size 2 * dp) = 3 microbatches. assert get_num_microbatches() == 3 - r1 = make_token_rollout( - torch.tensor([[1, 2, 3, tokenizer.eod]], dtype=torch.float).cuda(), - torch.tensor([[-0.2, -0.3, -3.2]]).cuda(), - torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(), - reward=3.14, - problem_id="2", - ) - r2 = make_token_rollout( - torch.tensor([[1, 2, 234, tokenizer.eod]], dtype=torch.float).cuda(), - torch.tensor([[-0.2, -0.3, -1.2]]), - torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(), - reward=0.14, - problem_id="2", - ) - rollouts = [[r1, r2] for _ in range(dp)] + def r1(): + return make_token_rollout( + torch.tensor([[1, 2, 3, tokenizer.eod]], dtype=torch.float).cuda(), + torch.tensor([[-0.2, -0.3, -3.2]]).cuda(), + torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(), + reward=3.14, + problem_id="2", + ) + + def r2(): + return make_token_rollout( + torch.tensor([[1, 2, 234, tokenizer.eod]], dtype=torch.float).cuda(), + torch.tensor([[-0.2, -0.3, -1.2]]), + torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(), + reward=0.14, + problem_id="2", + ) + + rollouts = [[r1(), r2()] for _ in range(dp)] + request_ledger = _ledger_for(*(r for g in rollouts for r in g)) data_iter, _, _ = rl_utils.prepare_data_for_update( - [model], {}, rollouts, tokenizer, sequence_packing=False, is_correction=False + [model], + {}, + rollouts, + tokenizer, + sequence_packing=False, + is_correction=False, + request_ledger=request_ledger, ) _, _, old_logprobs, _, _, _, _ = next(data_iter) @@ -677,8 +711,15 @@ def single(problem_id, reward): ) rollouts = [[single(str(i), float(i % 2)) for i in range(4)] for _ in range(dp)] + request_ledger = _ledger_for(*(r for g in rollouts for r in g)) rl_utils.prepare_data_for_update( - [model], {}, rollouts, tokenizer, sequence_packing=False, is_correction=False + [model], + {}, + rollouts, + tokenizer, + sequence_packing=False, + is_correction=False, + request_ledger=request_ledger, ) assert get_num_microbatches() == 1 @@ -808,6 +849,7 @@ def single(traj, reward): if scenario == "single_turn_only": group = [single([1, 2, 3, eod], 1.0), single([1, 2, eod], 0.0)] + request_ledger = _ledger_for(*group) else: # Cumulative per-turn lengths 4 then 7 -> turn 1 adds 3 tokens; trajectory length is # the full conversation (7), not 4 + 7 = 11. @@ -818,8 +860,11 @@ def single(traj, reward): problem_id="m", ) group = [multi, single([1, 2, 3, eod], 0.0)] + request_ledger = _ledger_for(*group) - stats = rl_utils.compute_group_stats([group], tokenizer, seq_len=8) + stats = rl_utils.compute_group_stats( + [group], tokenizer, seq_len=8, request_ledger=request_ledger + ) assert stats.turn_lens == expected_turn_lens assert stats.traj_lens == expected_traj_lens assert stats.num_turns == expected_num_turns @@ -1251,8 +1296,9 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders lst.append([sentinel, sentinel]) # the fully failed extra group for lst in (policy_epoch, kv_cache_epoch): for group in lst: - group.append([0]) # sentinel epoch stamp of a placeholder - lst.append([[0], [0]]) + # Placeholders pop no ledger records, so their epoch rows are empty. + group.append([]) + lst.append([[], []]) # Placeholders contribute no turns, and compute_group_stats already # excludes them from completed_epochs; the failed group adds empty # inner lists, which the group-level stats must skip, not crash on. @@ -1322,7 +1368,7 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders assert metrics["mean_completion_gap"] == 2.5 def test_compute_group_stats_excludes_placeholders_from_metric_fields(self): - def real_rollout(tokens, epoch, problem_id): + def real_rollout(tokens, problem_id): return TokenRollout( trajectory=[tokens], generation_mask=[[True] * len(tokens)], @@ -1330,9 +1376,7 @@ def real_rollout(tokens, epoch, problem_id): logprobs=[[0.0] * len(tokens)], env_id="swe", problem_id=problem_id, - policy_epoch=[[(0, epoch)]], - kv_cache_epoch=[[(0, epoch)]], - num_evictions=[0], + completion_ids=_mint_cids(1), ) def placeholder(): @@ -1343,26 +1387,26 @@ def placeholder(): logprobs=[], env_id="swe", problem_id="placeholder", - policy_epoch=[[(0, 0)]], - kv_cache_epoch=[[(0, 0)]], - num_evictions=[0], ) eod = MockTokenizer().eod - rollouts = [ - [ - real_rollout([1, 2, eod], epoch=5, problem_id="p0"), - real_rollout([1, 2, 3, eod], epoch=6, problem_id="p0"), - placeholder(), - ], - [placeholder(), placeholder(), placeholder()], - ] - stats = rl_utils.compute_group_stats(rollouts, MockTokenizer(), seq_len=16) + first = real_rollout([1, 2, eod], problem_id="p0") + second = real_rollout([1, 2, 3, eod], problem_id="p0") + rollouts = [[first, second, placeholder()], [placeholder(), placeholder(), placeholder()]] + # Placeholders have no turns (and no completion ids): they pop no records. + ledger = { + first.completion_ids[0]: _ledger_record(5), + second.completion_ids[0]: _ledger_record(6), + } + stats = rl_utils.compute_group_stats( + rollouts, MockTokenizer(), seq_len=16, request_ledger=ledger + ) # Per-rollout lists keep the placeholder entries: alignment with rewards # and num_turns is what lets prep_wandb_metrics mask them downstream. + # Placeholders join no records, so their epoch rows are empty. assert stats.num_turns == [[1, 1, 0], [0, 0, 0]] - assert stats.policy_epoch == [[[5], [6], [0]], [[0], [0], [0]]] + assert stats.policy_epoch == [[[5], [6], []], [[], [], []]] assert stats.traj_lens == [[3, 4, 0], [0, 0, 0]] # Per-turn lists exclude placeholders entirely: no sentinel epoch-0 stamp # in completed_epochs, no fake 0-length turn for all-placeholder groups. @@ -1371,7 +1415,7 @@ def placeholder(): # Rewards keep the placeholder zeros (they shape the GRPO baseline). assert stats.rewards == [[1.0, 1.0, 0.0], [0.0, 0.0, 0.0]] - # End to end: the sentinel epochs never reach the staleness metrics. + # End to end: the placeholder rows never reach the staleness metrics. metrics = rl_utils.prep_wandb_metrics( MagicMock(), stats.traj_lens, @@ -1385,9 +1429,121 @@ def placeholder(): num_evictions=stats.num_evictions, current_iteration=7, ) - assert metrics["max_policy_staleness"] == 2 # 7 - 5, not 7 - 0 + assert metrics["max_policy_staleness"] == 2 # 7 - 5, not a placeholder sentinel assert metrics["min_traj_length"] == 3 assert metrics["min_num_turns"] == 1 assert metrics["mean_completion_gap"] == np.mean([2, 1]) assert metrics["failed_rollouts/count"] == 4 assert np.isclose(metrics["failed_rollouts/ratio"], 4 / 6) + + def test_request_ledger_join(self): + """compute_group_stats joins finished-request records to rollouts by each + turn's completion id: records pop exactly once (identical token content + joins independently, text rollouts join nothing) and an unknown id is a + hard error. Downstream, prep_wandb_metrics filters staleness/eviction + telemetry to joined rollouts and emits none when nothing joined.""" + eod = MockTokenizer().eod + + def token_rollout(turns, masks, completion_ids=None): + return TokenRollout( + trajectory=turns, + reward=1.0, + generation_mask=masks, + logprobs=[[0.1] * len(t) for t in turns], + env_id='MEGAENV', + completion_ids=_mint_cids(len(turns)) if completion_ids is None else completion_ids, + ) + + # Multi-turn: each turn's stream is the full conversation prefix plus that + # turn's generation; the mask's False-prefix marks the prompt. + multi_turn = token_rollout( + [[1, 2, 3, eod], [1, 2, 3, eod, 4, 5, eod]], + [[False, True, True, True], [False] * 5 + [True, True]], + ) + # Identical sampled token content: distinct requests, distinct ids, + # independent records. + twin_a = token_rollout([[8, 9, eod]], [[False, True, True]]) + twin_b = token_rollout([[8, 9, eod]], [[False, True, True]]) + other = token_rollout([[4, 4, eod]], [[False, True, True]]) + # Text rollouts carry no token ids: they join nothing and need no records. + # (Groups are kept the same size: the advantage calculation the stats now + # run rejects ragged group shapes.) + text = Rollout(trajectory=["hello"], reward=1.0, env_id='MEGAENV') + text2 = Rollout(trajectory=["world"], reward=1.0, env_id='MEGAENV') + + ledger = { + multi_turn.completion_ids[0]: _ledger_record(7, num_evictions=1), + multi_turn.completion_ids[1]: _ledger_record(8), + twin_a.completion_ids[0]: _ledger_record(5), + twin_b.completion_ids[0]: _ledger_record(6), + other.completion_ids[0]: _ledger_record(9), + } + rollouts = [ + RolloutGroup(rollouts=[multi_turn, twin_a, text]), + RolloutGroup(rollouts=[twin_b, other, text2]), + ] + stats = rl_utils.compute_group_stats(rollouts, MockTokenizer(), 8, ledger) + + assert stats.policy_epoch == [[[7, 8], [5], []], [[6], [9], []]] + assert stats.kv_cache_epoch == [[[7, 8], [5], []], [[6], [9], []]] + assert stats.completed_epochs == [[7, 8, 5], [6, 9]] + assert stats.num_evictions == [[1, 0, 0], [0, 0, 0]] + assert not ledger # all records consumed + + # An id matching no record raises KeyError at the join. + unknown = token_rollout([[6, 6, eod]], [[False, True, True]]) + with pytest.raises(KeyError, match="cid-"): + rl_utils.compute_group_stats([RolloutGroup(rollouts=[unknown])], MockTokenizer(), 8, {}) + + # Text (non-token) rollouts join no finished-request records and yield + # empty epoch rows; staleness/eviction telemetry filters them out while + # keeping the rollout table row-aligned. + writer = MagicMock() + metrics = rl_utils.prep_wandb_metrics( + writer, + traj_lens=[[3, 7]], + turn_lens=[[3, 7]], + rewards=[[1.0, -1.0]], + num_turns=[[1, 1]], + advantages=[0.5], + policy_epoch=[[[4, 5], []]], + kv_cache_epoch=[[[4, 5], []]], + completed_epochs=[[5]], + num_evictions=[[2, 0]], + current_iteration=6, + ) + # Only the joined rollout counts: staleness [6-4], evictions [2]. + assert metrics["mean_policy_staleness"] == 2 + assert metrics["max_policy_staleness"] == 2 + assert metrics["total_eviction_count"] == 2 + assert metrics["max_num_evictions"] == 2 + rollout_table = next( + c for c in writer.Table.call_args_list if "traj_length" in c.kwargs.get("columns", []) + ) + assert rollout_table.kwargs["data"] == [(1.0, 3, 2, 2, 2, 1, 1)] + + # A window where nothing joined emits no staleness/eviction metrics at all. + metrics = rl_utils.prep_wandb_metrics( + MagicMock(), + traj_lens=[[3]], + turn_lens=[[3]], + rewards=[[1.0]], + num_turns=[[1]], + advantages=[0.5], + policy_epoch=[[[]]], + kv_cache_epoch=[[[]]], + completed_epochs=[[]], + num_evictions=[[0]], + current_iteration=6, + ) + assert metrics["mean_reward"] == 1.0 + assert "mean_policy_staleness" not in metrics + assert "total_eviction_count" not in metrics + assert "rollout_table" not in metrics + assert "per_token_table" not in metrics + + +def _ledger_record(epoch, num_evictions=0): + return FinishedRequestRecord( + policy_epoch=[(0, epoch)], kv_cache_epoch=[(0, epoch)], num_evictions=num_evictions + ) diff --git a/tests/unit_tests/rl/test_rollout_generation.py b/tests/unit_tests/rl/test_rollout_generation.py index ab1d960836c..909b4cdc6e8 100644 --- a/tests/unit_tests/rl/test_rollout_generation.py +++ b/tests/unit_tests/rl/test_rollout_generation.py @@ -44,9 +44,6 @@ async def base_generate(self, request): response=LLMChatMessage(role="assistant", content=prompt), raw_text=prompt, finish_reason="stop", - policy_epoch=[(0, 0)], - kv_cache_epoch=[(0, 0)], - num_evictions=0, ) finally: self.active_requests -= 1 @@ -88,12 +85,7 @@ async def build_rollout(episode): responses = episode.responses reward = float(responses[-1].response.content.removeprefix("t")) return Rollout( - trajectory=[r.raw_text for r in responses], - reward=reward, - env_id=self.env_id, - policy_epoch=[r.policy_epoch for r in responses], - kv_cache_epoch=[r.kv_cache_epoch for r in responses], - num_evictions=[r.num_evictions for r in responses], + trajectory=[r.raw_text for r in responses], reward=reward, env_id=self.env_id ) return GroupRolloutParams(run_episode=run_episode, build_rollout=build_rollout) @@ -552,7 +544,9 @@ def test_multi_env_layout_rejects_starving_batch_size(self): assert [a.num_groups for a in mt.rollout_allocations(8)] == [6, 2] -def make_response(epochs, prompt_length, total_len, content="resp", finish_reason="stop"): +def make_response( + prompt_length, total_len, content="resp", finish_reason="stop", completion_id=None +): return InferenceResponse( response=LLMChatMessage(role="assistant", content=content), raw_text=content, @@ -560,17 +554,15 @@ def make_response(epochs, prompt_length, total_len, content="resp", finish_reaso prompt_length=prompt_length, logprobs=[0.0] * (total_len - prompt_length), finish_reason=finish_reason, - policy_epoch=epochs, - kv_cache_epoch=epochs, - num_evictions=0, + completion_id=completion_id, ) # Conversation length -> response spec: length 1 is the first turn (the bare prompt), length 3 # the second (assistant reply + observation appended). TWO_TURN_SCRIPT = { - 1: dict(epochs=[(0, 5)], prompt_length=3, total_len=7, content="a0"), - 3: dict(epochs=[(0, 5)], prompt_length=6, total_len=11, content="a1"), + 1: dict(prompt_length=3, total_len=7, content="a0", completion_id="cc-1"), + 3: dict(prompt_length=6, total_len=11, content="a1", completion_id="cc-3"), } # Both two-turn termination modes (env-signaled done, max_turns exhausted) must produce this @@ -580,7 +572,7 @@ def make_response(epochs, prompt_length, total_len, content="resp", finish_reaso reward_conv=[("user", "hello"), ("assistant", "a0"), ("user", "obs0"), ("assistant", "a1")], rewarded=[("a1", "stop")], genmask_sums=[4, 5], - policy_epoch=[[(0, 5)], [(0, 5)]], + completion_ids=["cc-1", "cc-3"], ) @@ -639,13 +631,13 @@ class TestMultiTurnEpisode: pytest.param( 1, None, - {1: dict(epochs=[(0, 7)], prompt_length=2, total_len=6, content="only")}, + {1: dict(prompt_length=2, total_len=6, content="only", completion_id="cc-1")}, dict( seen_roles=[["user"]], reward_conv=[("user", "hello"), ("assistant", "only")], rewarded=[("only", "stop")], genmask_sums=[4], - policy_epoch=[[(0, 7)]], + completion_ids=["cc-1"], observation_turns=[], ), id="single_turn", @@ -704,6 +696,8 @@ async def _drain(): assert isinstance(rollout, TokenRollout) assert rollout.reward == 1.5 assert rollout.problem_id == "p0" + # Per-turn backend response ids ride onto the rollout for the ledger join. + assert rollout.completion_ids == expected["completion_ids"] # One trajectory entry per generated turn. assert len(rollout.trajectory) == len(expected["genmask_sums"]) # Each turn's inference request = prior conversation (reply + observation appended). @@ -714,9 +708,6 @@ async def _drain(): assert agent.rewarded == expected["rewarded"] # Per-turn generation masks cover exactly each turn's generated tokens. assert [sum(mask) for mask in rollout.generation_mask] == expected["genmask_sums"] - # Per-turn (engine-frame) staleness nesting is preserved. - assert rollout.policy_epoch == expected["policy_epoch"] - assert rollout.kv_cache_epoch == expected["policy_epoch"] # get_observation is consulted only when another generation is still possible -- never on # the final allowed turn. assert agent.observation_turns == expected["observation_turns"] From d6d48cc372974bc1f24e27fd88cbd0f1b73bf55f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Fri, 14 Aug 2026 00:53:49 +0000 Subject: [PATCH 032/101] fix(ci): prevent queue stalls and stale merge-group runs (#6523) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- .../cancel_superseded_merge_group_runs.py | 210 ++++++++++++++++++ .github/scripts/test_approve_test_queue.sh | 31 ++- ...test_cancel_superseded_merge_group_runs.py | 92 ++++++++ .github/workflows/cicd-approve-test-queue.yml | 15 +- .../cicd-cancel-superseded-merge-groups.yml | 35 +++ .github/workflows/cicd-main.yml | 3 + 6 files changed, 383 insertions(+), 3 deletions(-) create mode 100644 .github/scripts/cancel_superseded_merge_group_runs.py create mode 100644 .github/scripts/test_cancel_superseded_merge_group_runs.py create mode 100644 .github/workflows/cicd-cancel-superseded-merge-groups.yml diff --git a/.github/scripts/cancel_superseded_merge_group_runs.py b/.github/scripts/cancel_superseded_merge_group_runs.py new file mode 100644 index 00000000000..b71fbe62156 --- /dev/null +++ b/.github/scripts/cancel_superseded_merge_group_runs.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Cancel active workflow runs for superseded generations of one queued PR.""" + +import json +import os +import re +import time +import urllib.error +import urllib.parse +import urllib.request + +_ACTIVE_STATUSES = ("requested", "queued", "in_progress", "waiting", "pending") +_MERGE_GROUP_REF = re.compile(r"^gh-readonly-queue/.+/pr-(\d+)-[0-9a-f]+$", re.IGNORECASE) + + +class GitHubApiError(RuntimeError): + """Report a failed GitHub REST API request.""" + + def __init__(self, status: int, message: str) -> None: + super().__init__(message) + self.status = status + + +class GitHubClient: + """Minimal GitHub Actions REST API client.""" + + def __init__(self, api_url: str, repository: str, token: str) -> None: + owner, name = repository.split("/", maxsplit=1) + self._repository_path = f"repos/{urllib.parse.quote(owner)}/{urllib.parse.quote(name)}" + self._api_url = api_url.rstrip("/") + self._headers = { + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + } + + def _request(self, path: str, method: str = "GET") -> dict[str, object] | None: + for attempt in range(5): + request = urllib.request.Request( + f"{self._api_url}/{self._repository_path}/{path}", + data=b"" if method == "POST" else None, + headers=self._headers, + method=method, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + payload = response.read() + return json.loads(payload) if payload else None + except urllib.error.HTTPError as error: + details = error.read().decode("utf-8", errors="replace") + if error.code != 429 and error.code < 500: + raise GitHubApiError( + error.code, f"GitHub API {method} {path} failed: {details}" + ) from error + if attempt == 4: + raise GitHubApiError( + error.code, f"GitHub API {method} {path} failed after retries: {details}" + ) from error + delay = min(int(error.headers.get("Retry-After", 2**attempt)), 30) + except (TimeoutError, urllib.error.URLError) as error: + if attempt == 4: + raise RuntimeError(f"GitHub API {method} {path} failed after retries: {error}") from error + delay = 2**attempt + + print(f"GitHub API {method} {path} failed transiently; retrying in {delay}s") + time.sleep(delay) + + raise AssertionError("unreachable") + + def _list_runs(self, endpoint: str, parameters: dict[str, str]) -> list[dict[str, object]]: + runs: list[dict[str, object]] = [] + page = 1 + while True: + query = urllib.parse.urlencode({**parameters, "per_page": "100", "page": str(page)}) + response = self._request(f"{endpoint}?{query}") + if not isinstance(response, dict) or not isinstance(response.get("workflow_runs"), list): + raise RuntimeError(f"GitHub API returned an invalid workflow-runs response for {endpoint}") + + page_runs = response["workflow_runs"] + runs.extend(page_runs) + if len(page_runs) < 100: + return runs + page += 1 + + def list_trigger_workflow_runs( + self, workflow_id: int, head_branches: set[str] + ) -> list[dict[str, object]]: + """List merge-group runs that establish the ordering of active generations.""" + + runs: list[dict[str, object]] = [] + for head_branch in sorted(head_branches): + runs.extend( + self._list_runs( + f"actions/workflows/{workflow_id}/runs", + {"branch": head_branch, "event": "merge_group"}, + ) + ) + return runs + + def list_active_merge_group_runs(self) -> list[dict[str, object]]: + """List all non-terminal merge-group workflow runs in the repository.""" + + runs_by_id: dict[int, dict[str, object]] = {} + for status in _ACTIVE_STATUSES: + for run in self._list_runs("actions/runs", {"event": "merge_group", "status": status}): + runs_by_id[int(run["id"])] = run + return list(runs_by_id.values()) + + def cancel_run(self, run_id: int) -> None: + """Request cancellation of one workflow run.""" + + self._request(f"actions/runs/{run_id}/cancel", method="POST") + + +def _pr_number(head_branch: object) -> int | None: + match = _MERGE_GROUP_REF.fullmatch(str(head_branch)) + return int(match.group(1)) if match else None + + +def _run_order(run: dict[str, object]) -> tuple[str, int]: + return str(run["created_at"]), int(run["id"]) + + +def _select_superseded_runs( + active_runs: list[dict[str, object]], + trigger_workflow_runs: list[dict[str, object]], + current_run: dict[str, object], +) -> list[dict[str, object]]: + current_head = str(current_run["head_branch"]) + current_pr = _pr_number(current_head) + if current_pr is None: + raise ValueError(f"Unexpected merge-group head branch: {current_head}") + + generation_order: dict[str, tuple[str, int]] = {} + for run in trigger_workflow_runs: + head_branch = str(run.get("head_branch", "")) + if _pr_number(head_branch) != current_pr: + continue + run_order = _run_order(run) + previous_order = generation_order.get(head_branch) + if previous_order is None or run_order < previous_order: + generation_order[head_branch] = run_order + + current_order = _run_order(current_run) + generation_order[current_head] = min(generation_order.get(current_head, current_order), current_order) + if any(head != current_head and order > current_order for head, order in generation_order.items()): + print(f"A newer merge-group generation exists for PR #{current_pr}; this cleanup is stale") + return [] + + superseded = [] + for run in active_runs: + head_branch = str(run.get("head_branch", "")) + if head_branch == current_head or _pr_number(head_branch) != current_pr: + continue + + old_generation_order = generation_order.get(head_branch) + if old_generation_order is None: + print(f"Skipping run {run.get('id')} because its merge-group generation was not found") + continue + if old_generation_order < current_order: + superseded.append(run) + + return sorted(superseded, key=_run_order) + + +def main() -> None: + """Cancel active runs belonging to older merge-group generations of the triggering PR.""" + + current_run = { + "id": int(os.environ["TRIGGER_RUN_ID"]), + "head_branch": os.environ["TRIGGER_HEAD_BRANCH"], + "created_at": os.environ["TRIGGER_CREATED_AT"], + } + client = GitHubClient( + api_url=os.environ.get("GITHUB_API_URL", "https://api.github.com"), + repository=os.environ["GITHUB_REPOSITORY"], + token=os.environ["GITHUB_TOKEN"], + ) + + active_runs = client.list_active_merge_group_runs() + current_pr = _pr_number(current_run["head_branch"]) + if current_pr is None: + raise ValueError(f"Unexpected merge-group head branch: {current_run['head_branch']}") + active_heads = { + str(run.get("head_branch", "")) for run in active_runs if _pr_number(run.get("head_branch")) == current_pr + } + active_heads.add(str(current_run["head_branch"])) + trigger_runs = client.list_trigger_workflow_runs(int(os.environ["TRIGGER_WORKFLOW_ID"]), active_heads) + superseded_runs = _select_superseded_runs(active_runs, trigger_runs, current_run) + + canceled = 0 + for run in superseded_runs: + run_id = int(run["id"]) + try: + client.cancel_run(run_id) + except GitHubApiError as error: + if error.status == 409: + print(f"Run {run_id} became terminal before cancellation") + continue + raise + print(f"Canceled superseded {run.get('name', 'workflow')} run {run_id} ({run.get('head_branch')})") + canceled += 1 + + print(f"Canceled {canceled} superseded merge-group run(s)") + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/test_approve_test_queue.sh b/.github/scripts/test_approve_test_queue.sh index 9e091978651..605c8812a16 100755 --- a/.github/scripts/test_approve_test_queue.sh +++ b/.github/scripts/test_approve_test_queue.sh @@ -3,6 +3,13 @@ set -euo pipefail readonly WORKFLOW=${WORKFLOW:-.github/workflows/cicd-approve-test-queue.yml} readonly EXTERNAL_QUEUE=' - branch: all' +readonly WORKFLOW_RUN_EVENT=' workflow_run:' +readonly CICD_COMPLETION_TRIGGER=' workflows: ["CICD Megatron-LM"]' +readonly COMPLETED_ACTIVITY=' types: [completed]' +readonly CONCURRENCY_GROUP=' group: approve-test-queue' +readonly SERIALIZED_RUNS=' cancel-in-progress: false' +readonly INTERNAL_SERVICE_ACCOUNT=' INTERNAL_SERVICE_ACCOUNTS = {"svcnemo-autobot"}' +readonly INTERNAL_SERVICE_ACCOUNT_CHECK=' return login in INTERNAL_SERVICE_ACCOUNTS or any(' if [[ $(/usr/bin/grep -c -F "$EXTERNAL_QUEUE" "$WORKFLOW") -ne 1 ]]; then echo "Approve Test Queue must define exactly one global external worker" >&2 @@ -19,4 +26,26 @@ if [[ $(/usr/bin/grep -c -F 'if CONTRIBUTOR_TYPE == "external":' "$WORKFLOW") -n exit 1 fi -echo "Approve Test Queue uses one global external worker" +for expected in \ + "$WORKFLOW_RUN_EVENT" \ + "$CICD_COMPLETION_TRIGGER" \ + "$COMPLETED_ACTIVITY" \ + "$CONCURRENCY_GROUP" \ + "$SERIALIZED_RUNS"; do + if [[ $(/usr/bin/grep -c -F "$expected" "$WORKFLOW") -ne 1 ]]; then + echo "Approve Test Queue must wake on completed CICD runs without racing approvals: $expected" >&2 + exit 1 + fi +done + +if [[ $(/usr/bin/grep -c -F "$INTERNAL_SERVICE_ACCOUNT" "$WORKFLOW") -ne 1 ]]; then + echo "Approve Test Queue must classify svcnemo-autobot as internal without relying on the SSO export" >&2 + exit 1 +fi + +if [[ $(/usr/bin/grep -c -F "$INTERNAL_SERVICE_ACCOUNT_CHECK" "$WORKFLOW") -ne 1 ]]; then + echo "Approve Test Queue must apply its internal service-account allowlist" >&2 + exit 1 +fi + +echo "Approve Test Queue uses one global external worker, recognizes its service account, and wakes on completion" diff --git a/.github/scripts/test_cancel_superseded_merge_group_runs.py b/.github/scripts/test_cancel_superseded_merge_group_runs.py new file mode 100644 index 00000000000..678bedf9ae2 --- /dev/null +++ b/.github/scripts/test_cancel_superseded_merge_group_runs.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import unittest +from pathlib import Path + +from cancel_superseded_merge_group_runs import _pr_number, _select_superseded_runs + + +def _run(run_id: int, head_branch: str, created_at: str, name: str = "CICD Megatron-LM") -> dict[str, object]: + return { + "id": run_id, + "event": "merge_group", + "head_branch": head_branch, + "created_at": created_at, + "name": name, + } + + +class TestCancelSupersededMergeGroupRuns(unittest.TestCase): + def setUp(self) -> None: + self.old_head = "gh-readonly-queue/main/pr-123-a1d1234" + self.current_head = "gh-readonly-queue/main/pr-123-b2e5678" + self.current_run = _run(200, self.current_head, "2026-08-13T22:00:00Z") + self.trigger_runs = [ + self.current_run, + _run(100, self.old_head, "2026-08-13T21:00:00Z"), + ] + + def test_extracts_pr_number_only_from_merge_group_refs(self) -> None: + self.assertEqual(_pr_number(self.current_head), 123) + self.assertIsNone(_pr_number("pull-request/123")) + self.assertIsNone(_pr_number("feature/pr-123-not-a-merge-group")) + + def test_selects_all_active_runs_from_the_older_generation(self) -> None: + old_cicd = _run(100, self.old_head, "2026-08-13T21:00:00Z") + old_release = _run(250, self.old_head, "2026-08-13T22:01:00Z", name="Release") + current_release = _run(201, self.current_head, "2026-08-13T22:00:01Z", name="Release") + other_pr = _run(50, "gh-readonly-queue/main/pr-456-a7e1234", "2026-08-13T20:00:00Z") + + selected = _select_superseded_runs( + [old_cicd, old_release, self.current_run, current_release, other_pr], + self.trigger_runs, + self.current_run, + ) + + self.assertEqual([run["id"] for run in selected], [100, 250]) + + def test_stale_cleanup_does_not_cancel_a_newer_generation(self) -> None: + newer_head = "gh-readonly-queue/main/pr-123-c3f9abc" + newer_run = _run(300, newer_head, "2026-08-13T23:00:00Z") + + selected = _select_superseded_runs( + [self.current_run, newer_run], + [newer_run, *self.trigger_runs], + self.current_run, + ) + + self.assertEqual(selected, []) + + def test_rerun_does_not_change_the_generation_order(self) -> None: + old_rerun = _run(400, self.old_head, "2026-08-13T23:00:00Z") + + selected = _select_superseded_runs( + [old_rerun], + [old_rerun, *self.trigger_runs], + self.current_run, + ) + + self.assertEqual([run["id"] for run in selected], [400]) + + def test_skips_an_active_generation_without_trusted_ordering_data(self) -> None: + unknown_head = "gh-readonly-queue/main/pr-123-d4a1111" + unknown_run = _run(50, unknown_head, "2026-08-13T20:00:00Z") + + selected = _select_superseded_runs([unknown_run], self.trigger_runs, self.current_run) + + self.assertEqual(selected, []) + + def test_workflow_runs_from_the_default_branch_with_minimal_write_permission(self) -> None: + workflow = Path(".github/workflows/cicd-cancel-superseded-merge-groups.yml").read_text() + + self.assertIn('workflows: ["CICD Megatron-LM"]', workflow) + self.assertIn("types: [requested]", workflow) + self.assertIn("github.event.workflow_run.event == 'merge_group'", workflow) + self.assertIn("actions: write", workflow) + self.assertIn("ref: ${{ github.event.repository.default_branch }}", workflow) + self.assertIn("persist-credentials: false", workflow) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/cicd-approve-test-queue.yml b/.github/workflows/cicd-approve-test-queue.yml index b4b52aab20d..8d33a001bd1 100644 --- a/.github/workflows/cicd-approve-test-queue.yml +++ b/.github/workflows/cicd-approve-test-queue.yml @@ -17,8 +17,16 @@ name: Approve Test Queue on: schedule: - cron: "*/5 * * * *" # Runs every 5 minutes + # Scheduled events can be delayed, so refill released queue slots immediately. + workflow_run: + workflows: ["CICD Megatron-LM"] + types: [completed] workflow_dispatch: # Allows manual triggering +concurrency: + group: approve-test-queue + cancel-in-progress: false + jobs: approve-queue: runs-on: ubuntu-latest @@ -79,6 +87,7 @@ jobs: GITHUB_TOKEN = os.environ["GITHUB_TOKEN"] REPO = os.environ["GITHUB_REPOSITORY"] CONTRIBUTOR_TYPE = os.environ["CONTRIBUTOR_TYPE"] + INTERNAL_SERVICE_ACCOUNTS = {"svcnemo-autobot"} if CONTRIBUTOR_TYPE == "external": # Global limit across all branches — no division needed since we count globally. MAX_CONCURRENCY = int(os.environ["MAX_CONCURRENCY_EXTERNAL"]) @@ -131,10 +140,12 @@ jobs: return None def is_internal_contributor(pr_info): - """Return True if the PR author is a member of NVIDIA or NVIDIA-NeMo org (is_org_member).""" + """Return True for internal service accounts or NVIDIA organization members.""" login = pr_info.get("user", {}).get("login", "") org_roles = sso_users.get(login, {}).get("org_roles", []) - return any(role in ("NVIDIA:Member", "NVIDIA-NeMo:Member") for role in org_roles) + return login in INTERNAL_SERVICE_ACCOUNTS or any( + role in ("NVIDIA:Member", "NVIDIA-NeMo:Member") for role in org_roles + ) def get_pr_base_branch(workflow_run): """ diff --git a/.github/workflows/cicd-cancel-superseded-merge-groups.yml b/.github/workflows/cicd-cancel-superseded-merge-groups.yml new file mode 100644 index 00000000000..2aeebcaf0a4 --- /dev/null +++ b/.github/workflows/cicd-cancel-superseded-merge-groups.yml @@ -0,0 +1,35 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +name: Cancel Superseded Merge-Group Runs + +on: + workflow_run: + workflows: ["CICD Megatron-LM"] + types: [requested] + +permissions: + actions: write + contents: read + +jobs: + cancel-superseded-runs: + if: >- + github.repository == 'NVIDIA/Megatron-LM' && + github.event.workflow_run.event == 'merge_group' + runs-on: ubuntu-latest + steps: + # workflow_run is privileged, so execute only code from the trusted default branch. + - name: Checkout trusted repository code + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + ref: ${{ github.event.repository.default_branch }} + persist-credentials: false + + - name: Cancel superseded merge-group runs for this PR + env: + GITHUB_TOKEN: ${{ github.token }} + TRIGGER_CREATED_AT: ${{ github.event.workflow_run.created_at }} + TRIGGER_HEAD_BRANCH: ${{ github.event.workflow_run.head_branch }} + TRIGGER_RUN_ID: ${{ github.event.workflow_run.id }} + TRIGGER_WORKFLOW_ID: ${{ github.event.workflow_run.workflow_id }} + run: python3 .github/scripts/cancel_superseded_merge_group_runs.py diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 9312f3d4eb5..112425c9879 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -360,6 +360,9 @@ jobs: - name: Test queue approval topology run: .github/scripts/test_approve_test_queue.sh + - name: Test superseded merge-group cleanup + run: python3 .github/scripts/test_cancel_superseded_merge_group_runs.py + - name: Get PR info id: get-pr-info if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' From ad3ecd9861b2b2d8428c6843dd062126c012cd7c Mon Sep 17 00:00:00 2001 From: nvcsathe Date: Fri, 14 Aug 2026 02:42:04 +0000 Subject: [PATCH 033/101] Streaming tool call deltas support (#6226) Signed-off-by: Chaitra Sathe Co-authored-by: Shanmugam Ramasamy --- .../core/inference/engines/dynamic_engine.py | 10 + .../endpoints/chat_completions.py | 37 +- .../endpoints/completions.py | 2 + .../incremental_detokenizer.py | 5 +- .../openai_streaming.py | 335 ++++++++++++++++-- .../text/parsers/qwen3_coder_tool_parser.py | 1 + .../inference/engines/test_dynamic_engine.py | 17 +- .../inference/test_openai_streaming.py | 289 ++++++++++++++- 8 files changed, 655 insertions(+), 41 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 79a84981f04..22fe74525e9 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -2225,6 +2225,16 @@ def _try_send_streaming_partials(self) -> None: partial["new_log_probs"] = list( (request.generated_log_probs or [])[already:emit_end] ) + partial["new_top_n_logprobs"] = list( + (getattr(request, "generated_top_n_logprobs", None) or [])[already:] + ) + if already == 0 and not request.sampling_params.skip_prompt_log_probs: + partial["prompt_log_probs"] = list( + getattr(request, "prompt_log_probs", None) or [] + ) + partial["prompt_top_n_logprobs"] = list( + getattr(request, "prompt_top_n_logprobs", None) or [] + ) partials.append(partial) emit_lengths[rid] = emit_end diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index 23a55d0902d..ce1550aa55d 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -13,7 +13,7 @@ from megatron.core.tokenizers.text.parsers import PARSER_MAPPING from ..incremental_detokenizer import HuggingFaceFastIncrementalDetokenizer -from ..openai_streaming import openai_stream +from ..openai_streaming import StreamingChatParser, openai_stream logger = logging.getLogger(__name__) @@ -642,6 +642,40 @@ async def chat_completions(): streams = [ client.add_request_streaming(prompt_tokens, sampling_params) for _ in range(n) ] + chat_parsers = None + if parsers: + marker_prefixes = ( + tuple( + marker + for parser_name in parsers + for marker in getattr(PARSER_MAPPING[parser_name], "streaming_markers", ()) + ) + if tools_requested + else () + ) + + def parse_streaming_text(text): + parsed_text, metadata = apply_parsers( + text, + tools, + parsers, + tools_requested, + chat_template_kwargs=chat_template_kwargs, + ) + metadata["tool_calls"] = _maybe_filter_parallel_tool_calls( + metadata.get("tool_calls", []), parallel_tool_calls + ) + return parsed_text, metadata + + is_named_tool_choice = isinstance(tool_choice, dict) and "function" in tool_choice + chat_parsers = [ + StreamingChatParser( + parse_streaming_text, + marker_prefixes=marker_prefixes, + named_tool_choice=is_named_tool_choice, + ) + for _ in range(n) + ] include_usage = bool((req.get("stream_options") or {}).get("include_usage", False)) response = Response( openai_stream( @@ -651,6 +685,7 @@ async def chat_completions(): chat=True, return_log_probs=return_log_probs, include_usage=include_usage, + chat_parsers=chat_parsers, ), content_type="text/event-stream", ) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py index d5795b65c68..fa92988a10a 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py @@ -157,6 +157,8 @@ async def completions(): chat=False, return_log_probs=return_log_probs, include_usage=include_usage, + echo_prompts=prompts_as_strings if echo else None, + prompt_token_ids=prompts_as_tokens if echo else None, ), content_type="text/event-stream", ) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/incremental_detokenizer.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/incremental_detokenizer.py index 47b50e29399..a407047195f 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/incremental_detokenizer.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/incremental_detokenizer.py @@ -41,7 +41,10 @@ def __init__(self, tokenizer: Any, prompt_token_ids: list[int]) -> None: ) self._native_tokenizer = huggingface_tokenizer._tokenizer - self._skip_special_tokens = not getattr(tokenizer_wrapper, "include_special_tokens", True) + # Final inference decoding always strips special tokens. Streaming must + # make the same choice or its concatenated deltas can differ from + # ``generated_text`` in the non-streaming response. + self._skip_special_tokens = True self._decode_stream = self._new_decode_stream(prompt_token_ids) self._text_fragments: list[str] = [] self._text_length = 0 diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/openai_streaming.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/openai_streaming.py index 1d3aaaea7e2..88acb8962ff 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/openai_streaming.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/openai_streaming.py @@ -4,24 +4,39 @@ import asyncio import json +import logging import time import uuid from megatron.core.inference.inference_request import unwrap_serialized_tensors +logger = logging.getLogger(__name__) -def _token_logprobs(tokenizer, token_ids, log_probs, chat, start_offset): + +def _top_logprob_entries(top_logprobs): + if not isinstance(top_logprobs, dict): + return [] + return [ + {"token": str(token), "logprob": logprob, "bytes": list(str(token).encode("utf-8"))} + for token, logprob in top_logprobs.items() + ] + + +def _token_logprobs(tokenizer, token_ids, log_probs, top_log_probs, chat, start_offset): entries = [] offsets = [] offset = start_offset for i, token_id in enumerate(token_ids): token = tokenizer.detokenize([token_id]) + token_top_logprobs = ( + top_log_probs[i] if top_log_probs is not None and i < len(top_log_probs) else None + ) entries.append( { "token": token, "logprob": log_probs[i] if i < len(log_probs) else None, "bytes": list(token.encode("utf-8")), - "top_logprobs": [], + "top_logprobs": _top_logprob_entries(token_top_logprobs), } ) offsets.append(offset) @@ -31,7 +46,33 @@ def _token_logprobs(tokenizer, token_ids, log_probs, chat, start_offset): return { "tokens": [entry["token"] for entry in entries], "token_logprobs": [entry["logprob"] for entry in entries], - "top_logprobs": [None] * len(entries), + "top_logprobs": [ + top_log_probs[i] if top_log_probs is not None and i < len(top_log_probs) else None + for i in range(len(entries)) + ], + "text_offset": offsets, + } + + +def _prompt_logprobs(tokenizer, token_ids, log_probs, top_log_probs): + token_ids = list(token_ids or []) + tokens = [tokenizer.detokenize([token_id]) for token_id in token_ids] + token_logprobs = [None] + list(log_probs or []) + token_logprobs = token_logprobs[: len(tokens)] + token_logprobs.extend([None] * (len(tokens) - len(token_logprobs))) + prompt_top_logprobs = [None] + list(top_log_probs or []) + prompt_top_logprobs = prompt_top_logprobs[: len(tokens)] + prompt_top_logprobs.extend([None] * (len(tokens) - len(prompt_top_logprobs))) + + offsets = [] + offset = 0 + for token in tokens: + offsets.append(offset) + offset += len(token) + return { + "tokens": tokens, + "token_logprobs": token_logprobs, + "top_logprobs": prompt_top_logprobs, "text_offset": offsets, } @@ -42,6 +83,141 @@ def _finish_reason(result): return "length" if requested is not None and generated >= requested else "stop" +def _safe_content_prefix(content, markers): + """Hold text that may still become a parser control marker.""" + safe_end = len(content) + for marker in markers: + marker_index = content.find(marker) + if marker_index >= 0: + safe_end = min(safe_end, marker_index) + for prefix_length in range(1, min(len(marker), len(content)) + 1): + if content.endswith(marker[:prefix_length]): + safe_end = min(safe_end, len(content) - prefix_length) + return content[:safe_end] + + +class StreamingChatParser: + """Convert accumulated parser output into stable OpenAI chat deltas.""" + + def __init__(self, parse, *, marker_prefixes=(), named_tool_choice=False): + self._parse = parse + self._marker_prefixes = tuple(marker_prefixes) + self._named_tool_choice = named_tool_choice + # Retain emitted content/reasoning prefixes and stable per-tool IDs and emission flags. + self._content_sent = "" + self._reasoning_sent = "" + self._tool_arguments_sent = [] + self._tool_ids = [] + self._tool_names_sent = [] + # Track whether a tool delta was emitted so the final finish reason can reflect it. + self.tools_streamed = False + + def _append_state_for_tool(self): + self._tool_arguments_sent.append(False) + self._tool_ids.append(f"call_{uuid.uuid4().hex[:24]}") + self._tool_names_sent.append(False) + + @staticmethod + def _function(call): + function = call.get("function", {}) if isinstance(call, dict) else {} + return function if isinstance(function, dict) else {} + + def parse(self, text, *, finished=False): + """Parse accumulated text and return zero or more structured deltas.""" + try: + content, metadata = self._parse(text) + except Exception: + logger.exception("Failed to parse a streaming chat delta.") + return [] + + content = content or "" + metadata = metadata or {} + tool_calls = metadata.get("tool_calls") or [] + reasoning = metadata.get("reasoning") or "" + deltas = [] + + if reasoning.startswith(self._reasoning_sent): + reasoning_delta = reasoning[len(self._reasoning_sent) :] + if reasoning_delta: + deltas.append({"reasoning_content": reasoning_delta}) + self._reasoning_sent = reasoning + + # Buffer partial control-marker prefixes so parser syntax is not emitted as content. + safe_content = ( + content + if finished or tool_calls + else _safe_content_prefix(content, self._marker_prefixes) + ) + if safe_content.startswith(self._content_sent): + content_delta = safe_content[len(self._content_sent) :] + if content_delta: + deltas.append({"content": content_delta}) + self._content_sent = safe_content + + for index, call in enumerate(tool_calls): + while len(self._tool_ids) <= index: + self._append_state_for_tool() + + function = self._function(call) + function_name = function.get("name") + current_arguments = function.get("arguments", "") + if not isinstance(current_arguments, str): + current_arguments = json.dumps(current_arguments, ensure_ascii=False) + + if function_name and not self._tool_names_sent[index]: + deltas.append( + { + "tool_calls": [ + { + "index": index, + "id": self._tool_ids[index], + "type": "function", + "function": {"name": str(function_name)}, + } + ] + } + ) + self._tool_names_sent[index] = True + self.tools_streamed = True + + if self._tool_names_sent[index]: + call_is_complete = finished or index + 1 < len(tool_calls) + # Re-serialized arguments are not append-stable while XML parameters arrive. + # Emit the complete JSON object once instead of streaming corrupt prefixes. + if call_is_complete and not self._tool_arguments_sent[index]: + deltas.append( + { + "tool_calls": [ + {"index": index, "function": {"arguments": current_arguments}} + ] + } + ) + self._tool_arguments_sent[index] = True + + return deltas + + def finish_reason(self, result): + """Return the OpenAI finish reason for this parsed choice.""" + if self.tools_streamed and not self._named_tool_choice: + return "tool_calls" + return _finish_reason(result) + + +def _status_name(record): + status = record.get("status") if isinstance(record, dict) else None + return getattr(status, "name", str(status)).upper() + + +def _failure_message(record): + events = record.get("events") or [] + error_events = [ + event for event in events if event.get("type") in ("ERROR_NONTRANSIENT", "ERROR_TRANSIENT") + ] + if error_events: + return str(error_events[-1].get("payload", "Unknown error")) + return "Unknown inference error" + + async def openai_stream( streams, tokenizer, @@ -50,17 +226,37 @@ async def openai_stream( chat, return_log_probs=False, include_usage=False, + chat_parsers=None, + echo_prompts=None, + prompt_token_ids=None, ): """Yield SSE records for one or more inference streams.""" if len(streams) != len(incremental_detokenizers): raise ValueError("Each inference stream must have an incremental detokenizer.") + if chat_parsers is not None and len(streams) != len(chat_parsers): + raise ValueError("Each inference stream must have a streaming chat parser.") + if echo_prompts is not None and len(streams) != len(echo_prompts): + raise ValueError("Each inference stream must have an echo prompt.") + if prompt_token_ids is not None and len(streams) != len(prompt_token_ids): + raise ValueError("Each inference stream must have prompt token IDs.") response_id = f"chatcmpl-{uuid.uuid4().hex}" if chat else str(uuid.uuid4()) created = int(time.time()) queue = asyncio.Queue() states = [ - dict(tokens=[], log_probs=[], detokenizer=detokenizer, final=None) - for detokenizer in incremental_detokenizers + dict( + tokens=[], + log_probs=[], + top_log_probs=[], + detokenizer=detokenizer, + final=None, + role_sent=False, + echo_sent=False, + parser=chat_parsers[index] if chat_parsers is not None else None, + echo_prompt=echo_prompts[index] if echo_prompts is not None else None, + prompt_token_ids=prompt_token_ids[index] if prompt_token_ids is not None else [], + ) + for index, detokenizer in enumerate(incremental_detokenizers) ] async def pump(index, stream): @@ -87,58 +283,122 @@ def sse(choices, usage=None): return f"data: {json.dumps(payload)}\n\n" try: - if chat: - for index in range(len(streams)): - yield sse( - [ - { - "index": index, - "delta": {"role": "assistant", "content": ""}, - "logprobs": None, - "finish_reason": None, - } - ] - ) - remaining = len(streams) while remaining: index, item, error = await queue.get() if error is not None: yield f"data: {json.dumps({'error': {'message': str(error)}})}\n\n" - continue + return if item is None: remaining -= 1 continue state = states[index] + is_final = "final" in item if "partial" in item: partial = item["partial"] new_tokens = partial.get("new_tokens") or [] new_log_probs = partial.get("new_log_probs") or [] + new_top_log_probs = partial.get("new_top_n_logprobs") or [] + prompt_log_probs = partial.get("prompt_log_probs") or [] + prompt_top_log_probs = partial.get("prompt_top_n_logprobs") or [] else: - result = unwrap_serialized_tensors(item["final"]) + final_record = item["final"] + if _status_name(final_record) == "FAILED": + error_payload = {"error": {"message": _failure_message(final_record)}} + yield f"data: {json.dumps(error_payload)}\n\n" + return + result = unwrap_serialized_tensors(final_record) state["final"] = result already = len(state["tokens"]) new_tokens = (result.get("generated_tokens") or [])[already:] new_log_probs = (result.get("generated_log_probs") or [])[already:] + new_top_log_probs = (result.get("generated_top_n_logprobs") or [])[already:] + prompt_log_probs = result.get("prompt_log_probs") or [] + prompt_top_log_probs = result.get("prompt_top_n_logprobs") or [] - if not new_tokens: - continue - state["tokens"].extend(new_tokens) - state["log_probs"].extend(new_log_probs) - start_offset = state["detokenizer"].text_length - delta = state["detokenizer"].update(new_tokens) - choice = { - "index": index, - "logprobs": ( - _token_logprobs(tokenizer, new_tokens, new_log_probs, chat, start_offset) + if chat and not state["role_sent"]: + yield sse( + [ + { + "index": index, + "delta": {"role": "assistant", "content": ""}, + "logprobs": None, + "finish_reason": None, + } + ] + ) + state["role_sent"] = True + + if state["echo_prompt"] is not None and not state["echo_sent"]: + yield sse( + [ + { + "index": index, + "text": state["echo_prompt"], + "logprobs": ( + _prompt_logprobs( + tokenizer, + state["prompt_token_ids"], + prompt_log_probs, + prompt_top_log_probs, + ) + if return_log_probs + else None + ), + "finish_reason": None, + } + ] + ) + state["echo_sent"] = True + + if new_tokens: + state["tokens"].extend(new_tokens) + state["log_probs"].extend(new_log_probs) + state["top_log_probs"].extend(new_top_log_probs) + start_offset = state["detokenizer"].text_length + len(state["echo_prompt"] or "") + delta = state["detokenizer"].update(new_tokens) + logprobs = ( + _token_logprobs( + tokenizer, new_tokens, new_log_probs, new_top_log_probs, chat, start_offset + ) if return_log_probs else None - ), - "finish_reason": None, - } - choice["delta" if chat else "text"] = {"content": delta} if chat else delta - yield sse([choice]) + ) + else: + delta = "" + logprobs = None + + parser = state["parser"] + if parser is not None: + parsed_deltas = parser.parse(state["detokenizer"].text, finished=is_final) + logprobs_attached = False + for parsed_delta in parsed_deltas: + parsed_logprobs = None + if ( + not logprobs_attached + and logprobs is not None + and ( + parsed_delta.get("content") is not None + or parsed_delta.get("reasoning_content") is not None + ) + ): + parsed_logprobs = logprobs + logprobs_attached = True + yield sse( + [ + { + "index": index, + "delta": parsed_delta, + "logprobs": parsed_logprobs, + "finish_reason": None, + } + ] + ) + elif new_tokens: + choice = {"index": index, "logprobs": logprobs, "finish_reason": None} + choice["delta" if chat else "text"] = {"content": delta} if chat else delta + yield sse([choice]) prompt_tokens = completion_tokens = cached_token_count = 0 for index, state in enumerate(states): @@ -149,10 +409,13 @@ def sse(choices, usage=None): prompt_tokens = max(prompt_tokens, prompt_len) completion_tokens += len(result.get("generated_tokens") or []) cached_token_count = max(cached_token_count, result.get("num_cached_tokens", 0)) + parser = state["parser"] choice = { "index": index, "logprobs": None, - "finish_reason": _finish_reason(result), + "finish_reason": ( + parser.finish_reason(result) if parser is not None else _finish_reason(result) + ), "generation_token_ids": list(state["tokens"]), "generation_log_probs": list(state["log_probs"]), "generated_text": state["detokenizer"].text, diff --git a/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py b/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py index ace63a4b772..328d394f8c6 100644 --- a/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py +++ b/megatron/core/tokenizers/text/parsers/qwen3_coder_tool_parser.py @@ -267,6 +267,7 @@ class Qwen3CoderToolParser(BaseParser): """Parser for Qwen3 Coder style tool calls.""" implicit_reasoning_end_markers = ("",) + streaming_markers = ("", " tuple[str, dict[str, list[dict]]]: diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 0806cddf974..461fa762583 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -13,6 +13,7 @@ from typing import Dict, List, Optional, Tuple from unittest import mock +import msgpack import pytest import torch from tqdm import tqdm @@ -757,8 +758,13 @@ def test_streaming_partials_are_sent(): engine._partial_emit_lengths = {} request = types.SimpleNamespace( generated_tokens=[11, 12, 13], - generated_log_probs=None, - sampling_params=types.SimpleNamespace(streaming=True, return_log_probs=False), + generated_log_probs=[-0.1, -0.2, -0.3], + generated_top_n_logprobs=[{"eleven": -0.01}, {"twelve": -0.02}, {"thirteen": -0.03}], + prompt_log_probs=[-0.4], + prompt_top_n_logprobs=[{"prompt": -0.04}], + sampling_params=types.SimpleNamespace( + streaming=True, return_log_probs=True, skip_prompt_log_probs=False + ), ) engine.requests = {7: types.SimpleNamespace(record=[request])} engine.socket_for_receiving_requests = mock.Mock() @@ -767,6 +773,13 @@ def test_streaming_partials_are_sent(): engine.socket_for_receiving_requests.send.assert_called_once() assert engine._partial_emit_lengths == {7: 3} + payload = msgpack.unpackb( + engine.socket_for_receiving_requests.send.call_args.args[0], raw=False + ) + partial = payload[1][0] + assert partial["new_top_n_logprobs"] == request.generated_top_n_logprobs + assert partial["prompt_log_probs"] == request.prompt_log_probs + assert partial["prompt_top_n_logprobs"] == request.prompt_top_n_logprobs def test_streaming_partials_buffer_until_token_interval(): diff --git a/tests/unit_tests/inference/test_openai_streaming.py b/tests/unit_tests/inference/test_openai_streaming.py index c690ad5478e..827157d0719 100644 --- a/tests/unit_tests/inference/test_openai_streaming.py +++ b/tests/unit_tests/inference/test_openai_streaming.py @@ -9,12 +9,17 @@ from transformers import PreTrainedTokenizerFast from megatron.core.inference.async_stream import AsyncStream +from megatron.core.inference.text_generation_server.dynamic_text_gen_server.endpoints.chat_completions import ( + _sanitize_messages_for_template, +) from megatron.core.inference.text_generation_server.dynamic_text_gen_server.incremental_detokenizer import ( HuggingFaceFastIncrementalDetokenizer, ) from megatron.core.inference.text_generation_server.dynamic_text_gen_server.openai_streaming import ( + StreamingChatParser, openai_stream, ) +from megatron.core.tokenizers.text.parsers.qwen3_coder_tool_parser import Qwen3CoderToolParser class _Tokenizer: @@ -54,7 +59,16 @@ def _make_byte_level_fast_tokenizer(): @pytest.mark.asyncio async def test_openai_stream_emits_delta_chunks_and_terminal_metadata(): stream = AsyncStream(request_id=1, cancel=lambda: None) - stream.put({"partial": {"request_id": 1, "new_tokens": [1, 2], "new_log_probs": [-0.1, -0.2]}}) + stream.put( + { + "partial": { + "request_id": 1, + "new_tokens": [1, 2], + "new_log_probs": [-0.1, -0.2], + "new_top_n_logprobs": [{"a": -0.01}, {"b": -0.02}], + } + } + ) # Token 3 models a token completed before the engine's final reply and # therefore absent from its last partial frame. stream.put( @@ -63,6 +77,7 @@ async def test_openai_stream_emits_delta_chunks_and_terminal_metadata(): "prompt_tokens": [9, 9], "generated_tokens": [1, 2, 3], "generated_log_probs": [-0.1, -0.2, -0.3], + "generated_top_n_logprobs": [{"a": -0.01}, {"b": -0.02}, {"c": -0.03}], "num_cached_tokens": 2, "sampling_params": {"num_tokens_to_generate": 3}, } @@ -90,8 +105,10 @@ async def test_openai_stream_emits_delta_chunks_and_terminal_metadata(): assert "generated_text" not in first["choices"][0] assert "generated_length" not in first["choices"][0] assert first["choices"][0]["logprobs"]["token_logprobs"] == [-0.1, -0.2] + assert first["choices"][0]["logprobs"]["top_logprobs"] == [{"a": -0.01}, {"b": -0.02}] assert first["choices"][0]["logprobs"]["text_offset"] == [0, 1] assert reconciled["choices"][0]["text"] == "c" + assert reconciled["choices"][0]["logprobs"]["top_logprobs"] == [{"c": -0.03}] assert "generation_token_ids" not in reconciled["choices"][0] assert "generation_log_probs" not in reconciled["choices"][0] assert "generated_text" not in reconciled["choices"][0] @@ -124,6 +141,263 @@ def import_without_transformers(name, *args, **kwargs): HuggingFaceFastIncrementalDetokenizer(_Tokenizer(), []) +@pytest.mark.asyncio +async def test_openai_stream_echoes_completion_prompt_before_generated_text(): + stream = AsyncStream(request_id=1, cancel=lambda: None) + stream.put( + { + "partial": { + "request_id": 1, + "new_tokens": [1], + "new_log_probs": [-0.1], + "new_top_n_logprobs": [{"a": -0.01}], + "prompt_log_probs": [-0.4], + "prompt_top_n_logprobs": [{"z": -0.04}], + } + } + ) + stream.put( + { + "final": { + "prompt_tokens": [26, 26], + "generated_tokens": [1], + "generated_log_probs": [-0.1], + "generated_top_n_logprobs": [{"a": -0.01}], + "sampling_params": {"num_tokens_to_generate": 2}, + } + } + ) + stream.finish() + + records = [ + record + async for record in openai_stream( + [stream], + _Tokenizer(), + [_IncrementalDetokenizer()], + chat=False, + return_log_probs=True, + echo_prompts=["zz"], + prompt_token_ids=[[26, 26]], + ) + ] + payloads = [json.loads(record.removeprefix("data: ")) for record in records[:-1]] + + echoed, generated, finished = payloads + assert echoed["choices"][0]["text"] == "zz" + assert echoed["choices"][0]["logprobs"] == { + "tokens": ["z", "z"], + "token_logprobs": [None, -0.4], + "top_logprobs": [None, {"z": -0.04}], + "text_offset": [0, 1], + } + assert generated["choices"][0]["text"] == "a" + assert generated["choices"][0]["logprobs"]["text_offset"] == [2] + assert finished["choices"][0]["finish_reason"] == "stop" + assert records[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_openai_stream_surfaces_failed_final_without_success_terminator(): + stream = AsyncStream(request_id=1, cancel=lambda: None) + stream.put( + { + "final": { + "status": "FAILED", + "events": [{"type": "ERROR_NONTRANSIENT", "payload": "context length exceeded"}], + } + } + ) + stream.finish() + + records = [ + record + async for record in openai_stream( + [stream], _Tokenizer(), [_IncrementalDetokenizer()], chat=True + ) + ] + + assert records == ['data: {"error": {"message": "context length exceeded"}}\n\n'] + + +@pytest.mark.asyncio +async def test_openai_stream_preserves_chat_top_logprobs_with_parser(): + stream = AsyncStream(request_id=1, cancel=lambda: None) + stream.put( + { + "partial": { + "request_id": 1, + "new_tokens": [1], + "new_log_probs": [-0.1], + "new_top_n_logprobs": [{"a": -0.01}], + } + } + ) + stream.put( + { + "final": { + "prompt_tokens": [9], + "generated_tokens": [1], + "generated_log_probs": [-0.1], + "generated_top_n_logprobs": [{"a": -0.01}], + "sampling_params": {"num_tokens_to_generate": 2}, + } + } + ) + stream.finish() + parser = StreamingChatParser(lambda text: (text, {})) + + records = [ + record + async for record in openai_stream( + [stream], + _Tokenizer(), + [_IncrementalDetokenizer()], + chat=True, + return_log_probs=True, + chat_parsers=[parser], + ) + ] + payloads = [json.loads(record.removeprefix("data: ")) for record in records[:-1]] + + role, content, finished = payloads + assert role["choices"][0]["delta"] == {"role": "assistant", "content": ""} + assert content["choices"][0]["delta"] == {"content": "a"} + assert content["choices"][0]["logprobs"]["content"][0]["top_logprobs"] == [ + {"token": "a", "logprob": -0.01, "bytes": [97]} + ] + assert finished["choices"][0]["finish_reason"] == "stop" + assert records[-1] == "data: [DONE]\n\n" + + +def test_streaming_chat_parser_emits_structured_stable_tool_call_deltas(): + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ] + + def parse(text): + return Qwen3CoderToolParser.parse(text, tools=tools) + + parser = StreamingChatParser(parse, marker_prefixes=Qwen3CoderToolParser.streaming_markers) + model_output = ( + "Checking. " + "\nSF\n" + ) + deltas = [] + for end in range(1, len(model_output) + 1): + deltas.extend(parser.parse(model_output[:end])) + deltas.extend(parser.parse(model_output, finished=True)) + + assert "".join(delta.get("content", "") for delta in deltas) == "Checking. " + tool_deltas = [tool_call for delta in deltas for tool_call in delta.get("tool_calls", [])] + name_deltas = [ + tool_call + for tool_call in tool_deltas + if tool_call.get("function", {}).get("name") is not None + ] + assert len(name_deltas) == 1 + assert name_deltas[0]["index"] == 0 + assert name_deltas[0]["type"] == "function" + assert name_deltas[0]["id"].startswith("call_") + assert name_deltas[0]["function"] == {"name": "get_weather"} + + argument_text = "".join( + tool_call.get("function", {}).get("arguments", "") for tool_call in tool_deltas + ) + assert json.loads(argument_text) == {"city": "SF"} + assert all(tool_call["index"] == 0 for tool_call in tool_deltas) + assert ( + parser.finish_reason( + {"generated_tokens": [1], "sampling_params": {"num_tokens_to_generate": 2}} + ) + == "tool_calls" + ) + + +def test_streaming_chat_parser_handles_single_multi_turn_tool_call_request(): + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ] + request_payload = { + "stream": True, + "tools": tools, + "messages": [ + {"role": "user", "content": "What is the weather in SF?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_previous", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_previous", "content": "62 F"}, + {"role": "user", "content": "What about NYC?"}, + ], + } + + sanitized_messages = _sanitize_messages_for_template(request_payload["messages"]) + + assert sanitized_messages[1]["tool_calls"][0]["function"]["arguments"] == {"city": "SF"} + assert sanitized_messages[2] == { + "role": "tool", + "tool_call_id": "call_previous", + "content": "62 F", + } + assert sanitized_messages[3] == {"role": "user", "content": "What about NYC?"} + assert ( + request_payload["messages"][1]["tool_calls"][0]["function"]["arguments"] == '{"city": "SF"}' + ) + + parser = StreamingChatParser( + lambda text: Qwen3CoderToolParser.parse(text, tools=tools), + marker_prefixes=Qwen3CoderToolParser.streaming_markers, + ) + model_output = ( + "" + "\nNYC\n" + ) + deltas = [] + for end in range(1, len(model_output) + 1): + deltas.extend(parser.parse(model_output[:end])) + deltas.extend(parser.parse(model_output, finished=True)) + + tool_deltas = [tool_call for delta in deltas for tool_call in delta.get("tool_calls", [])] + assert ( + sum(tool_call.get("function", {}).get("name") == "get_weather" for tool_call in tool_deltas) + == 1 + ) + tool_call_ids = [tool_call["id"] for tool_call in tool_deltas if "id" in tool_call] + assert len(tool_call_ids) == 1 + assert tool_call_ids[0].startswith("call_") + assert all(tool_call["index"] == 0 for tool_call in tool_deltas) + argument_text = "".join( + tool_call.get("function", {}).get("arguments", "") for tool_call in tool_deltas + ) + assert json.loads(argument_text) == {"city": "NYC"} + assert ( + parser.finish_reason( + {"generated_tokens": [1], "sampling_params": {"num_tokens_to_generate": 2}} + ) + == "tool_calls" + ) + + def test_huggingface_fast_incremental_detokenizer_preserves_utf8_boundaries(): tokenizer = _make_byte_level_fast_tokenizer() huggingface_tokenizer = tokenizer._tokenizer.tokenizer @@ -154,6 +428,19 @@ def test_huggingface_fast_incremental_detokenizer_uses_prompt_context(): assert streamed_text == full_text[len(prompt_text) :] +def test_huggingface_fast_incremental_detokenizer_skips_special_tokens(): + tokenizer = _make_byte_level_fast_tokenizer() + huggingface_tokenizer = tokenizer._tokenizer.tokenizer + huggingface_tokenizer.add_special_tokens({"additional_special_tokens": [""]}) + token_ids = huggingface_tokenizer.encode("ab", add_special_tokens=False) + detokenizer = HuggingFaceFastIncrementalDetokenizer(tokenizer, []) + + streamed_text = "".join(detokenizer.update([token_id]) for token_id in token_ids) + + assert streamed_text == huggingface_tokenizer.decode(token_ids, skip_special_tokens=True) + assert streamed_text == "ab" + + def test_incremental_detokenizer_rejects_unsupported_tokenizer(): with pytest.raises( ValueError, match="Streaming is currently supported only for Hugging Face fast tokenizers" From 6c344290e54aa58a8be26d6af5ca2d46a3b00948 Mon Sep 17 00:00:00 2001 From: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:05:38 +0000 Subject: [PATCH 034/101] save tokenizer assets inside torch_dist checkpoint (#6280) Signed-off-by: dimapihtar Signed-off-by: Dmytro Pykhtar --- megatron/training/arguments.py | 2 +- megatron/training/checkpointing.py | 153 +++++++- megatron/training/training.py | 1 - .../models/test_gpt_hybrid_interop.py | 1 + .../test_global_metadata_reuse.py | 2 + .../test_layer_wise_optimizer.py | 2 + .../dist_checkpointing/test_local.py | 2 + .../dist_checkpointing/test_nonpersistent.py | 2 + .../dist_checkpointing/test_optimizer.py | 2 + .../test_pipeline_parallel_layout.py | 1 + .../dist_checkpointing/test_replication.py | 1 + .../test_save_tokenizer_assets.py | 370 ++++++++++++++++++ .../pipeline_parallel/test_pipeline_layout.py | 1 + tests/unit_tests/test_checkpointing.py | 3 + ...test_moe_single_grouped_weight_numerics.py | 1 + .../unit_tests/transformer/test_attention.py | 1 + .../test_multi_latent_attention.py | 1 + .../test_multi_token_prediction.py | 2 + 18 files changed, 545 insertions(+), 3 deletions(-) create mode 100644 tests/unit_tests/dist_checkpointing/test_save_tokenizer_assets.py diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 3497834ef93..e08a06ec150 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2869,7 +2869,7 @@ def _add_learning_rate_args(parser): def _add_checkpointing_args(parser): from megatron.training.config import CheckpointConfig - ckpt_factory = ArgumentGroupFactory(CheckpointConfig, exclude=["most_recent_k", "save_tokenizer_assets", "save_optim", "save_rng", "load_optim", "load_rng"]) + ckpt_factory = ArgumentGroupFactory(CheckpointConfig, exclude=["most_recent_k", "save_optim", "save_rng", "load_optim", "load_rng"]) group = ckpt_factory.build_group(parser, "checkpointing") group.add_argument('--no-save-optim', action='store_true', default=None, diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index a81a6c2dfe1..e7c1b2ec02c 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -37,11 +37,16 @@ TorchDistSaveShardedStrategy, get_async_strategy, ) -from megatron.core.msc_utils import maybe_msc +from megatron.core.msc_utils import maybe_msc, MultiStorageClientFeature from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.optimizer import DistributedOptimizer from megatron.core.rerun_state_machine import get_rerun_state_machine +from megatron.core.tokenizers import MegatronTokenizer from megatron.core.utils import get_pg_rank, get_pg_size, unwrap_model +from megatron.core._rank_utils import safe_get_rank as get_rank_safe +from megatron.training.argument_utils import _default_config_from_args +from megatron.training.config import TokenizerConfig +from megatron.training.global_vars import get_tokenizer from ..core.dist_checkpointing.utils import _clean_metadata_for_serialization from . import ft_integration, wandb_utils @@ -1057,6 +1062,17 @@ def iter_finalize_fn(): f'gtp_remat {gtp_remat_rank}/{gtp_remat_size_to_print}, ' f'p {pipeline_mp_rank}/{pp_size_to_print} ]' ) + # Save tokenizer files for torch_dist checkpoints (if enabled) + if ( + args.save_tokenizer_assets + and args.ckpt_format == 'torch_dist' + and iteration > 0 + ): + checkpoint_name = get_checkpoint_name( + save_dir, iteration=iteration, return_base_dir=True + ) + config = _default_config_from_args(TokenizerConfig, args) + save_tokenizer_assets(get_tokenizer(), config, checkpoint_name) if args.log_progress and args.async_save: append_to_progress_log( args.save, f'Saved async checkpoint\tIteration: {iteration}', barrier=False @@ -1205,6 +1221,141 @@ def progress_finalize_fn(): ft_integration.on_checkpointing_end(is_async_finalization=False) +def save_tokenizer_assets( + tokenizer: MegatronTokenizer, + config: TokenizerConfig, + checkpoint_path: str, +) -> None: + """Save tokenizer files to the checkpoint directory. + + Always saves tokenizer files to ensure checkpoints are self-contained + and portable. Handles both HuggingFace tokenizers and file-based tokenizers. + Compatible with MultiStorageClient for cloud storage support. + + Args: + tokenizer: The tokenizer instance to save. + config: Tokenizers config. + checkpoint_path: The checkpoint directory path. + """ + if tokenizer is None: + return + + # Only rank 0 saves tokenizer files + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + if rank != 0: + return + + def resolve_path(path_str: str) -> str: + """Resolve relative paths to absolute paths.""" + if not path_str: + return path_str + path_obj = Path(path_str) + if path_obj.is_absolute(): + return path_str + # Resolve relative to current working directory + return str(path_obj.resolve()) + + try: + # Check if MultiStorageClient is enabled + if MultiStorageClientFeature .is_enabled(): + msc = MultiStorageClientFeature .import_package() + checkpoint_path_obj = msc.Path(checkpoint_path) + tokenizer_dir = checkpoint_path_obj / "tokenizer" + tokenizer_dir.mkdir(parents=True, exist_ok=True) + use_msc = True + else: + tokenizer_dir = os.path.join(checkpoint_path, "tokenizer") + os.makedirs(tokenizer_dir, exist_ok=True) + use_msc = False + + tokenizer_type = config.tokenizer_type + + # Handle HuggingFace and Multimodal tokenizers + if tokenizer_type in ("HuggingFaceTokenizer", "MultimodalTokenizer"): + if use_msc: + import tempfile + + with tempfile.TemporaryDirectory() as tmp_dir: + if hasattr(tokenizer, "save_pretrained"): + tokenizer.save_pretrained(tmp_dir) + elif hasattr(tokenizer, "_tokenizer") and hasattr(tokenizer._tokenizer, "save_pretrained"): + tokenizer._tokenizer.save_pretrained(tmp_dir) + else: + logger.debug(f"{tokenizer_type} does not support save_pretrained(), skipping tokenizer save") + return + + logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") + for filename in os.listdir(tmp_dir): + src_path = os.path.join(tmp_dir, filename) + if os.path.isfile(src_path): + dest_path = tokenizer_dir / filename + with open(src_path, "rb") as src_f: + with msc.open(str(dest_path), "wb") as dest_f: + dest_f.write(src_f.read()) + else: + logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") + if hasattr(tokenizer, "save_pretrained"): + tokenizer.save_pretrained(tokenizer_dir) + elif hasattr(tokenizer, "_tokenizer") and hasattr(tokenizer._tokenizer, "save_pretrained"): + tokenizer._tokenizer.save_pretrained(tokenizer_dir) + return + + # Handle file-based tokenizers - resolve all paths + files_to_copy = [] + + if tokenizer_type in ("BertWordPieceLowerCase", "BertWordPieceCase"): + if config.vocab_file: + resolved_path = resolve_path(config.vocab_file) + files_to_copy.append(("vocab_file", resolved_path, "vocab.txt")) + + elif tokenizer_type == "GPT2BPETokenizer": + if config.vocab_file: + resolved_path = resolve_path(config.vocab_file) + files_to_copy.append(("vocab_file", resolved_path, "vocab.json")) + if config.merge_file: + resolved_path = resolve_path(config.merge_file) + files_to_copy.append(("merge_file", resolved_path, "merges.txt")) + + elif tokenizer_type in ("SentencePieceTokenizer", "GPTSentencePieceTokenizer", "Llama2Tokenizer"): + if config.tokenizer_model: + resolved_path = resolve_path(config.tokenizer_model) + files_to_copy.append(("tokenizer_model", resolved_path, "tokenizer.model")) + + elif tokenizer_type == "TikTokenizer": + if config.tokenizer_model: + resolved_path = resolve_path(config.tokenizer_model) + files_to_copy.append(("tokenizer_model", resolved_path, "tokenizer.json")) + + elif tokenizer_type == "NullTokenizer": + logger.debug(f"{tokenizer_type} requires no file artifacts") + return + + # Copy the files + if files_to_copy: + logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") + for config_attr, source_path, dest_filename in files_to_copy: + if source_path and os.path.exists(source_path): + if use_msc: + dest_path = tokenizer_dir / dest_filename + with open(source_path, "rb") as src_f: + with msc.open(str(dest_path), "wb") as dest_f: + dest_f.write(src_f.read()) + logger.debug(f"Copied {config_attr}: {source_path} -> {dest_path}") + else: + dest_path = os.path.join(tokenizer_dir, dest_filename) + shutil.copy2(source_path, dest_path) + logger.debug(f"Copied {config_attr}: {source_path} -> {dest_path}") + else: + logger.debug(f"{config_attr} not found at resolved path: {source_path}") + + except Exception as e: + if get_rank_safe() == 0: + logger.error(f"Failed to save tokenizer files: {e}") + import traceback + + logger.error(traceback.format_exc()) + + @_disable_gc() def _async_delete_checkpoint_impl( save_path, diff --git a/megatron/training/training.py b/megatron/training/training.py index 14d753ff6cd..11a9b946525 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -175,7 +175,6 @@ get_signal_handler, get_tensorboard_writer, get_timers, - get_tokenizer, get_wandb_writer, ) from .theoretical_memory_usage import report_theoretical_memory diff --git a/tests/unit_tests/dist_checkpointing/models/test_gpt_hybrid_interop.py b/tests/unit_tests/dist_checkpointing/models/test_gpt_hybrid_interop.py index dae973f4a7b..99d887b104e 100644 --- a/tests/unit_tests/dist_checkpointing/models/test_gpt_hybrid_interop.py +++ b/tests/unit_tests/dist_checkpointing/models/test_gpt_hybrid_interop.py @@ -657,6 +657,7 @@ def _run_gpt_to_hybrid_optimizer_load( ) with TempNamedDir(tmp_path_dist_ckpt / 'gpt_hybrid_opt_interop') as ckpt_dir: mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False with mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args): # Build a GPT model + distributed optimizer whose Adam moments are # seeded to random values, then save a full checkpoint. diff --git a/tests/unit_tests/dist_checkpointing/test_global_metadata_reuse.py b/tests/unit_tests/dist_checkpointing/test_global_metadata_reuse.py index 642d43b9d59..de027070182 100644 --- a/tests/unit_tests/dist_checkpointing/test_global_metadata_reuse.py +++ b/tests/unit_tests/dist_checkpointing/test_global_metadata_reuse.py @@ -33,6 +33,7 @@ def test_global_metadata_reuse(self, tmp_path_dist_ckpt, tp, pp): opt_param_scheduler = None mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False with ( TempNamedDir( tmp_path_dist_ckpt / "test_global_metadata_reuse" @@ -105,6 +106,7 @@ def test_no_global_metadata_reuse_on_different_parallelism(self, tmp_path_dist_c opt_param_scheduler = None mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False with ( TempNamedDir( tmp_path_dist_ckpt / "test_global_metadata_reuse" diff --git a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py index 42ef0a401ee..5c20be82e04 100644 --- a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py @@ -535,6 +535,7 @@ def test_optimizer_common_state_dict( with TempNamedDir(tmp_path_dist_ckpt / 'test_dp_sharding', sync=True) as ckpt_dir: mock_args = parse_args(ignore_unknown_args=True) mock_args.use_distributed_optimizer = False + mock_args.save_tokenizer_assets = False with mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args): # Initialize model and optimizer A if is_moe: @@ -628,6 +629,7 @@ def test_optimizer_common_state_dict_hybrid(self, tmp_path_dist_ckpt, tp, pp): tmp_path_dist_ckpt / 'test_optimizer_common_state_dict_hybrid', sync=True ) as ckpt_dir: mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False # Mirror the arg-parser's Muon path: ``use_distributed_optimizer`` # is flipped off and ``use_layer_wise_distributed_optimizer`` is # the surviving flag. diff --git a/tests/unit_tests/dist_checkpointing/test_local.py b/tests/unit_tests/dist_checkpointing/test_local.py index 8205bd8b36d..df27246fbf8 100644 --- a/tests/unit_tests/dist_checkpointing/test_local.py +++ b/tests/unit_tests/dist_checkpointing/test_local.py @@ -87,6 +87,7 @@ def test_sharded_tensors(self, tp, pp, use_torch_fsdp2): mock_args = parse_args(ignore_unknown_args=True) mock_args.no_save_optim = False mock_args.no_save_rng = True + mock_args.save_tokenizer_assets = False mock_args.use_torch_fsdp2 = use_torch_fsdp2 # Test save_local state_dict = generate_state_dict( @@ -270,6 +271,7 @@ def test_failed_save(self, caplog, tmp_path_dist_ckpt, tp, pp, use_ramdisk, asyn opt_param_scheduler = None mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False if use_ramdisk: tmp_path_dist_ckpt = Path("/dev/shm") diff --git a/tests/unit_tests/dist_checkpointing/test_nonpersistent.py b/tests/unit_tests/dist_checkpointing/test_nonpersistent.py index 0aab5573808..9db1977de11 100644 --- a/tests/unit_tests/dist_checkpointing/test_nonpersistent.py +++ b/tests/unit_tests/dist_checkpointing/test_nonpersistent.py @@ -38,6 +38,7 @@ def test_basic_save_load_scenarios(self, tmp_path_dist_ckpt, tp, pp): opt_param_scheduler = None mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False original_empty = torch.empty @@ -131,6 +132,7 @@ def test_basic_save_load_scenario(self, tmp_path_dist_ckpt, tp, pp): opt_param_scheduler = None mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False with ( TempNamedDir(tmp_path_dist_ckpt / "test_legacy") as legacy_ckpt_dir, mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args), diff --git a/tests/unit_tests/dist_checkpointing/test_optimizer.py b/tests/unit_tests/dist_checkpointing/test_optimizer.py index 7c319e0a14a..ce2d8aa13dd 100644 --- a/tests/unit_tests/dist_checkpointing/test_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_optimizer.py @@ -477,6 +477,7 @@ def test_optimizer_common_state_dict( with TempNamedDir(tmp_path_dist_ckpt / 'test_dp_sharding', sync=True) as ckpt_dir: mock_args = parse_args(ignore_unknown_args=True) mock_args.use_distributed_optimizer = True + mock_args.save_tokenizer_assets = False with mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args): # Initialize model and optimizer A if is_moe: @@ -552,6 +553,7 @@ def test_finetune_doesnt_load_optimizer( tmp_path_dist_ckpt / 'test_finetune_doesnt_load_optimizer', sync=True ) as ckpt_dir: mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False with mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp=src_tp_pp[0], pp=src_tp_pp[1]) init_checkpointing_mock_args(mock_args, ckpt_dir, False) diff --git a/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py b/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py index 4ed91aa2cb6..787482f2dbd 100644 --- a/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py +++ b/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py @@ -157,6 +157,7 @@ def create_args(): args.distrib_optim_fully_reshardable_mem_efficient = False args.phase_transition_iterations = None args.verify_integrity = False + args.save_tokenizer_assets = False yield args diff --git a/tests/unit_tests/dist_checkpointing/test_replication.py b/tests/unit_tests/dist_checkpointing/test_replication.py index 5ec9ad7ee53..eb0ff2b5ae1 100644 --- a/tests/unit_tests/dist_checkpointing/test_replication.py +++ b/tests/unit_tests/dist_checkpointing/test_replication.py @@ -94,6 +94,7 @@ def post_init(self, root_tmp_dir, tp, pp, async_save, algo, repl_groups): Utils.initialize_model_parallel(tp, pp) mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False with ( mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args), mock.patch('megatron.training.async_utils.get_args', new=lambda: mock_args), diff --git a/tests/unit_tests/dist_checkpointing/test_save_tokenizer_assets.py b/tests/unit_tests/dist_checkpointing/test_save_tokenizer_assets.py new file mode 100644 index 00000000000..585b7453873 --- /dev/null +++ b/tests/unit_tests/dist_checkpointing/test_save_tokenizer_assets.py @@ -0,0 +1,370 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Unit tests for tokenizer file saving in checkpointing.""" + +import os +import shutil +import tempfile +from unittest.mock import Mock, mock_open, patch + +import pytest + +from megatron.training.checkpointing import save_tokenizer_assets + + +class TestSaveTokenizerAssets: + """Test tokenizer file saving functionality.""" + + @pytest.fixture + def checkpoint_path_fixture(self): + """Create a temporary directory for checkpoint.""" + with tempfile.TemporaryDirectory() as temp_dir: + yield temp_dir + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + def test_save_tokenizer_assets_none_tokenizer( + self, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test that function returns early when tokenizer is None.""" + mock_dist_init.return_value = False + + config = Mock() + config.tokenizer_type = "SentencePieceTokenizer" + + # Should not raise error + save_tokenizer_assets(None, config, checkpoint_path_fixture) + + # Should not create tokenizer directory + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + assert not os.path.exists(tokenizer_dir) + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + def test_save_tokenizer_assets_non_rank0( + self, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test that non-rank-0 processes don't save files.""" + mock_dist_init.return_value = True + mock_get_rank.return_value = 1 # Not rank 0 + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "SentencePieceTokenizer" + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Should not create tokenizer directory + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + assert not os.path.exists(tokenizer_dir) + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_sentencepiece( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test saving SentencePiece tokenizer files.""" + mock_dist_init.return_value = False + + # Create a fake tokenizer model file + fake_model_path = os.path.join(checkpoint_path_fixture, "tokenizer.model") + with open(fake_model_path, "wb") as f: + f.write(b"fake sentencepiece model") + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "SentencePieceTokenizer" + config.tokenizer_model = fake_model_path + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Check that tokenizer directory was created + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + assert os.path.exists(tokenizer_dir) + + # Check that model file was copied + copied_file = os.path.join(tokenizer_dir, "tokenizer.model") + assert os.path.exists(copied_file) + + # Verify content + with open(copied_file, "rb") as f: + content = f.read() + assert content == b"fake sentencepiece model" + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_gpt2bpe( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test saving GPT2BPE tokenizer files (vocab + merges).""" + mock_dist_init.return_value = False + + # Create fake vocab and merge files + vocab_path = os.path.join(checkpoint_path_fixture, "vocab.json") + merge_path = os.path.join(checkpoint_path_fixture, "merges.txt") + + with open(vocab_path, "w") as f: + f.write('{"test": 0}') + with open(merge_path, "w") as f: + f.write("t e\ne s\n") + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "GPT2BPETokenizer" + config.vocab_file = vocab_path + config.merge_file = merge_path + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Check files were copied + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + assert os.path.exists(os.path.join(tokenizer_dir, "vocab.json")) + assert os.path.exists(os.path.join(tokenizer_dir, "merges.txt")) + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_bert( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test saving BERT tokenizer files.""" + mock_dist_init.return_value = False + + # Create fake vocab file + vocab_path = os.path.join(checkpoint_path_fixture, "vocab.txt") + with open(vocab_path, "w") as f: + f.write("[PAD]\n[UNK]\n[CLS]\n[SEP]\n") + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "BertWordPieceLowerCase" + config.vocab_file = vocab_path + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Check file was copied + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + copied_vocab = os.path.join(tokenizer_dir, "vocab.txt") + assert os.path.exists(copied_vocab) + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_tiktoken( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test saving TikToken tokenizer files.""" + mock_dist_init.return_value = False + + # Create fake tiktoken vocab file + vocab_path = os.path.join(checkpoint_path_fixture, "tokenizer.json") + with open(vocab_path, "w") as f: + f.write('{"vocab": {}}') + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "TikTokenizer" + config.tokenizer_model = vocab_path + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Check file was copied + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + assert os.path.exists(os.path.join(tokenizer_dir, "tokenizer.json")) + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_null_tokenizer( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test that NullTokenizer doesn't save files.""" + mock_dist_init.return_value = False + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "NullTokenizer" + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Tokenizer directory should be created but empty + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + assert os.path.exists(tokenizer_dir) + assert len(os.listdir(tokenizer_dir)) == 0 + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_relative_path(self, mock_logger, mock_get_rank, mock_dist_init): + """Test that relative paths are resolved correctly.""" + mock_dist_init.return_value = False + + with tempfile.TemporaryDirectory() as temp_dir: + # Create file using relative path + original_cwd = os.getcwd() + try: + os.chdir(temp_dir) + relative_path = "tokenizer.model" + with open(relative_path, "wb") as f: + f.write(b"test data") + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "SentencePieceTokenizer" + config.tokenizer_model = relative_path # Relative path + + checkpoint_path = os.path.join(temp_dir, "checkpoint") + os.makedirs(checkpoint_path, exist_ok=True) + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path) + + # Check that file was copied + tokenizer_dir = os.path.join(checkpoint_path, "tokenizer") + assert os.path.exists(os.path.join(tokenizer_dir, "tokenizer.model")) + finally: + os.chdir(original_cwd) + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_huggingface_with_instance( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test saving HuggingFace tokenizer files using provided instance.""" + mock_dist_init.return_value = False + + # Mock tokenizer with save_pretrained method + # Use spec_set to make hasattr work correctly + mock_inner_tokenizer = Mock(spec=["save_pretrained"]) + mock_tokenizer = Mock(spec=["_tokenizer"]) + mock_tokenizer._tokenizer = mock_inner_tokenizer + + config = Mock() + config.tokenizer_type = "HuggingFaceTokenizer" + config.tokenizer_model = "bert-base-uncased" # HF model ID + + # Pass tokenizer instance - uses it directly + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Check that save_pretrained was called + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + mock_tokenizer._tokenizer.save_pretrained.assert_called_once_with(tokenizer_dir) + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_all_sentencepiece_variants( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test all SentencePiece-based tokenizer types.""" + mock_dist_init.return_value = False + + # Create a fake tokenizer model file + fake_model_path = os.path.join(checkpoint_path_fixture, "tokenizer.model") + with open(fake_model_path, "wb") as f: + f.write(b"fake model data") + + tokenizer_types = ["SentencePieceTokenizer", "GPTSentencePieceTokenizer"] + + for tokenizer_type in tokenizer_types: + # Clear previous tokenizer directory + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + if os.path.exists(tokenizer_dir): + shutil.rmtree(tokenizer_dir) + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = tokenizer_type + config.tokenizer_model = fake_model_path + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Verify file was copied + copied_file = os.path.join(tokenizer_dir, "tokenizer.model") + assert os.path.exists(copied_file), f"Failed for {tokenizer_type}" + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_missing_file( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test handling of missing tokenizer files.""" + mock_dist_init.return_value = False + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "SentencePieceTokenizer" + config.tokenizer_model = "/nonexistent/path/tokenizer.model" + + # Should not raise error + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Should log debug message about missing file + assert mock_logger.debug.called + + @patch("megatron.training.checkpointing.MultiStorageClientFeature") + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_with_msc( + self, mock_logger, mock_get_rank, mock_dist_init, mock_msc_feature + ): + """Test saving tokenizer files with MultiStorageClient enabled.""" + mock_dist_init.return_value = False + mock_msc_feature.is_enabled.return_value = True + + # Mock MSC + mock_msc = Mock() + mock_path = Mock() + mock_tokenizer_dir = Mock() + mock_tokenizer_dir.mkdir = Mock() + mock_path.__truediv__ = Mock(return_value=mock_tokenizer_dir) + mock_msc.Path = Mock(return_value=mock_path) + mock_msc.open = mock_open() + mock_msc_feature.import_package.return_value = mock_msc + + with tempfile.TemporaryDirectory() as temp_dir: + # Create a fake tokenizer file + fake_model_path = os.path.join(temp_dir, "tokenizer.model") + with open(fake_model_path, "wb") as f: + f.write(b"test data") + + mock_tokenizer = Mock() + config = Mock() + config.tokenizer_type = "SentencePieceTokenizer" + config.tokenizer_model = fake_model_path + + checkpoint_path = os.path.join(temp_dir, "checkpoint") + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path) + + # Verify MSC methods were called + mock_msc_feature.is_enabled.assert_called() + mock_msc_feature.import_package.assert_called() + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.get_rank") + @patch("megatron.training.checkpointing.logger") + def test_save_tokenizer_assets_huggingface_save_pretrained( + self, mock_logger, mock_get_rank, mock_dist_init, checkpoint_path_fixture + ): + """Test that HuggingFace tokenizer's save_pretrained is called.""" + mock_dist_init.return_value = False + + # Mock tokenizer with save_pretrained on the wrapped _tokenizer + # Use spec to make hasattr work correctly + mock_inner_tokenizer = Mock(spec=["save_pretrained"]) + mock_tokenizer = Mock(spec=["_tokenizer"]) + mock_tokenizer._tokenizer = mock_inner_tokenizer + + config = Mock() + config.tokenizer_type = "HuggingFaceTokenizer" + config.tokenizer_model = "bert-base-uncased" + + save_tokenizer_assets(mock_tokenizer, config, checkpoint_path_fixture) + + # Verify save_pretrained was called with tokenizer directory + tokenizer_dir = os.path.join(checkpoint_path_fixture, "tokenizer") + mock_tokenizer._tokenizer.save_pretrained.assert_called_once_with(tokenizer_dir) diff --git a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py index 1c998181b50..aeabfe00944 100644 --- a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py +++ b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py @@ -221,6 +221,7 @@ def test_forward_vpp(create_args, tmp_path_dist_ckpt, tp_pp_vpp, pp_layout, is_m args.num_attention_heads = 8 # Ckpt format args.ckpt_format = "torch_dist" + args.save_tokenizer_assets = False set_args(args) def set_tp_pp_vpp(tp, pp, vpp=None, pp_layout=None, destroy_first=True): diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index 15ab11458ae..d9ae2b0f07e 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -334,6 +334,7 @@ def test_save_checkpoint(init_model_parallel, create_args, tmp_path_dist_ckpt, c with TempNamedDir(tmp_path_dist_ckpt / "test_save_checkpoint", sync=True) as save_dir: args.save = save_dir + args.save_tokenizer_assets = False set_args(args) save_checkpoint( @@ -370,6 +371,7 @@ def test_load_checkpoint( with TempNamedDir(tmp_path_dist_ckpt / "test_load_checkpoint", sync=True) as ckpt_dir: args.load = ckpt_dir args.save = ckpt_dir + args.save_tokenizer_assets = False set_args(args) # Create and save a checkpoint first. @@ -417,6 +419,7 @@ def test_dist_checkpoint_versioning(init_model_parallel, tmp_path_dist_ckpt, cre ) as ckpt_dir: args.load = ckpt_dir args.save = ckpt_dir + args.save_tokenizer_assets = False set_args(args) # Create and save a checkpoint first. diff --git a/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py b/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py index 3e1bdb244d4..0f4eb35c510 100644 --- a/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py +++ b/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py @@ -190,6 +190,7 @@ def create_test_args( args.overlap_grad_reduce = overlap_grad_reduce args.accumulate_allreduce_grads_in_fp32 = grad_reduce_in_fp32 args.ddp_bucket_size = 40960 + args.save_tokenizer_assets = False args.num_experts = 2 args.moe_layer_freq = 1 diff --git a/tests/unit_tests/transformer/test_attention.py b/tests/unit_tests/transformer/test_attention.py index c5f31666903..e82ae7e458a 100644 --- a/tests/unit_tests/transformer/test_attention.py +++ b/tests/unit_tests/transformer/test_attention.py @@ -515,6 +515,7 @@ def initialize_gpt_model( with TempNamedDir(tmp_path_dist_ckpt / 'test_parallel_attn', sync=True) as ckpt_dir: # Set argument mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False set_args(mock_args) # Initialize baseline model diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 95517f7176d..3131222f800 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -1531,6 +1531,7 @@ def initialize_gpt_model( with TempNamedDir(tmp_path_dist_ckpt / 'test_parallel_mla', sync=True) as ckpt_dir: # Set argument mock_args = parse_args(ignore_unknown_args=True) + mock_args.save_tokenizer_assets = False set_args(mock_args) # Initialize baseline model diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index ee1eb02267f..8852a8408b8 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -597,6 +597,7 @@ def create_test_args( args.no_save_optim = True args.no_load_optim = True args.no_load_rng = True + args.save_tokenizer_assets = False if HAVE_TE: # only use grouped gemm if there is TE args.moe_grouped_gemm = True @@ -1431,6 +1432,7 @@ def create_test_args( args.no_load_optim = True args.no_load_rng = True args.bf16 = True + args.save_tokenizer_assets = False # Unified pattern: "main/mtp/mtp" - main decoder "M*M*", MTP pattern "M*" with 2 depths args.hybrid_layer_pattern = "M*M*/M*/M*" From d347eef80c875ca66c8e9ff4a5e12d2936f614f3 Mon Sep 17 00:00:00 2001 From: Charlie Truong Date: Fri, 14 Aug 2026 14:25:45 -0500 Subject: [PATCH 035/101] ci: Do not upload docker oci artifact attestations (#6542) Signed-off-by: Charlie Truong --- .github/workflows/cicd-main.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 112425c9879..429632980a2 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -701,6 +701,8 @@ jobs: - name: Build and push uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f # v7.1.0 + env: + BUILDX_NO_DEFAULT_OCI_ARTIFACT: "true" with: file: ${{ steps.base-image.outputs.dockerfile }} push: true From 4dc5c9679f2809ee617bc8fdbe78a42ca81d52ab Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Fri, 14 Aug 2026 13:40:45 +0000 Subject: [PATCH 036/101] fix(ci): AUT-1427 skip Claude labeling for external authors (#6398) Signed-off-by: svcnemo-autobot --- .github/workflows/claude-complexity-label.yml | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/.github/workflows/claude-complexity-label.yml b/.github/workflows/claude-complexity-label.yml index d9b765fb3dd..f00eb482e40 100644 --- a/.github/workflows/claude-complexity-label.yml +++ b/.github/workflows/claude-complexity-label.yml @@ -5,9 +5,40 @@ on: types: [ready_for_review] jobs: + check-write-permission: + name: Check PR Author Permission + runs-on: ubuntu-latest + permissions: + contents: read + outputs: + has-write-permission: ${{ steps.permission.outputs.has-write-permission }} + env: + GH_TOKEN: ${{ github.token }} + REPO: ${{ github.repository }} + AUTHOR: ${{ github.event.pull_request.user.login }} + steps: + - name: Check repository permission + id: permission + run: | + set -euo pipefail + + permission=$(gh api "repos/$REPO/collaborators/$AUTHOR/permission" --jq '.permission') + case "$permission" in + admin|maintain|write) + has_write_permission=true + ;; + *) + has_write_permission=false + ;; + esac + + echo "has-write-permission=$has_write_permission" >> "$GITHUB_OUTPUT" + analyze_complexity: name: Analyze PR Complexity runs-on: ubuntu-latest + needs: check-write-permission + if: needs.check-write-permission.outputs.has-write-permission == 'true' permissions: contents: read pull-requests: read From c8b2a4a32dbd8da95f200208ef7a78d6b93ff533 Mon Sep 17 00:00:00 2001 From: Charlie Truong Date: Fri, 14 Aug 2026 15:34:00 +0000 Subject: [PATCH 037/101] docs: Add 0.19.0 to version picker (#6536) Signed-off-by: Charlie Truong --- docs/versions1.json | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/docs/versions1.json b/docs/versions1.json index b0bc489fa71..84942f53740 100644 --- a/docs/versions1.json +++ b/docs/versions1.json @@ -5,9 +5,20 @@ "url": "https://docs.nvidia.com/megatron-core/developer-guide/nightly/" }, { - "name": "0.17.0 (latest)", + "name": "0.19.0 (latest)", + "version": "0.19.0", + "url": "https://docs.nvidia.com/megatron-core/developer-guide/latest/", + "preferred": true + }, + { + "name": "0.18.0", + "version": "0.18.0", + "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.18.0/" + }, + { + "name": "0.17.0", "version": "0.17.0", - "url": "https://docs.nvidia.com/megatron-core/developer-guide/latest/" + "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.17.0/" }, { "name": "0.16.0", From 41ba4d94e9553452547d27f9f370b39640fc0e8f Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Fri, 14 Aug 2026 15:56:52 +0000 Subject: [PATCH 038/101] Miscellaneous fixes to the RL loop (#6044) Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Jorge Albericio Co-authored-by: Claude Fable 5 --- megatron/rl/agent/pass_at_evaluation_agent.py | 2 +- megatron/rl/rl_utils.py | 75 +++++++++- megatron/rl/sequence_packing_utils.py | 6 +- tests/unit_tests/rl/test_rl_agents.py | 33 +++++ tests/unit_tests/rl/test_rl_utils.py | 135 ++++++++++++++++++ .../rl/test_sequence_packing_utils.py | 19 ++- 6 files changed, 252 insertions(+), 18 deletions(-) diff --git a/megatron/rl/agent/pass_at_evaluation_agent.py b/megatron/rl/agent/pass_at_evaluation_agent.py index c04e1c2772f..09917df64e7 100644 --- a/megatron/rl/agent/pass_at_evaluation_agent.py +++ b/megatron/rl/agent/pass_at_evaluation_agent.py @@ -25,7 +25,7 @@ def pass_at_k(n_samples: int, n_correct: int, k: int) -> float: class PassAtEvaluationResult(RewardEvaluationResult): pass_at: dict[int, float] - response: list[str] | list[LLMChatMessage] + response: list[str | LLMChatMessage] reward: list[float] greedy_response: str | LLMChatMessage greedy_reward: float diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index bfa15cacd6e..6c77cf13694 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -10,6 +10,7 @@ import logging import math import os +import time from collections import Counter, defaultdict from contextlib import contextmanager, nullcontext from dataclasses import dataclass @@ -498,8 +499,12 @@ def align_unpacked_inference_logprobs( first_gen_tok = gen_masks_for_alignment.int().argmax(dim=1) - 1 # Align inference logprobs with old_logprobs - # Note: We use old_logprobs_for_data as template since it has correct shape - padded_inference_logprobs = old_logprobs_for_data.clone() + # Note: We use old_logprobs_for_data as template since it has correct shape. + # Inference logprobs ride in float32, matching pack_inference_logprobs: the + # IS-correction ratio is precision-sensitive, and under bf16 training + # old_logprobs arrives bf16 while the wire logprobs are float32 — cloning + # the template dtype would silently downcast them in the assignment below. + padded_inference_logprobs = old_logprobs_for_data.clone().float() # We need to align old_logprobs and inference logprobs as the latter are only for generations for i, inf_logprobs in enumerate(inference_logprobs): @@ -512,7 +517,9 @@ def align_unpacked_inference_logprobs( end_idx = min(first_gen_idx + len(inf_logprobs), padded_inference_logprobs.shape[1]) actual_len = end_idx - first_gen_idx if actual_len > 0: - padded_inference_logprobs[i, first_gen_idx:end_idx] = inf_logprobs[:actual_len] + padded_inference_logprobs[i, first_gen_idx:end_idx] = inf_logprobs[:actual_len].to( + padded_inference_logprobs.dtype + ) # Create truncated mask for statistics if old_logprobs_for_data.shape[1] + 1 < gen_masks_for_alignment.shape[1]: @@ -554,9 +561,43 @@ def get_agent(env_config_path): _INFERENCE_INTERFACE = None +def _eager_train_comm_warmup(): + """Run one tiny allreduce per training process group to initialize NCCL communicators.""" + group_getters = { + 'dp': partial(mpu.get_data_parallel_group), + 'dp_cp': partial(mpu.get_data_parallel_group, with_context_parallel=True), + 'cp': mpu.get_context_parallel_group, + 'mp': mpu.get_model_parallel_group, + 'tp_dp_cp': partial(mpu.get_tensor_and_data_parallel_group, with_context_parallel=True), + 'expert_dp': mpu.get_expert_data_parallel_group, + 'expert_tp': mpu.get_expert_tensor_parallel_group, + 'expert_mp': mpu.get_expert_model_parallel_group, + } + warmed = [] + start = time.perf_counter() + for name, getter in group_getters.items(): + try: + group = getter() + except AssertionError: + continue + if group is None: + continue + group_start = time.perf_counter() + dist.all_reduce(torch.zeros(1, device='cuda'), group=group) + torch.cuda.synchronize() + warmed.append(f'{name}={time.perf_counter() - group_start:.2f}s') + log_single_rank( + logger, + logging.INFO, + f'[EagerCommWarmup] initialized NCCL communicators in ' + f'{time.perf_counter() - start:.2f}s: {warmed}', + ) + + def get_inference_interface(args, loop, model): global _INFERENCE_INTERFACE if _INFERENCE_INTERFACE is None: + _eager_train_comm_warmup() _INFERENCE_INTERFACE = loop.run_until_complete( MegatronLocal.launch( model[0], @@ -1092,6 +1133,21 @@ def prep_wandb_metrics( ), } + # All-empty wave: every episode failed and was dropped (nothing to aggregate), + # or every surviving rollout is a zero-turn placeholder (rewards exist but the + # length/turn stats below would reduce over empty sequences). Keep the failure + # counters: they are the one signal such a wave still carries. + if ( + len(advantages) == 0 + or not rewards + or not any(keep for g in real_mask for keep in g) + ): + logger.warning( + "prep_wandb_metrics: empty wave (0 usable rollouts); " + "skipping rollout metrics this iteration." + ) + return failure_metrics + def _real(grouped): """Grouped per-rollout entries with placeholder (zero-turn) rollouts removed.""" return [ @@ -2151,7 +2207,18 @@ def evaluate_and_print_results_rl( tb_writer.add_scalar(k, v, iteration) wandb_writer = get_wandb_writer() if wandb_writer: - wandb_writer.log(eval_metrics, step=iteration) + if args.do_train: + wandb_writer.log(eval_metrics, step=iteration) + else: + # Without a training loop the eval may target an arbitrary (older) checkpoint, + # breaking wandb step monotonicity. + wandb_writer.define_metric('eval_only/*', step_metric='eval_checkpoint_iter') + wandb_writer.log( + { + 'eval_checkpoint_iter': iteration, + **{f'eval_only/{k}': v for k, v in eval_metrics.items()}, + } + ) logger.info( "Evaluation results:" + "".join([f"\n\t{k}: {v:0.4f}" for k, v in eval_metrics.items()]) diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index df94f633ff6..bc3f06ed587 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -606,9 +606,11 @@ def compute_packed_inference_logprobs_stats( # poisons the mismatch stats with spurious |p_old - 1| terms. mask = mask & filled_mask.to(mask.device) - # Ensure shapes match if mask.shape != old_logprobs.shape: - return + raise RuntimeError( + f"packed inference-logprob stats shape mismatch: loss-mask " + f"{tuple(mask.shape)} vs old_logprobs {tuple(old_logprobs.shape)}" + ) # Update group statistics using common helper update_inference_logprobs_group_stats( diff --git a/tests/unit_tests/rl/test_rl_agents.py b/tests/unit_tests/rl/test_rl_agents.py index 23f3a65e549..794b7335bc9 100644 --- a/tests/unit_tests/rl/test_rl_agents.py +++ b/tests/unit_tests/rl/test_rl_agents.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock import pytest +from pydantic import ValidationError def _mock_math_verify(): @@ -19,6 +20,8 @@ def _mock_math_verify(): from examples.rl.environments.countdown.countdown import compute_score as countdown_compute_score from examples.rl.environments.math.math_agent import MathAgent +from megatron.rl.agent.pass_at_evaluation_agent import PassAtEvaluationResult +from megatron.rl.inference import LLMChatMessage def _make_math_agent( @@ -105,3 +108,33 @@ def test_math_agent_compute_score( def test_countdown_compute_score(self, response, ground_truth, expected_reward): score = countdown_compute_score(response, ground_truth) assert score == pytest.approx(expected_reward) + + +def test_pass_at_result_accepts_mixed_responses(): + """Aggregated pass@k responses mix assistant messages with '' placeholders + for dropped/over-context rollouts; the result type validates element-wise.""" + msg = LLMChatMessage(role='assistant', content='the answer is 4') + result = PassAtEvaluationResult( + prompt='2+2?', + response=[msg, '', 'plain-string response'], + reward=[1.0, 0.0, 0.5], + pass_at={1: 0.5}, + greedy_response=msg, + greedy_reward=1.0, + ) + # Assert that the plain-string responses are not type-cast. + assert isinstance(result.response[0], LLMChatMessage) + assert result.response[0].content == 'the answer is 4' + assert result.response[1] == '' + assert result.response[2] == 'plain-string response' + + # Element-wise validation: invalid elements still fail. + with pytest.raises(ValidationError): + PassAtEvaluationResult( + prompt='x', + response=[{'not': 'valid'}], + reward=[0.0], + pass_at={1: 0.0}, + greedy_response='y', + greedy_reward=0.0, + ) diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 3d7eb2a1a21..dac3f963d4f 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -415,6 +415,102 @@ def test_megatron_rl_inference_mode_restores_training_cuda_graph_state( lang_module.eval.assert_called_once() lang_module.train.assert_called_once() + @pytest.mark.parametrize( + "do_train, skip_train, expect_train_branch", + [ + pytest.param(True, False, True, id="training"), + pytest.param(True, True, True, id="inference_only"), + pytest.param(False, False, False, id="eval_only"), + ], + ) + def test_evaluate_and_print_results_rl_wandb_step_handling( + self, monkeypatch, do_train, skip_train, expect_train_branch + ): + """Only eval-only mode (--skip-train --train-iters 0, do_train=False) takes the decoupled + eval_only/* wandb path. --skip-train alone still runs the RL loop with monotonically + advancing iteration, so its evals log against the global step (see PR #6043).""" + args = SimpleNamespace( + do_train=do_train, + skip_train=skip_train, + world_size=1, + cuda_graph_impl="none", + rl_offload_optimizer_during_inference=False, + langrl_env_config=None, + rl_prompts_per_eval=1, + rl_default_temperature=1.0, + rl_default_top_p=1.0, + rl_default_top_k=0, + seq_length=SEQ, + ) + response = SimpleNamespace(env_id="env", metrics=lambda: {"reward": [1.0, 0.5]}, results=[]) + agent = SimpleNamespace(run_evaluation=lambda request: [response]) + wandb_writer = MagicMock() + + monkeypatch.setattr(rl_utils, "get_args", lambda: args) + monkeypatch.setattr( + rl_utils, "megatron_rl_inference_mode", lambda *_a, **_k: nullcontext(MagicMock()) + ) + monkeypatch.setattr( + rl_utils, "get_nvtx_range", lambda: (lambda *args, **kwargs: nullcontext()) + ) + monkeypatch.setattr(rl_utils, "get_agent", lambda *_args: agent) + monkeypatch.setattr( + rl_utils, "get_asyncio_loop", lambda: SimpleNamespace(run_until_complete=lambda r: r) + ) + monkeypatch.setattr( + rl_utils, "EvaluationRequest", lambda **kwargs: SimpleNamespace(**kwargs) + ) + monkeypatch.setattr(rl_utils, "get_wandb_writer", lambda: wandb_writer) + monkeypatch.setattr(rl_utils, "lang_rl_log_dir", None) + monkeypatch.setattr(rl_utils.dist, "get_rank", lambda: 0) + monkeypatch.setattr( + rl_utils.dist, "gather_object", lambda obj, lst, dst=0: lst.__setitem__(0, obj) + ) + + rl_utils.evaluate_and_print_results_rl( + data_iterator=None, + model=[MagicMock()], + optimizer=MagicMock(), + iteration=7, + write_to_tensorboard=False, + ) + + if expect_train_branch: + wandb_writer.log.assert_called_once_with({"env_eval_mean_reward": 0.75}, step=7) + wandb_writer.define_metric.assert_not_called() + else: + wandb_writer.define_metric.assert_called_once_with( + "eval_only/*", step_metric="eval_checkpoint_iter" + ) + wandb_writer.log.assert_called_once_with( + {"eval_checkpoint_iter": 7, "eval_only/env_eval_mean_reward": 0.75} + ) + + @pytest.mark.parametrize("template_dtype", [torch.float32, torch.bfloat16]) + def test_align_unpacked_inference_logprobs_rides_in_float32(self, template_dtype): + """Under bf16 training old_logprobs arrives bf16 while wire inference + logprobs are float32; the alignment must not silently downcast them + into the template's dtype (the IS-correction ratio is + precision-sensitive), and the aligned result rides in float32 like + pack_inference_logprobs' packed twin.""" + old_logprobs = torch.full((2, 6), -1.0, dtype=template_dtype) + # Two generated tokens at positions 3,4: the logprob of the token at + # position p lives at index p-1, so the filled slots are 2,3. + gen_masks = torch.zeros((2, 7), dtype=torch.bool) + gen_masks[:, 3:5] = True + inference_logprobs = [torch.tensor([-0.5, -0.25], dtype=torch.float32) for _ in range(2)] + aligned = rl_utils.align_unpacked_inference_logprobs( + inference_logprobs=inference_logprobs, + old_logprobs_for_data=old_logprobs, + generation_masks=gen_masks, + group_stats=SimpleNamespace(), + ) + assert aligned.dtype == torch.float32 + assert aligned[0, 2].item() == pytest.approx(-0.5) + assert aligned[0, 3].item() == pytest.approx(-0.25) + # Slots without inference logprobs keep the template's values. + assert aligned[0, 0].item() == pytest.approx(-1.0) + @pytest.mark.parametrize( "share_config", [pytest.param(True, id="shared-config"), pytest.param(False, id="distinct-config")], @@ -1367,6 +1463,45 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders # mean_completion_gap = mean([6-5, 6-3, 6-5, 6-1]) = mean([1, 3, 1, 5]) = 2.5 assert metrics["mean_completion_gap"] == 2.5 + # Case where every episode has failed: the wave still reports the + # failure counters (the one signal it carries) and skips the tables. + writer = MagicMock() + metrics = rl_utils.prep_wandb_metrics( + writer, + traj_lens=[], + turn_lens=[], + rewards=[], + num_turns=[], + advantages=[], + policy_epoch=[], + kv_cache_epoch=[], + completed_epochs=[], + num_evictions=[], + current_iteration=6, + ) + assert metrics == {'failed_rollouts/count': 0, 'failed_rollouts/ratio': 0.0} + writer.Table.assert_not_called() + + # All-placeholder wave: rewards exist (placeholders carry 0.0) but every + # rollout is a zero-turn pad, so the masked stats have nothing to reduce + # over; the guard must still report the failure counters. + writer = MagicMock() + metrics = rl_utils.prep_wandb_metrics( + writer, + traj_lens=[[0, 0]], + turn_lens=[[]], + rewards=[[0.0, 0.0]], + num_turns=[[0, 0]], + advantages=[0.0, 0.0], + policy_epoch=[[[0], [0]]], + kv_cache_epoch=[[[0], [0]]], + completed_epochs=[[]], + num_evictions=[[0, 0]], + current_iteration=6, + ) + assert metrics == {'failed_rollouts/count': 2, 'failed_rollouts/ratio': 1.0} + writer.Table.assert_not_called() + def test_compute_group_stats_excludes_placeholders_from_metric_fields(self): def real_rollout(tokens, problem_id): return TokenRollout( diff --git a/tests/unit_tests/rl/test_sequence_packing_utils.py b/tests/unit_tests/rl/test_sequence_packing_utils.py index 687e3aef434..f5f56c8877d 100644 --- a/tests/unit_tests/rl/test_sequence_packing_utils.py +++ b/tests/unit_tests/rl/test_sequence_packing_utils.py @@ -389,7 +389,7 @@ def test_compute_packed_inference_logprobs_stats_with_mismatch(): def test_compute_packed_inference_logprobs_stats_shape_mismatch(): - """Test that function handles shape mismatch gracefully.""" + """Test that a shape mismatch raises instead of silently skipping the stats.""" # Mismatched shapes old_logprobs = torch.tensor([[-0.5, -0.3, -0.2]]) # 3 elements packed_inference_logprobs = torch.tensor([[-0.5, -0.3, -0.2]]) @@ -397,16 +397,13 @@ def test_compute_packed_inference_logprobs_stats_shape_mismatch(): group_stats = MockGroupStats() - # Should not raise, but stats should remain None due to shape mismatch - sequence_packing_utils.compute_packed_inference_logprobs_stats( - old_logprobs=old_logprobs, - packed_inference_logprobs=packed_inference_logprobs, - packed_loss_mask=packed_loss_mask, - group_stats=group_stats, - ) - - # Stats should remain None due to shape mismatch - assert group_stats.mean_piold_to_inf_prob is None + with pytest.raises(RuntimeError, match="shape mismatch"): + sequence_packing_utils.compute_packed_inference_logprobs_stats( + old_logprobs=old_logprobs, + packed_inference_logprobs=packed_inference_logprobs, + packed_loss_mask=packed_loss_mask, + group_stats=group_stats, + ) def test_packing_observability_metrics(): From cac6b7bc69dc24baffb4634b83b6e1c68ba97460 Mon Sep 17 00:00:00 2001 From: wdykas Date: Fri, 14 Aug 2026 16:32:51 +0000 Subject: [PATCH 039/101] Add disaggregated KV state handoff (#6222) Signed-off-by: William Dykas --- .../inference/contexts/kv_block_allocator.py | 21 + .../handlers.py | 63 ++ .../disaggregation/decode_admission.py | 136 +++ .../core/inference/disaggregation/engine.py | 16 + .../handoff_completion_tracker.py | 97 ++ .../disaggregation/inference_state_handoff.py | 968 ++++++++++++++++++ .../disaggregation/pending_handoff_imports.py | 44 + .../disaggregation/transfer_backends/nixl.py | 1 + .../core/inference/disaggregation/utils.py | 36 + .../core/inference/engines/dynamic_engine.py | 234 ++++- megatron/core/inference/headers.py | 4 + megatron/core/inference/inference_client.py | 102 +- megatron/core/inference/inference_request.py | 8 + megatron/core/inference/sampling_params.py | 1 + .../text_generation_controller.py | 84 +- .../contexts/test_kv_block_allocator.py | 26 + .../test_dynamic_engine_async_sched.py | 52 +- ...est_data_parallel_inference_coordinator.py | 10 +- .../test_disagg_handoff_lifecycle.py | 867 ++++++++++++++++ .../inference/test_inference_client.py | 21 + .../test_inference_client_streaming.py | 24 + .../test_text_generation_controller.py | 80 +- 22 files changed, 2855 insertions(+), 40 deletions(-) create mode 100644 megatron/core/inference/disaggregation/decode_admission.py create mode 100644 megatron/core/inference/disaggregation/engine.py create mode 100644 megatron/core/inference/disaggregation/handoff_completion_tracker.py create mode 100644 megatron/core/inference/disaggregation/inference_state_handoff.py create mode 100644 megatron/core/inference/disaggregation/pending_handoff_imports.py create mode 100644 tests/unit_tests/inference/test_disagg_handoff_lifecycle.py diff --git a/megatron/core/inference/contexts/kv_block_allocator.py b/megatron/core/inference/contexts/kv_block_allocator.py index 90043c9bad4..808b629919d 100644 --- a/megatron/core/inference/contexts/kv_block_allocator.py +++ b/megatron/core/inference/contexts/kv_block_allocator.py @@ -46,6 +46,12 @@ def __init__( self.on_blocks_deregistered: Optional[Callable] = None self._blocks_deregistered_observers: list[BlocksDeregisteredObserver] = [] + # Handoff blocks remain pinned until decode finishes pulling them. + # Pinning at request finish only happens on engines with KV transfer + # configured (setup_kv_transfer flips this on); other engines have no + # release path for the pins. + self.enable_handoff_pinning = False + assert ( 0 <= paused_limit <= pool_size - 2 ), "paused block limit must leave at least one usable block outside the limit" @@ -251,6 +257,21 @@ def release_memory_blocks(self, blocks: Tensor) -> None: self.block_bag[self.pool_avail : self.pool_avail + num_blocks] = blocks self.pool_avail += num_blocks + def retain_memory_blocks(self, block_ids: list[int]) -> None: + """Add one prefix-cache reference to each block. + + Args: + block_ids: Blocks retained by a new owner. + """ + assert self.enable_prefix_caching, "retaining KV blocks requires prefix caching" + if block_ids: + blocks = torch.tensor(block_ids, dtype=torch.int32, device='cpu') + unique_blocks, retain_counts = torch.unique(blocks, return_counts=True) + self.block_ref_counts[unique_blocks] += retain_counts.to( + dtype=self.block_ref_counts.dtype + ) + self.update_timestamps(unique_blocks) + def reset(self) -> None: """Reset the allocator to initial state. diff --git a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py index 2c1712a1de3..b514b5e62c7 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py @@ -135,6 +135,69 @@ def handle_submit_request(coordinator, sender_identity, payload): ) +@message_handler(Headers.SUBMIT_REQUEST_WITH_KV) +def handle_submit_request_with_kv(coordinator, sender_identity, payload): + """Route a client-supplied KV handoff to a decode engine.""" + + if sender_identity not in coordinator.known_clients: + logging.info( + "Received SUBMIT_REQUEST_WITH_KV from unknown client %s; ignoring.", sender_identity + ) + return + if len(payload) != 6: + logging.error( + "Coordinator: malformed SUBMIT_REQUEST_WITH_KV payload with %d fields", len(payload) - 1 + ) + return + + client_request_id, prompt, sampling_params, kv_meta, src_block_ids = payload[1:] + request_id = coordinator.next_request_id + coordinator.next_request_id += 1 + coordinator.request_id_to_client_id[request_id] = sender_identity + coordinator.request_id_to_client_request_id[request_id] = client_request_id + coordinator.client_request_to_request_id[(sender_identity, client_request_id)] = request_id + + if isinstance(prompt, torch.Tensor): + prompt = prompt.tolist() + elif not isinstance(prompt, (str, list)): + raise TypeError(f"unsupported prompt type {type(prompt).__name__}") + engine_payload = msgpack.packb( + [ + Headers.SUBMIT_REQUEST_WITH_KV.value, + request_id, + prompt, + sampling_params, + kv_meta, + src_block_ids, + ], + use_bin_type=True, + ) + + for _ in range(len(coordinator.identities_of_data_parallel_ranks)): + next_identity = coordinator.get_least_loaded_data_parallel_rank() + if coordinator._send_to_engine(next_identity, engine_payload): + break + else: + logging.error("Coordinator: no reachable engines for handoff request %d", request_id) + del coordinator.request_id_to_client_id[request_id] + del coordinator.request_id_to_client_request_id[request_id] + del coordinator.client_request_to_request_id[(sender_identity, client_request_id)] + return True + + coordinator.request_id_to_rank[request_id] = next_identity + coordinator._pending_counts[coordinator.identity_to_rank_index[next_identity]] += 1 + + +@message_handler(Headers.RELEASE_KV) +def handle_release_kv(coordinator, sender_identity, payload): + """Broadcast release of prefill blocks retained for a completed handoff.""" + + if sender_identity not in coordinator.known_clients: + logging.warning("Coordinator: ignoring RELEASE_KV from unknown client.") + return + coordinator._broadcast_to_engines([Headers.RELEASE_KV.value, int(payload[1])]) + + @message_handler( Headers.PAUSE, Headers.UNPAUSE, diff --git a/megatron/core/inference/disaggregation/decode_admission.py b/megatron/core/inference/disaggregation/decode_admission.py new file mode 100644 index 00000000000..9ec5c9f9474 --- /dev/null +++ b/megatron/core/inference/disaggregation/decode_admission.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Direct decode admission for imported prefill state.""" + +from __future__ import annotations + +import math + +import torch + + +def additional_decode_blocks(prompt_length: int, input_token_count: int, block_size: int) -> int: + """Return blocks needed beyond those containing the imported prompt.""" + + prompt_blocks = math.ceil(prompt_length / block_size) + blocks_after_input = math.ceil((prompt_length + input_token_count) / block_size) + return blocks_after_input - prompt_blocks + + +def can_admit_prefilled_decode(context, input_token_count: int) -> bool: + """Return whether an imported request can enter the live decode batch.""" + + if context.num_prefill_requests != 0 or context.chunked_prefill_request_id != -1: + raise RuntimeError("A decode-only engine cannot execute prefill requests") + if context.total_request_count >= context.max_requests: + return False + if context.active_token_count + input_token_count > context.max_tokens: + return False + return True + + +def admit_prefilled_decode( + context, + request, + prompt_block_ids: list[int], + continuation_block_ids: list[int], + input_tokens: list[int], +) -> None: + """Admit imported KV state directly as an active decode request. + + The imported prompt blocks already carry one allocator reference. This + function transfers that ownership to the live request and schedules the + sampled first token (plus any MTP proposals) as the next decode query. + """ + + input_token_count = len(input_tokens) + if not can_admit_prefilled_decode(context, input_token_count): + raise RuntimeError("Imported prefill state cannot enter the current decode batch") + + prompt_length = len(request.prompt_tokens) + expected_prompt_blocks = math.ceil(prompt_length / context.block_size_tokens) + if len(prompt_block_ids) != expected_prompt_blocks: + raise ValueError( + f"Expected {expected_prompt_blocks} imported prompt blocks, " + f"got {len(prompt_block_ids)}" + ) + expected_decode_blocks = additional_decode_blocks( + prompt_length, input_token_count, context.block_size_tokens + ) + if len(continuation_block_ids) != expected_decode_blocks: + raise ValueError( + f"Expected {expected_decode_blocks} decode continuation blocks, " + f"got {len(continuation_block_ids)}" + ) + if input_token_count != context.num_speculative_tokens + 1: + raise ValueError( + "Imported decode inputs must contain one sampled token plus the configured " + f"MTP proposals: expected {context.num_speculative_tokens + 1}, " + f"got {input_token_count}" + ) + if request.get_metadata_types() != context.request_metadata_types: + raise ValueError("Imported request metadata does not match the decode context") + + current_id = context.total_request_count + all_block_ids = prompt_block_ids + continuation_block_ids + row = context.request_to_kv_block_ids[current_id] + row.fill_(-1) + row[: len(all_block_ids)] = torch.tensor(all_block_ids, dtype=row.dtype, device=row.device) + + context.request_ids[current_id] = request.request_id + for metadata, metadata_type in zip(request.tracked_metadata, request.get_metadata_types()): + label, _ = metadata_type + value = metadata + if not isinstance(value, torch.Tensor): + value = torch.as_tensor( + value, + device=context.request_metadata[label].device, + dtype=context.request_metadata[label].dtype, + ) + context.request_metadata[label][current_id] = value + + context.request_kv_length_offsets[current_id] = prompt_length + context.request_query_lengths[current_id] = input_token_count + context.request_output_lengths[current_id] = ( + prompt_length + request.sampling_params.num_tokens_to_generate + ) + context.request_in_prefill_status_tensor[current_id] = 0 + context.request_kv_block_counts[current_id] = len(all_block_ids) + context.request_last_kv_block_id[current_id] = all_block_ids[-1] + context.request_last_kv_block_offset[current_id] = ( + prompt_length - 1 + input_token_count + ) % context.block_size_tokens + + token_start = context.active_token_count + token_end = token_start + input_token_count + token_positions = torch.arange( + prompt_length, + prompt_length + input_token_count, + dtype=context.token_to_pos_ids.dtype, + device=context.token_to_pos_ids.device, + ) + context.token_to_input_ids[token_start:token_end] = torch.tensor( + input_tokens, + dtype=context.token_to_input_ids.dtype, + device=context.token_to_input_ids.device, + ) + context.token_to_pos_ids[token_start:token_end] = token_positions + context.token_to_request_idx[token_start:token_end] = current_id + context.token_to_position_in_request[token_start:token_end] = token_positions + context.token_to_local_position_within_kv_block[token_start:token_end] = ( + token_positions % context.block_size_tokens + ) + block_ids_tensor = torch.tensor( + all_block_ids, + dtype=context.token_to_block_idx.dtype, + device=context.token_to_block_idx.device, + ) + context.token_to_block_idx[token_start:token_end] = block_ids_tensor[ + token_positions // context.block_size_tokens + ] + + context.active_token_count = token_end + context.total_request_count += 1 + + request.remaining_prompt_tokens = request.remaining_prompt_tokens.new_empty(0) + request.add_event_add_context() diff --git a/megatron/core/inference/disaggregation/engine.py b/megatron/core/inference/disaggregation/engine.py new file mode 100644 index 00000000000..411d163e83b --- /dev/null +++ b/megatron/core/inference/disaggregation/engine.py @@ -0,0 +1,16 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Dynamic inference engine composed with KV-cache handoff behavior.""" + +from megatron.core.inference.disaggregation.inference_state_handoff import ( + InferenceStateHandoffMixin, +) +from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine + + +class DisaggDynamicInferenceEngine(InferenceStateHandoffMixin, DynamicInferenceEngine): + """Dynamic inference engine with prefill/decode KV-cache handoff support. + + Used by both control planes: the Dynamo integration and the + coordinator-native 2-hop mode. + """ diff --git a/megatron/core/inference/disaggregation/handoff_completion_tracker.py b/megatron/core/inference/disaggregation/handoff_completion_tracker.py new file mode 100644 index 00000000000..ae5f6ac872a --- /dev/null +++ b/megatron/core/inference/disaggregation/handoff_completion_tracker.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU control-plane aggregation for model-parallel handoff completion.""" + +import socket +import struct +from typing import Any + +import torch.distributed as dist + +try: + import zmq +except ImportError: + from unittest.mock import MagicMock + + zmq = MagicMock() + HAVE_ZMQ = False +else: + HAVE_ZMQ = True + + +class HandoffCompletionTracker: + """Aggregate per-rank transfer results at the model-parallel coordinator.""" + + _REPORT_FORMAT = "!qi?" + + def __init__(self, zmq_context, process_group: dist.ProcessGroup, hostname: str | None = None): + if not HAVE_ZMQ: + raise ImportError("pyzmq is required for disaggregated KV handoff") + + self.rank = dist.get_rank(process_group) + self.world_size = dist.get_world_size(process_group) + self.is_coordinator = self.rank == 0 + self._reports: dict[int, dict[int, bool]] = {} + self._failure_notified: set[int] = set() + self._socket: Any = None + self.sockets = [] + + if self.world_size == 1: + return + + coordinator_rank = dist.get_process_group_ranks(process_group)[0] + if self.is_coordinator: + self._socket = zmq_context.socket(zmq.PULL) + self._socket.bind_to_random_port(f"tcp://{hostname or socket.gethostname()}") + address = self._socket.getsockopt_string(zmq.LAST_ENDPOINT) + dist.broadcast_object_list([address], src=coordinator_rank, group=process_group) + else: + address_holder = [None] + dist.broadcast_object_list(address_holder, src=coordinator_rank, group=process_group) + self._socket = zmq_context.socket(zmq.PUSH) + self._socket.connect(address_holder[0]) + self.sockets.append(self._socket) + + def report(self, request_id: int, failed: bool) -> None: + """Report this rank's terminal transfer result once for a request.""" + + if self.is_coordinator: + self._record(request_id, self.rank, failed) + return + self._socket.send(struct.pack(self._REPORT_FORMAT, request_id, self.rank, failed)) + + def drain_completed(self) -> list[tuple[int, bool]]: + """Return requests that failed or completed on every model-parallel rank.""" + + if not self.is_coordinator: + return [] + if self._socket is not None: + while True: + try: + request_id, rank, failed = struct.unpack( + self._REPORT_FORMAT, self._socket.recv(flags=zmq.NOBLOCK) + ) + self._record(request_id, rank, failed) + except zmq.Again: + break + + completed = [] + for request_id, reports in list(self._reports.items()): + failed = any(reports.values()) + if failed and request_id not in self._failure_notified: + completed.append((request_id, failed)) + self._failure_notified.add(request_id) + if len(reports) == self.world_size: + if request_id not in self._failure_notified: + completed.append((request_id, False)) + del self._reports[request_id] + self._failure_notified.discard(request_id) + return completed + + def _record(self, request_id: int, rank: int, failed: bool) -> None: + reports = self._reports.setdefault(request_id, {}) + previous = reports.setdefault(rank, failed) + if previous != failed: + raise RuntimeError( + f"Conflicting KV handoff results for request {request_id} from rank {rank}" + ) diff --git a/megatron/core/inference/disaggregation/inference_state_handoff.py b/megatron/core/inference/disaggregation/inference_state_handoff.py new file mode 100644 index 00000000000..fa989aad6f6 --- /dev/null +++ b/megatron/core/inference/disaggregation/inference_state_handoff.py @@ -0,0 +1,968 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Engine-side lifecycle for disaggregated prefill/decode KV-cache handoff.""" + +from __future__ import annotations + +import asyncio +import logging +from collections import deque +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict + +import torch + +from megatron.core.inference.disaggregation.decode_admission import ( + additional_decode_blocks, + admit_prefilled_decode, + can_admit_prefilled_decode, +) +from megatron.core.inference.disaggregation.handoff_completion_tracker import ( + HandoffCompletionTracker, +) +from megatron.core.inference.disaggregation.pending_handoff_imports import ( + DeferredKvHandoff, + PendingKvImport, +) +from megatron.core.inference.disaggregation.transfer_backends.base import ( + construct_kv_transfer_backend_class, +) +from megatron.core.inference.disaggregation.utils import ( + drop_transfer_prefix_blocks, + transfer_block_count, +) +from megatron.core.inference.inference_request import ( + DynamicInferenceEvent, + DynamicInferenceEventType, + Status, + compute_block_hashes_batched, +) +from megatron.core.utils import get_pg_rank, get_pg_size + +if TYPE_CHECKING: + from megatron.core.inference.inference_request import DynamicInferenceRequest + from megatron.core.inference.sampling_params import SamplingParams + + +@dataclass(frozen=True) +class _PreparedHandoffMetadata: + """Per-request metadata assembled once for a completed prefill batch.""" + + local_blocks: list[int] # Full prompt block table, including a partial tail. + kv_meta: Any + resume_tokens: list[int] + + +class InferenceStateHandoffMixin: + """Optional KV-cache handoff behavior composed into the dynamic engine.""" + + def _initialize_disaggregation_state(self) -> None: + """Initialize state without importing or constructing a transfer backend.""" + + self._pinned_handoff_blocks: Dict[int, list] = {} # Request ID -> pinned KV block IDs. + self._kv_transfer_agent = None + self._kv_peer_metas = None # KV descriptors for this PP stage's TP ranks. + self._pp_kv_peer_metas = None # Each PP stage's set of TP KV descriptors. + self._deferred_kv_handoffs = deque() + self._pending_kv_imports = deque() + self._quarantined_kv_imports: list[PendingKvImport] = [] + self._handoff_completion_tracker: HandoffCompletionTracker | None = None + self._handoff_completion_notifications: dict[int, bool] = {} # Request ID -> failed. + self._pending_kv_pushes: list = [] + self._kv_transfer_role: str | None = None + + def _setup_handoff_completion_tracking(self, hostname: str | None = None) -> None: + """Create the CPU path used to aggregate model-parallel transfer completion.""" + + self._handoff_completion_tracker = HandoffCompletionTracker( + self.zmq_context, self.pg_collection.mp, hostname + ) + self.zmq_sockets.extend(self._handoff_completion_tracker.sockets) + + def _drain_handoff_completion_notifications(self) -> list[tuple[int, bool]]: + """Collect decisions that the existing MP schedule broadcast must distribute. + + Side: decode MP coordinator; pull and push transport paths. + """ + + if self._handoff_completion_tracker is None: + return [] + return self._handoff_completion_tracker.drain_completed() + + def _record_handoff_completion_notification(self, request_id: int, failed: bool) -> None: + """Record the coordinator's shared admission decision on every MP rank. + + Side: decode engine; pull and push transport paths. + """ + + self._handoff_completion_notifications[request_id] = failed + + @property + def pending_kv_import_count(self) -> int: + """Number of decode requests waiting for capacity or transfer completion. + + Side: decode engine; pull and push transport paths. + """ + + return len(self._deferred_kv_handoffs) + len(self._pending_kv_imports) + + @property + def has_admittable_kv_import(self) -> bool: + """Whether a completed import can be handled at the next admission point.""" + + for pending in self._pending_kv_imports: + failed = self._handoff_completion_notifications.get(pending.request_id) + if failed is None: + continue + if failed or can_admit_prefilled_decode(self.context, len(pending.resume_tokens)): + return True + return False + + @property + def pending_kv_push_count(self) -> int: + """Number of prefill sends waiting for transport completion. + + Side: prefill engine; push transport path only. + """ + + return len(self._pending_kv_pushes) + + def _reset_pending_kv_imports(self) -> None: + """Drain and release pending handoff transfers before an engine reset.""" + + unsafe_pushes = [] + if not hasattr(self, "_pending_kv_pushes"): + self._pending_kv_pushes = [] + for request_id, handles in self._pending_kv_pushes: + if not self._wait_for_transfer_handles(*handles): + unsafe_pushes.append((request_id, handles)) + self._pending_kv_pushes = unsafe_pushes + + if not hasattr(self, "_pending_kv_imports"): + self._pending_kv_imports = deque() + if not hasattr(self, "_deferred_kv_handoffs"): + self._deferred_kv_handoffs = deque() + while self._deferred_kv_handoffs: + handoff = self._deferred_kv_handoffs.popleft() + if not handoff.future.done(): + handoff.future.cancel() + unsafe = deque() + while self._pending_kv_imports: + pending = self._pending_kv_imports.popleft() + safe_to_release = pending.destinations_safe and self._wait_for_transfer_handles( + *self._pending_transfer_handles(pending) + ) + if safe_to_release: + self._release_pending_kv_import(pending) + if not pending.future.done(): + pending.future.cancel() + else: + unsafe.append(pending) + self._pending_kv_imports = unsafe + quarantined = [] + for pending in self._quarantined_kv_imports: + safe_to_release = pending.destinations_safe and self._wait_for_transfer_handles( + *self._pending_transfer_handles(pending) + ) + if safe_to_release: + self._release_pending_kv_import(pending) + else: + quarantined.append(pending) + self._quarantined_kv_imports = quarantined + if unsafe_pushes or unsafe or quarantined: + raise RuntimeError( + "Cannot reset while KV handoff transfers may still access cache storage" + ) + if self._pinned_handoff_blocks: + raise RuntimeError( + "Cannot reset while KV handoff blocks remain pinned; wait for RELEASE_KV" + ) + self._handoff_completion_notifications.clear() + + def schedule_waiting_requests(self) -> None: + """Reject prompt scheduling on a dedicated disaggregated decode engine. + + Side: decode engine; pull and push transport paths. + """ + + if self._kv_transfer_role == "decode": + # Async scheduling reaches this point after consuming the previous + # forward's logits. The legacy path does not retain logits across steps. + self._admit_pending_kv_imports() + if self.waiting_request_ids: + raise RuntimeError( + "A disaggregated decode engine cannot schedule prompt prefill requests" + ) + super().schedule_waiting_requests() + + def setup_kv_transfer(self, role: str, backend: str = "nixl") -> None: + """Bring up the KV transfer agents for this engine. + + This method must be called collectively by every model-parallel rank in + the engine because each rank participates in PP and TP metadata gathers. + + Args: + role: "prefill" or "decode"; used to name the local transfer agent. + backend: transfer backend name, resolved through the explicit + registry ("nixl"; "nccl" selects the two-sided push family). + """ + if role not in ("prefill", "decode"): + raise ValueError(f"KV transfer role must be 'prefill' or 'decode', got {role!r}") + if self.context.is_hybrid_model: + raise RuntimeError( + "Hybrid models require recurrent-state handoff in addition to KV-cache handoff" + ) + self._kv_transfer_role = role + backend_cls = construct_kv_transfer_backend_class(backend) + + # Prefill output blocks stay pinned until the peer finishes reading + # them. Decode requests consume imports but do not produce another + # handoff, so pinning their completed blocks would retain cache state + # without a downstream release acknowledgement. + allocator = self.context.kv_block_allocator + assert allocator.enable_prefix_caching, ( + "KV handoff requires prefix caching on both prefill and decode " + "engines (--inference-dynamic-batching-prefix-caching)." + ) + allocator.enable_handoff_pinning = role == "prefill" + + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + + # TP topology, so a peer at a different TP can re-shard our KV heads. + # KV heads under GQA == num_query_groups (falls back to attention heads). + model_config = self.controller.inference_wrapped_model.model.config + num_kv_heads_global = model_config.num_query_groups or model_config.num_attention_heads + tp_size = get_pg_size(self.pg_collection.tp) + tp_rank = get_pg_rank(self.pg_collection.tp) + + # Compute this PP rank's global attention-layer range. + pp_size = get_pg_size(self.pg_collection.pp) + pp_rank = get_pg_rank(self.pg_collection.pp) + local_num_layers = self.context.num_attention_layers + + if pp_size > 1 and torch.distributed.is_initialized(): + layer_counts: list = [None] * pp_size + torch.distributed.all_gather_object( + layer_counts, local_num_layers, group=self.pg_collection.pp + ) + layer_start = sum(layer_counts[:pp_rank]) + num_layers_global = sum(layer_counts) + else: + layer_start = 0 + num_layers_global = local_num_layers + layer_end = layer_start + local_num_layers + + self._kv_transfer_agent = backend_cls( + agent_name=f"{role}-rank{rank}", + memory_buffer=self.context.memory_buffer, + expected_num_blocks=self.context.kv_block_allocator.pool_size, + tp_size=tp_size, + tp_rank=tp_rank, + num_kv_heads_global=num_kv_heads_global, + heads_per_partition=self.context.num_attention_heads_per_partition, + head_dim=self.context.hidden_size_per_attention_head, + tokens_per_block=self.context.block_size_tokens, + global_rank=rank, + pp_size=pp_size, + pp_rank=pp_rank, + num_layers_global=num_layers_global, + layer_start=layer_start, + layer_end=layer_end, + ) + + # Transport descriptors do not change between requests. Collect each + # stage's TP descriptors once and reuse them for every handoff. + self._kv_peer_metas = self._kv_transfer_agent.export_meta() + if torch.distributed.is_initialized() and tp_size > 1: + gathered_tp_metas: list = [None] * tp_size + torch.distributed.all_gather_object( + gathered_tp_metas, self._kv_peer_metas, group=self.pg_collection.tp + ) + self._kv_peer_metas = gathered_tp_metas + + self._pp_kv_peer_metas = [self._kv_peer_metas] + if torch.distributed.is_initialized() and pp_size > 1: + self._pp_kv_peer_metas = [None] * pp_size + torch.distributed.all_gather_object( + self._pp_kv_peer_metas, self._kv_peer_metas, group=self.pg_collection.pp + ) + + def push_handoff_kv(self, request_id: int, decode_metas: list) -> None: + """Push a pinned hand-off's KV to the decode + instance described by `decode_metas` (two-sided transports only). + + The decode posted its matching receives when SUBMIT_REQUEST_WITH_KV + arrived; the sends are reaped asynchronously and the pins stay until + the coordinator's RELEASE_KV. + + Side: prefill engine; push transport path only. + """ + block_ids = self._pinned_handoff_blocks.get(request_id) + if not block_ids: + logging.warning( + "SEND_KV for request %d with no pinned hand-off blocks; skipping", request_id + ) + return + kv_peer = {"tp_metas": list(decode_metas)} + handles = [self._kv_transfer_agent.begin_push_blocks(kv_peer, block_ids)] + self._pending_kv_pushes.append((request_id, handles)) + logging.info("DISAGG_PREFILL_PUSH request_id=%d blocks=%d", request_id, len(block_ids)) + + def _poll_pending_kv_pushes(self) -> int: + """Reap completed push sends; unfinished ones stay pending. + + Side: prefill engine; push transport path only. + """ + if not self._pending_kv_pushes: + return 0 + remaining = [] + reaped = 0 + for request_id, handles in self._pending_kv_pushes: + if all(h.poll() for h in handles): + reaped += 1 + else: + remaining.append((request_id, handles)) + self._pending_kv_pushes = remaining + return reaped + + def _prepare_handoff_metadata_batch( + self, + requests_and_blocks: list[tuple["DynamicInferenceRequest", list[int]]], + decode_tokens_by_request: Dict[int, list[int]], + ) -> dict[int, _PreparedHandoffMetadata]: + """Assemble metadata for all handoffs completed by one engine step. + + Side: prefill engine; pull and push transport paths. + """ + + handoffs = [ + (request, list(block_ids)) + for request, block_ids in requests_and_blocks + if request.sampling_params.do_kv_handoff + ] + if not handoffs: + return {} + if self._kv_peer_metas is None: + raise RuntimeError("KV handoff requested before transfer setup") + + candidate_blocks_by_request = {} + for request, block_ids in handoffs: + prompt_block_count = ( + len(request.prompt_tokens) + self.context.block_size_tokens - 1 + ) // self.context.block_size_tokens + candidate_blocks_by_request[request.request_id] = block_ids[:prompt_block_count] + + pp_size = get_pg_size(self.pg_collection.pp) + try: + static_pp_metas = self._pp_kv_peer_metas or [self._kv_peer_metas] + if len(static_pp_metas) != pp_size: + raise RuntimeError( + f"Expected static metadata for {pp_size} pipeline stages, " + f"got {len(static_pp_metas)}" + ) + + prepared = {} + for request, _ in handoffs: + candidate_blocks = candidate_blocks_by_request[request.request_id] + resume_tokens = list(decode_tokens_by_request.get(request.request_id, [])) + expected_resume_tokens = self.context.num_speculative_tokens + 1 + if len(resume_tokens) != expected_resume_tokens: + raise RuntimeError( + "Cannot create a decode-only handoff without one sampled token plus " + f"the configured MTP proposals: expected {expected_resume_tokens}, " + f"got {len(resume_tokens)}" + ) + if not candidate_blocks: + raise RuntimeError( + "Cannot create a decode-only handoff without the prompt KV state" + ) + block_ids = candidate_blocks + if pp_size > 1: + # Model-parallel ranks receive the same request stream and + # use synchronized KV cache capacities. + # Their deterministic allocators therefore assign the same + # physical IDs, while transport descriptors remain stage- and + # rank-specific metadata collected once during setup. + kv_meta: Any = { + "pp_metas": [ + {"tp_metas": stage_meta, "block_ids": block_ids} + for stage_meta in static_pp_metas + ] + } + else: + kv_meta = self._kv_peer_metas + + prepared[request.request_id] = _PreparedHandoffMetadata( + local_blocks=block_ids, kv_meta=kv_meta, resume_tokens=resume_tokens + ) + return prepared + except Exception: + for block_ids in candidate_blocks_by_request.values(): + self._release_pinned_handoff_blocks(block_ids) + raise + + def _capture_handoff_meta( + self, request: "DynamicInferenceRequest", prepared: _PreparedHandoffMetadata | None + ) -> None: + """Attach prepared transfer metadata and retain the request's blocks. + + Side: prefill engine; pull and push transport paths. + """ + + if prepared is None: + raise RuntimeError( + f"No handoff metadata was prepared for completed request {request.request_id}" + ) + + rid = request.request_id + block_ids = prepared.local_blocks + kv_meta = prepared.kv_meta + + self._pinned_handoff_blocks[rid] = list(block_ids) + + if isinstance(kv_meta, list): + kv_meta = {"tp_metas": kv_meta} + else: + # TP=1 caches one static metadata dictionary for the engine. Keep + # request-specific fields out of that shared object. + kv_meta = dict(kv_meta) + kv_meta["resume_tokens"] = prepared.resume_tokens + + request.disaggregated_params = { + "request_id": rid, + "block_ids": block_ids, + "kv_meta": kv_meta, + } + logging.info("DISAGG_PREFILL_HANDOFF request_id=%d pinned_blocks=%d", rid, len(block_ids)) + + def release_handoff_blocks(self, request_id: int) -> None: + """Release blocks pinned by a previous do_kv_handoff completion. + + Side: prefill engine; pull and push transport paths. + """ + block_ids = self._pinned_handoff_blocks.pop(request_id, None) + if not block_ids: + return + released = self._release_pinned_handoff_blocks(block_ids) + logging.info( + "DISAGG_PREFILL_RELEASE request_id=%d released_blocks=%d", request_id, released + ) + + def _release_pinned_handoff_blocks(self, block_ids: list) -> int: + """Release this request's ownership of its pinned handoff blocks. + + Side: prefill engine; pull and push transport paths. + """ + if not block_ids: + return 0 + allocator = self.context.kv_block_allocator + allocator.release_memory_blocks(torch.tensor(block_ids, dtype=torch.int32, device="cpu")) + return len(block_ids) + + def add_request_with_kv_handoff( + self, + request_id: int, + prompt: list, + sampling_params: "SamplingParams", + kv_meta: dict, + src_block_ids: list, + ) -> "asyncio.Future[DynamicInferenceRequest]": + """Start or capacity-queue a KV import and return its completion future. + + Side: decode engine; pull and push transport paths. A pull backend starts + the read here, while a push backend posts the matching receive. + """ + allocator = self.context.kv_block_allocator + if not allocator.enable_prefix_caching: + raise RuntimeError( + "add_request_with_kv_handoff requires " + "--inference-dynamic-batching-prefix-caching on the decode engine; " + "the prefill-skip path uses the prefix-cache match logic." + ) + + if self._kv_transfer_agent is None: + raise RuntimeError("KV handoff received without a transfer backend") + + resume_tokens = ( + [int(token) for token in kv_meta.get("resume_tokens", [])] + if isinstance(kv_meta, dict) + else [] + ) + expected_resume_tokens = self.context.num_speculative_tokens + 1 + if len(resume_tokens) != expected_resume_tokens: + raise RuntimeError( + "Decode-only handoff requires one sampled token plus the configured " + f"MTP proposals: expected {expected_resume_tokens}, got {len(resume_tokens)}" + ) + if sampling_params.return_log_probs or sampling_params.top_n_logprobs > 0: + raise NotImplementedError( + "Decode-only handoff does not yet transfer prompt or first-token log probabilities" + ) + + prompt_tensor = torch.tensor(prompt, dtype=torch.int64) + hashes = compute_block_hashes_batched(prompt_tensor, self.context.block_size_tokens) + num_blocks = transfer_block_count(kv_meta, src_block_ids) + expected_blocks = ( + len(prompt) + self.context.block_size_tokens - 1 + ) // self.context.block_size_tokens + if num_blocks != expected_blocks: + raise RuntimeError( + "Decode-only handoff requires every prompt KV block, including the partial tail: " + f"expected {expected_blocks}, got {num_blocks}" + ) + future = self._loop.create_future() + handoff = DeferredKvHandoff( + request_id=request_id, + prompt=prompt, + sampling_params=sampling_params, + kv_meta=kv_meta, + src_block_ids=list(src_block_ids), + hashes=hashes, + num_blocks=num_blocks, + future=future, + ) + # Preserve receive order under backpressure. This is required by NCCL's + # two-sided transport and also prevents a stream of small handoffs from + # starving an older, larger one. + if self._deferred_kv_handoffs: + self._deferred_kv_handoffs.append(handoff) + logging.debug( + "DISAGG_DECODE_CAPACITY_QUEUE request_id=%d queued=%d", + request_id, + len(self._deferred_kv_handoffs), + ) + return future + + started = self._try_start_kv_handoff_import(handoff) + if not started: + self._deferred_kv_handoffs.append(handoff) + logging.debug( + "DISAGG_DECODE_CAPACITY_QUEUE request_id=%d queued=%d", + request_id, + len(self._deferred_kv_handoffs), + ) + return future + + def _try_start_kv_handoff_import(self, handoff: DeferredKvHandoff) -> bool: + """Reserve cache state and start one import, or return without mutation. + + Side: decode engine; pull and push transport paths. + """ + + allocator = self.context.kv_block_allocator + cached_blocks = [] + if not self._kv_transfer_agent.is_push: + cached_blocks = self._find_cached_handoff_prefix(handoff.hashes, handoff.num_blocks) + + num_blocks_to_import = handoff.num_blocks - len(cached_blocks) + resume_tokens = ( + [int(token) for token in handoff.kv_meta.get("resume_tokens", [])] + if isinstance(handoff.kv_meta, dict) + else [] + ) + continuation_block_count = ( + additional_decode_blocks( + len(handoff.prompt), len(resume_tokens), self.context.block_size_tokens + ) + if resume_tokens + else 0 + ) + if not self._handoff_capacity_available( + num_blocks_to_import + continuation_block_count, cached_blocks + ): + return False + + allocator.retain_memory_blocks(cached_blocks) + allocated_blocks_tensor = allocator.allocate_memory_blocks( + num_blocks_to_import + continuation_block_count + ) + if allocated_blocks_tensor is None: + if cached_blocks: + allocator.release_memory_blocks( + torch.tensor(cached_blocks, dtype=torch.int32, device="cpu") + ) + raise RuntimeError("KV allocator capacity changed during handoff admission") + allocated_blocks = [int(block) for block in allocated_blocks_tensor.tolist()] + imported_blocks = allocated_blocks[:num_blocks_to_import] + continuation_blocks = allocated_blocks[num_blocks_to_import:] + local_blocks = cached_blocks + imported_blocks + handle = None + start_error = None + try: + transfer_meta, transfer_src_blocks = drop_transfer_prefix_blocks( + handoff.kv_meta, handoff.src_block_ids, len(cached_blocks) + ) + handle = self._kv_transfer_agent.begin_pull_blocks( + transfer_meta, transfer_src_blocks, imported_blocks + ) + except Exception as exc: + start_error = exc + + pending = PendingKvImport( + request_id=handoff.request_id, + prompt=handoff.prompt, + sampling_params=handoff.sampling_params, + local_blocks=local_blocks, + hashes=handoff.hashes, + cached_prefix_block_count=len(cached_blocks), + handle=handle, + future=handoff.future, + resume_tokens=resume_tokens, + continuation_blocks=continuation_blocks, + local_error=start_error, + destinations_safe=( + start_error is None or getattr(start_error, "transfer_destinations_safe", True) + ), + ) + self._pending_kv_imports.append(pending) + logging.debug( + "DISAGG_DECODE_PULL_SUBMIT request_id=%d prompt_tokens=%d " + "cached_blocks=%d imported_blocks=%d pending_imports=%d", + handoff.request_id, + len(handoff.prompt), + len(cached_blocks), + len(imported_blocks), + len(self._pending_kv_imports), + ) + self._loop.call_soon_threadsafe(self._loop.create_task, self._notify_cond_for_new_request()) + return True + + def _find_cached_handoff_prefix(self, hashes: list[int], num_blocks: int) -> list[int]: + """Find the contiguous handoff prefix already cached on decode. + + Side: decode engine; pull transport path only. + """ + + allocator = self.context.kv_block_allocator + cached_blocks = [] + for block_hash in hashes[:num_blocks]: + block_id = allocator.kv_hash_to_block_id.get(block_hash) + if block_id is None: + break + cached_blocks.append(int(block_id)) + return cached_blocks + + def _handoff_capacity_available(self, num_blocks: int, cached_blocks: list[int]) -> bool: + """Check capacity in the rank-local mirror of model-parallel allocator state. + + Side: decode engine; pull and push transport paths. + """ + + allocator = self.context.kv_block_allocator + potential_matched_count = 0 + if cached_blocks: + block_tensor = torch.tensor(cached_blocks, dtype=torch.int32, device="cpu") + potential_matched_count = int((allocator.block_ref_counts[block_tensor] == 0).sum()) + # Model-parallel ranks process the same request and allocator operations + # in lockstep, so capacity must remain mirrored without a per-request collective. + return allocator.is_memory_available( + num_blocks, potential_matched_count=potential_matched_count + ) + + def _drain_deferred_kv_handoffs(self) -> int: + """Start queued handoffs in FIFO order while the queue head fits. + + Side: decode engine; pull and push transport paths. + """ + + started_count = 0 + while self._deferred_kv_handoffs: + handoff = self._deferred_kv_handoffs[0] + started = self._try_start_kv_handoff_import(handoff) + if not started: + break + self._deferred_kv_handoffs.popleft() + started_count += 1 + logging.debug( + "DISAGG_DECODE_CAPACITY_ADMIT request_id=%d queued=%d", + handoff.request_id, + len(self._deferred_kv_handoffs), + ) + return started_count + + @staticmethod + def _pending_transfer_handles(pending: PendingKvImport) -> list: + """Return this decode import's active KV transfer handles. + + Side: decode engine; pull and push transport paths. + """ + + return [pending.handle] if pending.handle is not None else [] + + def _validate_decode_ready_handoff(self, pending: PendingKvImport) -> None: + """Validate that transferred state can start decode without prompt execution.""" + + expected_tokens = self.context.num_speculative_tokens + 1 + if len(pending.resume_tokens) != expected_tokens: + raise RuntimeError( + "Decode-only handoff is missing its sampled token or MTP proposals: " + f"expected {expected_tokens}, got {len(pending.resume_tokens)}" + ) + + def _finalize_kv_handoff_import(self, pending: PendingKvImport) -> None: + """Register transferred blocks and admit the decode request. + + Side: decode engine; pull and push transport paths. + """ + + allocator = self.context.kv_block_allocator + local_blocks = pending.local_blocks + cached_prefix_block_count = pending.cached_prefix_block_count + registration_end = min(len(local_blocks), len(pending.hashes)) + num_hashes_to_register = registration_end - cached_prefix_block_count + + if num_hashes_to_register > 0: + # The imported suffix extends any retained local prefix. Preserve + # that predecessor link in the allocator's parent-aware LRU forest. + parent_hashes = [ + pending.hashes[block_idx - 1] if block_idx > 0 else 0 + for block_idx in range(cached_prefix_block_count, registration_end) + ] + allocator.register_kv_block_hashes( + local_blocks[cached_prefix_block_count:registration_end], + pending.hashes[cached_prefix_block_count:registration_end], + parent_hashes=parent_hashes, + ) + + logging.debug( + "DISAGG_DECODE_IMPORT request_id=%d prompt_tokens=%d " + "cached_blocks=%d imported_blocks=%d hashes_registered=%d pending_imports=%d", + pending.request_id, + len(pending.prompt), + cached_prefix_block_count, + len(local_blocks) - cached_prefix_block_count, + num_hashes_to_register, + len(self._pending_kv_imports), + ) + + request_future = self.add_request( + pending.request_id, + pending.prompt, + pending.sampling_params, + precomputed_block_hashes=( + pending.hashes[:registration_end] if registration_end > 0 else None + ), + ) + + request = self.get_request(pending.request_id) + if request.status == Status.FAILED: + # add_request() reports validation failures through its own future. + self._release_pending_kv_import(pending) + else: + queued_request_id = self.waiting_request_ids[0] + if queued_request_id != pending.request_id: + raise RuntimeError( + "Decode-only admission expected the imported request at the waiting queue " + f"head: expected {pending.request_id}, got {queued_request_id}" + ) + self.waiting_request_ids.popleft() + self._validate_decode_ready_handoff(pending) + first_token = pending.resume_tokens[0] + if request.sampling_params.num_tokens_to_generate > 0: + request.generated_tokens.append(first_token) + if self.track_generated_token_events: + first_token_event = request.add_event_generated_token(first_token) + else: + first_token_event = DynamicInferenceEvent( + type=DynamicInferenceEventType.GENERATED_TOKEN, + payload={"token_id": first_token}, + ) + request.ttft = first_token_event.timestamp - request.event_add_engine.timestamp + + stop_word_hit = False + if request.stop_word_ids: + stop_word_hit, _ = self._check_stop_words_for_request_post_append(request) + if first_token == request.sampling_params.termination_id or stop_word_hit: + request.sampling_params.num_tokens_to_generate = len(request.generated_tokens) + + request.num_cached_tokens = len(pending.prompt) + if len(request.generated_tokens) >= request.sampling_params.num_tokens_to_generate: + self._release_pending_kv_import(pending) + self._complete_handoff_request_without_forward(pending.request_id) + else: + admit_prefilled_decode( + self.context, + request, + local_blocks, + pending.continuation_blocks, + pending.resume_tokens, + ) + pending.local_blocks = [] + pending.continuation_blocks = [] + if self.use_coordinator and self.is_mp_coordinator: + self._try_send_streaming_partials() + + def _relay_result(src: asyncio.Future) -> None: + """Forward decode completion to the handoff future.""" + + if pending.future.done(): + return + if src.cancelled(): + pending.future.cancel() + return + exc = src.exception() + if exc is not None: + pending.future.set_exception(exc) + else: + pending.future.set_result(src.result()) + + request_future.add_done_callback(_relay_result) + + def _complete_handoff_request_without_forward(self, request_id: int) -> None: + """Complete an imported request whose transferred token already satisfies it.""" + + request_entry = self.requests.pop(request_id) + request = request_entry.record[-1] + request.generated_length = len(request.generated_tokens) + request.status = Status.COMPLETED + request.add_event_finish() + self.finished_request_count += 1 + + if self.use_coordinator and self.is_mp_coordinator: + self._send_request_records_to_coordinator([request_entry.record]) + self._partial_emit_lengths.pop(request_id, None) + elif not self.use_coordinator: + if request.prompt is None: + request.prompt = self.controller.detokenize( + self.controller.tokenizer, request.prompt_tokens.tolist(), remove_EOD=False + ) + request.generated_text = self.controller.detokenize( + self.controller.tokenizer, + request.generated_tokens, + remove_EOD=not request.sampling_params.detokenize_stop_sequence, + ) + + request_entry.future.set_result(request_entry.record) + + def _release_pending_kv_import(self, pending: PendingKvImport) -> None: + """Release storage owned by an unadmitted decode import. + + Side: decode engine; pull and push transport paths. + """ + + owned_blocks = pending.local_blocks + pending.continuation_blocks + if owned_blocks: + block_tensor = torch.tensor(owned_blocks, dtype=torch.int32, device="cpu") + self.context.kv_block_allocator.release_memory_blocks(block_tensor) + pending.local_blocks = [] + pending.continuation_blocks = [] + + @staticmethod + def _wait_for_transfer_handles(*handles) -> bool: + """Wait for known handles; return false if any may still be active.""" + + safe_to_release = True + for handle in handles: + if handle is None: + continue + try: + handle.wait() + except TimeoutError: + safe_to_release = False + except Exception: + # NIXL reports transfer errors only after all segments belonging + # to the handle have reached a terminal state. + pass + return safe_to_release + + def _report_completed_kv_imports(self) -> None: + """Report locally terminal imports without synchronizing the compute ranks. + + Side: decode engine; pull and push transport paths. + """ + + for pending in self._pending_kv_imports: + request_id = pending.request_id + if pending.terminal_state_reported: + continue + failed = pending.local_error is not None + if not failed: + try: + if not all(handle.poll() for handle in self._pending_transfer_handles(pending)): + continue + except Exception as exc: + pending.local_error = exc + failed = True + + if ( + self._handoff_completion_tracker is None + or self._handoff_completion_tracker.world_size == 1 + ): + if get_pg_size(self.pg_collection.mp) != 1: + raise RuntimeError( + "Model-parallel KV handoff requires coordinator completion tracking" + ) + self._handoff_completion_notifications[request_id] = failed + else: + self._handoff_completion_tracker.report(request_id, failed) + pending.terminal_state_reported = True + + def _poll_pending_kv_imports(self) -> int: + """Progress imports without mutating the active decode batch. + + Side: decode engine; pull and push transport paths. + """ + + self._drain_deferred_kv_handoffs() + if not self._pending_kv_imports: + return 0 + self._report_completed_kv_imports() + return int(self.has_admittable_kv_import) + + def _admit_pending_kv_imports(self) -> int: + """Handle completed imports at an engine scheduling admission point. + + Side: decode engine; pull and push transport paths. + """ + + ready = 0 + remaining = deque() + while self._pending_kv_imports: + pending = self._pending_kv_imports.popleft() + failed = self._handoff_completion_notifications.get(pending.request_id) + if failed is None: + remaining.append(pending) + continue + try: + if failed: + raise pending.local_error or RuntimeError( + "KV handoff transfer failed on a model-parallel peer" + ) + self._validate_decode_ready_handoff(pending) + if not can_admit_prefilled_decode(self.context, len(pending.resume_tokens)): + remaining.append(pending) + continue + self._handoff_completion_notifications.pop(pending.request_id) + self._finalize_kv_handoff_import(pending) + ready += 1 + except Exception as exc: + self._handoff_completion_notifications.pop(pending.request_id, None) + # A peer can fail before posting its half of a two-sided transfer. + # Do not wait on this rank's unmatched operation; quarantine its + # destination unless its own handle already reached a terminal state. + safe_to_release = pending.destinations_safe and pending.terminal_state_reported + safe_to_release = safe_to_release and self._wait_for_transfer_handles( + *self._pending_transfer_handles(pending) + ) + if safe_to_release: + self._release_pending_kv_import(pending) + else: + self._quarantined_kv_imports.append(pending) + logging.error( + "Quarantining request %d cache storage after an incomplete handoff", + pending.request_id, + ) + if not pending.future.done(): + pending.future.set_exception(exc) + logging.exception("DISAGG_DECODE_PULL_FAILED request_id=%d", pending.request_id) + if failed: + continue + remaining.extend(self._pending_kv_imports) + self._pending_kv_imports = remaining + raise + self._pending_kv_imports = remaining + if ready: + self._loop.call_soon_threadsafe( + self._loop.create_task, self._notify_cond_for_new_request() + ) + return ready diff --git a/megatron/core/inference/disaggregation/pending_handoff_imports.py b/megatron/core/inference/disaggregation/pending_handoff_imports.py new file mode 100644 index 00000000000..d2fbe12663f --- /dev/null +++ b/megatron/core/inference/disaggregation/pending_handoff_imports.py @@ -0,0 +1,44 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""State records for KV-cache imports awaiting completion.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import Any, List + +from megatron.core.inference.sampling_params import SamplingParams + + +@dataclass(kw_only=True) +class DeferredKvHandoff: + """Decode handoff waiting for local cache capacity before transfer starts.""" + + request_id: int + prompt: list + sampling_params: SamplingParams + kv_meta: dict + src_block_ids: List[int] + hashes: List[int] + num_blocks: int + future: asyncio.Future + + +@dataclass(kw_only=True) +class PendingKvImport: + """Decode request waiting for an asynchronous KV-cache import.""" + + request_id: int + prompt: list + sampling_params: SamplingParams + local_blocks: List[int] + hashes: List[int] + cached_prefix_block_count: int + handle: Any + future: asyncio.Future + resume_tokens: List[int] = field(default_factory=list) # Sampled token, then MTP proposals. + continuation_blocks: List[int] = field(default_factory=list) # Empty KV for resume writes. + local_error: Exception | None = None # Exact local error, if this rank failed. + destinations_safe: bool = True # Whether allocated blocks may return to the pool. + terminal_state_reported: bool = False # Avoid repeat reports and premature block reuse. diff --git a/megatron/core/inference/disaggregation/transfer_backends/nixl.py b/megatron/core/inference/disaggregation/transfer_backends/nixl.py index 2f2bdd6b6a7..36da14a4654 100644 --- a/megatron/core/inference/disaggregation/transfer_backends/nixl.py +++ b/megatron/core/inference/disaggregation/transfer_backends/nixl.py @@ -213,6 +213,7 @@ class NixlTransferBackend: """ name = "nixl" + is_push = False def __init__( self, diff --git a/megatron/core/inference/disaggregation/utils.py b/megatron/core/inference/disaggregation/utils.py index 779207e40f6..8e17d4e9096 100644 --- a/megatron/core/inference/disaggregation/utils.py +++ b/megatron/core/inference/disaggregation/utils.py @@ -55,3 +55,39 @@ def transfer_block_count(peer_meta: Any, src_block_ids: List[int]) -> int: records = transfer_peer_records(peer_meta, src_block_ids) return len(records[0][1]) if records else 0 + + +def drop_transfer_prefix_blocks( + peer_meta: Any, src_block_ids: List[int], prefix_count: int +) -> Tuple[Any, List[int]]: + """Return transfer metadata with a sequence prefix removed. + + Used when the decode engine reuses cached prefix blocks and only pulls the uncached suffix. + """ + + if prefix_count == 0: + return peer_meta, list(src_block_ids) + if prefix_count < 0: + raise ValueError("transfer prefix count must be non-negative") + + def trim_metadata(value: Any) -> Any: + if isinstance(value, list): + return [trim_metadata(entry) for entry in value] + if not isinstance(value, dict): + return value + + trimmed = dict(value) + if "block_ids" in trimmed: + block_ids = list(trimmed["block_ids"]) + if prefix_count > len(block_ids): + raise ValueError( + f"cannot remove {prefix_count} blocks from a {len(block_ids)}-block peer" + ) + trimmed["block_ids"] = block_ids[prefix_count:] + for key in ("pp_metas", "tp_metas"): + if key in trimmed: + trimmed[key] = trim_metadata(trimmed[key]) + return trimmed + + trimmed_src_blocks = list(src_block_ids[prefix_count:]) + return trim_metadata(peer_meta), trimmed_src_blocks diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 22fe74525e9..de6b342c89a 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -302,6 +302,7 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen # Throw a cudagraph-admission warning if deferred for > max_sequence_length steps. # The floor value of 100 avoids warnings in test configs where max_sequence_length < 100. self._cg_admission_warn_after = max(100, self.context.max_sequence_length) + self._initialize_disaggregation_state() # Initialize engine. self.reset() @@ -341,9 +342,90 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen # Create cuda graphs. self.create_cuda_graphs() + def _initialize_disaggregation_state(self) -> None: + """Hook overridden by the KV-handoff engine composition.""" + + def _reset_pending_kv_imports(self) -> None: + """Hook overridden by the KV-handoff engine composition.""" + + @property + def pending_kv_import_count(self) -> int: + """Number of decode requests awaiting a KV import (none here).""" + return 0 + + @property + def has_admittable_kv_import(self) -> bool: + """Whether a completed KV import is eligible for admission (false here).""" + return False + + def _poll_pending_kv_imports(self) -> int: + return 0 + + def _admit_pending_kv_imports(self) -> int: + return 0 + + def _setup_handoff_completion_tracking(self, hostname: str | None = None) -> None: + """Hook overridden by the KV-handoff engine composition.""" + + def _drain_handoff_completion_notifications(self) -> list[tuple[int, bool]]: + """Hook overridden by the KV-handoff engine composition.""" + return [] + + def _record_handoff_completion_notification(self, request_id: int, failed: bool) -> None: + """Hook overridden by the KV-handoff engine composition.""" + self._raise_kv_handoff_not_enabled("KV handoff completion notification") + + def _prepare_handoff_metadata_batch( + self, requests_and_blocks: list[tuple], decode_tokens_by_request: Dict[int, list[int]] + ) -> dict: + """Hook overridden by the KV-handoff engine composition.""" + if any(request.sampling_params.do_kv_handoff for request, _ in requests_and_blocks): + self._raise_kv_handoff_not_enabled("KV handoff completion") + return {} + + def _capture_handoff_meta(self, request, prepared) -> None: + self._raise_kv_handoff_not_enabled("KV handoff completion") + + def _release_pinned_handoff_blocks(self, block_ids: list) -> int: + return 0 + + def setup_kv_transfer(self, role: str, backend: str = "nixl") -> None: + """Raising stub; the hand-off engine composition overrides it.""" + self._raise_kv_handoff_not_enabled("KV transfer setup") + + def push_handoff_kv(self, request_id: int, decode_metas: list) -> None: + """Raising stub; the hand-off engine composition overrides it.""" + self._raise_kv_handoff_not_enabled("SEND_KV") + + def _poll_pending_kv_pushes(self) -> int: + return 0 + + @property + def pending_kv_push_count(self) -> int: + """Number of prefill sends awaiting completion (none here).""" + return 0 + + def add_request_with_kv_handoff( + self, request_id, prompt, sampling_params, kv_meta, src_block_ids + ) -> asyncio.Future[DynamicInferenceRequest]: + """Raising stub; the hand-off engine composition overrides it.""" + self._raise_kv_handoff_not_enabled("SUBMIT_REQUEST_WITH_KV") + + def release_handoff_blocks(self, request_id: int) -> None: + """Raising stub; the hand-off engine composition overrides it.""" + self._raise_kv_handoff_not_enabled("RELEASE_KV") + + @staticmethod + def _raise_kv_handoff_not_enabled(operation: str) -> None: + raise RuntimeError( + f"{operation} requires KV handoff, but it is not enabled. " + "Use DisaggDynamicInferenceEngine with KV transfer configured." + ) + def reset(self) -> None: """Reset by removing all requests and reset all state.""" + self._reset_pending_kv_imports() self.context.reset() # Request state. @@ -353,6 +435,8 @@ def reset(self) -> None: self.requests: Dict[int, RequestEntry] = {} self.waiting_request_ids = deque() + if hasattr(self, "_pinned_handoff_blocks"): + self._pinned_handoff_blocks.clear() self.failed_request_ids = [] # Generated token count already streamed for each request. self._partial_emit_lengths: Dict[int, int] = {} @@ -747,6 +831,8 @@ async def start_listening_to_data_parallel_coordinator( self.zmq_sockets += [self.model_parallel_subscriber_socket] + self._setup_handoff_completion_tracking(hostname) + torch.distributed.barrier(mp_group) # initialize zmq-based EP communicator @@ -978,6 +1064,28 @@ async def _notify_cond_for_new_request(self): async with self._cond: self._cond.notify_all() + def _send_request_records_to_coordinator( + self, records: List[DynamicInferenceRequestRecord] + ) -> None: + """Send completed or failed request records from the MP coordinator.""" + + merged_requests = [record.merge() for record in records] + if self.local_metadata_ledger_enabled: + # Failed requests are sent immediately but remain in the engine until the + # next bookkeeping pass. Index only completed requests as they are dropped. + for merged in merged_requests: + if merged.status == Status.FAILED: + continue + assert ( + merged.uid not in self.local_metadata_ledger + ), f"finished-request ledger: duplicate uid {merged.uid!r}" + self.local_metadata_ledger[merged.uid] = FinishedRequestRecord.from_request(merged) + payload = msgpack.packb( + [Headers.ENGINE_REPLY.value, [request.serialize() for request in merged_requests]], + use_bin_type=True, + ) + self.socket_for_receiving_requests.send(payload) + def _handle_failed_request(self, request_id: int): """Handle a failed request by sending the reply immediately. @@ -1013,11 +1121,7 @@ def _handle_failed_request(self, request_id: int): # Send the reply immediately, because it may never get a chance to be sent again. if self.use_coordinator and self.is_mp_coordinator: - payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [request_entry.record.merge().serialize()]], - use_bin_type=True, - ) - self.socket_for_receiving_requests.send(payload) + self._send_request_records_to_coordinator([request_entry.record]) elif not self.use_coordinator: if request.prompt is None: request.prompt = self.controller.tokenizer.detokenize( @@ -1189,6 +1293,7 @@ def add_request( request_id: int, prompt: Union[str, List[int], Tensor], sampling_params: Optional[SamplingParams] = None, + precomputed_block_hashes: Optional[List[int]] = None, ) -> asyncio.Future[DynamicInferenceRequest]: """Add request to inference context. @@ -1196,6 +1301,9 @@ def add_request( request_id (int): Unique ID of request. prompt (Union[str, Tensor]): Prompt as either a text string or token IDs. sampling_params (Optional[SamplingParams]): Sampling parameters for the request. + precomputed_block_hashes (Optional[List[int]]): Prefix-cache hashes already + computed for the prompt's complete blocks. Values must match + ``compute_block_hashes_batched(prompt_tokens, block_size_tokens)``. Return: Returns an asyncio `Future[DynamicInferenceRequest]` for the user to wait on. @@ -1238,6 +1346,7 @@ def add_request( sampling_params=sampling_params, block_size_tokens=self.context.block_size_tokens, enable_prefix_caching=self.context.enable_prefix_caching, + precomputed_block_hashes=precomputed_block_hashes or [], ) # Add request. @@ -1257,6 +1366,8 @@ def post_process_requests( pre_fwd_active_token_count: Optional[int] = None, pre_fwd_step_count: Optional[int] = None, finished_routing_block_ids: Optional[Dict[int, list[int]]] = None, + finished_handoff_block_ids: Optional[Dict[int, list[int]]] = None, + finished_handoff_decode_tokens: Optional[Dict[int, list[int]]] = None, ) -> Tuple[List[DynamicInferenceRequest], List[DynamicInferenceRequest]]: """ Handles post-processing for requests after a step. @@ -1279,6 +1390,9 @@ def post_process_requests( finished_routing_block_ids: (Dict[int, List[int]]): Block IDs for finished requests, saved before update_requests released them. Used for per-block routing reconstruction. + finished_handoff_block_ids: Prompt KV block IDs retained for state handoff. + finished_handoff_decode_tokens: First sampled token and optional MTP proposals + needed to resume directly from imported prefill state on decode. Returns: A list of active requests and completed requests as `DynamicInferenceRequest` objects @@ -1310,8 +1424,21 @@ def post_process_requests( if self.num_speculative_tokens > 0 and accepted_tokens is not None: self._spec_steps += 1 + # Convert the step's request IDs once, then batch-prepare handoff metadata for + # finished requests using the KV blocks retained before context cleanup. + request_id_list = request_ids.tolist() + handoff_blocks_by_request = finished_handoff_block_ids or {} + prepared_handoff_metadata = self._prepare_handoff_metadata_batch( + [ + (self.get_request(request_id), handoff_blocks_by_request.get(request_id, [])) + for request_id in request_id_list + if request_id in finished_request_ids + ], + finished_handoff_decode_tokens or {}, + ) + for req_idx, (request_id, tokens, accepted_tokens_list, request_log_probs) in enumerate( - zip(request_ids.tolist(), sample.tolist(), accepted_tokens_iter, log_probs_iter) + zip(request_id_list, sample.tolist(), accepted_tokens_iter, log_probs_iter) ): # Ensure tokens is always a list for consistent handling @@ -1446,6 +1573,16 @@ def post_process_requests( request.generated_length = len(request.generated_tokens) request.status = Status.COMPLETED request.add_event_finish() + # Keep handoff blocks only when the request needs them. + handoff_blocks = handoff_blocks_by_request.get(request_id, []) + if request.sampling_params.do_kv_handoff: + self._capture_handoff_meta( + request, prepared_handoff_metadata.get(request_id) + ) + # A prefill-role engine may also serve regular requests; release the + # temporary block references when no handoff was requested. + elif handoff_blocks: + self._release_pinned_handoff_blocks(handoff_blocks) finished_entry = self.requests.pop(request_id) finished_request = finished_entry.record[-1] finished_request.generated_length = len(finished_request.generated_tokens) @@ -1745,6 +1882,8 @@ def _should_run_async_sched_overlap(self) -> bool: # lifecycle bookkeeping before preparing the next batch. if not self.context.can_prepare_requests(): return False + if self.has_admittable_kv_import: + return False if not self.waiting_request_ids: return True @@ -2280,6 +2419,8 @@ async def async_bookkeep( log_probs = step_result["log_probs"] top_n_logprobs = step_result.get("top_n_logprobs", None) finished_routing_block_ids = step_result.get("finished_routing_block_ids", None) + finished_handoff_block_ids = step_result.get("finished_handoff_block_ids", None) + finished_handoff_decode_tokens = step_result.get("finished_handoff_decode_tokens", None) cuda_graph_request_count = step_result["cuda_graph_request_count"] # Add paused events. @@ -2301,6 +2442,8 @@ async def async_bookkeep( pre_fwd_active_token_count=context_state.get("active_token_count"), pre_fwd_step_count=context_state.get("step_count"), finished_routing_block_ids=finished_routing_block_ids, + finished_handoff_block_ids=finished_handoff_block_ids, + finished_handoff_decode_tokens=finished_handoff_decode_tokens, ) else: @@ -2348,21 +2491,7 @@ async def async_bookkeep( ] if records_to_send: nvtx_range_push("coordinator_communication") - merged_requests = [r.merge() for r in records_to_send] - if self.local_metadata_ledger_enabled: - # Index every finished request's extra metadata before it is dropped. - for merged in merged_requests: - assert ( - merged.uid not in self.local_metadata_ledger - ), f"finished-request ledger: duplicate uid {merged.uid!r}" - self.local_metadata_ledger[merged.uid] = FinishedRequestRecord.from_request( - merged - ) - payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [m.serialize() for m in merged_requests]], - use_bin_type=True, - ) - self.socket_for_receiving_requests.send(payload) + self._send_request_records_to_coordinator(records_to_send) nvtx_range_pop("coordinator_communication") # Stream newly generated tokens for active requests. Finished @@ -2672,6 +2801,12 @@ def schedule_requests(self) -> int: nvtx_range_push("drain_zmq_socket") all_messages = [] if self.is_mp_coordinator: + all_messages.extend( + msgpack.packb( + [Headers.KV_HANDOFF_COMPLETE.value, request_id, failed], use_bin_type=True + ) + for request_id, failed in self._drain_handoff_completion_notifications() + ) while True: try: # Receive messages in a non-blocking way. @@ -2700,6 +2835,24 @@ def schedule_requests(self) -> int: nvtx_range_push("add_request") self.add_request(request_id, prompt, sampling_params) nvtx_range_pop("add_request") + elif header == Headers.SUBMIT_REQUEST_WITH_KV: + # Decode-side KV import. + request_id, prompt, sampling_params, kv_meta, src_block_ids = data[1:] + sampling_params = SamplingParams.deserialize(sampling_params) + nvtx_range_push("add_request_with_kv_handoff") + self.add_request_with_kv_handoff( + request_id, prompt, sampling_params, kv_meta, src_block_ids + ) + nvtx_range_pop("add_request_with_kv_handoff") + elif header == Headers.RELEASE_KV: + # Coordinator-broadcast release. Unknown request ids are no-ops. + self.release_handoff_blocks(int(data[1])) + elif header == Headers.SEND_KV: + # Push transport: send a pinned hand-off's KV to the decode + # instance the coordinator picked. + self.push_handoff_kv(int(data[1]), data[2]) + elif header == Headers.KV_HANDOFF_COMPLETE: + self._record_handoff_completion_notification(int(data[1]), bool(data[2])) elif header == Headers.ABORT_REQUEST: request_id = int(data[1]) entry = self.requests.get(request_id) @@ -2728,6 +2881,9 @@ def schedule_requests(self) -> int: # Control signal: queue for second pass. self._pending_signals.append(message) + self._poll_pending_kv_imports() + self._poll_pending_kv_pushes() + if new_generation_epoch is not None: self._generation_epoch = new_generation_epoch # Stamp all active requests with the new epoch. @@ -2837,10 +2993,25 @@ async def run_engine(self, *, loop: Optional[asyncio.AbstractEventLoop] = None): and ( self.context.get_active_request_count() > 0 or self.waiting_request_ids + or self.pending_kv_import_count > 0 + or self.pending_kv_push_count > 0 ) ) ) - await self.async_step() + self._poll_pending_kv_imports() + self._poll_pending_kv_pushes() + if ( + self.context.get_active_request_count() > 0 + or self.waiting_request_ids + or self.has_admittable_kv_import + ): + await self.async_step() + else: + # Reached when there is no model work to step but a handoff transfer is + # still pending. Handles expose polling rather than an async completion + # callback, so yield briefly to avoid busy-spinning while keeping decode + # admission responsive when the transfer completes. + await asyncio.sleep(0.001) except asyncio.CancelledError: pass @@ -2930,9 +3101,12 @@ async def run_engine_with_coordinator( self.schedule_requests() if self.state in (EngineState.RUNNING, EngineState.PAUSING): - local_pending = self.context.get_active_request_count() + len( - self.waiting_request_ids + local_schedulable = ( + self.context.get_active_request_count() + + len(self.waiting_request_ids) + + int(self.has_admittable_kv_import) ) + local_pending_imports = self.pending_kv_import_count if self.disable_ep_consensus: # Skip the EP consensus all-reduce; act on local state only. # NOTE: even with no consensus we must still participate in EP @@ -2944,8 +3118,12 @@ async def run_engine_with_coordinator( await self._world_barrier() self.state = EngineState.PAUSED self._state_events[EngineState.PAUSED].set() - elif local_pending > 0: + elif local_schedulable > 0: await self.async_step() + elif self.ep_world_size == 1 and local_pending_imports > 0: + # No model work is ready; poll the network transfer without + # spending a dummy forward while waiting for decode admission. + await asyncio.sleep(0.001) else: self.step_start_event.record() nvtx_range_push("EP-dummy-forward") @@ -2973,7 +3151,7 @@ async def run_engine_with_coordinator( # In the worst case, this delays pausing by 20 steps which is around # 200-400 milliseconds. self._last_ep_consensus = await self._ep_establish_consensus( - local_pending, signal_consensus=(self.state == EngineState.PAUSING) + local_schedulable, signal_consensus=(self.state == EngineState.PAUSING) ) global_work, all_pausing = self._last_ep_consensus self._ep_consensus_loop_counter += 1 @@ -2985,7 +3163,7 @@ async def run_engine_with_coordinator( self._state_events[EngineState.PAUSED].set() elif global_work > 0: # At least one EP peer has work: all must participate. - if local_pending > 0: + if local_schedulable > 0: await self.async_step() else: # Dummy forward to participate in the EP collective. @@ -2999,7 +3177,7 @@ async def run_engine_with_coordinator( self.context.prefix_cache_lru_clock += 1 else: # No work, but not all pausing: idle. - await asyncio.sleep(0.02) + await asyncio.sleep(0.001 if local_pending_imports > 0 else 0.02) elif self.state == EngineState.PAUSED: await asyncio.sleep(0.02) diff --git a/megatron/core/inference/headers.py b/megatron/core/inference/headers.py index 35f9505b5fe..72fdb0a5415 100644 --- a/megatron/core/inference/headers.py +++ b/megatron/core/inference/headers.py @@ -29,6 +29,9 @@ class Request(IntEnum): ENGINE_REPLY = 21 ENGINE_REPLY_PARTIAL = 22 ABORT_REQUEST = 23 + SUBMIT_REQUEST_WITH_KV = 24 # Decode-side KV import. + RELEASE_KV = 25 # Free pinned handoff blocks. + SEND_KV = 26 # Tell the prefill engine to send a handoff's KV. class Control(IntEnum): @@ -55,6 +58,7 @@ class Transport(IntEnum): """Low-level transport framing.""" TP_BROADCAST = 80 + KV_HANDOFF_COMPLETE = 81 # Internal MP-wide decode admission decision. # All header categories. To add a new class of headers, define a new IntEnum in diff --git a/megatron/core/inference/inference_client.py b/megatron/core/inference/inference_client.py index 84cd763844b..199018ce3a6 100644 --- a/megatron/core/inference/inference_client.py +++ b/megatron/core/inference/inference_client.py @@ -112,12 +112,89 @@ def add_request( request_id = self.next_request_id self.next_request_id += 1 payload = [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params.serialize()] - payload_serialized = msgpack.packb(payload, use_bin_type=True) - self.socket.send(payload_serialized) - assert request_id not in self.completion_futures - self.completion_futures[request_id] = asyncio.get_running_loop().create_future() - self.request_submission_times[request_id] = time.perf_counter() - return self.completion_futures[request_id] + return self._submit_request(payload, request_id) + + def _make_kv_handoff_request( + self, + prompt: Union[str, List[int]], + sampling_params: SamplingParams, + kv_meta: dict, + src_block_ids: List[int], + ) -> tuple[int, list]: + """Allocate an ID and build a decode request carrying remote KV metadata.""" + request_id = self.next_request_id + self.next_request_id += 1 + payload = [ + Headers.SUBMIT_REQUEST_WITH_KV.value, + request_id, + prompt, + sampling_params.serialize(), + kv_meta, + list(src_block_ids), + ] + return request_id, payload + + def add_request_with_kv_handoff( + self, + prompt: Union[str, List[int]], + sampling_params: SamplingParams, + kv_meta: dict, + src_block_ids: List[int], + ) -> asyncio.Future: + """Submit a request with remote KV metadata. + + The decode engine allocates local blocks, pulls the KV from the + prefill peer described by ``kv_meta``, then begins generation. + + Args: + prompt: A string or list of token IDs. + sampling_params: Sampling parameters for the decode request. + kv_meta: Metadata identifying the remote KV buffers. + src_block_ids: Remote block IDs containing the request's KV state. + + Returns: + asyncio.Future: A future that resolves to the completed request. + """ + request_id, payload = self._make_kv_handoff_request( + prompt, sampling_params, kv_meta, src_block_ids + ) + return self._submit_request(payload, request_id) + + def add_request_with_kv_handoff_streaming( + self, + prompt: Union[str, List[int]], + sampling_params: SamplingParams, + kv_meta: dict, + src_block_ids: List[int], + ) -> AsyncStream[dict]: + """Submit a streaming request with remote KV metadata. + + Returns the same per-step partial/final iterator as + :meth:`add_request_streaming`. + + Args: + prompt: A string or list of token IDs. + sampling_params: Sampling parameters for the decode request. + kv_meta: Metadata identifying the remote KV buffers. + src_block_ids: Remote block IDs containing the request's KV state. + + Returns: + AsyncStream[dict]: Per-step partial and final reply frames. + """ + sampling_params.streaming = True + request_id, payload = self._make_kv_handoff_request( + prompt, sampling_params, kv_meta, src_block_ids + ) + return self._submit_stream(payload, request_id) + + def release_handoff(self, request_id: int) -> None: + """Tell the coordinator to release the KV blocks pinned for `request_id`. + + Fire-and-forget. The coordinator broadcasts RELEASE_KV to every engine; + engines without that request_id ignore the message. + """ + payload = [Headers.RELEASE_KV.value, int(request_id)] + self.socket.send(msgpack.packb(payload, use_bin_type=True)) def abort_request(self, request_id: int) -> None: """Cancel an in-flight request and close its local response stream.""" @@ -162,6 +239,19 @@ def add_request_streaming( request_id = self.next_request_id self.next_request_id += 1 payload = [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params.serialize()] + return self._submit_stream(payload, request_id) + + def _submit_request(self, payload: list, request_id: int) -> asyncio.Future: + """Send a prepared request and register its completion future.""" + self.socket.send(msgpack.packb(payload, use_bin_type=True)) + assert request_id not in self.completion_futures + future = asyncio.get_running_loop().create_future() + self.completion_futures[request_id] = future + self.request_submission_times[request_id] = time.perf_counter() + return future + + def _submit_stream(self, payload: list, request_id: int) -> AsyncStream[dict]: + """Send a prepared streaming request and register its response stream.""" self.socket.send(msgpack.packb(payload, use_bin_type=True)) stream = AsyncStream( request_id, functools.partial(self.abort_request, request_id), loop=self._loop diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index d13859c07b6..799c15b6d43 100644 --- a/megatron/core/inference/inference_request.py +++ b/megatron/core/inference/inference_request.py @@ -393,6 +393,11 @@ class DynamicInferenceRequest(InferenceRequest): # Computed field - not passed by caller precomputed_block_hashes: List[int] = field(default_factory=list) + # KV handoff metadata describing this request's pinned prefill state. + # Used by decode-side pulls and prefill-side pushes. + # Shape: {"request_id", "block_ids", "kv_meta"}. + disaggregated_params: Optional[dict] = None + def __post_init__(self): self.sampling_params = copy.deepcopy(self.sampling_params) if self.prompt_tokens is not None: @@ -748,6 +753,8 @@ def merge_lists(key): policy_epoch = self.requests[-1].policy_epoch kv_cache_epoch = self.requests[-1].kv_cache_epoch + # Preserve KV handoff metadata when merging request segments. + disaggregated_params = self.requests[-1].disaggregated_params # Merged request. request = DynamicInferenceRequest( @@ -775,6 +782,7 @@ def merge_lists(key): enable_prefix_caching=self.requests[0].enable_prefix_caching, precomputed_block_hashes=self.requests[0].precomputed_block_hashes, num_cached_tokens=self.requests[0].num_cached_tokens, + disaggregated_params=disaggregated_params, ) return request diff --git a/megatron/core/inference/sampling_params.py b/megatron/core/inference/sampling_params.py index 2f7d5bb4551..1e3a43f0e89 100644 --- a/megatron/core/inference/sampling_params.py +++ b/megatron/core/inference/sampling_params.py @@ -40,6 +40,7 @@ class SamplingParams: return_prompt_tokens: bool = False streaming: bool = False # Emit incremental ENGINE_REPLY_PARTIAL frames. streaming_interval: int = 1 # Minimum unsent tokens per ENGINE_REPLY_PARTIAL. + do_kv_handoff: bool = False # Pin KV blocks and expose metadata for peer transfer. def __post_init__(self): """Validate parameters and maintain backward compatibility. diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index 8ff0050b257..ee6cbe14933 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -5,7 +5,7 @@ import copy import functools from collections import defaultdict -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, OrderedDict, Tuple, Union import numpy as np @@ -175,6 +175,8 @@ class _AsyncScheduleRequestResult: survivor_idxs: Optional[Tensor] = None newly_paused_request_ids: Optional[Tensor] = None evict_request_ids: Optional[Tensor] = None + finished_handoff_block_ids: Dict[int, List[int]] = field(default_factory=dict) + finished_handoff_decode_tokens: Dict[int, List[int]] = field(default_factory=dict) @dataclass @@ -1899,6 +1901,41 @@ def _apply_stop_word_finished_ids( if request_id in stop_word_finished_ids: active_request_mask[idx] = 0 + def _collect_finished_handoff_state( + self, + finished_idxs: Tensor, + sampled_tokens_cpu: Tensor, + sampled_mtp_tokens_cpu: Optional[Tensor], + ) -> Tuple[Dict[int, List[int]], Dict[int, List[int]]]: + """Retain completed prompt blocks and capture tokens needed to resume decode.""" + + context = self.inference_wrapped_model.inference_context + allocator = context.kv_block_allocator + finished_block_ids: Dict[int, List[int]] = {} + decode_tokens_by_request: Dict[int, List[int]] = {} + if not allocator.enable_handoff_pinning or finished_idxs.numel() == 0: + return finished_block_ids, decode_tokens_by_request + + for finished_idx in finished_idxs.tolist(): + request_id = int(context.request_ids[finished_idx].item()) + blocks = context.request_to_kv_block_ids[finished_idx] + valid_blocks = [int(block) for block in blocks.tolist() if block != -1] + if valid_blocks: + finished_block_ids[request_id] = valid_blocks + # Retain across context cleanup. For an exclusively owned block: + # active=1, retain=2, request cleanup=1, coordinator RELEASE_KV=0. + allocator.retain_memory_blocks(valid_blocks) + + active_idx = finished_idx - context.paused_request_count + decode_tokens = [int(sampled_tokens_cpu[active_idx].item())] + if sampled_mtp_tokens_cpu is not None: + decode_tokens.extend( + int(token) for token in sampled_mtp_tokens_cpu[:, active_idx].tolist() + ) + decode_tokens_by_request[request_id] = decode_tokens + + return finished_block_ids, decode_tokens_by_request + def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: """Update the dynamic inference context after sampling. @@ -1951,6 +1988,12 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: finished_idxs = ( torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count ) + chunked_prefill_idx = context.get_index_of_chunked_prefill_request(safe=True) + if chunked_prefill_idx >= 0: + # The provisional length mask can mark a partial prefill chunk as + # finished. update_requests() keeps it active; keep completion-side + # routing and handoff bookkeeping consistent with that decision. + finished_idxs = finished_idxs[finished_idxs != chunked_prefill_idx] finished_request_ids = context.request_ids[finished_idxs] # Save block IDs for finished requests before update_requests releases them. @@ -1964,6 +2007,14 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: if valid: finished_routing_block_ids[req_id] = valid + # Retain finished prefill blocks before request cleanup releases them; + # the handoff path owns this reference until the decode transfer completes. + finished_handoff_block_ids, finished_handoff_decode_tokens = ( + self._collect_finished_handoff_state( + finished_idxs, sampled_tokens_cpu, sampled_mtp_tokens_cpu + ) + ) + # Clone needed: update_requests mutates next_tokens in-place via tensor_swap, # which would corrupt the reused buffer. new_sample_copy = sampled_tokens_cpu.clone() @@ -1984,6 +2035,8 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: # D2H sync when the engine later calls sample.tolist(). "sample": sampled_tokens_cpu, "finished_routing_block_ids": finished_routing_block_ids, + "finished_handoff_block_ids": finished_handoff_block_ids, + "finished_handoff_decode_tokens": finished_handoff_decode_tokens, **(update_result or {}), } @@ -2693,6 +2746,15 @@ def _run_async_sched_resolve( ) range_pop() + finished_idxs = ( + torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count + ) + finished_handoff_block_ids, finished_handoff_decode_tokens = ( + self._collect_finished_handoff_state( + finished_idxs, sampled_tokens_cpu, sample_result.sampled_mtp_tokens_cpu_view + ) + ) + # Resolve CPU request lifecycle state. range_push("resolve_requests") resolved_finished_request_ids, survivor_idxs = context.resolve_requests(active_request_mask) @@ -2710,6 +2772,8 @@ def _run_async_sched_resolve( active_request_ids=active_request_ids, finished_request_ids=finished_request_ids, survivor_idxs=survivor_idxs, + finished_handoff_block_ids=finished_handoff_block_ids, + finished_handoff_decode_tokens=finished_handoff_decode_tokens, ) def _run_async_sched_update_requests( @@ -2737,6 +2801,15 @@ def _run_async_sched_update_requests( self._build_async_sched_request_state(sampled_tokens_cpu, resolved_sequence_lengths) ) + finished_idxs = ( + torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count + ) + finished_handoff_block_ids, finished_handoff_decode_tokens = ( + self._collect_finished_handoff_state( + finished_idxs, sampled_tokens_cpu, sample_result.sampled_mtp_tokens_cpu_view + ) + ) + mutable_sampled_tokens_cpu = sampled_tokens_cpu.clone() mutable_sampled_mtp_tokens_cpu = ( sample_result.sampled_mtp_tokens_cpu_view.clone() @@ -2758,6 +2831,8 @@ def _run_async_sched_update_requests( finished_request_ids=finished_request_ids, newly_paused_request_ids=update_result.get("newly_paused_request_ids"), evict_request_ids=update_result.get("evict_request_ids"), + finished_handoff_block_ids=finished_handoff_block_ids, + finished_handoff_decode_tokens=finished_handoff_decode_tokens, ) def _build_async_sched_step_result( @@ -2799,6 +2874,8 @@ def _build_async_sched_step_result( "finished_request_ids": request_result.finished_request_ids, "sample": request_result.sampled_tokens_cpu, "finished_routing_block_ids": {}, + "finished_handoff_block_ids": request_result.finished_handoff_block_ids, + "finished_handoff_decode_tokens": request_result.finished_handoff_decode_tokens, "newly_paused_request_ids": request_result.newly_paused_request_ids, "evict_request_ids": request_result.evict_request_ids, "accepted_tokens": request_result.accepted_tokens_cpu, @@ -2885,7 +2962,10 @@ async def _run_async_sched_step_no_overlap( decode_only = DecodeOnly(consumed=consumed_decode_only, launched=launched_decode_only) if not had_pending_forward: - assert active_request_count > 0, "Async no-overlap admission did not add a request." + if active_request_count == 0: + # An admission callback may resolve queued work without adding + # anything that requires a model forward. + return DynamicBatchControllerStepResult(decode_only=decode_only) return DynamicBatchControllerStepResult(decode_only=decode_only, primer_only=True) if log_probs_transfer is not None: diff --git a/tests/unit_tests/inference/contexts/test_kv_block_allocator.py b/tests/unit_tests/inference/contexts/test_kv_block_allocator.py index 6b9ada8cf5e..9025a22c836 100644 --- a/tests/unit_tests/inference/contexts/test_kv_block_allocator.py +++ b/tests/unit_tests/inference/contexts/test_kv_block_allocator.py @@ -212,6 +212,32 @@ def test_prefix_caching_allocate_and_hash_registration(): assert small.is_memory_available(5) is False +def test_retain_memory_blocks_adds_one_reference_per_owner(): + allocator = KVBlockAllocator( + _make_context(), + pool_size=8, + paused_limit=0, + enable_prefix_caching=True, + prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.REF_ZERO, + ) + block = allocator.allocate_memory_blocks(1) + block_id = int(block[0]) + allocator.register_kv_block_hashes([block_id], [101]) + + allocator.retain_memory_blocks([block_id]) + allocator.retain_memory_blocks([block_id]) + assert allocator.block_ref_counts[block_id].item() == 3 + + allocator.retain_memory_blocks([block_id, block_id]) + assert allocator.block_ref_counts[block_id].item() == 5 + + allocator.release_memory_blocks(block) + allocator.release_memory_blocks(block) + allocator.release_memory_blocks(torch.tensor([block_id, block_id, block_id])) + assert allocator.block_ref_counts[block_id].item() == 0 + assert 101 not in allocator.kv_hash_to_block_id + + @pytest.mark.parametrize( "policy", [PrefixCachingEvictionPolicy.REF_ZERO, PrefixCachingEvictionPolicy.LRU] ) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py b/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py index d72e59275d0..01d8629582f 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py @@ -9,6 +9,10 @@ from megatron.core.inference.config import AsyncScheduleMode from megatron.core.inference.contexts.dynamic_context import DynamoHelper +from megatron.core.inference.disaggregation.engine import DisaggDynamicInferenceEngine +from megatron.core.inference.disaggregation.inference_state_handoff import ( + InferenceStateHandoffMixin, +) from megatron.core.inference.engines import DynamicInferenceEngine from megatron.core.inference.engines.dynamic_engine import EngineState, _get_decode_only_log_state from megatron.core.inference.sampling_params import SamplingParams @@ -18,13 +22,17 @@ ) -def _make_engine(async_sched_mode=AsyncScheduleMode.ASYNC, **overrides): - engine = DynamicInferenceEngine.__new__(DynamicInferenceEngine) +def _make_engine( + async_sched_mode=AsyncScheduleMode.ASYNC, engine_cls=DynamicInferenceEngine, **overrides +): + engine = engine_cls.__new__(engine_cls) context = SimpleNamespace( config=SimpleNamespace(async_sched_mode=async_sched_mode), is_hybrid_model=False, enable_prefix_caching=False, num_prefill_requests=0, + total_request_count=0, + max_requests=8, can_prepare_requests=mock.Mock(return_value=True), active_token_count=0, max_tokens=8, @@ -175,6 +183,23 @@ def test_async_sched_overlap_probe_routes_schedulable_chunk_to_no_overlap(): assert not engine._should_run_async_sched_overlap() +def test_ready_handoff_uses_safe_no_overlap_admission_point(): + """A completed import cannot join the batch before pending logits are consumed.""" + + engine = _make_engine(engine_cls=DisaggDynamicInferenceEngine) + engine._initialize_disaggregation_state() + engine.waiting_request_ids = deque() + engine._pending_kv_imports.append(SimpleNamespace(request_id=7, resume_tokens=[55])) + engine._handoff_completion_notifications[7] = False + + assert not engine._should_run_async_sched_overlap() + + # Keep overlap enabled while the active batch is full; the normal lifecycle + # boundary will select no-overlap once capacity becomes available. + engine.context.total_request_count = engine.context.max_requests + assert engine._should_run_async_sched_overlap() + + @pytest.mark.parametrize( "mode, run_async_overlap, decode_only, primer_only, expected_schedule_calls, " "expected_nvtx_range", @@ -364,3 +389,26 @@ def test_async_bookkeep_uses_consumed_chunked_prefill_request_id(): def test_get_decode_only_log_state(mode, decode_only, expected): """Console logging reports transitions and colors the latest available phase.""" assert _get_decode_only_log_state(mode, decode_only) == expected + + +def test_base_engine_rejects_kv_handoff_commands(): + engine = DynamicInferenceEngine.__new__(DynamicInferenceEngine) + + assert InferenceStateHandoffMixin not in DynamicInferenceEngine.mro() + assert engine.pending_kv_import_count == 0 + with pytest.raises(RuntimeError, match="SUBMIT_REQUEST_WITH_KV"): + engine.add_request_with_kv_handoff(1, [], SamplingParams(), {}, []) + with pytest.raises(RuntimeError, match="RELEASE_KV"): + engine.release_handoff_blocks(1) + + +def test_disagg_engine_resolves_handoff_methods_from_mixin(): + assert DisaggDynamicInferenceEngine.mro()[:3] == [ + DisaggDynamicInferenceEngine, + InferenceStateHandoffMixin, + DynamicInferenceEngine, + ] + assert ( + DisaggDynamicInferenceEngine.add_request_with_kv_handoff + is InferenceStateHandoffMixin.add_request_with_kv_handoff + ) diff --git a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py index 3d95f158c32..55c0f50a83d 100644 --- a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py +++ b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py @@ -19,7 +19,10 @@ from megatron.core.inference.data_parallel_inference_coordinator import ( DataParallelInferenceCoordinator, ) -from megatron.core.inference.data_parallel_inference_coordinator.handlers import handle_engine_reply +from megatron.core.inference.data_parallel_inference_coordinator.handlers import ( + HANDLERS, + handle_engine_reply, +) from megatron.core.inference.engines.async_zmq_communicator import AsyncZMQCommunicator from megatron.core.inference.engines.dynamic_engine import ( DynamicInferenceEngine, @@ -50,6 +53,11 @@ DEFAULT_PORT = 46581 +def test_coordinator_registers_client_kv_handoff_handlers(): + assert Headers.SUBMIT_REQUEST_WITH_KV in HANDLERS + assert Headers.RELEASE_KV in HANDLERS + + class DummyTokenizer: """Dummy tokenizer.""" diff --git a/tests/unit_tests/inference/test_disagg_handoff_lifecycle.py b/tests/unit_tests/inference/test_disagg_handoff_lifecycle.py new file mode 100644 index 00000000000..d6988ba49d6 --- /dev/null +++ b/tests/unit_tests/inference/test_disagg_handoff_lifecycle.py @@ -0,0 +1,867 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Cache ownership and capacity tests for disaggregated state handoff.""" + +import asyncio +from collections import deque +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from megatron.core.inference.config import PrefixCachingEvictionPolicy +from megatron.core.inference.contexts.kv_block_allocator import KVBlockAllocator +from megatron.core.inference.disaggregation.decode_admission import ( + additional_decode_blocks, + admit_prefilled_decode, + can_admit_prefilled_decode, +) +from megatron.core.inference.disaggregation.handoff_completion_tracker import ( + HandoffCompletionTracker, +) +from megatron.core.inference.disaggregation.inference_state_handoff import ( + InferenceStateHandoffMixin, +) +from megatron.core.inference.disaggregation.pending_handoff_imports import ( + DeferredKvHandoff, + PendingKvImport, +) +from megatron.core.inference.inference_request import ( + DynamicInferenceRequest, + DynamicInferenceRequestRecord, + Status, + compute_block_hashes_batched, +) +from megatron.core.inference.sampling_params import SamplingParams + + +class _PendingHandle: + def poll(self): + return False + + def wait(self): + return None + + +class _TransferAgent: + def __init__(self): + self.calls = [] + self.is_push = False + + def begin_pull_blocks(self, peer_meta, src_block_ids, dst_block_ids): + self.calls.append((peer_meta, list(src_block_ids), list(dst_block_ids))) + return _PendingHandle() + + +class _KvAllocator: + enable_prefix_caching = True + + def __init__(self): + self.next_block = 10 + self.capacity_available = True + self.releases = [] + self.registered_parent_hashes = [] + self.block_ref_counts = torch.zeros(256, dtype=torch.int32) + self.kv_hash_to_block_id = {} + + def is_memory_available(self, _count, potential_matched_count=0): + assert potential_matched_count >= 0 + return self.capacity_available + + def allocate_memory_blocks(self, count): + blocks = torch.arange(self.next_block, self.next_block + count, dtype=torch.int32) + self.next_block += count + self.block_ref_counts[blocks] = 1 + return blocks + + def release_memory_blocks(self, blocks): + assert torch.all(self.block_ref_counts[blocks] > 0) + self.block_ref_counts[blocks] -= 1 + self.releases.append(blocks.tolist()) + + def retain_memory_blocks(self, block_ids): + if block_ids: + self.block_ref_counts[torch.tensor(block_ids, dtype=torch.int64)] += 1 + + def register_kv_block_hashes(self, block_ids, block_hashes, parent_hashes=None): + self.kv_hash_to_block_id.update(zip(block_hashes, block_ids)) + self.registered_parent_hashes.extend(parent_hashes or []) + + def update_timestamps(self, block_ids): + return None + + +class _SchedulerHarness: + def schedule_waiting_requests(self): + if not self.waiting_request_ids: + return + request_id = self.waiting_request_ids[0] + block_ids = torch.tensor(self.blocks_to_bind[request_id], dtype=torch.int32) + self.context.kv_block_allocator.block_ref_counts[block_ids] += 1 + if request_id in self.partial_admissions: + self.get_request(request_id).finished_chunk_token_count += 4 + else: + self.waiting_request_ids.popleft() + + +class _HandoffHarness(InferenceStateHandoffMixin, _SchedulerHarness): + def __init__(self, loop): + self._loop = loop + self._initialize_disaggregation_state() + self.context = SimpleNamespace( + block_size_tokens=4, + num_speculative_tokens=0, + num_prefill_requests=0, + chunked_prefill_request_id=-1, + total_request_count=0, + max_requests=8, + active_token_count=0, + max_tokens=8, + is_hybrid_model=False, + kv_block_allocator=_KvAllocator(), + memory_buffer=torch.empty(1), + ) + self._kv_transfer_agent = _TransferAgent() + self.pg_collection = SimpleNamespace(tp=None, pp=None, mp=None) + self.waiting_request_ids = deque() + self.requests = {} + self.blocks_to_bind = {} + self.precomputed_hashes = {} + self.partial_admissions = set() + self.track_generated_token_events = False + self.use_coordinator = False + self.is_mp_coordinator = False + + async def _notify_cond_for_new_request(self): + return None + + def add_request(self, request_id, prompt, sampling_params, precomputed_block_hashes=None): + request = DynamicInferenceRequest( + request_id=request_id, + prompt_tokens=torch.tensor(prompt), + sampling_params=sampling_params, + precomputed_block_hashes=precomputed_block_hashes or [], + ) + request.add_event_add_engine() + self.requests[request_id] = request + self.precomputed_hashes[request_id] = precomputed_block_hashes + self.waiting_request_ids.append(request_id) + return self._loop.create_future() + + def get_request(self, request_id): + return self.requests[request_id] + + def _check_stop_words_for_request_post_append(self, request): + for stop_word_ids in request.stop_word_ids or []: + if request.generated_tokens[-len(stop_word_ids) :] == stop_word_ids: + return True, 0 + return False, 0 + + +def _meta(request_id): + return {"request_id": request_id, "resume_tokens": [99]} + + +def _drain_loop(loop): + loop.run_until_complete(asyncio.sleep(0)) + + +def _pending_import(engine, request_id, block_id, block_hash): + return PendingKvImport( + request_id=request_id, + prompt=[1, 2, 3, 4], + sampling_params=SamplingParams(num_tokens_to_generate=2), + local_blocks=[block_id], + hashes=[block_hash], + cached_prefix_block_count=0, + handle=None, + future=engine._loop.create_future(), + resume_tokens=[99], + ) + + +def _completion_tracker(world_size=2): + tracker = object.__new__(HandoffCompletionTracker) + tracker.rank = 0 + tracker.world_size = world_size + tracker.is_coordinator = True + tracker._reports = {} + tracker._failure_notified = set() + tracker._socket = None + return tracker + + +@pytest.fixture +def handoff_loop(): + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +def test_prefilled_decode_admission_uses_imported_kv_without_prompt_tokens(): + metadata_types = DynamicInferenceRequest.get_metadata_types() + context = SimpleNamespace( + block_size_tokens=4, + num_speculative_tokens=2, + num_prefill_requests=0, + chunked_prefill_request_id=-1, + total_request_count=0, + max_requests=2, + active_token_count=0, + max_tokens=8, + request_metadata_types=metadata_types, + request_metadata={label: torch.empty(2, dtype=dtype) for label, dtype in metadata_types}, + request_to_kv_block_ids=torch.full((2, 4), -1, dtype=torch.int32), + request_ids=torch.full((2,), -1, dtype=torch.int64), + request_kv_length_offsets=torch.zeros(2, dtype=torch.int32), + request_query_lengths=torch.zeros(2, dtype=torch.int32), + request_output_lengths=torch.zeros(2, dtype=torch.int32), + request_in_prefill_status_tensor=torch.ones(2, dtype=torch.int32), + request_kv_block_counts=torch.zeros(2, dtype=torch.int32), + request_last_kv_block_id=torch.full((2,), -1, dtype=torch.int32), + request_last_kv_block_offset=torch.zeros(2, dtype=torch.int32), + token_to_input_ids=torch.zeros(8, dtype=torch.int64), + token_to_pos_ids=torch.zeros(8, dtype=torch.int64), + token_to_request_idx=torch.zeros(8, dtype=torch.int32), + token_to_position_in_request=torch.zeros(8, dtype=torch.int64), + token_to_local_position_within_kv_block=torch.zeros(8, dtype=torch.int32), + token_to_block_idx=torch.full((8,), -1, dtype=torch.int32), + ) + request = DynamicInferenceRequest( + request_id=7, + prompt_tokens=torch.arange(3), + sampling_params=SamplingParams(num_tokens_to_generate=8, termination_id=99), + ) + + context.chunked_prefill_request_id = 9 + with pytest.raises(RuntimeError, match="decode-only engine"): + can_admit_prefilled_decode(context, 3) + context.chunked_prefill_request_id = -1 + assert additional_decode_blocks(3, 3, 4) == 1 + admit_prefilled_decode(context, request, [10], [11], [40, 41, 42]) + + assert context.total_request_count == 1 + assert context.num_prefill_requests == 0 + assert context.active_token_count == 3 + assert context.request_kv_length_offsets[0].item() == 3 + assert context.request_query_lengths[0].item() == 3 + assert context.request_to_kv_block_ids[0, :2].tolist() == [10, 11] + assert context.token_to_input_ids[:3].tolist() == [40, 41, 42] + assert context.token_to_pos_ids[:3].tolist() == [3, 4, 5] + assert context.token_to_block_idx[:3].tolist() == [10, 11, 11] + assert request.remaining_prompt_length == 0 + + +def test_completed_kv_handoff_enters_decode_without_waiting_queue(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine.context.num_speculative_tokens = 0 + engine.use_coordinator = False + engine.is_mp_coordinator = False + request_future = handoff_loop.create_future() + request = DynamicInferenceRequest( + request_id=7, + prompt_tokens=torch.arange(4), + sampling_params=SamplingParams(num_tokens_to_generate=3, termination_id=99), + ) + request.add_event_add_engine() + + def add_request(request_id, prompt, sampling_params, precomputed_block_hashes=None): + assert request_id == request.request_id + engine.requests[request_id] = request + engine.waiting_request_ids.append(request_id) + return request_future + + engine.add_request = add_request + pending = PendingKvImport( + request_id=7, + prompt=[0, 1, 2, 3], + sampling_params=request.sampling_params, + local_blocks=[10], + continuation_blocks=[11], + hashes=[101], + cached_prefix_block_count=0, + handle=None, + future=handoff_loop.create_future(), + resume_tokens=[55], + ) + + admission = ( + "megatron.core.inference.disaggregation.inference_state_handoff.admit_prefilled_decode" + ) + with mock.patch(admission) as admit: + engine._finalize_kv_handoff_import(pending) + + assert request.generated_tokens == [55] + assert request.num_cached_tokens == 4 + assert list(engine.waiting_request_ids) == [] + assert pending.local_blocks == [] + assert pending.continuation_blocks == [] + admit.assert_called_once_with(engine.context, request, [10], [11], [55]) + + +def test_setup_pins_handoff_outputs_only_on_prefill(): + class _Backend: + def __init__(self, **_kwargs): + pass + + def export_meta(self): + return {"transport": "test"} + + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + allocator = SimpleNamespace( + enable_prefix_caching=True, enable_handoff_pinning=False, pool_size=8 + ) + engine.context = SimpleNamespace( + kv_block_allocator=allocator, + memory_buffer=torch.empty(2, 1, 8, 4, 1, 1), + is_hybrid_model=False, + num_attention_layers=1, + num_attention_heads_per_partition=1, + hidden_size_per_attention_head=1, + block_size_tokens=4, + ) + model_config = SimpleNamespace(num_query_groups=1, num_attention_heads=1) + engine.controller = SimpleNamespace( + inference_wrapped_model=SimpleNamespace(model=SimpleNamespace(config=model_config)) + ) + engine.pg_collection = SimpleNamespace(tp=None, pp=None, mp=None) + + backend_factory = ( + "megatron.core.inference.disaggregation.inference_state_handoff." + "construct_kv_transfer_backend_class" + ) + pg_size = "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_size" + pg_rank = "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_rank" + with ( + mock.patch(backend_factory, return_value=_Backend), + mock.patch(pg_size, return_value=1), + mock.patch(pg_rank, return_value=0), + ): + engine.setup_kv_transfer("decode") + assert not allocator.enable_handoff_pinning + engine.setup_kv_transfer("prefill") + assert allocator.enable_handoff_pinning + + +def test_reset_cancels_capacity_queued_handoffs(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine.context.kv_block_allocator.capacity_available = False + future = engine.add_request_with_kv_handoff( + 3, [3] * 4, SamplingParams(num_tokens_to_generate=2), _meta(3), [103] + ) + + engine._reset_pending_kv_imports() + + assert future.cancelled() + assert engine.pending_kv_import_count == 0 + + +def test_reset_waits_for_pending_prefill_pushes(handoff_loop): + engine = _HandoffHarness(handoff_loop) + handle = mock.Mock() + engine._pending_kv_pushes = [(7, [handle])] + + engine._reset_pending_kv_imports() + + handle.wait.assert_called_once_with() + assert not engine._pending_kv_pushes + + +def test_reset_rejects_an_active_prefill_push(handoff_loop): + engine = _HandoffHarness(handoff_loop) + handle = mock.Mock() + handle.wait.side_effect = TimeoutError + engine._pending_kv_pushes = [(7, [handle])] + + with pytest.raises(RuntimeError, match="may still access cache storage"): + engine._reset_pending_kv_imports() + + assert engine._pending_kv_pushes == [(7, [handle])] + + +def test_reset_rejects_pinned_prefill_blocks(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine._pinned_handoff_blocks[7] = [10] + + with pytest.raises(RuntimeError, match="KV handoff blocks remain pinned"): + engine._reset_pending_kv_imports() + + assert engine._pinned_handoff_blocks == {7: [10]} + + +def test_reset_does_not_release_unsafe_import_destinations(handoff_loop): + engine = _HandoffHarness(handoff_loop) + block_id = int(engine.context.kv_block_allocator.allocate_memory_blocks(1)[0]) + pending = _pending_import(engine, 4, block_id, 104) + pending.destinations_safe = False + engine._pending_kv_imports.append(pending) + + with pytest.raises(RuntimeError, match="may still access cache storage"): + engine._reset_pending_kv_imports() + + assert engine.context.kv_block_allocator.releases == [] + assert list(engine._pending_kv_imports) == [pending] + + +def test_prefill_handoff_pins_protect_blocks_and_restore_capacity(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine.context.prefix_cache_lru_clock = 0 + allocator = KVBlockAllocator( + engine.context, + pool_size=7, + paused_limit=0, + enable_prefix_caching=True, + prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.LRU, + ) + engine.context.kv_block_allocator = allocator + baseline_capacity = allocator.get_allocatable_count() + + blocks = allocator.allocate_memory_blocks(4).tolist() + first_handoff = blocks[:2] + second_handoff = blocks[2:] + registered_blocks = [first_handoff[0], second_handoff[0]] + partial_tail_blocks = [first_handoff[1], second_handoff[1]] + allocator.register_kv_block_hashes(registered_blocks, [101, 201], [0, 0]) + + for request_id, request_blocks in ((11, first_handoff), (12, second_handoff)): + allocator.retain_memory_blocks(request_blocks) + allocator.release_memory_blocks(torch.tensor(request_blocks, dtype=torch.int32)) + engine._pinned_handoff_blocks[request_id] = request_blocks + + assert allocator.block_ref_counts[blocks].tolist() == [1, 1, 1, 1] + assert int(allocator.get_evictable_block_count()) == 0 + assert allocator.get_allocatable_count() == baseline_capacity - 4 + assert allocator.evict_lru_blocks(1) is False + assert allocator.allocate_memory_blocks(3) is None + + engine.release_handoff_blocks(11) + assert allocator.get_allocatable_count() == baseline_capacity - 2 + assert allocator.block_ref_counts[first_handoff].tolist() == [0, 0] + assert allocator.block_hashes[first_handoff[0]].item() == 101 + assert first_handoff[1] in allocator.block_bag[: allocator.pool_avail].tolist() + + engine.release_handoff_blocks(11) + engine.release_handoff_blocks(999) + assert allocator.get_allocatable_count() == baseline_capacity - 2 + + engine.release_handoff_blocks(12) + assert allocator.get_allocatable_count() == baseline_capacity + assert allocator.block_ref_counts[blocks].tolist() == [0, 0, 0, 0] + assert allocator.block_hashes[registered_blocks].tolist() == [101, 201] + free_blocks = allocator.block_bag[: allocator.pool_avail].tolist() + assert all(block in free_blocks for block in partial_tail_blocks) + + +def test_peer_failure_quarantines_an_unfinished_local_transfer(handoff_loop): + engine = _HandoffHarness(handoff_loop) + block_id = int(engine.context.kv_block_allocator.allocate_memory_blocks(1)[0]) + pending = _pending_import(engine, 4, block_id, 104) + pending.handle = _PendingHandle() + engine._pending_kv_imports.append(pending) + engine._record_handoff_completion_notification(4, failed=True) + + engine._poll_pending_kv_imports() + engine._admit_pending_kv_imports() + + assert isinstance(pending.future.exception(), RuntimeError) + assert engine.context.kv_block_allocator.releases == [] + assert not engine._pending_kv_imports + assert engine._quarantined_kv_imports == [pending] + + +def test_completed_handoff_keeps_notification_while_decode_is_full(handoff_loop): + engine = _HandoffHarness(handoff_loop) + block_id = int(engine.context.kv_block_allocator.allocate_memory_blocks(1)[0]) + pending = _pending_import(engine, 4, block_id, 104) + pending.terminal_state_reported = True + engine._pending_kv_imports.append(pending) + engine._record_handoff_completion_notification(4, failed=False) + engine.context.total_request_count = engine.context.max_requests + + with mock.patch.object(engine, "_finalize_kv_handoff_import") as finalize: + engine._poll_pending_kv_imports() + engine._admit_pending_kv_imports() + + finalize.assert_not_called() + assert list(engine._pending_kv_imports) == [pending] + assert engine._handoff_completion_notifications == {4: False} + + engine.context.total_request_count = 0 + engine._admit_pending_kv_imports() + _drain_loop(handoff_loop) + + finalize.assert_called_once_with(pending) + assert not engine._pending_kv_imports + assert not engine._handoff_completion_notifications + + +def test_transfer_polling_defers_batch_mutation_to_scheduling(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine._kv_transfer_role = "decode" + block_id = int(engine.context.kv_block_allocator.allocate_memory_blocks(1)[0]) + pending = _pending_import(engine, 4, block_id, 104) + pending.terminal_state_reported = True + engine._pending_kv_imports.append(pending) + engine._record_handoff_completion_notification(4, failed=False) + + with mock.patch.object(engine, "_finalize_kv_handoff_import") as finalize: + assert engine._poll_pending_kv_imports() == 1 + finalize.assert_not_called() + + engine.schedule_waiting_requests() + _drain_loop(handoff_loop) + + finalize.assert_called_once_with(pending) + + +@pytest.mark.parametrize( + "num_tokens_to_generate, termination_id, stop_word_ids, expected_tokens", + [(0, -1, None, []), (1, -1, None, [55]), (3, 55, None, [55]), (3, -1, [[55]], [55])], + ids=["sequence-limit", "generation-limit", "termination-token", "stop-word"], +) +def test_handoff_finishes_without_an_extra_decode_step( + handoff_loop, num_tokens_to_generate, termination_id, stop_word_ids, expected_tokens +): + engine = _HandoffHarness(handoff_loop) + blocks = engine.context.kv_block_allocator.allocate_memory_blocks(2).tolist() + request_future = handoff_loop.create_future() + request = DynamicInferenceRequest( + request_id=7, + prompt_tokens=torch.arange(4), + sampling_params=SamplingParams( + num_tokens_to_generate=num_tokens_to_generate, termination_id=termination_id + ), + ) + request.add_event_add_engine() + request.stop_word_ids = stop_word_ids + + def add_request(request_id, _prompt, _sampling_params, precomputed_block_hashes=None): + engine.requests[request_id] = request + engine.waiting_request_ids.append(request_id) + return request_future + + engine.add_request = add_request + pending = PendingKvImport( + request_id=7, + prompt=[0, 1, 2, 3], + sampling_params=request.sampling_params, + local_blocks=[blocks[0]], + continuation_blocks=[blocks[1]], + hashes=[101], + cached_prefix_block_count=0, + handle=None, + future=handoff_loop.create_future(), + resume_tokens=[55], + ) + + with ( + mock.patch( + "megatron.core.inference.disaggregation.inference_state_handoff." + "admit_prefilled_decode" + ) as admit, + mock.patch.object(engine, "_complete_handoff_request_without_forward") as complete, + ): + engine._finalize_kv_handoff_import(pending) + + admit.assert_not_called() + complete.assert_called_once_with(7) + assert request.generated_tokens == expected_tokens + assert engine.context.kv_block_allocator.block_ref_counts[blocks].tolist() == [0, 0] + assert pending.local_blocks == [] + assert pending.continuation_blocks == [] + + +def test_immediate_handoff_completion_resolves_request_future(handoff_loop): + engine = object.__new__(InferenceStateHandoffMixin) + request = DynamicInferenceRequest( + request_id=7, + prompt_tokens=torch.arange(4), + sampling_params=SamplingParams(num_tokens_to_generate=1), + generated_tokens=[55], + ) + request.add_event_add_engine() + record = DynamicInferenceRequestRecord.from_request(request) + request_future = handoff_loop.create_future() + engine.requests = {7: SimpleNamespace(record=record, future=request_future)} + engine.finished_request_count = 0 + engine.use_coordinator = False + engine.is_mp_coordinator = False + engine.controller = SimpleNamespace( + tokenizer=object(), detokenize=mock.Mock(side_effect=["prompt", "answer"]) + ) + + engine._complete_handoff_request_without_forward(7) + + assert request.status == Status.COMPLETED + assert request.generated_length == 1 + assert request.generated_text == "answer" + assert request_future.result() is record + assert engine.finished_request_count == 1 + assert 7 not in engine.requests + + +def test_handoff_completion_waits_for_all_ranks_but_failure_is_immediate(): + tracker = _completion_tracker() + tracker._record(7, rank=0, failed=False) + assert tracker.drain_completed() == [] + + tracker._record(7, rank=1, failed=False) + assert tracker.drain_completed() == [(7, False)] + + tracker._record(8, rank=1, failed=True) + assert tracker.drain_completed() == [(8, True)] + tracker._record(8, rank=0, failed=False) + assert tracker.drain_completed() == [] + assert 8 not in tracker._reports + + +def test_handoff_metadata_batches_completed_requests_across_pipeline(monkeypatch): + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + pp_group = object() + engine.pg_collection = SimpleNamespace(tp=None, pp=pp_group, mp=pp_group) + engine.context = SimpleNamespace( + block_size_tokens=4, num_speculative_tokens=0, memory_buffer=torch.empty(1) + ) + engine._kv_peer_metas = {"global_rank": 0} + engine._pp_kv_peer_metas = [{"global_rank": 0}, {"global_rank": 1}] + requests = [ + SimpleNamespace( + request_id=7, + prompt_tokens=torch.arange(8), + sampling_params=SamplingParams(do_kv_handoff=True), + disaggregated_params=None, + ), + SimpleNamespace( + request_id=8, + prompt_tokens=torch.arange(12), + sampling_params=SamplingParams(do_kv_handoff=True), + disaggregated_params=None, + ), + ] + local_blocks = [[10, 11], [12, 13, 14]] + pg_size = "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_size" + with ( + mock.patch(pg_size, return_value=2), + mock.patch( + "torch.distributed.all_gather_into_tensor", + side_effect=AssertionError("handoff metadata must not run a collective"), + ), + ): + prepared = engine._prepare_handoff_metadata_batch( + zip(requests, local_blocks), {7: [71], 8: [81]} + ) + for request in requests: + engine._capture_handoff_meta(request, prepared[request.request_id]) + + assert requests[0].disaggregated_params == { + "request_id": 7, + "block_ids": [10, 11], + "kv_meta": { + "resume_tokens": [71], + "pp_metas": [ + {"tp_metas": {"global_rank": 0}, "block_ids": [10, 11]}, + {"tp_metas": {"global_rank": 1}, "block_ids": [10, 11]}, + ], + }, + } + assert requests[1].disaggregated_params["kv_meta"]["pp_metas"][1]["block_ids"] == [12, 13, 14] + + +def test_handoff_keeps_partial_block_for_exact_decode_resume(): + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + engine.pg_collection = SimpleNamespace(tp=None, pp=None, mp=None) + engine.context = SimpleNamespace(block_size_tokens=4, num_speculative_tokens=0) + engine._kv_peer_metas = {"global_rank": 0} + engine._release_pinned_handoff_blocks = mock.Mock() + request = SimpleNamespace( + request_id=7, + prompt_tokens=torch.arange(10), + sampling_params=SamplingParams(do_kv_handoff=True), + disaggregated_params=None, + ) + + with mock.patch( + "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_size", return_value=1 + ): + prepared = engine._prepare_handoff_metadata_batch([(request, [10, 11, 12])], {7: [99]}) + engine._capture_handoff_meta(request, prepared[7]) + + assert request.disaggregated_params["block_ids"] == [10, 11, 12] + assert request.disaggregated_params["kv_meta"]["resume_tokens"] == [99] + engine._release_pinned_handoff_blocks.assert_not_called() + + +def test_handoff_metadata_survives_request_serialization(): + metadata = {"request_id": 7, "block_ids": [10, 11], "kv_meta": {"transport": "nixl"}} + request = DynamicInferenceRequest( + request_id=7, + prompt_tokens=torch.arange(8), + sampling_params=SamplingParams(num_tokens_to_generate=0), + disaggregated_params=metadata, + ) + + restored = DynamicInferenceRequest.deserialize(request.serialize()) + + assert restored.disaggregated_params == metadata + + +def test_nixl_handoff_reuses_decode_cached_prefix(handoff_loop): + engine = _HandoffHarness(handoff_loop) + prompt = [1] * 12 + hashes = compute_block_hashes_batched(torch.tensor(prompt), engine.context.block_size_tokens) + cached = engine.context.kv_block_allocator.allocate_memory_blocks(2) + engine.context.kv_block_allocator.release_memory_blocks(cached) + engine.context.kv_block_allocator.kv_hash_to_block_id.update(zip(hashes[:2], cached.tolist())) + + kv_meta = {"request_id": 5, "resume_tokens": [99]} + engine.add_request_with_kv_handoff( + 5, prompt, SamplingParams(num_tokens_to_generate=2), kv_meta, [100, 101, 102] + ) + _drain_loop(handoff_loop) + + pending = engine._pending_kv_imports[0] + assert engine._kv_transfer_agent.calls == [(kv_meta, [102], [12])] + assert pending.local_blocks == [10, 11, 12] + assert pending.cached_prefix_block_count == 2 + with mock.patch( + "megatron.core.inference.disaggregation.inference_state_handoff.admit_prefilled_decode" + ) as admit: + engine._finalize_kv_handoff_import(pending) + assert engine.precomputed_hashes[5] == hashes + assert engine.context.kv_block_allocator.registered_parent_hashes == [hashes[1]] + assert engine.get_request(5).generated_tokens == [99] + assert 5 not in engine.waiting_request_ids + admit.assert_called_once_with(engine.context, engine.get_request(5), [10, 11, 12], [13], [99]) + + +def test_decode_handoff_defers_until_kv_capacity_is_available(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine.context.kv_block_allocator.capacity_available = False + + kv_meta = {"request_id": 8, "resume_tokens": [99]} + future = engine.add_request_with_kv_handoff( + 8, [1] * 8, SamplingParams(num_tokens_to_generate=2), kv_meta, [100, 101] + ) + + assert engine.pending_kv_import_count == 1 + assert len(engine._deferred_kv_handoffs) == 1 + assert not engine._pending_kv_imports + assert not engine._kv_transfer_agent.calls + assert not future.done() + + engine.context.kv_block_allocator.capacity_available = True + engine._poll_pending_kv_imports() + _drain_loop(handoff_loop) + + assert not engine._deferred_kv_handoffs + assert len(engine._pending_kv_imports) == 1 + assert engine._kv_transfer_agent.calls == [(kv_meta, [100, 101], [10, 11])] + + +def test_handoff_submission_failure_is_reported_to_model_parallel_coordinator(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine._kv_transfer_agent.begin_pull_blocks = mock.Mock( + side_effect=RuntimeError("local handoff submission failed") + ) + engine._handoff_completion_tracker = mock.Mock(world_size=2) + future = handoff_loop.create_future() + handoff = DeferredKvHandoff( + request_id=8, + prompt=[1] * 8, + sampling_params=SamplingParams(num_tokens_to_generate=2), + kv_meta={"request_id": 8, "resume_tokens": [99]}, + src_block_ids=[100, 101], + hashes=compute_block_hashes_batched(torch.tensor([1] * 8), 4), + num_blocks=2, + future=future, + ) + + assert engine._try_start_kv_handoff_import(handoff) + _drain_loop(handoff_loop) + engine._poll_pending_kv_imports() + engine._handoff_completion_tracker.report.assert_called_once_with(8, True) + + engine._record_handoff_completion_notification(8, failed=True) + engine._admit_pending_kv_imports() + + assert isinstance(future.exception(), RuntimeError) + assert engine.context.kv_block_allocator.releases == [[10, 11, 12]] + assert not engine._pending_kv_imports + + +def test_nixl_handoff_trims_pipeline_stage_block_lists(handoff_loop): + engine = _HandoffHarness(handoff_loop) + prompt = [2] * 12 + hashes = compute_block_hashes_batched(torch.tensor(prompt), engine.context.block_size_tokens) + cached = engine.context.kv_block_allocator.allocate_memory_blocks(1) + engine.context.kv_block_allocator.release_memory_blocks(cached) + engine.context.kv_block_allocator.kv_hash_to_block_id[hashes[0]] = int(cached[0]) + kv_meta = { + "resume_tokens": [99], + "pp_metas": [ + {"tp_metas": {"rank": 0}, "block_ids": [100, 101, 102]}, + {"tp_metas": {"rank": 1}, "block_ids": [200, 201, 202]}, + ], + } + + engine.add_request_with_kv_handoff( + 6, prompt, SamplingParams(num_tokens_to_generate=2), kv_meta, [100, 101, 102] + ) + _drain_loop(handoff_loop) + + submitted_meta, src_blocks, dst_blocks = engine._kv_transfer_agent.calls[0] + assert src_blocks == [101, 102] + assert dst_blocks == [11, 12] + assert [stage["block_ids"] for stage in submitted_meta["pp_metas"]] == [[101, 102], [201, 202]] + + +def test_nixl_handoff_trims_per_rank_block_lists(handoff_loop): + engine = _HandoffHarness(handoff_loop) + prompt = [2] * 8 + hashes = compute_block_hashes_batched(torch.tensor(prompt), engine.context.block_size_tokens) + cached = engine.context.kv_block_allocator.allocate_memory_blocks(1) + engine.context.kv_block_allocator.release_memory_blocks(cached) + engine.context.kv_block_allocator.kv_hash_to_block_id[hashes[0]] = int(cached[0]) + kv_meta = { + "resume_tokens": [99], + "tp_metas": [{"rank": 0, "block_ids": [100, 101]}, {"rank": 1, "block_ids": [200, 201]}], + } + + engine.add_request_with_kv_handoff( + 6, prompt, SamplingParams(num_tokens_to_generate=2), kv_meta, [100, 101] + ) + _drain_loop(handoff_loop) + + submitted_meta, src_blocks, dst_blocks = engine._kv_transfer_agent.calls[0] + assert src_blocks == [101] + assert dst_blocks == [11] + assert [meta["block_ids"] for meta in submitted_meta["tp_metas"]] == [[101], [201]] + + +def test_nccl_handoff_does_not_filter_source_push(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine._kv_transfer_agent.is_push = True + prompt = [3] * 8 + hashes = compute_block_hashes_batched(torch.tensor(prompt), engine.context.block_size_tokens) + engine.context.kv_block_allocator.kv_hash_to_block_id[hashes[0]] = 4 + + kv_meta = {"request_id": 9, "resume_tokens": [99]} + engine.add_request_with_kv_handoff( + 9, prompt, SamplingParams(num_tokens_to_generate=2), kv_meta, [100, 101] + ) + _drain_loop(handoff_loop) + + assert engine._kv_transfer_agent.calls == [(kv_meta, [100, 101], [10, 11])] + + +def test_decode_role_rejects_prompt_scheduling(handoff_loop): + engine = _HandoffHarness(handoff_loop) + engine._kv_transfer_role = "decode" + engine.waiting_request_ids.append(3) + + with pytest.raises(RuntimeError, match="cannot schedule prompt prefill"): + engine.schedule_waiting_requests() diff --git a/tests/unit_tests/inference/test_inference_client.py b/tests/unit_tests/inference/test_inference_client.py index e28bb8b8c4c..71e84d2c16d 100644 --- a/tests/unit_tests/inference/test_inference_client.py +++ b/tests/unit_tests/inference/test_inference_client.py @@ -118,3 +118,24 @@ async def test_inference_client_connect_handshake_rejects_unexpected_reply(): fake_socket.recv.return_value = msgpack.packb([Headers.STOP.value], use_bin_type=True) with pytest.raises(AssertionError): client._connect_with_inference_coordinator() + + +async def test_add_request_with_kv_handoff_returns_future(): + """Non-streaming handoffs use the normal final-reply future path.""" + client, _, fake_socket = _make_client() + params = SamplingParams(temperature=0.5) + + future = client.add_request_with_kv_handoff([1, 2, 3], params, {"agent": "prefill"}, [10, 11]) + + assert isinstance(future, asyncio.Future) + assert params.streaming is False + assert client.completion_futures == {0: future} + assert client.streams == {} + payload = msgpack.unpackb(fake_socket.send.call_args.args[0], raw=False) + assert payload[0] == Headers.SUBMIT_REQUEST_WITH_KV.value + assert payload[1] == 0 + assert payload[2] == [1, 2, 3] + assert payload[3]["streaming"] is False + assert payload[4] == {"agent": "prefill"} + assert payload[5] == [10, 11] + future.cancel() diff --git a/tests/unit_tests/inference/test_inference_client_streaming.py b/tests/unit_tests/inference/test_inference_client_streaming.py index 5d587ffae2f..29fdcc7497c 100644 --- a/tests/unit_tests/inference/test_inference_client_streaming.py +++ b/tests/unit_tests/inference/test_inference_client_streaming.py @@ -151,3 +151,27 @@ async def test_stream_close_sends_abort_and_terminates_iterator(): assert iterator.request_id not in client.streams assert iterator.request_id not in client.request_submission_times assert [item async for item in iterator] == [] + + +async def test_add_request_with_kv_handoff_streaming_returns_stream(): + """Streaming handoffs opt into partial replies and register an AsyncStream.""" + client, fake_socket = _make_client() + params = SamplingParams(temperature=0.5) + + stream = client.add_request_with_kv_handoff_streaming( + [1, 2, 3], params, {"agent": "prefill"}, [10, 11] + ) + + assert isinstance(stream, AsyncStream) + assert params.streaming is True + assert client.streams == {0: stream} + assert client.completion_futures == {} + payload = msgpack.unpackb(fake_socket.send.call_args.args[0], raw=False) + assert payload[0] == Headers.SUBMIT_REQUEST_WITH_KV.value + assert payload[1] == 0 + assert payload[2] == [1, 2, 3] + assert payload[3]["streaming"] is True + assert payload[4] == {"agent": "prefill"} + assert payload[5] == [10, 11] + + await stream.aclose() diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py index f6af8a66880..efb4b412798 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py +++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py @@ -282,6 +282,7 @@ def _make_async_sched_context(total_request_count=2, paused_request_count=0): max_requests=metadata_len, max_tokens=32, request_query_lengths=torch.ones(metadata_len, dtype=torch.int32), + kv_block_allocator=SimpleNamespace(enable_handoff_pinning=False), ) context.is_decode_only = mock.Mock(side_effect=lambda: context.num_prefill_requests == 0) return context @@ -1007,6 +1008,50 @@ def test_build_async_sched_request_state_keeps_partial_chunk_active(): assert finished_request_ids.tolist() == [11] +def test_legacy_bookkeeping_keeps_partial_chunk_active(): + """Legacy bookkeeping must not finish a request midway through its prompt.""" + context = _make_async_sched_context(total_request_count=1) + context.chunked_prefill_request_id = 10 + context.get_index_of_chunked_prefill_request = mock.Mock(return_value=0) + context.get_active_sequence_lengths.return_value = torch.tensor([3]) + context.get_max_sequence_lengths.return_value = torch.tensor([4]) + context.kv_block_allocator = SimpleNamespace(block_routing=False, enable_handoff_pinning=True) + context.request_to_kv_block_ids = torch.tensor([[1, 2, -1]], dtype=torch.int32) + controller = _make_async_sched_controller(context) + controller._sampled_tokens_cuda[0] = 7 + + result = controller._dynamic_step_context_bookkeeping() + + assert result["finished_request_ids"].numel() == 0 + assert result["finished_handoff_block_ids"] == {} + context.update_requests.assert_called_once() + + +def test_async_bookkeeping_retains_finished_handoff_state(): + context = _make_async_sched_context(total_request_count=3, paused_request_count=1) + context.get_max_sequence_lengths.return_value = torch.tensor([10, 10]) + context.kv_block_allocator = SimpleNamespace( + enable_handoff_pinning=True, retain_memory_blocks=mock.Mock() + ) + context.request_to_kv_block_ids = torch.tensor( + [[8, 9, -1], [10, 11, -1], [12, 13, -1]], dtype=torch.int32 + ) + controller = _make_async_sched_controller(context) + sample_result = SimpleNamespace( + sampled_tokens_cpu_view=torch.tensor([99, 7]), + sampled_mtp_tokens_cpu_view=None, + accepted_tokens_cpu_view=None, + ) + + result = controller._run_async_sched_update_requests( + sample_result, resolved_sequence_lengths=torch.tensor([3, 3]) + ) + + assert result.finished_handoff_block_ids == {11: [10, 11]} + assert result.finished_handoff_decode_tokens == {11: [99]} + context.kv_block_allocator.retain_memory_blocks.assert_called_once_with([10, 11]) + + @pytest.mark.parametrize( "termination_ids, stop_word_finished_ids", [([99, 99, 99], set()), ([99, 2, 99], set()), ([99, 99, 99], {11})], @@ -1034,7 +1079,9 @@ def test_run_async_sched_resolve_compacts_without_forward_sync( controller._compact_async_sched_logits = mock.Mock() sample_result = SimpleNamespace( - sampled_tokens_cpu_view=sample_tokens, accepted_tokens_cpu_view=None + sampled_tokens_cpu_view=sample_tokens, + sampled_mtp_tokens_cpu_view=None, + accepted_tokens_cpu_view=None, ) result = controller._run_async_sched_resolve( sample_result, context.get_active_sequence_lengths() + 1 @@ -1113,6 +1160,8 @@ def test_async_sched_step_overlap_order(): survivor_idxs=torch.tensor([0, 2]), newly_paused_request_ids=None, evict_request_ids=None, + finished_handoff_block_ids={}, + finished_handoff_decode_tokens={}, ) ) @@ -1240,6 +1289,8 @@ def test_async_sched_step_yields_after_resolution_outside_inference_mode(): survivor_idxs=torch.tensor([0]), newly_paused_request_ids=None, evict_request_ids=None, + finished_handoff_block_ids={}, + finished_handoff_decode_tokens={}, ) ) observed = [] @@ -1290,6 +1341,27 @@ def admit_request(): assert context.async_sched_step_count == 0 +def test_async_sched_no_overlap_admission_can_finish_without_launching_a_primer(): + """An admission callback may resolve its request without adding model work.""" + + context = _make_async_sched_context(total_request_count=0) + context.active_token_count = 0 + controller = _make_async_sched_controller(context) + controller._async_sched_logits = AsyncScheduleLogitsState() + admit_request = mock.Mock() + controller._run_async_sched_forward_primer = mock.Mock() + + result = asyncio.run( + controller._run_async_sched_step_no_overlap(schedule_waiting_requests=admit_request) + ) + + assert result == DynamicBatchControllerStepResult( + decode_only=DecodeOnly(consumed=None, launched=None) + ) + admit_request.assert_called_once_with() + controller._run_async_sched_forward_primer.assert_not_called() + + def test_run_async_sched_update_requests_preserves_pre_update_output(): """Legacy bookkeeping may reorder working samples without corrupting step output.""" context = _make_async_sched_context(total_request_count=3, paused_request_count=1) @@ -1360,6 +1432,8 @@ def test_async_sched_no_overlap_updates_before_admission( survivor_idxs=None, newly_paused_request_ids=torch.tensor([10]), evict_request_ids=torch.tensor([11]), + finished_handoff_block_ids={}, + finished_handoff_decode_tokens={}, ) input_ids = torch.empty(1, dtype=torch.int64) position_ids = torch.empty(1, dtype=torch.int64) @@ -1464,6 +1538,8 @@ def test_async_sched_no_overlap_finishes_with_matching_ep_base_forward(): survivor_idxs=None, newly_paused_request_ids=None, evict_request_ids=None, + finished_handoff_block_ids={}, + finished_handoff_decode_tokens={}, ) controller._run_async_sched_sample = mock.Mock(return_value=sample_result) controller._synchronize_async_sched_event = mock.Mock() @@ -1514,6 +1590,8 @@ def test_async_sched_mtp_overlap_step_order(): survivor_idxs=torch.tensor([2, 1]), newly_paused_request_ids=None, evict_request_ids=None, + finished_handoff_block_ids={}, + finished_handoff_decode_tokens={}, ) input_ids = torch.empty(9, dtype=torch.int64) position_ids = torch.empty(9, dtype=torch.int64) From 814300c879b12df55d64b97739649777bfc25b3a Mon Sep 17 00:00:00 2001 From: Jiangfei Duan <30710061+JF-D@users.noreply.github.com> Date: Fri, 14 Aug 2026 20:49:21 +0000 Subject: [PATCH 040/101] Fix hidden checkpoint parameter handling in local CUDA graphs (#6373) Signed-off-by: Jiangfei Duan --- megatron/core/transformer/cuda_graphs.py | 28 ++++- .../transformer/test_cuda_graphs.py | 110 ++++++++++++++++++ 2 files changed, 133 insertions(+), 5 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index df7a1dc2c9a..e800d2bfee3 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -864,6 +864,8 @@ def backward(ctx, *grads): runner.bwd_graph_replay_complete_event.record(torch.cuda.current_stream()) for param in runner.params_to_backprop: param._cudagraph_wgrad_ready_event = runner.bwd_graph_replay_complete_event + if hasattr(param, 'grad_added_to_main_grad'): + param.grad_added_to_main_grad = True runner.status = _GraphStatus.FWD_READY # Update FP8 scale factors if needed @@ -880,8 +882,6 @@ def backward(ctx, *grads): with torch.cuda.stream(gtp_rs_stream): for param in params: param.grad = None - if hasattr(param, 'grad_added_to_main_grad'): - param.grad_added_to_main_grad = True hook = getattr(param, '_grad_accum_hook', None) if hook is not None: hook() @@ -933,6 +933,7 @@ def __init__( self.fp4_runtime_enabled = None self.deallocate_pipeline_outputs = False self.num_warmup_steps = 0 + self.needs_recompute_param_discovery = False self.use_stream = False self.gtp_remat = False self.fwd_side_streams = [] @@ -964,6 +965,9 @@ def __init__( self.backward_retain_grad = self.base_module.config.cuda_graph_retain_backward_graph self.deallocate_pipeline_outputs = self.base_module.config.deallocate_pipeline_outputs self.num_warmup_steps = self.base_module.config.cuda_graph_warmup_steps + self.needs_recompute_param_discovery = ( + self.grad_enabled and self.base_module.config.recompute_granularity is not None + ) self.fp8_enabled = self.base_module.config.fp8 is not None self.fp4_enabled = self.base_module.config.fp4 is not None self.fp8_runtime_enabled = None @@ -1280,7 +1284,12 @@ def _resolve_input_buffer(ten): with ctx: # warmup again as case graph capture mode may execute a different codepath _set_warmup_start() - for _ in range(self.num_warmup_steps): + + # Recompute parameter discovery needs one warmup output: graph warmup bypasses + # CheckpointWithoutOutput, making parameters hidden by its nested backward reachable + # from that output's autograd graph. Force this pass even when configured warmup is 0. + num_warmup_steps = max(self.num_warmup_steps, int(self.needs_recompute_param_discovery)) + for _ in range(num_warmup_steps): with self.get_quantization_context(): def clone_ten(ten): @@ -1389,7 +1398,10 @@ def clone_ten(ten): # Preserve only forward buffers whose lifetime crosses into backward capture. self._weakref_forward_buffers(preserve_forward_to_backward_lifetimes=True) - self.params_to_backprop = self.get_connected_params(fwd_graph_outputs) + if self.needs_recompute_param_discovery: + self.params_to_backprop = self.get_connected_params(warmup_outputs) + else: + self.params_to_backprop = self.get_connected_params(fwd_graph_outputs) self.num_dgrads = len(self.fwd_graph_input_surface) self.fwd_graph_input_surface = self.fwd_graph_input_surface + self.params_to_backprop @@ -1467,8 +1479,14 @@ def create_bwd_graph(self): 1 for i in self.fwd_graph_input_surface[: self.num_dgrads] if i.requires_grad ) for param, wgrad in zip(self.params_to_backprop, grad_inputs[n_act_grads:]): - if wgrad is not None and not getattr(param, 'grad_added_to_main_grad', False): + grad_already_accumulated = getattr(param, 'grad_added_to_main_grad', False) + if wgrad is not None and not grad_already_accumulated: param.main_grad.add_(wgrad) + if param.grad is not None and not grad_already_accumulated: + # A nested backward can accumulate directly into param.grad instead of + # returning that contribution through the outer autograd.grad call. + param.main_grad.add_(param.grad) + param.grad_added_to_main_grad = True # GTP cross-graph RS overlap, two phases: # Phase 1 — drain AG, fence runner_stream past ag_stream's tail, diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 10297072664..933e1850b52 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -8,6 +8,7 @@ import torch from transformer_engine.pytorch.fp8 import check_fp8_support +from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.enums import ModelType from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_decoder_block_spec, @@ -28,6 +29,7 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import ( HAVE_TE, + CheckpointWithoutOutput, initialize_rng_tracker, model_parallel_cuda_manual_seed, ) @@ -36,6 +38,7 @@ TECudaGraphHelper, _CudagraphGlobalRecord, create_cudagraphs, + delete_cuda_graphs, ) from megatron.core.transformer.enums import ( AttnBackend, @@ -1643,6 +1646,25 @@ def my_op(self, x): return self.linear(x) +class _CheckpointDependencyModule(MegatronModule): + """A visible projection consuming the output of a checkpointed child module.""" + + def __init__(self, config): + super().__init__(config) + self.checkpointed = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.projection = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=False) + + def forward(self, x): + return self.my_op(x) + + def my_op(self, x): + checkpoint = CheckpointWithoutOutput() + hidden = checkpoint.checkpoint(self.checkpointed, x) + output = self.projection(hidden) + checkpoint.discard_output_and_register_recompute(output) + return output + + class _SimpleNonModule: """non-nn.Module base_module for testing the function_name= form of `CudaGraphManager`.""" @@ -1661,6 +1683,94 @@ def _make_simple_non_module(config): return _SimpleNonModule(config) +class TestCheckpointParameterDiscovery: + """Local-CG discovery and DDP readiness for a checkpointed nested module.""" + + def setup_method(self, method): + initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) + Utils.initialize_model_parallel() + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + if _CudagraphGlobalRecord.cudagraph_created: + delete_cuda_graphs() + else: + _CudagraphGlobalRecord.cudagraph_record = [] + _CudagraphGlobalRecord.cudagraph_inference_record = [] + CudaGraphManager.global_mempool = None + torch.cuda.set_stream(torch.cuda.default_stream()) + Utils.destroy_model_parallel() + + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.5.0")), + reason="use_te_rng_tracker requires TransformerEngine version >= 1.5", + ) + def test_nested_checkpoint_parameter_gradients_and_ddp_readiness(self): + if not torch.distributed.is_initialized() or torch.distributed.get_world_size() < 2: + pytest.skip("test requires at least two data-parallel ranks") + + config = TransformerConfig( + num_layers=1, + hidden_size=32, + num_attention_heads=1, + use_cpu_initialization=True, + cuda_graph_impl="local", + cuda_graph_warmup_steps=0, + recompute_granularity="selective", + recompute_modules=["layernorm"], + ) + torch.manual_seed(123) + reference = _CheckpointDependencyModule(config).cuda() + module = _CheckpointDependencyModule(config).cuda() + module.load_state_dict(reference.state_dict()) + + manager = CudaGraphManager( + config, base_module=module, function_name="my_op", need_backward=True + ) + ddp_model = DistributedDataParallel( + config, + DistributedDataParallelConfig(overlap_grad_reduce=True, bucket_size=1_000_000), + module, + ) + + # Distinct inputs make the expected DDP result the average of different local gradients. + torch.manual_seed(1000 + ddp_model.dp_group.rank()) + test_input = torch.randn(4, config.hidden_size, device="cuda", requires_grad=True) + reference_output = reference.my_op(test_input.detach().clone().requires_grad_(True)) + reference_output_value = reference_output.detach().clone() + reference_output.sum().backward() + reference_wgrads = { + name: param.grad.detach().clone() for name, param in reference.named_parameters() + } + for grad in reference_wgrads.values(): + torch.distributed.all_reduce(grad, group=ddp_model.dp_group) + grad.div_(ddp_model.dp_group.size()) + del reference_output + + ddp_model.zero_grad_buffer() + record_input = test_input.detach().clone().requires_grad_(True) + ddp_model(record_input).sum().backward() + ddp_model.finish_grad_sync() + ddp_model.zero_grad_buffer() + create_cudagraphs() + + runner = manager.cudagraph_runners[0] + assert any(param is module.checkpointed.weight for param in runner.params_to_backprop) + + # The first replay uses DDP's recorded ready counts; the second verifies per-step reset. + for _ in range(2): + ddp_model.zero_grad_buffer() + replay_output = ddp_model(test_input.detach().clone().requires_grad_(True)) + replay_output.sum().backward() + ddp_model.finish_grad_sync() + torch.cuda.synchronize() + + torch.testing.assert_close(replay_output, reference_output_value) + for name, param in module.named_parameters(): + assert param.grad is None + torch.testing.assert_close(param.main_grad, reference_wgrads[name]) + + class TestInlineCaptureManager: """Tests for CudaGraphManager with inline_capture, function_name, eager, and cache_key.""" From 78ad10898de3c49ef0a988e88a062310bad13087 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Fri, 14 Aug 2026 21:25:56 +0000 Subject: [PATCH 041/101] Track DBuffer allocation stream (#6441) Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/dbuffer.py | 27 +++++++++++++++++-- .../experimental/parameter_group.py | 6 +++++ .../distributed/mfsdp_v2/test_dbuffer.py | 6 ++++- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 1761232a848..2f4e622d878 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -67,6 +67,10 @@ class DBuffer: layout: GlobalLayout offset: int local_buffer: torch.Tensor + # Record the allocation stream so DBuffer users can check that uses are joined back to it + # before deleting the buffer. See https://github.com/NVIDIA/Megatron-LM/pull/6187 and + # https://docs.pytorch.org/docs/2.13/generated/torch.Tensor.record_stream.html. + allocation_stream: torch.cuda.Stream | None def __init__( self, @@ -100,6 +104,11 @@ def __init__( self.offset, local_numel = self.layout.get_local_range(self.mesh, self.placements) self.local_buffer = torch.empty(local_numel, dtype=dtype, device=device) + self.allocation_stream = ( + torch.cuda.current_stream(self.local_buffer.device) + if self.local_buffer.is_cuda + else None + ) @property def dtype(self) -> torch.dtype: @@ -155,6 +164,8 @@ def from_local( mesh: DeviceMesh, placements: Iterable[Placement], tensor_shapes: Iterable[Shape], + *, + allocation_stream: torch.cuda.Stream | None, ) -> "DBuffer": """Create a DBuffer from an existing local buffer. @@ -165,6 +176,7 @@ def from_local( mesh: Device mesh whose dimensions correspond to ``placements``. placements: Per-mesh-axis DBuffer placements. tensor_shapes: Global shapes for each logical tensor in this buffer. + allocation_stream: CUDA stream that allocated ``local_buffer``, or ``None`` for CPU. Returns: A DBuffer that reuses ``local_buffer`` without allocating storage. @@ -195,6 +207,7 @@ def from_local( buffer.layout = layout buffer.offset = offset buffer.local_buffer = local_buffer + buffer.allocation_stream = allocation_stream return buffer @classmethod @@ -340,7 +353,11 @@ def redistribute( "Replicate -> Partial redistribute does not support an out buffer." ) return DBuffer.from_local( - self.local_buffer, self.mesh, new_placements, self.layout.tensor_shapes + self.local_buffer, + self.mesh, + new_placements, + self.layout.tensor_shapes, + allocation_stream=self.allocation_stream, ) raise NotImplementedError( "Unsupported DBuffer placement transition on axis " @@ -450,7 +467,13 @@ def scatter( raise RuntimeError("scatter() destination is not contained in the source local buffer.") local_slice = self.local_buffer.narrow(0, local_buffer_offset, destination_numel) if out is None: - return DBuffer.from_local(local_slice, self.mesh, placements, self.layout.tensor_shapes) + return DBuffer.from_local( + local_slice, + self.mesh, + placements, + self.layout.tensor_shapes, + allocation_stream=self.allocation_stream, + ) out.local_buffer.copy_(local_slice) return out diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 59bd7d7b569..622d1f7908c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -375,6 +375,9 @@ def reduce_partial_gradients( # DP-outer-Partial accumulation placement -- a metadata relabel for HSDP, # and a fresh reduce-scattered buffer for HFSDP in the future. if self.main_grad.placements != self._accumulation_placements: + assert self.main_grad.allocation_stream == ( + torch.cuda.current_stream(self.main_grad.device) + ) self.main_grad = self.main_grad.redistribute(self._accumulation_placements) can_reduce_into_main_grad = ( @@ -406,6 +409,9 @@ def reduce_partial_gradients( if is_last_microbatch: # Finalize the deferred DP-outer reduction (all-reduce for HSDP, # reduce-scatter for HFSDP) before binding the sharded parameter grads. + assert self.main_grad.allocation_stream == ( + torch.cuda.current_stream(self.main_grad.device) + ) self.main_grad = self.main_grad.redistribute(self.main_weight.placements) # Make each sharded parameter's .grad consistent with the final main_grad. diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py b/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py index febb95d770d..8f8721b0a10 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py @@ -240,7 +240,11 @@ def test_from_local_reuses_required_local_buffer(distributed_setup): local_buffer = replicated_buffer.local_buffer.narrow(0, offset, local_numel) sharded_buffer = DBuffer.from_local( - local_buffer, mesh, iter([Flat()]), replicated_buffer.layout.tensor_shapes + local_buffer, + mesh, + iter([Flat()]), + replicated_buffer.layout.tensor_shapes, + allocation_stream=replicated_buffer.allocation_stream, ) assert sharded_buffer.placements == (Flat(),) From 883533e049497ec129d6c2055f3c1c944fefe460 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Fri, 14 Aug 2026 22:53:55 +0000 Subject: [PATCH 042/101] [Main] Support device-init grouped linear module too with when not using TE opfuser (#6000) Signed-off-by: Zhongbo Zhu --- .../core/extensions/transformer_engine.py | 13 + megatron/core/transformer/moe/experts.py | 116 +++- megatron/core/transformer/moe/moe_utils.py | 30 +- .../core/transformer/moe/token_dispatcher.py | 5 +- .../core/transformer/transformer_config.py | 46 +- .../models/test_hybrid_moe_model.py | 1 + .../transformer/moe/test_grouped_mlp.py | 278 +++++++- ...test_grouped_tensor_dispatcher_numerics.py | 599 ++++++++++++++++++ ...test_moe_single_grouped_weight_numerics.py | 51 +- 9 files changed, 1080 insertions(+), 59 deletions(-) create mode 100644 tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 8dd339f3cef..6df3049d870 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2338,6 +2338,10 @@ def sharded_state_dict( if HAVE_TE and is_te_min_version("1.9.0.dev0"): + _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR = ( + "use_grouped_tensor" in inspect.signature(te.pytorch.GroupedLinear.__init__).parameters + ) + class TEGroupedLinear(te.pytorch.GroupedLinear): """ Wrapper for the Transformer-Engine's `GroupedLinear` layer. @@ -2452,6 +2456,14 @@ def __init__( config, "moe_single_grouped_bias", False ) + if _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR: + extra_kwargs["use_grouped_tensor"] = config.moe_use_grouped_tensor + elif config.moe_use_grouped_tensor and not config.use_transformer_engine_op_fuser: + raise RuntimeError( + "moe_use_grouped_tensor=True requires a Transformer Engine GroupedLinear " + "that exposes the use_grouped_tensor argument." + ) + self.te_quant_params: Optional[TEQuantizationParams] = None quant_config = get_quant_config_or_none(name, config.quant_recipe) self.finish_init(quant_config) @@ -2959,6 +2971,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): ) else: + _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR = False TEGroupedLinear = None # type: ignore[assignment, misc] TEColumnParallelGroupedLinear = None # type: ignore[assignment, misc] TERowParallelGroupedLinear = None # type: ignore[assignment, misc] diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 35817f05e6a..721dc475c6d 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -301,25 +301,74 @@ def __init__( self.config.moe_mlp_glu_interleave_size, ) - if self.config.fp8 or self.config.fp4: - assert HAVE_TE, "FP8 and FP4 requires TE." - align_size = 256 if self._with_fused_impl else None + self._use_grouped_tensor = self.config.moe_use_grouped_tensor + if self.config.fp8 or self.config.fp4 or self._use_grouped_tensor: + assert HAVE_TE, "Quantized or TE grouped-tensor GroupedMLP execution requires TE." + align_size = ( + get_align_size_for_quantization(self.config) if self._use_grouped_tensor else None + ) self.quantization_padding = Fp8Padding(self.num_local_experts, align_size=align_size) self.quantization_unpadding = Fp8Unpadding( self.num_local_experts, align_size=align_size ) + @staticmethod + def _apply_packed_bias(intermediate_parallel, packed_bias, tokens_per_expert, permuted_probs): + """Apply a packed expert bias without reading token counts on the host.""" + # TODO: get rid of the .float() by having fused kernel compute in FP32 + shape = intermediate_parallel.shape + hidden_size = shape[-1] + output_dtype = intermediate_parallel.dtype + flat_output = intermediate_parallel.view(-1, hidden_size).float() + flat_probs = permuted_probs.reshape(-1, 1).float() + + if tokens_per_expert.device != packed_bias.device: + raise ValueError("Packed MoE bias and tokens_per_expert must be on the same device.") + + # Permutation stores tokens contiguously by expert. Repeat bias row e by that expert's + # token count to create one bias row per permuted token: + # + # packed_bias = [bias_e0, bias_e1] + # tokens_per_expert = [ 2, 1] + # bias_per_token = [bias_e0, bias_e0, bias_e1] + # + # output_size avoids a stream synchronization to compute sum(tokens_per_expert). + # Cast before repeating so both forward arithmetic and repeat_interleave's backward + # reduction are computed in FP32. Autograd casts the final parameter gradient once. + bias_per_token = torch.repeat_interleave( + packed_bias.float(), tokens_per_expert, dim=0, output_size=flat_output.size(0) + ) + return (flat_output + bias_per_token * flat_probs).view(shape).to(output_dtype) + @staticmethod def _apply_bias(intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs): if bias_parallel is None: return intermediate_parallel + + # CUDA-graph-safe packed path. With single_grouped_bias=True, TE returns one packed + # GroupedTensor [num_experts, hidden_size]. The grouped-tensor backend also provides + # tokens_per_expert as a tensor on the same device. + if isinstance(bias_parallel, torch.Tensor) and isinstance(tokens_per_expert, torch.Tensor): + return TEGroupedMLP._apply_packed_bias( + intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs + ) + + # Eager-only CPU-metadata path. The legacy contract returns List[Tensor[hidden_size]], + # and torch.split plus the Python zip below require concrete host token counts. A packed + # bias paired with Python counts also uses this compatibility path. Converting a tensor + # with .tolist() synchronizes and copies device data to the host, so this path must never + # be included in a CUDA graph. + if isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = tokens_per_expert.tolist() + shape = intermediate_parallel.shape + flat_output = intermediate_parallel.view(-1, shape[-1]) return ( torch.cat( [ t + b * p for t, b, p in zip( - torch.split(intermediate_parallel.view(-1, shape[-1]), tokens_per_expert), + torch.split(flat_output, tokens_per_expert), bias_parallel, torch.split(permuted_probs, tokens_per_expert), ) @@ -655,9 +704,21 @@ def _fused_forward( # Apply padding if needed unpadded_tokens_per_expert = None + # Some dispatchers have already padded each expert's token segment before the tokens + # reach this module: + # * router padding changes the routing map before dispatch; + # * HybridEP/NCCL-EP pad as part of their fused dispatch/permute operation; + # * DeepEP can pad in the fused local permutation after communication. + # Padding those tensors again would insert a second set of dummy tokens and make + # tokens_per_expert disagree with the already-permuted token buffer, so skip the local + # Fp8Padding fallback in those cases. if skip_routed_expert_padding(self.config): pass - elif self.config.fp8 or self.config.fp4: + # Regular AllToAll normally reaches this branch because its permutation does not insert + # the padding needed by the fused grouped-MLP contract. FP8/FP4 require recipe-specific + # alignment, while the TE operation-fuser grouped-tensor path currently uses 256-token + # expert segments. + elif self.config.fp8 or self.config.fp4 or self._use_grouped_tensor: tokens_per_expert = tokens_per_expert.tolist() unpadded_tokens_per_expert = tokens_per_expert permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( @@ -668,7 +729,22 @@ def _fused_forward( ) permuted_probs = permuted_probs.squeeze(-1) tokens_per_expert = torch.tensor( - tokens_per_expert, dtype=torch.int, device=permuted_probs.device + tokens_per_expert, dtype=torch.int64, device=permuted_probs.device + ) + + if self._use_grouped_tensor: + if not isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = torch.tensor( + tokens_per_expert, dtype=torch.int64, device=permuted_local_hidden_states.device + ) + else: + tokens_per_expert = tokens_per_expert.to( + device=permuted_local_hidden_states.device, dtype=torch.int64, non_blocking=True + ) + else: + raise RuntimeError( + "The Transformer Engine operation-fuser MoE path requires " + "moe_use_grouped_tensor=True." ) # if the number of tokens is 0, pad the hidden states to 256 @@ -787,11 +863,23 @@ def forward( # Apply padding if needed unpadded_tokens_per_expert = None - tokens_per_expert: list[int] = tokens_per_expert.tolist() permuted_probs = permuted_probs.unsqueeze(-1) + # The token buffer may already contain per-expert padding when padding was performed + # before expert compute: + # * router padding modified the routing map before dispatch; + # * HybridEP/NCCL-EP fused padding into dispatch/permute; + # * DeepEP fused padding into its post-communication local permutation. + # In those cases tokens_per_expert already describes the padded expert segments. Running + # Fp8Padding again would change the segment lengths without matching the existing token + # layout, so this module must leave both tensors unchanged. if skip_routed_expert_padding(self.config): pass - elif self.config.fp8 or self.config.fp4: + # Regular AllToAll normally supplies unpadded expert segments and therefore uses this + # explicit fallback. FP8/FP4 need their recipe-specific alignment. MCore currently also + # applies its common aligned-segment contract to the GroupedTensor backend so quantized + # grouped execution receives supported shapes + elif self.config.fp8 or self.config.fp4 or self._use_grouped_tensor: + tokens_per_expert = tokens_per_expert.tolist() unpadded_tokens_per_expert = tokens_per_expert permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( permuted_local_hidden_states, tokens_per_expert @@ -800,6 +888,18 @@ def forward( permuted_probs, unpadded_tokens_per_expert ) + if self._use_grouped_tensor: + if not isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = torch.tensor( + tokens_per_expert, dtype=torch.int64, device=permuted_local_hidden_states.device + ) + else: + tokens_per_expert = tokens_per_expert.to( + device=permuted_local_hidden_states.device, dtype=torch.int64, non_blocking=True + ) + elif isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = tokens_per_expert.tolist() + if self.config.moe_apply_probs_on_input: assert ( self.config.moe_router_topk == 1 diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 1d0768c4a97..355dcd9cd83 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1385,8 +1385,11 @@ def forward( inp = inp.view(-1, inp_shape[-1]) if te_general_gemm is not None and router_dtype != torch.float64: - output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=bias) - output = output[0] + # cuBLASLt's non-FP8 bias epilogue expects bias and output to have the same + # dtype. Router parameters may be BF16 while router logits are FP32, so cast the + # small bias vector before passing it to TE. + gemm_bias = bias.to(router_dtype) if bias is not None else None + output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=gemm_bias)[0] elif bias is None: output = torch.mm(inp.to(router_dtype), weight.to(router_dtype).t()) else: @@ -1464,22 +1467,33 @@ def get_align_size_for_quantization(config: TransformerConfig) -> int: Returns: int: The alignment size for quantization. """ - # CUTLASS kernel for grouped GEMM assumes 256 alignment. - if config.use_transformer_engine_op_fuser: + # TE's grouped-tensor and fused grouped-MLP kernels require 256-token alignment. + if config.use_transformer_engine_op_fuser or config.moe_use_grouped_tensor: return 256 if config.fp8: return get_fp8_align_size(config.fp8_recipe) if config.fp4: return get_fp4_align_size(config.fp4_recipe) - # Only FP8 or FP4 requires padding. Defaults to 0. + # Legacy high-precision grouped GEMM does not require padding. Defaults to 0. return 0 +def _deepep_permute_pads_grouped_tensor_input(config: TransformerConfig) -> bool: + """Whether DeepEP fused permutation pads input for TE grouped-tensor GEMM.""" + return ( + config.moe_use_grouped_tensor + and config.moe_token_dispatcher_type == "flex" + and config.moe_flex_dispatcher_backend == "deepep" + and config.moe_permute_fusion + and fused_permute_and_pad_with_probs is not None + ) + + def skip_routed_expert_padding(config: TransformerConfig) -> bool: """Whether the expert module should skip quantization padding. - Returns True when padding is already applied by the router or the - HybridEP / NCCL-EP dispatcher. + Returns True when padding is already applied by the router, the HybridEP / NCCL-EP + dispatcher, or DeepEP's fused permutation kernel. """ if config.moe_router_padding_for_quantization: return True @@ -1488,6 +1502,8 @@ def skip_routed_expert_padding(config: TransformerConfig) -> bool: "ncclep", ): return True + if _deepep_permute_pads_grouped_tensor_input(config): + return True return False diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 8fc4acedfef..73fbafc9d71 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1138,8 +1138,9 @@ def dispatch( "HybridEP only supports float32 probs, please set --moe-router-dtype=fp32" ) self.token_probs = self.token_probs.float() # downcast or upcast - if self.config.fp8 or self.config.fp4: - self.pad_multiple = get_align_size_for_quantization(self.config) + align_size = get_align_size_for_quantization(self.config) + if align_size > 0: + self.pad_multiple = align_size if self._padded_num_tokens is not None and hidden_states.shape[0] < self._padded_num_tokens: pad_rows = self._padded_num_tokens - hidden_states.shape[0] hidden_states = torch.cat( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2a35f8b585f..80f7bca5c49 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -834,21 +834,29 @@ class TransformerConfig(ModelParallelConfig): """ moe_grouped_gemm: bool = False - """When there are multiple experts per rank, compress multiple local (potentially small) gemms - in a single kernel launch to improve the utilization and performance by leveraging the Grouped - GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). + """Use grouped GEMM to execute multiple local MoE experts together. + + The concrete implementation is selected by Transformer Engine. Set + ``moe_use_grouped_tensor=True`` to use its CUDA-graph-safe GroupedTensor path. + """ + + moe_use_grouped_tensor: bool = False + """Use Transformer Engine's native GroupedTensor path for grouped MoE GEMMs. + + This path uses padded expert segments and CUDA split metadata so it can be captured in CUDA + graphs. Enabling the Transformer Engine operation fuser also enables this option. """ moe_single_grouped_weight: bool = False """When using TE GroupedLinear for MoE experts, store expert weights as a single grouped parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and - ``use_transformer_engine_op_fuser=True``. + ``moe_use_grouped_tensor=True``. """ moe_single_grouped_bias: bool = False """When using TE GroupedLinear for MoE experts, store expert biases as a single grouped - parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` - and ``add_bias_linear=True``.""" + parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``, + ``moe_use_grouped_tensor=True``, and ``add_bias_linear=True``.""" moe_aux_loss_coeff: Union[float, List[float]] = 0.0 """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. @@ -1321,6 +1329,12 @@ def __post_init__(self): self.experimental_attention_variant ) + if self.use_transformer_engine_op_fuser and self.moe_grouped_gemm: + self.moe_use_grouped_tensor = True + + if self.moe_use_grouped_tensor and not self.moe_grouped_gemm: + raise ValueError("moe_use_grouped_tensor=True requires moe_grouped_gemm=True.") + # When fp32 residual connections are enabled, pipeline parallel communication must # use fp32 to match the dtype of the residual stream between pipeline stages. if self.fp32_residual_connection and self.pipeline_dtype is not None: @@ -1634,15 +1648,12 @@ def __post_init__(self): "(--fp4-param-gather). Without FP4 parameter gather, Transformer Engine " "uses a split-quantize fallback that is being deprecated." ) - if not self.use_transformer_engine_op_fuser: - raise ValueError( - "moe_single_grouped_weight requires " - "use_transformer_engine_op_fuser=True. The non-op-fuser TE GroupedLinear " - "path splits the grouped parameter into per-expert tensors and does not " - "support single-grouped-weight training." - ) + if not self.moe_use_grouped_tensor: + raise ValueError("moe_single_grouped_weight requires moe_use_grouped_tensor=True.") if self.moe_single_grouped_bias and not self.add_bias_linear: raise ValueError("moe_single_grouped_bias requires add_bias_linear=True.") + if self.moe_single_grouped_bias and not self.moe_use_grouped_tensor: + raise ValueError("moe_single_grouped_bias requires moe_use_grouped_tensor=True.") if self.moe_enable_deepep: if self.moe_token_dispatcher_type != "flex": @@ -1670,6 +1681,12 @@ def __post_init__(self): "moe_flex_dispatcher_backend='ncclep' requires " "moe_token_dispatcher_type='flex'." ) + if self.moe_use_grouped_tensor and not self.use_transformer_engine_op_fuser: + raise ValueError( + "moe_use_grouped_tensor=True without use_transformer_engine_op_fuser is " + "not yet supported with the NCCL-EP dispatcher. Use the TE op-fuser path " + "or select the alltoall, DeepEP, or HybridEP dispatcher." + ) # moe_deepep_num_sms / moe_hybridep_num_sms are deprecated and unified into # moe_flex_dispatcher_num_sms. If either is set, route it (an explicit @@ -1756,10 +1773,11 @@ def __post_init__(self): if ( self.moe_flex_dispatcher_backend == "hybridep" and not self.use_transformer_engine_op_fuser + and not self.moe_use_grouped_tensor ): raise ValueError( "moe_expert_rank_capacity_factor with the 'hybridep' backend requires " - "use_transformer_engine_op_fuser to be enabled." + "use_transformer_engine_op_fuser=True or moe_use_grouped_tensor=True." ) if self.cpu_offloading and ( diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 6f3688b4385..e5b617cb931 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -187,6 +187,7 @@ "moe_flex_dispatcher_num_sms": None, "moe_grad_scale_func": None, "moe_grouped_gemm": True, + "moe_use_grouped_tensor": False, "moe_hybridep_num_sms": None, "moe_hybridep_num_sms_preprocessing": 108, "moe_hybridep_num_blocks_permute": None, diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 6b3531b090e..8cc1d96db2a 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -31,11 +31,41 @@ def test_op_fuser_transformer_config_args_are_exposed(): _add_network_size_args(parser) args = parser.parse_args( - ["--use-transformer-engine-op-fuser", "--moe-mlp-glu-interleave-size", "16"] + [ + "--use-transformer-engine-op-fuser", + "--moe-mlp-glu-interleave-size", + "16", + "--moe-use-grouped-tensor", + ] ) assert args.use_transformer_engine_op_fuser is True assert args.moe_mlp_glu_interleave_size == 16 + assert args.moe_use_grouped_tensor is True + + +def test_op_fuser_enables_grouped_tensor(): + config = TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + num_moe_experts=2, + moe_grouped_gemm=True, + use_transformer_engine_op_fuser=True, + ) + + assert config.moe_use_grouped_tensor is True + + +def test_grouped_tensor_requires_grouped_gemm(): + with pytest.raises(ValueError, match="requires moe_grouped_gemm=True"): + TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + num_moe_experts=2, + moe_use_grouped_tensor=True, + ) def test_clamped_swiglu_allows_te_op_fuser(): @@ -188,8 +218,12 @@ def __call__(self, *args): moe_router_padding_for_quantization=False, moe_token_dispatcher_type=None, moe_flex_dispatcher_backend=None, + moe_use_grouped_tensor=True, moe_paged_stash=False, ) + module._use_grouped_tensor = True + module.quantization_padding = lambda tensor, token_counts: (tensor, token_counts) + module.quantization_unpadding = lambda tensor, token_counts: tensor module._fused_ops = None module.linear_fc2 = SimpleNamespace(use_bias=fc2_bias) fused_ops = FakeFusedOps() @@ -203,11 +237,11 @@ def __call__(self, *args): torch.testing.assert_close(output, torch.ones_like(hidden_states)) assert module._fused_ops[0] is fused_ops assert fused_ops.args[0] is hidden_states - assert fused_ops.args[1] is tokens_per_expert - assert fused_ops.args[2] is probs - assert fused_ops.args[3] is tokens_per_expert + torch.testing.assert_close(fused_ops.args[1], tokens_per_expert) + torch.testing.assert_close(fused_ops.args[2], probs) + torch.testing.assert_close(fused_ops.args[3], tokens_per_expert) if fc2_bias: - assert fused_ops.args[4] is probs + torch.testing.assert_close(fused_ops.args[4], probs) else: assert len(fused_ops.args) == 4 @@ -237,6 +271,20 @@ def test_apply_bias_combines_per_expert_bias_and_probs(): assert output.dtype == intermediate.dtype +def test_apply_bias_combines_packed_grouped_bias_and_accumulates_gradient(): + intermediate = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + packed_bias = torch.tensor([[10.0, 20.0], [100.0, 200.0]], requires_grad=True) + tokens_per_expert = torch.tensor([2, 1], dtype=torch.int64) + permuted_probs = torch.tensor([0.25, 0.5, 1.5]) + expected = torch.tensor([[3.5, 7.0], [8.0, 14.0], [155.0, 306.0]]) + + output = TEGroupedMLP._apply_bias(intermediate, packed_bias, tokens_per_expert, permuted_probs) + output.sum().backward() + + torch.testing.assert_close(output, expected) + torch.testing.assert_close(packed_bias.grad, torch.tensor([[0.75, 0.75], [1.5, 1.5]])) + + def test_make_fused_impl_pre_forward_hook_dispatches_submodule_hooks(): module = TEGroupedMLP.__new__(TEGroupedMLP) torch.nn.Module.__init__(module) @@ -1171,7 +1219,8 @@ def test_gpu_make_fused_ops_constructs_with_real_te(self): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.internal - def test_gpu_fused_path_scales_fc2_bias(self): + @pytest.mark.parametrize("single_grouped_bias", (False, True)) + def test_gpu_fused_path_scales_fc2_bias(self, monkeypatch, single_grouped_bias): """FC2 bias and its gradients must use the per-token router probability.""" try: from transformer_engine.pytorch.ops import GroupedLinear @@ -1181,6 +1230,8 @@ def test_gpu_fused_path_scales_fc2_bias(self): if "scale_bias" not in inspect.signature(GroupedLinear.__init__).parameters: pytest.skip("Installed TE op fuser GroupedLinear lacks `scale_bias` support") + if single_grouped_bias: + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") Utils.destroy_model_parallel() Utils.initialize_model_parallel(1, 1) @@ -1202,6 +1253,7 @@ def test_gpu_fused_path_scales_fc2_bias(self): moe_router_topk=1, moe_grouped_gemm=True, use_transformer_engine_op_fuser=True, + moe_single_grouped_bias=single_grouped_bias, ) _set_random_seed(seed_=123, data_parallel_random_init=False) submodules = get_submodules( @@ -1219,9 +1271,19 @@ def test_gpu_fused_path_scales_fc2_bias(self): for linear in (experts.linear_fc1, experts.linear_fc2): for expert_idx in range(self.num_experts): getattr(linear, f"weight{expert_idx}").zero_() - getattr(linear, f"bias{expert_idx}").zero_() - experts.linear_fc2.bias0.fill_(2.0) - experts.linear_fc2.bias1.fill_(4.0) + if not single_grouped_bias: + getattr(linear, f"bias{expert_idx}").zero_() + if single_grouped_bias: + linear.bias.rowwise_data.zero_() + if single_grouped_bias: + packed_fc2_bias = experts.linear_fc2.bias.rowwise_data.view( + self.num_experts, self.hidden_size + ) + packed_fc2_bias[0].fill_(2.0) + packed_fc2_bias[1].fill_(4.0) + else: + experts.linear_fc2.bias0.fill_(2.0) + experts.linear_fc2.bias1.fill_(4.0) hidden_states = torch.zeros( 3, self.hidden_size, dtype=torch.bfloat16, device="cuda", requires_grad=True @@ -1245,15 +1307,201 @@ def test_gpu_fused_path_scales_fc2_bias(self): torch.tensor([2.0, 2.0, 4.0], dtype=torch.bfloat16, device="cuda") * self.hidden_size ) torch.testing.assert_close(probs.grad, expected_prob_grad) - torch.testing.assert_close( - experts.linear_fc2.bias0.grad, - torch.ones_like(experts.linear_fc2.bias0) * probs[:2].detach().sum(), + if single_grouped_bias: + assert experts.linear_fc2.bias.grad is not None + packed_dbias = experts.linear_fc2.bias.grad.view(self.num_experts, self.hidden_size) + torch.testing.assert_close( + packed_dbias[0], torch.ones_like(packed_dbias[0]) * probs[:2].detach().sum() + ) + torch.testing.assert_close( + packed_dbias[1], torch.ones_like(packed_dbias[1]) * probs[2:].detach().sum() + ) + assert experts._fused_ops[0][2].bias is experts.linear_fc2.bias + else: + torch.testing.assert_close( + experts.linear_fc2.bias0.grad, + torch.ones_like(experts.linear_fc2.bias0) * probs[:2].detach().sum(), + ) + torch.testing.assert_close( + experts.linear_fc2.bias1.grad, + torch.ones_like(experts.linear_fc2.bias1) * probs[2:].detach().sum(), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.internal + @pytest.mark.parametrize("use_op_fuser", (False, True), ids=("module", "op-fuser")) + @pytest.mark.parametrize( + "single_grouped_weight,single_grouped_bias", + ((True, False), (False, True), (True, True)), + ids=("single-weight", "single-bias", "single-weight-and-bias"), + ) + def test_gpu_single_grouped_parent_gradient_parity( + self, monkeypatch, use_op_fuser, single_grouped_weight, single_grouped_bias + ): + """Single grouped parents must receive the same gradients as discrete parameters. + + This is an integration test over the real MCore TEGroupedMLP wrapper. It covers both + gradient ownership mechanisms: + + * the module path returns the packed FC2 bias to ``_apply_packed_bias``, where normal + PyTorch autograd must update the registered grouped parent; + * the op-fuser path reattaches the same wrapper parameters to TE op shells, whose custom + backward must return packed wgrad/dbias in the corresponding parent slots. + + Comparing gradients directly is intentional. A forward or dgrad-only check would not + catch a disconnected grouped parent that the optimizer can never update. + """ + try: + from transformer_engine.pytorch.module import GroupedLinear as ModuleGroupedLinear + from transformer_engine.pytorch.ops import GroupedLinear as OpGroupedLinear + except ImportError: + pytest.skip("Required TE GroupedLinear APIs are not available") + import inspect + + module_parameters = inspect.signature(ModuleGroupedLinear.__init__).parameters + op_parameters = inspect.signature(OpGroupedLinear.__init__).parameters + if ( + "use_grouped_tensor" not in module_parameters + or "single_grouped_bias" not in module_parameters + or "single_grouped_bias" not in op_parameters + ): + pytest.skip("Installed TE lacks native single grouped bias support") + + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + def build_experts(single_weight, single_bias): + config = TransformerConfig( + num_layers=1, + hidden_size=self.hidden_size, + num_attention_heads=4, + num_moe_experts=self.num_experts, + use_cpu_initialization=False, + add_bias_linear=True, + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=False, + bias_dropout_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="sinkhorn", + moe_router_topk=1, + moe_grouped_gemm=True, + moe_use_grouped_tensor=True, + use_transformer_engine_op_fuser=use_op_fuser, + moe_single_grouped_weight=single_weight, + moe_single_grouped_bias=single_bias, + ) + submodules = get_submodules( + get_gpt_layer_with_transformer_engine_submodules( + self.num_experts, moe_grouped_gemm=True + ).mlp + ) + assert isinstance(submodules, MoESubmodules) + layer = MoELayer(config, submodules) + layer = Float16Module(layer.config, layer).module + layer.cuda() + assert isinstance(layer.experts, TEGroupedMLP) + return layer.experts + + def copy_linear_params(reference_linear, target_linear): + reference_weights = torch.stack( + [ + getattr(reference_linear, f"weight{idx}").detach() + for idx in range(self.num_experts) + ] + ) + reference_biases = torch.stack( + [ + getattr(reference_linear, f"bias{idx}").detach() + for idx in range(self.num_experts) + ] + ) + with torch.no_grad(): + if target_linear.single_grouped_weight: + target_linear.weight.rowwise_data.view_as(reference_weights).copy_( + reference_weights + ) + else: + for idx in range(self.num_experts): + getattr(target_linear, f"weight{idx}").copy_(reference_weights[idx]) + + if target_linear.single_grouped_bias: + target_linear.bias.rowwise_data.view_as(reference_biases).copy_( + reference_biases + ) + else: + for idx in range(self.num_experts): + getattr(target_linear, f"bias{idx}").copy_(reference_biases[idx]) + + def packed_grad(linear, name): + if getattr(linear, f"single_grouped_{name}"): + grad = getattr(linear, name).grad + assert grad is not None, f"Grouped {name} parent did not receive a gradient" + return grad.float() + grads = [getattr(linear, f"{name}{idx}").grad for idx in range(self.num_experts)] + assert all(grad is not None for grad in grads) + return torch.stack(grads).float() + + torch.manual_seed(1234) + reference = build_experts(False, False) + torch.manual_seed(5678) + target = build_experts(single_grouped_weight, single_grouped_bias) + copy_linear_params(reference.linear_fc1, target.linear_fc1) + copy_linear_params(reference.linear_fc2, target.linear_fc2) + + tokens_per_expert = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + num_tokens = int(tokens_per_expert.sum().item()) + base_input = 0.1 * torch.randn( + num_tokens, self.hidden_size, dtype=torch.bfloat16, device="cuda" ) - torch.testing.assert_close( - experts.linear_fc2.bias1.grad, - torch.ones_like(experts.linear_fc2.bias1) * probs[2:].detach().sum(), + base_probs = torch.rand(num_tokens, dtype=torch.bfloat16, device="cuda") + grad_output = 0.1 * torch.randn( + num_tokens, self.hidden_size, dtype=torch.bfloat16, device="cuda" ) + reference_input = base_input.detach().clone().requires_grad_(True) + reference_probs = base_probs.detach().clone().requires_grad_(True) + reference_output, _ = reference(reference_input, tokens_per_expert, reference_probs) + reference_output.backward(grad_output) + + target_input = base_input.detach().clone().requires_grad_(True) + target_probs = base_probs.detach().clone().requires_grad_(True) + target_output, _ = target(target_input, tokens_per_expert, target_probs) + target_output.backward(grad_output) + + tolerances = {"rtol": 1e-2, "atol": 1e-2} + torch.testing.assert_close(target_output, reference_output, **tolerances) + torch.testing.assert_close(target_input.grad, reference_input.grad, **tolerances) + torch.testing.assert_close(target_probs.grad, reference_probs.grad, **tolerances) + + for target_linear, reference_linear in ( + (target.linear_fc1, reference.linear_fc1), + (target.linear_fc2, reference.linear_fc2), + ): + torch.testing.assert_close( + packed_grad(target_linear, "weight"), + packed_grad(reference_linear, "weight"), + **tolerances, + ) + torch.testing.assert_close( + packed_grad(target_linear, "bias"), + packed_grad(reference_linear, "bias"), + **tolerances, + ) + + if use_op_fuser: + ops = target._fused_ops[0] + fc1_op = ops[0] + fc2_op = ops[2] + # The fused-op shells must register MCore's original parent parameters, not + # detached copies or member views, so autograd and the optimizer update the same objects. + if single_grouped_weight: + assert fc1_op.weight is target.linear_fc1.weight + assert fc2_op.weight is target.linear_fc2.weight + if single_grouped_bias: + assert fc1_op.bias is target.linear_fc1.bias + assert fc2_op.bias is target.linear_fc2.bias + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.internal def test_gpu_fused_path_loss_decreases(self): diff --git a/tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py b/tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py new file mode 100644 index 00000000000..0e60e64c809 --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py @@ -0,0 +1,599 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Distributed MoE coverage for the TE grouped-tensor expert path. + +The dispatchers place expert-token padding at different boundaries: + +* All-to-All returns unpadded expert segments and TEGroupedMLP pads them before FC1. +* DeepEP communicates first, then its local fused permutation pads expert segments. +* HybridEP fuses communication, permutation, and expert-segment padding. +* NCCL-EP returns aligned expert segments from fused dispatch, like HybridEP. Its non-op-fuser + grouped-tensor integration is not enabled yet, so its parity and lifecycle cases remain skipped. + +The numerical tests compare each grouped-tensor configuration with the old discrete-parameter, +CPU-split path on the same dispatcher. The lifecycle tests inspect the actual expert-compute +boundary to ensure padding rows are zero and the dispatcher-specific inverse removes them. +""" + +import inspect +import os +from typing import Dict + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core import config as mcore_config +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_submodules, +) +from megatron.core.transformer.module import Float16Module +from megatron.core.transformer.moe.fused_a2a import ( + HAVE_DEEP_EP, + HAVE_HYBRIDEP, + reset_hybrid_ep_buffer, +) +from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules +from megatron.core.transformer.moe.moe_utils import fused_permute_and_pad_with_probs +from megatron.core.transformer.spec_utils import get_submodules +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training.initialize import _set_random_seed +from tests.unit_tests.test_utilities import Utils + +pytestmark = [ + pytest.mark.internal, + pytest.mark.launch_on_gb200, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), +] + + +_ALIGN_SIZE = 256 +_HIDDEN_SIZE = 256 +_MOE_FFN_HIDDEN_SIZE = 256 +_NUM_LOCAL_EXPERTS = 2 +_NUM_LOCAL_TOKENS = 128 +_TOLERANCES = {"rtol": 1e-2, "atol": 1e-2} +_NCCL_EP_GROUPED_TENSOR_UNSUPPORTED_REASON = ( + "NCCL-EP support for the non-op-fuser grouped-tensor expert path is not implemented yet" +) + +_PARAMETER_LAYOUTS = ( + pytest.param(False, False, False, id="discrete-weight-no-bias"), + pytest.param(True, False, False, id="single-weight-no-bias"), + pytest.param(False, True, False, id="discrete-weight-discrete-bias"), + pytest.param(True, True, False, id="single-weight-discrete-bias"), + pytest.param(False, True, True, id="discrete-weight-single-bias"), + pytest.param(True, True, True, id="single-weight-single-bias"), +) + + +def _require_test_environment(dispatcher: str) -> int: + """Validate runtime support and return the EP world size used by the test.""" + world_size = torch.distributed.get_world_size() + if world_size < 2: + pytest.skip("DeepEP/HybridEP parity requires at least two distributed ranks") + if dispatcher == "deepep" and not HAVE_DEEP_EP: + pytest.skip("DeepEP is not available") + if dispatcher == "hybridep" and not HAVE_HYBRIDEP: + pytest.skip("HybridEP is not available") + if dispatcher == "deepep" and fused_permute_and_pad_with_probs is None: + pytest.skip("DeepEP grouped-tensor padding requires TE fused permute-and-pad") + + try: + from transformer_engine.pytorch.module import GroupedLinear + except ImportError: + pytest.skip("Transformer Engine GroupedLinear is not available") + parameters = inspect.signature(GroupedLinear.__init__).parameters + if "use_grouped_tensor" not in parameters or "single_grouped_bias" not in parameters: + pytest.skip("Installed TE lacks native grouped-tensor parameter support") + return world_size + + +def _dispatcher_options(dispatcher: str) -> Dict[str, object]: + """Return the production padding configuration for one dispatcher.""" + if dispatcher == "alltoall": + return { + "moe_token_dispatcher_type": "alltoall", + "moe_flex_dispatcher_backend": None, + "moe_permute_fusion": True, + } + if dispatcher == "deepep": + # DeepEP communicates first. Its fused local permutation then groups and pads tokens. + return { + "moe_token_dispatcher_type": "flex", + "moe_flex_dispatcher_backend": "deepep", + "moe_permute_fusion": True, + } + if dispatcher == "hybridep": + # HybridEP owns permutation and padding inside its fused dispatch/combine kernels. + return { + "moe_token_dispatcher_type": "flex", + "moe_flex_dispatcher_backend": "hybridep", + "moe_permute_fusion": True, + } + if dispatcher == "ncclep": + # NCCL-EP dispatch packs aligned expert segments, but the module grouped-tensor path is + # intentionally config-rejected until its end-to-end numerics and padding are validated. + return { + "moe_token_dispatcher_type": "flex", + "moe_flex_dispatcher_backend": "ncclep", + "moe_permute_fusion": True, + "moe_expert_rank_capacity_factor": 8.0, + # The lifecycle assertions inspect the dynamically narrowed expert buffer. Static + # NCCL-EP intentionally exposes the full receive-capacity buffer instead. + "moe_ncclep_static_shape": False, + } + raise ValueError(f"Unknown dispatcher {dispatcher!r}") + + +def _build_moe_layer( + dispatcher: str, + *, + ep_size: int, + use_grouped_tensor: bool, + single_grouped_weight: bool, + use_bias: bool, + single_grouped_bias: bool, +) -> MoELayer: + """Build a small real TE MoE layer without using the TE operation fuser.""" + options = _dispatcher_options(dispatcher) + transformer_config = TransformerConfig( + num_layers=1, + hidden_size=_HIDDEN_SIZE, + num_attention_heads=8, + num_moe_experts=ep_size * _NUM_LOCAL_EXPERTS, + moe_ffn_hidden_size=_MOE_FFN_HIDDEN_SIZE, + use_cpu_initialization=False, + add_bias_linear=use_bias, + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=False, + bias_dropout_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="none", + moe_router_topk=2, + moe_aux_loss_coeff=0.0, + moe_router_dtype="fp32", + moe_grouped_gemm=True, + moe_use_grouped_tensor=use_grouped_tensor, + moe_single_grouped_weight=single_grouped_weight, + moe_single_grouped_bias=single_grouped_bias, + use_transformer_engine_op_fuser=False, + tensor_model_parallel_size=1, + expert_model_parallel_size=ep_size, + sequence_parallel=False, + **options, + ) + submodules = get_submodules( + get_gpt_layer_with_transformer_engine_submodules( + num_experts=transformer_config.num_moe_experts, moe_grouped_gemm=True + ).mlp + ) + assert isinstance(submodules, MoESubmodules) + layer = MoELayer(transformer_config, submodules) + layer = Float16Module(layer.config, layer).module + layer.cuda() + layer.set_layer_number(0) + return layer + + +def _copy_linear_parameters(reference, target) -> None: + """Copy discrete expert parameters into either a discrete or packed target layout.""" + for parameter_name in ("weight", "bias"): + if parameter_name == "bias" and not reference.use_bias: + continue + reference_parts = torch.stack( + [ + getattr(reference, f"{parameter_name}{idx}").detach() + for idx in range(reference.num_gemms) + ] + ) + target_is_grouped = getattr(target, f"single_grouped_{parameter_name}") + if target_is_grouped: + grouped_parameter = getattr(target, parameter_name) + grouped_parameter.rowwise_data.view_as(reference_parts).copy_(reference_parts) + else: + for idx, part in enumerate(reference_parts): + getattr(target, f"{parameter_name}{idx}").copy_(part) + + +@torch.no_grad() +def _copy_layer_parameters(reference: MoELayer, target: MoELayer) -> None: + """Give reference and target identical router and expert parameters.""" + target_parameters = dict(target.named_parameters()) + for name, parameter in reference.named_parameters(): + if not name.startswith("experts."): + target_parameters[name].copy_(parameter) + _copy_linear_parameters(reference.experts.linear_fc1, target.experts.linear_fc1) + _copy_linear_parameters(reference.experts.linear_fc2, target.experts.linear_fc2) + + +def _canonical_gradient(linear, parameter_name: str) -> torch.Tensor: + """Return expert gradients as one [experts, ...] FP32 tensor for either layout.""" + if getattr(linear, f"single_grouped_{parameter_name}"): + gradient = getattr(linear, parameter_name).grad + assert gradient is not None, f"Grouped {parameter_name} parent has no gradient" + return gradient.reshape(linear.num_gemms, -1).float() + + gradients = [getattr(linear, f"{parameter_name}{idx}").grad for idx in range(linear.num_gemms)] + assert all(gradient is not None for gradient in gradients) + return torch.stack([gradient.reshape(-1) for gradient in gradients]).float() + + +def _run_forward_backward( + layer: MoELayer, base_input: torch.Tensor, grad_output: torch.Tensor +) -> Dict[str, torch.Tensor]: + """Run one MoE step and collect values sensitive to dispatch and parameter layout.""" + layer.zero_grad(set_to_none=True) + hidden_states = base_input.detach().clone().requires_grad_(True) + output, _ = layer(hidden_states) + output.backward(grad_output) + + result = { + "output": output.detach(), + "input_grad": hidden_states.grad.detach(), + "router_grad": layer.router.weight.grad.detach(), + "fc1_weight_grad": _canonical_gradient(layer.experts.linear_fc1, "weight"), + "fc2_weight_grad": _canonical_gradient(layer.experts.linear_fc2, "weight"), + } + if layer.config.add_bias_linear: + result["fc1_bias_grad"] = _canonical_gradient(layer.experts.linear_fc1, "bias") + result["fc2_bias_grad"] = _canonical_gradient(layer.experts.linear_fc2, "bias") + return result + + +def _run_numerical_parity_case( + dispatcher: str, *, single_grouped_weight: bool, use_bias: bool, single_grouped_bias: bool +) -> None: + """Compare grouped-tensor execution with the old path on the same dispatcher.""" + ep_size = _require_test_environment(dispatcher) + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, expert_model_parallel_size=ep_size + ) + mcore_config.ENABLE_EXPERIMENTAL = True + + _set_random_seed(seed_=1234, data_parallel_random_init=False) + reference = _build_moe_layer( + dispatcher, + ep_size=ep_size, + use_grouped_tensor=False, + single_grouped_weight=False, + use_bias=use_bias, + single_grouped_bias=False, + ) + target = _build_moe_layer( + dispatcher, + ep_size=ep_size, + use_grouped_tensor=True, + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + _copy_layer_parameters(reference, target) + + base_input = torch.randn( + _NUM_LOCAL_TOKENS, 1, _HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda" + ) + grad_output = torch.randn_like(base_input) + + reference_result = _run_forward_backward(reference, base_input, grad_output) + target_result = _run_forward_backward(target, base_input, grad_output) + + assert target.experts._use_grouped_tensor + assert not reference.experts._use_grouped_tensor + assert reference_result.keys() == target_result.keys() + for name in reference_result: + torch.testing.assert_close( + target_result[name], + reference_result[name], + msg=lambda message, value=name: f"{dispatcher} {value} mismatch: {message}", + **_TOLERANCES, + ) + + +def _make_padding_mask( + real_tokens_per_expert: torch.Tensor, padded_tokens_per_expert: torch.Tensor +) -> torch.Tensor: + """Build a mask for the padding suffix in every expert-major segment. + + For example, using an alignment of four for readability: + + ``real_tokens_per_expert = tensor([3, 2])`` + ``padded_tokens_per_expert = tensor([4, 4])`` + + The packed expert-major rows are ``[e0 real x3][e0 pad x1][e1 real x2][e1 pad x2]``, + so this function returns ``[F, F, F, T, F, F, T, T]``. Indexing the packed hidden states or + probabilities with that mask selects only the synthetic rows that must contain exact zeros. + """ + masks = [] + for real_count, padded_count in zip( + real_tokens_per_expert.cpu().tolist(), padded_tokens_per_expert.cpu().tolist() + ): + # Each expert contributes a False prefix for real rows followed by a True padding suffix. + masks.append(torch.zeros(real_count, dtype=torch.bool, device="cuda")) + masks.append(torch.ones(padded_count - real_count, dtype=torch.bool, device="cuda")) + return torch.cat(masks) + + +def _infer_real_tokens_per_expert( + padded_probs: torch.Tensor, padded_tokens_per_expert: torch.Tensor +) -> torch.Tensor: + """Infer real expert counts from the zero suffix in every padded probability segment. + + This dropless test uses FP32 softmax router probabilities, so every routed token has a + nonzero probability. All padding implementations use exact zeros and append them after the + real rows in each expert-major segment. This gives one dispatcher-independent source of truth + without inspecting backend-specific routing metadata. + """ + real_counts = [] + offset = 0 + for padded_count in padded_tokens_per_expert.cpu().tolist(): + segment = padded_probs[offset : offset + padded_count] + nonzero_rows = segment != 0 + real_count = int(nonzero_rows.sum().item()) + + # Padding must be one contiguous suffix. Interspersed zero rows would preserve the total + # nonzero count while still violating the expert-major layout expected by grouped GEMM. + assert torch.all(nonzero_rows[:real_count]) + assert not torch.any(nonzero_rows[real_count:]) + real_counts.append(real_count) + offset += padded_count + + assert offset == padded_probs.numel() + return torch.tensor(real_counts, dtype=torch.int64, device=padded_probs.device) + + +def _install_padding_probes(layer: MoELayer, monkeypatch): + """Observe the tensors crossing each padding boundary without changing execution. + + ``register_forward_pre_hook`` runs immediately before the selected module's ``forward``. + The hook receives the module and the tuple of positional arguments that forward is about to + consume. Returning ``None`` leaves those arguments unchanged, so these hooks are read-only + probes rather than replacements for any production operation. + + There are two relevant boundaries. ``layer.experts`` sees what the token dispatcher hands to + TEGroupedMLP, while ``linear_fc1`` sees the final padded tensors and CUDA expert counts that + TE's grouped-tensor GEMM actually consumes. They are different for AllToAll, where TEGroupedMLP + owns padding, but identical for DeepEP and HybridEP, whose fused dispatch paths already pad. + """ + captured = {"padding_calls": 0, "unpadding_calls": 0} + + def capture_dispatcher_input(_module, args): + # TEGroupedMLP.forward(hidden, tokens_per_expert, permuted_probs) is about to run. + # DeepEP and HybridEP have already padded at this boundary, so preserve their router + # probabilities for the generic real-row inference below. Detach so the test does not + # retain the graph, and clone in case downstream computation reuses the input storage. + captured["dispatcher_probs"] = args[2].detach().clone() + + def capture_fc1_input(_module, args): + # GroupedLinear.forward(hidden, m_splits, ...) is about to run. This is the authoritative + # view of the rows and device-side split tensor presented to the grouped GEMM. + captured["padded_hidden"] = args[0].detach().clone() + captured["padded_counts"] = args[1].detach().clone() + + # Pre-hooks observe module inputs before either module can transform them. The returned hook + # handles need not be retained because each test owns this layer and executes one forward. + layer.experts.register_forward_pre_hook(capture_dispatcher_input) + layer.experts.linear_fc1.register_forward_pre_hook(capture_fc1_input) + + def capture_padding(_module, args, output): + # quantization_padding is called once for hidden states and once for router probabilities + # in the AllToAll path. A forward hook is used here because the padded tensor is its output. + captured["padding_calls"] += 1 + padded_tensor = output[0] + # Probability padding receives [tokens, 1], whereas hidden-state padding receives + # [tokens, hidden_size]. Preserve the padded probabilities for an exact-zero assertion. + if args[0].shape[-1] == 1: + captured["padded_probs"] = padded_tensor.detach().clone().reshape(-1) + + def capture_unpadding(_module, _args, output): + # AllToAll uses TEGroupedMLP's explicit unpadding module after expert compute. Recording + # the call proves that its locally inserted padding reaches the matching removal path. + captured["unpadding_calls"] += 1 + + layer.experts.quantization_padding.register_forward_hook(capture_padding) + layer.experts.quantization_unpadding.register_forward_hook(capture_unpadding) + + # All dispatchers share this final restoration API. + original_combine_postprocess = layer.token_dispatcher.combine_postprocess + + def capture_combine_postprocess(hidden_states, *args, **kwargs): + output = original_combine_postprocess(hidden_states, *args, **kwargs) + captured["restored_shape"] = output.shape + return output + + monkeypatch.setattr(layer.token_dispatcher, "combine_postprocess", capture_combine_postprocess) + + return captured + + +def _run_padding_lifecycle_case(dispatcher: str, monkeypatch) -> None: + """Verify exact zero padding, 256 alignment, no double-padding, and unpadding. + + This test follows one real MoE forward through dispatch, expert compute, and restoration. It + independently reconstructs the number of real tokens assigned to each local expert, then + compares that metadata with the padded CUDA ``m_splits`` observed directly at FC1. Numerical + parity with the legacy backend is tested separately; this case focuses on padding ownership + and the structural contract required by TE's grouped-tensor kernels. + """ + # EP spans the whole torchrun world so every tested dispatcher performs real communication. + # TP remains one to keep the independently reconstructed local-expert counts unambiguous. + ep_size = _require_test_environment(dispatcher) + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, expert_model_parallel_size=ep_size + ) + mcore_config.ENABLE_EXPERIMENTAL = True + + # Use the strictest parameter layout. If packed weight and packed bias reach the native + # grouped-tensor path correctly, discrete parameter layouts use the same padding lifecycle. + _set_random_seed(seed_=1357, data_parallel_random_init=False) + layer = _build_moe_layer( + dispatcher, + ep_size=ep_size, + use_grouped_tensor=True, + single_grouped_weight=True, + use_bias=True, + single_grouped_bias=True, + ) + # Install observers before the forward so they capture dispatcher output, FC1 input, and the + # common dispatcher boundary that returns to the original token layout. + captured = _install_padding_probes(layer, monkeypatch) + + torch.manual_seed(9753) + hidden_states = torch.randn( + _NUM_LOCAL_TOKENS, 1, _HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + output, _ = layer(hidden_states) + + padded_counts = captured["padded_counts"] + + # TE's grouped-tensor API requires device-resident int64 splits. Every expert segment must be + # represented by its m_split, and the physical FC1 input must contain their total. + assert padded_counts.device.type == "cuda" + assert padded_counts.dtype == torch.int64 + assert torch.all(padded_counts % _ALIGN_SIZE == 0) + assert captured["padded_hidden"].shape[0] == padded_counts.sum().item() + + if dispatcher == "alltoall": + # All-to-All returns real expert rows; TEGroupedMLP pads hidden states and probs itself. + assert captured["padding_calls"] == 2 + assert captured["unpadding_calls"] == 1 + else: + # DeepEP and HybridEP already return padded rows. TEGroupedMLP must not pad them again. + assert captured["padding_calls"] == 0 + assert captured["unpadding_calls"] == 0 + # For fused dispatchers, probabilities are already padded when they enter TEGroupedMLP, + # so the experts pre-hook is the correct observation point for the zero check below. + captured["padded_probs"] = captured["dispatcher_probs"].reshape(-1) + + # Router probabilities provide a common representation across all dispatchers: real routed + # rows are nonzero and padded rows are an exact-zero suffix. + real_counts = _infer_real_tokens_per_expert(captured["padded_probs"], padded_counts) + expected_padded_counts = ((real_counts + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + torch.testing.assert_close(padded_counts, expected_padded_counts, rtol=0, atol=0) + + # Expert-major layout is [expert 0 real][expert 0 pad][expert 1 real][expert 1 pad]... + # Build that exact mask and require both hidden states and routing probabilities to use + # numerical zero for every synthetic row. Merely allocating the right shape is insufficient. + padding_mask = _make_padding_mask(real_counts, padded_counts) + # An expert may already have a 256-aligned token count and legitimately need no padding. + # Boolean indexing with an empty mask is valid; otherwise these checks inspect every pad row. + assert not torch.any(captured["padded_hidden"][padding_mask]) + assert not torch.any(captured["padded_probs"][padding_mask]) + + # Regardless of where a backend removes padding, the common dispatcher postprocess contract + # must restore exactly the shape that entered this MoE layer. + assert captured["restored_shape"] == hidden_states.shape + + # The public MoE contract is unchanged by internal alignment. Run backward as a final check + # that padding/unpadding preserved a connected, finite autograd path to the original input. + assert output.shape == hidden_states.shape + output.float().square().mean().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() + + +class TestGroupedTensorDispatcherNumerics: + """Distributed numerical and padding coverage for grouped-tensor dispatchers.""" + + def setup_method(self, method): + if not torch.distributed.is_available() or Utils.world_size < 2: + pytest.skip("Distributed dispatcher tests must be launched with torchrun") + self._old_single_param_env = os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") + self._previous_experimental = mcore_config.ENABLE_EXPERIMENTAL + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1" + Utils.initialize_distributed() + + def teardown_method(self, method): + try: + mcore_config.ENABLE_EXPERIMENTAL = self._previous_experimental + reset_hybrid_ep_buffer() + Utils.destroy_model_parallel() + finally: + if self._old_single_param_env is None: + os.environ.pop("NVTE_GROUPED_LINEAR_SINGLE_PARAM", None) + else: + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = self._old_single_param_env + + @pytest.mark.parametrize( + "single_grouped_weight,use_bias,single_grouped_bias", _PARAMETER_LAYOUTS + ) + @pytest.mark.timeout(180) + def test_alltoall_grouped_tensor_moe_parity( + self, single_grouped_weight, use_bias, single_grouped_bias + ): + """All-to-All grouped-tensor MoE forward/backward matches its legacy expert path.""" + _run_numerical_parity_case( + "alltoall", + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + + @pytest.mark.parametrize( + "single_grouped_weight,use_bias,single_grouped_bias", _PARAMETER_LAYOUTS + ) + @pytest.mark.timeout(180) + def test_deepep_grouped_tensor_moe_parity( + self, single_grouped_weight, use_bias, single_grouped_bias + ): + """DeepEP grouped-tensor MoE forward/backward matches its legacy expert path.""" + _run_numerical_parity_case( + "deepep", + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + + @pytest.mark.parametrize( + "single_grouped_weight,use_bias,single_grouped_bias", _PARAMETER_LAYOUTS + ) + @pytest.mark.timeout(180) + def test_hybridep_grouped_tensor_moe_parity( + self, single_grouped_weight, use_bias, single_grouped_bias + ): + """HybridEP grouped-tensor MoE forward/backward matches its legacy expert path.""" + _run_numerical_parity_case( + "hybridep", + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + + @pytest.mark.skip(reason=_NCCL_EP_GROUPED_TENSOR_UNSUPPORTED_REASON) + @pytest.mark.parametrize( + "single_grouped_weight,use_bias,single_grouped_bias", _PARAMETER_LAYOUTS + ) + @pytest.mark.timeout(180) + def test_ncclep_grouped_tensor_moe_parity( + self, single_grouped_weight, use_bias, single_grouped_bias + ): + """NCCL-EP grouped-tensor parity coverage reserved for future enablement.""" + _run_numerical_parity_case( + "ncclep", + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + + @pytest.mark.timeout(180) + def test_alltoall_grouped_tensor_padding_lifecycle(self, monkeypatch): + """All-to-All explicitly pads in TEGroupedMLP and removes it before combine.""" + _run_padding_lifecycle_case("alltoall", monkeypatch) + + @pytest.mark.timeout(180) + def test_deepep_grouped_tensor_padding_lifecycle(self, monkeypatch): + """DeepEP fused local permutation pads, and local unpermute removes those rows.""" + _run_padding_lifecycle_case("deepep", monkeypatch) + + @pytest.mark.timeout(180) + def test_hybridep_grouped_tensor_padding_lifecycle(self, monkeypatch): + """HybridEP fused dispatch pads, and fused combine returns the original token shape.""" + _run_padding_lifecycle_case("hybridep", monkeypatch) + + @pytest.mark.skip(reason=_NCCL_EP_GROUPED_TENSOR_UNSUPPORTED_REASON) + @pytest.mark.timeout(180) + def test_ncclep_grouped_tensor_padding_lifecycle(self, monkeypatch): + """NCCL-EP padding lifecycle coverage reserved for future enablement.""" + _run_padding_lifecycle_case("ncclep", monkeypatch) diff --git a/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py b/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py index 0f4eb35c510..c4a50be0432 100644 --- a/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py +++ b/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py @@ -54,8 +54,12 @@ _TE_GROUPED_LINEAR_SUPPORTS_SINGLE_PARAM = ( "single_grouped_weight" in inspect.signature(TEGroupedLinear.__init__).parameters ) + _TE_GROUPED_LINEAR_SUPPORTS_USE_GROUPED_TENSOR = ( + "use_grouped_tensor" in inspect.signature(TEGroupedLinear.__init__).parameters + ) except (ImportError, AttributeError): _TE_GROUPED_LINEAR_SUPPORTS_SINGLE_PARAM = False + _TE_GROUPED_LINEAR_SUPPORTS_USE_GROUPED_TENSOR = False pytestmark = [ pytest.mark.internal, @@ -68,6 +72,10 @@ not _TE_GROUPED_LINEAR_SUPPORTS_SINGLE_PARAM, reason="Installed TE GroupedLinear does not expose single_grouped_weight", ), + pytest.mark.skipif( + not _TE_GROUPED_LINEAR_SUPPORTS_USE_GROUPED_TENSOR, + reason="Installed TE GroupedLinear does not expose use_grouped_tensor", + ), ] @@ -195,6 +203,7 @@ def create_test_args( args.num_experts = 2 args.moe_layer_freq = 1 args.moe_grouped_gemm = True + args.moe_use_grouped_tensor = True args.moe_single_grouped_weight = single_weight args.moe_token_dispatcher_type = "alltoall" args.moe_router_topk = 1 @@ -684,7 +693,22 @@ def test_mxfp8_single_weight_torch_dist_checkpoint_matches_discrete_baseline( self.assert_all_ranks_passed(local_passed, local_error) - @pytest.mark.parametrize("precision", ["bf16", "mxfp8", "nvfp4"]) + @pytest.mark.parametrize( + "precision", + [ + "bf16", + "mxfp8", + pytest.param( + "nvfp4", + marks=pytest.mark.skip( + reason=( + "NVFP4 single grouped weights are not supported by the " + "TransformerEngine native grouped-tensor path yet." + ) + ), + ), + ], + ) @pytest.mark.parametrize("gradient_accumulation_fusion", [False, True]) def test_single_grouped_weight_parity_with_primary_param_gather( self, precision, gradient_accumulation_fusion @@ -727,17 +751,18 @@ def test_single_grouped_weight_parity_without_primary_param_gather( use_transformer_engine_op_fuser=True, ) - def test_single_grouped_weight_parity_module_grouped_linear(self): - """Single grouped weights require the TE op-fuser execution path.""" - args = self.create_test_args( - precision="bf16", - primary_param_gather=False, - single_weight=True, - gradient_accumulation_fusion=False, + @pytest.mark.parametrize( + "precision,primary_param_gather", [("bf16", False), ("mxfp8", False), ("mxfp8", True)] + ) + @pytest.mark.parametrize("gradient_accumulation_fusion", [False, True]) + def test_single_grouped_weight_parity_module_grouped_linear( + self, precision, primary_param_gather, gradient_accumulation_fusion + ): + """Compare native TE GroupedLinear single and discrete parameter layouts.""" + _skip_if_unsupported(precision) + self.run_parity_case( + precision=precision, + primary_param_gather=primary_param_gather, + gradient_accumulation_fusion=gradient_accumulation_fusion, use_transformer_engine_op_fuser=False, ) - with pytest.raises( - ValueError, - match="moe_single_grouped_weight requires use_transformer_engine_op_fuser=True", - ): - core_transformer_config_from_args(args) From 0615ac96b23ddea86c1e2bcdc855229e6b16e325 Mon Sep 17 00:00:00 2001 From: Madan kumar Date: Fri, 14 Aug 2026 22:56:56 +0000 Subject: [PATCH 043/101] fix(datasets): use len(lengths) for default modes in multimodal add_document (#5726) Signed-off-by: winklemad Co-authored-by: Antoni-Joan Solergibert Co-authored-by: Guihong Li Co-authored-by: Tom Long Co-authored-by: Philip Petrakian --- megatron/core/datasets/indexed_dataset.py | 2 +- tests/unit_tests/data/test_builder.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/megatron/core/datasets/indexed_dataset.py b/megatron/core/datasets/indexed_dataset.py index ce216c8b855..0be6db45d9d 100644 --- a/megatron/core/datasets/indexed_dataset.py +++ b/megatron/core/datasets/indexed_dataset.py @@ -974,7 +974,7 @@ def add_document( self.sequence_lengths.extend(lengths) self.document_indices.append(len(self.sequence_lengths)) if self.multimodal: - self.sequence_modes.extend(modes if modes is not None else [0] * lengths) + self.sequence_modes.extend(modes if modes is not None else [0] * len(lengths)) def end_document(self) -> None: """Finalize the document, for use with IndexedDatasetBuilder.add_item""" diff --git a/tests/unit_tests/data/test_builder.py b/tests/unit_tests/data/test_builder.py index 190b9bd1e0b..374731de331 100644 --- a/tests/unit_tests/data/test_builder.py +++ b/tests/unit_tests/data/test_builder.py @@ -66,6 +66,27 @@ def create_file_prefixes(tokenizer, number_of_files, maximum_number_of_documents return file_prefixes +def test_multimodal_builder_add_document_default_modes(): + # add_document documents modes as defaulting to None. On a multimodal builder + # that default must assign one mode (0) per item, matching the sequence + # lengths - not crash. Regression test for `[0] * lengths` (list * list). + with tempfile.TemporaryDirectory() as temp_dir: + builder = IndexedDatasetBuilder( + os.path.join(temp_dir, "mm.bin"), dtype=numpy.int32, multimodal=True + ) + + builder.add_document(numpy.array([1, 2, 3, 4, 5], dtype=numpy.int32), [3, 2]) + builder.add_document(numpy.array([6, 7], dtype=numpy.int32), [2]) + + assert builder.sequence_lengths == [3, 2, 2] + assert builder.sequence_modes == [0, 0, 0] + assert len(builder.sequence_modes) == len(builder.sequence_lengths) + + # Explicit modes are still honored. + builder.add_document(numpy.array([8], dtype=numpy.int32), [1], modes=[4]) + assert builder.sequence_modes == [0, 0, 0, 4] + + def do_setup(odir): paths = defaultdict(list) From a4417a5f7b23fd59056fd98f81a283263fa1dc64 Mon Sep 17 00:00:00 2001 From: Ming Liu <95666491+mlzoo@users.noreply.github.com> Date: Fri, 14 Aug 2026 22:58:38 +0000 Subject: [PATCH 044/101] Fix argument order in multi-chunk MoE upcycling state dict conversion (#5442) Signed-off-by: svcnvidia-nemo-ci Co-authored-by: Ming Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Co-authored-by: svcnvidia-nemo-ci --- .../core/transformer/moe/upcycling_utils.py | 9 ++- .../transformer/moe/test_upcycling.py | 55 ++++++++++++++++++- 2 files changed, 58 insertions(+), 6 deletions(-) diff --git a/megatron/core/transformer/moe/upcycling_utils.py b/megatron/core/transformer/moe/upcycling_utils.py index 4f9d2b7d835..1bd3f2a1a26 100644 --- a/megatron/core/transformer/moe/upcycling_utils.py +++ b/megatron/core/transformer/moe/upcycling_utils.py @@ -1,5 +1,6 @@ # Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. -""" Helpers for converting a dense model to a MoE model in runtime """ +"""Helpers for converting a dense model to a MoE model in runtime""" + import copy from enum import Enum @@ -114,10 +115,10 @@ def _convert_to_moe_state_dict(moe_model, dense_model): transformations for weights and biases specific to the MoE architecture. Args: - state_dict (dict): The dense model's state_dict. moe_model (nn.Module): The MoE model instance from which to get the submodule and state_dict, must be a model without FP16 and/or DDP wrapper. + dense_model (nn.Module): The dense model instance. Returns: dict: The converted MoE model state_dict, ready for use in the MoE architecture. @@ -320,9 +321,7 @@ def upcycle_state_dict(moe_model, dense_model): else: assert len(moe_model) == len(dense_model) for i in range(len(moe_model)): - state_dict['model%d' % i] = _convert_to_moe_state_dict( - dense_model[i].state_dict(), moe_model[i] - ) + state_dict['model%d' % i] = _convert_to_moe_state_dict(moe_model[i], dense_model[i]) return state_dict diff --git a/tests/unit_tests/transformer/moe/test_upcycling.py b/tests/unit_tests/transformer/moe/test_upcycling.py index feb9c9b9d2f..a522505171c 100644 --- a/tests/unit_tests/transformer/moe/test_upcycling.py +++ b/tests/unit_tests/transformer/moe/test_upcycling.py @@ -75,7 +75,7 @@ def model_provider( transformer_layer_spec=layer_spec_fn( args.num_experts, args.moe_grouped_gemm, args.qk_layernorm ), - vocab_size=args.vocal_size, + vocab_size=args.padded_vocab_size, max_sequence_length=args.max_position_embeddings, pre_process=pre_process, post_process=post_process, @@ -247,6 +247,59 @@ def test_upcycling_Local(self, tp_ep, granularity, grouped_gemm, swiglu, squared moe_logits, dense_logits, rtol=1e-01, atol=1e-01 ), "The output of moe model do not match the output of dense model." + @pytest.mark.parametrize( + ('tp_ep', 'granularity', 'grouped_gemm', 'swiglu', 'squared_relu'), + [pytest.param((1, 1), 1, False, False, False)], + ) + def test_upcycling_multi_model_chunks( + self, tp_ep, granularity, grouped_gemm, swiglu, squared_relu + ): + # Cover the virtual-pipeline path of ``upcycle_state_dict`` where ``moe_model`` + # and ``dense_model`` hold more than one model chunk. The single-chunk tests above + # never exercise the ``len(moe_model) > 1`` branch. + tp = tp_ep[0] + ep = tp_ep[1] + args = create_test_args(tp, grouped_gemm, swiglu, squared_relu, use_te=False) + set_args(args) + + torch.manual_seed(_SEED) + Utils.initialize_model_parallel(tensor_model_parallel_size=tp) + + dense_model, optimizer, opt_param_scheduler = setup_model_and_optimizer( + ModelType.encoder_or_decoder, model_provider + ) + dense_model = unwrap_model(dense_model) + set_bias_value(dense_model) + + Utils.destroy_model_parallel() + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp, expert_model_parallel_size=ep + ) + set_upcycling_args(ep, granularity, num_experts=2) + moe_model = unwrap_model(get_model(model_provider, ModelType.encoder_or_decoder)) + + # Emulate multiple model chunks by repeating the single chunk; this drives the + # ``len(moe_model) > 1`` branch, which previously crashed with an ``AttributeError`` + # because the converter arguments were swapped and a state_dict was passed in place + # of a module. + moe_chunks = [moe_model[0], moe_model[0]] + dense_chunks = [dense_model[0], dense_model[0]] + state_dict = upcycling_utils.upcycle_state_dict(moe_chunks, dense_chunks) + + assert set(state_dict.keys()) == {'model0', 'model1'} + # The multi-chunk branch must agree with the single-chunk branch for the same chunk. + single = upcycling_utils.upcycle_state_dict([moe_model[0]], [dense_model[0]]) + for i in range(len(moe_chunks)): + chunk = state_dict['model%d' % i] + assert chunk.keys() == single['model'].keys() + for k in chunk: + # ``_extra_state`` entries are not tensors; the local linear layers return + # ``None`` for them to stay compatible with the TE state dict. + if k.endswith('_extra_state'): + continue + assert torch.equal(chunk[k], single['model'][k]), f"Value mismatch for key {k}" + moe_model[0].load_state_dict(chunk, strict=True) + @pytest.mark.skipif( not HAVE_TE or not is_te_min_version("2.1.0"), reason="grouped_gemm requires TransformerEngine >= 2.1.0", From 669d6ed864ff6085565b3d6f1412381d1c11b847 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Sat, 15 Aug 2026 00:59:54 +0000 Subject: [PATCH 045/101] ci: AUT-1408 replace cicd PAT with GitHub App tokens (#6360) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: svcnemo-autobot Signed-off-by: oliver könig Co-authored-by: oliver könig --- .github/actions/action.yml | 3 - .github/scripts/trigger_workflow_and_wait.py | 458 ++++++++++++++++++ .github/workflows/cicd-main.yml | 139 +++--- .../test_trigger_workflow_and_wait.py | 223 +++++++++ 4 files changed, 761 insertions(+), 62 deletions(-) create mode 100755 .github/scripts/trigger_workflow_and_wait.py create mode 100644 tests/test_utils/python_scripts/test_trigger_workflow_and_wait.py diff --git a/.github/actions/action.yml b/.github/actions/action.yml index 97f6c17c8fd..f339659d483 100644 --- a/.github/actions/action.yml +++ b/.github/actions/action.yml @@ -42,9 +42,6 @@ inputs: model: description: Model to launch required: false - PAT: - description: "GitHub Personal Access Token" - required: true scope: description: "Test scope (e.g. mr-github, mr-github-slim)" required: false diff --git a/.github/scripts/trigger_workflow_and_wait.py b/.github/scripts/trigger_workflow_and_wait.py new file mode 100755 index 00000000000..031fc2cf879 --- /dev/null +++ b/.github/scripts/trigger_workflow_and_wait.py @@ -0,0 +1,458 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +"""Trigger and wait for Megatron-Bridge validation with refreshable App credentials.""" + +from __future__ import annotations + +import base64 +import json +import os +import re +import subprocess +import sys +import time +import urllib.error +import urllib.parse +import urllib.request +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime +from typing import Any + +API_ROOT = "https://api.github.com" +TARGET_OWNER = "NVIDIA-NeMo" +TARGET_REPO = "Megatron-Bridge" +TARGET_WORKFLOW = "cicd-main.yml" +TARGET_BASE_REF = "main" +ALLOWED_CALLER = "NVIDIA/Megatron-LM" +ALLOWED_SUITES = {"L1", "unit-only"} +TOKEN_REFRESH_MARGIN_MS = 10 * 60 * 1000 +DISCOVERY_TIMEOUT_MS = 10 * 60 * 1000 +POLL_INTERVAL_MS = 60 * 1000 + +JsonObject = dict[str, Any] +ApiRequest = Callable[[str, str, str, JsonObject | None], JsonObject | None] +Now = Callable[[], int] +Sleep = Callable[[int], None] + + +@dataclass(frozen=True) +class Config: + """Validated inputs for one MBridge orchestration attempt.""" + + app_id: str + private_key: str + mcore_ref: str + test_suite: str + poll_timeout_seconds: int + caller_repository: str + run_id: str + run_attempt: str + server_url: str + + +@dataclass(frozen=True) +class Credential: + """One scoped, expiring installation token.""" + + token: str + expires_at_ms: int + + +class ApiError(RuntimeError): + """A bounded GitHub API failure.""" + + def __init__(self, method: str, path: str, status: int, response: str) -> None: + super().__init__( + f"GitHub API {method} {path} failed with HTTP {status}: {response[:1000]}" + ) + self.status = status + + +def _base64_url(value: bytes) -> str: + return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii") + + +def create_app_jwt(app_id: str, private_key: str, now_ms: int | None = None) -> str: + """Create a short-lived GitHub App JWT using the runner's OpenSSL.""" + + current_ms = int(time.time() * 1000) if now_ms is None else now_ms + issued_at = current_ms // 1000 - 60 + header = _base64_url(json.dumps({"alg": "RS256", "typ": "JWT"}).encode()) + payload = _base64_url( + json.dumps( + {"iat": issued_at, "exp": issued_at + 9 * 60, "iss": app_id} + ).encode() + ) + signing_input = f"{header}.{payload}".encode() + private_key_fd = os.memfd_create("github-app-private-key", os.MFD_CLOEXEC) + try: + os.write(private_key_fd, private_key.encode()) + os.lseek(private_key_fd, 0, os.SEEK_SET) + signature = subprocess.run( + ["openssl", "dgst", "-sha256", "-sign", f"/proc/self/fd/{private_key_fd}"], + input=signing_input, + check=True, + capture_output=True, + pass_fds=(private_key_fd,), + ).stdout + finally: + os.close(private_key_fd) + return f"{header}.{payload}.{_base64_url(signature)}" + + +def _parse_timeout(raw_timeout: str) -> int: + try: + value = float(raw_timeout) + except ValueError as error: + raise ValueError( + "poll_timeout_seconds must be an integer between 60 and 14400" + ) from error + if not value.is_integer() or value < 60 or value > 4 * 60 * 60: + raise ValueError("poll_timeout_seconds must be an integer between 60 and 14400") + return int(value) + + +def read_config(env: dict[str, str] | None = None) -> Config: + """Read and validate the bounded workflow-call contract.""" + + values = os.environ if env is None else env + config = Config( + app_id=values.get("BOT_ID", ""), + private_key=values.get("BOT_KEY", ""), + mcore_ref=values.get("MBRIDGE_MCORE_REF", ""), + test_suite=values.get("MBRIDGE_TEST_SUITE", ""), + caller_repository=values.get("GITHUB_REPOSITORY", ""), + poll_timeout_seconds=_parse_timeout( + values.get("MBRIDGE_POLL_TIMEOUT_SECONDS", "12600") + ), + run_id=values.get("GITHUB_RUN_ID", ""), + run_attempt=values.get("GITHUB_RUN_ATTEMPT", ""), + server_url=values.get("GITHUB_SERVER_URL", "https://github.com"), + ) + validate_config(config) + return config + + +def validate_config(config: Config) -> None: + """Fail closed on any untrusted caller-controlled input.""" + + if re.fullmatch(r"\d+", config.app_id) is None: + raise ValueError("The GitHub App ID must be numeric") + if re.search(r"-----BEGIN (?:RSA )?PRIVATE KEY-----", config.private_key) is None: + raise ValueError("The GitHub App private key is missing or malformed") + if re.fullmatch(r"[0-9a-f]{40}", config.mcore_ref) is None: + raise ValueError("mcore_ref must be a full lowercase commit SHA") + if config.test_suite not in ALLOWED_SUITES: + raise ValueError(f"Unsupported MBridge test suite: {config.test_suite}") + if config.caller_repository != ALLOWED_CALLER: + raise ValueError(f"This workflow only accepts calls from {ALLOWED_CALLER}") + if not 60 <= config.poll_timeout_seconds <= 4 * 60 * 60: + raise ValueError("poll_timeout_seconds must be an integer between 60 and 14400") + if ( + re.fullmatch(r"\d+", config.run_id) is None + or re.fullmatch(r"\d+", config.run_attempt) is None + ): + raise ValueError("GitHub run ID and attempt must be numeric") + if config.server_url != "https://github.com": + raise ValueError("Only github.com callers are supported") + + +def testing_branch(config: Config) -> str: + return f"mcore-testing-{config.run_id}-{config.run_attempt}" + + +def triggered_by(config: Config) -> str: + return ( + f"{config.server_url}/{config.caller_repository}/actions/runs/{config.run_id}" + ) + + +def api_request( + token: str, method: str, path: str, body: JsonObject | None = None +) -> JsonObject | None: + """Make one bounded GitHub REST request.""" + + data = None if body is None else json.dumps(body).encode() + request = urllib.request.Request( + f"{API_ROOT}{path}", + data=data, + method=method, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=60) as response: + response_text = response.read().decode() + return json.loads(response_text) if response_text else None + except urllib.error.HTTPError as error: + response_text = error.read().decode(errors="replace") + raise ApiError(method, path, error.code, response_text) from error + + +def _parse_timestamp(value: str) -> int: + return int(datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() * 1000) + + +def create_token_minter( + config: Config, + *, + request: ApiRequest = api_request, + now: Now = lambda: int(time.time() * 1000), +) -> Callable[[dict[str, str]], Credential]: + """Return a minter that scopes every token to one operation.""" + + installation_id: int | None = None + + def mint_token(permissions: dict[str, str]) -> Credential: + nonlocal installation_id + jwt = create_app_jwt(config.app_id, config.private_key, now()) + if installation_id is None: + installation = request( + jwt, "GET", f"/repos/{TARGET_OWNER}/{TARGET_REPO}/installation", None + ) + if installation is None: + raise RuntimeError("GitHub returned an empty installation response") + installation_id = int(installation["id"]) + response = request( + jwt, + "POST", + f"/app/installations/{installation_id}/access_tokens", + {"repositories": [TARGET_REPO], "permissions": permissions}, + ) + if response is None: + raise RuntimeError("GitHub returned an empty token response") + return Credential( + token=str(response["token"]), + expires_at_ms=_parse_timestamp(str(response["expires_at"])), + ) + + return mint_token + + +def runs_path(branch: str) -> str: + query = urllib.parse.urlencode( + {"event": "workflow_dispatch", "branch": branch, "per_page": "100"} + ) + workflow = urllib.parse.quote(TARGET_WORKFLOW, safe="") + return ( + f"/repos/{TARGET_OWNER}/{TARGET_REPO}/actions/workflows/{workflow}/runs?{query}" + ) + + +def select_dispatched_run( + runs: list[JsonObject], old_run_ids: set[int], branch: str, started_at_ms: int +) -> JsonObject | None: + candidates = [ + run + for run in runs + if int(run["id"]) not in old_run_ids + and run.get("event") == "workflow_dispatch" + and run.get("head_branch") == branch + and _parse_timestamp(str(run["created_at"])) >= started_at_ms - 2 * 60 * 1000 + ] + return min( + candidates, + key=lambda run: _parse_timestamp(str(run["created_at"])), + default=None, + ) + + +def discover_run( + token: str, + old_run_ids: set[int], + branch: str, + started_at_ms: int, + *, + request: ApiRequest = api_request, + now: Now = lambda: int(time.time() * 1000), + sleep: Sleep = lambda milliseconds: time.sleep(milliseconds / 1000), +) -> int: + deadline = now() + DISCOVERY_TIMEOUT_MS + while now() < deadline: + response = request(token, "GET", runs_path(branch), None) + if response is None: + raise RuntimeError("GitHub returned an empty workflow-runs response") + run = select_dispatched_run( + list(response["workflow_runs"]), old_run_ids, branch, started_at_ms + ) + if run is not None: + return int(run["id"]) + sleep(POLL_INTERVAL_MS) + raise TimeoutError( + "Timed out waiting for the dispatched MBridge workflow run to appear" + ) + + +def poll_run( + run_id: int, + timeout_seconds: int, + mint_token: Callable[[dict[str, str]], Credential], + *, + request: ApiRequest = api_request, + now: Now = lambda: int(time.time() * 1000), + sleep: Sleep = lambda milliseconds: time.sleep(milliseconds / 1000), +) -> None: + deadline = now() + timeout_seconds * 1000 + credential: Credential | None = None + while True: + if ( + credential is None + or credential.expires_at_ms - now() <= TOKEN_REFRESH_MARGIN_MS + ): + credential = mint_token({"actions": "read"}) + run = request( + credential.token, + "GET", + f"/repos/{TARGET_OWNER}/{TARGET_REPO}/actions/runs/{run_id}", + None, + ) + if run is None: + raise RuntimeError("GitHub returned an empty workflow-run response") + print(f"MBridge workflow status={run['status']} conclusion={run['conclusion']}") + if run["status"] == "completed": + if run["conclusion"] != "success": + raise RuntimeError( + f"MBridge workflow {run_id} concluded with {run['conclusion']}" + ) + return + remaining = deadline - now() + if remaining <= 0: + raise TimeoutError( + f"Timed out waiting for MBridge workflow {run_id} to complete" + ) + sleep(min(POLL_INTERVAL_MS, remaining)) + + +def delete_branch( + mint_token: Callable[[dict[str, str]], Credential], + branch: str, + *, + request: ApiRequest = api_request, +) -> None: + credential = mint_token({"contents": "write"}) + branch_name = urllib.parse.quote(branch, safe="") + try: + request( + credential.token, + "DELETE", + f"/repos/{TARGET_OWNER}/{TARGET_REPO}/git/refs/heads/{branch_name}", + None, + ) + except ApiError as error: + if error.status != 404: + raise + + +def orchestrate( + config: Config, + *, + request: ApiRequest = api_request, + now: Now = lambda: int(time.time() * 1000), + sleep: Sleep = lambda milliseconds: time.sleep(milliseconds / 1000), + mint_token: Callable[[dict[str, str]], Credential] | None = None, +) -> None: + """Create, dispatch, monitor, and clean up one isolated MBridge run.""" + + token_minter = mint_token or create_token_minter(config, request=request, now=now) + branch = testing_branch(config) + branch_created = False + failure: Exception | None = None + try: + branch_credential = token_minter({"contents": "write"}) + base_ref = request( + branch_credential.token, + "GET", + f"/repos/{TARGET_OWNER}/{TARGET_REPO}/git/ref/heads/{TARGET_BASE_REF}", + None, + ) + if base_ref is None: + raise RuntimeError("GitHub returned an empty base-ref response") + request( + branch_credential.token, + "POST", + f"/repos/{TARGET_OWNER}/{TARGET_REPO}/git/refs", + {"ref": f"refs/heads/{branch}", "sha": base_ref["object"]["sha"]}, + ) + branch_created = True + print(f"Created temporary MBridge branch {branch}") + + trigger_credential = token_minter({"actions": "write"}) + prior_runs = request(trigger_credential.token, "GET", runs_path(branch), None) + if prior_runs is None: + raise RuntimeError("GitHub returned an empty workflow-runs response") + old_run_ids = {int(run["id"]) for run in prior_runs["workflow_runs"]} + started_at_ms = now() + workflow = urllib.parse.quote(TARGET_WORKFLOW, safe="") + request( + trigger_credential.token, + "POST", + f"/repos/{TARGET_OWNER}/{TARGET_REPO}/actions/workflows/{workflow}/dispatches", + { + "ref": branch, + "inputs": { + "mcore_ref": config.mcore_ref, + "test_suite": config.test_suite, + "triggered_by": triggered_by(config), + }, + }, + ) + run_id = discover_run( + trigger_credential.token, + old_run_ids, + branch, + started_at_ms, + request=request, + now=now, + sleep=sleep, + ) + print(f"Discovered MBridge workflow run {run_id}") + poll_run( + run_id, + config.poll_timeout_seconds, + token_minter, + request=request, + now=now, + sleep=sleep, + ) + except Exception as error: # noqa: BLE001 - cleanup must run for every operational failure + failure = error + finally: + if branch_created: + try: + delete_branch(token_minter, branch, request=request) + print(f"Deleted temporary MBridge branch {branch}") + except Exception as cleanup_error: # noqa: BLE001 - preserve the primary failure + if failure is None: + failure = cleanup_error + else: + print( + f"MBridge branch cleanup also failed: {cleanup_error}", + file=sys.stderr, + ) + if failure is not None: + raise failure + + +def main() -> None: + orchestrate(read_config()) + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 429632980a2..90fb157313e 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -49,14 +49,22 @@ jobs: issues: write pull-requests: write env: - GITHUB_TOKEN: ${{ secrets.PAT }} REPO: ${{ github.repository }} DISABLE_EXTERNAL_CONTRIBUTOR: ${{ vars.DISABLE_EXTERNAL_CONTRIBUTOR }} steps: + - name: Generate GitHub App token + id: app-token + uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + with: + app-id: ${{ vars.BOT_ID }} + private-key: ${{ secrets.BOT_KEY }} + permission-contents: read + - name: Checkout repository uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: - token: ${{ env.GITHUB_TOKEN }} + token: ${{ steps.app-token.outputs.token }} + persist-credentials: false - name: Get PR info id: get-pr-info @@ -71,9 +79,38 @@ jobs: github_token: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }} sso_users_filename: ${{ vars.SSO_USERS_FILENAME }} + - name: Generate NVIDIA-NeMo membership token + id: nvidia-nemo-members-token + if: | + startsWith(github.ref, 'refs/heads/pull-request/') + && github.event_name == 'push' + && vars.DISABLE_EXTERNAL_CONTRIBUTOR == 'true' && steps.check-sso.outputs.is_member != 'true' + uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + with: + app-id: ${{ vars.BOT_ID }} + private-key: ${{ secrets.BOT_KEY }} + owner: NVIDIA-NeMo + permission-members: read + + - name: Generate NVIDIA membership token + id: nvidia-members-token + if: | + startsWith(github.ref, 'refs/heads/pull-request/') + && github.event_name == 'push' + && vars.DISABLE_EXTERNAL_CONTRIBUTOR == 'true' && steps.check-sso.outputs.is_member != 'true' + uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + with: + app-id: ${{ vars.BOT_ID }} + private-key: ${{ secrets.BOT_KEY }} + owner: NVIDIA + permission-members: read + - name: Set maintainer status id: check-membership env: + GITHUB_TOKEN: ${{ steps.app-token.outputs.token }} + NVIDIA_MEMBERS_TOKEN: ${{ steps.nvidia-members-token.outputs.token }} + NVIDIA_NEMO_MEMBERS_TOKEN: ${{ steps.nvidia-nemo-members-token.outputs.token }} IS_MAIN_BRANCH: ${{ github.ref == 'refs/heads/main' }} IS_MERGE_GROUP: ${{ github.event_name == 'merge_group' }} SCHEDULED_JOB: ${{ github.event_name == 'schedule' }} @@ -104,7 +141,7 @@ jobs: API_URL="https://api.github.com/orgs/NVIDIA-NeMo/members/$PR_AUTHOR" ORG_NVIDIA_NEMO_MEMBERSHIP_RESPONSE=$(curl -s -o /dev/null -w "%{http_code}" -L \ -H "Accept: application/vnd.github+json" \ - -H "Authorization: Bearer $GITHUB_TOKEN" \ + -H "Authorization: Bearer $NVIDIA_NEMO_MEMBERS_TOKEN" \ -H "X-GitHub-Api-Version: 2022-11-28" \ $API_URL) @@ -112,7 +149,7 @@ jobs: API_URL="https://api.github.com/orgs/NVIDIA/members/$PR_AUTHOR" ORG_NVIDIA_MEMBERSHIP_RESPONSE=$(curl -s -o /dev/null -w "%{http_code}" -L \ -H "Accept: application/vnd.github+json" \ - -H "Authorization: Bearer $GITHUB_TOKEN" \ + -H "Authorization: Bearer $NVIDIA_MEMBERS_TOKEN" \ -H "X-GitHub-Api-Version: 2022-11-28" \ $API_URL) @@ -177,11 +214,19 @@ jobs: fi echo "sha=${SHA}" | tee -a "$GITHUB_OUTPUT" + - name: Generate GitHub App token + id: app-token + uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + with: + app-id: ${{ vars.BOT_ID }} + private-key: ${{ secrets.BOT_KEY }} + permission-pull-requests: read + - name: Configure id: configure shell: bash -x -e -u -o pipefail {0} env: - GH_TOKEN: ${{ secrets.PAT }} + GH_TOKEN: ${{ steps.app-token.outputs.token }} IS_CI_WORKLOAD: ${{ needs.pre-flight.outputs.is_ci_workload }} IS_MERGE_GROUP: ${{ needs.pre-flight.outputs.is_merge_group }} EVENT_NAME: ${{ github.event_name }} @@ -214,7 +259,7 @@ jobs: SCOPE=L0; N_REPEAT=5; LIGHTWEIGHT=false fi - if [ "$HAS_MBRIDGE" == "true" || $IS_MERGE_GROUP == "true" ]; then + if [ "$HAS_MBRIDGE" == "true" ] || [ "$IS_MERGE_GROUP" == "true" ]; then MBRIDGE_SUITE="L1" else MBRIDGE_SUITE="unit-only" @@ -458,6 +503,8 @@ jobs: && needs.cicd-wait-in-queue.result != 'cancelled' && needs.cicd-parse-downstream-testing.result != 'cancelled' && vars.ENABLE_CICD_MBRIDGE_TESTING == 'true' + && github.event_name == 'push' + && startsWith(github.ref, 'refs/heads/pull-request/') && needs.configure.outputs.run_mbridge == 'true' && needs.pre-flight.outputs.docs_only == 'false' && ( @@ -467,55 +514,19 @@ jobs: || needs.pre-flight.outputs.is_merge_group == 'true' ) && !cancelled() + timeout-minutes: 300 steps: - - name: Get PR info - id: get-pr-info - if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' - uses: nv-gha-runners/get-pr-info@090577647b8ddc4e06e809e264f7881650ecdccf # main - - - name: Checkout MBridge and create testing branch + - name: Checkout Megatron-LM uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - with: - ref: main - repository: NVIDIA-NeMo/Megatron-Bridge - path: megatron-bridge - token: ${{ secrets.PAT }} - - name: Create testing branch - env: - MBRIDGE_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} - run: | - cd megatron-bridge - git fetch origin main - git checkout -b ${{ env.MBRIDGE_BRANCH_NAME }} origin/main - git push origin ${{ env.MBRIDGE_BRANCH_NAME }} --force - - - name: Trigger MBridge tests - uses: convictional/trigger-workflow-and-wait@f69fa9eedd3c62a599220f4d5745230e237904be # v1.6.5 - env: - MBRIDGE_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} - with: - owner: NVIDIA-NeMo - repo: Megatron-Bridge - workflow_file_name: cicd-main.yml - github_token: ${{ secrets.PAT }} - ref: ${{ env.MBRIDGE_BRANCH_NAME }} - wait_interval: 60 - propagate_failure: true - client_payload: | - { - "mcore_ref": "${{ needs.configure.outputs.sha }}", - "test_suite": "${{ needs.cicd-parse-downstream-testing.outputs.mbridge-test-suite }}", - "triggered_by": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}" - } - - - name: Delete testing branch - if: always() + - name: Trigger and wait for MBridge tests env: - MBRIDGE_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} - run: | - cd megatron-bridge - git push origin --delete ${{ env.MBRIDGE_BRANCH_NAME }} + BOT_ID: ${{ vars.COMMUNITY_BOT_ID }} + BOT_KEY: ${{ secrets.COMMUNITY_BOT_KEY }} + MBRIDGE_MCORE_REF: ${{ needs.configure.outputs.sha }} + MBRIDGE_TEST_SUITE: ${{ needs.cicd-parse-downstream-testing.outputs.mbridge-test-suite }} + MBRIDGE_POLL_TIMEOUT_SECONDS: "12600" + run: python3 -u .github/scripts/trigger_workflow_and_wait.py cicd-mbridge-testing-notify: runs-on: ubuntu-latest @@ -717,8 +728,6 @@ jobs: tags: | ${{ matrix.registry }}/megatron-lm:${{ steps.cache_keys.outputs.key }}-${{ matrix.cloud }} ${{ matrix.registry }}/megatron-lm:${{ needs.configure.outputs.sha }}-${{ matrix.cloud }} - secrets: | - GH_TOKEN=${{ secrets.PAT }} cicd-parse-unit-tests: runs-on: ubuntu-latest @@ -799,7 +808,6 @@ jobs: tag: latest timeout: ${{ matrix.timeout || 30 }} is_unit_test: "true" - PAT: ${{ secrets.PAT }} container-image: ${{ env.container-registry }}/megatron-lm:${{ needs.configure.outputs.sha }}-aws-h100 sha: ${{ needs.configure.outputs.sha }} @@ -887,7 +895,6 @@ jobs: tag: latest timeout: ${{ matrix.timeout || 30 }} is_unit_test: "true" - PAT: ${{ secrets.PAT }} container-image: ${{ env.container-registry-gb200 }}/megatron-lm:${{ needs.configure.outputs.sha }}-gb-gpu platform: dgx_gb200 sha: ${{ needs.configure.outputs.sha }} @@ -1044,7 +1051,6 @@ jobs: tag: latest timeout: ${{ matrix.timeout || 30 }} is_unit_test: "false" - PAT: ${{ secrets.PAT }} container-image: ${{ env.container-registry }}/megatron-lm:${{ needs.configure.outputs.sha }}-aws-h100 scope: ${{ needs.configure.outputs.scope }} n_repeat: ${{ needs.configure.outputs.n_repeat }} @@ -1145,7 +1151,6 @@ jobs: tag: latest timeout: ${{ matrix.timeout || 30 }} is_unit_test: "false" - PAT: ${{ secrets.PAT }} container-image: ${{ env.container-registry-gb200 }}/megatron-lm:${{ needs.configure.outputs.sha }}-gb-gpu scope: ${{ needs.configure.outputs.scope }} n_repeat: ${{ needs.configure.outputs.n_repeat }} @@ -1285,10 +1290,18 @@ jobs: && !cancelled() && github.repository == 'NVIDIA/Megatron-LM' steps: + - name: Generate GitHub App token + id: app-token + uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + with: + app-id: ${{ vars.BOT_ID }} + private-key: ${{ secrets.BOT_KEY }} + permission-statuses: write + - name: Generate fake coverage report uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8 with: - github-token: ${{ secrets.PAT }} + github-token: ${{ steps.app-token.outputs.token }} script: | await github.rest.repos.createCommitStatus({ owner: context.repo.owner, @@ -1372,10 +1385,18 @@ jobs: PR_NUMBER=$(echo "${{ github.event.merge_group.head_ref }}" | sed -n 's/.*\/pr-\([0-9]*\)-.*/\1/p') echo "pr_number=$PR_NUMBER" >> $GITHUB_OUTPUT + - name: Generate GitHub App token + id: app-token + uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + with: + app-id: ${{ vars.BOT_ID }} + private-key: ${{ secrets.BOT_KEY }} + permission-issues: write + - name: Comment on PR with action run URL uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8 with: - github-token: ${{ secrets.PAT }} + github-token: ${{ steps.app-token.outputs.token }} script: | const prNumber = ${{ steps.get-pr-number.outputs.pr_number }}; const runUrl = `https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}`; diff --git a/tests/test_utils/python_scripts/test_trigger_workflow_and_wait.py b/tests/test_utils/python_scripts/test_trigger_workflow_and_wait.py new file mode 100644 index 00000000000..98185fd3a1f --- /dev/null +++ b/tests/test_utils/python_scripts/test_trigger_workflow_and_wait.py @@ -0,0 +1,223 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import sys +from pathlib import Path +from unittest import mock + +import pytest + +MODULE_PATH = Path(__file__).parents[3] / ".github/scripts/trigger_workflow_and_wait.py" +SPEC = importlib.util.spec_from_file_location("trigger_workflow_and_wait", MODULE_PATH) +assert SPEC is not None and SPEC.loader is not None +MODULE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) + + +def config(**overrides): + values = { + "app_id": "123", + "private_key": "-----BEGIN PRIVATE KEY-----\nfixture\n-----END PRIVATE KEY-----", + "mcore_ref": "a" * 40, + "test_suite": "L1", + "poll_timeout_seconds": 12600, + "caller_repository": "NVIDIA/Megatron-LM", + "run_id": "456", + "run_attempt": "2", + "server_url": "https://github.com", + } + values.update(overrides) + return MODULE.Config(**values) + + +def test_read_config_accepts_workflow_contract(): + result = MODULE.read_config( + { + "BOT_ID": "123", + "BOT_KEY": "-----BEGIN PRIVATE KEY-----\nfixture\n-----END PRIVATE KEY-----", + "MBRIDGE_MCORE_REF": "a" * 40, + "MBRIDGE_TEST_SUITE": "L1", + "MBRIDGE_POLL_TIMEOUT_SECONDS": "12600", + "GITHUB_REPOSITORY": "NVIDIA/Megatron-LM", + "GITHUB_RUN_ID": "456", + "GITHUB_RUN_ATTEMPT": "2", + "GITHUB_SERVER_URL": "https://github.com", + } + ) + assert result == config() + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"app_id": "not-numeric"}, "App ID"), + ({"private_key": "missing"}, "private key"), + ({"mcore_ref": "main"}, "full lowercase commit SHA"), + ({"test_suite": "unsupported"}, "Unsupported"), + ({"poll_timeout_seconds": 1}, "between 60 and 14400"), + ({"caller_repository": "attacker/fork"}, "only accepts calls"), + ({"run_id": "not-numeric"}, "run ID and attempt"), + ({"server_url": "https://example.test"}, "Only github.com"), + ], +) +def test_validate_config_rejects_untrusted_inputs(override, message): + with pytest.raises(ValueError, match=message): + MODULE.validate_config(config(**override)) + + +def test_create_app_jwt_is_short_lived_and_invokes_openssl(): + completed = mock.Mock(stdout=b"signature") + with mock.patch.object(MODULE.subprocess, "run", return_value=completed) as run: + jwt = MODULE.create_app_jwt("123", config().private_key, now_ms=1_800_000) + assert jwt.count(".") == 2 + run.assert_called_once() + assert run.call_args.args[0][:4] == ["openssl", "dgst", "-sha256", "-sign"] + assert run.call_args.kwargs["pass_fds"] + assert run.call_args.args[0][4].startswith("/proc/self/fd/") + + +def test_poll_run_refreshes_actions_read_tokens(): + credentials = iter( + [MODULE.Credential("token-1", 15 * 60 * 1000), MODULE.Credential("token-2", 60 * 60 * 1000)] + ) + minted = [] + + def mint(permissions): + minted.append(permissions) + return next(credentials) + + now = iter([0, 10 * 60 * 1000, 10 * 60 * 1000, 10 * 60 * 1000]) + responses = iter( + [ + {"status": "in_progress", "conclusion": None}, + {"status": "completed", "conclusion": "success"}, + ] + ) + MODULE.poll_run( + 77, + 12600, + mint, + request=lambda *_: next(responses), + now=lambda: next(now), + sleep=lambda _: None, + ) + assert minted == [{"actions": "read"}, {"actions": "read"}] + + +def test_poll_run_fails_closed_on_timeout(): + credential = MODULE.Credential("token", 60 * 60 * 1000) + now = iter([0, 0, 1_000, 1_000]) + with pytest.raises(TimeoutError, match="workflow 77"): + MODULE.poll_run( + 77, + 1, + lambda _: credential, + request=lambda *_: {"status": "in_progress", "conclusion": None}, + now=lambda: next(now), + sleep=lambda _: None, + ) + + +def test_orchestrate_scopes_tokens_and_cleans_up(): + minted = [] + requests = [] + credentials = iter( + [ + MODULE.Credential("branch", 60 * 60 * 1000), + MODULE.Credential("trigger", 60 * 60 * 1000), + MODULE.Credential("poll", 60 * 60 * 1000), + MODULE.Credential("cleanup", 60 * 60 * 1000), + ] + ) + + def mint(permissions): + minted.append(permissions) + return next(credentials) + + run_queries = 0 + + def request(token, method, path, body): + nonlocal run_queries + requests.append((token, method, path, body)) + if path.endswith("/git/ref/heads/main"): + return {"object": {"sha": "b" * 40}} + if path.endswith("/runs?event=workflow_dispatch&branch=mcore-testing-456-2&per_page=100"): + count = sum(1 for request_entry in requests if request_entry[2] == path) + if count == 1: + return {"workflow_runs": []} + return { + "workflow_runs": [ + { + "id": 77, + "event": "workflow_dispatch", + "head_branch": "mcore-testing-456-2", + "created_at": "1970-01-01T00:00:01Z", + } + ] + } + if path.endswith("/actions/runs/77"): + return {"status": "completed", "conclusion": "success"} + return None + + times = iter([1_000, 1_000, 1_000, 1_000, 1_000, 1_000, 1_000]) + MODULE.orchestrate( + config(), request=request, now=lambda: next(times), sleep=lambda _: None, mint_token=mint + ) + assert minted == [ + {"contents": "write"}, + {"actions": "write"}, + {"actions": "read"}, + {"contents": "write"}, + ] + assert any( + method == "DELETE" and path.endswith("/mcore-testing-456-2") + for _, method, path, _ in requests + ) + + +def test_orchestrate_preserves_test_failure_when_cleanup_fails(): + credentials = iter( + [ + MODULE.Credential("branch", 60 * 60 * 1000), + MODULE.Credential("trigger", 60 * 60 * 1000), + MODULE.Credential("poll", 60 * 60 * 1000), + MODULE.Credential("cleanup", 60 * 60 * 1000), + ] + ) + + run_queries = 0 + + def request(token, method, path, body): + nonlocal run_queries + if path.endswith("/git/ref/heads/main"): + return {"object": {"sha": "b" * 40}} + if "/runs?" in path: + run_queries += 1 + if run_queries == 1: + return {"workflow_runs": []} + return { + "workflow_runs": [ + { + "id": 77, + "event": "workflow_dispatch", + "head_branch": "mcore-testing-456-2", + "created_at": "1970-01-01T00:00:01Z", + } + ] + } + if path.endswith("/actions/runs/77"): + return {"status": "completed", "conclusion": "failure"} + if method == "DELETE": + raise RuntimeError("cleanup failed") + return None + + with pytest.raises(RuntimeError, match="concluded with failure"): + MODULE.orchestrate( + config(), + request=request, + now=lambda: 1_000, + sleep=lambda _: None, + mint_token=lambda _: next(credentials), + ) From 51f6a39fc8b6010fc062ac32fa6f637b97c4bf0e Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Sat, 15 Aug 2026 01:15:09 +0000 Subject: [PATCH 046/101] fix: AUT-1501 abort when SSO list download fails (#6516) Signed-off-by: svcnemo-autobot --- .github/actions/check-nvidia-sso-membership/action.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/actions/check-nvidia-sso-membership/action.yml b/.github/actions/check-nvidia-sso-membership/action.yml index e96863e2264..e27aaae2aad 100644 --- a/.github/actions/check-nvidia-sso-membership/action.yml +++ b/.github/actions/check-nvidia-sso-membership/action.yml @@ -58,7 +58,7 @@ runs: echo "ERROR: Failed to download ${{ inputs.sso_users_filename }} from github-audits release" echo "sso_file_available=false" >> $GITHUB_OUTPUT echo "user_count=0" >> $GITHUB_OUTPUT - exit 0 + exit 1 } # Verify file was downloaded and is valid JSON @@ -66,7 +66,7 @@ runs: echo "ERROR: ${{ inputs.sso_users_filename }} file not found after download" echo "sso_file_available=false" >> $GITHUB_OUTPUT echo "user_count=0" >> $GITHUB_OUTPUT - exit 0 + exit 1 fi # Validate JSON structure @@ -74,7 +74,7 @@ runs: echo "ERROR: ${{ inputs.sso_users_filename }} is not a valid JSON object" echo "sso_file_available=false" >> $GITHUB_OUTPUT echo "user_count=0" >> $GITHUB_OUTPUT - exit 0 + exit 1 fi USER_COUNT=$(jq 'length' ${{ inputs.sso_users_filename }}) From e79110c67b6095326a98f210e8adc5afc905b07d Mon Sep 17 00:00:00 2001 From: janEbert Date: Sat, 15 Aug 2026 13:10:10 +0000 Subject: [PATCH 047/101] Fix missing golden values updates (#6528) Signed-off-by: janEbert --- .../golden_values_dev_dgx_h100.json | 201 ++- .../golden_values_dev_dgx_gb200.json | 449 ++++--- .../golden_values_lts_dgx_a100.json | 461 ++++--- .../golden_values_dev_dgx_gb200.json | 449 ++++--- .../golden_values_lts_dgx_a100.json | 485 ++++--- .../golden_values_dev_dgx_gb200.json | 497 ++++--- .../golden_values_dev_dgx_gb200.json | 785 ++++++------ .../golden_values_dev_dgx_gb200.json | 1137 ++++++++--------- 8 files changed, 2228 insertions(+), 2236 deletions(-) diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/golden_values_dev_dgx_h100.json index 3550f734923..e9e60ded687 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/hybrid/hybrid_flextron_nightly_tp2_pp1_ep2_dgx_h100_1N8G/golden_values_dev_dgx_h100.json @@ -4,26 +4,26 @@ "end_step": 20, "step_interval": 1, "values": { - "1": 10.96772, - "2": 10.94878, - "3": 10.13223, - "4": 10.10273, - "5": 9.77116, - "6": 9.74409, - "7": 9.67572, - "8": 9.54443, - "9": 9.47498, - "10": 9.46493, - "11": 9.0426, - "12": 9.29951, - "13": 9.30226, - "14": 9.27311, - "15": 9.10353, - "16": 8.79531, - "17": 8.51664, - "18": 8.77952, - "19": 8.71592, - "20": 8.71692 + "1": 10.97032, + "2": 10.94795, + "3": 10.14519, + "4": 10.13125, + "5": 9.77793, + "6": 9.76375, + "7": 9.68447, + "8": 9.56926, + "9": 9.48546, + "10": 9.49063, + "11": 9.0613, + "12": 9.32517, + "13": 9.35135, + "14": 9.29727, + "15": 9.14824, + "16": 8.82513, + "17": 8.52714, + "18": 8.81542, + "19": 8.73224, + "20": 8.76141 } }, "num-zeros": { @@ -31,26 +31,26 @@ "end_step": 20, "step_interval": 1, "values": { - "1": 47889940.0, - "2": 49012016.0, - "3": 48644404.0, - "4": 48106204.0, - "5": 58246720.0, - "6": 63126404.0, - "7": 77135048.0, - "8": 67132312.0, - "9": 67750168.0, - "10": 64553460.0, - "11": 77738728.0, - "12": 64587012.0, - "13": 64661168.0, - "14": 61613476.0, - "15": 61997576.0, - "16": 71691368.0, - "17": 69936240.0, - "18": 70742744.0, - "19": 67630016.0, - "20": 67347280.0 + "1": 131775872.0, + "2": 90954752.0, + "3": 90855608.0, + "4": 55755820.0, + "5": 81253648.0, + "6": 122200864.0, + "7": 130715608.0, + "8": 113168032.0, + "9": 106898848.0, + "10": 102486976.0, + "11": 134026600.0, + "12": 117554208.0, + "13": 120946232.0, + "14": 109072088.0, + "15": 104583256.0, + "16": 125987144.0, + "17": 120386480.0, + "18": 119443640.0, + "19": 115415224.0, + "20": 113935544.0 } }, "mem-allocated-bytes": { @@ -58,26 +58,26 @@ "end_step": 20, "step_interval": 1, "values": { - "1": 1872031232.0, - "2": 1872031232.0, - "3": 1872031232.0, - "4": 1872031232.0, - "5": 1872031232.0, - "6": 1872031232.0, - "7": 1872031232.0, - "8": 1872031232.0, - "9": 1872031232.0, - "10": 1872031232.0, - "11": 1872031232.0, - "12": 1872031232.0, - "13": 1872031232.0, - "14": 1872031232.0, - "15": 1872031232.0, - "16": 1872031232.0, - "17": 1872031232.0, - "18": 1872031232.0, - "19": 1872031232.0, - "20": 1872031232.0 + "1": 1873604096.0, + "2": 1873604096.0, + "3": 1873604096.0, + "4": 1873604096.0, + "5": 1873604096.0, + "6": 1873604096.0, + "7": 1873604096.0, + "8": 1873604096.0, + "9": 1873604096.0, + "10": 1873604096.0, + "11": 1873604096.0, + "12": 1873604096.0, + "13": 1873604096.0, + "14": 1873604096.0, + "15": 1873604096.0, + "16": 1873604096.0, + "17": 1873604096.0, + "18": 1873604096.0, + "19": 1873604096.0, + "20": 1873604096.0 } }, "mem-max-allocated-bytes": { @@ -85,53 +85,52 @@ "end_step": 20, "step_interval": 1, "values": { - "1": 2476529664.0, - "2": 2827837440.0, - "3": 2827837440.0, - "4": 2847291904.0, - "5": 2886460416.0, - "6": 2886460416.0, - "7": 2938775552.0, - "8": 2938775552.0, - "9": 2938775552.0, - "10": 2938775552.0, - "11": 2938775552.0, - "12": 2938775552.0, - "13": 2938775552.0, - "14": 2938775552.0, - "15": 2938775552.0, - "16": 2938775552.0, - "17": 2938775552.0, - "18": 2938775552.0, - "19": 2938775552.0, - "20": 2938775552.0 + "1": 2478305280.0, + "2": 2860624896.0, + "3": 2860624896.0, + "4": 2930518016.0, + "5": 2943496192.0, + "6": 2943496192.0, + "7": 2943496192.0, + "8": 2943496192.0, + "9": 2943496192.0, + "10": 2943496192.0, + "11": 2954289152.0, + "12": 2954289152.0, + "13": 2954289152.0, + "14": 2954289152.0, + "15": 2954289152.0, + "16": 2954289152.0, + "17": 2955154944.0, + "18": 2955154944.0, + "19": 2955154944.0, + "20": 2955356160.0 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 20, "step_interval": 1, "values": { - "1": "nan", - "2": 43.50173, - "3": 1.91918, - "4": 0.33128, - "5": 0.32613, - "6": 0.30839, - "7": 0.30715, - "8": 0.29241, - "9": 0.28352, - "10": 0.2695, - "11": 0.27092, - "12": 0.27945, - "13": 0.26134, - "14": 0.28396, - "15": 0.2654, - "16": 0.2759, - "17": 0.26684, - "18": 0.25266, - "19": 0.25121, - "20": 0.25109 + "2": 43.33046, + "3": 1.94062, + "4": 0.37601, + "5": 0.35056, + "6": 0.32995, + "7": 0.31079, + "8": 0.30807, + "9": 0.32086, + "10": 0.30641, + "11": 0.28508, + "12": 0.29118, + "13": 0.28659, + "14": 0.30051, + "15": 0.27696, + "16": 0.28535, + "17": 0.28785, + "18": 0.27908, + "19": 0.28446, + "20": 0.27351 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_gb200.json index e17135639b3..e662fca9b5b 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_gb200.json @@ -6,54 +6,54 @@ "values": { "1": 10.92072, "2": 10.91167, - "3": 10.92066, - "4": 10.91184, - "5": 10.92404, - "6": 10.91113, - "7": 10.90393, - "8": 10.90985, - "9": 10.91959, - "10": 10.91758, - "11": 10.90389, - "12": 10.90344, - "13": 10.89447, - "14": 10.89157, - "15": 10.88138, - "16": 10.87501, - "17": 10.87674, - "18": 10.8606, - "19": 10.868, - "20": 10.82424, - "21": 10.80928, - "22": 10.79632, - "23": 10.79123, - "24": 10.75749, - "25": 10.76202, - "26": 10.75071, - "27": 10.7414, - "28": 10.67372, - "29": 10.64549, - "30": 10.62348, - "31": 10.61147, - "32": 10.59676, - "33": 10.57395, - "34": 10.54145, - "35": 10.53203, - "36": 10.5234, - "37": 10.48714, - "38": 10.49435, - "39": 10.45409, - "40": 10.44144, - "41": 10.42824, - "42": 10.41826, - "43": 10.38934, - "44": 10.35343, - "45": 10.37024, - "46": 10.33167, - "47": 10.31245, - "48": 10.27749, - "49": 10.27081, - "50": 10.27222 + "3": 10.92063, + "4": 10.91109, + "5": 10.92358, + "6": 10.91096, + "7": 10.90398, + "8": 10.91005, + "9": 10.91975, + "10": 10.91819, + "11": 10.90492, + "12": 10.90256, + "13": 10.89411, + "14": 10.89146, + "15": 10.88102, + "16": 10.87423, + "17": 10.8758, + "18": 10.86035, + "19": 10.86778, + "20": 10.82365, + "21": 10.80909, + "22": 10.79674, + "23": 10.79169, + "24": 10.75821, + "25": 10.76254, + "26": 10.74964, + "27": 10.74103, + "28": 10.67304, + "29": 10.64482, + "30": 10.62363, + "31": 10.6108, + "32": 10.59641, + "33": 10.57382, + "34": 10.54019, + "35": 10.5322, + "36": 10.52253, + "37": 10.48684, + "38": 10.49374, + "39": 10.45374, + "40": 10.44023, + "41": 10.42655, + "42": 10.41693, + "43": 10.38734, + "44": 10.35149, + "45": 10.36826, + "46": 10.32915, + "47": 10.31053, + "48": 10.2757, + "49": 10.26805, + "50": 10.2702 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 18555.0, - "2": 18999.0, - "3": 18950.0, - "4": 18669.0, - "5": 18777.0, - "6": 18519.0, - "7": 18807.0, - "8": 18670.0, - "9": 19307.0, - "10": 19430.0, - "11": 19062.0, - "12": 18644.0, - "13": 19076.0, - "14": 18936.0, - "15": 18541.0, - "16": 18430.0, - "17": 18751.0, - "18": 18698.0, - "19": 19066.0, - "20": 18677.0, - "21": 18789.0, - "22": 18783.0, - "23": 18643.0, - "24": 18816.0, - "25": 18303.0, - "26": 18784.0, - "27": 18692.0, - "28": 18297.0, - "29": 18282.0, - "30": 18205.0, - "31": 18931.0, - "32": 19019.0, - "33": 18449.0, - "34": 18535.0, - "35": 18869.0, - "36": 18390.0, - "37": 18979.0, - "38": 18529.0, - "39": 19178.0, - "40": 19142.0, - "41": 18772.0, - "42": 18360.0, - "43": 19127.0, - "44": 18528.0, - "45": 20133.0, - "46": 19760.0, - "47": 19924.0, - "48": 19649.0, - "49": 21808.0, - "50": 19795.0 + "1": 36032.0, + "2": 37013.0, + "3": 36506.0, + "4": 36647.0, + "5": 36566.0, + "6": 36106.0, + "7": 36662.0, + "8": 36674.0, + "9": 36838.0, + "10": 37375.0, + "11": 37017.0, + "12": 36039.0, + "13": 36808.0, + "14": 36548.0, + "15": 36220.0, + "16": 36072.0, + "17": 36754.0, + "18": 36616.0, + "19": 36918.0, + "20": 36426.0, + "21": 36332.0, + "22": 36730.0, + "23": 36348.0, + "24": 36554.0, + "25": 35486.0, + "26": 36223.0, + "27": 35971.0, + "28": 35883.0, + "29": 35562.0, + "30": 35549.0, + "31": 36626.0, + "32": 36672.0, + "33": 35601.0, + "34": 36086.0, + "35": 37007.0, + "36": 35586.0, + "37": 36542.0, + "38": 35690.0, + "39": 37082.0, + "40": 37310.0, + "41": 36599.0, + "42": 35449.0, + "43": 37576.0, + "44": 36227.0, + "45": 39150.0, + "46": 37880.0, + "47": 38734.0, + "48": 38832.0, + "49": 42289.0, + "50": 39404.0 } }, "mem-allocated-bytes": { @@ -120,54 +120,54 @@ "values": { "1": 1027095552.0, "2": 1027093504.0, - "3": 1027091456.0, - "4": 1027091968.0, + "3": 1027092480.0, + "4": 1027091456.0, "5": 1027093504.0, "6": 1027094528.0, - "7": 1027095040.0, - "8": 1027094016.0, - "9": 1027095040.0, + "7": 1027094016.0, + "8": 1027093504.0, + "9": 1027093504.0, "10": 1027092992.0, "11": 1027093504.0, - "12": 1027092992.0, - "13": 1027091456.0, - "14": 1027092992.0, + "12": 1027093504.0, + "13": 1027090944.0, + "14": 1027092480.0, "15": 1027094528.0, "16": 1027094016.0, - "17": 1027095552.0, + "17": 1027096064.0, "18": 1027096576.0, - "19": 1027095040.0, - "20": 1027092992.0, + "19": 1027093504.0, + "20": 1027093504.0, "21": 1027095552.0, "22": 1027091968.0, - "23": 1027091456.0, - "24": 1027097600.0, - "25": 1027092480.0, - "26": 1027094528.0, + "23": 1027095552.0, + "24": 1027094528.0, + "25": 1027093504.0, + "26": 1027095040.0, "27": 1027092480.0, - "28": 1027094528.0, - "29": 1027091968.0, + "28": 1027095040.0, + "29": 1027094016.0, "30": 1027097600.0, "31": 1027093504.0, - "32": 1027090944.0, - "33": 1027098112.0, + "32": 1027091968.0, + "33": 1027096576.0, "34": 1027096064.0, "35": 1027094528.0, - "36": 1027096064.0, - "37": 1027094528.0, - "38": 1027096064.0, + "36": 1027096576.0, + "37": 1027092992.0, + "38": 1027095040.0, "39": 1027094016.0, "40": 1027096064.0, "41": 1027093504.0, - "42": 1027098112.0, + "42": 1027096576.0, "43": 1027093504.0, - "44": 1027093504.0, + "44": 1027092992.0, "45": 1027095552.0, - "46": 1027093504.0, + "46": 1027091968.0, "47": 1027096064.0, - "48": 1027093504.0, - "49": 1027096576.0, - "50": 1027096064.0 + "48": 1027092992.0, + "49": 1027095552.0, + "50": 1027094016.0 } }, "mem-max-allocated-bytes": { @@ -177,111 +177,110 @@ "values": { "1": 3058330112.0, "2": 3298373120.0, - "3": 3298764288.0, - "4": 3299075072.0, - "5": 3300025344.0, - "6": 3300025344.0, - "7": 3300025344.0, - "8": 3300025344.0, - "9": 3300025344.0, - "10": 3300025344.0, - "11": 3300025344.0, - "12": 3300025344.0, - "13": 3300025344.0, - "14": 3300025344.0, - "15": 3300025344.0, - "16": 3300025344.0, - "17": 3300025344.0, - "18": 3300135424.0, - "19": 3300135424.0, - "20": 3300135424.0, - "21": 3300135424.0, - "22": 3300135424.0, - "23": 3300135424.0, - "24": 3300135424.0, - "25": 3300135424.0, - "26": 3300135424.0, - "27": 3300135424.0, - "28": 3300135424.0, - "29": 3300135424.0, - "30": 3300135424.0, - "31": 3300135424.0, - "32": 3300135424.0, - "33": 3300135424.0, - "34": 3300452352.0, - "35": 3300452352.0, - "36": 3300452352.0, - "37": 3300452352.0, - "38": 3300452352.0, - "39": 3300452352.0, - "40": 3300525056.0, - "41": 3300525056.0, - "42": 3301999616.0, - "43": 3301999616.0, - "44": 3301999616.0, - "45": 3301999616.0, - "46": 3301999616.0, - "47": 3301999616.0, - "48": 3301999616.0, - "49": 3301999616.0, - "50": 3301999616.0 + "3": 3298938368.0, + "4": 3298957824.0, + "5": 3299935232.0, + "6": 3299935232.0, + "7": 3299935232.0, + "8": 3299935232.0, + "9": 3299935232.0, + "10": 3299935232.0, + "11": 3299935232.0, + "12": 3299935232.0, + "13": 3299935232.0, + "14": 3299935232.0, + "15": 3299935232.0, + "16": 3299935232.0, + "17": 3299935232.0, + "18": 3299971584.0, + "19": 3299971584.0, + "20": 3299971584.0, + "21": 3300268544.0, + "22": 3300268544.0, + "23": 3300268544.0, + "24": 3300268544.0, + "25": 3300268544.0, + "26": 3300268544.0, + "27": 3300268544.0, + "28": 3300268544.0, + "29": 3300268544.0, + "30": 3300268544.0, + "31": 3300268544.0, + "32": 3300275712.0, + "33": 3300275712.0, + "34": 3300275712.0, + "35": 3300275712.0, + "36": 3300275712.0, + "37": 3300275712.0, + "38": 3300275712.0, + "39": 3300275712.0, + "40": 3300545536.0, + "41": 3300545536.0, + "42": 3301613056.0, + "43": 3301613056.0, + "44": 3301637632.0, + "45": 3301637632.0, + "46": 3301637632.0, + "47": 3301637632.0, + "48": 3301637632.0, + "49": 3301637632.0, + "50": 3301637632.0 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 50, "step_interval": 1, "values": { - "1": "nan", - "2": 5.08499, - "3": 0.32344, - "4": 0.26578, - "5": 0.26136, - "6": 0.26073, - "7": 0.25791, - "8": 0.26042, - "9": 0.2894, - "10": 0.317, - "11": 0.3188, - "12": 0.32736, - "13": 0.31992, - "14": 0.31394, - "15": 0.30917, - "16": 0.31272, - "17": 0.3125, - "18": 0.31561, - "19": 0.31717, - "20": 0.30933, - "21": 0.31495, - "22": 0.31815, - "23": 0.31889, - "24": 0.32024, - "25": 0.3175, - "26": 0.31981, - "27": 0.31518, - "28": 0.32058, - "29": 0.31851, - "30": 0.32621, - "31": 0.31894, - "32": 0.31952, - "33": 0.31826, - "34": 0.3143, - "35": 0.31728, - "36": 0.31983, - "37": 0.31754, - "38": 0.31813, - "39": 0.32013, - "40": 0.31775, - "41": 0.3177, - "42": 0.31806, - "43": 0.31839, - "44": 0.3185, - "45": 0.31489, - "46": 0.31404, - "47": 0.3105, - "48": 0.31484, - "49": 0.31734, - "50": 0.31583 + "2": 6.28528, + "3": 0.31761, + "4": 0.27043, + "5": 0.26527, + "6": 0.26539, + "7": 0.26148, + "8": 0.265, + "9": 0.26114, + "10": 0.26356, + "11": 0.25992, + "12": 0.7058, + "13": 0.25379, + "14": 0.25867, + "15": 0.72393, + "16": 0.25422, + "17": 0.2511, + "18": 0.72454, + "19": 0.25117, + "20": 0.2542, + "21": 0.25116, + "22": 0.25053, + "23": 0.25184, + "24": 0.70444, + "25": 0.25188, + "26": 0.25183, + "27": 0.254, + "28": 0.25276, + "29": 0.25374, + "30": 0.25477, + "31": 0.25084, + "32": 0.25364, + "33": 0.25172, + "34": 0.2559, + "35": 0.24913, + "36": 0.71441, + "37": 0.26234, + "38": 0.25097, + "39": 0.2543, + "40": 0.25198, + "41": 0.25096, + "42": 0.25123, + "43": 0.2525, + "44": 0.25335, + "45": 0.2506, + "46": 0.25256, + "47": 0.25168, + "48": 0.25675, + "49": 0.25156, + "50": 0.24879 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_lts_dgx_a100.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_lts_dgx_a100.json index 838b50e3d9e..f3e02b30505 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_lts_dgx_a100.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_lts_dgx_a100.json @@ -6,54 +6,54 @@ "values": { "1": 10.94539, "2": 10.93781, - "3": 10.93988, - "4": 10.93542, - "5": 10.94061, - "6": 10.92945, - "7": 10.93511, - "8": 10.93552, - "9": 10.93297, - "10": 10.92843, - "11": 10.93667, - "12": 10.93013, - "13": 10.93018, - "14": 10.92297, - "15": 10.91273, - "16": 10.89804, - "17": 10.90266, - "18": 10.89344, - "19": 10.89235, - "20": 10.8356, - "21": 10.83076, - "22": 10.82914, - "23": 10.8111, - "24": 10.78639, - "25": 10.78685, - "26": 10.76881, - "27": 10.76703, - "28": 10.71292, - "29": 10.66853, - "30": 10.63948, - "31": 10.64276, - "32": 10.61112, - "33": 10.57981, - "34": 10.56137, - "35": 10.5585, - "36": 10.52917, - "37": 10.50576, - "38": 10.50299, - "39": 10.47449, - "40": 10.45267, - "41": 10.4227, - "42": 10.40413, - "43": 10.39432, - "44": 10.35641, - "45": 10.37876, - "46": 10.32832, - "47": 10.3069, - "48": 10.28659, - "49": 10.26877, - "50": 10.26785 + "3": 10.9397, + "4": 10.93551, + "5": 10.9406, + "6": 10.9298, + "7": 10.93519, + "8": 10.93512, + "9": 10.93213, + "10": 10.92859, + "11": 10.93725, + "12": 10.93058, + "13": 10.93, + "14": 10.92223, + "15": 10.9122, + "16": 10.89697, + "17": 10.90252, + "18": 10.89393, + "19": 10.89204, + "20": 10.8349, + "21": 10.83161, + "22": 10.82927, + "23": 10.81156, + "24": 10.78511, + "25": 10.78632, + "26": 10.76689, + "27": 10.76635, + "28": 10.7117, + "29": 10.66898, + "30": 10.63886, + "31": 10.64231, + "32": 10.61117, + "33": 10.57889, + "34": 10.561, + "35": 10.55781, + "36": 10.52984, + "37": 10.50502, + "38": 10.50146, + "39": 10.47347, + "40": 10.45073, + "41": 10.42162, + "42": 10.40328, + "43": 10.39256, + "44": 10.35454, + "45": 10.37698, + "46": 10.3265, + "47": 10.30373, + "48": 10.28469, + "49": 10.2659, + "50": 10.26523 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 19375.0, - "2": 19327.0, - "3": 19211.0, - "4": 19543.0, - "5": 19271.0, - "6": 18913.0, - "7": 19722.0, - "8": 19225.0, - "9": 19509.0, - "10": 19667.0, - "11": 19423.0, - "12": 18664.0, - "13": 19112.0, - "14": 19234.0, - "15": 18910.0, - "16": 18987.0, - "17": 19328.0, - "18": 19451.0, - "19": 19613.0, - "20": 19074.0, - "21": 19178.0, - "22": 19098.0, - "23": 19117.0, - "24": 19272.0, - "25": 18744.0, - "26": 18972.0, - "27": 18941.0, - "28": 18736.0, - "29": 18664.0, - "30": 18286.0, - "31": 18789.0, - "32": 18884.0, - "33": 18512.0, - "34": 18793.0, - "35": 18540.0, - "36": 18586.0, - "37": 18667.0, - "38": 18246.0, - "39": 18480.0, - "40": 19065.0, - "41": 18332.0, - "42": 18256.0, - "43": 19136.0, - "44": 18646.0, - "45": 20341.0, - "46": 19744.0, - "47": 19705.0, - "48": 20120.0, - "49": 21420.0, - "50": 19996.0 + "1": 36962.0, + "2": 37109.0, + "3": 37373.0, + "4": 37302.0, + "5": 36946.0, + "6": 35914.0, + "7": 37793.0, + "8": 36384.0, + "9": 37338.0, + "10": 37740.0, + "11": 37084.0, + "12": 36015.0, + "13": 37331.0, + "14": 37383.0, + "15": 37017.0, + "16": 36521.0, + "17": 36794.0, + "18": 36925.0, + "19": 37355.0, + "20": 37133.0, + "21": 36720.0, + "22": 37145.0, + "23": 37342.0, + "24": 37133.0, + "25": 35841.0, + "26": 37086.0, + "27": 36940.0, + "28": 36425.0, + "29": 35839.0, + "30": 35502.0, + "31": 36894.0, + "32": 36579.0, + "33": 35924.0, + "34": 35750.0, + "35": 36642.0, + "36": 35743.0, + "37": 35948.0, + "38": 35710.0, + "39": 36207.0, + "40": 36812.0, + "41": 36175.0, + "42": 35843.0, + "43": 37560.0, + "44": 35836.0, + "45": 39555.0, + "46": 38335.0, + "47": 38031.0, + "48": 38990.0, + "49": 42019.0, + "50": 39433.0 } }, "mem-allocated-bytes": { @@ -121,53 +121,53 @@ "1": 880363520.0, "2": 880364544.0, "3": 880365056.0, - "4": 880365056.0, + "4": 880364032.0, "5": 880366592.0, "6": 880361984.0, - "7": 880364032.0, - "8": 880365568.0, - "9": 880365056.0, - "10": 880363008.0, - "11": 880362496.0, - "12": 880363008.0, - "13": 880363520.0, + "7": 880363520.0, + "8": 880363520.0, + "9": 880365568.0, + "10": 880362496.0, + "11": 880363008.0, + "12": 880365056.0, + "13": 880364032.0, "14": 880363520.0, - "15": 880367616.0, - "16": 880365056.0, + "15": 880368128.0, + "16": 880365568.0, "17": 880365056.0, - "18": 880363008.0, + "18": 880365056.0, "19": 880364544.0, - "20": 880366592.0, - "21": 880363520.0, - "22": 880365568.0, - "23": 880362496.0, - "24": 880365568.0, - "25": 880365568.0, + "20": 880365568.0, + "21": 880364544.0, + "22": 880365056.0, + "23": 880363520.0, + "24": 880366592.0, + "25": 880364544.0, "26": 880367616.0, - "27": 880362496.0, - "28": 880363520.0, - "29": 880365568.0, - "30": 880363520.0, - "31": 880364032.0, - "32": 880365056.0, + "27": 880363008.0, + "28": 880364544.0, + "29": 880366080.0, + "30": 880363008.0, + "31": 880365568.0, + "32": 880363520.0, "33": 880366080.0, - "34": 880365568.0, + "34": 880364544.0, "35": 880363008.0, "36": 880365568.0, - "37": 880365568.0, - "38": 880366080.0, + "37": 880364544.0, + "38": 880364544.0, "39": 880363520.0, - "40": 880365056.0, + "40": 880364544.0, "41": 880365568.0, - "42": 880365568.0, + "42": 880364544.0, "43": 880364544.0, - "44": 880365568.0, - "45": 880364544.0, - "46": 880363008.0, + "44": 880366592.0, + "45": 880365056.0, + "46": 880362496.0, "47": 880363008.0, - "48": 880364032.0, - "49": 880364032.0, - "50": 880359936.0 + "48": 880366080.0, + "49": 880365568.0, + "50": 880360448.0 } }, "mem-max-allocated-bytes": { @@ -175,113 +175,112 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 2910779392.0, + "1": 2910778880.0, "2": 3151056896.0, "3": 3151056896.0, - "4": 3151056896.0, - "5": 3151663104.0, - "6": 3151663104.0, - "7": 3151663104.0, - "8": 3151663104.0, - "9": 3151663104.0, - "10": 3151663104.0, - "11": 3151663104.0, - "12": 3151663104.0, - "13": 3151663104.0, - "14": 3151663104.0, - "15": 3151936512.0, - "16": 3151936512.0, - "17": 3151936512.0, - "18": 3151936512.0, - "19": 3151936512.0, - "20": 3151936512.0, - "21": 3151936512.0, - "22": 3151936512.0, - "23": 3151936512.0, - "24": 3151936512.0, - "25": 3151936512.0, - "26": 3152254976.0, - "27": 3152254976.0, - "28": 3152254976.0, - "29": 3152254976.0, - "30": 3152254976.0, - "31": 3152254976.0, - "32": 3153094144.0, - "33": 3153094144.0, - "34": 3153094144.0, - "35": 3153094144.0, - "36": 3153094144.0, - "37": 3153094144.0, - "38": 3153094144.0, - "39": 3153094144.0, - "40": 3153094144.0, - "41": 3153094144.0, - "42": 3153094144.0, - "43": 3153094144.0, - "44": 3153094144.0, - "45": 3153094144.0, - "46": 3153094144.0, - "47": 3153094144.0, - "48": 3153094144.0, - "49": 3153094144.0, - "50": 3153094144.0 + "4": 3151080448.0, + "5": 3151697408.0, + "6": 3151697408.0, + "7": 3151697408.0, + "8": 3151697408.0, + "9": 3151697408.0, + "10": 3151697408.0, + "11": 3151697408.0, + "12": 3151697408.0, + "13": 3152274432.0, + "14": 3152274432.0, + "15": 3152390144.0, + "16": 3152390144.0, + "17": 3152390144.0, + "18": 3152390144.0, + "19": 3152390144.0, + "20": 3152390144.0, + "21": 3152390144.0, + "22": 3152390144.0, + "23": 3152390144.0, + "24": 3152390144.0, + "25": 3152390144.0, + "26": 3152390144.0, + "27": 3152390144.0, + "28": 3152390144.0, + "29": 3152390144.0, + "30": 3152390144.0, + "31": 3152390144.0, + "32": 3152975360.0, + "33": 3152975360.0, + "34": 3152975360.0, + "35": 3152975360.0, + "36": 3152975360.0, + "37": 3152975360.0, + "38": 3152975360.0, + "39": 3152975360.0, + "40": 3152975360.0, + "41": 3152975360.0, + "42": 3152975360.0, + "43": 3152975360.0, + "44": 3152975360.0, + "45": 3152975360.0, + "46": 3152975360.0, + "47": 3152975360.0, + "48": 3152975360.0, + "49": 3152975360.0, + "50": 3152975360.0 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 50, "step_interval": 1, "values": { - "1": "nan", - "2": 3.77561, - "3": 0.31056, - "4": 0.29762, - "5": 0.28993, - "6": 0.28778, - "7": 0.28831, - "8": 0.28834, - "9": 0.28927, - "10": 0.28725, - "11": 0.28853, - "12": 0.28466, - "13": 0.2851, - "14": 0.28902, - "15": 0.28576, - "16": 0.28514, - "17": 0.28508, - "18": 0.28838, - "19": 0.2956, - "20": 0.28887, - "21": 0.30039, - "22": 0.28563, - "23": 0.28711, - "24": 0.3062, - "25": 0.29335, - "26": 0.28872, - "27": 0.28681, - "28": 0.28829, - "29": 0.28748, - "30": 0.28833, - "31": 0.28959, - "32": 0.28698, - "33": 0.28666, - "34": 0.28781, - "35": 0.28701, - "36": 0.28742, - "37": 0.28598, - "38": 0.28718, - "39": 0.28723, - "40": 0.28749, - "41": 0.28732, - "42": 0.28845, - "43": 0.28871, - "44": 0.28768, - "45": 0.28982, - "46": 0.28854, - "47": 0.28803, - "48": 0.28909, - "49": 0.28845, - "50": 0.28839 + "2": 3.87433, + "3": 0.79787, + "4": 0.29842, + "5": 0.29697, + "6": 0.29864, + "7": 0.29825, + "8": 0.2982, + "9": 0.29558, + "10": 0.29421, + "11": 0.29148, + "12": 0.29442, + "13": 0.29148, + "14": 0.28921, + "15": 0.28988, + "16": 0.29585, + "17": 0.30217, + "18": 0.29118, + "19": 0.29216, + "20": 0.29178, + "21": 0.29109, + "22": 0.29017, + "23": 0.29101, + "24": 0.2916, + "25": 0.29162, + "26": 0.29153, + "27": 0.29174, + "28": 0.29111, + "29": 0.29061, + "30": 0.29096, + "31": 0.2915, + "32": 0.29181, + "33": 0.29104, + "34": 0.29255, + "35": 0.29096, + "36": 0.29143, + "37": 0.29399, + "38": 0.29182, + "39": 0.29224, + "40": 0.29166, + "41": 0.29146, + "42": 0.29273, + "43": 0.29245, + "44": 0.291, + "45": 0.29197, + "46": 0.2921, + "47": 0.29168, + "48": 0.29201, + "49": 0.2909, + "50": 0.29225 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_gb200.json index d7084a96169..a6203d8eadb 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_gb200.json @@ -6,54 +6,54 @@ "values": { "1": 10.92072, "2": 10.91167, - "3": 10.92066, - "4": 10.91184, - "5": 10.92404, - "6": 10.91113, - "7": 10.90393, - "8": 10.90985, - "9": 10.91959, - "10": 10.91758, - "11": 10.90389, - "12": 10.90344, - "13": 10.89447, - "14": 10.89157, - "15": 10.88138, - "16": 10.87501, - "17": 10.87674, - "18": 10.8606, - "19": 10.868, - "20": 10.82424, - "21": 10.80928, - "22": 10.79632, - "23": 10.79123, - "24": 10.75749, - "25": 10.76202, - "26": 10.75071, - "27": 10.7414, - "28": 10.67372, - "29": 10.64549, - "30": 10.62348, - "31": 10.61147, - "32": 10.59676, - "33": 10.57395, - "34": 10.54145, - "35": 10.53203, - "36": 10.5234, - "37": 10.48714, - "38": 10.49435, - "39": 10.45409, - "40": 10.44144, - "41": 10.42824, - "42": 10.41826, - "43": 10.38934, - "44": 10.35343, - "45": 10.37024, - "46": 10.33167, - "47": 10.31245, - "48": 10.27749, - "49": 10.27081, - "50": 10.27222 + "3": 10.92063, + "4": 10.91109, + "5": 10.92358, + "6": 10.91096, + "7": 10.90398, + "8": 10.91005, + "9": 10.91975, + "10": 10.91819, + "11": 10.90492, + "12": 10.90256, + "13": 10.89411, + "14": 10.89146, + "15": 10.88102, + "16": 10.87423, + "17": 10.8758, + "18": 10.86035, + "19": 10.86778, + "20": 10.82365, + "21": 10.80909, + "22": 10.79674, + "23": 10.79169, + "24": 10.75821, + "25": 10.76254, + "26": 10.74964, + "27": 10.74103, + "28": 10.67304, + "29": 10.64482, + "30": 10.62363, + "31": 10.6108, + "32": 10.59641, + "33": 10.57382, + "34": 10.54019, + "35": 10.5322, + "36": 10.52253, + "37": 10.48684, + "38": 10.49374, + "39": 10.45374, + "40": 10.44023, + "41": 10.42655, + "42": 10.41693, + "43": 10.38734, + "44": 10.35149, + "45": 10.36826, + "46": 10.32915, + "47": 10.31053, + "48": 10.2757, + "49": 10.26805, + "50": 10.2702 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 18555.0, - "2": 18999.0, - "3": 18950.0, - "4": 18669.0, - "5": 18777.0, - "6": 18519.0, - "7": 18807.0, - "8": 18670.0, - "9": 19307.0, - "10": 19430.0, - "11": 19062.0, - "12": 18644.0, - "13": 19076.0, - "14": 18936.0, - "15": 18541.0, - "16": 18430.0, - "17": 18751.0, - "18": 18698.0, - "19": 19066.0, - "20": 18677.0, - "21": 18789.0, - "22": 18783.0, - "23": 18643.0, - "24": 18816.0, - "25": 18303.0, - "26": 18784.0, - "27": 18692.0, - "28": 18297.0, - "29": 18282.0, - "30": 18205.0, - "31": 18931.0, - "32": 19019.0, - "33": 18449.0, - "34": 18535.0, - "35": 18869.0, - "36": 18390.0, - "37": 18979.0, - "38": 18529.0, - "39": 19178.0, - "40": 19142.0, - "41": 18772.0, - "42": 18360.0, - "43": 19127.0, - "44": 18528.0, - "45": 20133.0, - "46": 19760.0, - "47": 19924.0, - "48": 19649.0, - "49": 21808.0, - "50": 19795.0 + "1": 36032.0, + "2": 37013.0, + "3": 36506.0, + "4": 36647.0, + "5": 36566.0, + "6": 36106.0, + "7": 36662.0, + "8": 36674.0, + "9": 36838.0, + "10": 37375.0, + "11": 37017.0, + "12": 36039.0, + "13": 36808.0, + "14": 36548.0, + "15": 36220.0, + "16": 36072.0, + "17": 36754.0, + "18": 36616.0, + "19": 36918.0, + "20": 36426.0, + "21": 36332.0, + "22": 36730.0, + "23": 36348.0, + "24": 36554.0, + "25": 35486.0, + "26": 36223.0, + "27": 35971.0, + "28": 35883.0, + "29": 35562.0, + "30": 35549.0, + "31": 36626.0, + "32": 36672.0, + "33": 35601.0, + "34": 36086.0, + "35": 37007.0, + "36": 35586.0, + "37": 36542.0, + "38": 35690.0, + "39": 37082.0, + "40": 37310.0, + "41": 36599.0, + "42": 35449.0, + "43": 37576.0, + "44": 36227.0, + "45": 39150.0, + "46": 37880.0, + "47": 38734.0, + "48": 38832.0, + "49": 42289.0, + "50": 39404.0 } }, "mem-allocated-bytes": { @@ -120,54 +120,54 @@ "values": { "1": 1027095552.0, "2": 1027093504.0, - "3": 1027091456.0, - "4": 1027091968.0, + "3": 1027092480.0, + "4": 1027091456.0, "5": 1027093504.0, "6": 1027094528.0, - "7": 1027095040.0, - "8": 1027094016.0, - "9": 1027095040.0, + "7": 1027094016.0, + "8": 1027093504.0, + "9": 1027093504.0, "10": 1027092992.0, "11": 1027093504.0, - "12": 1027092992.0, - "13": 1027091456.0, - "14": 1027092992.0, + "12": 1027093504.0, + "13": 1027090944.0, + "14": 1027092480.0, "15": 1027094528.0, "16": 1027094016.0, - "17": 1027095552.0, + "17": 1027096064.0, "18": 1027096576.0, - "19": 1027095040.0, - "20": 1027092992.0, + "19": 1027093504.0, + "20": 1027093504.0, "21": 1027095552.0, "22": 1027091968.0, - "23": 1027091456.0, - "24": 1027097600.0, - "25": 1027092480.0, - "26": 1027094528.0, + "23": 1027095552.0, + "24": 1027094528.0, + "25": 1027093504.0, + "26": 1027095040.0, "27": 1027092480.0, - "28": 1027094528.0, - "29": 1027091968.0, + "28": 1027095040.0, + "29": 1027094016.0, "30": 1027097600.0, "31": 1027093504.0, - "32": 1027090944.0, - "33": 1027098112.0, + "32": 1027091968.0, + "33": 1027096576.0, "34": 1027096064.0, "35": 1027094528.0, - "36": 1027096064.0, - "37": 1027094528.0, - "38": 1027096064.0, + "36": 1027096576.0, + "37": 1027092992.0, + "38": 1027095040.0, "39": 1027094016.0, "40": 1027096064.0, "41": 1027093504.0, - "42": 1027098112.0, + "42": 1027096576.0, "43": 1027093504.0, - "44": 1027093504.0, + "44": 1027092992.0, "45": 1027095552.0, - "46": 1027093504.0, + "46": 1027091968.0, "47": 1027096064.0, - "48": 1027093504.0, - "49": 1027096576.0, - "50": 1027096064.0 + "48": 1027092992.0, + "49": 1027095552.0, + "50": 1027094016.0 } }, "mem-max-allocated-bytes": { @@ -177,111 +177,110 @@ "values": { "1": 3058330112.0, "2": 3298373120.0, - "3": 3299813888.0, - "4": 3299813888.0, - "5": 3300025344.0, - "6": 3300025344.0, - "7": 3300025344.0, - "8": 3300025344.0, - "9": 3300025344.0, - "10": 3300025344.0, - "11": 3300025344.0, - "12": 3300025344.0, - "13": 3300025344.0, - "14": 3300025344.0, - "15": 3300025344.0, - "16": 3300025344.0, - "17": 3300025344.0, - "18": 3300135424.0, - "19": 3300135424.0, - "20": 3300135424.0, - "21": 3300973056.0, - "22": 3300973056.0, - "23": 3300973056.0, - "24": 3300973056.0, - "25": 3300973056.0, - "26": 3300973056.0, - "27": 3300973056.0, - "28": 3300973056.0, - "29": 3300973056.0, - "30": 3300973056.0, - "31": 3300973056.0, - "32": 3300973056.0, - "33": 3300973056.0, - "34": 3300973056.0, - "35": 3300973056.0, - "36": 3300973056.0, - "37": 3300973056.0, - "38": 3300973056.0, - "39": 3300973056.0, - "40": 3300973056.0, - "41": 3300973056.0, - "42": 3301952512.0, - "43": 3301952512.0, - "44": 3301952512.0, - "45": 3301952512.0, - "46": 3301952512.0, - "47": 3301952512.0, - "48": 3301952512.0, - "49": 3301952512.0, - "50": 3301952512.0 + "3": 3298938368.0, + "4": 3298957824.0, + "5": 3299935232.0, + "6": 3299935232.0, + "7": 3299935232.0, + "8": 3299935232.0, + "9": 3299935232.0, + "10": 3299935232.0, + "11": 3299935232.0, + "12": 3299935232.0, + "13": 3299935232.0, + "14": 3299935232.0, + "15": 3299935232.0, + "16": 3299935232.0, + "17": 3299935232.0, + "18": 3299971584.0, + "19": 3299971584.0, + "20": 3299971584.0, + "21": 3300268544.0, + "22": 3300268544.0, + "23": 3300268544.0, + "24": 3300268544.0, + "25": 3300268544.0, + "26": 3300268544.0, + "27": 3300268544.0, + "28": 3300268544.0, + "29": 3300268544.0, + "30": 3300268544.0, + "31": 3300268544.0, + "32": 3300275712.0, + "33": 3300275712.0, + "34": 3300275712.0, + "35": 3300275712.0, + "36": 3300275712.0, + "37": 3300275712.0, + "38": 3300275712.0, + "39": 3300275712.0, + "40": 3300545536.0, + "41": 3300545536.0, + "42": 3301613056.0, + "43": 3301613056.0, + "44": 3301637632.0, + "45": 3301637632.0, + "46": 3301637632.0, + "47": 3301637632.0, + "48": 3301637632.0, + "49": 3301637632.0, + "50": 3301637632.0 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 50, "step_interval": 1, "values": { - "1": "nan", - "2": 5.09505, - "3": 0.37144, - "4": 0.33646, - "5": 0.32892, - "6": 0.31965, - "7": 0.32563, - "8": 0.32577, - "9": 0.32429, - "10": 0.32403, - "11": 0.31694, - "12": 0.31735, - "13": 0.35415, - "14": 0.47485, - "15": 0.51791, - "16": 0.31929, - "17": 0.34876, - "18": 0.45504, - "19": 0.36122, - "20": 0.33659, - "21": 0.59492, - "22": 0.31743, - "23": 0.32651, - "24": 0.30634, - "25": 0.28312, - "26": 0.26693, - "27": 0.26144, - "28": 0.25899, - "29": 0.25409, - "30": 0.25775, - "31": 0.25227, - "32": 0.26194, - "33": 0.27509, - "34": 0.25455, - "35": 0.25264, - "36": 0.25091, - "37": 0.25181, - "38": 0.25531, - "39": 0.25406, - "40": 0.25889, - "41": 0.25316, - "42": 0.25117, - "43": 0.25122, - "44": 0.25049, - "45": 0.24793, - "46": 0.24796, - "47": 0.25122, - "48": 0.25116, - "49": 0.24844, - "50": 0.2517 + "2": 6.38561, + "3": 0.36861, + "4": 0.32134, + "5": 0.32224, + "6": 0.3204, + "7": 0.32184, + "8": 0.3464, + "9": 0.35016, + "10": 0.35249, + "11": 0.34499, + "12": 0.33662, + "13": 0.3405, + "14": 0.3387, + "15": 0.33414, + "16": 0.33537, + "17": 0.33614, + "18": 1.02806, + "19": 0.35635, + "20": 0.34369, + "21": 0.73393, + "22": 0.33582, + "23": 0.32712, + "24": 0.31365, + "25": 0.30207, + "26": 0.30527, + "27": 0.91245, + "28": 0.29693, + "29": 0.28094, + "30": 0.28223, + "31": 0.27862, + "32": 0.26837, + "33": 0.27548, + "34": 0.29592, + "35": 0.31612, + "36": 0.99505, + "37": 0.32506, + "38": 0.31285, + "39": 0.31492, + "40": 0.31174, + "41": 0.31603, + "42": 0.31393, + "43": 0.31082, + "44": 0.31619, + "45": 0.85935, + "46": 0.31478, + "47": 0.31605, + "48": 0.3246, + "49": 0.31906, + "50": 0.31812 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts_etp1_ep4/golden_values_lts_dgx_a100.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts_etp1_ep4/golden_values_lts_dgx_a100.json index 611a07e2b83..ef16659f9c1 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts_etp1_ep4/golden_values_lts_dgx_a100.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts_etp1_ep4/golden_values_lts_dgx_a100.json @@ -6,54 +6,54 @@ "values": { "1": 10.9302, "2": 10.91691, - "3": 10.92189, - "4": 10.9142, - "5": 10.92371, - "6": 10.91339, - "7": 10.92299, - "8": 10.92205, - "9": 10.91978, - "10": 10.92093, - "11": 10.91423, - "12": 10.91313, - "13": 10.9027, - "14": 10.89272, - "15": 10.88045, - "16": 10.86859, - "17": 10.87268, - "18": 10.86173, - "19": 10.85772, - "20": 10.77961, - "21": 10.78492, - "22": 10.75544, - "23": 10.74492, - "24": 10.71118, - "25": 10.70861, - "26": 10.68985, - "27": 10.66132, - "28": 10.60998, - "29": 10.56139, - "30": 10.5372, - "31": 10.53704, - "32": 10.50615, - "33": 10.46766, - "34": 10.44278, - "35": 10.44029, - "36": 10.41373, - "37": 10.39529, - "38": 10.39286, - "39": 10.35169, - "40": 10.34009, - "41": 10.30991, - "42": 10.29138, - "43": 10.28249, - "44": 10.24452, - "45": 10.27186, - "46": 10.22035, - "47": 10.20817, - "48": 10.18251, - "49": 10.16969, - "50": 10.17769 + "3": 10.92196, + "4": 10.91423, + "5": 10.92272, + "6": 10.91355, + "7": 10.92271, + "8": 10.92185, + "9": 10.92021, + "10": 10.92131, + "11": 10.91421, + "12": 10.91373, + "13": 10.90212, + "14": 10.89162, + "15": 10.88053, + "16": 10.86797, + "17": 10.87285, + "18": 10.86226, + "19": 10.8578, + "20": 10.77998, + "21": 10.78618, + "22": 10.75518, + "23": 10.74489, + "24": 10.71143, + "25": 10.70865, + "26": 10.68928, + "27": 10.66141, + "28": 10.61008, + "29": 10.56136, + "30": 10.53753, + "31": 10.53634, + "32": 10.50503, + "33": 10.46717, + "34": 10.44123, + "35": 10.43983, + "36": 10.4129, + "37": 10.39414, + "38": 10.39145, + "39": 10.35015, + "40": 10.33833, + "41": 10.30812, + "42": 10.28972, + "43": 10.28036, + "44": 10.24249, + "45": 10.26982, + "46": 10.21796, + "47": 10.20654, + "48": 10.18073, + "49": 10.16738, + "50": 10.17578 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 16515.0, - "2": 16687.0, - "3": 17072.0, - "4": 16843.0, - "5": 16250.0, - "6": 16224.0, - "7": 17373.0, - "8": 16300.0, - "9": 16750.0, - "10": 16957.0, - "11": 16650.0, - "12": 16103.0, - "13": 17107.0, - "14": 16984.0, - "15": 16262.0, - "16": 16161.0, - "17": 16530.0, - "18": 16663.0, - "19": 16837.0, - "20": 16455.0, - "21": 16801.0, - "22": 16552.0, - "23": 16543.0, - "24": 16944.0, - "25": 16044.0, - "26": 16907.0, - "27": 16562.0, - "28": 16489.0, - "29": 16349.0, - "30": 16530.0, - "31": 17203.0, - "32": 17391.0, - "33": 17239.0, - "34": 17162.0, - "35": 17628.0, - "36": 17811.0, - "37": 17985.0, - "38": 16934.0, - "39": 18416.0, - "40": 19262.0, - "41": 17827.0, - "42": 18251.0, - "43": 19492.0, - "44": 18531.0, - "45": 20958.0, - "46": 19964.0, - "47": 20020.0, - "48": 20817.0, - "49": 22662.0, - "50": 20875.0 + "1": 31826.0, + "2": 32260.0, + "3": 32832.0, + "4": 32678.0, + "5": 31720.0, + "6": 31490.0, + "7": 33597.0, + "8": 31520.0, + "9": 32121.0, + "10": 33148.0, + "11": 32472.0, + "12": 30509.0, + "13": 32620.0, + "14": 32548.0, + "15": 31380.0, + "16": 31580.0, + "17": 32081.0, + "18": 32017.0, + "19": 32252.0, + "20": 32052.0, + "21": 32042.0, + "22": 32636.0, + "23": 32178.0, + "24": 32545.0, + "25": 31285.0, + "26": 32616.0, + "27": 32487.0, + "28": 31851.0, + "29": 32213.0, + "30": 32377.0, + "31": 33795.0, + "32": 34476.0, + "33": 33833.0, + "34": 33631.0, + "35": 34939.0, + "36": 34581.0, + "37": 35588.0, + "38": 33759.0, + "39": 36170.0, + "40": 37194.0, + "41": 35254.0, + "42": 35553.0, + "43": 38452.0, + "44": 36716.0, + "45": 41189.0, + "46": 39193.0, + "47": 38846.0, + "48": 40882.0, + "49": 44711.0, + "50": 40679.0 } }, "mem-allocated-bytes": { @@ -120,54 +120,54 @@ "values": { "1": 1364415488.0, "2": 1364417536.0, - "3": 1364402176.0, - "4": 1364398592.0, - "5": 1364404736.0, + "3": 1364402688.0, + "4": 1364399616.0, + "5": 1364405760.0, "6": 1364394496.0, - "7": 1364425728.0, - "8": 1364408832.0, - "9": 1364413440.0, - "10": 1364403200.0, - "11": 1364400640.0, - "12": 1364401664.0, - "13": 1364397056.0, - "14": 1364410368.0, - "15": 1364410880.0, - "16": 1364399616.0, - "17": 1364394496.0, - "18": 1364395520.0, - "19": 1364397056.0, - "20": 1364392448.0, - "21": 1364393984.0, - "22": 1364391936.0, - "23": 1364383744.0, - "24": 1364384768.0, - "25": 1364383744.0, - "26": 1364376064.0, - "27": 1364375040.0, - "28": 1364373504.0, - "29": 1364370432.0, - "30": 1364362240.0, - "31": 1364374528.0, - "32": 1364361728.0, - "33": 1364359680.0, - "34": 1364368896.0, - "35": 1364368896.0, - "36": 1364355072.0, - "37": 1364369408.0, - "38": 1364370432.0, - "39": 1364371456.0, - "40": 1364363264.0, + "7": 1364424192.0, + "8": 1364407808.0, + "9": 1364412928.0, + "10": 1364402176.0, + "11": 1364401152.0, + "12": 1364402176.0, + "13": 1364396032.0, + "14": 1364409856.0, + "15": 1364411904.0, + "16": 1364401152.0, + "17": 1364393984.0, + "18": 1364393984.0, + "19": 1364396032.0, + "20": 1364392960.0, + "21": 1364395008.0, + "22": 1364392960.0, + "23": 1364384768.0, + "24": 1364386304.0, + "25": 1364382720.0, + "26": 1364375552.0, + "27": 1364374016.0, + "28": 1364372480.0, + "29": 1364371456.0, + "30": 1364363264.0, + "31": 1364372480.0, + "32": 1364362752.0, + "33": 1364359168.0, + "34": 1364368384.0, + "35": 1364369408.0, + "36": 1364353024.0, + "37": 1364368384.0, + "38": 1364371968.0, + "39": 1364372992.0, + "40": 1364362240.0, "41": 1364362240.0, - "42": 1364368384.0, + "42": 1364370432.0, "43": 1364346880.0, "44": 1364358144.0, - "45": 1364359168.0, - "46": 1364368384.0, - "47": 1364364288.0, - "48": 1364359680.0, - "49": 1364357632.0, - "50": 1364362752.0 + "45": 1364358656.0, + "46": 1364367872.0, + "47": 1364363776.0, + "48": 1364359168.0, + "49": 1364358144.0, + "50": 1364361728.0 } }, "mem-max-allocated-bytes": { @@ -175,113 +175,112 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 3401542656.0, - "2": 3970463744.0, - "3": 3970463744.0, - "4": 3970463744.0, - "5": 3970463744.0, - "6": 3970463744.0, - "7": 3977521152.0, - "8": 3977521152.0, - "9": 3977521152.0, - "10": 3977521152.0, - "11": 3977521152.0, - "12": 3977521152.0, - "13": 3977521152.0, - "14": 3977521152.0, - "15": 3977521152.0, - "16": 3977521152.0, - "17": 3977521152.0, - "18": 3977521152.0, - "19": 3977521152.0, - "20": 3977521152.0, - "21": 3977521152.0, - "22": 3977521152.0, - "23": 3977521152.0, - "24": 3977521152.0, - "25": 3977521152.0, - "26": 3977521152.0, - "27": 3977521152.0, - "28": 3977521152.0, - "29": 3977521152.0, - "30": 3977521152.0, - "31": 3977521152.0, - "32": 3977521152.0, - "33": 3977521152.0, - "34": 3977521152.0, - "35": 3977521152.0, - "36": 3977521152.0, - "37": 3977521152.0, - "38": 3977521152.0, - "39": 3977521152.0, - "40": 3977521152.0, - "41": 3977521152.0, - "42": 3977521152.0, - "43": 3977521152.0, - "44": 3977521152.0, - "45": 3977521152.0, - "46": 3977521152.0, - "47": 3977521152.0, - "48": 3977521152.0, - "49": 3977521152.0, - "50": 3977521152.0 + "1": 3401542144.0, + "2": 3970463232.0, + "3": 3970463232.0, + "4": 3970463232.0, + "5": 3970463232.0, + "6": 3970463232.0, + "7": 3977527296.0, + "8": 3977527296.0, + "9": 3977527296.0, + "10": 3977527296.0, + "11": 3977527296.0, + "12": 3977527296.0, + "13": 3977527296.0, + "14": 3977527296.0, + "15": 3977527296.0, + "16": 3977527296.0, + "17": 3977527296.0, + "18": 3977527296.0, + "19": 3977527296.0, + "20": 3977527296.0, + "21": 3977527296.0, + "22": 3977527296.0, + "23": 3977527296.0, + "24": 3977527296.0, + "25": 3977527296.0, + "26": 3977527296.0, + "27": 3977527296.0, + "28": 3977527296.0, + "29": 3977527296.0, + "30": 3977527296.0, + "31": 3977527296.0, + "32": 3977527296.0, + "33": 3977527296.0, + "34": 3977527296.0, + "35": 3977527296.0, + "36": 3977527296.0, + "37": 3977527296.0, + "38": 3977527296.0, + "39": 3977527296.0, + "40": 3977527296.0, + "41": 3977527296.0, + "42": 3977527296.0, + "43": 3977527296.0, + "44": 3977527296.0, + "45": 3977527296.0, + "46": 3977527296.0, + "47": 3977527296.0, + "48": 3977527296.0, + "49": 3977527296.0, + "50": 3977527296.0 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 50, "step_interval": 1, "values": { - "1": "nan", - "2": 6.70023, - "3": 0.32344, - "4": 0.30744, - "5": 0.29991, - "6": 0.30172, - "7": 0.31327, - "8": 0.31174, - "9": 0.30644, - "10": 0.30374, - "11": 0.30977, - "12": 0.304, - "13": 0.30551, - "14": 0.30644, - "15": 0.30406, - "16": 0.30192, - "17": 0.30152, - "18": 0.3017, - "19": 0.30193, - "20": 0.30286, - "21": 0.30159, - "22": 0.30034, - "23": 0.30046, - "24": 0.31202, - "25": 0.30151, - "26": 0.30114, - "27": 0.29956, - "28": 0.3043, - "29": 0.29759, - "30": 0.30726, - "31": 0.29734, - "32": 0.30861, - "33": 0.29846, - "34": 0.2981, - "35": 0.30189, - "36": 0.29782, - "37": 0.30569, - "38": 0.29862, - "39": 0.2985, - "40": 0.29799, - "41": 0.29885, - "42": 0.30133, - "43": 0.29873, - "44": 0.29726, - "45": 0.29716, - "46": 0.29841, - "47": 0.29866, - "48": 0.31148, - "49": 0.29928, - "50": 0.2999 + "2": 5.42077, + "3": 0.32259, + "4": 0.30586, + "5": 0.30517, + "6": 0.29841, + "7": 0.29964, + "8": 0.298, + "9": 0.29723, + "10": 0.29721, + "11": 0.29654, + "12": 0.29707, + "13": 0.29795, + "14": 0.297, + "15": 0.29669, + "16": 0.29822, + "17": 0.29879, + "18": 0.29833, + "19": 0.29859, + "20": 0.29806, + "21": 0.29827, + "22": 0.29773, + "23": 0.29817, + "24": 0.29816, + "25": 0.29894, + "26": 0.30349, + "27": 0.29733, + "28": 0.3063, + "29": 0.30398, + "30": 0.3057, + "31": 0.29826, + "32": 0.30672, + "33": 0.2988, + "34": 0.29646, + "35": 0.30687, + "36": 0.29853, + "37": 0.30571, + "38": 0.30075, + "39": 0.29823, + "40": 0.29822, + "41": 0.30263, + "42": 0.30053, + "43": 0.29871, + "44": 0.31154, + "45": 0.30022, + "46": 0.2988, + "47": 0.29965, + "48": 0.31055, + "49": 0.29988, + "50": 0.29851 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_gb200.json index e73c620b7ce..28edeaf540a 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_gb200.json @@ -6,54 +6,54 @@ "values": { "1": 10.91963, "2": 10.9166, - "3": 10.9199, - "4": 10.92456, - "5": 10.92484, - "6": 10.91858, - "7": 10.91748, - "8": 10.91501, - "9": 10.91647, - "10": 10.91536, - "11": 10.90479, - "12": 10.91195, - "13": 10.89944, - "14": 10.89559, - "15": 10.87847, - "16": 10.86591, - "17": 10.8659, - "18": 10.86335, - "19": 10.86334, - "20": 10.79958, - "21": 10.77882, - "22": 10.7628, - "23": 10.75541, - "24": 10.72183, - "25": 10.72015, - "26": 10.70582, - "27": 10.68006, - "28": 10.61348, - "29": 10.57442, - "30": 10.54775, - "31": 10.54743, - "32": 10.5213, - "33": 10.48753, - "34": 10.46188, - "35": 10.45971, - "36": 10.44025, - "37": 10.40376, - "38": 10.40842, - "39": 10.36701, - "40": 10.35543, - "41": 10.33079, - "42": 10.32796, - "43": 10.30255, - "44": 10.26445, - "45": 10.28208, - "46": 10.24743, - "47": 10.22724, - "48": 10.19664, - "49": 10.18999, - "50": 10.18976 + "3": 10.91966, + "4": 10.92474, + "5": 10.92395, + "6": 10.91844, + "7": 10.91719, + "8": 10.91409, + "9": 10.91707, + "10": 10.91608, + "11": 10.90425, + "12": 10.91127, + "13": 10.89938, + "14": 10.89585, + "15": 10.87884, + "16": 10.86543, + "17": 10.8666, + "18": 10.86348, + "19": 10.86381, + "20": 10.79934, + "21": 10.77907, + "22": 10.76266, + "23": 10.75518, + "24": 10.72151, + "25": 10.72004, + "26": 10.70557, + "27": 10.6795, + "28": 10.61301, + "29": 10.57418, + "30": 10.54671, + "31": 10.54758, + "32": 10.52069, + "33": 10.48632, + "34": 10.4612, + "35": 10.45829, + "36": 10.43895, + "37": 10.40269, + "38": 10.40676, + "39": 10.36513, + "40": 10.35316, + "41": 10.32927, + "42": 10.32571, + "43": 10.30056, + "44": 10.26248, + "45": 10.27971, + "46": 10.24474, + "47": 10.22489, + "48": 10.19366, + "49": 10.18774, + "50": 10.18691 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 16620.0, - "2": 16534.0, - "3": 17006.0, - "4": 16534.0, - "5": 16267.0, - "6": 16166.0, - "7": 16933.0, - "8": 16024.0, - "9": 16624.0, - "10": 16958.0, - "11": 16623.0, - "12": 16120.0, - "13": 17017.0, - "14": 16597.0, - "15": 16421.0, - "16": 16353.0, - "17": 16636.0, - "18": 16637.0, - "19": 16771.0, - "20": 16728.0, - "21": 16427.0, - "22": 16650.0, - "23": 16365.0, - "24": 16946.0, - "25": 16200.0, - "26": 16721.0, - "27": 16943.0, - "28": 16484.0, - "29": 16382.0, - "30": 16691.0, - "31": 17163.0, - "32": 17233.0, - "33": 17217.0, - "34": 17108.0, - "35": 17498.0, - "36": 17007.0, - "37": 17802.0, - "38": 17142.0, - "39": 17994.0, - "40": 18639.0, - "41": 17597.0, - "42": 17438.0, - "43": 18861.0, - "44": 17989.0, - "45": 20471.0, - "46": 19369.0, - "47": 19580.0, - "48": 19495.0, - "49": 22083.0, - "50": 20470.0 + "1": 32041.0, + "2": 32358.0, + "3": 32675.0, + "4": 32160.0, + "5": 31880.0, + "6": 31245.0, + "7": 32888.0, + "8": 31582.0, + "9": 32646.0, + "10": 32514.0, + "11": 32110.0, + "12": 31259.0, + "13": 32590.0, + "14": 32489.0, + "15": 31684.0, + "16": 31660.0, + "17": 32137.0, + "18": 32390.0, + "19": 32324.0, + "20": 32443.0, + "21": 32586.0, + "22": 32841.0, + "23": 32808.0, + "24": 33158.0, + "25": 31556.0, + "26": 33177.0, + "27": 33251.0, + "28": 32766.0, + "29": 32621.0, + "30": 33422.0, + "31": 34127.0, + "32": 33950.0, + "33": 34103.0, + "34": 34742.0, + "35": 34756.0, + "36": 33987.0, + "37": 35253.0, + "38": 34322.0, + "39": 36233.0, + "40": 38974.0, + "41": 36008.0, + "42": 34754.0, + "43": 37855.0, + "44": 39006.0, + "45": 40981.0, + "46": 161246.0, + "47": 2135725.0, + "48": 41251.0, + "49": 2140295.0, + "50": 2137377.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 1561525248.0, - "2": 1562172928.0, - "3": 1562373632.0, - "4": 1561779712.0, - "5": 1561811456.0, - "6": 1561242112.0, - "7": 1561976320.0, - "8": 1561222656.0, - "9": 1563140608.0, - "10": 1561320960.0, - "11": 1561682432.0, - "12": 1561222656.0, - "13": 1561222656.0, - "14": 1562254848.0, - "15": 1561312768.0, - "16": 1561992704.0, - "17": 1562004992.0, - "18": 1561914880.0, - "19": 1562181120.0, - "20": 1561401856.0, - "21": 1561222656.0, - "22": 1562348032.0, - "23": 1561468416.0, - "24": 1561222656.0, - "25": 1561222656.0, - "26": 1561222656.0, - "27": 1561222656.0, - "28": 1561951744.0, - "29": 1561513472.0, - "30": 1561222656.0, - "31": 1561559552.0, - "32": 1561222656.0, - "33": 1561524736.0, - "34": 1561222656.0, - "35": 1561222656.0, - "36": 1561222656.0, - "37": 1561669120.0, - "38": 1561282048.0, - "39": 1561837056.0, - "40": 1562086912.0, - "41": 1561222656.0, - "42": 1562263040.0, - "43": 1561263616.0, - "44": 1561871872.0, - "45": 1561865728.0, - "46": 1561222656.0, - "47": 1561570816.0, - "48": 1562082816.0, - "49": 1561452032.0, - "50": 1562226176.0 + "1": 1560094720.0, + "2": 1559330304.0, + "3": 1559751168.0, + "4": 1559330304.0, + "5": 1559330304.0, + "6": 1559783936.0, + "7": 1559846400.0, + "8": 1559330304.0, + "9": 1559837184.0, + "10": 1559448064.0, + "11": 1560685056.0, + "12": 1559608832.0, + "13": 1559330304.0, + "14": 1559330304.0, + "15": 1559330304.0, + "16": 1559330304.0, + "17": 1559330304.0, + "18": 1560493568.0, + "19": 1559330304.0, + "20": 1559745024.0, + "21": 1559783936.0, + "22": 1559330304.0, + "23": 1559330304.0, + "24": 1559330304.0, + "25": 1559330304.0, + "26": 1560300032.0, + "27": 1559330304.0, + "28": 1560271360.0, + "29": 1559695872.0, + "30": 1559330304.0, + "31": 1559629312.0, + "32": 1559330304.0, + "33": 1560422912.0, + "34": 1559456256.0, + "35": 1560415744.0, + "36": 1559330304.0, + "37": 1560976896.0, + "38": 1560811008.0, + "39": 1559330304.0, + "40": 1559649792.0, + "41": 1559330304.0, + "42": 1560197632.0, + "43": 1559469568.0, + "44": 1560458752.0, + "45": 1559330304.0, + "46": 1559330304.0, + "47": 1560821248.0, + "48": 1561056768.0, + "49": 1560043008.0, + "50": 1559330304.0 } }, "mem-max-allocated-bytes": { @@ -175,113 +175,112 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 3478863360.0, - "2": 4057549824.0, - "3": 4057549824.0, - "4": 4057549824.0, - "5": 4064312320.0, - "6": 4064312320.0, - "7": 4064312320.0, - "8": 4064312320.0, - "9": 4064312320.0, - "10": 4064312320.0, - "11": 4064312320.0, - "12": 4064312320.0, - "13": 4064312320.0, - "14": 4064312320.0, - "15": 4064312320.0, - "16": 4064312320.0, - "17": 4064312320.0, - "18": 4064312320.0, - "19": 4064312320.0, - "20": 4064312320.0, - "21": 4064312320.0, - "22": 4064312320.0, - "23": 4064312320.0, - "24": 4064312320.0, - "25": 4064312320.0, - "26": 4064312320.0, - "27": 4064312320.0, - "28": 4064312320.0, - "29": 4064312320.0, - "30": 4064312320.0, - "31": 4064312320.0, - "32": 4064312320.0, - "33": 4064312320.0, - "34": 4064312320.0, - "35": 4064312320.0, - "36": 4064312320.0, - "37": 4064312320.0, - "38": 4064312320.0, - "39": 4064312320.0, - "40": 4064312320.0, - "41": 4064312320.0, - "42": 4064312320.0, - "43": 4064312320.0, - "44": 4064312320.0, - "45": 4064312320.0, - "46": 4064312320.0, - "47": 4064312320.0, - "48": 4064312320.0, - "49": 4064312320.0, - "50": 4064312320.0 + "1": 3478912512.0, + "2": 4055036928.0, + "3": 4055036928.0, + "4": 4055036928.0, + "5": 4060747776.0, + "6": 4060747776.0, + "7": 4060747776.0, + "8": 4060747776.0, + "9": 4060747776.0, + "10": 4060747776.0, + "11": 4060747776.0, + "12": 4060747776.0, + "13": 4060747776.0, + "14": 4060747776.0, + "15": 4060747776.0, + "16": 4060747776.0, + "17": 4060747776.0, + "18": 4060747776.0, + "19": 4060747776.0, + "20": 4060747776.0, + "21": 4060747776.0, + "22": 4060747776.0, + "23": 4060747776.0, + "24": 4060747776.0, + "25": 4060747776.0, + "26": 4060747776.0, + "27": 4060747776.0, + "28": 4060747776.0, + "29": 4060747776.0, + "30": 4060747776.0, + "31": 4060747776.0, + "32": 4060747776.0, + "33": 4060747776.0, + "34": 4060747776.0, + "35": 4060747776.0, + "36": 4060747776.0, + "37": 4060747776.0, + "38": 4060747776.0, + "39": 4060747776.0, + "40": 4060747776.0, + "41": 4060747776.0, + "42": 4060747776.0, + "43": 4060747776.0, + "44": 4060747776.0, + "45": 4060747776.0, + "46": 4060747776.0, + "47": 4060747776.0, + "48": 4060747776.0, + "49": 4060747776.0, + "50": 4060747776.0 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 50, "step_interval": 1, "values": { - "1": "nan", - "2": 7.28192, - "3": 0.48787, - "4": 0.45201, - "5": 0.43575, - "6": 0.44125, - "7": 0.45147, - "8": 0.43509, - "9": 0.43968, - "10": 0.43228, - "11": 0.44373, - "12": 0.44276, - "13": 0.44406, - "14": 0.43766, - "15": 0.43286, - "16": 0.44948, - "17": 0.44794, - "18": 0.45565, - "19": 0.45045, - "20": 0.4531, - "21": 0.46157, - "22": 0.45742, - "23": 0.45152, - "24": 0.45695, - "25": 0.46644, - "26": 0.4541, - "27": 0.4502, - "28": 0.4593, - "29": 0.45941, - "30": 0.45501, - "31": 0.46525, - "32": 0.48077, - "33": 0.43971, - "34": 0.44932, - "35": 0.44807, - "36": 0.43603, - "37": 0.4341, - "38": 0.43907, - "39": 0.43214, - "40": 0.44553, - "41": 0.43817, - "42": 0.42521, - "43": 0.42868, - "44": 0.42273, - "45": 0.4405, - "46": 0.46649, - "47": 0.44854, - "48": 0.44657, - "49": 0.42855, - "50": 0.44345 + "2": 8.76848, + "3": 0.51847, + "4": 0.44451, + "5": 0.42561, + "6": 0.41192, + "7": 0.40268, + "8": 0.40629, + "9": 0.39737, + "10": 0.38313, + "11": 0.39238, + "12": 0.84424, + "13": 0.36903, + "14": 0.35807, + "15": 0.80727, + "16": 0.34037, + "17": 0.89468, + "18": 0.3416, + "19": 0.33901, + "20": 0.35218, + "21": 0.35179, + "22": 0.79253, + "23": 0.33727, + "24": 0.3354, + "25": 0.34272, + "26": 0.33425, + "27": 0.84289, + "28": 0.34575, + "29": 0.34756, + "30": 0.35201, + "31": 0.34147, + "32": 0.34441, + "33": 0.35069, + "34": 0.35702, + "35": 0.36375, + "36": 0.35846, + "37": 0.35239, + "38": 0.35852, + "39": 0.35471, + "40": 0.35162, + "41": 0.34809, + "42": 0.3433, + "43": 0.35044, + "44": 0.33572, + "45": 0.34394, + "46": 0.34517, + "47": 0.35124, + "48": 0.3525, + "49": 0.34863, + "50": 0.34969 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_gb200.json index 9c55b67895a..b95c36e24d6 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_gb200.json @@ -7,103 +7,103 @@ "1": 10.92137, "2": 10.91202, "3": 10.91768, - "4": 10.91863, - "5": 10.90423, - "6": 10.9067, - "7": 10.91437, - "8": 10.91056, - "9": 10.91447, - "10": 10.90704, - "11": 10.89659, - "12": 10.8959, - "13": 10.90064, - "14": 10.89352, - "15": 10.87253, - "16": 10.86144, - "17": 10.87413, - "18": 10.85919, - "19": 10.86552, - "20": 10.78225, - "21": 10.77925, - "22": 10.76543, - "23": 10.75797, - "24": 10.72021, - "25": 10.72236, - "26": 10.71659, - "27": 10.68375, - "28": 10.62333, - "29": 10.58113, - "30": 10.56502, - "31": 10.55937, - "32": 10.5483, - "33": 10.51334, - "34": 10.487, - "35": 10.47893, - "36": 10.45972, - "37": 10.43079, - "38": 10.4355, - "39": 10.40227, - "40": 10.37985, - "41": 10.36324, - "42": 10.33753, - "43": 10.32309, - "44": 10.28593, - "45": 10.31097, - "46": 10.26822, - "47": 10.25636, - "48": 10.2077, - "49": 10.2086, - "50": 10.20888, - "51": 10.20917, - "52": 10.16628, - "53": 10.17255, - "54": 10.13289, - "55": 10.10853, - "56": 10.13283, - "57": 10.12157, - "58": 10.13408, - "59": 10.08188, - "60": 10.10019, - "61": 10.04681, - "62": 10.01985, - "63": 10.09447, - "64": 10.04943, - "65": 10.02338, - "66": 10.04778, - "67": 10.02637, - "68": 9.98521, - "69": 10.00765, - "70": 9.98465, - "71": 10.01262, - "72": 9.99392, - "73": 9.9866, - "74": 9.96469, - "75": 9.93791, - "76": 9.97185, - "77": 9.96314, - "78": 9.91583, - "79": 9.92237, - "80": 9.93257, - "81": 9.96107, - "82": 9.8946, - "83": 9.86433, - "84": 9.79835, - "85": 9.78983, - "86": 9.89001, - "87": 9.90939, - "88": 9.88672, - "89": 9.82926, - "90": 9.82374, - "91": 9.83424, - "92": 9.82048, - "93": 9.75668, - "94": 9.83611, - "95": 9.82852, - "96": 9.80839, - "97": 9.75439, - "98": 9.78617, - "99": 9.82596, - "100": 9.72104 + "4": 10.91867, + "5": 10.90429, + "6": 10.9062, + "7": 10.914, + "8": 10.91006, + "9": 10.91461, + "10": 10.90711, + "11": 10.89655, + "12": 10.89605, + "13": 10.90084, + "14": 10.89353, + "15": 10.87232, + "16": 10.86136, + "17": 10.87416, + "18": 10.85925, + "19": 10.86561, + "20": 10.78245, + "21": 10.7792, + "22": 10.7655, + "23": 10.75784, + "24": 10.72029, + "25": 10.72272, + "26": 10.71634, + "27": 10.68397, + "28": 10.62345, + "29": 10.58126, + "30": 10.56533, + "31": 10.5593, + "32": 10.54822, + "33": 10.51337, + "34": 10.48685, + "35": 10.47871, + "36": 10.45967, + "37": 10.43094, + "38": 10.43582, + "39": 10.40263, + "40": 10.37971, + "41": 10.36341, + "42": 10.33811, + "43": 10.3231, + "44": 10.28605, + "45": 10.31118, + "46": 10.26833, + "47": 10.25651, + "48": 10.20784, + "49": 10.20864, + "50": 10.20894, + "51": 10.20907, + "52": 10.1664, + "53": 10.17258, + "54": 10.1329, + "55": 10.10858, + "56": 10.13295, + "57": 10.12151, + "58": 10.13389, + "59": 10.08174, + "60": 10.10018, + "61": 10.04694, + "62": 10.01976, + "63": 10.09419, + "64": 10.04926, + "65": 10.02317, + "66": 10.04748, + "67": 10.02609, + "68": 9.98481, + "69": 10.00719, + "70": 9.98435, + "71": 10.01225, + "72": 9.99337, + "73": 9.98595, + "74": 9.96408, + "75": 9.93716, + "76": 9.97108, + "77": 9.96223, + "78": 9.91482, + "79": 9.92103, + "80": 9.93141, + "81": 9.95995, + "82": 9.89348, + "83": 9.86307, + "84": 9.79718, + "85": 9.78856, + "86": 9.88919, + "87": 9.90846, + "88": 9.88568, + "89": 9.82832, + "90": 9.82266, + "91": 9.83293, + "92": 9.81925, + "93": 9.75539, + "94": 9.83457, + "95": 9.82715, + "96": 9.80685, + "97": 9.75264, + "98": 9.78463, + "99": 9.82436, + "100": 9.71945 } }, "num-zeros": { @@ -111,106 +111,106 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 2666.0, - "2": 2644.0, - "3": 2751.0, - "4": 2603.0, - "5": 2552.0, - "6": 2425.0, - "7": 2606.0, - "8": 2505.0, - "9": 2490.0, - "10": 2506.0, - "11": 2668.0, - "12": 2586.0, - "13": 2519.0, - "14": 2483.0, - "15": 2434.0, - "16": 2528.0, - "17": 2682.0, - "18": 2604.0, - "19": 2510.0, - "20": 2434.0, - "21": 2425.0, - "22": 2551.0, - "23": 2511.0, - "24": 2495.0, - "25": 2503.0, - "26": 2631.0, - "27": 2652.0, - "28": 2669.0, - "29": 2733.0, - "30": 2613.0, - "31": 2775.0, - "32": 2808.0, - "33": 2750.0, - "34": 2788.0, - "35": 2882.0, - "36": 2782.0, - "37": 3001.0, - "38": 2950.0, - "39": 2915.0, - "40": 3165.0, - "41": 3101.0, - "42": 2987.0, - "43": 3038.0, - "44": 3017.0, - "45": 3336.0, - "46": 3207.0, - "47": 3368.0, - "48": 3472.0, - "49": 3487.0, - "50": 3453.0, - "51": 3438.0, - "52": 3540.0, - "53": 3438.0, - "54": 7447.0, - "55": 3892.0, - "56": 4010.0, - "57": 5332.0, - "58": 529109.0, - "59": 6516.0, - "60": 3780.0, - "61": 4406.0, - "62": 4225.0, - "63": 5252.0, - "64": 4844.0, - "65": 9249.0, - "66": 5773.0, - "67": 4636.0, - "68": 37715.0, - "69": 7172.0, - "70": 528692.0, - "71": 4227.0, - "72": 26929.0, - "73": 7106.0, - "74": 36832.0, - "75": 28664.0, - "76": 556371.0, - "77": 557196.0, - "78": 36991.0, - "79": 529101.0, - "80": 554041.0, - "81": 529463.0, - "82": 529423.0, - "83": 529040.0, - "84": 528140.0, - "85": 528770.0, - "86": 529266.0, - "87": 583382.0, - "88": 529195.0, - "89": 529009.0, - "90": 529255.0, - "91": 528899.0, - "92": 528696.0, - "93": 529174.0, - "94": 528741.0, - "95": 528662.0, - "96": 531201.0, - "97": 529233.0, - "98": 529714.0, - "99": 40294.0, - "100": 560372.0 + "1": 6670.0, + "2": 6580.0, + "3": 6653.0, + "4": 6526.0, + "5": 6462.0, + "6": 6209.0, + "7": 6831.0, + "8": 6551.0, + "9": 6643.0, + "10": 6568.0, + "11": 6679.0, + "12": 6265.0, + "13": 6608.0, + "14": 6484.0, + "15": 6336.0, + "16": 6452.0, + "17": 6592.0, + "18": 6571.0, + "19": 6523.0, + "20": 6191.0, + "21": 6672.0, + "22": 6529.0, + "23": 6372.0, + "24": 6646.0, + "25": 6532.0, + "26": 6596.0, + "27": 6675.0, + "28": 6831.0, + "29": 7222.0, + "30": 6842.0, + "31": 7049.0, + "32": 7290.0, + "33": 7229.0, + "34": 7344.0, + "35": 7383.0, + "36": 7250.0, + "37": 7828.0, + "38": 7533.0, + "39": 7790.0, + "40": 7960.0, + "41": 8186.0, + "42": 7584.0, + "43": 8105.0, + "44": 7859.0, + "45": 8684.0, + "46": 8605.0, + "47": 8515.0, + "48": 8762.0, + "49": 9131.0, + "50": 8830.0, + "51": 10301.0, + "52": 9384.0, + "53": 8941.0, + "54": 17166.0, + "55": 9783.0, + "56": 10246.0, + "57": 12994.0, + "58": 1060368.0, + "59": 15622.0, + "60": 59852.0, + "61": 11305.0, + "62": 18030.0, + "63": 1060940.0, + "64": 1059913.0, + "65": 1063020.0, + "66": 17265.0, + "67": 58014.0, + "68": 137665.0, + "69": 1064689.0, + "70": 1108015.0, + "71": 19901.0, + "72": 1103930.0, + "73": 1063927.0, + "74": 1170080.0, + "75": 1105577.0, + "76": 3156838.0, + "77": 2163555.0, + "78": 1122743.0, + "79": 2107511.0, + "80": 2157105.0, + "81": 2108230.0, + "82": 2108359.0, + "83": 2106919.0, + "84": 1104784.0, + "85": 2106376.0, + "86": 2107704.0, + "87": 2215262.0, + "88": 2107733.0, + "89": 2107132.0, + "90": 1109659.0, + "91": 2106948.0, + "92": 1109424.0, + "93": 2106782.0, + "94": 2106957.0, + "95": 1111714.0, + "96": 1122793.0, + "97": 2107258.0, + "98": 2107901.0, + "99": 1125984.0, + "100": 2113039.0 } }, "mem-allocated-bytes": { @@ -330,208 +330,207 @@ "3": 1152957440.0, "4": 1152957440.0, "5": 1152957440.0, - "6": 1153038848.0, - "7": 1153038848.0, - "8": 1153038848.0, - "9": 1153038848.0, - "10": 1153038848.0, - "11": 1153038848.0, - "12": 1153038848.0, - "13": 1153038848.0, - "14": 1153038848.0, - "15": 1153038848.0, - "16": 1153038848.0, - "17": 1153038848.0, - "18": 1154730496.0, - "19": 1154730496.0, - "20": 1154730496.0, - "21": 1154730496.0, - "22": 1154730496.0, - "23": 1154730496.0, - "24": 1154730496.0, - "25": 1154730496.0, - "26": 1154730496.0, - "27": 1154730496.0, - "28": 1154730496.0, - "29": 1154730496.0, - "30": 1154730496.0, - "31": 1154730496.0, - "32": 1154730496.0, - "33": 1154730496.0, - "34": 1154730496.0, - "35": 1154730496.0, - "36": 1154730496.0, - "37": 1155892736.0, - "38": 1155892736.0, - "39": 1155892736.0, - "40": 1155892736.0, - "41": 1155892736.0, - "42": 1155892736.0, - "43": 1155892736.0, - "44": 1155892736.0, - "45": 1155892736.0, - "46": 1155892736.0, - "47": 1155892736.0, - "48": 1155892736.0, - "49": 1155892736.0, - "50": 1155892736.0, - "51": 1155892736.0, - "52": 1155892736.0, - "53": 1155892736.0, - "54": 1155892736.0, - "55": 1155892736.0, - "56": 1155892736.0, - "57": 1155892736.0, - "58": 1155892736.0, - "59": 1155892736.0, - "60": 1155892736.0, - "61": 1155892736.0, - "62": 1155892736.0, - "63": 1155892736.0, - "64": 1155892736.0, - "65": 1155892736.0, - "66": 1155892736.0, - "67": 1155892736.0, - "68": 1155892736.0, - "69": 1155892736.0, - "70": 1155892736.0, - "71": 1155892736.0, - "72": 1155892736.0, - "73": 1155892736.0, - "74": 1155892736.0, - "75": 1155892736.0, - "76": 1155892736.0, - "77": 1155892736.0, - "78": 1155892736.0, - "79": 1155892736.0, - "80": 1155892736.0, - "81": 1155892736.0, - "82": 1155892736.0, - "83": 1155892736.0, - "84": 1155892736.0, - "85": 1155892736.0, - "86": 1155892736.0, - "87": 1155892736.0, - "88": 1155892736.0, - "89": 1155892736.0, - "90": 1155892736.0, - "91": 1155892736.0, - "92": 1155892736.0, - "93": 1155892736.0, - "94": 1155892736.0, - "95": 1155892736.0, - "96": 1155892736.0, - "97": 1155892736.0, - "98": 1155892736.0, - "99": 1155892736.0, - "100": 1155892736.0 + "6": 1153049088.0, + "7": 1153049088.0, + "8": 1153049088.0, + "9": 1153049088.0, + "10": 1153049088.0, + "11": 1153049088.0, + "12": 1153049088.0, + "13": 1153049088.0, + "14": 1153049088.0, + "15": 1153049088.0, + "16": 1153049088.0, + "17": 1153049088.0, + "18": 1153221120.0, + "19": 1153221120.0, + "20": 1153221120.0, + "21": 1153221120.0, + "22": 1153221120.0, + "23": 1153649664.0, + "24": 1153649664.0, + "25": 1153649664.0, + "26": 1153649664.0, + "27": 1153682432.0, + "28": 1153682432.0, + "29": 1153682432.0, + "30": 1153682432.0, + "31": 1153682432.0, + "32": 1153682432.0, + "33": 1153682432.0, + "34": 1154268672.0, + "35": 1154268672.0, + "36": 1154268672.0, + "37": 1155870208.0, + "38": 1155870208.0, + "39": 1155870208.0, + "40": 1155870208.0, + "41": 1155870208.0, + "42": 1155870208.0, + "43": 1155870208.0, + "44": 1155870208.0, + "45": 1155870208.0, + "46": 1155870208.0, + "47": 1155870208.0, + "48": 1155870208.0, + "49": 1155870208.0, + "50": 1155870208.0, + "51": 1155870208.0, + "52": 1155870208.0, + "53": 1155870208.0, + "54": 1155870208.0, + "55": 1155870208.0, + "56": 1155870208.0, + "57": 1155870208.0, + "58": 1155870208.0, + "59": 1155870208.0, + "60": 1155870208.0, + "61": 1155870208.0, + "62": 1155870208.0, + "63": 1155870208.0, + "64": 1155870208.0, + "65": 1155870208.0, + "66": 1155870208.0, + "67": 1155870208.0, + "68": 1155870208.0, + "69": 1155870208.0, + "70": 1155870208.0, + "71": 1155870208.0, + "72": 1155870208.0, + "73": 1155870208.0, + "74": 1155870208.0, + "75": 1155870208.0, + "76": 1155870208.0, + "77": 1155870208.0, + "78": 1155870208.0, + "79": 1155870208.0, + "80": 1155870208.0, + "81": 1155870208.0, + "82": 1155870208.0, + "83": 1155870208.0, + "84": 1155870208.0, + "85": 1155870208.0, + "86": 1155870208.0, + "87": 1155870208.0, + "88": 1155870208.0, + "89": 1155870208.0, + "90": 1155870208.0, + "91": 1155870208.0, + "92": 1155870208.0, + "93": 1155870208.0, + "94": 1155870208.0, + "95": 1155870208.0, + "96": 1155870208.0, + "97": 1155870208.0, + "98": 1155870208.0, + "99": 1155870208.0, + "100": 1155870208.0 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 100, "step_interval": 1, "values": { - "1": "nan", - "2": 11.10133, - "3": 1.00423, - "4": 0.87576, - "5": 0.85174, - "6": 0.89147, - "7": 0.85652, - "8": 0.81429, - "9": 0.78782, - "10": 0.77925, - "11": 0.77442, - "12": 0.76661, - "13": 0.75464, - "14": 0.73604, - "15": 0.75554, - "16": 0.73924, - "17": 0.72982, - "18": 0.73443, - "19": 0.80729, - "20": 0.86288, - "21": 0.8837, - "22": 0.87191, - "23": 0.86056, - "24": 0.83907, - "25": 0.83018, - "26": 0.82874, - "27": 0.81482, - "28": 0.80051, - "29": 0.80103, - "30": 0.79464, - "31": 0.78839, - "32": 0.78643, - "33": 0.7887, - "34": 0.7918, - "35": 0.78506, - "36": 0.77353, - "37": 0.77466, - "38": 0.84635, - "39": 0.86024, - "40": 0.8233, - "41": 0.78652, - "42": 0.76746, - "43": 0.75677, - "44": 0.76688, - "45": 0.75709, - "46": 0.7512, - "47": 0.76635, - "48": 0.74707, - "49": 0.75704, - "50": 0.74817, - "51": 1.53778, - "52": 1.40548, - "53": 0.8928, - "54": 0.89804, - "55": 0.89234, - "56": 0.88819, - "57": 0.87448, - "58": 0.84218, - "59": 0.81327, - "60": 0.80856, - "61": 0.78612, - "62": 0.77474, - "63": 0.75246, - "64": 0.74889, - "65": 0.76004, - "66": 1.23226, - "67": 0.76409, - "68": 0.76577, - "69": 1.18807, - "70": 0.74565, - "71": 0.73788, - "72": 0.74486, - "73": 0.83307, - "74": 0.85133, - "75": 0.83092, - "76": 0.78755, - "77": 0.76405, - "78": 0.75659, - "79": 0.74733, - "80": 0.74014, - "81": 0.75757, - "82": 0.73935, - "83": 0.74015, - "84": 0.74067, - "85": 0.73626, - "86": 0.74291, - "87": 0.7382, - "88": 0.73417, - "89": 0.73046, - "90": 0.73499, - "91": 0.73824, - "92": 0.73677, - "93": 0.74129, - "94": 0.74578, - "95": 0.737, - "96": 1.17233, - "97": 0.74753, - "98": 0.74302, - "99": 0.73302, - "100": 0.734 + "2": 11.62031, + "3": 0.82946, + "4": 0.75366, + "5": 0.74335, + "6": 0.75463, + "7": 0.7365, + "8": 0.7246, + "9": 0.73593, + "10": 0.7335, + "11": 0.71203, + "12": 0.72555, + "13": 0.74312, + "14": 1.13535, + "15": 0.71681, + "16": 0.71209, + "17": 0.71788, + "18": 0.70937, + "19": 0.70144, + "20": 1.16809, + "21": 0.70724, + "22": 0.70409, + "23": 0.71193, + "24": 1.1743, + "25": 0.72109, + "26": 0.71703, + "27": 0.70948, + "28": 0.72083, + "29": 0.73183, + "30": 0.72532, + "31": 0.72889, + "32": 0.75893, + "33": 0.7539, + "34": 1.23456, + "35": 1.21961, + "36": 0.77501, + "37": 0.76934, + "38": 0.76401, + "39": 0.7514, + "40": 0.77279, + "41": 0.77099, + "42": 0.7833, + "43": 0.77533, + "44": 1.17538, + "45": 0.78388, + "46": 0.78653, + "47": 0.77219, + "48": 0.75972, + "49": 0.78431, + "50": 0.77345, + "51": 1.11606, + "52": 0.84396, + "53": 0.76433, + "54": 0.76189, + "55": 0.76688, + "56": 0.7712, + "57": 0.76937, + "58": 0.77147, + "59": 0.78428, + "60": 0.76924, + "61": 0.77311, + "62": 0.7757, + "63": 0.85294, + "64": 0.76871, + "65": 0.76006, + "66": 0.77454, + "67": 0.78211, + "68": 0.75672, + "69": 0.75847, + "70": 0.75526, + "71": 0.74527, + "72": 0.75072, + "73": 0.75816, + "74": 0.75776, + "75": 0.7608, + "76": 0.75893, + "77": 0.7581, + "78": 0.76254, + "79": 0.7611, + "80": 0.77079, + "81": 0.77364, + "82": 0.77324, + "83": 0.77199, + "84": 0.87537, + "85": 0.94416, + "86": 0.89911, + "87": 0.86544, + "88": 0.84094, + "89": 0.79274, + "90": 0.76926, + "91": 0.77717, + "92": 0.7729, + "93": 0.76403, + "94": 0.78437, + "95": 0.76069, + "96": 0.76582, + "97": 0.77904, + "98": 0.77605, + "99": 0.78215, + "100": 0.78286 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json index 93c18afda0e..356655a2f72 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json @@ -7,103 +7,103 @@ "1": 10.98117, "2": 10.9645, "3": 10.96986, - "4": 10.96038, - "5": 10.9701, - "6": 10.96815, - "7": 10.96507, - "8": 10.97054, - "9": 10.96141, - "10": 10.96424, - "11": 10.95664, - "12": 10.94737, - "13": 10.94819, - "14": 10.95025, - "15": 10.91641, - "16": 10.90025, - "17": 10.89894, - "18": 10.90045, - "19": 10.88544, - "20": 10.8187, - "21": 10.77645, - "22": 10.77489, - "23": 10.76234, + "4": 10.96007, + "5": 10.96998, + "6": 10.96837, + "7": 10.96536, + "8": 10.97077, + "9": 10.96083, + "10": 10.96459, + "11": 10.95716, + "12": 10.94789, + "13": 10.94736, + "14": 10.95035, + "15": 10.91734, + "16": 10.90035, + "17": 10.89933, + "18": 10.90087, + "19": 10.88555, + "20": 10.81877, + "21": 10.77584, + "22": 10.77439, + "23": 10.76215, "24": 10.74094, - "25": 10.71574, - "26": 10.70618, - "27": 10.66985, - "28": 10.60725, - "29": 10.57586, - "30": 10.54319, - "31": 10.53935, - "32": 10.48648, - "33": 10.45188, - "34": 10.44968, - "35": 10.41723, - "36": 10.40441, - "37": 10.38008, - "38": 10.3728, - "39": 10.33083, - "40": 10.3386, - "41": 10.29305, - "42": 10.24315, - "43": 10.23829, - "44": 10.19158, - "45": 10.22758, - "46": 10.16368, - "47": 10.1506, - "48": 10.11899, - "49": 10.11135, - "50": 10.09763, - "51": 10.11138, - "52": 10.05432, - "53": 10.03179, - "54": 10.01418, - "55": 9.96809, - "56": 10.01525, - "57": 9.99335, - "58": 9.99662, - "59": 9.93395, - "60": 9.96146, - "61": 9.91376, - "62": 9.85848, - "63": 9.96287, - "64": 9.90765, - "65": 9.87405, - "66": 9.89551, - "67": 9.87298, - "68": 9.79949, - "69": 9.84522, - "70": 9.82568, - "71": 9.84955, - "72": 9.83187, - "73": 9.79639, - "74": 9.7821, - "75": 9.72976, - "76": 9.79558, - "77": 9.79604, - "78": 9.75162, - "79": 9.73342, - "80": 9.75354, - "81": 9.78965, - "82": 9.71909, - "83": 9.65341, - "84": 9.61892, - "85": 9.57955, - "86": 9.72668, - "87": 9.73079, - "88": 9.73212, - "89": 9.6346, - "90": 9.62591, - "91": 9.66704, - "92": 9.63591, - "93": 9.53928, - "94": 9.64991, - "95": 9.6367, - "96": 9.62983, - "97": 9.54251, - "98": 9.58549, - "99": 9.63024, - "100": 9.52511 + "25": 10.71575, + "26": 10.70576, + "27": 10.67012, + "28": 10.60658, + "29": 10.57561, + "30": 10.54346, + "31": 10.53971, + "32": 10.48642, + "33": 10.45191, + "34": 10.44938, + "35": 10.41762, + "36": 10.40467, + "37": 10.38053, + "38": 10.37325, + "39": 10.33101, + "40": 10.33885, + "41": 10.29378, + "42": 10.24359, + "43": 10.23841, + "44": 10.19204, + "45": 10.228, + "46": 10.16383, + "47": 10.15072, + "48": 10.11961, + "49": 10.1119, + "50": 10.09834, + "51": 10.11172, + "52": 10.05509, + "53": 10.03212, + "54": 10.01448, + "55": 9.96858, + "56": 10.01565, + "57": 9.99419, + "58": 9.99714, + "59": 9.93475, + "60": 9.96191, + "61": 9.91382, + "62": 9.85877, + "63": 9.96335, + "64": 9.908, + "65": 9.87403, + "66": 9.8959, + "67": 9.87286, + "68": 9.79981, + "69": 9.84586, + "70": 9.82596, + "71": 9.84984, + "72": 9.8318, + "73": 9.79652, + "74": 9.78204, + "75": 9.72998, + "76": 9.79584, + "77": 9.79644, + "78": 9.75151, + "79": 9.73337, + "80": 9.75376, + "81": 9.78971, + "82": 9.71907, + "83": 9.65377, + "84": 9.61876, + "85": 9.57915, + "86": 9.72619, + "87": 9.73085, + "88": 9.73229, + "89": 9.63454, + "90": 9.62564, + "91": 9.66692, + "92": 9.63586, + "93": 9.53873, + "94": 9.64989, + "95": 9.63616, + "96": 9.62938, + "97": 9.54237, + "98": 9.58505, + "99": 9.62999, + "100": 9.52477 } }, "num-zeros": { @@ -111,106 +111,106 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 22985550.0, - "2": 22866696.0, - "3": 22718816.0, - "4": 22793004.0, - "5": 22800288.0, - "6": 22758748.0, - "7": 22889372.0, - "8": 22616932.0, - "9": 22770608.0, - "10": 22482312.0, - "11": 22767860.0, - "12": 22646624.0, - "13": 23376188.0, - "14": 23020892.0, - "15": 22728500.0, - "16": 22844224.0, - "17": 22956374.0, - "18": 23025448.0, - "19": 23121808.0, - "20": 22737748.0, - "21": 22939010.0, - "22": 22975312.0, - "23": 22636400.0, - "24": 22885738.0, - "25": 22646644.0, - "26": 23036360.0, - "27": 22820198.0, - "28": 23031728.0, - "29": 23007826.0, - "30": 22977996.0, - "31": 22931432.0, - "32": 22671932.0, - "33": 22753836.0, - "34": 23115324.0, - "35": 22764116.0, - "36": 22708324.0, - "37": 23140486.0, - "38": 22991004.0, - "39": 23017608.0, - "40": 22766856.0, - "41": 23101476.0, - "42": 22699692.0, - "43": 23019168.0, - "44": 22716684.0, - "45": 22868822.0, - "46": 22742870.0, - "47": 22872520.0, - "48": 22852580.0, - "49": 22908216.0, - "50": 22654296.0, - "51": 22713688.0, - "52": 22833222.0, - "53": 22987634.0, - "54": 22807092.0, - "55": 22950776.0, - "56": 22669972.0, - "57": 23234460.0, - "58": 22699660.0, - "59": 22861762.0, - "60": 23046726.0, - "61": 22688404.0, - "62": 22743736.0, - "63": 22644196.0, - "64": 23031364.0, - "65": 23243098.0, - "66": 22705374.0, - "67": 22986432.0, - "68": 22949096.0, - "69": 23193052.0, - "70": 22837850.0, - "71": 22749418.0, - "72": 23154708.0, - "73": 23168010.0, - "74": 22969342.0, - "75": 22902846.0, - "76": 22713620.0, - "77": 23010912.0, - "78": 23009874.0, - "79": 22845022.0, - "80": 22957728.0, - "81": 22849644.0, - "82": 22745748.0, - "83": 22741216.0, - "84": 23135212.0, - "85": 22945266.0, - "86": 23107608.0, - "87": 22369124.0, - "88": 22564568.0, - "89": 22737612.0, - "90": 22782104.0, - "91": 22940644.0, - "92": 22680084.0, - "93": 22647824.0, - "94": 23168470.0, - "95": 22701628.0, - "96": 22866924.0, - "97": 22852120.0, - "98": 22897152.0, - "99": 22645312.0, - "100": 23029176.0 + "1": 22986500.0, + "2": 22867754.0, + "3": 22719740.0, + "4": 22793872.0, + "5": 22801148.0, + "6": 22759688.0, + "7": 22890320.0, + "8": 22617806.0, + "9": 22771422.0, + "10": 22483220.0, + "11": 22768770.0, + "12": 22647588.0, + "13": 23377208.0, + "14": 23021848.0, + "15": 22729496.0, + "16": 22845148.0, + "17": 22957352.0, + "18": 23026256.0, + "19": 23122762.0, + "20": 22738628.0, + "21": 22939938.0, + "22": 22976138.0, + "23": 22637284.0, + "24": 22886688.0, + "25": 22647576.0, + "26": 23037160.0, + "27": 22821264.0, + "28": 23032636.0, + "29": 23008834.0, + "30": 22978970.0, + "31": 22932444.0, + "32": 22672836.0, + "33": 22754776.0, + "34": 23116420.0, + "35": 22765096.0, + "36": 22709354.0, + "37": 23141482.0, + "38": 22992028.0, + "39": 23018736.0, + "40": 22768956.0, + "41": 23103932.0, + "42": 22700624.0, + "43": 23023730.0, + "44": 23766364.0, + "45": 22871416.0, + "46": 22743956.0, + "47": 23922412.0, + "48": 22855166.0, + "49": 23957920.0, + "50": 23704088.0, + "51": 23763576.0, + "52": 22833932.0, + "53": 24037402.0, + "54": 23857020.0, + "55": 24001264.0, + "56": 23719884.0, + "57": 24286134.0, + "58": 23750276.0, + "59": 22867408.0, + "60": 24097470.0, + "61": 23739132.0, + "62": 23794220.0, + "63": 24743926.0, + "64": 23037320.0, + "65": 23248866.0, + "66": 23756504.0, + "67": 22988692.0, + "68": 22957062.0, + "69": 23201288.0, + "70": 22842814.0, + "71": 22753444.0, + "72": 23158108.0, + "73": 23170850.0, + "74": 22975384.0, + "75": 22905316.0, + "76": 22717736.0, + "77": 23012708.0, + "78": 23012448.0, + "79": 22847056.0, + "80": 22959264.0, + "81": 22851702.0, + "82": 22747688.0, + "83": 22744370.0, + "84": 23136788.0, + "85": 22947836.0, + "86": 23109100.0, + "87": 23419368.0, + "88": 22565890.0, + "89": 22739148.0, + "90": 23831870.0, + "91": 22941742.0, + "92": 22681372.0, + "93": 22648836.0, + "94": 23169778.0, + "95": 22702866.0, + "96": 22868248.0, + "97": 22853186.0, + "98": 22898314.0, + "99": 22646596.0, + "100": 23030464.0 } }, "mem-allocated-bytes": { @@ -218,105 +218,105 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 776184832.0, - "2": 792302080.0, - "3": 776188416.0, - "4": 803877376.0, - "5": 807465472.0, - "6": 801239552.0, + "1": 773817344.0, + "2": 773820928.0, + "3": 783290880.0, + "4": 806515200.0, + "5": 806515200.0, + "6": 806515200.0, "7": 801239552.0, - "8": 801239552.0, + "8": 807072256.0, "9": 801239552.0, - "10": 803877376.0, + "10": 804308480.0, "11": 806515200.0, - "12": 801239552.0, - "13": 803877376.0, - "14": 807175680.0, - "15": 807137792.0, - "16": 808035840.0, - "17": 801239552.0, + "12": 806515200.0, + "13": 801239552.0, + "14": 801239552.0, + "15": 806515200.0, + "16": 801239552.0, + "17": 804414976.0, "18": 801239552.0, - "19": 801239552.0, + "19": 804465664.0, "20": 801239552.0, - "21": 808100864.0, - "22": 807793152.0, + "21": 801239552.0, + "22": 806990848.0, "23": 801239552.0, - "24": 804506112.0, + "24": 801239552.0, "25": 801239552.0, - "26": 807669248.0, - "27": 801239552.0, + "26": 804096000.0, + "27": 804142592.0, "28": 801239552.0, - "29": 803877376.0, - "30": 808099328.0, - "31": 801239552.0, - "32": 804414976.0, - "33": 801239552.0, - "34": 801239552.0, - "35": 806515200.0, - "36": 801239552.0, + "29": 807285248.0, + "30": 801239552.0, + "31": 807185920.0, + "32": 804123136.0, + "33": 807465472.0, + "34": 804414976.0, + "35": 801239552.0, + "36": 807117312.0, "37": 801239552.0, - "38": 803877376.0, - "39": 807889408.0, + "38": 806812160.0, + "39": 804086272.0, "40": 801239552.0, - "41": 803877376.0, - "42": 806515200.0, + "41": 801239552.0, + "42": 806621696.0, "43": 801239552.0, - "44": 801239552.0, - "45": 806515200.0, - "46": 806515200.0, - "47": 801239552.0, - "48": 806515200.0, - "49": 806515200.0, - "50": 806515200.0, - "51": 803877376.0, - "52": 807052800.0, + "44": 806515200.0, + "45": 804414976.0, + "46": 801239552.0, + "47": 807562752.0, + "48": 804414976.0, + "49": 801239552.0, + "50": 801239552.0, + "51": 801239552.0, + "52": 801239552.0, "53": 801239552.0, - "54": 801239552.0, - "55": 807037440.0, - "56": 803877376.0, - "57": 801239552.0, - "58": 806515200.0, - "59": 804414976.0, - "60": 803877376.0, + "54": 803877376.0, + "55": 801239552.0, + "56": 806515200.0, + "57": 803877376.0, + "58": 801239552.0, + "59": 806515200.0, + "60": 801239552.0, "61": 801239552.0, - "62": 801239552.0, + "62": 804414976.0, "63": 801239552.0, - "64": 804414976.0, + "64": 801239552.0, "65": 803877376.0, "66": 801239552.0, - "67": 801239552.0, - "68": 801239552.0, - "69": 803877376.0, - "70": 804414976.0, - "71": 801239552.0, + "67": 807842304.0, + "68": 806882816.0, + "69": 803884544.0, + "70": 801239552.0, + "71": 807052800.0, "72": 801239552.0, "73": 801239552.0, - "74": 806515200.0, - "75": 803877376.0, - "76": 801239552.0, + "74": 807236608.0, + "75": 806515200.0, + "76": 804555264.0, "77": 801239552.0, - "78": 803877376.0, + "78": 801239552.0, "79": 803877376.0, - "80": 803877376.0, - "81": 801239552.0, - "82": 807052800.0, - "83": 803877376.0, - "84": 801239552.0, - "85": 801239552.0, + "80": 801239552.0, + "81": 807938560.0, + "82": 804414976.0, + "83": 801239552.0, + "84": 807137792.0, + "85": 803877376.0, "86": 801239552.0, - "87": 807052800.0, - "88": 806515200.0, - "89": 804082176.0, - "90": 801239552.0, + "87": 807353856.0, + "88": 804414976.0, + "89": 801239552.0, + "90": 803877376.0, "91": 801239552.0, - "92": 807152128.0, - "93": 803877376.0, + "92": 803877376.0, + "93": 807052800.0, "94": 801239552.0, - "95": 804414976.0, + "95": 801239552.0, "96": 801239552.0, "97": 801239552.0, - "98": 807473664.0, - "99": 804164096.0, + "98": 803877376.0, + "99": 803877376.0, "100": 801239552.0 } }, @@ -325,106 +325,106 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 934987264.0, - "2": 1151992320.0, - "3": 1151992320.0, - "4": 1151992320.0, - "5": 1151992320.0, - "6": 1151992320.0, - "7": 1151992320.0, - "8": 1151992320.0, - "9": 1151992320.0, - "10": 1151992320.0, - "11": 1151992320.0, - "12": 1151992320.0, - "13": 1151992320.0, - "14": 1151992320.0, - "15": 1151992320.0, - "16": 1151992320.0, - "17": 1151992320.0, - "18": 1151992320.0, - "19": 1151992320.0, - "20": 1151992320.0, - "21": 1151992320.0, - "22": 1151992320.0, - "23": 1151992320.0, - "24": 1151992320.0, - "25": 1151992320.0, - "26": 1151992320.0, - "27": 1151992320.0, - "28": 1151992320.0, - "29": 1151992320.0, - "30": 1151992320.0, - "31": 1151992320.0, - "32": 1151992320.0, - "33": 1151992320.0, - "34": 1151992320.0, - "35": 1151992320.0, - "36": 1151992320.0, - "37": 1151992320.0, - "38": 1151992320.0, - "39": 1151992320.0, - "40": 1151992320.0, - "41": 1151992320.0, - "42": 1151992320.0, - "43": 1151992320.0, - "44": 1151992320.0, - "45": 1151992320.0, - "46": 1151992320.0, - "47": 1151992320.0, - "48": 1151992320.0, - "49": 1151992320.0, - "50": 1151992320.0, - "51": 1151992320.0, - "52": 1151992320.0, - "53": 1151992320.0, - "54": 1151992320.0, - "55": 1151992320.0, - "56": 1151992320.0, - "57": 1151992320.0, - "58": 1151992320.0, - "59": 1151992320.0, - "60": 1151992320.0, - "61": 1151992320.0, - "62": 1151992320.0, - "63": 1151992320.0, - "64": 1151992320.0, - "65": 1151992320.0, - "66": 1151992320.0, - "67": 1151992320.0, - "68": 1151992320.0, - "69": 1151992320.0, - "70": 1151992320.0, - "71": 1151992320.0, - "72": 1151992320.0, - "73": 1151992320.0, - "74": 1151992320.0, - "75": 1151992320.0, - "76": 1151992320.0, - "77": 1151992320.0, - "78": 1151992320.0, - "79": 1151992320.0, - "80": 1151992320.0, - "81": 1151992320.0, - "82": 1151992320.0, - "83": 1151992320.0, - "84": 1151992320.0, - "85": 1151992320.0, - "86": 1151992320.0, - "87": 1151992320.0, - "88": 1151992320.0, - "89": 1151992320.0, - "90": 1151992320.0, - "91": 1151992320.0, - "92": 1151992320.0, - "93": 1151992320.0, - "94": 1151992320.0, - "95": 1151992320.0, - "96": 1151992320.0, - "97": 1151992320.0, - "98": 1151992320.0, - "99": 1151992320.0, - "100": 1151992320.0 + "1": 934613504.0, + "2": 1146764288.0, + "3": 1147428864.0, + "4": 1147428864.0, + "5": 1147428864.0, + "6": 1147428864.0, + "7": 1147428864.0, + "8": 1147428864.0, + "9": 1147428864.0, + "10": 1147428864.0, + "11": 1147428864.0, + "12": 1147428864.0, + "13": 1147428864.0, + "14": 1147428864.0, + "15": 1147428864.0, + "16": 1147428864.0, + "17": 1147428864.0, + "18": 1147428864.0, + "19": 1147428864.0, + "20": 1147428864.0, + "21": 1147428864.0, + "22": 1147428864.0, + "23": 1147428864.0, + "24": 1147428864.0, + "25": 1147428864.0, + "26": 1147428864.0, + "27": 1147428864.0, + "28": 1147428864.0, + "29": 1147428864.0, + "30": 1147428864.0, + "31": 1147428864.0, + "32": 1147428864.0, + "33": 1147428864.0, + "34": 1147428864.0, + "35": 1147428864.0, + "36": 1147428864.0, + "37": 1147428864.0, + "38": 1147428864.0, + "39": 1147428864.0, + "40": 1147428864.0, + "41": 1147428864.0, + "42": 1147428864.0, + "43": 1147428864.0, + "44": 1147428864.0, + "45": 1147428864.0, + "46": 1147428864.0, + "47": 1147428864.0, + "48": 1147428864.0, + "49": 1147428864.0, + "50": 1147428864.0, + "51": 1147428864.0, + "52": 1147428864.0, + "53": 1147428864.0, + "54": 1147428864.0, + "55": 1147428864.0, + "56": 1147428864.0, + "57": 1147428864.0, + "58": 1147428864.0, + "59": 1147428864.0, + "60": 1147428864.0, + "61": 1147428864.0, + "62": 1147428864.0, + "63": 1147428864.0, + "64": 1147428864.0, + "65": 1147428864.0, + "66": 1147428864.0, + "67": 1147428864.0, + "68": 1147428864.0, + "69": 1147428864.0, + "70": 1147428864.0, + "71": 1147428864.0, + "72": 1147428864.0, + "73": 1147428864.0, + "74": 1147428864.0, + "75": 1147428864.0, + "76": 1147428864.0, + "77": 1147428864.0, + "78": 1147428864.0, + "79": 1147428864.0, + "80": 1147428864.0, + "81": 1147428864.0, + "82": 1147428864.0, + "83": 1147428864.0, + "84": 1147428864.0, + "85": 1147428864.0, + "86": 1147428864.0, + "87": 1147428864.0, + "88": 1147428864.0, + "89": 1147428864.0, + "90": 1147428864.0, + "91": 1147428864.0, + "92": 1147428864.0, + "93": 1147428864.0, + "94": 1147428864.0, + "95": 1147428864.0, + "96": 1147428864.0, + "97": 1147428864.0, + "98": 1147428864.0, + "99": 1147428864.0, + "100": 1147428864.0 } }, "mtp_1 loss": { @@ -434,211 +434,210 @@ "values": { "1": 10.91979, "2": 10.92264, - "3": 10.92209, - "4": 10.91955, - "5": 10.90835, - "6": 10.92272, - "7": 10.93087, - "8": 10.92437, - "9": 10.91985, - "10": 10.92389, - "11": 10.91598, - "12": 10.90802, - "13": 10.92185, - "14": 10.91605, - "15": 10.91495, - "16": 10.91889, - "17": 10.90071, - "18": 10.91147, - "19": 10.91668, - "20": 10.8999, - "21": 10.90635, - "22": 10.90341, - "23": 10.90177, - "24": 10.89122, - "25": 10.8824, - "26": 10.89437, - "27": 10.88841, - "28": 10.87614, - "29": 10.86752, - "30": 10.87704, - "31": 10.86234, - "32": 10.86731, - "33": 10.86105, - "34": 10.85919, - "35": 10.85696, - "36": 10.86409, - "37": 10.84378, - "38": 10.85715, - "39": 10.84402, - "40": 10.85552, - "41": 10.84961, - "42": 10.84625, - "43": 10.84286, - "44": 10.82915, - "45": 10.84334, - "46": 10.83246, - "47": 10.83934, - "48": 10.8254, - "49": 10.83159, - "50": 10.8209, - "51": 10.82446, - "52": 10.81983, - "53": 10.80681, - "54": 10.80125, - "55": 10.79716, - "56": 10.79536, - "57": 10.80256, - "58": 10.80015, - "59": 10.77957, - "60": 10.77884, - "61": 10.77508, - "62": 10.76253, - "63": 10.77603, - "64": 10.761, - "65": 10.75846, - "66": 10.75776, - "67": 10.74768, - "68": 10.7305, - "69": 10.73663, - "70": 10.73712, - "71": 10.72532, - "72": 10.73045, - "73": 10.71784, - "74": 10.70335, - "75": 10.69014, - "76": 10.70558, - "77": 10.68498, - "78": 10.6753, - "79": 10.66888, - "80": 10.67872, - "81": 10.67143, - "82": 10.65162, - "83": 10.63075, - "84": 10.61692, - "85": 10.60406, - "86": 10.63193, - "87": 10.62773, - "88": 10.62836, - "89": 10.5979, - "90": 10.59492, - "91": 10.60603, - "92": 10.58196, - "93": 10.55322, - "94": 10.58501, - "95": 10.57284, - "96": 10.56957, - "97": 10.55499, - "98": 10.55773, - "99": 10.54915, - "100": 10.52235 + "3": 10.9221, + "4": 10.91993, + "5": 10.90869, + "6": 10.92286, + "7": 10.931, + "8": 10.92453, + "9": 10.92022, + "10": 10.9236, + "11": 10.91587, + "12": 10.90821, + "13": 10.92167, + "14": 10.91611, + "15": 10.91537, + "16": 10.91895, + "17": 10.90102, + "18": 10.9115, + "19": 10.91676, + "20": 10.90004, + "21": 10.90603, + "22": 10.90305, + "23": 10.90172, + "24": 10.89078, + "25": 10.88234, + "26": 10.89373, + "27": 10.88813, + "28": 10.87636, + "29": 10.86783, + "30": 10.87714, + "31": 10.86239, + "32": 10.86703, + "33": 10.86125, + "34": 10.85896, + "35": 10.85674, + "36": 10.86428, + "37": 10.84354, + "38": 10.85707, + "39": 10.84396, + "40": 10.85579, + "41": 10.84922, + "42": 10.84627, + "43": 10.84337, + "44": 10.829, + "45": 10.8435, + "46": 10.83292, + "47": 10.83933, + "48": 10.82565, + "49": 10.83153, + "50": 10.82111, + "51": 10.82452, + "52": 10.8198, + "53": 10.80685, + "54": 10.80143, + "55": 10.7972, + "56": 10.79556, + "57": 10.80273, + "58": 10.80037, + "59": 10.77951, + "60": 10.77882, + "61": 10.77525, + "62": 10.76244, + "63": 10.7759, + "64": 10.76146, + "65": 10.75832, + "66": 10.75802, + "67": 10.74779, + "68": 10.73056, + "69": 10.7368, + "70": 10.73718, + "71": 10.72525, + "72": 10.73049, + "73": 10.71786, + "74": 10.70345, + "75": 10.69033, + "76": 10.70576, + "77": 10.6848, + "78": 10.67536, + "79": 10.66892, + "80": 10.67858, + "81": 10.67138, + "82": 10.65165, + "83": 10.63062, + "84": 10.61682, + "85": 10.60375, + "86": 10.63172, + "87": 10.62801, + "88": 10.62831, + "89": 10.59767, + "90": 10.59488, + "91": 10.60607, + "92": 10.58199, + "93": 10.553, + "94": 10.58477, + "95": 10.57294, + "96": 10.56937, + "97": 10.55493, + "98": 10.55738, + "99": 10.54871, + "100": 10.52187 } }, "iteration-time": { - "start_step": 1, + "start_step": 2, "end_step": 100, "step_interval": 1, "values": { - "1": "nan", - "2": 35.74123, - "3": 2.18315, - "4": 5.01562, - "5": 1.03725, - "6": 1.02347, - "7": 1.02133, - "8": 1.02003, - "9": 1.01236, - "10": 1.02816, - "11": 1.02785, - "12": 1.0306, - "13": 3.07844, - "14": 1.08648, - "15": 1.90378, - "16": 1.42103, - "17": 1.03321, - "18": 1.03208, - "19": 1.08999, - "20": 1.12218, - "21": 1.13644, - "22": 1.12541, - "23": 1.11741, - "24": 1.14956, - "25": 1.16864, - "26": 1.15669, - "27": 1.13971, - "28": 1.12588, - "29": 1.13278, - "30": 1.13947, - "31": 1.12834, - "32": 1.14288, - "33": 1.1523, - "34": 1.14619, - "35": 1.14037, - "36": 1.13045, - "37": 1.13097, - "38": 1.12346, - "39": 1.12979, - "40": 1.12916, - "41": 1.14953, - "42": 1.14055, - "43": 1.12561, - "44": 1.13109, - "45": 1.14685, - "46": 1.12948, - "47": 1.15187, - "48": 1.13531, - "49": 1.09506, - "50": 1.32392, - "51": 1.3855, - "52": 1.32904, - "53": 1.33371, - "54": 1.33298, - "55": 1.36204, - "56": 1.36401, - "57": 1.38403, - "58": 1.36486, - "59": 1.38035, - "60": 1.38823, - "61": 1.38756, - "62": 1.39908, - "63": 1.38566, - "64": 1.36931, - "65": 1.39594, - "66": 1.37573, - "67": 1.38916, - "68": 1.39546, - "69": 1.36946, - "70": 1.39074, - "71": 1.39027, - "72": 1.39378, - "73": 1.3835, - "74": 1.37947, - "75": 1.38571, - "76": 1.36867, - "77": 1.38108, - "78": 1.38407, - "79": 1.39463, - "80": 1.396, - "81": 1.36373, - "82": 1.3735, - "83": 1.37647, - "84": 1.39929, - "85": 1.38807, - "86": 1.38388, - "87": 1.3898, - "88": 1.39568, - "89": 1.3878, - "90": 1.37603, - "91": 1.38357, - "92": 1.38586, - "93": 1.3676, - "94": 1.36346, - "95": 1.36686, - "96": 1.37126, - "97": 1.37851, - "98": 1.38658, - "99": 1.38347, - "100": 1.39137 + "2": 36.95995, + "3": 2.25128, + "4": 4.84102, + "5": 1.07708, + "6": 1.13272, + "7": 1.14333, + "8": 1.15024, + "9": 1.12844, + "10": 1.14333, + "11": 1.12443, + "12": 1.11755, + "13": 1.13044, + "14": 1.13118, + "15": 1.13964, + "16": 1.14241, + "17": 1.14914, + "18": 1.16541, + "19": 1.15392, + "20": 1.13862, + "21": 1.14563, + "22": 1.15857, + "23": 1.13109, + "24": 1.2463, + "25": 1.27553, + "26": 1.23883, + "27": 1.19841, + "28": 1.17881, + "29": 1.16399, + "30": 1.14686, + "31": 1.15004, + "32": 1.16511, + "33": 1.15369, + "34": 1.15654, + "35": 1.16587, + "36": 1.16761, + "37": 1.15285, + "38": 1.17031, + "39": 1.1845, + "40": 1.16815, + "41": 1.20312, + "42": 1.27, + "43": 1.21109, + "44": 1.18879, + "45": 1.17708, + "46": 1.16366, + "47": 1.17611, + "48": 1.17039, + "49": 1.17102, + "50": 1.17962, + "51": 1.3918, + "52": 1.237, + "53": 1.29956, + "54": 1.23425, + "55": 1.16788, + "56": 1.16657, + "57": 1.1691, + "58": 1.15581, + "59": 1.18635, + "60": 1.16628, + "61": 1.17024, + "62": 1.35794, + "63": 1.38504, + "64": 1.36071, + "65": 1.35403, + "66": 1.36469, + "67": 1.38767, + "68": 1.4093, + "69": 1.37622, + "70": 1.35787, + "71": 1.3634, + "72": 1.36398, + "73": 1.37404, + "74": 1.37641, + "75": 1.39726, + "76": 1.3754, + "77": 1.38181, + "78": 1.39532, + "79": 1.3611, + "80": 1.35324, + "81": 1.37605, + "82": 1.3567, + "83": 1.35862, + "84": 1.38332, + "85": 1.3776, + "86": 1.36853, + "87": 1.38434, + "88": 1.36442, + "89": 1.35665, + "90": 1.36918, + "91": 1.35432, + "92": 1.33636, + "93": 1.32397, + "94": 1.32647, + "95": 1.3132, + "96": 1.31839, + "97": 1.29889, + "98": 1.23952, + "99": 1.20057, + "100": 1.36829 } } } \ No newline at end of file From e60e42fa1e82c2b9e2fb4145f93e446394447bda Mon Sep 17 00:00:00 2001 From: "achyuthan.s" <113010327+Achyuthan-S@users.noreply.github.com> Date: Sat, 15 Aug 2026 21:54:01 +0000 Subject: [PATCH 048/101] Support HFSDP deferred DP-outer gradient reduction (#5972) Signed-off-by: Achyuthan Sivasankar Signed-off-by: Jingyue Wu Signed-off-by: svcnvidia-nemo-ci Co-authored-by: Jingyue Wu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: svcnvidia-nemo-ci --- .../experimental/parameter_group.py | 43 +++-- .../distributed/mfsdp_v2/test_fully_shard.py | 161 ++++++++++++++++++ 2 files changed, 194 insertions(+), 10 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 622d1f7908c..aefcd15289d 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -356,12 +356,13 @@ def reduce_partial_gradients( ) -> None: """Reduce a packed partial gradient buffer into sharded parameter gradients. - For HSDP main_grad rests DP-outer-Partial (Partial where main_weight is - Replicate) between microbatches, accumulating each backward through the - standard zero_grad contract; the last microbatch reduces the DP-outer axes, - finalizing main_grad to main_weight's placements so ``.grad`` is the fully - reduced gradient before ``optimizer.step()``. With every axis Flat (plain - DP) main_grad already rests finalized. + For HSDP/HFSDP main_grad rests DP-outer-Partial between microbatches, + accumulating each backward through the standard zero_grad contract; the + last microbatch reduces the DP-outer axes, finalizing main_grad to + main_weight's placements (all-reduce to Replicate for HSDP, reduce-scatter + to Flat for HFSDP) so ``.grad`` is the fully reduced gradient before + ``optimizer.step()``. With every axis Flat (plain DP) main_grad already + rests finalized. """ assert self.main_grad is not None @@ -371,14 +372,36 @@ def reduce_partial_gradients( has_sharded_grads = self._has_sharded_grads() # A non-accumulation main_grad means the previous step finalized it; this - # only happens on the first microbatch. Redistribute it back to the - # DP-outer-Partial accumulation placement -- a metadata relabel for HSDP, - # and a fresh reduce-scattered buffer for HFSDP in the future. + # only happens on the first microbatch. Restore it to the DP-outer-Partial + # accumulation placement. HSDP's finalize keeps the buffer size (Replicate + # on DP-outer), so relabel it in place; HFSDP's finalize reduce-scattered + # DP-outer to a smaller optimizer-sharded buffer, so allocate a fresh one + # (zeroed only when we accumulate into it, i.e. sharded grads are still set). if self.main_grad.placements != self._accumulation_placements: assert self.main_grad.allocation_stream == ( torch.cuda.current_stream(self.main_grad.device) ) - self.main_grad = self.main_grad.redistribute(self._accumulation_placements) + reset_axis = changed_mesh_axis(self.main_grad.placements, self._accumulation_placements) + assert reset_axis is not None # the placements differ, so an axis changed + if isinstance(self.main_grad.placements[reset_axis], Replicate): + # HSDP: Replicate -> Partial changes only metadata and reuses the tensor. + self.main_grad = self.main_grad.redistribute(self._accumulation_placements) + else: + # HFSDP: main_grad was reduce-scattered to the optimizer shard, too small + # to hold the accumulation, so re-allocate. This runs inside the + # reduce_scatter stream context (see FsdpModule._reduce_gradient_groups), + # so the buffer is allocated on that stream and stays race-safe. Zero it + # only when we accumulate (set_to_none=False); with set_to_none=True the + # reduction below overwrites it via out=. + self.main_grad = DBuffer( + mesh=self.mesh, + placements=self._accumulation_placements, + tensor_shapes=self.main_weight.layout.tensor_shapes, + dtype=self.main_grad.dtype, + device=self.main_weight.device, + ) + if has_sharded_grads: + self.main_grad.local_buffer.zero_() can_reduce_into_main_grad = ( not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 1886945f3be..131afdbd0c8 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -162,6 +162,19 @@ def _hsdp_placements() -> Placements: ) +def _hfsdp_placements() -> Placements: + """HFSDP: params replicated across DP-outer (axis 0) for compute but the + optimizer sharded across it, all sharded within DP-inner (axis 1). main_grad + rests [Partial, Flat] between microbatches and is reduce-scattered to + [Flat, Flat] (the optimizer placement) on the last microbatch.""" + return Placements( + dp_axes=[0, 1], + parameter=[Replicate(), Flat()], + gradient=[Partial(dist.ReduceOp.AVG), Flat()], + optimizer=[Flat(), Flat()], + ) + + def _mb(num_bytes: int) -> str: return f"{num_bytes / 1024**2:.2f} MB" @@ -348,6 +361,90 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]: ) +@pytest.mark.parametrize("set_to_none", [True, False]) +@pytest.mark.parametrize("num_microbatches", [1, 3]) +def test_hfsdp_losses_match_baseline(distributed_setup, num_microbatches, set_to_none): + """HFSDP (optimizer sharded across DP-outer too) training should match single-rank SGD. + + Like HSDP, gradients reduce-scatter within DP-inner every backward and + accumulate into main_grad. Unlike HSDP, the last-microbatch DP-outer reduction + is a reduce-scatter (not an all-reduce) that finalizes main_grad to the + optimizer's [Flat, Flat] placement, shrinking the buffer; the next step's reset + therefore allocates a fresh [Partial, Flat] accumulation buffer. Every rank + sees identical data, so the averaged gradient equals the single-rank gradient + and losses must match. Both ``zero_grad`` modes are covered. + """ + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + # world_size=2 gives a 2x1 mesh: DP-inner is trivial but the DP-outer + # reduce-scatter finalize and the fresh-buffer reset still run and converge. + if world_size % 2 != 0: + pytest.skip("This test requires an even number of ranks for a 2-D DP mesh.") + + outer_size = 2 + inner_size = world_size // outer_size + mesh = init_device_mesh( + device.type, (outer_size, inner_size), mesh_dim_names=("dp_outer", "dp_inner") + ) + torch.manual_seed(1234) + dim = 8 + baseline = MultiChildModel(dim=dim, num_children=2).to(device) + model = MultiChildModel(dim=dim, num_children=2).to(device) + model.load_state_dict(baseline.state_dict()) + + # Shard the child layers, then the model, so the children share a root context + # and reduce through the overlap path instead of as independent roots. + with fully_shard_context(device=device) as context: + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=_hfsdp_placements()) + fully_shard(model, mesh=mesh, placements=_hfsdp_placements()) + baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + # HFSDP's optimizer placement [Flat, Flat] differs from the parameter placement + # [Replicate, Flat], so main_weight and model_weight are distinct buffers and the + # compute weight is stale until the step post-hook registered here refreshes it. + # HSDP needs no wrapper: its two placements match, so the buffers alias. + fully_shard_optimizer(optimizer) + + micro_batch_size = 2 + x = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + target = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + microbatches = tuple(zip(x.unbind(), target.unbind())) + + def train(model, optimizer, log_prefix) -> list[torch.Tensor]: + losses = [] + for step in range(5): + optimizer.zero_grad(set_to_none=set_to_none) + + for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): + is_last = microbatch_index == num_microbatches - 1 + with microbatch(context, is_last=is_last): + loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) + (loss / num_microbatches).backward() + losses.append(loss.detach()) + logger.debug( + "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s", + log_prefix, + rank, + step, + microbatch_index, + loss, + ) + + optimizer.step() + return losses + + baseline_losses = train(baseline, baseline_optimizer, "Baseline") + sharded_losses = train(model, optimizer, "HFSDP") + + torch.testing.assert_close( + torch.stack(sharded_losses), + torch.stack(baseline_losses), + msg="HFSDP losses did not match baseline losses.", + ) + + def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup): """HSDP reduce-scatters DP-inner every microbatch but all-reduces DP-outer once. @@ -411,6 +508,70 @@ def train_one_step() -> None: ) +def test_hfsdp_reduce_scatters_dp_outer_on_last_microbatch(distributed_setup): + """HFSDP finalizes with a reduce-scatter, not an all-reduce, on the last microbatch. + + Because the optimizer is sharded across DP-outer, the DP-outer finalize is a + reduce-scatter like the per-microbatch DP-inner reduction -- so there are no + all-reduces at all, and the reduce-scatter count is (num_microbatches + 1) per + parameter group: one DP-inner reduce-scatter every microbatch plus one DP-outer + reduce-scatter on the last. This asserts on kernel counts only, not numerics. + """ + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 4 or world_size % 2 != 0: + pytest.skip("This test requires an even number of at least 4 ranks for a 2-D DP mesh.") + + outer_size = 2 + inner_size = world_size // outer_size + mesh = init_device_mesh( + device.type, (outer_size, inner_size), mesh_dim_names=("dp_outer", "dp_inner") + ) + torch.manual_seed(1234) + dim = 8 + num_children = 2 + model = MultiChildModel(dim=dim, num_children=num_children).to(device) + with fully_shard_context(device=device) as context: + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=_hfsdp_placements()) + fully_shard(model, mesh=mesh, placements=_hfsdp_placements()) + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + + num_microbatches = 3 + micro_batch_size = 2 + x = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + target = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + microbatches = tuple(zip(x.unbind(), target.unbind())) + + def train_one_step() -> None: + optimizer.zero_grad(set_to_none=True) + for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): + is_last = microbatch_index == num_microbatches - 1 + with microbatch(context, is_last=is_last): + loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) + (loss / num_microbatches).backward() + optimizer.step() + + train_one_step() + torch.cuda.synchronize(device) + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + train_one_step() + torch.cuda.synchronize(device) + + reduce_scatter_kernels = collect_linked_kernels(prof, _REDUCE_SCATTER_OP_NAME_SUBSTRING) + allreduce_kernels = collect_linked_kernels(prof, _ALLREDUCE_OP_NAME_SUBSTRING) + # HFSDP reduce-scatters the DP-outer axis, so it never all-reduces. + assert not allreduce_kernels, [event.name for event in prof.events()] + # Per group (each child layer plus the root bias): one DP-inner reduce-scatter + # every microbatch plus one DP-outer reduce-scatter on the last microbatch. + expected = (num_microbatches + 1) * (num_children + 1) + assert len(reduce_scatter_kernels) == expected, ( + f"Expected {expected} reduce-scatters ((num_microbatches + 1) x (num_children + 1)), " + f"got {len(reduce_scatter_kernels)}." + ) + + def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): """An outer FsdpModule owns direct parameters but not nested child FsdpModule parameters.""" world_size = distributed_setup.world_size From 9e96f458b5cd7e6bd0871e368b359660b32a135c Mon Sep 17 00:00:00 2001 From: Jiancong Zhu Date: Sun, 16 Aug 2026 02:17:26 +0000 Subject: [PATCH 049/101] Fix redundant contiguous() call in save_dist_checkpoint_full (#5751) Signed-off-by: Chasing1020 --- tools/checkpoint/dist_checkpoint_io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/checkpoint/dist_checkpoint_io.py b/tools/checkpoint/dist_checkpoint_io.py index 33f0814cc5a..4b3db6de184 100644 --- a/tools/checkpoint/dist_checkpoint_io.py +++ b/tools/checkpoint/dist_checkpoint_io.py @@ -240,7 +240,7 @@ def save_dist_checkpoint_full( raw_state_dict = OrderedDict() for bare_key, tensor in model_state_dict.items(): full_key = f"{model_prefix}{bare_key}" if model_prefix else bare_key - raw_state_dict[full_key] = tensor.contiguous() if tensor.is_contiguous() else tensor.contiguous() + raw_state_dict[full_key] = tensor.contiguous() writer = FileSystemWriter(save_dir) dcp.save( From c95b0da7ded84e26e0b7065aed4fcc03d9ce8fcd Mon Sep 17 00:00:00 2001 From: Antoni-Joan Solergibert Date: Sun, 16 Aug 2026 11:24:19 +0000 Subject: [PATCH 050/101] Cache the FullyParallel save distribution on disk (--ckpt-pg-tensors-cache-path) (#5553) Signed-off-by: asolergi-nv --- .../core/dist_checkpointing/exchange_utils.py | 276 ++++++++++++++++-- .../strategies/fully_parallel.py | 36 ++- megatron/training/arguments.py | 19 ++ megatron/training/checkpointing.py | 4 + megatron/training/config/training_config.py | 19 ++ .../dist_checkpointing/models/common.py | 81 +++++ .../models/test_gpt_model.py | 19 ++ .../models/test_moe_experts.py | 30 ++ .../test_pg_distribution_cache.py | 198 +++++++++++++ 9 files changed, 653 insertions(+), 29 deletions(-) create mode 100644 tests/unit_tests/dist_checkpointing/test_pg_distribution_cache.py diff --git a/megatron/core/dist_checkpointing/exchange_utils.py b/megatron/core/dist_checkpointing/exchange_utils.py index 45eb3d6f2fd..38628e2e416 100644 --- a/megatron/core/dist_checkpointing/exchange_utils.py +++ b/megatron/core/dist_checkpointing/exchange_utils.py @@ -3,6 +3,8 @@ """Utilities for exchanging data between ranks.""" import logging +import os +import pickle from collections import defaultdict from functools import reduce from itertools import zip_longest @@ -172,39 +174,49 @@ def distribute_shards_to_ranks( return shard_to_saving_rank -def determine_main_replica_uniform_distribution( - sharded_state_dict: ShardedStateDict, - parallelization_group: torch.distributed.ProcessGroup, - ignore_groups: bool = False, -) -> Optional[ShardDistribution]: - """Computes the save distribution. +# Prefix of the per-parallelization-group cache files written/read by the +# PG-collective caching feature (see `determine_main_replica_uniform_distribution`). +PG_DIST_CACHE_FILE_PREFIX = "pg_dist" - Should be used in conjunction with `distribute_main_replicas_with_precomputed_distribution` - which applies the computed save distribution. +# Process-global in-memory cache of the per-group distributions, keyed by the +# resolved cache file path (which uniquely identifies a (cache dir, group) pair). +# Each entry holds the full {ignore_groups: ShardDistribution} map, so a single +# entry serves both the save and load variants. This memoization makes +# `determine_main_replica_uniform_distribution` idempotent within a job: +# * CREATE runs the gather + write at most once per group (subsequent calls, +# e.g. later saves or the second `replicate_local_replicas` call, hit memory +# and skip the collective and the file write), and +# * READ opens each group's file at most once per job. +# It is checked *before* the collective so every rank of a group hits or misses +# in lockstep, preserving collective symmetry. +_PG_DIST_CACHE: Dict[str, Dict[bool, "ShardDistribution"]] = {} - We rely on the fact that the assignment algorithm is deterministic on all ranks, - so there is no extra communication needed after metadata exchange. - Args: - sharded_state_dict (ShardedStateDict): state dict to compute the distribution of - parallelization_group (ProcessGroup): distribution will be computed - within this process group - ignore_groups (bool, optional): whether the distribution defines groups. - This option is primarily used during loading, as it ensures that all replicas, - including non-main ones, are loaded by this parallelization group - Defaults to False. +def _gather_shards_metadata( + sharded_state_dict: ShardedStateDict, parallelization_group: torch.distributed.ProcessGroup +) -> List[List[ShardedTensor]]: + """Runs the single expensive collective behind the load/save distribution. - Returns (ShardDistribution, optional): distribution that can be used to apply the - parallelization. Returns None if the process_group is trivial (1 rank) + Every rank in the group contributes the metadata (no tensor data) of the + ``ShardedTensor``s it holds; the result is the list of every rank's shard + metadata, identical on all ranks of the group. This is the + ``all_gather_object`` that the caching feature exists to eliminate on + subsequent jobs. + Decision: data is stripped with ``without_data()`` before the gather. Only + the metadata (key, offsets, shape, dtype, replica_id) drives the + distribution algorithm, and gathering it keeps the pickled payload tiny. + + Args: + sharded_state_dict (ShardedStateDict): state dict whose local shards + are advertised to the group. + parallelization_group (ProcessGroup): group to gather within. + + Returns: + List[List[ShardedTensor]]: per-rank lists of data-less ShardedTensors. """ from ..utils import get_pg_size - if parallelization_group is None: - parallelization_group = torch.distributed.group.WORLD - group_size = get_pg_size(group=parallelization_group) - if group_size <= 1: - return local_shards = list( sh_base for sh_base in nested_values(sharded_state_dict) @@ -212,11 +224,38 @@ def determine_main_replica_uniform_distribution( ) local_shards_no_data = [ten.without_data() for ten in local_shards] - all_shards = [None] * get_pg_size(group=parallelization_group) + all_shards: List[Optional[List[ShardedTensor]]] = [None] * get_pg_size( + group=parallelization_group + ) torch.distributed.all_gather_object( all_shards, local_shards_no_data, group=parallelization_group ) + return cast(List[List[ShardedTensor]], all_shards) + + +def _build_shard_distribution( + all_shards: List[List[ShardedTensor]], ignore_groups: bool +) -> ShardDistribution: + """Computes a ShardDistribution from already-gathered shard metadata. + + This is the pure, deterministic post-gather logic of + ``determine_main_replica_uniform_distribution`` factored out so that: + * it can be run twice (for both ``ignore_groups`` variants) from a single + gather when creating the cache, and + * it has no dependence on the local rank, which is what makes the result + identical on every rank of the group and therefore safe to persist and + reload across jobs. + Args: + all_shards (List[List[ShardedTensor]]): gathered per-rank shard metadata + (output of ``_gather_shards_metadata``). + ignore_groups (bool): if True (load mode) include shards whose main + replica lives outside this group; if False (save mode) restrict to + shards with a main replica in this group. + + Returns: + ShardDistribution: the computed distribution. + """ shard_to_ranks = defaultdict(list) shard_to_size = {} shard_to_metadata = {} @@ -255,6 +294,191 @@ def determine_main_replica_uniform_distribution( ) +def _pg_dist_cache_file_path( + cache_path: str, parallelization_group: torch.distributed.ProcessGroup +) -> str: + """Returns the cache file path holding the distributions for a given group. + + The file is keyed by the minimum global rank of the parallelization group. + Decision rationale: + * Parallelization groups partition the world, so ``min(global ranks)`` is + a unique, collision-free identifier of each group. + * It is computed from purely local metadata + (``get_process_group_ranks``), so the filename is derivable without any + collective — essential for the zero-communication read path. + * Every member of a group maps to the same file (one file per group rather + than one per rank): far fewer files at scale and a single read per rank, + with duplicate reads absorbed by the FS page cache. + + Args: + cache_path (str): directory holding the cache files. + parallelization_group (ProcessGroup): the group to key on. + + Returns: + str: absolute-or-relative path of the group's cache file. + """ + ranks = torch.distributed.get_process_group_ranks(parallelization_group) + return os.path.join(cache_path, f"{PG_DIST_CACHE_FILE_PREFIX}_{min(ranks)}.pkl") + + +def _load_pg_dist_cache(cache_file: str) -> Dict[bool, ShardDistribution]: + """Reads both pre-computed ShardDistributions for a group from disk. + + Returns the full ``{ignore_groups: ShardDistribution}`` map (not just one + variant) so the caller can memoize both in the process-global cache and serve + later save/load calls without re-opening the file. + + Performs no existence/validity checks by design (see the feature docs): the + user opts into the cache only when they guarantee the config and world size + match the run that created it, trading safety for the lowest possible + latency. + + Args: + cache_file (str): resolved path of the group's cache file. + + Returns: + Dict[bool, ShardDistribution]: the {ignore_groups: distribution} map. + """ + with open(cache_file, "rb") as f: + return pickle.load(f) + + +def _create_pg_dist_cache( + all_shards: List[List[ShardedTensor]], + cache_path: str, + parallelization_group: torch.distributed.ProcessGroup, +) -> Dict[bool, ShardDistribution]: + """Builds both distribution variants and persists them for the group. + + Both the save (``ignore_groups=False``) and load (``ignore_groups=True``) + distributions are derived from the *same* gathered metadata, so a single + collective suffices to populate a cache usable by every later load and save + — even if the creating job only ever loads (e.g. ``exit_after_loading_ckpt``). + + Only the group's writer rank (global rank == ``min`` of the group) writes, + keeping create-time IO to one file per group. The write is atomic + (temp file + ``os.replace``) so a concurrent/interrupted create can never + leave a torn file that a reader would later trust blindly. + + Args: + all_shards (List[List[ShardedTensor]]): gathered shard metadata. + cache_path (str): directory to write the cache file into. + parallelization_group (ProcessGroup): group being cached. + + Returns: + Dict[bool, ShardDistribution]: the {ignore_groups: distribution} map, + so the caller can return the variant it needs without recomputing. + """ + both: Dict[bool, ShardDistribution] = { + ignore_groups: _build_shard_distribution(all_shards, ignore_groups) + for ignore_groups in (False, True) + } + ranks = torch.distributed.get_process_group_ranks(parallelization_group) + if torch.distributed.get_rank() == min(ranks): + os.makedirs(cache_path, exist_ok=True) + path = _pg_dist_cache_file_path(cache_path, parallelization_group) + tmp_path = f"{path}.tmp.{torch.distributed.get_rank()}" + with open(tmp_path, "wb") as f: + pickle.dump(both, f) + os.replace(tmp_path, path) + return both + + +def determine_main_replica_uniform_distribution( + sharded_state_dict: ShardedStateDict, + parallelization_group: torch.distributed.ProcessGroup, + ignore_groups: bool = False, + pg_cache_path: Optional[str] = None, + pg_cache_create: bool = False, +) -> Optional[ShardDistribution]: + """Computes (or loads from cache) the load/save shard distribution. + + Should be used in conjunction with `distribute_main_replicas_with_precomputed_distribution` + which applies the computed save distribution. + + We rely on the fact that the assignment algorithm is deterministic on all ranks, + so there is no extra communication needed after metadata exchange. + + PG-collective caching: + The metadata exchange (`all_gather_object`) is by far the dominant cost of + this function at scale. Its result is a pure deterministic function of the + sharded state dict structure and the parallel layout, so for a fixed config + and world size it is identical across jobs and can be cached on disk: + + * ``pg_cache_path`` set, ``pg_cache_create=False`` (READ): skip the + collective entirely and load the pre-computed distribution for this + group and mode from disk. + * ``pg_cache_path`` set, ``pg_cache_create=True`` (CREATE): run the + collective once, build *both* the save and load distributions from the + single gather, persist them, and return the requested one. + * both unset (default): the original behaviour, unchanged. + + Results are additionally memoized in a process-global cache (``_PG_DIST_CACHE``) + keyed by the group's cache file, so within a single job each group is created + (gather + write) or read (file open) at most once — later calls (subsequent + saves, the second ``replicate_local_replicas`` pass, or the save wrapper + reusing the load wrapper's entry) are served from memory. + + No existence/validity checks are performed on the cache by design — the user + opts in only when guaranteeing a matching config and world size. + + Args: + sharded_state_dict (ShardedStateDict): state dict to compute the distribution of + parallelization_group (ProcessGroup): distribution will be computed + within this process group + ignore_groups (bool, optional): whether the distribution defines groups. + This option is primarily used during loading, as it ensures that all replicas, + including non-main ones, are loaded by this parallelization group + Defaults to False. + pg_cache_path (str, optional): directory of the PG-distribution cache. + When set, enables the caching feature. Defaults to None (disabled). + pg_cache_create (bool, optional): when True (and ``pg_cache_path`` set), + compute the distribution normally and write the cache; when False + (and ``pg_cache_path`` set), read the cache and skip the collective. + Defaults to False. + + Returns (ShardDistribution, optional): distribution that can be used to apply the + parallelization. Returns None if the process_group is trivial (1 rank) + + """ + from ..utils import get_pg_size + + if parallelization_group is None: + parallelization_group = torch.distributed.group.WORLD + group_size = get_pg_size(group=parallelization_group) + if group_size <= 1: + return + + if pg_cache_path is not None: + # The resolved file path uniquely identifies this (cache dir, group) and + # keys the process-global memo. Checking it *first* makes both READ and + # CREATE idempotent within a job: a second call for the same group (a + # later save, the second `replicate_local_replicas` call, or the save + # wrapper reusing what the load wrapper already cached) returns from + # memory with no collective and no file IO. The check is identical on all + # ranks of the group, so it never desynchronizes the collective below. + cache_file = _pg_dist_cache_file_path(pg_cache_path, parallelization_group) + if cache_file in _PG_DIST_CACHE: + return _PG_DIST_CACHE[cache_file][bool(ignore_groups)] + + # READ path: no collective; open the file once, then memoize both variants. + if not pg_cache_create: + both = _load_pg_dist_cache(cache_file) + _PG_DIST_CACHE[cache_file] = both + return both[bool(ignore_groups)] + + # CREATE path: gather once, build & persist both variants, then memoize + # so no subsequent call for this group gathers or writes again. + all_shards = _gather_shards_metadata(sharded_state_dict, parallelization_group) + both = _create_pg_dist_cache(all_shards, pg_cache_path, parallelization_group) + _PG_DIST_CACHE[cache_file] = both + return both[bool(ignore_groups)] + + # Default path: unchanged behaviour. + all_shards = _gather_shards_metadata(sharded_state_dict, parallelization_group) + return _build_shard_distribution(all_shards, ignore_groups) + + @torch.no_grad() @debug_time(f"exchange_loaded_tensors_gather_rounds", logger) def exchange_loaded_tensors_gather_rounds( diff --git a/megatron/core/dist_checkpointing/strategies/fully_parallel.py b/megatron/core/dist_checkpointing/strategies/fully_parallel.py index 657942cadb0..bc4c36f99a9 100644 --- a/megatron/core/dist_checkpointing/strategies/fully_parallel.py +++ b/megatron/core/dist_checkpointing/strategies/fully_parallel.py @@ -76,6 +76,8 @@ def __init__( backend: str = "torch_dist", version: int = 1, validate_access_integrity: bool = True, + pg_cache_path: Optional[str] = None, + pg_cache_create: bool = False, ): """ """ self.base_strategy = strategy @@ -91,6 +93,11 @@ def __init__( # --no-ckpt-load-validate-sharding-integrity skips *both* save-side # validation collectives, not just the one in save_preprocess. self.validate_access_integrity = validate_access_integrity + # PG-collective caching (see determine_main_replica_uniform_distribution). + # When pg_cache_path is set and pg_cache_create is False, the save + # distribution is read from disk and the all_gather_object is skipped. + self.pg_cache_path = pg_cache_path + self.pg_cache_create = pg_cache_create self.cached_distribution: Optional[ShardDistribution] = None @@ -131,13 +138,25 @@ def apply_saving_parallelization(self, sharded_state_dict: ShardedStateDict) -> else: logger.debug(f'Apply save parallelization') precomputed_distribution = determine_main_replica_uniform_distribution( - sharded_state_dict, self.parallelization_group + sharded_state_dict, + self.parallelization_group, + pg_cache_path=self.pg_cache_path, + pg_cache_create=self.pg_cache_create, ) distribute_main_replicas_with_precomputed_distribution( sharded_state_dict, self.parallelization_group, precomputed_distribution ) - if self.cached_distribution is None and self.validate_access_integrity: + # In cache-read mode the user vouches for a matching config/world size, + # so we also skip this first-save integrity check — it triggers its own + # collective (determine_global_metadata), which would defeat the purpose + # of skipping the distribution all_gather. + pg_cache_read = self.pg_cache_path is not None and not self.pg_cache_create + if ( + self.cached_distribution is None + and not pg_cache_read + and self.validate_access_integrity + ): # First time applying the parallelization validate_sharding_integrity(determine_global_metadata(sharded_state_dict)[1]) if self.do_cache_distribution: @@ -179,6 +198,8 @@ def __init__( do_cache_distribution: bool = False, exchange_algo: str = 'broadcast', per_rank_object_load: bool = False, + pg_cache_path: Optional[str] = None, + pg_cache_create: bool = False, ): self.base_strategy = strategy if parallelization_group is None: @@ -193,6 +214,11 @@ def __init__( # with a WORLD-wide `all_gather_object`. Opt-in; defaults to the legacy # gather-based exchange. See `load` for correctness rationale. self.per_rank_object_load = per_rank_object_load + # PG-collective caching (see determine_main_replica_uniform_distribution). + # When pg_cache_path is set and pg_cache_create is False, the load + # distribution is read from disk and the all_gather_object is skipped. + self.pg_cache_path = pg_cache_path + self.pg_cache_create = pg_cache_create self.cached_distribution: Optional[ShardDistribution] = None self.cached_global_metadata: Optional[Metadata] = None @@ -433,7 +459,11 @@ def apply_loading_parallelization( else: logger.debug(f'Apply load parallelization') precomputed_distribution = determine_main_replica_uniform_distribution( - sharded_state_dict, self.parallelization_group, True + sharded_state_dict, + self.parallelization_group, + True, + pg_cache_path=self.pg_cache_path, + pg_cache_create=self.pg_cache_create, ) distribute_main_replicas_with_precomputed_distribution( diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e08a06ec150..1d49434b24e 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1723,6 +1723,25 @@ def validate_args(args, defaults={}): "Other algorithms cannot guarantee numerical stability yet." ) + # The PG-distribution cache stores one file per parallelization group, keyed by + # that group's minimum global rank, and each file holds both the save and the + # load distribution. If save and load used different parallelization groups, + # two distinct groups could map to the same file (they can share a minimum + # global rank), so one side would silently read a distribution computed for the + # other group's layout. Require both sides to use the same group. + if args.ckpt_pg_tensors_cache_path is not None: + assert ( + args.ckpt_fully_parallel_save_process_group + == args.ckpt_fully_parallel_load_process_group + ), ( + "--ckpt-pg-tensors-cache-path requires --ckpt-fully-parallel-save-process-group and " + "--ckpt-fully-parallel-load-process-group to be identical, but got " + f"save='{args.ckpt_fully_parallel_save_process_group}' and " + f"load='{args.ckpt_fully_parallel_load_process_group}'. The cache is keyed by the " + "parallelization group, so mixing groups can read a distribution computed for a " + "different layout." + ) + if args.load_main_params_from_ckpt: assert args.no_load_optim, '--load-main-params-from-ckpt must be used with --no-load-optim.' diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index e7c1b2ec02c..00ec96f9a41 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -853,6 +853,8 @@ def save_checkpoint( process_group, args.ckpt_assume_constant_structure, validate_access_integrity=args.ckpt_load_validate_sharding_integrity, + pg_cache_path=getattr(args, 'ckpt_pg_tensors_cache_path', None), + pg_cache_create=getattr(args, 'ckpt_pg_tensors_cache_create', False), ) # Allow opting out of save-side sharding validation entirely (even on # the first save) via --no-ckpt-load-validate-sharding-integrity. This @@ -1905,6 +1907,8 @@ def _load_global_dist_base_checkpoint( per_rank_object_load=getattr( args, 'ckpt_fully_parallel_load_per_rank_objects', False ), + pg_cache_path=getattr(args, 'ckpt_pg_tensors_cache_path', None), + pg_cache_create=getattr(args, 'ckpt_pg_tensors_cache_create', False), ) if checkpointing_context is not None: checkpointing_context['load_strategy'] = load_strategy diff --git a/megatron/training/config/training_config.py b/megatron/training/config/training_config.py index a379039ff20..ecc4736cb36 100644 --- a/megatron/training/config/training_config.py +++ b/megatron/training/config/training_config.py @@ -571,6 +571,25 @@ class CheckpointConfig: ckpt_assume_constant_structure: bool = False """Assume the checkpoint structure is constant across saves to enable optimizations.""" + ckpt_pg_tensors_cache_path: Optional[str] = None + """Directory of the parallelization-group distribution cache for fully parallel + save/load of distributed checkpoints. When set, the expensive + ``all_gather_object`` in ``determine_main_replica_uniform_distribution`` is + replaced by a single per-group file read (the load and save distributions are + loaded from this directory). Only safe when the config and world size match the + run that created the cache (see ``--ckpt-pg-tensors-cache-create``); no + existence/validity checks are performed, for the lowest possible latency. + Default (None) preserves the original collective-based behaviour.""" + + ckpt_pg_tensors_cache_create: bool = False + """Create (rather than read) the parallelization-group distribution cache at + ``--ckpt-pg-tensors-cache-path``. Set this for a single run with a matching + config/world size: the collective runs as usual but both the save and load + distributions are derived from that single gather and written to disk, one + file per parallelization group. Subsequent runs set only + ``--ckpt-pg-tensors-cache-path`` (leave this False) to skip the collective. + Default: False.""" + ckpt_load_validate_sharding_integrity: bool = True """Whether to validate sharding access integrity when loading *and saving* a distributed checkpoint. When True (default), each tensor shard is checked to be accessed exactly once as diff --git a/tests/unit_tests/dist_checkpointing/models/common.py b/tests/unit_tests/dist_checkpointing/models/common.py index d18d99d0b11..2f68a61148f 100644 --- a/tests/unit_tests/dist_checkpointing/models/common.py +++ b/tests/unit_tests/dist_checkpointing/models/common.py @@ -141,6 +141,87 @@ def common_test_parallel_reconfiguration_e2e( Utils.destroy_model_parallel() +def common_test_pg_distribution_cache_e2e( + initialize_model_fn, tmp_path_dist_ckpt, parallelization_group, sharded_state_dict_fn=None +): + """Save/load a real sharded model with and without the PG-distribution cache. + + The cache only changes *how* the fully-parallel save/load distribution is + obtained, so it must be invisible end to end: the checkpoint written through it + and the state dict loaded through it have to match the ones produced without it. + + Both halves use the same parallel layout, which is the only regime the cache is + valid for (it is keyed by the parallelization group), and it mirrors how the + feature is used in practice: one job creates the cache, a later job with the + same config reuses it and skips the distribution all_gather. + """ + from megatron.core.dist_checkpointing import exchange_utils + + if sharded_state_dict_fn is None: + sharded_state_dict_fn = lambda model: model.sharded_state_dict() + + def _save(model, ckpt_dir, pg_cache_path=None, pg_cache_create=False): + save_strategy = FullyParallelSaveStrategyWrapper( + TorchDistSaveShardedStrategy(), + parallelization_group, + True, + pg_cache_path=pg_cache_path, + pg_cache_create=pg_cache_create, + ) + save(sharded_state_dict_fn(model), ckpt_dir, save_strategy) + + def _load(model, ckpt_dir, pg_cache_path=None): + load_strategy = FullyParallelLoadStrategyWrapper( + TorchDistLoadShardedStrategy(), parallelization_group, pg_cache_path=pg_cache_path + ) + state_dict, missing_keys, unexpected_keys = load( + sharded_state_dict_fn(model), ckpt_dir, load_strategy, strict=StrictHandling.RETURN_ALL + ) + # Potential mismatch is because of extra states which is ok + assert all('_extra_state' in k for k in missing_keys) + assert all('_extra_state' in k for k in unexpected_keys) + return state_dict + + with ( + TempNamedDir(tmp_path_dist_ckpt / 'pg_cache_dir', sync=True) as cache_dir, + TempNamedDir(tmp_path_dist_ckpt / 'pg_cache_ckpt_plain', sync=True) as ckpt_dir_plain, + TempNamedDir(tmp_path_dist_ckpt / 'pg_cache_ckpt_cached', sync=True) as ckpt_dir_cached, + ): + cache_path = str(cache_dir) + model_A = initialize_model_fn(1) + + # Reference save, then the same save with the cache being created. + exchange_utils._PG_DIST_CACHE.clear() + _save(model_A, ckpt_dir_plain) + exchange_utils._PG_DIST_CACHE.clear() + _save(model_A, ckpt_dir_cached, pg_cache_path=cache_path, pg_cache_create=True) + torch.distributed.barrier() + + # Load both back into freshly initialized models (seed 2, so the loaded + # values can only come from the checkpoints). The cached load drops the + # in-process memo first so it genuinely reads the distribution from disk. + model_B_plain = initialize_model_fn(2) + exchange_utils._PG_DIST_CACHE.clear() + loaded_plain = _load(model_B_plain, ckpt_dir_plain) + + model_B_cached = initialize_model_fn(2) + exchange_utils._PG_DIST_CACHE.clear() + loaded_cached = _load(model_B_cached, ckpt_dir_cached, pg_cache_path=cache_path) + + # The cached load must return exactly what the uncached load returns. + diffs = diff(loaded_plain, loaded_cached) + assert not any(map(bool, diffs)), diffs + + # ... and both checkpoints must hold the same tensors on disk. + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(1, 1) + diffs = diff(load_plain_tensors(ckpt_dir_plain), load_plain_tensors(ckpt_dir_cached)) + assert not any(map(bool, diffs)), diffs + + exchange_utils._PG_DIST_CACHE.clear() + Utils.destroy_model_parallel() + + def common_test_state_dict_comparison(initialize_model_fn, tmp_path_dist_ckpt): tp = 2 pp = 4 diff --git a/tests/unit_tests/dist_checkpointing/models/test_gpt_model.py b/tests/unit_tests/dist_checkpointing/models/test_gpt_model.py index e18d3b4683b..cd91ee9cffe 100644 --- a/tests/unit_tests/dist_checkpointing/models/test_gpt_model.py +++ b/tests/unit_tests/dist_checkpointing/models/test_gpt_model.py @@ -20,6 +20,7 @@ from megatron.core.utils import is_te_min_version from tests.unit_tests.dist_checkpointing.models.common import ( common_test_parallel_reconfiguration_e2e, + common_test_pg_distribution_cache_e2e, common_test_simple_sharded_state_dict_save_load, common_test_state_dict_comparison, common_test_vocab_size_padding_change, @@ -157,6 +158,24 @@ def test_parallel_reconfiguration_e2e( metadata={'singleton_local_shards': singleton_local_shards}, ) + @pytest.mark.parametrize(('tp', 'pp'), [(1, 1), (2, 1), (1, 2), (2, 2)]) + def test_pg_distribution_cache_e2e(self, tmp_path_dist_ckpt, tp, pp): + """A real sharded GPT checkpoint must round-trip identically through the + PG-distribution cache (--ckpt-pg-tensors-cache-path). + + Covers tensor and pipeline parallelism: each pipeline stage has its own + data-parallel group, hence its own cache file, so the cached distribution + must stay stage-local. + """ + Utils.initialize_model_parallel(tp, pp) + common_test_pg_distribution_cache_e2e( + lambda seed: initialize_gpt_model( + seed, gpt_te_spec, tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp + ), + tmp_path_dist_ckpt, + ps.get_data_parallel_group(with_context_parallel=True), + ) + def test_state_dict_comparison(self, tmp_path_dist_ckpt): common_test_state_dict_comparison(initialize_gpt_model, tmp_path_dist_ckpt) diff --git a/tests/unit_tests/dist_checkpointing/models/test_moe_experts.py b/tests/unit_tests/dist_checkpointing/models/test_moe_experts.py index 126968a3af9..d6f66e0ce86 100644 --- a/tests/unit_tests/dist_checkpointing/models/test_moe_experts.py +++ b/tests/unit_tests/dist_checkpointing/models/test_moe_experts.py @@ -29,6 +29,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_te_min_version from tests.unit_tests.dist_checkpointing import TempNamedDir +from tests.unit_tests.dist_checkpointing.models.common import common_test_pg_distribution_cache_e2e from tests.unit_tests.test_utilities import Utils fp8_available, reason_for_no_fp8 = check_fp8_support() @@ -132,6 +133,35 @@ def setup_method(self, method): def teardown_method(self, method): Utils.destroy_model_parallel() + @pytest.mark.parametrize('save_load_process_group', ['dp', 'ep_dp']) + @pytest.mark.parametrize(('tp', 'pp', 'ep'), [(1, 1, 1), (1, 1, 2), (1, 2, 2), (2, 1, 2)]) + @pytest.mark.parametrize('expert_type', expert_type) + def test_pg_distribution_cache_e2e( + self, tmp_path_dist_ckpt, tp, pp, ep, expert_type, save_load_process_group + ): + """MoE experts must round-trip identically through the PG-distribution cache. + + Experts are the case the cache targets: they dominate the shard count and can + be distributed over either parallelization group that + --ckpt-fully-parallel-{save,load}-process-group selects, so both 'dp' and + 'ep_dp' are covered here, across tensor, pipeline and expert parallelism. + """ + Utils.initialize_model_parallel(tp, pp, expert_model_parallel_size=ep) + if save_load_process_group == 'ep_dp': + parallelization_group = parallel_state.get_expert_data_parallel_group() + else: + parallelization_group = parallel_state.get_data_parallel_group( + with_context_parallel=True + ) + common_test_pg_distribution_cache_e2e( + lambda seed: initialize_expert_layer(seed, expert_type=expert_type), + tmp_path_dist_ckpt, + parallelization_group, + sharded_state_dict_fn=lambda model: model.sharded_state_dict( + prefix=f'{parallel_state.get_pipeline_model_parallel_rank()}.' + ), + ) + @pytest.mark.internal @pytest.mark.parametrize( "use_fpsl,src_tp_pp_ep_etp,dest_tp_pp_ep_etp,use_glu", diff --git a/tests/unit_tests/dist_checkpointing/test_pg_distribution_cache.py b/tests/unit_tests/dist_checkpointing/test_pg_distribution_cache.py new file mode 100644 index 00000000000..78027d666bf --- /dev/null +++ b/tests/unit_tests/dist_checkpointing/test_pg_distribution_cache.py @@ -0,0 +1,198 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for the PG-distribution disk cache (``--ckpt-pg-tensors-cache-path``). + +The FullyParallel save/load distribution is a deterministic function of the +sharded-state-dict structure and the parallel layout, computed via a single +``all_gather_object`` (``_gather_shards_metadata``). The cache persists that +result so later jobs can: + * CREATE (``pg_cache_create=True``): gather once and write the distribution. + * READ (``pg_cache_create=False``): load it from disk and skip the gather. + +These tests run at world_size=2 (TP=2) but make no world-size-specific +assertions, so they hold at any size. +""" + +import os + +import torch + +from megatron.core.dist_checkpointing import ShardedTensor, exchange_utils +from megatron.core.dist_checkpointing.exchange_utils import ( + _pg_dist_cache_file_path, + determine_main_replica_uniform_distribution, +) +from megatron.core.dist_checkpointing.strategies.fully_parallel import ( + FullyParallelLoadStrategyWrapper, + FullyParallelSaveStrategyWrapper, +) +from megatron.core.dist_checkpointing.strategies.torch import ( + TorchDistLoadShardedStrategy, + TorchDistSaveShardedStrategy, +) +from tests.unit_tests.dist_checkpointing import TempNamedDir +from tests.unit_tests.test_utilities import Utils + + +def _state_dict(): + # Fully-replicated tensors of different sizes; varying replica_id makes + # different ranks the main replica, so the save distribution is non-trivial. + return { + f"sd_key{i}": ShardedTensor.from_rank_offsets( + f"key{i}", torch.ones(10 * (i + 1)), replica_id=(Utils.rank + i) % Utils.world_size + ) + for i in range(4) + } + + +def _rank_marked_state_dict(): + """Replicated shards whose *content* identifies the rank holding them. + + Replicas normally carry identical data, which makes the elected main replica + unobservable in the loaded values. Marking each rank's copy with its own rank + makes the save distribution visible end-to-end: the loaded value tells you + which rank the distribution elected to write that shard. + """ + return { + f"sd_key{i}": ShardedTensor.from_rank_offsets( + f"key{i}", + torch.full((10 * (i + 1),), float(Utils.rank)), + replica_id=(Utils.rank + i) % Utils.world_size, + ) + for i in range(4) + } + + +class TestPgDistributionCache: + def setup_method(self, method): + Utils.destroy_model_parallel() + exchange_utils._PG_DIST_CACHE.clear() + + def teardown_method(self, method): + Utils.destroy_model_parallel() + exchange_utils._PG_DIST_CACHE.clear() + + def test_create_writes_cache_and_read_matches(self, tmp_path_dist_ckpt): + Utils.initialize_model_parallel(2, 1) + group = torch.distributed.group.WORLD + with TempNamedDir(tmp_path_dist_ckpt / "pg_cache") as cache_dir: + cache_path = str(cache_dir) + + # CREATE: runs the gather once and writes the cache file. + created = determine_main_replica_uniform_distribution( + _state_dict(), group, pg_cache_path=cache_path, pg_cache_create=True + ) + torch.distributed.barrier() + assert os.path.exists(_pg_dist_cache_file_path(cache_path, group)) + + # Drop the in-process memo so the next call truly reads from disk. + exchange_utils._PG_DIST_CACHE.clear() + + # READ: loads the distribution from disk. + read = determine_main_replica_uniform_distribution( + _state_dict(), group, pg_cache_path=cache_path, pg_cache_create=False + ) + + # The deterministic distribution decisions round-trip exactly. + assert read.main_rank_for_shard == created.main_rank_for_shard + assert read.all_ranks_for_shard == created.all_ranks_for_shard + + def test_read_path_skips_the_gather_collective(self, tmp_path_dist_ckpt): + Utils.initialize_model_parallel(2, 1) + group = torch.distributed.group.WORLD + with TempNamedDir(tmp_path_dist_ckpt / "pg_cache_skip") as cache_dir: + cache_path = str(cache_dir) + + created = determine_main_replica_uniform_distribution( + _state_dict(), group, pg_cache_path=cache_path, pg_cache_create=True + ) + torch.distributed.barrier() + exchange_utils._PG_DIST_CACHE.clear() + + # Make the metadata-exchange collective explode; the READ path must + # never reach it. + orig_gather = exchange_utils._gather_shards_metadata + + def _boom(*args, **kwargs): + raise AssertionError("READ path must not run the all_gather_object collective") + + exchange_utils._gather_shards_metadata = _boom + try: + read = determine_main_replica_uniform_distribution( + _state_dict(), group, pg_cache_path=cache_path, pg_cache_create=False + ) + finally: + exchange_utils._gather_shards_metadata = orig_gather + + assert read is not None + assert read.main_rank_for_shard == created.main_rank_for_shard + + def test_checkpoint_roundtrip_matches_without_cache(self, tmp_path_dist_ckpt): + """A real save+load with the cache must return exactly what it does without it. + + The cache only replaces *how* the save/load distribution is obtained, so a + checkpoint written and read back through it must be indistinguishable from + the default path. This exercises the full FullyParallel save -> load + round-trip, which is what would break if a cached distribution were applied + to the wrong shards. + """ + Utils.initialize_model_parallel(2, 1) + group = torch.distributed.group.WORLD + + # Reference (no cache): plain save + load. + with TempNamedDir(tmp_path_dist_ckpt / "rt_plain", sync=True) as plain_dir: + FullyParallelSaveStrategyWrapper(TorchDistSaveShardedStrategy(), group).save( + _rank_marked_state_dict(), plain_dir + ) + torch.distributed.barrier() + plain_loaded = FullyParallelLoadStrategyWrapper( + TorchDistLoadShardedStrategy(), group + ).load(_rank_marked_state_dict(), plain_dir) + + exchange_utils._PG_DIST_CACHE.clear() + + # Same round-trip, but the distribution is written to and then read back + # from the on-disk cache. + with TempNamedDir(tmp_path_dist_ckpt / "rt_cache_dir", sync=True) as cache_dir: + cache_path = str(cache_dir) + with TempNamedDir(tmp_path_dist_ckpt / "rt_cached", sync=True) as cached_ckpt_dir: + # CREATE: this save computes and persists the distribution. + FullyParallelSaveStrategyWrapper( + TorchDistSaveShardedStrategy(), + group, + pg_cache_path=cache_path, + pg_cache_create=True, + ).save(_rank_marked_state_dict(), cached_ckpt_dir) + torch.distributed.barrier() + + # The create-mode save must have persisted the distribution, and + # dropping the in-process memo forces the load to read that file. + assert os.path.exists(_pg_dist_cache_file_path(cache_path, group)) + exchange_utils._PG_DIST_CACHE.clear() + + # READ: the load takes the distribution from disk. Fail loudly if it + # falls back to the collective, which would make this test vacuous. + orig_gather = exchange_utils._gather_shards_metadata + + def _boom(*args, **kwargs): + raise AssertionError("cached load must not run the all_gather_object") + + exchange_utils._gather_shards_metadata = _boom + try: + cached_loaded = FullyParallelLoadStrategyWrapper( + TorchDistLoadShardedStrategy(), group, pg_cache_path=cache_path + ).load(_rank_marked_state_dict(), cached_ckpt_dir) + finally: + exchange_utils._gather_shards_metadata = orig_gather + + assert cached_loaded.keys() == plain_loaded.keys() + for key, expected in plain_loaded.items(): + actual = cached_loaded[key] + assert torch.equal(actual, expected), f"mismatch for {key}" + # Guard against a vacuous comparison: each loaded shard must be a + # rank-marked tensor of the expected shape (a single rank's copy). + for i, key in enumerate(sorted(plain_loaded)): + shard = plain_loaded[key] + assert shard.shape == (10 * (i + 1),) + assert torch.equal(shard, torch.full_like(shard, shard[0].item())) + assert 0 <= shard[0].item() < Utils.world_size From b164d7408f0ef5091861f245bf15063f304121bc Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Sun, 16 Aug 2026 15:24:23 +0000 Subject: [PATCH 051/101] cp: build: bump transformer-engine to 2.17.1 (#6433) Signed-off-by: svcnemo-autobot --- pyproject.toml | 4 ++-- uv.lock | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 26c0fbc33be..954c699107f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -208,7 +208,7 @@ requires-dist = [] [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.18.0+e7c550c5" +version = "2.17.1+4329ff84" requires-dist = [ # Cap below 2.14: pydantic 2.14 breaks langchain_core's module-level # RunnablePassthrough() instantiation, which is imported transitively in @@ -234,7 +234,7 @@ flash_mla = [ { git = "https://github.com/deepseek-ai/FlashMLA", rev = "nv_dev" }, ] deep_gemm = { git = "https://github.com/deepseek-ai/DeepGEMM.git", rev = "714dd1a4a980f7937a74343d19a8eba4fe321480" } -transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "e7c550c5f80636cf841a8204b1d6f85a5f3f28b7" } +transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "4329ff84bfbdaa778a33cba02a15fb0807c64689" } nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "e3935393a290aed1822af52139b4b8ee270fed1f" } emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" } fast-hadamard-transform = { git = "https://github.com/Dao-AILab/fast-hadamard-transform.git", rev = "f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" } diff --git a/uv.lock b/uv.lock index 3fc0c7b0308..f521105d3bc 100644 --- a/uv.lock +++ b/uv.lock @@ -32,7 +32,7 @@ version = "1.0.0+9edee0c" [[manifest.dependency-metadata]] name = "transformer-engine" -version = "2.18.0+e7c550c5" +version = "2.17.1+4329ff84" requires-dist = ["pydantic<2.14", "importlib-metadata>=1.0", "packaging", "torch>=2.1", "einops", "onnxscript", "onnx", "nvdlfw-inspect"] [[package]] @@ -2348,7 +2348,7 @@ requires-dist = [ { name = "tiktoken", marker = "extra == 'training'" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", marker = "extra == 'dev'" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=e7c550c5f80636cf841a8204b1d6f85a5f3f28b7" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689" }, { name = "transformers", marker = "extra == 'mlm'" }, { name = "transformers", marker = "extra == 'training'" }, { name = "wandb", marker = "extra == 'mlm'" }, @@ -5253,8 +5253,8 @@ wheels = [ [[package]] name = "transformer-engine" -version = "2.18.0+e7c550c5" -source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=e7c550c5f80636cf841a8204b1d6f85a5f3f28b7#e7c550c5f80636cf841a8204b1d6f85a5f3f28b7" } +version = "2.17.1+4329ff84" +source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689#4329ff84bfbdaa778a33cba02a15fb0807c64689" } dependencies = [ { name = "einops" }, { name = "importlib-metadata" }, From c374e5b65e32b498763e3666b40490314b554cb7 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Sun, 16 Aug 2026 23:23:14 +0000 Subject: [PATCH 052/101] Autotune RL lag (#4127) Signed-off-by: Teodor-Dumitru Ene --- megatron/rl/README.md | 72 ++++++++++++++++++- megatron/rl/agent/rollout_pipeline.py | 13 ++-- megatron/rl/inference/megatron.py | 39 ++++++++-- megatron/rl/rl_utils.py | 6 ++ megatron/training/arguments.py | 43 +++++++++-- .../gpt_grpo_basic_function/model_config.yaml | 1 + .../model_config.yaml | 1 + .../model_config.yaml | 1 + .../model_config.yaml | 1 + tests/unit_tests/rl/test_rl_utils.py | 21 ++++++ 10 files changed, 181 insertions(+), 17 deletions(-) diff --git a/megatron/rl/README.md b/megatron/rl/README.md index 0c86c151d97..e82505fd181 100644 --- a/megatron/rl/README.md +++ b/megatron/rl/README.md @@ -37,7 +37,77 @@ Below we describe the different conceptual components and how they divide respon - Provides the endpoint the `AgenticEnvironment` uses to run `.generate(prompt, **generation_args)`. - Can take many forms (e.g. Megatron, OpenAI, HF) and supports many configuration options. +## Off-Policy Generation + +Pure on-policy RL runs rollout generation and training sequentially: each training step has to wait for fresh rollouts from the current policy. This significantly underutilizes hardware. + +Off-policy generation overlaps the two, whether that's in the time dimension (by switching resources between training and inference), in the space dimension (by dividing resources between training and inference), or both. +Note that whether we are overlapping in the time dimension (like Python's asyncio) or the space dimension (like Python's threading), the core behavior is almost the same. +While training is busy with step *N*, the inference engine keeps generating rollouts that will be consumed by steps *N+1*, *N+2*, … The rollouts come from slightly older policy weights, so they are *stale* by some amount. We call this the **collection lag**: + +> A lag of *L* means the rollouts consumed by the current training step were generated *L* training steps earlier, i.e. by a policy that is *L* updates behind the one we're training now. + +Higher lag means more overlap (throughput) at the cost of more staleness (sample quality / stability). Finding the right point on that tradeoff is the main thing this knob exists for. + +### Turning it on + +Off-policy generation requires `--rl-partial-rollouts`. The lag itself is set with `--rl-generation-lag`: + +``` +--rl-partial-rollouts --rl-generation-lag 2 +``` + +Useful values: +- `L = -1`: the minimum — a single unit of generation work in flight at a time (effectively serial, mostly useful for testing; only distinguishable from `L = 0` at submission granularity `G` or `R`, see below). +- `L = 0`: fill exactly one training step's worth of generation capacity. The inference engine is never ahead of the trainer. +- `L = 1, 2, …`: the engine runs this many steps ahead. Higher L = more throughput, more staleness. Fractional values are allowed. + +### Autotuning + +If you omit `--rl-generation-lag`, the system sets it automatically, and prints what it picked at startup. +Alternatively, if you provide your own `--rl-generation-lag`, a startup warning will inform you whether *your* `--rl-generation-lag` oversubscribes the inference engine — +i.e. whether you're asking for more in-flight work than the engine can actually serve. +This is why it is important to allow `L` to turn negative. Negative values of `L` occur when there are so few inference resources that even fully-synchronous RL will oversubscribe the engine. +Given such information, users can utilize the formula below to appropriately scale their inference resources: + +``` +max_effective_lag = DP * engine.max_requests / (G * P) - 1 +``` + +where `DP` is the inference data-parallel size, `G` is `--grpo-group-size`, `P` is `--grpo-prompts-per-step`, and `engine.max_requests` is the per-rank request capacity of the inference engine (derived from KV-cache memory, or overridden via `--inference-dynamic-batching-max-requests`). + +### Submission and consumption granularity + +Two further flags shape the overlap by picking the pipeline's *units of work*: a single rollout (`R`), a prompt group (`G` — one prompt's `--grpo-group-size` samples), or a whole training batch (`B` — all `P` groups of a step): + +- `--rl-submission-granularity {R,G,B}` sets the unit the engine's capacity is parceled out in. A unit occupies its share of the engine from admission until it completes; only then is that capacity released for new work. +- `--rl-consumption-granularity {G,B}` sets the unit the trainer waits for: `G` consumes prompt groups in completion order, `B` consumes whole batches in submission order. Consumption can't be finer than submission (`B`-submission forces `B`-consumption), and per-rollout consumption is impossible under GRPO — advantages are relative to the group's rewards, so the group is the atomic training unit. + +Writing modes as submission/consumption, the supported modes are `B/B` (the default), `G/G`, `G/B`, `R/G`, and `R/B`. Their behavior is governed by one fact: a unit is only finished when its *slowest* rollout is finished. With long-tailed generation lengths, a unit's expected completion time therefore exceeds the mean rollout time by a *tail factor* `τ = E[slowest rollout in unit] / E[rollout]`, which grows with unit size: `τ_rollout = 1 ≤ τ_group ≤ τ_batch`. Two laws follow (for an engine right-sized to the lag gate — which is exactly what autotuning arranges): + +- **Utilization is set by submission alone:** `U = 1 / τ_submit`. A coarse unit holds engine slots while its stragglers finish, idling the engine a fraction `1 − 1/τ_submit` of the time; finer submission backfills freed slots immediately, and `R`-submission keeps the engine fully busy. +- **Staleness is set by the ratio:** mean first-token staleness `≈ (τ_consume / τ_submit) · (1 + L)`. Consuming coarser than you submit multiplies staleness by every tail factor spanned in between, while buying *zero* utilization — so at fixed submission the matched mode dominates (`G/G` over `G/B`, `R/G` over `R/B`). + +Consumption granularity also sets the staleness *distribution*. Under `G`-consumption a group's staleness scales with its own completion time — long groups always train proportionally staler. `B/B` is the special corner: releasing capacity only at consumption phase-locks generation to the training clock, so every group in every step trains at *exactly* `L` versions of staleness — deterministic, and one version below the `(1 + L)` that release-on-completion modes pay. At `L = 0`, `B/B` is fully synchronous on-policy training. Its price is the idle engine (`1 − 1/τ_batch`), which it uniquely does *not* convert into run-ahead staleness. + +Rules of thumb: +- `B/B` (default): the on-policy corner — smallest and deterministic staleness, lowest throughput; best for debugging, reproducibility, and ablations where lag-variance would confound. +- `G/G`: `τ_batch/τ_group`× the utilization of `B/B`, for one extra version of mean staleness paid as a length-correlated spread. +- `R/G`: the only route to a fully-utilized engine, at `τ_group`× the staleness of `G/G`. On its own a par trade; with importance-sampling correction (`--rl-inference-logprobs-is-correction`) it becomes a win — corrected staleness is a statistical cost, while an idle engine is unrecoverable wall-clock. +- `G/B`, `R/B`: dominated — same utilization as their matched mode with strictly more staleness; use only if you specifically need batch-ordered consumption together with streaming submission. + +### How it maps to the internal knobs + +Under the hood, the lag controls how many submission units ("tasks") are in flight simultaneously: + +``` +tasks = max(1, round((L + 1) * P * G / unit)) # in-flight units of generation work +rollouts_in_flight = tasks * unit # total concurrent inference requests +``` + +where `unit` is the number of rollouts per submission unit (`P * G` for `B`, `G` for `G`, `1` for `R`). You generally don't need to think about this — pick an `L` (or let autotune pick one), or set `--rl-max-inflight-requests` to cap the total in-flight rollout requests directly (`(L + 1) * P * G`; mutually exclusive with `--rl-generation-lag`) — then check the gate capacity in the startup log to confirm the resulting task count is what you expected. + ## Examples See `examples/rl` for demonstrations of: - Custom `InferenceInterface` endpoints -- Example `Agents` \ No newline at end of file +- Example `Agents` diff --git a/megatron/rl/agent/rollout_pipeline.py b/megatron/rl/agent/rollout_pipeline.py index 35712ce6992..c111fa29258 100644 --- a/megatron/rl/agent/rollout_pipeline.py +++ b/megatron/rl/agent/rollout_pipeline.py @@ -225,15 +225,16 @@ def __init__( self.gran_policy = _GranularityConfig.from_request( request, [allocation.num_groups for allocation in self.allocations] ) + # Lag may be fractional or negative (>= -1): clamp and round the slot math at granularity. self.gate = _SubmissionGate( - capacity=parallel_generation_tasks - * self.gran_policy.units_per_batch(), + capacity=max( + 1, round(parallel_generation_tasks * self.gran_policy.units_per_batch()) + ), submission=self.gran_policy.submission, ) - self.num_infer_workers = ( - parallel_generation_tasks - * self.gran_policy.num_groups_per_batch - * request.rollouts_per_group + rollouts_per_batch = self.gran_policy.num_groups_per_batch * request.rollouts_per_group + self.num_infer_workers = self.gate.capacity * ( + rollouts_per_batch // self.gran_policy.units_per_batch() ) if not request.streaming: self.num_infer_workers = min( diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index 6f6610f6289..fa5116f1e5a 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -19,8 +19,9 @@ from megatron.core.inference.inference_client import InferenceClient from megatron.core.inference.inference_request import FinishedRequestRecord from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.utils import log_single_rank +from megatron.core.utils import get_pg_size, log_single_rank from megatron.training.global_vars import get_args, get_tokenizer +from megatron.training.utils import print_rank_0 from ..inference.inference_interface import ( InferenceRequest, @@ -34,6 +35,7 @@ logger = logging.getLogger(__name__) logging.getLogger("httpx").setLevel(logging.WARNING) + class MegatronLocal(InferenceServer, ReturnsTokens, ReturnsRaw): """Interface to use MCoreEngine directly as an inference engine.""" @@ -108,6 +110,37 @@ async def launch(cls, model: GPTModel, **kwargs): inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(model=model) inference_engine.local_metadata_ledger_enabled = True + if args.rl_partial_rollouts: + # Resolve args.rl_generation_lag against the engine's request capacity: + # autotune it when unset, or report how the requested lag compares. + dp_size = get_pg_size(inference_engine.pg_collection.dp) + max_requests = inference_engine.context.max_requests + G = args.grpo_group_size + P = args.grpo_prompts_per_step + max_effective_groups = dp_size * max_requests // G + max_effective_lag = max_effective_groups / P - 1 + if args.rl_generation_lag is None: + args.rl_generation_lag = max_effective_lag + print_rank_0( + f"Autotuned rl-generation-lag={max_effective_lag:.2f} " + f"(DP={dp_size}, max_requests={max_requests}, G={G}, P={P}).") + else: + print_rank_0( + f"Using rl-generation-lag={args.rl_generation_lag} " + f"(max effective lag={max_effective_lag:.2f}; " + f"DP={dp_size}, max_requests={max_requests}, G={G}, P={P}).") + groups_in_flight = (args.rl_generation_lag + 1) * P + if groups_in_flight > max_effective_groups + 1e-6: + print_rank_0( + f"WARNING: {groups_in_flight:.1f} groups in flight oversubscribes the " + f"inference engine (max effective lag is {max_effective_lag:.2f}). " + f"Additional run-ahead beyond that point has no benefit.") + if max_effective_lag < 0: + print_rank_0( + f"WARNING: max effective lag is {max_effective_lag:.2f} (negative) — the " + f"inference engine cannot hold even one training step's worth of rollouts " + f"({max_effective_groups} groups < P={P}). Even fully-synchronous GRPO would " + f"oversubscribe. Consider scaling up inference resources.") dp_addr = await inference_engine.start_listening_to_data_parallel_coordinator( inference_coordinator_port=41521, launch_inference_coordinator=True, ) @@ -137,9 +170,7 @@ async def launch(cls, model: GPTModel, **kwargs): ) concurrency_limit = ( - args.grpo_prompts_per_step - * args.grpo_group_size - * (args.rl_generation_lag + 1) + get_pg_size(inference_engine.pg_collection.dp) * inference_engine.context.max_requests ) custom_limits = httpx.Limits( max_connections=concurrency_limit, diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 6c77cf13694..cd233eeab86 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -648,6 +648,12 @@ def get_rollout_generator( request=request, parallel_generation_tasks=generation_lag + 1, ) + log_single_rank( + logger, + logging.INFO, + f"Rollout pipeline gate: {_ROLLOUT_PIPELINE.gate.capacity} " + f"{submission_granularity}-granularity slots (generation lag {generation_lag}).", + ) _ROLLOUT_GENERATOR = _ROLLOUT_PIPELINE.run() return _ROLLOUT_GENERATOR diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 1d49434b24e..0e0c6664860 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -526,9 +526,29 @@ def validate_args(args, defaults={}): submit_rollouts_at_rollout_granularity = ( args.rl_submission_granularity == "R" ) - if args.rl_generation_lag > 0: - assert args.rl_partial_rollouts, \ - "--rl-generation-lag requires --rl-partial-rollouts." + if args.rl_max_inflight_requests is not None: + requests_per_batch = args.grpo_prompts_per_step * args.grpo_group_size + assert args.rl_generation_lag is None, \ + "--rl-generation-lag and --rl-max-inflight-requests are mutually exclusive." + assert args.rl_max_inflight_requests >= 1, \ + f"--rl-max-inflight-requests ({args.rl_max_inflight_requests}) must be >= 1." + if args.rl_max_inflight_requests > requests_per_batch: + assert args.rl_partial_rollouts, \ + f"--rl-max-inflight-requests above one training batch " \ + f"({requests_per_batch} requests) requires --rl-partial-rollouts." + # Total in-flight requests = (lag + 1) trainer batches of P * G requests each. + args.rl_generation_lag = args.rl_max_inflight_requests / requests_per_batch - 1 + if args.rl_generation_lag is None: + # With --rl-partial-rollouts the lag is autotuned from engine capacity + # at inference launch; otherwise generation is fully synchronous. + if not args.rl_partial_rollouts: + args.rl_generation_lag = 0 + else: + assert args.rl_generation_lag >= -1, \ + f"--rl-generation-lag ({args.rl_generation_lag}) must be >= -1." + if args.rl_generation_lag > 0: + assert args.rl_partial_rollouts, \ + "--rl-generation-lag requires --rl-partial-rollouts." if submit_rollouts_at_rollout_granularity: assert ( args.rl_partial_rollouts @@ -2595,10 +2615,21 @@ def _add_rl_args(parser): help="Number of GRPO groups (G in the paper).") group.add_argument('--grpo-group-size', type=int, default=2, help="Number of samples per a GRPO group.") - group.add_argument('--rl-generation-lag', type=int, default=0, - help='Number of trainer batches of rollout generation lag to allow. ' + group.add_argument('--rl-generation-lag', type=float, default=None, + help='Number of trainer batches of rollout generation lag to allow ' 'The number of in-flight trainer batches is this value plus one. ' - 'Requires --rl-partial-rollouts when greater than 0.') + 'May be fractional or negative; the minimum of -1 keeps a single unit ' + 'of generation work in flight. If omitted, the lag is autotuned to the ' + 'inference engine\'s request capacity when --rl-partial-rollouts is ' + 'set, and is 0 otherwise. ' + 'Requires --rl-partial-rollouts when greater than 0. ' + 'Mutually exclusive with --rl-max-inflight-requests.') + group.add_argument('--rl-max-inflight-requests', type=int, default=None, + help='Maximum number of inference requests RL generation may keep inflight: ' + 'equivalent to (--rl-generation-lag + 1) training batches ' + 'of grpo_prompts_per_step * grpo_group_size requests each. ' + 'Requires --rl-partial-rollouts when above one training batch; ' + 'mutually exclusive with --rl-generation-lag.') # TODO: Refactor these string literals back to an enum after the megatron.training refactor. group.add_argument('--rl-submission-granularity', type=str, default="B", diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_basic_function/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_basic_function/model_config.yaml index f052a0bca75..c90b99d531c 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_basic_function/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_basic_function/model_config.yaml @@ -86,6 +86,7 @@ MODEL_ARGS: --grpo-entropy-term-weight: 0.0 --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_basic_function/env_config.yaml --rl-partial-rollouts: true + --rl-generation-lag: 0 --use-checkpoint-args: true --dist-ckpt-strictness: log_unexpected --perform-rl-step: true diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/model_config.yaml index 7ef22862c00..1c2f94c37c1 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/model_config.yaml @@ -62,6 +62,7 @@ MODEL_ARGS: --grpo-entropy-term-weight: 0.0 --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/env_config.yaml --rl-partial-rollouts: true + --rl-generation-lag: 0 --lr: 0.000001 --lr-warmup-samples: 0 --clip-grad: 1.0 diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest_github/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest_github/model_config.yaml index e12eea86625..5cd530ed79e 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest_github/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest_github/model_config.yaml @@ -62,6 +62,7 @@ MODEL_ARGS: --grpo-entropy-term-weight: 0.0 --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest_github/env_config.yaml --rl-partial-rollouts: true + --rl-generation-lag: 0 --lr: 0.000001 --lr-warmup-samples: 0 --clip-grad: 1.0 diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml index 8b4a14b9836..b2ffb9633f5 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml @@ -84,6 +84,7 @@ MODEL_ARGS: --grpo-entropy-term-weight: 0.0 --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/env_config.yaml --rl-partial-rollouts: true + --rl-generation-lag: 0 --use-checkpoint-args: true --dist-ckpt-strictness: log_unexpected --perform-rl-step: true diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index dac3f963d4f..cb6103a6d9c 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -256,6 +256,7 @@ def __init__(self, agent, request, parallel_generation_tasks): captured["agent"] = agent captured["request"] = request captured["parallel_generation_tasks"] = parallel_generation_tasks + self.gate = SimpleNamespace(capacity=parallel_generation_tasks) def run(self): return rollout_generator @@ -296,6 +297,26 @@ def run(self): "--rl-generation-lag requires --rl-partial-rollouts", id="lag_requires_partial_rollouts", ), + pytest.param({"rl_generation_lag": -2}, "must be >= -1", id="lag_below_minimum"), + pytest.param( + { + "rl_generation_lag": 1, + "rl_max_inflight_requests": 64, + "rl_partial_rollouts": True, + }, + "mutually exclusive", + id="lag_and_max_inflight_exclusive", + ), + pytest.param( + {"rl_max_inflight_requests": 128}, + "requires --rl-partial-rollouts", + id="max_inflight_above_batch_requires_partial_rollouts", + ), + pytest.param( + {"rl_max_inflight_requests": 0, "rl_partial_rollouts": True}, + "must be >= 1", + id="max_inflight_below_minimum", + ), pytest.param( {"rl_submission_granularity": "R"}, "Rollout submission granularity requires streaming grouped rollouts", From d262dbed3d48a0b6a22b1f0ddc58c2cede2b9149 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Mon, 17 Aug 2026 03:06:04 +0000 Subject: [PATCH 053/101] fix(gtp+mtp): zero the placeholder wgrad when zero_out_wgrad is set (#6565) Signed-off-by: Shiqing Fan --- .../generalized_tensor_parallelism.py | 14 ++- .../test_gtp_mtp.py | 106 ++++++++++++++++++ 2 files changed, 115 insertions(+), 5 deletions(-) diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index ff84be77a38..f423365fd2b 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -1765,15 +1765,19 @@ def register_grad_accum_hook( def _handle_megatron_grad_accum(param): """Handle megatron DDP and gradient-accumulation fusion. - Do NOT set param.grad before calling the hook — the hook checks param.grad and would - accumulate it into main_grad if zero_out_wgrad is True, corrupting it with a dummy. - Returns a cached dummy wgrad; sync callers use it as the graph-safe grad, async drains - discard it. + discard it. It is zeroed when DDP will accumulate it — see below. """ if hasattr(param, "grad_added_to_main_grad"): param.grad_added_to_main_grad = True - dummy_grad = get_dummy_wgrad(list(param.main_grad.shape), param.dtype) + # This dummy becomes param.grad, and DDP accumulates it when zero_out_wgrad is set + # (DistributedDataParallel._make_backward_post_hook in distributed_data_parallel.py). + # get_dummy_wgrad returns a SHARED reused buffer, so it must be zeroed in that case + # or it injects whatever it last held into the gradient. Same pairing as layers.py. + if getattr(param, "zero_out_wgrad", False): + dummy_grad = get_dummy_wgrad(list(param.main_grad.shape), param.dtype, zero=True) + else: + dummy_grad = get_dummy_wgrad(list(param.main_grad.shape), param.dtype) if getattr(param, "_grad_accum_hook", None) is not None: param._grad_accum_hook() diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_mtp.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_mtp.py index fae23a29c66..c262446474b 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_mtp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_mtp.py @@ -511,6 +511,100 @@ def counting_register(self, param, *a, **k): ) +# Value written into the shared dummy-wgrad buffer to make the leak observable. Any non-zero +# value works; a large one keeps it far outside the range of a real gradient. +DUMMY_POISON = 12345.0 + + +def _worker_dummy_wgrad_not_leaked(rank, world_size, port, repeated_layer=False): + """A GTP weight's gradient must not depend on the shared dummy-wgrad buffer's contents. + + GTP returns a PLACEHOLDER from backward, not a gradient -- the real wgrad already went into + main_grad via reduce-scatter. But it becomes ``param.grad``, and DDP adds ``param.grad`` into + main_grad when ``zero_out_wgrad`` is set (``_make_backward_post_hook``), which MTP does to the + shared embedding weight. ``get_dummy_wgrad`` hands back ONE reused buffer, and the shared + weight gets N = (1 main pass + mtp_num_layers) x gradient-accumulation steps backward passes + per iteration, so:: + + want: S1 + S2 + ... + SN Si = call i's reduce-scattered gradient + got: S1 + S2 + ... + SN + N*D D = whatever that shared buffer last held + + D is near-zero while the buffer is fresh and real data once traffic has passed through it -- + hence the original failure needing both MTP and gradient accumulation, and vanishing whenever + an extra allocation shifted the buffer. This test POISONS it with a known value instead, and + requires the gradients to come out bit-identical. It runs a single step, so N = 1 + 2 = 3 and + reverting the fix shifts the gradient by ~3x the poison. + """ + from transformer_engine.pytorch.module.base import get_dummy_wgrad + + from megatron.core import parallel_state as ps + from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + GTP_CONFIG, + reset_gtp_state, + ) + from megatron.core.tensor_parallel.gtp_api import classify_gtp_remat_chains + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + def _train_one_step_with_poison(poison): + """Build a fresh MTP model, poison the shared dummy, run fwd+bwd, return {name: grad}.""" + reset_gtp_state() # reset_gtp_globals is an autouse FIXTURE; this is what it defers to + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=world_size + ) + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + + model = _build_mtp_gpt_model(repeated_layer, moe=False) + classify_gtp_remat_chains([model]) + for p in model.parameters(): + p.main_grad = torch.zeros(p.shape, dtype=torch.float32, device='cuda') + + tagged = [n for n, p in model.named_parameters() if getattr(p, "zero_out_wgrad", False)] + # Fail loudly rather than pass vacuously: with nothing tagged, DDP never reads the + # placeholder and this test would be green no matter what GTP returns. + assert tagged, ( + "no parameter has zero_out_wgrad set -- MTP should tag the shared embedding weight. " + "Without it this test cannot observe the placeholder at all." + ) + # get_dummy_wgrad keys on main_grad's shape, so this dirties the exact buffer GTP will + # hand back from _handle_megatron_grad_accum. + for _, p in model.named_parameters(): + if getattr(p, "zero_out_wgrad", False): + get_dummy_wgrad(list(p.main_grad.shape), p.dtype).fill_(poison) + + _forward_backward(model) + + # The model here is unwrapped, so replay DDP's accumulation rule explicitly -- see + # DistributedDataParallel._make_backward_post_hook. This is what makes the test able to + # observe the placeholder at all: drop it and the test passes even with the bug present. + for p in model.parameters(): + if p.grad is not None and ( + not getattr(p, "grad_added_to_main_grad", False) + or getattr(p, "zero_out_wgrad", False) + ): + p.main_grad.add_(p.grad.to(p.main_grad.dtype)) + + return _gathered_main_grads(model), tagged + + saved_pad = GTP_CONFIG.pad_for_alignment + try: + # Shards stay exactly 1/gtp of the weight so _gathered_main_grads reconstructs it. + GTP_CONFIG.pad_for_alignment = 0 + clean, tagged = _train_one_step_with_poison(0.0) + dirty, _ = _train_one_step_with_poison(DUMMY_POISON) + finally: + GTP_CONFIG.pad_for_alignment = saved_pad + + for name in clean: + delta = (dirty[name] - clean[name]).abs().max().item() + assert delta == 0.0, ( + f"{name}: gradient changed by {delta} when the shared dummy-wgrad buffer was " + f"poisoned with {DUMMY_POISON}. GTP's placeholder grad leaked into the gradient; it " + f"must be zeroed whenever zero_out_wgrad is set (tagged here: {tagged})." + ) + + class TestGTPMTP: @pytest.mark.parametrize("moe", [False, True], ids=["dense", "moe"]) @pytest.mark.parametrize("repeated_layer", [False, True]) @@ -562,3 +656,15 @@ def test_mtp_shared_weight_ddp_grad_ready_counts(self, repeated_layer): if torch.cuda.device_count() < 4: pytest.skip("Requires 4 CUDA devices") _run_distributed(_worker_ddp_grad_ready_counts, 4, repeated_layer) + + @pytest.mark.parametrize("repeated_layer", [False, True]) + def test_mtp_dummy_wgrad_does_not_leak_into_grad(self, repeated_layer): + """MTP sets zero_out_wgrad on the shared embedding, which makes DDP ADD param.grad. + + GTP's placeholder grad must therefore be zeroed, or the shared dummy buffer's stale + contents land in main_grad -- silently at first, then as a NaN once enough backward + passes have dirtied it. + """ + if torch.cuda.device_count() < 4: + pytest.skip("Requires 4 CUDA devices") + _run_distributed(_worker_dummy_wgrad_not_leaked, 4, repeated_layer) From 34c3741f564c6793aac969c0a1204a85cb4efe70 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 17 Aug 2026 10:05:11 +0000 Subject: [PATCH 054/101] Enable context parallelism for RL training (#5882) Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Jorge Albericio --- .../core/extensions/transformer_engine.py | 1 + megatron/core/packed_seq_params.py | 1 + megatron/rl/rl_utils.py | 177 +++++++++++++++-- megatron/rl/sequence_packing_utils.py | 123 +++++++++--- tests/unit_tests/rl/test_context_parallel.py | 187 ++++++++++++++++++ tests/unit_tests/rl/test_rl_utils.py | 118 +++++------ .../rl/test_sequence_packing_utils.py | 156 ++++++++++++++- train_rl.py | 10 +- 8 files changed, 670 insertions(+), 103 deletions(-) create mode 100644 tests/unit_tests/rl/test_context_parallel.py diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 6df3049d870..9393ad825c3 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2175,6 +2175,7 @@ def __init__( self.kept_packed_seq_params.discard("total_tokens") self.kept_packed_seq_params.discard("seq_idx") self.kept_packed_seq_params.discard("tokens_per_sample") + self.kept_packed_seq_params.discard("cp_scatter_cache") if get_te_version() < PkgVersion("2.2.0"): self.kept_packed_seq_params.discard("pad_between_seqs") diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 1c26af56244..52e7dbaa516 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -26,6 +26,7 @@ class PackedSeqParams: seq_idx: Tensor = None tokens_per_sample: int = None pad_between_seqs: bool = None + cp_scatter_cache: object = None def __post_init__(self): """Pre-compute seq_idx for Mamba mixer CUDA graph compatibility. diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index cd233eeab86..302d8f86e89 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -21,6 +21,14 @@ import numpy as np import torch import torch.distributed as dist +import torch.distributed.nn + +try: + import transformer_engine_torch as tex + + HAVE_TEX = True +except ImportError: + HAVE_TEX = False import yaml from torch.utils.data import DataLoader, TensorDataset from torch.utils.tensorboard import SummaryWriter @@ -831,6 +839,121 @@ def selective_log_softmax(logits, index): return per_token_logps +@dataclass +class _CPScatterCache: + """Per-bin CP partition artifacts, derived once from (packed_seq_params, cp_group). + + Attributes: + seq_len: Full (pre-partition) sequence length of the bin. + cp_group: Process group the partition was derived for. + cp_packed_seq_params: Copy of the bin's params annotated for the local CP forward. + partition_index: Full-sequence positions owned by this rank [seq_len // cp_size]. + inverse_gather_perm: Used to restore full-sequence order after the all-gather [seq_len]. + """ + + seq_len: int + cp_group: torch.distributed.ProcessGroup + cp_packed_seq_params: PackedSeqParams + partition_index: torch.Tensor + inverse_gather_perm: torch.Tensor + + +def _build_cp_scatter_cache( + packed_seq_params: PackedSeqParams, seq_len: int, cp_group: torch.distributed.ProcessGroup +) -> _CPScatterCache: + """Derive the CP partition for one bin: annotated params copy + partition indices.""" + assert HAVE_TEX, ( + "Context-parallel RL logprobs require Transformer Engine to compute THD indices." + ) + cp_rank = torch.distributed.get_rank(cp_group) + cp_size = torch.distributed.get_world_size(cp_group) + + # TE's THD context parallelism splits every sequence into 2*cp_size chunks to balance load: + # rank r owns the chunks at `r` and at `2*cp_size-1-r`. + assert seq_len % (2 * cp_size) == 0, ( + f"Sequence length {seq_len} must be divisible by 2*context_parallel_size." + ) + + # Shallow-copy so we do not mutate the main sequence params. + cp_packed_seq_params = copy.copy(packed_seq_params) + cp_packed_seq_params.cp_scatter_cache = None + cp_packed_seq_params.cp_group = cp_group + cp_packed_seq_params.local_cp_size = cp_size + # THD CP expects to have padded seqlens; if we don't need to pad, use the unpadded version. + if cp_packed_seq_params.cu_seqlens_q_padded is None: + cp_packed_seq_params.cu_seqlens_q_padded = cp_packed_seq_params.cu_seqlens_q + if cp_packed_seq_params.cu_seqlens_kv_padded is None: + cp_packed_seq_params.cu_seqlens_kv_padded = cp_packed_seq_params.cu_seqlens_kv + + # Rank-major concatenation of every rank's partition indices: + # a permutation of range(seq_len) whose position j is owned by rank j // local_len. + gather_perm = torch.cat( + [ + tex.thd_get_partitioned_indices( + cp_packed_seq_params.cu_seqlens_q_padded, seq_len, cp_size, rank + ) + for rank in range(cp_size) + ] + ).long() + local_len = seq_len // cp_size + partition_index = gather_perm[cp_rank * local_len : (cp_rank + 1) * local_len] + + return _CPScatterCache( + seq_len=seq_len, + cp_group=cp_group, + cp_packed_seq_params=cp_packed_seq_params, + partition_index=partition_index, + inverse_gather_perm=torch.argsort(gather_perm), + ) + + +def _scatter_for_context_parallel( + tokens: torch.Tensor, + position_ids: torch.Tensor, + packed_seq_params: PackedSeqParams, + cp_group: torch.distributed.ProcessGroup, +) -> tuple: + """Prepare local inputs for one context-parallel rank. + + Formats the bins to allow for the use of Transformer Engine's `thd_get_partitioned_indices`. + + Args: + tokens: Full token tensor [batch, seq_len]. + position_ids: Full position-id tensor [batch, seq_len]. + packed_seq_params: PackedSeqParams for the full bin. All slots are divisible by 2*cp_size. + cp_group: Context-parallel process group. + + Returns: + (local_tokens, local_position_ids, cp_packed_seq_params, local_labels, cp_scatter) + where every local tensor has sequence length `seq_len // cp_size` and `cp_scatter` + is this bin's _CPScatterCache. + """ + seq_len = tokens.shape[1] + # The partition depends only on the bin layout and CP topology. Derive it once per bin. + cp_scatter = packed_seq_params.cp_scatter_cache + if cp_scatter is None: + cp_scatter = _build_cp_scatter_cache(packed_seq_params, seq_len, cp_group) + packed_seq_params.cp_scatter_cache = cp_scatter + assert cp_scatter.cp_group is cp_group and cp_scatter.seq_len == seq_len, ( + "Cached CP partition was derived for a different (cp_group, seq_len)." + ) + + # We do not need logprobs for the n+1 token. + tokens_shifted = torch.cat([tokens[:, 1:], tokens[:, -1:]], dim=1) + + local_tokens = tokens.index_select(1, cp_scatter.partition_index) + local_position_ids = position_ids.index_select(1, cp_scatter.partition_index) + local_labels = tokens_shifted.index_select(1, cp_scatter.partition_index) + + return ( + local_tokens, + local_position_ids, + cp_scatter.cp_packed_seq_params, + local_labels, + cp_scatter, + ) + + def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=False, packed_seq_params=None): """Get sequence logprobs from their token ids. @@ -838,17 +961,14 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa model: model to predict with. tokens: inputs for which we want to get logprobs. position_ids: position ids that come with tokens. - attention_mask: attention mask that comes with tokens. no_grad: whether to run in no_grad mode. + sequence_packing: whether tokens holds packed bins. packed_seq_params: Optional PackedSeqParams for sequence packing with TE. When provided with qkv_format='thd', the input tokens are sliced to remove padding before the forward pass, and outputs are padded back. - packed_seq_len: Optional length of the packed sequence (excluding padding). - Required when packed_seq_params is provided to avoid CPU-GPU synchronization. Returns: - Logprobs of input sequences. - + Next-token logprobs from `forward()` over the input sequences [batch, seq_len - 1]. """ args = get_args() @@ -866,15 +986,20 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa ) else: cu_seqlens = torch.tensor([0, tokens.shape[1]], dtype=torch.int32, device=tokens.device) + # Make sure to omit `total_tokens` to prevent `seq_idx` from being auto-computed. + # That would cause a sequence packing kernel to be incorrectly used. packed_seq_params = PackedSeqParams( qkv_format='thd', cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, max_seqlen_q=tokens.shape[1], max_seqlen_kv=tokens.shape[1], - total_tokens=tokens.shape[1], + pad_between_seqs=False, ) + pg_collection = get_attr_wrapped_model(model, "pg_collection") + cp_group = pg_collection.cp + cp_size = torch.distributed.get_world_size(cp_group) if cp_group is not None else 1 nvtx_range = get_nvtx_range() with nvtx_range("rl/get-logprobs", time=True): @@ -887,28 +1012,52 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa flash_decode = model.config.flash_decode set_model_config_attribute(model, "flash_decode", False) fp32_output = not (args.fp16 or args.bf16) - with torch.no_grad() if no_grad else nullcontext(): - logits_or_hidden_states = model( + + if cp_size > 1: + # Scatter: each rank processes seq_len // cp_size tokens. + tokens_in, position_ids_in, packed_seq_params_in, local_labels, cp_scatter = ( + _scatter_for_context_parallel(tokens, position_ids, packed_seq_params, cp_group) + ) + else: + tokens_in, position_ids_in, packed_seq_params_in = ( tokens, position_ids, + packed_seq_params, + ) + + with torch.no_grad() if no_grad else nullcontext(): + logits_or_hidden_states = model( + tokens_in, + position_ids_in, attention_mask_for_forward, - packed_seq_params=packed_seq_params, + packed_seq_params=packed_seq_params_in, runtime_gather_output=True, fp32_output=fp32_output, ) + set_model_config_attribute(model, "flash_decode", flash_decode) - pg_collection = get_attr_wrapped_model(model, "pg_collection") pp_group = pg_collection.pp if not is_pp_last_stage(pp_group): return logits_or_hidden_states - else: - logits = logits_or_hidden_states - with nvtx_range("rl/log-softmax", time=True): + + logits = logits_or_hidden_states + with nvtx_range("rl/log-softmax", time=True): + if cp_size > 1: + local_logprobs = selective_log_softmax(logits, local_labels) + # Differentiable all-gather so training pass can backprop through the reassembly. + # With no-grad, this acts as a plain all_gather. + gathered = torch.distributed.nn.functional.all_gather( + local_logprobs.contiguous(), group=cp_scatter.cp_group + ) + full = torch.cat(gathered, dim=1).index_select(1, cp_scatter.inverse_gather_perm) + # Drop the dummy boundary position appended by the label shift. + logprobs = full[:, :-1] + else: # We do not need logprobs for the n+1 token. logprobs = selective_log_softmax(logits[:, :-1, :], tokens[:, 1:]) - return logprobs + return logprobs def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[int]]) -> np.ndarray: diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index bc3f06ed587..bcd066207c2 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -6,8 +6,8 @@ from typing import List, Dict, Any, Tuple, Optional from torch.utils.data import DataLoader, TensorDataset from dataclasses import dataclass, field -from megatron.core.utils import log_single_rank -from megatron.training.global_vars import get_args, get_tokenizer +from megatron.core.utils import log_single_rank, round_up_to_nearest_multiple +from megatron.training.global_vars import get_tokenizer from megatron.training.utils import get_nvtx_range from megatron.core.packed_seq_params import PackedSeqParams from megatron.core import mpu @@ -391,8 +391,6 @@ def get_default_packed_seq_params(seq_length: int, max_sequences_per_bin: int, d PackedSeqParams configured as a single unpacked sequence. """ - args = get_args() - # Pad to the maximum number of sequences in the bin for the attention kernel. # We add 2 to account for the initial 0 and the final bin_size. cu_seqlens = torch.full( @@ -409,6 +407,7 @@ def get_default_packed_seq_params(seq_length: int, max_sequences_per_bin: int, d max_seqlen_q=seq_length, max_seqlen_kv=seq_length, total_tokens=seq_length, + pad_between_seqs=False, ) def create_packed_seq_params(packing_context: PackingContext): @@ -424,6 +423,7 @@ def create_packed_seq_params(packing_context: PackingContext): bin_size=bin_size, max_sequences_per_bin=max_sequences_per_bin, device=device, + seq_length_multiple=packing_context.packer.seq_length_multiple, ) cached_packed_seq_params.append(params) return cached_packed_seq_params @@ -433,7 +433,8 @@ def create_packed_seq_params_for_bin( bin_idx: int, bin_size: int, max_sequences_per_bin: int, - device: torch.device + device: torch.device, + seq_length_multiple: int = 1, ) -> Optional[PackedSeqParams]: """Create PackedSeqParams for a single bin to enable proper attention masking in TE. @@ -447,6 +448,9 @@ def create_packed_seq_params_for_bin( bin_size: Size of the bin (padded sequence length) max_sequences_per_bin: Maximum number of sequences per bin device: Device to create tensors on + seq_length_multiple: Slot alignment the packer used (2*cp_size under context + parallelism). When > 1, every padded slot length is validated against it — + TE's THD partitioner requires slots to split evenly across CP ranks. Returns: PackedSeqParams with cu_seqlens set for proper attention masking (or None if empty) @@ -459,20 +463,57 @@ def create_packed_seq_params_for_bin( # Get actual sequence lengths for sequences in this bin seq_lengths_in_bin = [packing_info.seq_lengths[idx] for idx in seq_indices] + assert len(seq_lengths_in_bin) <= max_sequences_per_bin, ( + f"bin {bin_idx} holds {len(seq_lengths_in_bin)} sequences, exceeding " + f"max_sequences_per_bin={max_sequences_per_bin}" + ) + # Aligned placement offsets from the packer: [start_0, ..., start_{n-1}, padded_end]. + seq_starts = packing_info.seq_starts[bin_idx] + assert len(seq_starts) == len(seq_lengths_in_bin) + 1, ( + f"seq_starts for bin {bin_idx} must contain one start per sequence plus the padded " + f"end sentinel (got {len(seq_starts)} entries for {len(seq_lengths_in_bin)} sequences)" + ) + padded_end = seq_starts[-1] + + # Layout boundaries (cu_seqlens_padded): where each sequence's reserved slot + # begins and ends in the bin. The packer places sequences at + # seq_length_multiple-aligned offsets; the trailing bin padding + # [padded_end, bin_size) forms a final ghost slot. + cu_seqlens_padded_list = np.append(np.asarray(seq_starts), bin_size) + + if seq_length_multiple > 1: + slot_lengths = np.diff(cu_seqlens_padded_list) + assert np.all(slot_lengths % seq_length_multiple == 0), ( + f"bin {bin_idx} padded slot lengths {slot_lengths.tolist()} must each be " + f"divisible by seq_length_multiple={seq_length_multiple} for TE's THD " + f"context-parallel partitioning" + ) - # Build cumulative sequence lengths for actual sequences - # cu_seqlens should be [0, len(seq1), len(seq1)+len(seq2), ..., total_actual_len] - cu_seqlens_list = np.append(np.cumsum([0] + seq_lengths_in_bin), bin_size) + # Actual cumulative lengths (cu_seqlens): real token count per slot, so TE + # masks each sequence's pad tail. The trailing ghost slot keeps its full + # size (its pad tokens self-attend), matching the legacy behaviour — with + # seq_length_multiple == 1 both arrays equal the legacy single array. + cu_seqlens_list = np.cumsum([0] + seq_lengths_in_bin + [bin_size - padded_end]) + + # Pad both arrays to a fixed size by repeating their last boundary (creates + # zero-length ghost sequences). This ensures a fixed tensor size for CUDA + # graph compatibility. We add 2 to account for the initial 0 and the final + # boundary. + fixed_size = max_sequences_per_bin + 2 + cu_seqlens_list = np.pad(cu_seqlens_list, (0, fixed_size - len(cu_seqlens_list)), mode='edge') + cu_seqlens_padded_list = np.pad( + cu_seqlens_padded_list, (0, fixed_size - len(cu_seqlens_padded_list)), mode='edge' + ) cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int32, device=device) - - # Pad cu_seqlens to bin_size by repeating the last value (creates zero-length ghost sequences) - # This ensures a fixed tensor size for CUDA graph compatibility - # We add 2 to account for the initial 0 and the final bin_size. - if len(cu_seqlens) < max_sequences_per_bin + 2: - out = cu_seqlens.new_full((max_sequences_per_bin + 2,), bin_size) - out[:len(cu_seqlens)] = cu_seqlens - cu_seqlens = out + if seq_length_multiple > 1: + cu_seqlens_padded = torch.tensor(cu_seqlens_padded_list, dtype=torch.int32, device=device) + else: + assert np.array_equal(cu_seqlens_padded_list, cu_seqlens_list), ( + f"bin {bin_idx}: padded layout {cu_seqlens_padded_list.tolist()} deviates from " + f"actual boundaries {cu_seqlens_list.tolist()} with seq_length_multiple == 1" + ) + cu_seqlens_padded = None max_seqlen = bin_size @@ -480,11 +521,12 @@ def create_packed_seq_params_for_bin( qkv_format='thd', cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, - cu_seqlens_q_padded=None, - cu_seqlens_kv_padded=None, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, max_seqlen_q=max_seqlen, max_seqlen_kv=max_seqlen, total_tokens=bin_size, + pad_between_seqs=seq_length_multiple > 1, ) @@ -622,12 +664,33 @@ def compute_packed_inference_logprobs_stats( class SequencePacker: - """Packs multiple sequences into bins to minimize padding and improve GPU utilization.""" + """Packs multiple sequences into bins to minimize padding and improve GPU utilization. + + Args: + bin_size: Size of each bin (padded sequence length). + pad_token: Padding token id. + max_sequences_per_bin: Maximum number of sequences per bin. + seq_length_multiple: Align sequences within a bin to multiples of this value: + the slot a sequence occupies is its length rounded up to the next multiple. + Context parallelism requires every padded sequence length to be a multiple of 2*cp_size. + """ - def __init__(self, bin_size: int, pad_token: int, max_sequences_per_bin: int = 16): + def __init__( + self, + bin_size: int, + pad_token: int, + max_sequences_per_bin: int = 16, + seq_length_multiple: int = 1, + ): + assert seq_length_multiple >= 1 + assert bin_size % seq_length_multiple == 0, ( + f"bin_size ({bin_size}) must be divisible by seq_length_multiple " + f"({seq_length_multiple})" + ) self.bin_size = bin_size self.pad_token = pad_token self.max_sequences_per_bin = max_sequences_per_bin + self.seq_length_multiple = seq_length_multiple def pack_sequences( self, trajs: torch.Tensor, generation_masks: Optional[torch.Tensor] = None @@ -651,11 +714,11 @@ def pack_sequences( current_bin_indices = [] current_bin_length = 0 - # Pack sequences into bins + # Pack sequences into bins. sequences_per_bin = [] for idx in sorted_indices: seq = sequences[idx] - seq_len = len(seq) + seq_len = round_up_to_nearest_multiple(len(seq), self.seq_length_multiple) if ( current_bin_length + seq_len <= self.bin_size @@ -720,7 +783,9 @@ def pack_sequences( for seq_idx in seq_indices: seq_to_bin_idx[seq_idx] = bin_idx - # Fill bins + # Fill bins. Each sequence is placed at a seq_length_multiple-aligned + # offset; the gap up to the next aligned offset stays pad tokens with a + # zero loss mask. for bin_idx, (bin_seqs, seq_indices) in enumerate(zip(bins, bin_seq_indices)): seq_starts = [] current_pos = 0 @@ -729,7 +794,9 @@ def pack_sequences( start = current_pos end = start + len(seq) seq_starts.append(start) - current_pos = end + current_pos = start + round_up_to_nearest_multiple( + len(seq), self.seq_length_multiple + ) # Pack sequence packed_sequences[bin_idx, start:end] = seq @@ -961,11 +1028,16 @@ def _gather(data): inference_logprobs = _gather(inference_logprobs) with nvtx_range("rl/pack-sequences", time=True): - # Create packer with max sequences per bin limit to prevent extreme imbalance + # Create packer with max sequences per bin limit to prevent extreme imbalance. + # Context parallelism requires every sequence's padded slot in the bin to be + # divisible by 2*cp_size so TE's THD ring attention can split each sequence + # evenly across CP ranks (see _scatter_for_context_parallel in rl_utils). + cp_size = mpu.get_context_parallel_world_size() packer = SequencePacker( bin_size=bin_size, pad_token=tokenizer.pad, max_sequences_per_bin=max_sequences_per_bin, + seq_length_multiple=2 * cp_size if cp_size > 1 else 1, ) # Pack sequences with generation masks @@ -1010,6 +1082,7 @@ def _gather(data): bin_size=bin_size, max_sequences_per_bin=max_sequences_per_bin, device=packed_trajs.device, + seq_length_multiple=packer.seq_length_multiple, ) for bin_idx in range(len(packed_trajs)) ] diff --git a/tests/unit_tests/rl/test_context_parallel.py b/tests/unit_tests/rl/test_context_parallel.py new file mode 100644 index 00000000000..0c71379e3f3 --- /dev/null +++ b/tests/unit_tests/rl/test_context_parallel.py @@ -0,0 +1,187 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch + +from megatron.core import mpu +from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer import TransformerConfig +from megatron.rl import rl_utils +from megatron.rl.sequence_packing_utils import SequencePacker, create_packed_seq_params_for_bin +from megatron.training.arguments import parse_args, validate_args +from megatron.training.global_vars import destroy_global_vars, set_global_variables +from tests.unit_tests.test_utilities import Utils + +VOCAB = 128 +MAX_SEQUENCES_PER_BIN = 4 + +pytestmark = pytest.mark.skipif( + not rl_utils.HAVE_TEX, reason="CP logprobs require Transformer Engine THD indices" +) + +CP_SIZES = [pytest.param(cp, id=f"cp{cp}") for cp in (2, 4, 8) if cp <= Utils.world_size] + + +class PositionWiseModel: + """Deterministic model whose logits at each index depend only on that index's + token and position id (tables are identical on every rank).""" + + def __init__(self, max_position: int, device: torch.device): + self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() + self.config = TransformerConfig( + num_attention_heads=8, num_layers=8, pipeline_dtype=torch.bfloat16 + ) + generator = torch.Generator().manual_seed(1234) + self.token_logits = torch.randn(VOCAB, VOCAB, generator=generator).to(device) + self.position_logits = torch.randn(max_position, VOCAB, generator=generator).to(device) + + def __call__(self, tokens, position_ids, attention_mask, **kwargs): + return self.token_logits[tokens] + self.position_logits[position_ids] + + +def _set_minimal_args(): + destroy_global_vars() + destroy_num_microbatches_calculator() + args = parse_args(ignore_unknown_args=True) + args.num_layers = 8 + args.num_attention_heads = 8 + args.vocab_size = VOCAB + args.hidden_size = 128 + args.max_position_embeddings = 256 + args.seq_length = 256 + args.micro_batch_size = 1 + args.wandb_project = None + args = validate_args(args) + set_global_variables(args, False) + + +@pytest.fixture +def initialize_context_parallel(request, monkeypatch): + """Initialize/destroy model parallel with the requested CP size (tp = pp = 1).""" + monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1") + monkeypatch.setenv("WANDB_MODE", "disabled") + monkeypatch.setenv("LOG_TO_WANDB", "false") + Utils.initialize_model_parallel(context_parallel_size=request.param) + _set_minimal_args() + yield request.param + Utils.destroy_model_parallel() + destroy_global_vars() + destroy_num_microbatches_calculator() + + +def _pack_bin(lengths, seq_length_multiple, bin_size, device): + """Pack variable-length sequences into one bin and build its PackedSeqParams.""" + generator = torch.Generator().manual_seed(7) + seqs = torch.zeros(len(lengths), max(lengths), dtype=torch.long) # 0 is the pad token + for i, length in enumerate(lengths): + seqs[i, :length] = torch.randint(1, VOCAB, (length,), generator=generator) + packer = SequencePacker( + bin_size=bin_size, + pad_token=0, + max_sequences_per_bin=MAX_SEQUENCES_PER_BIN, + seq_length_multiple=seq_length_multiple, + ) + packed_trajs, packed_position_ids, _, packing_info = packer.pack_sequences(seqs) + assert packed_trajs.shape[0] == 1, "test data must fit a single bin" + packed_seq_params = create_packed_seq_params_for_bin( + packing_info=packing_info, + bin_idx=0, + bin_size=bin_size, + max_sequences_per_bin=MAX_SEQUENCES_PER_BIN, + device=device, + seq_length_multiple=seq_length_multiple, + ) + return packed_trajs[:1].to(device), packed_position_ids[:1].to(device), packed_seq_params + + +class TestContextParallelLogprobs: + + @pytest.mark.parametrize("initialize_context_parallel", CP_SIZES, indirect=True) + @pytest.mark.parametrize( + "packed", [pytest.param(False, id="single_sequence"), pytest.param(True, id="packed_bin")] + ) + @pytest.mark.parametrize( + "no_grad", [pytest.param(True, id="no_grad"), pytest.param(False, id="grad")] + ) + def test_cp_logprobs_match_single_rank_reference( + self, initialize_context_parallel, packed, no_grad + ): + """get_logprobs under CP must return the exact cp_size == 1 result on every rank.""" + cp_size = initialize_context_parallel + slot = 2 * cp_size # TE splits every padded slot into 2*cp_size chunks + device = torch.device("cuda") + if packed: + # Footprints of 3/2/1 slots: pad gaps after each sequence and a + # trailing ghost slot of 2 more slots at the end of the bin. + bin_size, lengths = 8 * slot, [2 * slot + 1, slot + 3, slot - 2] + else: + bin_size, lengths = 4 * slot, [4 * slot] + tokens, position_ids, packed_seq_params = _pack_bin(lengths, slot, bin_size, device) + # The packer must reserve every slot (incl. ghosts) on the 2*cp_size grid. + boundaries = packed_seq_params.cu_seqlens_q_padded + assert boundaries[-1].item() == bin_size + assert torch.all((boundaries[1:] - boundaries[:-1]) % slot == 0) + + model = PositionWiseModel(max_position=bin_size, device=device) + if not no_grad: + model.token_logits.requires_grad_() + + logprobs = rl_utils.get_logprobs( + model, + tokens, + position_ids, + no_grad=no_grad, + sequence_packing=True, + packed_seq_params=packed_seq_params, + ) + + with torch.no_grad(): + full_logits = model(tokens, position_ids, None) + reference = rl_utils.selective_log_softmax(full_logits[:, :-1, :], tokens[:, 1:]) + # The tables are identical on every rank, so matching the local reference + # on each rank also means all CP ranks returned the same logprobs. + torch.testing.assert_close(logprobs, reference) + # The CP scatter must annotate a copy, not the caller's params. + assert packed_seq_params.cp_group is None + assert packed_seq_params.local_cp_size is None + + # A second call must hit the memoized per-bin partition and return the same reference. + cp_scatter = packed_seq_params.cp_scatter_cache + logprobs_again = rl_utils.get_logprobs( + model, + tokens, + position_ids, + no_grad=True, + sequence_packing=True, + packed_seq_params=packed_seq_params, + ) + assert packed_seq_params.cp_scatter_cache is cp_scatter + torch.testing.assert_close(logprobs_again, reference) + + if not no_grad: + # Gradients must flow back through the differentiable logprob all-gather. + logprobs.sum().backward() + grad = model.token_logits.grad + assert grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 + + @pytest.mark.parametrize("initialize_context_parallel", CP_SIZES[:1], indirect=True) + def test_scatter_rejects_indivisible_seq_len(self, initialize_context_parallel): + """Unaligned inputs (bypassing the packer) must be rejected, not mis-partitioned.""" + cp_size = initialize_context_parallel + seq_len = 2 * cp_size + 1 + cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda") + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=seq_len, + max_seqlen_kv=seq_len, + total_tokens=seq_len, + ) + tokens = torch.zeros(1, seq_len, dtype=torch.long, device="cuda") + with pytest.raises(AssertionError, match="divisible"): + rl_utils._scatter_for_context_parallel( + tokens, tokens.clone(), packed_seq_params, mpu.get_context_parallel_group() + ) diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index cb6103a6d9c..51281bdd3d6 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -30,7 +30,6 @@ from megatron.core.transformer import TransformerConfig from megatron.core.transformer.cuda_graphs import ( CudaGraphManager, - _CudagraphGlobalRecord, create_cudagraphs, delete_cuda_graphs, ) @@ -129,7 +128,7 @@ def __init__(self, config, layer_config): super().__init__() self.config = config self.layer = DummyConfigModule(layer_config) - self.pg_collection = SimpleNamespace(pp=object()) + self.pg_collection = SimpleNamespace(pp=object(), cp=None) self.config_values_during_forward = None def forward(self, tokens, position_ids, attention_mask, **kwargs): @@ -164,8 +163,8 @@ def transition_cudagraph_scope(self, mode): def initialize_model_parallel(request, monkeypatch): """Fixture to initialize and destroy model parallel. - Parameters are passed via request.param as a tuple: (tp, pp) - Skips if world_size < tp * pp. + Parameters are passed via request.param as a tuple: (tp, pp) or (tp, pp, cp) + Skips if world_size < tp * pp * cp. """ monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1") monkeypatch.setenv("WANDB_MODE", "disabled") @@ -173,11 +172,14 @@ def initialize_model_parallel(request, monkeypatch): initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) - tp, pp = request.param + tp, pp = request.param[:2] + cp = request.param[2] if len(request.param) > 2 else 1 world_size = Utils.world_size - Utils.initialize_model_parallel(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp) + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp, context_parallel_size=cp + ) model_parallel_cuda_manual_seed(123) - dp = world_size // (tp * pp) + dp = world_size // (tp * pp * cp) yield world_size, dp, tp, pp Utils.destroy_model_parallel() destroy_global_vars() @@ -195,6 +197,16 @@ def cleanup_global_state(): class TestRLUtils: """Test class for RL utilities.""" + def teardown_method(self, method): + # Runs even when a test fails: leaked capture state otherwise wedges the + # session-end distributed barrier into the 60-minute job timeout. + delete_cuda_graphs() + if torch.cuda.is_available(): + torch.cuda.synchronize() + CudaGraphManager.global_mempool = None + CudaGraphManager.fwd_mempools = None + CudaGraphManager.bwd_mempools = None + def create_test_args(self, **kwargs): destroy_global_vars() destroy_num_microbatches_calculator() @@ -1250,24 +1262,36 @@ def test_gpt_logprobs(self, initialize_model_parallel): @pytest.mark.parametrize( "initialize_model_parallel", - [pytest.param((1, 1), id="tp1-pp1")], + [ + pytest.param((1, 1), id="tp1-pp1"), + pytest.param((2, 1), id="tp2-pp1"), + pytest.param((1, 2), id="tp1-pp2"), + pytest.param( + (1, 1, 2), + id="tp1-pp1-cp2", + marks=pytest.mark.skip( + reason="PackedSeqParams.cp_group (a ProcessGroup) is rejected by the " + "CUDA-graph arg type whitelist; fixed and unskipped in a follow-up PR" + ), + ), + ], indirect=["initialize_model_parallel"], ) def test_get_logprobs_cuda_graphs(self, initialize_model_parallel): - """Test that get_logprobs reuses CUDA graphs created during training forward pass. + """Training CUDA graphs recorded and reused through get_logprobs. - This test verifies that rl_utils.get_logprobs can reuse CUDA graphs by: - 1. Running a training-style forward pass on some model to record CUDA graph runners. - 2. Creating the CUDA graphs. - 3. Running `get_logprobs` to verify it reuses the same forward graph from training. + The training-style pass goes through get_logprobs exactly like train_rl.py's + forward_step, so the recorded runners see the shapes production sees — under CP + that is the seq // cp local scatter and the explicit pad_between_seqs flag, and + capture includes TE's ring-P2P context-parallel attention. The reference-logprobs + pass must then replay those graphs without recording new runners. """ - - num_layers = 2 - world_size, dp, tp, pp = initialize_model_parallel + cp = world_size // (dp * tp * pp) self.create_test_args( tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp, + context_parallel_size=cp, rl_training_cuda_graphs=True, cuda_graph_impl="local", bf16=True, @@ -1276,12 +1300,16 @@ def test_get_logprobs_cuda_graphs(self, initialize_model_parallel): # Create a model with training CUDA graphs enabled transformer_config = TransformerConfig( - num_layers=num_layers, + num_layers=2, hidden_size=64, num_attention_heads=4, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + context_parallel_size=cp, use_cpu_initialization=True, cuda_graph_impl="local", bf16=True, + pipeline_dtype=torch.bfloat16, ) model = GPTModel( config=transformer_config, @@ -1296,37 +1324,28 @@ def test_get_logprobs_cuda_graphs(self, initialize_model_parallel): for param in wrapped_model.parameters(): param.main_grad = torch.zeros_like(param) - # Create test inputs (batch_size=1 required for thd format with sequence packing) + # batch_size=1 required for thd format with sequence packing; seq_length divisible + # by 2 * cp, as TE's THD partitioning requires. batch_size = 1 - seq_length = 16 + seq_length = 32 tokens = torch.randint(0, 256, (batch_size, seq_length), dtype=torch.long).cuda() position_ids = torch.arange(seq_length).unsqueeze(0).expand(batch_size, -1).cuda() - # Create packed_seq_params for dummy data - packed_seq_params = get_default_packed_seq_params( - seq_length=seq_length, max_sequences_per_bin=4, device=tokens.device - ) - - # Run a single training forward pass to record cudagraphs - output = wrapped_model( - tokens, - position_ids, - attention_mask=None, - packed_seq_params=packed_seq_params, - runtime_gather_output=True, - fp32_output=False, + # Training-style forward pass: like train_rl.py's forward_step it goes through + # get_logprobs (which scatters to seq // cp under CP), recording the runners. + output = rl_utils.get_logprobs( + wrapped_model, tokens, position_ids=position_ids, no_grad=False, sequence_packing=True ) # Run backward to reset runner status from BWD_READY back to FWD_READY - # This is needed because get_logprobs runs in no_grad mode and expects FWD_READY - loss = output.sum() - loss.backward() + output.sum().backward() # Collect all CudaGraphManager instances and their runners cudagraph_managers = [] for module in wrapped_model.modules(): if hasattr(module, 'cudagraph_manager') and module.cudagraph_manager is not None: cudagraph_managers.append(module.cudagraph_manager) + assert cudagraph_managers, "Expected CUDA graph managers on the transformer layers" # Record runner count before creating graphs runners_before = {id(mgr): len(mgr.cudagraph_runners) for mgr in cudagraph_managers} @@ -1337,8 +1356,8 @@ def test_get_logprobs_cuda_graphs(self, initialize_model_parallel): for mgr in cudagraph_managers: for runner in mgr.cudagraph_runners: assert runner.fwd_graph is not None, ( - f"Expected runner to have fwd_graph created after create_cudagraphs(), " - f"but fwd_graph is None" + "Expected runner to have fwd_graph created after create_cudagraphs(), " + "but fwd_graph is None" ) # Now test `get_logprobs`; this should reuse the existing CUDA graphs @@ -1356,27 +1375,12 @@ def test_get_logprobs_cuda_graphs(self, initialize_model_parallel): f"but got {count_after}. `get_logprobs` should not create new runners." ) - # Verify outputs are valid + # Verify outputs are valid; only the last pipeline stage returns logprobs. assert output is not None, "Training forward pass should return valid output" - assert logprobs is not None, "get_logprobs should return valid output" - - # Destroy all captured graphs deterministically - for l in model.decoder.layers: - for runner in getattr(l.cudagraph_manager, "cudagraph_runners", []): - # Safely delete both graphs if present - if hasattr(runner, "fwd_graph"): - del runner.fwd_graph - if hasattr(runner, "bwd_graph"): - del runner.bwd_graph - - # Ensure all pending work is complete and graph destruction runs now - torch.cuda.synchronize() - - _CudagraphGlobalRecord.cudagraph_created = False - _CudagraphGlobalRecord.cudagraph_record = [] - CudaGraphManager.global_mempool = None - CudaGraphManager.fwd_mempools = None - CudaGraphManager.bwd_mempools = None + if is_pp_last_stage(model.pg_collection.pp): + assert logprobs.shape == (batch_size, seq_length - 1) + else: + assert logprobs is not None, "get_logprobs should return valid output" @pytest.mark.parametrize( "initialize_model_parallel", diff --git a/tests/unit_tests/rl/test_sequence_packing_utils.py b/tests/unit_tests/rl/test_sequence_packing_utils.py index f5f56c8877d..85fbdb150fb 100644 --- a/tests/unit_tests/rl/test_sequence_packing_utils.py +++ b/tests/unit_tests/rl/test_sequence_packing_utils.py @@ -1,10 +1,12 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import dataclasses from unittest.mock import patch import pytest import torch +from megatron.core.packed_seq_params import PackedSeqParams from megatron.rl import rl_utils, sequence_packing_utils from megatron.training import arguments, global_vars @@ -461,7 +463,8 @@ def test_cu_seqlens_size(num_sequences): packing_info = sequence_packing_utils.PackingInfo( bin_seq_indices=[list(range(num_sequences))], - seq_starts={0: []}, + # Back-to-back starts plus the padded end sentinel, as pack_sequences records them. + seq_starts={0: [i * seq_len for i in range(num_sequences + 1)]}, seq_lengths=seq_lengths, seq_to_bin_idx=[0] * num_sequences, packing_algo='fifo', @@ -530,3 +533,154 @@ def side_eff( # Iterator is local, batch size is global assert expected_bs == actual_bs + + +@pytest.mark.parametrize( + "multiple, bin_size, expected_starts, expected_cu_padded, expected_cu", + [ + # Back-to-back placement: the layout equals the actual boundaries, so the + # padded fields stay None (keeps TE off its pad_between_seqs GPU sync, + # which would break CUDA graph capture). + pytest.param(1, 16, [0, 7, 12], None, [0, 7, 12, 16, 16, 16], id="legacy_back_to_back"), + # 4-aligned placement (cp=2): lengths 7 and 5 reserve footprints of 8; + # cu_seqlens_*_padded is the slot grid (with the trailing ghost slot), + # cu_seqlens_* the real token counts (the ghost keeps its full size). + pytest.param( + 4, 24, [0, 8, 16], [0, 8, 16, 24, 24, 24], [0, 7, 12, 20, 20, 20], id="aligned_for_cp" + ), + ], +) +def test_sequence_packing_alignment_pipeline( + multiple, bin_size, expected_starts, expected_cu_padded, expected_cu +): + """pack_sequences places each sequence at a seq_length_multiple-aligned offset + (gaps stay pad tokens with a zero loss mask) and create_packed_seq_params_for_bin + turns that placement into the cu_seqlens / cu_seqlens_padded pair, materializing + the padded boundaries only when alignment makes them differ.""" + tokenizer = MockTokenizer() + seq_a = torch.cat([torch.arange(1, 6), torch.full((2,), tokenizer.pad)]) # length 5 + seq_b = torch.arange(10, 17) # length 7, placed first (packer sorts by length) + packer = sequence_packing_utils.SequencePacker( + bin_size=bin_size, + pad_token=tokenizer.pad, + max_sequences_per_bin=4, + seq_length_multiple=multiple, + ) + + packed_trajs, packed_position_ids, packed_loss_mask, packing_info = packer.pack_sequences( + torch.stack([seq_a, seq_b]) + ) + + assert packed_trajs.shape == (1, bin_size) + assert packing_info.bin_seq_indices[0] == [1, 0] + # Aligned starts plus the padded end sentinel. + assert packing_info.seq_starts[0] == expected_starts + + # Data at the aligned starts; everything outside the two sequences + # (alignment gaps + bin tail) stays pad tokens with a zero loss mask. + start_a = expected_starts[1] + assert torch.equal(packed_trajs[0, 0:7], seq_b) + assert torch.equal(packed_trajs[0, start_a : start_a + 5], seq_a[:5]) + is_data = torch.zeros(bin_size, dtype=torch.bool) + is_data[0:7] = True + is_data[start_a : start_a + 5] = True + assert torch.all(packed_trajs[0, ~is_data] == tokenizer.pad) + assert torch.all(packed_loss_mask[0, ~is_data] == 0) + # Position ids restart at each start. + assert packed_position_ids[0, start_a] == 0 + + params = sequence_packing_utils.create_packed_seq_params_for_bin( + packing_info=packing_info, + bin_idx=0, + bin_size=bin_size, + max_sequences_per_bin=4, + device=torch.device('cpu'), + seq_length_multiple=multiple, + ) + if expected_cu_padded is None: + assert params.cu_seqlens_q_padded is None + assert params.cu_seqlens_kv_padded is None + assert params.pad_between_seqs is False + else: + assert params.cu_seqlens_q_padded.tolist() == expected_cu_padded + assert params.cu_seqlens_kv_padded.tolist() == expected_cu_padded + assert params.pad_between_seqs is True + assert params.cu_seqlens_q.tolist() == expected_cu + + +def test_sequence_packing_aligned_fit_accounting(): + """Fit decisions must use the reserved (rounded) footprint, not the raw length.""" + tokenizer = MockTokenizer() + packer = sequence_packing_utils.SequencePacker( + bin_size=8, pad_token=tokenizer.pad, seq_length_multiple=4 + ) + # Three sequences of length 3 (footprint 4): only two fit per 8-token bin. + sequences = torch.arange(1, 10).reshape(3, 3) + packed_trajs, _, _, packing_info = packer.pack_sequences(sequences) + assert packed_trajs.shape[0] == 2 + assert sorted(len(idxs) for idxs in packing_info.bin_seq_indices) == [1, 2] + + +def test_sequence_packer_rejects_indivisible_bin(): + with pytest.raises(AssertionError, match="divisible"): + sequence_packing_utils.SequencePacker(bin_size=10, pad_token=0, seq_length_multiple=4) + + +def test_default_packed_seq_params_signature_matches_bin_params(): + """The empty-padding-bin fallback (get_default_packed_seq_params) must be structurally + interchangeable with real bins' params — same None-ness, shape, and dtype per field — + or CUDA graph replay rejects the microbatches that fall back to it.""" + tokenizer = MockTokenizer() + packer = sequence_packing_utils.SequencePacker( + bin_size=16, pad_token=tokenizer.pad, max_sequences_per_bin=4 + ) + _, _, _, packing_info = packer.pack_sequences(torch.arange(1, 11).reshape(2, 5)) + bin_params = sequence_packing_utils.create_packed_seq_params_for_bin( + packing_info=packing_info, + bin_idx=0, + bin_size=16, + max_sequences_per_bin=4, + device=torch.device('cpu'), + ) + default_params = sequence_packing_utils.get_default_packed_seq_params( + seq_length=16, max_sequences_per_bin=4, device=torch.device('cpu') + ) + + for field in dataclasses.fields(PackedSeqParams): + bin_value = getattr(bin_params, field.name) + default_value = getattr(default_params, field.name) + assert type(default_value) is type(bin_value), field.name + if isinstance(bin_value, torch.Tensor): + assert default_value.shape == bin_value.shape, field.name + assert default_value.dtype == bin_value.dtype, field.name + + +@pytest.mark.parametrize( + "override, match", + [ + pytest.param({'max_sequences_per_bin': 1}, "max_sequences_per_bin", id="overfull_bin"), + pytest.param({'seq_length_multiple': 4}, "divisible", id="misaligned_slots"), + pytest.param({'seq_starts': []}, "sentinel", id="missing_seq_starts"), + ], +) +def test_create_packed_seq_params_rejects_invalid_bin(override, match): + """Bins violating construction contracts must fail loudly: more sequences than the + fixed-size cu_seqlens can hold, padded slots misaligned to seq_length_multiple + (TE's THD context-parallel partitioning contract), or missing seq_starts.""" + tokenizer = MockTokenizer() + packer = sequence_packing_utils.SequencePacker( + bin_size=16, pad_token=tokenizer.pad, max_sequences_per_bin=4 + ) + _, _, _, packing_info = packer.pack_sequences(torch.arange(1, 11).reshape(2, 5)) + kwargs = dict( + packing_info=packing_info, + bin_idx=0, + bin_size=16, + max_sequences_per_bin=4, + device=torch.device('cpu'), + ) + if 'seq_starts' in override: + packing_info.seq_starts[0] = override.pop('seq_starts') + kwargs.update(override) + with pytest.raises(AssertionError, match=match): + sequence_packing_utils.create_packed_seq_params_for_bin(**kwargs) diff --git a/train_rl.py b/train_rl.py index 17eec357b15..04ec7bc66c1 100644 --- a/train_rl.py +++ b/train_rl.py @@ -9,7 +9,6 @@ from gpt_builders import gpt_builder from hybrid_builders import hybrid_builder -from megatron.core import mpu from megatron.core.enums import ModelType from megatron.core.models.gpt import GPTModel from megatron.core.parallel_state import is_pipeline_last_stage @@ -140,9 +139,6 @@ def loss_func( masked_truncated_from_above = torch.sum(loss_mask_flat * truncated_from_above_flat) masked_truncated_from_below = torch.sum(loss_mask_flat * truncated_from_below_flat) - if args.context_parallel_size > 1: - torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group()) - # Check individual rank losses are not NaN prior to DP all-reduce. rerun_state_machine = get_rerun_state_machine() if args.check_for_nan_in_loss_and_grad: @@ -188,7 +184,7 @@ def loss_func( # Note: This information needs to be determined in forward_step where we have access to the batch data # The loss_func doesn't have direct access to this information - return (loss[0] * args.context_parallel_size, total_tokens.int(), output_dict) + return (loss[0].clone(), total_tokens.int(), output_dict) def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): @@ -268,13 +264,15 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): ) else: cu_seqlens = torch.tensor([0, tokens.shape[1]], dtype=torch.int32, device=tokens.device) + # Make sure to omit `total_tokens` to prevent `seq_idx` from being auto-computed. + # That would cause a sequence packing kernel to be incorrectly used. packed_seq_params = PackedSeqParams( qkv_format='thd', cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, max_seqlen_q=tokens.shape[1], max_seqlen_kv=tokens.shape[1], - total_tokens=tokens.shape[1], + pad_between_seqs=False, ) # Clear RoPE cache to avoid inference tensor errors From 32e7edaecd610da90ff1668130d95e1e8275736b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Mon, 17 Aug 2026 19:45:17 +0200 Subject: [PATCH 055/101] docs(version-picker): preserve 0.19.0 and add patch releases (#6578) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- docs/versions1.json | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/versions1.json b/docs/versions1.json index 84942f53740..926d43d16ee 100644 --- a/docs/versions1.json +++ b/docs/versions1.json @@ -10,16 +10,31 @@ "url": "https://docs.nvidia.com/megatron-core/developer-guide/latest/", "preferred": true }, + { + "name": "0.18.2", + "version": "0.18.2", + "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.18.2/" + }, { "name": "0.18.0", "version": "0.18.0", "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.18.0/" }, + { + "name": "0.17.1", + "version": "0.17.1", + "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.17.1/" + }, { "name": "0.17.0", "version": "0.17.0", "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.17.0/" }, + { + "name": "0.16.1", + "version": "0.16.1", + "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.16.1/" + }, { "name": "0.16.0", "version": "0.16.0", From 62a617338b8143d2911b36234022c2e55704a12e Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 17 Aug 2026 16:56:52 +0000 Subject: [PATCH 056/101] test(mfsdp): cover v2 training-step peak memory (#6551) Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 69 ++++++++----------- 1 file changed, 27 insertions(+), 42 deletions(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 131afdbd0c8..2bb6d92a7ba 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -608,8 +608,8 @@ def test_tied_child_parameters_allocate_one_physical_weight(distributed_setup): assert len(list(model.parameters())) == 1 -def test_forward_peak_memory_bounds_in_flight_child_all_gathers(distributed_setup): - """Forward peak memory should stay below three live child all-gathers.""" +def test_training_step_peak_memory_bounds_full_size_buffers(distributed_setup): + """A training step should stay below five full-size child buffers.""" rank = distributed_setup.rank world_size = distributed_setup.world_size device = distributed_setup.device @@ -619,38 +619,39 @@ def test_forward_peak_memory_bounds_in_flight_child_all_gathers(distributed_setu mesh = init_device_mesh(device.type, (world_size,)) dim = 4096 dtype = torch.bfloat16 - model = MultiChildModel(dim=dim, num_children=4).to(dtype=dtype, device=device) + model = MultiChildModel(dim=dim, num_children=8).to(dtype=dtype) placements = _flat_placements() policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) with fully_shard_context(device=device): for layer in model.layers: fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + fully_shard_optimizer(optimizer) x = torch.randn(2, dim, device=device, dtype=dtype) - with torch.no_grad(): - model(x) - torch.cuda.synchronize(device) - torch.cuda.empty_cache() + def train_step() -> None: + optimizer.zero_grad(set_to_none=True) + model(x).float().sum().backward() + optimizer.step() + + child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() resting_allocated = torch.cuda.memory_allocated(device) torch.cuda.reset_peak_memory_stats(device) - with torch.no_grad(): - model(x) - torch.cuda.synchronize(device) + train_step() peak_delta = torch.cuda.max_memory_allocated(device) - resting_allocated - child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() - bound_nbytes = 3 * child_weight_nbytes + # Backward keeps the current child and one prefetched child unsharded. The current + # child also has a full wgrad until it is copied into a full reduce-scatter input, + # for a four-full-child-buffer peak. Allow one additional buffer for cuBLAS + # workspace, allocator granularity, and small temporaries. + bound_nbytes = (4 + 1) * child_weight_nbytes - # A parent forward should keep one previous child unsharded until its compute - # stream consumer is safe, plus the current child being unsharded. The bound - # is looser than two child weights to avoid coupling this test to CUDA - # allocator granularity and small temporary buffers, while still catching - # delayed releases piling up across the four child layers. assert peak_delta < bound_nbytes, ( - "FSDP forward peak memory exceeded the in-flight all-gather bound: " + "FSDP training-step peak memory exceeded the full-size-buffer bound: " f"rank={rank}, peak_delta={_mb(peak_delta)}, " - f"three_child_weights={_mb(bound_nbytes)}" + f"five_full_child_buffers={_mb(bound_nbytes)}" ) @@ -671,7 +672,6 @@ def test_deleted_model_releases_fsdp_storage(distributed_setup): x = torch.ones(1, 8192, dtype=torch.bfloat16, device=device) output = model(x) del output, x, model - torch.cuda.synchronize(device) assert torch.cuda.memory_allocated(device) - allocated_before < 1024**2 @@ -696,14 +696,10 @@ def test_root_forward_returns_to_resting_memory(distributed_setup): fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) x = torch.randn(2, dim, device=device, dtype=dtype) - torch.cuda.synchronize(device) - torch.cuda.empty_cache() resting_allocated = torch.cuda.memory_allocated(device) - with torch.no_grad(): - output = model(x) + output = model(x) del output - torch.cuda.synchronize(device) allocated_after_forward = torch.cuda.memory_allocated(device) extra_allocated = allocated_after_forward - resting_allocated child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() @@ -737,13 +733,10 @@ def test_root_backward_returns_to_resting_memory(distributed_setup): x = torch.randn(2, dim, device=device, dtype=dtype, requires_grad=True) output = model(x) loss = output.float().square().mean() - torch.cuda.synchronize(device) - torch.cuda.empty_cache() allocated_before_backward = torch.cuda.memory_allocated(device) loss.backward() del loss, output - torch.cuda.synchronize(device) allocated_after_backward = torch.cuda.memory_allocated(device) extra_allocated = allocated_after_backward - allocated_before_backward child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() @@ -948,8 +941,7 @@ def test_backward_averages_across_dp_and_accumulates_across_calls(distributed_se mesh = init_device_mesh(device.type, (world_size,)) model = nn.Linear(1, world_size, bias=False).to(device) - with torch.no_grad(): - model.weight.fill_(1.0) + nn.init.constant_(model.weight, 1.0) with fully_shard_context(device=device): fully_shard(model, mesh=mesh, placements=_flat_placements()) @@ -973,8 +965,7 @@ def test_next_forward_uses_optimizer_updated_weights(distributed_setup): mesh = init_device_mesh(device.type, (world_size,)) model = nn.Linear(1, world_size, bias=False, dtype=torch.bfloat16).to(device) - with torch.no_grad(): - model.weight.fill_(1.0) + nn.init.constant_(model.weight, 1.0) with fully_shard_context(device=device): fully_shard( @@ -1106,9 +1097,8 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): mesh = init_device_mesh(device.type, (world_size,)) model = nn.Sequential(nn.Linear(4, 4, bias=False), nn.Linear(4, 4, bias=False)) - with torch.no_grad(): - model[0].weight.fill_(2.0) - model[1].weight.fill_(3.0) + nn.init.constant_(model[0].weight, 2.0) + nn.init.constant_(model[1].weight, 3.0) x = torch.ones(1, 4) expected_output = model(x).to(device) @@ -1141,9 +1131,8 @@ def test_meta_parameters_shard_to_mesh_device(distributed_setup): with fully_shard_context(device=device): fully_shard(model, mesh=mesh, placements=_flat_placements()) - with torch.no_grad(): - model[0].weight.fill_(2.0) - model[1].weight.fill_(3.0) + nn.init.constant_(model[0].weight, 2.0) + nn.init.constant_(model[1].weight, 3.0) # The exposed parameters update FP32 main weights, while forward uses separate BF16 # model weights. This simulates load_checkpoint() until # https://github.com/NVIDIA/Megatron-LM/pull/6024 lands and syncs after loading. @@ -1206,13 +1195,11 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten x = torch.randn(batch, dim, device=device, dtype=dtype) torch.cuda.reset_peak_memory_stats(device) train_steps(baseline, baseline_optimizer, x) - torch.cuda.synchronize(device) baseline_peak = torch.cuda.max_memory_allocated(device) del baseline_optimizer del baseline del x - torch.cuda.empty_cache() torch.manual_seed(4321) model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) @@ -1227,12 +1214,10 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten ), ) optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) - torch.cuda.empty_cache() x = torch.randn(batch, dim, device=device, dtype=dtype) torch.cuda.reset_peak_memory_stats(device) train_steps(model, optimizer, x) - torch.cuda.synchronize(device) sharded_peak = torch.cuda.max_memory_allocated(device) logger.info( "FSDP peak memory: rank=%s, baseline=%s, sharded=%s", From afb55ddaf9cce95033d17570dee775efcf169f56 Mon Sep 17 00:00:00 2001 From: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:48:06 +0000 Subject: [PATCH 057/101] Complete explicit GTP RNG seeding (#6569) Signed-off-by: ykarnati --- megatron/core/tensor_parallel/random.py | 10 +++- megatron/training/initialize.py | 26 ++++++++++- .../test_gtp_custom_pgs.py | 46 +++++++++++++++++++ 3 files changed, 78 insertions(+), 4 deletions(-) diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index a9619ea4819..22aaf515a8f 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -454,6 +454,8 @@ def model_parallel_cuda_manual_seed( gtp_remat_rank: Optional[int] = None, egtp_remat_rank: Optional[int] = None, force_reset_rng: bool = False, + gtp_remat_world_size: Optional[int] = None, + egtp_remat_world_size: Optional[int] = None, ): """Initialize model parallel cuda seed. @@ -482,6 +484,10 @@ def model_parallel_cuda_manual_seed( gtp_remat_rank = get_gtp_weight_remat_rank() if egtp_remat_rank is None: egtp_remat_rank = get_expert_gtp_weight_remat_rank() + if gtp_remat_world_size is None: + gtp_remat_world_size = get_gtp_weight_remat_world_size() + if egtp_remat_world_size is None: + egtp_remat_world_size = get_expert_gtp_weight_remat_world_size() # 2718 is just for fun and any POSITIVE value will work. offset = seed + 2718 tensor_model_parallel_seed = offset + tp_rank @@ -506,10 +512,10 @@ def model_parallel_cuda_manual_seed( # must draw DIFFERENT values (everything above is identical across peers by design). The 65536 # stride keeps these disjoint from the tp/ep/etp seeds. Added only when the axis is active, so # non-GTP runs keep a byte-identical tracker set (and checkpoint rng payload). - if get_gtp_weight_remat_world_size() > 1: + if gtp_remat_world_size > 1: gtp_remat_seed = tensor_model_parallel_seed + 65536 * (1 + gtp_remat_rank) _CUDA_RNG_STATE_TRACKER.add(_GTP_REMAT_RNG_TRACKER_NAME, gtp_remat_seed) - if get_expert_gtp_weight_remat_world_size() > 1: + if egtp_remat_world_size > 1: egtp_remat_seed = expert_parallel_seed + 32768 + 65536 * (1 + egtp_remat_rank) _CUDA_RNG_STATE_TRACKER.add(_EXPERT_GTP_REMAT_RNG_TRACKER_NAME, egtp_remat_seed) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index c3eb2cccb2e..61374a4ae08 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -26,7 +26,13 @@ from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( enable_batch_invariant_mode, ) -from megatron.core.utils import get_pg_rank, get_te_version, is_te_min_version, is_torch_min_version +from megatron.core.utils import ( + get_pg_rank, + get_pg_size, + get_te_version, + is_te_min_version, + is_torch_min_version, +) from megatron.training import ( get_adlr_autoresume, get_args, @@ -426,10 +432,12 @@ def _set_random_seed( tp_group: Optional[torch.distributed.ProcessGroup] = None, ep_group: Optional[torch.distributed.ProcessGroup] = None, etp_group: Optional[torch.distributed.ProcessGroup] = None, + gtp_remat_group: Optional[torch.distributed.ProcessGroup] = None, + egtp_remat_group: Optional[torch.distributed.ProcessGroup] = None, ): """Set random seed for reproducability. - The optional pp/dp/tp/ep/etp groups let a caller without an initialized mpu + The optional parallel groups let a caller without an initialized mpu (e.g. a disjoint-grid run) supply the parallel ranks explicitly; each falls back to the mpu group when None. """ @@ -448,6 +456,16 @@ def _set_random_seed( tp_rank = get_pg_rank(tp_group) if tp_group is not None else None ep_rank = get_pg_rank(ep_group) if ep_group is not None else None etp_rank = get_pg_rank(etp_group) if etp_group is not None else None + gtp_remat_rank = get_pg_rank(gtp_remat_group) if gtp_remat_group is not None else None + egtp_remat_rank = ( + get_pg_rank(egtp_remat_group) if egtp_remat_group is not None else None + ) + gtp_remat_world_size = ( + get_pg_size(gtp_remat_group) if gtp_remat_group is not None else None + ) + egtp_remat_world_size = ( + get_pg_size(egtp_remat_group) if egtp_remat_group is not None else None + ) tensor_parallel.model_parallel_cuda_manual_seed( seed, te_rng_tracker, @@ -456,6 +474,10 @@ def _set_random_seed( tp_rank=tp_rank, ep_rank=ep_rank, etp_rank=etp_rank, + gtp_remat_rank=gtp_remat_rank, + egtp_remat_rank=egtp_remat_rank, + gtp_remat_world_size=gtp_remat_world_size, + egtp_remat_world_size=egtp_remat_world_size, ) else: raise ValueError("Seed ({}) should be a positive integer.".format(seed_)) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py index a649c1ac36d..d86da2d54d7 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py @@ -279,6 +279,47 @@ def _worker_partial_pgs_fall_back_to_mpu(rank, world_size, port): ps.initialize_model_parallel() +def _worker_seed_custom_gtp_groups_without_mpu(rank, world_size, port): + """Explicit GTP groups must seed their trackers without MPU groups.""" + from megatron.core import parallel_state as ps + from megatron.core.tensor_parallel.random import ( + get_cuda_rng_tracker, + get_gtp_remat_rng_tracker_name, + ) + from megatron.training.initialize import _set_random_seed + + singleton_group = None + for group_rank in range(world_size): + group = dist.new_group([group_rank]) + if rank == group_rank: + singleton_group = group + + pairing_groups = {} + for name, pairs in _PAIRINGS.items(): + for pair in pairs: + group = dist.new_group(pair) + if rank in pair: + pairing_groups[name] = group + + ps.destroy_model_parallel() + try: + _set_random_seed( + 42, + pp_group=singleton_group, + dp_group=singleton_group, + tp_group=singleton_group, + ep_group=singleton_group, + etp_group=singleton_group, + gtp_remat_group=pairing_groups["adjacent"], + egtp_remat_group=pairing_groups["strided"], + ) + states = get_cuda_rng_tracker().get_states() + assert get_gtp_remat_rng_tracker_name() in states + assert get_gtp_remat_rng_tracker_name(is_expert=True) in states + finally: + ps.initialize_model_parallel() + + class TestGTPCustomProcessGroups: def test_custom_gtp_pg_collection_matches_mpu(self): """A permuted-but-equivalent gtp_remat group must give identical fwd/bwd results.""" @@ -289,3 +330,8 @@ def test_pg_collection_without_gtp_remat_falls_back_to_mpu(self): """Omitting gtp_remat must fall back to the MPU group, not silently disable sharding.""" _requires_multi_gpu(WORLD) _run_distributed(_worker_partial_pgs_fall_back_to_mpu, WORLD) + + def test_custom_gtp_groups_seed_without_mpu(self): + """Explicit GTP groups must not depend on MPU rank or world-size state.""" + _requires_multi_gpu(WORLD) + _run_distributed(_worker_seed_custom_gtp_groups_without_mpu, WORLD) From 43c0e027e493cfb71d5c3b8f67800f2cb894b819 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Mon, 17 Aug 2026 18:08:16 +0000 Subject: [PATCH 058/101] fix(core): warn on invalid RANK env values (#4993) Signed-off-by: Minh Vu --- megatron/core/_rank_utils.py | 27 +++++++++---------- .../fsdp/src/megatron_fsdp/utils.py | 6 +---- .../test_mcore_tensor_parallelism_detect.py | 9 +++++++ tests/unit_tests/test_utils.py | 26 +++++++++++++++++- 4 files changed, 47 insertions(+), 21 deletions(-) diff --git a/megatron/core/_rank_utils.py b/megatron/core/_rank_utils.py index 5d9200e526e..017aa3772de 100644 --- a/megatron/core/_rank_utils.py +++ b/megatron/core/_rank_utils.py @@ -28,21 +28,18 @@ def safe_get_rank() -> int: return torch.distributed.get_rank() # If torch.distributed is not initialized, try to read environment variables. - try: - if "RANK" in os.environ: - return int(os.environ["RANK"]) - - slurm_rank = resolve_slurm_rank() - if slurm_rank is not None: - return slurm_rank - - warnings.warn( - "Could not determine rank from torch.distributed, RANK, or SLURM_PROCID. " - "Defaulting to rank 0." - ) - return 0 - except (ValueError, TypeError): - return 0 + if "RANK" in os.environ: + return int(os.environ["RANK"]) + + slurm_rank = resolve_slurm_rank() + if slurm_rank is not None: + return slurm_rank + + warnings.warn( + "Could not determine rank from torch.distributed, RANK, or SLURM_PROCID. " + "Defaulting to rank 0." + ) + return 0 def safe_get_world_size() -> int: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py index 8fd56795065..88c2257e716 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py @@ -472,11 +472,7 @@ def safe_get_rank() -> int: return torch.distributed.get_rank() # If torch.distributed is not initialized, try to read environment variables. - try: - return int(os.environ.get("RANK", 0)) - except (ValueError, TypeError): - # Return rank 0 regardless of the actual rank. - return 0 + return int(os.environ.get("RANK", 0)) def log_single_rank(logger_: logging.Logger, level: int, msg: str, *args, rank: int = 0, **kwargs): diff --git a/tests/unit_tests/distributed/mfsdp_v1/test_mcore_tensor_parallelism_detect.py b/tests/unit_tests/distributed/mfsdp_v1/test_mcore_tensor_parallelism_detect.py index f548609ad94..0fc76f1afee 100644 --- a/tests/unit_tests/distributed/mfsdp_v1/test_mcore_tensor_parallelism_detect.py +++ b/tests/unit_tests/distributed/mfsdp_v1/test_mcore_tensor_parallelism_detect.py @@ -1,5 +1,6 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +import pytest import torch from torch import nn @@ -87,6 +88,14 @@ def test_safe_get_rank_should_fall_back_to_rank_env_if_distributed_is_not_initia assert safe_get_rank() == 7 +def test_safe_get_rank_should_raise_value_error_if_rank_env_is_invalid(monkeypatch): + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + monkeypatch.setenv("RANK", "not-an-int") + + with pytest.raises(ValueError): + safe_get_rank() + + class DummyConfig: # Just enough attributes for __init__ to run if needed in future tests. gradient_accumulation_fusion = False diff --git a/tests/unit_tests/test_utils.py b/tests/unit_tests/test_utils.py index 0178258cd45..773a74a17f8 100644 --- a/tests/unit_tests/test_utils.py +++ b/tests/unit_tests/test_utils.py @@ -14,7 +14,7 @@ import megatron.core.utils as util import megatron.training.utils as training_util from megatron.core import config -from megatron.core._rank_utils import safe_get_world_size +from megatron.core._rank_utils import safe_get_rank, safe_get_world_size from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_with_transformer_engine_submodules, @@ -577,6 +577,30 @@ def test_invalid_world_size_env_var(self, mock_is_initialized): safe_get_world_size() +class TestGetRankSafe: + """Test safe_get_rank function.""" + + @patch("torch.distributed.is_initialized") + @patch.dict(os.environ, {"RANK": "9"}) + def test_uninitialized_torch_distributed_with_rank_env_var(self, mock_is_initialized): + """Test safe_get_rank when torch.distributed is not initialized but RANK env var exists.""" + mock_is_initialized.return_value = False + + result = safe_get_rank() + + assert result == 9 + mock_is_initialized.assert_called_once() + + @patch("torch.distributed.is_initialized") + @patch.dict(os.environ, {"RANK": "invalid"}) + def test_invalid_rank_env_var_should_raise_value_error(self, mock_is_initialized): + """Test safe_get_rank with invalid RANK environment variable.""" + mock_is_initialized.return_value = False + + with pytest.raises(ValueError): + safe_get_rank() + + class TestGetLocalRankPreinit: """Test get_local_rank_preinit function.""" From 86380223f23ddc2c8ad10f6bec30705ab5cf36f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Mon, 17 Aug 2026 18:25:31 +0000 Subject: [PATCH 059/101] build(container): support non-root CI execution (#6421) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- docker/Dockerfile.ci.dev | 77 ++++++++++++++++++- tests/test_utils/recipes/gb200/gpt-1node.yaml | 8 +- .../recipes/gb200/gpt-dynamic-inference.yaml | 8 +- .../recipes/gb200/gpt-perf-dp4.yaml | 8 +- tests/test_utils/recipes/gb200/gpt-perf.yaml | 8 +- tests/test_utils/recipes/gb200/gpt.yaml | 8 +- .../recipes/gb200/hybrid-perf-ep4.yaml | 8 +- .../test_utils/recipes/gb200/hybrid-perf.yaml | 8 +- tests/test_utils/recipes/gb200/moe-1node.yaml | 8 +- .../recipes/gb200/moe-dynamic-inference.yaml | 8 +- tests/test_utils/recipes/gb200/moe.yaml | 8 +- .../recipes/gb200/nemotron-1node.yaml | 8 +- tests/test_utils/recipes/gb200/nemotron.yaml | 8 +- .../test_utils/recipes/gb200/unit-tests.yaml | 8 +- tests/test_utils/recipes/h100/bert.yaml | 8 +- .../recipes/h100/ckpt_converter.yaml | 8 +- .../recipes/h100/determinism-perf.yaml | 8 +- tests/test_utils/recipes/h100/flextron.yaml | 8 +- .../gpt-dynamic-inference-cuda-graphs.yaml | 8 +- ...pt-dynamic-inference-with-coordinator.yaml | 8 +- .../recipes/h100/gpt-dynamic-inference.yaml | 8 +- tests/test_utils/recipes/h100/gpt-grads.yaml | 8 +- tests/test_utils/recipes/h100/gpt-grpo.yaml | 8 +- .../h100/gpt-inference-server-smoke.yaml | 8 +- tests/test_utils/recipes/h100/gpt-nemo.yaml | 8 +- .../recipes/h100/gpt-offline-inference.yaml | 8 +- .../test_utils/recipes/h100/gpt-perf-dp8.yaml | 8 +- tests/test_utils/recipes/h100/gpt-perf.yaml | 8 +- .../recipes/h100/gpt-static-inference.yaml | 8 +- tests/test_utils/recipes/h100/gpt.yaml | 8 +- .../recipes/h100/hybrid-perf-ep8.yaml | 8 +- .../test_utils/recipes/h100/hybrid-perf.yaml | 8 +- ...ba-dynamic-inference-with-coordinator.yaml | 8 +- .../recipes/h100/mamba-dynamic-inference.yaml | 8 +- .../recipes/h100/mamba-static-inference.yaml | 8 +- tests/test_utils/recipes/h100/mamba.yaml | 8 +- tests/test_utils/recipes/h100/mimo.yaml | 8 +- .../recipes/h100/module_performance.yaml | 8 +- ...oe-dynamic-inference-with-coordinator.yaml | 8 +- .../recipes/h100/moe-dynamic-inference.yaml | 8 +- tests/test_utils/recipes/h100/moe-grpo.yaml | 8 +- .../recipes/h100/moe-static-inference.yaml | 8 +- tests/test_utils/recipes/h100/moe.yaml | 8 +- .../recipes/h100/multimodal-llava.yaml | 8 +- tests/test_utils/recipes/h100/t5.yaml | 8 +- tests/test_utils/recipes/h100/unit-tests.yaml | 8 +- 46 files changed, 255 insertions(+), 182 deletions(-) diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index bd342e613b1..f92148059d5 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -6,10 +6,11 @@ FROM ${FROM_IMAGE_NAME} AS main ENV PIP_CONSTRAINT="" ENV DEBIAN_FRONTEND=noninteractive ARG UV_VERSION=0.7.2 +ARG TARGETARCH ARG YQ_VERSION=4.44.1 # TE 2.18's NCCL EP build requires NCCL 2.30.7+; the 26.06 base uses CUDA 13.3. ARG NCCL_VERSION=2.30.7-1+cuda13.3 -ENV PATH="/root/.local/bin:$PATH" +ENV PATH="/usr/local/bin:$PATH" ARG UV_PROJECT_ENVIRONMENT=/opt/venv ENV UV_PROJECT_ENVIRONMENT=${UV_PROJECT_ENVIRONMENT} ENV VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT @@ -30,7 +31,19 @@ RUN bash -ex <<"EOF" esac wget https://github.com/mikefarah/yq/releases/download/v${YQ_VERSION}/yq_linux_${YQ_ARCH} -O /usr/local/bin/yq chmod a+x /usr/local/bin/yq - curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh + case "${TARGETARCH}" in \ + "amd64") UV_ARCH=x86_64; UV_SHA256=cfaab1b5166a6439ff66f020333d3a12bbdf622deee3b510718283e8f06c9de7 ;; \ + "arm64") UV_ARCH=aarch64; UV_SHA256=2872fdf4785666575d129ba90590c44e6508e22584735f3e7e8a30d773dfc3db ;; \ + *) echo "Unsupported TARGETARCH for uv: ${TARGETARCH}" >&2; exit 1 ;; \ + esac + UV_ARCHIVE="uv-${UV_ARCH}-unknown-linux-gnu.tar.gz" + curl -fLSs -o "/tmp/${UV_ARCHIVE}" \ + "https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/${UV_ARCHIVE}" + echo "${UV_SHA256} /tmp/${UV_ARCHIVE}" | sha256sum -c - + tar -xzf "/tmp/${UV_ARCHIVE}" -C /tmp + install -m 0755 "/tmp/uv-${UV_ARCH}-unknown-linux-gnu/uv" /usr/local/bin/uv + install -m 0755 "/tmp/uv-${UV_ARCH}-unknown-linux-gnu/uvx" /usr/local/bin/uvx + rm -rf "/tmp/${UV_ARCHIVE}" "/tmp/uv-${UV_ARCH}-unknown-linux-gnu" EOF RUN --mount=type=bind,source=docker/common/install_nccl.sh,target=/opt/install_nccl.sh \ @@ -126,6 +139,64 @@ EOF COPY assets/ /opt/data/ ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python +ARG RUNTIME_UID=65532 +ARG RUNTIME_GID=65532 +RUN if ! getent group "${RUNTIME_GID}" >/dev/null; then \ + groupadd --gid "${RUNTIME_GID}" nemo-runtime; \ + fi && \ + if ! getent passwd "${RUNTIME_UID}" >/dev/null; then \ + useradd --no-log-init --uid "${RUNTIME_UID}" --gid "${RUNTIME_GID}" \ + --create-home --home-dir /home/nemo-runtime --shell /bin/bash nemo-runtime; \ + fi && \ + install -d -o "${RUNTIME_UID}" -g "${RUNTIME_GID}" \ + /home/nemo-runtime \ + /home/gitlab-runner/data/gpt3_data \ + /opt/megatron-lm \ + /tmp/regent-cache \ + /tmp/regent-cache/pip \ + /tmp/regent-cache/uv \ + /tmp/regent-cache/xdg \ + /workspace/data/cache && \ + venv_site_packages="$("${UV_PROJECT_ENVIRONMENT}/bin/python" -c \ + 'import site; print(site.getsitepackages()[0])')" && \ + touch "${venv_site_packages}/.nonroot-uninstall-probe" && \ + find "${UV_PROJECT_ENVIRONMENT}" -type d \ + -exec chown "${RUNTIME_UID}:${RUNTIME_GID}" {} + + +USER ${RUNTIME_UID}:${RUNTIME_GID} + +# The CI harness materializes an immutable test commit here rather than over +# the dependency-build checkout in /workspace. Runtime callers select this UID +# and set HOME/cache variables explicitly; the published image remains root by +# default. Prove that the non-root contract works before restoring that default. +# The cache root is probed for directory creation, not just file writes: a +# caller pointing XDG_CACHE_HOME or HF_HOME below it materializes new +# subdirectories, which a root-owned parent rejects even when every pre-created +# leaf is writable. +RUN export HOME=/home/nemo-runtime \ + PIP_CACHE_DIR=/tmp/regent-cache/pip \ + XDG_CACHE_HOME=/tmp/regent-cache/xdg \ + UV_CACHE_DIR=/tmp/regent-cache/uv && \ + test "$(id -u)" = "${RUNTIME_UID}" && \ + test -w "${HOME}" && \ + test -w "${PIP_CACHE_DIR}" && \ + test -w "${XDG_CACHE_HOME}" && \ + test -w "${UV_CACHE_DIR}" && \ + test -x /usr/local/bin/uv && \ + test -x /opt/venv/bin/python && \ + venv_site_packages="$(python -c 'import site; print(site.getsitepackages()[0])')" && \ + rm "${venv_site_packages}/.nonroot-uninstall-probe" && \ + touch "${venv_site_packages}/.nonroot-write-probe" && \ + rm "${venv_site_packages}/.nonroot-write-probe" && \ + touch /workspace/data/cache/.nonroot-write-probe && \ + rm /workspace/data/cache/.nonroot-write-probe && \ + touch /opt/megatron-lm/.nonroot-write-probe && \ + rm /opt/megatron-lm/.nonroot-write-probe && \ + mkdir /tmp/regent-cache/.nonroot-mkdir-probe && \ + rmdir /tmp/regent-cache/.nonroot-mkdir-probe + +USER root + ##### For NVIDIANS only ##### FROM main AS jet ARG JET_API_VERSION @@ -145,4 +216,6 @@ RUN --mount=type=secret,id=JET_INDEX_URLS \ uv pip install --no-cache-dir --upgrade "setuptools>=80" uv pip install --no-cache-dir --upgrade $JET_INDEX_URLS "jet-client~=4.0" EOF + +USER root ### diff --git a/tests/test_utils/recipes/gb200/gpt-1node.yaml b/tests/test_utils/recipes/gb200/gpt-1node.yaml index 3485dc41d25..2465b68a6d6 100644 --- a/tests/test_utils/recipes/gb200/gpt-1node.yaml +++ b/tests/test_utils/recipes/gb200/gpt-1node.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail : "${{RUN_ID:=pr-$$}}" diff --git a/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml b/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml index 4287eea3bae..81624c5c5ba 100644 --- a/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml +++ b/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml b/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml index fa0539c77d9..50dfae839e0 100644 --- a/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml +++ b/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml @@ -18,8 +18,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -29,10 +29,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/gb200/gpt-perf.yaml b/tests/test_utils/recipes/gb200/gpt-perf.yaml index 14717ed1654..6dcdeb88d78 100644 --- a/tests/test_utils/recipes/gb200/gpt-perf.yaml +++ b/tests/test_utils/recipes/gb200/gpt-perf.yaml @@ -20,8 +20,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -31,10 +31,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/gb200/gpt.yaml b/tests/test_utils/recipes/gb200/gpt.yaml index 1142861ea20..aaf8f4e6876 100644 --- a/tests/test_utils/recipes/gb200/gpt.yaml +++ b/tests/test_utils/recipes/gb200/gpt.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail : "${{RUN_ID:=pr-$$}}" diff --git a/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml b/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml index 99f9ed67af6..ab4d2f1b04c 100644 --- a/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml +++ b/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml @@ -18,8 +18,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -29,10 +29,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/gb200/hybrid-perf.yaml b/tests/test_utils/recipes/gb200/hybrid-perf.yaml index 1f17e5d1590..2a25ef04032 100644 --- a/tests/test_utils/recipes/gb200/hybrid-perf.yaml +++ b/tests/test_utils/recipes/gb200/hybrid-perf.yaml @@ -19,8 +19,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -30,10 +30,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/gb200/moe-1node.yaml b/tests/test_utils/recipes/gb200/moe-1node.yaml index 330e911c6c5..fca78124824 100644 --- a/tests/test_utils/recipes/gb200/moe-1node.yaml +++ b/tests/test_utils/recipes/gb200/moe-1node.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail : "${{RUN_ID:=pr-$$}}" diff --git a/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml b/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml index d3fc99b193d..39bd46086ee 100644 --- a/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml +++ b/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/gb200/moe.yaml b/tests/test_utils/recipes/gb200/moe.yaml index 79799201386..ff6f155c3e1 100644 --- a/tests/test_utils/recipes/gb200/moe.yaml +++ b/tests/test_utils/recipes/gb200/moe.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail : "${{RUN_ID:=pr-$$}}" diff --git a/tests/test_utils/recipes/gb200/nemotron-1node.yaml b/tests/test_utils/recipes/gb200/nemotron-1node.yaml index ee61a7246cd..b39bc6a4af2 100644 --- a/tests/test_utils/recipes/gb200/nemotron-1node.yaml +++ b/tests/test_utils/recipes/gb200/nemotron-1node.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail : "${{RUN_ID:=pr-$$}}" diff --git a/tests/test_utils/recipes/gb200/nemotron.yaml b/tests/test_utils/recipes/gb200/nemotron.yaml index 9aae75eaebf..f41a1e877f6 100644 --- a/tests/test_utils/recipes/gb200/nemotron.yaml +++ b/tests/test_utils/recipes/gb200/nemotron.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -37,10 +37,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail : "${{RUN_ID:=pr-$$}}" diff --git a/tests/test_utils/recipes/gb200/unit-tests.yaml b/tests/test_utils/recipes/gb200/unit-tests.yaml index 0b462f95cf5..c972254463c 100644 --- a/tests/test_utils/recipes/gb200/unit-tests.yaml +++ b/tests/test_utils/recipes/gb200/unit-tests.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -32,10 +32,10 @@ spec: --backwards-commit $MCORE_BACKWARDS_COMMIT \ --repo $MCORE_REPO - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/bert.yaml b/tests/test_utils/recipes/h100/bert.yaml index b0e35676169..d5e366ffacc 100644 --- a/tests/test_utils/recipes/h100/bert.yaml +++ b/tests/test_utils/recipes/h100/bert.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -37,10 +37,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/ckpt_converter.yaml b/tests/test_utils/recipes/h100/ckpt_converter.yaml index 03719103548..bc726d9df0f 100644 --- a/tests/test_utils/recipes/h100/ckpt_converter.yaml +++ b/tests/test_utils/recipes/h100/ckpt_converter.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/determinism-perf.yaml b/tests/test_utils/recipes/h100/determinism-perf.yaml index d1d14bf0749..2545fd973f8 100644 --- a/tests/test_utils/recipes/h100/determinism-perf.yaml +++ b/tests/test_utils/recipes/h100/determinism-perf.yaml @@ -12,8 +12,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -23,10 +23,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/h100/flextron.yaml b/tests/test_utils/recipes/h100/flextron.yaml index 783d6b1476b..9864a566b32 100644 --- a/tests/test_utils/recipes/h100/flextron.yaml +++ b/tests/test_utils/recipes/h100/flextron.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml b/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml index 23a67467176..a9336dd036c 100644 --- a/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml +++ b/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml index 55a07a36090..bbc02dac20c 100644 --- a/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml +++ b/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml b/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml index f41fc3cdca4..6045b35c219 100644 --- a/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml +++ b/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt-grads.yaml b/tests/test_utils/recipes/h100/gpt-grads.yaml index 39c92c20c9a..5c7cf5a8f8b 100644 --- a/tests/test_utils/recipes/h100/gpt-grads.yaml +++ b/tests/test_utils/recipes/h100/gpt-grads.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt-grpo.yaml b/tests/test_utils/recipes/h100/gpt-grpo.yaml index 82d9d44b9a9..9e333cfebe3 100644 --- a/tests/test_utils/recipes/h100/gpt-grpo.yaml +++ b/tests/test_utils/recipes/h100/gpt-grpo.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml b/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml index 9998e9a4b74..d1ed10b2f36 100644 --- a/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml +++ b/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml @@ -13,8 +13,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -34,10 +34,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/h100/gpt-nemo.yaml b/tests/test_utils/recipes/h100/gpt-nemo.yaml index 75fbbfa5c58..d9c56bd52fa 100644 --- a/tests/test_utils/recipes/h100/gpt-nemo.yaml +++ b/tests/test_utils/recipes/h100/gpt-nemo.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -37,10 +37,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt-offline-inference.yaml b/tests/test_utils/recipes/h100/gpt-offline-inference.yaml index 472245ed476..9e1804c6e67 100644 --- a/tests/test_utils/recipes/h100/gpt-offline-inference.yaml +++ b/tests/test_utils/recipes/h100/gpt-offline-inference.yaml @@ -13,8 +13,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -34,10 +34,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN ls cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml b/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml index 834920b2995..b1c6f43ca83 100644 --- a/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml +++ b/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml @@ -17,8 +17,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -28,10 +28,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/h100/gpt-perf.yaml b/tests/test_utils/recipes/h100/gpt-perf.yaml index 2b0f922b9c8..be4656addcf 100644 --- a/tests/test_utils/recipes/h100/gpt-perf.yaml +++ b/tests/test_utils/recipes/h100/gpt-perf.yaml @@ -16,8 +16,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -27,10 +27,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/h100/gpt-static-inference.yaml b/tests/test_utils/recipes/h100/gpt-static-inference.yaml index 8331f048721..66771761421 100644 --- a/tests/test_utils/recipes/h100/gpt-static-inference.yaml +++ b/tests/test_utils/recipes/h100/gpt-static-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/gpt.yaml b/tests/test_utils/recipes/h100/gpt.yaml index 10f2bebf722..c2a1f17eea3 100644 --- a/tests/test_utils/recipes/h100/gpt.yaml +++ b/tests/test_utils/recipes/h100/gpt.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml b/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml index a4e501912f1..2b64f923e3e 100644 --- a/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml +++ b/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml @@ -16,8 +16,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -27,10 +27,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/h100/hybrid-perf.yaml b/tests/test_utils/recipes/h100/hybrid-perf.yaml index 9efbb27b740..3e37ec95d7c 100644 --- a/tests/test_utils/recipes/h100/hybrid-perf.yaml +++ b/tests/test_utils/recipes/h100/hybrid-perf.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -26,10 +26,10 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail cd /opt/megatron-lm diff --git a/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml index e58d70fe70c..53e912131ee 100644 --- a/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml +++ b/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml index bf3c6e26d67..f5e6b7f8d8c 100644 --- a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml +++ b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/mamba-static-inference.yaml b/tests/test_utils/recipes/h100/mamba-static-inference.yaml index 8450dba539f..29cfb2e0a67 100644 --- a/tests/test_utils/recipes/h100/mamba-static-inference.yaml +++ b/tests/test_utils/recipes/h100/mamba-static-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index bea183b36c2..0a338d91b61 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/mimo.yaml b/tests/test_utils/recipes/h100/mimo.yaml index 26f68251d55..4aea76c6b19 100644 --- a/tests/test_utils/recipes/h100/mimo.yaml +++ b/tests/test_utils/recipes/h100/mimo.yaml @@ -20,8 +20,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -42,10 +42,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/module_performance.yaml b/tests/test_utils/recipes/h100/module_performance.yaml index 63530d825e6..19d58d03ef8 100644 --- a/tests/test_utils/recipes/h100/module_performance.yaml +++ b/tests/test_utils/recipes/h100/module_performance.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml index 035927b2e17..b51e5188cdb 100644 --- a/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml +++ b/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml b/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml index ca81d0fbba1..933d1d25551 100644 --- a/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml +++ b/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/moe-grpo.yaml b/tests/test_utils/recipes/h100/moe-grpo.yaml index a15bb590fe2..3793815b8d1 100644 --- a/tests/test_utils/recipes/h100/moe-grpo.yaml +++ b/tests/test_utils/recipes/h100/moe-grpo.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/moe-static-inference.yaml b/tests/test_utils/recipes/h100/moe-static-inference.yaml index 96578b3d1c2..a7f63591ad7 100644 --- a/tests/test_utils/recipes/h100/moe-static-inference.yaml +++ b/tests/test_utils/recipes/h100/moe-static-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/moe.yaml b/tests/test_utils/recipes/h100/moe.yaml index 4cbcd298b23..7620eed364b 100644 --- a/tests/test_utils/recipes/h100/moe.yaml +++ b/tests/test_utils/recipes/h100/moe.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,10 +36,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/multimodal-llava.yaml b/tests/test_utils/recipes/h100/multimodal-llava.yaml index 90b23f438c2..f37a81d6648 100644 --- a/tests/test_utils/recipes/h100/multimodal-llava.yaml +++ b/tests/test_utils/recipes/h100/multimodal-llava.yaml @@ -17,8 +17,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -39,10 +39,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/t5.yaml b/tests/test_utils/recipes/h100/t5.yaml index 648afe9b34f..7fd369b2070 100644 --- a/tests/test_utils/recipes/h100/t5.yaml +++ b/tests/test_utils/recipes/h100/t5.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,10 +35,10 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls diff --git a/tests/test_utils/recipes/h100/unit-tests.yaml b/tests/test_utils/recipes/h100/unit-tests.yaml index e3c8e3bea85..c37edef8dd3 100644 --- a/tests/test_utils/recipes/h100/unit-tests.yaml +++ b/tests/test_utils/recipes/h100/unit-tests.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> /root/.netrc - trap 'rm -f /root/.netrc; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" + trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -32,10 +32,10 @@ spec: --backwards-commit $MCORE_BACKWARDS_COMMIT \ --repo $MCORE_REPO - rm -f /root/.netrc + rm -f "${HOME}/.netrc" unset RO_API_TOKEN script: |- - rm -f /root/.netrc + rm -f "${{HOME}}/.netrc" unset RO_API_TOKEN set -euo pipefail ls From 51270a3adb2f56c800d0bb161446169bbf7dbd54 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Mon, 17 Aug 2026 19:00:25 +0000 Subject: [PATCH 060/101] Add flash-linear-attention GDP kernels (#6538) Signed-off-by: Keshav Santhanam --- LICENSE | 25 ++ megatron/core/ssm/gated_delta_product.py | 17 +- megatron/core/ssm/ops/__init__.py | 1 + megatron/core/ssm/ops/common/__init__.py | 5 +- megatron/core/ssm/ops/gdp/__init__.py | 26 ++ megatron/core/ssm/ops/gdp/chunk.py | 178 +++++++++++ megatron/core/ssm/ops/gdp/chunk_h.py | 289 ++++++++++++++++++ megatron/core/ssm/ops/gdp/chunk_o.py | 215 +++++++++++++ megatron/core/ssm/ops/gdp/common.py | 250 ++++++++++++++++ megatron/core/ssm/ops/gdp/cumsum.py | 124 ++++++++ megatron/core/ssm/ops/gdp/fused_recurrent.py | 198 ++++++++++++ megatron/core/ssm/ops/gdp/scaled_dot_kkt.py | 141 +++++++++ megatron/core/ssm/ops/gdp/solve_tril.py | 300 +++++++++++++++++++ megatron/core/ssm/ops/gdp/wy_fast.py | 166 ++++++++++ 14 files changed, 1928 insertions(+), 7 deletions(-) create mode 100644 megatron/core/ssm/ops/gdp/__init__.py create mode 100644 megatron/core/ssm/ops/gdp/chunk.py create mode 100644 megatron/core/ssm/ops/gdp/chunk_h.py create mode 100644 megatron/core/ssm/ops/gdp/chunk_o.py create mode 100644 megatron/core/ssm/ops/gdp/common.py create mode 100644 megatron/core/ssm/ops/gdp/cumsum.py create mode 100644 megatron/core/ssm/ops/gdp/fused_recurrent.py create mode 100644 megatron/core/ssm/ops/gdp/scaled_dot_kkt.py create mode 100644 megatron/core/ssm/ops/gdp/solve_tril.py create mode 100644 megatron/core/ssm/ops/gdp/wy_fast.py diff --git a/LICENSE b/LICENSE index 4570d9283c4..1f1ee0241bd 100644 --- a/LICENSE +++ b/LICENSE @@ -325,3 +325,28 @@ LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +-------------------------------------------------------------------------------- +LICENSE FOR fla-org/flash-linear-attention. + +MIT License + +Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index a9d939b6a15..22754aa981e 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -73,12 +73,19 @@ try: from fla.ops.gated_delta_product import chunk_gated_delta_product - from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule HAVE_FLA = True except ImportError: HAVE_FLA = False +# Dynamic-batching inference runs the in-tree fork of these kernels rather than +# the pip `flash-linear-attention` ones. The fork is forward-only, so training +# and the static-batching path (which shares the training body) keep calling +# upstream, which owns the backward pass. +from megatron.core.ssm.ops.gdp import ( + chunk_gated_delta_product_varlen, + fused_recurrent_gated_delta_rule_update, +) logger = logging.getLogger(__name__) @@ -670,7 +677,7 @@ def ssm_decode( # for the scatter below; the padding rows' outputs are never scattered back. initial_state = ssm_state[batch_indices.clamp(min=0)] - core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule( + core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule_update( query, key, value, @@ -778,17 +785,17 @@ def ssm_prefill( query = query.repeat_interleave(rep, dim=2) key = key.repeat_interleave(rep, dim=2) - core_attn_out, last_recurrent_state = chunk_gated_delta_product( + core_attn_out, last_recurrent_state = chunk_gated_delta_product_varlen( query, key, value, g=g, beta=beta, + num_householder=M, + cu_seqlens=cu_seqlens, initial_state=None, output_final_state=True, - num_householder=M, use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, ) # Write per-request final SSM states into the cache for subsequent decode. diff --git a/megatron/core/ssm/ops/__init__.py b/megatron/core/ssm/ops/__init__.py index 2a42df6bef9..84386bc7675 100644 --- a/megatron/core/ssm/ops/__init__.py +++ b/megatron/core/ssm/ops/__init__.py @@ -8,6 +8,7 @@ gather/scatter, autotune determinism). * `mamba2` -- the Mamba2 state-space duality kernels, adapted from vLLM and state-spaces/mamba. +* `gdp` -- the Gated Delta Product kernels, forked from flash-linear-attention. The re-exports below are the historical top-level entry points and are kept so existing callers do not have to spell out the subpackage. diff --git a/megatron/core/ssm/ops/common/__init__.py b/megatron/core/ssm/ops/common/__init__.py index f968191a055..bc5f879dd55 100644 --- a/megatron/core/ssm/ops/common/__init__.py +++ b/megatron/core/ssm/ops/common/__init__.py @@ -5,6 +5,7 @@ Nothing here is specific to a recurrence. The short causal convolution that front-runs the recurrence, the slot-indexed state gather/scatter used by prefix caching, and the autotune-configuration filtering that makes the kernels -deterministic are all common to Mamba2 and its relatives. Recurrence-specific -kernels live in the sibling `mamba2` package. +deterministic are all common to Mamba2, Gated Delta Product and their +relatives. Recurrence-specific kernels live in the sibling `mamba2` and `gdp` +packages. """ diff --git a/megatron/core/ssm/ops/gdp/__init__.py b/megatron/core/ssm/ops/gdp/__init__.py new file mode 100644 index 00000000000..4d5468bbb8d --- /dev/null +++ b/megatron/core/ssm/ops/gdp/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# The modules in this package are forked from the Gated Delta Product kernels +# in flash-linear-attention v0.5.1 +# (https://github.com/fla-org/flash-linear-attention), licensed under the MIT +# license. See the LICENSE file in this directory. + +"""Inference-only fork of the Gated Delta Product kernels from +`flash-linear-attention `_ +(v0.5.1). + +Training and the static-batching inference path call the pip +`flash-linear-attention` kernels, which own the backward pass. Only the +dynamic-batching decode and prefill steps route here: + +* `fused_recurrent_gated_delta_rule_update` -- decode. +* `chunk_gated_delta_product_varlen` -- prefill. + +Both entry points are forward-only. +""" + +from .chunk import chunk_gated_delta_product_varlen +from .fused_recurrent import fused_recurrent_gated_delta_rule_update + +__all__ = ["chunk_gated_delta_product_varlen", "fused_recurrent_gated_delta_rule_update"] diff --git a/megatron/core/ssm/ops/gdp/chunk.py b/megatron/core/ssm/ops/gdp/chunk.py new file mode 100644 index 00000000000..58c89eeeddc --- /dev/null +++ b/megatron/core/ssm/ops/gdp/chunk.py @@ -0,0 +1,178 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/gated_delta_product/chunk.py` in flash-linear-attention v0.5.1 +# (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Chunked Gated Delta Product prefill. + +The stage sequence: + +1. Interleave the log decays across the Householder copies, then take the + within-chunk cumulative sum on both the unexpanded and the expanded stream, + converting to base 2 on the way (the kernels exponentiate with `exp2`). +2. Build the WY representation: `beta * K K^T`, then invert `I + A`, then + recompute the `w` and `u` factors. +3. Sweep the inter-chunk state recurrence, emitting the chunk-boundary states. +4. Combine those with the within-chunk attention to get the outputs. + +Two chunkings are in play throughout. The queries and the outputs live on the +token stream as written; the keys, values and betas live on the +Householder-expanded stream, whose sequences are `M` times longer. The second +chunking is not a rescaling of the first, because `ceil(L*M/64)` is not +`M*ceil(L/64)` unless `L` is a multiple of the chunk size. +""" + +import torch + +from .chunk_h import chunk_gated_delta_product_fwd_h +from .chunk_o import chunk_gated_delta_product_fwd_o +from .common import CHUNK_SIZE, RCP_LN2, l2norm_fwd, prepare_chunk_indices +from .cumsum import chunk_local_cumsum +from .scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from .solve_tril import solve_tril +from .wy_fast import recompute_w_u_fwd + + +def chunk_gated_delta_product_varlen( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + num_householder: int, + cu_seqlens: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Variable-length chunked Gated Delta Product forward pass. + + Args: + q: Queries `[1, T, H, K]`. + k: Keys `[1, T*M, H, K]` (Householder-expanded). + v: Values `[1, T*M, H, V]` (Householder-expanded). + g: Log decays `[1, T, H]`. + beta: Betas `[1, T*M, H]`. + num_householder: Number of Householder copies `M`. + cu_seqlens: Sequence boundaries over the unexpanded stream, `[N+1]`. + scale: Score scale; defaults to `K ** -0.5`. + initial_state: Starting state `[N, H, K, V]`, or `None` for zeros. + output_final_state: Whether to return the final state. + use_qk_l2norm_in_kernel: Whether to L2-normalize `q` and `k` first. + + Returns `(o, final_state)` with `o` shaped `[1, T, H, V]`. + """ + B, T, H, K = q.shape + V = v.shape[-1] + assert q.dtype != torch.float32, "the chunked GDP kernels require bf16/fp16 inputs" + assert B == 1, f"varlen prefill expects a single packed sequence, got batch {B}" + assert k.shape == (B, T * num_householder, H, K), f"unexpected key shape {tuple(k.shape)}" + assert v.shape == (B, T * num_householder, H, V), f"unexpected value shape {tuple(v.shape)}" + assert beta.shape == (B, T * num_householder, H), f"unexpected beta shape {tuple(beta.shape)}" + if g is not None: + assert g.shape == (B, T, H), f"unexpected decay shape {tuple(g.shape)}" + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + "The number of initial states is expected to be equal to the number of input " + f"sequences, i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = K**-0.5 + + # The kernels index with raw pointer arithmetic and would read garbage from + # a strided input, so force every tensor argument contiguous. + q, k, v, beta = q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous() + if g is not None: + g = g.contiguous() + if initial_state is not None: + initial_state = initial_state.contiguous() + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + + cu_seqlens_dp = cu_seqlens * num_householder + # Both chunkings are derived here, once, and threaded through every stage + # so the sub-kernels do not each re-derive (and re-synchronize on) them. + chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) + chunk_indices_dp = prepare_chunk_indices(cu_seqlens_dp, CHUNK_SIZE) + + if g is not None: + # The decay applies to the first Householder copy of each token; the + # remaining copies are decay-free (zeros in log space). + g_interleaved = g.new_zeros(B, T, num_householder, H, dtype=torch.float32) + g_interleaved[:, :, 0] = g + g_interleaved = g_interleaved.view(B, T * num_householder, H).contiguous() + # The chunked kernels exponentiate in base 2, so the decays are + # converted out of natural-log space here rather than in every kernel. + g = chunk_local_cumsum( + g, + chunk_size=CHUNK_SIZE, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + output_dtype=torch.float32, + chunk_indices=chunk_indices, + ) + g_interleaved = chunk_local_cumsum( + g_interleaved, + chunk_size=CHUNK_SIZE, + scale=RCP_LN2, + cu_seqlens=cu_seqlens_dp, + output_dtype=torch.float32, + chunk_indices=chunk_indices_dp, + ) + else: + g_interleaved = None + + # WY representation of the (inverse) transition matrix. u is the new v. + A = chunk_scaled_dot_kkt_fwd( + k=k, + g=g_interleaved, + beta=beta, + cu_seqlens=cu_seqlens_dp, + chunk_size=CHUNK_SIZE, + output_dtype=torch.float32, + chunk_indices=chunk_indices_dp, + ) + A = solve_tril( + A=A, cu_seqlens=cu_seqlens_dp, chunk_indices=chunk_indices_dp, output_dtype=k.dtype + ) + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g_interleaved, + cu_seqlens=cu_seqlens_dp, + chunk_indices=chunk_indices_dp, + ) + + h, v_new, final_state = chunk_gated_delta_product_fwd_h( + k=k, + w=w, + u=u, + g=g_interleaved, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens_dp, + num_householder=num_householder, + chunk_size=CHUNK_SIZE, + chunk_indices=chunk_indices, + ) + o = chunk_gated_delta_product_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=CHUNK_SIZE, + num_householder=num_householder, + chunk_indices=chunk_indices, + ) + return o.to(q.dtype), final_state diff --git a/megatron/core/ssm/ops/gdp/chunk_h.py b/megatron/core/ssm/ops/gdp/chunk_h.py new file mode 100644 index 00000000000..e922b9061a4 --- /dev/null +++ b/megatron/core/ssm/ops/gdp/chunk_h.py @@ -0,0 +1,289 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/gated_delta_product/chunk_deltaproduct_h.py` in +# flash-linear-attention v0.5.1 (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Inter-chunk state recurrence for the Gated Delta Product. + +Carries the matrix-valued state forward across chunks, emitting the state at +each chunk boundary (for the intra-chunk output kernel) and the corrected +values `v_new`. The Householder expansion shows up here as the stride between +state stores: the state is only checkpointed once every `num_householder` +expanded chunks, because that is one chunk of the original token stream. + +The key dimension is handled in fixed 64-wide registers, which is why `K` is +capped at 256. +""" + +import torch + +from .common import HAVE_TRITON, exp2, prepare_chunk_indices, prepare_chunk_offsets, tl, triton + + +@triton.heuristics( + { + 'USE_G': lambda args: args['g'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + for BV in [32, 64] + ], + key=['H', 'K', 'V', 'BT', 'USE_G'], +) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_product_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + num_householder: tl.constexpr, # number of delta products + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """Sweep one sequence's chunks, carrying the `[K, V]` state in registers.""" + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * tl.cdiv(T // num_householder, BT) + + # [BK, BV] + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * H + i_h) * K * V + v += (bos * H + i_h) * V + k += (bos * H + i_h) * K + w += (bos * H + i_h) * K + if SAVE_NEW_VALUE: + v_new += (bos * H + i_h) * V + stride_v = H * V + stride_h = H * K * V + stride_k = H * K + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + # load initial state + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + # main recurrence + for i_t in range(NT): + if i_t % num_householder == 0: + i_t_true = i_t // num_householder + p_h1 = tl.make_block_ptr( + h + i_t_true * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr( + h + i_t_true * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr( + h + i_t_true * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr( + h + i_t_true * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v_new = tl.zeros([BT, BV], dtype=tl.float32) + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype)) + b_v_new = -b_v_new + tl.load(p_v, boundary_check=(0, 1)) + + if SAVE_NEW_VALUE: + p_v_new = tl.make_block_ptr( + v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + tl.store(p_v_new, b_v_new.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v_new = b_v_new * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + b_h1 = b_h1 * b_g_last + if K > 64: + b_h2 = b_h2 * b_g_last + if K > 128: + b_h3 = b_h3 * b_g_last + if K > 192: + b_h4 = b_h4 * b_g_last + b_v_new = b_v_new.to(k.dtype.element_ty) + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.dot(b_k, b_v_new) + if K > 64: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.dot(b_k, b_v_new) + if K > 128: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.dot(b_k, b_v_new) + if K > 192: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.dot(b_k, b_v_new) + # epilogue + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_product_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.Tensor | None = None, + num_householder: int = 1, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the inter-chunk state recurrence. + + Args: + k: Keys `[B, T, H, K]` on the Householder-expanded token stream. + w: WY key factor, shaped like `k`. + u: WY value factor `[B, T, H, V]`. + g: Within-chunk cumulative log2 decays on the expanded stream, or `None`. + initial_state: Starting state `[N, H, K, V]`, or `None` for zeros. + output_final_state: Whether to return the final state. + chunk_size: Chunk length. + save_new_value: Whether to emit the corrected values `v_new`. + cu_seqlens: Sequence boundaries `[N+1]` on the expanded stream. + num_householder: Number of Householder copies `M`. + chunk_indices: Chunk descriptors for the *unexpanded* stream. Derived + from `cu_seqlens // num_householder` when omitted, which + synchronizes on the device. + + Returns `(h, v_new, final_state)`. `h` holds the state at each unexpanded + chunk boundary. + """ + assert HAVE_TRITON, "chunk_gated_delta_product_fwd_h requires Triton" + B, T, H, K, V = *k.shape, u.shape[-1] + assert T % num_householder == 0, "T must be divisible by num_householder" + T_true = T // num_householder + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens // num_householder, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T_true, BT), None + else: + N = len(cu_seqlens) - 1 + NT = len(chunk_indices) + chunk_offsets = prepare_chunk_offsets(cu_seqlens // num_householder, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + h = k.new_empty(B, NT, H, K, V) + final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + v_new = torch.empty_like(u) if save_new_value else None + + def grid(meta): + return (triton.cdiv(V, meta['BV']), N * H) + + chunk_gated_delta_product_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + num_householder=num_householder, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return h, v_new, final_state diff --git a/megatron/core/ssm/ops/gdp/chunk_o.py b/megatron/core/ssm/ops/gdp/chunk_o.py new file mode 100644 index 00000000000..437d5ea9115 --- /dev/null +++ b/megatron/core/ssm/ops/gdp/chunk_o.py @@ -0,0 +1,215 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/gated_delta_product/chunk_deltaproduct_o.py` in +# flash-linear-attention v0.5.1 (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Intra-chunk output for the Gated Delta Product. + +Combines the inter-chunk state from `chunk_h` with the causal within-chunk +attention. The queries live on the unexpanded token stream while the keys and +values live on the Householder-expanded one, so the within-chunk term is +accumulated over the `num_householder` copies, each strided `H*K` (or `H*V`) +apart. +""" + +import torch + +from .common import ( + HAVE_TRITON, + IS_NVIDIA_HOPPER, + check_shared_mem, + exp2, + prepare_chunk_indices, + tl, + triton, +) + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics( + { + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT'], +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_o( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + num_householder: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """Emit one chunk of output from the chunk-boundary state plus local attention.""" + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * num_householder * H + i_h) * K + v += (bos * num_householder * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K * V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o += tl.dot(b_q, b_h) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + b_m = tl.where(m_A, exp2(b_g[:, None] - b_g[None, :]), 0) + b_o = b_o * exp2(b_g)[:, None] + else: + b_m = ((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)).to(tl.float32) + + for i_dp in range(num_householder): + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr( + k + i_dp * H * K, + (K, T), + (1, num_householder * H * K), + (i_k * BK, i_t * BT), + (BK, BT), + (0, 1), + ) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A += tl.dot(b_q, b_k) + b_A = b_A * b_m + p_v = tl.make_block_ptr( + v + i_dp * H * V, + (T, V), + (H * V * num_householder, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_o += tl.dot(b_A.to(b_v.dtype), b_v) + b_o = b_o * scale + p_o = tl.make_block_ptr(o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_product_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + num_householder: int = 1, + chunk_indices: torch.Tensor | None = None, +) -> torch.Tensor: + """Compute the Gated Delta Product outputs. + + Args: + q: Queries `[B, T, H, K]` on the unexpanded token stream. + k: Keys `[B, T*M, H, K]` on the Householder-expanded stream. + v: Corrected values `[B, T*M, H, V]` from `chunk_h`. + h: State at each chunk boundary, from `chunk_h`. + g: Within-chunk cumulative log2 decays `[B, T, H]`, or `None`. + scale: Score scale. + cu_seqlens: Sequence boundaries `[N+1]` on the unexpanded stream. + chunk_size: Chunk length. + num_householder: Number of Householder copies `M`. + chunk_indices: Precomputed chunk descriptors for the unexpanded stream. + Derived from `cu_seqlens` when omitted, which synchronizes on the + device. + + Returns the outputs `[B, T, H, V]`. + """ + assert HAVE_TRITON, "chunk_gated_delta_product_fwd_o requires Triton" + assert ( + q.shape[1] * num_householder == k.shape[1] + ), "q.shape[1] * num_householder must be equal to k.shape[1]" + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + o = v.new_empty(B, T, H, V).fill_(-float('inf')) + + def grid(meta): + return (triton.cdiv(V, meta['BV']), NT, B * H) + + chunk_fwd_kernel_o[grid]( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T=T, + num_householder=num_householder, + H=H, + K=K, + V=V, + BT=BT, + ) + return o diff --git a/megatron/core/ssm/ops/gdp/common.py b/megatron/core/ssm/ops/gdp/common.py new file mode 100644 index 00000000000..f41848d740b --- /dev/null +++ b/megatron/core/ssm/ops/gdp/common.py @@ -0,0 +1,250 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/utils.py`, `fla/ops/utils/op.py`, `fla/ops/utils/index.py` and +# `fla/modules/l2norm.py` in flash-linear-attention v0.5.1 +# (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Shared helpers for the Gated Delta Product kernels. + +The kernel modules in this package reference a handful of names that do not +belong to any one of them: the hardware-capability probes that select autotune +configurations, the chunk-descriptor builders, `exp` / `exp2`, and the L2 +normalization applied to the queries and keys. Keeping them here makes the +package self-contained, with no `fla` import at run time. + +Take care when changing the probes: they choose which autotune configurations +exist, so editing one changes which kernel variants get benchmarked and picked. +""" + +import functools +import os + +import torch + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + from unittest.mock import MagicMock + + from megatron.core.utils import null_decorator + + triton = MagicMock() + triton.jit = null_decorator + tl = MagicMock() + HAVE_TRITON = False + + +# The Gated Delta Product chunk length is fixed at 64: `solve_tril` +# merges 16x16 blocks up to 64x64, the WY representation is built on 64-wide +# blocks, and the h kernel stores one state block per `num_householder` +# expanded chunks. +CHUNK_SIZE = 64 + +# 1/ln(2), used to convert natural-log decays into the base-2 space the chunked +# kernels exponentiate in. Best fp32 approximation (hex 0x3FB8AA3B). +RCP_LN2 = 1.4426950216 + + +# ---------------------------------------------------------------------------- +# Hardware capability probes (`fla.utils`). +# ---------------------------------------------------------------------------- +def _is_nvidia() -> bool: + return torch.cuda.is_available() and torch.version.hip is None + + +IS_NVIDIA = _is_nvidia() +IS_NVIDIA_HOPPER = IS_NVIDIA and ( + 'NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] >= 9 +) +IS_NVIDIA_BLACKWELL = IS_NVIDIA and torch.cuda.get_device_capability()[0] in (10, 12) +IS_TMA_SUPPORTED = ( + IS_NVIDIA + and torch.cuda.get_device_capability(0)[0] >= 9 + and os.environ.get('FLA_USE_TMA', '0') == '1' + and ( + hasattr(tl, '_experimental_make_tensor_descriptor') or hasattr(tl, 'make_tensor_descriptor') + ) +) + +if hasattr(tl, '_experimental_make_tensor_descriptor'): + # Triton 3.3.x + make_tensor_descriptor = tl._experimental_make_tensor_descriptor +elif hasattr(tl, 'make_tensor_descriptor'): + # Triton 3.4.x and later + make_tensor_descriptor = tl.make_tensor_descriptor +else: + + @triton.jit + def make_tensor_descriptor(base, shape, strides, block_shape, _builder=None): + """No-op stand-in for Triton builds without TMA support.""" + return None + + +# Shared-memory thresholds per architecture, from `fla.utils.Backend`. +_SHARED_MEM_BY_ARCH = { + 'ada': 101376, # RTX 4090 + 'ampere': 166912, # A100 + 'hopper': 232448, # H100 + 'none': 102400, # default +} + + +@functools.cache +def check_shared_mem(arch: str = "none") -> bool: + """Whether the current device has at least `arch`'s shared memory budget.""" + try: + return torch.cuda.get_device_properties( + 0 + ).shared_memory_per_block_optin >= _SHARED_MEM_BY_ARCH.get( + arch.lower(), _SHARED_MEM_BY_ARCH['none'] + ) + except Exception: + return False + + +# ---------------------------------------------------------------------------- +# `fla.ops.utils.op` +# ---------------------------------------------------------------------------- +@triton.jit +def exp(x): + """Exponentiate in fp32 regardless of the input dtype.""" + return tl.exp(x.to(tl.float32)) + + +@triton.jit +def exp2(x): + """Base-2 exponentiate in fp32 regardless of the input dtype.""" + return tl.math.exp2(x.to(tl.float32)) + + +# ---------------------------------------------------------------------------- +# `fla.ops.utils.index` +# ---------------------------------------------------------------------------- +def _segmented_arange(counts: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Expand per-segment counts into flat per-slot index tensors. + + Given segment sizes `counts = [c0, c1, ...]`, return two 1-D tensors of + length `counts.sum()` labelling every slot with its segment and its + position within that segment. For `counts = [2, 3]`:: + + seg_id = [0, 0, 1, 1, 1] + intra_idx = [0, 1, 0, 1, 2] + """ + seg_id = torch.repeat_interleave( + torch.arange(counts.numel(), device=counts.device, dtype=counts.dtype), counts + ) + seg_start = torch.nn.functional.pad(counts.cumsum(0), (1, 0))[:-1] + intra_idx = ( + torch.arange(seg_id.shape[0], device=seg_id.device, dtype=counts.dtype) - seg_start[seg_id] + ) + return seg_id, intra_idx + + +def prepare_chunk_indices( + cu_seqlens: torch.Tensor, chunk_size: int, cu_seqlens_cpu: torch.Tensor | None = None +) -> torch.Tensor: + """Flattened `(sequence, chunk-within-sequence)` pairs, one per chunk. + + Note that this synchronizes on the device when the counts live on the GPU: + `repeat_interleave` reads `counts.sum()` on the host to size its output. + Passing `cu_seqlens_cpu` avoids the sync. + """ + src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + counts = triton.cdiv(torch.diff(src), chunk_size) + seg_id, intra_idx = _segmented_arange(counts) + return torch.stack([seg_id, intra_idx], 1).to(cu_seqlens) + + +def prepare_chunk_offsets(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor: + """Per-sequence prefix sum of chunk counts, with a leading zero.""" + lens = torch.diff(cu_seqlens) + return torch.nn.functional.pad(triton.cdiv(lens, chunk_size), (1, 0), value=0).cumsum(-1) + + +# ---------------------------------------------------------------------------- +# `fla.modules.l2norm` +# ---------------------------------------------------------------------------- +_BT_LIST = [8, 16, 32, 64, 128] + + +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16, 32]], + key=["D"], +) +@triton.jit +def l2norm_fwd_kernel1(x, y, rstd, eps, D, BD: tl.constexpr): + """Row-per-program L2 normalization, used when `D > 512`.""" + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + cols = tl.arange(0, BD) + mask = cols < D + + b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) + b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x) + eps) + b_y = b_x * b_rstd + tl.store(y + cols, b_y, mask=mask) + tl.store(rstd + i_t, b_rstd) + + +@triton.autotune( + configs=[ + triton.Config({"BT": BT}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16] + for BT in _BT_LIST + ], + key=["D", "NB"], +) +@triton.jit(do_not_specialize=["T"]) +def l2norm_fwd_kernel( + x, y, rstd, eps, T, D: tl.constexpr, BD: tl.constexpr, NB: tl.constexpr, BT: tl.constexpr +): + """Block-of-rows L2 normalization, used when `D <= 512`.""" + i_t = tl.program_id(0) + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps) + b_y = b_x * b_rstd[:, None] + + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + +def l2norm_fwd(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + """Row-wise L2 normalization over the last dimension, computed in fp32. + + The kernel also writes `rstd`, which only a backward pass would consume, so + it is not returned. + """ + x_shape_og = x.shape + x = x.view(-1, x.shape[-1]) + y = torch.empty_like(x) + assert y.stride(-1) == 1 + T, D = x.shape[0], x.shape[-1] + # Less than 64KB per feature: enqueue the fused kernel. + max_fused_size = 65536 // x.element_size() + BD = min(max_fused_size, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer doesn't support feature dim >= 64KB.") + + rstd = torch.empty((T,), dtype=torch.float32, device=x.device) + if D <= 512: + # Tolerate a wide range of T before recompiling, to limit autotuning. + NB = triton.cdiv(T, 2048 * 32) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]),) + + l2norm_fwd_kernel[grid](x=x, y=y, rstd=rstd, eps=eps, T=T, D=D, BD=BD, NB=NB) + else: + l2norm_fwd_kernel1[(T,)](x=x, y=y, rstd=rstd, eps=eps, D=D, BD=BD) + return y.view(x_shape_og) diff --git a/megatron/core/ssm/ops/gdp/cumsum.py b/megatron/core/ssm/ops/gdp/cumsum.py new file mode 100644 index 00000000000..e71d5f71ce7 --- /dev/null +++ b/megatron/core/ssm/ops/gdp/cumsum.py @@ -0,0 +1,124 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/utils/cumsum.py` in flash-linear-attention v0.5.1 +# (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Within-chunk cumulative sum of the scalar log decays. + +Only the scalar, non-reversed, variable-length path is provided -- the one the +Gated Delta Product prefill calls. +""" + +import torch + +from .common import HAVE_TRITON, prepare_chunk_indices, tl, triton + + +@triton.heuristics( + { + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + } +) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], + key=['B', 'H', 'BT', 'IS_VARLEN', 'REVERSE'], +) +@triton.jit(do_not_specialize=['T']) +def chunk_local_cumsum_scalar_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + """Cumulative sum within each chunk of a `[B, T, H]` scalar sequence.""" + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + p_s = tl.make_block_ptr(s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + else: + p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_o = tl.cumsum(b_s, axis=0) + if REVERSE: + b_z = tl.sum(b_s, axis=0) + b_o = -b_o + b_z[None] + b_s + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + + +def chunk_local_cumsum( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + scale: float = None, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, + chunk_indices: torch.Tensor | None = None, +) -> torch.Tensor: + """Cumulative sum of `g` within each chunk of length `chunk_size`. + + Args: + g: Scalar sequence `[B, T, H]` (or `[B, H, T]` when `head_first`). + chunk_size: Chunk length; must be a power of two. + reverse: Accumulate from the end of each chunk instead of the start. + scale: Optional scale applied to the result. + cu_seqlens: Sequence boundaries `[N+1]` for variable-length input. + head_first: Whether `g` is laid out head-major. + output_dtype: Result dtype; `None` keeps `g`'s dtype. + chunk_indices: Precomputed chunk descriptors. Derived from `cu_seqlens` + when omitted, which synchronizes on the device. + + Returns the cumulative sums, shaped like `g`. + """ + assert HAVE_TRITON, "chunk_local_cumsum requires Triton" + if head_first: + B, H, T = g.shape + else: + B, T, H = g.shape + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2" + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + chunk_local_cumsum_scalar_kernel[(NT, B * H)]( + s=g_org, + o=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return g diff --git a/megatron/core/ssm/ops/gdp/fused_recurrent.py b/megatron/core/ssm/ops/gdp/fused_recurrent.py new file mode 100644 index 00000000000..7a1510184b7 --- /dev/null +++ b/megatron/core/ssm/ops/gdp/fused_recurrent.py @@ -0,0 +1,198 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/gated_delta_rule/fused_recurrent.py` in flash-linear-attention +# v0.5.1 (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Fused recurrent Gated Delta Rule step, used by the decode path. + +Gated Delta Product decode reaches this kernel by folding the `M` Householder +copies into the sequence dimension, so a single decode token becomes an +`M`-length sequence with the query placed on the last copy and the decay on the +first; the caller slices the answer back out. +""" + +import torch + +from .common import HAVE_TRITON, exp, tl, triton + + +@triton.heuristics( + { + 'USE_G': lambda args: args['g'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + } +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + beta, + o, + h0, + ht, + cu_seqlens, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """Walk one sequence token by token, carrying the `[K, V]` state.""" + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if USE_G: + p_g = g + bos * HV + i_hv + if IS_BETA_HEADWISE: + p_beta = beta + bos * HV + i_hv + else: + p_beta = beta + (bos * HV + i_hv) * V + o_v + + p_o = o + (bos * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in tl.range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta).to(tl.float32) + else: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + + if USE_G: + b_g = tl.load(p_g).to(tl.float32) + b_h *= exp(b_g) + + b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0)) + b_h += b_k[:, None] * b_v + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H * K + p_k += H * K + p_v += HV * V + if USE_G: + p_g += HV + p_beta += HV * (1 if IS_BETA_HEADWISE else V) + p_o += HV * V + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_update( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the recurrent Gated Delta Rule forward pass. + + Args: + q: Queries `[B, T, H, K]`. + k: Keys `[B, T, H, K]`. + v: Values `[B, T, HV, V]`. + g: Log decays `[B, T, HV]`, or `None`. + beta: Betas `[B, T, HV]` (head-wise) or `[B, T, HV, V]`. + scale: Score scale; defaults to `K ** -0.5`. + initial_state: Starting state `[N, HV, K, V]`, or `None` for zeros. + output_final_state: Whether to return the final state. + use_qk_l2norm_in_kernel: Whether to L2-normalize `q` and `k` in-kernel. + cu_seqlens: Sequence boundaries `[N+1]` for variable-length input. + + Returns `(o, final_state)` with `o` shaped like `v`. + """ + assert HAVE_TRITON, "fused_recurrent_gated_delta_rule_update requires Triton" + # The kernel indexes with raw pointer arithmetic and would read garbage from + # a strided input, so force every tensor argument contiguous. + q, k, v, beta = q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous() + if g is not None: + g = g.contiguous() + if initial_state is not None: + initial_state = initial_state.contiguous() + + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK = triton.next_power_of_2(K) + BV = min(8, triton.next_power_of_2(V)) + NV = triton.cdiv(V, BV) + if scale is None: + scale = K**-0.5 + + o = torch.empty_like(v) + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) if output_final_state else None + + fused_recurrent_gated_delta_rule_fwd_kernel[(NV, N * HV)]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim != v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=1, + num_stages=3, + ) + return o, final_state diff --git a/megatron/core/ssm/ops/gdp/scaled_dot_kkt.py b/megatron/core/ssm/ops/gdp/scaled_dot_kkt.py new file mode 100644 index 00000000000..ea31aa295d3 --- /dev/null +++ b/megatron/core/ssm/ops/gdp/scaled_dot_kkt.py @@ -0,0 +1,141 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/common/chunk_scaled_dot_kkt.py` in flash-linear-attention v0.5.1 +# (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Strictly-lower-triangular `beta * K K^T` per chunk. + +This is the first half of building the WY representation of the chunk's +transition matrix; `solve_tril` inverts `I + A` afterwards. +""" + +import torch + +from .common import HAVE_TRITON, exp2, prepare_chunk_indices, tl, triton + + +@triton.heuristics( + { + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'HV', 'K', 'BT', 'IS_VARLEN'], +) +@triton.jit(do_not_specialize=['T']) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + """Compute one chunk's `beta * K K^T`, masked to below the diagonal.""" + i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_b = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * H + i_h // (HV // H)) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A += tl.dot(b_k, tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A *= exp2(b_g_diff) + b_A *= b_b[:, None] + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + p_A = tl.make_block_ptr( + A + (bos * HV + i_h) * BT, (T, BT), (BT * HV, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_scaled_dot_kkt_fwd( + k: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32, + chunk_indices: torch.Tensor | None = None, +) -> torch.Tensor: + """Compute `beta * K K^T` per chunk, decayed by `g` and masked below the diagonal. + + Args: + k: Keys `[B, T, H, K]`, where `H` is the number of query/key heads. + beta: Betas `[B, T, HV]`, where `HV` is the number of value/output heads. + For GVA, `H < HV` and `HV % H == 0`; otherwise `H == HV`. + g: Within-chunk cumulative log2 decays `[B, T, HV]`, or `None`. + cu_seqlens: Sequence boundaries `[N+1]` for variable-length input. + chunk_size: Chunk length. + output_dtype: Result dtype. + chunk_indices: Precomputed chunk descriptors. Derived from `cu_seqlens` + when omitted, which synchronizes on the device. + + Returns `[B, T, HV, BT]`, the per-chunk lower-triangular block. + """ + assert HAVE_TRITON, "chunk_scaled_dot_kkt_fwd requires Triton" + B, T, H, K, HV = *k.shape, beta.shape[2] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + A = torch.empty(B, T, HV, BT, device=k.device, dtype=output_dtype) + chunk_scaled_dot_kkt_fwd_kernel[(NT, B * HV)]( + k=k, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + ) + return A diff --git a/megatron/core/ssm/ops/gdp/solve_tril.py b/megatron/core/ssm/ops/gdp/solve_tril.py new file mode 100644 index 00000000000..31624082170 --- /dev/null +++ b/megatron/core/ssm/ops/gdp/solve_tril.py @@ -0,0 +1,300 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/utils/solve_tril.py` in flash-linear-attention v0.5.1 +# (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Inverse of `I + A` for a strictly lower-triangular, chunk-blocked `A`. + +Only the 64x64 path is provided; the Gated Delta Product kernels are wired to a +chunk length of 64 throughout (see `CHUNK_SIZE`). + +The 64x64 inverse is assembled from four 16x16 diagonal blocks inverted by +forward substitution, then merged with the off-diagonal blocks via the standard +block-triangular identity. +""" + +import os + +import torch + +from .common import ( + HAVE_TRITON, + IS_TMA_SUPPORTED, + make_tensor_descriptor, + prepare_chunk_indices, + tl, + triton, +) + +FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee') +assert FLA_TRIL_PRECISION in [ + 'ieee', + 'tf32', + 'tf32x3', +], f"FLA_TRIL_PRECISION must be one of 'ieee', 'tf32', or 'tf32x3', but got {FLA_TRIL_PRECISION}" +DOT_PRECISION_AUTOTUNE_LIST = ( + ["ieee"] if not IS_TMA_SUPPORTED else list({"ieee", FLA_TRIL_PRECISION}) +) + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4, 5] + for DOT_PRECISION in DOT_PRECISION_AUTOTUNE_LIST + ], + key=['H', 'BT', 'IS_VARLEN'], +) +@triton.jit(do_not_specialize=['T']) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + """Invert one 64x64 unit-lower-triangular block from its 16x16 sub-blocks.""" + i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_A_33 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_A_44 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) + b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) + b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) + b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 = tl.where(o_i < i, b_a_11, 0.0) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 = tl.where(o_i < i - 16, b_a_22, 0.0) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + for i in range(32 + 2, min(48, T - i_t * BT)): + b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) + b_a_33 = tl.where(o_i < i - 32, b_a_33, 0.0) + b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) + b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) + for i in range(48 + 2, min(64, T - i_t * BT)): + b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) + b_a_44 = tl.where(o_i < i - 48, b_a_44, 0.0) + b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) + b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) + b_Ai_11 += m_I + b_Ai_22 += m_I + b_Ai_33 += m_I + b_Ai_44 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_A_31 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_A_32 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_A_41 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_A_42 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_A_43 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) + b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) + b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) + b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) + b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) + b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32) + b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32) + b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32) + b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32) + + b_Ai_21 = -tl.dot( + tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), + b_Ai_11, + input_precision=DOT_PRECISION, + ) + b_Ai_32 = -tl.dot( + tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), + b_Ai_22, + input_precision=DOT_PRECISION, + ) + b_Ai_43 = -tl.dot( + tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), + b_Ai_33, + input_precision=DOT_PRECISION, + ) + + b_Ai_31 = -tl.dot( + b_Ai_33, + tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_42 = -tl.dot( + b_Ai_44, + tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_41 = -tl.dot( + b_Ai_44, + tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_Ai_33 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_Ai_44 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_Ai_31 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_Ai_32 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_Ai_41 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_Ai_42 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_Ai_43 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + tl.store( + p_Ai_11, + b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_33, + b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_44, + b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_31, + b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_32, + b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_41, + b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_42, + b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_43, + b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne")) + + +def solve_tril( + A: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float, +) -> torch.Tensor: + """Compute the inverse of `I + A` for a strictly lower-triangular `A`. + + Args: + A: `[B, T, H, 64]`; `A.triu() == 0` is assumed, not checked. + cu_seqlens: Sequence boundaries `[N+1]` for variable-length input. + chunk_indices: Precomputed chunk descriptors. Derived from `cu_seqlens` + when omitted, which synchronizes on the device. + output_dtype: Result dtype; `None` keeps `A`'s dtype. + + Returns `(I + A)^-1`, shaped like `A`. + """ + assert HAVE_TRITON, "solve_tril requires Triton" + assert A.shape[-1] == 64, "only the 64x64 inverse is supported" + output_dtype = A.dtype if output_dtype is None else output_dtype + + B, T, H, BT = A.shape + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + + Ai = torch.zeros_like(A, dtype=output_dtype) + merge_16x16_to_64x64_inverse_kernel[NT, B * H]( + A=A, + Ai=Ai, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + USE_TMA=IS_TMA_SUPPORTED, + ) + return Ai diff --git a/megatron/core/ssm/ops/gdp/wy_fast.py b/megatron/core/ssm/ops/gdp/wy_fast.py new file mode 100644 index 00000000000..dadf2f550ca --- /dev/null +++ b/megatron/core/ssm/ops/gdp/wy_fast.py @@ -0,0 +1,166 @@ +# Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Forked from `fla/ops/gated_delta_rule/wy_fast.py` in flash-linear-attention v0.5.1 +# (https://github.com/fla-org/flash-linear-attention). +# +# Licensed under the MIT license; see the LICENSE file in this directory. + +"""Recompute the `w` and `u` factors of the WY representation. + +Given the inverted transition block `A` from `solve_tril`, this produces the +two per-chunk factors the state recurrence consumes: `u`, the effective values, +and `w`, the effective keys carrying the decay. +""" + +import torch + +from .common import HAVE_TRITON, exp2, prepare_chunk_indices, tl, triton + + +@triton.heuristics( + { + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """Apply the inverted transition block to the betas-scaled keys and values.""" + i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr( + A + (bos * HV + i_h) * BT, (T, BT), (HV * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + p_u = tl.make_block_ptr( + u + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + p_g = tl.make_block_ptr(g + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + b_g = exp2(tl.load(p_g, boundary_check=(0,))) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * H + i_h // (HV // H)) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_w = tl.make_block_ptr( + w + (bos * HV + i_h) * K, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + if USE_G: + b_kb *= b_g[:, None] + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + g: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Recompute the WY factors `w` and `u` for every chunk. + + Args: + k: Keys `[B, T, H, K]`, where `H` is the number of query/key heads. + v: Values `[B, T, HV, V]`, where `HV` is the number of value/output heads. + For GVA, `H < HV` and `HV % H == 0`; otherwise `H == HV`. + beta: Betas `[B, T, HV]`. + A: Inverted transition blocks `[B, T, HV, BT]` from `solve_tril`. + g: Within-chunk cumulative log2 decays `[B, T, HV]`, or `None`. + cu_seqlens: Sequence boundaries `[N+1]` for variable-length input. + chunk_indices: Precomputed chunk descriptors. Derived from `cu_seqlens` + when omitted, which synchronizes on the device. + + Returns `(w, u)`, shaped `[B, T, HV, K]` and like `v` respectively. + """ + assert HAVE_TRITON, "recompute_w_u_fwd requires Triton" + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] + BT = A.shape[-1] + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = 64 + BV = 64 + + w = k.new_empty(B, T, HV, K) + u = torch.empty_like(v) + recompute_w_u_fwd_kernel[(NT, B * HV)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u From fcc331916ae17c5f31a2ca1eff7e44ab31bb3d0c Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 17 Aug 2026 20:37:06 +0000 Subject: [PATCH 061/101] Re-enable group filtering in RL (#5824) Signed-off-by: Teodor-Dumitru Ene --- megatron/rl/agent/rollout_pipeline.py | 139 +++++++++++------- megatron/rl/rl_utils.py | 3 + .../unit_tests/rl/test_rollout_generation.py | 114 ++++++++++++-- 3 files changed, 196 insertions(+), 60 deletions(-) diff --git a/megatron/rl/agent/rollout_pipeline.py b/megatron/rl/agent/rollout_pipeline.py index c111fa29258..d515e972831 100644 --- a/megatron/rl/agent/rollout_pipeline.py +++ b/megatron/rl/agent/rollout_pipeline.py @@ -85,7 +85,7 @@ def units_per_batch(self) -> int: @staticmethod def _validate(request: GroupedRolloutRequest, num_groups_per_env: list[int]) -> None: - """Reject invalid granularity, layout, and filter combinations. + """Reject invalid granularity and layout combinations. Args: request: Grouped rollout request to check. @@ -104,11 +104,6 @@ def _validate(request: GroupedRolloutRequest, num_groups_per_env: list[int]) -> assert ( sum(num_groups_per_env) == request.num_groups ), "The sum of groups per environment must equal the total number of groups requested." - assert not request.filter_groups_with_same_reward, ( - "filter_groups_with_same_reward is not currently supported: dropped groups " - "are not regenerated, so non-streaming callers receive fewer groups than " - "requested and batch-order consumers stall on incomplete batches." - ) class _SubmissionGate: @@ -246,6 +241,12 @@ def __init__( self.assemble_queue = asyncio_Queue() self.output_queue = asyncio_Queue() + # Track regenerated groups and tasks for proper shutdown. + self._next_regen_group_id = -1 + self._prepare_done = False + self._groups_in_flight = 0 + self._regen_tasks: set[asyncio.Task] = set() + # Buffers of partial results. self._assemble_pending: dict[int, list[_InferredItem]] = {} self._consume_pending: dict[int, list[RolloutGroup]] = {} @@ -259,6 +260,7 @@ def __init__( self.prepared_count = 0 self.inferred_count = 0 self.assembled_count = 0 + self.filtered_count = 0 self.yielded_count = 0 self.prepared_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) self.assembled_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) @@ -300,9 +302,38 @@ async def run(self) -> AsyncIterator[RolloutGroup]: + ("" if failure is not None else " (no stage exception was recovered)") ) from failure finally: - for task in tasks: + regen_tasks = tuple(self._regen_tasks) + for task in (*tasks, *regen_tasks): task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.gather(*tasks, *regen_tasks, return_exceptions=True) + + async def _submit_group_to_infer_queue( + self, *, group_id: int, batch_id: int, index_in_batch: int + ) -> None: + """Enqueue one group's inference items, acquiring per-rollout gate slots.""" + env_index = self.gran_policy.env_of_index(index_in_batch) + agent = self.allocations[env_index].agent + params: GroupRolloutParams = await agent.prepare_group_rollout(self.request) + self.prepared_groups_per_env[env_index] += 1 + + for rollout_idx in range(self.request.rollouts_per_group): + await self.gate.acquire_for("R") + item = _InferWorkItem( + group_id=group_id, + rollout_idx=rollout_idx, + batch_id=batch_id, + index_in_batch=index_in_batch, + params=params, + env_index=env_index, + prepared_at=time.monotonic(), + ) + await self.infer_queue.put(item) + self.prepared_count += 1 + + def _maybe_close_intake(self) -> None: + """Shut down infer_queue once no work can ever be submitted again.""" + if self._prepare_done and self._groups_in_flight <= 0: + self.infer_queue.shutdown() async def stage_prepare(self) -> None: """Generate gated inference work items.""" @@ -317,28 +348,18 @@ async def stage_prepare(self) -> None: batch_id = group_id // self.gran_policy.num_groups_per_batch for index_in_batch in range(self.gran_policy.num_groups_per_batch): - env_index = self.gran_policy.env_of_index(index_in_batch) await self.gate.acquire_for("G") - agent = self.allocations[env_index].agent - params: GroupRolloutParams = await agent.prepare_group_rollout(self.request) - self.prepared_groups_per_env[env_index] += 1 - - for rollout_idx in range(self.request.rollouts_per_group): - await self.gate.acquire_for("R") - item = _InferWorkItem( - group_id=group_id, - rollout_idx=rollout_idx, - batch_id=batch_id, - index_in_batch=index_in_batch, - params=params, - env_index=env_index, - prepared_at=time.monotonic(), - ) - await self.infer_queue.put(item) - self.prepared_count += 1 + self._groups_in_flight += 1 + await self._submit_group_to_infer_queue( + group_id=group_id, batch_id=batch_id, index_in_batch=index_in_batch + ) group_id += 1 - finally: + except BaseException: self.infer_queue.shutdown() + raise + finally: + self._prepare_done = True + self._maybe_close_intake() async def stage_infer(self) -> None: """Run a persistent pool of inference workers, spawned once per pipeline.""" @@ -381,7 +402,7 @@ async def _infer_one(self, item: _InferWorkItem) -> None: ) async def stage_assemble(self) -> None: - """Build complete rollout groups from inferred items.""" + """Build complete rollout groups, and regenerate filtered groups.""" pending = self._assemble_pending try: while True: @@ -404,33 +425,49 @@ async def stage_assemble(self) -> None: first = completed[0] self.assembled_count += 1 self.assembled_groups_per_env[first.item.env_index] += 1 - # NOTE: this filter is currently non-functional dead code: - # _GranularityConfig._validate rejects filter_groups_with_same_reward - # at pipeline construction, so `keep` is always True. Kept for a - # future PR that regenerates dropped groups instead of - # under-delivering to the caller. That PR must also release the - # gate slot on the drop path: G/B slots free on consumption, - # and a dropped group never reaches stage_consume, so its slot - # (and eventually its batch's) would leak permanently. - keep = ( - not self.request.filter_groups_with_same_reward - or np.std([rollout.reward for rollout in rollouts]) > 1e-6 + group = RolloutGroup( + rollouts=rollouts, + batch_id=first.item.batch_id, + index_in_batch=first.item.index_in_batch, ) - if keep: - output_enqueued_at = time.monotonic() - self._output_enqueued_at[ - (first.item.batch_id, first.item.index_in_batch) - ] = output_enqueued_at - await self.output_queue.put( - RolloutGroup( - rollouts=rollouts, - batch_id=first.item.batch_id, - index_in_batch=first.item.index_in_batch, - ) - ) + if self._should_drop(group): + self.filtered_count += 1 + task = asyncio.create_task(self._regenerate_group(group)) + self._regen_tasks.add(task) + task.add_done_callback(self._regen_tasks.discard) + continue + self._output_enqueued_at[(group.batch_id, group.index_in_batch)] = ( + time.monotonic() + ) + await self.output_queue.put(group) + self._groups_in_flight -= 1 + self._maybe_close_intake() finally: self.output_queue.shutdown() + def _should_drop(self, group: RolloutGroup) -> bool: + """A group with zero reward variance carries no learning signal.""" + if not self.request.filter_groups_with_same_reward: + return False + return np.std([rollout.reward for rollout in group.rollouts]) <= 1e-6 + + @trace_async_exceptions(verbose=True) + async def _regenerate_group(self, dropped: RolloutGroup) -> None: + """Resubmit a replacement group for a dropped group's batch slot.""" + group_id = self._next_regen_group_id + self._next_regen_group_id -= 1 + try: + await self._submit_group_to_infer_queue( + group_id=group_id, + batch_id=dropped.batch_id, + index_in_batch=dropped.index_in_batch, + ) + except asyncio_QueueShutDown: + # Intake closed mid-regeneration (teardown or other failure). + self.gate.release_for("G") + self._groups_in_flight -= 1 + self._maybe_close_intake() + def _record_output_dwell(self, group: RolloutGroup) -> None: """Record how long a group sat in output_queue before being yielded.""" key = (group.batch_id, group.index_in_batch) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 302d8f86e89..5e0a66b3f0d 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -1472,6 +1472,7 @@ def _collect_rollout_pipeline_metrics() -> dict: "rollout_pipeline_output_queue_size": pipeline.output_queue.qsize(), "rollout_pipeline_assemble_pending_groups": len(pipeline._assemble_pending), "rollout_pipeline_consume_pending_groups": len(pipeline._consume_pending), + "rollout_pipeline_regen_pending_groups": len(pipeline._regen_tasks), "rollout_pipeline_gate_capacity": gate.capacity, "rollout_pipeline_gate_held": gate.held, "rollout_pipeline_gate_utilization": ( @@ -1484,6 +1485,7 @@ def _collect_rollout_pipeline_metrics() -> dict: "rollout_pipeline_prepared_count": pipeline.prepared_count, "rollout_pipeline_inferred_count": pipeline.inferred_count, "rollout_pipeline_assembled_count": pipeline.assembled_count, + "rollout_pipeline_filtered_count": pipeline.filtered_count, "rollout_pipeline_yielded_count": pipeline.yielded_count, }) for name, samples in ( @@ -1523,6 +1525,7 @@ def _collect_rollout_pipeline_metrics() -> dict: pipeline.prepared_count = 0 pipeline.inferred_count = 0 pipeline.assembled_count = 0 + pipeline.filtered_count = 0 pipeline.yielded_count = 0 pipeline.prepared_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) pipeline.assembled_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) diff --git a/tests/unit_tests/rl/test_rollout_generation.py b/tests/unit_tests/rl/test_rollout_generation.py index 909b4cdc6e8..1efaf22806e 100644 --- a/tests/unit_tests/rl/test_rollout_generation.py +++ b/tests/unit_tests/rl/test_rollout_generation.py @@ -1,9 +1,11 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import asyncio +import itertools from contextlib import aclosing from unittest.mock import MagicMock +import numpy as np import pytest from pydantic import Field, ValidationError @@ -91,6 +93,28 @@ async def build_rollout(episode): return GroupRolloutParams(run_episode=run_episode, build_rollout=build_rollout) +class FilteringMockGenerator(MockGenerator): + """Mock generator whose first `num_degenerate` prepared groups have zero reward variance.""" + + def __init__(self, num_degenerate=0, **kwargs): + super().__init__(**kwargs) + self.num_degenerate = num_degenerate + + async def prepare_group_rollout(self, request): + idx = self._call_count + params = await super().prepare_group_rollout(request) + degenerate = idx < self.num_degenerate + rollout_counter = itertools.count() + base_build = params.build_rollout + + async def build_rollout(episode): + rollout = await base_build(episode) + rollout.reward = 0.0 if degenerate else float(next(rollout_counter)) + return rollout + + return GroupRolloutParams(run_episode=params.run_episode, build_rollout=build_rollout) + + class CountingRewardAgent(RewardOnlyAgent): """Minimal RewardOnlyAgent: prompts t0, t1, ... and reward = echoed index.""" @@ -302,27 +326,99 @@ def test_grouped_rollout_request_rejects_unknown_granularity(self, field): @pytest.mark.asyncio @pytest.mark.parametrize( - "num_groups, submission_granularity, consumption_granularity", + "streaming, submission_granularity, consumption_granularity, num_degenerate", [ - pytest.param(1, "B", "B", id="num_groups_1_batch"), - pytest.param(4, "G", "G", id="num_groups_gt_1_group"), - pytest.param(4, "R", "B", id="num_groups_gt_1_rollout"), + pytest.param(False, "B", "B", 3, id="batch_submission"), + pytest.param(False, "G", "B", 3, id="group_submission"), + pytest.param(False, "R", "B", 3, id="rollout_submission"), + pytest.param(False, "G", "G", 3, id="group_consumption"), + pytest.param(False, "R", "G", 3, id="rollout_submission_group_consume"), + pytest.param(False, "B", "B", 9, id="cascading_regeneration"), + pytest.param(True, "B", "B", 3, id="streaming_batch_consume"), + pytest.param(True, "G", "G", 3, id="streaming_group_consume"), ], ) - async def test_filter_groups_with_same_reward_rejected( - self, num_groups, submission_granularity, consumption_granularity + async def test_filter_groups_and_regenerate( + self, streaming, submission_granularity, consumption_granularity, num_degenerate ): - gen = MockGenerator() + num_groups = 4 + gen = FilteringMockGenerator(num_degenerate=num_degenerate) request = GroupedRolloutRequest( num_groups=num_groups, rollouts_per_group=2, inference_interface=MockInferenceInterface(), filter_groups_with_same_reward=True, + streaming=streaming, submission_granularity=submission_granularity, consumption_granularity=consumption_granularity, ) - with pytest.raises(AssertionError, match="filter_groups_with_same_reward"): - RolloutPipeline(gen, request, parallel_generation_tasks=8) + pipeline = RolloutPipeline(gen, request, parallel_generation_tasks=8) + + expected_count = 2 * num_groups if streaming else num_groups + groups, filtered_counts = [], [] + async with aclosing(pipeline.run()) as it: + async for group in it: + groups.append(group) + filtered_counts.append(pipeline.filtered_count) + if streaming and len(groups) >= expected_count: + break + + assert len(groups) == expected_count + # Every delivered group carries reward signal. + for group in groups: + assert np.std([rollout.reward for rollout in group]) > 1e-6 + if consumption_granularity == "B": + # Batches complete and arrive in submission order despite drops. + assert [g.batch_id for g in groups] == sorted(g.batch_id for g in groups) + for batch_start in range(0, expected_count, num_groups): + batch = groups[batch_start : batch_start + num_groups] + assert sorted(g.index_in_batch for g in batch) == list(range(num_groups)) + if not streaming: + # One extra prepare per dropped group, and no under-delivery. + assert gen.prepare_group_rollout_calls == num_groups + num_degenerate + assert filtered_counts[-1] == num_degenerate + + @pytest.mark.asyncio + async def test_pending_regeneration_does_not_block_delivery(self): + """A regeneration pinned on R-slot acquires must not stall delivery of ready groups. + + Rollouts run for minutes in production: with gate capacity 2, group 2's + gated rollouts hog both R slots, so group 0's regeneration sits inside + its slot acquires while group 1 is already assembled and deliverable. + """ + events = {"t2": asyncio.Event()} + + class GatedInference(MockInferenceInterface): + async def base_generate(self, request): + event = events.get(request.prompt[0].content) + if event is not None: + await event.wait() + return await super().base_generate(request) + + gen = FilteringMockGenerator(num_degenerate=1) + request = GroupedRolloutRequest( + num_groups=3, + rollouts_per_group=2, + inference_interface=GatedInference(), + filter_groups_with_same_reward=True, + submission_granularity="R", + consumption_granularity="G", + ) + pipeline = RolloutPipeline(gen, request, parallel_generation_tasks=1 / 3) + assert pipeline.gate.capacity == 2 + async with aclosing(pipeline.run()) as it: + # wait_for turns the pre-fix failure mode (the regeneration awaited + # inline on the delivery path, delivering nothing) into a test failure. + group = await asyncio.wait_for(anext(it), timeout=10) + await _flush() + assert [rollout.trajectory[0] for rollout in group.rollouts] == ["t1", "t1"] + assert pipeline.filtered_count == 1 + assert len(pipeline._regen_tasks) == 1 + assert gen.prepare_group_rollout_calls == 4 + events["t2"].set() + rest = [await asyncio.wait_for(anext(it), timeout=10) for _ in range(2)] + assert sorted(g.rollouts[0].trajectory[0] for g in rest) == ["t2", "t3"] + assert not pipeline._regen_tasks @pytest.mark.asyncio @pytest.mark.parametrize( From dd8c825b3ad59ac15f70c7d075e1527cf8e29215 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Mon, 17 Aug 2026 21:55:57 +0000 Subject: [PATCH 062/101] Instrument Megatron-LM with OpenTelemetry traces and metrics (#6438) Signed-off-by: Deepak Narayanan --- docs/index.md | 13 + docs/user-guide/observability/index.md | 15 - docs/user-guide/observability/metrics.md | 20 +- docs/user-guide/observability/span-groups.md | 33 +- .../blended_megatron_dataset_builder.py | 10 + .../distributed/distributed_data_parallel.py | 8 + .../fine_grained_activation_offload.py | 13 +- .../pipeline_parallel/p2p_communication.py | 9 + megatron/core/pipeline_parallel/schedules.py | 7 + megatron/core/ssm/mamba_layer.py | 52 +- megatron/core/telemetry/README.md | 17 +- .../core/transformer/transformer_layer.py | 50 +- megatron/training/arguments.py | 27 + megatron/training/async_utils.py | 101 +- megatron/training/checkpointing.py | 58 +- megatron/training/global_vars.py | 211 +++ megatron/training/training.py | 1281 +++++++++++++---- pretrain_hybrid.py | 73 +- pretrain_vlm.py | 21 +- pyproject.toml | 8 +- uv.lock | 132 +- 21 files changed, 1737 insertions(+), 422 deletions(-) diff --git a/docs/index.md b/docs/index.md index 3995c70217d..13443ab42c1 100644 --- a/docs/index.md +++ b/docs/index.md @@ -76,6 +76,19 @@ user-guide/features/megatron_rl user-guide/features/tokenizers ``` +```{toctree} +:maxdepth: 2 +:hidden: +:caption: Observability + +user-guide/observability/index +user-guide/observability/configuration +user-guide/observability/span-groups +user-guide/observability/metrics +user-guide/observability/pipeline-parallel +user-guide/observability/extending +``` + ```{toctree} :maxdepth: 1 :hidden: diff --git a/docs/user-guide/observability/index.md b/docs/user-guide/observability/index.md index e9239e0a692..89d0513d59b 100644 --- a/docs/user-guide/observability/index.md +++ b/docs/user-guide/observability/index.md @@ -1,7 +1,3 @@ ---- -orphan: true ---- - b --[g]--> c --[h]--> d --[sum]--> loss + """ + + def func_f(x): + return torch.nn.functional.gelu(x * 2 + 1) + + def func_g(x): + return x * 3 - 2 + + def func_h(x): + return torch.sigmoid(x) + x + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + b_ref = func_f(input_ref) + c_ref = func_g(b_ref) + d_ref = func_h(c_ref) + loss_ref = d_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + b = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_f, input_ckpt) + c = func_g(b) + d = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_h, c) + + loss_ckpt = d.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert b.untyped_storage().size() == 0, "b storage should be released" + assert d.untyped_storage().size() == 0, "d storage should be released" + assert c.untyped_storage().size() > 0, "c storage should NOT be released (not checkpointed)" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6), ( + f"Gradients mismatch with partial checkpoint!\n" + f"With manager: {grad_ckpt}\nReference: {grad_ref}" + ) + + def test_partial_checkpoint_with_tuple_output(self): + """ + Mimics HyperConnection's computation pattern with tuple outputs. + + - compute_mappings: checkpointed, returns tuple (h_pre, h_post, h_res) + - aggregate: NOT checkpointed + - apply_h_res: checkpointed + - apply_h_post: checkpointed + """ + + def compute_mappings(x): + h_pre = torch.sigmoid(x.mean(dim=-1, keepdim=True).expand_as(x)) + h_post = torch.tanh(x.sum(dim=-1, keepdim=True).expand_as(x)) + h_res = torch.relu(x) + return h_pre, h_post, h_res + + def aggregate(x, h_pre): + return x * h_pre + + def apply_h_res(h_res, residual): + return h_res + residual * 0.5 + + def apply_h_post(y, h_post): + return y * h_post + y + + x_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + residual_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + h_pre_ref, h_post_ref, h_res_ref = compute_mappings(x_ref) + agg_ref = aggregate(x_ref, h_pre_ref) + y_ref = torch.nn.functional.gelu(agg_ref) + mixed_ref = apply_h_res(h_res_ref, residual_ref) + output_ref = apply_h_post(y_ref, h_post_ref) + final_ref = output_ref + mixed_ref + loss_ref = final_ref.sum() + loss_ref.backward() + grad_x_ref = x_ref.grad.clone() + grad_residual_ref = residual_ref.grad.clone() + + x_ckpt = x_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + h_pre, h_post, h_res = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + compute_mappings, x_ckpt + ) + agg = aggregate(x_ckpt, h_pre) + y = torch.nn.functional.gelu(agg) + mixed = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + apply_h_res, h_res, residual_ckpt + ) + output = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(apply_h_post, y, h_post) + + final = output + mixed + loss_ckpt = final.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert h_pre.untyped_storage().size() == 0, "h_pre storage should be released" + assert h_post.untyped_storage().size() == 0, "h_post storage should be released" + assert h_res.untyped_storage().size() == 0, "h_res storage should be released" + assert mixed.untyped_storage().size() == 0, "mixed storage should be released" + assert output.untyped_storage().size() == 0, "output storage should be released" + + assert agg.untyped_storage().size() > 0, "agg storage should NOT be released" + assert y.untyped_storage().size() > 0, "y storage should NOT be released" + + loss_ckpt.backward() + grad_x_ckpt = x_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + assert torch.allclose( + grad_x_ckpt, grad_x_ref, atol=1e-6 + ), f"Gradients for x mismatch!\nWith manager: {grad_x_ckpt}\nReference: {grad_x_ref}" + assert torch.allclose(grad_residual_ckpt, grad_residual_ref, atol=1e-6), ( + f"Gradients for residual mismatch!\n" + f"With manager: {grad_residual_ckpt}\nReference: {grad_residual_ref}" + ) + + +# ============================================================================ +# Block-level mHC recompute coverage +# ============================================================================ +# +# These tests instantiate a full ``TransformerBlock`` with mHC enabled to +# exercise: +# * ``build_mhc_recompute_layer_plan`` (per-layer ``CheckpointWithoutOutputManager`` +# allocation, including the ``mhc_recompute_layer_num`` boundary case), +# * ``finalize_mhc_recompute_layer`` (manager finalization at block end), +# * the ``HyperConnectionModule.input_expand`` / ``output_contract`` calls +# in ``TransformerBlock.forward`` for ``pre_process`` / ``post_process`` +# stages. +# +# Single-process (no PP) so they can run on a single-GPU CI lane. + + +class TestTransformerBlockMHCRecompute: + """End-to-end ``TransformerBlock`` forward with mHC selective recompute.""" + + def setup_method(self, method): + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _make_mhc_block(num_layers, num_streams=4, mhc_recompute_layer_num=None): + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, + ) + from megatron.core.transformer.hyper_connection import HyperConnectionModule + from megatron.core.transformer.transformer_block import TransformerBlock + from megatron.core.transformer.transformer_config import TransformerConfig + from megatron.core.transformer.transformer_layer import HyperConnectionTransformerLayer + + config = TransformerConfig( + num_layers=num_layers, + hidden_size=64, + num_attention_heads=4, + use_cpu_initialization=True, + enable_mhc_connections=True, + mhc_num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + mhc_recompute_layer_num=mhc_recompute_layer_num, + recompute_granularity='selective', + recompute_modules=['mhc'], + hidden_dropout=0.0, + attention_dropout=0.0, + ) + spec = get_gpt_layer_with_transformer_engine_spec() + spec.module = HyperConnectionTransformerLayer + spec.submodules.self_attention_hyper_connection = HyperConnectionModule + spec.submodules.mlp_hyper_connection = HyperConnectionModule + return TransformerBlock(config, spec, pre_process=True, post_process=True).cuda(), config + + def _check_recompute_plan(self, block, expected_block_ends): + """Build the mHC recompute plan directly and check the boundary list.""" + block.train() + managers, ends = build_mhc_recompute_layer_plan( + num_layers=len(block.layers), + mhc_recompute_layer_num=block.config.mhc_recompute_layer_num, + use_mhc_recompute=True, + ) + assert len(managers) == len(block.layers) + assert ends == expected_block_ends, f"got {ends}, expected {expected_block_ends}" + # Layers in the same recompute block share a manager; new block → new manager. + last_was_end = True + last_mgr = None + for mgr, end in zip(managers, ends): + assert mgr is not None + if last_was_end: + assert mgr is not last_mgr, "new recompute block should get a new manager" + else: + assert mgr is last_mgr, "layers within a recompute block share a manager" + last_was_end = end + last_mgr = mgr + + def test_recompute_plan_no_layer_num(self): + """Without ``mhc_recompute_layer_num`` only the final layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4) + self._check_recompute_plan(block, expected_block_ends=[False, False, False, True]) + + def test_recompute_plan_with_layer_num(self): + """With ``mhc_recompute_layer_num=2`` every other layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4, mhc_recompute_layer_num=2) + self._check_recompute_plan(block, expected_block_ends=[False, True, False, True]) + + def test_recompute_plan_disabled(self): + """``use_mhc_recompute=False`` returns an all-None / all-False plan.""" + block, _ = self._make_mhc_block(num_layers=3) + managers, ends = build_mhc_recompute_layer_plan( + num_layers=len(block.layers), + mhc_recompute_layer_num=block.config.mhc_recompute_layer_num, + use_mhc_recompute=False, + ) + assert managers == [None, None, None] + assert ends == [False, False, False] + + def test_recompute_enablement_initialized(self): + """TransformerBlock records static mHC recompute enablement during initialization.""" + block, _ = self._make_mhc_block(num_layers=2) + assert block.mhc_recompute_enabled + + def test_block_forward_input_expand_output_contract(self): + """Forward exercises ``input_expand`` (pre) and ``output_contract`` (post).""" + block, config = self._make_mhc_block(num_layers=2, mhc_recompute_layer_num=2) + block.train() + + seq_len = 8 + batch_size = 2 + # Input is [s, b, hidden_size]; the block must expand to [s, b, n*hidden_size] + # internally, then contract back to [s, b, hidden_size] before final layernorm. + hidden_states = torch.randn( + seq_len, batch_size, config.hidden_size, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=torch.bool, device='cuda') + + out = block(hidden_states=hidden_states, attention_mask=attention_mask) + assert out.shape == hidden_states.shape, ( + f"output_contract should restore original shape, got {tuple(out.shape)} " + f"vs expected {tuple(hidden_states.shape)}" + ) + # Backward should flow through the recompute path without error. + out.sum().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 9f6a20ae453..8c4d1134a1c 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc @@ -19,6 +19,7 @@ from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec from megatron.core.tensor_parallel.random import ( HAVE_TE, + CheckpointWithoutOutputManager, initialize_rng_tracker, model_parallel_cuda_manual_seed, ) @@ -30,6 +31,7 @@ from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, MoETransformerLayer, TransformerLayer, TransformerLayerSubmodules, @@ -39,6 +41,44 @@ from tests.unit_tests.test_utilities import Utils +def _make_mhc_layer_spec(**kwargs): + """Build a layer spec with HyperConnectionModule submodules. + + The ``enable_hyper_connection`` kwarg on ``gpt_layer_specs`` is added by + the GPT-wiring follow-up split, so this helper patches the mHC submodules + directly to keep the unit tests self-contained for this split. + """ + from megatron.core.transformer.hyper_connection import HyperConnectionModule + + layer_spec = get_gpt_layer_with_transformer_engine_spec(**kwargs) + layer_spec.module = HyperConnectionTransformerLayer + layer_spec.submodules.self_attention_hyper_connection = HyperConnectionModule + layer_spec.submodules.mlp_hyper_connection = HyperConnectionModule + return layer_spec + + +def _make_mhc_config(hidden_size=64, num_streams=4, **extra): + """Build a TransformerConfig with common MHC defaults. + + Any default can be overridden via **extra + (e.g. ``_make_mhc_config(num_layers=8, recompute_modules=["core_attn", "mhc"])``). + """ + base = dict( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_mhc_connections=True, + mhc_num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + base.update(extra) + return TransformerConfig(**base) + + class TestParallelTransformerLayer: def setup_method(self, method): @@ -508,3 +548,796 @@ def test_deprecated_full_iteration_inference_scope_string_matches_new_granularit assert block.config.cuda_graph_modules == [] assert _no_layers_have_manager(block) _reset_cudagraph_state() + + +class TestTransformerLayerWithHyperConnectionRecompute: + """Test TransformerLayer with HyperConnection and MHC block recomputation.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_layer_with_hyper_connection( + self, hidden_size=64, num_streams=4, layer_number=1, **extra + ): + """Create a HyperConnectionTransformerLayer with hyper connection enabled.""" + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + recompute_modules=["core_attn", "mhc"], + recompute_granularity='selective', + **extra, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=layer_number + ) + layer.cuda() + return layer, config + + def test_forward_with_hyper_connection_recompute(self): + """ + Test that TransformerLayer forward works correctly with HyperConnection + and MHC block recomputation enabled. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() # Enable training mode for recomputation + + # Input shape: [seq_len, batch_size, n * hidden_size] for hyper connections + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Create manager for MHC block recomputation + manager = CheckpointWithoutOutputManager() + + # Forward pass with recompute manager + manager.is_last_layer_in_recompute_block = True + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Expected output shape {(seq_len, batch_size, n_channels)}, got {output.shape}" + + # Register unified recompute hook at block boundary. + manager.discard_all_outputs_and_register_unified_recompute(output) + + # Backward pass should work without error + loss = output.sum() + loss.backward() + + # Verify gradients exist + assert hidden_states.grad is not None, "Gradients should be computed for hidden_states" + assert hidden_states.grad.shape == hidden_states.shape + + def test_intermediate_layer_with_recompute(self): + """ + Test TransformerLayer as an intermediate layer (not last in block). + In this case, MLP BDA should also be checkpointed. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + manager = CheckpointWithoutOutputManager() + + # Forward pass - NOT the last layer in block + manager.is_last_layer_in_recompute_block = False + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == (seq_len, batch_size, n_channels) + + # Backward pass should work + loss = output.sum() + # For intermediate layers, we need to pass output to next layer + # Here we just register the recompute hook on output for testing + manager.discard_all_outputs_and_register_unified_recompute(loss) + + loss.backward() + + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + + def test_multiple_layers_chain_with_recompute(self): + """ + Test multiple TransformerLayers chained together with a single + CheckpointWithoutOutputManager, simulating TransformerBlock behavior. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + num_layers = 3 + + layers = [ + self._create_layer_with_hyper_connection( + hidden_size, num_streams, layer_number=i + 1, num_layers=num_layers + )[0] + for i in range(num_layers) + ] + + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Single manager for all layers (like TransformerBlock) + manager = CheckpointWithoutOutputManager() + + # Forward through all layers + h = hidden_states + for i, layer in enumerate(layers): + is_last = i == num_layers - 1 + manager.is_last_layer_in_recompute_block = is_last + h, _ = layer( + hidden_states=h, attention_mask=attention_mask, mhc_recompute_manager=manager + ) + if is_last: + manager.discard_all_outputs_and_register_unified_recompute(h) + + # Backward pass + loss = h.sum() + loss.backward() + + # Verify gradients + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + # Check that gradient is non-trivial (not all zeros) + assert hidden_states.grad.abs().sum() > 0 + + @pytest.mark.parametrize("norm_attr", ["input_layernorm", "pre_mlp_layernorm"]) + def test_residual_returning_layernorm_is_rejected(self, norm_attr): + """A norm returning (output, residual) must fail fast, not feed a tuple downstream. + + Base TransformerLayer accepts that contract and uses the returned residual; + mHC cannot, because its residual is the n-stream tensor captured before + hyper-connection aggregation. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, _ = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() + + class _ResidualReturningNorm(torch.nn.Module): + def __init__(self, inner): + super().__init__() + self.inner = inner + + def forward(self, x): + return self.inner(x), x + + setattr(layer, norm_attr, _ResidualReturningNorm(getattr(layer, norm_attr)).cuda()) + + n_channels = num_streams * hidden_size + hidden_states = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + with pytest.raises(ValueError, match=norm_attr): + layer(hidden_states=hidden_states, attention_mask=attention_mask) + + +class TestMHCRecomputeMemorySaving: + """Verify that 'mhc' in recompute_modules actually reduces peak GPU memory.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute, + recompute_block_size=2, + ): + """Run a full forward + backward pass and return (peak memory, output grad). + + When use_recompute=True, a new CheckpointWithoutOutputManager is created every + `recompute_block_size` layers, mirroring TransformerBlock's + build_mhc_recompute_layer_plan logic. + """ + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + num_layers=num_layers, + recompute_modules=["core_attn", "mhc"] if use_recompute else None, + recompute_granularity='selective' if use_recompute else None, + ) + layer_spec = _make_mhc_layer_spec() + layers = [ + HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=i + 1 + ).cuda() + for i in range(num_layers) + ] + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + + manager = CheckpointWithoutOutputManager() if use_recompute else None + + h = hidden_states + for i, layer in enumerate(layers): + is_last_in_block = (i == num_layers - 1) or ((i + 1) % recompute_block_size == 0) + kwargs = dict(hidden_states=h, attention_mask=attention_mask) + if manager is not None: + manager.is_last_layer_in_recompute_block = is_last_in_block + kwargs['mhc_recompute_manager'] = manager + h, _ = layer(**kwargs) + if manager is not None and is_last_in_block: + manager.discard_all_outputs_and_register_unified_recompute(h) + if i < num_layers - 1: + manager = CheckpointWithoutOutputManager() + + loss = h.sum() + loss.backward() + torch.cuda.synchronize() + + peak_mem = torch.cuda.max_memory_allocated() + grad = hidden_states.grad.clone() + + del layers, hidden_states, h, loss, manager + torch.cuda.empty_cache() + + return peak_mem, grad + + def test_recompute_reduces_peak_memory(self): + """Peak memory with recompute (block_size=2) should be lower than without.""" + num_layers = 8 + hidden_size = 128 + num_streams = 4 + seq_len = 64 + batch_size = 4 + + peak_no_recompute, _ = self._run_forward_backward( + num_layers, hidden_size, num_streams, seq_len, batch_size, use_recompute=False + ) + peak_recompute, _ = self._run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute=True, + recompute_block_size=2, + ) + + saving_pct = (peak_no_recompute - peak_recompute) / peak_no_recompute * 100 + + assert peak_recompute < peak_no_recompute, ( + f"Recompute should reduce peak memory, but got " + f"no_recompute={peak_no_recompute / 1e6:.1f}MB vs " + f"recompute={peak_recompute / 1e6:.1f}MB " + f"(saving={saving_pct:.1f}%)" + ) + + +class TestMHCWithCudaGraph: + """Test HyperConnectionTransformerLayer compatibility with CUDA graphs. + + CUDA graph capture requires static computation graphs and fixed tensor shapes. + These tests verify that the mHC layer properly supports the CUDA graph interface + defined in GraphableMegatronModule and TransformerLayer. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123, use_cudagraphable_rng=True, force_reset_rng=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer(self, hidden_size=64, num_streams=4, **extra_config): + config = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams, **extra_config) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_get_layer_static_inputs_shape_for_mhc(self): + """get_layer_static_inputs must return [s, b, n*C] for mHC layers. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. If the shape is [s, b, C] instead of [s, b, n*C], the graph + capture will produce a shape mismatch at the first hyper connection module. + """ + layer, config = self._create_mhc_layer() + seq_length = 32 + micro_batch_size = 2 + + static_inputs = layer.get_layer_static_inputs(seq_length, micro_batch_size) + hidden_states = static_inputs["hidden_states"] + + expected_hidden_dim = config.mhc_num_residual_streams * config.hidden_size + assert hidden_states.shape[-1] == expected_hidden_dim, ( + f"get_layer_static_inputs returns hidden dim {hidden_states.shape[-1]} " + f"but mHC expects {expected_hidden_dim} (n={config.mhc_num_residual_streams} * " + f"C={config.hidden_size}). " + f"HyperConnectionTransformerLayer must override get_layer_static_inputs." + ) + + def test_submodules_under_cudagraphs_includes_hyper_connection(self): + """_get_submodules_under_cudagraphs must include hyper connection modules. + + CUDA graph manual hooks are set up for parameters of submodules returned + by this method. Missing hyper connection modules means their parameters + (mapping_proj, alpha_*, bias) will not get proper pre-forward hooks during + graph replay, leading to stale parameter values. + """ + layer, config = self._create_mhc_layer( + cuda_graph_modules=[CudaGraphModule.attn, CudaGraphModule.mlp] + ) + + submodules = layer._get_submodules_under_cudagraphs() + + assert layer.self_attention_hyper_connection in submodules + assert layer.mlp_hyper_connection in submodules + + def test_forward_through_te_cuda_graph_capture_path(self): + """_te_cuda_graph_capture must produce correct output shapes for mHC. + + TE CUDA graph capture calls _te_cuda_graph_capture() during warmup. + For mHC layers, the input must be n-stream [s, b, n*C] and output must + also be [s, b, n*C]. + """ + layer, config = self._create_mhc_layer() + layer.eval() + + seq_len = 8 + batch_size = 2 + n_channels = config.mhc_num_residual_streams * config.hidden_size + + hidden_states = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + with torch.no_grad(): + outputs = layer._te_cuda_graph_capture( + hidden_states=hidden_states, attention_mask=attention_mask + ) + + if isinstance(outputs, tuple): + output = outputs[0] + else: + output = outputs + + assert output.shape == (seq_len, batch_size, n_channels), ( + f"_te_cuda_graph_capture output shape {output.shape} != " + f"expected {(seq_len, batch_size, n_channels)}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection(self): + """End-to-end CUDA graph capture and replay for forward+backward with mHC. + + Captures both the forward and backward pass of HyperConnectionTransformerLayer + into a torch.cuda.CUDAGraph and replays it with fresh input data, verifying + that the computation graph is fully static (capturable) and produces correct + output shapes and non-trivial gradients. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.mhc_num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream to trigger lazy allocations + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + out, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + # Set .grad to None so backward allocates fresh gradient tensors in the + # graph's private memory pool during capture. + layer.zero_grad(set_to_none=True) + static_input.grad = None + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + output.sum().backward() + + # Replay with new input data. + # Use no_grad because backward inside the captured graph already + # bumped the autograd version counter on static_input, making + # in-place copy_ illegal without disabling grad tracking. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Verify numerical consistency: graph replay should match eager execution + # with the same input and weights. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer(hidden_states=eager_input, attention_mask=attention_mask) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection_and_recompute(self): + """CUDA graph capture+replay for fwd+bwd with mHC and CheckpointWithoutOutputManager. + + When a CheckpointWithoutOutputManager is used, additional CheckpointWithoutOutput + objects are created for layernorm and hyper-connection operations. The + manager discards intermediate activations during forward (storage.resize_(0)) + and recomputes them during backward via a unified gradient hook. + This test verifies the full capture+replay still works correctly. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.mhc_num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream; fresh manager per iteration to avoid stale state. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + out, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=mgr, + ) + mgr.discard_all_outputs_and_register_unified_recompute(out) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + layer.zero_grad(set_to_none=True) + static_input.grad = None + + capture_mgr = CheckpointWithoutOutputManager() + capture_mgr.is_last_layer_in_recompute_block = True + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=capture_mgr, + ) + capture_mgr.discard_all_outputs_and_register_unified_recompute(output) + output.sum().backward() + + # Replay with new input data. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Numerical consistency: graph replay vs eager with the same input. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_mgr = CheckpointWithoutOutputManager() + eager_mgr.is_last_layer_in_recompute_block = True + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer( + hidden_states=eager_input, + attention_mask=attention_mask, + mhc_recompute_manager=eager_mgr, + ) + eager_mgr.discard_all_outputs_and_register_unified_recompute(eager_output) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_mcore_cudagraph_manager_with_mhc_recompute_manager(self): + """MCore CudaGraphManager must not crash on mhc_recompute_manager kwarg. + + When cuda_graph_impl="local" is set, HyperConnectionTransformerLayer.__call__ + runs first and pops mhc_recompute_manager off kwargs before + super().__call__ → MegatronModule.__call__ → CudaGraphManager.__call__, + which iterates over all kwargs to check supported types. + CheckpointWithoutOutputManager (used by mhc_recompute_manager) is not a + CUDA-graph-supported type. + + This test verifies that mhc_recompute_manager is properly extracted + from kwargs before the CudaGraphManager sees them, preventing the + AssertionError that would otherwise occur. + """ + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + assert hasattr( + layer, 'cudagraph_manager' + ), "Layer should have cudagraph_manager with cuda_graph_impl='local'" + + seq_len = 8 + batch_size = 2 + n_channels = config.mhc_num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + + output, context = layer( + hidden_states=hidden_states, attention_mask=attention_mask, mhc_recompute_manager=mgr + ) + + assert output.shape == (seq_len, batch_size, n_channels) + + def test_mcore_cudagraph_manager_without_mhc_recompute_manager(self): + """MCore CudaGraphManager path works when mhc_recompute_manager is None.""" + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.mhc_num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == (seq_len, batch_size, n_channels) + + +class TestMHCWithOffloading: + """Test HyperConnectionTransformerLayer with fine-grained activation offloading. + + Fine-grained activation offloading transfers specific activations (e.g., layernorm + inputs) to CPU during forward and reloads them during backward. These tests verify + that the mHC layer's multi-stream architecture works correctly with offloading. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer_with_offloading( + self, hidden_size=64, num_streams=4, offload_modules=None + ): + if offload_modules is None: + offload_modules = ["attn_norm", "mlp_norm"] + + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=offload_modules, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_forward_backward_with_offloading(self): + """Forward+backward should work with activation offloading enabled. + + This exercises the off_interface context manager around layernorms in + the mHC forward path, including the group_commit that commits the + offloading group for the aggregated 1-stream layernorm input. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + layer, config = self._create_mhc_layer_with_offloading() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.mhc_num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler( + pp_rank=0, vp_size=1, vp_stage=0, min_offloaded_tensor_size=0 + ) + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + + loss = output.sum() + loss.backward() + + assert hidden_states.grad is not None, "Gradients should flow through offloaded path" + assert hidden_states.grad.shape == hidden_states.shape + assert hidden_states.grad.abs().sum() > 0, "Gradients should be non-trivial" + + PipelineOffloadManager.reset_instance() + + def test_offloading_numerical_equivalence(self): + """Offloaded forward+backward must produce the same result as non-offloaded. + + Compares outputs and gradients between a layer with offloading disabled + vs enabled to ensure the offloading path does not corrupt activations. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + PipelineOffloadManager.reset_instance() + + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + torch.manual_seed(42) + input_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Run without offloading + config_no_offload = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams) + layer_spec = _make_mhc_layer_spec() + layer_no_offload = HyperConnectionTransformerLayer( + config_no_offload, layer_spec.submodules + ).cuda() + layer_no_offload.train() + + h1 = input_data.clone().detach().requires_grad_(True) + out1, _ = layer_no_offload(hidden_states=h1, attention_mask=attention_mask) + out1.sum().backward() + grad_no_offload = h1.grad.clone() + out1_detached = out1.detach().clone() + + # Run with offloading using the same weights + config_offload = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=["attn_norm", "mlp_norm"], + ) + layer_offload = HyperConnectionTransformerLayer( + config_offload, layer_spec.submodules + ).cuda() + layer_offload.load_state_dict(layer_no_offload.state_dict()) + layer_offload.train() + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler( + pp_rank=0, vp_size=1, vp_stage=0, min_offloaded_tensor_size=0 + ) + + h2 = input_data.clone().detach().requires_grad_(True) + out2, _ = layer_offload(hidden_states=h2, attention_mask=attention_mask) + out2.sum().backward() + grad_offload = h2.grad.clone() + + PipelineOffloadManager.reset_instance() + + assert torch.allclose(out1_detached, out2.detach(), atol=1e-5), ( + f"Forward outputs differ: max diff = " + f"{(out1_detached - out2.detach()).abs().max().item()}" + ) + assert torch.allclose(grad_no_offload, grad_offload, atol=1e-5), ( + f"Gradients differ: max diff = " + f"{(grad_no_offload - grad_offload).abs().max().item()}" + ) From d7e2c137351dc35560d82001dc5a64c1cd816e84 Mon Sep 17 00:00:00 2001 From: Ajay Date: Tue, 18 Aug 2026 00:08:49 +0000 Subject: [PATCH 066/101] ci: Update workflow to include NeMo RL testing (#6548) Signed-off-by: Ajay Balasa --- .github/workflows/cicd-main.yml | 108 ++++++++++++++++++++++++++++++++ skills/mcore-cicd/SKILL.md | 3 + 2 files changed, 111 insertions(+) diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 90fb157313e..ab24d5e3112 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -183,6 +183,7 @@ jobs: lts: ${{ steps.configure.outputs.lts }} mbridge_suite: ${{ steps.configure.outputs.mbridge_suite }} run_mbridge: ${{ steps.configure.outputs.run_mbridge }} + run_nemo_rl: ${{ steps.configure.outputs.run_nemo_rl }} dev: ${{ steps.configure.outputs.dev }} cadence: ${{ steps.configure.outputs.cadence }} cadence_bypass: ${{ steps.configure.outputs.cadence_bypass }} @@ -240,6 +241,7 @@ jobs: HAS_RUN_FUNCTIONAL=$(echo "$LABELS" | jq 'any(. == "Run functional tests")') HAS_LTS=$(echo "$LABELS" | jq 'any(. == "container::lts")') HAS_MBRIDGE=$(echo "$LABELS" | jq 'any(. == "Run MBridge tests")') + HAS_NEMO_RL=$(echo "$LABELS" | jq 'any(. == "Run NeMoRL tests")') if [ "$IS_MERGE_GROUP" == "true" ]; then SCOPE=L1; N_REPEAT=1; LIGHTWEIGHT=false @@ -278,6 +280,14 @@ jobs: RUN_MBRIDGE=false fi + # NeMo RL job gating: PR authors opt in by adding the `Run NeMoRL tests` + # label. Merge queue validation intentionally skips the downstream run. + if [ "$HAS_NEMO_RL" == "true" ] && [ "$IS_MERGE_GROUP" != "true" ]; then + RUN_NEMO_RL=true + else + RUN_NEMO_RL=false + fi + # Cadence: trigger-driven test selection axis (see filter_by_cadence # in tests/test_utils/python_scripts/recipe_parser.py). PR labels # `Run tests` and `Run functional tests` bypass the cadence filter so @@ -306,6 +316,7 @@ jobs: echo "lts=$HAS_LTS" | tee -a $GITHUB_OUTPUT echo "mbridge_suite=$MBRIDGE_SUITE" | tee -a $GITHUB_OUTPUT echo "run_mbridge=$RUN_MBRIDGE" | tee -a $GITHUB_OUTPUT + echo "run_nemo_rl=$RUN_NEMO_RL" | tee -a $GITHUB_OUTPUT echo "dev=$DEV" | tee -a $GITHUB_OUTPUT echo "cadence=$CADENCE_OUTPUT" | tee -a $GITHUB_OUTPUT echo "cadence_bypass=$CADENCE_BYPASS" | tee -a $GITHUB_OUTPUT @@ -333,6 +344,7 @@ jobs: | \`lts\` | \`$HAS_LTS\` | | \`dev\` | \`$DEV\` | | \`run_mbridge\` | \`$RUN_MBRIDGE\` | + | \`run_nemo_rl\` | \`$RUN_NEMO_RL\` | | \`mbridge_suite\` | \`$MBRIDGE_SUITE\` | | \`cadence\` | \`$CADENCE\` | | \`cadence_bypass\` | \`$CADENCE_BYPASS\` | @@ -370,6 +382,7 @@ jobs: - **\`dev\`**: uses the latest development container base image (default) - **\`cadence\`**: per-test trigger filter (recipe \`cadence:\` field). Recipes default to \`[pr, nightly, mergegroup]\`. - **\`run_mbridge\`**: whether to trigger the Megatron-Bridge downstream CI. Off for PR pushes by default; flip on by adding the _Run MBridge tests_ label. + - **\`run_nemo_rl\`**: whether to trigger the NeMo RL downstream CI. Off for merge queue validation; PR authors can opt in by adding the _Run NeMoRL tests_ label. SUMMARY linting: @@ -547,6 +560,99 @@ jobs: • Run: ${{ needs.cicd-mbridge-testing.result == 'failure' && format('cc ', secrets.SLACK_NEMO_MB_CODEOWNERS_GROUP_ID) || '' }} + cicd-nemo-rl-testing: + runs-on: ubuntu-latest + needs: + - pre-flight + - configure + - cicd-wait-in-queue + - cicd-parse-downstream-testing + # Skip downstream NeMo RL testing for docs-only and merge queue changes. + # PR authors opt in with the "Run NeMoRL tests" label. + if: | + needs.pre-flight.result != 'cancelled' + && needs.configure.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-parse-downstream-testing.result != 'cancelled' + && needs.configure.outputs.run_nemo_rl == 'true' + && github.event_name != 'merge_group' + && needs.pre-flight.outputs.docs_only == 'false' + && ( + success() + || needs.pre-flight.outputs.force_run_all == 'true' + ) + && !cancelled() + steps: + - name: Get PR info + id: get-pr-info + if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' + uses: nv-gha-runners/get-pr-info@090577647b8ddc4e06e809e264f7881650ecdccf # main + + - name: Checkout NeMo RL + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + ref: main + repository: NVIDIA-NeMo/RL + path: nemo-rl + token: ${{ secrets.PAT }} + + - name: Create testing branch + env: + NEMO_RL_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} + run: | + cd nemo-rl + git fetch origin main + git checkout -b ${{ env.NEMO_RL_BRANCH_NAME }} origin/main + git push origin ${{ env.NEMO_RL_BRANCH_NAME }} --force + + # Dispatch on a per-run branch so NeMo RL's concurrency policy does not + # cancel an unrelated MCore-triggered workflow running against `main`. + - name: Trigger NeMo RL tests + uses: convictional/trigger-workflow-and-wait@f69fa9eedd3c62a599220f4d5745230e237904be # v1.6.5 + env: + NEMO_RL_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} + with: + owner: NVIDIA-NeMo + repo: RL + workflow_file_name: cicd-main.yml + github_token: ${{ secrets.PAT }} + ref: ${{ env.NEMO_RL_BRANCH_NAME }} + wait_interval: 60 + propagate_failure: true + client_payload: | + { + "test_to_run": "L1", + "test_script_pattern": "L1_Functional_Tests_Megatron_4", + "test_platform": "h100", + "megatron_lm_commit": "${{ github.sha }}" + } + + - name: Delete testing branch + if: always() + env: + NEMO_RL_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} + run: | + cd nemo-rl + git push origin --delete ${{ env.NEMO_RL_BRANCH_NAME }} + + cicd-nemo-rl-testing-notify: + runs-on: ubuntu-latest + needs: [cicd-nemo-rl-testing] + # Notify on both success and failure of the NeMo RL downstream tests. + # Skipped/cancelled runs are intentionally not announced. + if: | + always() + && (needs.cicd-nemo-rl-testing.result == 'success' || needs.cicd-nemo-rl-testing.result == 'failure') + steps: + - name: Send Slack alert + uses: NVIDIA-NeMo/FW-CI-templates/.github/actions/send-slack-alert@209ac7913b0419a5ccbac47b02d00fbea4939243 # main + with: + webhook: ${{ secrets.SLACK_WH_MLM_RL_ALERTS }} + message: | + ${{ needs.cicd-nemo-rl-testing.result == 'success' && ':white_check_mark: *NeMo RL downstream tests passed*' || ':rotating_light: *NeMo RL downstream tests failed*' }} + • Trigger: `${{ github.event_name }}` on `${{ github.ref_name }}` + • Run: + cicd-compute-build-matrix: runs-on: ubuntu-latest needs: [is-not-external-contributor] @@ -1256,6 +1362,7 @@ jobs: and (.conclusion == "failure" or .conclusion == "cancelled") and .name != "merge-queue-notification" and .name != "cicd-mbridge-testing" + and .name != "cicd-nemo-rl-testing" )] | length ') || BAD_JOBS=0 @@ -1267,6 +1374,7 @@ jobs: and (.conclusion == "failure" or .conclusion == "cancelled") and .name != "merge-queue-notification" and .name != "cicd-mbridge-testing" + and .name != "cicd-nemo-rl-testing" ) | .name + " → " + .conclusion ' FAILED=true diff --git a/skills/mcore-cicd/SKILL.md b/skills/mcore-cicd/SKILL.md index df7efab03a6..ea2efd30766 100644 --- a/skills/mcore-cicd/SKILL.md +++ b/skills/mcore-cicd/SKILL.md @@ -23,6 +23,7 @@ For PR-label or trigger questions, lead with the exact values: explicitly asks for LTS validation; never add it on your own initiative, even for a container or dependency change.** - `Run MBridge tests` additionally triggers the MBridge L1 suite. +- `Run NeMoRL tests` additionally triggers NeMo RL's Megatron functional test suite. - ⚠️ **WARNING — destructive remote write.** `tools/trigger_internal_ci.py` **force-pushes the current branch** to the internal GitLab remote as `pull-request/`. Always run with `--dry-run` first and confirm the @@ -78,6 +79,7 @@ The CI pipeline reads PR labels to decide test scope, n_repeat, and container im |-------|--------| | **`container::lts`** | Build on the older long-term-support NGC PyTorch base instead of `dev`'s latest — a backward-compat check, not a different test set (combinable with any scope label) | | **`Run MBridge tests`** | Also triggers the MBridge L1 test suite | +| **`Run NeMoRL tests`** | Also triggers NeMo RL's Megatron functional test suite | ### Which label to attach when opening a PR @@ -92,6 +94,7 @@ The CI pipeline reads PR labels to decide test scope, n_repeat, and container im | Could affect training numerics (model arch, attention, optimizer, distributed, MoE routing) | `Run functional tests` | | Container or dependency changes (`docker/`, `pyproject.toml`, `uv.lock`) | `Run tests` (add `container::lts` **only if the user explicitly asks** to validate LTS) | | Touches MBridge integration | add `Run MBridge tests` | +| Could affect NeMo RL's Megatron integration | add `Run NeMoRL tests` | **Rule of thumb:** default to `Run tests`. Always use `Run functional tests` when the PR adds new test cases (golden values must be generated) or when the change could plausibly shift loss curves. From 574e04ebf802b3b3cfc90d97dcbfc0721e7313af Mon Sep 17 00:00:00 2001 From: Tom Long Date: Tue, 18 Aug 2026 00:32:24 +0000 Subject: [PATCH 067/101] Add data-schedule utilities and GPTDatasetConfig fields for sequence packing (#5903) Signed-off-by: ilml Co-authored-by: Claude Fable 5 --- megatron/core/datasets/data_schedule.py | 75 +-- megatron/core/datasets/data_schedule_utils.py | 615 ++++++++++++++++++ megatron/core/datasets/gpt_dataset.py | 37 +- 3 files changed, 683 insertions(+), 44 deletions(-) create mode 100644 megatron/core/datasets/data_schedule_utils.py diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 0f016473b6a..26d083c3037 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -5,6 +5,7 @@ import torch from megatron.core import parallel_state +from megatron.core.datasets.data_schedule_utils import _get_global_seqlens_and_ids from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler from megatron.core.process_groups_config import ProcessGroupCollection @@ -60,47 +61,16 @@ def get_global_seqlens(self, subsample_seqlens: torch.Tensor) -> List[int]: We find the number of subsamples each rank holds and then gather the sequence lengths of all subsamples from all ranks. - """ - # Collect the number of subsamples from all ranks - local_len = torch.tensor([subsample_seqlens.shape[0]], dtype=torch.int32).cuda() - dp_subsample_count = [torch.zeros_like(local_len) for _ in range(self.dp_group.size())] - torch.distributed.all_gather(dp_subsample_count, local_len, group=self.dp_group) - - # Find the max number of subsamples across all ranks and pad subsample_seqlens to max length - dp_subsample_counts = torch.stack(dp_subsample_count, dim=0).cpu().view(-1) - max_sub_samples = int(dp_subsample_counts.max().item()) - if local_len.item() < max_sub_samples: - subsample_seqlens_padded = torch.cat( - [ - subsample_seqlens, - torch.zeros(max_sub_samples - local_len.item(), dtype=torch.int32).cuda(), - ], - dim=0, - ) - else: - subsample_seqlens_padded = subsample_seqlens - - # Gather the subsample_seqlens from all ranks - seqlens_gathered = [ - torch.empty_like(subsample_seqlens_padded) for _ in range(self.dp_group.size()) - ] - torch.distributed.all_gather( - seqlens_gathered, subsample_seqlens_padded, group=self.dp_group + Delegates to ``data_schedule_utils._get_global_seqlens_and_ids``. The + shared helper returns ``offsets`` with one extra trailing entry (the + total subsample count); this method preserves the original contract of + returning per-rank start offsets only. + """ + _, _, offsets, seqlens_gathered = _get_global_seqlens_and_ids( + subsample_seqlens, self.dp_group ) - - # Trim each seqlens_gathered to the length of the correct sample - for dp_rank, seqlen in enumerate(seqlens_gathered): - seqlens_gathered[dp_rank] = seqlen[: dp_subsample_counts[dp_rank]] - - seqlens_gathered = torch.cat(seqlens_gathered, dim=0) - seqlens_gathered = seqlens_gathered.cpu().tolist() - - # Calculate the offsets to assign unique global ID to each subsample. - csum = torch.cumsum(dp_subsample_counts, dim=0, dtype=torch.int32) - offsets = torch.cat([torch.zeros(1, dtype=torch.int32), csum[:-1]], dim=0) - - return seqlens_gathered, offsets + return seqlens_gathered, offsets[:-1] def get_global_id_seqlens(self, num_local_subsamples, offsets, seqlens_gathered): """ @@ -108,6 +78,10 @@ def get_global_id_seqlens(self, num_local_subsamples, offsets, seqlens_gathered) We assign a unique global ID to each subsample. + Kept as a local implementation: the shared helper fuses this pure + indexing step with the collective gather, so delegating here would + re-run the all-gathers. + Returns: global_id_seqlens: list of (global_id, seqlen) tuples for scheduling. global_ids_this_rank: list of global IDs locally present on this rank. @@ -133,6 +107,16 @@ def _gid_to_src_rank(self, gid: int, offsets: List[int]) -> int: ) return hdp_rank + # NOTE: data_schedule_utils.reroute_samples_to_dcp_ranks is the sequence- + # packing generalization of this method. It is intentionally NOT used here + # because it is not behavior-preserving for the hybrid-CP path: + # * its empty-send fallback allocates a 1-element (not 0-element) buffer, + # which fails all_to_all_single split validation when a rank has + # nothing to send (reachable here); + # * it transports the "original_seq_len"/"padded_seq_len" metadata keys + # with count-based splits, keys the hybrid-CP batch does not carry; + # * its rank mapping applies "% dp_cp_group.size()" for PP support, + # while hybrid CP asserts pipeline_model_parallel_size == 1. def reroute_samples_to_hdp_ranks( self, batch, global_ids_this_rank, global_id_seqlens, sample_id_groups, offsets ): @@ -242,6 +226,13 @@ def _unpack_sample_by_key(key: str, recv_tensor: torch.Tensor): } return recv_sample_with_id + # NOTE: data_schedule_utils._unpack_batch is the sequence-packing variant + # of this method. It is intentionally NOT used here because it changes the + # sub-sample contract for the hybrid-CP path: it copies a fixed key list + # (dropping any extra dataset keys), squeezes 2-D tensors in place, and + # adds synthesized "original_seq_len"/"padded_seq_len" keys, which would + # alter the all-to-all key set in reroute_samples_to_hdp_ranks and the + # samples returned from __next__. def unpack_batch(self, batch): """ Unpacks the packed samples into a list of sub-samples. @@ -284,10 +275,8 @@ def __next__(self) -> Any: subsample_seqlens = torch.tensor(subsample_seqlens, dtype=torch.int32).cuda() subsample_seqlens = subsample_seqlens[subsample_seqlens != 0] - seqlens_gathered, offsets = self.get_global_seqlens(subsample_seqlens) - - global_id_seqlens, global_ids_this_rank = self.get_global_id_seqlens( - subsample_seqlens.shape[0], offsets, seqlens_gathered + global_id_seqlens, global_ids_this_rank, offsets, _ = _get_global_seqlens_and_ids( + subsample_seqlens, self.dp_group ) groups, sample_id_groups = self.cp_balancing_scheduler.get_groups_and_subsamples( diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py new file mode 100644 index 00000000000..e4dd6fb8bd0 --- /dev/null +++ b/megatron/core/datasets/data_schedule_utils.py @@ -0,0 +1,615 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from typing import Dict, List + +import torch + +from megatron.core.rerun_state_machine import RerunDataIterator + + +def _unpack_batch(batch: List[Dict[str, torch.Tensor]]) -> List[Dict[str, torch.Tensor]]: + """ + Unpacks the packed samples into a list of sub-samples. + Since each sub-sample may be routed to different DPxCP ranks, + we unpack the sample here to avoid unnecessarily transferring + the entire packed sample. + + Two mutually exclusive input shapes are accepted, and every sample in + ``batch`` must use the same one: + + * **Pre-packed** (e.g. :class:`SFTDataset`): each sample carries a + ``cu_seqlens`` tensor and the tokens of multiple sub-samples + concatenated together. We slice them apart and synthesize + ``original_seq_len`` / ``padded_seq_len`` from the cu_seqlens deltas. + + * **Already unpacked** (e.g. :class:`VarlenDataset`): each sample is a + single sub-sample that already carries ``padded_seq_len`` (and + usually ``original_seq_len``). We just normalize the leading batch + dimension introduced by the default collate_fn and return as-is. + + The shape is decided once for the whole batch and asserted per sample, so a + dataset that emits both keys cannot silently bypass the ``cu_seqlens`` + slicing below. + """ + if not batch: + return batch + + # Pick the input shape from the first sample, then validate every sample + # against it and normalize the collate dimension in the same pass. + is_unpacked = "padded_seq_len" in batch[0] + for i, sample in enumerate(batch): + assert ("padded_seq_len" in sample) == is_unpacked, ( + f"_unpack_batch got a mixed batch: sample {i} and sample 0 disagree on " + "whether they carry 'padded_seq_len' (already unpacked) or not (pre-packed)." + ) + assert ("cu_seqlens" in sample) != is_unpacked, ( + f"_unpack_batch: sample {i} must carry exactly one of 'padded_seq_len' " + "(already unpacked, e.g. VarlenDataset) or 'cu_seqlens' (pre-packed, " + "e.g. SFTDataset)." + ) + for key, value in sample.items(): + if value.ndim == 2: + # Drop the redundant batch dimension added by the default + # collate_fn in the pytorch dataloader. squeeze(0) is a silent + # no-op when the leading dimension is not 1, so assert on it + # instead of slicing along the batch dimension further down. + # The packing path installs an identity collate_fn (see + # build_pretraining_data_loader), which never adds this + # dimension in the first place and therefore supports + # micro_batch_size > 1; the default collate_fn only works here + # with micro_batch_size == 1. + assert value.shape[0] == 1, ( + f"_unpack_batch got '{key}' with shape {tuple(value.shape)}; the " + "packed-sequence path needs one sub-sample per collated entry. Use " + "micro_batch_size 1 with the default collate_fn, or an identity " + "collate_fn." + ) + sample[key] = value.squeeze(0) + + # Short-circuit for datasets that already emit one sub-sample per index. + if is_unpacked: + for sample in batch: + if "original_seq_len" not in sample: + sample["original_seq_len"] = sample["padded_seq_len"].clone() + return batch + + batch_unpacked = [] + device = batch[0]["tokens"].device + original_seq_lens = [] + padded_seq_lens = [] + for sample in batch: + for sub_sample in range(sample["cu_seqlens"].shape[0] - 1): + sub_sample_dict = {} + start_idx = sample["cu_seqlens"][sub_sample] + end_idx = sample["cu_seqlens"][sub_sample + 1] + if end_idx - start_idx == 0: + continue + for key in ["tokens", "labels", "loss_mask", "position_ids"]: + sub_sample_dict[key] = sample[key][start_idx:end_idx] + # Since sft_dataset.py does not provide cu_seqlens_original, + # we assume original_seq_len equals padded_seq_len here. + # Ideally the dataset should define the pre-padding seq_len. + seq_len = (end_idx - start_idx).item() + original_seq_lens.append(seq_len) + padded_seq_lens.append(seq_len) + batch_unpacked.append(sub_sample_dict) + + # Single H2D transfer for all seq lens + original_seq_lens_cuda = torch.tensor(original_seq_lens, device=device) + padded_seq_lens_cuda = torch.tensor(padded_seq_lens, device=device) + for i, sub_sample_dict in enumerate(batch_unpacked): + sub_sample_dict["original_seq_len"] = original_seq_lens_cuda[i : i + 1] + sub_sample_dict["padded_seq_len"] = padded_seq_lens_cuda[i : i + 1] + + return batch_unpacked + + +def _get_global_seqlens_and_ids(subsample_seqlens: torch.Tensor, dp_group): + """ + Gathers the sequence lengths of all subsamples from all DP ranks and calculates global IDs. + """ + # Collect the number of subsamples from all ranks + num_local_subsamples = subsample_seqlens.shape[0] + local_len = torch.tensor([num_local_subsamples], dtype=torch.int32).cuda() + dp_subsample_count = [torch.zeros_like(local_len) for _ in range(dp_group.size())] + torch.distributed.all_gather(dp_subsample_count, local_len, group=dp_group) + + # Find the max number of subsamples across all ranks and pad subsample_seqlens to max length + dp_subsample_counts = torch.stack(dp_subsample_count, dim=0).cpu().view(-1) + max_sub_samples = int(dp_subsample_counts.max().item()) + + if num_local_subsamples < max_sub_samples: + subsample_seqlens_padded = torch.cat( + [ + subsample_seqlens, + torch.zeros(max_sub_samples - num_local_subsamples, dtype=torch.int32).cuda(), + ], + dim=0, + ) + else: + subsample_seqlens_padded = subsample_seqlens + + # Gather the subsample_seqlens from all ranks + seqlens_gathered = [torch.empty_like(subsample_seqlens_padded) for _ in range(dp_group.size())] + torch.distributed.all_gather(seqlens_gathered, subsample_seqlens_padded, group=dp_group) + + # Trim each seqlens_gathered to the length of the correct sample + for dp_rank, seqlen in enumerate(seqlens_gathered): + seqlens_gathered[dp_rank] = seqlen[: dp_subsample_counts[dp_rank]] + + seqlens_gathered = torch.cat(seqlens_gathered, dim=0) + seqlens_gathered = seqlens_gathered.cpu().tolist() + + # Calculate the offsets to assign unique global ID to each subsample. + csum = torch.cumsum(dp_subsample_counts, dim=0, dtype=torch.int32) + offsets = torch.cat([torch.zeros(1, dtype=torch.int32), csum], dim=0) + + # Calculate global ID for each subsample + dp_rank = dp_group.rank() + global_ids = torch.arange(len(seqlens_gathered), dtype=torch.int32).cuda() + + # Create a list of (global_id, seqlen) tuples for scheduling + global_id_seqlens = [(i, seqlens_gathered[i]) for i in range(len(global_ids))] + + # Get the global IDs locally present on this rank + start_idx = offsets[dp_rank] + end_idx = offsets[dp_rank + 1] + + global_ids_this_rank = global_ids[start_idx:end_idx] + + return global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered + + +def _pack_sequences( + samples: List, padded_lengths: torch.Tensor, original_lengths: torch.Tensor, dev: torch.device +) -> Dict[str, torch.Tensor]: + """Pack multiple samples into a single packed sample.""" + + def _pack_tensors(tensors): + return torch.cat([t.reshape(-1) for t in tensors], dim=0) + + tokens = _pack_tensors([sample["tokens"] for sample in samples]) + labels = _pack_tensors([sample["labels"] for sample in samples]) + loss_mask = _pack_tensors([sample["loss_mask"] for sample in samples]) + position_ids = _pack_tensors([sample["position_ids"] for sample in samples]) + + new_sample = {} + new_sample["tokens"] = tokens + new_sample["labels"] = labels + new_sample["loss_mask"] = loss_mask + new_sample["position_ids"] = position_ids + + padded_lengths = padded_lengths.to(device=dev, dtype=torch.int32, non_blocking=True).reshape(-1) + cu_seqlens_padded = torch.empty(padded_lengths.numel() + 1, device=dev, dtype=torch.int32) + cu_seqlens_padded[0] = 0 + cu_seqlens_padded[1:] = torch.cumsum(padded_lengths, dim=0) + max_seqlen = torch.max(padded_lengths).to(dtype=torch.int32) + + new_sample["cu_seqlens_padded"] = cu_seqlens_padded + new_sample["max_seqlen"] = max_seqlen + + original_lengths = original_lengths.to( + device=dev, dtype=torch.int32, non_blocking=True + ).reshape(-1) + cu_seqlens = torch.empty(original_lengths.numel() + 1, device=dev, dtype=torch.int32) + cu_seqlens[0] = 0 + cu_seqlens[1:] = torch.cumsum(original_lengths, dim=0).reshape(-1) + new_sample["cu_seqlens"] = cu_seqlens + + return new_sample + + +def broadcast_tensor(item, src_rank, group) -> None: + """Broadcast a tensor from src_rank to all ranks in the group.""" + if item is not None: + torch.distributed.broadcast(item, src_rank, group=group) + + +def broadcast_to_pp_group( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + pp_group, + dev, +): + """ + Broadcast num_micro_batches, seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch and metadata to middle PP stages. + Before this broadcast, the new_samples on middle PP stages are None, + after this broadcast, the new_samples on middle PP stages contain the metadata but + without tokens, labels, loss_mask, position_ids. + + Who needs what: + + * **PP rank 0 and the last PP rank** both own a data iterator (only TP rank 0 + on the first and last PP stage does), so both run the whole schedule -> + reroute -> pack pipeline on the same input samples and independently end up + with complete ``new_samples``: tokens, labels, loss_mask, position_ids *and* + the packing metadata. Neither takes anything from this broadcast; the last + stage in particular must keep its own labels / loss_mask. + * **Middle PP stages** have no data iterator, so ``new_samples`` is None on + entry. They only need the packing metadata (max_seqlen / cu_seqlens / + cu_seqlens_padded) to rebuild the packed-sequence params, never the token + tensors. + + The last PP rank still takes part in the transfer because + ``torch.distributed.broadcast`` is a collective over ``pp_group``: every member + has to call it or the group deadlocks. It therefore receives the payload and + drops it, which is what the ``pp_group.rank() != pp_group.size() - 1`` guard + below implements. Filtering it out of the transfer itself would require a + separate "first + middle" process group, which is not worth an extra process + group for a payload of a few hundred bytes per global batch. + """ + + pp_src_rank = torch.distributed.get_process_group_ranks(pp_group)[0] + + # size() > 2 asks "does a middle PP stage exist at all": with 1 or 2 PP ranks + # every rank is a first and/or last stage and already owns its packed samples, + # so there is nobody to broadcast to. + if pp_group.size() > 2: + if pp_group.rank() == 0: + cu_seqlens_lengths = torch.tensor( + [sample["cu_seqlens"].numel() for sample in new_samples], + dtype=torch.float32, + device=dev, + ) + cu_seqlens_padded_lengths = torch.tensor( + [sample["cu_seqlens_padded"].numel() for sample in new_samples], + dtype=torch.float32, + device=dev, + ) + tensor_list = [ + torch.tensor( + [ + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ], + dtype=torch.float32, + device=dev, + ) + ] + for sample in new_samples: + tensor_list.append(sample["max_seqlen"].reshape(1)) + tensor_list.append(cu_seqlens_lengths) + tensor_list.append(cu_seqlens_padded_lengths) + for sample in new_samples: + tensor_list.append(sample["cu_seqlens"]) + tensor_list.append(sample["cu_seqlens_padded"]) + info_to_broadcast = torch.cat(tensor_list, dim=0).to(device=dev, dtype=torch.float32) + info_length_tensor = torch.tensor( + info_to_broadcast.shape[0], dtype=torch.int32, device=dev + ) + broadcast_tensor(info_length_tensor, pp_src_rank, pp_group) + broadcast_tensor(info_to_broadcast, pp_src_rank, pp_group) + else: + # Every non-source rank has to take part in the collective, including + # the last PP stage. + info_length_tensor = torch.tensor(0, dtype=torch.int32, device=dev) + broadcast_tensor(info_length_tensor, pp_src_rank, pp_group) + info_to_broadcast = torch.empty( + info_length_tensor.item(), dtype=torch.float32, device=dev + ) + broadcast_tensor(info_to_broadcast, pp_src_rank, pp_group) + if pp_group.rank() != pp_group.size() - 1: + # Middle PP stages receive the broadcasted info and unpack it. + # Cu-seqlens lengths are encoded explicitly so zero values inside + # the payload cannot be mistaken for tensor boundaries. + # The last PP stage deliberately falls through: it built its own + # new_samples from its own data iterator (with the labels and + # loss_mask this payload does not carry), so it discards what it + # just received rather than overwriting them. + num_micro_batches = int(info_to_broadcast[0].item()) + seqlen_sum_this_global_batch = info_to_broadcast[1].item() + seqlen_squared_sum_this_global_batch = info_to_broadcast[2].item() + + cursor = 3 + max_seqlens = info_to_broadcast[cursor : cursor + num_micro_batches] + cursor += num_micro_batches + cu_seqlens_lengths = info_to_broadcast[cursor : cursor + num_micro_batches].to( + torch.int64 + ) + cursor += num_micro_batches + cu_seqlens_padded_lengths = info_to_broadcast[ + cursor : cursor + num_micro_batches + ].to(torch.int64) + cursor += num_micro_batches + + new_samples = [] + for i in range(num_micro_batches): + cu_seqlens_len = int(cu_seqlens_lengths[i].item()) + cu_seqlens_padded_len = int(cu_seqlens_padded_lengths[i].item()) + new_sample = {} + new_sample["max_seqlen"] = max_seqlens[i].to(torch.int32) + new_sample["cu_seqlens"] = info_to_broadcast[ + cursor : cursor + cu_seqlens_len + ].to(torch.int32) + cursor += cu_seqlens_len + new_sample["cu_seqlens_padded"] = info_to_broadcast[ + cursor : cursor + cu_seqlens_padded_len + ].to(torch.int32) + cursor += cu_seqlens_padded_len + new_samples.append(new_sample) + + return ( + new_samples, + num_micro_batches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) + + +def broadcast_scalars(values: List, group, dev, dtype=torch.float32) -> List: + """ + Broadcast scalar values from rank 0 to all ranks in the group. + + Args: + values: List of scalar values to broadcast (only used on rank 0). + group: The process group to broadcast within. + dev: The device to use for the tensor. + dtype: The data type for the tensor. + + Returns: + List of broadcasted values. + """ + if group.size() <= 1: + return values + + src_rank = torch.distributed.get_process_group_ranks(group)[0] + num_values = len(values) + + if group.rank() == 0: + info_to_broadcast = torch.tensor(values, dtype=dtype, device=dev) + else: + info_to_broadcast = torch.zeros(num_values, dtype=dtype, device=dev) + + broadcast_tensor(info_to_broadcast, src_rank, group) + + if group.rank() != 0: + values = info_to_broadcast.cpu().tolist() + + return values + + +def create_data_iterator(new_samples, pp_group, tp_group, config): + """Handle virtual pipeline parallelism.""" + if ( + config.virtual_pipeline_model_parallel_size is not None + and config.virtual_pipeline_model_parallel_size > 1 + ): + vpp_size = config.virtual_pipeline_model_parallel_size + if tp_group.rank() == 0: + if pp_group.rank() == 0 or pp_group.rank() == pp_group.size() - 1: + metadata = [ + {k: sample[k] for k in ["max_seqlen", "cu_seqlens", "cu_seqlens_padded"]} + for sample in new_samples + ] + if pp_group.rank() == 0: + new_data_iterator = [RerunDataIterator(iter(new_samples))] + [ + RerunDataIterator(iter(metadata)) for _ in range(vpp_size - 1) + ] + else: + new_data_iterator = [ + RerunDataIterator(iter(metadata)) for _ in range(vpp_size - 1) + ] + [RerunDataIterator(iter(new_samples))] + else: + # on middle PP stages, the new_samples are the metadata + metadata = new_samples + new_data_iterator = [RerunDataIterator(iter(metadata)) for _ in range(vpp_size)] + else: + new_data_iterator = [None for _ in range(vpp_size)] + else: + new_data_iterator = RerunDataIterator(iter(new_samples)) if tp_group.rank() == 0 else None + + return new_data_iterator + + +def reroute_samples_to_dcp_ranks( + batch, + global_ids_this_rank, + global_id_seqlens, + sample_id_groups, + offsets, + dp_group, + tp_group, + dp_cp_group, + total_dcp_gpus, +): + """ + Reroutes the sub-samples to the correct rank after scheduling. + + For each key in the batch dict, we perform an all-to-all communication + to transfer the data to the correct ranks. + """ + + def _gid_to_src_rank(gid: int) -> int: + dp_src_rank = torch.bucketize(gid, offsets[1:] - 1) + dcp_rank = ( + torch.distributed.get_process_group_ranks(dp_group)[dp_src_rank] // tp_group.size() + ) % dp_cp_group.size() + return dcp_rank + + gid2local_id = {int(gid): i for i, gid in enumerate(global_ids_this_rank)} + dcp_rank = dp_cp_group.rank() + dp_ranks = torch.distributed.get_process_group_ranks(dp_group) + dp_ranks = [(r // tp_group.size()) % dp_cp_group.size() for r in dp_ranks] + + data_keys = batch[0].keys() + + # Create the send plan + combined_sample_id_groups: List[List[int]] = [[] for _ in range(total_dcp_gpus)] + for d in range(total_dcp_gpus): + for sample_id_group in sample_id_groups: + combined_sample_id_groups[d].extend(sample_id_group[d]) + for dest_rank in range(total_dcp_gpus): + combined_sample_id_groups[dest_rank].sort() + + send_ids_sorted = [ + gid for d in dp_ranks for gid in combined_sample_id_groups[d] if gid in global_ids_this_rank + ] + + send_num_split = [0] * total_dcp_gpus + send_lens_split = [0] * total_dcp_gpus + for dest_rank in range(total_dcp_gpus): + if dest_rank in dp_ranks: + send_seq_lens = [ + global_id_seqlens[gid][1] + for gid in combined_sample_id_groups[dest_rank] + if gid in global_ids_this_rank + ] + send_num_split[dest_rank] = len(send_seq_lens) + send_lens_split[dest_rank] = sum(send_seq_lens) + else: + send_lens_split[dest_rank] = 0 + + # Create the recv plan + recv_sample_id_groups = [[] for _ in range(total_dcp_gpus)] + for gid in combined_sample_id_groups[dcp_rank]: + src_rank = _gid_to_src_rank(gid) + recv_sample_id_groups[src_rank].append(gid) + + recv_lens_split = [0] * total_dcp_gpus + for src_rank in range(total_dcp_gpus): + recv_lens_split[src_rank] = sum( + [global_id_seqlens[gid][1] for gid in recv_sample_id_groups[src_rank]] + ) + + recv_ids_sorted = [gid for d in range(total_dcp_gpus) for gid in recv_sample_id_groups[d]] + recv_counts = [len(recv_sample_id_groups[d]) for d in range(total_dcp_gpus)] + + recv_samples = [{k: None for k in data_keys} for _ in range(sum(recv_counts))] + + def _pack_sample_by_key(key: str) -> torch.Tensor: + flattened_tensors = [] + for gid in send_ids_sorted: + t = batch[gid2local_id[gid]][key].to(torch.cuda.current_device(), non_blocking=True) + flattened_tensors.append(t.reshape(-1)) + return ( + torch.cat(flattened_tensors, dim=0) + if flattened_tensors + else torch.empty(1, device=torch.cuda.current_device(), dtype=batch[0][key].dtype) + ) + + def _unpack_sample_by_key(key: str, recv_tensor: torch.Tensor): + cursor = 0 + for i, gid in enumerate(recv_ids_sorted): + sample_len = ( + 1 if key in ["original_seq_len", "padded_seq_len"] else global_id_seqlens[gid][1] + ) + recv_samples[i][key] = recv_tensor[cursor : cursor + sample_len] + cursor += sample_len + + for key in data_keys: + output_split_sizes, input_split_sizes = ( + (recv_counts, send_num_split) + if key in ["original_seq_len", "padded_seq_len"] + else (recv_lens_split, send_lens_split) + ) + send_tensor = _pack_sample_by_key(key) + recv_tensor_size = sum(output_split_sizes) + recv_tensor = torch.empty( + recv_tensor_size, device=torch.cuda.current_device(), dtype=send_tensor.dtype + ) + torch.distributed.all_to_all_single( + output=recv_tensor, + input=send_tensor, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=dp_cp_group, + ) + _unpack_sample_by_key(key, recv_tensor) + + recv_sample_with_id = {recv_id: recv_samples[i] for i, recv_id in enumerate(recv_ids_sorted)} + return recv_sample_with_id + + +def build_packed_microbatches( + grouped_samples: List[List[Dict[str, torch.Tensor]]], dev: torch.device +) -> List[Dict[str, torch.Tensor]]: + """Build packed samples for each microbatch.""" + num_micro_batches = len(grouped_samples) + seg_starts: List[int] = [0] + original_lens_tensors = [] + padded_lens_tensors = [] + + for i in range(num_micro_batches): + samples = grouped_samples[i] + seg_starts.append(seg_starts[-1] + len(samples)) + original_lens_tensors.extend([s["original_seq_len"].reshape(-1) for s in samples]) + padded_lens_tensors.extend([s["padded_seq_len"].reshape(-1) for s in samples]) + + padded_lens_all_gpu = torch.cat(padded_lens_tensors, dim=0).to(dtype=torch.int32) + original_lens_all_gpu = torch.cat(original_lens_tensors, dim=0).to(dtype=torch.int32) + + new_samples: List[Dict[str, torch.Tensor]] = [] + for i in range(num_micro_batches): + samples = grouped_samples[i] + lens_padded = padded_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]] + lens_original = original_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]] + new_sample = _pack_sequences(samples, lens_padded, lens_original, dev) + new_samples.append(new_sample) + + return new_samples + + +def get_batch_and_global_seqlens(data_iterator, num_microbatches, dp_group): + """ + Get the batch and global sequence lengths. + Each DP rank loads the same number of sequences, so we need to gather the sequence + lengths from all ranks then we can schedule the sequences into groups. + Args: + data_iterator: The data iterator. + num_microbatches: The number of microbatches. + dp_group: The data parallel group. + + Returns: + batch (List[Dict[str, torch.Tensor]]): The sub-samples pulled from this rank's + ``data_iterator`` over ``num_microbatches`` steps, flattened and unpacked + (see :func:`_unpack_batch`). Every dict carries ``tokens`` / ``labels`` / + ``loss_mask`` / ``position_ids`` plus the ``original_seq_len`` and + ``padded_seq_len`` scalars used for scheduling. + global_id_seqlens (List[Tuple[int, int]]): ``(global_id, padded_seq_len)`` for + every sub-sample in the DP group, ordered by DP rank and then by local + index. Identical on all ranks; this is the scheduler's input. + global_ids_this_rank (torch.Tensor): int32 CUDA tensor holding the global IDs of + the sub-samples loaded by this rank, i.e. ``batch[i]`` has global ID + ``global_ids_this_rank[i]``. + offsets (torch.Tensor): int32 CPU tensor of shape ``[dp_size + 1]`` with the + exclusive prefix sum of the per-rank sub-sample counts, so DP rank ``r`` owns + global IDs ``offsets[r]:offsets[r + 1]``. Used by + :func:`reroute_samples_to_dcp_ranks` to map a global ID back to its source + rank. + seqlens_gathered (List[int]): Padded sequence length of every sub-sample in the + DP group, indexed by global ID (``seqlens_gathered[gid]`` equals + ``global_id_seqlens[gid][1]``). Handy for global-batch token counts such as + the FLOPs accounting. + """ + + batch_list = [next(data_iterator) for _ in range(num_microbatches)] + + batch = [] + for item in batch_list: + if isinstance(item, dict): + batch.append(item) + elif isinstance(item, list): + batch.extend(item) + else: + raise ValueError(f"Invalid item type: {type(item)}") + + # in sft_dataset.py, sequences are already packed before rescheduling, + # so we need to unpack them here and repack after rescheduling. + # This is only to adapt to the current megatron-lm sft_dataset. + # If you implement your own dataset, just have __getitem__ return List[Dict] + # and this step can be skipped. + batch = _unpack_batch(batch) + + subsample_seqlens = torch.cat([sample["padded_seq_len"] for sample in batch]).to( + dtype=torch.int32, device=torch.cuda.current_device() + ) + + global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered = ( + _get_global_seqlens_and_ids(subsample_seqlens, dp_group) + ) + + return batch, global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 166a47c084a..b3d43ba67b0 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -80,6 +80,35 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): inter_document_masking: bool = False """When True, return cu_seqlens marking document boundaries within each sample so that attention is restricted to individual documents.""" + sft_mock_dataset_config_json: Optional[str] = None + """JSON string (or path to a JSON file) configuring the mock SFT dataset's + sequence-length distribution. Two modes: + + * ``{"mode": "file", "path": "/path/to/seqlens.csv"}`` -- read the + per-sample sequence lengths from a CSV file. + * ``{"mode": "distribution", "type": "lognormal", "min_seq_len": 1024, + "max_seq_len": 8192, "mean_seq_len": 4096, "lognormal_sigma": 1.1}`` + -- draw lengths from a clipped lognormal distribution. + + Consumed by ``MockSFTLowLevelDataset`` in + ``megatron/training/datasets/sft_dataset.py``, which documents the fields. + """ + + varlen_mock_dataset_config_json: Optional[str] = None + """Mock-dataset config (same JSON schema as ``sft_mock_dataset_config_json``) + used by the ``--use-varlen-dataset`` path; kept separate so the varlen path + does not implicitly inherit SFT-specific knobs.""" + + varlen_sbhd_validation: bool = False + """When True, :class:`VarlenDataset.__getitem__` emits SBHD samples padded + to ``sequence_length`` (no ``cu_seqlens`` / ``original_seq_len`` / + ``padded_seq_len``), bypassing the packed-sequence path. Used to obtain a + SBHD reference run that mirrors the THD path's tokenization but skips all + packing — useful for THD numerical-correctness validation. + + NOTE: this is a debugging/verification knob, not a training feature. + TODO: drop this field once a functional test covers THD-vs-SBHD numerical + parity in CI.""" def __post_init__(self) -> None: """Do asserts and set fields post init""" @@ -91,6 +120,12 @@ def __post_init__(self) -> None: assert self.reset_attention_mask is not None assert self.eod_mask_loss is not None + if self.varlen_sbhd_validation: + assert not self.hybrid_context_parallel, ( + "--varlen-sbhd-validation is incompatible with " + "--hybrid-context-parallel (SBHD mode is not packed)." + ) + self.token_dtype_code = ( None if self.tokenizer.vocab_size is None @@ -145,7 +180,7 @@ def __init__( self.cached_loss_mask = None self.cached_position_ids = None - (self.document_index, self.sample_index, self.shuffle_index) = ( + self.document_index, self.sample_index, self.shuffle_index = ( self._build_document_sample_shuffle_indices() ) From 59ff1719da0622667f4d99ec5e1ce58c2a5cb4a3 Mon Sep 17 00:00:00 2001 From: yhgalaxy <130375321+yhgalaxy@users.noreply.github.com> Date: Tue, 18 Aug 2026 02:56:34 +0000 Subject: [PATCH 068/101] Enhance optimizer state loading with runtime overrides (#3720) Signed-off-by: yhgalaxy Co-authored-by: Guihong Li Co-authored-by: Xin Yao Co-authored-by: Maanu Grover --- megatron/training/checkpointing.py | 41 +++++++++++ tests/unit_tests/test_checkpointing.py | 99 ++++++++++++++++++++++++++ 2 files changed, 140 insertions(+) diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index bbb2ce5d9bc..040cdac4c9b 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -2948,6 +2948,47 @@ def load_model_state_dict(module, state_dict, strict: bool): opt_param_scheduler.load_state_dict(state_dict['lr_scheduler']) else: opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler']) + + # Optimizer state dict can overwrite per-group max_lr/min_lr values. + # If scheduler override is requested, re-apply runtime lr bounds from + # args so scheduler math does not stay pinned to checkpoint lr settings. + if getattr(args, "override_opt_param_scheduler", False): + if ( + optimizer is None + or getattr(optimizer, "is_stub_optimizer", False) + or opt_param_scheduler is None + ): + print_rank_0( + " > WARNING: --override-opt-param-scheduler is set, but optimizer or " + "opt_param_scheduler is not available; skipping override of scheduler " + "and optimizer param_group max_lr/min_lr." + ) + else: + for param_group in optimizer.param_groups: + if param_group.get("is_decoupled_lr", False): + max_lr = getattr(args, "decoupled_lr", None) + min_lr = getattr(args, "decoupled_min_lr", None) + else: + max_lr = args.lr + min_lr = args.min_lr + if max_lr is not None: + param_group["max_lr"] = max_lr + if min_lr is not None: + param_group["min_lr"] = min_lr + # Synchronize scheduler num_steps with consumed_train_samples + # to ensure lr calculation is based on current training progress + if opt_param_scheduler.num_steps != args.consumed_train_samples: + print_rank_0( + f" > WARNING: scheduler num_steps ({opt_param_scheduler.num_steps}) " + f"differs from consumed_train_samples ({args.consumed_train_samples}). " + f"Resetting scheduler num_steps to match consumed_train_samples." + ) + opt_param_scheduler.num_steps = args.consumed_train_samples + opt_param_scheduler.step(increment=0) + print_rank_0( + " > restored optimizer param_group max_lr/min_lr from runtime config " + "because --override-opt-param-scheduler is set" + ) except KeyError as e: print_rank_0( 'Unable to load optimizer from checkpoint {}. ' diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index d9ae2b0f07e..d5c625ca8ca 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -150,6 +150,31 @@ def test_maybe_save_dataloader_state_skips_empty_state_after_barriers(tmp_path): save.assert_not_called() +class MockOptParamScheduler(MockState): + def __init__(self, state_dict): + super().__init__(state_dict) + self.num_steps = state_dict.get("num_steps", 0) + self.step_calls = [] + + def load_state_dict(self, state_dict): + super().load_state_dict(state_dict) + self.num_steps = state_dict.get("num_steps", self.num_steps) + + def step(self, increment=1): + self.step_calls.append(increment) + + +class MockOptimizer(MockState): + def state_dict(self, is_loading=False): + state_dict = super().state_dict(is_loading=is_loading).copy() + state_dict["param_groups"] = [param_group.copy() for param_group in self.param_groups] + return state_dict + + def load_state_dict(self, state_dict): + super().load_state_dict(state_dict) + self.param_groups = [param_group.copy() for param_group in state_dict["param_groups"]] + + @pytest.mark.parametrize( ("checkpoint_args", "configured_num_householder", "expected_num_householder"), [(SimpleNamespace(gdp_num_householder=5), 3, 5), (SimpleNamespace(), 5, 3)], @@ -217,6 +242,7 @@ def create_args(): args.auto_detect_ckpt_format = False args.ckpt_convert_update_legacy_dist_opt_format = False args.ckpt_step = None + args.override_opt_param_scheduler = False args.swiglu = True args.num_experts = 1 args.verify_integrity = False @@ -407,6 +433,79 @@ def test_load_checkpoint( assert new_opt_param_scheduler.state_dict() == opt_param_scheduler.state_dict() +@pytest.mark.parametrize("ckpt_format", ["torch"]) +def test_load_checkpoint_override_opt_param_scheduler( + init_model_parallel, create_ckpt_load_args, tmp_path_dist_ckpt, ckpt_format +): + """Test override_opt_param_scheduler behavior during checkpoint load.""" + args = create_ckpt_load_args + args.ckpt_format = ckpt_format + args.use_distributed_optimizer = False + args.use_dist_ckpt = ckpt_format != "torch" + args.override_opt_param_scheduler = True + args.lr = 1.0 + args.min_lr = 0.1 + args.decoupled_lr = 0.5 + args.decoupled_min_lr = 0.05 + args.consumed_train_samples = 42 + + with TempNamedDir( + tmp_path_dist_ckpt / "test_load_checkpoint_override_opt_param_scheduler", sync=True + ) as ckpt_dir: + args.load = ckpt_dir + args.save = ckpt_dir + args.save_tokenizer_assets = False + set_args(args) + + # Create and save a checkpoint first. + iteration = 123 + config = TransformerConfig(num_layers=1, kv_channels=1) + model = MockModel(config) + + optimizer = MockOptimizer({"optimizer": "optimizer_state"}) + optimizer.param_groups = [ + {"is_decoupled_lr": False, "max_lr": -1.0, "min_lr": -1.0}, + {"is_decoupled_lr": True, "max_lr": -1.0, "min_lr": -1.0}, + ] + opt_param_scheduler = MockOptParamScheduler( + {"opt_param_scheduler": "scheduler_state", "num_steps": 3} + ) + num_floating_point_operations_so_far = 456 + + save_checkpoint( + iteration, [model], optimizer, opt_param_scheduler, num_floating_point_operations_so_far + ) + + # Create new model, optimizer, and scheduler instances to load into. + new_model = MockModel(config) + new_optimizer = MockOptimizer({"optimizer": "dummy1"}) + new_optimizer.param_groups = [ + {"is_decoupled_lr": False, "max_lr": -2.0, "min_lr": -2.0}, + {"is_decoupled_lr": True, "max_lr": -2.0, "min_lr": -2.0}, + ] + new_opt_param_scheduler = MockOptParamScheduler( + {"opt_param_scheduler": "dummy2", "num_steps": 0} + ) + + # Load checkpoint and verify runtime overrides are restored. + loaded_iter, loaded_flops = load_checkpoint( + [new_model], new_optimizer, new_opt_param_scheduler, strict=True + ) + assert loaded_iter == iteration + assert loaded_flops == num_floating_point_operations_so_far + assert new_optimizer.param_groups[0]["max_lr"] == args.lr + assert new_optimizer.param_groups[0]["min_lr"] == args.min_lr + assert new_optimizer.param_groups[1]["max_lr"] == args.decoupled_lr + assert new_optimizer.param_groups[1]["min_lr"] == args.decoupled_min_lr + assert new_opt_param_scheduler.num_steps == args.consumed_train_samples + assert new_opt_param_scheduler.step_calls[-1] == 0 + + # Ensure loading without optimizer/scheduler remains safe. + loaded_iter_none, loaded_flops_none = load_checkpoint([new_model], None, None, strict=True) + assert loaded_iter_none == iteration + assert loaded_flops_none == num_floating_point_operations_so_far + + def test_dist_checkpoint_versioning(init_model_parallel, tmp_path_dist_ckpt, create_ckpt_load_args): """Test distributed checkpoint versioning.""" args = create_ckpt_load_args From 35ae46c5bf61d8000793ddd1c67e946f344fea5d Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Tue, 18 Aug 2026 03:10:10 +0000 Subject: [PATCH 069/101] test(mfsdp): cover v2 training memory (#6567) Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 118 +++++++++++------- 1 file changed, 72 insertions(+), 46 deletions(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 2bb6d92a7ba..ab341521ee3 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -608,6 +608,56 @@ def test_tied_child_parameters_allocate_one_physical_weight(distributed_setup): assert len(list(model.parameters())) == 1 +@pytest.mark.parametrize("main_params_dtype", [torch.bfloat16, torch.float32]) +def test_persistent_sharded_storage(distributed_setup, main_params_dtype): + """FSDP should retain only its sharded weights and gradients at rest.""" + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(device.type, (world_size,)) + dim = 4096 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=8).to(dtype=dtype) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=main_params_dtype) + allocated_before = torch.cuda.memory_allocated(device) + with fully_shard_context(device=device): + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() + persistent_allocated = torch.cuda.memory_allocated(device) - allocated_before + if main_params_dtype == dtype: + # Model and main weights alias, leaving only one BF16 weight buffer and one + # BF16 main-gradient buffer per child. + assert all( + group.model_weight is group.main_weight + for layer in model.layers + for group in layer.parameter_groups + ) + expected_per_child_nbytes = 2 * child_weight_nbytes + else: + # FP32 main weights require a distinct buffer in addition to the BF16 model + # weight and BF16 main-gradient buffers. + main_weight_nbytes = dim * dim * torch.empty((), dtype=main_params_dtype).element_size() + expected_per_child_nbytes = 2 * child_weight_nbytes + main_weight_nbytes + + # All persistent buffers are sharded over the data-parallel group. Small bookkeeping + # allocations stay below 1 MiB. + expected_persistent_nbytes = len(model.layers) * expected_per_child_nbytes // world_size + assert ( + expected_persistent_nbytes <= persistent_allocated < expected_persistent_nbytes + 1024**2 + ), ( + "FSDP persistent memory does not match its sharded weight and gradient storage: " + f"rank={rank}, persistent_allocated={_mb(persistent_allocated)}, " + f"expected={_mb(expected_persistent_nbytes)}" + ) + + def test_training_step_peak_memory_bounds_full_size_buffers(distributed_setup): """A training step should stay below five full-size child buffers.""" rank = distributed_setup.rank @@ -676,8 +726,8 @@ def test_deleted_model_releases_fsdp_storage(distributed_setup): assert torch.cuda.memory_allocated(device) - allocated_before < 1024**2 -def test_root_forward_returns_to_resting_memory(distributed_setup): - """Root forward should release child all-gather storage before returning.""" +def test_fully_shard_returns_to_resting_memory(distributed_setup): + """Fully-sharded temporary storage should be released after forward and backward.""" rank = distributed_setup.rank world_size = distributed_setup.world_size device = distributed_setup.device @@ -696,56 +746,32 @@ def test_root_forward_returns_to_resting_memory(distributed_setup): fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) x = torch.randn(2, dim, device=device, dtype=dtype) - resting_allocated = torch.cuda.memory_allocated(device) - - output = model(x) - del output - allocated_after_forward = torch.cuda.memory_allocated(device) - extra_allocated = allocated_after_forward - resting_allocated - child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() - - assert extra_allocated < child_weight_nbytes, ( - "Root forward did not return to resting memory after draining child releases: " - f"rank={rank}, extra_allocated={_mb(extra_allocated)}, " - f"one_child_weight={_mb(child_weight_nbytes)}" - ) - -def test_root_backward_returns_to_resting_memory(distributed_setup): - """Root backward should release child all-gather storage before returning.""" - rank = distributed_setup.rank - world_size = distributed_setup.world_size - device = distributed_setup.device - if world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(device.type, (world_size,)) - dim = 4096 - dtype = torch.bfloat16 - model = MultiChildModel(dim=dim, num_children=2).to(dtype=dtype, device=device) - placements = _flat_placements() - policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) - with fully_shard_context(device=device): - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) - fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + def clear_cublas_workspaces_and_get_allocated_memory() -> int: + # PyTorch retains a cuBLAS workspace for each handle/stream pair. Clear those + # library caches so this measurement isolates FSDP-managed storage. + torch._C._cuda_clearCublasWorkspaces() + return torch.cuda.memory_allocated(device) + + resting_allocated = clear_cublas_workspaces_and_get_allocated_memory() + + def assert_returns_to_resting_memory(phase: str) -> None: + extra_allocated = clear_cublas_workspaces_and_get_allocated_memory() - resting_allocated + # The live output, activations, and root-owned bias gradient are small; unsharded + # parameter storage must be released. + assert extra_allocated < 1024**2, ( + f"Fully-sharded storage did not return to resting memory after {phase}: " + f"rank={rank}, extra_allocated={_mb(extra_allocated)}, " + "max_extra_allocated=1.00 MB" + ) - x = torch.randn(2, dim, device=device, dtype=dtype, requires_grad=True) output = model(x) - loss = output.float().square().mean() - allocated_before_backward = torch.cuda.memory_allocated(device) + assert_returns_to_resting_memory("forward") + loss = output.float().square().mean() loss.backward() del loss, output - allocated_after_backward = torch.cuda.memory_allocated(device) - extra_allocated = allocated_after_backward - allocated_before_backward - child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() - - assert extra_allocated < child_weight_nbytes, ( - "Root backward did not return to resting memory after draining child releases: " - f"rank={rank}, extra_allocated={_mb(extra_allocated)}, " - f"one_child_weight={_mb(child_weight_nbytes)}" - ) + assert_returns_to_resting_memory("backward") @pytest.mark.parametrize( From 5d50b16e95c4634206acf14ba7021765f848b207 Mon Sep 17 00:00:00 2001 From: Ajay Date: Tue, 18 Aug 2026 04:00:26 +0000 Subject: [PATCH 070/101] fix(ci): avoid JET expansion of HOME in setup scripts (#6599) Signed-off-by: Ajay Balasa --- tests/test_utils/recipes/gb200/gpt-1node.yaml | 6 +++--- tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml | 6 +++--- tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml | 6 +++--- tests/test_utils/recipes/gb200/gpt-perf.yaml | 6 +++--- tests/test_utils/recipes/gb200/gpt.yaml | 6 +++--- tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml | 6 +++--- tests/test_utils/recipes/gb200/hybrid-perf.yaml | 6 +++--- tests/test_utils/recipes/gb200/moe-1node.yaml | 6 +++--- tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml | 6 +++--- tests/test_utils/recipes/gb200/moe.yaml | 6 +++--- tests/test_utils/recipes/gb200/nemotron-1node.yaml | 6 +++--- tests/test_utils/recipes/gb200/nemotron.yaml | 6 +++--- tests/test_utils/recipes/gb200/unit-tests.yaml | 6 +++--- tests/test_utils/recipes/h100/bert.yaml | 6 +++--- tests/test_utils/recipes/h100/ckpt_converter.yaml | 6 +++--- tests/test_utils/recipes/h100/determinism-perf.yaml | 6 +++--- tests/test_utils/recipes/h100/flextron.yaml | 6 +++--- .../recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml | 6 +++--- .../h100/gpt-dynamic-inference-with-coordinator.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-grads.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-grpo.yaml | 6 +++--- .../test_utils/recipes/h100/gpt-inference-server-smoke.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-nemo.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-offline-inference.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-perf-dp8.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-perf.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt-static-inference.yaml | 6 +++--- tests/test_utils/recipes/h100/gpt.yaml | 6 +++--- tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml | 6 +++--- tests/test_utils/recipes/h100/hybrid-perf.yaml | 6 +++--- .../h100/mamba-dynamic-inference-with-coordinator.yaml | 6 +++--- tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml | 6 +++--- tests/test_utils/recipes/h100/mamba-static-inference.yaml | 6 +++--- tests/test_utils/recipes/h100/mamba.yaml | 6 +++--- tests/test_utils/recipes/h100/mimo.yaml | 6 +++--- tests/test_utils/recipes/h100/module_performance.yaml | 6 +++--- .../h100/moe-dynamic-inference-with-coordinator.yaml | 6 +++--- tests/test_utils/recipes/h100/moe-dynamic-inference.yaml | 6 +++--- tests/test_utils/recipes/h100/moe-grpo.yaml | 6 +++--- tests/test_utils/recipes/h100/moe-static-inference.yaml | 6 +++--- tests/test_utils/recipes/h100/moe.yaml | 6 +++--- tests/test_utils/recipes/h100/multimodal-llava.yaml | 6 +++--- tests/test_utils/recipes/h100/t5.yaml | 6 +++--- tests/test_utils/recipes/h100/unit-tests.yaml | 6 +++--- 45 files changed, 135 insertions(+), 135 deletions(-) diff --git a/tests/test_utils/recipes/gb200/gpt-1node.yaml b/tests/test_utils/recipes/gb200/gpt-1node.yaml index 2465b68a6d6..45ce564e6c1 100644 --- a/tests/test_utils/recipes/gb200/gpt-1node.yaml +++ b/tests/test_utils/recipes/gb200/gpt-1node.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml b/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml index 81624c5c5ba..4630403b381 100644 --- a/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml +++ b/tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml b/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml index 50dfae839e0..6b789751b4b 100644 --- a/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml +++ b/tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml @@ -18,8 +18,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -29,7 +29,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/gpt-perf.yaml b/tests/test_utils/recipes/gb200/gpt-perf.yaml index 6dcdeb88d78..af0db7de0e7 100644 --- a/tests/test_utils/recipes/gb200/gpt-perf.yaml +++ b/tests/test_utils/recipes/gb200/gpt-perf.yaml @@ -20,8 +20,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -31,7 +31,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/gpt.yaml b/tests/test_utils/recipes/gb200/gpt.yaml index aaf8f4e6876..50af10e9688 100644 --- a/tests/test_utils/recipes/gb200/gpt.yaml +++ b/tests/test_utils/recipes/gb200/gpt.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml b/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml index ab4d2f1b04c..e46eda4d002 100644 --- a/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml +++ b/tests/test_utils/recipes/gb200/hybrid-perf-ep4.yaml @@ -18,8 +18,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -29,7 +29,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/hybrid-perf.yaml b/tests/test_utils/recipes/gb200/hybrid-perf.yaml index 2a25ef04032..a5abe5d986e 100644 --- a/tests/test_utils/recipes/gb200/hybrid-perf.yaml +++ b/tests/test_utils/recipes/gb200/hybrid-perf.yaml @@ -19,8 +19,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -30,7 +30,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/moe-1node.yaml b/tests/test_utils/recipes/gb200/moe-1node.yaml index fca78124824..8dc11c2f1b6 100644 --- a/tests/test_utils/recipes/gb200/moe-1node.yaml +++ b/tests/test_utils/recipes/gb200/moe-1node.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml b/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml index 39bd46086ee..bff130b611b 100644 --- a/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml +++ b/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/moe.yaml b/tests/test_utils/recipes/gb200/moe.yaml index ff6f155c3e1..71c2e28015f 100644 --- a/tests/test_utils/recipes/gb200/moe.yaml +++ b/tests/test_utils/recipes/gb200/moe.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/nemotron-1node.yaml b/tests/test_utils/recipes/gb200/nemotron-1node.yaml index b39bc6a4af2..d4178a7d054 100644 --- a/tests/test_utils/recipes/gb200/nemotron-1node.yaml +++ b/tests/test_utils/recipes/gb200/nemotron-1node.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/nemotron.yaml b/tests/test_utils/recipes/gb200/nemotron.yaml index f41a1e877f6..7f6dfed7d3e 100644 --- a/tests/test_utils/recipes/gb200/nemotron.yaml +++ b/tests/test_utils/recipes/gb200/nemotron.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -37,7 +37,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/gb200/unit-tests.yaml b/tests/test_utils/recipes/gb200/unit-tests.yaml index c972254463c..f12c7776ace 100644 --- a/tests/test_utils/recipes/gb200/unit-tests.yaml +++ b/tests/test_utils/recipes/gb200/unit-tests.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -32,7 +32,7 @@ spec: --backwards-commit $MCORE_BACKWARDS_COMMIT \ --repo $MCORE_REPO - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/bert.yaml b/tests/test_utils/recipes/h100/bert.yaml index d5e366ffacc..bc7e29f2894 100644 --- a/tests/test_utils/recipes/h100/bert.yaml +++ b/tests/test_utils/recipes/h100/bert.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -37,7 +37,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/ckpt_converter.yaml b/tests/test_utils/recipes/h100/ckpt_converter.yaml index bc726d9df0f..11c6a97fba6 100644 --- a/tests/test_utils/recipes/h100/ckpt_converter.yaml +++ b/tests/test_utils/recipes/h100/ckpt_converter.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/determinism-perf.yaml b/tests/test_utils/recipes/h100/determinism-perf.yaml index 2545fd973f8..230d4435add 100644 --- a/tests/test_utils/recipes/h100/determinism-perf.yaml +++ b/tests/test_utils/recipes/h100/determinism-perf.yaml @@ -12,8 +12,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -23,7 +23,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/flextron.yaml b/tests/test_utils/recipes/h100/flextron.yaml index 9864a566b32..45515e547ba 100644 --- a/tests/test_utils/recipes/h100/flextron.yaml +++ b/tests/test_utils/recipes/h100/flextron.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml b/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml index a9336dd036c..af82ce3aaa4 100644 --- a/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml +++ b/tests/test_utils/recipes/h100/gpt-dynamic-inference-cuda-graphs.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml index bbc02dac20c..e07e5f08235 100644 --- a/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml +++ b/tests/test_utils/recipes/h100/gpt-dynamic-inference-with-coordinator.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml b/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml index 6045b35c219..cf22d0095e9 100644 --- a/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml +++ b/tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-grads.yaml b/tests/test_utils/recipes/h100/gpt-grads.yaml index 5c7cf5a8f8b..184d86fc074 100644 --- a/tests/test_utils/recipes/h100/gpt-grads.yaml +++ b/tests/test_utils/recipes/h100/gpt-grads.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-grpo.yaml b/tests/test_utils/recipes/h100/gpt-grpo.yaml index 9e333cfebe3..56e59af4ffc 100644 --- a/tests/test_utils/recipes/h100/gpt-grpo.yaml +++ b/tests/test_utils/recipes/h100/gpt-grpo.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml b/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml index d1ed10b2f36..eceead5089a 100644 --- a/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml +++ b/tests/test_utils/recipes/h100/gpt-inference-server-smoke.yaml @@ -13,8 +13,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -34,7 +34,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-nemo.yaml b/tests/test_utils/recipes/h100/gpt-nemo.yaml index d9c56bd52fa..2d77b29dfe5 100644 --- a/tests/test_utils/recipes/h100/gpt-nemo.yaml +++ b/tests/test_utils/recipes/h100/gpt-nemo.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -37,7 +37,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-offline-inference.yaml b/tests/test_utils/recipes/h100/gpt-offline-inference.yaml index 9e1804c6e67..3e9e96f2e2b 100644 --- a/tests/test_utils/recipes/h100/gpt-offline-inference.yaml +++ b/tests/test_utils/recipes/h100/gpt-offline-inference.yaml @@ -13,8 +13,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -34,7 +34,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml b/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml index b1c6f43ca83..f91aff5eef5 100644 --- a/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml +++ b/tests/test_utils/recipes/h100/gpt-perf-dp8.yaml @@ -17,8 +17,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -28,7 +28,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-perf.yaml b/tests/test_utils/recipes/h100/gpt-perf.yaml index be4656addcf..a584adc527f 100644 --- a/tests/test_utils/recipes/h100/gpt-perf.yaml +++ b/tests/test_utils/recipes/h100/gpt-perf.yaml @@ -16,8 +16,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -27,7 +27,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt-static-inference.yaml b/tests/test_utils/recipes/h100/gpt-static-inference.yaml index 66771761421..c47a360c695 100644 --- a/tests/test_utils/recipes/h100/gpt-static-inference.yaml +++ b/tests/test_utils/recipes/h100/gpt-static-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/gpt.yaml b/tests/test_utils/recipes/h100/gpt.yaml index c2a1f17eea3..2b4c0849b04 100644 --- a/tests/test_utils/recipes/h100/gpt.yaml +++ b/tests/test_utils/recipes/h100/gpt.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml b/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml index 2b64f923e3e..b0ae34e0d07 100644 --- a/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml +++ b/tests/test_utils/recipes/h100/hybrid-perf-ep8.yaml @@ -16,8 +16,8 @@ spec: script_setup: | unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -27,7 +27,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/hybrid-perf.yaml b/tests/test_utils/recipes/h100/hybrid-perf.yaml index 3e37ec95d7c..f1caebaacfa 100644 --- a/tests/test_utils/recipes/h100/hybrid-perf.yaml +++ b/tests/test_utils/recipes/h100/hybrid-perf.yaml @@ -15,8 +15,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT cd /opt rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm @@ -26,7 +26,7 @@ spec: git fetch origin $MCORE_MR_COMMIT git checkout $MCORE_MR_COMMIT git rev-parse HEAD - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml index 53e912131ee..b2892878f94 100644 --- a/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml +++ b/tests/test_utils/recipes/h100/mamba-dynamic-inference-with-coordinator.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml index f5e6b7f8d8c..3a1ad18c282 100644 --- a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml +++ b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/mamba-static-inference.yaml b/tests/test_utils/recipes/h100/mamba-static-inference.yaml index 29cfb2e0a67..4b2df8c2d41 100644 --- a/tests/test_utils/recipes/h100/mamba-static-inference.yaml +++ b/tests/test_utils/recipes/h100/mamba-static-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index 0a338d91b61..b4596e2d71a 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/mimo.yaml b/tests/test_utils/recipes/h100/mimo.yaml index 4aea76c6b19..b76cc643647 100644 --- a/tests/test_utils/recipes/h100/mimo.yaml +++ b/tests/test_utils/recipes/h100/mimo.yaml @@ -20,8 +20,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -42,7 +42,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/module_performance.yaml b/tests/test_utils/recipes/h100/module_performance.yaml index 19d58d03ef8..6af791ff3e9 100644 --- a/tests/test_utils/recipes/h100/module_performance.yaml +++ b/tests/test_utils/recipes/h100/module_performance.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml index b51e5188cdb..1c9eb18104c 100644 --- a/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml +++ b/tests/test_utils/recipes/h100/moe-dynamic-inference-with-coordinator.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml b/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml index 933d1d25551..12e2dcadf60 100644 --- a/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml +++ b/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/moe-grpo.yaml b/tests/test_utils/recipes/h100/moe-grpo.yaml index 3793815b8d1..04f3a482ae0 100644 --- a/tests/test_utils/recipes/h100/moe-grpo.yaml +++ b/tests/test_utils/recipes/h100/moe-grpo.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/moe-static-inference.yaml b/tests/test_utils/recipes/h100/moe-static-inference.yaml index a7f63591ad7..cc0776f0910 100644 --- a/tests/test_utils/recipes/h100/moe-static-inference.yaml +++ b/tests/test_utils/recipes/h100/moe-static-inference.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/moe.yaml b/tests/test_utils/recipes/h100/moe.yaml index 7620eed364b..599b9b5237f 100644 --- a/tests/test_utils/recipes/h100/moe.yaml +++ b/tests/test_utils/recipes/h100/moe.yaml @@ -14,8 +14,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -36,7 +36,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/multimodal-llava.yaml b/tests/test_utils/recipes/h100/multimodal-llava.yaml index f37a81d6648..dfd888e9797 100644 --- a/tests/test_utils/recipes/h100/multimodal-llava.yaml +++ b/tests/test_utils/recipes/h100/multimodal-llava.yaml @@ -17,8 +17,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -39,7 +39,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/t5.yaml b/tests/test_utils/recipes/h100/t5.yaml index 7fd369b2070..6e5dbe953b0 100644 --- a/tests/test_utils/recipes/h100/t5.yaml +++ b/tests/test_utils/recipes/h100/t5.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -35,7 +35,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" diff --git a/tests/test_utils/recipes/h100/unit-tests.yaml b/tests/test_utils/recipes/h100/unit-tests.yaml index c37edef8dd3..2e730853043 100644 --- a/tests/test_utils/recipes/h100/unit-tests.yaml +++ b/tests/test_utils/recipes/h100/unit-tests.yaml @@ -13,8 +13,8 @@ spec: set -euo pipefail unset https_proxy umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "${HOME}/.netrc" - trap 'rm -f "${HOME}/.netrc"; unset RO_API_TOKEN' EXIT + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -32,7 +32,7 @@ spec: --backwards-commit $MCORE_BACKWARDS_COMMIT \ --repo $MCORE_REPO - rm -f "${HOME}/.netrc" + rm -f "$HOME/.netrc" unset RO_API_TOKEN script: |- rm -f "${{HOME}}/.netrc" From 32756e3683b242b73114a41ed83658a9eab2e5cc Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 18 Aug 2026 04:55:55 +0000 Subject: [PATCH 071/101] RL: isolate placeholder rollouts from GRPO rewards (#6475) Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Jorge Albericio --- megatron/rl/agent/api.py | 7 ++ megatron/rl/rl_utils.py | 102 ++++++++++++++--- tests/unit_tests/rl/test_rl_utils.py | 159 ++++++++++++++++++++++++--- 3 files changed, 239 insertions(+), 29 deletions(-) diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index 310869d142e..cad2ab5739e 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -40,6 +40,9 @@ class GroupedRolloutRequest(Request): consumption_granularity: ConsumptionGranularity = "B" +KNOWN_ROLLOUT_STATUSES = ('ok', 'placeholder', 'masked', 'graded') + + class Rollout(AgentBaseModel): """Data for language-based Rollout.""" @@ -48,6 +51,8 @@ class Rollout(AgentBaseModel): reward: float = None env_id: str = '' problem_id: str | None = None + rollout_status: str = 'ok' + failure_reason: str | None = None class TokenRollout(AgentBaseModel): @@ -60,6 +65,8 @@ class TokenRollout(AgentBaseModel): env_id: str = '' problem_id: str | None = None completion_ids: list[str] = [] + rollout_status: str = 'ok' + failure_reason: str | None = None Rollouts = list[TokenRollout | Rollout] diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 5e0a66b3f0d..08e5b7b579e 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -73,6 +73,7 @@ unwrap_model, ) from megatron.rl.agent.api import ( + KNOWN_ROLLOUT_STATUSES, EvaluationRequest, EvaluationResponse, GroupedRolloutRequest, @@ -308,6 +309,8 @@ class RolloutStats: kv_cache_epoch: list[list[int]] completed_epochs: list[list[int]] num_evictions: list[list[int]] + rollout_statuses: list[list[str]] + failure_reasons: list[list[str | None]] # Runtime state container for RL-specific data that shouldn't be checkpointed @@ -1069,15 +1072,25 @@ def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[i # get [a,b] with 1.0 and [c,d,e] with 1.0 when doing updates. """ - rewards = np.array(rewards) + rewards = np.asarray(rewards, dtype=np.float64) + num_turns = np.asarray(num_turns, dtype=np.int64) + assert rewards.shape == num_turns.shape, ( + f"rewards and num_turns must have matching shape, got {rewards.shape} and {num_turns.shape}" + ) - num_turns = np.array(num_turns) # Each outer dimension of num_turns is a group. Sum of those gives total num_turns per group. # Let's use this to calculate advantage. # mean/std should be repeated based on group lens group_turns = num_turns.sum(axis=-1) - reward_means = rewards.mean(axis=1, keepdims=True).repeat(group_turns) - reward_stds = rewards.std(axis=1, keepdims=True).repeat(group_turns) + # Zero-turn rollouts are empty-trajectory placeholders for failed episodes and are masked out. + real_mask = num_turns > 0 + # Clamping to 1 avoids 0/0 for all-placeholder groups. + real_counts = np.maximum(real_mask.sum(axis=-1), 1) + group_means = np.where(real_mask, rewards, 0.0).sum(axis=-1) / real_counts + group_deviations = np.where(real_mask, rewards - group_means[:, np.newaxis], 0.0) + group_stds = np.sqrt((group_deviations**2).sum(axis=-1) / real_counts) + reward_means = group_means.repeat(group_turns) + reward_stds = group_stds.repeat(group_turns) # rewards are originally [g, group_size] # Making an assumption that all groups are of the same size! @@ -1117,6 +1130,8 @@ def compute_group_stats( all_kv_cache_epoch = [] all_completed_epochs = [] all_num_evictions = [] + all_rollout_statuses = [] + all_failure_reasons = [] for group in rollouts: group_rewards = [] group_traj_lengths = [] @@ -1126,6 +1141,8 @@ def compute_group_stats( group_kv_epoch = [] group_completed_epochs = [] group_num_evictions = [] + group_rollout_statuses = [] + group_failure_reasons = [] for rollout in group: if isinstance(rollout, TokenRollout): for turn_traj in rollout.trajectory: @@ -1177,10 +1194,18 @@ def compute_group_stats( ) group_completed_epochs.extend(record.policy_epoch[-1][1] for record in turn_records) group_num_evictions.append(sum(record.num_evictions for record in turn_records)) + rollout_status = rollout.rollout_status + # All empty-trajectory rollouts are placeholders. + if not rollout.trajectory and rollout_status == 'ok': + rollout_status = 'placeholder' + group_rollout_statuses.append(rollout_status) + group_failure_reasons.append(rollout.failure_reason) all_policy_epoch.append(group_policy_epoch) all_kv_cache_epoch.append(group_kv_epoch) all_completed_epochs.append(group_completed_epochs) all_num_evictions.append(group_num_evictions) + all_rollout_statuses.append(group_rollout_statuses) + all_failure_reasons.append(group_failure_reasons) traj_lens.append(group_traj_lengths) turn_lens.append(group_turn_lengths) env_ids.append(group[0].env_id) # All rollouts in a group share the env_id by design. @@ -1212,6 +1237,8 @@ def compute_group_stats( kv_cache_epoch=all_kv_cache_epoch, completed_epochs=all_completed_epochs, num_evictions=all_num_evictions, + rollout_statuses=all_rollout_statuses, + failure_reasons=all_failure_reasons, ) return stats @@ -1251,16 +1278,16 @@ def prep_wandb_metrics( completed_epochs: List[List[int]], num_evictions: List[List[int]], current_iteration: int, + rollout_statuses: List[List[str]] | None = None, + failure_reasons: List[List[str | None]] | None = None, example_group: list[TokenRollout | Rollout] | None = None, tokenizer: MegatronTokenizer | None = None, ): """Make a wandb-parseable dictionary of metrics for logging. - Zero-turn rollouts are a mark of placeholders (empty-trajectory pads for failed episodes). - Their 0.0 reward deliberately stays in the reward and group-mean/std aggregates: - it does affect training dynamics. - All other per-rollout field of theirs is masked out of stats. + Zero-turn rollouts are empty-trajectory placeholders for failed episodes. + Their rewards are excluded from GRPO group normalization; the stats reflect this. Args: wandb_writer: Wandb run to log to. @@ -1274,6 +1301,8 @@ def prep_wandb_metrics( completed_epochs: Grouped list of per-turn max policy epoch stamps. num_evictions: Grouped list of per-rollout number of evictions. current_iteration: Current training iteration. + rollout_statuses: Grouped adapter-stamped statuses; None means all 'ok'. + failure_reasons: Grouped failure-cause labels; None means all None. example_group: A list of rollouts of one group to log examples of trajectories. tokenizer: Tokenizer to untokenize trajectories for logging. """ @@ -1281,13 +1310,39 @@ def prep_wandb_metrics( real_mask = [[nt > 0 for nt in g] for g in num_turns] total_rollouts = sum(len(g) for g in num_turns) failed_rollouts = sum(not keep for g in real_mask for keep in g) - failure_metrics = { + rollout_metrics = { + 'rollout/count': total_rollouts, 'failed_rollouts/count': failed_rollouts, 'failed_rollouts/ratio': ( failed_rollouts / total_rollouts if total_rollouts else 0.0 ), } + def _safe_key(label): + """Sanitize a free-form label for use inside a wandb metric key.""" + return ''.join(c if c.isalnum() or c in ('_', '-') else '_' for c in label) + + if rollout_statuses is None: + rollout_statuses = [['ok'] * len(g) for g in rewards] + if failure_reasons is None: + failure_reasons = [[None] * len(g) for g in rewards] + status_counts = Counter(s for g in rollout_statuses for s in g) + status_counts.update({status: 0 for status in KNOWN_ROLLOUT_STATUSES}) + for status, count in sorted(status_counts.items()): + safe_status = _safe_key(status) + rollout_metrics[f'rollout/{safe_status}_count'] = count + rollout_metrics[f'rollout/{safe_status}_rate'] = ( + count / total_rollouts if total_rollouts else 0.0 + ) + reason_counts = Counter(r for g in failure_reasons for r in g if r) + rollout_metrics['failure_reason_table'] = wandb_writer.Table( + columns=['failure_reason', 'count', 'rate'], + data=[ + [reason, count, count / total_rollouts if total_rollouts else 0.0] + for reason, count in sorted(reason_counts.items()) + ], + ) + # All-empty wave: every episode failed and was dropped (nothing to aggregate), # or every surviving rollout is a zero-turn placeholder (rewards exist but the # length/turn stats below would reduce over empty sequences). Keep the failure @@ -1301,7 +1356,7 @@ def prep_wandb_metrics( "prep_wandb_metrics: empty wave (0 usable rollouts); " "skipping rollout metrics this iteration." ) - return failure_metrics + return rollout_metrics def _real(grouped): """Grouped per-rollout entries with placeholder (zero-turn) rollouts removed.""" @@ -1309,26 +1364,31 @@ def _real(grouped): [x for x, keep in zip(g, m) if keep] for g, m in zip(grouped, real_mask) ] - # Reward metrics include failures. All other metrics do not. Staleness and - # eviction telemetry filters through `joined` instead: placeholders pop no - # ledger records, so they have no epoch rows and never join. + # Raw delivery aggregates (mean_reward, rewards_hist) keep failures. All other metrics do + # not. Staleness and eviction telemetry filters through `joined` instead: placeholders pop + # no ledger records, so they have no epoch rows and never join. + rewards_real = _real(rewards) + table_rewards = [r for g in rewards_real for r in g] traj_lens_real = _real(traj_lens) num_turns_real = _real(num_turns) + # Mirror calculate_grpo_advantages: group means/stds over real rollouts only. group_table = wandb_writer.Table( columns=['group_means', 'group_stds'], - data=[[np.mean(g), np.std(g)] for g in rewards], + data=[[np.mean(g) if g else 0.0, np.std(g) if g else 0.0] for g in rewards_real], ) flat_rewards = [r for g in rewards for r in g] flat_traj_lens = [l for g in traj_lens for l in g] flat_num_evictions = [e for g in num_evictions for e in g] + flat_statuses = [s for g in rollout_statuses for s in g] flat_policy_epochs = [r for g in policy_epoch for r in g] flat_kv_epochs = [r for g in kv_cache_epoch for r in g] joined = [i for i, row in enumerate(flat_policy_epochs) if row] joined_rewards = [flat_rewards[i] for i in joined] joined_traj_lens = [flat_traj_lens[i] for i in joined] joined_num_evictions = [flat_num_evictions[i] for i in joined] + joined_statuses = [flat_statuses[i] for i in joined] per_rollout_policy_epochs = [flat_policy_epochs[i] for i in joined] per_rollout_kv_epochs = [flat_kv_epochs[i] for i in joined] # Per-rollout staleness (oldest token) @@ -1373,10 +1433,11 @@ def _real(grouped): 'max_num_turns': max(max(g) for g in num_turns_real if g), 'min_num_turns': min(min(g) for g in num_turns_real if g), 'mean_reward': np.mean([np.mean(g) for g in rewards]), + 'valid_mean_reward': np.mean(table_rewards) if table_rewards else 0.0, 'mean_advantage': np.mean(advantages), 'nonzero_groups_ratio': np.count_nonzero(advantages) / len(advantages), - **failure_metrics, + **rollout_metrics, } # Staleness/eviction telemetry: joined rollouts only; skip when nothing joined. @@ -1387,6 +1448,7 @@ def _real(grouped): 'reward', 'traj_length', 'num_evictions', 'policy_staleness', 'kv_staleness', 'policy_last_token_staleness', 'kv_last_token_staleness', + 'rollout_status', ], data=list(zip( joined_rewards, @@ -1396,6 +1458,7 @@ def _real(grouped): rollout_kv_staleness, rollout_policy_last_token_staleness, rollout_kv_last_token_staleness, + joined_statuses, )), ), # NOTE: This table can get very large (one row per token across all rollouts). @@ -1429,6 +1492,7 @@ def _real(grouped): } + if example_group: if tokenizer is None: raise ValueError("If you provide an example group to log, you need to provide a tokenizer too.") @@ -1602,11 +1666,15 @@ def maybe_log_training_metrics( kv_cache_epoch = group_stats.kv_cache_epoch completed_epochs = group_stats.completed_epochs num_evictions = group_stats.num_evictions + rollout_statuses = group_stats.rollout_statuses + failure_reasons = group_stats.failure_reasons metrics = metrics | prep_wandb_metrics(wandb_writer=wandb_writer, traj_lens=traj_lens, turn_lens=turn_lens, rewards=rewards, num_turns=num_turns, advantages=advantages, policy_epoch=policy_epoch, kv_cache_epoch=kv_cache_epoch, completed_epochs=completed_epochs, - num_evictions=num_evictions, current_iteration=current_iteration) + num_evictions=num_evictions, current_iteration=current_iteration, + rollout_statuses=rollout_statuses, failure_reasons=failure_reasons) + env_stats = lambda cont, idx: [cont[i] for i in idx] group_turn_counts = [sum(nt) for nt in num_turns] @@ -1631,6 +1699,8 @@ def maybe_log_training_metrics( completed_epochs=env_stats(completed_epochs, env_idx), num_evictions=env_stats(num_evictions, env_idx), current_iteration=current_iteration, + rollout_statuses=env_stats(rollout_statuses, env_idx), + failure_reasons=env_stats(failure_reasons, env_idx), example_group=example_groups[env_id], tokenizer=tokenizer, ) diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 51281bdd3d6..52567a99901 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -952,6 +952,22 @@ def test_bounded_artifact_key(self): assert bounded == rl_utils._bounded_artifact_key(long_key) assert bounded != rl_utils._bounded_artifact_key(long_key + "y") + def test_placeholder_rewards_do_not_affect_advantage_normalization(self): + # Group of 3 with one zero-turn placeholder: the two real rollouts + # normalize against each other, so the placeholder's synthetic 0.0 + # reward manufactures no advantage. + advs = rl_utils.calculate_grpo_advantages([[0, 1, 0]], [[0, 1, 1]]) + torch.testing.assert_close( + torch.tensor(advs), torch.tensor([1.0, -1.0]), atol=3e-4, rtol=1e-5 + ) + + def test_all_placeholder_group_has_no_advantages(self): + assert rl_utils.calculate_grpo_advantages([[0, 0]], [[0, 0]]) == [] + + def test_advantage_shape_mismatch_raises(self): + with pytest.raises(AssertionError, match="matching shape"): + rl_utils.calculate_grpo_advantages([[1.0, 0.0]], [[1, 1, 1]]) + @pytest.mark.parametrize( "scenario, expected_turn_lens, expected_traj_lens, expected_num_turns", [ @@ -1404,6 +1420,9 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders # Per-turn max epoch stamps (when each turn completed) completed_epochs = [[5, 3], [5, 1]] num_evictions = [[0, 1], [0, 0]] + # "env/timeout" sits outside KNOWN_ROLLOUT_STATUSES and needs key sanitization. + rollout_statuses = [["graded", "ok"], ["env/timeout", "ok"]] + failure_reasons = [[None, None], [None, None]] current_iteration = 6 if inject_placeholders: for lst, sentinel in ( @@ -1420,6 +1439,12 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders # Placeholders pop no ledger records, so their epoch rows are empty. group.append([]) lst.append([[], []]) + for group in rollout_statuses: + group.append("placeholder") + rollout_statuses.append(["placeholder", "placeholder"]) + for group in failure_reasons: + group.append("http_500") + failure_reasons.append(["http_500", "http_500"]) # Placeholders contribute no turns, and compute_group_stats already # excludes them from completed_epochs; the failed group adds empty # inner lists, which the group-level stats must skip, not crash on. @@ -1439,12 +1464,51 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders completed_epochs=completed_epochs, num_evictions=num_evictions, current_iteration=current_iteration, + rollout_statuses=rollout_statuses, + failure_reasons=failure_reasons, ) assert metrics["failed_rollouts/count"] == (4 if inject_placeholders else 0) assert metrics["failed_rollouts/ratio"] == (0.5 if inject_placeholders else 0.0) - # Reward aggregates keep the placeholder zeros by design: group means - # become [2/3, 1/3, 0] instead of [1, 0.5]. + assert metrics["rollout/count"] == (8 if inject_placeholders else 4) + assert metrics["rollout/placeholder_count"] == (4 if inject_placeholders else 0) + assert metrics["rollout/placeholder_rate"] == (0.5 if inject_placeholders else 0.0) + # Statuses in KNOWN_ROLLOUT_STATUSES always emit a series, zeros included... + assert metrics["rollout/ok_count"] == 2 + assert metrics["rollout/ok_rate"] == (0.25 if inject_placeholders else 0.5) + assert metrics["rollout/masked_count"] == 0 + assert metrics["rollout/masked_rate"] == 0.0 + assert metrics["rollout/graded_count"] == 1 + # ...and statuses beyond the vocabulary are still counted, under a + # wandb-safe key. + assert metrics["rollout/env_timeout_count"] == 1 + assert metrics["rollout/env_timeout_rate"] == (0.125 if inject_placeholders else 0.25) + assert not any(k.startswith("rollout/failure_reason/") for k in metrics) + assert "failure_reason_table" in metrics + reason_table_calls = [ + c + for c in writer.Table.call_args_list + if c.kwargs.get("columns") == ["failure_reason", "count", "rate"] + ] + assert len(reason_table_calls) == 1 + expected_reason_rows = [["http_500", 4, 0.5]] if inject_placeholders else [] + assert reason_table_calls[0].kwargs["data"] == expected_reason_rows + # The raw delivery aggregate keeps the placeholder zeros (mean of group + # means becomes 1/3 instead of 0.75); the valid view masks them. The + # pair diverging is the infrastructure-failure signature. assert np.isclose(metrics["mean_reward"], 1 / 3 if inject_placeholders else 0.75) + assert metrics["valid_mean_reward"] == 0.75 + # Group means/stds mirror GRPO normalization: real rollouts only, 0.0 + # rows for all-placeholder groups. + group_table_calls = [ + c + for c in writer.Table.call_args_list + if c.kwargs.get("columns") == ["group_means", "group_stds"] + ] + assert len(group_table_calls) == 1 + expected_group_rows = [[1.0, 0.0], [0.5, 1.5]] + if inject_placeholders: + expected_group_rows.append([0.0, 0.0]) + assert group_table_calls[0].kwargs["data"] == expected_group_rows assert metrics["mean_advantage"] == 0.5 # The rollout table lists real rollouts only, in either case. rollout_table_calls = [ @@ -1488,8 +1552,8 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders # mean_completion_gap = mean([6-5, 6-3, 6-5, 6-1]) = mean([1, 3, 1, 5]) = 2.5 assert metrics["mean_completion_gap"] == 2.5 - # Case where every episode has failed: the wave still reports the - # failure counters (the one signal it carries) and skips the tables. + # Case where every episode has failed: the wave still reports the failure + # counters (now status series + reason table) and skips the rollout tables. writer = MagicMock() metrics = rl_utils.prep_wandb_metrics( writer, @@ -1504,8 +1568,14 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders num_evictions=[], current_iteration=6, ) - assert metrics == {'failed_rollouts/count': 0, 'failed_rollouts/ratio': 0.0} - writer.Table.assert_not_called() + assert metrics['rollout/count'] == 0 + assert metrics['failed_rollouts/count'] == 0 + assert metrics['failed_rollouts/ratio'] == 0.0 + # Unstamped wave: the known-status series still emit, zeroed. + assert metrics['rollout/placeholder_count'] == 0 + assert metrics['rollout/placeholder_rate'] == 0.0 + assert 'rollout_table' not in metrics and 'mean_reward' not in metrics + assert [c.kwargs['data'] for c in writer.Table.call_args_list] == [[]] # All-placeholder wave: rewards exist (placeholders carry 0.0) but every # rollout is a zero-turn pad, so the masked stats have nothing to reduce @@ -1523,9 +1593,22 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders completed_epochs=[[]], num_evictions=[[0, 0]], current_iteration=6, + rollout_statuses=[["placeholder", "placeholder"]], + failure_reasons=[["http_500", None]], ) - assert metrics == {'failed_rollouts/count': 2, 'failed_rollouts/ratio': 1.0} - writer.Table.assert_not_called() + assert metrics['rollout/count'] == 2 + assert metrics['failed_rollouts/count'] == 2 + assert metrics['failed_rollouts/ratio'] == 1.0 + assert metrics['rollout/placeholder_count'] == 2 + assert metrics['rollout/placeholder_rate'] == 1.0 + assert 'rollout_table' not in metrics and 'mean_reward' not in metrics + reason_calls = [ + c + for c in writer.Table.call_args_list + if c.kwargs.get("columns") == ["failure_reason", "count", "rate"] + ] + assert len(reason_calls) == 1 + assert reason_calls[0].kwargs["data"] == [["http_500", 1, 0.5]] def test_compute_group_stats_excludes_placeholders_from_metric_fields(self): def real_rollout(tokens, problem_id): @@ -1539,7 +1622,7 @@ def real_rollout(tokens, problem_id): completion_ids=_mint_cids(1), ) - def placeholder(): + def placeholder(rollout_status="ok", failure_reason=None): return TokenRollout( trajectory=[], generation_mask=[], @@ -1547,12 +1630,21 @@ def placeholder(): logprobs=[], env_id="swe", problem_id="placeholder", + rollout_status=rollout_status, + failure_reason=failure_reason, ) eod = MockTokenizer().eod first = real_rollout([1, 2, eod], problem_id="p0") second = real_rollout([1, 2, 3, eod], problem_id="p0") - rollouts = [[first, second, placeholder()], [placeholder(), placeholder(), placeholder()]] + rollouts = [ + [first, second, placeholder()], + [ + placeholder(rollout_status="masked", failure_reason="judge_unparseable"), + placeholder(failure_reason="http_500"), + placeholder(), + ], + ] # Placeholders have no turns (and no completion ids): they pop no records. ledger = { first.completion_ids[0]: _ledger_record(5), @@ -1572,12 +1664,26 @@ def placeholder(): # in completed_epochs, no fake 0-length turn for all-placeholder groups. assert stats.completed_epochs == [[5, 6], []] assert stats.turn_lens == [[3, 4], []] - # Rewards keep the placeholder zeros (they shape the GRPO baseline). + # Rewards keep the placeholder zeros: the lists stay rectangular for + # alignment; GRPO normalization masks them by num_turns downstream. assert stats.rewards == [[1.0, 1.0, 0.0], [0.0, 0.0, 0.0]] + # Which the advantages reflect: both real rollouts earned 1.0, so with + # the placeholder masked there is no contrast left in the group. + assert stats.advantages == [0.0, 0.0] + # Adapter stamps pass through; bare empties are inferred to 'placeholder'. + assert stats.rollout_statuses == [ + ["ok", "ok", "placeholder"], + ["masked", "placeholder", "placeholder"], + ] + assert stats.failure_reasons == [ + [None, None, None], + ["judge_unparseable", "http_500", None], + ] # End to end: the placeholder rows never reach the staleness metrics. + writer = MagicMock() metrics = rl_utils.prep_wandb_metrics( - MagicMock(), + writer, stats.traj_lens, stats.turn_lens, stats.rewards, @@ -1588,6 +1694,8 @@ def placeholder(): completed_epochs=stats.completed_epochs, num_evictions=stats.num_evictions, current_iteration=7, + rollout_statuses=stats.rollout_statuses, + failure_reasons=stats.failure_reasons, ) assert metrics["max_policy_staleness"] == 2 # 7 - 5, not a placeholder sentinel assert metrics["min_traj_length"] == 3 @@ -1595,6 +1703,31 @@ def placeholder(): assert metrics["mean_completion_gap"] == np.mean([2, 1]) assert metrics["failed_rollouts/count"] == 4 assert np.isclose(metrics["failed_rollouts/ratio"], 4 / 6) + # GRPO-facing reward views mask the placeholders; raw delivery keeps them. + assert metrics["valid_mean_reward"] == 1.0 + assert np.isclose(metrics["mean_reward"], 1 / 3) + group_rows = [ + c + for c in writer.Table.call_args_list + if c.kwargs.get("columns") == ["group_means", "group_stds"] + ][0].kwargs["data"] + assert group_rows == [[1.0, 0.0], [0.0, 0.0]] + # The stamped taxonomy decomposes the structural zero-turn count. + assert metrics["rollout/count"] == 6 + assert metrics["rollout/ok_count"] == 2 + assert metrics["rollout/placeholder_count"] == 3 + assert metrics["rollout/masked_count"] == 1 + assert metrics["rollout/graded_count"] == 0 + assert ( + metrics["rollout/placeholder_count"] + metrics["rollout/masked_count"] + == metrics["failed_rollouts/count"] + ) + reason_rows = [ + c + for c in writer.Table.call_args_list + if c.kwargs.get("columns") == ["failure_reason", "count", "rate"] + ][0].kwargs["data"] + assert reason_rows == [["http_500", 1, 1 / 6], ["judge_unparseable", 1, 1 / 6]] def test_request_ledger_join(self): """compute_group_stats joins finished-request records to rollouts by each @@ -1680,7 +1813,7 @@ def token_rollout(turns, masks, completion_ids=None): rollout_table = next( c for c in writer.Table.call_args_list if "traj_length" in c.kwargs.get("columns", []) ) - assert rollout_table.kwargs["data"] == [(1.0, 3, 2, 2, 2, 1, 1)] + assert rollout_table.kwargs["data"] == [(1.0, 3, 2, 2, 2, 1, 1, 'ok')] # A window where nothing joined emits no staleness/eviction metrics at all. metrics = rl_utils.prep_wandb_metrics( From ae5ec14dfcbe547a08816af1a9150ee15a583e03 Mon Sep 17 00:00:00 2001 From: Ajay Date: Tue, 18 Aug 2026 05:20:48 +0000 Subject: [PATCH 072/101] tests: Update cerno.yml with linear issue status (#6304) Signed-off-by: Ajay Balasa --- .gitlab/cerno.yml | 1 + tests/test_utils/test_ci_triage.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/.gitlab/cerno.yml b/.gitlab/cerno.yml index 6f723daa166..e00985e7d25 100644 --- a/.gitlab/cerno.yml +++ b/.gitlab/cerno.yml @@ -11,6 +11,7 @@ modules: reconcile_proposal: true team_key: MCORE project_template: "MCore CI Testing" + linear_issue_status: "Need Triage" enable_linear_open: true enable_linear_modify: true enable_linear_close: true diff --git a/tests/test_utils/test_ci_triage.py b/tests/test_utils/test_ci_triage.py index 3c3b3b7ed08..18b2d2a877e 100644 --- a/tests/test_utils/test_ci_triage.py +++ b/tests/test_utils/test_ci_triage.py @@ -121,11 +121,14 @@ def test_error_extraction_is_opt_in_for_generated_jobs( def test_cerno_hard_cutover_contract(): pipeline = yaml.safe_load(Path(".gitlab-ci.yml").read_text()) + config = yaml.safe_load(Path(".gitlab/cerno.yml").read_text()) triage = Path(".gitlab/stages/06.triage.yml").read_text() dockerfile = Path("docker/Dockerfile.linting").read_text() build_script = Path(".gitlab/scripts/build.sh").read_text() assert pipeline["variables"]["CERNO_CONFIG"] == ".gitlab/cerno.yml" + assert config["report_artifacts"] == {} + assert config["error_extraction"]["resolution_configs"] == {} assert triage.count("cerno-linear") == 3 assert triage.count("cerno-notify") == 2 assert triage.count('--config "${CERNO_CONFIG}"') == 3 @@ -404,6 +407,7 @@ def test_triage_config_selects_megatron_and_enables_write_actions(): "reconcile_proposal": True, "team_key": "MCORE", "project_template": "MCore CI Testing", + "linear_issue_status": "Need Triage", "enable_linear_open": True, "enable_linear_modify": True, "enable_linear_close": True, @@ -413,6 +417,7 @@ def test_triage_config_selects_megatron_and_enables_write_actions(): assert linear_write.write_gates(config) == { linear_ci.LINEAR_MODULE: {"open": True, "modify": True, "close": True} } + assert linear_write.issue_statuses(config) == {linear_ci.LINEAR_MODULE: "Need Triage"} def test_slack_followup_uses_upstream_detailed_and_execution_summaries(): From cfb8e46ffe08fa027c38ffbe3898e7d8937bd019 Mon Sep 17 00:00:00 2001 From: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:42:51 +0000 Subject: [PATCH 073/101] Enable GTP for heterogeneous MIMO training (#6570) Signed-off-by: ykarnati --- .../mimo/model_providers/nemotron_moe_vlm.py | 19 ++++++ .../mimo/model_providers/radio_encoder.py | 10 ++++ examples/mimo/pretrain_mimo.py | 22 +++---- .../run_hetero_nemotron_20l_mock_train.sh | 8 ++- examples/mimo/training/args.py | 34 +++++++++-- examples/mimo/training/data.py | 13 ++-- examples/mimo/training/grad_sync.py | 13 ++-- examples/mimo/training/runtime.py | 4 +- examples/mimo/training/topology.py | 60 ++++++++++++++----- megatron/core/models/mimo/model/base.py | 2 + megatron/core/models/mimo/optimizer.py | 29 ++------- .../models/mimo/test_mimo_hetero_grid_args.py | 27 +++++++++ .../models/mimo/test_mimo_mock_data.py | 9 ++- tests/unit_tests/test_mimo_hetero_topology.py | 40 +++++++++++++ 14 files changed, 218 insertions(+), 72 deletions(-) diff --git a/examples/mimo/model_providers/nemotron_moe_vlm.py b/examples/mimo/model_providers/nemotron_moe_vlm.py index 133bf9bc1d2..0979cd4e791 100644 --- a/examples/mimo/model_providers/nemotron_moe_vlm.py +++ b/examples/mimo/model_providers/nemotron_moe_vlm.py @@ -12,6 +12,7 @@ from examples.mimo.model_providers.radio_encoder import ( RADIO_ENCODER_MODULE_NAME, _base_config, + _disable_gtp, _make_dense_non_hybrid, add_radio_encoder_args, radio_vision_config, @@ -21,6 +22,7 @@ from megatron.core.activations import squared_relu from megatron.core.hyper_comm_grid import HyperCommGrid from megatron.core.hyper_comm_grid import _is_process_group_member as is_process_group_member +from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec from megatron.core.models.mamba.mamba_model import MambaModel from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY @@ -105,6 +107,22 @@ def nemotron_language_config( config.expert_tensor_parallel_size = expt_tp_size config.tensor_model_parallel_size = tp_size config.pipeline_model_parallel_size = pp_size + config.tensor_parallel_num_weight_shards, config.gtp_weight_remat_size = ( + resolve_tensor_parallel_weight_shards( + tp_size, + getattr(args, "tensor_parallel_num_weight_shards", None), + getattr(args, "gtp_weight_remat_size", 1), + ) + ) + config.expert_tensor_parallel_num_weight_shards, config.expert_gtp_weight_remat_size = ( + resolve_tensor_parallel_weight_shards( + expt_tp_size, + getattr(args, "expert_tensor_parallel_num_weight_shards", None), + getattr(args, "expert_gtp_weight_remat_size", 1), + shards_field="expert_tensor_parallel_num_weight_shards", + tp_field="expert_tensor_parallel_size", + ) + ) config.sequence_parallel = tp_size > 1 config.position_embedding_type = "none" return config @@ -142,6 +160,7 @@ def nemotron_projection_config( config.normalization = "RMSNorm" _make_dense_non_hybrid(config) # Projection inherits no MoE/Mamba/hybrid settings. config.tensor_model_parallel_size = tp_size + _disable_gtp(config) config.sequence_parallel = False return config diff --git a/examples/mimo/model_providers/radio_encoder.py b/examples/mimo/model_providers/radio_encoder.py index 9e0591cc7e7..f28dd5af6a3 100644 --- a/examples/mimo/model_providers/radio_encoder.py +++ b/examples/mimo/model_providers/radio_encoder.py @@ -83,6 +83,15 @@ def _make_dense_non_hybrid(config: TransformerConfig) -> None: config.use_fused_weighted_squared_relu = False +def _disable_gtp(config: TransformerConfig) -> None: + """Keep this module replicated across any LLM GTP axes.""" + config.tensor_parallel_num_weight_shards = config.tensor_model_parallel_size + config.gtp_weight_remat_size = 1 + expert_tp = config.expert_tensor_parallel_size or config.tensor_model_parallel_size + config.expert_tensor_parallel_num_weight_shards = expert_tp + config.expert_gtp_weight_remat_size = 1 + + def radio_vision_config(args: argparse.Namespace, tp_size: int, pp_size: int) -> TransformerConfig: """RADIO vision config: stock from-args base + RADIO-specific overrides.""" config = deepcopy(_base_config(args)) @@ -114,6 +123,7 @@ def radio_vision_config(args: argparse.Namespace, tp_size: int, pp_size: int) -> config.bf16 = bf16 config.tensor_model_parallel_size = tp_size config.pipeline_model_parallel_size = pp_size + _disable_gtp(config) config.sequence_parallel = False return config diff --git a/examples/mimo/pretrain_mimo.py b/examples/mimo/pretrain_mimo.py index ee521b188a5..9832cc806c8 100644 --- a/examples/mimo/pretrain_mimo.py +++ b/examples/mimo/pretrain_mimo.py @@ -43,20 +43,22 @@ def extra_args_provider(parser: argparse.ArgumentParser) -> argparse.ArgumentPar return parser +def _set_stock_parallel_args(args: argparse.Namespace) -> None: + # validate_args and the stock training loop consume these fields. + args.tensor_model_parallel_size = args.llm_tp + args.pipeline_model_parallel_size = args.llm_pp + args.context_parallel_size = args.llm_cp + args.expert_model_parallel_size = args.llm_ep + args.expert_tensor_parallel_size = args.llm_expt_tp or 1 + + def _parse_and_validate() -> argparse.Namespace: """Parse stock plus MIMO arguments and validate the disjoint module grids.""" args = parse_args(extra_args_provider) - validate_hetero_grid_args(args, args.world_size) + _, llm_size = validate_hetero_grid_args(args, args.world_size) + _set_stock_parallel_args(args) physical_world_size = args.world_size - # Stock validate_args sets data_parallel_size = world_size // (tp*pp*cp); feed the - # language module's world (llm_dp; stock tp/pp/cp stay 1, MIMO parallelism is in --llm-*) - # so it yields llm_dp. The physical world incl. encoder ranks is restored below. - args.world_size = ( - args.llm_dp - * args.tensor_model_parallel_size - * args.pipeline_model_parallel_size - * args.context_parallel_size - ) + args.world_size = llm_size try: validate_args(args, {"dataloader_type": "external"}) finally: diff --git a/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh b/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh index 787df50fef8..a0fd5be85a0 100755 --- a/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh +++ b/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh @@ -9,8 +9,8 @@ NUM_MICROBATCHES=${NUM_MICROBATCHES:-4} EVAL_INTERVAL=${EVAL_INTERVAL:-1} EVAL_ITERS=${EVAL_ITERS:-0} MICRO_BATCH_SIZE=1 -LLM_DP=2 -GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * LLM_DP)) +LLM_DP=1 +GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * LLM_DP * 2)) TORCHRUN_LOG_DIR=${TORCHRUN_LOG_DIR:-"${PWD}/logs/torchrun-$(date +%Y%m%d_%H%M%S)-$$"} mkdir -p "${TORCHRUN_LOG_DIR}" @@ -78,8 +78,10 @@ uv run --extra ssm python -m torch.distributed.run \ --llm-cp 1 \ --llm-pp 1 \ --llm-dp "${LLM_DP}" \ - --llm-ep 4 \ + --llm-ep 2 \ --llm-expt-tp 1 \ + --tensor-parallel-num-weight-shards 4 \ + --expert-tensor-parallel-num-weight-shards 2 \ --vocab-size 131072 \ --micro-batch-size "${MICRO_BATCH_SIZE}" \ --global-batch-size "${GLOBAL_BATCH_SIZE}" \ diff --git a/examples/mimo/training/args.py b/examples/mimo/training/args.py index e62bc31967b..cd89f4579d4 100644 --- a/examples/mimo/training/args.py +++ b/examples/mimo/training/args.py @@ -8,6 +8,7 @@ from typing import List from examples.mimo.training.topology import ModuleGridSpec +from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY @@ -61,6 +62,7 @@ def add_hetero_grid_args(parser: argparse.ArgumentParser) -> argparse.ArgumentPa def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tuple[int, int]: """Validate the disjoint hetero grid layout; returns ``(encoder_size, llm_size)``.""" + gtp_weight_remat_size, _ = resolve_hetero_gtp_degrees(args) if args.llm_cp != 1: raise ValueError("hetero MIMO training currently supports CP=1 only") @@ -76,7 +78,7 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl f"--num-experts ({num_experts}) must be divisible by --llm-ep ({args.llm_ep})" ) - llm_size = args.llm_tp * args.llm_cp * args.llm_pp * args.llm_dp + llm_size = args.llm_tp * gtp_weight_remat_size * args.llm_cp * args.llm_pp * args.llm_dp if args.llm_only: if getattr(args, "encoder_ddp_overlap", False): @@ -94,12 +96,14 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl ) return 0, llm_size - # Fan-out divisibility: the bridge splits (mbs * llm_dp) LLM lanes across + # Fan-out divisibility: the bridge splits every LLM data lane across # encoder_dp encoder lanes; the split must be exact. - if (args.micro_batch_size * args.llm_dp) % args.encoder_dp != 0: + llm_data_parallel_size = args.llm_dp * gtp_weight_remat_size + if (args.micro_batch_size * llm_data_parallel_size) % args.encoder_dp != 0: raise ValueError( - "--micro-batch-size * --llm-dp must be divisible by --encoder-dp " - f"(got {args.micro_batch_size} * {args.llm_dp} % {args.encoder_dp} != 0)" + "--micro-batch-size * --llm-dp * GTP must be divisible by --encoder-dp " + f"(got {args.micro_batch_size} * {args.llm_dp} * " + f"{gtp_weight_remat_size} % {args.encoder_dp} != 0)" ) encoder_size = args.encoder_tp * args.encoder_dp @@ -126,6 +130,7 @@ def build_module_grid_specs( ) -> List[ModuleGridSpec]: """Map grid args to the ModuleGridSpec list create_topology consumes.""" encoder_size, llm_size = validate_hetero_grid_args(args, world_size) + gtp_weight_remat_size, expert_gtp_weight_remat_size = resolve_hetero_gtp_degrees(args) language_grid_spec = ModuleGridSpec( name=MIMO_LANGUAGE_MODULE_KEY, @@ -134,8 +139,10 @@ def build_module_grid_specs( cp=args.llm_cp, pp=args.llm_pp, ep=args.llm_ep, + gtp_remat=gtp_weight_remat_size, rank_offset=args.llm_offset, expt_tp=args.llm_expt_tp or 1, + expt_gtp_remat=expert_gtp_weight_remat_size, ) if args.llm_only: @@ -154,6 +161,23 @@ def build_module_grid_specs( return [encoder_grid_spec, language_grid_spec] +def resolve_hetero_gtp_degrees(args: argparse.Namespace) -> tuple[int, int]: + """Return the language dense and expert GTP degrees.""" + _, gtp_weight_remat_size = resolve_tensor_parallel_weight_shards( + args.llm_tp, + getattr(args, "tensor_parallel_num_weight_shards", None), + getattr(args, "gtp_weight_remat_size", 1), + ) + _, expert_gtp_weight_remat_size = resolve_tensor_parallel_weight_shards( + args.llm_expt_tp or 1, + getattr(args, "expert_tensor_parallel_num_weight_shards", None), + getattr(args, "expert_gtp_weight_remat_size", 1), + shards_field="expert_tensor_parallel_num_weight_shards", + tp_field="expert_tensor_parallel_size", + ) + return gtp_weight_remat_size, expert_gtp_weight_remat_size + + def _num_experts(args: argparse.Namespace) -> int: """Resolve MoE expert count from the stock --num-experts arg.""" value = getattr(args, "num_experts", None) diff --git a/examples/mimo/training/data.py b/examples/mimo/training/data.py index 9139711662b..986c521874e 100644 --- a/examples/mimo/training/data.py +++ b/examples/mimo/training/data.py @@ -291,8 +291,12 @@ def build_train_valid_test_data_loaders( raise ValueError(f"unsupported dataset provider: {args.dataset_provider}") encoder_name = _encoder_name(topology) - if encoder_name is not None and (args.micro_batch_size * args.llm_dp) % args.encoder_dp: - raise ValueError("micro_batch_size * llm_dp must be divisible by encoder_dp") + llm_data_parallel_size = args.llm_dp * args.gtp_weight_remat_size + if ( + encoder_name is not None + and (args.micro_batch_size * llm_data_parallel_size) % args.encoder_dp + ): + raise ValueError("micro_batch_size * llm_dp * GTP must be divisible by encoder_dp") language_grid = topology.grids[MIMO_LANGUAGE_MODULE_KEY] language_pgc = topology.module_pgs[MIMO_LANGUAGE_MODULE_KEY] @@ -312,7 +316,7 @@ def build_train_valid_test_data_loaders( if encoder_needs_data and language_needs_data: raise ValueError("the external DataLoader adapter requires non-colocated module grids") if encoder_needs_data: - encoder_mbs = args.micro_batch_size * args.llm_dp // args.encoder_dp + encoder_mbs = args.micro_batch_size * llm_data_parallel_size // args.encoder_dp return _build_split_loaders( args, batch_size=encoder_mbs, @@ -340,7 +344,8 @@ def _build_split_loaders( encoder_name: Optional[str], ) -> tuple[DataLoader, DataLoader, DataLoader]: """Build split-local datasets with deterministic module/DP/split seeds.""" - base_seed = args.seed + module_seed_offset + get_pg_rank(pg_collection.dp) + data_group = pg_collection.dp_cp_gtp_remat or pg_collection.dp + base_seed = args.seed + module_seed_offset + get_pg_rank(data_group) common = _mock_loader_kwargs(args, encoder_name) return tuple( _build_mock_vlm_dataloader( diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py index 2a06a4b8188..bff506c074e 100644 --- a/examples/mimo/training/grad_sync.py +++ b/examples/mimo/training/grad_sync.py @@ -62,7 +62,7 @@ def _is_pg_member(pg) -> bool: def _is_token_source_rank(language_pg) -> bool: """Whether this rank is on the LLM (last PP stage, TP rank 0) coordinate that sums - the global token count over DP/CP. + the global token count over all data lanes. Sourcing from this single coordinate avoids double-counting across TP/PP replicas. The _is_pg_member guards short-circuit encoder-grid ranks (non-member pp/tp groups) @@ -101,18 +101,17 @@ def _token_source_global_rank(language_grid) -> int: def _global_token_count(num_tokens, language_pg, src_global_rank) -> float: """Total non-padded tokens in the global batch, visible on every rank. - Only the LLM token-source rank computes the count by summing over the LLM DP/CP - group; it then broadcasts that N_global from its global rank to every rank in the + Only the LLM token-source ranks compute the count by summing over all LLM data lanes; + the source then broadcasts N_global from its global rank to every rank in the world (including the non-colocated encoder grid, where ``language_pg`` is None) so both modules divide by the same per-token mean. """ global_num_tokens = torch.zeros(1, dtype=torch.float32, device="cuda") if _is_token_source_rank(language_pg): - # Collective over DP/CP: every (pp_last, tp0) rank participates so the all-reduce - # does not hang; only DP/CP rank 0 keeps the result and is the broadcast root. + data_group = language_pg.dp_cp_gtp_remat or language_pg.dp_cp token_count = num_tokens.to(dtype=torch.float32).sum().view(1) - dist.all_reduce(token_count, group=language_pg.dp_cp, op=dist.ReduceOp.SUM) - if dist.get_rank(group=language_pg.dp_cp) == 0: + dist.all_reduce(token_count, group=data_group, op=dist.ReduceOp.SUM) + if dist.get_rank(group=data_group) == 0: global_num_tokens.copy_(token_count) dist.broadcast(global_num_tokens, src=src_global_rank) return float(global_num_tokens.item()) diff --git a/examples/mimo/training/runtime.py b/examples/mimo/training/runtime.py index 6a5235aba13..8b46f1fb5f2 100644 --- a/examples/mimo/training/runtime.py +++ b/examples/mimo/training/runtime.py @@ -40,7 +40,7 @@ def configure_module_rng( so disjoint modules (and stages) get independent RNG state. Caller invokes once per active module on this rank. """ - for _required in ("pp", "dp", "tp", "ep", "expt_tp"): + for _required in ("pp", "dp", "tp", "ep", "expt_tp", "gtp_remat", "expt_gtp_remat"): assert ( getattr(pg_collection, _required, None) is not None ), f"pg_collection passed to configure_module_rng must define {_required}" @@ -52,6 +52,8 @@ def configure_module_rng( tp_group=pg_collection.tp, ep_group=pg_collection.ep, etp_group=pg_collection.expt_tp, + gtp_remat_group=pg_collection.gtp_remat, + egtp_remat_group=pg_collection.expt_gtp_remat, ) diff --git a/examples/mimo/training/topology.py b/examples/mimo/training/topology.py index 60e473f58ff..8146972da00 100644 --- a/examples/mimo/training/topology.py +++ b/examples/mimo/training/topology.py @@ -33,23 +33,26 @@ class ModuleGridSpec: cp: int = 1 pp: int = 1 ep: int = 1 + gtp_remat: int = 1 rank_offset: int = 0 # Experts default to TP=1 (set explicitly for MoE); intentionally not Megatron's etp=tp default. expt_tp: int = 1 + expt_gtp_remat: int = 1 dp: int = field(init=False) expt_dp: int = field(init=False) def __post_init__(self) -> None: - dense = self.tp * self.cp * self.pp + dense = self.tp * self.gtp_remat * self.cp * self.pp if self.num_ranks % dense != 0: raise ValueError( - f"num_ranks ({self.num_ranks}) must be divisible by tp*cp*pp ({dense})" + f"num_ranks ({self.num_ranks}) must be divisible by tp*gtp_remat*cp*pp ({dense})" ) self.dp = self.num_ranks // dense - expert = self.expt_tp * self.ep * self.pp + expert = self.expt_tp * self.ep * self.expt_gtp_remat * self.pp if self.num_ranks % expert != 0: raise ValueError( - f"num_ranks ({self.num_ranks}) must be divisible by expt_tp*ep*pp ({expert})" + f"num_ranks ({self.num_ranks}) must be divisible by " + f"expt_tp*ep*expt_gtp_remat*pp ({expert})" ) self.expt_dp = self.num_ranks // expert @@ -118,27 +121,45 @@ def create_topology(specs: list[ModuleGridSpec]) -> HeteroTopology: def _build_grid(spec: ModuleGridSpec) -> HyperCommGrid: """Create a dense grid plus its expert view and the process groups MIMO needs.""" grid = HyperCommGrid( - shape=[spec.tp, spec.cp, spec.dp, spec.pp], - dim_names=["tp", "cp", "dp", "pp"], + shape=[spec.tp, spec.gtp_remat, spec.cp, spec.dp, spec.pp], + dim_names=["tp", "gtp_remat", "cp", "dp", "pp"], rank_offset=spec.rank_offset, backend="nccl", ) # Expert factorization over the same rank span; pp is shared with the base view. grid.register_view( _EXPERT_VIEW, - shape=[spec.expt_tp, spec.ep, spec.expt_dp, spec.pp], - dim_names=["expt_tp", "ep", "expt_dp", "pp"], + shape=[spec.expt_tp, spec.ep, spec.expt_gtp_remat, spec.expt_dp, spec.pp], + dim_names=["expt_tp", "ep", "expt_gtp_remat", "expt_dp", "pp"], shared_dims=["pp"], ) try: for dims in ( - ["tp"], ["cp"], ["pp"], ["dp"], - ["dp", "cp"], ["tp", "cp"], ["tp", "pp"], - ["tp", "dp"], ["tp", "dp", "cp"], ["tp", "cp", "dp", "pp"], + ["tp"], + ["gtp_remat"], + ["cp"], + ["pp"], + ["dp"], + ["dp", "cp"], + ["gtp_remat", "dp", "cp"], + ["tp", "cp"], + ["tp", "gtp_remat", "pp"], + ["tp", "gtp_remat", "dp"], + ["tp", "gtp_remat", "dp", "cp"], + ["tp", "gtp_remat", "cp", "dp", "pp"], ): grid.create_pg(dims) - for dims in (["ep"], ["expt_tp"], ["expt_dp"], ["expt_tp", "ep"], ["expt_tp", "ep", "pp"]): + for dims in ( + ["ep"], + ["expt_tp"], + ["expt_gtp_remat"], + ["expt_dp"], + ["expt_tp", "ep"], + ["expt_tp", "ep", "pp"], + ["expt_tp", "ep", "expt_gtp_remat", "pp"], + ["expt_gtp_remat", "expt_dp"], + ): grid.create_pg(dims, view=_EXPERT_VIEW) except Exception: grid.destroy() @@ -193,18 +214,25 @@ def pg_collection_from_grid( pgc.pp = grid.get_pg("pp") pgc.dp = grid.get_pg("dp") pgc.dp_cp = grid.get_pg(["dp", "cp"]) + pgc.dp_cp_gtp_remat = grid.get_pg(["gtp_remat", "dp", "cp"]) pgc.intra_dp_cp = pgc.dp_cp + pgc.gtp_remat = grid.get_pg("gtp_remat") pgc.tp_cp = grid.get_pg(["tp", "cp"]) - pgc.tp_dp = grid.get_pg(["tp", "dp"]) - pgc.tp_dp_cp = grid.get_pg(["tp", "dp", "cp"]) - pgc.mp = grid.get_pg(["tp", "pp"]) - pgc.intra_dist_opt = grid.get_pg(["tp", "cp", "dp", "pp"]) + pgc.tp_dp = grid.get_pg(["tp", "gtp_remat", "dp"]) + pgc.tp_dp_cp = grid.get_pg(["tp", "gtp_remat", "dp", "cp"]) + pgc.mp = grid.get_pg(["tp", "gtp_remat", "pp"]) + pgc.intra_dist_opt = grid.get_pg(["tp", "gtp_remat", "cp", "dp", "pp"]) pgc.ep = grid.get_pg("ep", view=_EXPERT_VIEW) pgc.expt_tp = grid.get_pg("expt_tp", view=_EXPERT_VIEW) + pgc.expt_gtp_remat = grid.get_pg("expt_gtp_remat", view=_EXPERT_VIEW) pgc.expt_dp = grid.get_pg("expt_dp", view=_EXPERT_VIEW) + pgc.expt_dp_gtp_remat = grid.get_pg(["expt_gtp_remat", "expt_dp"], view=_EXPERT_VIEW) pgc.intra_expt_dp = pgc.expt_dp pgc.tp_ep = grid.get_pg(["expt_tp", "ep"], view=_EXPERT_VIEW) pgc.tp_ep_pp = grid.get_pg(["expt_tp", "ep", "pp"], view=_EXPERT_VIEW) + pgc.tp_ep_pp_with_egtp_remat = grid.get_pg( + ["expt_tp", "ep", "expt_gtp_remat", "pp"], view=_EXPERT_VIEW + ) pgc.embd = None pgc.pos_embd = None if is_language: diff --git a/megatron/core/models/mimo/model/base.py b/megatron/core/models/mimo/model/base.py index d5b70826bb5..854f6abd227 100644 --- a/megatron/core/models/mimo/model/base.py +++ b/megatron/core/models/mimo/model/base.py @@ -549,6 +549,8 @@ def _attach_modality_split_sizes( language_grid = grid_map[MIMO_LANGUAGE_MODULE_KEY] encoder_dp = encoder_grid.shape[encoder_grid.dim_names.index("dp")] language_dp = language_grid.shape[language_grid.dim_names.index("dp")] + if "gtp_remat" in language_grid.dim_names: + language_dp *= language_grid.shape[language_grid.dim_names.index("gtp_remat")] assert encoder_dp <= language_dp, ( f"Bridge fan-out split metadata with non-uniform per-sample sizes " f"requires encoder DP <= LM DP (got encoder='{encoder_name}' " diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py index 598f6c883af..2eb72ceee51 100644 --- a/megatron/core/models/mimo/optimizer.py +++ b/megatron/core/models/mimo/optimizer.py @@ -16,6 +16,7 @@ from megatron.core.optimizer.optimizer import MegatronOptimizer from megatron.core.optimizer.optimizer_config import OptimizerConfig from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.utils import unwrap_model if TYPE_CHECKING: from megatron.core.hyper_comm_grid import HyperCommGrid @@ -339,29 +340,6 @@ def _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple: return (pg_collection.tp.rank(), pg_collection.pp.rank(), pg_collection.dp.rank()) -_EXPERT_VIEW = "expert" - - -def _get_pg_collection_for_optimizer(grid) -> ProcessGroupCollection: - """Derive the optimizer's ProcessGroupCollection from a populated HyperCommGrid. - - Dense groups come from the base view; expert-parallel groups (tp_ep_pp, expt_dp) come from - the grid's dedicated expert view -- expert parallelism is always factored into a separate - view (expt_tp/ep/expt_dp), never the base view. All groups must be pre-created on the grid. - """ - pg = ProcessGroupCollection() - pg.dp = grid.get_pg("dp") - pg.dp_cp = grid.get_pg(["dp", "cp"]) - pg.tp = grid.get_pg("tp") - pg.pp = grid.get_pg("pp") - pg.mp = grid.get_pg(["tp", "pp"]) - pg.tp_ep_pp = grid.get_pg(["expt_tp", "ep", "pp"], view=_EXPERT_VIEW) - pg.expt_dp = grid.get_pg("expt_dp", view=_EXPERT_VIEW) - # Distributed-optimizer grad-stats group spans the dense shards (mirrors the topology PGC). - pg.intra_dist_opt = grid.get_pg(["tp", "cp", "dp", "pp"]) - return pg - - def get_mimo_optimizer(mimo_model: "MimoModel", config: OptimizerConfig) -> MimoOptimizer: """Create optimizer for MimoModel with heterogeneous parallelism.""" from megatron.core.optimizer import get_megatron_optimizer @@ -386,7 +364,10 @@ def get_mimo_optimizer(mimo_model: "MimoModel", config: OptimizerConfig) -> Mimo module = mimo_model.modality_submodules[module_name] if module is not None: - pg_collection = _get_pg_collection_for_optimizer(grid) + pg_collection = getattr(unwrap_model(module), 'pg_collection', None) + assert ( + pg_collection is not None + ), f"Module '{module_name}' must own a ProcessGroupCollection for optimizer setup" assert ( not hasattr(module, 'ddp_config') or module.ddp_config is None diff --git a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py index fb923c2b37d..8f6ac5435e2 100644 --- a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py +++ b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py @@ -62,6 +62,33 @@ def test_canonical_layout_validates_and_maps_specs(): assert language_grid_spec.expt_tp == 1 +def test_gtp_layout_validates_and_maps_weight_shard_axes(): + args = _layout_8gpu_20l( + llm_dp=1, + llm_ep=2, + tensor_parallel_num_weight_shards=4, + expert_tensor_parallel_num_weight_shards=2, + ) + + assert validate_hetero_grid_args(args, WORLD_SIZE_8) == (4, 4) + assert not hasattr(args, "gtp_weight_remat_size") + assert not hasattr(args, "expert_gtp_weight_remat_size") + + _, language_grid_spec = build_module_grid_specs( + args, WORLD_SIZE_8, encoder_module_name="radio_encoder" + ) + assert language_grid_spec.gtp_remat == 2 + assert language_grid_spec.dp == 1 + assert language_grid_spec.expt_gtp_remat == 2 + assert language_grid_spec.expt_dp == 1 + + +def test_weight_shards_must_divide_language_tp(): + args = _layout_8gpu_20l(tensor_parallel_num_weight_shards=3) + with pytest.raises(ValueError, match="must be divisible"): + validate_hetero_grid_args(args, WORLD_SIZE_8) + + def test_overlapping_spans_raise(): # llm-offset 2 makes llm ranks {2,3,4,5} overlap encoder ranks {0,1,2,3}. args = _layout_8gpu_20l(llm_offset=2) diff --git a/tests/unit_tests/models/mimo/test_mimo_mock_data.py b/tests/unit_tests/models/mimo/test_mimo_mock_data.py index 6ff70d318ac..873ae995b34 100644 --- a/tests/unit_tests/models/mimo/test_mimo_mock_data.py +++ b/tests/unit_tests/models/mimo/test_mimo_mock_data.py @@ -26,6 +26,7 @@ def _args(): dataset_provider="mock", micro_batch_size=2, llm_dp=2, + gtp_weight_remat_size=1, encoder_dp=1, seq_length=8, image_seq_length=4, @@ -46,10 +47,14 @@ def _args(): def _topology(*, language_rank, encoder_rank=None): encoder = RADIO_ENCODER_MODULE_NAME grids = {"language": _grid(language_rank)} - pgs = {"language": SimpleNamespace(pp=_group(size=3), dp=_group(rank=0, size=2))} + pgs = { + "language": SimpleNamespace( + pp=_group(size=3), dp=_group(rank=0, size=2), dp_cp_gtp_remat=None + ) + } if encoder_rank is not None: grids[encoder] = _grid(encoder_rank) - pgs[encoder] = SimpleNamespace(pp=_group(), dp=_group(rank=1, size=2)) + pgs[encoder] = SimpleNamespace(pp=_group(), dp=_group(rank=1, size=2), dp_cp_gtp_remat=None) return SimpleNamespace(grids=grids, module_pgs=pgs) diff --git a/tests/unit_tests/test_mimo_hetero_topology.py b/tests/unit_tests/test_mimo_hetero_topology.py index 7d68e107bf7..e7e8551b5a0 100644 --- a/tests/unit_tests/test_mimo_hetero_topology.py +++ b/tests/unit_tests/test_mimo_hetero_topology.py @@ -25,6 +25,21 @@ def _specs(): ] +def _gtp_specs(): + return [ + ModuleGridSpec(name=ENCODER, num_ranks=4, tp=2, rank_offset=0), + ModuleGridSpec( + name=MIMO_LANGUAGE_MODULE_KEY, + num_ranks=4, + tp=2, + gtp_remat=2, + ep=2, + expt_gtp_remat=2, + rank_offset=4, + ), + ] + + class TestModuleGridSpecResolution: def test_derived_dims_resolve_to_concrete_ints(self): # num_ranks=4,tp=2 with default expt_tp=1: dp=2, expt_dp=4. @@ -87,6 +102,31 @@ def test_pgc_group_sizes(self): finally: topo.destroy() + def test_gtp_pgc_group_sizes(self): + topo = create_topology(_gtp_specs()) + try: + rank = dist.get_rank() + pgc = ( + topo.module_pgs[ENCODER] if rank < 4 else topo.module_pgs[MIMO_LANGUAGE_MODULE_KEY] + ) + if rank < 4: + assert pgc.gtp_remat.size() == 1 + assert pgc.expt_gtp_remat.size() == 1 + else: + assert pgc.tp.size() == 2 + assert pgc.gtp_remat.size() == 2 + assert pgc.dp.size() == 1 + assert pgc.dp_cp_gtp_remat.size() == 2 + assert pgc.mp.size() == 4 + assert pgc.expt_tp.size() == 1 + assert pgc.ep.size() == 2 + assert pgc.expt_gtp_remat.size() == 2 + assert pgc.expt_dp.size() == 1 + assert pgc.expt_dp_gtp_remat.size() == 2 + assert pgc.tp_ep_pp_with_egtp_remat.size() == 4 + finally: + topo.destroy() + def test_embedding_groups(self): # Language grid is tp=2,pp=2 at ranks 4-7: each PP group is [first,last] (size 2), # so first/last-stage ranks get a 2-rank .embd and the first stage gets .pos_embd. From dcb2bfd97f82d2bf32e344dc1d785c63b52f5f07 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Tue, 18 Aug 2026 05:51:31 +0000 Subject: [PATCH 074/101] test(mfsdp): separate overlap test (#6591) Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 148 +------------- .../distributed/mfsdp_v2/test_overlap.py | 188 ++++++++++++++++++ 2 files changed, 190 insertions(+), 146 deletions(-) create mode 100644 tests/unit_tests/distributed/mfsdp_v2/test_overlap.py diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index ab341521ee3..967c18056fd 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -8,7 +8,7 @@ import torch import torch.distributed as dist from torch import nn -from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor from torch.profiler import ProfilerActivity, profile from torch.utils.checkpoint import checkpoint @@ -25,10 +25,7 @@ ) from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module import FsdpModule from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy -from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import ( - collect_linked_kernels, - events_overlap, -) +from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import collect_linked_kernels logger = logging.getLogger(__name__) @@ -181,10 +178,8 @@ def _mb(num_bytes: int) -> str: # CPU ops that a device event chains up to via cpu_parent, used to attribute the device # work to its enclosing collective or matmul operation. -_ALL_GATHER_OP_NAME_SUBSTRING = "allgather" _REDUCE_SCATTER_OP_NAME_SUBSTRING = "reduce_scatter" _ALLREDUCE_OP_NAME_SUBSTRING = "allreduce" -_GEMM_OP_NAME_SUBSTRING = "aten::mm" @pytest.mark.parametrize("num_microbatches", [1, 3]) @@ -774,145 +769,6 @@ def assert_returns_to_resting_memory(phase: str) -> None: assert_returns_to_resting_memory("backward") -@pytest.mark.parametrize( - "use_symmetric_memory", - [ - # Both variants' all-gathers are launch-timing sensitive, but default-CTA - # kernels occupy more compute CTAs, making the profiler overlap count less - # stable across ranks (see - # https://github.com/NVIDIA/Megatron-LM/actions/runs/31615942188). - # The symmetric-memory variant uses zero-CTA all-gather kernels, so its overlap - # measurement is more stable and remains enabled. - pytest.param(False, marks=(pytest.mark.flaky, pytest.mark.flaky_in_dev)), - pytest.param(True), - ], - ids=["default", "symmetric_memory"], -) -def test_overlaps_communication_and_compute(distributed_setup, use_symmetric_memory): - """Forward and backward communication should overlap GEMM compute.""" - world_size = distributed_setup.world_size - device = distributed_setup.device - if world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - # A large hidden size keeps the per-layer GEMMs long enough that the - # collectives reliably overlap them. The overlap count is otherwise - # launch-bound: the host issues kernels with gaps (amplified by CI's - # `coverage run` wrapper), so with short GEMMs a collective can land in a - # gap between GEMMs instead of running alongside one, making the count jitter - # run to run. At dim=16384 the GEMMs dominate that launch jitter and the - # overlap becomes deterministic. (dim=8192 was flaky under coverage.) - dim = 16384 - num_children = 4 - dtype = torch.bfloat16 - - # new_group requires a default process group. Initialize it here so this test works - # in isolation. Do not eagerly initialize it with device_id in the shared fixture: - # that can hang teardown after communicator splits; see - # https://github.com/pytorch/pytorch/issues/190396. - if not dist.is_initialized(): - dist.init_process_group(backend="nccl") - - if use_symmetric_memory: - # Dedicated communicator with NCCL's zero-CTA policy. cta_policy is a - # per-communicator property, so scoping it to this group leaves the rest of the - # bucket on default-CTA symmetric-memory kernels (test_symmetric_memory.py asserts - # ncclSymk all-gather kernel counts, which zero-CTA would turn into copy-engine - # memcpys). This 1-D group models the DP (FSDP) sub-mesh that mfsdp is handed in - # production: with EP/TP the full device mesh is multi-dimensional, but mfsdp - # requires an all-FSDP mesh (see experimental/module.py) and never sees the TP/EP - # axes, so only the DP communicator needs the zero-CTA policy. - zero_cta_options = dist.ProcessGroupNCCL.Options() - zero_cta_options.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO - dp_group = dist.new_group(backend="nccl", pg_options=zero_cta_options) - # NCCL window registration can fail when symmetric-memory rendezvous is the first - # operation on a communicator, so initialize this communicator explicitly. - dist.barrier(group=dp_group, device_ids=[device.index]) - else: - dp_group = dist.new_group(backend="nccl") - - mesh = DeviceMesh.from_group(dp_group, device.type) - model = MultiChildModel(dim=dim, num_children=num_children).to(device=device, dtype=dtype) - placements = _flat_placements() - policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) - with fully_shard_context(device=device, use_symmetric_memory=use_symmetric_memory): - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) - - x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True) - - def train_one_iteration() -> None: - model.zero_grad(set_to_none=True) - model(x).sum().backward() - - train_one_iteration() - torch.cuda.synchronize(device) - - with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: - train_one_iteration() - # Synchronize inside the profiler context so in-flight device kernels - # complete and get recorded before the profiler stops on __exit__. - # Synchronizing after the context would finalize the trace first and - # drop the CUDA events. - torch.cuda.synchronize(device) - - gemm_kernels = collect_linked_kernels(prof, _GEMM_OP_NAME_SUBSTRING) - # Each child Linear runs one forward and two backward matmuls. aten::mm may also - # launch auxiliary kernels, so check only the matmul lower bound. - assert len(gemm_kernels) >= 3 * num_children, ( - f"Expected at least {3 * num_children} kernels linked to GEMMs, got " - f"{len(gemm_kernels)}: " - f"{[kernel.name for kernel in gemm_kernels]}" - ) - - allgather_kernels = collect_linked_kernels(prof, _ALL_GATHER_OP_NAME_SUBSTRING) - reduce_scatter_kernels = collect_linked_kernels(prof, _REDUCE_SCATTER_OP_NAME_SUBSTRING) - # Each child layer does a forward and a backward all-gather and one - # reduce-scatter. Zero-CTA moves the all-gather to copy-engine memcpys, so it - # should not emit all-gather kernels. - expected_allgather_kernel_count = 0 if use_symmetric_memory else 2 * num_children - assert len(allgather_kernels) == expected_allgather_kernel_count, ( - f"Expected {expected_allgather_kernel_count} all-gather kernels, got " - f"{len(allgather_kernels)}: {[kernel.name for kernel in allgather_kernels]}" - ) - assert len(reduce_scatter_kernels) == num_children, ( - f"Expected {num_children} reduce-scatter kernels, got " - f"{len(reduce_scatter_kernels)}: {[kernel.name for kernel in reduce_scatter_kernels]}" - ) - - allgather_streams = {kernel.device_resource_id for kernel in allgather_kernels} - reduce_scatter_streams = {kernel.device_resource_id for kernel in reduce_scatter_kernels} - gemm_streams = {kernel.device_resource_id for kernel in gemm_kernels} - if allgather_kernels: - assert len(allgather_streams) == 1 - assert len(reduce_scatter_streams) == 1 - assert allgather_streams.isdisjoint(reduce_scatter_streams) - assert allgather_streams.isdisjoint(gemm_streams) - assert reduce_scatter_streams.isdisjoint(gemm_streams) - - allgather_overlap_count = sum( - any(events_overlap(kernel, gemm) for gemm in gemm_kernels) for kernel in allgather_kernels - ) - reduce_scatter_overlap_count = sum( - any(events_overlap(kernel, gemm) for gemm in gemm_kernels) - for kernel in reduce_scatter_kernels - ) - expected_allgather_overlap = 2 * (num_children - 1) - expected_reduce_scatter_overlap = num_children - 1 - if not use_symmetric_memory: - assert allgather_overlap_count >= expected_allgather_overlap, ( - f"Expected at least {expected_allgather_overlap} all-gathers to " - f"overlap compute, got {allgather_overlap_count}/{len(allgather_kernels)}." - ) - assert reduce_scatter_overlap_count >= expected_reduce_scatter_overlap, ( - f"Expected at least {expected_reduce_scatter_overlap} reduce-scatters to overlap " - f"compute, got {reduce_scatter_overlap_count}/{len(reduce_scatter_kernels)}." - ) - - # Release the dedicated communicator so it does not leak into the shared session. - dist.destroy_process_group(dp_group) - - def test_parameterless_parent_with_child_modules_trains(distributed_setup): """A parent with no unowned parameters should still root trainable child FsdpModules.""" world_size = distributed_setup.world_size diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py b/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py new file mode 100644 index 00000000000..26326962a6b --- /dev/null +++ b/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py @@ -0,0 +1,188 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Communication-overlap tests for the minimal Megatron-FSDP path.""" + +import pytest +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.device_mesh import DeviceMesh +from torch.profiler import ProfilerActivity, profile + +from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( + Flat, + Placements, + fully_shard, + fully_shard_context, +) +from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy +from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import ( + collect_linked_kernels, + events_overlap, +) + + +class MultiChildModel(nn.Module): + """Model with direct parameters and multiple child FsdpModules.""" + + def __init__(self, dim: int, num_children: int) -> None: + super().__init__() + self.bias = nn.Parameter(torch.ones(dim)) + self.layers = nn.ModuleList([nn.Linear(dim, dim, bias=False) for _ in range(num_children)]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run through every child layer with a root-owned bias.""" + x = x + self.bias + for layer in self.layers: + x = torch.relu(layer(x)) + return x + + +def _flat_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) + + +# CPU ops that a device event chains up to via cpu_parent, used to attribute the device +# work to its enclosing collective or matmul operation. +_ALL_GATHER_OP_NAME_SUBSTRING = "allgather" +_REDUCE_SCATTER_OP_NAME_SUBSTRING = "reduce_scatter" +_GEMM_OP_NAME_SUBSTRING = "aten::mm" + + +@pytest.mark.parametrize( + "use_symmetric_memory", + [ + # Both variants' all-gathers are launch-timing sensitive, but default-CTA + # kernels occupy more compute CTAs, making the profiler overlap count less + # stable across ranks (see + # https://github.com/NVIDIA/Megatron-LM/actions/runs/31615942188). + # The symmetric-memory variant uses zero-CTA all-gather kernels, so its overlap + # measurement is more stable and remains enabled. + pytest.param(False, marks=(pytest.mark.flaky, pytest.mark.flaky_in_dev)), + pytest.param(True), + ], + ids=["default", "symmetric_memory"], +) +def test_overlaps_communication_and_compute(distributed_setup, use_symmetric_memory): + """Forward and backward communication should overlap GEMM compute.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + # A large hidden size keeps the per-layer GEMMs long enough that the + # collectives reliably overlap them. The overlap count is otherwise + # launch-bound: the host issues kernels with gaps (amplified by CI's + # `coverage run` wrapper), so with short GEMMs a collective can land in a + # gap between GEMMs instead of running alongside one, making the count jitter + # run to run. At dim=16384 the GEMMs dominate that launch jitter and the + # overlap becomes deterministic. (dim=8192 was flaky under coverage.) + dim = 16384 + num_children = 4 + dtype = torch.bfloat16 + + # new_group requires a default process group. Initialize it here so this test works + # in isolation. Do not eagerly initialize it with device_id in the shared fixture: + # that can hang teardown after communicator splits; see + # https://github.com/pytorch/pytorch/issues/190396. + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + + if use_symmetric_memory: + # Dedicated communicator with NCCL's zero-CTA policy. cta_policy is a + # per-communicator property, so scoping it to this group leaves the rest of the + # bucket on default-CTA symmetric-memory kernels (test_symmetric_memory.py asserts + # ncclSymk all-gather kernel counts, which zero-CTA would turn into copy-engine + # memcpys). This 1-D group models the DP (FSDP) sub-mesh that mfsdp is handed in + # production: with EP/TP the full device mesh is multi-dimensional, but mfsdp + # requires an all-FSDP mesh (see experimental/module.py) and never sees the TP/EP + # axes, so only the DP communicator needs the zero-CTA policy. + zero_cta_options = dist.ProcessGroupNCCL.Options() + zero_cta_options.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO + dp_group = dist.new_group(backend="nccl", pg_options=zero_cta_options) + # NCCL window registration can fail when symmetric-memory rendezvous is the first + # operation on a communicator, so initialize this communicator explicitly. + dist.barrier(group=dp_group, device_ids=[device.index]) + else: + dp_group = dist.new_group(backend="nccl") + + mesh = DeviceMesh.from_group(dp_group, device.type) + model = MultiChildModel(dim=dim, num_children=num_children).to(device=device, dtype=dtype) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) + with fully_shard_context(device=device, use_symmetric_memory=use_symmetric_memory): + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True) + + def train_one_iteration() -> None: + model.zero_grad(set_to_none=True) + model(x).sum().backward() + + train_one_iteration() + torch.cuda.synchronize(device) + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + train_one_iteration() + # Synchronize inside the profiler context so in-flight device kernels + # complete and get recorded before the profiler stops on __exit__. + # Synchronizing after the context would finalize the trace first and + # drop the CUDA events. + torch.cuda.synchronize(device) + + gemm_kernels = collect_linked_kernels(prof, _GEMM_OP_NAME_SUBSTRING) + # Each child Linear runs one forward and two backward matmuls. aten::mm may also + # launch auxiliary kernels, so check only the matmul lower bound. + assert len(gemm_kernels) >= 3 * num_children, ( + f"Expected at least {3 * num_children} kernels linked to GEMMs, got " + f"{len(gemm_kernels)}: " + f"{[kernel.name for kernel in gemm_kernels]}" + ) + + allgather_kernels = collect_linked_kernels(prof, _ALL_GATHER_OP_NAME_SUBSTRING) + reduce_scatter_kernels = collect_linked_kernels(prof, _REDUCE_SCATTER_OP_NAME_SUBSTRING) + # Each child layer does a forward and a backward all-gather and one + # reduce-scatter. Zero-CTA moves the all-gather to copy-engine memcpys, so it + # should not emit all-gather kernels. + expected_allgather_kernel_count = 0 if use_symmetric_memory else 2 * num_children + assert len(allgather_kernels) == expected_allgather_kernel_count, ( + f"Expected {expected_allgather_kernel_count} all-gather kernels, got " + f"{len(allgather_kernels)}: {[kernel.name for kernel in allgather_kernels]}" + ) + assert len(reduce_scatter_kernels) == num_children, ( + f"Expected {num_children} reduce-scatter kernels, got " + f"{len(reduce_scatter_kernels)}: {[kernel.name for kernel in reduce_scatter_kernels]}" + ) + + allgather_streams = {kernel.device_resource_id for kernel in allgather_kernels} + reduce_scatter_streams = {kernel.device_resource_id for kernel in reduce_scatter_kernels} + gemm_streams = {kernel.device_resource_id for kernel in gemm_kernels} + if allgather_kernels: + assert len(allgather_streams) == 1 + assert len(reduce_scatter_streams) == 1 + assert allgather_streams.isdisjoint(reduce_scatter_streams) + assert allgather_streams.isdisjoint(gemm_streams) + assert reduce_scatter_streams.isdisjoint(gemm_streams) + + allgather_overlap_count = sum( + any(events_overlap(kernel, gemm) for gemm in gemm_kernels) for kernel in allgather_kernels + ) + reduce_scatter_overlap_count = sum( + any(events_overlap(kernel, gemm) for gemm in gemm_kernels) + for kernel in reduce_scatter_kernels + ) + expected_allgather_overlap = 2 * (num_children - 1) + expected_reduce_scatter_overlap = num_children - 1 + if not use_symmetric_memory: + assert allgather_overlap_count >= expected_allgather_overlap, ( + f"Expected at least {expected_allgather_overlap} all-gathers to " + f"overlap compute, got {allgather_overlap_count}/{len(allgather_kernels)}." + ) + assert reduce_scatter_overlap_count >= expected_reduce_scatter_overlap, ( + f"Expected at least {expected_reduce_scatter_overlap} reduce-scatters to overlap " + f"compute, got {reduce_scatter_overlap_count}/{len(reduce_scatter_kernels)}." + ) + + # Release the dedicated communicator so it does not leak into the shared session. + dist.destroy_process_group(dp_group) From bd72caf8a4afbb1789897789db261cfcfeea5eae Mon Sep 17 00:00:00 2001 From: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:07:47 +0000 Subject: [PATCH 075/101] GDP perf optimizations (#6553) Signed-off-by: Xiaowei Ren Signed-off-by: hongbinl Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: hongbinl --- .../fine_grained_activation_offload.py | 179 +++++- megatron/core/ssm/gated_delta_product.py | 601 ++++++++++++------ megatron/core/tensor_parallel/random.py | 8 +- .../core/transformer/transformer_config.py | 47 +- megatron/training/training.py | 5 +- .../models/test_hybrid_moe_model.py | 2 + ...test_fine_grained_activation_offloading.py | 285 +++++++++ tests/unit_tests/ssm/test_causal_conv1d.py | 80 +++ .../ssm/test_gated_delta_product.py | 220 +++++++ 9 files changed, 1203 insertions(+), 224 deletions(-) create mode 100644 tests/unit_tests/ssm/test_gated_delta_product.py diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index a9e7195dddf..a60036adcf9 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import logging from collections import defaultdict, deque from contextlib import nullcontext from typing import Any, Dict, Optional, Tuple @@ -16,9 +17,12 @@ except ImportError: from megatron.core.telemetry.fallbacks import trace_fn as _otel_trace_fn +from megatron.core._rank_utils import log_single_rank from megatron.core.transformer.cuda_graphs import is_graph_capturing from megatron.core.utils import nvtx_range_pop, nvtx_range_push +logger = logging.getLogger(__name__) + def debug_rank(message): """Print debug message for a specific rank when DEBUG is enabled.""" @@ -46,34 +50,61 @@ def _te_do_not_offload(tensor): ) -def print_offload_summary_table(total_offload_bytes: Dict[str, int]): +def print_offload_summary_table( + total_offload_bytes: Dict[str, int], + local_duplicate_summary: Optional[Dict[str, Tuple[int, int]]] = None, +): """ - Print an ASCII table summarizing offload bytes across all ranks. + Print offload bytes and warn about duplicate copies across all ranks. - Gathers offload data from all ranks and prints a formatted table on rank 0, - with rows representing ranks and columns representing groups. + Gathers both summaries in one object collective and reports them on rank 0. + The table has rows representing ranks and columns representing groups. + + The reported bytes are every byte copied to CPU, so a redundant copy of an + already-offloaded storage is counted like any other copy. Args: total_offload_bytes: Dict mapping group names to offload bytes for this rank. + local_duplicate_summary: Dict mapping group names to redundant-copy counts + and bytes for this rank. """ # pylint: disable=bad-builtin assert torch.distributed.is_initialized() rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() - # Gather all group names across ranks - local_names = list(total_offload_bytes.keys()) - all_names_list = [None] * world_size - torch.distributed.all_gather_object(all_names_list, local_names) - all_group_names = sorted(set(name for names in all_names_list for name in names)) - - # Gather offload bytes from all ranks: each rank sends a list of bytes per group - local_bytes = [total_offload_bytes.get(name, 0) for name in all_group_names] - all_bytes_list = [None] * world_size - torch.distributed.all_gather_object(all_bytes_list, local_bytes) + local_summary = { + "offload_bytes": dict(total_offload_bytes), + "duplicate_summary": dict(local_duplicate_summary or {}), + } + all_summaries = [None] * world_size + torch.distributed.all_gather_object(all_summaries, local_summary) - # Print ASCII table on rank 0 if rank == 0: + details = [] + for summary_rank, summary in enumerate(all_summaries): + for name, (copy_count, duplicate_bytes) in sorted(summary["duplicate_summary"].items()): + details.append( + f"rank {summary_rank}, {name}: {copy_count} redundant copies, " + f"{duplicate_bytes / (1024 * 1024):.2f} MB" + ) + if details: + log_single_rank( + logger, + logging.WARNING, + "Fine-grained activation offloading is copying the same tensor storage to CPU " + "more than once within an offload group, wasting host memory and PCIe bandwidth " + f"({', '.join(details)}). Offloading proceeds with the duplicated copies.", + ) + + all_group_names = sorted( + {name for summary in all_summaries for name in summary["offload_bytes"]} + ) + all_bytes_list = [ + [summary["offload_bytes"].get(name, 0) for name in all_group_names] + for summary in all_summaries + ] + # Calculate column widths col_width = max(12, max((len(name) for name in all_group_names), default=8) + 2) rank_col_width = max(6, len(f"Rank {world_size - 1}") + 2) @@ -364,6 +395,9 @@ def __init__(self, name): self.offload = True self.total_offload_bytes = 0 self.total_tensor_count = 0 + # Warmup-only bookkeeping for redundant copies within this group. + self.duplicate_storage_tensor_count = 0 + self.duplicate_storage_bytes = 0 # Using memory pool is for the compatibility with cuda graph. # Shapes of tensors for MoE activation offload groups are not known in advance, # so we do not use CPU pool for them. @@ -397,10 +431,15 @@ def wait_reload_event(self, stream): stream.wait_event(self._reload_event) def update_offload_info(self, tensor): - """Update the offload information.""" + """Update the offload information with the tensor actually copied to CPU.""" self.total_offload_bytes += tensor.numel() * tensor.element_size() self.total_tensor_count += 1 + def set_duplicate_storage_info(self, tensor_count, duplicate_bytes): + """Set redundant-copy statistics calculated for this group.""" + self.duplicate_storage_tensor_count = tensor_count + self.duplicate_storage_bytes = duplicate_bytes + class PipelineOffloadManager: """ @@ -628,6 +667,17 @@ def post_warmup_callback(self): group.offload = False else: break + # Warn about source bytes copied to CPU more than once within a single group. + # Once one view takes the whole-storage path, every other copy from that + # storage in the group is redundant. The copies are left in place -- this + # is only a diagnostic. This summary is gathered with the offload bytes + # below so post-warmup reporting needs only one object collective. + local_duplicate_summary = defaultdict(lambda: [0, 0]) + for chunk in self._cached_chunks_forward: + for group in chunk.offload_groups: + if group.offload and group.duplicate_storage_tensor_count > 0: + local_duplicate_summary[group._name][0] += group.duplicate_storage_tensor_count + local_duplicate_summary[group._name][1] += group.duplicate_storage_bytes # Dump the offload information total_tensor_count = {} total_offload_bytes = {} @@ -650,7 +700,10 @@ def post_warmup_callback(self): # Cache summary for downstream consumers (e.g., unit tests). self._offload_summary_bytes = dict(total_offload_bytes) self._offload_summary_total_bytes = int(sum(total_offload_bytes.values())) - print_offload_summary_table(total_offload_bytes) + print_offload_summary_table( + total_offload_bytes, + {name: tuple(summary) for name, summary in local_duplicate_summary.items()}, + ) def push(self, handler): """Add a chunk handler to the backward queue.""" @@ -831,13 +884,43 @@ class ChunkOffloadHandler: Manages tensor groups, coordinates asynchronous GPU-CPU transfers, and handles synchronization. """ + # Minimum fraction of the underlying storage bytes a non-contiguous view must + # cover before offloading the storage instead of gathering the view: the + # storage path trades D2H/H2D bytes on the copy engines for skipping the + # gather kernel and its temporary GPU buffer, which only pays off when the + # view spans most of the storage. + BASE_OFFLOAD_MIN_COVERAGE = 0.5 + @_otel_trace_fn('activation_offload', 'megatron.activation.offload') def offload(self, src_tensor, pin_memory=True, use_cpu_pool=True): - """Offload.""" + """Offload. + + A non-contiguous view that covers most of its underlying storage is + offloaded as the flat storage (a plain memcpy plus the view's + (size, stride, storage_offset) metadata) instead of being gathered into a + temporary contiguous copy. The base is resolved through + ``untyped_storage()`` rather than ``._base`` because autograd hands the + saved-tensor hooks a detach()-ed alias (e.g. CheckpointWithoutOutput's + saved input), which drops the view metadata but not the storage. + """ debug_rank("--------offload") + view_meta = None if not src_tensor.is_contiguous(): - src_tensor = src_tensor.contiguous() + storage = src_tensor.untyped_storage() + element_size = src_tensor.element_size() + covered_bytes = src_tensor.numel() * element_size + if ( + storage.nbytes() % element_size == 0 + and covered_bytes >= self.BASE_OFFLOAD_MIN_COVERAGE * storage.nbytes() + ): + view_meta = (src_tensor.size(), src_tensor.stride(), src_tensor.storage_offset()) + # Flat alias of the full storage; contiguous by construction. + src_tensor = torch.empty(0, dtype=src_tensor.dtype, device=src_tensor.device).set_( + storage + ) + else: + src_tensor = src_tensor.contiguous() if use_cpu_pool: cpu_backup = self.cpu_tensor_pool.allocate(src_tensor.shape, dtype=src_tensor.dtype) @@ -847,14 +930,14 @@ def offload(self, src_tensor, pin_memory=True, use_cpu_pool=True): ) cpu_backup.copy_(src_tensor, non_blocking=pin_memory) - state = (src_tensor.device, cpu_backup, use_cpu_pool) + state = (src_tensor.device, cpu_backup, use_cpu_pool, view_meta) return state @_otel_trace_fn('activation_offload', 'megatron.activation.reload') def reload(self, state, non_blocking=None): """Reload.""" debug_rank("------reload") - dev, cpu_backup, use_cpu_pool = state + dev, cpu_backup, use_cpu_pool, view_meta = state if non_blocking is None: non_blocking = cpu_backup.is_pinned() gpu_tensor = torch.empty( @@ -863,6 +946,9 @@ def reload(self, state, non_blocking=None): gpu_tensor.copy_(cpu_backup, non_blocking=non_blocking) if use_cpu_pool: self.cpu_tensor_pool.free(cpu_backup) + if view_meta is not None: + size, stride, storage_offset = view_meta + gpu_tensor = gpu_tensor.as_strided(size, stride, storage_offset) return gpu_tensor def __init__( @@ -1010,15 +1096,36 @@ def bulk_offload_group(self, group_to_offload): nvtx_msg = "activation offloading " + group_to_offload._name nvtx_range_push(nvtx_msg) with torch.cuda.stream(self.d2h_stream): + # Warmup-only accounting local to this offload group. Every tensor in + # the group is alive here, so device + data_ptr identifies its storage. + storage_records = ( + defaultdict( + lambda: { + "storage_bytes": 0, + "transfer_count": 0, + "full_storage_count": 0, + "transferred_bytes": 0, + } + ) + if self.is_warmup + else None + ) for tensor_tag, tensor_on_device in group_to_offload._tensors.items(): if self.tensor_need_offloading_checker(tensor_on_device): state = self.offload( tensor_on_device, use_cpu_pool=group_to_offload.use_cpu_pool ) + # Account the bytes actually copied to CPU (the base storage + # for view offloads). if self.is_warmup: - group_to_offload.update_offload_info(tensor_on_device) + group_to_offload.update_offload_info(state[1]) + self._record_offload_transfer(tensor_on_device, state, storage_records) + # record_stream on the view marks the shared storage + # allocation, so this also covers base-storage offloads. tensor_on_device.record_stream(self.d2h_stream) group_to_offload.push_tensor(tensor_tag, state) + if self.is_warmup: + self._set_duplicate_storage_info(group_to_offload, storage_records) group_to_offload.record_offload_event(self.d2h_stream) nvtx_range_pop(nvtx_msg) # Under full-iteration CG capture, the main stream may not wait on d2h @@ -1030,6 +1137,34 @@ def bulk_offload_group(self, group_to_offload): self._offload_pending_by_name[gname].append(group_to_offload._offload_event) self._drain_offload_pending(gname) + @staticmethod + def _record_offload_transfer(tensor_on_device, state, storage_records): + """Record one transfer for group-local duplicate-byte accounting.""" + _, cpu_backup, _, view_meta = state + storage = tensor_on_device.untyped_storage() + storage_key = (tensor_on_device.device, storage.data_ptr(), storage.nbytes()) + record = storage_records[storage_key] + record["storage_bytes"] = storage.nbytes() + record["transfer_count"] += 1 + record["full_storage_count"] += int(view_meta is not None) + record["transferred_bytes"] += cpu_backup.numel() * cpu_backup.element_size() + + @staticmethod + def _set_duplicate_storage_info(group, storage_records): + """Calculate exact redundant bytes for storages copied in full by this group.""" + duplicate_tensor_count = 0 + duplicate_bytes = 0 + for record in storage_records.values(): + # Without a full-storage copy, finding overlap between arbitrary + # strided logical views requires real storage-address analysis. + if record["full_storage_count"] == 0: + continue + storage_duplicate_bytes = record["transferred_bytes"] - record["storage_bytes"] + if storage_duplicate_bytes > 0: + duplicate_tensor_count += record["transfer_count"] - 1 + duplicate_bytes += storage_duplicate_bytes + group.set_duplicate_storage_info(duplicate_tensor_count, duplicate_bytes) + def get_max_deduplicated_groups(self): """Get the maximum number of deduplicated groups.""" count_modules = [] diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 22754aa981e..e6c8c945a8c 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -4,21 +4,27 @@ # This source code is licensed under the Apache license found in the # LICENSE file in the root directory of this source tree. +import inspect import logging import math from dataclasses import dataclass, replace +from functools import partial from typing import List, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F +from megatron.core import tensor_parallel from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory from megatron.core.inference.contexts import BaseInferenceContext, DynamicInferenceContext from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( tensor_masked_update, ) +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gdp_context_parallel import GDPContextParallel @@ -72,12 +78,20 @@ HAVE_EINOPS = False try: + from fla.modules.l2norm import l2_norm from fla.ops.gated_delta_product import chunk_gated_delta_product HAVE_FLA = True except ImportError: HAVE_FLA = False +try: + from gdp_attn import chunk_gated_delta_product as cutedsl_chunk_gated_delta_product + + HAVE_CUTEDSL_GDP = True +except ImportError: + HAVE_CUTEDSL_GDP = False + # Dynamic-batching inference runs the in-tree fork of these kernels rather than # the pip `flash-linear-attention` ones. The fork is forward-only, so training # and the static-batching path (which shares the training body) keep calling @@ -90,6 +104,20 @@ logger = logging.getLogger(__name__) +def _kernel_accepts_kwarg(kernel, name: str) -> bool: + """Return True if `kernel` explicitly declares a keyword argument called `name`.""" + try: + parameters = inspect.signature(kernel).parameters + except (TypeError, ValueError): + # Builtins / C extensions without introspectable signatures. + return False + parameter = parameters.get(name) + return parameter is not None and parameter.kind in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + + class ExtendedRMSNorm(RMSNormGated): """ RMSNormGated with sharded state dict. @@ -183,7 +211,10 @@ def __init__( "MambaSSM is not installed. Please install it with `pip install mamba-ssm`." ) - if not HAVE_FLA: + if config.gdp_cutedsl_kernel: + if not HAVE_CUTEDSL_GDP: + raise ImportError("gdp_attn (CuTeDSL GatedDeltaProduct) is not installed") + elif not HAVE_FLA: raise ImportError("FLA is not installed") super().__init__(config) @@ -199,6 +230,23 @@ def __init__( self.num_householder = config.gdp_num_householder + # Select the chunked gated delta product kernel once. The CuTeDSL and FLA + # implementations share the same keyword API but are imported under distinct + # names so neither shadows the other when both packages are installed. + self.gdp_kernel = ( + cutedsl_chunk_gated_delta_product + if config.gdp_cutedsl_kernel + else chunk_gated_delta_product + ) + + # Newer CuTeDSL kernel builds accept num_chunk_states_to_recompute; probe the + # signature once so older builds that predate the argument still work. + self.gdp_kernel_extra_kwargs = {} + if _kernel_accepts_kwarg(self.gdp_kernel, "num_chunk_states_to_recompute"): + self.gdp_kernel_extra_kwargs["num_chunk_states_to_recompute"] = ( + config.gdp_num_chunk_states_to_recompute + ) + self.config = config self.d_model = d_model self.d_conv = d_conv @@ -220,6 +268,21 @@ def __init__( self.pp_layer_offset = pp_layer_offset self.cached_batch_size = None + self.recompute_in_proj = ( + self.config.recompute_granularity == "selective" + and "gdp_in_proj" in self.config.recompute_modules + ) + + self.recompute_qkv = ( + self.config.recompute_granularity == "selective" + and "gdp_qkv" in self.config.recompute_modules + ) + + self.offload_gdp_qkv = ( + self.config.fine_grained_activation_offloading + and "gdp_qkv" in self.config.offload_modules + ) + tp_size = self.pg_collection.tp.size() self.nheads_local_tp = self.nheads // tp_size @@ -391,10 +454,21 @@ def forward( packed_seq_params=None, ): """Run the gated delta product mixer on hidden states.""" - seq_len, batch_size, dim = hidden_states.shape + inference_context = deprecate_inference_params(inference_context, inference_params) + + _, batch_size, _ = hidden_states.shape conv_state, ssm_state = None, None if inference_context is not None: + # Only decode is actually incompatible: it runs + # fused_recurrent_gated_delta_rule, which needs the batched layout, while + # _prepare_qkv and _compute_gating emit the CuTeDSL varlen layout whenever + # gdp_cutedsl_kernel is set. Checked once here, ahead of the static/dynamic + # split, so a run fails at the entry rather than part way through + # generation: both batching modes prefill and then decode the same request. + assert ( + not self.config.gdp_cutedsl_kernel + ), "gdp_cutedsl_kernel is only supported for training, not inference." if inference_context.is_dynamic_batching(): ok, reason = check_fla_sequence_packing_support() assert ok, reason @@ -415,30 +489,191 @@ def forward( # The states are updated in place. return self._static_decode(hidden_states, conv_state, ssm_state) - # Build cu_seqlens for the chunked recurrence (FLA) when running with - # packed (THD) sequences on the training/prefill path. - cu_seqlens_packed = None if packed_seq_params is not None: # ``hidden_states`` is [seq_len, batch, dim]; THD requires batch=1. assert batch_size == 1, "Packed sequences require batch=1 (THD/varlen format)." - cu_seqlens_packed = get_cu_seqlens(packed_seq_params) + y = self._gdp_chunk_forward( + hidden_states, + conv_state=conv_state, + ssm_state=ssm_state, + packed_seq_params=packed_seq_params, + ) + + out, out_bias = self.out_proj(y) + + return out, out_bias + + def _packed_metadata(self, packed_seq_params, pack_length): + """Return the ``(seq_idx, cu_seqlens)`` pair describing packed (THD) sequences. + + ``seq_idx`` is built from ``pack_length`` rather than the pre-projection sequence + length so that it matches the conv1d's input after in_proj's SP all-gather and + ``pre_conv_ssm``'s CP all-to-all, regardless of SP/CP/TP upstream-slicing. + Mirrors mamba_mixer.py, which calls _create_packed_seq_idx after the same gather + points. Both are ``None`` when sequences are not packed. + """ + if packed_seq_params is None: + return None, None + return build_packed_seq_idx(packed_seq_params, pack_length), get_cu_seqlens( + packed_seq_params + ) + + def _gdp_chunk_forward( + self, hidden_states, conv_state=None, ssm_state=None, packed_seq_params=None + ): + """Chunked-kernel forward, shared by training and static-batching prefill: input + projection, causal conv, QKV preparation, and the chunked gated delta product + kernel, with optional in_proj/QKV recompute and fine-grained activation + offloading of the causal conv input. + + Prefill passes conv_state and ssm_state so the trailing conv window and the final + recurrent state are cached for the decode steps. + """ + if self.recompute_in_proj: + # Checkpoint the input projection and its preprocessing, discard the z, VKQ, + # and ba outputs after the forward pass, and recompute them in the backward + # pass to save memory. Only the (much smaller) hidden_states input is kept. + # + # Quantized replay is safe for the same reason as in MLASelfAttention: + # CheckpointWithoutOutput records the forward recipe and amax state and + # replays under the recorded fp8_autocast, and the only quantized op inside + # _in_proj_preprocess is in_proj itself. Everything downstream of it (causal + # conv, gated delta product kernel) runs unquantized. + quantization = self.config.fp8 or self.config.fp4 + in_proj_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) + # ``packed_seq_params`` is bound rather than passed through ``checkpoint`` so + # that only tensors reach ``ctx.save_for_backward``. + z, VKQ, ba = in_proj_checkpoint.checkpoint( + partial(self._in_proj_preprocess, packed_seq_params=packed_seq_params), + hidden_states, + ) + else: + z, VKQ, ba = self._in_proj_preprocess( + hidden_states, packed_seq_params=packed_seq_params + ) + + seq_idx_packed, cu_seqlens_packed = self._packed_metadata(packed_seq_params, VKQ.shape[1]) + + # The offload group captures the tensors saved for backward inside the causal + # conv and QKV preparation: the checkpoint's saved input VKQ when QKV recompute + # is enabled, or the conv/QKV-prep internal saves otherwise. + qkv_manager = off_interface(self.offload_gdp_qkv, VKQ, "gdp_qkv") + with qkv_manager as VKQ: + if self.recompute_qkv: + # Checkpoint the causal conv and QKV preparation, discard the + # query/key/value outputs after the forward pass, and recompute them + # in the backward pass to save memory. + qkv_checkpoint = tensor_parallel.CheckpointWithoutOutput() + query, key, value = qkv_checkpoint.checkpoint( + partial( + self._prepare_qkv, + conv_state=conv_state, + seq_idx=seq_idx_packed, + l2_norm_in_kernel=self.config.gdp_cutedsl_kernel, + ), + VKQ, + ) + else: + # _prepare_qkv also seeds conv_state during prefill + query, key, value = self._prepare_qkv( + VKQ, + conv_state=conv_state, + seq_idx=seq_idx_packed, + l2_norm_in_kernel=self.config.gdp_cutedsl_kernel, + ) + + beta, g = self._compute_gating(ba) + + core_attn_out, last_recurrent_state = self._run_gdp_kernel( + query, + key, + value, + g, + beta, + VKQ, + output_final_state=(ssm_state is not None), + cu_seqlens=cu_seqlens_packed, + ) + + if ssm_state is not None: + ssm_state.copy_(last_recurrent_state) + + if self.recompute_qkv: + qkv_checkpoint.discard_output_and_register_recompute(core_attn_out) + + y = self._postprocess(core_attn_out, z, packed_seq_params=packed_seq_params) + + # Commit the offload group downstream of core_attn_out: the commit node's + # backward, which waits on the reload event, then runs before the recompute + # hook on core_attn_out unpacks the saved conv input. + y = qkv_manager.group_offload(y, forced_released_tensors=[]) + + if self.recompute_in_proj: + # Hook on the mixer output so the recompute runs at the very start of this + # mixer's backward: z is needed by the output norm, and VKQ is needed by the + # causal conv backward or by the QKV recompute hook on core_attn_out, all of + # which come later in the backward pass. + in_proj_checkpoint.discard_output_and_register_recompute(y) + + return y + + def _run_gdp_kernel(self, query, key, value, g, beta, VKQ, output_final_state, cu_seqlens=None): + """Run the selected chunked gated delta product kernel and return + (core_attn_out, final_state). + + For the CuTeDSL path the kernel expects packed varlen sequences described by + cu_seqlens and applies the query/key L2 norm itself, and its output is collapsed + to (b l) h p and restored here; the FLA path uses the batched layout and has + already applied the L2 norm in _prepare_qkv. ``cu_seqlens`` is the packed-sequence + layout when THD packing is active; otherwise the CuTeDSL path synthesizes the + uniform-length equivalent and the FLA path leaves it unset.""" + if self.config.gdp_cutedsl_kernel and cu_seqlens is None: + # query/key/value are already collapsed to (b l) ..., so read the batch and + # sequence dimensions from VKQ, which is still in (b, l, d) layout. + b, l = VKQ.shape[0], VKQ.shape[1] + cu_seqlens = torch.arange(0, (b + 1) * l, l, device=VKQ.device, dtype=torch.int32) + + core_attn_out, final_state = self.gdp_kernel( + q=query, + k=key, + v=value, + g=g, + beta=beta, + num_householder=self.num_householder, + initial_state=None, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=self.config.gdp_cutedsl_kernel, + cu_seqlens=cu_seqlens, + **self.gdp_kernel_extra_kwargs, + ) + if self.config.gdp_cutedsl_kernel: + core_attn_out = rearrange( + core_attn_out, "(b l) (h p) -> b l h p", b=VKQ.shape[0], p=self.headdim + ) + return core_attn_out, final_state + + def _in_proj_preprocess(self, hidden_states, packed_seq_params=None): + """Run the input projection, gather its output across CP ranks, switch to + (b, l, d) layout, and split it into the z, VKQ, and ba groups. + + Kept as a single function so that it can be checkpointed as one unit: the in_proj + output and the intermediate CP-gathered/transposed copies of it are then freed + inside the checkpoint, leaving only hidden_states saved for the backward pass. + """ zVKQba, _ = self.in_proj(hidden_states) zVKQba = self.cp.pre_conv_ssm(zVKQba, packed_seq_params=packed_seq_params) - # Build seq_idx *after* in_proj's SP all-gather and pre_conv_ssm's CP - # all-to-all. ``zVKQba.shape[0]`` is now the true pack_length, so the - # helper produces a seq_idx matching the conv1d's input length - # regardless of SP/CP/TP upstream-slicing. Mirrors mamba_mixer.py - # which calls _create_packed_seq_idx after the same gather points. - seq_idx_packed = None - if packed_seq_params is not None: - seq_idx_packed = build_packed_seq_idx(packed_seq_params, zVKQba.shape[0]) + return self._preprocess(zVKQba) + def _preprocess(self, zVKQba): + """Switch the (l, b, proj_dim) input projection to the batch-first layout the + causal conv and the kernels expect, and split it into the z, VKQ, and ba groups. + """ zVKQba = rearrange(zVKQba, "l b d -> b l d").contiguous() - z, VKQ, ba = torch.split( + return torch.split( zVKQba, [ self.cp.d_inner_local_tpcp, @@ -449,35 +684,72 @@ def forward( dim=-1, ) - # ``causal_conv1d_fn`` expects a ``[B, D, L]`` tensor. - # But the expected memory layout varies depending on whether seq_idx is set. - if seq_idx_packed is None: - # Default path: channels-first contiguous, stride(2) == 1. - VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() - else: - # ``causal_conv1d_fn(seq_idx=...)`` requires channels-last memory but [B, D, L] - # logical shape. This keeps the channels contiguous in memory. - VKQ = VKQ.contiguous() - VKQ = rearrange(VKQ, "b l d -> b d l") - - if conv_state is not None: - # Static-batching prefill: seed the conv state from the prompt's tail. - # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv - # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. - conv_state.copy_(F.pad(VKQ, (self.d_conv - VKQ.shape[-1], 0))) # Update state (B D W) - # causal_conv1d uses seq_idx_packed to reset the convolution boundaries - VKQ = causal_conv1d_fn( - x=VKQ, - weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), - bias=self.cp.get_conv1d_bias(), - activation=self.activation, - seq_idx=seq_idx_packed, - ) + def _prepare_qkv( + self, + x, + conv_state=None, + seq_idx=None, + is_decode=False, + conv_state_indices=None, + l2_norm_in_kernel=False, + ): + """Run the causal conv on the VKQ slice and split/reshape it into query, key, + and value. + + x: (b, l, d). Keep the transposes to the conv layout inside this function so that + the checkpointed subgraph takes the (b, l, d) slice as input and hands back a + gradient that is already contiguous in (b, l, d) for the split backward's concat. + + Decode passes ``is_decode`` to step the cached conv state one token at a time + with ``causal_conv1d_update`` (l == 1), selecting per-request cache rows with + ``conv_state_indices``. - VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() + ``l2_norm_in_kernel`` leaves the query/key L2 norm to the caller's kernel, which + every caller whose kernel normalizes internally must set: the CuTeDSL path + (``use_qk_l2norm_in_kernel``), ``ssm_prefill``'s varlen kernel, and decode's fused + recurrent kernel, which normalizes after the householder zeros are interleaved in. + """ + if is_decode: + # Indexed conv update: reads/writes the per-request conv state rows selected + # by ``conv_state_indices``, in place. Unlike ``causal_conv1d_fn`` below, the + # Triton ``causal_conv1d_update`` takes ``x`` as (batch, seq_len, dim) -- a + # 2-D (batch, dim) input is unsqueezed at dim 1 -- so decode keeps the + # (b, l, d) layout rather than transposing to [B, D, L]. Here l == 1. + x = causal_conv1d_update( + x, + conv_state, + rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), + self.cp.get_conv1d_bias(), + self.activation, + conv_state_indices=conv_state_indices, + ) + else: + # ``causal_conv1d_fn`` expects a ``[B, D, L]`` tensor in channels-last memory, + # which is also what it requires when ``seq_idx`` is set. ``x`` is a view into + # the channels-last ``zVKQba``, so the transpose alone already gives + # stride(1) == 1 and no copy is needed on either path. + x = rearrange(x, "b l d -> b d l") + if conv_state is not None: + # Static-batching prefill: seed the conv state from the prompt's tail. + # If we just take x[:, :, -self.d_conv :], it errors if seqlen < d_conv. + # Instead F.pad pads with zeros if seqlen < d_conv, and truncates otherwise. + conv_state.copy_(F.pad(x, (self.d_conv - x.shape[-1], 0))) # state (B D W) + if causal_conv1d_fn is None: + seqlen = x.size(2) + x = self.act(self.cp.conv1d(x)[..., :seqlen]) + else: + # causal_conv1d uses seq_idx to reset the convolution boundaries + x = causal_conv1d_fn( + x=x, + weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), + bias=self.cp.get_conv1d_bias(), + activation=self.activation, + seq_idx=seq_idx, + ) + x = rearrange(x, "b d l -> b l d") value, key, query = torch.split( - VKQ, + x, [ self.cp.d_inner_local_tpcp * self.num_householder, self.cp.ngroups_local_tpcp * self.d_state * self.num_householder, @@ -486,62 +758,76 @@ def forward( dim=-1, ) + if self.config.gdp_cutedsl_kernel: + # value is not GQA-expanded, so flatten it straight to the kernel layout. + # Keep query/key's group axis (g) separate so it can be GQA-expanded below; + # they are collapsed to the flat (b l) (...) layout afterwards. + value = rearrange( + value, "b l (m h p) -> (b l) (m h p)", m=self.num_householder, p=self.headdim + ) + key = rearrange(key, "b l (m g n) -> b l m g n", m=self.num_householder, n=self.d_state) + query = rearrange(query, "b l (g n) -> b l g n", n=self.d_state) + else: + value = rearrange( + value, "b l (m h p) -> b (l m) h p", m=self.num_householder, p=self.headdim + ).contiguous() + key = rearrange( + key, "b l (m g n) -> b (l m) g n", m=self.num_householder, n=self.d_state + ).contiguous() + query = rearrange(query, "b l (g n) -> b l g n", n=self.d_state).contiguous() + + if not l2_norm_in_kernel: + # Apply the L2 norm here so that it falls inside the QKV recompute boundary. + query = l2_norm(query) + key = l2_norm(key) + + if self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp > 1: + query = query.repeat_interleave( + self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=-2 + ) + key = key.repeat_interleave( + self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=-2 + ) + + if self.config.gdp_cutedsl_kernel: + # Collapse batch/seq and the feature axes into the flat layout the kernel wants. + key = rearrange(key, "b l m g n -> (b l) (m g n)") + query = rearrange(query, "b l g n -> (b l) (g n)") + + return query, key, value + + def _compute_gating(self, ba): + """Compute the beta and g gating tensors from the ba slice.""" b, a = torch.split( ba, [self.cp.nheads_local_tpcp * self.num_householder, self.cp.nheads_local_tpcp], dim=-1, ) - z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).contiguous() - value = rearrange( - value, "b l (m h p) -> b (l m) h p", m=self.num_householder, p=self.headdim - ).contiguous() - key = rearrange( - key, "b l (m g n) -> b (l m) g n", m=self.num_householder, n=self.d_state - ).contiguous() - query = rearrange(query, "b l (g n) -> b l g n", n=self.d_state).contiguous() - b, a = b.contiguous(), a.contiguous() beta = b.sigmoid() - beta = rearrange(beta, "b l (m h) -> b (l m) h", m=self.num_householder).contiguous() # If the model is loaded in fp16, without the .float() here, A might be -inf g = -self.cp.get_A_log().float().exp() * F.softplus(a.float() + self.cp.get_dt_bias()) - if self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp > 1: - query = query.repeat_interleave( - self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=2 - ) - key = key.repeat_interleave( - self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=2 - ) - - core_attn_out, last_recurrent_state = chunk_gated_delta_product( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=(ssm_state is not None), - num_householder=self.num_householder, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens_packed, - ) + if self.config.gdp_cutedsl_kernel: + beta = rearrange(beta, "b l d -> (b l) d") + g = rearrange(g, "b l h -> (b l) h") + else: + beta = rearrange(beta, "b l (m h) -> b (l m) h", m=self.num_householder).contiguous() - if ssm_state is not None: - ssm_state.copy_(last_recurrent_state) + return beta, g + def _postprocess(self, core_attn_out, z, packed_seq_params=None): + """Switch back to (l, b, d) layout, scatter across CP ranks, and apply the + gated output norm.""" y = rearrange(core_attn_out, "b l h p -> l b (h p)").contiguous() y = self.cp.post_conv_ssm(y, packed_seq_params=packed_seq_params) if self.rmsnorm: - z = rearrange(z, "b l h p -> l b (h p)").contiguous() + z = rearrange(z, "b l d -> l b d").contiguous() z = self.cp.post_conv_ssm(z, packed_seq_params=packed_seq_params) y = self.norm(y, z) - - out, out_bias = self.out_proj(y) - - return out, out_bias + return y # ================================================================== # Static / eager inference @@ -549,9 +835,9 @@ def forward( # ``_static_decode`` implements legacy static-batching decode. It is # deliberately kept separate from the dynamic inference hooks below so that # static-batching bookkeeping does not pollute the interface defined by - # ``SSMDynamicInferenceMixin``. Static-batching prefill shares the training - # body in ``forward``, which seeds the conv/SSM state when the caches are - # present. Mirrors ``MambaMixer._static_decode``. + # ``SSMDynamicInferenceMixin``. Static-batching prefill instead shares + # ``_gdp_chunk_forward`` with training, seeding the conv/SSM state when the + # caches are present. Mirrors ``MambaMixer._static_decode``. # ================================================================== def _static_decode( self, hidden_states: torch.Tensor, conv_state: torch.Tensor, ssm_state: torch.Tensor @@ -606,64 +892,31 @@ def ssm_decode( assert ( intermediate_conv_state is None and intermediate_ssm_state is None ), "GDP decode does not support speculative decoding yet" - zVKQba = zVKQba.squeeze(1) # [n, proj_dim] - M = self.num_householder - z, VKQ, ba = torch.split( - zVKQba, - [ - self.cp.d_inner_local_tpcp, - self.cp.d_inner_local_tpcp * M - + (M + 1) * self.cp.ngroups_local_tpcp * self.d_state, - self.cp.nheads_local_tpcp * (M + 1), - ], - dim=-1, - ) + # Keep the length-1 sequence dimension so the shared helpers apply: with l == 1 + # their (b, l, ...) reshapes collapse to exactly the layouts the fla recurrent + # kernel wants here, i.e. "b (l m) h p" is "n m h p" and "b l g n" is "n 1 g n". + # ``_preprocess`` takes the sequence-first layout that in_proj produces, so hand + # it back the (1, n, proj_dim) view; both transposes are free at l == 1. + z, VKQ, ba = self._preprocess(zVKQba.transpose(0, 1)) - # Indexed conv update into the per-request state rows (``batch_indices`` - # is None for static batching, where the cache is already in order). - VKQ = causal_conv1d_update( + query, key, value = self._prepare_qkv( VKQ, - conv_state, - rearrange(self.conv1d.weight, "d 1 w -> d w"), - self.conv1d.bias, - self.activation, + conv_state=conv_state, + is_decode=True, conv_state_indices=batch_indices, + l2_norm_in_kernel=True, ) - value, key, query = torch.split( - VKQ, - [ - self.cp.d_inner_local_tpcp * M, - self.cp.ngroups_local_tpcp * self.d_state * M, - self.cp.ngroups_local_tpcp * self.d_state, - ], - dim=-1, - ) - b, a = torch.split(ba, [self.cp.nheads_local_tpcp * M, self.cp.nheads_local_tpcp], dim=-1) - - # Reshape to the fla layout with batch=n requests, seq length 1, and the - # householder copies folded into the sequence dimension (static path, l=1). - value = rearrange(value, "n (m h p) -> n m h p", m=M, p=self.headdim).contiguous() - key = rearrange(key, "n (m g s) -> n m g s", m=M, s=self.d_state).contiguous() - query = rearrange(query, "n (g s) -> n 1 g s", s=self.d_state).contiguous() - z = rearrange(z, "n (h p) -> n 1 h p", p=self.headdim).contiguous() - beta = rearrange(b.sigmoid(), "n (m h) -> n m h", m=M).contiguous() - g = -self.cp.get_A_log().float().exp() * F.softplus(a.float() + self.cp.get_dt_bias()) - g = rearrange(g, "n h -> n 1 h").contiguous() + beta, g = self._compute_gating(ba) - if self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp > 1: - rep = self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp - query = query.repeat_interleave(rep, dim=2) - key = key.repeat_interleave(rep, dim=2) - - # Interleave the (length-1) query / decay with householder zeros so the - # recurrent kernel sees an (1 * M)-length sequence (matches static decode). - g_new = g.new_zeros(g.shape[0], g.shape[1], M, g.shape[2]) + # Interleave the (length-1) query / decay with householder zeros so the recurrent + # kernel sees an (1 * num_householder)-length sequence (matches static decode). + g_new = g.new_zeros(g.shape[0], g.shape[1], self.num_householder, g.shape[2]) g_new[:, :, 0] = g g = rearrange(g_new, "n t m h -> n (t m) h") query_new = query.new_zeros( - query.shape[0], query.shape[1], M, query.shape[2], query.shape[3] + query.shape[0], query.shape[1], self.num_householder, query.shape[2], query.shape[3] ) query_new[:, :, -1] = query query = rearrange(query_new, "n t m h d -> n (t m) h d") @@ -687,7 +940,9 @@ def ssm_decode( output_final_state=True, use_qk_l2norm_in_kernel=True, ) - core_attn_out = rearrange(core_attn_out, "n (t m) h d -> n t m h d", m=M)[ + core_attn_out = rearrange( + core_attn_out, "n (t m) h d -> n t m h d", m=self.num_householder + )[ ..., -1, :, : ].contiguous() # [n, 1, h, d] @@ -697,11 +952,11 @@ def ssm_decode( # Scatter updated states back into the cache (skips -1 padding slots). tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) - y = rearrange(core_attn_out, "n t h p -> n t (h p)").contiguous() # [n, 1, d_inner] - if self.rmsnorm: - z = rearrange(z, "n t h p -> n t (h p)").contiguous() - y = self.norm(y, z) - return y + # ``_postprocess`` returns the sequence-first layout, so transpose back to the + # batch-first [n, seq_len, d_inner] this method contracts to return; the + # transpose is free at l == 1. post_conv_ssm inside it is a no-op here: decode + # only runs at cp_size == 1. + return self._postprocess(core_attn_out, z).transpose(0, 1) def ssm_prefill( self, @@ -722,76 +977,36 @@ def ssm_prefill( seq_idx = metadata.seq_idx cu_seqlens = metadata.cu_seqlens batch_indices = metadata.batch_indices_prefill - M = self.num_householder - # l b d -> b l d - zVKQba = rearrange(zVKQba, "l b d -> b l d").contiguous() - - z, VKQ, ba = torch.split( - zVKQba, - [ - self.cp.d_inner_local_tpcp, - self.cp.d_inner_local_tpcp * M - + (M + 1) * self.cp.ngroups_local_tpcp * self.d_state, - self.cp.nheads_local_tpcp * (M + 1), - ], - dim=-1, - ) + z, VKQ, ba = self._preprocess(zVKQba) # Capture per-request final conv states (before the conv consumes the - # inputs) and write them into the prefill requests' cache rows. + # inputs) and write them into the prefill requests' cache rows. Done here + # rather than via _prepare_qkv's conv_state argument, which seeds a single + # static-batching window with F.pad instead of one row per request. conv_varlen_states = causal_conv1d_varlen_states( VKQ.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] ) tensor_masked_update(conv_state, batch_indices, conv_varlen_states) - # Maintain channels-last memory layout so causal_conv1d_fn can use seq_idx. - VKQ = VKQ.transpose(1, 2) - - seqlen = VKQ.size(2) - if causal_conv1d_fn is None: - VKQ = self.act(self.cp.conv1d(VKQ)[..., :seqlen]) - else: - VKQ = causal_conv1d_fn( - x=VKQ, - weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), - bias=self.cp.get_conv1d_bias(), - activation=self.activation, - seq_idx=seq_idx, - ) - VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() - value, key, query = torch.split( - VKQ, - [ - self.cp.d_inner_local_tpcp * M, - self.cp.ngroups_local_tpcp * self.d_state * M, - self.cp.ngroups_local_tpcp * self.d_state, - ], - dim=-1, - ) - b, a = torch.split(ba, [self.cp.nheads_local_tpcp * M, self.cp.nheads_local_tpcp], dim=-1) - - # batch = 1 packed sequence of length T; householder folded into seq. - value = rearrange(value, "b l (m h p) -> b (l m) h p", m=M, p=self.headdim).contiguous() - key = rearrange(key, "b l (m g s) -> b (l m) g s", m=M, s=self.d_state).contiguous() - query = rearrange(query, "b l (g s) -> b l g s", s=self.d_state).contiguous() - z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).contiguous() - beta = rearrange(b.sigmoid(), "b l (m h) -> b (l m) h", m=M).contiguous() - g = -self.cp.get_A_log().float().exp() * F.softplus(a.float() + self.cp.get_dt_bias()) - g = g.contiguous() + query, key, value = self._prepare_qkv(VKQ, seq_idx=seq_idx, l2_norm_in_kernel=True) - if self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp > 1: - rep = self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp - query = query.repeat_interleave(rep, dim=2) - key = key.repeat_interleave(rep, dim=2) + beta, g = self._compute_gating(ba) + # batch = 1 packed sequence of length T; householder folded into seq. The shapes + # and the GQA repeat_interleave come from _prepare_qkv and _compute_gating above + # -- inference asserts gdp_cutedsl_kernel is off, so both take their FLA branch, + # which emits exactly the layout this kernel wants. The query/key L2 norm is left + # to the kernel (l2_norm_in_kernel above), so it is applied exactly once. This + # calls the in-tree fork rather than _run_gdp_kernel because dynamic batching is + # forward-only. core_attn_out, last_recurrent_state = chunk_gated_delta_product_varlen( query, key, value, g=g, beta=beta, - num_householder=M, + num_householder=self.num_householder, cu_seqlens=cu_seqlens, initial_state=None, output_final_state=True, @@ -801,11 +1016,9 @@ def ssm_prefill( # Write per-request final SSM states into the cache for subsequent decode. tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) - y = rearrange(core_attn_out, "b l h p -> l b (h p)").contiguous() - if self.rmsnorm: - z = rearrange(z, "b l h p -> l b (h p)").contiguous() - y = self.norm(y, z) - return y + # post_conv_ssm inside _postprocess is a no-op here: dynamic inference asserts + # cp_size == 1. + return self._postprocess(core_attn_out, z) def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None): """ diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index ef70b11b852..ac523378c34 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -705,7 +705,13 @@ def _save_args_to_ctx(ctx, args): continue non_tensor_entries.append((index, arg)) - ctx.save_for_backward(*detach_variable(tuple(tensor_args))) + # Save the raw tensors (as torch.utils.checkpoint does) rather than + # detach_variable()-ed copies: detaching here creates leaf tensors that require + # grad, and autograd's SavedVariable keeps such leaves alive until backward even + # when saved-tensor hooks (e.g. fine-grained activation offload) pack them away, + # pinning the input storage on GPU for the whole forward-backward interval. + # _load_args_from_ctx() detaches the unpacked tensors before they are reused. + ctx.save_for_backward(*tensor_args) ctx._non_tensor_entries = tuple(non_tensor_entries) ctx._total_args_count = len(args) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 087c41284d1..e783a017056 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -553,7 +553,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "gdn_norm_out", "mhc". + "shared_experts", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", "mhc". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -563,12 +563,15 @@ class TransformerConfig(ModelParallelConfig): "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. + "gdp_in_proj": recompute the GatedDeltaProduct input projection and its CP gather/split + preprocessing. + "gdp_qkv": recompute the GatedDeltaProduct causal conv and QKV preparation. "mhc": recompute HyperConnection intermediate activations via CheckpointWithoutOutput + CheckpointWithoutOutputManager. Requires enable_mhc_connections=True. Cannot be used with "mlp". - "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", and "mhc" use - output-discarding checkpointing, - "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. + "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", and + "mhc" use output-discarding checkpointing, "core_attn", "mlp", "moe", and + "shared_experts" use normal checkpointing. """ #################### @@ -1249,6 +1252,17 @@ class TransformerConfig(ModelParallelConfig): ) """Controls usage of the memory efficient path for Mamba layers.""" + gdp_cutedsl_kernel: bool = False + """Whether to use the CuTeDSL kernel for the GatedDeltaProduct mixer.""" + + gdp_num_chunk_states_to_recompute: int = 2 + """Checkpoint-coarsening ratio N in [0, 64] for the CuTeDSL GatedDeltaProduct kernel. + N=0 checkpoints every chunk state for the backward pass (dense, no recompute). N>=1 + stores one checkpoint per group of N+1 chunks (1/(N+1) the checkpoint memory) and the + backward recomputes each group's N missing chunk states, trading recompute time for + activation memory monotonically in N (sweet spot N=2..3). Only honored by kernel + builds that expose the num_chunk_states_to_recompute argument.""" + mlp_chunks_for_prefill: int = 1 """The number of chunks along the sequence dimension to use for MLP computation during prefill.""" @@ -1285,7 +1299,7 @@ class TransformerConfig(ModelParallelConfig): offload_modules: Optional[list[str]] = field(default_factory=list) """The submodules to offload its input. choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", - "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp". + "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp", "gdp_qkv". "attn_norm": offload the input of the normalization in the attention part. "qkv_linear": offload the input of the qkv linear part. "core_attn": offload the input of the core attention part. @@ -1294,6 +1308,8 @@ class TransformerConfig(ModelParallelConfig): "expert_fc1": offload the input of the expert fc1 part. "moe_act": offload the input of the moe act part. "fused_group_mlp": offload the input of the whole fused grouped MLP. + "gdp_qkv": offload the input of the causal conv and QKV preparation in the + GatedDeltaProduct mixer. """ min_offloaded_tensor_size: int = 1024 * 1024 """The minimum size of the tensor to be offloaded.""" @@ -1883,6 +1899,8 @@ def __post_init__(self): "moe", "shared_experts", "gdn_norm_out", + "gdp_in_proj", + "gdp_qkv", "mhc", } invalid_modules = set(self.recompute_modules) - allowed_modules @@ -2068,6 +2086,7 @@ def __post_init__(self): "attn_norm", "mlp_norm", "qkv_linear", + "gdp_qkv", } invalid_modules = set(self.offload_modules) - allowed_modules assert not invalid_modules, ( @@ -2080,6 +2099,17 @@ def __post_init__(self): "because the input of attn_proj is the output of core_attn, " "which is needed in core_attn.backward()." ) + if ( + "gdp_qkv" in self.offload_modules + and self.recompute_granularity == "selective" + and "gdp_in_proj" in self.recompute_modules + ): + raise ValueError( + "gdp_qkv cannot be set in offload_modules together with gdp_in_proj in " + "recompute_modules, because gdp_in_proj discards the input of the causal " + "conv and rematerializes it at the start of the mixer backward, leaving " + "nothing for the gdp_qkv offload group to keep on the host." + ) if self.recompute_granularity == "selective" and "moe" in self.recompute_modules: offload_inside_moe = {"moe_act", "expert_fc1", "fused_group_mlp"} & set( self.offload_modules @@ -2110,6 +2140,13 @@ def __post_init__(self): "fused_group_mlp offloads the whole fused grouped MLP and cannot be " f"combined with expert_fc1 or moe_act. Remove: {moe_partial_offload}" ) + + if self.gdp_cutedsl_kernel: + assert 0 <= self.gdp_num_chunk_states_to_recompute <= 64, ( + "gdp_num_chunk_states_to_recompute must be in range [0, 64], got " + f"{self.gdp_num_chunk_states_to_recompute}." + ) + if self.gtp_remat_opt_in_modules: _allowed_gtp_remat_opt_in_modules = {"moe_latent_proj"} invalid = set(self.gtp_remat_opt_in_modules) - _allowed_gtp_remat_opt_in_modules diff --git a/megatron/training/training.py b/megatron/training/training.py index 1284942886d..37ee411e019 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -4526,8 +4526,9 @@ def trace_handler(p): prof.step() elif iteration == args.profile_step_start: torch.cuda.check_error(torch.cuda.cudart().cudaProfilerStart()) - nsys_nvtx_context = torch.autograd.profiler.emit_nvtx(record_shapes=args.record_shapes) - nsys_nvtx_context.__enter__() + if args.record_shapes: + nsys_nvtx_context = torch.autograd.profiler.emit_nvtx(record_shapes=args.record_shapes) + nsys_nvtx_context.__enter__() # Fault-tolerance heartbeat at the top of the loop -- uninstrumented # main-thread work that sits in the post-checkpoint gap alongside the diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index b3ea98434c4..037dbae3ee4 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -135,6 +135,8 @@ "fused_residual_rmsnorm": False, "fused_single_qkv_rope": False, "gated_linear_unit": False, + "gdp_cutedsl_kernel": False, + "gdp_num_chunk_states_to_recompute": 2, "gdp_num_householder": 3, "gtp_remat_opt_in_modules": [], "gtp_weight_remat_size": 1, diff --git a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py index dce35bdd4fa..67bf73616d2 100644 --- a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py +++ b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc +import logging import os from contextlib import nullcontext from typing import Dict, List, Optional, Tuple @@ -8,12 +9,18 @@ import pytest import torch +from megatron.core._rank_utils import safe_get_rank from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.pipeline_parallel.fine_grained_activation_offload import ChunkOffloadHandler from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface, ) +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + OffloadTensorGroup, + OffloadTensorPool, + PipelineOffloadManager, +) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig @@ -75,6 +82,284 @@ def test_chunk_offload_handler_respects_tensor_opt_out_flags(): assert not handler.tensor_need_offloading_checker(tensor) +def _make_chunk_handler_for_offload_reload(): + handler = ChunkOffloadHandler.__new__(ChunkOffloadHandler) + handler.cpu_tensor_pool = OffloadTensorPool(device="cpu", pin_memory=True) + return handler + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +def test_chunk_offload_handler_offloads_base_storage_for_covering_views(): + handler = _make_chunk_handler_for_offload_reload() + + base = torch.randn(64, 96, device="cuda") + view = base[:, :80] # non-contiguous last-dim slice covering >50% of the storage + assert not view.is_contiguous() + + state = handler.offload(view) + _, cpu_backup, _, view_meta = state + # The full flat storage is offloaded, not a gathered copy of the view. + assert cpu_backup.shape == (base.numel(),) + assert view_meta == (view.size(), view.stride(), view.storage_offset()) + + reloaded = handler.reload(state) + torch.cuda.synchronize() + assert reloaded.shape == view.shape + assert reloaded.stride() == view.stride() + assert torch.equal(reloaded, view) + # The pool backup was returned after the reload. + assert handler.cpu_tensor_pool._stats["current_in_use"] == 0 + + # Autograd hands the saved-tensor hooks a detach()-ed alias (e.g. the + # checkpoint's saved input), which has no ._base; the storage-based + # resolution must still take the base path. + detached = base[:, :80].detach() + state_d = handler.offload(detached) + assert state_d[3] is not None + reloaded_d = handler.reload(state_d) + torch.cuda.synchronize() + assert torch.equal(reloaded_d, detached) + assert handler.cpu_tensor_pool._stats["current_in_use"] == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +def test_chunk_offload_handler_gathers_low_coverage_views(): + handler = _make_chunk_handler_for_offload_reload() + + base = torch.randn(64, 96, device="cuda") + view = base[:, :8] # covers well under BASE_OFFLOAD_MIN_COVERAGE of base + + state = handler.offload(view) + _, cpu_backup, _, view_meta = state + # Low-coverage views fall back to the contiguous gather. + assert view_meta is None + assert cpu_backup.shape == view.shape + + reloaded = handler.reload(state) + torch.cuda.synchronize() + assert reloaded.is_contiguous() + assert torch.equal(reloaded, view) + + +def _make_warmup_chunk(groups: List["OffloadTensorGroup"]) -> ChunkOffloadHandler: + """Build a warmup chunk handler that already owns ``groups``.""" + handler = ChunkOffloadHandler( + min_offloaded_tensor_size=1, + cpu_tensor_pool=OffloadTensorPool(device="cpu", pin_memory=True), + ) + handler.offload_groups = list(groups) + handler._max_group_size = len(groups) + return handler + + +def _run_post_warmup_callback(chunk: ChunkOffloadHandler) -> None: + """Drive post_warmup_callback over a single hand-built chunk.""" + manager = PipelineOffloadManager.__new__(PipelineOffloadManager) + manager._is_warmup = True + manager._cached_chunks_forward = [chunk] + manager._cached_chunks_backward = [chunk] + manager._offload_margin = 0 + manager._pp_rank = 0 + manager._delta_offload_bytes_across_pp_ranks = 0 + manager._activation_offload_fraction = 1.0 + manager.post_warmup_callback() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +@pytest.mark.skipif( + Utils.world_size < 2, reason="Rank-0 duplicate aggregation requires at least two ranks." +) +def test_post_warmup_callback_warns_on_duplicated_storage_offload(caplog, monkeypatch): + from megatron.core.pipeline_parallel import fine_grained_activation_offload as off_module + + # post_warmup_callback gathers the per-rank duplicate summaries, so this test needs a + # live process group. The skipif above reads the launcher's world size rather than + # torch.distributed.get_world_size(), which raises when no group exists yet. + Utils.initialize_model_parallel() + + try: + rank = safe_get_rank() + duplicate_rank = 1 + base = torch.randn(1024, 512, device="cuda") + storage_bytes = base.numel() * base.element_size() + # Two strided views that each cover >50% of the shared storage, so both take + # the whole-storage offload path and copy the same bytes to CPU twice. Only + # rank 1 creates this duplication, exercising the rank-0 aggregation path. + views = (base[:, :400], base[:, 112:]) if rank == duplicate_rank else (base[:, :400],) + assert all(not view.is_contiguous() for view in views) + + # post_warmup_callback disables the last group of each name (offload margin), + # so build two same-named groups and assert on the surviving one. + groups = [] + for group_idx in range(2): + group = OffloadTensorGroup("core_attn") + for tensor_idx, view in enumerate(views): + group.push_tensor((group_idx, tensor_idx), view) + groups.append(group) + + chunk = _make_warmup_chunk(groups) + for group in groups: + chunk.bulk_offload_group(group) + torch.cuda.synchronize() + + # The second view of each group on rank 1 is redundant; other ranks have no + # local duplication to report. + for group in groups: + assert group.duplicate_storage_tensor_count == int(rank == duplicate_rank) + assert group.duplicate_storage_bytes == (storage_bytes if rank == duplicate_rank else 0) + + real_all_gather_object = torch.distributed.all_gather_object + all_gather_object_calls = 0 + + def counting_all_gather_object(*args, **kwargs): + nonlocal all_gather_object_calls + all_gather_object_calls += 1 + return real_all_gather_object(*args, **kwargs) + + monkeypatch.setattr(torch.distributed, "all_gather_object", counting_all_gather_object) + with caplog.at_level(logging.WARNING, logger=off_module.__name__): + _run_post_warmup_callback(chunk) + + assert all_gather_object_calls == 1 + + # The margin rule keeps only the first group offloaded. Rank 0 gathers and + # reports the duplication even though it occurred only on rank 1. + if rank == 0: + assert "copying the same tensor storage to CPU more than once" in caplog.text + assert ( + f"rank {duplicate_rank}, core_attn: 1 redundant copies, " + f"{storage_bytes / (1024 * 1024):.2f} MB" in caplog.text + ) + assert "Offloading proceeds with the duplicated copies." in caplog.text + else: + assert "copying the same tensor storage to CPU more than once" not in caplog.text + # The duplicated offloading is diagnosed, not suppressed. + assert groups[0].offload + assert groups[0].total_tensor_count == len(views) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +def test_bulk_offload_group_flags_disjoint_views_taking_the_whole_storage_path(): + base = torch.randn(1024, 512, device="cuda") + storage_bytes = base.numel() * base.element_size() + # Disjoint halves, but each covers exactly BASE_OFFLOAD_MIN_COVERAGE of the + # storage, so both are offloaded as the full buffer: the views not sharing a + # single element does not make the second copy any less redundant. + group = OffloadTensorGroup("core_attn") + group.push_tensor((1, 0), base[:, :256]) + group.push_tensor((1, 1), base[:, 256:]) + + chunk = _make_warmup_chunk([group]) + chunk.bulk_offload_group(group) + torch.cuda.synchronize() + + assert all(state[3] is not None for state in group._tensors.values()) + assert group.total_offload_bytes == 2 * storage_bytes + assert group.duplicate_storage_tensor_count == 1 + assert group.duplicate_storage_bytes == storage_bytes + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +@pytest.mark.parametrize("full_storage_first", [False, True]) +def test_bulk_offload_group_counts_full_storage_and_logical_view_duplicates(full_storage_first): + base = torch.randn(1024, 512, device="cuda") + storage_bytes = base.numel() * base.element_size() + logical_view = base[:128] + logical_view_bytes = logical_view.numel() * logical_view.element_size() + full_storage_view = base[:, :400] + tensors = ( + (full_storage_view, logical_view) + if full_storage_first + else (logical_view, full_storage_view) + ) + + group = OffloadTensorGroup("core_attn") + for tensor_idx, tensor in enumerate(tensors): + group.push_tensor((1, tensor_idx), tensor) + + chunk = _make_warmup_chunk([group]) + chunk.bulk_offload_group(group) + torch.cuda.synchronize() + + assert group.total_offload_bytes == storage_bytes + logical_view_bytes + assert group.duplicate_storage_tensor_count == 1 + assert group.duplicate_storage_bytes == logical_view_bytes + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +def test_bulk_offload_group_does_not_flag_views_copying_their_own_bytes(caplog, monkeypatch): + from megatron.core.pipeline_parallel import fine_grained_activation_offload as off_module + + base = torch.randn(1024, 512, device="cuda") + storage_bytes = base.numel() * base.element_size() + # Disjoint contiguous slices copy their own bytes, together exactly the + # storage once. + group = OffloadTensorGroup("core_attn") + group.push_tensor((1, 0), base[:512]) + group.push_tensor((1, 1), base[512:]) + # Low-coverage strided views are gathered, so they too copy only their own + # elements even though the gathered spans interleave in the storage. + other_group = OffloadTensorGroup("core_attn") + other_group.push_tensor((2, 0), base[:, :8]) + other_group.push_tensor((2, 1), base[:, 8:16]) + + chunk = _make_warmup_chunk([group, other_group]) + chunk.bulk_offload_group(group) + chunk.bulk_offload_group(other_group) + torch.cuda.synchronize() + + assert group.total_offload_bytes == storage_bytes + for offload_group in (group, other_group): + assert all(state[3] is None for state in offload_group._tensors.values()) + assert offload_group.duplicate_storage_tensor_count == 0 + assert offload_group.duplicate_storage_bytes == 0 + + monkeypatch.setattr(off_module, "print_offload_summary_table", lambda *args, **kwargs: None) + with caplog.at_level(logging.WARNING, logger=off_module.__name__): + _run_post_warmup_callback(chunk) + assert "more than once" not in caplog.text + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +def test_bulk_offload_group_does_not_analyze_logical_view_overlap(): + base = torch.randn(1024, 512, device="cuda") + half_bytes = base[:512].numel() * base.element_size() + # General overlap analysis for logical views is intentionally out of scope; + # duplicate-byte accounting is exact only when the storage is copied in full. + group = OffloadTensorGroup("core_attn") + group.push_tensor((1, 0), base[:512]) + group.push_tensor((1, 1), base[:512]) + + chunk = _make_warmup_chunk([group]) + chunk.bulk_offload_group(group) + torch.cuda.synchronize() + + assert group.total_offload_bytes == 2 * half_bytes + assert group.duplicate_storage_tensor_count == 0 + assert group.duplicate_storage_bytes == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.") +def test_duplicate_storage_accounting_is_scoped_to_one_group(): + base = torch.randn(1024, 512, device="cuda") + groups = [] + for group_idx, view in enumerate((base[:, :400], base[:, 112:])): + group = OffloadTensorGroup("core_attn") + group.push_tensor((group_idx, 0), view) + groups.append(group) + + chunk = _make_warmup_chunk(groups) + for group in groups: + chunk.bulk_offload_group(group) + torch.cuda.synchronize() + + for group in groups: + assert group.duplicate_storage_tensor_count == 0 + assert group.duplicate_storage_bytes == 0 + + def _build_gpt_model( *, seed: int, diff --git a/tests/unit_tests/ssm/test_causal_conv1d.py b/tests/unit_tests/ssm/test_causal_conv1d.py index 6818a4e2026..c51515d712b 100644 --- a/tests/unit_tests/ssm/test_causal_conv1d.py +++ b/tests/unit_tests/ssm/test_causal_conv1d.py @@ -71,3 +71,83 @@ def test_causal_conv1d_cp_matches_full_sequence(): torch.testing.assert_close(bias_local.grad, bias_ref.grad, rtol=rtol, atol=atol) finally: Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not HAVE_CAUSAL_CONV1D, reason="causal_conv1d is not installed") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_causal_conv1d_channel_contiguous_matches_sequence_contiguous(dtype): + """Verify forward and backward equivalence across causal_conv1d layouts. + + Both inputs have shape (batch, dim, seqlen) and identical values. Only + their strides differ: a unit dim stride selects the channel-last kernel, + while a unit seqlen stride selects the standard kernel. + """ + torch.manual_seed(42) + batch, seq_len, dim, width = 2, 128, 64, 4 + padding = 16 + storage_dim = dim + 2 * padding + + # Slice a larger allocation so the channel-last input is non-dense and + # has gaps between consecutive sequence elements. + channel_last_storage = torch.randn(batch, seq_len, storage_dim, device="cuda", dtype=dtype) + x_channel_contiguous = ( + channel_last_storage[:, :, padding : padding + dim] + .transpose(1, 2) + .detach() + .requires_grad_() + ) + x_sequence_contiguous = x_channel_contiguous.detach().contiguous().requires_grad_() + + assert x_channel_contiguous.shape == x_sequence_contiguous.shape == (batch, dim, seq_len) + assert torch.equal(x_channel_contiguous, x_sequence_contiguous) + assert x_channel_contiguous.stride() == (seq_len * storage_dim, 1, storage_dim) + assert x_sequence_contiguous.stride() == (dim * seq_len, seq_len, 1) + + weight = torch.randn(dim, width, device="cuda", dtype=torch.float32, requires_grad=True) + bias = torch.randn(dim, device="cuda", dtype=torch.float32, requires_grad=True) + + out_channel_contiguous = causal_conv1d_fn(x_channel_contiguous, weight, bias, activation="silu") + out_sequence_contiguous = causal_conv1d_fn( + x_sequence_contiguous, weight, bias, activation="silu" + ) + + assert out_channel_contiguous.stride() == (seq_len * dim, 1, dim) + assert out_sequence_contiguous.stride() == (dim * seq_len, seq_len, 1) + + # The forward kernels should be bitwise identical. + assert torch.equal(out_channel_contiguous, out_sequence_contiguous) + + grad_channel_contiguous = torch.randn_like(out_channel_contiguous) + grad_sequence_contiguous = grad_channel_contiguous.contiguous() + assert torch.equal(grad_channel_contiguous, grad_sequence_contiguous) + assert grad_channel_contiguous.stride() == (seq_len * dim, 1, dim) + assert grad_sequence_contiguous.stride() == (dim * seq_len, seq_len, 1) + + dx_channel, dweight_channel, dbias_channel = torch.autograd.grad( + out_channel_contiguous, + (x_channel_contiguous, weight, bias), + grad_outputs=grad_channel_contiguous, + ) + dx_sequence, dweight_sequence, dbias_sequence = torch.autograd.grad( + out_sequence_contiguous, + (x_sequence_contiguous, weight, bias), + grad_outputs=grad_sequence_contiguous, + ) + + assert dx_channel.stride() == (seq_len * dim, 1, dim) + assert dx_sequence.stride() == (dim * seq_len, seq_len, 1) + + # dx is returned in the input dtype, so BF16 needs looser tolerances. Weight + # and bias gradients are accumulated in FP32; their small differences come + # from the kernels using different parallel reduction orders. + input_grad_rtol, input_grad_atol = (3e-4, 1e-3) if dtype == torch.float32 else (1e-2, 5e-2) + param_grad_rtol, param_grad_atol = 1e-3, 1e-3 + + torch.testing.assert_close(dx_channel, dx_sequence, rtol=input_grad_rtol, atol=input_grad_atol) + torch.testing.assert_close( + dweight_channel, dweight_sequence, rtol=param_grad_rtol, atol=param_grad_atol + ) + torch.testing.assert_close( + dbias_channel, dbias_sequence, rtol=param_grad_rtol, atol=param_grad_atol + ) diff --git a/tests/unit_tests/ssm/test_gated_delta_product.py b/tests/unit_tests/ssm/test_gated_delta_product.py new file mode 100644 index 00000000000..fa562d9bc50 --- /dev/null +++ b/tests/unit_tests/ssm/test_gated_delta_product.py @@ -0,0 +1,220 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Parity tests for GatedDeltaProduct selective recompute. + +`gdp_in_proj` recompute is meant to be an exact replay of the input projection and its +preprocessing: the forward pass is unaffected, and the backward pass rebuilds the same +subgraph from the same saved input. So enabling it must not move the numerics at all. +Any drift would mean the replay observed a different weight, a different RNG state, or +a stale storage after the output-discarding checkpoint shared the recomputed one back. + +The one exception is `conv1d.weight`, see NONDETERMINISTIC_PARAMS below: its gradient is +not reproducible run to run, for reasons that have nothing to do with recompute. +""" + +import pytest +import torch + +from megatron.core.fp8_utils import get_fp8_context +from megatron.core.models.hybrid.hybrid_layer_specs import gdp_stack_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_product import ( + HAVE_CUTEDSL_GDP, + HAVE_EINOPS, + HAVE_FLA, + HAVE_MAMBA_SSM, + GatedDeltaProductMixer, + causal_conv1d_fn, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from transformer_engine.pytorch.fp8 import check_fp8_support + + fp8_available, reason_for_no_fp8 = check_fp8_support() +except ImportError: + fp8_available, reason_for_no_fp8 = False, "transformer-engine is not installed" + +HAVE_GDP_DEPS = ( + HAVE_MAMBA_SSM + and HAVE_EINOPS + and (HAVE_FLA or HAVE_CUTEDSL_GDP) + and causal_conv1d_fn is not None +) + +# The mixer places several divisibility constraints on these: d_inner = num_heads * +# head_dim must divide evenly by the group count, and the in_proj output width +# (d_inner * 4 + num_groups * state_dim * 4 + num_heads * 4 = 2064) must stay a +# multiple of 16 so the FP8 GEMM is legal. +HIDDEN_SIZE = 256 +MAMBA_NUM_HEADS = 4 +MAMBA_HEAD_DIM = 64 +MAMBA_NUM_GROUPS = 2 +MAMBA_STATE_DIM = 128 + + +def _make_config(recompute_modules, fp8_recipe=None): + """Build a GDP config; an empty recompute_modules disables selective recompute.""" + return TransformerConfig( + num_layers=1, + hidden_size=HIDDEN_SIZE, + num_attention_heads=MAMBA_NUM_HEADS, + mamba_num_heads=MAMBA_NUM_HEADS, + mamba_head_dim=MAMBA_HEAD_DIM, + mamba_num_groups=MAMBA_NUM_GROUPS, + mamba_state_dim=MAMBA_STATE_DIM, + params_dtype=torch.bfloat16, + bf16=True, + recompute_granularity="selective" if recompute_modules else None, + recompute_modules=list(recompute_modules), + fp8="hybrid" if fp8_recipe else None, + fp8_recipe=fp8_recipe if fp8_recipe else "delayed", + ) + + +def _build_mixer(config, state_dict=None): + """Build a GatedDeltaProductMixer, optionally loading a shared state dict.""" + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + mixer = GatedDeltaProductMixer( + config, + gdp_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + config.hidden_size, + layer_number=1, + pg_collection=pg_collection, + ).cuda() + if state_dict is not None: + mixer.load_state_dict(state_dict) + return mixer + + +def _run(mixer, hidden_states, grads): + """Run one forward/backward and collect the output and every gradient.""" + x = hidden_states.clone().requires_grad_(True) + with get_fp8_context(mixer.config): + out, _ = mixer(x) + out.backward(grads) + param_grads = {name: param.grad.clone() for name, param in mixer.named_parameters()} + return out.detach().clone(), x.grad.clone(), param_grads + + +# `causal_conv1d_fn` accumulates its weight gradient across blocks with atomicAdd on the +# channels-last layout the mixer feeds it, so the reduction order -- and with it the +# last-place rounding -- varies from launch to launch. Two runs of the *same* mixer can +# therefore disagree on this tensor. It is unrelated to recompute and to FP8: it +# reproduces with recompute disabled, in plain BF16, and the conv sees the identical +# layout and dtypes either way. Everything else here must still match bitwise, so scope +# the tolerance to this one tensor rather than loosening the whole comparison. +NONDETERMINISTIC_PARAMS = ("conv1d.weight",) + + +def _assert_grad_matches(name, recomputed_grad, base_grad): + """Compare one gradient: bitwise, unless the kernel that produced it is unstable.""" + if name in NONDETERMINISTIC_PARAMS: + # Reordering shifts an element by a few ulps of its own value, so the budget is + # relative: BF16 carries 8 mantissa bits, so 2**-6 is four ulps. atol only has to + # cover elements too close to zero for rtol to mean anything, and is derived from + # the tensor's magnitude so it cannot go stale if these tests change shape. + atol = float(base_grad.abs().max()) * 2**-11 + torch.testing.assert_close( + recomputed_grad, base_grad, rtol=2**-6, atol=atol, msg=f"{name} gradient drifted" + ) + else: + torch.testing.assert_close( + recomputed_grad, base_grad, rtol=0, atol=0, msg=f"{name} gradient drifted" + ) + + +def _assert_replay_matches(baseline, recomputed): + """Compare a (output, input grad, param grads) triple; exact but for the conv weight.""" + base_out, base_input_grad, base_param_grads = baseline + recomp_out, recomp_input_grad, recomp_param_grads = recomputed + + torch.testing.assert_close(recomp_out, base_out, rtol=0, atol=0, msg="output drifted") + torch.testing.assert_close( + recomp_input_grad, base_input_grad, rtol=0, atol=0, msg="input gradient drifted" + ) + assert set(recomp_param_grads) == set(base_param_grads) + for name, base_grad in base_param_grads.items(): + _assert_grad_matches(name, recomp_param_grads[name], base_grad) + + +@pytest.mark.skipif(not HAVE_GDP_DEPS, reason="requires mamba-ssm, einops, causal-conv1d and FLA") +@pytest.mark.parametrize("recompute_qkv", [False, True]) +def test_in_proj_recompute_parity(recompute_qkv): + """`gdp_in_proj` recompute must be an exact replay, alone and with `gdp_qkv`. + + The two are stacked deliberately: `gdp_qkv` saves VKQ, which `gdp_in_proj` discards, + so this pins the ordering contract between the two recompute hooks -- the in_proj + replay on the mixer output has to rematerialize VKQ before the QKV replay on + core_attn_out consumes it. + """ + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1) + model_parallel_cuda_manual_seed(123) + try: + modules = ["gdp_in_proj"] + (["gdp_qkv"] if recompute_qkv else []) + baseline = _build_mixer(_make_config([])) + recomputed = _build_mixer(_make_config(modules), state_dict=baseline.state_dict()) + assert recomputed.recompute_in_proj and not baseline.recompute_in_proj + assert recomputed.recompute_qkv == recompute_qkv + + sequence_length, micro_batch_size = 256, 2 + hidden_states = torch.randn( + (sequence_length, micro_batch_size, HIDDEN_SIZE), dtype=torch.bfloat16, device="cuda" + ) + grads = torch.randn_like(hidden_states) + + _assert_replay_matches( + _run(baseline, hidden_states, grads), _run(recomputed, hidden_states, grads) + ) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not HAVE_GDP_DEPS, reason="requires mamba-ssm, einops, causal-conv1d and FLA") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("fp8_recipe", ["tensorwise"]) +def test_fp8_in_proj_recompute_parity(fp8_recipe): + """The replay must stay exact under FP8, where in_proj is the quantized op. + + CheckpointWithoutOutput records the forward recipe and amax state and replays under + the recorded fp8_autocast, so the replayed GEMM has to see the same scales as the + original. Drift here would mean the recompute re-derived them from an amax history + that had already moved on. + """ + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1) + model_parallel_cuda_manual_seed(123) + try: + baseline = _build_mixer(_make_config([], fp8_recipe=fp8_recipe)) + recomputed = _build_mixer( + _make_config(["gdp_in_proj"], fp8_recipe=fp8_recipe), state_dict=baseline.state_dict() + ) + assert recomputed.recompute_in_proj and not baseline.recompute_in_proj + + # FP8 GEMMs need the token dimension aligned; keep it a multiple of 128. + sequence_length, micro_batch_size = 256, 2 + hidden_states = torch.randn( + (sequence_length, micro_batch_size, HIDDEN_SIZE), dtype=torch.bfloat16, device="cuda" + ) + grads = torch.randn_like(hidden_states) + + _assert_replay_matches( + _run(baseline, hidden_states, grads), _run(recomputed, hidden_states, grads) + ) + finally: + Utils.destroy_model_parallel() + + +def test_in_proj_recompute_rejects_qkv_offload(): + """gdp_in_proj leaves nothing for the gdp_qkv offload group, so the pair is rejected.""" + with pytest.raises(ValueError, match="gdp_qkv cannot be set in offload_modules"): + TransformerConfig( + num_layers=1, + hidden_size=HIDDEN_SIZE, + num_attention_heads=MAMBA_NUM_HEADS, + recompute_granularity="selective", + recompute_modules=["gdp_in_proj"], + fine_grained_activation_offloading=True, + offload_modules=["gdp_qkv"], + ) From c9676f41e257ae6ab2dee4a31be6567e108fc03e Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 18 Aug 2026 12:49:50 +0000 Subject: [PATCH 076/101] RL: compress per-token staleness metrics (#6376) Signed-off-by: Teodor-Dumitru Ene --- megatron/rl/rl_utils.py | 233 +++++++++++++++++++++------ tests/unit_tests/rl/test_rl_utils.py | 168 ++++++++++++++++--- 2 files changed, 336 insertions(+), 65 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 08e5b7b579e..754441f5a50 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -16,7 +16,7 @@ from dataclasses import dataclass from functools import partial from pathlib import Path -from typing import Any, AsyncIterator, Dict, Iterator, List, Optional +from typing import Any, AsyncIterator, Dict, Iterator, List, NamedTuple, Optional, TypeAlias import numpy as np import torch @@ -287,6 +287,38 @@ def verify_model_weights_swap( inf_core.train() +class EpochBoundary(NamedTuple): + """RLE boundary: `epoch` applies from `start_token_index` until the next boundary.""" + + start_token_index: int + epoch: int + + +RolloutEpochBoundaries: TypeAlias = list[list[EpochBoundary]] +"""Per-turn lists of epoch boundaries over the turn's cumulative token sequence.""" + + +class EpochSegment(NamedTuple): + """A run of `token_count` consecutive tokens sharing one epoch.""" + + epoch: int + token_count: int + + +class AlignedEpochSegment(NamedTuple): + """A run of `token_count` tokens whose policy and KV-cache epochs are both constant.""" + + policy_epoch: int + kv_cache_epoch: int + token_count: int + + +RolloutEpochSegments: TypeAlias = list[EpochSegment] +"""One rollout's (epoch, token_count) segments, covering its trajectory in order.""" + +GroupedEpochSegments: TypeAlias = list[list[RolloutEpochSegments]] +"""Per-group, per-rollout epoch segments.""" + @dataclass(slots=True) class RolloutStats: @@ -305,8 +337,8 @@ class RolloutStats: min_inf_prob: None | float max_inf_prob: None | float mean_inf_prob: None | float - policy_epoch: list[list[int]] - kv_cache_epoch: list[list[int]] + policy_epoch_segments: GroupedEpochSegments + kv_cache_epoch_segments: GroupedEpochSegments completed_epochs: list[list[int]] num_evictions: list[list[int]] rollout_statuses: list[list[str]] @@ -1100,6 +1132,49 @@ def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[i return ((rewards - reward_means) / (1e-4 + reward_stds)).tolist() +def expand_epoch_segments( + per_turn_boundaries: RolloutEpochBoundaries, + per_turn_token_counts: list[int], +) -> RolloutEpochSegments: + """Expand RLE (start_token_index, epoch) boundaries into (epoch, token_count) segments.""" + segments: RolloutEpochSegments = [] + for boundaries, total_len in zip(per_turn_boundaries, per_turn_token_counts, strict=True): + if not boundaries: + continue + for idx, (start, epoch) in enumerate(boundaries): + end = boundaries[idx + 1][0] if idx + 1 < len(boundaries) else total_len + count = end - start + if count > 0: + segments.append(EpochSegment(epoch=epoch, token_count=count)) + return segments + + +def merge_epoch_segments( + policy_segments: RolloutEpochSegments, kv_segments: RolloutEpochSegments +) -> Iterator[AlignedEpochSegment]: + """Yield (policy_epoch, kv_cache_epoch, token_count) runs over shared token positions.""" + pol_idx = kv_idx = 0 + pol_left = policy_segments[0].token_count if policy_segments else 0 + kv_left = kv_segments[0].token_count if kv_segments else 0 + while pol_idx < len(policy_segments) and kv_idx < len(kv_segments): + take = min(pol_left, kv_left) + yield AlignedEpochSegment( + policy_epoch=policy_segments[pol_idx].epoch, + kv_cache_epoch=kv_segments[kv_idx].epoch, + token_count=take, + ) + pol_left -= take + kv_left -= take + if pol_left == 0: + pol_idx += 1 + if pol_idx < len(policy_segments): + pol_left = policy_segments[pol_idx].token_count + if kv_left == 0: + kv_idx += 1 + if kv_idx < len(kv_segments): + kv_left = kv_segments[kv_idx].token_count + + def compute_group_stats( rollouts: GroupedRollouts, tokenizer: MegatronTokenizer, @@ -1126,8 +1201,8 @@ def compute_group_stats( env_ids = [] group_reward_ids = [] num_turns = [] # num_turns per traj - all_policy_epoch = [] - all_kv_cache_epoch = [] + all_policy_epoch_segments = [] + all_kv_cache_epoch_segments = [] all_completed_epochs = [] all_num_evictions = [] all_rollout_statuses = [] @@ -1137,8 +1212,8 @@ def compute_group_stats( group_traj_lengths = [] group_turn_lengths = [] group_num_turns = [] - group_policy_epoch = [] - group_kv_epoch = [] + group_policy_epoch_segments = [] + group_kv_epoch_segments = [] group_completed_epochs = [] group_num_evictions = [] group_rollout_statuses = [] @@ -1186,12 +1261,25 @@ def compute_group_stats( if isinstance(rollout, TokenRollout) else [] ) - group_policy_epoch.append( - [epoch for record in turn_records for _, epoch in record.policy_epoch] - ) - group_kv_epoch.append( - [epoch for record in turn_records for _, epoch in record.kv_cache_epoch] - ) + for record in turn_records: + assert record.policy_epoch, "Request record has no policy_epoch data" + assert record.kv_cache_epoch, "Request record has no kv_cache_epoch data" + if turn_records: + cumulative_turn_lens = [len(t) for t in rollout.trajectory] + group_policy_epoch_segments.append( + expand_epoch_segments( + [record.policy_epoch for record in turn_records], cumulative_turn_lens + ) + ) + group_kv_epoch_segments.append( + expand_epoch_segments( + [record.kv_cache_epoch for record in turn_records], cumulative_turn_lens + ) + ) + else: + # Unjoined rollouts (placeholders, text rollouts) pop no records: empty rows. + group_policy_epoch_segments.append([]) + group_kv_epoch_segments.append([]) group_completed_epochs.extend(record.policy_epoch[-1][1] for record in turn_records) group_num_evictions.append(sum(record.num_evictions for record in turn_records)) rollout_status = rollout.rollout_status @@ -1200,8 +1288,8 @@ def compute_group_stats( rollout_status = 'placeholder' group_rollout_statuses.append(rollout_status) group_failure_reasons.append(rollout.failure_reason) - all_policy_epoch.append(group_policy_epoch) - all_kv_cache_epoch.append(group_kv_epoch) + all_policy_epoch_segments.append(group_policy_epoch_segments) + all_kv_cache_epoch_segments.append(group_kv_epoch_segments) all_completed_epochs.append(group_completed_epochs) all_num_evictions.append(group_num_evictions) all_rollout_statuses.append(group_rollout_statuses) @@ -1233,8 +1321,8 @@ def compute_group_stats( min_inf_prob=None, max_inf_prob=None, mean_inf_prob=None, - policy_epoch=all_policy_epoch, - kv_cache_epoch=all_kv_cache_epoch, + policy_epoch_segments=all_policy_epoch_segments, + kv_cache_epoch_segments=all_kv_cache_epoch_segments, completed_epochs=all_completed_epochs, num_evictions=all_num_evictions, rollout_statuses=all_rollout_statuses, @@ -1273,8 +1361,8 @@ def prep_wandb_metrics( rewards: List[List[float]], num_turns: List[List[int]], advantages: List[float], - policy_epoch: List[List[List[int]]], - kv_cache_epoch: List[List[List[int]]], + policy_epoch_segments: GroupedEpochSegments, + kv_cache_epoch_segments: GroupedEpochSegments, completed_epochs: List[List[int]], num_evictions: List[List[int]], current_iteration: int, @@ -1296,8 +1384,8 @@ def prep_wandb_metrics( rewards: Grouped list of rewards. num_turns: Grouped list of number of turns in the trajectories. Zero means failure. advantages: Flattened list of advantages. - policy_epoch: Grouped list of per-token policy epoch stamps. - kv_cache_epoch: Grouped list of per-token KV cache epoch stamps. + policy_epoch_segments: Grouped list of per-rollout (epoch, token_count) segments. + kv_cache_epoch_segments: Grouped list of per-rollout (epoch, token_count) segments. completed_epochs: Grouped list of per-turn max policy epoch stamps. num_evictions: Grouped list of per-rollout number of evictions. current_iteration: Current training iteration. @@ -1382,8 +1470,8 @@ def _real(grouped): flat_traj_lens = [l for g in traj_lens for l in g] flat_num_evictions = [e for g in num_evictions for e in g] flat_statuses = [s for g in rollout_statuses for s in g] - flat_policy_epochs = [r for g in policy_epoch for r in g] - flat_kv_epochs = [r for g in kv_cache_epoch for r in g] + flat_policy_epochs = [r for g in policy_epoch_segments for r in g] + flat_kv_epochs = [r for g in kv_cache_epoch_segments for r in g] joined = [i for i, row in enumerate(flat_policy_epochs) if row] joined_rewards = [flat_rewards[i] for i in joined] joined_traj_lens = [flat_traj_lens[i] for i in joined] @@ -1391,15 +1479,58 @@ def _real(grouped): joined_statuses = [flat_statuses[i] for i in joined] per_rollout_policy_epochs = [flat_policy_epochs[i] for i in joined] per_rollout_kv_epochs = [flat_kv_epochs[i] for i in joined] - # Per-rollout staleness (oldest token) - rollout_policy_staleness = [current_iteration - r[0] for r in per_rollout_policy_epochs] - rollout_kv_staleness = [current_iteration - r[0] for r in per_rollout_kv_epochs] + # Per-rollout staleness (oldest token). Epoch rows are (epoch, token_count) segments; + # a segment's first/last epoch is its first/last token's epoch. + rollout_policy_staleness = [current_iteration - r[0].epoch for r in per_rollout_policy_epochs] + rollout_kv_staleness = [current_iteration - r[0].epoch for r in per_rollout_kv_epochs] # Per-rollout staleness (newest token) - rollout_policy_last_token_staleness = [current_iteration - r[-1] for r in per_rollout_policy_epochs] - rollout_kv_last_token_staleness = [current_iteration - r[-1] for r in per_rollout_kv_epochs] - # Per-token staleness - per_token_policy_staleness = [current_iteration - e for r in per_rollout_policy_epochs for e in r] - per_token_kv_staleness = [current_iteration - e for r in per_rollout_kv_epochs for e in r] + rollout_policy_last_token_staleness = [ + current_iteration - r[-1].epoch for r in per_rollout_policy_epochs + ] + rollout_kv_last_token_staleness = [ + current_iteration - r[-1].epoch for r in per_rollout_kv_epochs + ] + # Exact token-weighted per-rollout average staleness. + rollout_policy_avg_staleness = [ + current_iteration + - sum(s.epoch * s.token_count for s in r) / sum(s.token_count for s in r) + for r in per_rollout_policy_epochs + ] + rollout_kv_avg_staleness = [ + current_iteration + - sum(s.epoch * s.token_count for s in r) / sum(s.token_count for s in r) + for r in per_rollout_kv_epochs + ] + # Token-weighted within-rollout staleness dispersion. + rollout_policy_staleness_std = [ + ( + sum(s.token_count * (current_iteration - s.epoch - m) ** 2 for s in r) + / sum(s.token_count for s in r) + ) + ** 0.5 + for r, m in zip(per_rollout_policy_epochs, rollout_policy_avg_staleness) + ] + rollout_kv_staleness_std = [ + ( + sum(s.token_count * (current_iteration - s.epoch - m) ** 2 for s in r) + / sum(s.token_count for s in r) + ) + ** 0.5 + for r, m in zip(per_rollout_kv_epochs, rollout_kv_avg_staleness) + ] + # Per-token staleness as (rollout, policy, kv, token_count) rows. + per_token_staleness_rows = [ + ( + rollout_idx, + current_iteration - aligned.policy_epoch, + current_iteration - aligned.kv_cache_epoch, + aligned.token_count, + ) + for rollout_idx, (pol_row, kv_row) in enumerate( + zip(per_rollout_policy_epochs, per_rollout_kv_epochs, strict=True) + ) + for aligned in merge_epoch_segments(pol_row, kv_row) + ] metrics = { 'group_means_hist': wandb_writer.plot.histogram( @@ -1448,6 +1579,8 @@ def _real(grouped): 'reward', 'traj_length', 'num_evictions', 'policy_staleness', 'kv_staleness', 'policy_last_token_staleness', 'kv_last_token_staleness', + 'policy_avg_staleness', 'kv_avg_staleness', + 'policy_staleness_std', 'kv_staleness_std', 'rollout_status', ], data=list(zip( @@ -1458,13 +1591,18 @@ def _real(grouped): rollout_kv_staleness, rollout_policy_last_token_staleness, rollout_kv_last_token_staleness, + rollout_policy_avg_staleness, + rollout_kv_avg_staleness, + rollout_policy_staleness_std, + rollout_kv_staleness_std, joined_statuses, )), ), - # NOTE: This table can get very large (one row per token across all rollouts). + # NOTE: rows are not individual tokens, but instead compressed data. + # The same information is contained, but it takes less space. 'per_token_table': wandb_writer.Table( - columns=['policy_staleness', 'kv_staleness'], - data=list(zip(per_token_policy_staleness, per_token_kv_staleness)), + columns=['rollout', 'policy_staleness', 'kv_staleness', 'token_count'], + data=per_token_staleness_rows, ), 'mean_policy_staleness': np.mean(rollout_policy_staleness), 'max_policy_staleness': max(rollout_policy_staleness), @@ -1481,13 +1619,17 @@ def _real(grouped): 'total_eviction_count': sum(joined_num_evictions), 'max_num_evictions': max(joined_num_evictions), 'mean_completion_gap': np.mean([current_iteration - s for g in completed_epochs for s in g]), - 'per_token_policy_staleness_hist': wandb_writer.plot.histogram( - wandb_writer.Table(columns=['staleness'], data=[[s] for s in per_token_policy_staleness]), - 'staleness', 'Per-Token Policy Staleness' + 'rollout_avg_policy_staleness_hist': wandb_writer.plot.histogram( + wandb_writer.Table( + columns=['staleness'], data=[[s] for s in rollout_policy_avg_staleness] + ), + 'staleness', 'Per-Rollout Token-Weighted Avg Policy Staleness' ), - 'per_token_kv_staleness_hist': wandb_writer.plot.histogram( - wandb_writer.Table(columns=['staleness'], data=[[s] for s in per_token_kv_staleness]), - 'staleness', 'Per-Token KV Cache Staleness' + 'rollout_avg_kv_staleness_hist': wandb_writer.plot.histogram( + wandb_writer.Table( + columns=['staleness'], data=[[s] for s in rollout_kv_avg_staleness] + ), + 'staleness', 'Per-Rollout Token-Weighted Avg KV Cache Staleness' ), } @@ -1662,8 +1804,8 @@ def maybe_log_training_metrics( rewards = group_stats.rewards num_turns = group_stats.num_turns advantages = group_stats.advantages - policy_epoch = group_stats.policy_epoch - kv_cache_epoch = group_stats.kv_cache_epoch + policy_epoch_segments = group_stats.policy_epoch_segments + kv_cache_epoch_segments = group_stats.kv_cache_epoch_segments completed_epochs = group_stats.completed_epochs num_evictions = group_stats.num_evictions rollout_statuses = group_stats.rollout_statuses @@ -1671,7 +1813,8 @@ def maybe_log_training_metrics( metrics = metrics | prep_wandb_metrics(wandb_writer=wandb_writer, traj_lens=traj_lens, turn_lens=turn_lens, rewards=rewards, num_turns=num_turns, advantages=advantages, - policy_epoch=policy_epoch, kv_cache_epoch=kv_cache_epoch, completed_epochs=completed_epochs, + policy_epoch_segments=policy_epoch_segments, + kv_cache_epoch_segments=kv_cache_epoch_segments, completed_epochs=completed_epochs, num_evictions=num_evictions, current_iteration=current_iteration, rollout_statuses=rollout_statuses, failure_reasons=failure_reasons) @@ -1694,8 +1837,8 @@ def maybe_log_training_metrics( rewards=env_stats(rewards, env_idx), num_turns=env_stats(num_turns, env_idx), advantages=env_advantages, - policy_epoch=env_stats(policy_epoch, env_idx), - kv_cache_epoch=env_stats(kv_cache_epoch, env_idx), + policy_epoch_segments=env_stats(policy_epoch_segments, env_idx), + kv_cache_epoch_segments=env_stats(kv_cache_epoch_segments, env_idx), completed_epochs=env_stats(completed_epochs, env_idx), num_evictions=env_stats(num_evictions, env_idx), current_iteration=current_iteration, diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 52567a99901..69fd60dd765 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -1414,9 +1414,16 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders rewards = [[1, 1], [-1, 2]] num_turns = [[42, 2], [10, 8]] advantages = [0, 1] - # Per-token epoch stamps, grouped by group then rollout - policy_epoch = [[[4, 5], [2, 3]], [[5], [0, 1]]] - kv_cache_epoch = [[[4, 5], [3, 4]], [[5], [1, 2]]] + # Per-rollout (epoch, token_count) segments, grouped by group then rollout + seg = rl_utils.EpochSegment + policy_epoch_segments = [ + [[seg(4, 1), seg(5, 2)], [seg(2, 2), seg(3, 1)]], + [[seg(5, 1)], [seg(0, 3), seg(1, 1)]], + ] + kv_cache_epoch_segments = [ + [[seg(4, 2), seg(5, 1)], [seg(3, 1), seg(4, 2)]], + [[seg(5, 1)], [seg(1, 2), seg(2, 2)]], + ] # Per-turn max epoch stamps (when each turn completed) completed_epochs = [[5, 3], [5, 1]] num_evictions = [[0, 1], [0, 0]] @@ -1434,7 +1441,7 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders for group in lst: group.append(sentinel) lst.append([sentinel, sentinel]) # the fully failed extra group - for lst in (policy_epoch, kv_cache_epoch): + for lst in (policy_epoch_segments, kv_cache_epoch_segments): for group in lst: # Placeholders pop no ledger records, so their epoch rows are empty. group.append([]) @@ -1459,8 +1466,8 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders rewards, num_turns, advantages, - policy_epoch=policy_epoch, - kv_cache_epoch=kv_cache_epoch, + policy_epoch_segments=policy_epoch_segments, + kv_cache_epoch_segments=kv_cache_epoch_segments, completed_epochs=completed_epochs, num_evictions=num_evictions, current_iteration=current_iteration, @@ -1519,6 +1526,9 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders assert len(rollout_table_calls) == 1 rows = rollout_table_calls[0].kwargs["data"] assert [r[3] for r in rows] == [2, 4, 1, 6] # policy_staleness column + # policy_avg_staleness column is token-weighted: rollout 4 covers epochs + # (0 x3 tokens, 1 x1 token) -> 6 - 1/4 = 5.75. + assert rows[3][7] == 5.75 assert metrics["nonzero_groups_ratio"] == 0.5 assert metrics["max_traj_length"] == 3 assert metrics["min_traj_length"] == 1 @@ -1562,8 +1572,8 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders rewards=[], num_turns=[], advantages=[], - policy_epoch=[], - kv_cache_epoch=[], + policy_epoch_segments=[], + kv_cache_epoch_segments=[], completed_epochs=[], num_evictions=[], current_iteration=6, @@ -1588,8 +1598,8 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders rewards=[[0.0, 0.0]], num_turns=[[0, 0]], advantages=[0.0, 0.0], - policy_epoch=[[[0], [0]]], - kv_cache_epoch=[[[0], [0]]], + policy_epoch_segments=[[[], []]], + kv_cache_epoch_segments=[[[], []]], completed_epochs=[[]], num_evictions=[[0, 0]], current_iteration=6, @@ -1658,7 +1668,8 @@ def placeholder(rollout_status="ok", failure_reason=None): # and num_turns is what lets prep_wandb_metrics mask them downstream. # Placeholders join no records, so their epoch rows are empty. assert stats.num_turns == [[1, 1, 0], [0, 0, 0]] - assert stats.policy_epoch == [[[5], [6], []], [[], [], []]] + seg = rl_utils.EpochSegment + assert stats.policy_epoch_segments == [[[seg(5, 3)], [seg(6, 4)], []], [[], [], []]] assert stats.traj_lens == [[3, 4, 0], [0, 0, 0]] # Per-turn lists exclude placeholders entirely: no sentinel epoch-0 stamp # in completed_epochs, no fake 0-length turn for all-placeholder groups. @@ -1689,8 +1700,8 @@ def placeholder(rollout_status="ok", failure_reason=None): stats.rewards, stats.num_turns, stats.advantages, - policy_epoch=stats.policy_epoch, - kv_cache_epoch=stats.kv_cache_epoch, + policy_epoch_segments=stats.policy_epoch_segments, + kv_cache_epoch_segments=stats.kv_cache_epoch_segments, completed_epochs=stats.completed_epochs, num_evictions=stats.num_evictions, current_iteration=7, @@ -1777,8 +1788,14 @@ def token_rollout(turns, masks, completion_ids=None): ] stats = rl_utils.compute_group_stats(rollouts, MockTokenizer(), 8, ledger) - assert stats.policy_epoch == [[[7, 8], [5], []], [[6], [9], []]] - assert stats.kv_cache_epoch == [[[7, 8], [5], []], [[6], [9], []]] + assert stats.policy_epoch_segments == [ + [[(7, 4), (8, 7)], [(5, 3)], []], + [[(6, 3)], [(9, 3)], []], + ] + assert stats.kv_cache_epoch_segments == [ + [[(7, 4), (8, 7)], [(5, 3)], []], + [[(6, 3)], [(9, 3)], []], + ] assert stats.completed_epochs == [[7, 8, 5], [6, 9]] assert stats.num_evictions == [[1, 0, 0], [0, 0, 0]] assert not ledger # all records consumed @@ -1791,6 +1808,7 @@ def token_rollout(turns, masks, completion_ids=None): # Text (non-token) rollouts join no finished-request records and yield # empty epoch rows; staleness/eviction telemetry filters them out while # keeping the rollout table row-aligned. + seg = rl_utils.EpochSegment writer = MagicMock() metrics = rl_utils.prep_wandb_metrics( writer, @@ -1799,8 +1817,8 @@ def token_rollout(turns, masks, completion_ids=None): rewards=[[1.0, -1.0]], num_turns=[[1, 1]], advantages=[0.5], - policy_epoch=[[[4, 5], []]], - kv_cache_epoch=[[[4, 5], []]], + policy_epoch_segments=[[[seg(4, 1), seg(5, 1)], []]], + kv_cache_epoch_segments=[[[seg(4, 1), seg(5, 1)], []]], completed_epochs=[[5]], num_evictions=[[2, 0]], current_iteration=6, @@ -1813,7 +1831,21 @@ def token_rollout(turns, masks, completion_ids=None): rollout_table = next( c for c in writer.Table.call_args_list if "traj_length" in c.kwargs.get("columns", []) ) - assert rollout_table.kwargs["data"] == [(1.0, 3, 2, 2, 2, 1, 1, 'ok')] + (row,) = rollout_table.kwargs["data"] + assert dict(zip(rollout_table.kwargs["columns"], row)) == { + "reward": 1.0, + "traj_length": 3, + "num_evictions": 2, + "policy_staleness": 2, + "kv_staleness": 2, + "policy_last_token_staleness": 1, + "kv_last_token_staleness": 1, + "policy_avg_staleness": 1.5, + "kv_avg_staleness": 1.5, + "policy_staleness_std": 0.5, + "kv_staleness_std": 0.5, + "rollout_status": 'ok', + } # A window where nothing joined emits no staleness/eviction metrics at all. metrics = rl_utils.prep_wandb_metrics( @@ -1823,8 +1855,8 @@ def token_rollout(turns, masks, completion_ids=None): rewards=[[1.0]], num_turns=[[1]], advantages=[0.5], - policy_epoch=[[[]]], - kv_cache_epoch=[[[]]], + policy_epoch_segments=[[[]]], + kv_cache_epoch_segments=[[[]]], completed_epochs=[[]], num_evictions=[[0]], current_iteration=6, @@ -1835,6 +1867,102 @@ def token_rollout(turns, masks, completion_ids=None): assert "rollout_table" not in metrics assert "per_token_table" not in metrics + def test_epoch_segment_helpers(self): + expand, merge = rl_utils.expand_epoch_segments, rl_utils.merge_epoch_segments + seg, aligned = rl_utils.EpochSegment, rl_utils.AlignedEpochSegment + # expand: boundaries -> (epoch, token_count). Multi-turn boundaries index + # each turn's full CUMULATIVE sequence; boundary-less turns contribute nothing. + assert expand([[(0, 3), (4, 4)]], [6]) == [seg(3, 4), seg(4, 2)] + assert expand([[(0, 5)], [(0, 5), (4, 6)]], [4, 7]) == [seg(5, 4), seg(5, 4), seg(6, 3)] + assert expand([[], [(0, 2)]], [3, 5]) == [seg(2, 5)] + # Per-turn lists of mismatched lengths raise instead of silently dropping turns. + with pytest.raises(ValueError): + expand([[(0, 9)]], []) + # merge: (policy, kv, token_count) runs aligned by token position, not an + # index-wise zip of the two segment lists; a tail covered by only one + # stream (disagreeing totals) is dropped rather than misattributed. + assert list(merge([seg(10, 4)], [seg(9, 2), seg(8, 2)])) == [ + aligned(10, 9, 2), + aligned(10, 8, 2), + ] + assert list(merge([seg(7, 3)], [seg(6, 1)])) == [aligned(7, 6, 1)] + assert list(merge([], [seg(1, 2)])) == [] + + def test_per_token_staleness_is_token_weighted(self): + # One turn, 10 tokens: 9 at epoch 1 + 1 at epoch 5. Per-segment accounting + # scores both halves equally; per-token, the stale run must dominate. + eod = MockTokenizer().eod + rollout = TokenRollout( + trajectory=[[7] * 9 + [eod]], + generation_mask=[[False] + [True] * 9], + logprobs=[[0.0] * 10], + reward=1.0, + env_id="w", + problem_id="p", + completion_ids=_mint_cids(1), + ) + ledger = { + rollout.completion_ids[0]: FinishedRequestRecord( + policy_epoch=[(0, 1), (9, 5)], kv_cache_epoch=[(0, 5)], num_evictions=0 + ) + } + stats = rl_utils.compute_group_stats( + [[rollout]], MockTokenizer(), seq_len=16, request_ledger=ledger + ) + seg = rl_utils.EpochSegment + assert stats.policy_epoch_segments == [[[seg(1, 9), seg(5, 1)]]] + assert stats.kv_cache_epoch_segments == [[[seg(5, 10)]]] + + writer = MagicMock() + metrics = rl_utils.prep_wandb_metrics( + writer, + stats.traj_lens, + stats.turn_lens, + stats.rewards, + stats.num_turns, + stats.advantages, + policy_epoch_segments=stats.policy_epoch_segments, + kv_cache_epoch_segments=stats.kv_cache_epoch_segments, + completed_epochs=stats.completed_epochs, + num_evictions=stats.num_evictions, + current_iteration=6, + ) + # First/last-token scalars: a segment's first/last epoch is its first/last token's. + assert metrics["max_policy_staleness"] == 5 # 6 - 1 + assert metrics["min_policy_last_token_staleness"] == 1 # 6 - 5 + tables = [c.kwargs for c in writer.Table.call_args_list] + # per_token_table rows are (rollout, policy, kv, token_count) runs merged + # over positions; 'rollout' indexes rollout_table's rows. + assert [t["data"] for t in tables if "token_count" in t.get("columns", [])] == [ + [(0, 5, 1, 9), (0, 1, 1, 1)] + ] + # rollout_table carries the within-rollout dispersion: token-weighted avg + # 6 - (1*9 + 5*1)/10 = 4.6 and std sqrt((9*0.4^2 + 3.6^2)/10) = 1.2 for + # policy; kv is a single epoch, so avg 1.0 and std 0. + rollout_tbl = next(t for t in tables if "policy_staleness_std" in t.get("columns", [])) + (rollout_row,) = rollout_tbl["data"] + row_by_col = dict(zip(rollout_tbl["columns"], rollout_row)) + assert row_by_col.pop("rollout_status") == 'ok' + assert row_by_col == pytest.approx( + { + "reward": 1.0, + "traj_length": 10, + "num_evictions": 0, + "policy_staleness": 5, + "kv_staleness": 1, + "policy_last_token_staleness": 1, + "kv_last_token_staleness": 1, + "policy_avg_staleness": 4.6, + "kv_avg_staleness": 1.0, + "policy_staleness_std": 1.2, + "kv_staleness_std": 0.0, + } + ) + # Histogram feed = the same exact token-weighted per-rollout averages + # (a per-segment mean would say 3.0 for policy). + hists = [t["data"] for t in tables if t.get("columns") == ["staleness"]] + assert [[4.6]] in hists and [[1.0]] in hists + def _ledger_record(epoch, num_evictions=0): return FinishedRequestRecord( From 3126ff2e3c73c18a52130bf73f51fb9dba252b4c Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 18 Aug 2026 13:21:29 +0000 Subject: [PATCH 077/101] Guard against divide/multiply by 0 in RL task distribution (#4362) Signed-off-by: Teodor-Dumitru Ene --- megatron/rl/agent/api.py | 1 - megatron/rl/agent/rollout_pipeline.py | 53 ++-- megatron/rl/agent/weighted_multi_task.py | 183 ++++++------ megatron/rl/rl_utils.py | 66 ++--- .../rl/server/agent/fastapi_env_server.py | 1 - megatron/training/arguments.py | 10 +- tests/unit_tests/rl/test_rl_utils.py | 27 +- .../unit_tests/rl/test_rollout_generation.py | 267 +++++++++++------- 8 files changed, 335 insertions(+), 273 deletions(-) diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index cad2ab5739e..8307015b240 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -35,7 +35,6 @@ class GroupedRolloutRequest(Request): inference_interface: InferenceInterface validation: bool = False filter_groups_with_same_reward: bool = False - streaming: bool = False submission_granularity: SubmissionGranularity = "B" consumption_granularity: ConsumptionGranularity = "B" diff --git a/megatron/rl/agent/rollout_pipeline.py b/megatron/rl/agent/rollout_pipeline.py index d515e972831..0fec043186b 100644 --- a/megatron/rl/agent/rollout_pipeline.py +++ b/megatron/rl/agent/rollout_pipeline.py @@ -89,7 +89,7 @@ def _validate(request: GroupedRolloutRequest, num_groups_per_env: list[int]) -> Args: request: Grouped rollout request to check. - num_groups_per_env: Proposed per-env group layout. + num_groups_per_env: Constant per-env group layout. """ assert ( GRANULARITY_RANK[request.consumption_granularity] @@ -231,10 +231,6 @@ def __init__( self.num_infer_workers = self.gate.capacity * ( rollouts_per_batch // self.gran_policy.units_per_batch() ) - if not request.streaming: - self.num_infer_workers = min( - self.num_infer_workers, request.num_groups * request.rollouts_per_group - ) # Core queues. self.infer_queue = asyncio_Queue() @@ -292,15 +288,10 @@ async def run(self) -> AsyncIterator[RolloutGroup]: ), None, ) - expected_end = ( - not self.request.streaming - and self.yielded_count == self.request.num_groups - ) - if failure is not None or not expected_end: - raise RuntimeError( - "RolloutPipeline output stream ended: a pipeline stage died" - + ("" if failure is not None else " (no stage exception was recovered)") - ) from failure + raise RuntimeError( + "RolloutPipeline output stream ended: a pipeline stage died" + + ("" if failure is not None else " (no stage exception was recovered)") + ) from failure finally: regen_tasks = tuple(self._regen_tasks) for task in (*tasks, *regen_tasks): @@ -335,17 +326,38 @@ def _maybe_close_intake(self) -> None: if self._prepare_done and self._groups_in_flight <= 0: self.infer_queue.shutdown() + def assert_no_inflight_rollouts(self) -> None: + """Verify no rollouts are buffered or in-flight at an iteration boundary for lag=0.""" + buffered = { + "infer_queue": self.infer_queue.qsize(), + "assemble_queue": self.assemble_queue.qsize(), + "output_queue": self.output_queue.qsize(), + "assemble_pending": sum(len(items) for items in self._assemble_pending.values()), + "consume_pending": sum(len(groups) for groups in self._consume_pending.values()), + "regen_pending": sum(1 for task in self._regen_tasks if not task.done()), + } + assert not any(buffered.values()), ( + f"The rollout pipeline has buffered rollouts at iteration boundary: {buffered}. " + "The generator has run ahead under a stale policy." + ) + # Filtered groups consume prepared rollouts without ever being yielded. + in_flight = self.prepared_count - ( + (self.yielded_count + self.filtered_count) * self.request.rollouts_per_group + ) + assert in_flight == 0, ( + f"The rollout pipeline prepared {self.prepared_count} rollout(s) but yielded " + f"{self.yielded_count} and filtered {self.filtered_count} group(s) of " + f"{self.request.rollouts_per_group}; " + f"{in_flight} rollout(s) in flight at iteration boundary." + ) + async def stage_prepare(self) -> None: """Generate gated inference work items.""" - assert ( - self.request.streaming - or self.request.num_groups % self.gran_policy.num_groups_per_batch == 0 - ), "non-streaming requires num_groups to be a multiple of num_groups_per_batch" group_id = 0 + batch_id = 0 try: - while self.request.streaming or group_id < self.request.num_groups: + while True: await self.gate.acquire_for("B") - batch_id = group_id // self.gran_policy.num_groups_per_batch for index_in_batch in range(self.gran_policy.num_groups_per_batch): await self.gate.acquire_for("G") @@ -354,6 +366,7 @@ async def stage_prepare(self) -> None: group_id=group_id, batch_id=batch_id, index_in_batch=index_in_batch ) group_id += 1 + batch_id += 1 except BaseException: self.infer_queue.shutdown() raise diff --git a/megatron/rl/agent/weighted_multi_task.py b/megatron/rl/agent/weighted_multi_task.py index a11e93fd7a7..67f47ea7f54 100644 --- a/megatron/rl/agent/weighted_multi_task.py +++ b/megatron/rl/agent/weighted_multi_task.py @@ -50,7 +50,6 @@ def __init__(self, agent_configs: list[AgentConfig]): # Initialize all sub-agents self.agents = [] - self.weights = [] self.agent_configs = agent_configs # Store the configs for later use # Calculate total weight only among non-evaluation agents @@ -59,14 +58,36 @@ def __init__(self, agent_configs: list[AgentConfig]): raise ValueError("Total weight of non-evaluation agents must be positive") for config in agent_configs: - # Initialize the agent with its arguments - agent = config.agent_type(**config.agent_args) - self.agents.append(agent) - # Only normalize weights for non-evaluation agents - if config.evaluation_only: - self.weights.append(0.0) - else: - self.weights.append(config.weight / total_weight) + self.agents.append(config.agent_type(**config.agent_args)) + + # Weight-0 entries exist so the launcher boots their servers ("weight 0 = never + # sampled"); exclude them here so the min-one-group bump cannot revive them. + self._rollout_agents = [] + self._rollout_env_ids = [] + self._rollout_weights = [] + self._rollout_config_indices = [] + for idx, (agent, config) in enumerate(zip(self.agents, agent_configs)): + env_id = getattr(agent, "env_id", None) or f"agent_{idx}" + if config.evaluation_only or config.weight <= 0.0: + if not config.evaluation_only: + logger.info( + "WeightedMultiTask: env %s has weight 0 and is excluded " + "from rollout generation.", + env_id, + ) + continue + self._rollout_agents.append(agent) + self._rollout_env_ids.append(env_id) + self._rollout_weights.append(config.weight / total_weight) + self._rollout_config_indices.append(idx) + duplicates = { + env_id for env_id in self._rollout_env_ids if self._rollout_env_ids.count(env_id) > 1 + } + if duplicates: + raise ValueError( + f"Duplicate env_ids among weighted environments: {sorted(duplicates)}; " + "per-env layout and metrics require unique names." + ) @classmethod def from_config(cls, config: list[dict[str, Any]]) -> 'WeightedMultiTask': @@ -98,93 +119,89 @@ def from_config(cls, config: list[dict[str, Any]]) -> 'WeightedMultiTask': return cls(agent_configs) - def _distribute_counts(self, total_count: int, distribute_remainder: bool = True) -> list[int]: - """Helper method to distribute counts according to weights. + @staticmethod + def _round_shares(targets: list[float], total: int) -> list[int]: + """Round fractional targets to integers, awarding the shortfall to the largest residuals.""" + counts = [int(t) for t in targets] + by_residual = sorted( + range(len(targets)), key=lambda i: targets[i] - counts[i], reverse=True + ) + for i in by_residual[: total - sum(counts)]: + counts[i] += 1 + return counts - This implementation ensures the most balanced distribution possible while - maintaining the relative proportions specified by weights. + def _quantized_counts(self, total: int) -> list[int]: + """Quantize weights into integer counts summing to `total`, at least one per weighted env. - Args: - total_count: Total number of items to distribute - distribute_remainder: Whether to distribute the remainder of the counts to the agents with the largest fractional parts + Raises ValueError when total is smaller than the number of weighted envs. + Note that this does not operate on eval-only tasks; those are pre-filtered in `__init__`. + """ + num_envs = len(self._rollout_weights) + if total < num_envs: + raise ValueError( + f"{num_envs} weighted environments cannot fit into {total} slots; " + "increase the batch or request size." + ) + exact = [weight * total for weight in self._rollout_weights] + counts = self._round_shares(exact, total) + # Round zero shares up to one, taking from the most over-served env. + while 0 in counts: + zero = counts.index(0) + donor = max( + (i for i in range(num_envs) if counts[i] >= 2), + key=lambda i: counts[i] - exact[i], + ) + counts[zero] += 1 + counts[donor] -= 1 + return counts - Returns: - List of counts for each agent, summing to total_count + def _distribute_counts(self, total_count: int) -> list[int]: + """Split a count across weighted agents by weight (largest remainder, min one each). + + Returns a per-agent list summing to total_count, 0 only for evaluation-only + and zero-weight agents; raises when total_count < the number of weighted envs. """ - # Filter out evaluation-only agents for rollout distribution - rollout_weights = [ - w for w, config in zip(self.weights, self.agent_configs) if not config.evaluation_only - ] - if not rollout_weights: - raise ValueError("No non-evaluation agents available for rollout generation") - - # Calculate exact fractional counts - exact_counts = [total_count * w for w in rollout_weights] - - # Get integer part of each count - base_counts = [int(count) for count in exact_counts] - remaining = total_count - sum(base_counts) - - if distribute_remainder: - # Sort indices by fractional parts to distribute remaining counts - # to those with largest fractional parts first - fractional_parts = [count - int(count) for count in exact_counts] - indices = list(range(len(rollout_weights))) - indices.sort(key=lambda i: fractional_parts[i], reverse=True) - - # Distribute remaining counts - for i in range(remaining): - base_counts[indices[i]] += 1 - - # Map back to original indices, skipping evaluation-only agents - final_counts = [] - rollout_idx = 0 - for config in self.agent_configs: - if config.evaluation_only: - final_counts.append(0) - else: - final_counts.append(base_counts[rollout_idx]) - rollout_idx += 1 - - return final_counts + shares = self._quantized_counts(total_count) + counts = [0] * len(self.agent_configs) + for idx, share in zip(self._rollout_config_indices, shares): + counts[idx] = share + return counts def rollout_allocations(self, num_groups: int) -> list[EnvAllocation]: - """Constant per-batch allocation for each weighted env, in env order.""" - counts = self._distribute_counts(num_groups) - env_ids = [ - getattr(agent, "env_id", None) or f"agent_{idx}" - for idx, agent in enumerate(self.agents) - ] - starved = [ - env_ids[idx] - for idx, count in enumerate(counts) - if count == 0 and self.weights[idx] > 0 - ] - if starved: - raise ValueError( - f"num_groups={num_groups} is too small to give every weighted env a group " - f"per batch (starved envs: {starved}); increase the trainer batch size." + """Constant per-batch allocation for each weighted env, in env order. + + Weights that cannot be realized as an integer split of the batch are rounded with a warning. + """ + for agent in self._rollout_agents: + if not isinstance(agent, GroupedRolloutGenerator): + raise TypeError(f"Agent of type {type(agent)} does not support grouped rollouts") + + counts = self._quantized_counts(num_groups) + exact = [weight * num_groups for weight in self._rollout_weights] + if any(abs(count - target) > 1e-9 for count, target in zip(counts, exact)): + logger.warning( + "WeightedMultiTask weights changed to fit num_groups=%d: %s", + num_groups, + ", ".join( + f"{eid}: {weight:g} -> {count}/{num_groups}" + for eid, weight, count in zip( + self._rollout_env_ids, self._rollout_weights, counts + ) + ), ) - for agent, count in zip(self.agents, counts): - if count > 0 and not isinstance(agent, GroupedRolloutGenerator): - raise TypeError( - f"Agent of type {type(agent)} does not support grouped rollouts" - ) - # Snapshot for metric logging; read back by rl_utils. - self.latest_distribution = { - "env_ids": env_ids, - "agent_groups": list(counts), - "num_groups": num_groups, - } logger.info( "WeightedMultiTask layout: num_groups=%d per_agent=%s", num_groups, - ", ".join(f"{eid}(groups={c})" for eid, c in zip(env_ids, counts)), + ", ".join( + f"{eid}(groups={c}, weight={w:g})" + for eid, c, w in zip(self._rollout_env_ids, counts, self._rollout_weights) + ), ) return [ EnvAllocation(agent=agent, env_id=env_id, num_groups=count) - for agent, env_id, count in zip(self.agents, env_ids, counts) - if count > 0 + for agent, env_id, count in zip( + self._rollout_agents, self._rollout_env_ids, counts + ) ] async def prepare_group_rollout( diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 754441f5a50..f1642095ef9 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -660,7 +660,6 @@ def get_rollout_generator( n_prompts: int, samples_per_group: int, *, - streaming: bool, generation_args: dict[str, Any], filter_groups_with_same_reward: bool, submission_granularity: SubmissionGranularity, @@ -674,10 +673,9 @@ def get_rollout_generator( The async iterator produced by RolloutPipeline.run(). """ global _ROLLOUT_GENERATOR, _ROLLOUT_PIPELINE - if not streaming or _ROLLOUT_GENERATOR is None: + if _ROLLOUT_GENERATOR is None: request = GroupedRolloutRequest( num_groups=n_prompts, - streaming=streaming, rollouts_per_group=samples_per_group, inference_interface=inference_interface, generation_args=generation_args, @@ -768,7 +766,6 @@ def get_environment_rollouts( inference_interface, n_prompts, samples_per_group, - streaming=args.rl_partial_rollouts, generation_args={ 'temperature': args.rl_default_temperature, 'max_tokens': args.inference_max_seq_length, @@ -799,12 +796,7 @@ def get_environment_rollouts( if torch.are_deterministic_algorithms_enabled(): rollouts.sort(key=lambda group: group[0].problem_id if group and group[0].problem_id else "") if not args.rl_partial_rollouts: - while True: - try: - loop.run_until_complete(anext(rollout_generator)) - assert False, "Unexpected group left in generator." - except StopAsyncIteration: - break + _ROLLOUT_PIPELINE.assert_no_inflight_rollouts() else: # Just set up space to collect the rollouts rollouts = [[None for _ in range(samples_per_group)] for _ in range(n_prompts)] @@ -1660,7 +1652,7 @@ def _collect_rollout_pipeline_metrics() -> dict: """Snapshot pipeline instrumentation into wandb-loggable scalars. Reads the RolloutPipeline held by get_rollout_generator: queue sizes, gate state, - per-stage dwell times, and rate counters, plus per-env group counters for multi-env agents. + per-stage dwell times, and rate counters, plus per-env group counters. Returns: Metric name -> value dict; empty when no pipeline exists yet. @@ -1668,7 +1660,6 @@ def _collect_rollout_pipeline_metrics() -> dict: if _ROLLOUT_PIPELINE is None: return {} pipeline = _ROLLOUT_PIPELINE - dist = getattr(pipeline.agent, "latest_distribution", None) metrics: dict = {} gate = pipeline.gate metrics.update({ @@ -1706,23 +1697,23 @@ def _collect_rollout_pipeline_metrics() -> dict: metrics[f"rollout_pipeline_max_{name}_s"] = float(arr.max()) metrics[f"rollout_pipeline_p50_{name}_s"] = float(np.percentile(arr, 50)) metrics[f"rollout_pipeline_p99_{name}_s"] = float(np.percentile(arr, 99)) - # Per-env group counters, mapped from env_index to env_id via the multi-env layout. - if dist: - active_env_ids = [ - env_id - for env_id, groups in zip(dist["env_ids"], dist["agent_groups"]) - if groups > 0 - ] - for env_index, env_id in enumerate(active_env_ids): - metrics[f"{env_id}_prepared_groups"] = ( - pipeline.prepared_groups_per_env[env_index] - ) - metrics[f"{env_id}_assembled_groups"] = ( - pipeline.assembled_groups_per_env[env_index] - ) - metrics[f"{env_id}_yielded_groups"] = ( - pipeline.yielded_groups_per_env[env_index] - ) + # Per-env metrics, in env layout order (the pipeline arrays are env-indexed; + # weighted env_ids are unique by construction). + for env_index, allocation in enumerate(pipeline.allocations): + metrics[f"{allocation.env_id}_prepared_groups"] = ( + pipeline.prepared_groups_per_env[env_index] + ) + metrics[f"{allocation.env_id}_assembled_groups"] = ( + pipeline.assembled_groups_per_env[env_index] + ) + metrics[f"{allocation.env_id}_yielded_groups"] = ( + pipeline.yielded_groups_per_env[env_index] + ) + metrics[f"{allocation.env_id}_agent_groups"] = allocation.num_groups + # The realized weight: the constant share of each batch the env actually owns. + metrics[f"{allocation.env_id}_weight"] = ( + allocation.num_groups / pipeline.request.num_groups + ) # Reset accumulators; queue sizes and gate held are point-in-time. pipeline.infer_queue_dwell = [] pipeline.engine_dwell = [] @@ -1733,23 +1724,14 @@ def _collect_rollout_pipeline_metrics() -> dict: pipeline.assembled_count = 0 pipeline.filtered_count = 0 pipeline.yielded_count = 0 - pipeline.prepared_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) - pipeline.assembled_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) - pipeline.yielded_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) + num_envs = len(pipeline.gran_policy.num_groups_per_env) + pipeline.prepared_groups_per_env = [0] * num_envs + pipeline.assembled_groups_per_env = [0] * num_envs + pipeline.yielded_groups_per_env = [0] * num_envs gate.prepare_blocked_seconds = 0.0 gate.acquire_calls = 0 gate.release_calls = 0 - # WeightedMultiTask per-batch group distribution. - if dist: - # An env_id can appear more than once in the config (e.g. an active - # entry plus an evaluation-only twin with zero weight). Sum per - # env_id so the zero twin does not overwrite the active entry. - per_env: dict = {} - for env_id, groups in zip(dist["env_ids"], dist["agent_groups"]): - per_env[env_id] = per_env.get(env_id, 0) + groups - for env_id, groups in per_env.items(): - metrics[f"{env_id}_agent_groups"] = groups return metrics diff --git a/megatron/rl/server/agent/fastapi_env_server.py b/megatron/rl/server/agent/fastapi_env_server.py index 1bcd3a21fa1..b4a8e42e747 100644 --- a/megatron/rl/server/agent/fastapi_env_server.py +++ b/megatron/rl/server/agent/fastapi_env_server.py @@ -141,7 +141,6 @@ async def get_grouped_rollouts( assert ( request.submission_granularity != "R" ), "FastAPIEnvServer does not support rollout submission granularity" - assert not request.streaming, "FastAPIEnvServer does not support group rollout streaming" payload = request.model_dump() payload["inference_interface"] = request.inference_interface.model_dump() async with httpx.AsyncClient() as client: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 73c4242e886..3c9984a6b96 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -523,9 +523,6 @@ def validate_args(args, defaults={}): "installed. See https://github.com/fzyzcjy/torch_memory_saver." ) - submit_rollouts_at_rollout_granularity = ( - args.rl_submission_granularity == "R" - ) if args.rl_max_inflight_requests is not None: requests_per_batch = args.grpo_prompts_per_step * args.grpo_group_size assert args.rl_generation_lag is None, \ @@ -549,10 +546,9 @@ def validate_args(args, defaults={}): if args.rl_generation_lag > 0: assert args.rl_partial_rollouts, \ "--rl-generation-lag requires --rl-partial-rollouts." - if submit_rollouts_at_rollout_granularity: - assert ( - args.rl_partial_rollouts - ), "Rollout submission granularity requires streaming grouped rollouts." + assert args.rl_submission_granularity == "B" or args.rl_partial_rollouts, \ + f"--rl-submission-granularity {args.rl_submission_granularity} requires " \ + "--rl-partial-rollouts." assert args.rl_consumption_granularity != "R", \ "--rl-consumption-granularity R is not currently supported." assert not ( diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 69fd60dd765..f3f43119ad2 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -238,30 +238,28 @@ def test_rl_granularity_defaults(self): assert not hasattr(args, "rl_parallel_generation_tasks") @pytest.mark.parametrize( - "rl_partial_rollouts, submission_granularity", + "submission_granularity", [ - pytest.param(False, "B", id="non_streaming_batch"), - pytest.param(True, "B", id="streaming_batch"), - pytest.param(True, "G", id="streaming_group"), - pytest.param(True, "R", id="streaming_rollout"), + pytest.param("B", id="batch"), + pytest.param("G", id="group"), + pytest.param("R", id="rollout"), ], ) def test_get_rollout_generator_keeps_num_groups_at_trainer_batch_size( - self, monkeypatch, rl_partial_rollouts, submission_granularity + self, monkeypatch, submission_granularity ): """Regression for the removed ``num_groups=1`` streaming override. Previously ``get_rollout_generator`` forced ``num_groups`` to 1 whenever it streamed with a non-batch submission granularity, collapsing the per-env - group layout so some environments received zero groups (now a loud - ``ValueError`` from ``rollout_allocations``). ``num_groups`` must stay at - the trainer batch size (``n_prompts``) regardless of streaming or + group layout so some environments received zero groups. ``num_groups`` + must stay at the trainer batch size (``n_prompts``) regardless of submission granularity. """ n_prompts = 8 captured = {} rollout_generator = object() - agent = object() + agent = MagicMock() class FakePipeline: def __init__(self, agent, request, parallel_generation_tasks): @@ -283,7 +281,6 @@ def run(self): ReturnsRaw(), n_prompts, 4, - streaming=rl_partial_rollouts, generation_args={'temperature': 1.0, 'max_tokens': 128, 'top_p': 1.0, 'top_k': 0}, filter_groups_with_same_reward=False, submission_granularity=submission_granularity, @@ -295,7 +292,6 @@ def run(self): assert result is rollout_generator assert captured["agent"] is agent assert captured["request"].num_groups == n_prompts - assert captured["request"].streaming == rl_partial_rollouts # The gate depth handed to the pipeline is lag + 1 trainer batches, # independent of submission granularity. assert captured["parallel_generation_tasks"] == generation_lag + 1 @@ -331,9 +327,14 @@ def run(self): ), pytest.param( {"rl_submission_granularity": "R"}, - "Rollout submission granularity requires streaming grouped rollouts", + "--rl-submission-granularity R requires --rl-partial-rollouts", id="rollout_submission_requires_partial_rollouts", ), + pytest.param( + {"rl_submission_granularity": "G", "rl_consumption_granularity": "G"}, + "--rl-submission-granularity G requires --rl-partial-rollouts", + id="group_submission_requires_partial_rollouts", + ), pytest.param( {"rl_consumption_granularity": "R"}, "--rl-consumption-granularity R is not currently supported", diff --git a/tests/unit_tests/rl/test_rollout_generation.py b/tests/unit_tests/rl/test_rollout_generation.py index 1efaf22806e..9e423810a17 100644 --- a/tests/unit_tests/rl/test_rollout_generation.py +++ b/tests/unit_tests/rl/test_rollout_generation.py @@ -178,7 +178,6 @@ async def test_group_submission_stalls_until_consumption( num_groups=num_groups, rollouts_per_group=1, inference_interface=MockInferenceInterface(), - streaming=True, submission_granularity="G", consumption_granularity=consumption_granularity, ) @@ -207,7 +206,6 @@ async def test_batch_submission_releases_once_per_batch(self): num_groups=2, rollouts_per_group=1, inference_interface=MockInferenceInterface(), - streaming=True, submission_granularity="B", consumption_granularity="B", ) @@ -326,20 +324,25 @@ def test_grouped_rollout_request_rejects_unknown_granularity(self, field): @pytest.mark.asyncio @pytest.mark.parametrize( - "streaming, submission_granularity, consumption_granularity, num_degenerate", + ( + "submission_granularity, consumption_granularity, num_degenerate, " + "parallel_generation_tasks" + ), [ - pytest.param(False, "B", "B", 3, id="batch_submission"), - pytest.param(False, "G", "B", 3, id="group_submission"), - pytest.param(False, "R", "B", 3, id="rollout_submission"), - pytest.param(False, "G", "G", 3, id="group_consumption"), - pytest.param(False, "R", "G", 3, id="rollout_submission_group_consume"), - pytest.param(False, "B", "B", 9, id="cascading_regeneration"), - pytest.param(True, "B", "B", 3, id="streaming_batch_consume"), - pytest.param(True, "G", "G", 3, id="streaming_group_consume"), + pytest.param("B", "B", 3, 1, id="batch_submission"), + pytest.param("G", "B", 3, 8, id="group_submission"), + pytest.param("R", "B", 3, 8, id="rollout_submission"), + pytest.param("G", "G", 3, 8, id="group_consumption"), + pytest.param("R", "G", 3, 8, id="rollout_submission_group_consume"), + pytest.param("B", "B", 9, 1, id="cascading_regeneration"), ], ) async def test_filter_groups_and_regenerate( - self, streaming, submission_granularity, consumption_granularity, num_degenerate + self, + submission_granularity, + consumption_granularity, + num_degenerate, + parallel_generation_tasks, ): num_groups = 4 gen = FilteringMockGenerator(num_degenerate=num_degenerate) @@ -348,20 +351,24 @@ async def test_filter_groups_and_regenerate( rollouts_per_group=2, inference_interface=MockInferenceInterface(), filter_groups_with_same_reward=True, - streaming=streaming, submission_granularity=submission_granularity, consumption_granularity=consumption_granularity, ) - pipeline = RolloutPipeline(gen, request, parallel_generation_tasks=8) + pipeline = RolloutPipeline( + gen, request, parallel_generation_tasks=parallel_generation_tasks + ) - expected_count = 2 * num_groups if streaming else num_groups - groups, filtered_counts = [], [] + expected_count = 2 * num_groups + groups = [] async with aclosing(pipeline.run()) as it: async for group in it: groups.append(group) - filtered_counts.append(pipeline.filtered_count) - if streaming and len(groups) >= expected_count: + if len(groups) >= expected_count: break + if submission_granularity == "B" and parallel_generation_tasks == 1: + # lag=0: the boundary is quiescent, checked before close() drops + # shutdown sentinels into the queues on Python < 3.13. + pipeline.assert_no_inflight_rollouts() assert len(groups) == expected_count # Every delivered group carries reward signal. @@ -373,10 +380,10 @@ async def test_filter_groups_and_regenerate( for batch_start in range(0, expected_count, num_groups): batch = groups[batch_start : batch_start + num_groups] assert sorted(g.index_in_batch for g in batch) == list(range(num_groups)) - if not streaming: - # One extra prepare per dropped group, and no under-delivery. - assert gen.prepare_group_rollout_calls == num_groups + num_degenerate - assert filtered_counts[-1] == num_degenerate + if submission_granularity == "B" and parallel_generation_tasks == 1: + # Every dropped group cost exactly one extra prepare (no under-delivery). + assert gen.prepare_group_rollout_calls == expected_count + num_degenerate + assert pipeline.filtered_count == num_degenerate @pytest.mark.asyncio async def test_pending_regeneration_does_not_block_delivery(self): @@ -414,39 +421,37 @@ async def base_generate(self, request): assert [rollout.trajectory[0] for rollout in group.rollouts] == ["t1", "t1"] assert pipeline.filtered_count == 1 assert len(pipeline._regen_tasks) == 1 - assert gen.prepare_group_rollout_calls == 4 + # t0-t2 plus the pinned regeneration, plus the streaming prepare of + # the next batch's first group (pinned on the same full gate). + assert gen.prepare_group_rollout_calls == 5 events["t2"].set() - rest = [await asyncio.wait_for(anext(it), timeout=10) for _ in range(2)] - assert sorted(g.rollouts[0].trajectory[0] for g in rest) == ["t2", "t3"] + # Completion order may interleave the perpetual stream's next-batch + # groups; keep pulling (bounded) until both stragglers land. + rest = [] + for _ in range(5): + rest.append(await asyncio.wait_for(anext(it), timeout=10)) + if {(0, 0), (0, 2)} <= {(g.batch_id, g.index_in_batch) for g in rest}: + break + # The gated group (slot 2) and the regenerated replacement (slot 0) both deliver. + assert {(0, 0), (0, 2)} <= {(g.batch_id, g.index_in_batch) for g in rest} assert not pipeline._regen_tasks @pytest.mark.asyncio @pytest.mark.parametrize( ( - "num_slow_calls, streaming, num_groups, submission_granularity, " + "num_slow_calls, num_groups, submission_granularity, " "consumption_granularity, expected_count, expected_batch_ids, " "expected_trajectories" ), [ - pytest.param(0, False, 8, "B", "B", 8, None, None, id="non_batched"), - pytest.param( - 0, False, 4, "B", "B", 4, None, None, id="non_streaming_fewer_than_parallel" - ), + pytest.param(0, 8, "B", "B", 8, None, None, id="single_batch"), + pytest.param(0, 4, "B", "B", 4, None, None, id="fewer_groups_than_parallel"), pytest.param( - 4, - True, - 2, - "B", - "B", - 8, - [0, 0, 1, 1, 2, 2, 3, 3], - None, - id="batched_submission_order", + 4, 2, "B", "B", 8, [0, 0, 1, 1, 2, 2, 3, 3], None, id="batched_submission_order" ), - pytest.param(0, True, 1, "G", "B", 10, None, None, id="streaming"), + pytest.param(0, 1, "G", "B", 10, None, None, id="streaming"), pytest.param( 4, - True, 1, "G", "G", @@ -457,7 +462,6 @@ async def base_generate(self, request): ), pytest.param( 4, - True, 1, "G", "B", @@ -471,7 +475,6 @@ async def base_generate(self, request): async def test_grouped_rollout_generation( self, num_slow_calls, - streaming, num_groups, submission_granularity, consumption_granularity, @@ -484,7 +487,6 @@ async def test_grouped_rollout_generation( num_groups=num_groups, rollouts_per_group=1, inference_interface=MockInferenceInterface(num_slow_calls=num_slow_calls), - streaming=streaming, submission_granularity=submission_granularity, consumption_granularity=consumption_granularity, ) @@ -492,7 +494,7 @@ async def test_grouped_rollout_generation( groups = [] async for group in RolloutPipeline(gen, request, parallel_generation_tasks=8).run(): groups.append(group) - if request.streaming and len(groups) >= expected_count: + if len(groups) >= expected_count: break assert len(groups) == expected_count @@ -513,7 +515,6 @@ async def test_rollout_submission_granularity_limits_inference_concurrency(self) num_groups=2, rollouts_per_group=2, inference_interface=inference_interface, - streaming=True, submission_granularity="R", consumption_granularity="B", ) @@ -534,66 +535,96 @@ async def test_rollout_submission_granularity_limits_inference_concurrency(self) @pytest.mark.asyncio @pytest.mark.parametrize( - "submission_granularity, consumption_granularity", - [pytest.param("B", "B", id="batch_batch"), pytest.param("G", "G", id="group_group")], + ( + "submission_granularity, consumption_granularity, num_slow_calls, " + "parallel_generation_tasks, windows" + ), + [ + pytest.param("B", "B", 0, 1, 2, id="batch_batch"), + pytest.param("G", "G", 0, 1, 2, id="group_group"), + # Balanced G under out-of-order completion and a depth-2 gate. + pytest.param("G", "G", 2, 2, 3, id="group_group_out_of_order"), + ], ) - async def test_weighted_multi_task(self, submission_granularity, consumption_granularity): - configs = [ - AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), - AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), - ] - mt = WeightedMultiTask(configs) - + async def test_weighted_multi_task( + self, + submission_granularity, + consumption_granularity, + num_slow_calls, + parallel_generation_tasks, + windows, + ): + """Routing and consumption keep every trainer-batch window at the exact env mix.""" + mt = WeightedMultiTask( + [ + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), + ] + ) request = GroupedRolloutRequest( num_groups=4, rollouts_per_group=1, - inference_interface=MockInferenceInterface(), - streaming=False, + inference_interface=MockInferenceInterface(num_slow_calls=num_slow_calls), submission_granularity=submission_granularity, consumption_granularity=consumption_granularity, ) - groups = [] - pipeline = RolloutPipeline(mt, request, parallel_generation_tasks=1) - async for group in pipeline.run(): - groups.append(group) + pipeline = RolloutPipeline(mt, request, parallel_generation_tasks=parallel_generation_tasks) + gen = pipeline.run() + groups = [await anext(gen) for _ in range(windows * 4)] - assert len(groups) == 4 - # Weights 3:1 → env "a" owns 3 batch slots, env "b" owns 1; the single - # pipeline routes preparation and generation to the owning sub-agent. + assert [a.num_groups for a in pipeline.allocations] == [3, 1] + # Weights 3:1 → env "a" owns 3 of every 4 batch slots; the pipeline routes + # each slot to the owning sub-agent regardless of completion order. env_ids = [g[0].env_id for g in groups] - assert sorted(env_ids) == ["a", "a", "a", "b"] - assert [agent.prepare_group_rollout_calls for agent in mt.agents] == [3, 1] - assert [agent.get_rollout_response_calls for agent in mt.agents] == [3, 1] - assert mt.latest_distribution["agent_groups"] == [3, 1] - # The pipeline drains fully: every gate slot is released at exhaustion. - assert pipeline.gate.held == 0 + for start in range(0, windows * 4, 4): + assert sorted(env_ids[start : start + 4]) == ["a", "a", "a", "b"] + if consumption_granularity == "B": + # With depth-1 gating and consumed-release, nothing is buffered or in flight. + pipeline.assert_no_inflight_rollouts() @pytest.mark.asyncio - async def test_group_consumption_balances_each_batch(self): - """Balanced G: every trainer-batch window holds each env's exact share.""" + async def test_lag0_streaming_matches_non_streaming_boundaries(self): + """lag=0 (B/B, depth-1 gate): each iteration of the persistent stream is exactly + one batch, generated entirely after the previous boundary — the old + non-streaming per-iteration contract, enforced by assert_no_inflight_rollouts.""" configs = [ - AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=1.0), AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), ] mt = WeightedMultiTask(configs) + request = GroupedRolloutRequest( + num_groups=4, + rollouts_per_group=2, + inference_interface=MockInferenceInterface(), + submission_granularity="B", + consumption_granularity="B", + ) + pipeline = RolloutPipeline(mt, request, parallel_generation_tasks=1) + gen = pipeline.run() + for iteration in range(3): + groups = [await anext(gen) for _ in range(4)] + # Exactly this iteration's batch, whole and in order. + assert [g.batch_id for g in groups] == [iteration] * 4 + # Nothing of the next batch has even been prepared: everything the + # next iteration consumes is generated after this boundary. + assert sum(a.prepare_group_rollout_calls for a in mt.agents) == (iteration + 1) * 4 + pipeline.assert_no_inflight_rollouts() + @pytest.mark.asyncio + async def test_assert_no_inflight_rollouts_detects_run_ahead(self): + """With lag>0 the gate legitimately runs ahead; the boundary checker must fire.""" request = GroupedRolloutRequest( num_groups=4, rollouts_per_group=1, - inference_interface=MockInferenceInterface(num_slow_calls=2), - streaming=True, - submission_granularity="G", - consumption_granularity="G", + inference_interface=MockInferenceInterface(), + submission_granularity="B", + consumption_granularity="B", ) - groups = [] - async for group in RolloutPipeline(mt, request, parallel_generation_tasks=2).run(): - groups.append(group) - if len(groups) >= 12: - break - - for start in range(0, 12, 4): - env_ids = [g[0].env_id for g in groups[start : start + 4]] - assert sorted(env_ids) == ["a", "a", "a", "b"] + pipeline = RolloutPipeline(MockGenerator(), request, parallel_generation_tasks=2) + gen = pipeline.run() + [await anext(gen) for _ in range(4)] + with pytest.raises(AssertionError, match="The rollout pipeline"): + pipeline.assert_no_inflight_rollouts() @pytest.mark.asyncio @pytest.mark.parametrize( @@ -614,30 +645,54 @@ async def test_consumption_finer_than_submission_rejected( with pytest.raises(AssertionError, match="no finer"): RolloutPipeline(gen, request, parallel_generation_tasks=1) - def test_multi_env_layout_rejects_starving_batch_size(self): - """The layout raises rather than silently starving a weighted env.""" + @pytest.mark.parametrize( + "weights, num_groups, expected_layout, warns", + [ + # 8 groups cannot realize 1:2 exactly; quantized with a warning. + pytest.param([1.0, 2.0], 8, [3, 5], True, id="quantized"), + # A weight below 1/num_groups keeps one group per batch. + pytest.param([0.01, 0.99], 8, [1, 7], True, id="zero_share_rounded_up"), + pytest.param([3.0, 1.0], 8, [6, 2], False, id="exact"), + pytest.param([1.0, 1.0, 1.0], 3, [1, 1, 1], False, id="one_group_each"), + # Only an env count exceeding the batch size is infeasible. + pytest.param([1.0, 1.0, 1.0], 2, None, False, id="too_many_envs"), + ], + ) + def test_multi_env_layout(self, caplog, weights, num_groups, expected_layout, warns): + """Weights quantize to a constant split (warned); weight-0 and eval-only envs take + no slot — the min-one-group bump must not revive boot-only entries.""" configs = [ - AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), - AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": f"e{i}"}, weight=w) + for i, w in enumerate(weights) ] + configs.insert( + 1, AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "boot"}, weight=0.0) + ) + configs.append( + AgentConfig( + agent_type=MockGenerator, + agent_args={"env_id": "eval"}, + weight=1.0, + evaluation_only=True, + ) + ) mt = WeightedMultiTask(configs) - with pytest.raises(ValueError, match="starved"): - mt.rollout_allocations(1) - assert [a.num_groups for a in mt.rollout_allocations(8)] == [6, 2] - - # Evaluation-only envs take no groups and never count as starved. - mt = WeightedMultiTask( - configs - + [ - AgentConfig( - agent_type=MockGenerator, - agent_args={"env_id": "c"}, - weight=1.0, - evaluation_only=True, - ) - ] + assert mt._rollout_env_ids == [f"e{i}" for i in range(len(weights))] + if expected_layout is None: + with pytest.raises(ValueError, match="cannot fit"): + mt.rollout_allocations(num_groups) + with pytest.raises(ValueError, match="cannot fit"): + mt._distribute_counts(num_groups) + return + # The split is identical on every call. + assert [[a.num_groups for a in mt.rollout_allocations(num_groups)] for _ in range(3)] == [ + expected_layout + ] * 3 + assert warns == any("weights changed" in message for message in caplog.messages) + # Config-order distribution: the boot (index 1) and eval (last) slots pinned to 0. + assert mt._distribute_counts(num_groups) == ( + expected_layout[:1] + [0] + expected_layout[1:] + [0] ) - assert [a.num_groups for a in mt.rollout_allocations(8)] == [6, 2] def make_response( From 74a48940c309fe78b6f2d8591abeb0ae2ea6b6c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Tue, 18 Aug 2026 21:19:55 +0200 Subject: [PATCH 078/101] fix(ci): route scoped tokens through community bot (#6632) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- .github/workflows/cicd-main.yml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index ab24d5e3112..6162f67d239 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -87,8 +87,8 @@ jobs: && vars.DISABLE_EXTERNAL_CONTRIBUTOR == 'true' && steps.check-sso.outputs.is_member != 'true' uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 with: - app-id: ${{ vars.BOT_ID }} - private-key: ${{ secrets.BOT_KEY }} + app-id: ${{ vars.COMMUNITY_BOT_ID }} + private-key: ${{ secrets.COMMUNITY_BOT_KEY }} owner: NVIDIA-NeMo permission-members: read @@ -100,8 +100,8 @@ jobs: && vars.DISABLE_EXTERNAL_CONTRIBUTOR == 'true' && steps.check-sso.outputs.is_member != 'true' uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 with: - app-id: ${{ vars.BOT_ID }} - private-key: ${{ secrets.BOT_KEY }} + app-id: ${{ vars.COMMUNITY_BOT_ID }} + private-key: ${{ secrets.COMMUNITY_BOT_KEY }} owner: NVIDIA permission-members: read @@ -1402,8 +1402,8 @@ jobs: id: app-token uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 with: - app-id: ${{ vars.BOT_ID }} - private-key: ${{ secrets.BOT_KEY }} + app-id: ${{ vars.COMMUNITY_BOT_ID }} + private-key: ${{ secrets.COMMUNITY_BOT_KEY }} permission-statuses: write - name: Generate fake coverage report @@ -1497,8 +1497,8 @@ jobs: id: app-token uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 with: - app-id: ${{ vars.BOT_ID }} - private-key: ${{ secrets.BOT_KEY }} + app-id: ${{ vars.COMMUNITY_BOT_ID }} + private-key: ${{ secrets.COMMUNITY_BOT_KEY }} permission-issues: write - name: Comment on PR with action run URL From 0980313cb6caacfa6a0f14c5038c7d3aaeb7f963 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Tue, 18 Aug 2026 21:40:13 +0200 Subject: [PATCH 079/101] fix(tests): restrict result extraction to rank zero (#6620) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- tests/functional_tests/shell_test_utils/run_ci_test.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/functional_tests/shell_test_utils/run_ci_test.sh b/tests/functional_tests/shell_test_utils/run_ci_test.sh index b58c37d7e23..f6cd272135b 100644 --- a/tests/functional_tests/shell_test_utils/run_ci_test.sh +++ b/tests/functional_tests/shell_test_utils/run_ci_test.sh @@ -386,7 +386,7 @@ for i in $(seq 1 $N_REPEAT); do SKIP_PYTEST=1 fi - if [[ ${SKIP_PYTEST:-0} != 1 || "$TEST_TYPE" == "release" ]]; then + if [[ "$NODE_RANK" -eq 0 && (${SKIP_PYTEST:-0} != 1 || "$TEST_TYPE" == "release") ]]; then # Save run results export PYTHONPATH=$ROOT_DIR if [[ "$TEST_TYPE" == "release" ]]; then @@ -438,7 +438,7 @@ for i in $(seq 1 $N_REPEAT); do ALLOW_NONDETERMINISTIC_ALGO_ARG="--allow-nondeterministic-algo" fi - if [[ "$SLURM_NODEID" -eq 0 ]]; then + if [[ "$NODE_RANK" -eq 0 ]]; then echo "Running pytest checks against golden values" # For pretraining jobs From 54c7918b9a012626dc0a65c76e6c79086647056f Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Tue, 18 Aug 2026 17:19:07 +0000 Subject: [PATCH 080/101] test(mfsdp): separate memory tests (#6568) Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 242 --------------- .../distributed/mfsdp_v2/test_memory.py | 283 ++++++++++++++++++ 2 files changed, 283 insertions(+), 242 deletions(-) create mode 100644 tests/unit_tests/distributed/mfsdp_v2/test_memory.py diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 967c18056fd..dcb154cbcd8 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -86,18 +86,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -class ElementwiseModel(nn.Module): - """Small activation path over a large FSDP-managed weight.""" - - def __init__(self, dim: int) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(dim, dim)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Apply the first weight row to an activation tensor.""" - return torch.relu(x + self.weight[0]) - - class TiedLM(nn.Module): """Tiny language model with shared input and output embedding weights.""" @@ -172,10 +160,6 @@ def _hfsdp_placements() -> Placements: ) -def _mb(num_bytes: int) -> str: - return f"{num_bytes / 1024**2:.2f} MB" - - # CPU ops that a device event chains up to via cpu_parent, used to attribute the device # work to its enclosing collective or matmul operation. _REDUCE_SCATTER_OP_NAME_SUBSTRING = "reduce_scatter" @@ -603,172 +587,6 @@ def test_tied_child_parameters_allocate_one_physical_weight(distributed_setup): assert len(list(model.parameters())) == 1 -@pytest.mark.parametrize("main_params_dtype", [torch.bfloat16, torch.float32]) -def test_persistent_sharded_storage(distributed_setup, main_params_dtype): - """FSDP should retain only its sharded weights and gradients at rest.""" - rank = distributed_setup.rank - world_size = distributed_setup.world_size - device = distributed_setup.device - if world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(device.type, (world_size,)) - dim = 4096 - dtype = torch.bfloat16 - model = MultiChildModel(dim=dim, num_children=8).to(dtype=dtype) - placements = _flat_placements() - policy = MixedPrecisionPolicy(main_params_dtype=main_params_dtype) - allocated_before = torch.cuda.memory_allocated(device) - with fully_shard_context(device=device): - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) - fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) - - child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() - persistent_allocated = torch.cuda.memory_allocated(device) - allocated_before - if main_params_dtype == dtype: - # Model and main weights alias, leaving only one BF16 weight buffer and one - # BF16 main-gradient buffer per child. - assert all( - group.model_weight is group.main_weight - for layer in model.layers - for group in layer.parameter_groups - ) - expected_per_child_nbytes = 2 * child_weight_nbytes - else: - # FP32 main weights require a distinct buffer in addition to the BF16 model - # weight and BF16 main-gradient buffers. - main_weight_nbytes = dim * dim * torch.empty((), dtype=main_params_dtype).element_size() - expected_per_child_nbytes = 2 * child_weight_nbytes + main_weight_nbytes - - # All persistent buffers are sharded over the data-parallel group. Small bookkeeping - # allocations stay below 1 MiB. - expected_persistent_nbytes = len(model.layers) * expected_per_child_nbytes // world_size - assert ( - expected_persistent_nbytes <= persistent_allocated < expected_persistent_nbytes + 1024**2 - ), ( - "FSDP persistent memory does not match its sharded weight and gradient storage: " - f"rank={rank}, persistent_allocated={_mb(persistent_allocated)}, " - f"expected={_mb(expected_persistent_nbytes)}" - ) - - -def test_training_step_peak_memory_bounds_full_size_buffers(distributed_setup): - """A training step should stay below five full-size child buffers.""" - rank = distributed_setup.rank - world_size = distributed_setup.world_size - device = distributed_setup.device - if world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(device.type, (world_size,)) - dim = 4096 - dtype = torch.bfloat16 - model = MultiChildModel(dim=dim, num_children=8).to(dtype=dtype) - placements = _flat_placements() - policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) - with fully_shard_context(device=device): - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) - fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) - - optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - fully_shard_optimizer(optimizer) - x = torch.randn(2, dim, device=device, dtype=dtype) - - def train_step() -> None: - optimizer.zero_grad(set_to_none=True) - model(x).float().sum().backward() - optimizer.step() - - child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() - resting_allocated = torch.cuda.memory_allocated(device) - torch.cuda.reset_peak_memory_stats(device) - train_step() - peak_delta = torch.cuda.max_memory_allocated(device) - resting_allocated - - # Backward keeps the current child and one prefetched child unsharded. The current - # child also has a full wgrad until it is copied into a full reduce-scatter input, - # for a four-full-child-buffer peak. Allow one additional buffer for cuBLAS - # workspace, allocator granularity, and small temporaries. - bound_nbytes = (4 + 1) * child_weight_nbytes - - assert peak_delta < bound_nbytes, ( - "FSDP training-step peak memory exceeded the full-size-buffer bound: " - f"rank={rank}, peak_delta={_mb(peak_delta)}, " - f"five_full_child_buffers={_mb(bound_nbytes)}" - ) - - -def test_deleted_model_releases_fsdp_storage(distributed_setup): - """Deleting an FSDP model should release its persistent storage.""" - world_size = distributed_setup.world_size - device = distributed_setup.device - - mesh = init_device_mesh(device.type, (world_size,)) - # Earlier tests may retain process-global CUDA allocations such as the - # CuBLAS workspace. Capture them before creating this model, so the test - # only detects storage retained by the deleted FSDP model itself. - allocated_before = torch.cuda.memory_allocated(device) - model = ElementwiseModel(dim=8192).to(dtype=torch.bfloat16, device=device) - with fully_shard_context(device=device): - fully_shard(model, mesh=mesh, placements=_flat_placements()) - - x = torch.ones(1, 8192, dtype=torch.bfloat16, device=device) - output = model(x) - del output, x, model - - assert torch.cuda.memory_allocated(device) - allocated_before < 1024**2 - - -def test_fully_shard_returns_to_resting_memory(distributed_setup): - """Fully-sharded temporary storage should be released after forward and backward.""" - rank = distributed_setup.rank - world_size = distributed_setup.world_size - device = distributed_setup.device - if world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - - mesh = init_device_mesh(device.type, (world_size,)) - dim = 4096 - dtype = torch.bfloat16 - model = MultiChildModel(dim=dim, num_children=2).to(dtype=dtype, device=device) - placements = _flat_placements() - policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) - with fully_shard_context(device=device): - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) - fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) - - x = torch.randn(2, dim, device=device, dtype=dtype) - - def clear_cublas_workspaces_and_get_allocated_memory() -> int: - # PyTorch retains a cuBLAS workspace for each handle/stream pair. Clear those - # library caches so this measurement isolates FSDP-managed storage. - torch._C._cuda_clearCublasWorkspaces() - return torch.cuda.memory_allocated(device) - - resting_allocated = clear_cublas_workspaces_and_get_allocated_memory() - - def assert_returns_to_resting_memory(phase: str) -> None: - extra_allocated = clear_cublas_workspaces_and_get_allocated_memory() - resting_allocated - # The live output, activations, and root-owned bias gradient are small; unsharded - # parameter storage must be released. - assert extra_allocated < 1024**2, ( - f"Fully-sharded storage did not return to resting memory after {phase}: " - f"rank={rank}, extra_allocated={_mb(extra_allocated)}, " - "max_extra_allocated=1.00 MB" - ) - - output = model(x) - assert_returns_to_resting_memory("forward") - - loss = output.float().square().mean() - loss.backward() - del loss, output - assert_returns_to_resting_memory("backward") - - def test_parameterless_parent_with_child_modules_trains(distributed_setup): """A parent with no unowned parameters should still root trainable child FsdpModules.""" world_size = distributed_setup.world_size @@ -1049,63 +867,3 @@ def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup): assert group.main_grad is not None assert group._unsharded_model_weight is not None assert group._unsharded_model_weight.local_buffer.untyped_storage().nbytes() == 0 - - -def test_fully_shard_reduces_peak_training_memory(distributed_setup): - """Per-layer FSDP should reduce peak CUDA memory during training.""" - rank = distributed_setup.rank - world_size = distributed_setup.world_size - device = distributed_setup.device - if world_size < 2: - pytest.skip("This test requires at least 2 ranks.") - mesh = init_device_mesh(device.type, (world_size,)) - dim = 1024 - layers = 16 - batch = 8 - steps = 2 - dtype = torch.bfloat16 - - def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Tensor) -> None: - for _ in range(steps): - optimizer.zero_grad(set_to_none=True) - model(x).sum().backward() - optimizer.step() - - torch.manual_seed(4321) - baseline = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) - baseline_optimizer = torch.optim.AdamW(baseline.parameters(), lr=0.01) - x = torch.randn(batch, dim, device=device, dtype=dtype) - torch.cuda.reset_peak_memory_stats(device) - train_steps(baseline, baseline_optimizer, x) - baseline_peak = torch.cuda.max_memory_allocated(device) - - del baseline_optimizer - del baseline - del x - - torch.manual_seed(4321) - model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) - with fully_shard_context(device=device): - for layer in model: - fully_shard( - layer, - mesh=mesh, - placements=_flat_placements(), - mixed_precision_policy=MixedPrecisionPolicy( - main_params_dtype=dtype, main_grads_dtype=dtype - ), - ) - optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) - - x = torch.randn(batch, dim, device=device, dtype=dtype) - torch.cuda.reset_peak_memory_stats(device) - train_steps(model, optimizer, x) - sharded_peak = torch.cuda.max_memory_allocated(device) - logger.info( - "FSDP peak memory: rank=%s, baseline=%s, sharded=%s", - rank, - _mb(baseline_peak), - _mb(sharded_peak), - ) - - assert sharded_peak < baseline_peak diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_memory.py b/tests/unit_tests/distributed/mfsdp_v2/test_memory.py new file mode 100644 index 00000000000..6e0e2b2e389 --- /dev/null +++ b/tests/unit_tests/distributed/mfsdp_v2/test_memory.py @@ -0,0 +1,283 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Memory-accounting tests for Megatron-FSDP.""" + +import logging + +import pytest +import torch +from torch import nn +from torch.distributed.device_mesh import init_device_mesh + +from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( + Flat, + Placements, + fully_shard, + fully_shard_context, + fully_shard_optimizer, +) +from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy + +logger = logging.getLogger(__name__) + + +class MultiChildModel(nn.Module): + """Model with direct parameters and multiple child FsdpModules.""" + + def __init__(self, dim: int, num_children: int) -> None: + super().__init__() + self.bias = nn.Parameter(torch.ones(dim)) + self.layers = nn.ModuleList([nn.Linear(dim, dim, bias=False) for _ in range(num_children)]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run through every child layer with a root-owned bias.""" + x = x + self.bias + for layer in self.layers: + x = torch.relu(layer(x)) + return x + + +class ElementwiseModel(nn.Module): + """Small activation path over a large FSDP-managed weight.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(dim, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply the first weight row to an activation tensor.""" + return torch.relu(x + self.weight[0]) + + +def _flat_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) + + +def _mb(num_bytes: int) -> str: + return f"{num_bytes / 1024**2:.2f} MB" + + +@pytest.mark.parametrize("main_params_dtype", [torch.bfloat16, torch.float32]) +def test_persistent_sharded_storage(distributed_setup, main_params_dtype): + """FSDP should retain only its sharded weights and gradients at rest.""" + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(device.type, (world_size,)) + dim = 4096 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=8).to(dtype=dtype) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=main_params_dtype) + allocated_before = torch.cuda.memory_allocated(device) + with fully_shard_context(device=device): + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() + persistent_allocated = torch.cuda.memory_allocated(device) - allocated_before + if main_params_dtype == dtype: + # Model and main weights alias, leaving only one BF16 weight buffer and one + # BF16 main-gradient buffer per child. + assert all( + group.model_weight is group.main_weight + for layer in model.layers + for group in layer.parameter_groups + ) + expected_per_child_nbytes = 2 * child_weight_nbytes + else: + # FP32 main weights require a distinct buffer in addition to the BF16 model + # weight and BF16 main-gradient buffers. + main_weight_nbytes = dim * dim * torch.empty((), dtype=main_params_dtype).element_size() + expected_per_child_nbytes = 2 * child_weight_nbytes + main_weight_nbytes + + # All persistent buffers are sharded over the data-parallel group. Small bookkeeping + # allocations stay below 1 MiB. + expected_persistent_nbytes = len(model.layers) * expected_per_child_nbytes // world_size + assert ( + expected_persistent_nbytes <= persistent_allocated < expected_persistent_nbytes + 1024**2 + ), ( + "FSDP persistent memory does not match its sharded weight and gradient storage: " + f"rank={rank}, persistent_allocated={_mb(persistent_allocated)}, " + f"expected={_mb(expected_persistent_nbytes)}" + ) + + +def test_training_step_peak_memory_bounds_full_size_buffers(distributed_setup): + """A training step should stay below five full-size child buffers.""" + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(device.type, (world_size,)) + dim = 4096 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=8).to(dtype=dtype) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) + with fully_shard_context(device=device): + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + fully_shard_optimizer(optimizer) + x = torch.randn(2, dim, device=device, dtype=dtype) + + def train_step() -> None: + optimizer.zero_grad(set_to_none=True) + model(x).float().sum().backward() + optimizer.step() + + child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() + resting_allocated = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + train_step() + peak_delta = torch.cuda.max_memory_allocated(device) - resting_allocated + + # Backward keeps the current child and one prefetched child unsharded. The current + # child also has a full wgrad until it is copied into a full reduce-scatter input, + # for a four-full-child-buffer peak. Allow one additional buffer for cuBLAS + # workspace, allocator granularity, and small temporaries. + bound_nbytes = (4 + 1) * child_weight_nbytes + + assert peak_delta < bound_nbytes, ( + "FSDP training-step peak memory exceeded the full-size-buffer bound: " + f"rank={rank}, peak_delta={_mb(peak_delta)}, " + f"five_full_child_buffers={_mb(bound_nbytes)}" + ) + + +def test_deleted_model_releases_fsdp_storage(distributed_setup): + """Deleting an FSDP model should release its persistent storage.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + + mesh = init_device_mesh(device.type, (world_size,)) + # Earlier tests may retain process-global CUDA allocations such as the + # CuBLAS workspace. Capture them before creating this model, so the test + # only detects storage retained by the deleted FSDP model itself. + allocated_before = torch.cuda.memory_allocated(device) + model = ElementwiseModel(dim=8192).to(dtype=torch.bfloat16, device=device) + with fully_shard_context(device=device): + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + x = torch.ones(1, 8192, dtype=torch.bfloat16, device=device) + output = model(x) + del output, x, model + + assert torch.cuda.memory_allocated(device) - allocated_before < 1024**2 + + +def test_fully_shard_returns_to_resting_memory(distributed_setup): + """Fully-sharded temporary storage should be released after forward and backward.""" + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(device.type, (world_size,)) + dim = 4096 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=2).to(dtype=dtype, device=device) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) + with fully_shard_context(device=device): + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + x = torch.randn(2, dim, device=device, dtype=dtype) + + def clear_cublas_workspaces_and_get_allocated_memory() -> int: + # PyTorch retains a cuBLAS workspace for each handle/stream pair. Clear those + # library caches so this measurement isolates FSDP-managed storage. + torch._C._cuda_clearCublasWorkspaces() + return torch.cuda.memory_allocated(device) + + resting_allocated = clear_cublas_workspaces_and_get_allocated_memory() + + def assert_returns_to_resting_memory(phase: str) -> None: + extra_allocated = clear_cublas_workspaces_and_get_allocated_memory() - resting_allocated + # The live output, activations, and root-owned bias gradient are small; unsharded + # parameter storage must be released. + assert extra_allocated < 1024**2, ( + f"Fully-sharded storage did not return to resting memory after {phase}: " + f"rank={rank}, extra_allocated={_mb(extra_allocated)}, " + "max_extra_allocated=1.00 MB" + ) + + output = model(x) + assert_returns_to_resting_memory("forward") + + loss = output.float().square().mean() + loss.backward() + del loss, output + assert_returns_to_resting_memory("backward") + + +def test_fully_shard_reduces_peak_training_memory(distributed_setup): + """Per-layer FSDP should reduce peak CUDA memory during training.""" + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + mesh = init_device_mesh(device.type, (world_size,)) + dim = 1024 + layers = 16 + batch = 8 + steps = 2 + dtype = torch.bfloat16 + + def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Tensor) -> None: + for _ in range(steps): + optimizer.zero_grad(set_to_none=True) + model(x).sum().backward() + optimizer.step() + + torch.manual_seed(4321) + baseline = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) + baseline_optimizer = torch.optim.AdamW(baseline.parameters(), lr=0.01) + x = torch.randn(batch, dim, device=device, dtype=dtype) + torch.cuda.reset_peak_memory_stats(device) + train_steps(baseline, baseline_optimizer, x) + baseline_peak = torch.cuda.max_memory_allocated(device) + + del baseline_optimizer + del baseline + del x + + torch.manual_seed(4321) + model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) + with fully_shard_context(device=device): + for layer in model: + fully_shard( + layer, + mesh=mesh, + placements=_flat_placements(), + mixed_precision_policy=MixedPrecisionPolicy( + main_params_dtype=dtype, main_grads_dtype=dtype + ), + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) + + x = torch.randn(batch, dim, device=device, dtype=dtype) + torch.cuda.reset_peak_memory_stats(device) + train_steps(model, optimizer, x) + sharded_peak = torch.cuda.max_memory_allocated(device) + logger.info( + "FSDP peak memory: rank=%s, baseline=%s, sharded=%s", + rank, + _mb(baseline_peak), + _mb(sharded_peak), + ) + + assert sharded_peak < baseline_peak From d5a019acccfad6afc3af90da2327a78d510e9c51 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 18 Aug 2026 18:34:44 +0000 Subject: [PATCH 081/101] Support PackedSeqParams CP CG as in #2282 (#6571) Signed-off-by: Teodor-Dumitru Ene --- megatron/core/transformer/cuda_graphs.py | 1 + tests/unit_tests/rl/test_rl_utils.py | 9 +-------- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index e800d2bfee3..b22eddc328e 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -338,6 +338,7 @@ def _check_supported_type(meta): _SUPPORTED_TYPES = { torch.Tensor, + torch.distributed.ProcessGroup, type(None), bool, int, diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index f3f43119ad2..915e1d9fc7f 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -1283,14 +1283,7 @@ def test_gpt_logprobs(self, initialize_model_parallel): pytest.param((1, 1), id="tp1-pp1"), pytest.param((2, 1), id="tp2-pp1"), pytest.param((1, 2), id="tp1-pp2"), - pytest.param( - (1, 1, 2), - id="tp1-pp1-cp2", - marks=pytest.mark.skip( - reason="PackedSeqParams.cp_group (a ProcessGroup) is rejected by the " - "CUDA-graph arg type whitelist; fixed and unskipped in a follow-up PR" - ), - ), + pytest.param((1, 1, 2), id="tp1-pp1-cp2"), ], indirect=["initialize_model_parallel"], ) From aaa4dfd9c4820c4c217ec2a5633c753753663979 Mon Sep 17 00:00:00 2001 From: Robert Kirby Date: Tue, 18 Aug 2026 19:11:42 +0000 Subject: [PATCH 082/101] Use named module logger in dynamic inference engine (#6001) Signed-off-by: Robert Kirby --- .../core/inference/engines/dynamic_engine.py | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index de6b342c89a..d429a8141eb 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -107,6 +107,9 @@ except ImportError: HAVE_PSUTIL = False +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + DEPRECATED_ARGS = [ "enable_cuda_graph", "random_seed", @@ -313,7 +316,7 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen # Configure wandb to use separate step counter for inference metrics (only once) if self.logging_step_interval > 0 and self.metrics_writer is not None: - logging.info( + logger.info( f"\033[1;93m[INFERENCE]\033[0m " f"\033[1;95mLogging inference metrics to wandb (rank {self.rank})\033[0m" ) @@ -527,9 +530,9 @@ def create_cuda_graphs(self, reset_context: bool = True): # Pool-scoped baselines for the per-iteration deltas. prev_pool_reserved, prev_pool_alloc = _cuda_graph_mempool_bytes() - logging.info("> dynamic_engine.py: building cuda graphs for ") + logger.info("> dynamic_engine.py: building cuda graphs for ") for graph in context.cuda_graph_batch_dimensions_list: - logging.info(graph) + logger.info(graph) # Enable inference dispatcher for EP during graph capture model_config = controller.inference_wrapped_model.model.config @@ -576,7 +579,7 @@ def create_cuda_graphs(self, reset_context: bool = True): if HAVE_TQDM: tbar.set_description(tbar_str) else: - logging.info( + logger.info( f"{tbar_idx}/{len(context.cuda_graph_batch_dimensions_list)}. {tbar_str}" ) @@ -626,7 +629,7 @@ def create_cuda_graphs(self, reset_context: bool = True): # This isolates pool growth from process-wide scratch churn (KV cache, # NCCL workspaces, etc.) that pollutes `torch.cuda.memory_stats()`. pool_reserved, pool_alloc = _cuda_graph_mempool_bytes() - logging.info( + logger.info( " [graph %d/%d] %s | pool reserved=%s (Δiter=%s) " "pool allocated=%s (Δiter=%s)", tbar_idx + 1, len(context.cuda_graph_batch_dimensions_list), @@ -639,7 +642,7 @@ def create_cuda_graphs(self, reset_context: bool = True): prev_pool_reserved, prev_pool_alloc = pool_reserved, pool_alloc if mtp_warmup_enabled and mtp_seen_batch_sizes: - logging.info("> MTP CUDA graph warmup: %d batch size(s)", len(mtp_seen_batch_sizes)) + logger.info("> MTP CUDA graph warmup: %d batch size(s)", len(mtp_seen_batch_sizes)) # Memory usage. time_end = time.time() @@ -652,7 +655,7 @@ def create_cuda_graphs(self, reset_context: bool = True): "pool_reserved_bytes": final_pool_reserved, "pool_allocated_bytes": final_pool_alloc, } - logging.info( + logger.info( "> built cuda graph(s) in %.2f sec. " "Mempool: reserved %s, allocated %s. " "Process-wide delta: allocated %s, reserved %s.", @@ -778,7 +781,7 @@ async def start_listening_to_data_parallel_coordinator( # Check if the port number is not inference_coordinator_port actual_port = int(dp_addr.rsplit(":", 1)[-1]) if inference_coordinator_port != None and actual_port != inference_coordinator_port: - logging.warning( + logger.warning( f"Requested InferenceCoordinator port {inference_coordinator_port} " f"but got port {actual_port} instead. This happens if the request port " f"is already in use." @@ -860,8 +863,8 @@ async def start_listening_to_data_parallel_coordinator( await await_process_call( coordinator_ready_event.wait, self.inference_coordinator_process ) - logging.info("Inference co-ordinator is ready to receive requests!") - logging.info(f"Data parallel coordinator can be found at {dp_addr}") + logger.info("Inference co-ordinator is ready to receive requests!") + logger.info(f"Data parallel coordinator can be found at {dp_addr}") # Finally run the engine infinite loop. loop = get_asyncio_loop(loop) @@ -924,7 +927,7 @@ def suspend_resume_ctx(key: str, *, unified_memory_level: int) -> None: f"res {end_mem_res / 1024**3:.1f} gb", ) ) - logging.info( + logger.info( f"[rank {rank_str}] dynamic engine {key}, " f"unified {unified_memory_level}, " f"{dir_str} " @@ -1039,7 +1042,7 @@ def resume(self): add_time = time.time() - add_time # Print inner timing (must be outside context manager above for correct formatting). - logging.info( + logger.info( " > " + ", ".join( ( @@ -1986,7 +1989,7 @@ def _register_cg_wait(self, req) -> None: """ req.cg_wait_iters += 1 if req.cg_wait_iters % self._cg_admission_warn_after == 0: - logging.warning( + logger.warning( "request %d has been deferred by CG-aware admission for %d steps — " "possible starvation (strict=%s, active P=%d D=%d tok=%d)", req.request_id, @@ -2676,7 +2679,7 @@ async def async_bookkeep( ) if color_decode_only: output_str = f"\033[94m{output_str}\033[0m" - logging.info(output_str) + logger.info(output_str) nvtx_range_pop("console_logging") @@ -3208,7 +3211,7 @@ async def run_engine_with_coordinator( elif self.state == EngineState.STOPPING: await self._world_barrier() if self.rank == 0: - logging.info("Stopping engine.") + logger.info("Stopping engine.") break finally: From aeae5d74e7fe029d10ac7ebffae6eef5c6a4c08a Mon Sep 17 00:00:00 2001 From: Tom Long Date: Tue, 18 Aug 2026 19:26:12 +0000 Subject: [PATCH 083/101] Correct the split-PR skill's stacking rules for copy-pr-bot CI (#6593) Signed-off-by: ilml Co-authored-by: Claude Opus 5 (1M context) --- skills/mcore-split-pr/SKILL.md | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/skills/mcore-split-pr/SKILL.md b/skills/mcore-split-pr/SKILL.md index 28c86f9a172..38dd13760ff 100644 --- a/skills/mcore-split-pr/SKILL.md +++ b/skills/mcore-split-pr/SKILL.md @@ -28,11 +28,24 @@ workflow: separate PR just to reduce reviewer groups. - If PR B depends on symbols renamed in PR A, call out the dependency and put backward-compatible aliases, re-exports, or shims in PR A when needed. -- When creating dependent PRs, set the dependent PR's GitHub base/diffbase to - `pull-request/`, not the base PR author's branch. -- Before merging a base PR, retarget each dependent PR back to `main` and - refresh it against `main`; otherwise GitHub may automatically close the - dependent PR, losing approvals and review discussion. +- [GitHub's standard stacked-PR flow](https://docs.github.com/en/pull-requests/how-tos/create-pull-requests/creating-stacked-pull-requests) + — push each branch to the upstream repo and base each PR on the previous + branch — does not work here: contributors cannot push branches to + `NVIDIA/Megatron-LM`, and a PR's base must be an upstream branch. The only + upstream refs containing a fork PR's commits are the `pull-request/` + mirrors that copy-pr-bot creates, so stacking goes through them. +- Create every PR with base `main`; the `pull-request/` mirror refs do not + exist until a vetter comments `/ok to test ` (copy-pr-bot). Once + the mirror exists, stack a dependent PR with + `gh pr edit --base pull-request/`. +- Never merge a PR while its base is a `pull-request/*` ref: the squash lands + in the bot's scratch ref, not `main`, and the PR ends up MERGED and + unreopenable. Retarget to `main` first. +- Retarget each dependent PR back to `main` as soon as its base PR is + approved: copy-pr-bot deletes `pull-request/` when PR N merges or + closes, and GitHub then auto-closes (unreopenably) every PR based on that + ref. After the base PR squash-merges, rebase dependents onto the new `main` + and force-push. - Wait for user approval before execution. - Execution creates draft PRs from the right base, applies file-scoped diffs with `git diff upstream/main.. -- | git apply`, pushes @@ -73,7 +86,7 @@ For each new PR: 1. Create a new branch from the appropriate local base (`main`, or a dependency PR's branch). 2. Extract the relevant changes: `git diff upstream/main.. -- | git apply`. 3. Stage, commit with a clear message, and push to the user's fork. -4. Create the PR as a **draft** (per repo contributing guidelines). For dependent PRs, set the GitHub base/diffbase to `pull-request/`. +4. Create the PR as a **draft** with base `main` (per repo contributing guidelines). Retarget dependent PRs to `pull-request/` only after a vetter's `/ok to test` has created that mirror ref. 5. If the original PR needs to be narrowed in scope, confirm with the user before force-pushing. 6. Report all PR URLs when done. @@ -81,7 +94,6 @@ For each new PR: - Always create PRs as **drafts** and push to the user's fork, never directly to upstream. - Backward-compatible changes (aliases, re-exports, deprecation shims) should go in the first PR so subsequent PRs can depend on them. -- Dependent PRs should target `pull-request/` while stacked, then be retargeted and refreshed to `main` before the base PR is merged. - Test files should go with the production code they test, not in a separate PR. - Prefer a single clean commit per split PR over replaying the original commit history. - If a file is hard to categorize (e.g., it touches two groups), ask the user which PR it should go in. From 5f9b2bb48c13d04803630a926cf71624dbd4722a Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Wed, 19 Aug 2026 00:16:25 +0000 Subject: [PATCH 084/101] Update copy-pr-bot.yaml [skip ci] --- .github/copy-pr-bot.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml index 9b33b89739d..43c02a2a2b0 100644 --- a/.github/copy-pr-bot.yaml +++ b/.github/copy-pr-bot.yaml @@ -1,4 +1,4 @@ enabled: true auto_sync_draft: false auto_sync_ready: true -trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "CarlosGomes98", "ChenhanYu", "Connor-XY", "DanialTaheri", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JF-D", "JRD971000", "Leili", "Mellonta", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "WanZzzzzz", "Wohox", "YangFei1990", "ZhiyuLi-Nvidia", "adistomar", "ahmadki", "aklife97", "alokpathy", "ananthsub", "anlthms", "aroshanghias-nvd", "ashehper", "asolergi-nv", "athitten", "balasaajay", "buptzyb", "chtruong814", "cjld", "cspades", "cuichenx", "deepakn94", "desh2608", "dimapihtar", "dingqingy-nv", "duncanriach", "ehosseiniasl", "erhoo82", "ericharper", "fanshiqing", "faradawn", "fitsumreda", "freewym", "frsun-nvda", "gautham-kollu", "gdengk", "goelarushi", "guihong-nv", "guyueh1", "hexinw-nvidia", "huvunvidia", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "niyunsheng", "ntajbakhsh", "nvcsathe", "parthmannan", "philipcmonk", "prajwal1210", "pthombre", "rapatel", "rhewett-nv", "rogerwaleffe", "sajadn", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "sharathts", "sheliang-nv", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sraman-rgb", "sudhakarsingh27", "svcnvidia-nemo-ci", "tdene", "theothermike", "thomasdhc", "tomlifu", "trintamaki", "tylerpoon", "vasunvidia", "wdykas", "wujingyue", "xiaoyao0115", "xuantengh", "xuwchen", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yqwangustc", "yueshen2016", "yuzhongw-nvidia", "zhehuaichen", "zhongbozhu"] +trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "CarlosGomes98", "ChenhanYu", "Connor-XY", "DanialTaheri", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JF-D", "JRD971000", "Leili", "Mellonta", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "WanZzzzzz", "Wohox", "YangFei1990", "ZhiyuLi-Nvidia", "adistomar", "ahmadki", "aklife97", "alokpathy", "ananthsub", "anlthms", "aroshanghias-nvd", "ashehper", "asolergi-nv", "athitten", "balasaajay", "buptzyb", "chtruong814", "cjld", "cspades", "cuichenx", "deepakn94", "desh2608", "dimapihtar", "dingqingy-nv", "duncanriach", "ehosseiniasl", "erhoo82", "ericharper", "fanshiqing", "faradawn", "fitsumreda", "freewym", "frsun-nvda", "gautham-kollu", "gdengk", "goelarushi", "guihong-nv", "guyueh1", "hexinw-nvidia", "huvunvidia", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwangjunahn", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "niyunsheng", "ntajbakhsh", "nvcsathe", "parthmannan", "philipcmonk", "prajwal1210", "pthombre", "rapatel", "rhewett-nv", "rogerwaleffe", "sajadn", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "sharathts", "sheliang-nv", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sraman-rgb", "sudhakarsingh27", "svcnvidia-nemo-ci", "tdene", "theothermike", "thomasdhc", "tomlifu", "trintamaki", "tylerpoon", "vasunvidia", "wdykas", "wujingyue", "xiaoyao0115", "xuantengh", "xuwchen", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yqwangustc", "yueshen2016", "yuzhongw-nvidia", "zhehuaichen", "zhongbozhu"] From 342faa43cd4c5bcd3297b59bf3cd96aaa3d1d3a7 Mon Sep 17 00:00:00 2001 From: Cory Ye <44509866+cspades@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:26:09 +0000 Subject: [PATCH 085/101] [MFSDPv1] Fix gradient sharding / allocation schedule with ZeRO-2, and fix wgrad accumulate fusion for ZeRO-1. (#6357) Signed-off-by: Cory Ye --- .../distributed/fsdp/mcore_fsdp_adapter.py | 6 +- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 69 ++++++++++++++----- 2 files changed, 56 insertions(+), 19 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 0b1c8a9d5f2..64cced2e593 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -169,7 +169,11 @@ def __init__( if fsdp_unit_modules is not None: self.fsdp_unit_modules = fsdp_unit_modules else: - if self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params": + # FSDP unit modules control the granularity of FSDP communications. + # "optim": Reduce-scatter communication groups on the final microbatch. + # "optim_grads": Additionally, RS communication groups on all microbatches. + # "optim_grads_params": RS & AG communication groups on all microbatches. + if self.ddp_config.data_parallel_sharding_strategy != "no_shard": self.fsdp_unit_modules = [TransformerLayer, MoETransformerLayer, MambaLayer] else: self.fsdp_unit_modules = [] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index f2df87ff256..f6ce9bba260 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -505,10 +505,12 @@ def all_gather_and_wait_parameters_ready( # un-shard of the FSDP params. param.__fsdp_param__ = True # Transformer Engine accumulates gradient on top of the `main_grad` - # buffer when gradient accumulation fusion in enabled. But with FSDP, - # we want to overwrite the `main_grad` which is enabled by this - # attribute. - param.overwrite_main_grad = True + # buffer when gradient accumulation fusion in enabled. When sharding + # the gradient, we want to overwrite the allocated gradient buffer. + param.overwrite_main_grad = self.data_parallel_sharding_strategy in [ + "optim_grads", + "optim_grads_params", + ] def _register_fsdp_hooks(self, root_module): """Register necessary hooks for Fully Sharded Data Parallel (FSDP) execution on the model. @@ -752,13 +754,25 @@ def _process_post_backward_gradients(param_list): for param in param_list: _grad_acc(param) + # Only reduce if we are sharding gradients, or are on the final microbatch. + # If is_last_microbatch is not specified, then we should reduce gradients + # if model_auto_sync is enabled, otherwise wait until is_last_microbatch + # is specified by the user, context manager, or FW before reduction. grad_reduce_every_bprop = self.data_parallel_sharding_strategy in [ "optim_grads", "optim_grads_params", ] is_last_microbatch = getattr(self, "is_last_microbatch", False) - if grad_reduce_every_bprop or is_last_microbatch or self.model_auto_sync: + if grad_reduce_every_bprop or ( + # "no_shard" and "optim" can have overlapped or non-overlapped + # gradient reduction. Non-overlapped gradient reduction happens + # in start_grad_sync() or synchronize_gradient_reduce(). + # Must only happen once on the final microbatch prior to optimization, + # or else unsharded gradient buffers will be redundantly reduced. + self.ddp_config.overlap_grad_reduce + and (is_last_microbatch or self.model_auto_sync) + ): # Launch asynchronous reduce-scatter of gradients before the optimizer # step. This requires a later call to finish_grad_sync() to wait for # completion. @@ -858,17 +872,24 @@ def _root_post_backward(*unused): for param in ordered_params: _grad_acc(param) - # Reduce the remaining gradients. - grad_reduce_every_bprop = self.data_parallel_sharding_strategy in [ - "optim_grads", - "optim_grads_params", - ] # Only reduce if we are sharding gradients, or are on the final microbatch. # If is_last_microbatch is not specified, then we should reduce gradients # if model_auto_sync is enabled, otherwise wait until is_last_microbatch # is specified by the user, context manager, or FW before reduction. + grad_reduce_every_bprop = self.data_parallel_sharding_strategy in [ + "optim_grads", + "optim_grads_params", + ] is_last_microbatch = getattr(self, "is_last_microbatch", False) - if grad_reduce_every_bprop or is_last_microbatch or self.model_auto_sync: + if grad_reduce_every_bprop or ( + # "no_shard" and "optim" can have overlapped or non-overlapped + # gradient reduction. Non-overlapped gradient reduction happens + # in start_grad_sync() or synchronize_gradient_reduce(). + # Must only happen once on the final microbatch prior to optimization, + # or else unsharded gradient buffers will be redundantly reduced. + self.ddp_config.overlap_grad_reduce + and (is_last_microbatch or self.model_auto_sync) + ): self.grad_reduce_pipeline.reduce_gradients( ordered_params, suggested_queue_capacity=self.suggested_RS_queue_capacity, @@ -954,6 +975,8 @@ def _root_pre_backward(module: nn.Module, *unused, skip_backward_hook: bool = Fa @torch.compiler.disable def _post_forward(module: nn.Module, input: Any, output: Any): + assert self.data_parallel_sharding_strategy == "optim_grads_params" + # When composed with module-hook-based activation recomputation, the # post-backward hook is responsible for resharding the module parameters # after the forward pass. In this case, the resharding is performed lazily. @@ -1065,12 +1088,15 @@ def _register_pre_backward_param_unshard_hook(module): if isinstance(module, tuple(fsdp_unit_modules)): fsdp_modules.append(module) - # Register the forward post-hook to reshard FSDP unit module parameters - # after the forward pass, except when recomputing forward activations, - # in which case we skip resharding for the subsequent backward pass. - self.forward_hooks[f"release module {name} parameters"] = ( - module.register_forward_hook(_post_forward, prepend=False) - ) + + if self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params": + # Register the forward post-hook to reshard FSDP unit module parameters + # after the forward pass, except when recomputing forward activations, + # in which case we skip resharding for the subsequent backward pass. + # Only relevant when using ZeRO-3 / FSDP. + self.forward_hooks[f"release module {name} parameters"] = ( + module.register_forward_hook(_post_forward, prepend=False) + ) _register_pre_backward_param_unshard_hook(module) elif ( @@ -1330,7 +1356,14 @@ def synchronize_gradient_reduce(self): self.grad_reduce_pipeline.wait_for_previous_grad_reduce(0) self.grad_reduce_pipeline.reset() else: - # Synchronous gradient all-reduce when sharding optimizer state or not sharding. + assert self.data_parallel_sharding_strategy not in [ + "optim_grads", + "optim_grads_params", + ], ( + "Cannot sync-reduce gradients for optim_grads / optim_grads_params!" + "Ensure that overlap_grad_reduce=True for ZeRO-2 and ZeRO-3." + ) + # Synchronous gradient reduction when sharding optimizer state or not sharding. self.start_grad_sync() def attach_grad_to_optimizer_state(self): From e5542e85f294fc215f2e08b3c7e17213f0ffae93 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Tue, 18 Aug 2026 22:48:59 +0000 Subject: [PATCH 086/101] fix(gtp+grad-reduce-in-fp32): emit the wgrad in main_grad's dtype so the reduce-scatter honours FP32 (#6624) Signed-off-by: Shiqing Fan --- megatron/core/tensor_parallel/layers.py | 48 ++++++++----- megatron/training/arguments.py | 17 +++++ .../test_gtp_basics.py | 69 +++++++++++++++++++ 3 files changed, 117 insertions(+), 17 deletions(-) diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index e842e5d9f39..a90b6e2b92e 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -551,6 +551,25 @@ def linear_with_frozen_weight( return LinearWithFrozenWeight.apply(*args) +def _wgrad_gemm(out, grad_output, total_input): + """Weight-gradient GEMM into ``out``, which may be wider than the inputs (bf16 -> fp32). + + Returns ``out``, filled with the weight gradient. + """ + # Import here to avoid circular import + from megatron.core.extensions.transformer_engine import te_general_gemm + + if te_general_gemm is not None: + # torch.matmul cannot widen via out=, so TE's GEMM does the mixed-precision output. + te_general_gemm( + total_input, grad_output, out_dtype=out.dtype, layout="NT", out=out, grad=True + ) + else: + # matmul rejects an out= of a different dtype, so land in the compute dtype and cast. + out.copy_(grad_output.t().matmul(total_input)) + return out + + class LinearWithGradAccumulationAndAsyncCommunication(torch.autograd.Function): """See linear_with_grad_accumulation_and_async_allreduce""" @@ -689,23 +708,7 @@ def backward(ctx, grad_output): # In case of Megatron-FSDP, need to create main grad buffers in-place if hasattr(weight, "__fsdp_param__"): weight.main_grad = weight.get_main_grad() - # Import here to avoid circular import - from megatron.core.extensions.transformer_engine import te_general_gemm - - if te_general_gemm is not None: - # Use TE general_gemm to support mixed-precision output - # (e.g. bf16 input -> fp32 main_grad) which torch.matmul - # does not support via the out= parameter. - te_general_gemm( - total_input, - grad_output, - out_dtype=weight.main_grad.dtype, - layout="NT", - out=weight.main_grad, - grad=True, - ) - else: - torch.matmul(grad_output.t(), total_input, out=weight.main_grad) + _wgrad_gemm(weight.main_grad, grad_output, total_input) else: if weight.main_grad.dtype == torch.float32: fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32( @@ -752,6 +755,17 @@ def backward(ctx, grad_output): weight.grad_added_to_main_grad = True else: grad_weight = None + elif ctx.gtp_remat_size > 1 and sharded_weight.main_grad.dtype != total_input.dtype: + # Fusion is off for GTP (main_grad is SHARDED, the GEMM output is not), so the wgrad + # takes three hops, and get_wgrad_tensor() types it from main_grad: + # + # epilogue RS accum + # --accumulate-allreduce-grads-in-fp32 fp32 ---> fp32 ---> fp32 <- here + # --grad-reduce-in-bf16 bf16 ---> bf16 ---> bf16 <- else + # + # This branch widens the epilogue to main_grad's dtype, so the RS no longer rounds + # across ranks before the fp32 accum sees the value. + grad_weight = _wgrad_gemm(sharded_weight.get_wgrad_tensor(), grad_output, total_input) else: grad_weight = grad_output.t().matmul(total_input) grad_bias = grad_output.sum(dim=0) if use_bias else None diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 3c9984a6b96..2f619468b86 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1197,6 +1197,23 @@ def validate_args(args, defaults={}): elif not args.accumulate_allreduce_grads_in_fp32 and args.main_grads_dtype == torch.float32: args.accumulate_allreduce_grads_in_fp32 = True print_rank_0('accumulate and all-reduce gradients in fp32 for bfloat16 data type.') + + if args.accumulate_allreduce_grads_in_fp32: + # The FP32-accumulation reduce-scatters only exist to recover FP32 accumulation from a + # lower-precision wire dtype. With FP32 main_grads the plain ring reduce-scatter already + # sums in FP32 on both axes, so they buy no precision and cost one extra + # unsharded-wgrad-sized scratch buffer per in-flight reduce-scatter. + for arg_name in ( + 'ddp_reduce_scatter_with_fp32_accumulation', + 'gtp_remat_reduce_scatter_with_fp32_accumulation', + ): + if getattr(args, arg_name, False): + setattr(args, arg_name, False) + warn_rank_0( + f"Setting args.{arg_name} to False since " + "--accumulate-allreduce-grads-in-fp32 already reduces in fp32" + ) + if args.cuda_graph_impl == "full_iteration": assert not args.check_for_nan_in_loss_and_grad, \ "--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-impl=full_iteration for training." diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py index 85e9b942d30..2c1b6d8903d 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py @@ -306,6 +306,75 @@ def test_forward_backward(self): _run_distributed(_worker_layernorm_linear, 4) +# --------------------------------------------------------------------------- +# Megatron linear: the wgrad handed to the reduce-scatter carries main_grad's dtype +# --------------------------------------------------------------------------- + + +def _worker_wgrad_rs_dtype(rank, world_size, port): + """With an FP32 main_grad, the gtp_remat reduce-scatter must run in FP32, not the compute dtype. + + Guards a SILENT failure: a bf16 wgrad makes the RS round at every rank, which leaves gradients + finite and merely less precise, so the value-based tests here pass either way. + + Covers Megatron's linear (layers.py); TE types its wgrad from main_grad via grad_buffer(). + """ + from megatron.core.tensor_parallel.generalized_tensor_parallelism import GTPShardedParam + from megatron.core.tensor_parallel.gtp_api import wrap_module_params_gtp + from megatron.core.tensor_parallel.layers import ( + linear_with_grad_accumulation_and_async_allreduce, + ) + + torch.manual_seed(0) + batch, in_f, out_f = 16, 64, 128 # out_f % (16*world_size)==0 -> no padding + layer = nn.Linear(in_f, out_f, bias=False, dtype=torch.bfloat16, device="cuda") + wrap_module_params_gtp(layer, ["weight"], dist.new_group(list(range(world_size)))) + layer.weight.main_grad = torch.zeros(layer.weight.shape, dtype=torch.float32, device="cuda") + + seen = [] + captured = [] + original = GTPShardedParam.wgrad_reduce_scatter + + def recording_rs(self, wgrad, nvtx_label=None): + seen.append(wgrad.dtype) + captured.append(wgrad.detach().float().clone()) + return original(self, wgrad, nvtx_label=nvtx_label) + + try: + GTPShardedParam.wgrad_reduce_scatter = recording_rs + inp = torch.randn(batch, in_f, dtype=torch.bfloat16, device="cuda", requires_grad=True) + linear_with_grad_accumulation_and_async_allreduce( + input=inp, + weight=layer.weight, + bias=None, + gradient_accumulation_fusion=True, + allreduce_dgrad=False, + sequence_parallel=False, + gtp_remat_size=world_size, + ).sum().backward() + finally: + GTPShardedParam.wgrad_reduce_scatter = original + + # Every rank asserts: the tally is rank-local. + assert seen, "wgrad_reduce_scatter was never reached -- the weight is not GTP-sharded" + assert set(seen) == {torch.float32}, ( + f"gtp_remat reduce-scatter ran in {set(seen)} with an FP32 main_grad. The wgrad must carry " + f"main_grad's dtype, or the cross-rank sum rounds before main_grad ever sees it." + ) + # The widened branch is the only caller of _wgrad_gemm here, so pin the VALUE too: a wrong + # layout/operand order would still be FP32 and still reduce-scatter cleanly. grad_output is + # all ones (loss = .sum()), so every wgrad row is just the input column sum -- computed in + # FP32 without a GEMM, the reference carries no BF16 rounding and the tolerance stays tight. + expected = inp.detach().float().sum(dim=0).expand(out_f, in_f) + torch.testing.assert_close(captured[0], expected, rtol=1e-5, atol=1e-4) + + +class TestGTPWgradRSDtype: + def test_wgrad_rs_follows_main_grad_dtype(self): + _requires_multi_gpu(4) + _run_distributed(_worker_wgrad_rs_dtype, 4) + + # --------------------------------------------------------------------------- # GroupedLinear forward/backward smoke test # --------------------------------------------------------------------------- From bf51cf46897374e62859b328798effcfb0ef05fe Mon Sep 17 00:00:00 2001 From: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:02:40 +0000 Subject: [PATCH 087/101] Honor high-priority streams in MIMO topology (#6604) Signed-off-by: ykarnati --- examples/mimo/pretrain_mimo.py | 2 +- examples/mimo/training/topology.py | 101 ++++++++++++------ tests/unit_tests/test_mimo_hetero_topology.py | 78 +++++++++++++- 3 files changed, 144 insertions(+), 37 deletions(-) diff --git a/examples/mimo/pretrain_mimo.py b/examples/mimo/pretrain_mimo.py index 9832cc806c8..bfe900d9a53 100644 --- a/examples/mimo/pretrain_mimo.py +++ b/examples/mimo/pretrain_mimo.py @@ -86,7 +86,7 @@ def main() -> None: # generic over any number of encoder grids in the topology. encoder_name = provider.encoder_module_names[0] if provider.encoder_module_names else None specs = build_module_grid_specs(args, args.world_size, encoder_name) - topology = create_topology(specs) + topology = create_topology(specs, args.high_priority_stream_groups) communicator = provider.build_communicator(args, topology) diff --git a/examples/mimo/training/topology.py b/examples/mimo/training/topology.py index 8146972da00..8634fb0b496 100644 --- a/examples/mimo/training/topology.py +++ b/examples/mimo/training/topology.py @@ -83,7 +83,9 @@ def destroy(self) -> None: grid.destroy() -def create_topology(specs: list[ModuleGridSpec]) -> HeteroTopology: +def create_topology( + specs: list[ModuleGridSpec], high_priority_stream_groups: Optional[list[str]] = None +) -> HeteroTopology: """Build every module's grid, PGC, and embedding groups. Exactly one spec must be named ``MIMO_LANGUAGE_MODULE_KEY`` (the language module); specs must @@ -103,11 +105,15 @@ def create_topology(specs: list[ModuleGridSpec]) -> HeteroTopology: module_pgs: dict[str, ProcessGroupCollection] = {} try: for spec in specs: - grids[spec.name] = _build_grid(spec) + grids[spec.name] = _build_grid(spec, high_priority_stream_groups) _validate_grid_layout(grids) for name, grid in grids.items(): - module_pgs[name] = pg_collection_from_grid(grid, is_language=(name == language_name)) + module_pgs[name] = pg_collection_from_grid( + grid, + is_language=(name == language_name), + high_priority_stream_groups=high_priority_stream_groups, + ) schedule_pg_collection = build_schedule_pg_collection(grids, module_pgs, language_name) return HeteroTopology( @@ -118,7 +124,9 @@ def create_topology(specs: list[ModuleGridSpec]) -> HeteroTopology: raise -def _build_grid(spec: ModuleGridSpec) -> HyperCommGrid: +def _build_grid( + spec: ModuleGridSpec, high_priority_stream_groups: Optional[list[str]] = None +) -> HyperCommGrid: """Create a dense grid plus its expert view and the process groups MIMO needs.""" grid = HyperCommGrid( shape=[spec.tp, spec.gtp_remat, spec.cp, spec.dp, spec.pp], @@ -134,39 +142,50 @@ def _build_grid(spec: ModuleGridSpec) -> HyperCommGrid: shared_dims=["pp"], ) + def create_pg(dims: list[str], pg_name: str, *, view: Optional[str] = None) -> None: + grid.create_pg( + dims, view=view, pg_options=_get_pg_options(pg_name, high_priority_stream_groups) + ) + try: - for dims in ( - ["tp"], - ["gtp_remat"], - ["cp"], - ["pp"], - ["dp"], - ["dp", "cp"], - ["gtp_remat", "dp", "cp"], - ["tp", "cp"], - ["tp", "gtp_remat", "pp"], - ["tp", "gtp_remat", "dp"], - ["tp", "gtp_remat", "dp", "cp"], - ["tp", "gtp_remat", "cp", "dp", "pp"], - ): - grid.create_pg(dims) - for dims in ( - ["ep"], - ["expt_tp"], - ["expt_gtp_remat"], - ["expt_dp"], - ["expt_tp", "ep"], - ["expt_tp", "ep", "pp"], - ["expt_tp", "ep", "expt_gtp_remat", "pp"], - ["expt_gtp_remat", "expt_dp"], - ): - grid.create_pg(dims, view=_EXPERT_VIEW) + create_pg(["tp"], "tp") + create_pg(["gtp_remat"], "gtp_remat") + create_pg(["cp"], "cp") + create_pg(["pp"], "pp") + create_pg(["dp"], "dp") + create_pg(["dp", "cp"], "dp_cp") + create_pg(["gtp_remat", "dp", "cp"], "gtp_remat_dp_cp") + create_pg(["tp", "cp"], "tp_cp") + create_pg(["tp", "gtp_remat", "pp"], "mp") + create_pg(["tp", "gtp_remat", "dp"], "tp_dp") + create_pg(["tp", "gtp_remat", "dp", "cp"], "tp_dp_cp") + create_pg(["tp", "gtp_remat", "cp", "dp", "pp"], "intra_dist_opt_instance") + + create_pg(["ep"], "ep", view=_EXPERT_VIEW) + create_pg(["expt_tp"], "ep_tp", view=_EXPERT_VIEW) + create_pg(["expt_gtp_remat"], "expt_gtp_remat", view=_EXPERT_VIEW) + create_pg(["expt_dp"], "ep_dp", view=_EXPERT_VIEW) + create_pg(["expt_tp", "ep"], "tp_ep_mp", view=_EXPERT_VIEW) + create_pg(["expt_tp", "ep", "pp"], "tp_ep_pp", view=_EXPERT_VIEW) + create_pg( + ["expt_tp", "ep", "expt_gtp_remat", "pp"], "tp_ep_gtp_remat_pp", view=_EXPERT_VIEW + ) + create_pg(["expt_gtp_remat", "expt_dp"], "ep_gtp_remat_dp", view=_EXPERT_VIEW) except Exception: grid.destroy() raise return grid +def _get_pg_options( + group_name: str, high_priority_stream_groups: Optional[list[str]] +) -> Optional[dist.ProcessGroupNCCL.Options]: + """Return high-priority NCCL options for a stock parallel-state group name.""" + if group_name not in (high_priority_stream_groups or ()): + return None + return dist.ProcessGroupNCCL.Options(is_high_priority_stream=True) + + def _validate_grid_layout(grids: dict[str, HyperCommGrid]) -> None: """Assert grids tile the world disjointly (non-colocated) XOR fully share ranks (colocated), with no gaps. Colocated-vs-not is decided via the core ``RankRole.build`` path. @@ -202,7 +221,9 @@ def _validate_grid_layout(grids: dict[str, HyperCommGrid]) -> None: def pg_collection_from_grid( - grid: HyperCommGrid, is_language: bool = False + grid: HyperCommGrid, + is_language: bool = False, + high_priority_stream_groups: Optional[list[str]] = None, ) -> ProcessGroupCollection: """Adapt a populated ``HyperCommGrid`` into a ``ProcessGroupCollection``. @@ -236,19 +257,29 @@ def pg_collection_from_grid( pgc.embd = None pgc.pos_embd = None if is_language: - _build_language_embedding_groups(grid, pgc) + _build_language_embedding_groups(grid, pgc, high_priority_stream_groups) return pgc -def _build_language_embedding_groups(grid: HyperCommGrid, pgc: ProcessGroupCollection) -> None: +def _build_language_embedding_groups( + grid: HyperCommGrid, + pgc: ProcessGroupCollection, + high_priority_stream_groups: Optional[list[str]] = None, +) -> None: """Set this rank's word/position embedding groups, mirroring parallel_state. Creation is collective: every grid rank calls ``new_group`` for each PP tuple. """ own_pp_ranks = tuple(dist.get_process_group_ranks(pgc.pp)) if pgc.pp is not None else () for pp_ranks in grid.get_rank_enum("pp"): - emb_group = dist.new_group(ranks=default_embedding_ranks(list(pp_ranks))) - pos_group = dist.new_group(ranks=default_position_embedding_ranks(list(pp_ranks))) + emb_group = dist.new_group( + ranks=default_embedding_ranks(list(pp_ranks)), + pg_options=_get_pg_options("embd", high_priority_stream_groups), + ) + pos_group = dist.new_group( + ranks=default_position_embedding_ranks(list(pp_ranks)), + pg_options=_get_pg_options("pos_embd", high_priority_stream_groups), + ) if tuple(pp_ranks) == own_pp_ranks: if _is_process_group_member(emb_group): pgc.embd = emb_group diff --git a/tests/unit_tests/test_mimo_hetero_topology.py b/tests/unit_tests/test_mimo_hetero_topology.py index e7e8551b5a0..1d9e555c6b2 100644 --- a/tests/unit_tests/test_mimo_hetero_topology.py +++ b/tests/unit_tests/test_mimo_hetero_topology.py @@ -5,11 +5,18 @@ Layout under test: encoder grid tp=2,dp=2 at ranks 0-3, language grid tp=2,pp=2 at ranks 4-7. """ +from types import SimpleNamespace + import pytest import torch import torch.distributed as dist -from examples.mimo.training.topology import ModuleGridSpec, _validate_grid_layout, create_topology +from examples.mimo.training.topology import ( + ModuleGridSpec, + _get_pg_options, + _validate_grid_layout, + create_topology, +) from megatron.core.hyper_comm_grid import HyperCommGrid from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY from megatron.core.parallel_state import default_embedding_ranks @@ -18,6 +25,15 @@ ENCODER = "images" +@pytest.fixture +def mock_nccl_options(monkeypatch): + monkeypatch.setattr( + dist, + "ProcessGroupNCCL", + SimpleNamespace(Options=lambda **kwargs: SimpleNamespace(**kwargs)), + ) + + def _specs(): return [ ModuleGridSpec(name=ENCODER, num_ranks=4, tp=2, rank_offset=0), @@ -63,6 +79,66 @@ def test_indivisible_expert_raises(self): ModuleGridSpec(name=ENCODER, num_ranks=4, tp=2, ep=3, expt_tp=2) +@pytest.mark.parametrize( + ("group_name", "requested", "enabled"), + [ + ("tp", ["tp"], True), + ("gtp_remat", ["gtp_remat"], True), + ("ep", ["ep"], True), + ("expt_gtp_remat", ["expt_gtp_remat"], True), + ("mp", ["tp"], False), + ("cp", ["cp"], True), + ("tp", None, False), + ], +) +def test_high_priority_pg_options(mock_nccl_options, group_name, requested, enabled): + options = _get_pg_options(group_name, requested) + + assert (options is not None) is enabled + if options is not None: + assert options.is_high_priority_stream + + +def test_high_priority_options_are_wired_to_each_module_grid(monkeypatch, mock_nccl_options): + class FakeGrid: + instances = [] + + def __init__(self, *args, **kwargs): + self.calls = [] + self.instances.append(self) + + def register_view(self, *args, **kwargs): + pass + + def create_pg(self, dims, view=None, pg_options=None): + self.calls.append((tuple(dims), view, pg_options)) + + monkeypatch.setattr("examples.mimo.training.topology.HyperCommGrid", FakeGrid) + monkeypatch.setattr("examples.mimo.training.topology._validate_grid_layout", lambda grids: None) + monkeypatch.setattr( + "examples.mimo.training.topology.pg_collection_from_grid", + lambda grid, is_language, high_priority_stream_groups: object(), + ) + monkeypatch.setattr( + "examples.mimo.training.topology.build_schedule_pg_collection", + lambda grids, module_pgs, language_name: object(), + ) + + create_topology(_specs(), ["tp", "gtp_remat_dp_cp", "ep_tp", "tp_ep_gtp_remat_pp"]) + + expected = { + (("tp",), None), + (("gtp_remat", "dp", "cp"), None), + (("expt_tp",), "expert"), + (("expt_tp", "ep", "expt_gtp_remat", "pp"), "expert"), + } + assert len(FakeGrid.instances) == 2 + for grid in FakeGrid.instances: + enabled = {(dims, view) for dims, view, options in grid.calls if options is not None} + assert enabled == expected + assert all(options.is_high_priority_stream for _, _, options in grid.calls if options) + + @pytest.mark.skipif(torch.cuda.device_count() < 8, reason="requires 8 GPUs") class TestHeteroTopology: def setup_method(self, method): From 0961f12e6f5f917b593a8281c74042f4189ff753 Mon Sep 17 00:00:00 2001 From: Charlie Truong Date: Tue, 18 Aug 2026 22:38:19 -0500 Subject: [PATCH 088/101] ci: Remove scheduled CI workflow (#6640) Signed-off-by: Charlie Truong --- .github/workflows/cicd-main.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 6162f67d239..6716359e78c 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -14,8 +14,6 @@ name: CICD Megatron-LM on: - schedule: - - cron: 0 0 * * * push: branches: - "pull-request/[0-9]+" From e7d58db4bc9a608a3eed9b9a4dbdf309d904b699 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Wed, 19 Aug 2026 05:38:31 +0200 Subject: [PATCH 089/101] fix(ci): request pull request permission for queue comments (#6641) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- .github/workflows/cicd-main.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 6716359e78c..9e4eaef7969 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -1497,7 +1497,7 @@ jobs: with: app-id: ${{ vars.COMMUNITY_BOT_ID }} private-key: ${{ secrets.COMMUNITY_BOT_KEY }} - permission-issues: write + permission-pull-requests: write - name: Comment on PR with action run URL uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8 From 29361d65827db246f79b996e5205b10f77b2c913 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Tue, 18 Aug 2026 23:38:43 -0400 Subject: [PATCH 090/101] Remove core teams from FSDP owners (#6623) Signed-off-by: Philip Petrakian --- .github/CODEOWNERS | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index d939d68729c..7482add19c1 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -24,7 +24,7 @@ megatron/core/datasets/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/datasets megatron/core/tokenizers/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/tokenizers megatron/core/distributed/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/data-parallelism -megatron/core/distributed/fsdp/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/megatron-fsdp +megatron/core/distributed/fsdp/ @NVIDIA/megatron-fsdp megatron/core/transformer/fsdp_dtensor_checkpoint.py @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/megatron-fsdp From c452d5a7ee3301f279c6dce6398b2739f1681f77 Mon Sep 17 00:00:00 2001 From: Tom Long Date: Wed, 19 Aug 2026 00:10:48 +0000 Subject: [PATCH 091/101] Add THD-format per-rank batch fetch with padding mask for sequence packing (#6625) Signed-off-by: ilml Co-authored-by: Claude Fable 5 --- megatron/core/datasets/data_schedule.py | 248 +++++++++++++++- tests/unit_tests/test_sequence_packing.py | 331 ++++++++++++++++++++++ 2 files changed, 576 insertions(+), 3 deletions(-) create mode 100644 tests/unit_tests/test_sequence_packing.py diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 26d083c3037..fa0026f3937 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -1,14 +1,71 @@ # Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional import torch from megatron.core import parallel_state -from megatron.core.datasets.data_schedule_utils import _get_global_seqlens_and_ids +from megatron.core.datasets.data_schedule_utils import ( + _get_global_seqlens_and_ids, + broadcast_scalars, + broadcast_tensor, +) +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler from megatron.core.process_groups_config import ProcessGroupCollection +try: + # Register the TE CUDA kernels + import transformer_engine # pylint: disable=unused-import + + # Alias the PyTorch wrapper so we can call tex.* APIs + import transformer_engine_torch as tex +except ImportError: + # TE isn't installed or the torch wrapper is missing + tex = None + + +def _build_thd_padding_mask( + cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.Tensor +) -> torch.Tensor: + """Build a 1D THD padding mask from scheduler sequence metadata.""" + assert cu_seqlens.dim() == 1 + assert cu_seqlens_padded.dim() == 1 + assert cu_seqlens.numel() == cu_seqlens_padded.numel() + + total_tokens = int(cu_seqlens_padded[-1].item()) + if total_tokens == 0: + return torch.empty((0,), dtype=torch.bool, device=cu_seqlens.device) + + num_sequences = cu_seqlens.numel() - 1 + if num_sequences <= 0: + return torch.ones((total_tokens,), dtype=torch.bool, device=cu_seqlens.device) + + positions = torch.arange( + total_tokens, dtype=cu_seqlens_padded.dtype, device=cu_seqlens_padded.device + ) + seq_indices = torch.searchsorted(cu_seqlens_padded[1:].contiguous(), positions, right=True) + + valid_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).clamp(min=0) + valid_ends = cu_seqlens_padded[:-1] + valid_lengths + return positions >= valid_ends[seq_indices] + + +def _sanitize_thd_padding_values(batch: Dict[str, Any], padding_mask: torch.Tensor) -> None: + """Replace padded token-like slots with safe neutral values in-place.""" + assert padding_mask.dim() == 1 + pad_values = {'tokens': 0, 'labels': 0, 'loss_mask': 0.0, 'position_ids': 0} + for key, pad_value in pad_values.items(): + tensor = batch.get(key) + if tensor is None: + continue + assert tensor.dim() == 1, f"{key} must be 1D before CP slicing, got {tensor.dim()}D" + assert tensor.numel() == padding_mask.numel(), ( + f"{key} length ({tensor.numel()}) must match padding_mask length " + f"({padding_mask.numel()}) before CP slicing." + ) + batch[key] = tensor.masked_fill(padding_mask, pad_value) + class HybridCPDataLoaderWrapper: """ @@ -58,7 +115,6 @@ def get_global_seqlens(self, subsample_seqlens: torch.Tensor) -> List[int]: Gathers the sequence lengths of all subsamples from all DP ranks. Each DP rank loads the same number of microbatches but each microbatch may have a different number of subsamples. - We find the number of subsamples each rank holds and then gather the sequence lengths of all subsamples from all ranks. @@ -288,3 +344,189 @@ def __next__(self) -> Any: batch, global_ids_this_rank, global_id_seqlens, sample_id_groups, offsets ) return samples_this_rank_with_id, sample_id_groups + + +def get_batch_on_this_rank_for_sequence_packing( + data_iterator, + vpp_size: Optional[int] = None, + mtp_on_this_rank: bool = False, + vp_stage: Optional[int] = None, + pg_collection: Optional[ProcessGroupCollection] = None, +): + """ + Get a batch of data for sequence packing. + Args: + data_iterator (Iterator): The data iterator to get the batch from. + mtp_on_this_rank (bool): Whether to use multi-token prediction. + vp_stage (Optional[int]): The stage of the pipeline. + Returns: + tuple of (tokens, labels, loss_mask, attention_mask, position_ids, + packed_seq_params, padding_mask) + """ + + if pg_collection is None: + tp_group = parallel_state.get_tensor_model_parallel_group() + pp_group = parallel_state.get_pipeline_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + else: + tp_group = pg_collection.tp + pp_group = pg_collection.pp + cp_group = pg_collection.cp + + tp_src_rank = torch.distributed.get_process_group_ranks(tp_group)[0] + + is_tp_rank_0 = tp_group.rank() == 0 + is_first_stage = pp_group.rank() == 0 and (vp_stage is None or vp_stage == 0) + is_last_stage = pp_group.rank() == pp_group.size() - 1 and ( + vp_stage is None or vp_stage == vpp_size - 1 + ) + + is_first_or_last_stage = is_first_stage or is_last_stage + dev = torch.cuda.current_device() + + # data_iterator should return a batch including the following keys. + batch_keys = ['cu_seqlens', 'cu_seqlens_padded', 'max_seqlen'] + if is_first_stage: + batch_keys.append('tokens') + batch_keys.append('position_ids') + if is_last_stage: + batch_keys.append('labels') + batch_keys.append('loss_mask') + + # Get a batch from data_iterator or create an emtpy batch. + if is_tp_rank_0: + assert data_iterator is not None + batch = next(data_iterator) + for key in batch_keys: + assert key in batch, f"{key} is missing in current batch." + else: + assert data_iterator is None, "Non TP 0 rank should not have data_iterator" + batch = {} + + # Build padding_mask before CP slicing while tensors still have the full + # packed length represented by cu_seqlens_padded[-1]. + if is_tp_rank_0: + batch['padding_mask'] = _build_thd_padding_mask( + batch['cu_seqlens'], batch['cu_seqlens_padded'] + ) + _sanitize_thd_padding_values(batch, batch['padding_mask']) + + # Partition padding_mask for context parallel on every PP stage. Partition + # token-like tensors only on stages that own them. + if is_tp_rank_0: + cp_size = cp_group.size() + cp_rank = cp_group.rank() + # If cp_size == 1, no need to do further processing. + if cp_size > 1: + # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as + # cu_seqlens to get the correct result. + # TODO: Revert this workaround once TE fixes the issue. + cu_seqlens = batch["cu_seqlens_padded"] + total_tokens = int(cu_seqlens[-1].item()) + assert ( + tex is not None + ), "Transformer Engine is required to use Context Parallel with THD format data." + index = tex.thd_get_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) + cp_slice_keys = ['padding_mask'] + if is_first_or_last_stage: + cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask']) + for key in cp_slice_keys: + batch[key] = batch[key].index_select(0, index) + + # Broadcast the receive-buffer shapes inside the TP group: + # - cu_seqlen_size is needed to allocate cu_seqlens / cu_seqlens_padded on non TP 0 ranks. + # - total_tokens is needed because padding_mask is prepared on every PP stage, and + # tokens/labels/loss_mask/position_ids use the same length on stages that own them. + shapes = ( + [batch['cu_seqlens'].size(0), batch['padding_mask'].size(0)] if is_tp_rank_0 else [0, 0] + ) + cu_seqlen_size, total_tokens = broadcast_scalars(shapes, tp_group, dev, dtype=torch.int32) + + # Step1: Prepare "tokens", "position_ids" on all ranks. + if is_first_stage or mtp_on_this_rank: + if is_tp_rank_0: + assert batch['tokens'].dtype == torch.int64 + assert batch['position_ids'].dtype == torch.int64 + batch['tokens'] = batch['tokens'].view(1, total_tokens) + batch['position_ids'] = batch['position_ids'].view(1, total_tokens) + else: + batch['tokens'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + batch['position_ids'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + else: + # Non first stage rank doesn't need tokens and position_ids. + batch['tokens'] = None + batch['position_ids'] = None + + # Step2: Prepare "labels", "loss_mask" on all ranks. + if is_last_stage: + if is_tp_rank_0: + assert batch['labels'].dtype == torch.int64 + assert batch['loss_mask'].dtype == torch.float32 + batch['labels'] = batch['labels'].view(1, total_tokens) + batch['loss_mask'] = batch['loss_mask'].view(1, total_tokens) + else: + batch['labels'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev) + batch['loss_mask'] = torch.empty([1, total_tokens], dtype=torch.float32, device=dev) + else: + # Non last stage rank doesn't need labels and loss_mask. + batch['labels'] = None + batch['loss_mask'] = None + + # Step3: Prepare "padding_mask" on all TP ranks. + if is_tp_rank_0: + assert batch['padding_mask'].dtype == torch.bool + batch['padding_mask'] = batch['padding_mask'].view(1, total_tokens) + else: + batch['padding_mask'] = torch.empty([1, total_tokens], dtype=torch.bool, device=dev) + + # Step4: Prepare "cu_seqlens", "cu_seqlens_padded", "max_seqlen" on all ranks. + if is_tp_rank_0: + assert batch['cu_seqlens'].dtype == torch.int32 + assert batch['cu_seqlens_padded'].dtype == torch.int32 + assert batch['cu_seqlens'].dim() == 1 + assert batch['cu_seqlens_padded'].dim() == 1 + if type(batch['max_seqlen']) == int: + batch['max_seqlen'] = torch.tensor(batch['max_seqlen'], dtype=torch.int32, device=dev) + else: + assert batch['max_seqlen'].dtype == torch.int32 + assert batch['max_seqlen'].numel() == 1 + else: + batch['cu_seqlens'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev) + batch['cu_seqlens_padded'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev) + batch['max_seqlen'] = torch.empty(1, dtype=torch.int32, device=dev) + + # Broadcast batch inside TP group. + broadcast_tensor(batch['tokens'], tp_src_rank, tp_group) + broadcast_tensor(batch['position_ids'], tp_src_rank, tp_group) + broadcast_tensor(batch['labels'], tp_src_rank, tp_group) + broadcast_tensor(batch['loss_mask'], tp_src_rank, tp_group) + broadcast_tensor(batch['padding_mask'], tp_src_rank, tp_group) + broadcast_tensor(batch['cu_seqlens'], tp_src_rank, tp_group) + broadcast_tensor(batch['cu_seqlens_padded'], tp_src_rank, tp_group) + broadcast_tensor(batch['max_seqlen'], tp_src_rank, tp_group) + + # Extract the data from batch after broadcasting. + tokens = batch['tokens'] + position_ids = batch['position_ids'] + labels = batch['labels'] + loss_mask = batch['loss_mask'] + padding_mask = batch['padding_mask'] + cu_seqlens = batch['cu_seqlens'] + cu_seqlens_padded = batch['cu_seqlens_padded'] + max_seqlen = batch['max_seqlen'].item() + + # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as cu_seqlens to + # get the correct result. + # TODO: Revert this workaround once TE fixes the issue. + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_padded, + cu_seqlens_kv=cu_seqlens_padded, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + ) + + # "attention_mask" is not valid for sequence packing, so set it to None. + return tokens, labels, loss_mask, None, position_ids, packed_seq_params, padding_mask diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py new file mode 100644 index 00000000000..99520e1c4ca --- /dev/null +++ b/tests/unit_tests/test_sequence_packing.py @@ -0,0 +1,331 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.datasets.data_schedule import ( + _build_thd_padding_mask, + _sanitize_thd_padding_values, + get_batch_on_this_rank_for_sequence_packing, +) +from megatron.training.global_vars import unset_global_variables +from tests.unit_tests.test_utilities import Utils + + +def test_scheduler_thd_padding_mask_from_cu_seqlens(): + cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32) + cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32) + + padding_mask = _build_thd_padding_mask(cu_seqlens, cu_seqlens_padded) + + assert torch.equal( + padding_mask, torch.tensor([False, False, False, True, False, False, True, True]) + ) + + +def test_scheduler_sanitizes_thd_padding_values(): + padding_mask = torch.tensor([False, False, True, False, True]) + batch = { + 'tokens': torch.tensor([11, 12, -1, 21, -1], dtype=torch.int64), + 'labels': torch.tensor([12, 13, -1, 22, -1], dtype=torch.int64), + 'loss_mask': torch.ones(5, dtype=torch.float32), + 'position_ids': torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64), + } + + _sanitize_thd_padding_values(batch, padding_mask) + + assert torch.equal(batch['tokens'], torch.tensor([11, 12, 0, 21, 0])) + assert torch.equal(batch['labels'], torch.tensor([12, 13, 0, 22, 0])) + assert torch.equal(batch['loss_mask'], torch.tensor([1.0, 1.0, 0.0, 1.0, 0.0])) + assert torch.equal(batch['position_ids'], torch.tensor([0, 1, 0, 0, 0])) + + +class MockVariableLengthSequencePackingDataIterator: + """ + Mock data iterator for testing get_batch_on_this_rank_for_sequence_packing. + + Generates variable-length (THD format) packed sequences with deterministic + data for verification across parallel ranks. + """ + + def __init__( + self, + total_seq_length: int, + sequence_lengths: list, + local_cp_size: int = None, + device: str = "cuda", + seed: int = 42, + ): + """ + Args: + total_seq_length: Total length of packed sequences + sequence_lengths: List of individual sequence lengths (variable-length). + If None, generates random variable lengths. + device: Device to create tensors on + seed: Random seed for reproducibility + """ + self.total_seq_length = total_seq_length + self.sequence_lengths = sequence_lengths + self.local_cp_size = local_cp_size + self.device = device + self.seed = seed + assert ( + sum(self.sequence_lengths) == total_seq_length + ), f"Sequence lengths sum {sum(self.sequence_lengths)} != total {total_seq_length}" + + def __iter__(self): + """Interface for the data iterator.""" + return self + + def __next__(self): + """Generate a mock batch with variable-length THD format.""" + dev = self.device + torch.manual_seed(self.seed) + torch.cuda.manual_seed(self.seed) + + tokens = torch.randint(0, 16384, (self.total_seq_length,), dtype=torch.int64, device=dev) + + # Create position_ids that reset for each sequence (THD format) + position_ids = [] + for seq_len in self.sequence_lengths: + position_ids.extend(range(seq_len)) + position_ids = torch.tensor(position_ids, dtype=torch.int64, device=dev) + + # Labels are tokens shifted by 1 for easy verification + labels = tokens + 1 + + # Loss mask: 1.0 for all positions except padding (none here) + loss_mask = torch.ones(self.total_seq_length, dtype=torch.float32, device=dev) + + # Create cu_seqlens for variable-length packed sequences + cu_seqlens = [0] + for seq_len in self.sequence_lengths: + cu_seqlens.append(cu_seqlens[-1] + seq_len) + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=dev) + cu_seqlens_padded = cu_seqlens.clone() + + max_seqlen = torch.tensor([max(self.sequence_lengths)], dtype=torch.int32, device=dev) + + batch = { + "tokens": tokens, + "position_ids": position_ids, + "labels": labels, + "loss_mask": loss_mask, + "cu_seqlens": cu_seqlens, + "cu_seqlens_padded": cu_seqlens_padded, + "max_seqlen": max_seqlen, + } + + if not ( + parallel_state.is_pipeline_first_stage(ignore_virtual=True) + or parallel_state.is_pipeline_last_stage(ignore_virtual=True) + ): + batch["tokens"] = None + batch["position_ids"] = None + batch["labels"] = None + batch["loss_mask"] = None + + if self.local_cp_size is not None: + batch["local_cp_size"] = torch.tensor( + [self.local_cp_size], dtype=torch.int32, device=dev + ) + + return batch + + +def _gather_tensor_from_tp_group(tensor): + """Gather tensors from all TP ranks for comparison.""" + assert tensor is not None, "Tensor should not be None" + tp_size = parallel_state.get_tensor_model_parallel_world_size() + gathered = [torch.zeros_like(tensor) for _ in range(tp_size)] + torch.distributed.all_gather( + gathered, tensor, group=parallel_state.get_tensor_model_parallel_group() + ) + return gathered + + +def _gather_tensor_from_all_ranks(tensor): + """Gather tensors from all PP ranks for comparison.""" + assert tensor is not None, "Tensor should not be None" + if type(tensor) is int: + tensor = torch.tensor(tensor, dtype=torch.int32, device=torch.cuda.current_device()) + gathered = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())] + torch.distributed.all_gather(gathered, tensor) + return gathered + + +@pytest.mark.parametrize( + ("tp", "pp", "cp"), + [ + (1, 1, 1), # Basic case: no parallelism + (2, 1, 1), # Tensor parallel only + (1, 2, 1), # Pipeline parallel only + (2, 2, 1), # TP + PP + (1, 1, 2), # CP only + (2, 1, 2), # TP + CP + (1, 2, 2), # PP + CP + (1, 4, 1), # Has middle pp stage + ], +) +def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp): + """ + Test get_batch_on_this_rank_for_sequence_packing function with variable-length THD format. + + This test verifies: + 1. TP ranks: All ranks within a TP group receive identical data after broadcast + 2. PP ranks: Middle PP ranks have the same packed_seq_params as first/last stages + 3. CP ranks: Data is correctly partitioned with proper shape and values + 4. Variable-length (THD) format: Different sequence lengths are handled correctly + """ + args = SimpleNamespace() + args.tensor_model_parallel_size = tp + args.pipeline_model_parallel_size = pp + args.context_parallel_size = cp + args.virtual_pipeline_model_parallel_size = None + args.data_parallel_size = 8 // (tp * pp * cp) + args.seq_length = 8192 + + # Skip invalid configurations + if args.data_parallel_size < 1: + raise ValueError(f"Invalid config: tp={tp}, pp={pp}, cp={cp} exceeds world size 8") + + # Initialize model parallel + Utils.initialize_model_parallel(tp, pp, None, context_parallel_size=cp) + + try: + # Create mock data iterator with variable-length sequences + # Only TP rank 0 needs the iterator; other TP ranks pass None + tp_rank = parallel_state.get_tensor_model_parallel_rank() + if tp_rank == 0: + # Use deterministic seed based on DP rank so same data within TP/PP/CP group + dp_rank = parallel_state.get_data_parallel_rank() + sequence_lengths = [1024, 2048, 512, 1536, 3072] + assert ( + sum(sequence_lengths) == args.seq_length + ), f"Sequence lengths sum {sum(sequence_lengths)} != total {args.seq_length}" + data_iterator = iter( + MockVariableLengthSequencePackingDataIterator( + total_seq_length=args.seq_length, + sequence_lengths=sequence_lengths, # Variable lengths, sum=8192 + seed=42 + dp_rank, # Same seed within PP/CP group + ) + ) + else: + # Non-TP-rank-0 ranks don't need the iterator + data_iterator = None + + # Call the function under test + result = get_batch_on_this_rank_for_sequence_packing( + data_iterator=data_iterator, mtp_on_this_rank=False, vp_stage=None + ) + + # Unpack the result. Scheduler THD always returns padding_mask. + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = ( + result + ) + + # Get parallel state info + tp_rank = parallel_state.get_tensor_model_parallel_rank() + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + cp_rank = parallel_state.get_context_parallel_rank() + is_first_stage = parallel_state.is_pipeline_first_stage(ignore_virtual=True) + is_last_stage = parallel_state.is_pipeline_last_stage(ignore_virtual=True) + is_first_or_last = is_first_stage or is_last_stage + + assert padding_mask is not None + assert padding_mask.dtype == torch.bool + assert padding_mask.dim() == 2 + assert padding_mask.size(0) == 1 + assert not padding_mask.any(), "Mock data has no per-sequence padding." + + # ===================================================================== + # TEST 1: Verify data based on pipeline stage + # ===================================================================== + if is_first_stage: + assert tokens is not None, "First stage should have tokens" + assert position_ids is not None, "First stage should have position_ids" + assert tokens.dim() == 2, "Tokens should be 2D (batch, seq)" + assert position_ids.dim() == 2, "Position IDs should be 2D (batch, seq)" + assert tokens.size(0) == 1, "batch should be 1 in THD format" + assert position_ids.size(0) == 1, "batch should be 1 in THD format" + else: + assert tokens is None, "Non-first stage should not have tokens" + assert position_ids is None, "Non-first stage should not have position_ids" + + if is_last_stage: + assert labels is not None, "Last stage should have labels" + assert loss_mask is not None, "Last stage should have loss_mask" + assert labels.dim() == 2, "Labels should be 2D (batch, seq)" + assert loss_mask.dim() == 2, "Loss mask should be 2D (batch, seq)" + assert labels.size(0) == 1, "batch should be 1 in THD format" + assert loss_mask.size(0) == 1, "batch should be 1 in THD format" + else: + assert labels is None, "Non-last stage should not have labels" + assert loss_mask is None, "Non-last stage should not have loss_mask" + + # ===================================================================== + # TEST 2: Verify all ranks have consistent packed_seq_params + # ===================================================================== + assert packed_seq_params is not None + assert packed_seq_params.qkv_format == "thd" + + test_keys = [ + "cu_seqlens_q", + "cu_seqlens_q_padded", + "max_seqlen_q", + "cu_seqlens_kv", + "cu_seqlens_kv_padded", + "max_seqlen_kv", + ] + for key in test_keys: + tensor = getattr(packed_seq_params, key) + assert tensor is not None + gathered_tensor = _gather_tensor_from_all_ranks(tensor) + for i in range(1, len(gathered_tensor)): + assert torch.equal( + gathered_tensor[0], gathered_tensor[i] + ), f"Rank 0 and rank {i} have different {key}" + + # ===================================================================== + # TEST 3: Verify TP ranks receive identical data after broadcast + # ===================================================================== + if tp > 1: + test_tensors = [] + if is_first_stage: + test_tensors.extend([tokens, position_ids]) + if is_last_stage: + test_tensors.extend([labels, loss_mask]) + + for tensor in test_tensors: + gathered_tensors = _gather_tensor_from_tp_group(tensor) + for i in range(1, tp): + assert torch.equal( + gathered_tensors[0], gathered_tensors[i] + ), f"TP rank 0 and rank {i} have different data" + + # ===================================================================== + # TEST 4: Verify CP partitioning + # ===================================================================== + if cp > 1: + # With CP, the sequence should be partitioned + expected_seq_len = args.seq_length // cp + + if is_first_stage: + actual_seq_len = tokens.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned tokens have wrong shape: {actual_seq_len} != {expected_seq_len}" + + # Verify labels only if all CP ranks are at last stage + if is_last_stage: + actual_seq_len = labels.shape[1] + assert ( + actual_seq_len == expected_seq_len + ), f"CP partitioned labels have wrong shape: {actual_seq_len} != {expected_seq_len}" + + finally: + Utils.destroy_model_parallel() + unset_global_variables() From 3ad8c2a6b13d9841d6cc6123f791af446e4916bc Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 19 Aug 2026 00:13:35 +0000 Subject: [PATCH 092/101] Share prefill in sequence packing (#5887) Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Vitaly Kurin Co-authored-by: Jorge Albericio --- megatron/rl/agent/api.py | 1 + megatron/rl/rl_utils.py | 322 +++++++++------ megatron/rl/sequence_packing_utils.py | 41 +- tests/unit_tests/rl/test_rl_utils.py | 369 +++++++++++++----- .../rl/test_sequence_packing_utils.py | 36 ++ 5 files changed, 545 insertions(+), 224 deletions(-) diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index 8307015b240..9859d0ac4b2 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -64,6 +64,7 @@ class TokenRollout(AgentBaseModel): env_id: str = '' problem_id: str | None = None completion_ids: list[str] = [] + generation_cap: int | None = None rollout_status: str = 'ok' failure_reason: str | None = None diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index f1642095ef9..4018badc342 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -16,7 +16,17 @@ from dataclasses import dataclass from functools import partial from pathlib import Path -from typing import Any, AsyncIterator, Dict, Iterator, List, NamedTuple, Optional, TypeAlias +from typing import ( + Any, + AsyncIterator, + Dict, + Iterator, + List, + NamedTuple, + Optional, + Sequence, + TypeAlias, +) import numpy as np import torch @@ -536,10 +546,7 @@ def align_unpacked_inference_logprobs( Returns: Aligned inference logprobs tensor """ - # Get first occurrence of a generation token - # In get_logprobs() we chop off the first token -> the generation mask is shifted by one gen_masks_for_alignment = generation_masks - first_gen_tok = gen_masks_for_alignment.int().argmax(dim=1) - 1 # Align inference logprobs with old_logprobs # Note: We use old_logprobs_for_data as template since it has correct shape. @@ -548,21 +555,24 @@ def align_unpacked_inference_logprobs( # old_logprobs arrives bf16 while the wire logprobs are float32 — cloning # the template dtype would silently downcast them in the assignment below. padded_inference_logprobs = old_logprobs_for_data.clone().float() + width = padded_inference_logprobs.shape[1] + + # Inference logprobs cover only generated tokens; multi-turn rollouts have multiple regions. + # Scatter each logprob: the logprob of the token at position p lives at index p-1. - # We need to align old_logprobs and inference logprobs as the latter are only for generations for i, inf_logprobs in enumerate(inference_logprobs): - if not gen_masks_for_alignment[i].any(): - # No generation tokens; nothing to align. - continue - first_gen_idx = int(first_gen_tok[i]) - # We subtract -1 here because we append eod token on the train side, and we do not - # get it from the inference. For the eod token, we reuse old_logprobs value. - end_idx = min(first_gen_idx + len(inf_logprobs), padded_inference_logprobs.shape[1]) - actual_len = end_idx - first_gen_idx - if actual_len > 0: - padded_inference_logprobs[i, first_gen_idx:end_idx] = inf_logprobs[:actual_len].to( - padded_inference_logprobs.dtype - ) + gen_positions = gen_masks_for_alignment[i].nonzero(as_tuple=True)[0] + if gen_positions.numel() == 0: + continue # No generation tokens; nothing to align. + target = gen_positions - 1 + # If inference returns fewer logprobs than generated tokens (e.g. an eod appended on + # the train side), the trailing generated slots keep their old_logprobs value. + n = min(len(inf_logprobs), target.numel()) + target = target[:n] + in_range = (target >= 0) & (target < width) + padded_inference_logprobs[i, target[in_range]] = inf_logprobs[:n][in_range].to( + padded_inference_logprobs.dtype + ) # Create truncated mask for statistics if old_logprobs_for_data.shape[1] + 1 < gen_masks_for_alignment.shape[1]: @@ -1087,13 +1097,29 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa return logprobs -def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[int]]) -> np.ndarray: - """Calculate GRPO advantages from rewards/num_turns. +def calculate_grpo_advantages( + rewards: list[list[float]], num_turns: list[list[int]] +) -> list[float]: + """Calculate GRPO advantages from rewards. + + Each multi-turn episode is trained as a single trajectory with a single advantage. + The result is flattened group-major, matching the order of `[r for g in rollouts for r in g]`. + + Args: + rewards: One inner list per group, one entry per rollout in that group: + `rewards[g][i]` is the reward of rollout `i` in group `g`. + All groups are assumed to be the same size. Shape: (num_groups, group_size). + e.g. rewards = [[-1, 1], [4, 4]] # 2 groups, 2 rollouts each + num_turns: Same shape as `rewards`; turns per rollout. Zero marks a placeholder. + + Returns: + A flat list with one advantage per rollout, in group-major order + (group 0's rollouts, then group 1's, ...); aligned with `[r for g in rollouts for r in g]`. - For multiturn rollouts, the logic is a bit more involved. - # For training, we'll be turning each turn into a trajectory with the same reward - # within a trajectory, e.g. if [[a,b],[c,d,e]] trajectory has reward 1.0, we will - # get [a,b] with 1.0 and [c,d,e] with 1.0 when doing updates. + Example: + calculate_grpo_advantages([[-1, 1], [4, 4]], [[1, 1], [1, 1]]) ~= [-1.0, 1.0, 0.0, 0.0] + # approximate: stds are damped by the +1e-4 epsilon + # group 0 (mean 0, std 1) -> [-1, +1]; group 1 (zero variance) -> [0, 0] """ rewards = np.asarray(rewards, dtype=np.float64) @@ -1102,26 +1128,21 @@ def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[i f"rewards and num_turns must have matching shape, got {rewards.shape} and {num_turns.shape}" ) - # Each outer dimension of num_turns is a group. Sum of those gives total num_turns per group. - # Let's use this to calculate advantage. - # mean/std should be repeated based on group lens - group_turns = num_turns.sum(axis=-1) - # Zero-turn rollouts are empty-trajectory placeholders for failed episodes and are masked out. + # rewards are [num_groups, group_size]; making an assumption that all groups are the same size! + # @vitalyk: this will go away when we start sending env-based sample reqs. + # Zero-turn rollouts are empty-trajectory placeholders for failed episodes: they are masked + # out of the group statistics and pinned to advantage 0.0 (they train no rows). real_mask = num_turns > 0 # Clamping to 1 avoids 0/0 for all-placeholder groups. - real_counts = np.maximum(real_mask.sum(axis=-1), 1) - group_means = np.where(real_mask, rewards, 0.0).sum(axis=-1) / real_counts - group_deviations = np.where(real_mask, rewards - group_means[:, np.newaxis], 0.0) - group_stds = np.sqrt((group_deviations**2).sum(axis=-1) / real_counts) - reward_means = group_means.repeat(group_turns) - reward_stds = group_stds.repeat(group_turns) - - # rewards are originally [g, group_size] - # Making an assumption that all groups are of the same size! - # @vitalyk: this will go away when we start sending env-based sample reqs. - rewards = rewards.flatten().repeat(num_turns.flatten()) + real_counts = np.maximum(real_mask.sum(axis=-1, keepdims=True), 1) + reward_means = np.where(real_mask, rewards, 0.0).sum(axis=-1, keepdims=True) / real_counts + reward_stds = np.sqrt( + (np.where(real_mask, rewards - reward_means, 0.0) ** 2).sum(axis=-1, keepdims=True) + / real_counts + ) - return ((rewards - reward_means) / (1e-4 + reward_stds)).tolist() + advantages = (rewards - reward_means) / (1e-4 + reward_stds) + return np.where(real_mask, advantages, 0.0).flatten().tolist() def expand_epoch_segments( @@ -1167,6 +1188,19 @@ def merge_epoch_segments( kv_left = kv_segments[kv_idx].token_count +def single_turn_termination_ok( + rollout: TokenRollout, turn_traj: Sequence[int], seq_len: int, eod: int +) -> bool: + """Whether a single-turn trajectory ends in a known-legitimate way.""" + if len(turn_traj) == seq_len: + return True + if len(turn_traj) > 0 and turn_traj[-1] == eod: + return True + if rollout.generation_cap is None or not rollout.generation_mask: + return False + return int(sum(rollout.generation_mask[-1])) == rollout.generation_cap + + def compute_group_stats( rollouts: GroupedRollouts, tokenizer: MegatronTokenizer, @@ -1222,12 +1256,20 @@ def compute_group_stats( f"Rollout too long: {len(turn_traj)} > {seq_len} " f"(last token {turn_traj[-1]})\n{detokenized_traj}" ) - # A single-turn completion can only end in eod or be truncated at seq_len. + # A single-turn completion can only end in eod, be truncated at seq_len, + # or stop exactly at the agent's stamped output cap. # Multi-turn agents can additionally end a turn on a tool-call boundary. if len(rollout.trajectory) == 1: - assert len(turn_traj) == seq_len or turn_traj[-1] == tokenizer.eod, ( - f"Single-turn rollout under seq_length must end in eod: " - f"len={len(turn_traj)} last={turn_traj[-1]}\n{detokenized_traj}" + assert single_turn_termination_ok( + rollout, turn_traj, seq_len, tokenizer.eod + ), ( + f"Single-turn rollout under seq_length must end in eod or exactly " + f"at its generation cap: " + f"env={rollout.env_id} " + f"len={len(turn_traj)} " + f"last={turn_traj[-1] if turn_traj else None} " + f"generated={int(sum(rollout.generation_mask[-1])) if rollout.generation_mask else None} " + f"cap={rollout.generation_cap}\n{detokenized_traj}" ) else: lang_rl_log( @@ -1802,7 +1844,8 @@ def maybe_log_training_metrics( env_stats = lambda cont, idx: [cont[i] for i in idx] - group_turn_counts = [sum(nt) for nt in num_turns] + # advantages is flattened with ONE entry per rollout; it is sliced by per-group rollout counts + group_rollout_counts = [len(g) for g in rewards] for env_id in set(group_stats.env_ids): env_idx = [i for i, eidx in enumerate(group_stats.env_ids) if eidx == env_id] @@ -1810,8 +1853,8 @@ def maybe_log_training_metrics( # Advantages are flattened, we need to be more careful with those. env_advantages = [] for i in env_idx: - st = sum(group_turn_counts[:i]) - end = st + group_turn_counts[i] + st = sum(group_rollout_counts[:i]) + end = st + group_rollout_counts[i] env_advantages.extend(advantages[st:end]) env_metrics = prep_wandb_metrics(wandb_writer=wandb_writer, traj_lens=env_stats(traj_lens, env_idx), @@ -1845,30 +1888,59 @@ def maybe_log_training_metrics( wandb_writer.log(metrics, step=current_iteration) -PAD_TURN_UNIT = (None, -1) +PAD_ROW = None """Special filler entry to pad rollout turns for logprobs calculation. Used to equalize per-rank trajectory counts without contributing to the loss.""" +def _to_list(x): + """Normalize a token / mask / logprob sequence (list or 1-D tensor) to a plain list.""" + return x.tolist() if torch.is_tensor(x) else list(x) + + +def _is_collapsible_rollout(rollout) -> bool: + """Whether a rollout can be trained as one combined sequence. + + A multi-turn rollout collapses only if each turn's token_ids is an exact token-prefix of + the next (so the per-turn generated regions are disjoint in the final frame) and the + per-turn logprob counts match the generated-token counts (so the concatenated inference + logprobs line up with the union mask). + Single-turn and raw rollouts are trivially collapsible. + Non-collapsible rollouts must fall back to one row per turn so that logprobs are not misaligned. + """ + if not isinstance(rollout, TokenRollout) or len(rollout.trajectory) == 1: + return True + turns = [_to_list(t) for t in rollout.trajectory] + for prev, cur in zip(turns, turns[1:]): + if len(prev) > len(cur) or cur[: len(prev)] != prev: + return False + # Make sure that every turn's logprobs count matches the generated token count exactly. + # (the final turn may be short by one: the EOD) + for k, (mask, lps) in enumerate(zip(rollout.generation_mask, rollout.logprobs)): + short = int(sum(_to_list(mask))) - len(_to_list(lps)) + if short != 0 and not (k == len(rollout.logprobs) - 1 and short == 1): + return False + return True + + def prepare_trajectories( - rollout_turns: list[tuple[TokenRollout | Rollout | None, int]], - tokenizer: MegatronTokenizer, - seq_length: int, - skip_bos_token: bool, + rows: list[TokenRollout | Rollout | None], tokenizer: MegatronTokenizer, seq_length: int, skip_bos_token: bool ): """Pad trajectories and extract the generation masks. + Each (possibly multi-turn) rollout becomes one training row: the last turn's token_ids + already contain the whole conversation, so we train it once with a multi-region + generation mask (the union of every turn's generated span) rather than re-encoding the + prefix once per turn. + Args: - rollout_turns: (rollout, turn_idx) pairs; each pair becomes one trajectory. + rows: one entry per training row: a rollout, or PAD_ROW (None) for inert padding. tokenizer: Tokenizer to get the padding token and potentially tokenize. seq_length: Maximum sequence length to pad to. Returns: Trajectories, their generation masks, and per-row inference logprobs (unpadded tensor per real row, None per PAD row). - - Raises: - ValueError: """ # Track counts for each environment ID env_id_counts = Counter() @@ -1906,46 +1978,62 @@ def prepare_trajectories( trajs = [] generation_masks = [] inference_logprobs = [] - for rollout, turn_idx in rollout_turns: - if rollout is None: - # PAD_TURN_UNIT: inert filler so all DP ranks hold the same number of - # trajectories. All pad tokens, nothing generated, no inference logprobs. + single_turn_rows = [] # for the single-turn-only "one eod" sanity check below + for row in rows: + if row is None: + # PAD_ROW: inert filler so all DP ranks hold the same number of trajectories. + # All pad tokens, nothing generated, no inference logprobs. trajs.append([tokenizer.pad] * seq_length) generation_masks.append([False] * seq_length) inference_logprobs.append(None) + single_turn_rows.append(False) continue - # traj, gen mask and logprobs are per-turn lists on the rollout; - # single-turn environments just have single-element lists. - # We assume that all the structs above have the same lengths (number of turns). - trajectory = ( - copy.deepcopy(rollout.trajectory[turn_idx]) - if isinstance(rollout, TokenRollout) - else tokenizer.tokenize(rollout.trajectory)[turn_idx] - ) - inf_logprobs = rollout.logprobs[turn_idx] - generation_mask = ( - copy.deepcopy(rollout.generation_mask[turn_idx]) - if isinstance(rollout, TokenRollout) - else None - ) + + if isinstance(row, tuple): + # Fallback row: one turn of a rollout whose turns are NOT a prefix chain, + # so it could not be safely collapsed. Train this turn on its own. It is a + # multi-turn-origin turn, hence exempt from the single-turn "must end in eod" check. + rollout, turn_idx = row + trajectory = _to_list(rollout.trajectory[turn_idx]) + generation_mask = _to_list(rollout.generation_mask[turn_idx]) + inf_logprobs = _to_list(rollout.logprobs[turn_idx]) + single_turn_rows.append(False) + elif isinstance(row, TokenRollout): + # Collapse the (possibly multi-turn) rollout into ONE training sequence. The last + # turn's token_ids already contain the whole conversation (each turn re-encodes the + # prior context). The generation mask is the union of every turn's generated span; + # inference logprobs are the per-turn logprobs concatenated in turn order. + rollout = row + trajectory = _to_list(rollout.trajectory[-1]) + generation_mask = [False] * len(trajectory) + for turn_mask in rollout.generation_mask: + for pos, generated in enumerate(_to_list(turn_mask)): + if generated: + generation_mask[pos] = True + inf_logprobs = [lp for turn in rollout.logprobs for lp in _to_list(turn)] + single_turn_rows.append(len(rollout.trajectory) == 1) + else: + # Raw (non-token) rollout: tokenize the final turn; there is no per-token mask. + rollout = row + trajectory = tokenizer.tokenize(rollout.trajectory[-1]) + generation_mask = None + inf_logprobs = None + single_turn_rows.append(len(rollout.trajectory) == 1) + length = len(trajectory) assert length <= seq_length, "Rollout too long, how did this happen?" - + # Single-turn completions under seq_length must end in eod or exactly at their stamped generation cap (enforced upstream in + # compute_group_stats); multi-turn turns may stop on a tool-call boundary. if length < seq_length: trajectory.extend([tokenizer.pad] * (seq_length - length)) if generation_mask: - generation_mask.extend([False] * (seq_length - length)) + generation_mask = generation_mask + [False] * (seq_length - length) trajs.append(trajectory) generation_masks.append(generation_mask) + inference_logprobs.append(torch.Tensor(inf_logprobs) if inf_logprobs else None) - if inf_logprobs is not None: - inf_logprobs_tensor = torch.Tensor(inf_logprobs) - # Don't pad individual logprobs here - padding happens later if needed - inference_logprobs.append(inf_logprobs_tensor) - else: - inference_logprobs.append(None) - - if turn_idx == 0: + # A non-prefix rollout expands to one tuple row per turn; count the rollout once. + if not isinstance(row, tuple) or row[1] == 0: env_id_counts[rollout.env_id] += 1 if torch.distributed.is_initialized(): @@ -1959,7 +2047,7 @@ def prepare_trajectories( # Some sanity checks regarding the tokenization. Pad units start with the pad # token rather than bos, so the bos-equality check only applies to real rows. real_rows = torch.tensor( - [rollout is not None for rollout, _ in rollout_turns], dtype=torch.bool + [rollout is not None for rollout in rows], dtype=torch.bool ) if not skip_bos_token: assert ( @@ -1972,9 +2060,13 @@ def prepare_trajectories( assert ( tokenizer.bos is None or (trajs[:, 1] != tokenizer.bos).all() ), "Second token should not be bos" + # A multi-turn row legitimately contains one eod per turn (each turn can end on eod), + # so the "single eod" sanity check only applies to single-turn rows. + single_turn = torch.tensor(single_turn_rows, dtype=torch.bool) assert ( - (trajs * generation_masks.int() == tokenizer.eod).sum(axis=1) <= 1 - ).all(), "Only one eod per trajectory in generated tokens." + (trajs[single_turn] * generation_masks[single_turn].int() == tokenizer.eod).sum(axis=1) + <= 1 + ).all(), "Only one eod per trajectory in single-turn generated tokens." # TODO(rkirby): # We should avoid the tokenizer pad token being the same as the eod token for proper loss masking, # But now the deepseek tokenizer has the pad token set to eod, we need to handle this. @@ -2102,8 +2194,6 @@ def prepare_data_for_update( with nvtx_range("rl/prepare-data-for-update", time=True): with nvtx_range("rl/compute-group-stats", time=True): group_stats = compute_group_stats(rollouts, tokenizer, args.seq_length, request_ledger) - # TODO(vitalyk): why do we need global_advantages here? go inside packing - advantages = global_advantages = torch.tensor(group_stats.advantages, dtype=dtype).cuda() # Now split the rollouts across the data parallel ranks for training # This needs to be done at this point because we are about to calculate logprobs @@ -2125,47 +2215,55 @@ def prepare_data_for_update( samples_ratio_per_step = args.global_batch_size / (args.grpo_prompts_per_step * args.grpo_group_size) assert samples_ratio_per_step <= 1, "You cannot use more data than you sampled." - # Multi-turn rollouts contribute one trainable trajectory per turn, and turn counts vary. - # Flatten to single turns and split the turns across DP ranks. - - # advantages is already one entry per turn, so it is sliced with the same range. - rollout_turns = [ - (rollout, turn_idx) - for rollout in rollouts - for turn_idx in range(len(rollout.trajectory)) - ] - if not rollout_turns: - raise RuntimeError( - f"prepare_data_for_update: 0 usable trajectories from {len(rollouts)} rollout(s). " - "All rollouts have empty trajectories." - ) + # Build training rows. A rollout collapses to one combined row when its turns form a + # prefix chain (the common case: each turn re-encodes the prior conversation); if not, + # it falls back to one row per turn so logprobs are never misaligned. group_stats.advantages + # has one entry per rollout (group-major, matching `rollouts`); expand it to one per row. + rows = [] + row_advantages = [] + for rollout, adv in zip(rollouts, group_stats.advantages): + if not rollout.trajectory: + # Zero-turn placeholder (failed episode): nothing to train. + continue + if _is_collapsible_rollout(rollout): + rows.append(rollout) + row_advantages.append(adv) + else: + rows.extend((rollout, turn_idx) for turn_idx in range(len(rollout.trajectory))) + row_advantages.extend(adv for _ in range(len(rollout.trajectory))) + if not rows: + raise RuntimeError("prepare_data_for_update: 0 usable rollouts; all groups were empty.") + # TODO(vitalyk): why do we need global_advantages here? go inside packing + advantages = global_advantages = torch.tensor(row_advantages, dtype=dtype).cuda() data_parallel_world_size = mpu.get_data_parallel_world_size() - # The total turn count is data-dependent, so it needs to be padded. + # The rollout count is data-dependent (oversampling, group filtering), so it need not + # divide evenly across ranks and microbatches. Pad to the next multiple of + # (DP size * micro_batch_size) with inert rows (zero advantage, no loss contribution). pad_to_multiple = data_parallel_world_size * args.micro_batch_size - if pad_n := -len(rollout_turns) % pad_to_multiple: - rollout_turns = rollout_turns + [PAD_TURN_UNIT] * pad_n + if pad_n := -len(rows) % pad_to_multiple: + rows = rows + [PAD_ROW] * pad_n advantages = global_advantages = torch.cat( [advantages, torch.zeros(pad_n, dtype=advantages.dtype, device=advantages.device)] ) - total_turns_sampled = len(rollout_turns) + total_rows = len(rows) - has_inference_logprobs = any(isinstance(rollout, TokenRollout) for rollout, _ in rollout_turns) + has_inference_logprobs = any(isinstance(row, (tuple, TokenRollout)) for row in rows) if data_parallel_world_size > 0: - data_split_size = len(rollout_turns) // data_parallel_world_size + data_split_size = len(rows) // data_parallel_world_size data_split_range = ( mpu.get_data_parallel_rank() * data_split_size, (mpu.get_data_parallel_rank() + 1) * data_split_size, ) - rollout_turns = rollout_turns[data_split_range[0] : data_split_range[1]] + rows = rows[data_split_range[0] : data_split_range[1]] advantages = advantages[data_split_range[0] : data_split_range[1]] # First we calculate them on a global level and then we split and recalculate on a local level. # Sequence packing and reporting needs it global but non-packing wants it local. with nvtx_range("rl/prepare-trajectories", time=True): trajs, generation_masks, inference_logprobs = prepare_trajectories( - rollout_turns, tokenizer, args.seq_length, args.rl_skip_bos_token, + rows, tokenizer, args.seq_length, args.rl_skip_bos_token, ) if not has_inference_logprobs: inference_logprobs = None @@ -2352,7 +2450,7 @@ def prepare_data_for_update( reconfigure_num_microbatches_calculator( rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0, - global_batch_size=math.ceil(samples_ratio_per_step*total_turns_sampled), + global_batch_size=math.ceil(samples_ratio_per_step*total_rows), micro_batch_size=args.micro_batch_size, decrease_batch_size_if_needed=args.decrease_batch_size_if_needed, data_parallel_size=mpu.get_data_parallel_world_size(), diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index bcd066207c2..b9a0696dc6e 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -567,6 +567,7 @@ def pack_inference_logprobs( seq_to_local_bin[seq_idx] = local_bin_idx # Align and pack inference logprobs based on generation masks + device = packed_inference_logprobs.device for seq_idx in range(len(inference_logprobs)): if seq_idx not in seq_to_local_bin: continue # Skip sequences not on this rank @@ -578,30 +579,26 @@ def pack_inference_logprobs( seq_pos_in_bin = seq_positions.index(seq_idx) seq_start = packing_info.seq_starts[local_bin_idx][seq_pos_in_bin] - # Get generation mask for this sequence to find where generation starts - gen_mask = generation_masks[seq_idx] - # Find first generation token (accounting for the shift in get_logprobs) - first_gen_idx = gen_mask.int().argmax().item() - 1 - # Get the inference logprobs for this sequence - if isinstance(inference_logprobs[seq_idx], torch.Tensor): - seq_inf_logprobs = inference_logprobs[seq_idx] - else: + if not isinstance(inference_logprobs[seq_idx], torch.Tensor): continue # Skip if no inference logprobs - - # Calculate where to place inference logprobs in the packed tensor - # The inference logprobs start at the first generated token position - pack_start = seq_start + first_gen_idx - pack_end = min( - pack_start + len(seq_inf_logprobs), seq_start + packing_info.seq_lengths[seq_idx] - 1 - ) - actual_len = pack_end - pack_start - - if actual_len > 0 and pack_end <= bin_size - 1: - packed_inference_logprobs[local_bin_idx, pack_start:pack_end] = seq_inf_logprobs[ - :actual_len - ] - filled_mask[local_bin_idx, pack_start:pack_end] = True + seq_inf_logprobs = inference_logprobs[seq_idx].to(device) + + # Scatter logprobs to their generated token slots. + # Generation can span multiple regions for multi-turn rollouts; place each logprob at the + # slot for its token: a token at position `p` occupies packed slot `seq_start + (p - 1)`. + gen_positions = generation_masks[seq_idx].nonzero(as_tuple=True)[0].to(device) + if gen_positions.numel() == 0: + continue + target = seq_start + gen_positions - 1 + n = min(seq_inf_logprobs.numel(), target.numel()) + target = target[:n] + # Keep within this sequence's logprob span in the bin. + seq_lp_end = seq_start + packing_info.seq_lengths[seq_idx] - 1 + in_range = (target >= seq_start) & (target < seq_lp_end) + target_in_range = target[in_range] + packed_inference_logprobs[local_bin_idx, target_in_range] = seq_inf_logprobs[:n][in_range] + filled_mask[local_bin_idx, target_in_range] = True return packed_inference_logprobs, filled_mask diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 915e1d9fc7f..f6465caa9a6 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -221,6 +221,7 @@ def create_test_args(self, **kwargs): args.wandb_project = None args.micro_batch_size = 1 + args.eval_micro_batch_size = 1 for key, value in kwargs.items(): setattr(args, key, value) @@ -229,6 +230,34 @@ def create_test_args(self, **kwargs): set_global_variables(args, False) return args + @pytest.mark.parametrize( + "traj_len, last_token, cap, generated, ok", + [ + pytest.param(8, 99, None, 8, True, id="ends_in_eod"), + pytest.param(16, 7, None, 8, True, id="truncated_at_seq_len"), + pytest.param(8, 7, 5, 5, True, id="stopped_exactly_at_cap"), + pytest.param(8, 7, 5, 4, False, id="under_cap_is_corruption"), + pytest.param(8, 7, None, 8, False, id="uncapped_stays_strict"), + pytest.param(0, None, 5, 0, False, id="empty_turn_rejected"), + ], + ) + def test_single_turn_termination_ok(self, traj_len, last_token, cap, generated, ok): + """The eod tripwire admits exactly-at-cap stops (verified from the + generation mask, not a label) and stays strict everywhere else — the + 2026-07-30 workplace_assistant rollout (prompt + exactly + max_output_tokens_per_step generated, no eod) is the admit case.""" + eod, seq_len = 99, 16 + traj = [1] * (traj_len - 1) + [last_token] if traj_len else [] + mask = [False] * (traj_len - generated) + [True] * generated + rollout = TokenRollout( + trajectory=[traj], + generation_mask=[mask], + reward=0.0, + env_id="cap-test", + generation_cap=cap, + ) + assert rl_utils.single_turn_termination_ok(rollout, traj, seq_len, eod) is ok + def test_rl_granularity_defaults(self): args = self.create_test_args(perform_rl_step=True, grpo_prompts_per_step=8) @@ -695,8 +724,8 @@ def test_grpo_loss_truncation(self, current, old, expected_above, expected_below indirect=["initialize_model_parallel"], ) def test_prepare_data_for_update(self, initialize_model_parallel): - """Logprobs path runs; single-turn EOD guard holds; multi-turn turns are split, - padded to a DP*microbatch multiple, and size the microbatch calculator by turns.""" + """Logprobs path runs; single-turn EOD guard holds; multi-turn rollouts collapse to one + row each, padded to a DP*microbatch multiple to size the microbatch calculator.""" world_size, dp, tp, pp = initialize_model_parallel # Here I assume that we will be consuming all data in one step. group_size = 2 @@ -731,16 +760,15 @@ def test_prepare_data_for_update(self, initialize_model_parallel): request_ledger={}, ) - # Multi-turn rollouts with uneven turn counts: a turn may stop on a tool-call - # boundary (short, no eod) and is accepted. 2 + 3 turns per group * dp groups - # = 5*dp turns, padded up to the next multiple of micro_batch_size*dp (= 2*dp) - # -> 6*dp turns. With samples_ratio 1 the calculator is sized by total turns - # (6*dp -> 3 microbatches), not by rollout count (the pre-fix bug gave 1). + # Multi-turn rollouts with uneven turn counts: each collapses to ONE combined row + # (the final turn) regardless of turn count, so a group of 2 rollouts yields 2 rows. + # token_ids are cumulative and each turn generates the newly-added tokens (no overlap), + # so the combined generation mask / logprobs are well-formed. def mt1(): return make_token_rollout( [[1, 2, 3], [1, 2, 3, 4]], - [[-0.1, -0.2], [-0.3, -0.4]], - [[False, True, True], [False, False, True, True]], + [[-0.1, -0.2], [-0.3]], + [[False, True, True], [False, False, False, True]], problem_id="1", ) @@ -767,8 +795,9 @@ def mt2(): is_correction=False, request_ledger=mt_ledger, ) - # 5*dp turns padded to 6*dp; 6*dp / (micro_batch_size 2 * dp) = 3 microbatches. - assert get_num_microbatches() == 3 + # 2 rollouts/group * dp groups = 2*dp rows (already a multiple of micro_batch_size*dp); + # 2*dp / (micro_batch_size 2 * dp) = 1 microbatch. + assert get_num_microbatches() == 1 def r1(): return make_token_rollout( @@ -810,27 +839,16 @@ def r2(): [pytest.param((1, 1), id="tp1-pp1")], indirect=["initialize_model_parallel"], ) - def test_prepare_data_for_update_oversampling(self, initialize_model_parallel): - """Oversampling (ratio < 1) consumes a fraction of the (padded) turn count per step: - the microbatch calculator is sized by ceil(ratio * total turns), not the full batch.""" + @pytest.mark.parametrize("scenario", ["fallback_and_pad", "oversampling", "empty_rollout"]) + def test_prepare_data_for_update_row_accounting(self, initialize_model_parallel, scenario): + """Rows = one per rollout, except non-prefix rollouts fall back to one row per turn + and zero-turn placeholders (failed episodes) contribute no row at all. + The (padded) row count -- not the rollout count -- sizes the microbatch calculator, and + oversampling (ratio < 1) consumes a fraction of the rows per step.""" world_size, dp, tp, pp = initialize_model_parallel tokenizer = MockTokenizer() model = MockModel() - # ratio = global_batch_size/(prompts*group) = 2*dp/(dp*4) = 0.5. - # 4*dp single-turn turns (already a multiple of 2*dp); ceil(0.5 * 4*dp) = 2*dp; - # 2*dp / (2 * dp) = 1 microbatch. - self.create_test_args( - micro_batch_size=2, - seq_length=4, - curr_iteration=1, - tensor_model_parallel_size=tp, - pipeline_model_parallel_size=pp, - global_batch_size=dp * 2, - grpo_prompts_per_step=dp, - grpo_group_size=4, - ) - def single(problem_id, reward): return make_token_rollout( [[1, 2, 3, tokenizer.eod]], @@ -840,7 +858,70 @@ def single(problem_id, reward): problem_id=problem_id, ) - rollouts = [[single(str(i), float(i % 2)) for i in range(4)] for _ in range(dp)] + if scenario == "fallback_and_pad": + # group = [collapsible single-turn, NON-prefix 2-turn] -> 1 + 2 = 3 rows/group. + # 3*dp rows padded up to the next multiple of micro_batch_size*dp (= 2*dp) -> 4*dp; + # 4*dp / (2 * dp) = 2 microbatches. + self.create_test_args( + micro_batch_size=2, + seq_length=4, + curr_iteration=1, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + global_batch_size=dp * 2, + grpo_prompts_per_step=dp, + grpo_group_size=2, + ) + + def non_prefix(): + return make_token_rollout( + [[1, 2, 3, tokenizer.eod], [7, 8, 9, tokenizer.eod]], + [[-0.1, -0.2, -0.3], [-0.4, -0.5, -0.6]], + [[False, True, True, True], [False, True, True, True]], + reward=0.0, + problem_id="x", + ) + + rollouts = [[single("a", 1.0), non_prefix()] for _ in range(dp)] + expected_microbatches = 2 + elif scenario == "empty_rollout": + # group = [2 single-turn, 1 zero-turn placeholder (failed episode)] -> 2 rows/group: + # the placeholder is skipped (keeping it would crash indexing trajectory[-1], and + # counting it would give 3*dp rows -> 3 microbatches). + # 2*dp rows / (micro_batch_size 1 * dp) -> 2 microbatches. + # micro_batch_size must be 1: the microbatch calculator is first built from the + # raw args, and global_batch_size (3*dp) must be divisible by mbs*dp. + self.create_test_args( + micro_batch_size=1, + seq_length=4, + curr_iteration=1, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + global_batch_size=dp * 3, + grpo_prompts_per_step=dp, + grpo_group_size=3, + ) + rollouts = [ + [single("a", 1.0), single("b", 0.0), make_token_rollout([], [], [], reward=0.0)] + for _ in range(dp) + ] + expected_microbatches = 2 + else: # oversampling: ratio = global_batch_size/(prompts*group) = 2*dp/(dp*4) = 0.5. + # 4*dp single-turn rows (already a multiple of 2*dp); ceil(0.5 * 4*dp) = 2*dp; + # 2*dp / (2 * dp) = 1 microbatch. + self.create_test_args( + micro_batch_size=2, + seq_length=4, + curr_iteration=1, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + global_batch_size=dp * 2, + grpo_prompts_per_step=dp, + grpo_group_size=4, + ) + rollouts = [[single(str(i), float(i % 2)) for i in range(4)] for _ in range(dp)] + expected_microbatches = 1 + request_ledger = _ledger_for(*(r for g in rollouts for r in g)) rl_utils.prepare_data_for_update( [model], @@ -851,77 +932,88 @@ def single(problem_id, reward): is_correction=False, request_ledger=request_ledger, ) - assert get_num_microbatches() == 1 + assert get_num_microbatches() == expected_microbatches - @pytest.mark.parametrize("num_turns", [1, 2]) - def test_prepare_trajectories(self, num_turns): - """Each (rollout, turn_idx) unit becomes one padded training row (a rollout with T - turns contributes T rows); PAD_TURN_UNIT entries become inert all-pad rows with no - generated tokens and no inference logprobs (DP count equalization).""" + def test_prepare_trajectories(self): + """Every row kind becomes one padded training row in a single batch: a single-turn + rollout (one contiguous generated region), a multi-turn rollout collapsed to its final + sequence (multi-region mask with a gap at the observation, per-turn logprobs + concatenated), the per-turn fallback rows of a NON-prefix rollout (each trained on its own + tokens/mask), and an inert PAD_ROW.""" seq_length = 8 self.create_test_args(rl_skip_bos_token=False, micro_batch_size=1, seq_length=seq_length) tokenizer = MockTokenizer() eod, pad = tokenizer.eod, tokenizer.pad - r1 = make_token_rollout( - [[1, 2, 3, eod]] * num_turns, - [[0.1, 0.2, 0.3, 0.35]] * num_turns, - [[False, True, True, True]] * num_turns, + # Single-turn rollout: one contiguous generated region. + single = make_token_rollout( + [[1, 2, 3, eod]], + [[0.1, 0.2, 0.3]], + [[False, True, True, True]], reward=3.14, problem_id="1", ) - r2 = make_token_rollout( - [[4, 5, 6, 7, eod]] * num_turns, - [[0.4, 0.5, 0.6, 0.7, 0.75]] * num_turns, - [[False, True, True, True, True]] * num_turns, + # Multi-turn prefix chain: turn 0 generates [2, 3, eod]; an observation (token 5) is + # appended; turn 1 generates [6, eod]. token_ids are cumulative, so this collapses to the + # final sequence with TWO generated regions (a gap at the observation) and per-turn + # logprobs concatenated. + multi = make_token_rollout( + [[1, 2, 3, eod], [1, 2, 3, eod, 5, 6, eod]], + [[0.4, 0.5, 0.6], [0.7, 0.75]], + [[False, True, True, True], [False, False, False, False, False, True, True]], reward=0.14, problem_id="2", ) - r3 = make_token_rollout( - [[8, 9, eod]] * num_turns, - [[0.8, 0.9, 0.95]] * num_turns, - [[False, True, True]] * num_turns, - reward=2.71, + # Non-prefix rollout: turn 1 is not a prefix-extension of turn 0, so it is trained as + # per-turn fallback rows (rollout, turn_idx) -- each on its own tokens/mask. + non_prefix = make_token_rollout( + [[1, 2, 3, eod], [7, 8, 9, eod]], + [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], + [[False, True, True, True], [False, True, True, True]], problem_id="3", ) - turn_units = [ - (rollout, turn_idx) - for rollout in [r1, r2, r3] - for turn_idx in range(len(rollout.trajectory)) - ] + [rl_utils.PAD_TURN_UNIT] + rows = [single, multi, (non_prefix, 0), (non_prefix, 1), rl_utils.PAD_ROW] trajs, genmask, inference_logprobs = rl_utils.prepare_trajectories( - turn_units, tokenizer, seq_length, skip_bos_token=False + rows, tokenizer, seq_length, skip_bos_token=False ) expected_trajs = torch.tensor( - [[1, 2, 3, eod] + [pad] * 4, [4, 5, 6, 7, eod] + [pad] * 3, [8, 9, eod] + [pad] * 5], + [ + [1, 2, 3, eod] + [pad] * 4, + [1, 2, 3, eod, 5, 6, eod] + [pad] * 1, + [1, 2, 3, eod] + [pad] * 4, + [7, 8, 9, eod] + [pad] * 4, + [pad] * 8, + ], dtype=torch.long, device=trajs.device, - ).repeat_interleave(num_turns, dim=0) - expected_trajs = torch.cat( - [expected_trajs, torch.full((1, seq_length), pad, dtype=expected_trajs.dtype)] ) assert torch.equal(trajs, expected_trajs) expected_genmask = torch.tensor( [ - [False, True, True, True] + [False] * 4, - [False, True, True, True, True] + [False] * 3, - [False, True, True] + [False] * 5, + [False, True, True, True, False, False, False, False], + # union of both turns' generated spans, with a gap at the observation (pos 4) + [False, True, True, True, False, True, True, False], + [False, True, True, True, False, False, False, False], + [False, True, True, True, False, False, False, False], + [False] * 8, # inert PAD_ROW ], dtype=torch.bool, device=genmask.device, - ).repeat_interleave(num_turns, dim=0) - expected_genmask = torch.cat( - [expected_genmask, torch.zeros((1, seq_length), dtype=torch.bool)] ) assert torch.equal(genmask, expected_genmask) - # Per-row list: unpadded tensor per real row, None for the pad unit. (Packing-mode + # Per-row list: unpadded tensor per real row, None for the pad row. (Packing-mode # densification happens at the call site via _pad_nonnull_with_zeros, tested separately.) - expected_logprobs = [[0.1, 0.2, 0.3, 0.35], [0.4, 0.5, 0.6, 0.7, 0.75], [0.8, 0.9, 0.95]] - expected_logprobs = [el for el in expected_logprobs for _ in range(num_turns)] + [None] + expected_logprobs = [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6, 0.7, 0.75], + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + None, + ] assert len(inference_logprobs) == len(expected_logprobs) for got, exp in zip(inference_logprobs, expected_logprobs): if exp is None: @@ -931,16 +1023,90 @@ def test_prepare_trajectories(self, num_turns): torch.testing.assert_close(got, exp_t, rtol=0, atol=0) @pytest.mark.parametrize( - "num_turns, expected", + "trajectory, logprobs, generation_mask, collapsible", + [ + pytest.param([[1, 2, 3]], [[0.1, 0.2]], [[False, True, True]], True, id="single_turn"), + pytest.param( + [[1, 2], [1, 2, 3, 4]], + [[0.1], [0.2, 0.3]], + [[False, True], [False, False, True, True]], + True, + id="prefix_chain", + ), + pytest.param( + [[1, 2, 3, 4], [5, 6, 7, 8]], + [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]], + [[False, True, True, True], [False, True, True, True]], + False, + id="non_prefix", + ), + pytest.param( + [[1, 2, 3], [1, 2]], + [[0.1, 0.2], [0.1]], + [[False, True, True], [False, True]], + False, + id="later_turn_shorter", + ), + pytest.param( + [[1, 2], [1, 2, 3, 4]], + [[0.1], [0.2]], + [[False, True], [False, False, True, True]], + True, + id="final_turn_short_one_eod", + ), + pytest.param( + [[1, 2], [1, 2, 3, 4]], + [[0.1], [0.2, 0.3]], + [[True, True], [False, False, True, True]], + False, + id="nonfinal_turn_short_one", + ), + pytest.param( + [[1, 2], [1, 2, 3, 4]], + [[0.1], [0.2]], + [[False, True], [False, True, True, True]], + False, + id="final_turn_short_two", + ), + pytest.param( + [[1, 2], [1, 2, 3, 4]], + [[0.1], [0.2, 0.3]], + [[False, True], [False, False, False, True]], + False, + id="extra_logprobs", + ), + ], + ) + def test_is_collapsible_rollout(self, trajectory, logprobs, generation_mask, collapsible): + """A rollout collapses to one combined row only when each turn's token_ids is an exact + prefix of the next AND each turn's logprob count matches its generated-token count + (the final turn may be short one: train-side-appended eod); otherwise it must fall + back to one row per turn.""" + rollout = make_token_rollout(trajectory, logprobs, generation_mask) + assert rl_utils._is_collapsible_rollout(rollout) is collapsible + + @pytest.mark.parametrize( + "rewards, num_turns, expected", [ - pytest.param([[1, 1], [1, 1]], [-1.0, 1.0, 0.0, 0.0], id="single_turn"), - # A rollout's group advantage is repeated once per turn. - pytest.param([[2, 1], [1, 3]], [-1.0, -1.0, 1.0, 0.0, 0.0, 0.0, 0.0], id="multi_turn"), + pytest.param( + [[-1, 1], [4, 4]], + [[1, 1], [1, 1]], + [-1.0, 1.0, 0.0, 0.0], + id="normalized_and_zero_variance", + ), + pytest.param([[2, 2, 2]], [[1, 2, 3]], [0.0, 0.0, 0.0], id="all_equal"), + pytest.param( + [[0, 1, 0]], [[0, 1, 1]], [0.0, 1.0, -1.0], id="placeholder_masked_from_stats" + ), + pytest.param([[0, 0]], [[0, 0]], [0.0, 0.0], id="all_placeholder_group"), ], ) - def test_advantage_calculation(self, num_turns, expected): - advs = rl_utils.calculate_grpo_advantages([[-1, 1], [4, 4]], num_turns) - torch.testing.assert_close(torch.tensor(advs), torch.tensor(expected), atol=1e-4, rtol=1e-5) + def test_calculate_grpo_advantages(self, rewards, num_turns, expected): + """One group-normalized advantage per rollout (independent of turn counts), flattened + group-major; zero-turn placeholders are masked out of the group statistics and pinned + to advantage 0.0, so their synthetic rewards manufacture no advantage.""" + advs = rl_utils.calculate_grpo_advantages(rewards, num_turns) + torch.testing.assert_close(torch.tensor(advs), torch.tensor(expected), atol=3e-4, rtol=1e-5) def test_bounded_artifact_key(self): short = "env_rollout_table" @@ -953,18 +1119,6 @@ def test_bounded_artifact_key(self): assert bounded == rl_utils._bounded_artifact_key(long_key) assert bounded != rl_utils._bounded_artifact_key(long_key + "y") - def test_placeholder_rewards_do_not_affect_advantage_normalization(self): - # Group of 3 with one zero-turn placeholder: the two real rollouts - # normalize against each other, so the placeholder's synthetic 0.0 - # reward manufactures no advantage. - advs = rl_utils.calculate_grpo_advantages([[0, 1, 0]], [[0, 1, 1]]) - torch.testing.assert_close( - torch.tensor(advs), torch.tensor([1.0, -1.0]), atol=3e-4, rtol=1e-5 - ) - - def test_all_placeholder_group_has_no_advantages(self): - assert rl_utils.calculate_grpo_advantages([[0, 0]], [[0, 0]]) == [] - def test_advantage_shape_mismatch_raises(self): with pytest.raises(AssertionError, match="matching shape"): rl_utils.calculate_grpo_advantages([[1.0, 0.0]], [[1, 1, 1]]) @@ -1037,6 +1191,39 @@ def test_pad_nonnull_with_zeros(self, lengths, max_len, expected_shape): assert padded.shape == expected_shape assert (padded == 0).all() # zero inputs and zero-filled padding/None rows + def test_align_unpacked_inference_logprobs(self): + """Scatter positions: a generated token at position p lands at index p-1, hopping + observation gaps; a short logprob list leaves trailing generated slots at old_logprobs + (train-side eod); all-False (PAD) rows are untouched and never dereference their + (None) list entry.""" + width = 7 + # Distinct value per cell so any wrong-position write breaks exact equality. + old = -0.5 - 0.01 * torch.arange(4 * width, dtype=torch.float32).reshape(4, width) + # Masks are seq_length (width + 1) wide, as in the pipeline: exercises the stats clip. + masks = torch.tensor( + [ + [False, True, True, False, True, True, False, False], # multi-region, gap at 3 + [False, True, True, True, False, False, False, False], # contiguous, 3 gen slots + [True, True, False, False, False, False, False, False], # gen at 0: target -1 + [False] * 8, # PAD row + ] + ) + logprobs = [ + torch.tensor([-0.11, -0.12, -0.21, -0.22]), # two turns' logprobs, concatenated + torch.tensor([-0.31, -0.32]), # one short: the eod slot keeps old_logprobs + torch.tensor([-0.41, -0.42]), # first pairs with the dropped target -1 + None, # PAD: must not be touched + ] + expected = old.clone() + expected[0, [0, 1, 3, 4]] = torch.tensor([-0.11, -0.12, -0.21, -0.22]) + expected[1, [0, 1]] = torch.tensor([-0.31, -0.32]) # index 2 keeps old (ratio 1) + expected[2, 0] = -0.42 + + aligned = rl_utils.align_unpacked_inference_logprobs( + logprobs, old, masks, SimpleNamespace() + ) + torch.testing.assert_close(aligned, expected, rtol=0, atol=0) + @pytest.mark.parametrize( "initialize_model_parallel", [ @@ -1451,7 +1638,8 @@ def test_prep_wandb_metrics(self, initialize_model_parallel, inject_placeholders # inner lists, which the group-level stats must skip, not crash on. turn_lens.append([]) completed_epochs.append([]) - # advantages stay [0, 1]: zero-turn rollouts emit no advantage entries. + # advantages stay [0, 1]: prep treats them as an opaque flat vector + # (production pins placeholder entries to 0.0; alignment is not prep's concern). writer = MagicMock() metrics = rl_utils.prep_wandb_metrics( writer, @@ -1672,9 +1860,10 @@ def placeholder(rollout_status="ok", failure_reason=None): # Rewards keep the placeholder zeros: the lists stay rectangular for # alignment; GRPO normalization masks them by num_turns downstream. assert stats.rewards == [[1.0, 1.0, 0.0], [0.0, 0.0, 0.0]] - # Which the advantages reflect: both real rollouts earned 1.0, so with - # the placeholder masked there is no contrast left in the group. - assert stats.advantages == [0.0, 0.0] + # Which the advantages reflect: one entry per rollout, placeholders pinned to 0.0; + # both real rollouts earned 1.0, so with the placeholders masked from the group + # statistics there is no contrast left anywhere. + assert stats.advantages == [0.0] * 6 # Adapter stamps pass through; bare empties are inferred to 'placeholder'. assert stats.rollout_statuses == [ ["ok", "ok", "placeholder"], diff --git a/tests/unit_tests/rl/test_sequence_packing_utils.py b/tests/unit_tests/rl/test_sequence_packing_utils.py index 85fbdb150fb..95d8559ab80 100644 --- a/tests/unit_tests/rl/test_sequence_packing_utils.py +++ b/tests/unit_tests/rl/test_sequence_packing_utils.py @@ -408,6 +408,42 @@ def test_compute_packed_inference_logprobs_stats_shape_mismatch(): ) +def test_pack_inference_logprobs_multi_region(): + """pack_inference_logprobs scatters each logprob to its generated token's slot, including + multi-region (multi-turn) masks with gaps, and reduces to contiguous placement otherwise. + The logprob of a token at sequence position p lands at packed index seq_start + p - 1.""" + bin_size = 8 + # Bin 0 holds two length-4 sequences: seq 0 at offset 0, seq 1 at offset 4. + packing_info = sequence_packing_utils.PackingInfo( + bin_seq_indices=[[0, 1]], + seq_starts={0: [0, 4]}, + seq_lengths=[4, 4], + seq_to_bin_idx=[0, 0], + packing_algo='fifo', + ) + # seq 0: MULTI-region (generated at positions 1 and 3, gap at 2 — e.g. an observation). + # seq 1: single contiguous region (generated at positions 1, 2). + generation_masks = torch.tensor( + [ + [False, True, False, True, False, False, False, False], + [False, True, True, False, False, False, False, False], + ] + ) + inference_logprobs = [torch.tensor([0.1, 0.2]), torch.tensor([0.3, 0.4])] + + packed, filled_mask = sequence_packing_utils.pack_inference_logprobs( + inference_logprobs, packing_info, generation_masks, bin_size + ) + + # seq 0 (start 0): positions 1,3 -> indices 0,2 (gap index 1 stays 0). + # seq 1 (start 4): positions 1,2 -> indices 4,5. + expected = torch.tensor([[0.1, 0.0, 0.2, 0.0, 0.3, 0.4, 0.0]]) + assert packed.shape == (1, bin_size - 1) + torch.testing.assert_close(packed, expected, rtol=0, atol=0) + # The filled mask marks exactly the scattered targets. + assert filled_mask.tolist() == [[True, False, True, False, True, True, False]] + + def test_packing_observability_metrics(): """Test various observability metrics related to sequence packing.""" From 540f4203efd66588319c45e0a265a0afa4854082 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 19 Aug 2026 00:50:58 +0000 Subject: [PATCH 093/101] Add opt-in MFSDP unified communication stream (#6563) Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 14 +++++++++-- .../src/megatron_fsdp/experimental/module.py | 23 ++++++++++++++---- .../distributed/mfsdp_v2/test_memory.py | 24 ++++++++++++------- .../distributed/mfsdp_v2/test_overlap.py | 18 +++++++++++--- 4 files changed, 61 insertions(+), 18 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index a5dab7de9d3..63354a7a403 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -32,7 +32,10 @@ @contextmanager def fully_shard_context( - device: torch.device | None = None, *, use_symmetric_memory: bool = False + device: torch.device | None = None, + *, + use_symmetric_memory: bool = False, + unify_communication_stream: bool = False, ) -> Iterator[FsdpContext]: """Construct FSDP modules that share runtime streams and prefetch orders. @@ -44,6 +47,9 @@ def fully_shard_context( the current CUDA device. use_symmetric_memory: Allocate communication staging buffers from PyTorch's NCCL symmetric-memory pool. + unify_communication_stream: Whether all-gathers and reduce-scatters share one + communication stream to reduce peak transient memory. See + https://github.com/NVIDIA/Megatron-LM/issues/6471. """ if _FSDP_CONTEXT.get() is not None: raise RuntimeError("fully_shard_context does not support nesting.") @@ -52,7 +58,11 @@ def fully_shard_context( if device.type != "cuda": raise ValueError(f"fully_shard_context requires a CUDA device, got {device}.") - context = FsdpContext(device=device, use_symmetric_memory=use_symmetric_memory) + context = FsdpContext( + device=device, + use_symmetric_memory=use_symmetric_memory, + unify_communication_stream=unify_communication_stream, + ) token = _FSDP_CONTEXT.set(context) try: yield context diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index 4c410899285..dbc0df69481 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -44,30 +44,43 @@ class FsdpContext: # from ``main_weight``, has placements different from ``Placements.optimizer``. is_last_microbatch: bool use_symmetric_memory: bool + unify_communication_stream: bool # Static orders used to drive all-gather prefetch. We may want to switch to # capturing runtime order if static module order proves too fragile. Each # FsdpModule tracks its own materialized state via ``FsdpModule._unshard_event``. forward_order: IndexedOrder["FsdpModule"] backward_order: IndexedOrder["FsdpModule"] - def __init__(self, device: torch.device, use_symmetric_memory: bool = False) -> None: + def __init__( + self, + device: torch.device, + use_symmetric_memory: bool = False, + unify_communication_stream: bool = False, + ) -> None: """Create rank-local runtime state for FSDP modules on ``device``. Args: device: Device on which this context schedules communication. use_symmetric_memory: Whether modules constructed in this context allocate communication staging buffers from PyTorch's NCCL symmetric-memory pool. + unify_communication_stream: Whether all-gathers and reduce-scatters share one + communication stream to reduce peak transient memory. """ self.is_last_microbatch = True self.use_symmetric_memory = use_symmetric_memory + self.unify_communication_stream = unify_communication_stream self.forward_order = IndexedOrder() self.backward_order = IndexedOrder() # Construction-only; empty after finalization. self._registered_modules: list[FsdpModule] = [] self._is_finalized = False - with torch.cuda.device(device): - self.allgather_stream = torch.cuda.Stream() - self.reduce_scatter_stream = torch.cuda.Stream() + self.allgather_stream = torch.cuda.Stream(device) + if unify_communication_stream: + # A unified stream lets an all-gather reuse the storage released by a + # preceding reduce-scatter. + self.reduce_scatter_stream = self.allgather_stream + else: + self.reduce_scatter_stream = torch.cuda.Stream(device) def register_module(self, module: "FsdpModule") -> None: """Register a module constructed in this context.""" @@ -386,8 +399,8 @@ def pre_backward(self) -> None: def post_backward(self) -> None: """Reduce gradients and return parameters to their sharded resting state.""" - self._reduce_gradient_groups() self._reshard_parameter_groups() + self._reduce_gradient_groups() self.phase = FsdpModule.Phase.RESTING torch.cuda.nvtx.range_pop() diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_memory.py b/tests/unit_tests/distributed/mfsdp_v2/test_memory.py index 6e0e2b2e389..78e427788cc 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_memory.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_memory.py @@ -107,8 +107,13 @@ def test_persistent_sharded_storage(distributed_setup, main_params_dtype): ) -def test_training_step_peak_memory_bounds_full_size_buffers(distributed_setup): - """A training step should stay below five full-size child buffers.""" +@pytest.mark.parametrize( + "unify_communication_stream", [False, True], ids=["separate_streams", "unified_stream"] +) +def test_training_step_peak_memory_bounds_full_size_buffers( + distributed_setup, unify_communication_stream +): + """A training step should stay within its full-size-buffer bound.""" rank = distributed_setup.rank world_size = distributed_setup.world_size device = distributed_setup.device @@ -121,7 +126,7 @@ def test_training_step_peak_memory_bounds_full_size_buffers(distributed_setup): model = MultiChildModel(dim=dim, num_children=8).to(dtype=dtype) placements = _flat_placements() policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) - with fully_shard_context(device=device): + with fully_shard_context(device=device, unify_communication_stream=unify_communication_stream): for layer in model.layers: fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) @@ -142,15 +147,18 @@ def train_step() -> None: peak_delta = torch.cuda.max_memory_allocated(device) - resting_allocated # Backward keeps the current child and one prefetched child unsharded. The current - # child also has a full wgrad until it is copied into a full reduce-scatter input, - # for a four-full-child-buffer peak. Allow one additional buffer for cuBLAS - # workspace, allocator granularity, and small temporaries. - bound_nbytes = (4 + 1) * child_weight_nbytes + # child also has a full wgrad until it is copied into a full reduce-scatter input. + # With separate streams, their allocation cannot reuse the released full-weight + # storage, for a four-buffer peak. With a unified stream, the release precedes the + # allocation on that stream, reducing the peak to three. Allow one additional buffer + # for cuBLAS workspace, allocator granularity, and small temporaries. + full_buffer_bound = 4 if unify_communication_stream else 5 + bound_nbytes = full_buffer_bound * child_weight_nbytes assert peak_delta < bound_nbytes, ( "FSDP training-step peak memory exceeded the full-size-buffer bound: " f"rank={rank}, peak_delta={_mb(peak_delta)}, " - f"five_full_child_buffers={_mb(bound_nbytes)}" + f"{full_buffer_bound}_full_child_buffers={_mb(bound_nbytes)}" ) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py b/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py index 26326962a6b..7d83526442f 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py @@ -63,7 +63,12 @@ def _flat_placements() -> Placements: ], ids=["default", "symmetric_memory"], ) -def test_overlaps_communication_and_compute(distributed_setup, use_symmetric_memory): +@pytest.mark.parametrize( + "unify_communication_stream", [False, True], ids=["separate_streams", "unified_stream"] +) +def test_overlaps_communication_and_compute( + distributed_setup, use_symmetric_memory, unify_communication_stream +): """Forward and backward communication should overlap GEMM compute.""" world_size = distributed_setup.world_size device = distributed_setup.device @@ -110,7 +115,11 @@ def test_overlaps_communication_and_compute(distributed_setup, use_symmetric_mem model = MultiChildModel(dim=dim, num_children=num_children).to(device=device, dtype=dtype) placements = _flat_placements() policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) - with fully_shard_context(device=device, use_symmetric_memory=use_symmetric_memory): + with fully_shard_context( + device=device, + use_symmetric_memory=use_symmetric_memory, + unify_communication_stream=unify_communication_stream, + ): for layer in model.layers: fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) @@ -160,8 +169,11 @@ def train_one_iteration() -> None: gemm_streams = {kernel.device_resource_id for kernel in gemm_kernels} if allgather_kernels: assert len(allgather_streams) == 1 + if unify_communication_stream: + assert allgather_streams == reduce_scatter_streams + else: + assert allgather_streams.isdisjoint(reduce_scatter_streams) assert len(reduce_scatter_streams) == 1 - assert allgather_streams.isdisjoint(reduce_scatter_streams) assert allgather_streams.isdisjoint(gemm_streams) assert reduce_scatter_streams.isdisjoint(gemm_streams) From 1d693644948831636ca06aff9633928265313742 Mon Sep 17 00:00:00 2001 From: Maanu Grover Date: Wed, 19 Aug 2026 00:50:58 +0000 Subject: [PATCH 094/101] tests: allow trusted TE pickle extra state in affected tests (#6602) Signed-off-by: Maanu Grover --- .../model_config.yaml | 1 + .../model_config.yaml | 1 + .../model_config.yaml | 1 + .../model_config.yaml | 1 + .../model_config.yaml | 1 + .../model_config.yaml | 1 + 6 files changed, 6 insertions(+) diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml index b32eb7e76dc..b05134ba37d 100644 --- a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_chunked_prefill/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_chunked_prefill/model_config.yaml index 35f0b7951e7..78037a8e520 100644 --- a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_chunked_prefill/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_chunked_prefill/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_prefix_caching/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_prefix_caching/model_config.yaml index 69fba5c46cc..6750b36db94 100644 --- a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_prefix_caching/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_prefix_caching/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 diff --git a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml index 2e9157cbcc2..0107499bc01 100644 --- a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1 THROUGHPUT_START_STEP: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring diff --git a/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml index 013eee2e49f..f1c8d39d63f 100644 --- a/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 diff --git a/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml index 93532c83707..080c34b9520 100644 --- a/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 From afa707eb1c83220d18ac4d93bbd158822a06371e Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 19 Aug 2026 03:23:49 +0000 Subject: [PATCH 095/101] Reflect #4440 onto GPT models, not just hybrid (#6476) Signed-off-by: Teodor-Dumitru Ene --- megatron/core/models/gpt/gpt_model.py | 41 ++++++++++++++++++- .../inference/engines/test_dynamic_engine.py | 13 +++--- 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index ad5f46c1f97..b0e6db63079 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -29,6 +29,7 @@ from megatron.core.quantization.utils import get_quant_config_or_none from megatron.core.tensor_parallel import gather_from_sequence_parallel_region from megatron.core.transformer.enums import InferenceCudaGraphScope, ModelType +from megatron.core.transformer.module import GraphableMegatronModule from megatron.core.transformer.moe.paged_stash import paged_stash_init_chunk_handler from megatron.core.transformer.multi_token_prediction import ( MultiTokenPredictionBlock, @@ -48,7 +49,7 @@ logger = logging.getLogger(__name__) -class GPTModel(LanguageModule): +class GPTModel(LanguageModule, GraphableMegatronModule): """GPT Transformer language model. Args: @@ -237,6 +238,8 @@ def __init__( pg_collection=self.pg_collection, vp_stage=vp_stage, ) + if hasattr(self, 'cudagraph_manager') and hasattr(self.decoder, 'cudagraph_manager'): + del self.decoder.cudagraph_manager if self.mtp_process: self.mtp = MultiTokenPredictionBlock( @@ -522,6 +525,42 @@ def preprocess_for_paged_stash(self): vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage ) + def _should_call_local_cudagraph(self, *args, **kwargs): + """ + Check if we should call the local cudagraph path. + """ + if ( + InferenceMode.is_active() + and hasattr(self, 'cudagraph_manager') + and ( + kwargs.get('inference_context') is not None + or kwargs.get('inference_params') is not None + ) + and self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block + ): + if kwargs['inference_context'].is_static_batching(): + using_cuda_graph = kwargs['inference_context'].is_decode_only() + else: + using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step() + + if using_cuda_graph: + return True + return False + + def __call__(self, *args, **kwargs): + if self._should_call_local_cudagraph(*args, **kwargs): + return super().__call__(*args, **kwargs)[0] + return super().__call__(*args, **kwargs) + + def create_mcore_cudagraph_manager(self, config): + """ + Create the cudagraph manager for the full iteration inference scope + """ + if config.inference_cuda_graph_scope == InferenceCudaGraphScope.block: + from megatron.core.transformer.cuda_graphs import CudaGraphManager + + self.cudagraph_manager = CudaGraphManager(config) + def forward( self, input_ids: Tensor, diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 461fa762583..093e4062cc8 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -214,6 +214,7 @@ def _assert_inference_cuda_graphs_disabled(env) -> None: assert env.engine.cuda_graph_impl == "full_iteration" assert env.engine.inference_cuda_graph_scope == InferenceCudaGraphScope.none assert env.engine.capture_stats is None + assert not hasattr(model, 'cudagraph_manager') assert not hasattr(model.decoder, 'cudagraph_manager') for layer in model.decoder.layers: assert not hasattr(layer, 'cudagraph_manager') @@ -858,11 +859,10 @@ def test_simple(self, model_provider, num_cuda_graphs, inference_cuda_graph_scop assert env.engine.context.cuda_graph_batch_dimensions_list model = env.engine.controller.inference_wrapped_model.model if inference_cuda_graph_scope == InferenceCudaGraphScope.block: - # hybrid models attach cudagraph_manager to the model; others attach to the decoder - if model_provider == "hybrid": - assert model.cudagraph_manager.cudagraph_runners - else: - assert model.decoder.cudagraph_manager.cudagraph_runners + # GPT and hybrid models both own the block-scope graph at model level; + # GPT removes the decoder's fallback manager at construction. + assert model.cudagraph_manager.cudagraph_runners + assert not hasattr(model.decoder, 'cudagraph_manager') else: # check if cudagraph runners are created at the layer level for layer in model.decoder.layers: @@ -972,7 +972,8 @@ def test_deprecated_full_iteration_inference_scope_matches_new_flag_runtime_beha model = env.engine.controller.inference_wrapped_model.model assert model.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block assert model.config.cuda_graph_modules == [] - assert model.decoder.cudagraph_manager.cudagraph_runners + assert model.cudagraph_manager.cudagraph_runners + assert not hasattr(model.decoder, 'cudagraph_manager') for layer in model.decoder.layers: assert not hasattr(layer, 'cudagraph_manager') From cab7116fcf6b89b8e1deeab6bbaceedd2c4b78f0 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 19 Aug 2026 09:11:34 +0000 Subject: [PATCH 096/101] chore: rotate oncall schedule --- .github/oncall_schedule.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/oncall_schedule.json b/.github/oncall_schedule.json index 622a830a73e..093166018f2 100644 --- a/.github/oncall_schedule.json +++ b/.github/oncall_schedule.json @@ -1,8 +1,4 @@ [ - { - "user": "maanug-nv", - "date": "2026-08-12" - }, { "user": "Phlip79", "date": "2026-08-19" @@ -46,5 +42,9 @@ { "user": "maanug-nv", "date": "2026-10-28" + }, + { + "user": "Phlip79", + "date": "2026-11-04" } ] From 01dc20b4a570391cf95183986861829d820553fd Mon Sep 17 00:00:00 2001 From: Tong Liu Date: Wed, 19 Aug 2026 08:12:51 +0000 Subject: [PATCH 097/101] Support no-shard, ZeRO-1, and ZeRO-2 in MFSDP v2 (#6137) Signed-off-by: tongliu Signed-off-by: Jingyue Wu Signed-off-by: Tong Liu Co-authored-by: Jingyue Wu Co-authored-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/dbuffer.py | 18 ++- .../src/megatron_fsdp/experimental/module.py | 1 + .../experimental/parameter_group.py | 148 +++++++++--------- .../distributed/mfsdp_v2/test_fully_shard.py | 74 ++++++--- .../distributed/mfsdp_v2/test_memory.py | 126 ++++++++++++++- .../distributed/mfsdp_v2/test_overlap.py | 106 ++++++++++--- 6 files changed, 344 insertions(+), 129 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 2f4e622d878..127dd6ab798 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -120,6 +120,13 @@ def device(self) -> torch.device: """Device of the local buffer.""" return self.local_buffer.device + @property + def is_symmetric_memory(self) -> bool: + """Whether the local buffer is allocated from symmetric memory.""" + return hasattr(symm_mem, "is_symm_mem_tensor") and symm_mem.is_symm_mem_tensor( + self.local_buffer + ) + def reallocate_storage(self) -> None: """Restore the local buffer's backing storage to its logical size.""" self._resize_storage(self.local_buffer.numel()) @@ -375,6 +382,10 @@ def allgather(self, mesh_axis: int, *, out: "DBuffer | None" = None) -> "DBuffer placements[mesh_axis] = Replicate() _validate_placements(placements) out = self._create_or_validate_out(out, placements=placements) + # Symmetric-memory registration is scoped to the collective's process + # group, so rendezvous the output on the same mesh axis as the all-gather. + if out.is_symmetric_memory: + out.rendezvous(mesh_axis) dist.all_gather_into_tensor( output_tensor=out.local_buffer, input_tensor=self.local_buffer, @@ -416,10 +427,7 @@ def reduce_scatter( reduce_op = partial_placement.reduce_op # Symmetric-memory MFSDP requires this detector, but ordinary DBuffer # reductions remain supported on older PyTorch versions that lack it. - is_symm_mem = hasattr(symm_mem, "is_symm_mem_tensor") and symm_mem.is_symm_mem_tensor( - self.local_buffer - ) - if is_symm_mem: + if self.is_symmetric_memory: self.rendezvous(axis) # NCCL symmetric-memory reduce-scatter selects its symmetric kernel # for SUM. Preserve the placement's AVG semantics by scaling the @@ -432,7 +440,7 @@ def reduce_scatter( op=reduce_op, group=self.mesh.get_group(axis), ) - if is_symm_mem and partial_placement.reduce_op == dist.ReduceOp.AVG: + if self.is_symmetric_memory and partial_placement.reduce_op == dist.ReduceOp.AVG: out.local_buffer.div_(self.mesh.size(axis)) return out diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index dbc0df69481..1fae3dbfa2e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -195,6 +195,7 @@ def __init__( mesh=mesh, placements=placements, mixed_precision_policy=mixed_precision_policy, + allgather_stream=context.allgather_stream, reduce_scatter_stream=context.reduce_scatter_stream, grad_divisor=grad_divisor, use_symmetric_memory=use_symmetric_memory, diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index aefcd15289d..60cb8f6a571 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -43,10 +43,10 @@ def get_containing_parameter_group(parameter: nn.Parameter) -> "FsdpParameterGro def sync_model_weights_from_main_weights(parameters: Iterable[nn.Parameter]) -> None: - """Refresh MFSDP compute weights for parameter groups represented by ``parameters``. + """Sync MFSDP compute weights for parameter groups represented by ``parameters``. Parameters outside the experimental MFSDP path are ignored. A parameter group - may own multiple parameters, but its compute-weight buffer is refreshed once. + may own multiple parameters, but its compute-weight buffer is synced once. """ seen_parameter_groups = set() for parameter in parameters: @@ -93,6 +93,7 @@ def __init__( mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy, + allgather_stream: torch.cuda.Stream, reduce_scatter_stream: torch.cuda.Stream, grad_divisor: int = 1, use_symmetric_memory: bool = False, @@ -105,6 +106,7 @@ def __init__( mesh: Device mesh used for all DBuffer storage in this version. placements: Parameter, gradient, and optimizer placements. mixed_precision_policy: Precision policy for main weights and gradients. + allgather_stream: Stream used to allocate model weights when a dtype cast is required. reduce_scatter_stream: Stream on which to allocate the main-gradient buffer. use_symmetric_memory: Allocate communication staging buffers from PyTorch's NCCL symmetric-memory pool. @@ -123,6 +125,10 @@ def __init__( model_weight_placements = tuple(placements.parameter) main_grad_placements = tuple(placements.gradient) main_weight_placements = tuple(placements.optimizer) + self._model_weight_placements = model_weight_placements + # main_grad rests here (DP-outer-Partial for HSDP) between microbatches and + # is finalized to main_weight's placements after the last microbatch. + self._main_grad_placements = main_grad_placements # Python dicts preserve insertion order, so parameter_to_fqns and # fsdp_parameters define the same stable DBuffer tensor order. @@ -158,6 +164,22 @@ def __init__( self._symm_mem_pool = symm_mem.get_mem_pool(self.main_weight.device) else: self._symm_mem_pool = None + if main_weight_dtype == self.dtype: + self.model_weight = self.main_weight + else: + # Record the all-gather stream as model_weight's allocation stream to allow + # its uses to be joined back to it before the buffer is deleted. + with torch.cuda.stream(allgather_stream): + self.model_weight = DBuffer( + mesh=self.mesh, + placements=main_weight_placements, + tensor_shapes=tensor_shapes, + dtype=self.dtype, + device=self.main_weight.device, + ) + # Cast into the preallocated model_weight on the current stream without + # replacing its storage or its all-gather allocation stream. + self.main_weight.cast(self.model_weight.dtype, out=self.model_weight) with self._symmetric_memory_context(): self._unsharded_model_weight = DBuffer( @@ -167,16 +189,6 @@ def __init__( dtype=self.dtype, device=self.main_weight.device, ) - if main_weight_dtype == self.dtype and main_weight_placements == model_weight_placements: - self.model_weight = self.main_weight - else: - self.model_weight = DBuffer( - mesh=self.mesh, - placements=model_weight_placements, - tensor_shapes=tensor_shapes, - dtype=self.dtype, - device=self.main_weight.device, - ) self.main_grad = None if self.requires_grad: @@ -198,9 +210,6 @@ def __init__( "main_grad is built from main_weight tensor shapes on the same mesh, " "and DBuffer layouts are deterministic from those shapes and mesh size." ) - # main_grad rests here (DP-outer-Partial for HSDP) between microbatches and - # is finalized to main_weight's placements after the last microbatch. - self._accumulation_placements = main_grad_placements fsdp_parameters: list[FsdpParameter] = [] main_grad_dtype = self.main_grad.dtype if self.main_grad is not None else None for index, (parameter, fqns) in enumerate(parameter_to_fqns.items()): @@ -231,11 +240,8 @@ def __init__( ) self.fsdp_parameters = tuple(fsdp_parameters) - # Compute weights must be initialized before the first forward; subsequent - # refreshes happen from the FSDP optimizer's post-step hook. - self.sync_model_weight_from_main_weight() - self._switch_to_sharded_parameters() self._unsharded_model_weight.release_storage() + self._switch_to_sharded_parameters() def _symmetric_memory_context(self): if self._symm_mem_pool is None: @@ -260,42 +266,45 @@ def _switch_to_unsharded_parameters(self) -> None: def sync_model_weight_from_main_weight(self) -> None: """Refresh compute weights from optimizer weights.""" - if self.main_weight is self.model_weight: - return - - if self.main_weight.placements == self.model_weight.placements: - self.main_weight.cast(self.model_weight.dtype, out=self.model_weight) - return - - # main_weight is typically the higher-precision optimizer dtype, while - # model_weight is the lower-precision compute dtype. Cast before redistributing - # so cross-rank communication moves the smaller compute-dtype payload. - self.main_weight.cast(self.model_weight.dtype).redistribute( - self.model_weight.placements, out=self.model_weight - ) + allgather_stream = self.model_weight.allocation_stream + assert allgather_stream is not None + current_stream = torch.cuda.current_stream(self.model_weight.device) + allgather_stream.wait_stream(current_stream) + with torch.cuda.stream(allgather_stream): + self.model_weight = self.main_weight.cast(self.model_weight.dtype) + # CUDA graph capture requires every forked stream to rejoin the capture + # stream before capture ends. + current_stream.wait_stream(allgather_stream) def unshard_parameters(self) -> None: """Install full parameters for local compute.""" - with self._symmetric_memory_context(): - self._unsharded_model_weight.reallocate_storage() - # This buffer backs unsharded Parameters whose views may be saved by autograd. - # Autograd records a tensor's version counter when saving it for backward, and - # in-place writes like the out= redistribution below increment that counter even - # under no_grad. Without preserving it, backward can fail with "modified by an - # inplace operation" even though FSDP only materialized internal storage. - gather_axis = changed_mesh_axis( - self.model_weight.placements, self._unsharded_model_weight.placements - ) - if gather_axis is None: - raise RuntimeError("FSDP parameter unshard requires a changed placement axis.") - with torch.autograd._unsafe_preserve_version_counter( - self._unsharded_model_weight.local_buffer - ): - if self._symm_mem_pool is not None: - self._unsharded_model_weight.rendezvous(gather_axis) - self.model_weight.redistribute( - self._unsharded_model_weight.placements, out=self._unsharded_model_weight - ) + # In ZeRO-1, the post-step cast leaves model_weight sharded. Only the first + # microbatch sees placements different from the configured model placements + # and restores the replicated model weight. + if self.model_weight.placements != self._model_weight_placements: + with self._symmetric_memory_context(): + # Allocate the restored destination in symmetric memory when enabled so the + # redistribution can use the faster symmetric-memory all-gather path. + self.model_weight = self.model_weight.redistribute(self._model_weight_placements) + if self.model_weight.placements == self._unsharded_model_weight.placements: + unsharded_model_weight = self.model_weight + else: + with self._symmetric_memory_context(): + self._unsharded_model_weight.reallocate_storage() + # This buffer backs unsharded Parameters whose views may be saved by autograd. + # Autograd records a tensor's version counter when saving it for backward, and + # in-place writes like the out= redistribution below increment that counter even + # under no_grad. Without preserving it, backward can fail with "modified by an + # inplace operation" even though FSDP only materialized internal storage. + with torch.autograd._unsafe_preserve_version_counter( + self._unsharded_model_weight.local_buffer + ): + self.model_weight.redistribute( + self._unsharded_model_weight.placements, out=self._unsharded_model_weight + ) + unsharded_model_weight = self._unsharded_model_weight + for index, fsdp_parameter in enumerate(self.fsdp_parameters): + fsdp_parameter.unsharded.data = unsharded_model_weight.get_local_tensor(index) self._switch_to_unsharded_parameters() def reshard_parameters(self) -> None: @@ -377,15 +386,15 @@ def reduce_partial_gradients( # on DP-outer), so relabel it in place; HFSDP's finalize reduce-scattered # DP-outer to a smaller optimizer-sharded buffer, so allocate a fresh one # (zeroed only when we accumulate into it, i.e. sharded grads are still set). - if self.main_grad.placements != self._accumulation_placements: + if self.main_grad.placements != self._main_grad_placements: assert self.main_grad.allocation_stream == ( torch.cuda.current_stream(self.main_grad.device) ) - reset_axis = changed_mesh_axis(self.main_grad.placements, self._accumulation_placements) + reset_axis = changed_mesh_axis(self.main_grad.placements, self._main_grad_placements) assert reset_axis is not None # the placements differ, so an axis changed if isinstance(self.main_grad.placements[reset_axis], Replicate): # HSDP: Replicate -> Partial changes only metadata and reuses the tensor. - self.main_grad = self.main_grad.redistribute(self._accumulation_placements) + self.main_grad = self.main_grad.redistribute(self._main_grad_placements) else: # HFSDP: main_grad was reduce-scattered to the optimizer shard, too small # to hold the accumulation, so re-allocate. This runs inside the @@ -395,7 +404,7 @@ def reduce_partial_gradients( # reduction below overwrites it via out=. self.main_grad = DBuffer( mesh=self.mesh, - placements=self._accumulation_placements, + placements=self._main_grad_placements, tensor_shapes=self.main_weight.layout.tensor_shapes, dtype=self.main_grad.dtype, device=self.main_weight.device, @@ -403,27 +412,20 @@ def reduce_partial_gradients( if has_sharded_grads: self.main_grad.local_buffer.zero_() - can_reduce_into_main_grad = ( + if can_reduce_into_main_grad := ( not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype - ) - reduce_axis = changed_mesh_axis(partial_grad.placements, self.main_grad.placements) - if reduce_axis is None: - raise RuntimeError("FSDP gradient reduction requires a changed placement axis.") - if self._symm_mem_pool is not None: - partial_grad.rendezvous(reduce_axis) - # Divide this backward's contribution, not the accumulated total: with plain - # all-Flat DP every backward is a last microbatch, so main_grad accumulates - # across microbatches below and a scale applied to the running sum would - # compound. Dividing before the deferred DP-outer reduction is equivalent because - # both that reduction and this scale are linear. - if can_reduce_into_main_grad: + ): partial_grad.redistribute(self.main_grad.placements, out=self.main_grad) - if self.grad_divisor != 1: - self.main_grad.local_buffer.div_(self.grad_divisor) + reduced_grad = self.main_grad else: reduced_grad = partial_grad.redistribute(self.main_grad.placements) - if self.grad_divisor != 1: - reduced_grad.local_buffer.div_(self.grad_divisor) + + # Scale this backward's contribution before accumulating it so repeated + # backwards do not repeatedly scale the running total. + if self.grad_divisor != 1: + reduced_grad.local_buffer.div_(self.grad_divisor) + + if reduced_grad is not self.main_grad: if has_sharded_grads: self.main_grad.local_buffer.add_(reduced_grad.local_buffer) else: diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index dcb154cbcd8..ed02b4be048 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -135,6 +135,28 @@ def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) +def _no_shard_placements() -> Placements: + return Placements( + dp_axes=[0], + parameter=[Replicate()], + gradient=[Partial(dist.ReduceOp.AVG)], + optimizer=[Replicate()], + ) + + +def _zero1_placements() -> Placements: + return Placements( + dp_axes=[0], + parameter=[Replicate()], + gradient=[Partial(dist.ReduceOp.AVG)], + optimizer=[Flat()], + ) + + +def _zero2_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Replicate()], gradient=[Flat()], optimizer=[Flat()]) + + def _hsdp_placements() -> Placements: """HSDP: params/optimizer replicated across DP-outer (axis 0), sharded within DP-inner (axis 1). main_grad rests [Partial, Flat] between microbatches and is @@ -166,9 +188,16 @@ def _hfsdp_placements() -> Placements: _ALLREDUCE_OP_NAME_SUBSTRING = "allreduce" +@pytest.mark.parametrize( + "placements_factory", + [_no_shard_placements, _zero1_placements, _zero2_placements, _flat_placements], + ids=["no_shard", "zero1", "zero2", "zero3"], +) @pytest.mark.parametrize("num_microbatches", [1, 3]) -def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatches): - """Minimal per-module FSDP training should match single-rank SGD.""" +def test_fully_shard_sgd_losses_match_baseline( + distributed_setup, num_microbatches, placements_factory +): + """Every supported sharding strategy should match single-rank SGD.""" rank = distributed_setup.rank world_size = distributed_setup.world_size device = distributed_setup.device @@ -176,16 +205,18 @@ def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatch pytest.skip("This test requires at least 2 ranks.") mesh = init_device_mesh(device.type, (world_size,)) + placements = placements_factory() torch.manual_seed(1234) baseline = TinyModel().to(device) model = TinyModel().to(device) model.load_state_dict(baseline.state_dict()) - with fully_shard_context(device=device): - fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) - fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device) as context: + fully_shard(model.fc1, mesh=mesh, placements=placements) + fully_shard(model.fc2, mesh=mesh, placements=placements) baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + fully_shard_optimizer(optimizer) micro_batch_size = 2 x = torch.randn(num_microbatches, micro_batch_size, 8, device=device) @@ -197,19 +228,19 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]: for step in range(5): optimizer.zero_grad() - for microbatch, (microbatch_x, microbatch_target) in enumerate(microbatches): - loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) - losses.append(loss.detach()) - logger.debug( - "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s", - log_prefix, - rank, - step, - microbatch, - loss, - ) - - (loss / num_microbatches).backward() + for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): + with microbatch(context, is_last=microbatch_index == num_microbatches - 1): + loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) + losses.append(loss.detach()) + logger.debug( + "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s", + log_prefix, + rank, + step, + microbatch_index, + loss, + ) + (loss / num_microbatches).backward() optimizer.step() return losses @@ -643,12 +674,13 @@ def test_backward_averages_across_dp_and_accumulates_across_calls(distributed_se model = nn.Linear(1, world_size, bias=False).to(device) nn.init.constant_(model.weight, 1.0) - with fully_shard_context(device=device): + with fully_shard_context(device=device) as context: fully_shard(model, mesh=mesh, placements=_flat_placements()) x = torch.full((1, 1), float(rank + 1), device=device) - model(x).sum().backward() - model(x).sum().backward() + with microbatch(context, is_last=False): + model(x).sum().backward() + model(x).sum().backward() assert isinstance(model.weight.grad, DTensor) local_grad = model.weight.grad.to_local() diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_memory.py b/tests/unit_tests/distributed/mfsdp_v2/test_memory.py index 78e427788cc..0c2c65685be 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_memory.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_memory.py @@ -6,12 +6,15 @@ import pytest import torch +import torch.distributed as dist from torch import nn from torch.distributed.device_mesh import init_device_mesh from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, + Partial, Placements, + Replicate, fully_shard, fully_shard_context, fully_shard_optimizer, @@ -53,6 +56,19 @@ def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) +def _zero1_placements() -> Placements: + return Placements( + dp_axes=[0], + parameter=[Replicate()], + gradient=[Partial(dist.ReduceOp.AVG)], + optimizer=[Flat()], + ) + + +def _zero2_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Replicate()], gradient=[Flat()], optimizer=[Flat()]) + + def _mb(num_bytes: int) -> str: return f"{num_bytes / 1024**2:.2f} MB" @@ -162,6 +178,97 @@ def train_step() -> None: ) +def test_zero1_memory_matches_sharded_optimizer_and_replicated_weight(distributed_setup): + """First-step peak memory should match the expected ZeRO-1 training allocations.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + dim = 4096 + num_tokens = 256 + dtype = torch.bfloat16 + x = torch.ones(num_tokens, dim, device=device, dtype=dtype) + model = ElementwiseModel(dim).to(device=device, dtype=dtype) + mesh = init_device_mesh(device.type, (world_size,)) + placements = _zero1_placements() + with fully_shard_context(device=device): + fully_shard(model, mesh=mesh, placements=placements) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01, foreach=False) + fully_shard_optimizer(optimizer) + (parameter_group,) = model.parameter_groups + + allocated_before_training = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + + loss = model(x).sum() + loss.backward() + optimizer.step() + assert parameter_group.model_weight.placements == (Flat(),) + + actual_optimizer_size = sum( + state["exp_avg"].to_local().nbytes + state["exp_avg_sq"].to_local().nbytes + for state in optimizer.state.values() + ) + actual_training_growth = torch.cuda.max_memory_allocated(device) - allocated_before_training + + # Theoretical sharded Adam size. + shard_numel = dim * dim // world_size + fp32_size = torch.empty((), dtype=torch.float32).element_size() + bf16_size = torch.empty((), dtype=dtype).element_size() + theoretical_optimizer_size = 2 * shard_numel * fp32_size + + # ElementwiseModel itself is the only layer in this test model. + num_layers = 1 + full_bf16_weight_size = dim * dim * bf16_size + sharded_bf16_weight_size = full_bf16_weight_size // world_size + sharded_fp32_weight_size = shard_numel * fp32_size + activation_size = num_tokens * dim * bf16_size + # Backward peak growth: + # - Unsharded model weight: dim * dim * bf16_size. + backward_model_weight_size = full_bf16_weight_size + # - Full parameter gradient: dim * dim * bf16_size. + backward_parameter_gradient_size = full_bf16_weight_size + # - Unreduced partial-gradient buffer: dim * dim * bf16_size. + backward_partial_gradient_size = full_bf16_weight_size + # - Earlier-layer activations: (num_layers - 1) * num_tokens * dim * bf16_size. + # This one-layer model has no stacked activation at its backward peak. + backward_stacked_activation_size = (num_layers - 1) * activation_size + # - Released baseline model-weight shard: -(dim * dim * bf16_size / world_size). + released_sharded_model_weight_size = sharded_bf16_weight_size + theoretical_backward_growth = ( + backward_model_weight_size + + backward_parameter_gradient_size + + backward_partial_gradient_size + + backward_stacked_activation_size + - released_sharded_model_weight_size + ) + + # First-Adam-step peak growth: + # - Model weight grows from sharded to full BF16: full - sharded BF16 weight. + optimizer_model_weight_growth = full_bf16_weight_size - sharded_bf16_weight_size + # - BF16 main gradient shrinks from full to sharded: -(full - sharded BF16 weight). + optimizer_main_gradient_reduction = full_bf16_weight_size - sharded_bf16_weight_size + # - Casted gradient: shard_numel * fp32_size. + optimizer_casted_gradient_size = sharded_fp32_weight_size + # - Optimizer states exp_avg and exp_avg_sq: 2 * shard_numel * fp32_size. + optimizer_state_size = 2 * sharded_fp32_weight_size + # - Two FP32 buffers that briefly coexist while Adam computes the + # bias-corrected denominator: 2 * shard_numel * fp32_size. + optimizer_temporary_size = 2 * sharded_fp32_weight_size + theoretical_optimizer_growth = ( + optimizer_model_weight_growth + - optimizer_main_gradient_reduction + + optimizer_casted_gradient_size + + optimizer_state_size + + optimizer_temporary_size + ) + theoretical_training_growth = max(theoretical_backward_growth, theoretical_optimizer_growth) + + assert actual_optimizer_size == theoretical_optimizer_size + assert abs(actual_training_growth - theoretical_training_growth) < 1024**2 + + def test_deleted_model_releases_fsdp_storage(distributed_setup): """Deleting an FSDP model should release its persistent storage.""" world_size = distributed_setup.world_size @@ -231,8 +338,16 @@ def assert_returns_to_resting_memory(phase: str) -> None: assert_returns_to_resting_memory("backward") -def test_fully_shard_reduces_peak_training_memory(distributed_setup): - """Per-layer FSDP should reduce peak CUDA memory during training.""" +@pytest.mark.parametrize( + "placements_factory", + [ + pytest.param(_zero1_placements, id="zero1"), + pytest.param(_zero2_placements, id="zero2"), + pytest.param(_flat_placements, id="zero3"), + ], +) +def test_fully_shard_reduces_peak_training_memory(distributed_setup, placements_factory): + """Per-layer FSDP should reduce peak CUDA memory for each sharding strategy.""" rank = distributed_setup.rank world_size = distributed_setup.world_size device = distributed_setup.device @@ -270,7 +385,7 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten fully_shard( layer, mesh=mesh, - placements=_flat_placements(), + placements=placements_factory(), mixed_precision_policy=MixedPrecisionPolicy( main_params_dtype=dtype, main_grads_dtype=dtype ), @@ -288,4 +403,7 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten _mb(sharded_peak), ) - assert sharded_peak < baseline_peak + assert sharded_peak < baseline_peak, ( + f"Expected FSDP to reduce peak training memory on rank {rank}: " + f"baseline={_mb(baseline_peak)}, sharded={_mb(sharded_peak)}" + ) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py b/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py index 7d83526442f..3838f611d8c 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_overlap.py @@ -11,9 +11,13 @@ from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, + Partial, Placements, + Replicate, fully_shard, fully_shard_context, + fully_shard_optimizer, + microbatch, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import ( @@ -42,6 +46,19 @@ def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) +def _zero1_placements() -> Placements: + return Placements( + dp_axes=[0], + parameter=[Replicate()], + gradient=[Partial(dist.ReduceOp.AVG)], + optimizer=[Flat()], + ) + + +def _zero2_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Replicate()], gradient=[Flat()], optimizer=[Flat()]) + + # CPU ops that a device event chains up to via cpu_parent, used to attribute the device # work to its enclosing collective or matmul operation. _ALL_GATHER_OP_NAME_SUBSTRING = "allgather" @@ -49,6 +66,14 @@ def _flat_placements() -> Placements: _GEMM_OP_NAME_SUBSTRING = "aten::mm" +@pytest.mark.parametrize( + "placements_factory", + [ + pytest.param(_zero1_placements, id="zero1"), + pytest.param(_zero2_placements, id="zero2"), + pytest.param(_flat_placements, id="zero3"), + ], +) @pytest.mark.parametrize( "use_symmetric_memory", [ @@ -67,9 +92,9 @@ def _flat_placements() -> Placements: "unify_communication_stream", [False, True], ids=["separate_streams", "unified_stream"] ) def test_overlaps_communication_and_compute( - distributed_setup, use_symmetric_memory, unify_communication_stream + distributed_setup, placements_factory, use_symmetric_memory, unify_communication_stream ): - """Forward and backward communication should overlap GEMM compute.""" + """ZeRO-1/2/3 communication should overlap GEMM compute.""" world_size = distributed_setup.world_size device = distributed_setup.device if world_size < 2: @@ -85,6 +110,7 @@ def test_overlaps_communication_and_compute( dim = 16384 num_children = 4 dtype = torch.bfloat16 + placements = placements_factory() # new_group requires a default process group. Initialize it here so this test works # in isolation. Do not eagerly initialize it with device_id in the shared fixture: @@ -113,27 +139,32 @@ def test_overlaps_communication_and_compute( mesh = DeviceMesh.from_group(dp_group, device.type) model = MultiChildModel(dim=dim, num_children=num_children).to(device=device, dtype=dtype) - placements = _flat_placements() policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) with fully_shard_context( device=device, use_symmetric_memory=use_symmetric_memory, unify_communication_stream=unify_communication_stream, - ): + ) as context: for layer in model.layers: fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) - x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01, foreach=False) + fully_shard_optimizer(optimizer) + x = torch.randn(8192, dim, device=device, dtype=dtype, requires_grad=True) - def train_one_iteration() -> None: - model.zero_grad(set_to_none=True) - model(x).sum().backward() + def train_one_step() -> None: + """Run one optimizer step consuming two microbatches.""" + optimizer.zero_grad(set_to_none=True) + for microbatch_index in range(2): + with microbatch(context, is_last=microbatch_index == 1): + model(x).sum().backward() + optimizer.step() - train_one_iteration() + train_one_step() torch.cuda.synchronize(device) with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: - train_one_iteration() + train_one_step() # Synchronize inside the profiler context so in-flight device kernels # complete and get recorded before the profiler stops on __exit__. # Synchronizing after the context would finalize the trace first and @@ -141,26 +172,51 @@ def train_one_iteration() -> None: torch.cuda.synchronize(device) gemm_kernels = collect_linked_kernels(prof, _GEMM_OP_NAME_SUBSTRING) - # Each child Linear runs one forward and two backward matmuls. aten::mm may also - # launch auxiliary kernels, so check only the matmul lower bound. - assert len(gemm_kernels) >= 3 * num_children, ( - f"Expected at least {3 * num_children} kernels linked to GEMMs, got " + # Each child Linear runs one forward and two backward matmuls per microbatch. + # aten::mm may also launch auxiliary kernels, so check only the matmul lower bound. + expected_gemm_count = 6 * num_children + assert len(gemm_kernels) >= expected_gemm_count, ( + f"Expected at least {expected_gemm_count} kernels linked to GEMMs, got " f"{len(gemm_kernels)}: " f"{[kernel.name for kernel in gemm_kernels]}" ) allgather_kernels = collect_linked_kernels(prof, _ALL_GATHER_OP_NAME_SUBSTRING) reduce_scatter_kernels = collect_linked_kernels(prof, _REDUCE_SCATTER_OP_NAME_SUBSTRING) - # Each child layer does a forward and a backward all-gather and one - # reduce-scatter. Zero-CTA moves the all-gather to copy-engine memcpys, so it - # should not emit all-gather kernels. - expected_allgather_kernel_count = 0 if use_symmetric_memory else 2 * num_children + # ZeRO-1/2 all-gather once per optimizer step. ZeRO-1 reduces only the second + # microbatch, while ZeRO-2 reduces both. ZeRO-3 gathers for every forward and + # backward and reduces every microbatch. + expected_collective_counts = { + _zero1_placements: (num_children, num_children, num_children - 1, num_children - 1), + _zero2_placements: ( + num_children, + 2 * num_children, + num_children - 1, + 2 * (num_children - 1), + ), + _flat_placements: ( + 4 * num_children, + 2 * num_children, + 4 * (num_children - 1), + 2 * (num_children - 1), + ), + } + ( + expected_allgather_count, + expected_reduce_scatter_count, + expected_allgather_overlap_count, + expected_reduce_scatter_overlap_count, + ) = expected_collective_counts[placements_factory] + + expected_allgather_kernel_count = 0 if use_symmetric_memory else expected_allgather_count + # Zero-CTA moves the all-gather to copy-engine memcpys, so it should not emit + # all-gather kernels. assert len(allgather_kernels) == expected_allgather_kernel_count, ( f"Expected {expected_allgather_kernel_count} all-gather kernels, got " f"{len(allgather_kernels)}: {[kernel.name for kernel in allgather_kernels]}" ) - assert len(reduce_scatter_kernels) == num_children, ( - f"Expected {num_children} reduce-scatter kernels, got " + assert len(reduce_scatter_kernels) == expected_reduce_scatter_count, ( + f"Expected {expected_reduce_scatter_count} reduce-scatter kernels, got " f"{len(reduce_scatter_kernels)}: {[kernel.name for kernel in reduce_scatter_kernels]}" ) @@ -184,15 +240,13 @@ def train_one_iteration() -> None: any(events_overlap(kernel, gemm) for gemm in gemm_kernels) for kernel in reduce_scatter_kernels ) - expected_allgather_overlap = 2 * (num_children - 1) - expected_reduce_scatter_overlap = num_children - 1 if not use_symmetric_memory: - assert allgather_overlap_count >= expected_allgather_overlap, ( - f"Expected at least {expected_allgather_overlap} all-gathers to " + assert allgather_overlap_count == expected_allgather_overlap_count, ( + f"Expected exactly {expected_allgather_overlap_count} all-gathers to " f"overlap compute, got {allgather_overlap_count}/{len(allgather_kernels)}." ) - assert reduce_scatter_overlap_count >= expected_reduce_scatter_overlap, ( - f"Expected at least {expected_reduce_scatter_overlap} reduce-scatters to overlap " + assert reduce_scatter_overlap_count == expected_reduce_scatter_overlap_count, ( + f"Expected exactly {expected_reduce_scatter_overlap_count} reduce-scatters to overlap " f"compute, got {reduce_scatter_overlap_count}/{len(reduce_scatter_kernels)}." ) From f481e6361520dcac1554891a6ae83b353eb1d91b Mon Sep 17 00:00:00 2001 From: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:52:04 +0000 Subject: [PATCH 098/101] Optimize MIMO pipeline communication (#6606) Signed-off-by: ykarnati --- .../pipeline_parallel/bridge_communicator.py | 50 ++++++-- .../multimodule_communicator.py | 4 + megatron/core/pipeline_parallel/schedules.py | 11 +- .../test_bridge_communicator.py | 67 ++++++++++ .../test_multimodule_communicator.py | 5 + .../pipeline_parallel/test_schedules.py | 120 ++++++++++++++++++ 6 files changed, 246 insertions(+), 11 deletions(-) diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py index 94b14317842..15f2bf5c524 100644 --- a/megatron/core/pipeline_parallel/bridge_communicator.py +++ b/megatron/core/pipeline_parallel/bridge_communicator.py @@ -384,7 +384,7 @@ def send_forward(self, tensor_to_send: torch.Tensor): self.current_rank, dest_rank, ) - dist.send(tensor_split, dst=dest_rank, group=self.bridge_pg) + self._run_batched_payload_p2p(tensor_splits, rank_info.send_to_ranks, op="send") def recv_forward(self) -> torch.Tensor: """Receive forward activation tensor. @@ -433,8 +433,14 @@ def recv_forward(self) -> torch.Tensor: dtype=self.comm_dtype, requires_grad=True, ) - dist.recv(tensor_to_recv, src=src_rank, group=self.bridge_pg) - if logger.isEnabledFor(logging.DEBUG): + received_tensors_list.append(tensor_to_recv) + self._run_batched_payload_p2p( + received_tensors_list, rank_info.recv_from_ranks, op="recv" + ) + if logger.isEnabledFor(logging.DEBUG): + for src_rank, tensor_to_recv in zip( + rank_info.recv_from_ranks, received_tensors_list + ): logger.debug( "[Bridge Communicator] [receive_forward] Rank %s " "received tensor from src rank %s shape %s sum %s", @@ -443,7 +449,6 @@ def recv_forward(self) -> torch.Tensor: tensor_to_recv.shape, tensor_to_recv.sum(), ) - received_tensors_list.append(tensor_to_recv) aggregated_tensor = torch.cat(received_tensors_list, dim=self._batch_dim) if logger.isEnabledFor(logging.DEBUG): logger.debug( @@ -527,7 +532,6 @@ def send_backward(self, grad_tensor: torch.Tensor): self._communicate_shapes(tensor_to_send_prev=tensor_splits) if num_receives > 0: for src_rank, tensor_split in zip(rank_info.recv_from_ranks, tensor_splits): - # Send the gradient split back to the source rank if logger.isEnabledFor(logging.DEBUG): logger.debug( "[Bridge Communicator] [send_backward] Rank %s " @@ -537,7 +541,7 @@ def send_backward(self, grad_tensor: torch.Tensor): tensor_split.shape, tensor_split.sum(), ) - dist.send(tensor_split, dst=src_rank, group=self.bridge_pg) + self._run_batched_payload_p2p(tensor_splits, rank_info.recv_from_ranks, op="send") def recv_backward(self) -> torch.Tensor: """Receive backward gradient tensor. @@ -579,8 +583,12 @@ def recv_backward(self) -> torch.Tensor: grad_tensor = torch.empty( grad_shape, device=torch.cuda.current_device(), dtype=self.comm_dtype ) - dist.recv(grad_tensor, src=dest_rank, group=self.bridge_pg) - if logger.isEnabledFor(logging.DEBUG): + received_gradients_list.append(grad_tensor) + self._run_batched_payload_p2p( + received_gradients_list, rank_info.send_to_ranks, op="recv" + ) + if logger.isEnabledFor(logging.DEBUG): + for dest_rank, grad_tensor in zip(rank_info.send_to_ranks, received_gradients_list): logger.debug( "[Bridge Communicator] [receive_backward] Rank %s " "received gradient from dest rank %s shape %s sum %s", @@ -589,7 +597,6 @@ def recv_backward(self) -> torch.Tensor: grad_tensor.shape, grad_tensor.sum(), ) - received_gradients_list.append(grad_tensor) # Concatenate received gradients aggregated_gradient = torch.cat(received_gradients_list, dim=self._batch_dim) @@ -925,6 +932,31 @@ def send_backward_recv_forward( ) return received_activation + def _run_batched_payload_p2p( + self, tensors: List[torch.Tensor], peers: List[int], *, op: str + ) -> None: + """Launch one payload operation per peer in a single P2P batch.""" + if len(tensors) != len(peers): + raise ValueError( + f"expected one payload tensor per peer, got {len(tensors)} tensors " + f"for {len(peers)} peers" + ) + if op == "send": + p2p_op = dist.isend + elif op == "recv": + p2p_op = dist.irecv + else: + raise ValueError(f"unsupported bridge P2P operation: {op}") + + ops = [ + dist.P2POp(p2p_op, tensor, peer, self.bridge_pg) for tensor, peer in zip(tensors, peers) + ] + if not ops: + return + requests = dist.batch_isend_irecv(ops) + for request in requests: + request.wait() + def _communicate_shapes( self, tensor_to_send_next: Optional[torch.Tensor] = None, diff --git a/megatron/core/pipeline_parallel/multimodule_communicator.py b/megatron/core/pipeline_parallel/multimodule_communicator.py index c27d3f753e2..d5d2b76a843 100644 --- a/megatron/core/pipeline_parallel/multimodule_communicator.py +++ b/megatron/core/pipeline_parallel/multimodule_communicator.py @@ -199,6 +199,10 @@ def is_pp_first_stage(self): return True return False + def is_module_pp_first_stage(self, module_name: str) -> bool: + """Return True if the current rank is the module's first PP stage.""" + return self.rank_module_map[module_name].pp_rank == 0 + @property def is_pp_last_stage(self): """Return True if the current rank has the absolute last stage in the overall model. diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index aa3640f32ae..fdbc8a9a60c 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -2255,6 +2255,13 @@ def enable_grad_sync(): disable_grad_sync() + grad_sync_first_stage = p2p_communicator.is_pp_first_stage + if isinstance(p2p_communicator, MultiModulePipelineCommunicator): + grad_sync_first_stage = all( + p2p_communicator.is_module_pp_first_stage(module_name) + for module_name in p2p_communicator.rank_module_map + ) + # Compute number of warmup microbatches. num_warmup_microbatches = p2p_communicator.total_stages - p2p_communicator.current_stage - 1 num_warmup_microbatches = min(num_warmup_microbatches, num_microbatches) @@ -2412,7 +2419,7 @@ def enable_grad_sync(): # Enable grad sync for the last microbatch in the batch if the full # backward pass completes in the 1F1B stage. if num_warmup_microbatches == 0 and last_iteration: - if config.grad_sync_func is None or p2p_communicator.is_pp_first_stage: + if config.grad_sync_func is None or grad_sync_first_stage: enable_grad_sync() input_tensor_grad = backward_func( @@ -2439,7 +2446,7 @@ def enable_grad_sync(): # pipeline stages do grad reduction during pipeline # bubble. if i == num_warmup_microbatches - 1: - if config.grad_sync_func is None or p2p_communicator.is_pp_first_stage: + if config.grad_sync_func is None or grad_sync_first_stage: enable_grad_sync() input_tensor = input_tensors.pop(0) diff --git a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py index 4eaa11b8c24..1d3f1f12989 100644 --- a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py +++ b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py @@ -270,6 +270,73 @@ def test_as_per_peer_tensors_rejects_count_mismatch(self): with pytest.raises(ValueError, match="expected 3 tensors for shape communication, got 2"): BridgeCommunicator._as_per_peer_tensors(tensors, expected_count=3) + @pytest.mark.parametrize( + "op, expected_callable", + [("send", torch.distributed.isend), ("recv", torch.distributed.irecv)], + ) + def test_batched_payload_launches_all_peers_before_waiting( + self, monkeypatch, op, expected_callable + ): + bridge = self._bridge() + bridge.bridge_pg = object() + peers = [17, 5, 29] + tensors = [torch.empty((3, 2)), torch.empty((0, 2)), torch.empty((7, 2))] + created_ops = [] + calls = [] + + class FakeOp: + def __init__(self, p2p_op, tensor, peer, group): + self.op = p2p_op + self.tensor = tensor + self.peer = peer + self.group = group + created_ops.append(self) + calls.append(("construct", peer)) + + class FakeWork: + def __init__(self, index): + self.index = index + + def wait(self): + assert len(created_ops) == len(peers) + calls.append(("wait", self.index)) + + def fake_batch(ops): + assert list(ops) == created_ops + calls.append(("launch", tuple(item.peer for item in ops))) + return [FakeWork(index) for index in range(len(ops))] + + monkeypatch.setattr(torch.distributed, "P2POp", FakeOp) + monkeypatch.setattr(torch.distributed, "batch_isend_irecv", fake_batch) + + bridge._run_batched_payload_p2p(tensors, peers, op=op) + + assert calls == [ + ("construct", 17), + ("construct", 5), + ("construct", 29), + ("launch", (17, 5, 29)), + ("wait", 0), + ("wait", 1), + ("wait", 2), + ] + assert [item.op for item in created_ops] == [expected_callable] * len(peers) + assert all(item.tensor is tensor for item, tensor in zip(created_ops, tensors)) + assert [item.peer for item in created_ops] == peers + assert all(item.group is bridge.bridge_pg for item in created_ops) + + def test_batched_payload_rejects_invalid_mapping_before_launch(self, monkeypatch): + bridge = self._bridge() + bridge.bridge_pg = object() + monkeypatch.setattr( + torch.distributed, + "batch_isend_irecv", + lambda _ops: (_ for _ in ()).throw(AssertionError("invalid batch was launched")), + ) + + with pytest.raises(ValueError, match="one payload tensor per peer"): + bridge._run_batched_payload_p2p([torch.empty(1)], [4, 5], op="send") + class TestBridgeCommunicator: diff --git a/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py index e7026c714ea..29766a292f1 100644 --- a/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py +++ b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py @@ -79,6 +79,11 @@ def test_multimodule_communicator_init(self): assert mllm_comm.config == config assert mllm_comm.current_rank == dist.get_rank() + for module_name, rank_module_info in mllm_comm.rank_module_map.items(): + assert mllm_comm.is_module_pp_first_stage(module_name) == ( + rank_module_info.pp_rank == 0 + ) + def test_compute_total_pipeline_stages(self): """Test compute_total_pipeline_stages for overall chain and until specific ranks.""" diff --git a/tests/unit_tests/pipeline_parallel/test_schedules.py b/tests/unit_tests/pipeline_parallel/test_schedules.py index 92db675d193..616f7d1e7fe 100644 --- a/tests/unit_tests/pipeline_parallel/test_schedules.py +++ b/tests/unit_tests/pipeline_parallel/test_schedules.py @@ -15,6 +15,7 @@ from megatron.core import ModelParallelConfig from megatron.core.distributed.finalize_model_grads import finalize_model_grads from megatron.core.hyper_comm_grid import HyperCommGrid +from megatron.core.pipeline_parallel.multimodule_communicator import MultiModulePipelineCommunicator from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection @@ -569,6 +570,125 @@ def set_input_tensor(input_tensor): Utils.destroy_model_parallel() +@pytest.mark.parametrize( + "communicator_base,module_first_stage,pipeline_stages", + [ + (P2PCommunicator, None, 1), + (MultiModulePipelineCommunicator, True, 1), + (MultiModulePipelineCommunicator, True, 2), + ], +) +def test_schedule_enables_grad_sync_on_first_stage( + monkeypatch, communicator_base, module_first_stage, pipeline_stages +): + events = [] + + @contextmanager + def no_sync(): + events.append("enter_no_sync") + try: + yield + finally: + events.append("exit_no_sync") + + config = SimpleNamespace( + overlap_p2p_comm=False, + variable_seq_lengths=True, + finalize_model_grads_func=None, + timers=None, + no_sync_func=no_sync, + num_microbatches_with_partial_activation_checkpoints=None, + deallocate_pipeline_outputs=False, + grad_sync_func=lambda _parameters: events.append("grad_sync"), + calculate_per_token_loss=False, + ) + model = SimpleNamespace(config=config, parameters=lambda: []) + + is_multimodule = communicator_base is MultiModulePipelineCommunicator + + class FakeCommunicator(communicator_base): + total_stages = pipeline_stages + current_stage = 0 + is_pp_first_stage = not is_multimodule + is_pp_last_stage = True + + def __init__(self): + self.config = config + if is_multimodule: + self.rank_module_map = {"llm": SimpleNamespace()} + + def is_module_pp_first_stage(self, _module_name): + return module_first_stage + + @staticmethod + def recv_forward(*_args): + return {"llm": None} if is_multimodule else None + + @staticmethod + def send_forward_recv_backward(*_args): + return None + + @staticmethod + def send_forward(*_args): + return None + + @staticmethod + def recv_backward(*_args): + return None + + @staticmethod + def send_backward(*_args): + return None + + monkeypatch.setattr( + schedule, + "forward_step", + lambda *_args, **_kwargs: ( + {"llm": torch.tensor(1.0)} if is_multimodule else torch.tensor(1.0), + torch.tensor(0), + ), + ) + monkeypatch.setattr( + schedule, "backward_step", lambda *_args, **_kwargs: events.append("backward") or None + ) + monkeypatch.setattr( + schedule, + "backward_step_multimodule", + lambda *_args, **_kwargs: events.append("backward") or {"llm": None}, + ) + monkeypatch.setattr(schedule, "deallocate_output_tensor", lambda *_args: None) + real_zeros = torch.zeros + monkeypatch.setattr( + schedule.torch, + "zeros", + lambda *args, **kwargs: real_zeros( + *args, **{key: value for key, value in kwargs.items() if key != "device"} + ), + ) + + if is_multimodule: + pg_collection = SimpleNamespace( + has_language_model=lambda: False, language_model_module_name=None + ) + else: + pg_collection = ProcessGroupCollection() + pg_collection.tp = SimpleNamespace(size=lambda: 1) + pg_collection.cp = SimpleNamespace(size=lambda: 1) + + schedule.forward_backward_pipelining_without_interleaving( + forward_step_func=None, + data_iterator=None, + model=model, + num_microbatches=1, + seq_length=1, + micro_batch_size=1, + p2p_communicator=FakeCommunicator(), + pg_collection=pg_collection, + ) + + assert events == ["enter_no_sync", "exit_no_sync", "backward"] + + def test_forward_backward_func_with_pipeline_parallel(mocker): from megatron.core.pipeline_parallel import get_forward_backward_func From 731b791469004f8fdcb896e65d610d6b1b0ebb32 Mon Sep 17 00:00:00 2001 From: wdykas Date: Wed, 19 Aug 2026 11:17:50 +0000 Subject: [PATCH 099/101] Add hybrid SSM state handoff (#6584) Signed-off-by: William Dykas --- .../attention_context/mamba_metadata.py | 47 ++- .../inference/contexts/dynamic_context.py | 45 ++- .../disaggregation/decode_admission.py | 16 +- .../core/inference/disaggregation/engine.py | 4 +- .../disaggregation/inference_state_handoff.py | 359 ++++++++++++++--- .../disaggregation/pending_handoff_imports.py | 11 +- .../disaggregation/transfer_backends/base.py | 36 +- .../disaggregation/transfer_backends/nccl.py | 5 + .../core/inference/engines/dynamic_engine.py | 26 +- megatron/core/inference/inference_request.py | 1 + .../text_generation_controller.py | 27 +- .../attention_metadata/test_mamba_metadata.py | 44 +++ .../contexts/test_dynamic_context.py | 57 ++- .../contexts/test_dynamic_prefix_caching.py | 42 +- .../test_disagg_handoff_lifecycle.py | 361 +++++++++++++++++- .../inference/test_kv_transfer_backends.py | 21 + .../test_text_generation_controller.py | 26 ++ 17 files changed, 993 insertions(+), 135 deletions(-) diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index da48faced90..fd4bfacc47d 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -155,10 +155,7 @@ def reset(self) -> None: self.reset_varlen_metadata() - # Re-initialize the free slot pool - self.mamba_state_free_slots = torch.arange( - self.max_requests, dtype=torch.int32, device='cpu' - ) + torch.arange(self.max_requests, out=self.mamba_state_free_slots) self.mamba_state_free_slot_count = self.max_requests def reset_varlen_metadata(self) -> None: @@ -750,8 +747,27 @@ def allocate_slot(self) -> Optional[int]: self.mamba_state_free_slot_count -= 1 mamba_idx = self.mamba_state_free_slots[self.mamba_state_free_slot_count] + return int(mamba_idx) + + def detach_state_slot(self, request_idx: int) -> int: + """Detach and return a request's live state slot without freeing it.""" + + mamba_idx = int(self.request_to_mamba_state_idx[request_idx].item()) + if mamba_idx < 0: + raise RuntimeError(f"Request index {request_idx} has no live Mamba state slot") + self.request_to_mamba_state_idx[request_idx] = -1 return mamba_idx + def free_slot(self, mamba_idx: int) -> None: + """Return one unbound slot to the live Mamba state pool.""" + + if not 0 <= mamba_idx < self.max_requests: + raise ValueError(f"Mamba state slot {mamba_idx} is outside the live state pool") + if self.mamba_state_free_slot_count >= self.max_requests: + raise RuntimeError("Cannot free a Mamba state slot when the pool is already full") + self.mamba_state_free_slots[self.mamba_state_free_slot_count] = mamba_idx + self.mamba_state_free_slot_count += 1 + def batch_allocate_slots(self, num_slots: int) -> Optional[torch.Tensor]: """ Allocates new slots for the given number of requests in the Mamba state buffers. @@ -769,7 +785,19 @@ def batch_allocate_slots(self, num_slots: int) -> Optional[torch.Tensor]: self.mamba_state_free_slot_count : self.mamba_state_free_slot_count + num_slots ] - return mamba_idx + return mamba_idx.clone() + + def _return_slots(self, mamba_indices: torch.Tensor) -> None: + """Return live state slots to the free-slot stack.""" + + if mamba_indices.numel() == 0: + return + start = self.mamba_state_free_slot_count + end = start + mamba_indices.numel() + if end > self.max_requests: + raise RuntimeError("Mamba state free-slot pool overflow") + self.mamba_state_free_slots[start:end] = mamba_indices.to(torch.int32) + self.mamba_state_free_slot_count = end def free_slots(self, request_indices: torch.Tensor) -> None: """ @@ -783,14 +811,7 @@ def free_slots(self, request_indices: torch.Tensor) -> None: # Filter out any invalid indices (e.g., -1) mamba_indices_to_free = mamba_indices_to_free[mamba_indices_to_free != -1] - num_to_free = len(mamba_indices_to_free) - - if num_to_free > 0: - # Add the freed indices back to the free slot pool - start_idx = self.mamba_state_free_slot_count - end_idx = start_idx + num_to_free - self.mamba_state_free_slots[start_idx:end_idx] = mamba_indices_to_free - self.mamba_state_free_slot_count = end_idx + self._return_slots(mamba_indices_to_free) # Invalidate the Mamba state index for the finished requests self.request_to_mamba_state_idx[request_indices] = -1 diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 4a5a2836e5e..d06176aef96 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -901,17 +901,14 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC and prefix_caching_mamba_gb is not None and prefix_caching_mamba_gb > 0 ): + assert self.mamba_slot_allocator is not None prefix_cache_bytes = int(prefix_caching_mamba_gb * 1024**3) - # Mirror the split done in _allocate_mamba_cache so this preview - # matches what is actually allocated: the "scratch" buffers - # (intermediate_ssm_out/intermediate_conv_out) are reserved from the - # budget first, then the rest sizes the "durable" cache - # (ssm_states/conv_states). mamba_bytes_per_req is the shared - # per-slot footprint of both. + # The allocator contains the PP-synchronized durable capacity. + # Scratch remains stage-local because it is temporary forward + # storage and does not participate in prefix-cache matching. scratch_slots = self.max_mamba_intermediate_states_per_step scratch_bytes = scratch_slots * mamba_bytes_per_req - durable_slots = (prefix_cache_bytes - scratch_bytes) // mamba_bytes_per_req - durable_slots = max(durable_slots, 0) + durable_slots = self.mamba_slot_allocator.max_slots log_lines += [ f" Mamba prefix cache:", f" budget: {get_mem_size_str(prefix_cache_bytes)}", @@ -1826,7 +1823,33 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: scratch_slots = self.max_mamba_intermediate_states_per_step scratch_bytes = scratch_slots * per_slot_bytes max_slots = (total_bytes - scratch_bytes) // per_slot_bytes # durable slots + local_max_slots = max_slots + + # Prefix-cache state is replicated across pipeline stages. Use the + # smallest stage-local capacity so identical cache operations produce + # identical block-to-slot mappings on every stage. A recurrent prefix + # is executable only when all stages retain its state, so additional + # slots on an individual stage would not increase usable cache capacity. + if get_pg_size(self.pipeline_parallel_group) > 1: + max_slots_tensor = torch.tensor( + max_slots, dtype=torch.int64, device=torch.cuda.current_device() + ) + torch.distributed.all_reduce( + max_slots_tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.pipeline_parallel_group, + ) + max_slots = int(max_slots_tensor.item()) + if max_slots < 1: + if local_max_slots >= 1: + raise ValueError( + f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " + f"has room for {local_max_slots} durable slots on this pipeline stage, but " + f"another stage has room for fewer than one. Cache capacity is mirrored to " + f"the minimum across stages; increase prefix_caching_mamba_gb, reduce " + f"max_tokens, or rebalance the Mamba layers." + ) raise ValueError( f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " f"is too small. The CUDA-graph extraction scratch reserves " @@ -3074,10 +3097,14 @@ def check_availability(self, req: DynamicInferenceRequest) -> Tuple[bool, bool, Check if the request can be added to the context. """ # Note that for hybrid models checking the total request count is sufficient - # because we allocate a single set of Mamba state tensors for each request + # because we allocate a single set of Mamba state tensors for each request. request_can_be_added = ( self.total_request_count < self.max_requests and self.paused_request_count == 0 ) + # A handoff is the exception: it keeps its live recurrent-state slot after + # its request row is removed, so both capacities must be checked independently. + if self.is_hybrid_model and self.kv_block_allocator.enable_handoff_pinning: + request_can_be_added &= self.mamba_metadata.mamba_state_free_slot_count > 0 matched_block_ids, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length = ( self._compute_prefix_match(req, req.remaining_prompt_length) diff --git a/megatron/core/inference/disaggregation/decode_admission.py b/megatron/core/inference/disaggregation/decode_admission.py index 9ec5c9f9474..1a436677811 100644 --- a/megatron/core/inference/disaggregation/decode_admission.py +++ b/megatron/core/inference/disaggregation/decode_admission.py @@ -35,12 +35,14 @@ def admit_prefilled_decode( prompt_block_ids: list[int], continuation_block_ids: list[int], input_tokens: list[int], + ssm_state_idx: int | None = None, ) -> None: - """Admit imported KV state directly as an active decode request. + """Admit imported KV/SSM state directly as an active decode request. The imported prompt blocks already carry one allocator reference. This - function transfers that ownership to the live request and schedules the - sampled first token (plus any MTP proposals) as the next decode query. + function transfers that ownership to the live request, binds the exact + post-prompt recurrent state, and schedules the sampled first token (plus + any MTP proposals) as the next decode query. """ input_token_count = len(input_tokens) @@ -71,6 +73,10 @@ def admit_prefilled_decode( if request.get_metadata_types() != context.request_metadata_types: raise ValueError("Imported request metadata does not match the decode context") + if context.is_hybrid_model: + if ssm_state_idx is None: + raise ValueError("A hybrid decode request requires an exact imported SSM state") + current_id = context.total_request_count all_block_ids = prompt_block_ids + continuation_block_ids row = context.request_to_kv_block_ids[current_id] @@ -129,6 +135,10 @@ def admit_prefilled_decode( token_positions // context.block_size_tokens ] + if context.is_hybrid_model: + assert ssm_state_idx is not None + context.mamba_metadata.request_to_mamba_state_idx[current_id] = ssm_state_idx + context.active_token_count = token_end context.total_request_count += 1 diff --git a/megatron/core/inference/disaggregation/engine.py b/megatron/core/inference/disaggregation/engine.py index 411d163e83b..ec65e219f62 100644 --- a/megatron/core/inference/disaggregation/engine.py +++ b/megatron/core/inference/disaggregation/engine.py @@ -1,6 +1,6 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -"""Dynamic inference engine composed with KV-cache handoff behavior.""" +"""Dynamic inference engine composed with the KV/state hand-off behavior.""" from megatron.core.inference.disaggregation.inference_state_handoff import ( InferenceStateHandoffMixin, @@ -9,7 +9,7 @@ class DisaggDynamicInferenceEngine(InferenceStateHandoffMixin, DynamicInferenceEngine): - """Dynamic inference engine with prefill/decode KV-cache handoff support. + """Dynamic inference engine with prefill/decode state hand-off support. Used by both control planes: the Dynamo integration and the coordinator-native 2-hop mode. diff --git a/megatron/core/inference/disaggregation/inference_state_handoff.py b/megatron/core/inference/disaggregation/inference_state_handoff.py index fa989aad6f6..500f13fec82 100644 --- a/megatron/core/inference/disaggregation/inference_state_handoff.py +++ b/megatron/core/inference/disaggregation/inference_state_handoff.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Engine-side lifecycle for disaggregated prefill/decode KV-cache handoff.""" +"""Engine-side lifecycle for disaggregated prefill/decode state handoff.""" from __future__ import annotations @@ -23,7 +23,9 @@ from megatron.core.inference.disaggregation.pending_handoff_imports import ( DeferredKvHandoff, PendingKvImport, + PendingSSMImport, ) +from megatron.core.inference.disaggregation.ssm_reshard import SSMShardLayout, SSMStateDims from megatron.core.inference.disaggregation.transfer_backends.base import ( construct_kv_transfer_backend_class, ) @@ -43,26 +45,33 @@ from megatron.core.inference.inference_request import DynamicInferenceRequest from megatron.core.inference.sampling_params import SamplingParams +_SSM_STATE_KINDS = ("conv", "recurrent") + @dataclass(frozen=True) class _PreparedHandoffMetadata: """Per-request metadata assembled once for a completed prefill batch.""" local_blocks: list[int] # Full prompt block table, including a partial tail. + local_ssm_slot: int | None kv_meta: Any + ssm_meta: Any resume_tokens: list[int] class InferenceStateHandoffMixin: - """Optional KV-cache handoff behavior composed into the dynamic engine.""" + """Optional KV/SSM handoff behavior composed into the dynamic engine.""" def _initialize_disaggregation_state(self) -> None: """Initialize state without importing or constructing a transfer backend.""" self._pinned_handoff_blocks: Dict[int, list] = {} # Request ID -> pinned KV block IDs. + self._pinned_handoff_ssm_slots: Dict[int, int] = {} # Request ID -> detached live slot. self._kv_transfer_agent = None self._kv_peer_metas = None # KV descriptors for this PP stage's TP ranks. self._pp_kv_peer_metas = None # Each PP stage's set of TP KV descriptors. + self._ssm_transfer_agents = {} + self._pp_ssm_peer_metas = None # Each PP stage's TP SSM descriptors. self._deferred_kv_handoffs = deque() self._pending_kv_imports = deque() self._quarantined_kv_imports: list[PendingKvImport] = [] @@ -173,9 +182,9 @@ def _reset_pending_kv_imports(self) -> None: raise RuntimeError( "Cannot reset while KV handoff transfers may still access cache storage" ) - if self._pinned_handoff_blocks: + if self._pinned_handoff_blocks or self._pinned_handoff_ssm_slots: raise RuntimeError( - "Cannot reset while KV handoff blocks remain pinned; wait for RELEASE_KV" + "Cannot reset while handoff state remains pinned; wait for RELEASE_KV" ) self._handoff_completion_notifications.clear() @@ -209,9 +218,11 @@ def setup_kv_transfer(self, role: str, backend: str = "nixl") -> None: if role not in ("prefill", "decode"): raise ValueError(f"KV transfer role must be 'prefill' or 'decode', got {role!r}") if self.context.is_hybrid_model: - raise RuntimeError( - "Hybrid models require recurrent-state handoff in addition to KV-cache handoff" - ) + if role == "decode" and self.context.mamba_slot_allocator is not None: + raise RuntimeError( + "A hybrid decode handoff writes directly into live SSM state; " + "do not configure prefix_caching_mamba_gb on the decode engine" + ) self._kv_transfer_role = role backend_cls = construct_kv_transfer_backend_class(backend) @@ -235,21 +246,24 @@ def setup_kv_transfer(self, role: str, backend: str = "nixl") -> None: tp_size = get_pg_size(self.pg_collection.tp) tp_rank = get_pg_rank(self.pg_collection.tp) - # Compute this PP rank's global attention-layer range. + # Compute this PP rank's global attention- and SSM-layer ranges. pp_size = get_pg_size(self.pg_collection.pp) pp_rank = get_pg_rank(self.pg_collection.pp) local_num_layers = self.context.num_attention_layers + local_num_ssm_layers = self.context.num_mamba_layers if pp_size > 1 and torch.distributed.is_initialized(): layer_counts: list = [None] * pp_size torch.distributed.all_gather_object( - layer_counts, local_num_layers, group=self.pg_collection.pp + layer_counts, (local_num_layers, local_num_ssm_layers), group=self.pg_collection.pp ) - layer_start = sum(layer_counts[:pp_rank]) - num_layers_global = sum(layer_counts) + layer_start = sum(counts[0] for counts in layer_counts[:pp_rank]) + num_layers_global = sum(counts[0] for counts in layer_counts) + ssm_layer_start = sum(counts[1] for counts in layer_counts[:pp_rank]) else: layer_start = 0 num_layers_global = local_num_layers + ssm_layer_start = 0 layer_end = layer_start + local_num_layers self._kv_transfer_agent = backend_cls( @@ -270,25 +284,103 @@ def setup_kv_transfer(self, role: str, backend: str = "nixl") -> None: layer_end=layer_end, ) + # Both roles transfer directly between live recurrent-state slots. Prefill + # detaches its source slot until decode acknowledges transfer completion. + if self.context.is_hybrid_model: + conv_states = self.context.mamba_conv_states + recurrent_states = self.context.mamba_ssm_states + state_slot_count = self.context.max_requests + + conv_shape = conv_states.shape + ssm_shape = recurrent_states.shape + nheads_local = int(ssm_shape[-3]) + nheads_global = nheads_local * tp_size + configured_nheads = model_config.mamba_num_heads + if configured_nheads is not None and configured_nheads != nheads_global: + raise ValueError( + f"SSM state shape implies {nheads_global} global heads, " + f"but model config specifies {configured_nheads}" + ) + ssm_dims = SSMStateDims( + nheads=nheads_global, + headdim=int(ssm_shape[-2]), + d_state=int(ssm_shape[-1]), + ngroups=int(model_config.mamba_num_groups), + d_conv=int(conv_shape[-1]), + ) + ssm_layout = SSMShardLayout( + global_rank=rank, + tp_size=tp_size, + tp_rank=tp_rank, + layer_start=ssm_layer_start, + num_layers=local_num_ssm_layers, + dims=ssm_dims, + ) + if int(conv_shape[-2]) != ssm_layout.conv_dim_local: + raise ValueError( + "SSM conv state shape does not match the model TP layout: " + f"{conv_shape[-2]} vs {ssm_layout.conv_dim_local}" + ) + state_specs = { + "conv": (conv_states, conv_states.shape[-2], conv_states.shape[-1]), + "recurrent": ( + recurrent_states, + recurrent_states.shape[-3], + recurrent_states.shape[-2] * recurrent_states.shape[-1], + ), + } + for state_kind, (memory_buffer, width, state_dim) in state_specs.items(): + self._ssm_transfer_agents[state_kind] = ( + self._kv_transfer_agent.new_registered_buffer( + agent_name=f"{role}-ssm-{state_kind}-rank{rank}", + memory_buffer=memory_buffer, + expected_num_blocks=state_slot_count, + heads_per_partition=width, + head_dim=state_dim, + tokens_per_block=1, + ssm_layout=ssm_layout, + ssm_state_kind=state_kind, + ) + ) + ssm_peer_metas = { + state_kind: agent.export_meta() + for state_kind, agent in self._ssm_transfer_agents.items() + } + else: + ssm_peer_metas = {} + # Transport descriptors do not change between requests. Collect each # stage's TP descriptors once and reuse them for every handoff. - self._kv_peer_metas = self._kv_transfer_agent.export_meta() + local_peer_metas = {"kv": self._kv_transfer_agent.export_meta(), "ssm": ssm_peer_metas} if torch.distributed.is_initialized() and tp_size > 1: gathered_tp_metas: list = [None] * tp_size torch.distributed.all_gather_object( - gathered_tp_metas, self._kv_peer_metas, group=self.pg_collection.tp + gathered_tp_metas, local_peer_metas, group=self.pg_collection.tp ) - self._kv_peer_metas = gathered_tp_metas + ssm_state_kinds = [set(entry["ssm"]) for entry in gathered_tp_metas] + if any(kinds != ssm_state_kinds[0] for kinds in ssm_state_kinds[1:]): + raise RuntimeError("SSM handoff agents are not configured on every TP rank") + self._kv_peer_metas = [entry["kv"] for entry in gathered_tp_metas] + tp_ssm_peer_metas = { + state_kind: [entry["ssm"][state_kind] for entry in gathered_tp_metas] + for state_kind in ssm_state_kinds[0] + } + else: + self._kv_peer_metas = local_peer_metas["kv"] + tp_ssm_peer_metas = dict(local_peer_metas["ssm"]) - self._pp_kv_peer_metas = [self._kv_peer_metas] + local_stage_metas = {"kv": self._kv_peer_metas, "ssm": tp_ssm_peer_metas} + gathered_stage_metas = [local_stage_metas] if torch.distributed.is_initialized() and pp_size > 1: - self._pp_kv_peer_metas = [None] * pp_size + gathered_stage_metas = [None] * pp_size torch.distributed.all_gather_object( - self._pp_kv_peer_metas, self._kv_peer_metas, group=self.pg_collection.pp + gathered_stage_metas, local_stage_metas, group=self.pg_collection.pp ) + self._pp_kv_peer_metas = [entry["kv"] for entry in gathered_stage_metas] + self._pp_ssm_peer_metas = [entry["ssm"] for entry in gathered_stage_metas] def push_handoff_kv(self, request_id: int, decode_metas: list) -> None: - """Push a pinned hand-off's KV to the decode + """Push a pinned hand-off's KV and live SSM state to the decode instance described by `decode_metas` (two-sided transports only). The decode posted its matching receives when SUBMIT_REQUEST_WITH_KV @@ -299,14 +391,42 @@ def push_handoff_kv(self, request_id: int, decode_metas: list) -> None: """ block_ids = self._pinned_handoff_blocks.get(request_id) if not block_ids: + if self._ssm_transfer_agents and request_id in self._pinned_handoff_ssm_slots: + raise RuntimeError( + f"Handoff request {request_id} has detached SSM state but no pinned KV blocks" + ) logging.warning( "SEND_KV for request %d with no pinned hand-off blocks; skipping", request_id ) return + + ssm_pushes = [] + if self._ssm_transfer_agents: + slot = self._pinned_handoff_ssm_slots.get(request_id) + if slot is None: + raise RuntimeError(f"No detached SSM state for handoff request {request_id}") + for state_kind in _SSM_STATE_KINDS: + try: + peer_metas = [entry["ssm"][state_kind] for entry in decode_metas] + except (KeyError, TypeError) as error: + raise RuntimeError( + f"Decode metadata is missing {state_kind} SSM transfer state" + ) from error + ssm_pushes.append((self._ssm_transfer_agents[state_kind], peer_metas, [slot])) + kv_peer = {"tp_metas": list(decode_metas)} handles = [self._kv_transfer_agent.begin_push_blocks(kv_peer, block_ids)] + # The source live slot remains detached until RELEASE_KV, so request + # cleanup cannot reuse it while these sends are active. + for agent, peer_metas, slots in ssm_pushes: + handles.append(agent.begin_push_blocks({"tp_metas": peer_metas}, slots)) self._pending_kv_pushes.append((request_id, handles)) - logging.info("DISAGG_PREFILL_PUSH request_id=%d blocks=%d", request_id, len(block_ids)) + logging.info( + "DISAGG_PREFILL_PUSH request_id=%d blocks=%d ssm=%d", + request_id, + len(block_ids), + len(handles) - 1, + ) def _poll_pending_kv_pushes(self) -> int: """Reap completed push sends; unfinished ones stay pending. @@ -327,7 +447,7 @@ def _poll_pending_kv_pushes(self) -> int: def _prepare_handoff_metadata_batch( self, - requests_and_blocks: list[tuple["DynamicInferenceRequest", list[int]]], + requests_and_state: list[tuple["DynamicInferenceRequest", list[int], int | None]], decode_tokens_by_request: Dict[int, list[int]], ) -> dict[int, _PreparedHandoffMetadata]: """Assemble metadata for all handoffs completed by one engine step. @@ -335,35 +455,37 @@ def _prepare_handoff_metadata_batch( Side: prefill engine; pull and push transport paths. """ - handoffs = [ - (request, list(block_ids)) - for request, block_ids in requests_and_blocks - if request.sampling_params.do_kv_handoff - ] - if not handoffs: - return {} - if self._kv_peer_metas is None: - raise RuntimeError("KV handoff requested before transfer setup") - - candidate_blocks_by_request = {} - for request, block_ids in handoffs: + handoffs = [] + for request, block_ids, ssm_slot in requests_and_state: + if not request.sampling_params.do_kv_handoff: + continue prompt_block_count = ( len(request.prompt_tokens) + self.context.block_size_tokens - 1 ) // self.context.block_size_tokens - candidate_blocks_by_request[request.request_id] = block_ids[:prompt_block_count] + handoffs.append((request, list(block_ids[:prompt_block_count]), ssm_slot)) + if not handoffs: + return {} - pp_size = get_pg_size(self.pg_collection.pp) try: + if self._kv_peer_metas is None: + raise RuntimeError("KV handoff requested before transfer setup") + + pp_size = get_pg_size(self.pg_collection.pp) static_pp_metas = self._pp_kv_peer_metas or [self._kv_peer_metas] + static_pp_ssm_metas = self._pp_ssm_peer_metas or [{}] if len(static_pp_metas) != pp_size: raise RuntimeError( f"Expected static metadata for {pp_size} pipeline stages, " f"got {len(static_pp_metas)}" ) + if len(static_pp_ssm_metas) != pp_size: + raise RuntimeError( + f"Expected static SSM metadata for {pp_size} pipeline stages, " + f"got {len(static_pp_ssm_metas)}" + ) prepared = {} - for request, _ in handoffs: - candidate_blocks = candidate_blocks_by_request[request.request_id] + for request, candidate_blocks, ssm_slot in handoffs: resume_tokens = list(decode_tokens_by_request.get(request.request_id, [])) expected_resume_tokens = self.context.num_speculative_tokens + 1 if len(resume_tokens) != expected_resume_tokens: @@ -376,10 +498,15 @@ def _prepare_handoff_metadata_batch( raise RuntimeError( "Cannot create a decode-only handoff without the prompt KV state" ) + if self._ssm_transfer_agents and ssm_slot is None: + raise RuntimeError( + "Cannot create a hybrid decode-only handoff without the exact final " + "SSM state" + ) block_ids = candidate_blocks if pp_size > 1: # Model-parallel ranks receive the same request stream and - # use synchronized KV cache capacities. + # use synchronized KV and recurrent-state cache capacities. # Their deterministic allocators therefore assign the same # physical IDs, while transport descriptors remain stage- and # rank-specific metadata collected once during setup. @@ -392,15 +519,59 @@ def _prepare_handoff_metadata_batch( else: kv_meta = self._kv_peer_metas + ssm_meta = None + if self._ssm_transfer_agents: + stage_state_metas = [ + self._ssm_stage_transfer_meta(static_meta, ssm_slot) + for static_meta in static_pp_ssm_metas + ] + ssm_meta = {} + for state_kind in _SSM_STATE_KINDS: + if pp_size > 1: + ssm_meta[state_kind] = { + "pp_metas": [ + {"tp_metas": stage_meta[state_kind]} + for stage_meta in stage_state_metas + ] + } + else: + ssm_meta[state_kind] = stage_state_metas[0][state_kind] + prepared[request.request_id] = _PreparedHandoffMetadata( - local_blocks=block_ids, kv_meta=kv_meta, resume_tokens=resume_tokens + local_blocks=block_ids, + local_ssm_slot=ssm_slot, + kv_meta=kv_meta, + ssm_meta=ssm_meta, + resume_tokens=resume_tokens, ) return prepared except Exception: - for block_ids in candidate_blocks_by_request.values(): + # The controller temporarily transfers ownership for every finished + # request on a prefill engine, including regular requests. Return all + # of it if metadata preparation aborts before the per-request loop. + for _, block_ids, ssm_slot in requests_and_state: self._release_pinned_handoff_blocks(block_ids) + self._release_pinned_handoff_ssm_slot(ssm_slot) raise + @staticmethod + def _ssm_stage_transfer_meta(static_metas: dict, selected_slot: int) -> dict: + """Attach the final recurrent-state slot to one stage's descriptors.""" + + state_metas = {} + selected_slots = [selected_slot] + for state_kind in _SSM_STATE_KINDS: + state_static_metas = static_metas.get(state_kind) + if state_static_metas is None: + raise RuntimeError(f"Missing static {state_kind} SSM handoff metadata") + if isinstance(state_static_metas, list): + state_metas[state_kind] = [ + {**meta, "block_ids": selected_slots} for meta in state_static_metas + ] + else: + state_metas[state_kind] = {**state_static_metas, "block_ids": selected_slots} + return state_metas + def _capture_handoff_meta( self, request: "DynamicInferenceRequest", prepared: _PreparedHandoffMetadata | None ) -> None: @@ -417,23 +588,33 @@ def _capture_handoff_meta( rid = request.request_id block_ids = prepared.local_blocks kv_meta = prepared.kv_meta + ssm_meta = prepared.ssm_meta self._pinned_handoff_blocks[rid] = list(block_ids) + if prepared.local_ssm_slot is not None: + self._pinned_handoff_ssm_slots[rid] = prepared.local_ssm_slot if isinstance(kv_meta, list): kv_meta = {"tp_metas": kv_meta} else: # TP=1 caches one static metadata dictionary for the engine. Keep - # request-specific fields out of that shared object. + # request-specific SSM metadata out of that shared object. kv_meta = dict(kv_meta) kv_meta["resume_tokens"] = prepared.resume_tokens + if ssm_meta is not None: + kv_meta["ssm"] = ssm_meta request.disaggregated_params = { "request_id": rid, "block_ids": block_ids, "kv_meta": kv_meta, } - logging.info("DISAGG_PREFILL_HANDOFF request_id=%d pinned_blocks=%d", rid, len(block_ids)) + logging.info( + "DISAGG_PREFILL_HANDOFF request_id=%d pinned_blocks=%d ssm_blocks=%d", + rid, + len(block_ids), + int(ssm_meta is not None), + ) def release_handoff_blocks(self, request_id: int) -> None: """Release blocks pinned by a previous do_kv_handoff completion. @@ -441,11 +622,16 @@ def release_handoff_blocks(self, request_id: int) -> None: Side: prefill engine; pull and push transport paths. """ block_ids = self._pinned_handoff_blocks.pop(request_id, None) - if not block_ids: + ssm_slot = self._pinned_handoff_ssm_slots.pop(request_id, None) + if not block_ids and ssm_slot is None: return - released = self._release_pinned_handoff_blocks(block_ids) + released = self._release_pinned_handoff_blocks(block_ids or []) + self._release_pinned_handoff_ssm_slot(ssm_slot) logging.info( - "DISAGG_PREFILL_RELEASE request_id=%d released_blocks=%d", request_id, released + "DISAGG_PREFILL_RELEASE request_id=%d released_blocks=%d ssm=%d", + request_id, + released, + int(ssm_slot is not None), ) def _release_pinned_handoff_blocks(self, block_ids: list) -> int: @@ -459,6 +645,12 @@ def _release_pinned_handoff_blocks(self, block_ids: list) -> int: allocator.release_memory_blocks(torch.tensor(block_ids, dtype=torch.int32, device="cpu")) return len(block_ids) + def _release_pinned_handoff_ssm_slot(self, ssm_slot: int | None) -> None: + """Release a prefill live-state slot after its handoff ownership ends.""" + + if ssm_slot is not None: + self.context.mamba_metadata.free_slot(ssm_slot) + def add_request_with_kv_handoff( self, request_id: int, @@ -480,6 +672,25 @@ def add_request_with_kv_handoff( "the prefill-skip path uses the prefix-cache match logic." ) + ssm_meta = kv_meta.get("ssm") if isinstance(kv_meta, dict) else None + local_has_ssm = bool(self._ssm_transfer_agents) + if self.context.is_hybrid_model and not local_has_ssm: + raise RuntimeError("Hybrid decode received a handoff before SSM transfer setup") + if local_has_ssm and ssm_meta is None: + raise RuntimeError( + "Decode has SSM state transfer agents but the handoff contains no " + "SSM metadata; prefill and decode must use the same hybrid model" + ) + if local_has_ssm: + if not isinstance(ssm_meta, dict): + raise RuntimeError("SSM handoff metadata must be a mapping") + missing_state_kinds = [kind for kind in _SSM_STATE_KINDS if kind not in ssm_meta] + if missing_state_kinds: + raise RuntimeError( + f"SSM handoff metadata is missing state kinds {missing_state_kinds}" + ) + if not self.context.is_hybrid_model and ssm_meta is not None: + raise RuntimeError("Transformer decode received SSM metadata from a hybrid prefill") if self._kv_transfer_agent is None: raise RuntimeError("KV handoff received without a transfer backend") @@ -560,6 +771,7 @@ def _try_start_kv_handoff_import(self, handoff: DeferredKvHandoff) -> bool: if isinstance(handoff.kv_meta, dict) else [] ) + ssm_meta = handoff.kv_meta.get("ssm") if isinstance(handoff.kv_meta, dict) else None continuation_block_count = ( additional_decode_blocks( len(handoff.prompt), len(resume_tokens), self.context.block_size_tokens @@ -571,6 +783,8 @@ def _try_start_kv_handoff_import(self, handoff: DeferredKvHandoff) -> bool: num_blocks_to_import + continuation_block_count, cached_blocks ): return False + if ssm_meta and self.context.mamba_metadata.mamba_state_free_slot_count < 1: + return False allocator.retain_memory_blocks(cached_blocks) allocated_blocks_tensor = allocator.allocate_memory_blocks( @@ -586,8 +800,20 @@ def _try_start_kv_handoff_import(self, handoff: DeferredKvHandoff) -> bool: imported_blocks = allocated_blocks[:num_blocks_to_import] continuation_blocks = allocated_blocks[num_blocks_to_import:] local_blocks = cached_blocks + imported_blocks + owned_blocks_tensor = torch.tensor( + local_blocks + continuation_blocks, dtype=torch.int32, device="cpu" + ) + handle = None start_error = None + ssm_import = None + if ssm_meta: + try: + ssm_import = self._reserve_ssm_handoff_import() + except Exception: + allocator.release_memory_blocks(owned_blocks_tensor) + raise + try: transfer_meta, transfer_src_blocks = drop_transfer_prefix_blocks( handoff.kv_meta, handoff.src_block_ids, len(cached_blocks) @@ -595,6 +821,8 @@ def _try_start_kv_handoff_import(self, handoff: DeferredKvHandoff) -> bool: handle = self._kv_transfer_agent.begin_pull_blocks( transfer_meta, transfer_src_blocks, imported_blocks ) + if ssm_import is not None: + self._start_ssm_handoff_import(handoff.request_id, ssm_meta, ssm_import) except Exception as exc: start_error = exc @@ -607,6 +835,7 @@ def _try_start_kv_handoff_import(self, handoff: DeferredKvHandoff) -> bool: cached_prefix_block_count=len(cached_blocks), handle=handle, future=handoff.future, + ssm=ssm_import, resume_tokens=resume_tokens, continuation_blocks=continuation_blocks, local_error=start_error, @@ -682,12 +911,15 @@ def _drain_deferred_kv_handoffs(self) -> int: @staticmethod def _pending_transfer_handles(pending: PendingKvImport) -> list: - """Return this decode import's active KV transfer handles. + """Return this decode import's active KV and SSM transfer handles. Side: decode engine; pull and push transport paths. """ - return [pending.handle] if pending.handle is not None else [] + handles = [pending.handle] + if pending.ssm is not None: + handles.extend(pending.ssm.handles) + return [handle for handle in handles if handle is not None] def _validate_decode_ready_handoff(self, pending: PendingKvImport) -> None: """Validate that transferred state can start decode without prompt execution.""" @@ -698,6 +930,9 @@ def _validate_decode_ready_handoff(self, pending: PendingKvImport) -> None: "Decode-only handoff is missing its sampled token or MTP proposals: " f"expected {expected_tokens}, got {len(pending.resume_tokens)}" ) + if self.context.is_hybrid_model: + if pending.ssm is None: + raise RuntimeError("Hybrid decode-only handoff is missing transferred SSM state") def _finalize_kv_handoff_import(self, pending: PendingKvImport) -> None: """Register transferred blocks and admit the decode request. @@ -786,7 +1021,9 @@ def _finalize_kv_handoff_import(self, pending: PendingKvImport) -> None: local_blocks, pending.continuation_blocks, pending.resume_tokens, + ssm_state_idx=(pending.ssm.live_slot if pending.ssm is not None else None), ) + pending.ssm = None pending.local_blocks = [] pending.continuation_blocks = [] if self.use_coordinator and self.is_mp_coordinator: @@ -846,6 +1083,9 @@ def _release_pending_kv_import(self, pending: PendingKvImport) -> None: self.context.kv_block_allocator.release_memory_blocks(block_tensor) pending.local_blocks = [] pending.continuation_blocks = [] + if pending.ssm is not None: + self.context.mamba_metadata.free_slot(pending.ssm.live_slot) + pending.ssm = None @staticmethod def _wait_for_transfer_handles(*handles) -> bool: @@ -966,3 +1206,30 @@ def _admit_pending_kv_imports(self) -> int: self._loop.create_task, self._notify_cond_for_new_request() ) return ready + + def _reserve_ssm_handoff_import(self) -> PendingSSMImport: + """Reserve the live request slot that receives exact SSM state.""" + + if not self._ssm_transfer_agents: + raise RuntimeError( + "Received SSM handoff state but this decode engine has no " + "SSM transfer agents. Ensure KV transfer was initialized on a " + "hybrid decode engine." + ) + live_slot = self.context.mamba_metadata.allocate_slot() + if live_slot is None: + raise RuntimeError("Live SSM slot capacity changed during handoff admission") + return PendingSSMImport(handles=[], live_slot=int(live_slot)) + + def _start_ssm_handoff_import( + self, request_id: int, ssm_meta: dict, pending: PendingSSMImport + ) -> None: + """Post transfers into slots already reserved for one handoff.""" + + pending.handles.clear() + for state_kind in _SSM_STATE_KINDS: + agent = self._ssm_transfer_agents[state_kind] + pending.handles.append( + agent.begin_pull_blocks(ssm_meta[state_kind], [], [pending.live_slot]) + ) + logging.debug("DISAGG_DECODE_SSM_IMPORT_SUBMIT request_id=%d ssm_blocks=%d", request_id, 1) diff --git a/megatron/core/inference/disaggregation/pending_handoff_imports.py b/megatron/core/inference/disaggregation/pending_handoff_imports.py index d2fbe12663f..0ae4e214185 100644 --- a/megatron/core/inference/disaggregation/pending_handoff_imports.py +++ b/megatron/core/inference/disaggregation/pending_handoff_imports.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""State records for KV-cache imports awaiting completion.""" +"""State records for KV-cache and SSM-state imports awaiting completion.""" from __future__ import annotations @@ -25,6 +25,14 @@ class DeferredKvHandoff: future: asyncio.Future +@dataclass(kw_only=True) +class PendingSSMImport: + """Exact SSM state being transferred into a reserved live request slot.""" + + handles: list[Any] + live_slot: int + + @dataclass(kw_only=True) class PendingKvImport: """Decode request waiting for an asynchronous KV-cache import.""" @@ -37,6 +45,7 @@ class PendingKvImport: cached_prefix_block_count: int handle: Any future: asyncio.Future + ssm: PendingSSMImport | None = None resume_tokens: List[int] = field(default_factory=list) # Sampled token, then MTP proposals. continuation_blocks: List[int] = field(default_factory=list) # Empty KV for resume writes. local_error: Exception | None = None # Exact local error, if this rank failed. diff --git a/megatron/core/inference/disaggregation/transfer_backends/base.py b/megatron/core/inference/disaggregation/transfer_backends/base.py index fa8df492e44..a93d93a3a7f 100644 --- a/megatron/core/inference/disaggregation/transfer_backends/base.py +++ b/megatron/core/inference/disaggregation/transfer_backends/base.py @@ -109,20 +109,28 @@ def compute_buffer_geometry( ) shape = list(memory_buffer.shape) - candidates = [i for i, dim in enumerate(shape) if dim == expected_num_blocks] - if not candidates: - raise RuntimeError( - f"{backend_name}: no axis in memory_buffer shape {shape} matches " - f"expected_num_blocks={expected_num_blocks}. Layout is unrecognized; " - "bug in caller or new Megatron tensor shape." - ) - if len(candidates) > 1: - raise RuntimeError( - f"{backend_name}: ambiguous blocks axis in shape {shape} " - f"(expected_num_blocks={expected_num_blocks} matches multiple axes " - f"{candidates}). Caller must pass a more distinctive value." - ) - blocks_axis = candidates[0] + if ssm_layout is not None: + blocks_axis = 1 + if len(shape) <= blocks_axis or shape[blocks_axis] != expected_num_blocks: + raise RuntimeError( + f"{backend_name}: SSM state shape {shape} does not have " + f"expected_num_blocks={expected_num_blocks} on slot axis 1" + ) + else: + candidates = [i for i, dim in enumerate(shape) if dim == expected_num_blocks] + if not candidates: + raise RuntimeError( + f"{backend_name}: no axis in memory_buffer shape {shape} matches " + f"expected_num_blocks={expected_num_blocks}. Layout is unrecognized; " + "bug in caller or new Megatron tensor shape." + ) + if len(candidates) > 1: + raise RuntimeError( + f"{backend_name}: ambiguous blocks axis in shape {shape} " + f"(expected_num_blocks={expected_num_blocks} matches multiple axes " + f"{candidates}). Caller must pass a more distinctive value." + ) + blocks_axis = candidates[0] elements_per_slice = 1 for dim in shape[blocks_axis + 1 :]: diff --git a/megatron/core/inference/disaggregation/transfer_backends/nccl.py b/megatron/core/inference/disaggregation/transfer_backends/nccl.py index 96350fb54bf..1242b63b5dd 100644 --- a/megatron/core/inference/disaggregation/transfer_backends/nccl.py +++ b/megatron/core/inference/disaggregation/transfer_backends/nccl.py @@ -162,6 +162,11 @@ def __init__( list(memory_buffer.shape), ) + def new_registered_buffer(self, **kwargs) -> "NcclTransferBackend": + """Create a transfer backend for another state buffer.""" + + return type(self)(**kwargs) + def export_meta(self) -> Dict[str, Any]: """The shared geometry schema plus this rank's NCCL address.""" meta = export_geometry_meta(self._geometry, self._ssm_layout) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index d429a8141eb..61ade21cbb3 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -379,10 +379,10 @@ def _record_handoff_completion_notification(self, request_id: int, failed: bool) self._raise_kv_handoff_not_enabled("KV handoff completion notification") def _prepare_handoff_metadata_batch( - self, requests_and_blocks: list[tuple], decode_tokens_by_request: Dict[int, list[int]] + self, requests_and_state: list[tuple], decode_tokens_by_request: Dict[int, list[int]] ) -> dict: """Hook overridden by the KV-handoff engine composition.""" - if any(request.sampling_params.do_kv_handoff for request, _ in requests_and_blocks): + if any(request.sampling_params.do_kv_handoff for request, *_ in requests_and_state): self._raise_kv_handoff_not_enabled("KV handoff completion") return {} @@ -392,6 +392,9 @@ def _capture_handoff_meta(self, request, prepared) -> None: def _release_pinned_handoff_blocks(self, block_ids: list) -> int: return 0 + def _release_pinned_handoff_ssm_slot(self, ssm_slot: int | None) -> None: + return None + def setup_kv_transfer(self, role: str, backend: str = "nixl") -> None: """Raising stub; the hand-off engine composition overrides it.""" self._raise_kv_handoff_not_enabled("KV transfer setup") @@ -440,6 +443,7 @@ def reset(self) -> None: self.waiting_request_ids = deque() if hasattr(self, "_pinned_handoff_blocks"): self._pinned_handoff_blocks.clear() + self._pinned_handoff_ssm_slots.clear() self.failed_request_ids = [] # Generated token count already streamed for each request. self._partial_emit_lengths: Dict[int, int] = {} @@ -1370,6 +1374,7 @@ def post_process_requests( pre_fwd_step_count: Optional[int] = None, finished_routing_block_ids: Optional[Dict[int, list[int]]] = None, finished_handoff_block_ids: Optional[Dict[int, list[int]]] = None, + finished_handoff_ssm_slots: Optional[Dict[int, int]] = None, finished_handoff_decode_tokens: Optional[Dict[int, list[int]]] = None, ) -> Tuple[List[DynamicInferenceRequest], List[DynamicInferenceRequest]]: """ @@ -1394,6 +1399,7 @@ def post_process_requests( finished requests, saved before update_requests released them. Used for per-block routing reconstruction. finished_handoff_block_ids: Prompt KV block IDs retained for state handoff. + finished_handoff_ssm_slots: Live SSM slots detached for state handoff. finished_handoff_decode_tokens: First sampled token and optional MTP proposals needed to resume directly from imported prefill state on decode. @@ -1431,9 +1437,14 @@ def post_process_requests( # finished requests using the KV blocks retained before context cleanup. request_id_list = request_ids.tolist() handoff_blocks_by_request = finished_handoff_block_ids or {} + handoff_ssm_slots_by_request = finished_handoff_ssm_slots or {} prepared_handoff_metadata = self._prepare_handoff_metadata_batch( [ - (self.get_request(request_id), handoff_blocks_by_request.get(request_id, [])) + ( + self.get_request(request_id), + handoff_blocks_by_request.get(request_id, []), + handoff_ssm_slots_by_request.get(request_id), + ) for request_id in request_id_list if request_id in finished_request_ids ], @@ -1583,9 +1594,12 @@ def post_process_requests( request, prepared_handoff_metadata.get(request_id) ) # A prefill-role engine may also serve regular requests; release the - # temporary block references when no handoff was requested. - elif handoff_blocks: + # temporary state ownership when no handoff was requested. + elif handoff_blocks or request_id in handoff_ssm_slots_by_request: self._release_pinned_handoff_blocks(handoff_blocks) + self._release_pinned_handoff_ssm_slot( + handoff_ssm_slots_by_request.get(request_id) + ) finished_entry = self.requests.pop(request_id) finished_request = finished_entry.record[-1] finished_request.generated_length = len(finished_request.generated_tokens) @@ -2423,6 +2437,7 @@ async def async_bookkeep( top_n_logprobs = step_result.get("top_n_logprobs", None) finished_routing_block_ids = step_result.get("finished_routing_block_ids", None) finished_handoff_block_ids = step_result.get("finished_handoff_block_ids", None) + finished_handoff_ssm_slots = step_result.get("finished_handoff_ssm_slots", None) finished_handoff_decode_tokens = step_result.get("finished_handoff_decode_tokens", None) cuda_graph_request_count = step_result["cuda_graph_request_count"] @@ -2446,6 +2461,7 @@ async def async_bookkeep( pre_fwd_step_count=context_state.get("step_count"), finished_routing_block_ids=finished_routing_block_ids, finished_handoff_block_ids=finished_handoff_block_ids, + finished_handoff_ssm_slots=finished_handoff_ssm_slots, finished_handoff_decode_tokens=finished_handoff_decode_tokens, ) diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index 799c15b6d43..895f3973cc1 100644 --- a/megatron/core/inference/inference_request.py +++ b/megatron/core/inference/inference_request.py @@ -396,6 +396,7 @@ class DynamicInferenceRequest(InferenceRequest): # KV handoff metadata describing this request's pinned prefill state. # Used by decode-side pulls and prefill-side pushes. # Shape: {"request_id", "block_ids", "kv_meta"}. + # Hybrid models may add kv_meta["ssm"] for recurrent state. disaggregated_params: Optional[dict] = None def __post_init__(self): diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index ee6cbe14933..23d522fd6ca 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -176,6 +176,7 @@ class _AsyncScheduleRequestResult: newly_paused_request_ids: Optional[Tensor] = None evict_request_ids: Optional[Tensor] = None finished_handoff_block_ids: Dict[int, List[int]] = field(default_factory=dict) + finished_handoff_ssm_slots: Dict[int, int] = field(default_factory=dict) finished_handoff_decode_tokens: Dict[int, List[int]] = field(default_factory=dict) @@ -1906,15 +1907,16 @@ def _collect_finished_handoff_state( finished_idxs: Tensor, sampled_tokens_cpu: Tensor, sampled_mtp_tokens_cpu: Optional[Tensor], - ) -> Tuple[Dict[int, List[int]], Dict[int, List[int]]]: - """Retain completed prompt blocks and capture tokens needed to resume decode.""" + ) -> Tuple[Dict[int, List[int]], Dict[int, int], Dict[int, List[int]]]: + """Preserve completed prompt state and capture tokens needed to resume decode.""" context = self.inference_wrapped_model.inference_context allocator = context.kv_block_allocator finished_block_ids: Dict[int, List[int]] = {} + finished_ssm_slots: Dict[int, int] = {} decode_tokens_by_request: Dict[int, List[int]] = {} if not allocator.enable_handoff_pinning or finished_idxs.numel() == 0: - return finished_block_ids, decode_tokens_by_request + return finished_block_ids, finished_ssm_slots, decode_tokens_by_request for finished_idx in finished_idxs.tolist(): request_id = int(context.request_ids[finished_idx].item()) @@ -1925,6 +1927,13 @@ def _collect_finished_handoff_state( # Retain across context cleanup. For an exclusively owned block: # active=1, retain=2, request cleanup=1, coordinator RELEASE_KV=0. allocator.retain_memory_blocks(valid_blocks) + if context.is_hybrid_model: + # Transfer ownership out of the finished request before normal + # cleanup; the slot stays absent from the free-slot stack until + # the prefill engine receives RELEASE_KV. + finished_ssm_slots[request_id] = context.mamba_metadata.detach_state_slot( + finished_idx + ) active_idx = finished_idx - context.paused_request_count decode_tokens = [int(sampled_tokens_cpu[active_idx].item())] @@ -1934,7 +1943,7 @@ def _collect_finished_handoff_state( ) decode_tokens_by_request[request_id] = decode_tokens - return finished_block_ids, decode_tokens_by_request + return finished_block_ids, finished_ssm_slots, decode_tokens_by_request def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: """Update the dynamic inference context after sampling. @@ -2009,7 +2018,7 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: # Retain finished prefill blocks before request cleanup releases them; # the handoff path owns this reference until the decode transfer completes. - finished_handoff_block_ids, finished_handoff_decode_tokens = ( + finished_handoff_block_ids, finished_handoff_ssm_slots, finished_handoff_decode_tokens = ( self._collect_finished_handoff_state( finished_idxs, sampled_tokens_cpu, sampled_mtp_tokens_cpu ) @@ -2036,6 +2045,7 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: "sample": sampled_tokens_cpu, "finished_routing_block_ids": finished_routing_block_ids, "finished_handoff_block_ids": finished_handoff_block_ids, + "finished_handoff_ssm_slots": finished_handoff_ssm_slots, "finished_handoff_decode_tokens": finished_handoff_decode_tokens, **(update_result or {}), } @@ -2749,7 +2759,7 @@ def _run_async_sched_resolve( finished_idxs = ( torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count ) - finished_handoff_block_ids, finished_handoff_decode_tokens = ( + finished_handoff_block_ids, finished_handoff_ssm_slots, finished_handoff_decode_tokens = ( self._collect_finished_handoff_state( finished_idxs, sampled_tokens_cpu, sample_result.sampled_mtp_tokens_cpu_view ) @@ -2773,6 +2783,7 @@ def _run_async_sched_resolve( finished_request_ids=finished_request_ids, survivor_idxs=survivor_idxs, finished_handoff_block_ids=finished_handoff_block_ids, + finished_handoff_ssm_slots=finished_handoff_ssm_slots, finished_handoff_decode_tokens=finished_handoff_decode_tokens, ) @@ -2804,7 +2815,7 @@ def _run_async_sched_update_requests( finished_idxs = ( torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count ) - finished_handoff_block_ids, finished_handoff_decode_tokens = ( + finished_handoff_block_ids, finished_handoff_ssm_slots, finished_handoff_decode_tokens = ( self._collect_finished_handoff_state( finished_idxs, sampled_tokens_cpu, sample_result.sampled_mtp_tokens_cpu_view ) @@ -2832,6 +2843,7 @@ def _run_async_sched_update_requests( newly_paused_request_ids=update_result.get("newly_paused_request_ids"), evict_request_ids=update_result.get("evict_request_ids"), finished_handoff_block_ids=finished_handoff_block_ids, + finished_handoff_ssm_slots=finished_handoff_ssm_slots, finished_handoff_decode_tokens=finished_handoff_decode_tokens, ) @@ -2875,6 +2887,7 @@ def _build_async_sched_step_result( "sample": request_result.sampled_tokens_cpu, "finished_routing_block_ids": {}, "finished_handoff_block_ids": request_result.finished_handoff_block_ids, + "finished_handoff_ssm_slots": request_result.finished_handoff_ssm_slots, "finished_handoff_decode_tokens": request_result.finished_handoff_decode_tokens, "newly_paused_request_ids": request_result.newly_paused_request_ids, "evict_request_ids": request_result.evict_request_ids, diff --git a/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py b/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py index 43a7d57d8ba..4f4f56d45fb 100644 --- a/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py +++ b/tests/unit_tests/inference/contexts/attention_metadata/test_mamba_metadata.py @@ -40,6 +40,50 @@ def test_decode_indices_dtype(self, dtype): assert metadata._batch_indices_decode_buffer.dtype == dtype + def test_free_unbound_live_slot(self): + metadata = MambaMetadata(max_requests=2, max_tokens=4, max_intermediate_count=1) + + slot = int(metadata.allocate_slot()) + assert metadata.mamba_state_free_slot_count == 1 + + metadata.free_slot(slot) + + assert metadata.mamba_state_free_slot_count == 2 + assert int(metadata.allocate_slot()) == slot + + def test_allocated_slots_do_not_alias_free_slot_stack(self): + metadata = MambaMetadata(max_requests=3, max_tokens=4, max_intermediate_count=1) + + slot_to_release = metadata.allocate_slot() + allocated_slot = metadata.allocate_slot() + metadata.free_slot(slot_to_release) + + assert isinstance(allocated_slot, int) + assert allocated_slot == 1 + + metadata.reset() + slot_to_release = metadata.allocate_slot() + allocated_slots = metadata.batch_allocate_slots(2) + metadata.free_slot(slot_to_release) + + assert torch.equal(allocated_slots, torch.tensor([0, 1], dtype=torch.int32)) + + def test_detached_live_slot_survives_request_cleanup(self): + metadata = MambaMetadata(max_requests=2, max_tokens=4, max_intermediate_count=1) + slot = int(metadata.allocate_slot()) + metadata.request_to_mamba_state_idx[0] = slot + + assert metadata.detach_state_slot(0) == slot + metadata.free_slots(torch.tensor([0], dtype=torch.int64)) + + assert metadata.mamba_state_free_slot_count == 1 + assert int(metadata.allocate_slot()) != slot + + metadata.free_slot(slot) + + assert metadata.mamba_state_free_slot_count == 1 + assert int(metadata.allocate_slot()) == slot + def _run_update_test( self, metadata: MambaMetadata, diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 48abd7a05c8..b7789575014 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -1941,7 +1941,7 @@ def expected_log_probs(logits, active_id_and_counts): @rounder_override(64) def test_pipeline_parallel_uneven_layers(self): """ - Test that DynamicInferenceContext synchronizes the total block count across + Test that DynamicInferenceContext synchronizes cache capacities across pipeline stages when they have unequal layer counts. """ pp_size = 2 @@ -1986,29 +1986,60 @@ def test_pipeline_parallel_uneven_layers(self): block_size_tokens=16, max_tokens=1024, unified_memory_level=0, + enable_prefix_caching=True, + prefix_caching_mamba_gb=0.05, + mamba_inference_state_config=mamba_inference_state_config, ), ) - # Collect the total block counts on each rank (CUDA needed for NCCL all_gather) - local_total_blocks = torch.tensor( - [context.kv_block_allocator.pool_size], device='cuda', dtype=torch.long + # Collect cache capacities on each rank (CUDA needed for NCCL all_gather). + local_capacities = torch.tensor( + [context.kv_block_allocator.pool_size, context.mamba_slot_allocator.max_slots], + device='cuda', + dtype=torch.long, ) - gathered_block_counts = [torch.zeros_like(local_total_blocks) for _ in range(pp_size)] + gathered_capacities = [torch.zeros_like(local_capacities) for _ in range(pp_size)] torch.distributed.all_gather( - gathered_block_counts, - local_total_blocks, + gathered_capacities, + local_capacities, group=parallel_state.get_pipeline_model_parallel_group(), ) - all_counts = [t.item() for t in gathered_block_counts] + all_capacities = [tuple(t.tolist()) for t in gathered_capacities] - # Verify that there is only 1 unique value across all ranks - unique_counts = set(all_counts) + # Both allocators must remain mirrored across pipeline stages. + unique_capacities = set(all_capacities) assert ( - len(unique_counts) == 1 - ), f"Block counts were not synchronized across ranks. Gathered: {all_counts}" + len(unique_capacities) == 1 + ), f"Cache capacities were not synchronized across ranks. Gathered: {all_capacities}" self._restore_model_parallel() + @pytest.mark.internal + def test_mamba_cache_error_identifies_limiting_pipeline_stage(self): + context = object.__new__(DynamicInferenceContext) + context.mamba_conv_states_shape = (1,) + context.mamba_ssm_states_shape = (1,) + context.mamba_conv_states_dtype = torch.float32 + context.mamba_ssm_states_dtype = torch.float32 + context.num_mamba_layers = 1 + context.max_mamba_intermediate_states_per_step = 1 + context.pipeline_parallel_group = object() + + def reduce_to_remote_capacity(tensor, **_kwargs): + tensor.fill_(0) + + get_pg_size = "megatron.core.inference.contexts.dynamic_context.get_pg_size" + with ( + mock.patch(get_pg_size, return_value=2), + mock.patch.object( + torch.distributed, "all_reduce", side_effect=reduce_to_remote_capacity + ), + pytest.raises(ValueError, match="another stage has room for fewer than one") as error, + ): + context._allocate_mamba_cache(32 / 1024**3) + + assert "room for 3 durable slots on this pipeline stage" in str(error.value) + @pytest.mark.internal @rounder_override(64) def test_uneven_decoder_pp_layer_map_matches_get_num_layers_to_build(self): @@ -3330,7 +3361,7 @@ def test_chunked_prefill_prefix_caching_from_hidden_state(self): prefix_skip = 2 * bs - 1 eff_chunk = chunk_length - prefix_skip - (_, _, _, _, prefix_skip, eff_chunk) = ctx._compute_prefix_match(req2, chunk_length) + _, _, _, _, prefix_skip, eff_chunk = ctx._compute_prefix_match(req2, chunk_length) expected_active = tokens_before_chunk_2 + eff_chunk assert ctx.active_token_count == expected_active diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index a3a01f8ce4c..4ada5a4c570 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -14,6 +14,9 @@ MambaSlotAllocator, MambaSlotCapacityError, ) +from megatron.core.inference.disaggregation.inference_state_handoff import ( + InferenceStateHandoffMixin, +) from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine from megatron.core.inference.inference_request import ( DynamicInferenceRequest, @@ -814,7 +817,7 @@ def test_hybrid_memory_only(self): req2 = self._req(ctx, prompt.clone(), request_id=2) # no prefill skipping - (matched, _, _, _, prefix_skip, eff_chunk) = ctx._compute_prefix_match(req2, len(prompt)) + matched, _, _, _, prefix_skip, eff_chunk = ctx._compute_prefix_match(req2, len(prompt)) assert len(matched) == 3 and prefix_skip == 0 and eff_chunk == len(prompt) ctx.add_request(req2) @@ -958,6 +961,31 @@ def test_mamba_cache_budget_too_small_raises(self): with pytest.raises(ValueError, match="prefix cache budget"): self._mctx(prefix_caching_mamba_gb=1e-5) + @pytest.mark.internal + def test_hybrid_admission_resumes_after_handoff_releases_live_slot(self): + ctx = self._mctx(max_requests=1, rounder=1) + ctx.kv_block_allocator.enable_handoff_pinning = True + slot = ctx.mamba_metadata.allocate_slot() + assert slot is not None + ctx.mamba_metadata.request_to_mamba_state_idx[0] = slot + ctx.mamba_metadata.detach_state_slot(0) + request = self._req(ctx, self._prompt(ctx.block_size_tokens)) + + request_available, _, _ = ctx.check_availability(request) + + assert not request_available + + engine = InferenceStateHandoffMixin() + engine.context = ctx + engine._initialize_disaggregation_state() + engine._pinned_handoff_ssm_slots[7] = slot + engine.release_handoff_blocks(7) + + request_available, _, _ = ctx.check_availability(request) + + assert request_available + assert ctx.mamba_metadata.mamba_state_free_slot_count == 1 + @pytest.mark.internal def test_mamba_prefill_skip_and_zero_prefill(self): # mamba match limits prefill skip @@ -970,7 +998,7 @@ def test_mamba_prefill_skip_and_zero_prefill(self): self._mamba_allocate_and_register(ctx, self._block_ids(ctx, 0, 3)[:1]) req2 = self._req(ctx, prompt.clone(), request_id=2) req2._mamba_num_matched_blocks = 1 - (matched, _, _, _, prefix_skip, eff_chunk) = ctx._compute_prefix_match(req2, len(prompt)) + matched, _, _, _, prefix_skip, eff_chunk = ctx._compute_prefix_match(req2, len(prompt)) assert len(matched) == 3 and prefix_skip == bs and eff_chunk == len(prompt) - bs # no mamba match means no skip @@ -979,7 +1007,7 @@ def test_mamba_prefill_skip_and_zero_prefill(self): ctx2.add_request(self._req(ctx2, p2.clone())) req2b = self._req(ctx2, p2.clone(), request_id=2) req2b._mamba_num_matched_blocks = 0 - (m2, _, _, _, ps2, ec2) = ctx2._compute_prefix_match(req2b, len(p2)) + m2, _, _, _, ps2, ec2 = ctx2._compute_prefix_match(req2b, len(p2)) assert len(m2) == 3 and ps2 == 0 and ec2 == len(p2) # zero prefill for hybrid (mamba-cached, block-aligned) @@ -989,7 +1017,7 @@ def test_mamba_prefill_skip_and_zero_prefill(self): self._mamba_allocate_and_register(ctx3, self._block_ids(ctx3, 0, 3)) req3 = self._req(ctx3, p3.clone(), request_id=2) req3._mamba_num_matched_blocks = 3 - (m3, _, _, _, ps3, ec3) = ctx3._compute_prefix_match(req3, len(p3)) + m3, _, _, _, ps3, ec3 = ctx3._compute_prefix_match(req3, len(p3)) assert len(m3) == 3 and ps3 == 2 * bs and ec3 == bs # KV-only prefix skip with non-block-aligned prompt: all 3 full blocks @@ -1001,7 +1029,7 @@ def test_mamba_prefill_skip_and_zero_prefill(self): req4a = self._req(ctx4, p4.clone()) ctx4.add_request(req4a) req4b = self._req(ctx4, p4.clone(), request_id=2) - (m4, _, _, _, ps4, ec4) = ctx4._compute_prefix_match(req4b, len(p4)) + m4, _, _, _, ps4, ec4 = ctx4._compute_prefix_match(req4b, len(p4)) assert len(m4) == 3 and ps4 == 3 * bs4 and ec4 == tail ctx4.add_request(req4b) @@ -1070,7 +1098,7 @@ def test_mamba_intermediate_offsets(self): self._mamba_allocate_and_register(ctx, self._block_ids(ctx, 0, 4)[:2]) req2 = self._req(ctx, prompt.clone(), request_id=2) req2._mamba_num_matched_blocks = 2 - (matched, _, _, overall, prefix_skip, _) = ctx._compute_prefix_match(req2, len(prompt)) + matched, _, _, overall, prefix_skip, _ = ctx._compute_prefix_match(req2, len(prompt)) # Copy block IDs to slot 1 so compute_and_store_offsets can resolve EOS block ctx.request_to_kv_block_ids[1] = ctx.request_to_kv_block_ids[0] msa.compute_and_store_offsets( @@ -1923,7 +1951,7 @@ def test_prefill_computed_and_skipped_counters(self): # request 2 shares the first 4 blocks, adds 2 new blocks req2 = self._req(ctx, self._prompt(bs * 6), request_id=2) - (matched, _, _, _, prefix_skip, _) = ctx._compute_prefix_match(req2, bs * 6) + matched, _, _, _, prefix_skip, _ = ctx._compute_prefix_match(req2, bs * 6) assert len(matched) == 4 and prefix_skip == bs * 4 ctx.add_request(req2) diff --git a/tests/unit_tests/inference/test_disagg_handoff_lifecycle.py b/tests/unit_tests/inference/test_disagg_handoff_lifecycle.py index d6988ba49d6..f9255872fbf 100644 --- a/tests/unit_tests/inference/test_disagg_handoff_lifecycle.py +++ b/tests/unit_tests/inference/test_disagg_handoff_lifecycle.py @@ -26,6 +26,7 @@ from megatron.core.inference.disaggregation.pending_handoff_imports import ( DeferredKvHandoff, PendingKvImport, + PendingSSMImport, ) from megatron.core.inference.inference_request import ( DynamicInferenceRequest, @@ -92,6 +93,25 @@ def update_timestamps(self, block_ids): return None +class _MambaMetadata: + def __init__(self, available): + self.mamba_state_free_slot_count = available + self.next_slot = 20 + self.freed = [] + + def allocate_slot(self): + if self.mamba_state_free_slot_count == 0: + return None + slot = self.next_slot + self.next_slot += 1 + self.mamba_state_free_slot_count -= 1 + return slot + + def free_slot(self, slot): + self.freed.append(slot) + self.mamba_state_free_slot_count += 1 + + class _SchedulerHarness: def schedule_waiting_requests(self): if not self.waiting_request_ids: @@ -106,7 +126,7 @@ def schedule_waiting_requests(self): class _HandoffHarness(InferenceStateHandoffMixin, _SchedulerHarness): - def __init__(self, loop): + def __init__(self, loop, *, hybrid=False, available=0): self._loop = loop self._initialize_disaggregation_state() self.context = SimpleNamespace( @@ -118,11 +138,15 @@ def __init__(self, loop): max_requests=8, active_token_count=0, max_tokens=8, - is_hybrid_model=False, + is_hybrid_model=hybrid, kv_block_allocator=_KvAllocator(), + mamba_slot_allocator=None, + mamba_metadata=_MambaMetadata(available) if hybrid else None, memory_buffer=torch.empty(1), ) self._kv_transfer_agent = _TransferAgent() + if hybrid: + self._ssm_transfer_agents = {"conv": _TransferAgent(), "recurrent": _TransferAgent()} self.pg_collection = SimpleNamespace(tp=None, pp=None, mp=None) self.waiting_request_ids = deque() self.requests = {} @@ -160,7 +184,11 @@ def _check_stop_words_for_request_post_append(self, request): def _meta(request_id): - return {"request_id": request_id, "resume_tokens": [99]} + return { + "request_id": request_id, + "resume_tokens": [99], + "ssm": {"conv": {"request_id": request_id}, "recurrent": {"request_id": request_id}}, + } def _drain_loop(loop): @@ -199,7 +227,7 @@ def handoff_loop(): loop.close() -def test_prefilled_decode_admission_uses_imported_kv_without_prompt_tokens(): +def test_prefilled_decode_admission_uses_exact_ssm_state_without_prompt_tokens(): metadata_types = DynamicInferenceRequest.get_metadata_types() context = SimpleNamespace( block_size_tokens=4, @@ -210,6 +238,7 @@ def test_prefilled_decode_admission_uses_imported_kv_without_prompt_tokens(): max_requests=2, active_token_count=0, max_tokens=8, + is_hybrid_model=True, request_metadata_types=metadata_types, request_metadata={label: torch.empty(2, dtype=dtype) for label, dtype in metadata_types}, request_to_kv_block_ids=torch.full((2, 4), -1, dtype=torch.int32), @@ -228,6 +257,9 @@ def test_prefilled_decode_admission_uses_imported_kv_without_prompt_tokens(): token_to_local_position_within_kv_block=torch.zeros(8, dtype=torch.int32), token_to_block_idx=torch.full((8,), -1, dtype=torch.int32), ) + context.mamba_metadata = SimpleNamespace( + request_to_mamba_state_idx=torch.full((2,), -1, dtype=torch.int32) + ) request = DynamicInferenceRequest( request_id=7, prompt_tokens=torch.arange(3), @@ -239,7 +271,7 @@ def test_prefilled_decode_admission_uses_imported_kv_without_prompt_tokens(): can_admit_prefilled_decode(context, 3) context.chunked_prefill_request_id = -1 assert additional_decode_blocks(3, 3, 4) == 1 - admit_prefilled_decode(context, request, [10], [11], [40, 41, 42]) + admit_prefilled_decode(context, request, [10], [11], [40, 41, 42], 10) assert context.total_request_count == 1 assert context.num_prefill_requests == 0 @@ -250,11 +282,12 @@ def test_prefilled_decode_admission_uses_imported_kv_without_prompt_tokens(): assert context.token_to_input_ids[:3].tolist() == [40, 41, 42] assert context.token_to_pos_ids[:3].tolist() == [3, 4, 5] assert context.token_to_block_idx[:3].tolist() == [10, 11, 11] + assert context.mamba_metadata.request_to_mamba_state_idx[0].item() == 10 assert request.remaining_prompt_length == 0 -def test_completed_kv_handoff_enters_decode_without_waiting_queue(handoff_loop): - engine = _HandoffHarness(handoff_loop) +def test_completed_exact_ssm_handoff_enters_decode_without_waiting_queue(handoff_loop): + engine = _HandoffHarness(handoff_loop, hybrid=True) engine.context.num_speculative_tokens = 0 engine.use_coordinator = False engine.is_mp_coordinator = False @@ -283,6 +316,7 @@ def add_request(request_id, prompt, sampling_params, precomputed_block_hashes=No cached_prefix_block_count=0, handle=None, future=handoff_loop.create_future(), + ssm=PendingSSMImport(handles=[], live_slot=20), resume_tokens=[55], ) @@ -297,7 +331,8 @@ def add_request(request_id, prompt, sampling_params, precomputed_block_hashes=No assert list(engine.waiting_request_ids) == [] assert pending.local_blocks == [] assert pending.continuation_blocks == [] - admit.assert_called_once_with(engine.context, request, [10], [11], [55]) + admit.assert_called_once_with(engine.context, request, [10], [11], [55], ssm_state_idx=20) + assert pending.ssm is None def test_setup_pins_handoff_outputs_only_on_prefill(): @@ -321,6 +356,7 @@ def export_meta(self): num_attention_heads_per_partition=1, hidden_size_per_attention_head=1, block_size_tokens=4, + num_mamba_layers=0, ) model_config = SimpleNamespace(num_query_groups=1, num_attention_heads=1) engine.controller = SimpleNamespace( @@ -345,11 +381,132 @@ def export_meta(self): assert allocator.enable_handoff_pinning +def test_handoff_roles_use_live_ssm_buffers_and_decode_rejects_durable_cache(): + class _Backend: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.instances.append(self) + + def export_meta(self): + return {"transport": "test"} + + def new_registered_buffer(self, **kwargs): + return type(self)(**kwargs) + + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + engine.context = SimpleNamespace( + kv_block_allocator=SimpleNamespace( + enable_prefix_caching=True, enable_handoff_pinning=False, pool_size=8 + ), + memory_buffer=torch.empty(2, 1, 8, 4, 1, 1), + is_hybrid_model=True, + num_attention_layers=1, + num_mamba_layers=1, + num_attention_heads_per_partition=1, + hidden_size_per_attention_head=1, + block_size_tokens=4, + mamba_slot_allocator=object(), + ) + model_config = SimpleNamespace( + num_query_groups=1, num_attention_heads=1, mamba_num_heads=2, mamba_num_groups=1 + ) + engine.controller = SimpleNamespace( + inference_wrapped_model=SimpleNamespace(model=SimpleNamespace(config=model_config)) + ) + engine.pg_collection = SimpleNamespace(tp=None, pp=None, mp=None) + + backend_factory = ( + "megatron.core.inference.disaggregation.inference_state_handoff." + "construct_kv_transfer_backend_class" + ) + pg_size = "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_size" + pg_rank = "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_rank" + with ( + mock.patch(backend_factory, return_value=_Backend), + mock.patch(pg_size, return_value=1), + mock.patch(pg_rank, return_value=0), + ): + with pytest.raises(RuntimeError, match="directly into live SSM state"): + engine.setup_kv_transfer("decode") + + _Backend.instances.clear() + engine.context.mamba_slot_allocator = None + engine.context.max_requests = 8 + engine.context.mamba_conv_states = torch.empty(1, 8, 18, 3) + engine.context.mamba_ssm_states = torch.empty(1, 8, 2, 4, 5) + engine.setup_kv_transfer("decode") + + assert _Backend.instances[1].kwargs["memory_buffer"] is engine.context.mamba_conv_states + assert _Backend.instances[2].kwargs["memory_buffer"] is engine.context.mamba_ssm_states + assert [backend.kwargs["expected_num_blocks"] for backend in _Backend.instances[1:]] == [8, 8] + + _Backend.instances.clear() + engine.context.mamba_slot_allocator = object() + with ( + mock.patch(backend_factory, return_value=_Backend), + mock.patch(pg_size, return_value=1), + mock.patch(pg_rank, return_value=0), + ): + engine.setup_kv_transfer("prefill") + + assert _Backend.instances[1].kwargs["memory_buffer"] is engine.context.mamba_conv_states + assert _Backend.instances[2].kwargs["memory_buffer"] is engine.context.mamba_ssm_states + assert [backend.kwargs["expected_num_blocks"] for backend in _Backend.instances[1:]] == [8, 8] + + +def test_capacity_miss_defers_before_any_transfer(handoff_loop): + engine = _HandoffHarness(handoff_loop, hybrid=True, available=0) + future = engine.add_request_with_kv_handoff( + 7, [1, 2, 3, 4, 5], SamplingParams(num_tokens_to_generate=2), _meta(7), [100, 101] + ) + + assert not future.done() + assert engine.pending_kv_import_count == 1 + assert [item.request_id for item in engine._deferred_kv_handoffs] == [7] + assert not engine._pending_kv_imports + assert not engine._kv_transfer_agent.calls + assert not engine._ssm_transfer_agents["conv"].calls + assert engine.context.kv_block_allocator.releases == [] + + engine.context.mamba_metadata.mamba_state_free_slot_count = 1 + assert engine._poll_pending_kv_imports() == 0 + _drain_loop(handoff_loop) + + assert not engine._deferred_kv_handoffs + assert len(engine._pending_kv_imports) == 1 + assert len(engine._kv_transfer_agent.calls) == 1 + assert len(engine._ssm_transfer_agents["conv"].calls) == 1 + assert len(engine._ssm_transfer_agents["recurrent"].calls) == 1 + assert engine._pending_kv_imports[0].ssm.live_slot == 20 + assert engine.context.mamba_metadata.mamba_state_free_slot_count == 0 + + +def test_releasing_pending_hybrid_import_returns_live_slot(handoff_loop): + engine = _HandoffHarness(handoff_loop, hybrid=True, available=1) + engine.add_request_with_kv_handoff( + 7, [1, 2, 3, 4], SamplingParams(num_tokens_to_generate=2), _meta(7), [100] + ) + _drain_loop(handoff_loop) + + pending = engine._pending_kv_imports[0] + engine._release_pending_kv_import(pending) + + assert pending.ssm is None + assert engine.context.mamba_metadata.freed == [20] + assert engine.context.mamba_metadata.mamba_state_free_slot_count == 1 + + def test_reset_cancels_capacity_queued_handoffs(handoff_loop): engine = _HandoffHarness(handoff_loop) - engine.context.kv_block_allocator.capacity_available = False future = engine.add_request_with_kv_handoff( - 3, [3] * 4, SamplingParams(num_tokens_to_generate=2), _meta(3), [103] + 3, + [3] * 4, + SamplingParams(num_tokens_to_generate=2), + {"request_id": 3, "resume_tokens": [99]}, + [103], ) engine._reset_pending_kv_imports() @@ -385,12 +542,22 @@ def test_reset_rejects_pinned_prefill_blocks(handoff_loop): engine = _HandoffHarness(handoff_loop) engine._pinned_handoff_blocks[7] = [10] - with pytest.raises(RuntimeError, match="KV handoff blocks remain pinned"): + with pytest.raises(RuntimeError, match="handoff state remains pinned"): engine._reset_pending_kv_imports() assert engine._pinned_handoff_blocks == {7: [10]} +def test_release_handoff_returns_detached_prefill_ssm_slot(handoff_loop): + engine = _HandoffHarness(handoff_loop, hybrid=True) + engine._pinned_handoff_ssm_slots[7] = 3 + + engine.release_handoff_blocks(7) + + assert engine.context.mamba_metadata.freed == [3] + assert engine._pinned_handoff_ssm_slots == {} + + def test_reset_does_not_release_unsafe_import_destinations(handoff_loop): engine = _HandoffHarness(handoff_loop) block_id = int(engine.context.kv_block_allocator.allocate_memory_blocks(1)[0]) @@ -627,6 +794,7 @@ def test_handoff_metadata_batches_completed_requests_across_pipeline(monkeypatch ) engine._kv_peer_metas = {"global_rank": 0} engine._pp_kv_peer_metas = [{"global_rank": 0}, {"global_rank": 1}] + engine._pp_ssm_peer_metas = [{}, {}] requests = [ SimpleNamespace( request_id=7, @@ -651,7 +819,8 @@ def test_handoff_metadata_batches_completed_requests_across_pipeline(monkeypatch ), ): prepared = engine._prepare_handoff_metadata_batch( - zip(requests, local_blocks), {7: [71], 8: [81]} + [(request, blocks, None) for request, blocks in zip(requests, local_blocks)], + {7: [71], 8: [81]}, ) for request in requests: engine._capture_handoff_meta(request, prepared[request.request_id]) @@ -670,12 +839,49 @@ def test_handoff_metadata_batches_completed_requests_across_pipeline(monkeypatch assert requests[1].disaggregated_params["kv_meta"]["pp_metas"][1]["block_ids"] == [12, 13, 14] -def test_handoff_keeps_partial_block_for_exact_decode_resume(): +def test_handoff_metadata_error_releases_all_finished_state(): + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + engine.pg_collection = SimpleNamespace(tp=None, pp=None, mp=None) + engine.context = SimpleNamespace(block_size_tokens=4, num_speculative_tokens=0) + engine._kv_peer_metas = None + engine._release_pinned_handoff_blocks = mock.Mock() + engine._release_pinned_handoff_ssm_slot = mock.Mock() + regular = SimpleNamespace( + request_id=7, + prompt_tokens=[1, 2, 3, 4], + sampling_params=SamplingParams(do_kv_handoff=False), + ) + handoff = SimpleNamespace( + request_id=8, prompt_tokens=[1, 2, 3, 4], sampling_params=SamplingParams(do_kv_handoff=True) + ) + + with ( + mock.patch( + "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_size", + return_value=1, + ), + pytest.raises(RuntimeError, match="transfer setup"), + ): + engine._prepare_handoff_metadata_batch( + [(regular, [10], 3), (handoff, [11], 4)], decode_tokens_by_request={} + ) + + assert engine._release_pinned_handoff_blocks.call_args_list == [ + mock.call([10]), + mock.call([11]), + ] + assert engine._release_pinned_handoff_ssm_slot.call_args_list == [mock.call(3), mock.call(4)] + + +def test_hybrid_handoff_keeps_partial_block_for_exact_decode_resume(): engine = object.__new__(InferenceStateHandoffMixin) engine._initialize_disaggregation_state() engine.pg_collection = SimpleNamespace(tp=None, pp=None, mp=None) engine.context = SimpleNamespace(block_size_tokens=4, num_speculative_tokens=0) engine._kv_peer_metas = {"global_rank": 0} + engine._pp_ssm_peer_metas = [{"conv": {"global_rank": 0}, "recurrent": {"global_rank": 0}}] + engine._ssm_transfer_agents = {"conv": mock.Mock(), "recurrent": mock.Mock()} engine._release_pinned_handoff_blocks = mock.Mock() request = SimpleNamespace( request_id=7, @@ -687,11 +893,12 @@ def test_handoff_keeps_partial_block_for_exact_decode_resume(): with mock.patch( "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_size", return_value=1 ): - prepared = engine._prepare_handoff_metadata_batch([(request, [10, 11, 12])], {7: [99]}) + prepared = engine._prepare_handoff_metadata_batch([(request, [10, 11, 12], 9)], {7: [99]}) engine._capture_handoff_meta(request, prepared[7]) assert request.disaggregated_params["block_ids"] == [10, 11, 12] assert request.disaggregated_params["kv_meta"]["resume_tokens"] == [99] + assert engine._pinned_handoff_ssm_slots == {7: 9} engine._release_pinned_handoff_blocks.assert_not_called() @@ -735,7 +942,9 @@ def test_nixl_handoff_reuses_decode_cached_prefix(handoff_loop): assert engine.context.kv_block_allocator.registered_parent_hashes == [hashes[1]] assert engine.get_request(5).generated_tokens == [99] assert 5 not in engine.waiting_request_ids - admit.assert_called_once_with(engine.context, engine.get_request(5), [10, 11, 12], [13], [99]) + admit.assert_called_once_with( + engine.context, engine.get_request(5), [10, 11, 12], [13], [99], ssm_state_idx=None + ) def test_decode_handoff_defers_until_kv_capacity_is_available(handoff_loop): @@ -865,3 +1074,125 @@ def test_decode_role_rejects_prompt_scheduling(handoff_loop): with pytest.raises(RuntimeError, match="cannot schedule prompt prefill"): engine.schedule_waiting_requests() + + +def test_push_handoff_uses_final_pinned_blocks_ssm_slot(): + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + engine._kv_transfer_agent = mock.Mock() + engine._ssm_transfer_agents = {"conv": mock.Mock(), "recurrent": mock.Mock()} + engine._pinned_handoff_blocks[7] = [20, 21] + engine._pinned_handoff_ssm_slots[7] = 3 + decode_metas = [ + {"ssm": {"conv": {"agent": "decode-conv"}, "recurrent": {"agent": "decode-ssm"}}} + ] + + engine.push_handoff_kv(7, decode_metas) + + engine._kv_transfer_agent.begin_push_blocks.assert_called_once_with( + {"tp_metas": decode_metas}, [20, 21] + ) + engine._ssm_transfer_agents["conv"].begin_push_blocks.assert_called_once_with( + {"tp_metas": [{"agent": "decode-conv"}]}, [3] + ) + engine._ssm_transfer_agents["recurrent"].begin_push_blocks.assert_called_once_with( + {"tp_metas": [{"agent": "decode-ssm"}]}, [3] + ) + + +def test_push_handoff_rejects_ssm_state_without_kv_blocks(): + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + engine._kv_transfer_agent = mock.Mock() + engine._ssm_transfer_agents = {"conv": mock.Mock(), "recurrent": mock.Mock()} + engine._pinned_handoff_ssm_slots[7] = 3 + + with pytest.raises(RuntimeError, match="detached SSM state but no pinned KV blocks"): + engine.push_handoff_kv(7, []) + + engine._kv_transfer_agent.begin_push_blocks.assert_not_called() + engine._ssm_transfer_agents["conv"].begin_push_blocks.assert_not_called() + engine._ssm_transfer_agents["recurrent"].begin_push_blocks.assert_not_called() + + +def test_push_handoff_validates_ssm_metadata_before_posting_sends(): + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + engine._kv_transfer_agent = mock.Mock() + engine._ssm_transfer_agents = {"conv": mock.Mock(), "recurrent": mock.Mock()} + engine._pinned_handoff_blocks[7] = [20] + engine._pinned_handoff_ssm_slots[7] = 3 + + with pytest.raises(RuntimeError, match="recurrent SSM transfer state"): + engine.push_handoff_kv(7, [{"ssm": {"conv": {"agent": "decode"}}}]) + + engine._kv_transfer_agent.begin_push_blocks.assert_not_called() + engine._ssm_transfer_agents["conv"].begin_push_blocks.assert_not_called() + + +def test_hybrid_handoff_uses_mirrored_slots_without_request_collectives(): + engine = object.__new__(InferenceStateHandoffMixin) + engine._initialize_disaggregation_state() + tp_group = object() + pp_group = object() + engine.pg_collection = SimpleNamespace(tp=tp_group, pp=pp_group, mp=object()) + engine._kv_peer_metas = [{"rank": "s0t0"}, {"rank": "s0t1"}] + engine._pp_kv_peer_metas = [engine._kv_peer_metas, [{"rank": "s1t0"}, {"rank": "s1t1"}]] + engine._ssm_transfer_agents = {"conv": mock.Mock(), "recurrent": mock.Mock()} + + def stage_ssm_metas(stage): + return { + state: [{"rank": f"s{stage}t{tp}"} for tp in range(2)] + for state in ("conv", "recurrent") + } + + engine._pp_ssm_peer_metas = [stage_ssm_metas(0), stage_ssm_metas(1)] + engine.context = SimpleNamespace( + block_size_tokens=4, num_speculative_tokens=0, memory_buffer=torch.empty(1) + ) + requests = [ + SimpleNamespace( + request_id=2, + prompt_tokens=[0] * 16, + sampling_params=SamplingParams(do_kv_handoff=True), + disaggregated_params=None, + ), + SimpleNamespace( + request_id=3, + prompt_tokens=[0] * 12, + sampling_params=SamplingParams(do_kv_handoff=True), + disaggregated_params=None, + ), + ] + local_blocks = [[10, 11, 12, 13], [14, 15, 16]] + live_slots = [103, 112] + + pg_size = "megatron.core.inference.disaggregation.inference_state_handoff.get_pg_size" + with ( + mock.patch(pg_size, return_value=2), + mock.patch( + "torch.distributed.all_gather_into_tensor", + side_effect=AssertionError("handoff metadata must not run a collective"), + ), + ): + prepared = engine._prepare_handoff_metadata_batch( + [ + (request, blocks, slot) + for request, blocks, slot in zip(requests, local_blocks, live_slots) + ], + {2: [90], 3: [91]}, + ) + for request in requests: + engine._capture_handoff_meta(request, prepared[request.request_id]) + + assert engine._pinned_handoff_ssm_slots == {2: 103, 3: 112} + first_ssm = requests[0].disaggregated_params["kv_meta"]["ssm"] + assert [ + [tp_meta["block_ids"] for tp_meta in stage["tp_metas"]] + for stage in first_ssm["conv"]["pp_metas"] + ] == [[[103], [103]], [[103], [103]]] + assert ( + requests[0].disaggregated_params["kv_meta"] + is not requests[1].disaggregated_params["kv_meta"] + ) + assert all("ssm" not in meta for meta in engine._kv_peer_metas) diff --git a/tests/unit_tests/inference/test_kv_transfer_backends.py b/tests/unit_tests/inference/test_kv_transfer_backends.py index c3bbfa60825..7558428bcfa 100644 --- a/tests/unit_tests/inference/test_kv_transfer_backends.py +++ b/tests/unit_tests/inference/test_kv_transfer_backends.py @@ -135,6 +135,27 @@ def test_ssm_geometry_uses_conv_and_recurrent_state_names(): ) +def test_ssm_geometry_uses_explicit_live_slot_axis(): + layout = SSMShardLayout( + global_rank=0, + tp_size=1, + tp_rank=0, + layer_start=0, + num_layers=3, + dims=SSMStateDims(nheads=2, headdim=4, d_state=5, ngroups=1, d_conv=3), + ) + geometry = base.compute_buffer_geometry( + torch.zeros(3, 3, 2, 4, 5), + expected_num_blocks=3, + backend_name="test", + heads_per_partition=2, + ssm_layout=layout, + ssm_state_kind="recurrent", + ) + + assert geometry.blocks_axis == 1 + + def test_nixl_direct_backend_exports_metadata_with_fake_agent(monkeypatch): from megatron.core.inference.disaggregation.transfer_backends import nixl as nixl_mod diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py index efb4b412798..b24a987be23 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py +++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py @@ -1048,10 +1048,31 @@ def test_async_bookkeeping_retains_finished_handoff_state(): ) assert result.finished_handoff_block_ids == {11: [10, 11]} + assert result.finished_handoff_ssm_slots == {} assert result.finished_handoff_decode_tokens == {11: [99]} context.kv_block_allocator.retain_memory_blocks.assert_called_once_with([10, 11]) +def test_finished_hybrid_handoff_detaches_live_ssm_slot(): + context = _make_async_sched_context(total_request_count=2) + context.is_hybrid_model = True + context.kv_block_allocator = SimpleNamespace( + enable_handoff_pinning=True, retain_memory_blocks=mock.Mock() + ) + context.request_to_kv_block_ids = torch.tensor([[10, 11, -1], [12, 13, -1]], dtype=torch.int32) + context.mamba_metadata = SimpleNamespace(detach_state_slot=mock.Mock(return_value=7)) + controller = _make_async_sched_controller(context) + + blocks, ssm_slots, decode_tokens = controller._collect_finished_handoff_state( + torch.tensor([1]), torch.tensor([91, 92]), None + ) + + assert blocks == {11: [12, 13]} + assert ssm_slots == {11: 7} + assert decode_tokens == {11: [92]} + context.mamba_metadata.detach_state_slot.assert_called_once_with(1) + + @pytest.mark.parametrize( "termination_ids, stop_word_finished_ids", [([99, 99, 99], set()), ([99, 2, 99], set()), ([99, 99, 99], {11})], @@ -1161,6 +1182,7 @@ def test_async_sched_step_overlap_order(): newly_paused_request_ids=None, evict_request_ids=None, finished_handoff_block_ids={}, + finished_handoff_ssm_slots={}, finished_handoff_decode_tokens={}, ) ) @@ -1290,6 +1312,7 @@ def test_async_sched_step_yields_after_resolution_outside_inference_mode(): newly_paused_request_ids=None, evict_request_ids=None, finished_handoff_block_ids={}, + finished_handoff_ssm_slots={}, finished_handoff_decode_tokens={}, ) ) @@ -1433,6 +1456,7 @@ def test_async_sched_no_overlap_updates_before_admission( newly_paused_request_ids=torch.tensor([10]), evict_request_ids=torch.tensor([11]), finished_handoff_block_ids={}, + finished_handoff_ssm_slots={}, finished_handoff_decode_tokens={}, ) input_ids = torch.empty(1, dtype=torch.int64) @@ -1539,6 +1563,7 @@ def test_async_sched_no_overlap_finishes_with_matching_ep_base_forward(): newly_paused_request_ids=None, evict_request_ids=None, finished_handoff_block_ids={}, + finished_handoff_ssm_slots={}, finished_handoff_decode_tokens={}, ) controller._run_async_sched_sample = mock.Mock(return_value=sample_result) @@ -1591,6 +1616,7 @@ def test_async_sched_mtp_overlap_step_order(): newly_paused_request_ids=None, evict_request_ids=None, finished_handoff_block_ids={}, + finished_handoff_ssm_slots={}, finished_handoff_decode_tokens={}, ) input_ids = torch.empty(9, dtype=torch.int64) From 19deef67f910c96c213f33b33b30277be8b94d6d Mon Sep 17 00:00:00 2001 From: Carlos Gomes Date: Wed, 19 Aug 2026 15:13:03 +0000 Subject: [PATCH 100/101] Checkpoint shared storage when inputs are returned (#4497) Signed-off-by: CarlosGomes98 --- .../core/extensions/transformer_engine.py | 1 + .../core/models/gpt/fine_grained_callables.py | 12 +-- megatron/core/tensor_parallel/random.py | 38 +++++++-- megatron/core/transformer/torch_norm.py | 6 +- .../core/transformer/transformer_layer.py | 51 +++++------- .../unit_tests/tensor_parallel/test_random.py | 80 +++++++++++++++++++ 6 files changed, 139 insertions(+), 49 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 9393ad825c3..c63f0c75dc2 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1071,6 +1071,7 @@ def __new__( **_get_extra_te_kwargs(config), ) + instance.returns_residual = use_fused_residual return cast(LayerNormInterface, instance) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index ba1fd3af499..fd64f3d044f 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -103,7 +103,9 @@ def forward_func( mlp_norm_manager = off_interface(layer.offload_mlp_norm, hidden_states, "mlp_norm") node.layer_state.mlp_norm_manager = mlp_norm_manager if layer.recompute_pre_mlp_layernorm: - layer.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + layer.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + retain_input_tensors=layer._pre_mlp_layernorm_returns_residual + ) with mlp_norm_manager as hidden_states: pre_mlp_layernorm_output = layer.pre_mlp_norm_checkpoint.checkpoint( apply_module(layer.pre_mlp_layernorm), hidden_states @@ -117,13 +119,7 @@ def forward_func( # When using fused residual norm (e.g. TEFusedResidualRMSNorm), # the layernorm returns (normalized_output, residual). Unpack # and use the fused residual for the downstream BDA connection. - if isinstance(pre_mlp_layernorm_output, tuple): - if len(pre_mlp_layernorm_output) != 2: - raise ValueError( - f"When the output of pre_mlp_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(pre_mlp_layernorm_output)}" - ) + if layer._pre_mlp_layernorm_returns_residual: pre_mlp_layernorm_output, hidden_states = pre_mlp_layernorm_output shared_expert_output = layer.mlp.shared_experts_compute(pre_mlp_layernorm_output) diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index ac523378c34..7915f045bda 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -832,8 +832,7 @@ def add_checkpoint(self, ckpt): def discard_all_outputs_and_register_unified_recompute(self, hook_tensor): """Discard all checkpoint outputs to save memory and register unified recompute hook.""" for ckpt in self.checkpoints: - for output in ckpt.outputs: - output.untyped_storage().resize_(0) + ckpt._discard_outputs() # Register unified recompute hook if hook_tensor.requires_grad: @@ -860,7 +859,7 @@ class CheckpointWithoutOutput(object): discarded output tensors are directly saved in the following modules for backward computation. """ - def __init__(self, fp8=False, ckpt_manager=None): + def __init__(self, fp8=False, ckpt_manager=None, retain_input_tensors=False): """ Initialize CheckpointWithoutOutput. @@ -878,9 +877,12 @@ def __init__(self, fp8=False, ckpt_manager=None): checkpoint() will auto-register to the manager, and discard_output_and_register_recompute() will only discard output without registering individual hooks. + retain_input_tensors: Whether outputs sharing storage with checkpoint inputs + should be retained when discarding outputs. """ self.fp8 = fp8 is not None self.ckpt_manager = ckpt_manager + self.retain_input_tensors = retain_input_tensors self.run_function = None self.fwd_cpu_rng_state = None self.fwd_cuda_rng_state = None @@ -907,6 +909,11 @@ def checkpoint(self, run_function: Callable[[Unpack[_Ts]], _R], *args: Unpack[_T self.rng_states = _get_all_rng_states() + if self.retain_input_tensors: + self._saved_input_ptrs = { + t.untyped_storage().data_ptr() for t in args if isinstance(t, torch.Tensor) + } + outputs = CheckpointWithoutOutputFunction.apply(run_function, self, *args) self.outputs = outputs if isinstance(self.outputs, torch.Tensor): @@ -964,13 +971,30 @@ def _recompute(self, _): # - No tensor version-counter bump (no autograd complaint) share_storage = _get_share_storage() for output, recomputation_output in zip(self.outputs, outputs): - share_storage(output, recomputation_output) + if ( + output.untyped_storage().data_ptr() + != recomputation_output.untyped_storage().data_ptr() + ): + share_storage(output, recomputation_output) self.ctx.outputs = outputs self.ctx.inputs = inputs self.outputs = None self.ctx = None + def _discard_outputs(self): + """Release output storage, preserving outputs that alias retained inputs.""" + if self.retain_input_tensors: + # Skip outputs whose storage is shared with a saved input — freeing those + # would destroy the data needed for recomputation (e.g. TE.ops.Sequential + # operations with MakeExtraOutput). + for output in self.outputs: + if output.untyped_storage().data_ptr() not in self._saved_input_ptrs: + output.untyped_storage().resize_(0) + else: + for output in self.outputs: + output.untyped_storage().resize_(0) + def discard_output_and_register_recompute(self, hook_tensor): """ Release the output tensor storages and register the recompute function as a grad hook of @@ -987,10 +1011,8 @@ def discard_output_and_register_recompute(self, hook_tensor): if self.ckpt_manager is not None or is_graph_warmup(): return - # use resize to release the output tensor memory and still keep the metadata in the tensors. - # the metadata is still needed for backward - for output in self.outputs: - output.untyped_storage().resize_(0) + # Release output tensor memory while keeping metadata for backward. + self._discard_outputs() # register the recomputation as a backward hook, when the the gradient of the hook_tensor # is computed, the recomputation will be triggered. The hook_tensor should be selected diff --git a/megatron/core/transformer/torch_norm.py b/megatron/core/transformer/torch_norm.py index 5948ae600f9..d961b214d43 100644 --- a/megatron/core/transformer/torch_norm.py +++ b/megatron/core/transformer/torch_norm.py @@ -11,6 +11,8 @@ class LayerNormInterface(Protocol): """Interface that all LayerNorm implementations should follow.""" + returns_residual: bool + def forward(self, x: torch.Tensor, /) -> torch.Tensor: """Forward method for a LayerNorm implementation.""" ... @@ -66,7 +68,9 @@ def __new__( else: raise Exception("Only LayerNorm, RMSNorm and L2Norm are currently supported") - return norm_cls(normalized_shape=hidden_size, eps=eps) + instance = norm_cls(normalized_shape=hidden_size, eps=eps) + instance.returns_residual = False + return instance class L2Norm(torch.nn.Module, LayerNormInterface): diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 0a8f4c8408a..11f04c73b27 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -370,6 +370,9 @@ def __init__( hidden_size=self.config.hidden_size, eps=self.config.layernorm_epsilon, ) + self._input_layernorm_returns_residual = getattr( + self.input_layernorm, 'returns_residual', False + ) attention_optional_kwargs = {} if config.context_parallel_size > 1 and config.cp_comm_type is not None: @@ -403,6 +406,9 @@ def __init__( hidden_size=self.config.hidden_size, eps=self.config.layernorm_epsilon, ) + self._pre_cross_attn_layernorm_returns_residual = getattr( + self.pre_cross_attn_layernorm, 'returns_residual', False + ) # [Module 5: CrossAttention] self.cross_attention = build_module( @@ -422,6 +428,9 @@ def __init__( hidden_size=self.config.hidden_size, eps=self.config.layernorm_epsilon, ) + self._pre_mlp_layernorm_returns_residual = getattr( + self.pre_mlp_layernorm, 'returns_residual', False + ) # [Module 8: MLP block] # import here to avoid circular import from megatron.core.extensions.transformer_engine import TEFusedMLP @@ -702,7 +711,9 @@ def _run_input_layernorm(self, hidden_states): ) self._input_layernorm_checkpoint_active = self.recompute_input_layernorm if self._input_layernorm_checkpoint_active: - self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + retain_input_tensors=self._input_layernorm_returns_residual + ) with self.attn_norm_manager as hidden_states: input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( apply_module(self.input_layernorm), hidden_states @@ -711,13 +722,7 @@ def _run_input_layernorm(self, hidden_states): with self.attn_norm_manager as hidden_states: input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) - if isinstance(input_layernorm_output, tuple): - if len(input_layernorm_output) != 2: - raise ValueError( - f"When the output of input_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(input_layernorm_output)}" - ) + if self._input_layernorm_returns_residual: input_layernorm_output, residual = input_layernorm_output else: residual = hidden_states @@ -763,14 +768,7 @@ def _run_cross_attention(self, hidden_states, context, context_mask, inference_c """Optional pre-cross-attn layernorm + cross-attention + bda block.""" pre_cross_attn_layernorm_output = apply_module(self.pre_cross_attn_layernorm)(hidden_states) - if isinstance(pre_cross_attn_layernorm_output, tuple): - if len(pre_cross_attn_layernorm_output) != 2: - raise ValueError( - f"When the output of pre_cross_attn_layernorm_output " - f"is a tuple, it is expected to have 2 elements " - f"(output, residual), but " - f"got {len(pre_cross_attn_layernorm_output)}" - ) + if self._pre_cross_attn_layernorm_returns_residual: pre_cross_attn_layernorm_output, residual = pre_cross_attn_layernorm_output else: residual = hidden_states @@ -820,7 +818,9 @@ def forward(self, *args, **kwargs): def _forward_pre_mlp_layernorm(self, hidden_states: Tensor): self.mlp_norm_manager = self.off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") if self.recompute_pre_mlp_layernorm: - self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + retain_input_tensors=self._pre_mlp_layernorm_returns_residual + ) with self.mlp_norm_manager as hidden_states: pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( apply_module(self.pre_mlp_layernorm), hidden_states @@ -882,16 +882,9 @@ def _pre_mlp_layernorm_and_residual(self, hidden_states): """ pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) - if isinstance(pre_mlp_layernorm_output, tuple): - if len(pre_mlp_layernorm_output) != 2: - raise ValueError( - f"When the output of pre_mlp_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(pre_mlp_layernorm_output)}" - ) + if self._pre_mlp_layernorm_returns_residual: pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output else: - # Residual connection. residual = hidden_states if self.config.fp32_residual_connection: @@ -1449,13 +1442,7 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): if not self.is_moe_layer: return residual, None, None, None hidden_states = apply_module(self.pre_mlp_layernorm)(residual) - if isinstance(hidden_states, tuple): - if len(hidden_states) != 2: - raise ValueError( - f"When the output of pre_mlp_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(hidden_states)}" - ) + if self._pre_mlp_layernorm_returns_residual: hidden_states, residual = hidden_states shared_expert_output = self.mlp.shared_experts_compute(hidden_states) diff --git a/tests/unit_tests/tensor_parallel/test_random.py b/tests/unit_tests/tensor_parallel/test_random.py index 873e6a30f00..e603ffd4abd 100644 --- a/tests/unit_tests/tensor_parallel/test_random.py +++ b/tests/unit_tests/tensor_parallel/test_random.py @@ -5,6 +5,7 @@ from megatron.core.tensor_parallel.random import ( CheckpointWithoutOutput, + CheckpointWithoutOutputManager, CudaRNGStatesTracker, checkpoint, convert_cuda_rng_state, @@ -296,3 +297,82 @@ def checkpoint_forward(input_, weight): assert torch.allclose(weight1.grad, weight2.grad) finally: Utils.destroy_model_parallel() + + +class _FusedResidualNorm(torch.autograd.Function): + """Mimics TEFusedResidualRMSNorm: returns (norm(input), input) where the + residual shares storage with the input (the MakeExtraOutput pattern). + Backward combines grad_norm and grad_residual into a single grad_input.""" + + @staticmethod + def forward(ctx, x, weight): + rms = x.float().pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() + normed = (x.float() * rms).to(x.dtype) * weight + ctx.save_for_backward(x, weight, rms) + return normed, x + + @staticmethod + def backward(ctx, grad_normed, grad_residual): + x, weight, rms = ctx.saved_tensors + x_float = x.float() + normed_float = x_float * rms + grad_normed_float = (grad_normed * weight).float() + d = x.shape[-1] + grad_x = ( + grad_normed_float * rms + - normed_float * (grad_normed_float * normed_float).sum(-1, keepdim=True) / d + ) + grad_x = grad_x.to(x.dtype) + grad_residual + grad_weight = (grad_normed * normed_float.to(x.dtype)).sum( + dim=tuple(range(grad_normed.ndim - 1)) + ) + return grad_x, grad_weight + + +@pytest.mark.parametrize("use_manager", [False, True]) +def test_checkpoint_without_output_retain_input_tensors(use_manager): + """CheckpointWithoutOutput with retain_input_tensors=True must not free + outputs that share storage with inputs (the fused residual norm pattern).""" + + hidden = 32 + + def fused_norm(x): + return _FusedResidualNorm.apply(x, weight) + + def normal_forward(x): + normed, residual = fused_norm(x) + return normed + residual + + def checkpoint_forward(x): + manager = CheckpointWithoutOutputManager() if use_manager else None + ckpt = CheckpointWithoutOutput(ckpt_manager=manager, retain_input_tensors=True) + normed, residual = ckpt.checkpoint(fused_norm, x) + y = normed + residual + if manager is None: + ckpt.discard_output_and_register_recompute(y) + else: + manager.discard_all_outputs_and_register_unified_recompute(y) + return y + + Utils.initialize_model_parallel() + try: + weight = torch.randn(hidden, requires_grad=True) + input_ref = torch.randn((4, hidden), requires_grad=True) + + input1 = input_ref.detach().clone().requires_grad_(True) + output1 = normal_forward(input1) + + input2 = input_ref.detach().clone().requires_grad_(True) + weight.grad = None + output2 = checkpoint_forward(input2) + assert torch.allclose(output1, output2) + + grad = torch.randn_like(output1) + output1.backward(grad) + ref_input_grad = input1.grad.clone() + + weight.grad = None + output2.backward(grad) + assert torch.allclose(ref_input_grad, input2.grad) + finally: + Utils.destroy_model_parallel() From 4a534df611eeb12250fad0b09a1a9d2eccf83fec Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Fri, 16 Oct 2026 05:30:30 +0000 Subject: [PATCH 101/101] fix: reconcile main slice 19deef67 with dev for nightly sync (2026-10-13) Repair commit for the 2026-10-13 main-to-dev sync merge 1eaec475. - inference: restore main #5301 pre-bound socket support (the sock parameter) in start_text_gen_server. The slice's #6513 makes MegatronLLM.serve() and MegatronAsyncLLM.serve() pass sock=, but dev's server lacked #5301, so serve() raised TypeError on the primary rank. Add tests that bind both serve() calls to the real signature and cover the pre-bound socket path. - inference: restore main #5348/#6584 mamba-cache scratch reservation and budget-too-small / PP-limiting-stage errors in DynamicInferenceContext (the dev warn-and-disable path was an artifact of revert #7283), and restore test_mamba_cache_budget_too_small_raises. - megatron_fsdp experimental: keep the symmetric-memory all-gather buffer rendezvous working on the dev torch build, which lacks symm_mem.is_symm_mem_tensor. - functional moe_perf: restore the per-iteration CUDA RNG tracker state reset from #7094. - unit test doubles: align with merged APIs (RecordingCheckpoint retain_input_tensors, paged-stash config cuda_graph fields, profile_ranks in test_training, contextmanager import and the dev pipeline_p2p_fixed_shape config field in test_schedules). - unit test isolation: reset the global CUDA graph record after the dev TestMHCWithCudaGraph local-impl tests and before main #6432's TestTransformerLayerCudaGraphManagers, so its create_cudagraphs() does not capture stale mHC layers recorded under the cudagraphable RNG tracker. - unit test teardown: Utils.destroy_model_parallel resets the TE FP8 global state and the HybridEP buffer that cache process groups, and retires the groups parallel_state created; a new autouse conftest fixture destroys them once the test finishes. They are not destroyed inside destroy_model_parallel because a test may re-initialize and keep using a group it captured earlier (test_upcycling_TE reuses its ProcessGroupCollection). Before this, each test leaked its NCCL communicators. The tests/unit_tests/transformer bucket then crossed vm.max_map_count (65530) at test_vision_cuda_graphs, on both this branch and dev c75a269e. Before destroying the retired groups, the fixture clears DTensor's per-thread sharding-propagation caches (Python and native). They key on value-equal DeviceMeshes, so a later test that builds a mesh with the same layout would otherwise get placements on the earlier mesh, whose group is destroyed (test_gpt_hybrid_interop fsdp4-dense). - rl unit tests: a bucket-scoped autouse conftest fixture runs gc.collect() before the retired groups are destroyed. test_get_logprobs_cuda_graphs leaves CUDA graphs that captured those communicators reachable only through reference cycles, and ncclCommDestroy waits for them to be freed, so the teardown hung. Collecting suite-wide instead changed object lifetimes in other buckets, so the collection stays in tests/unit_tests/rl. Signed-off-by: svcnemo-autobot --- .../src/megatron_fsdp/experimental/dbuffer.py | 28 +++- .../experimental/parameter_group.py | 10 +- .../inference/contexts/dynamic_context.py | 57 ++++++-- .../text_generation_server.py | 36 +++-- .../test_cases/common/moe_perf/__main__.py | 11 +- tests/unit_tests/conftest.py | 7 + .../contexts/test_dynamic_prefix_caching.py | 18 ++- .../test_serve_text_gen_server.py | 134 ++++++++++++++++++ .../pipeline_parallel/test_schedules.py | 2 + tests/unit_tests/rl/conftest.py | 25 ++++ tests/unit_tests/test_training.py | 1 + tests/unit_tests/test_utilities.py | 59 ++++++++ .../moe/test_paged_stash_runner.py | 7 +- .../transformer/test_transformer_layer.py | 6 +- 14 files changed, 363 insertions(+), 38 deletions(-) create mode 100644 tests/unit_tests/inference/high_level_api/test_serve_text_gen_server.py create mode 100644 tests/unit_tests/rl/conftest.py diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 30b78c6282c..831ef5087e4 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -14,6 +14,8 @@ """Distributed tensor buffers for Megatron-FSDP.""" +import contextlib +import contextvars import dataclasses from collections.abc import Iterable @@ -35,6 +37,29 @@ class _OwnedRange: buffer_relative_offset: int +# True while DBuffers are allocated from the NCCL symmetric-memory pool; see +# symmetric_memory_allocation. +_ALLOCATING_FROM_SYMM_MEM: contextvars.ContextVar[bool] = contextvars.ContextVar( + "megatron_fsdp_allocating_from_symm_mem", default=False +) + + +@contextlib.contextmanager +def symmetric_memory_allocation(pool: torch.cuda.MemPool): + """Allocate DBuffers created in this scope from the symmetric-memory ``pool``. + + DBuffers constructed inside the scope record ``allocated_from_symm_mem``, so + ``DBuffer.is_symmetric_memory`` still recognizes them, and their collectives still + rendezvous, on PyTorch builds without ``symm_mem.is_symm_mem_tensor``. + """ + token = _ALLOCATING_FROM_SYMM_MEM.set(True) + try: + with torch.cuda.use_mem_pool(pool): + yield + finally: + _ALLOCATING_FROM_SYMM_MEM.reset(token) + + def _validate_placements(placements: Iterable[Placement]) -> None: """Validate DBuffer placements form a supported contiguous local layout.""" seen_flat = False @@ -71,7 +96,7 @@ class DBuffer: # before deleting the buffer. See https://github.com/NVIDIA/Megatron-LM/pull/6187 and # https://docs.pytorch.org/docs/2.13/generated/torch.Tensor.record_stream.html. allocation_stream: torch.cuda.Stream | None - # Set by owners that allocate ``local_buffer`` from the NCCL symmetric-memory pool. Only + # Whether ``local_buffer`` was allocated inside ``symmetric_memory_allocation``. Only # consulted on PyTorch builds without ``symm_mem.is_symm_mem_tensor``. allocated_from_symm_mem: bool = False @@ -107,6 +132,7 @@ def __init__( self.offset, local_numel = self.layout.get_local_range(self.mesh, self.placements) self.local_buffer = torch.empty(local_numel, dtype=dtype, device=device) + self.allocated_from_symm_mem = _ALLOCATING_FROM_SYMM_MEM.get() self.allocation_stream = ( torch.cuda.current_stream(self.local_buffer.device) if self.local_buffer.is_cuda diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index cb47045632f..b05dbd7fa3c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -26,7 +26,7 @@ from torch.distributed import DeviceMesh from ..mixed_precision import MixedPrecisionPolicy -from .dbuffer import DBuffer +from .dbuffer import DBuffer, symmetric_memory_allocation from .placement import Partial, Placements, Replicate, changed_mesh_axis _CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group" @@ -244,7 +244,7 @@ def __init__( def _symmetric_memory_context(self): if self._symm_mem_pool is None: return nullcontext() - return torch.cuda.use_mem_pool(self._symm_mem_pool) + return symmetric_memory_allocation(self._symm_mem_pool) def _set_module_parameter(self, fqns: tuple[str, ...], parameter: nn.Parameter) -> None: owning_module = self._owning_module() @@ -331,17 +331,13 @@ def allocate_partial_grad_buffer(self) -> DBuffer: raise RuntimeError(f"Missing gradient for FSDP parameter {fsdp_parameter.fqns!r}.") grads.append(fsdp_parameter.unsharded.grad) with self._symmetric_memory_context(): - partial_grad = DBuffer( + return DBuffer( mesh=self.mesh, placements=[Partial(dist.ReduceOp.AVG)] * self.mesh.ndim, tensor_shapes=tuple(grad.shape for grad in grads), dtype=grads[0].dtype, device=grads[0].device, ) - # Lets DBuffer.reduce_scatter select the symmetric-memory SUM path on PyTorch builds - # without symm_mem.is_symm_mem_tensor. - partial_grad.allocated_from_symm_mem = self._symm_mem_pool is not None - return partial_grad def copy_gradients_to_partial_buffer(self, partial_grad: DBuffer) -> None: """Pack full local gradients into an existing reduce-scatter input buffer.""" diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 79c4797a971..2786e970751 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -897,11 +897,18 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC assert self.mamba_slot_allocator is not None prefix_cache_bytes = int(prefix_caching_mamba_gb * 1024**3) # The allocator contains the PP-synchronized durable capacity. - prefix_cache_slots = self.mamba_slot_allocator.max_slots + # Scratch remains stage-local because it is temporary forward + # storage and does not participate in prefix-cache matching. + scratch_slots = self.max_mamba_intermediate_states_per_step + scratch_bytes = scratch_slots * mamba_bytes_per_req + durable_slots = self.mamba_slot_allocator.max_slots log_lines += [ f" Mamba prefix cache:", f" budget: {get_mem_size_str(prefix_cache_bytes)}", - f" slots: {prefix_cache_slots}", + f" extraction_scratch: {scratch_slots} slots " + f"({get_mem_size_str(scratch_bytes)})", + f" durable_slots: {durable_slots} " + f"({get_mem_size_str(durable_slots * mamba_bytes_per_req)})", f" per_slot: {get_mem_size_str(mamba_bytes_per_req)}", ] @@ -1755,7 +1762,22 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: ssm_size = _math.prod(self.mamba_ssm_states_shape) * self.mamba_ssm_states_dtype.itemsize per_slot_bytes = self.num_mamba_layers * (conv_size + ssm_size) total_bytes = int(mamba_gb * 1024**3) - max_slots = total_bytes // per_slot_bytes + + # MambaSlotAllocator allocates two GPU buffer families with the same + # per-slot footprint, both of which must fit in this budget: + # - "durable" cache: self.ssm_states / self.conv_states, sized to + # `max_slots` slots (computed below). + # - "scratch" buffers: self.intermediate_ssm_out / self.intermediate_conv_out, + # fixed CUDA-graph-safe staging for intermediate-state + # extraction, sized to the per-step token-budget cap + # `scratch_slots` = max_mamba_intermediate_states_per_step. + # The scratch is not part of the durable cache but consumes the same + # per-slot bytes, so reserve it from the budget up front before sizing the + # durable cache; otherwise total usage silently exceeds mamba_gb (and can + # OOM) when scratch_slots > max_slots. + scratch_slots = self.max_mamba_intermediate_states_per_step + scratch_bytes = scratch_slots * per_slot_bytes + max_slots = (total_bytes - scratch_bytes) // per_slot_bytes # durable slots local_max_slots = max_slots # Prefix-cache state is replicated across pipeline stages. Use the @@ -1775,13 +1797,23 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: max_slots = int(max_slots_tensor.item()) if max_slots < 1: - logging.warning( - "Mamba cache budget (%.3f GB) too small for even 1 slot " - "(need %.3f GB per slot). Mamba caching disabled.", - mamba_gb, - per_slot_bytes / 1024**3, + if local_max_slots >= 1: + raise ValueError( + f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " + f"has room for {local_max_slots} durable slots on this pipeline stage, but " + f"another stage has room for fewer than one. Cache capacity is mirrored to " + f"the minimum across stages; increase prefix_caching_mamba_gb, reduce " + f"max_tokens, or rebalance the Mamba layers." + ) + raise ValueError( + f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " + f"is too small. The CUDA-graph extraction scratch reserves " + f"{scratch_bytes / 1024**3:.4g} GB ({scratch_slots} slots x " + f"{per_slot_bytes / 1024:.1f} KB/slot), leaving room for " + f"fewer than one durable cache slot. Increase prefix_caching_mamba_gb " + f"to at least {(scratch_bytes + per_slot_bytes) / 1024**3:.4g} GB, or " + f"reduce max_tokens." ) - return self.mamba_slot_allocator = MambaSlotAllocator( context=self, @@ -1797,9 +1829,14 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: ) logging.info( - "Mamba prefix cache: %d slots (%.3f GB), per-slot %.1f KB", + "Mamba prefix cache: %d durable slots (%.3f GB) + %d scratch slots " + "(%.3f GB) = %.3f GB total within %.3f GB budget, per-slot %.1f KB", max_slots, max_slots * per_slot_bytes / 1024**3, + scratch_slots, + scratch_bytes / 1024**3, + (max_slots + scratch_slots) * per_slot_bytes / 1024**3, + mamba_gb, per_slot_bytes / 1024, ) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py index 8f91644afdd..af12bc5bbc1 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. import asyncio import logging @@ -150,6 +150,7 @@ def start_text_gen_server( verbose: bool = False, num_replicas: int = 4, hostname: Optional[str] = None, + sock: Optional[socket.socket] = None, ): """Start the text generation server.""" global _SERVER_PROCESSES @@ -159,18 +160,29 @@ def start_text_gen_server( logger.warning("Text gen server processes are already running.") return - _SHARED_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - if hasattr(socket, 'SO_REUSEPORT'): - try: - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - except OSError: - pass + # The caller may pass in a socket it has already bound ahead of time. + if sock is not None: + bound_port = sock.getsockname()[1] + if bound_port == 0: + raise ValueError( + "socket must be bound to a real port before being passed to start_text_gen_server" + ) + _SHARED_SOCKET = sock + server_port = bound_port + _SHARED_SOCKET.setblocking(False) + else: + _SHARED_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + + if hasattr(socket, 'SO_REUSEPORT'): + try: + _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + except OSError: + pass - bind_address = hostname if hostname is not None else "0.0.0.0" - _SHARED_SOCKET.bind((bind_address, server_port)) - _SHARED_SOCKET.setblocking(False) + bind_address = hostname if hostname is not None else "0.0.0.0" + _SHARED_SOCKET.bind((bind_address, server_port)) + _SHARED_SOCKET.setblocking(False) _SHARED_SOCKET.set_inheritable(True) fd = _SHARED_SOCKET.fileno() diff --git a/tests/functional_tests/test_cases/common/moe_perf/__main__.py b/tests/functional_tests/test_cases/common/moe_perf/__main__.py index a722121f079..f63ec3d5ce1 100644 --- a/tests/functional_tests/test_cases/common/moe_perf/__main__.py +++ b/tests/functional_tests/test_cases/common/moe_perf/__main__.py @@ -19,9 +19,9 @@ from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_with_transformer_engine_submodules, ) +from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP, HAVE_HYBRIDEP from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules -from megatron.core.transformer.moe.moe_utils import RandomSTE from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import nvtx_range_pop, nvtx_range_push @@ -206,9 +206,14 @@ def _benchmark_moe_layer(layer: MoELayer, case: MoEPerformanceCase): generator=generator, ) input_tensor.requires_grad_(True) + initial_rng_states = get_cuda_rng_tracker().get_states() for iteration in range(WARMUP_ITERS + MEASURE_ITERS): - if RandomSTE.generator is not None: - RandomSTE.generator.manual_seed(RandomSTE.generator.initial_seed()) + get_cuda_rng_tracker().set_states( + { + name: state.clone() if isinstance(state, torch.Tensor) else state.clone_state() + for name, state in initial_rng_states.items() + } + ) if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.barrier() nvtx_iter_msg = f"({case.name}) iteration {iteration}" diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index 4620c6a57fc..2d09781fe19 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -68,6 +68,13 @@ def cleanup(): torch.distributed.destroy_process_group() +@pytest.fixture(scope="function", autouse=True) +def destroy_retired_process_groups(): + """Destroy the model-parallel process groups a test retired once it finishes.""" + yield + Utils.destroy_retired_process_groups() + + @pytest.fixture(scope="function", autouse=True) def set_env(): if is_te_min_version("1.3"): diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index 57da62ead8d..77104af9304 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -936,6 +936,17 @@ def test_mamba_cache_lifecycle(self): and ctx8.mamba_slot_allocator.has_state(bids8[0]) ) + @pytest.mark.internal + def test_mamba_cache_budget_too_small_raises(self): + # The CUDA-graph extraction scratch (sized to the per-step token-budget + # cap, max_mamba_intermediate_states_per_step) is reserved from + # prefix_caching_mamba_gb before the durable cache is sized. A budget too + # small to fit the scratch plus at least one durable slot is a hard + # configuration error, not a silent over-allocation (which previously + # could OOM at startup). + with pytest.raises(ValueError, match="prefix cache budget"): + self._mctx(prefix_caching_mamba_gb=1e-5) + @pytest.mark.internal def test_hybrid_admission_resumes_after_handoff_releases_live_slot(self): ctx = self._mctx(max_requests=1, rounder=1) @@ -1533,8 +1544,11 @@ def test_allocate_slots_batch(self): assert mixed_slots[0] == pre_slot assert msa3.free_count == free_before3 - 2 # only 2 new - # Eviction: exhaust free pool, verify eviction fires and returns valid slots - ctx4 = self._mctx(prefix_caching_mamba_gb=0.001) + # Eviction: exhaust free pool, verify eviction fires and returns valid slots. + # Budget must cover the CUDA-graph extraction scratch (3 * max_requests + # slots) plus the durable cache; a budget too small to fit the scratch now + # raises (see test_mamba_cache_budget_too_small_raises). + ctx4 = self._mctx(prefix_caching_mamba_gb=0.01) msa4 = ctx4.mamba_slot_allocator total_slots = msa4.max_slots ctx4.add_request(self._req(ctx4, self._prompt(bs * 4))) diff --git a/tests/unit_tests/inference/high_level_api/test_serve_text_gen_server.py b/tests/unit_tests/inference/high_level_api/test_serve_text_gen_server.py new file mode 100644 index 00000000000..5811760f61f --- /dev/null +++ b/tests/unit_tests/inference/high_level_api/test_serve_text_gen_server.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Contract tests between the high-level ``serve()`` APIs and +``start_text_gen_server``. + +``test_apis.py`` replaces ``start_text_gen_server`` with a ``**kwargs`` lambda, +which accepts any keyword. These tests bind both ``serve()`` calls to the real +signature with ``create_autospec`` and exercise the pre-bound socket path, so a +caller and the server cannot drift apart unnoticed. +""" + +import socket +from unittest.mock import MagicMock, create_autospec + +import pytest + +from megatron.core.inference.apis.async_llm import MegatronAsyncLLM +from megatron.core.inference.apis.llm import MegatronLLM +from megatron.core.inference.apis.serve_config import ServeConfig + +tgs = pytest.importorskip( + "megatron.core.inference.text_generation_server.dynamic_text_gen_server" + ".text_generation_server" +) + + +def _make_primary_instance(cls): + """Build a coordinator-mode primary-rank instance without running the + constructor's engine/runtime setup.""" + obj = cls.__new__(cls) + obj._engine = MagicMock() + obj._context = MagicMock() + obj._controller = MagicMock() + obj._use_coordinator = True + obj._is_primary_rank = True + obj._loop_manager = None + obj._coord_runtime = MagicMock() + obj._coord_runtime.coord_addr = "tcp://coord:5555" + obj._shutdown_called = False + obj._serve_started = False + return obj + + +@pytest.fixture +def autospec_start(monkeypatch): + """Replace ``start_text_gen_server`` with a mock that enforces its real signature.""" + import torch.distributed as dist + + monkeypatch.setattr(dist, "get_rank", lambda: 0) + start = create_autospec(tgs.start_text_gen_server) + monkeypatch.setattr(tgs, "start_text_gen_server", start) + return start + + +@pytest.fixture +def fake_mp(monkeypatch): + """Record frontend replicas instead of spawning processes, and isolate the + module-level server globals.""" + fake = MagicMock() + fake.Process.return_value.is_alive.return_value = False + monkeypatch.setattr(tgs, "mp", fake) + monkeypatch.setattr(tgs, "_SERVER_PROCESSES", []) + monkeypatch.setattr(tgs, "_SHARED_SOCKET", None) + return fake + + +class TestServeMatchesStartTextGenServer: + """Both ``serve()`` implementations call the server with accepted arguments.""" + + @pytest.mark.parametrize("sock", [None, MagicMock(spec=socket.socket)]) + def test_sync_serve(self, autospec_start, sock): + llm = _make_primary_instance(MegatronLLM) + llm.serve(ServeConfig(port=1234, sock=sock), blocking=False) + assert llm._serve_started is True + autospec_start.assert_called_once() + kwargs = autospec_start.call_args.kwargs + assert kwargs["coordinator_addr"] == "tcp://coord:5555" + assert kwargs["server_port"] == 1234 + assert kwargs["sock"] is sock + + @pytest.mark.asyncio + @pytest.mark.parametrize("sock", [None, MagicMock(spec=socket.socket)]) + async def test_async_serve(self, autospec_start, sock): + llm = _make_primary_instance(MegatronAsyncLLM) + await llm.serve(ServeConfig(port=1234, sock=sock), blocking=False) + assert llm._serve_started is True + autospec_start.assert_called_once() + kwargs = autospec_start.call_args.kwargs + assert kwargs["coordinator_addr"] == "tcp://coord:5555" + assert kwargs["server_port"] == 1234 + assert kwargs["sock"] is sock + + +class TestPreboundSocket: + """``start_text_gen_server(sock=...)`` serves on the caller's bound socket.""" + + def test_replicas_share_the_prebound_socket(self, fake_mp): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + bound_port = sock.getsockname()[1] + tgs.start_text_gen_server( + coordinator_addr="tcp://coord:5555", + tokenizer=MagicMock(), + rank=0, + server_port=9999, + num_replicas=2, + sock=sock, + ) + assert tgs._SHARED_SOCKET is sock + assert sock.getblocking() is False + assert sock.get_inheritable() is True + assert fake_mp.Process.call_count == 2 + for call in fake_mp.Process.call_args_list: + args = call.kwargs["args"] + # (coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname) + assert args[3] == bound_port + assert args[6] == sock.fileno() + + tgs.stop_text_gen_server() + assert tgs._SHARED_SOCKET is None + assert sock.fileno() == -1 + + def test_unbound_socket_is_rejected(self, fake_mp): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + with pytest.raises(ValueError, match="bound to a real port"): + tgs.start_text_gen_server( + coordinator_addr="tcp://coord:5555", + tokenizer=MagicMock(), + rank=0, + server_port=0, + sock=sock, + ) + fake_mp.Process.assert_not_called() + assert tgs._SERVER_PROCESSES == [] diff --git a/tests/unit_tests/pipeline_parallel/test_schedules.py b/tests/unit_tests/pipeline_parallel/test_schedules.py index 3c76c6e8b54..31a93d8f3f5 100644 --- a/tests/unit_tests/pipeline_parallel/test_schedules.py +++ b/tests/unit_tests/pipeline_parallel/test_schedules.py @@ -1,6 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import os +from contextlib import contextmanager from types import SimpleNamespace import pytest @@ -375,6 +376,7 @@ def no_sync(): config = SimpleNamespace( overlap_p2p_comm=False, variable_seq_lengths=True, + pipeline_p2p_fixed_shape=False, finalize_model_grads_func=None, timers=None, no_sync_func=no_sync, diff --git a/tests/unit_tests/rl/conftest.py b/tests/unit_tests/rl/conftest.py new file mode 100644 index 00000000000..7a25d70d933 --- /dev/null +++ b/tests/unit_tests/rl/conftest.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import gc + +import pytest +import torch + +from tests.unit_tests.test_utilities import Utils + + +@pytest.fixture(scope="function", autouse=True) +def collect_cuda_graphs_before_group_teardown(): + """Free CUDA graphs held only by reference cycles before retired groups are destroyed. + + NCCL will not destroy a communicator while a CUDA graph that captured its + collectives is still alive: ncclCommDestroy waits for every such graph to be + freed. RL tests such as test_get_logprobs_cuda_graphs leave their graphs + reachable only through reference cycles, so the suite-wide + destroy_retired_process_groups fixture would hang without a collection first. + This fixture's teardown runs before that one's. It is scoped to this bucket so + other buckets keep their existing object lifetimes. + """ + yield + if Utils.retired_process_groups and torch.distributed.is_initialized(): + gc.collect() diff --git a/tests/unit_tests/test_training.py b/tests/unit_tests/test_training.py index c6260798e81..2e0bc6f345d 100644 --- a/tests/unit_tests/test_training.py +++ b/tests/unit_tests/test_training.py @@ -140,6 +140,7 @@ def test_training_log_resets_first_iteration_when_log_interval_is_one(monkeypatc mtp_num_layers=None, num_experts=None, perform_rl_step=False, + profile_ranks=[], record_memory_history=False, rl_profile=False, rl_use_sequence_packing=False, diff --git a/tests/unit_tests/test_utilities.py b/tests/unit_tests/test_utilities.py index 0ff22fefb5f..29b2224ad6c 100644 --- a/tests/unit_tests/test_utilities.py +++ b/tests/unit_tests/test_utilities.py @@ -46,6 +46,10 @@ class Utils: rank = int(os.environ.get('LOCAL_RANK', '0')) inited = False store = None + # Model-parallel process groups dropped by destroy_model_parallel(). The + # destroy_retired_process_groups fixture in conftest.py destroys them once the + # test finishes. + retired_process_groups = [] @staticmethod def initialize_distributed(): @@ -114,10 +118,65 @@ def destroy_model_parallel(): except Exception: Utils.inited = False return + # Drop module-level caches that hold model-parallel process groups before + # those groups are destroyed, so the next test cannot reuse an aborted + # communicator. + from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer + + reset_hybrid_ep_buffer() + try: + from transformer_engine.pytorch.fp8 import FP8GlobalStateManager + + FP8GlobalStateManager.reset() + except ImportError: + pass + # ps.destroy_model_parallel() only drops its references to the groups it + # created. Retire them so they are destroyed when the test finishes; + # otherwise every test leaks its NCCL communicators and their threads, and + # a long bucket exhausts the process's memory-map limit. They cannot be + # destroyed here: a test may re-initialize model parallelism and keep using + # a group it captured earlier, for example in a ProcessGroupCollection. + Utils.retired_process_groups.extend( + g for g in (ps._global_process_group_list or []) if g is not None + ) ps.destroy_model_parallel() Utils.inited = False torch.cuda.memory.empty_cache() + @staticmethod + def destroy_retired_process_groups(): + """Destroy the process groups retired by destroy_model_parallel().""" + groups, Utils.retired_process_groups = Utils.retired_process_groups, [] + if not groups or not torch.distributed.is_initialized(): + return + if torch.cuda.is_available(): + torch.cuda.synchronize() + Utils.clear_dtensor_sharding_caches() + live_groups = torch.distributed.distributed_c10d._world.pg_map + for group in groups: + if group in live_groups: + torch.distributed.destroy_process_group(group) + + @staticmethod + def clear_dtensor_sharding_caches(): + """Drop DTensor's cached sharding propagation results. + + DTensor caches output shardings keyed on value-equal DeviceMeshes, so a later + test that builds a mesh with the same layout can receive placements that still + reference an earlier test's mesh, and through it a destroyed process group. + Both the Python and the native caches are per thread; this clears the calling + thread's, which is where tests build and dispatch their DTensors. + """ + try: + from torch.distributed.tensor import DTensor + except ImportError: + return + propagate = DTensor._op_dispatcher.sharding_propagator.propagate_op_sharding + getattr(propagate, "cache", propagate).cache_clear() + clear_native_cache = getattr(torch._C, "_clear_DTensor_sharding_propagator_cache", None) + if clear_native_cache is not None: + clear_native_cache() + @staticmethod def initialize_model_parallel( tensor_model_parallel_size=1, diff --git a/tests/unit_tests/transformer/moe/test_paged_stash_runner.py b/tests/unit_tests/transformer/moe/test_paged_stash_runner.py index aee4203cb1b..b35e6e5ca2a 100644 --- a/tests/unit_tests/transformer/moe/test_paged_stash_runner.py +++ b/tests/unit_tests/transformer/moe/test_paged_stash_runner.py @@ -8,7 +8,12 @@ def _config(moe_paged_stash): - return SimpleNamespace(moe_paged_stash=moe_paged_stash, moe_expert_rank_capacity_factor=1.5) + return SimpleNamespace( + moe_paged_stash=moe_paged_stash, + moe_expert_rank_capacity_factor=1.5, + cuda_graph_impl="none", + cuda_graph_modules=[], + ) class _FakeTokenDispatcher: diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 61d61c9a09d..f9e6edea792 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -256,8 +256,8 @@ def forward(self, hidden_states, **_kwargs): discarded_outputs = [] class RecordingCheckpoint: - def __init__(self, fp8=False, ckpt_manager=None): - del fp8 + def __init__(self, fp8=False, ckpt_manager=None, retain_input_tensors=False): + del fp8, retain_input_tensors seen_managers.append(ckpt_manager) def checkpoint(self, function, *args): @@ -967,6 +967,7 @@ def setup_method(self, method): def teardown_method(self, method): Utils.destroy_model_parallel() + _reset_cudagraph_state() def _create_mhc_layer(self, hidden_size=64, num_streams=4, **extra_config): config = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams, **extra_config) @@ -1771,6 +1772,7 @@ def test_moe_router_synchronizes_host_outputs_and_reuses_event(): ) class TestTransformerLayerCudaGraphManagers: def setup_method(self, method): + _reset_cudagraph_state() initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) Utils.initialize_model_parallel(1, 1) model_parallel_cuda_manual_seed(123)