diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index bd598bb557a..e4df624a2cb 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -1,10 +1,28 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. from dataclasses import dataclass +from typing import Mapping, MutableMapping import torch import torch.distributed as dist from torch import Tensor +CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX = "_packed_seq_params_" + +PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS = ( + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", +) + +PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS = ( + "qkv_format", + "max_seqlen_q", + "max_seqlen_kv", + "local_cp_size", + "cp_group", +) + @dataclass class PackedSeqParams: @@ -65,3 +83,90 @@ def __post_init__(self): .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) + + +def _cuda_graph_packed_seq_params_key(field_name: str, prefix: str) -> str: + return f"{prefix}{field_name}" + + +def split_packed_seq_params_for_cuda_graph( + packed_seq_params: PackedSeqParams | None, prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX +) -> tuple[dict[str, Tensor | None], dict[str, object]]: + """Split ``PackedSeqParams`` into graph Tensor inputs and static metadata. + + Transformer Engine CUDA graph inputs must be tensors or ``None``. ``PackedSeqParams`` mixes + dynamic Tensor fields, such as cumulative sequence lengths, with static metadata, such as THD + format and max sequence lengths. This helper keeps only the fields TE attention consumes; + Mamba-only fields such as ``total_tokens`` and ``seq_idx`` stay outside this graph boundary. + """ + if packed_seq_params is None: + return {}, {} + + tensor_kwargs = {} + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS: + value = getattr(packed_seq_params, field_name) + if value is not None and not isinstance(value, Tensor): + raise TypeError( + f"PackedSeqParams.{field_name} must be a Tensor or None for CUDA graphs, " + f"got {type(value).__name__}." + ) + if value is not None: + tensor_kwargs[_cuda_graph_packed_seq_params_key(field_name, prefix)] = value + + static_metadata = {} + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS: + value = getattr(packed_seq_params, field_name) + if isinstance(value, Tensor): + raise TypeError( + f"PackedSeqParams.{field_name} is static CUDA graph metadata and must not be " + "a Tensor." + ) + static_metadata[field_name] = value + + return tensor_kwargs, static_metadata + + +def has_packed_seq_params_cuda_graph_kwargs( + kwargs: Mapping[str, object], prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX +) -> bool: + """Return whether ``kwargs`` contains flattened ``PackedSeqParams`` Tensor fields.""" + return any( + _cuda_graph_packed_seq_params_key(field_name, prefix) in kwargs + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS + ) + + +def build_packed_seq_params_from_cuda_graph_kwargs( + kwargs: MutableMapping[str, object], + static_metadata: Mapping[str, object] | None, + prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX, + remove_from_kwargs: bool = True, +) -> PackedSeqParams | None: + """Rebuild ``PackedSeqParams`` from flattened CUDA graph kwargs. + + Args: + kwargs: Graph kwargs that may contain flattened packed-sequence Tensor fields. + static_metadata: Non-Tensor metadata produced by + :func:`split_packed_seq_params_for_cuda_graph`. + prefix: Prefix used for flattened Tensor fields. + remove_from_kwargs: Whether to pop consumed flattened fields from ``kwargs``. + """ + packed_seq_params_kwargs = dict(static_metadata or {}) + found_tensor_field = False + for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS: + key = _cuda_graph_packed_seq_params_key(field_name, prefix) + if key not in kwargs: + continue + found_tensor_field = True + value = kwargs.pop(key) if remove_from_kwargs else kwargs[key] + if value is not None and not isinstance(value, Tensor): + raise TypeError( + f"Flattened PackedSeqParams field {key} must be a Tensor or None, " + f"got {type(value).__name__}." + ) + packed_seq_params_kwargs[field_name] = value + + if not packed_seq_params_kwargs and not found_tensor_field: + return None + + return PackedSeqParams(**packed_seq_params_kwargs) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 210f39fa217..d135ebca636 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -21,6 +21,7 @@ from torch.utils._pytree import tree_map as tree_map_pyt from megatron.core.num_microbatches_calculator import get_num_microbatches +from megatron.core.packed_seq_params import split_packed_seq_params_for_cuda_graph from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import ( CudaRNGStatesTracker, @@ -1668,6 +1669,31 @@ def _layer_is_graphable(layer, config): return False +def _add_packed_seq_params_to_te_cuda_graph_sample_kwargs( + layer, sample_kwargs, sample_packed_seq_params +): + """Add flattened ``PackedSeqParams`` Tensor inputs to TE graph sample kwargs.""" + if sample_packed_seq_params is None: + return + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph( + sample_packed_seq_params + ) + duplicate_keys = set(sample_kwargs) & set(tensor_kwargs) + assert not duplicate_keys, ( + "PackedSeqParams CUDA graph Tensor kwargs overlap with existing sample kwargs: " + f"{', '.join(sorted(duplicate_keys))}." + ) + assert hasattr(layer, '_set_te_cuda_graph_packed_seq_params_static_metadata'), ( + "Transformer layers using TE CUDA graph packed sequence samples must support " + "PackedSeqParams static metadata." + ) + layer._set_te_cuda_graph_packed_seq_params_static_metadata( + static_metadata, tensor_kwargs.keys() + ) + sample_kwargs.update(tensor_kwargs) + + class TECudaGraphHelper: """ Helper class to capture CUDA Graphs using TE make_graphed_callables(). @@ -1678,7 +1704,14 @@ class TECudaGraphHelper: """ def __init__( - self, model, config, seq_length, micro_batch_size, optimizers=[], pg_collection=None + self, + model, + config, + seq_length, + micro_batch_size, + optimizers=[], + pg_collection=None, + sample_packed_seq_params=None, ): assert HAVE_TE_GRAPHS, "CUDA Graphs are not supported without TE." assert ( @@ -1692,11 +1725,16 @@ def __init__( "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) self.model = model + assert sample_packed_seq_params is None or is_te_min_version("1.10.0"), ( + "TE CUDA graph packed_seq_params support requires Transformer Engine >= 1.10.0 " + "because packed-sequence Tensor fields are passed as keyword arguments." + ) self.config = config self.seq_length = seq_length self.micro_batch_size = micro_batch_size self.optimizers = optimizers self.pg_collection = pg_collection + self.sample_packed_seq_params = sample_packed_seq_params if self.pg_collection is None: self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() self.tp_group = self.pg_collection.tp @@ -1925,6 +1963,9 @@ def get_rotary_pos_emb(transformer_module, transformer_input): rotary_pos_emb = get_rotary_pos_emb(chunk_of_the_layer, hidden_states) if rotary_pos_emb is not None: static_inputs["rotary_pos_emb"] = rotary_pos_emb + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs( + layer, static_inputs, self.sample_packed_seq_params + ) _sample_kwargs = static_inputs elif contains_self_attn: _sample_args = ( diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index f6ea382077e..0a877648a05 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -16,7 +16,13 @@ from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.dist_checkpointing.utils import apply_prefix_mapping from megatron.core.inference.utils import InferenceMode -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import ( + CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX, + PackedSeqParams, + build_packed_seq_params_from_cuda_graph_kwargs, + has_packed_seq_params_cuda_graph_kwargs, + split_packed_seq_params_for_cuda_graph, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup, make_weakref from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType @@ -1135,6 +1141,110 @@ def _get_submodules_under_cudagraphs(self): submodules += [self.mlp.shared_experts] return submodules + def _set_te_cuda_graph_packed_seq_params_static_metadata( + self, static_metadata, tensor_kwarg_names=None + ): + """Store non-Tensor ``PackedSeqParams`` metadata used during TE graph capture.""" + self._te_cuda_graph_packed_seq_params_static_metadata = dict(static_metadata) + self._te_cuda_graph_packed_seq_params_tensor_kwarg_names = ( + None if tensor_kwarg_names is None else tuple(sorted(tensor_kwarg_names)) + ) + + def _get_te_cuda_graph_packed_seq_params_static_metadata(self): + """Return the static ``PackedSeqParams`` metadata used for this TE graph.""" + return getattr(self, '_te_cuda_graph_packed_seq_params_static_metadata', None) + + def _validate_te_cuda_graph_packed_seq_params_static_metadata(self, static_metadata): + """Validate that replay uses the same static packed-sequence contract as capture.""" + expected_static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata() + assert expected_static_metadata is not None, ( + "TE CUDA graph replay received packed_seq_params, but the graph was captured without " + "packed-sequence sample inputs. Recapture the graph with matching PackedSeqParams " + "static metadata." + ) + + mismatched_fields = [] + for field_name in sorted(set(expected_static_metadata) | set(static_metadata)): + expected_value = expected_static_metadata.get(field_name) + actual_value = static_metadata.get(field_name) + if expected_value is actual_value: + continue + if expected_value != actual_value: + mismatched_fields.append(field_name) + + assert not mismatched_fields, ( + "TE CUDA graph replay received PackedSeqParams with static metadata that differs " + "from capture. Recapture the graph for changed fields: " + f"{', '.join(mismatched_fields)}." + ) + + def _get_te_cuda_graph_packed_seq_params_tensor_kwarg_names(self): + """Return flattened ``PackedSeqParams`` Tensor kwargs used for this TE graph.""" + return getattr(self, '_te_cuda_graph_packed_seq_params_tensor_kwarg_names', None) + + def _validate_te_cuda_graph_packed_seq_params_tensor_kwargs(self, tensor_kwargs): + """Validate replay uses the same flattened Tensor field set as capture.""" + expected_names = self._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names() + if expected_names is None: + return + + expected_names = set(expected_names) + actual_names = set(tensor_kwargs) + missing_names = sorted(expected_names - actual_names) + extra_names = sorted(actual_names - expected_names) + assert not missing_names and not extra_names, ( + "TE CUDA graph replay received PackedSeqParams with Tensor fields that differ " + "from capture. Recapture the graph for missing fields " + f"{missing_names} and extra fields {extra_names}." + ) + + def _rebuild_te_cuda_graph_packed_seq_params(self, kwargs): + """Rebuild ``PackedSeqParams`` from flattened TE graph capture kwargs.""" + if not has_packed_seq_params_cuda_graph_kwargs(kwargs): + return + + assert kwargs.get('packed_seq_params') is None, ( + "PackedSeqParams must be passed either as flattened TE CUDA graph kwargs or as " + "packed_seq_params, but not both." + ) + static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata() + assert static_metadata is not None, ( + "Flattened PackedSeqParams Tensor fields require static metadata captured on the " + "TransformerLayer." + ) + tensor_kwargs = { + key: value + for key, value in kwargs.items() + if key.startswith(CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX) + } + self._validate_te_cuda_graph_packed_seq_params_tensor_kwargs(tensor_kwargs) + + kwargs['packed_seq_params'] = build_packed_seq_params_from_cuda_graph_kwargs( + kwargs, static_metadata + ) + + def _flatten_te_cuda_graph_packed_seq_params(self, kwargs): + """Flatten replay-time ``PackedSeqParams`` into Tensor kwargs for TE graphs.""" + packed_seq_params = kwargs.pop('packed_seq_params', None) + expected_static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata() + if packed_seq_params is None: + assert expected_static_metadata is None, ( + "TE CUDA graph was captured with packed_seq_params, so replay must also pass " + "packed_seq_params with matching static metadata." + ) + return + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + self._validate_te_cuda_graph_packed_seq_params_static_metadata(static_metadata) + self._validate_te_cuda_graph_packed_seq_params_tensor_kwargs(tensor_kwargs) + + duplicate_keys = set(kwargs) & set(tensor_kwargs) + assert not duplicate_keys, ( + "PackedSeqParams CUDA graph Tensor kwargs overlap with existing replay kwargs: " + f"{', '.join(sorted(duplicate_keys))}." + ) + kwargs.update(tensor_kwargs) + def _te_cuda_graph_capture(self, *args, **kwargs): """ CUDA Graph capture for this layer using TE interface. @@ -1155,6 +1265,8 @@ def _te_cuda_graph_capture(self, *args, **kwargs): hidden_states = kwargs.pop("hidden_states") hidden_states = self.off_interface.backward_record(hidden_states) kwargs["hidden_states"] = hidden_states + self._rebuild_te_cuda_graph_packed_seq_params(kwargs) + context = None if ( not self.config.cuda_graph_modules @@ -1197,7 +1309,8 @@ def _te_cuda_graph_replay(self, *args, **kwargs): CUDA graph replay for this layer and microbatch `self.current_microbatch` using TE interface. TransformerEngine versions>=1.10 allow keyword arguments with CUDA graph. However, CUDA graph accepts only Tensor inputs. - Hence, `inference_context` and `packed_seq_params` are excluded from input list. + Hence, `inference_context` is excluded from input list. `packed_seq_params` is split + into Tensor graph inputs and static metadata when attention is in the graph scope. """ context = None if ( @@ -1207,12 +1320,13 @@ def _te_cuda_graph_replay(self, *args, **kwargs): hidden_states, context = self._forward_attention(*args, **kwargs) args = (hidden_states,) kwargs = {} + else: + self._flatten_te_cuda_graph_packed_seq_params(kwargs) - assert (kwargs.get('inference_context') is None) and ( - kwargs.get('packed_seq_params') is None - ), ( + assert kwargs.get('inference_context') is None, ( "CUDA graph accepts only Tensor inputs. " - "inference_context and packed_seq_params are excluded from input list. " + "inference_context is excluded from input list; packed_seq_params must be " + "flattened into Tensor kwargs with matching static metadata. " "For inference cuda graph, please use cuda_graph_impl=local instead." ) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 76791720e9d..72967b0ebf3 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -746,6 +746,34 @@ def selective_log_softmax(logits, index): return per_token_logps +def get_rl_packed_seq_params_for_cuda_graph( + seq_length: int, + device: torch.device, + sequence_packing: bool = False, + max_sequences_per_bin: int = None, +) -> PackedSeqParams: + """Build RL ``PackedSeqParams`` used to keep CUDA graph signatures stable.""" + if sequence_packing: + assert max_sequences_per_bin is not None, ( + "max_sequences_per_bin is required when sequence_packing is enabled." + ) + return get_default_packed_seq_params( + seq_length=seq_length, + max_sequences_per_bin=max_sequences_per_bin, + device=device, + ) + + cu_seqlens = torch.tensor([0, seq_length], dtype=torch.int32, device=device) + return PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=seq_length, + max_seqlen_kv=seq_length, + total_tokens=seq_length, + ) + + def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=False, packed_seq_params=None): """Get sequence logprobs from their token ids. @@ -773,22 +801,12 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa # graph signature matches the training forward_step in train_rl.py. # This is necessary because reference logprobs steps will reuse the training forward graph. if packed_seq_params is None: - if sequence_packing: - packed_seq_params = get_default_packed_seq_params( - seq_length=tokens.shape[1], - max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, - device=tokens.device, - ) - else: - cu_seqlens = torch.tensor([0, tokens.shape[1]], dtype=torch.int32, device=tokens.device) - 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], - ) + packed_seq_params = get_rl_packed_seq_params_for_cuda_graph( + seq_length=tokens.shape[1], + device=tokens.device, + sequence_packing=sequence_packing, + max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, + ) nvtx_range = get_nvtx_range() diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index ff98b0a58e2..24aafafc512 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -7,7 +7,7 @@ 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.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( diff --git a/megatron/training/training.py b/megatron/training/training.py index f67397fe889..68a64078dc5 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3625,12 +3625,21 @@ def trace_handler(p): # Initialize CUDA Graphs helper. if args.cuda_graph_impl == "transformer_engine": + cuda_graph_sample_packed_seq_params = None + if has_rl_utils and args.perform_rl_step: + cuda_graph_sample_packed_seq_params = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=args.seq_length, + device=torch.device("cuda", torch.cuda.current_device()), + sequence_packing=args.rl_use_sequence_packing, + max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, + ) cuda_graph_helper = TECudaGraphHelper( model=model, config=config, seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, optimizers=[optimizer], + sample_packed_seq_params=cuda_graph_sample_packed_seq_params, ) # Run training iterations till done. diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index dd6c85b2125..06f67bc7038 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -167,6 +167,64 @@ def create_test_args(self, **kwargs): set_global_variables(args, False) return args + def test_get_rl_packed_seq_params_for_cuda_graph_without_sequence_packing(self): + params = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=False + ) + + assert params.qkv_format == 'thd' + assert params.max_seqlen_q == 8 + assert params.max_seqlen_kv == 8 + assert params.total_tokens == 8 + assert params.cu_seqlens_kv is params.cu_seqlens_q + assert params.cu_seqlens_q.dtype == torch.int32 + assert params.cu_seqlens_q.device.type == "cpu" + assert torch.equal(params.cu_seqlens_q, torch.tensor([0, 8], dtype=torch.int32)) + assert torch.equal(params.seq_idx, torch.zeros((1, 8), dtype=torch.int32)) + + def test_get_rl_packed_seq_params_for_cuda_graph_with_sequence_packing(self): + params = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=True, max_sequences_per_bin=3 + ) + + expected_cu_seqlens = torch.tensor([0, 8, 8, 8, 8], dtype=torch.int32) + assert params.qkv_format == 'thd' + assert params.max_seqlen_q == 8 + assert params.max_seqlen_kv == 8 + assert params.total_tokens == 8 + assert params.cu_seqlens_kv is params.cu_seqlens_q + assert params.cu_seqlens_q.dtype == torch.int32 + assert params.cu_seqlens_q.device.type == "cpu" + assert params.cu_seqlens_q.shape == (5,) + assert torch.equal(params.cu_seqlens_q, expected_cu_seqlens) + assert torch.equal(params.seq_idx, torch.zeros((1, 8), dtype=torch.int32)) + + def test_get_rl_packed_seq_params_for_cuda_graph_requires_max_sequences_per_bin(self): + with pytest.raises(AssertionError, match="max_sequences_per_bin is required"): + rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=True + ) + + def test_get_rl_packed_seq_params_for_cuda_graph_edge_cases(self): + # Parametrize seq_length of 1 (single-token boundary condition) + params_single = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=1, device=torch.device("cpu"), sequence_packing=False + ) + assert params_single.max_seqlen_q == 1 + assert params_single.max_seqlen_kv == 1 + assert params_single.total_tokens == 1 + assert torch.equal(params_single.cu_seqlens_q, torch.tensor([0, 1], dtype=torch.int32)) + + # Parametrize sequence packing with max_sequences_per_bin > 1 (e.g. 4) + params_multi = rl_utils.get_rl_packed_seq_params_for_cuda_graph( + seq_length=8, device=torch.device("cpu"), sequence_packing=True, max_sequences_per_bin=4 + ) + assert params_multi.max_seqlen_q == 8 + assert params_multi.max_seqlen_kv == 8 + assert params_multi.total_tokens == 8 + assert params_multi.cu_seqlens_q.shape == (6,) + assert torch.equal(params_multi.cu_seqlens_q, torch.tensor([0, 8, 8, 8, 8, 8], dtype=torch.int32)) + def test_rl_granularity_defaults(self): args = self.create_test_args(perform_rl_step=True, grpo_prompts_per_step=8) diff --git a/tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py b/tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py new file mode 100644 index 00000000000..b1030cf95e2 --- /dev/null +++ b/tests/unit_tests/transformer/test_packed_seq_params_cuda_graph.py @@ -0,0 +1,368 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.packed_seq_params import ( + CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX, + PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS, + PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS, + PackedSeqParams, + build_packed_seq_params_from_cuda_graph_kwargs, + has_packed_seq_params_cuda_graph_kwargs, + split_packed_seq_params_for_cuda_graph, +) +from megatron.core.transformer.cuda_graphs import ( + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs, +) +from megatron.core.transformer.transformer_layer import TransformerLayer + + +class _TransformerLayerCudaGraphStub: + _set_te_cuda_graph_packed_seq_params_static_metadata = ( + TransformerLayer._set_te_cuda_graph_packed_seq_params_static_metadata + ) + _get_te_cuda_graph_packed_seq_params_static_metadata = ( + TransformerLayer._get_te_cuda_graph_packed_seq_params_static_metadata + ) + _validate_te_cuda_graph_packed_seq_params_static_metadata = ( + TransformerLayer._validate_te_cuda_graph_packed_seq_params_static_metadata + ) + _get_te_cuda_graph_packed_seq_params_tensor_kwarg_names = ( + TransformerLayer._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names + ) + _validate_te_cuda_graph_packed_seq_params_tensor_kwargs = ( + TransformerLayer._validate_te_cuda_graph_packed_seq_params_tensor_kwargs + ) + _rebuild_te_cuda_graph_packed_seq_params = ( + TransformerLayer._rebuild_te_cuda_graph_packed_seq_params + ) + _flatten_te_cuda_graph_packed_seq_params = ( + TransformerLayer._flatten_te_cuda_graph_packed_seq_params + ) + + +def _make_packed_seq_params(): + cu_seqlens = torch.IntTensor([0, 4, 9, 16]) + cu_seqlens_padded = torch.IntTensor([0, 8, 12, 16]) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=8, + max_seqlen_kv=8, + local_cp_size=1, + ) + + +def test_split_packed_seq_params_for_cuda_graph_separates_tensors_from_metadata(): + packed_seq_params = _make_packed_seq_params() + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + + assert set(static_metadata) == set(PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS) + assert static_metadata == { + "qkv_format": "thd", + "max_seqlen_q": 8, + "max_seqlen_kv": 8, + "local_cp_size": 1, + "cp_group": None, + } + assert all(not isinstance(value, torch.Tensor) for value in static_metadata.values()) + + expected_tensor_fields = { + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + } + assert set(tensor_kwargs) == { + f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}{field}" for field in expected_tensor_fields + } + assert set(PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS) >= expected_tensor_fields + for value in tensor_kwargs.values(): + assert isinstance(value, torch.Tensor) + + +def test_has_packed_seq_params_cuda_graph_kwargs_detects_flattened_fields(): + tensor_kwargs, _ = split_packed_seq_params_for_cuda_graph(_make_packed_seq_params()) + + assert has_packed_seq_params_cuda_graph_kwargs(tensor_kwargs) + assert not has_packed_seq_params_cuda_graph_kwargs({"hidden_states": torch.ones(2, 1, 4)}) + assert build_packed_seq_params_from_cuda_graph_kwargs({}, None) is None + + +def test_build_packed_seq_params_from_cuda_graph_kwargs_pops_flattened_fields(): + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + kwargs = {"hidden_states": torch.ones(2, 1, 4), **tensor_kwargs} + + rebuilt = build_packed_seq_params_from_cuda_graph_kwargs(kwargs, static_metadata) + + assert set(kwargs) == {"hidden_states"} + assert rebuilt.qkv_format == "thd" + assert rebuilt.max_seqlen_q == 8 + assert rebuilt.max_seqlen_kv == 8 + assert rebuilt.local_cp_size == 1 + assert rebuilt.cp_group is None + assert rebuilt.total_tokens is None + assert rebuilt.seq_idx is None + assert torch.equal(rebuilt.cu_seqlens_q, packed_seq_params.cu_seqlens_q) + assert torch.equal(rebuilt.cu_seqlens_kv, packed_seq_params.cu_seqlens_kv) + assert torch.equal(rebuilt.cu_seqlens_q_padded, packed_seq_params.cu_seqlens_q_padded) + assert torch.equal(rebuilt.cu_seqlens_kv_padded, packed_seq_params.cu_seqlens_kv_padded) + + +def test_build_packed_seq_params_from_cuda_graph_kwargs_can_keep_kwargs_intact(): + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph( + _make_packed_seq_params() + ) + kwargs = dict(tensor_kwargs) + + build_packed_seq_params_from_cuda_graph_kwargs( + kwargs, static_metadata, remove_from_kwargs=False + ) + + assert kwargs == tensor_kwargs + + +def test_split_packed_seq_params_for_cuda_graph_rejects_static_tensor_metadata(): + packed_seq_params = _make_packed_seq_params() + packed_seq_params.max_seqlen_q = torch.IntTensor([8]) + + with pytest.raises(TypeError, match="max_seqlen_q"): + split_packed_seq_params_for_cuda_graph(packed_seq_params) + + +def test_split_packed_seq_params_for_cuda_graph_ignores_mamba_only_fields(): + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=torch.IntTensor([0, 2, 5]), + cu_seqlens_kv=torch.IntTensor([0, 2, 5]), + max_seqlen_q=3, + max_seqlen_kv=3, + total_tokens=5, + ) + assert packed_seq_params.seq_idx is not None + + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + + assert f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}seq_idx" not in tensor_kwargs + assert "total_tokens" not in static_metadata + + +def test_transformer_layer_rebuilds_flattened_cuda_graph_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + kwargs = {"hidden_states": torch.ones(2, 1, 4), **tensor_kwargs} + + layer._rebuild_te_cuda_graph_packed_seq_params(kwargs) + + assert set(kwargs) == {"hidden_states", "packed_seq_params"} + rebuilt = kwargs["packed_seq_params"] + assert rebuilt.qkv_format == "thd" + assert rebuilt.max_seqlen_q == 8 + assert rebuilt.max_seqlen_kv == 8 + assert torch.equal(rebuilt.cu_seqlens_q, packed_seq_params.cu_seqlens_q) + assert torch.equal(rebuilt.cu_seqlens_kv, packed_seq_params.cu_seqlens_kv) + + +def test_transformer_layer_flattens_replay_time_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + attention_mask = torch.zeros(1, 1, 16, 16, dtype=torch.bool) + kwargs = {"attention_mask": attention_mask, "packed_seq_params": packed_seq_params} + + layer._flatten_te_cuda_graph_packed_seq_params(kwargs) + + assert kwargs["attention_mask"] is attention_mask + assert "packed_seq_params" not in kwargs + assert set(tensor_kwargs).issubset(kwargs) + for key, value in tensor_kwargs.items(): + assert kwargs[key] is value + + +def test_transformer_layer_rejects_replay_without_captured_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + _, static_metadata = split_packed_seq_params_for_cuda_graph(_make_packed_seq_params()) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata) + + with pytest.raises(AssertionError, match="captured with packed_seq_params"): + layer._flatten_te_cuda_graph_packed_seq_params({"hidden_states": torch.ones(2, 1, 4)}) + + +def test_transformer_layer_rejects_changed_packed_seq_params_static_metadata(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + _, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata) + packed_seq_params.max_seqlen_q = 4 + + with pytest.raises(AssertionError, match="max_seqlen_q"): + layer._flatten_te_cuda_graph_packed_seq_params({"packed_seq_params": packed_seq_params}) + + +def test_transformer_layer_rejects_changed_packed_seq_params_tensor_fields(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + packed_seq_params.cu_seqlens_q_padded = None + + with pytest.raises(AssertionError, match="Tensor fields"): + layer._flatten_te_cuda_graph_packed_seq_params({"packed_seq_params": packed_seq_params}) + + +def test_transformer_layer_rejects_replay_with_overlapping_flattened_kwargs(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + existing_key = f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}cu_seqlens_q" + + with pytest.raises(AssertionError, match="overlap"): + layer._flatten_te_cuda_graph_packed_seq_params( + {existing_key: torch.IntTensor([0]), "packed_seq_params": packed_seq_params} + ) + + +def test_te_cuda_graph_sample_kwargs_include_flattened_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + expected_tensor_kwargs, expected_static_metadata = split_packed_seq_params_for_cuda_graph( + packed_seq_params + ) + attention_mask = torch.zeros(1, 1, 16, 16, dtype=torch.bool) + sample_kwargs = {"attention_mask": attention_mask} + + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs(layer, sample_kwargs, packed_seq_params) + + assert sample_kwargs["attention_mask"] is attention_mask + assert set(expected_tensor_kwargs).issubset(sample_kwargs) + for key, value in expected_tensor_kwargs.items(): + assert sample_kwargs[key] is value + assert layer._get_te_cuda_graph_packed_seq_params_static_metadata() == expected_static_metadata + assert layer._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names() == tuple( + sorted(expected_tensor_kwargs) + ) + + +def test_te_cuda_graph_sample_kwargs_noop_without_packed_seq_params(): + layer = _TransformerLayerCudaGraphStub() + attention_mask = torch.zeros(1, 1, 16, 16, dtype=torch.bool) + sample_kwargs = {"attention_mask": attention_mask} + + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs(layer, sample_kwargs, None) + + assert sample_kwargs == {"attention_mask": attention_mask} + assert layer._get_te_cuda_graph_packed_seq_params_static_metadata() is None + + +def test_te_cuda_graph_sample_kwargs_reject_overlapping_flattened_keys(): + layer = _TransformerLayerCudaGraphStub() + packed_seq_params = _make_packed_seq_params() + sample_kwargs = {f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}cu_seqlens_q": torch.IntTensor([0])} + + with pytest.raises(AssertionError, match="overlap"): + _add_packed_seq_params_to_te_cuda_graph_sample_kwargs( + layer, sample_kwargs, packed_seq_params + ) + + +def test_te_cuda_graph_partial_attn_only_flow(): + from megatron.core.transformer.enums import CudaGraphModule + + class _ConfigStub: + def __init__(self, cuda_graph_modules): + self.cuda_graph_modules = cuda_graph_modules + self.delay_offload_until_cuda_graph = False + + class _TestLayer(_TransformerLayerCudaGraphStub): + _te_cuda_graph_replay = TransformerLayer._te_cuda_graph_replay + + def __init__(self, cuda_graph_modules): + self.config = _ConfigStub(cuda_graph_modules) + self.attn_called = False + self.replay_impl_called = False + self.replay_impl_args = None + self.replay_impl_kwargs = None + self.replay_impl_context = None + + def _forward_attention(self, *args, **kwargs): + self.attn_called = True + return torch.ones(2, 1, 4) * 2.0, "attn_context" + + def _te_cuda_graph_replay_impl(self, args, kwargs, context): + self.replay_impl_called = True + self.replay_impl_args = args + self.replay_impl_kwargs = kwargs + self.replay_impl_context = context + return torch.ones(2, 1, 4) * 3.0 + + # Case 1: When CudaGraphModule.attn is captured + layer_attn = _TestLayer([CudaGraphModule.attn]) + packed_seq_params = _make_packed_seq_params() + tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params) + layer_attn._set_te_cuda_graph_packed_seq_params_static_metadata(static_metadata, tensor_kwargs) + + kwargs = {"packed_seq_params": packed_seq_params, "hidden_states": torch.ones(2, 1, 4)} + layer_attn._te_cuda_graph_replay(**kwargs) + + assert not layer_attn.attn_called + assert layer_attn.replay_impl_called + assert layer_attn.replay_impl_context is None + assert "packed_seq_params" not in layer_attn.replay_impl_kwargs + assert f"{CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX}cu_seqlens_q" in layer_attn.replay_impl_kwargs + + # Case 2: When CudaGraphModule.attn is NOT captured (e.g. only mlp is captured) + layer_mlp = _TestLayer([CudaGraphModule.mlp]) + + kwargs = {"packed_seq_params": packed_seq_params, "hidden_states": torch.ones(2, 1, 4)} + layer_mlp._te_cuda_graph_replay(**kwargs) + + assert layer_mlp.attn_called + assert layer_mlp.replay_impl_called + assert layer_mlp.replay_impl_context == "attn_context" + assert len(layer_mlp.replay_impl_args) == 1 + assert torch.equal(layer_mlp.replay_impl_args[0], torch.ones(2, 1, 4) * 2.0) + assert layer_mlp.replay_impl_kwargs == {} + + +def test_seq_idx_determinism_across_replays(): + cu_seqlens = torch.IntTensor([0, 3, 7, 10]) + cu_seqlens_padded = torch.IntTensor([0, 4, 8, 12]) + + params1 = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=4, + max_seqlen_kv=4, + total_tokens=10, + ) + + params2 = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=4, + max_seqlen_kv=4, + total_tokens=10, + ) + + assert params1.seq_idx is not None + assert params2.seq_idx is not None + assert torch.equal(params1.seq_idx, params2.seq_idx) + assert params1.seq_idx.shape == params2.seq_idx.shape + assert params1.seq_idx.dtype == torch.int32 + diff --git a/train_rl.py b/train_rl.py index acf54680f4a..8d5b66369d4 100644 --- a/train_rl.py +++ b/train_rl.py @@ -20,6 +20,7 @@ get_logprobs, get_rl_runtime_state, load_packed_data_by_index, + get_rl_packed_seq_params_for_cuda_graph, ) from megatron.training import get_args, get_timers, pretrain, print_rank_0 from megatron.training.utils import is_hybrid_model @@ -27,9 +28,6 @@ from megatron.training.argument_utils import gpt_config_from_args, hybrid_config_from_args, pretrain_cfg_container_from_args from model_provider import model_provider -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.rl.sequence_packing_utils import get_default_packed_seq_params - stimer = StragglerDetector() import logging @@ -260,22 +258,12 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): model_to_use = model[0] if isinstance(model, list) else model if packed_seq_params is None: - if args.rl_use_sequence_packing: - packed_seq_params = get_default_packed_seq_params( - seq_length=tokens.shape[1], - max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, - device=tokens.device, - ) - else: - cu_seqlens = torch.tensor([0, tokens.shape[1]], dtype=torch.int32, device=tokens.device) - 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], - ) + packed_seq_params = get_rl_packed_seq_params_for_cuda_graph( + seq_length=tokens.shape[1], + device=tokens.device, + sequence_packing=args.rl_use_sequence_packing, + max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin, + ) # Clear RoPE cache to avoid inference tensor errors try: