From 8892dc4dd2b391abc32bdc35e4e2bc2fb85b526d Mon Sep 17 00:00:00 2001 From: Li Tao Date: Thu, 20 Aug 2026 08:01:44 -0700 Subject: [PATCH 1/5] feat(mdp): greedy token-budget packing and fixed-shape THD batches Give MDP's decoder data path two independent, opt-in properties. --mdp-greedy-packing fills a fixed num_microbatches bins to a token budget of max_seqlen_per_dp_cp_rank x cp_size, consuming as many samples as that takes, instead of a fixed --micro-batch-size count. This decouples padding waste from MBS while keeping num_microbatches static, which the PP schedule, the VPP replay cursors, and per-layer CUDA-graph slot sizing all depend on. The grouping rule is DpBalancedScheduler's in-order greedy fill; MDP cannot use that scheduler itself, because it asserts on GPT-only sample keys, drops pixel_values / image_grid_thw, and reroutes samples across DP without pixel awareness, so that combination is now rejected with that reason. --micro-batch-size and --global-batch-size become pure bin-count knobs, and consumed_train_samples is reported from a real all-reduced count. --thd-static-packing makes the collator emit exactly max_seqlen_per_dp_cp_rank x cp_size rows with cu_seqlens* padded to thd_max_packed_sequences + 1 entries -- the same shape contract --sequence-packing-scheduler gives the decoder-only path. The tail is an append_dummy_seq sequence; the pre-tail cu_seqlens is emitted alongside so the FLOPs accumulator keeps excluding pad tokens. --mdp-mock-dataset-config-json gives the mock dataset a controllable sequence-length distribution, reusing --varlen-mock-dataset-config-json's schema and MockSFTLowLevelDataset, seeded so every rank rebuilds an identical scenario pool. The default pool is byte-identical. Verified on 4xGB300: with a degenerate min=max=mean=1024 distribution and a 4096 budget, greedy and greedy+static both reproduce the fixed-MBS=4 baseline loss and grad-norm trajectory bit-for-bit. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Li Tao --- examples/multimodal_dev/arguments.py | 34 ++ examples/multimodal_dev/data/mdp_mock.py | 57 ++- examples/multimodal_dev/data/mdp_scenarios.py | 115 +++++- examples/multimodal_dev/forward_step.py | 119 +++++- .../tests/test_flops_accounting.py | 8 +- .../tests/test_thd_static_packing.py | 341 ++++++++++++++++++ megatron/core/mdp/README.md | 1 + megatron/core/mdp/config.py | 75 ++++ megatron/core/mdp/integration.py | 16 + megatron/core/mdp/knowledge.md | 58 ++- megatron/core/mdp/packing.py | 214 +++++++++++ megatron/core/mdp/runtime.py | 68 ++++ megatron/core/model_parallel_config.py | 33 ++ megatron/core/packed_seq_params.py | 78 ++++ .../core/transformer/transformer_config.py | 7 + megatron/training/training.py | 43 +++ tests/unit_tests/mdp/test_config.py | 63 ++++ tests/unit_tests/mdp/test_packing.py | 157 ++++++++ 18 files changed, 1451 insertions(+), 36 deletions(-) create mode 100644 examples/multimodal_dev/tests/test_thd_static_packing.py create mode 100644 megatron/core/mdp/packing.py create mode 100644 tests/unit_tests/mdp/test_packing.py diff --git a/examples/multimodal_dev/arguments.py b/examples/multimodal_dev/arguments.py index 5716cb6d2ee..b35ef80c27d 100644 --- a/examples/multimodal_dev/arguments.py +++ b/examples/multimodal_dev/arguments.py @@ -168,6 +168,40 @@ def add_multimodal_args(parser): "pixel traffic." ), ) + group.add_argument( + "--mdp-greedy-packing", + action="store_true", + default=False, + help=( + "Fill each decoder microbatch to a token budget " + "(--max-seqlen-per-dp-cp-rank x CP) by consuming as many samples as " + "it takes, instead of a fixed --micro-batch-size count. " + "IMPORTANT: this REINTERPRETS --micro-batch-size and " + "--global-batch-size. They no longer describe what goes into a " + "microbatch; they only set the number of bins per iteration " + "(N = GBS / (MBS x DP)). The sample count per iteration then floats, " + "so --global-batch-size means 'N x token budget' and loss curves are " + "not iteration-by-iteration comparable against a fixed-GBS run. " + "Requires --max-seqlen-per-dp-cp-rank. Independent of " + "--thd-static-packing." + ), + ) + group.add_argument( + "--mdp-mock-dataset-config-json", + type=str, + default=None, + help=( + "Sequence-length distribution for the MDP mock dataset, as JSON or " + "a path to a JSON file. Same schema as " + "--varlen-mock-dataset-config-json, e.g. " + '\'{"mode":"distribution","type":"lognormal","min_seq_len":512,' + '"max_seq_len":4096,"mean_seq_len":2048,"lognormal_sigma":1.1}\'. ' + "A dedicated flag because --varlen-mock-dataset-config-json is only " + "honored under --use-varlen-dataset, which auto-sets the packing " + "scheduler MDP must not have. Unset keeps the built-in " + "[1000, 2000] uniform range." + ), + ) group.add_argument( "--mdp-debug-plan-payload-check", action="store_true", diff --git a/examples/multimodal_dev/data/mdp_mock.py b/examples/multimodal_dev/data/mdp_mock.py index 046063fe555..d0d888e020f 100644 --- a/examples/multimodal_dev/data/mdp_mock.py +++ b/examples/multimodal_dev/data/mdp_mock.py @@ -17,12 +17,13 @@ routing can be verified element by element. """ +import math from typing import Optional, Sequence import torch from torch.utils.data import Dataset -from examples.multimodal_dev.data.mdp_scenarios import build_scenarios +from examples.multimodal_dev.data.mdp_scenarios import build_scenarios, scenario_totals from examples.multimodal_dev.models.qwen35_vl.configuration import ( QWEN35_VL_IMAGE_TOKEN_ID, QWEN35_VL_VISION_START_TOKEN_ID, @@ -63,6 +64,10 @@ class MdpThdMockDataset(Dataset): ``(grids, text_chunk_lengths)`` with ``len(text_chunks) == len(grids) + 1`` (text before/between/after vision blocks; the text-only scenario uses one chunk). + length_config: Optional ``--mdp-mock-dataset-config-json`` payload + controlling the per-sample *total* token length distribution + (same schema as ``--varlen-mock-dataset-config-json``). Ignored + when ``scenarios`` is given. """ def __init__( @@ -76,6 +81,7 @@ def __init__( spatial_merge_size: int = 2, seed: int = 1234, scenarios: Optional[Sequence] = None, + length_config: Optional[dict] = None, ): self.num_samples = num_samples self.vocab_size = vocab_size @@ -85,7 +91,14 @@ def __init__( self.temporal_patch_size = temporal_patch_size self.spatial_merge_size = spatial_merge_size self.seed = seed - self.scenarios = tuple(scenarios) if scenarios is not None else _SCENARIOS + if scenarios is not None: + self.scenarios = tuple(scenarios) + elif length_config is not None: + # Rebuilt independently on every rank; build_scenarios is seeded + # from the fixed GENERATOR_SEED so the pools stay identical. + self.scenarios = build_scenarios(length_config=length_config) + else: + self.scenarios = _SCENARIOS self.pixel_dim = 3 * temporal_patch_size * patch_size * patch_size for grids, text_chunks in self.scenarios: expected_chunks = len(grids) + 1 if grids else 1 @@ -165,16 +178,54 @@ def _text(length): } +#: Extra margin on top of the computed greedy sample requirement. The mean +#: sample length only predicts the *average* bin occupancy; individual bins run +#: short or long, and the shortfall compounds across iterations. +GREEDY_SAMPLE_SAFETY = 1.5 + + +def _greedy_sample_scale(args, length_config): + """Scale factor for the synthetic dataset length under greedy packing. + + Megatron sizes the dataset as ``train_iters * global_batch_size`` samples. + Under ``--mdp-greedy-packing`` that is a *bin* count, not a sample count: + each bin swallows roughly ``token_budget / mean_sample_len`` samples instead + of exactly ``micro_batch_size``. Whenever the mean sample is shorter than + ``token_budget / micro_batch_size`` the stream runs dry mid-run. The mock + dataset is a pure function of the sample index, so enlarging it costs + nothing and changes no sample's content. + + Returns 1.0 when greedy packing is off, so the default path is unchanged. + """ + if not getattr(args, "mdp_greedy_packing", False): + return 1.0 + budget = int(args.max_seqlen_per_dp_cp_rank) * int(args.context_parallel_size) + cap = getattr(args, "thd_max_packed_sequences", None) + pool = build_scenarios(length_config=length_config) + mean_len = sum(scenario_totals(s)[0] for s in pool) / len(pool) + samples_per_bin = budget / mean_len + if cap: + samples_per_bin = min(samples_per_bin, float(cap)) + return max(1.0, samples_per_bin / int(args.micro_batch_size)) * GREEDY_SAMPLE_SAFETY + + def train_valid_test_datasets_provider(train_val_test_num_samples): """Provide MDP mock train / val / test datasets.""" from megatron.training import get_args args = get_args() + length_config = getattr(args, "mdp_mock_dataset_config_json", None) + if length_config is not None: + from megatron.training.datasets.utils import load_json_arg + + length_config = load_json_arg(length_config) kwargs = dict( vocab_size=getattr(args, "padded_vocab_size", 1024), image_token_id=getattr(args, "image_token_id", QWEN35_VL_IMAGE_TOKEN_ID), + length_config=length_config, ) + scale = _greedy_sample_scale(args, length_config) return tuple( - MdpThdMockDataset(num_samples=n, seed=1234 + split, **kwargs) + MdpThdMockDataset(num_samples=math.ceil(n * scale), seed=1234 + split, **kwargs) for split, n in enumerate(train_val_test_num_samples) ) diff --git a/examples/multimodal_dev/data/mdp_scenarios.py b/examples/multimodal_dev/data/mdp_scenarios.py index 1cef4d85f14..3c0fe31d968 100644 --- a/examples/multimodal_dev/data/mdp_scenarios.py +++ b/examples/multimodal_dev/data/mdp_scenarios.py @@ -34,6 +34,40 @@ TOTAL_TOKENS_RANGE = (1000, 2000) IMAGE_TOKENS_RANGE = (100, 500) TEXT_ONLY_EVERY = 5 # every 5th scenario is text-only (~1/5 of the pool) +MIN_TOTAL_TOKENS = 64 # floor for externally supplied length distributions +MAX_GRID_DRAWS = 200 # rejection-sampling cap; see _multimodal_scenario + + +def draw_total_token_lengths(length_config, count, seed=GENERATOR_SEED): + """``count`` per-sample total token lengths from a length distribution. + + ``length_config`` is the ``--mdp-mock-dataset-config-json`` payload, the same + schema ``--varlen-mock-dataset-config-json`` uses, so MDP and the + decoder-only reference can be pointed at identical distributions. + + Determinism matters more than usual here: every MDP rank rebuilds the + scenario pool independently, so a divergent draw becomes a collective hang, + not a wrong number. Two things guarantee it. ``MockSFTLowLevelDataset`` + seeds NumPy's global RNG with its own fixed class seed, so its length vector + is identical everywhere; and the *selection* out of that vector runs on a + private ``random.Random(seed)`` rather than on whatever global state the + process happens to hold. The global NumPy state is saved and restored so the + draw leaves no side effect on other datasets. + """ + import numpy as np + + from megatron.training.datasets.sft_dataset import MockSFTLowLevelDataset + + state = np.random.get_state() + try: + low_level = MockSFTLowLevelDataset(**dict(length_config)) + finally: + np.random.set_state(state) + lengths = low_level.sequence_lengths + rng = random.Random(seed) + return [ + max(MIN_TOTAL_TOKENS, int(lengths[rng.randrange(len(lengths))])) for _ in range(count) + ] def _item_tokens(grid, merge=SPATIAL_MERGE_SIZE): @@ -49,6 +83,11 @@ def _random_grid(rng, max_tokens): """ # ~25% of items are multi-frame video (t in 2-3), rest still images. t = rng.choice((2, 3)) if rng.random() < 0.25 else 1 + # A short budget cannot carry a 2x2 merged grid per frame. Clamp after the + # draw rather than before it, so the rng stream is unchanged and only + # otherwise-impossible cases behave differently. + while t > 1 and max_tokens // t < 4: + t -= 1 # Choose merged spatial extent (h/2, w/2) so tokens = t * mh * mw fills the # share as closely as the factorization allows without exceeding it. budget = max_tokens // t @@ -58,14 +97,34 @@ def _random_grid(rng, max_tokens): return (t, mh * SPATIAL_MERGE_SIZE, mw * SPATIAL_MERGE_SIZE) -def _multimodal_scenario(rng): - total_target = rng.randint(*TOTAL_TOKENS_RANGE) +def _image_token_bounds(total_target): + """Image-token window for a sample of ``total_target`` tokens. + + Identical to ``IMAGE_TOKENS_RANGE`` for the built-in [1000, 2000] pool; the + clamps only bind for the short samples an external length distribution can + produce, where a fixed 100-500 image budget would leave no room for text. + """ + high = max(4, min(IMAGE_TOKENS_RANGE[1], total_target // 2)) + low = max(4, min(IMAGE_TOKENS_RANGE[0], high)) + return low, high + + +def _multimodal_scenario(rng, total_target): + image_low, image_high = _image_token_bounds(total_target) # Rejection-sample grids until the image token total lands in range (the # grid factorization can undershoot a share, so one draw is not enough). + # The window is a *quality* constraint -- text fills whatever remains, so + # any draw yields an exact total. Short samples from an external length + # distribution leave a narrow window, so after MAX_GRID_DRAWS attempts any + # structurally valid draw is accepted rather than looping. + attempts = 0 while True: - image_target = rng.randint(*IMAGE_TOKENS_RANGE) - num_items = rng.randint(1, 3) + attempts += 1 + image_target = rng.randint(image_low, image_high) + # Every item needs at least a 2x2 merged grid, so the item count is + # bounded by the drawn budget (always 3 for the built-in pool). + num_items = rng.randint(1, max(1, min(3, image_target // 4))) grids = [] remaining = image_target for i in range(num_items): @@ -78,12 +137,14 @@ def _multimodal_scenario(rng): grids.append(_random_grid(rng, share)) remaining -= _item_tokens(grids[-1]) image_tokens = sum(_item_tokens(g) for g in grids) - if IMAGE_TOKENS_RANGE[0] <= image_tokens <= IMAGE_TOKENS_RANGE[1]: + # k items need k+1 text chunks, each at least 1 token. + text_total = total_target - image_tokens - num_items + if text_total < num_items + 1: + continue + if image_low <= image_tokens <= image_high or attempts >= MAX_GRID_DRAWS: break # Text fills the remainder: total = text + image_tokens + num_items sentinels. - text_total = total_target - image_tokens - num_items - # k items need k+1 text chunks, each at least 1 token. cuts = sorted(rng.sample(range(1, text_total), num_items)) bounds = [0] + cuts + [text_total] text_chunks = tuple(bounds[i + 1] - bounds[i] for i in range(num_items + 1)) @@ -91,8 +152,8 @@ def _multimodal_scenario(rng): return (tuple(grids), text_chunks) -def _text_only_scenario(rng): - return ((), (rng.randint(*TOTAL_TOKENS_RANGE),)) +def _text_only_scenario(rng, total_target): + return ((), (total_target,)) def scenario_totals(scenario, merge=SPATIAL_MERGE_SIZE): @@ -103,29 +164,47 @@ def scenario_totals(scenario, merge=SPATIAL_MERGE_SIZE): return total, image_tokens -def build_scenarios(pool_size=POOL_SIZE, seed=GENERATOR_SEED): +def build_scenarios(pool_size=POOL_SIZE, seed=GENERATOR_SEED, length_config=None): """Deterministic pool of ``pool_size`` scenarios. The first five entries deliberately cover every structural case the dataset contract needs -- multi-image, multi-frame video, variable grids, and a text-only sample -- so a caller that only looks at a short prefix still sees all of them. + + Args: + pool_size: Number of scenarios. + seed: Drives the private ``random.Random``; the pool is byte-identical + across processes and ranks for a given seed. + length_config: Optional ``--mdp-mock-dataset-config-json`` payload. When + given, each sample's *total* token length is drawn from that + distribution and the vision/text split is fitted to it; when + ``None``, totals are drawn uniformly from ``TOTAL_TOKENS_RANGE`` as + before (the default pool is unchanged). """ rng = random.Random(seed) + totals = ( + None + if length_config is None + else draw_total_token_lengths(length_config, pool_size, seed=seed) + ) pool = [] for i in range(pool_size): - if i % TEXT_ONLY_EVERY == TEXT_ONLY_EVERY - 1: - scenario = _text_only_scenario(rng) + text_only = i % TEXT_ONLY_EVERY == TEXT_ONLY_EVERY - 1 + # Drawn here rather than inside the scenario builders so both branches + # consume the same rng budget with and without a length distribution. + total_target = rng.randint(*TOTAL_TOKENS_RANGE) if totals is None else totals[i] + if text_only: + scenario = _text_only_scenario(rng, total_target) else: - scenario = _multimodal_scenario(rng) + scenario = _multimodal_scenario(rng, total_target) total, image_tokens = scenario_totals(scenario) - assert TOTAL_TOKENS_RANGE[0] <= total <= TOTAL_TOKENS_RANGE[1], ( - f"scenario {i}: total {total} outside {TOTAL_TOKENS_RANGE}" - ) + assert total == total_target, f"scenario {i}: total {total} != target {total_target}" grids = scenario[0] if grids: - assert IMAGE_TOKENS_RANGE[0] <= image_tokens <= IMAGE_TOKENS_RANGE[1], ( - f"scenario {i}: image tokens {image_tokens} outside {IMAGE_TOKENS_RANGE}" + _, image_high = _image_token_bounds(total_target) + assert 4 <= image_tokens <= image_high, ( + f"scenario {i}: image tokens {image_tokens} outside [4, {image_high}]" ) for t, h, w in grids: assert h % SPATIAL_MERGE_SIZE == 0 and w % SPATIAL_MERGE_SIZE == 0 diff --git a/examples/multimodal_dev/forward_step.py b/examples/multimodal_dev/forward_step.py index b0d7e138e5f..0e0560b74da 100644 --- a/examples/multimodal_dev/forward_step.py +++ b/examples/multimodal_dev/forward_step.py @@ -12,7 +12,7 @@ from examples.multimodal_dev.observability import nvtx_phase from megatron.core import mpu -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, build_static_thd_metadata from megatron.core.parallel_state import ( get_tensor_model_parallel_group, get_tensor_model_parallel_rank, @@ -146,7 +146,7 @@ def broadcast_data_batch(data, device="cuda"): # ------------------------------------------------------------------- -def accumulate_flops_stats(packed_seq_params) -> None: +def accumulate_flops_stats(packed_seq_params, real_cu_seqlens=None) -> None: """Feed one micro-batch's real ``cu_seqlens`` into the FLOPs accumulators. Called from the forward step -- once per micro-batch, on every rank of the @@ -169,6 +169,15 @@ def accumulate_flops_stats(packed_seq_params) -> None: intentionally not used. The reduction stays on device (no ``.item()``), so no host sync is added to the hot path. + ``real_cu_seqlens`` overrides ``cu_seqlens_q``. Under ``--thd-static-packing`` + with the ``append_dummy_seq`` tail policy (forced at CP>1), the static pad is + represented as an ordinary extra sequence and therefore lands in + ``cu_seqlens_q`` itself. Accumulating that would overstate ``sum(L)`` and + ``sum(L^2)`` by the padding fraction -- exactly the shape of a false speedup, + since it appears only on the padded side. The collator emits the pre-tail + vector as ``flops_cu_seqlens`` in that case. ``extend_last`` (CP=1) leaves + ``cu_seqlens_q`` untouched and needs no override. + Imported lazily: ``megatron.training.training`` pulls in the whole training stack, and this module is also imported by unit tests that never build it. """ @@ -176,7 +185,11 @@ def accumulate_flops_stats(packed_seq_params) -> None: # BSHD path: leave the accumulator untouched so # ``num_floating_point_operations`` keeps its closed-form defaults. return - cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q", None) + cu_seqlens = ( + real_cu_seqlens + if real_cu_seqlens is not None + else getattr(packed_seq_params, "cu_seqlens_q", None) + ) if cu_seqlens is None: return try: @@ -221,7 +234,7 @@ def accumulate_vision_flops_stats_from_items(vision_items) -> None: def _accumulate_workload_stats( - model, packed_seq_params, *, vision_items=None, image_grid_thw=None + model, packed_seq_params, *, vision_items=None, image_grid_thw=None, real_cu_seqlens=None ) -> None: """Report one micro-batch's decoder and vision work exactly once per rank. @@ -236,7 +249,7 @@ def _accumulate_workload_stats( if vp_stage not in (None, 0): return - accumulate_flops_stats(packed_seq_params) + accumulate_flops_stats(packed_seq_params, real_cu_seqlens=real_cu_seqlens) if vision_items is not None: accumulate_vision_flops_stats_from_items(vision_items) else: @@ -420,6 +433,37 @@ def pack_or_pad_batch( if use_packed_sequence: packed_batch: Dict[str, Any] = {} + # --thd-static-packing: emit a fixed-shape THD batch. Every microbatch + # becomes exactly `max_seqlen_per_dp_cp_rank * cp_size` rows with + # cu_seqlens* of `thd_max_packed_sequences + 1` entries, which is the + # contract MCore's THD CUDA-graph machinery expects from a packing + # scheduler. This collator works in GLOBAL (pre-CP-slice) coordinates: + # CP slicing happens later in models/base.py. + static_target_T = None + static_max_num_seqs = None + static_tail_policy = "extend_last" + try: + static_args = get_args() + except AssertionError: + static_args = None + if static_args is not None and getattr(static_args, "thd_static_packing", False): + static_target_T = int(static_args.max_seqlen_per_dp_cp_rank) * cp_size + static_max_num_seqs = int(static_args.thd_max_packed_sequences) + # append_dummy_seq, matching what --sequence-packing-scheduler + # produces. `extend_last` is *not* usable here even at CP=1: it + # leaves cu_seqlens_q ending at the real token count while the + # tensors are padded to target_T, and TE then returns a shorter + # attention output than the padded input (observed as a view + # mismatch in Attention._apply_output_gate). + # + # The cost is that the pad tail becomes an ordinary sequence in + # cu_seqlens_q, which would inflate the FLOPs accumulator; the + # pre-tail vector is therefore emitted separately (see + # accumulate_flops_stats). + static_tail_policy = ( + getattr(static_args, "thd_tail_padding_policy", None) or "append_dummy_seq" + ) + # Owner-sharded pixel reading: during MDP window capture of a # microbatch owned by another worker, skip pixel # materialization + H2D wholesale. All text tensors and vision item @@ -480,21 +524,35 @@ def pack_or_pad_batch( # routing in megatron.core to exclude padded tokens from aux loss, # z-loss, and expert-bias accumulation. total_tokens_padded = cu_seqlens_padded[-1] + # Physical row count of the emitted tensors. Under static packing it + # is the fixed target, so the tail beyond the pack is padding too. + physical_T = total_tokens_padded + if static_target_T is not None: + assert total_tokens_padded <= static_target_T, ( + f"Packed THD length ({total_tokens_padded}) exceeds the static " + f"target ({static_target_T}). Increase " + "--max-seqlen-per-dp-cp-rank, or lower the number of samples per " + "microbatch (--micro-batch-size, or the greedy token budget)." + ) + physical_T = static_target_T padding_mask_thd = torch.zeros( - total_tokens_padded, dtype=torch.bool, pin_memory=use_pinned + physical_T, dtype=torch.bool, pin_memory=use_pinned ) for i, real_seqlen in enumerate(seqlens_list): pad_start = cu_seqlens_padded[i] + real_seqlen pad_end = cu_seqlens_padded[i + 1] if pad_end > pad_start: padding_mask_thd[pad_start:pad_end] = True + if physical_T > total_tokens_padded: + padding_mask_thd[total_tokens_padded:] = True if use_pinned: # Single padded buffer per field; pad regions filled with the - # same values F.pad used, sample slices copied in place. + # same values F.pad used, sample slices copied in place. Sized to + # physical_T so static packing costs no second copy. def _packed_field(key, fill): out = torch.empty( - total_tokens_padded, dtype=batch[0][key].dtype, pin_memory=True + physical_T, dtype=batch[0][key].dtype, pin_memory=True ) out.fill_(fill) for i, sample in enumerate(batch): @@ -531,9 +589,16 @@ def _packed_field(key, fill): packed_batch["labels"] = labels_list[0].unsqueeze(0) packed_batch["loss_mask"] = loss_mask_list[0].unsqueeze(0) else: - packed_batch["input_ids"] = torch.concat(input_ids_list, dim=0).unsqueeze(0) - packed_batch["labels"] = torch.concat(labels_list, dim=0).unsqueeze(0) - packed_batch["loss_mask"] = torch.concat(loss_mask_list, dim=0).unsqueeze(0) + def _concat_field(pieces, fill): + packed = torch.concat(pieces, dim=0) + tail = physical_T - packed.shape[0] + if tail: + packed = F.pad(packed, (0, tail), value=fill) + return packed.unsqueeze(0) + + packed_batch["input_ids"] = _concat_field(input_ids_list, 0) + packed_batch["labels"] = _concat_field(labels_list, -100) + packed_batch["loss_mask"] = _concat_field(loss_mask_list, 0) packed_batch["padding_mask"] = padding_mask_thd.unsqueeze(0) if not suppress_pixels: if use_pinned and pixel_values_list: @@ -593,6 +658,35 @@ def _packed_field(key, fill): max_seqlen_q = int((cu_seqlens_padded_t[1:] - cu_seqlens_padded_t[:-1]).max().item()) total_tokens = int(cu_seqlens_padded_t[-1].item()) + pad_between_seqs = None + if static_target_T is not None: + cu_seqlens_t, cu_seqlens_padded_t, real_cu_seqlens_t = build_static_thd_metadata( + cu_seqlens_t, + cu_seqlens_padded_t, + target_len=static_target_T, + max_num_seqs=static_max_num_seqs, + tail_padding_policy=static_tail_policy, + cp_size=cp_size, + ) + # max_seqlen must be the padded static value: the tail belongs to a + # sequence now, and a stale (shorter) max silently produces wrong + # attention rather than a crash. + max_seqlen_q = static_target_T + total_tokens = static_target_T + # Must be batch-independent (that is the point of static shapes), so + # derive it from the alignment rather than from this batch's + # vectors: with divisible_by == 1 no sample is ever padded, so + # cu_seqlens and cu_seqlens_padded coincide and there is provably no + # gap between sequences. Saying True there is not free -- it makes + # FlashAttention ineligible and, when the fused cuDNN backend is not + # selected either, drops TE onto its unfused O(T^2) attention, which + # OOMs at these lengths. + pad_between_seqs = divisible_by > 1 + if real_cu_seqlens_t is not None: + # append_dummy_seq put the tail into cu_seqlens_q itself, which + # would inflate sum(L) and sum(L^2) in the FLOPs accumulator. + packed_batch["flops_cu_seqlens"] = real_cu_seqlens_t + packed_batch["packed_seq_params"] = PackedSeqParams( qkv_format="thd", cu_seqlens_q=cu_seqlens_t, @@ -602,6 +696,7 @@ def _packed_field(key, fill): max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_q, total_tokens=total_tokens, + pad_between_seqs=pad_between_seqs, ) return packed_batch @@ -753,6 +848,7 @@ def mdp_forward_step(runtime, data_iterator, model): model, record.decoder_packed_seq_params, vision_items=record.vision_items, + real_cu_seqlens=batch.get("flops_cu_seqlens"), ) vision_embeddings = None @@ -810,6 +906,7 @@ def forward_step(data_iterator, model): model, batch.get("packed_seq_params", None), image_grid_thw=batch.get("image_grid_thw", None), + real_cu_seqlens=batch.get("flops_cu_seqlens"), ) # ``pixel_values`` is the heavy vision tensor and is only consumed diff --git a/examples/multimodal_dev/tests/test_flops_accounting.py b/examples/multimodal_dev/tests/test_flops_accounting.py index b6674898af5..a96b395ec4b 100644 --- a/examples/multimodal_dev/tests/test_flops_accounting.py +++ b/examples/multimodal_dev/tests/test_flops_accounting.py @@ -23,7 +23,9 @@ def __init__(self, module): def test_workload_stats_report_once_per_physical_rank(monkeypatch, vp_stage, expected_calls): """Only the canonical VPP chunk reports a physical rank's micro-batch.""" calls = [] - monkeypatch.setattr(forward_step, "accumulate_flops_stats", lambda _: calls.append("decoder")) + monkeypatch.setattr( + forward_step, "accumulate_flops_stats", lambda *_, **__: calls.append("decoder") + ) monkeypatch.setattr( forward_step, "accumulate_vision_flops_stats_from_items", @@ -44,7 +46,9 @@ def test_workload_stats_native_path_uses_grid_metadata(monkeypatch): """The native path keeps reporting grid-based vision statistics.""" calls = [] grid = object() - monkeypatch.setattr(forward_step, "accumulate_flops_stats", lambda _: calls.append("decoder")) + monkeypatch.setattr( + forward_step, "accumulate_flops_stats", lambda *_, **__: calls.append("decoder") + ) monkeypatch.setattr( forward_step, "accumulate_vision_flops_stats_from_grids", diff --git a/examples/multimodal_dev/tests/test_thd_static_packing.py b/examples/multimodal_dev/tests/test_thd_static_packing.py new file mode 100644 index 00000000000..10326f03c2d --- /dev/null +++ b/examples/multimodal_dev/tests/test_thd_static_packing.py @@ -0,0 +1,341 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Fixed-shape THD batches (``--thd-static-packing``) and the MDP mock length +distribution. + +``pack_or_pad_batch`` ends with a TP-group broadcast, so these tests require +``torch.distributed`` to be initialised. Run via:: + + torchrun --nproc-per-node 1 -m pytest -q \\ + examples/multimodal_dev/tests/test_thd_static_packing.py +""" + +import os +import sys +import types + +import numpy as np +import pytest +import torch + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from examples.multimodal_dev import forward_step as fs +from examples.multimodal_dev.data import mdp_scenarios +from megatron.core.packed_seq_params import build_static_thd_metadata +from tests.unit_tests.test_utilities import Utils + +MAX_SEQLEN = 64 +MAX_PACKED = 8 + + +@pytest.fixture(scope="module", autouse=True) +def _init_model_parallel(): + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + yield + Utils.destroy_model_parallel() + + +def _fake_args(**overrides): + base = dict( + sequence_parallel=False, + mdp_enable=False, + thd_static_packing=True, + max_seqlen_per_dp_cp_rank=MAX_SEQLEN, + thd_max_packed_sequences=MAX_PACKED, + thd_tail_padding_policy=None, + image_token_id=248056, + vision_spatial_merge_size=2, + ) + base.update(overrides) + return types.SimpleNamespace(**base) + + +@pytest.fixture +def static_args(monkeypatch): + """Install a fake ``get_args`` returning a static-packing configuration.""" + + def _install(**overrides): + args = _fake_args(**overrides) + monkeypatch.setattr(fs, "get_args", lambda: args) + return args + + return _install + + +def _make_sample(seq_len, *, base=0, num_patches=4, pixel_dim=8, device="cuda"): + return { + "input_ids": torch.arange(seq_len, dtype=torch.long, device=device) + base, + "labels": torch.arange(seq_len, dtype=torch.long, device=device) + base + 100, + "loss_mask": torch.ones(seq_len, dtype=torch.float, device=device), + "pixel_values": torch.full((num_patches, pixel_dim), float(base), device=device), + "image_grid_thw": torch.tensor([[2, 4, 4]], dtype=torch.long, device=device), + } + + +# --------------------------------------------------------------------------- +# build_static_thd_metadata -- pure helper +# --------------------------------------------------------------------------- + + +class TestBuildStaticThdMetadata: + def test_extend_last_keeps_valid_boundaries(self): + cu = torch.tensor([0, 5, 12], dtype=torch.int32) + cu_padded = torch.tensor([0, 6, 14], dtype=torch.int32) + q, q_padded, real = build_static_thd_metadata( + cu, cu_padded, target_len=32, max_num_seqs=4, tail_padding_policy="extend_last" + ) + assert real is None # no FLOPs override needed + assert q.tolist() == [0, 5, 12, 12, 12] # valid vector untouched, then padded + assert q_padded.tolist() == [0, 6, 32, 32, 32] + assert q.numel() == q_padded.numel() == 5 + + def test_append_dummy_seq_returns_the_real_vector(self): + cu = torch.tensor([0, 5, 12], dtype=torch.int32) + cu_padded = torch.tensor([0, 6, 14], dtype=torch.int32) + q, q_padded, real = build_static_thd_metadata( + cu, cu_padded, target_len=32, max_num_seqs=4, tail_padding_policy="append_dummy_seq" + ) + assert real is not None and real.tolist() == [0, 5, 12] + assert q.tolist() == [0, 5, 12, 30, 30] # tail became an ordinary sequence + assert q_padded.tolist() == [0, 6, 14, 32, 32] + + def test_exact_fit_needs_no_tail(self): + cu = torch.tensor([0, 16, 32], dtype=torch.int32) + q, q_padded, real = build_static_thd_metadata( + cu, cu.clone(), target_len=32, max_num_seqs=3, tail_padding_policy="extend_last" + ) + assert real is None + assert q.tolist() == [0, 16, 32, 32] + + def test_overflowing_pack_names_the_flag(self): + cu = torch.tensor([0, 40], dtype=torch.int32) + with pytest.raises(AssertionError, match="max-seqlen-per-dp-cp-rank"): + build_static_thd_metadata( + cu, cu.clone(), target_len=32, max_num_seqs=4, tail_padding_policy="extend_last" + ) + + def test_too_many_sequences_names_the_flag(self): + cu = torch.tensor([0, 4, 8, 12, 16], dtype=torch.int32) + with pytest.raises(AssertionError, match="thd_max_packed_sequences"): + build_static_thd_metadata( + cu, cu.clone(), target_len=32, max_num_seqs=2, tail_padding_policy="extend_last" + ) + + def test_extend_last_is_rejected_under_cp(self): + cu = torch.tensor([0, 8], dtype=torch.int32) + with pytest.raises(AssertionError, match="before CP slicing"): + build_static_thd_metadata( + cu, + cu.clone(), + target_len=32, + max_num_seqs=4, + tail_padding_policy="extend_last", + cp_size=2, + ) + + +# --------------------------------------------------------------------------- +# pack_or_pad_batch -- fixed shapes +# --------------------------------------------------------------------------- + + +# The pinned fast path (--mdp-enable, TP=1) stages into pinned host buffers, so +# it requires the CPU sample tensors a real dataloader produces. +@pytest.mark.parametrize( + "mdp_enable, sample_device", + [(False, "cuda"), (True, "cpu")], + ids=["generic", "pinned"], +) +class TestStaticShapes: + def test_heterogeneous_batch_has_fixed_shapes(self, static_args, mdp_enable, sample_device): + static_args(mdp_enable=mdp_enable) + batch = [ + _make_sample(L, base=i * 1000, device=sample_device) + for i, L in enumerate([5, 13, 7]) + ] + packed = fs.pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + for key in ("input_ids", "labels", "loss_mask", "padding_mask"): + assert packed[key].shape == (1, MAX_SEQLEN), key + psp = packed["packed_seq_params"] + for name in ( + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + ): + assert getattr(psp, name).numel() == MAX_PACKED + 1, name + assert psp.max_seqlen_q == MAX_SEQLEN + assert psp.max_seqlen_kv == MAX_SEQLEN + assert psp.pad_between_seqs is True + + def test_shapes_do_not_depend_on_the_batch(self, static_args, mdp_enable, sample_device): + static_args(mdp_enable=mdp_enable) + shapes = set() + entries = set() + for lengths in ([3], [5, 13, 7], [11, 11, 11, 11, 2]): + batch = [ + _make_sample(L, base=i * 1000, device=sample_device) + for i, L in enumerate(lengths) + ] + packed = fs.pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + shapes.add(tuple(packed["input_ids"].shape)) + entries.add(int(packed["packed_seq_params"].cu_seqlens_q.numel())) + assert shapes == {(1, MAX_SEQLEN)} + assert entries == {MAX_PACKED + 1} + + def test_pad_positions_are_masked_out(self, static_args, mdp_enable, sample_device): + static_args(mdp_enable=mdp_enable) + lengths = [5, 13, 7] + batch = [ + _make_sample(L, base=i * 1000, device=sample_device) + for i, L in enumerate(lengths) + ] + packed = fs.pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + real = sum(lengths) + assert packed["loss_mask"][0, :real].sum().item() == real + assert packed["loss_mask"][0, real:].sum().item() == 0 + assert bool(packed["padding_mask"][0, real:].all()) + assert not bool(packed["padding_mask"][0, :real].any()) + assert bool((packed["labels"][0, real:] == -100).all()) + + def test_overflow_is_rejected_with_a_named_flag(self, static_args, mdp_enable, sample_device): + static_args(mdp_enable=mdp_enable) + batch = [ + _make_sample(MAX_SEQLEN, base=0, device=sample_device), + _make_sample(MAX_SEQLEN, base=1000, device=sample_device), + ] + with pytest.raises(AssertionError, match="max-seqlen-per-dp-cp-rank"): + fs.pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + +def test_static_packing_off_is_byte_identical(static_args): + """The opt-in branch must not perturb the default path.""" + batch = [_make_sample(L, base=i * 1000) for i, L in enumerate([5, 13, 7])] + reference = fs.pack_or_pad_batch( + [dict(s) for s in batch], use_packed_sequence=True, device="cuda" + ) + static_args(thd_static_packing=False) + actual = fs.pack_or_pad_batch( + [dict(s) for s in batch], use_packed_sequence=True, device="cuda" + ) + for key in ("input_ids", "labels", "loss_mask", "padding_mask"): + assert torch.equal(reference[key], actual[key]), key + assert torch.equal( + reference["packed_seq_params"].cu_seqlens_q, actual["packed_seq_params"].cu_seqlens_q + ) + assert reference["packed_seq_params"].max_seqlen_q == actual["packed_seq_params"].max_seqlen_q + assert "flops_cu_seqlens" not in actual + + +# --------------------------------------------------------------------------- +# FLOPs accounting +# --------------------------------------------------------------------------- + + +class TestFlopsAccounting: + def test_extend_last_reports_only_real_tokens(self, static_args): + static_args() + lengths = [5, 13, 7] + batch = [_make_sample(L, base=i * 1000) for i, L in enumerate(lengths)] + packed = fs.pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + # extend_last leaves cu_seqlens_q compact, so no override is emitted and + # the accumulator sees exactly the real tokens. + assert "flops_cu_seqlens" not in packed + cu = packed["packed_seq_params"].cu_seqlens_q + seqlens = (cu[1:] - cu[:-1]).tolist() + assert sum(seqlens) == sum(lengths) + + def test_dummy_tail_would_pollute_without_the_override(self): + """The override exists because the dummy tail lands in cu_seqlens_q.""" + lengths = [5, 13, 7] + cu = torch.tensor([0, 5, 18, 25], dtype=torch.int32) + q, _, real = build_static_thd_metadata( + cu, + cu.clone(), + target_len=MAX_SEQLEN, + max_num_seqs=MAX_PACKED, + tail_padding_policy="append_dummy_seq", + ) + polluted = (q[1:] - q[:-1]).clamp(min=0).sum().item() + honest = (real[1:] - real[:-1]).sum().item() + assert honest == sum(lengths) + assert polluted == MAX_SEQLEN # would overstate by the whole pad + + +# --------------------------------------------------------------------------- +# Mock length distribution +# --------------------------------------------------------------------------- + + +LOGNORMAL = { + "mode": "distribution", + "type": "lognormal", + "format": "thd", + "min_seq_len": 256, + "max_seq_len": 4096, + "mean_seq_len": 2048, + "lognormal_sigma": 1.1, +} +DEGENERATE = { + "mode": "distribution", + "type": "lognormal", + "format": "thd", + "min_seq_len": 1024, + "max_seq_len": 1024, + "mean_seq_len": 1024, + "lognormal_sigma": 1.1, +} + + +class TestMockLengthDistribution: + def test_default_pool_is_unchanged(self): + pool = mdp_scenarios.build_scenarios() + totals = [mdp_scenarios.scenario_totals(s)[0] for s in pool] + assert len(pool) == mdp_scenarios.POOL_SIZE + assert min(totals) >= mdp_scenarios.TOTAL_TOKENS_RANGE[0] + assert max(totals) <= mdp_scenarios.TOTAL_TOKENS_RANGE[1] + + def test_distribution_controls_the_totals(self): + pool = mdp_scenarios.build_scenarios(length_config=LOGNORMAL) + totals = [mdp_scenarios.scenario_totals(s)[0] for s in pool] + assert min(totals) >= LOGNORMAL["min_seq_len"] + assert max(totals) <= LOGNORMAL["max_seq_len"] + # Wider than the built-in [1000, 2000] window, which is the point. + assert max(totals) - min(totals) > 1000 + + def test_degenerate_distribution_is_constant(self): + pool = mdp_scenarios.build_scenarios(length_config=DEGENERATE) + totals = {mdp_scenarios.scenario_totals(s)[0] for s in pool} + assert totals == {1024} + + def test_pool_is_identical_across_ranks(self): + """Every MDP rank rebuilds the pool independently; a divergence hangs. + + Perturbing NumPy's global RNG between builds stands in for the different + global state each rank carries. + """ + first = mdp_scenarios.build_scenarios(length_config=LOGNORMAL) + np.random.seed(9871) + np.random.random(1000) + second = mdp_scenarios.build_scenarios(length_config=LOGNORMAL) + assert first == second + + def test_the_draw_leaves_no_global_numpy_side_effect(self): + np.random.seed(4242) + expected = np.random.random(4).tolist() + np.random.seed(4242) + mdp_scenarios.draw_total_token_lengths(LOGNORMAL, 8) + assert np.random.random(4).tolist() == expected + + def test_dataset_emits_the_requested_lengths(self): + from examples.multimodal_dev.data.mdp_mock import MdpThdMockDataset + + dataset = MdpThdMockDataset(num_samples=16, length_config=DEGENERATE) + for i in range(16): + assert dataset[i]["input_ids"].shape[0] == 1024 diff --git a/megatron/core/mdp/README.md b/megatron/core/mdp/README.md index 1c414deb87f..1e8822831fd 100644 --- a/megatron/core/mdp/README.md +++ b/megatron/core/mdp/README.md @@ -49,6 +49,7 @@ schedule model list. | `allocator.py` / `storage.py` | Single allocation point for MDP buffers; endpoint leaf storage | | `bridge.py` | One ledger + transport for pixels/embeddings/gradients | | `window.py` / `activation.py` | Iteration window with VPP replay cursors; forward handle, chunking, encoder THD params | +| `packing.py` | Greedy token-budget bin filling and the cross-iteration sample buffer | | `runtime.py` / `schedule.py` | Phase machine; schedule and finalizer wrappers | | `encoder.py` / `optimizer.py` | Encoder DDP over WORLD + ZeRO-1; composite optimizer with WORLD overflow union | | `checkpoint.py` | Weight-only torch_dist facade (`vision_model.*` with WORLD replica metadata) | diff --git a/megatron/core/mdp/config.py b/megatron/core/mdp/config.py index f08e23d05d6..128bce8b46c 100644 --- a/megatron/core/mdp/config.py +++ b/megatron/core/mdp/config.py @@ -48,6 +48,7 @@ class MdpConfig: debug_plan_payload_check: bool = False pixel_locality: bool = False overlap_window_capture: bool = False + greedy_packing: bool = False @dataclass(frozen=True) @@ -76,6 +77,27 @@ class MdpCompatibilityOptions: checkpoint_mode: str save_requested: bool load_requested: bool + sequence_parallel: bool = False + sequence_packing_scheduler: Optional[str] = None + thd_static_packing: bool = False + max_seqlen_per_dp_cp_rank: Optional[int] = None + thd_max_packed_sequences: Optional[int] = None + + +def thd_row_alignment(options: "MdpCompatibilityOptions") -> int: + """Row alignment the MDP collator pads each packed sample to. + + Mirrors ``pack_or_pad_batch``'s ``divisible_by`` (zigzag CP wants an even + per-rank split; SP additionally splits across TP). The greedy token budget + must be a multiple of this, or a full bin cannot be partitioned legally. + """ + if options.context_parallel_size > 1: + return ( + options.tensor_parallel_size * options.context_parallel_size * 2 + if options.sequence_parallel + else options.context_parallel_size * 2 + ) + return options.tensor_parallel_size if options.sequence_parallel else 1 def _reject(option: str, value: Any, condition: str, why: str, suggestion: str = "") -> None: @@ -150,6 +172,7 @@ def validate_mdp_config(config: MdpConfig, options: MdpCompatibilityOptions) -> "False", ) _validate_override_entries(config.vision_config_overrides) + _validate_packing(config, options) # --- parallel dimensions and rank mapping preconditions --- if options.rank_order != SUPPORTED_RANK_ORDER: @@ -305,6 +328,58 @@ def validate_mdp_config(config: MdpConfig, options: MdpCompatibilityOptions) -> ) +def _validate_packing(config: MdpConfig, options: MdpCompatibilityOptions) -> None: + """Reject packing configurations MDP cannot honor. + + ``--sequence-packing-scheduler`` is rejected outright, not merely untested: + ``training.py`` wraps the data iterator whenever it is set, and + ``DpBalancedScheduler.run`` then asserts on GPT-only sample keys, deletes + every key outside those six (dropping ``pixel_values`` / ``image_grid_thw``), + and reroutes samples across DP with an all-to-all that has no notion of + variable-size pixel payloads. Without this rejection the run dies deep inside + an assert about a missing ``tokens`` key. + """ + if options.sequence_packing_scheduler is not None: + _reject( + "sequence_packing_scheduler", + options.sequence_packing_scheduler, + "sequence_packing_scheduler is None", + "MCore's packing schedulers assert on GPT-only sample keys, drop the " + "pixel payload, and reroute samples across DP without pixel awareness. " + "MDP owns its packing (--mdp-greedy-packing).", + "None", + ) + if not config.greedy_packing: + return + if options.max_seqlen_per_dp_cp_rank is None: + _reject( + "max_seqlen_per_dp_cp_rank", + options.max_seqlen_per_dp_cp_rank, + "max_seqlen_per_dp_cp_rank is set when --mdp-greedy-packing is on", + "The greedy token budget is max_seqlen_per_dp_cp_rank x " + "context_parallel_size; there is no default for it.", + ) + alignment = thd_row_alignment(options) + budget = options.max_seqlen_per_dp_cp_rank * options.context_parallel_size + if budget % alignment != 0: + _reject( + "max_seqlen_per_dp_cp_rank", + options.max_seqlen_per_dp_cp_rank, + f"the greedy token budget ({budget}) is divisible by the collator row " + f"alignment ({alignment})", + "A bin filled to the budget must still split legally across CP/SP ranks; " + "discovering this inside TransformerEngine gives a far worse error.", + ) + if options.thd_max_packed_sequences is not None and options.thd_max_packed_sequences < 1: + _reject( + "thd_max_packed_sequences", + options.thd_max_packed_sequences, + "thd_max_packed_sequences >= 1", + "It caps the real sequences per greedy bin.", + "8", + ) + + def _validate_override_entries(overrides: Sequence) -> None: """Shared structural validation for vision config override entry sequences.""" seen = set() diff --git a/megatron/core/mdp/integration.py b/megatron/core/mdp/integration.py index 4e4aa6495c1..cef0913cfcf 100644 --- a/megatron/core/mdp/integration.py +++ b/megatron/core/mdp/integration.py @@ -29,6 +29,7 @@ SUPPORTED_RANK_ORDER, MdpCompatibilityOptions, MdpConfig, + thd_row_alignment, validate_mdp_config, ) from megatron.core.mdp.encoder import ( @@ -101,6 +102,7 @@ def mdp_config_from_args(args) -> MdpConfig: debug_plan_payload_check=getattr(args, "mdp_debug_plan_payload_check", False), pixel_locality=getattr(args, "mdp_pixel_locality", False), overlap_window_capture=getattr(args, "mdp_overlap_window_capture", False), + greedy_packing=getattr(args, "mdp_greedy_packing", False), ) @@ -154,6 +156,11 @@ def compatibility_options_from_args(args) -> MdpCompatibilityOptions: checkpoint_mode=getattr(args, "ckpt_format", "torch_dist"), save_requested=getattr(args, "save", None) is not None, load_requested=getattr(args, "load", None) is not None, + sequence_parallel=bool(getattr(args, "sequence_parallel", False)), + sequence_packing_scheduler=getattr(args, "sequence_packing_scheduler", None), + thd_static_packing=bool(getattr(args, "thd_static_packing", False)), + max_seqlen_per_dp_cp_rank=getattr(args, "max_seqlen_per_dp_cp_rank", None), + thd_max_packed_sequences=getattr(args, "thd_max_packed_sequences", None), ) @@ -226,6 +233,12 @@ def maybe_build_mdp_domain(*, args, model, optimizer, optimizer_config, ddp_conf else: params_dtype = torch.float32 allocator = DirectBufferAllocator() + compat = compatibility_options_from_args(args) + greedy_token_budget = ( + args.max_seqlen_per_dp_cp_rank * args.context_parallel_size + if mdp_config.greedy_packing + else None + ) _RUNTIME = MdpRuntime( config=mdp_config, rank_map=rank_map, @@ -245,6 +258,9 @@ def maybe_build_mdp_domain(*, args, model, optimizer, optimizer_config, ddp_conf hidden_size=args.hidden_size, params_dtype=params_dtype, num_vpp_chunks=len(model), + greedy_token_budget=greedy_token_budget, + greedy_max_num_seqs=getattr(args, "thd_max_packed_sequences", None), + greedy_row_alignment=thd_row_alignment(compat), ) logger.info( "MDP: runtime installed (outer_dp_rank=%d, worker_id=%s, endpoint=%d, " diff --git a/megatron/core/mdp/knowledge.md b/megatron/core/mdp/knowledge.md index 354be0c543d..e3b8700b21d 100644 --- a/megatron/core/mdp/knowledge.md +++ b/megatron/core/mdp/knowledge.md @@ -124,6 +124,7 @@ returns to `EMPTY`. | `storage.py` | Endpoint embedding leaves and lifecycle checks. | | `bridge.py` | Canonical ledger and `all_to_all_single` transport for all three payload phases. | | `window.py` | Whole-iteration capture, microbatch replay cursors, pixel ownership context. | +| `packing.py` | Greedy token-budget bin filling and the cross-iteration sample buffer (`--mdp-greedy-packing`). | | `activation.py` | Encoder forward handle, chunk output retention, multi-tensor backward. | | `encoder.py` | Encoder process groups, DDP/ZeRO-1 domain, gradient finalization. | | `runtime.py` | P0-P5 orchestration, prefetch handoff, per-iteration state and metrics. | @@ -165,7 +166,19 @@ The collator builds normal decoder tensors plus an MDP vision sidecar: - `vision_decoder_positions`: absolute image-token positions in the decoder's packed physical layout; - `pixel_values`: present only on the owner worker for that microbatch; -- `image_grid_thw`: present on all workers and used to derive item shapes. +- `image_grid_thw`: present on all workers and used to derive item shapes; +- `flops_cu_seqlens`: present only under `--thd-static-packing`; the + pre-tail-pad valid `cu_seqlens`, because the static tail is represented as an + ordinary dummy sequence and would otherwise inflate the FLOPs accumulator. + +Under `--thd-static-packing` the tail policy is `append_dummy_seq`, not +`extend_last`. `extend_last` leaves `cu_seqlens_q` ending at the real token +count while the tensors are padded to the static target; TE then returns a +shorter attention output than the padded input. `pad_between_seqs` is derived +from the collator's row alignment (`divisible_by > 1`), not hardcoded to +`True`: at TP=CP=1 no sample is ever padded, so there is provably no gap, and +claiming otherwise makes FlashAttention ineligible and can drop TE onto its +unfused O(T^2) backend. `MdpModelAdapter.get_batch` converts the model-specific batch into `CapturedMicrobatch`. Core MDP treats `model_payload` as opaque and consumes @@ -250,6 +263,46 @@ Primary flags: - `--mdp-plan-check-interval` - `--mdp-overlap-window-capture` - `--mdp-debug-plan-payload-check` +- `--mdp-greedy-packing` +- `--mdp-mock-dataset-config-json` + +Packing flags MDP consumes from the core config (all optional, all off by +default): + +- `--max-seqlen-per-dp-cp-rank` -- required by `--mdp-greedy-packing`; the + greedy token budget is this times `context_parallel_size`. +- `--thd-max-packed-sequences` -- caps real sequences per bin, and fixes the + `cu_seqlens` entry count under `--thd-static-packing`. +- `--thd-static-packing` -- the data path emits fixed-shape THD batches + (`T == max_seqlen_per_dp_cp_rank * cp_size`, `cu_seqlens*` of + `thd_max_packed_sequences + 1` entries). Requires + `--pad-packed-seq-alignment max`. Independent of `--mdp-greedy-packing`: all + four corners of the 2x2 are reachable. +- `--sequence-packing-scheduler` is **rejected** under MDP. It is not merely + untested: `training.py` wraps the data iterator whenever it is set, and + `DpBalancedScheduler.run` then asserts on GPT-only sample keys, deletes every + key outside those six (dropping `pixel_values` / `image_grid_thw`), and + reroutes samples across DP with an all-to-all that has no notion of + variable-size pixel payloads. MDP owns its packing instead. + +`--mdp-greedy-packing` **reinterprets** `--micro-batch-size` and +`--global-batch-size`: they no longer describe what goes into a microbatch, +only how many bins an iteration has (`N = GBS / (MBS x DP)`). Two consequences +to state in any comparison: + +- GBS means "N x token budget", so loss curves are not iteration-by-iteration + comparable against a fixed-GBS run; +- DP ranks consume different sample counts, so `consumed_train_samples` is + computed from a real all-reduced count + (`training._mdp_greedy_consumed_samples`) rather than the closed form, and + resume determinism weakens (the sample buffer is not checkpointed). + +The stream must be provisioned by **tokens**, not samples: an iteration eats +about `token_budget / mean_sample_len` samples per bin, so +`train_iters x GBS` under-provisions whenever the mean sample is shorter than +`token_budget / MBS`. The mock provider scales its dataset accordingly +(`mdp_mock._greedy_sample_scale`); a real dataset must be sized by the +operator. There is deliberately no pixel-sharding flag. Pixel owner sharding is part of the MDP definition in this baseline. @@ -265,7 +318,8 @@ Current major constraints: - bf16/fp16 mixed precision; - synchronous global `torch_dist` weight-only checkpointing; - no FSDP/HSDP, FP8, full-iteration CUDA graph, CPU activation offload, or - communication-overlap modes rejected by `validate_mdp_config`. + communication-overlap modes rejected by `validate_mdp_config`; +- no `--sequence-packing-scheduler`. Always read `validate_mdp_config` before relaxing a constraint. A validation change without corresponding runtime/test support is not an implementation. diff --git a/megatron/core/mdp/packing.py b/megatron/core/mdp/packing.py new file mode 100644 index 00000000000..f7208365b76 --- /dev/null +++ b/megatron/core/mdp/packing.py @@ -0,0 +1,214 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Greedy token-budget packing for the MDP decoder data path. + +Without this module a microbatch is exactly ``--micro-batch-size`` samples, so +the packed THD length ``T`` is whatever those samples happen to sum to. With it, +a microbatch is a *bin*: samples are appended while they fit a token budget and +a real-sequence cap, so ``T`` is bounded and the padding waste is decoupled from +``--micro-batch-size``. + +The grouping rule is the in-order greedy fill of +``DpBalancedScheduler.get_groups_and_subsamples`` +(``megatron/core/datasets/data_schedule.py``). MDP cannot use that scheduler -- +it asserts on GPT-only sample keys, drops ``pixel_values`` / +``image_grid_thw``, and reroutes samples across DP with an all-to-all that has +no notion of variable-size pixel payloads (see ``mdp/README.md``). So MDP reuses +the *logic*, not the path. + +Variant: a fixed number of bins per iteration (``num_microbatches``), with the +number of samples consumed floating. ``num_microbatches`` therefore stays +static, which the PP schedule, the VPP replay cursors, and per-layer CUDA graph +slot sizing all depend on. The cost is that ``--global-batch-size`` / +``--micro-batch-size`` stop describing sample counts and become pure bin-count +knobs; MDP already hard-requires ``calculate_per_token_loss=True``, so loss +normalization is unaffected by a varying sample count per iteration. +""" + +import threading +from typing import Any, Callable, Iterator, List, Optional + +from megatron.core.mdp.errors import MdpConfigurationError, MdpStateError + + +def decoder_sample_length(sample: Any) -> int: + """Token count of one decoder sample dict. + + ``input_ids`` is the MDP dataset contract (``examples/multimodal_dev/data``); + image slots are already materialized in it, so it is the packed row count. + """ + try: + return int(sample["input_ids"].shape[0]) + except (KeyError, TypeError, AttributeError) as exc: + raise MdpConfigurationError( + "MDP: greedy packing violates: every sample carries a 1-D 'input_ids' " + f"tensor (got {type(sample).__name__})." + ) from exc + + +def greedy_bin_sizes( + lengths: List[int], *, token_budget: int, max_num_seqs: Optional[int], num_bins: int +) -> List[int]: + """Number of samples taken into each of ``num_bins`` in-order greedy bins. + + Pure function over a length vector, so the grouping rule is testable without + a dataset. ``lengths`` must already be aligned (see + :class:`GreedySampleStream`). Raises when ``lengths`` cannot fill + ``num_bins`` bins. + """ + sizes = [] + cursor = 0 + for _ in range(num_bins): + taken = 0 + total = 0 + while cursor + taken < len(lengths): + length = lengths[cursor + taken] + if taken and total + length > token_budget: + break + if max_num_seqs is not None and taken >= max_num_seqs: + break + total += length + taken += 1 + if taken == 0: + raise MdpStateError( + f"MDP: greedy packing violates: {num_bins} non-empty bins per iteration " + f"(the sample stream ran out after {len(sizes)})." + ) + sizes.append(taken) + cursor += taken + return sizes + + +class GreedySampleStream: + """Wrap a microbatch-list iterator so ``next()`` returns one greedy bin. + + The underlying data iterator yields whole ``--micro-batch-size`` lists (the + dataloader uses an identity collate over a ``batch_sampler``), so this holds + a **sample buffer**: it pulls those lists and drains them sample by sample + into bins. The batch_sampler itself is untouched -- shrinking it to size 1 + would change sampler bookkeeping and the shuffle order. + + The buffer is training state that carries across iterations: after filling + an iteration's bins it usually holds a partial list, and dropping those + samples would silently skip data. It is **not** checkpointed, so a resume + under greedy packing restarts on a batch_sampler boundary and is only + approximately reproducible. + + Args: + iterator: The underlying iterator of sample-dict lists. + token_budget: Maximum aligned token count per bin, + ``max_seqlen_per_dp_cp_rank * cp_size``. + max_num_seqs: Maximum real sequences per bin + (``thd_max_packed_sequences``), or ``None`` for no cap. + align: Per-sample row alignment applied by the collator; a sample of + length ``L`` occupies ``ceil(L / align) * align`` rows in the pack, + and that is what is charged against the budget. + length_of: Extracts a sample's unaligned token count. + """ + + def __init__( + self, + iterator: Iterator, + *, + token_budget: int, + max_num_seqs: Optional[int] = None, + align: int = 1, + length_of: Callable[[Any], int], + ) -> None: + if token_budget <= 0: + raise MdpConfigurationError( + f"MDP: greedy packing violates: token_budget > 0 (got {token_budget}). " + "Set --max-seqlen-per-dp-cp-rank." + ) + if align < 1 or token_budget % align != 0: + raise MdpConfigurationError( + f"MDP: greedy packing violates: token_budget ({token_budget}) is " + f"divisible by the collator row alignment ({align}); otherwise a full " + "bin cannot be split legally across CP/SP ranks." + ) + self._iterator = iterator + self._token_budget = token_budget + self._max_num_seqs = max_num_seqs + self._align = align + self._length_of = length_of + self._buffer: List[Any] = [] + self._buffer_cursor = 0 + self._exhausted = False + self._consumed_samples = 0 + # --mdp-overlap-window-capture captures the next iteration's window on a + # background thread. Only one capture per iterator is ever in flight (the + # consumer joins the prefetch before capturing again), but the buffer is + # mutable state shared with that thread, so guard it rather than rely on + # the caller's ordering. + self._lock = threading.Lock() + + @property + def consumed_samples(self) -> int: + """Real samples drained into bins since construction.""" + return self._consumed_samples + + @property + def exhausted(self) -> bool: + """True once the underlying iterator has raised ``StopIteration``.""" + return self._exhausted and self._buffer_cursor >= len(self._buffer) + + def __iter__(self) -> "GreedySampleStream": + return self + + def _next_sample(self) -> Optional[Any]: + """One sample from the buffer, refilling from the iterator as needed.""" + while self._buffer_cursor >= len(self._buffer): + if self._exhausted: + return None + try: + self._buffer = list(next(self._iterator)) + except StopIteration: + self._exhausted = True + return None + self._buffer_cursor = 0 + sample = self._buffer[self._buffer_cursor] + self._buffer_cursor += 1 + return sample + + def _unread(self) -> None: + """Push the last sample back; it starts the next bin.""" + self._buffer_cursor -= 1 + + def _aligned_length(self, sample: Any) -> int: + length = int(self._length_of(sample)) + align = self._align + return ((length + align - 1) // align) * align + + def __next__(self) -> List[Any]: + """One greedy bin, or raise ``StopIteration`` at end of stream. + + A bin is closed when the next sample would exceed the token budget or + the real-sequence cap. At end of stream a partially filled bin is + returned as is -- correct, and never an *empty* pack. + """ + with self._lock: + bin_samples: List[Any] = [] + total = 0 + while True: + sample = self._next_sample() + if sample is None: + break + length = self._aligned_length(sample) + if not bin_samples and length > self._token_budget: + raise MdpStateError( + f"MDP: sample violates: aligned length ({length}) <= the greedy " + f"token budget ({self._token_budget}). Raise " + "--max-seqlen-per-dp-cp-rank or filter overlong samples." + ) + if bin_samples and total + length > self._token_budget: + self._unread() + break + if self._max_num_seqs is not None and len(bin_samples) >= self._max_num_seqs: + self._unread() + break + bin_samples.append(sample) + total += length + if not bin_samples: + raise StopIteration + self._consumed_samples += len(bin_samples) + return bin_samples diff --git a/megatron/core/mdp/runtime.py b/megatron/core/mdp/runtime.py index ab393dd2eb7..8525a0a887e 100644 --- a/megatron/core/mdp/runtime.py +++ b/megatron/core/mdp/runtime.py @@ -34,6 +34,7 @@ from megatron.core.mdp.encoder import EncoderDomain, finalize_encoder_grads from megatron.core.mdp.errors import MdpConfigurationError, MdpStateError from megatron.core.mdp.groups import MdpProcessGroups, broadcast_descriptors +from megatron.core.mdp.packing import GreedySampleStream, decoder_sample_length from megatron.core.mdp.observability import ( MdpIterationMetrics, nvtx_phase, @@ -77,6 +78,9 @@ def __init__( params_dtype: torch.dtype, num_vpp_chunks: int = 1, device: Optional[torch.device] = None, + greedy_token_budget: Optional[int] = None, + greedy_max_num_seqs: Optional[int] = None, + greedy_row_alignment: int = 1, ) -> None: self.config = config self.rank_map = rank_map @@ -122,6 +126,15 @@ def __init__( self._prefetch_thread = None self._prefetch_box: Optional[dict] = None self._prefetch_stream: Optional[torch.cuda.Stream] = None + # Greedy token-budget packing (--mdp-greedy-packing). One + # GreedySampleStream per underlying data iterator, so train and eval keep + # independent sample buffers: an eval window must never consume (or be + # consumed by) the training stream's leftovers. Keyed by iterator + # identity, which the training loop keeps stable for the whole run. + self._greedy_token_budget = greedy_token_budget + self._greedy_max_num_seqs = greedy_max_num_seqs + self._greedy_row_alignment = greedy_row_alignment + self._greedy_streams: dict = {} # ------------------------------------------------------------------ # Public API @@ -501,7 +514,62 @@ def consumed_num_tokens(self) -> Optional[torch.Tensor]: # Internals # ------------------------------------------------------------------ + @staticmethod + def _first_iterator(data_iterators): + if isinstance(data_iterators, (list, tuple)): + return data_iterators[0] if data_iterators else None + return data_iterators + + def _greedy_stream(self, data_iterators): + """The greedy sample stream for this data iterator, created on first use.""" + iterator = self._first_iterator(data_iterators) + stream = self._greedy_streams.get(id(iterator)) + if stream is None: + stream = GreedySampleStream( + iterator, + token_budget=self._greedy_token_budget, + max_num_seqs=self._greedy_max_num_seqs, + align=self._greedy_row_alignment, + length_of=decoder_sample_length, + ) + self._greedy_streams[id(iterator)] = stream + return stream + + def consumed_samples(self) -> Optional[int]: + """Real samples drained by greedy packing so far, or ``None`` when off. + + Sums every stream (train and eval) this rank owns. ``training.py`` reads + the delta per iteration because the closed form + ``dp x mbs x num_microbatches`` is wrong under greedy packing. + """ + if not self.config.greedy_packing: + return None + return sum(stream.consumed_samples for stream in self._greedy_streams.values()) + def _capture_window(self, data_iterators, num_microbatches: int) -> MdpIterationWindow: + if not self.config.greedy_packing: + return self._capture(data_iterators, num_microbatches) + stream = self._greedy_stream(data_iterators) + try: + return self._capture(stream, num_microbatches) + except MdpStateError as error: + if not stream.exhausted: + raise + # Greedy fills a fixed number of bins to a token budget, so an + # iteration eats roughly token_budget/mean_sample_len samples per + # bin, not micro_batch_size. Megatron provisions the sampler as + # train_iters x global_batch_size *samples*, which under-counts + # whenever the mean sample is shorter than the per-bin share. + raise MdpStateError( + f"{error} Under --mdp-greedy-packing the sample stream must be " + "provisioned by tokens, not by samples: each bin consumes about " + f"{self._greedy_token_budget} tokens' worth of samples, so raise " + "--train-samples / the dataset size (roughly by " + "token_budget / (mean_sample_len x micro_batch_size)), or lower " + "--max-seqlen-per-dp-cp-rank." + ) from error + + def _capture(self, data_iterators, num_microbatches: int) -> MdpIterationWindow: return MdpIterationWindow.capture( data_iterators, num_microbatches=num_microbatches, diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 8a943b3ef2b..e9fc0db9d47 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -135,6 +135,17 @@ class ModelParallelConfig: that value + 1 entries in both eager and CUDA Graph modes. """ + thd_static_packing: bool = False + """The data path emits THD batches already padded to max_seqlen_per_dp_cp_rank x + context_parallel_size tokens with cu_seqlens* padded to thd_max_packed_sequences + 1 + entries. Set this when packing is done outside --sequence-packing-scheduler (for + example by MDP's greedy token-budget packer). It is an explicit statement about the + shape contract of the incoming batches; it does not itself perform any packing. + Requires max_seqlen_per_dp_cp_rank, thd_max_packed_sequences, and + pad_packed_seq_alignment='max', and is mutually exclusive with + sequence_packing_scheduler. + """ + expert_model_parallel_size: int = 1 """Distributes Moe Experts across sub data parallel dimension.""" @@ -550,6 +561,28 @@ def __post_init__(self): f"{self.pad_packed_seq_alignment}." ) + if self.thd_static_packing: + if self.sequence_packing_scheduler is not None: + raise ValueError( + "thd_static_packing declares that the data path already emits " + "fixed-shape THD batches, so it is mutually exclusive with " + "sequence_packing_scheduler " + f"(got {self.sequence_packing_scheduler!r})." + ) + if self.max_seqlen_per_dp_cp_rank is None: + raise ValueError( + "thd_static_packing requires --max-seqlen-per-dp-cp-rank: it is the " + "per-DPxCP-rank token budget every microbatch is padded to." + ) + if self.pad_packed_seq_alignment not in ("max", self.max_seqlen_per_dp_cp_rank): + raise ValueError( + "thd_static_packing requires --pad-packed-seq-alignment='max' (or a " + "value equal to max_seqlen_per_dp_cp_rank=" + f"{self.max_seqlen_per_dp_cp_rank}), got " + f"{self.pad_packed_seq_alignment!r}: any other alignment produces " + "variable token counts." + ) + if self.sequence_parallel: if self.tensor_model_parallel_size <= 1: raise ValueError("Cannot use sequence parallelism without tensor parallelism") diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index a0957a39eab..2f259e57bf1 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -293,6 +293,84 @@ def extend_thd_padding_before_cp_slice( return cu_seqlens_padded, max_seqlen, global_target_len +def build_static_thd_metadata( + cu_seqlens: Tensor, + cu_seqlens_padded: Tensor, + *, + target_len: int, + max_num_seqs: int, + tail_padding_policy: Literal["append_dummy_seq", "extend_last"], + cp_size: int = 1, + cp_partition_mode: str = "zigzag", +) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """Pad already-packed *global* THD metadata to a fixed shape. + + For collators that pack outside ``--sequence-packing-scheduler`` and pad the + token-like tensors themselves (see ``thd_static_packing``). Operates on the + global, pre-CP-slice metadata, which is where ``extend_last`` must be + applied. + + Args: + cu_seqlens: Valid-token boundaries, ``num_samples + 1`` entries. + cu_seqlens_padded: Physical boundaries, ``num_samples + 1`` entries. + target_len: Global physical row count every batch is padded to + (``max_seqlen_per_dp_cp_rank * cp_size``). + max_num_seqs: ``thd_max_packed_sequences``; both tensors are padded to + ``max_num_seqs + 1`` entries. + tail_padding_policy: ``extend_last`` keeps the valid coordinates + untouched (CP=1 only); ``append_dummy_seq`` represents the tail as an + ordinary extra sequence, which also lands in ``cu_seqlens``. + cp_size: Context-parallel world size. + cp_partition_mode: ``zigzag`` or ``contiguous``. + + Returns: + ``(cu_seqlens, cu_seqlens_padded, real_cu_seqlens)``. ``real_cu_seqlens`` + is the pre-tail valid vector and is not ``None`` only when + ``append_dummy_seq`` polluted ``cu_seqlens`` -- FLOPs accounting must use + it instead, or the tail is counted as real tokens. + """ + actual_len = int(cu_seqlens_padded[-1].item()) + assert actual_len <= target_len, ( + f"Packed THD length ({actual_len}) exceeds the static target ({target_len}). " + "Increase --max-seqlen-per-dp-cp-rank, or reduce the number of samples per " + "microbatch so the pack fits." + ) + + real_cu_seqlens = None + if actual_len < target_len: + if tail_padding_policy == "extend_last": + assert cp_size == 1, ( + "thd_tail_padding_policy='extend_last' needs the global metadata " + "extended before CP slicing, which this collator does not do; use " + "'append_dummy_seq' with context parallelism." + ) + cu_seqlens_padded = _extend_last_padded_sequence(cu_seqlens_padded, target_len) + else: + dummy_seq_len = target_len - actual_len + if cp_size > 1 and cp_partition_mode == "zigzag": + assert dummy_seq_len % (2 * cp_size) == 0, ( + f"THD dummy padding length ({dummy_seq_len}) must be divisible by " + f"2 * context_parallel_size ({2 * cp_size}) for zigzag partitioning." + ) + real_cu_seqlens = cu_seqlens + if torch.equal(cu_seqlens, cu_seqlens_padded): + cu_seqlens = _append_dummy_seq(cu_seqlens, target_len) + else: + # Gaps already exist between real sequences; the dummy's valid and + # physical lengths are both exactly the new tail length. + cu_seqlens = _append_dummy_seq( + cu_seqlens, int(cu_seqlens[-1].item()) + dummy_seq_len + ) + cu_seqlens_padded = _append_dummy_seq(cu_seqlens_padded, target_len) + + target_entries = max_num_seqs + 1 + return ( + _pad_cu_seqlens(cu_seqlens, target_entries), + _pad_cu_seqlens(cu_seqlens_padded, target_entries), + real_cu_seqlens, + ) + + def _resolve_thd_padding_lengths( tokens: Optional[Tensor], labels: Optional[Tensor], diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 4091b647103..58261444a02 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -3464,6 +3464,13 @@ def _scope_to_str(s): f"({self.max_seqlen_per_dp_cp_rank}), got {self.pad_packed_seq_alignment}." ) + if self.thd_static_packing and self.thd_max_packed_sequences is None: + raise ValueError( + "thd_static_packing requires --thd-max-packed-sequences: it fixes the " + "cu_seqlens entry count (thd_max_packed_sequences + 1) that makes the " + "THD metadata shape static." + ) + # 'extend_last' THD tail padding with context parallelism requires the # global metadata to be extended before CP slicing, which only the # sequence-packing scheduler path performs. Restrict this combination. diff --git a/megatron/training/training.py b/megatron/training/training.py index 061211ce56e..bcc0302fe64 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -391,6 +391,41 @@ def print_datetime(string, override_timestamp=None): print_rank_0(f'[{string}] datetime: {time_str} ') +def _mdp_greedy_consumed_samples(args): + """Real samples this iteration under ``--mdp-greedy-packing``, summed over DP. + + Returns ``None`` when MDP greedy packing is off, so the caller keeps the + closed-form ``dp * mbs * num_microbatches``. Under greedy packing that form + is wrong: each rank fills a fixed number of bins to a token budget and + consumes however many samples that takes, which differs per iteration and + per rank. + + The MDP runtime keeps a cumulative per-rank counter; this returns the delta + since the previous call, all-reduced over the data-parallel group. One small + D2H per iteration on the logging path, next to the existing host-side + ``consumed_train_samples`` arithmetic. + """ + if not getattr(args, "mdp_enable", False) or not getattr(args, "mdp_greedy_packing", False): + return None + from megatron.core.mdp import integration as mdp_integration + + runtime = mdp_integration.get_runtime() + if runtime is None: + return None + consumed = runtime.consumed_samples() + if consumed is None: + return None + previous = getattr(args, "_mdp_prev_consumed_samples", 0) + args._mdp_prev_consumed_samples = consumed + delta = torch.tensor( + [consumed - previous], + dtype=torch.long, + device=torch.cuda.current_device() if torch.cuda.is_available() else 'cpu', + ) + torch.distributed.all_reduce(delta, group=mpu.get_data_parallel_group()) + return int(delta.item()) + + def update_seqlen_stats_from_cu_seqlens(cu_seqlens): """Add ``sum(L_i)`` and ``sum(L_i ** 2)`` from one micro-batch's REAL ``cu_seqlens``. @@ -4641,6 +4676,14 @@ def trace_handler(p): else: batch_size = _dp_world_size() * args.micro_batch_size * get_num_microbatches() iteration_sequences = batch_size + mdp_consumed = _mdp_greedy_consumed_samples(args) + if mdp_consumed is not None: + # --mdp-greedy-packing fills a fixed number of bins to a token + # budget, so the number of samples consumed floats per iteration + # AND per DP rank. The closed form above is simply wrong then; + # report the real all-reduced count instead. Same hook shape as + # rl_utils.get_iteration_sequence_count above. + iteration_sequences = mdp_consumed # Update consumed samples (always means sequences now) args.consumed_train_samples += iteration_sequences diff --git a/tests/unit_tests/mdp/test_config.py b/tests/unit_tests/mdp/test_config.py index c9ab58d945d..2c839f18d0e 100644 --- a/tests/unit_tests/mdp/test_config.py +++ b/tests/unit_tests/mdp/test_config.py @@ -244,3 +244,66 @@ def test_snapshot_reports_the_real_rank_order(): assert remapped_options.rank_order == "tp-cp-ep-pp-dp" with pytest.raises(MdpConfigurationError, match="rank_order"): validate_mdp_config(MdpConfig(enable=True), remapped_options) + + +# --------------------------------------------------------------------------- +# Packing: greedy token budget and the MCore scheduler rejection +# --------------------------------------------------------------------------- + + +def test_mcore_packing_scheduler_is_rejected(): + # Not merely untested: training.py wraps the data iterator whenever this is + # set, and DpBalancedScheduler.run then asserts on GPT-only sample keys and + # drops pixel_values / image_grid_thw. + with pytest.raises(MdpConfigurationError, match="sequence_packing_scheduler"): + validate_mdp_config( + MdpConfig(enable=True), _options(sequence_packing_scheduler="dp_balanced") + ) + + +def test_greedy_packing_requires_a_token_budget(): + with pytest.raises(MdpConfigurationError, match="max_seqlen_per_dp_cp_rank"): + validate_mdp_config(MdpConfig(enable=True, greedy_packing=True), _options()) + + +def test_greedy_packing_accepts_a_valid_budget(): + validate_mdp_config( + MdpConfig(enable=True, greedy_packing=True), + _options(max_seqlen_per_dp_cp_rank=8192, thd_max_packed_sequences=8), + ) + + +def test_greedy_budget_must_match_the_collator_row_alignment(): + # SP splits the packed rows across TP, so the budget must divide by TP. + with pytest.raises(MdpConfigurationError, match="row alignment"): + validate_mdp_config( + MdpConfig(enable=True, greedy_packing=True), + _options( + tensor_parallel_size=4, + sequence_parallel=True, + max_seqlen_per_dp_cp_rank=8190, + ), + ) + + +def test_greedy_packing_rejects_a_zero_sequence_cap(): + with pytest.raises(MdpConfigurationError, match="thd_max_packed_sequences"): + validate_mdp_config( + MdpConfig(enable=True, greedy_packing=True), + _options(max_seqlen_per_dp_cp_rank=8192, thd_max_packed_sequences=0), + ) + + +def test_greedy_packing_is_independent_of_static_packing(): + # Task 2 needs greedy + eager integer alignment (no static pad) as its + # honest baseline, so all four corners of the 2x2 must validate. + for greedy in (False, True): + for static in (False, True): + validate_mdp_config( + MdpConfig(enable=True, greedy_packing=greedy), + _options( + thd_static_packing=static, + max_seqlen_per_dp_cp_rank=8192, + thd_max_packed_sequences=8, + ), + ) diff --git a/tests/unit_tests/mdp/test_packing.py b/tests/unit_tests/mdp/test_packing.py new file mode 100644 index 00000000000..1f130965257 --- /dev/null +++ b/tests/unit_tests/mdp/test_packing.py @@ -0,0 +1,157 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Greedy token-budget packing. Pure compute: no distributed state, no CUDA.""" + +import pytest +import torch + +from megatron.core.mdp.errors import MdpConfigurationError, MdpStateError +from megatron.core.mdp.packing import ( + GreedySampleStream, + decoder_sample_length, + greedy_bin_sizes, +) + + +def _sample(length, tag=0): + return {"input_ids": torch.zeros(length, dtype=torch.long), "tag": tag} + + +def _microbatches(lengths, mbs): + """Emulate the dataloader: identity collate over an MBS-sized batch_sampler.""" + samples = [_sample(length, tag=i) for i, length in enumerate(lengths)] + return iter([samples[i : i + mbs] for i in range(0, len(samples), mbs)]) + + +def _bin_lengths(bins): + return [[int(s["input_ids"].shape[0]) for s in b] for b in bins] + + +# --------------------------------------------------------------------------- +# greedy_bin_sizes +# --------------------------------------------------------------------------- + + +def test_bins_respect_the_token_budget(): + lengths = [400, 400, 400, 100, 900] + sizes = greedy_bin_sizes(lengths, token_budget=1000, max_num_seqs=None, num_bins=2) + # 400+400 fits, +400 would be 1200 -> close. Second bin: 400+100+900 = 1400 -> 400+100. + assert sizes == [2, 2] + + +def test_bins_respect_the_sequence_cap(): + lengths = [10] * 20 + sizes = greedy_bin_sizes(lengths, token_budget=1000, max_num_seqs=4, num_bins=3) + assert sizes == [4, 4, 4] + + +def test_exactly_num_bins_are_produced(): + lengths = [100] * 50 + sizes = greedy_bin_sizes(lengths, token_budget=250, max_num_seqs=8, num_bins=5) + assert len(sizes) == 5 + assert all(size == 2 for size in sizes) + + +def test_no_bin_exceeds_the_budget(): + lengths = [137, 998, 5, 640, 640, 1, 512, 512, 511] + sizes = greedy_bin_sizes(lengths, token_budget=1024, max_num_seqs=8, num_bins=3) + cursor = 0 + for size in sizes: + assert sum(lengths[cursor : cursor + size]) <= 1024 + cursor += size + assert cursor <= len(lengths) + + +def test_exhausted_stream_is_a_named_error(): + with pytest.raises(MdpStateError, match="non-empty bins per iteration"): + greedy_bin_sizes([10, 10], token_budget=1000, max_num_seqs=None, num_bins=3) + + +# --------------------------------------------------------------------------- +# GreedySampleStream +# --------------------------------------------------------------------------- + + +def test_stream_drains_microbatch_lists_sample_by_sample(): + lengths = [300, 300, 300, 300, 300, 300] + stream = GreedySampleStream( + _microbatches(lengths, mbs=4), token_budget=900, length_of=decoder_sample_length + ) + bins = [next(stream), next(stream)] + # The 4-sample list is split across bins; the leftover carries forward. + assert _bin_lengths(bins) == [[300, 300, 300], [300, 300, 300]] + assert stream.consumed_samples == 6 + + +def test_leftovers_carry_across_iterations(): + lengths = [500] * 8 + stream = GreedySampleStream( + _microbatches(lengths, mbs=4), token_budget=1000, length_of=decoder_sample_length + ) + first_iteration = [next(stream), next(stream)] + second_iteration = [next(stream), next(stream)] + tags = [[s["tag"] for s in b] for b in first_iteration + second_iteration] + assert tags == [[0, 1], [2, 3], [4, 5], [6, 7]] + + +def test_alignment_is_charged_against_the_budget(): + # Aligned to 8, a 5-token sample occupies 8 rows: 3 fit in 24, not 4. + stream = GreedySampleStream( + _microbatches([5] * 8, mbs=8), + token_budget=24, + align=8, + length_of=decoder_sample_length, + ) + assert _bin_lengths([next(stream)]) == [[5, 5, 5]] + + +def test_sequence_cap_closes_the_bin(): + stream = GreedySampleStream( + _microbatches([10] * 10, mbs=10), + token_budget=10_000, + max_num_seqs=3, + length_of=decoder_sample_length, + ) + assert len(next(stream)) == 3 + + +def test_partial_bin_at_end_of_stream_then_stop(): + stream = GreedySampleStream( + _microbatches([400] * 3, mbs=3), token_budget=1000, length_of=decoder_sample_length + ) + assert len(next(stream)) == 2 + assert len(next(stream)) == 1 # partial, never empty + with pytest.raises(StopIteration): + next(stream) + + +def test_oversized_sample_names_the_flag(): + stream = GreedySampleStream( + _microbatches([5000], mbs=1), token_budget=1024, length_of=decoder_sample_length + ) + with pytest.raises(MdpStateError, match="max-seqlen-per-dp-cp-rank"): + next(stream) + + +def test_budget_must_be_divisible_by_the_row_alignment(): + with pytest.raises(MdpConfigurationError, match="row alignment"): + GreedySampleStream( + _microbatches([10], mbs=1), + token_budget=100, + align=8, + length_of=decoder_sample_length, + ) + + +def test_degenerate_distribution_reproduces_fixed_mbs(): + """min=max=mean=L with budget k*L packs exactly k samples per bin. + + This is the exact-equivalence configuration: greedy is then bit-identical + to today's fixed ``--micro-batch-size k``. + """ + L, k = 256, 4 + stream = GreedySampleStream( + _microbatches([L] * 32, mbs=1), token_budget=k * L, length_of=decoder_sample_length + ) + for _ in range(8): + assert len(next(stream)) == k From 4813e7ec61da21f01b27472d353bd249953b5ee6 Mon Sep 17 00:00:00 2001 From: Li Tao Date: Thu, 20 Aug 2026 08:19:25 -0700 Subject: [PATCH 2/5] fix(mdp): reserve a cu_seqlens slot for the static padding tail Strict-review follow-ups on the greedy/static packing branch. CRITICAL: the greedy sequence cap was thd_max_packed_sequences verbatim, but under --thd-static-packing the padding tail is appended to cu_seqlens as an ordinary dummy sequence, so a bin filled to the cap needs cap + 2 entries and overflows the cap + 1 capacity inside _pad_cu_seqlens. dp_balanced already reserves that slot (_get_scheduler_max_real_num_seqs); greedy now does too, via greedy_max_real_sequences(), and thd_max_packed_sequences < 2 is rejected under static packing. Latent until a bin happened to hit the cap. IMPORTANT: consumed_samples() summed every greedy stream, so an evaluation pass leaked into the next iteration's consumed_train_samples delta. Streams are now tagged train/eval at creation and eval is excluded. The one-iteration shift --mdp-overlap-window-capture introduces is documented. Also: drop greedy_bin_sizes (no runtime use path; its assertions moved onto GreedySampleStream), and build the mock scenario pool once per provider call instead of four times. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Li Tao --- examples/multimodal_dev/data/mdp_mock.py | 12 ++--- megatron/core/mdp/config.py | 28 ++++++++++-- megatron/core/mdp/integration.py | 4 +- megatron/core/mdp/packing.py | 33 -------------- megatron/core/mdp/runtime.py | 48 +++++++++++++------- tests/unit_tests/mdp/test_config.py | 26 +++++++++++ tests/unit_tests/mdp/test_packing.py | 57 +++++++++--------------- 7 files changed, 114 insertions(+), 94 deletions(-) diff --git a/examples/multimodal_dev/data/mdp_mock.py b/examples/multimodal_dev/data/mdp_mock.py index d0d888e020f..a9314c1a787 100644 --- a/examples/multimodal_dev/data/mdp_mock.py +++ b/examples/multimodal_dev/data/mdp_mock.py @@ -184,7 +184,7 @@ def _text(length): GREEDY_SAMPLE_SAFETY = 1.5 -def _greedy_sample_scale(args, length_config): +def _greedy_sample_scale(args, scenarios): """Scale factor for the synthetic dataset length under greedy packing. Megatron sizes the dataset as ``train_iters * global_batch_size`` samples. @@ -201,8 +201,7 @@ def _greedy_sample_scale(args, length_config): return 1.0 budget = int(args.max_seqlen_per_dp_cp_rank) * int(args.context_parallel_size) cap = getattr(args, "thd_max_packed_sequences", None) - pool = build_scenarios(length_config=length_config) - mean_len = sum(scenario_totals(s)[0] for s in pool) / len(pool) + mean_len = sum(scenario_totals(s)[0] for s in scenarios) / len(scenarios) samples_per_bin = budget / mean_len if cap: samples_per_bin = min(samples_per_bin, float(cap)) @@ -219,12 +218,15 @@ def train_valid_test_datasets_provider(train_val_test_num_samples): from megatron.training.datasets.utils import load_json_arg length_config = load_json_arg(length_config) + # Built once and shared: the pool is a deterministic function of the length + # config, and drawing it re-samples a million lognormal lengths. + scenarios = build_scenarios(length_config=length_config) kwargs = dict( vocab_size=getattr(args, "padded_vocab_size", 1024), image_token_id=getattr(args, "image_token_id", QWEN35_VL_IMAGE_TOKEN_ID), - length_config=length_config, + scenarios=scenarios, ) - scale = _greedy_sample_scale(args, length_config) + scale = _greedy_sample_scale(args, scenarios) return tuple( MdpThdMockDataset(num_samples=math.ceil(n * scale), seed=1234 + split, **kwargs) for split, n in enumerate(train_val_test_num_samples) diff --git a/megatron/core/mdp/config.py b/megatron/core/mdp/config.py index 128bce8b46c..f3630be9f0b 100644 --- a/megatron/core/mdp/config.py +++ b/megatron/core/mdp/config.py @@ -328,6 +328,23 @@ def validate_mdp_config(config: MdpConfig, options: MdpCompatibilityOptions) -> ) +def greedy_max_real_sequences(options: "MdpCompatibilityOptions") -> Optional[int]: + """Real sequences a greedy bin may hold, or ``None`` for no cap. + + ``thd_max_packed_sequences`` is the *final* static THD capacity. Under + ``--thd-static-packing`` the padding tail is represented as an ordinary + dummy sequence appended to ``cu_seqlens``, so one slot must be reserved for + it -- exactly what ``_get_scheduler_max_real_num_seqs`` does for + ``dp_balanced``. Without the reservation a bin filled to the cap overflows + the ``thd_max_packed_sequences + 1`` entry budget and dies inside + ``_pad_cu_seqlens``. + """ + cap = options.thd_max_packed_sequences + if cap is None: + return None + return int(cap) - 1 if options.thd_static_packing else int(cap) + + def _validate_packing(config: MdpConfig, options: MdpCompatibilityOptions) -> None: """Reject packing configurations MDP cannot honor. @@ -370,12 +387,17 @@ def _validate_packing(config: MdpConfig, options: MdpCompatibilityOptions) -> No "A bin filled to the budget must still split legally across CP/SP ranks; " "discovering this inside TransformerEngine gives a far worse error.", ) - if options.thd_max_packed_sequences is not None and options.thd_max_packed_sequences < 1: + minimum = 2 if options.thd_static_packing else 1 + if ( + options.thd_max_packed_sequences is not None + and options.thd_max_packed_sequences < minimum + ): _reject( "thd_max_packed_sequences", options.thd_max_packed_sequences, - "thd_max_packed_sequences >= 1", - "It caps the real sequences per greedy bin.", + f"thd_max_packed_sequences >= {minimum}", + "It caps the real sequences per greedy bin; under --thd-static-packing " + "one slot is reserved for the padding tail's dummy sequence.", "8", ) diff --git a/megatron/core/mdp/integration.py b/megatron/core/mdp/integration.py index cef0913cfcf..b28afd493c8 100644 --- a/megatron/core/mdp/integration.py +++ b/megatron/core/mdp/integration.py @@ -29,6 +29,7 @@ SUPPORTED_RANK_ORDER, MdpCompatibilityOptions, MdpConfig, + greedy_max_real_sequences, thd_row_alignment, validate_mdp_config, ) @@ -239,6 +240,7 @@ def maybe_build_mdp_domain(*, args, model, optimizer, optimizer_config, ddp_conf if mdp_config.greedy_packing else None ) + greedy_max_num_seqs = greedy_max_real_sequences(compat) _RUNTIME = MdpRuntime( config=mdp_config, rank_map=rank_map, @@ -259,7 +261,7 @@ def maybe_build_mdp_domain(*, args, model, optimizer, optimizer_config, ddp_conf params_dtype=params_dtype, num_vpp_chunks=len(model), greedy_token_budget=greedy_token_budget, - greedy_max_num_seqs=getattr(args, "thd_max_packed_sequences", None), + greedy_max_num_seqs=greedy_max_num_seqs, greedy_row_alignment=thd_row_alignment(compat), ) logger.info( diff --git a/megatron/core/mdp/packing.py b/megatron/core/mdp/packing.py index f7208365b76..d0b2bcb9036 100644 --- a/megatron/core/mdp/packing.py +++ b/megatron/core/mdp/packing.py @@ -46,39 +46,6 @@ def decoder_sample_length(sample: Any) -> int: ) from exc -def greedy_bin_sizes( - lengths: List[int], *, token_budget: int, max_num_seqs: Optional[int], num_bins: int -) -> List[int]: - """Number of samples taken into each of ``num_bins`` in-order greedy bins. - - Pure function over a length vector, so the grouping rule is testable without - a dataset. ``lengths`` must already be aligned (see - :class:`GreedySampleStream`). Raises when ``lengths`` cannot fill - ``num_bins`` bins. - """ - sizes = [] - cursor = 0 - for _ in range(num_bins): - taken = 0 - total = 0 - while cursor + taken < len(lengths): - length = lengths[cursor + taken] - if taken and total + length > token_budget: - break - if max_num_seqs is not None and taken >= max_num_seqs: - break - total += length - taken += 1 - if taken == 0: - raise MdpStateError( - f"MDP: greedy packing violates: {num_bins} non-empty bins per iteration " - f"(the sample stream ran out after {len(sizes)})." - ) - sizes.append(taken) - cursor += taken - return sizes - - class GreedySampleStream: """Wrap a microbatch-list iterator so ``next()`` returns one greedy bin. diff --git a/megatron/core/mdp/runtime.py b/megatron/core/mdp/runtime.py index 8525a0a887e..1e56108abe2 100644 --- a/megatron/core/mdp/runtime.py +++ b/megatron/core/mdp/runtime.py @@ -523,28 +523,44 @@ def _first_iterator(data_iterators): def _greedy_stream(self, data_iterators): """The greedy sample stream for this data iterator, created on first use.""" iterator = self._first_iterator(data_iterators) - stream = self._greedy_streams.get(id(iterator)) - if stream is None: - stream = GreedySampleStream( - iterator, - token_budget=self._greedy_token_budget, - max_num_seqs=self._greedy_max_num_seqs, - align=self._greedy_row_alignment, - length_of=decoder_sample_length, + entry = self._greedy_streams.get(id(iterator)) + if entry is None: + entry = ( + GreedySampleStream( + iterator, + token_budget=self._greedy_token_budget, + max_num_seqs=self._greedy_max_num_seqs, + align=self._greedy_row_alignment, + length_of=decoder_sample_length, + ), + # Evaluation runs forward_only; recorded so its consumption is + # kept out of consumed_train_samples. + self._forward_only, ) - self._greedy_streams[id(iterator)] = stream - return stream + self._greedy_streams[id(iterator)] = entry + return entry[0] def consumed_samples(self) -> Optional[int]: - """Real samples drained by greedy packing so far, or ``None`` when off. - - Sums every stream (train and eval) this rank owns. ``training.py`` reads - the delta per iteration because the closed form - ``dp x mbs x num_microbatches`` is wrong under greedy packing. + """Real samples drained by greedy *training* packing, or ``None`` when off. + + ``training.py`` reads the delta per iteration because the closed form + ``dp x mbs x num_microbatches`` is wrong under greedy packing. Evaluation + streams are excluded: they consume their own samples, and folding them in + would charge an eval pass to the next training iteration. + + Under ``--mdp-overlap-window-capture`` the count is shifted by one + iteration -- the prefetch thread drains iteration i+1's samples during + iteration i, and the final prefetch is captured but never consumed. + Per-iteration values are therefore approximate with overlap on; the + running total is not. """ if not self.config.greedy_packing: return None - return sum(stream.consumed_samples for stream in self._greedy_streams.values()) + return sum( + stream.consumed_samples + for stream, forward_only in self._greedy_streams.values() + if not forward_only + ) def _capture_window(self, data_iterators, num_microbatches: int) -> MdpIterationWindow: if not self.config.greedy_packing: diff --git a/tests/unit_tests/mdp/test_config.py b/tests/unit_tests/mdp/test_config.py index 2c839f18d0e..ba510fad057 100644 --- a/tests/unit_tests/mdp/test_config.py +++ b/tests/unit_tests/mdp/test_config.py @@ -12,6 +12,7 @@ MdpCompatibilityOptions, MdpConfig, apply_vision_config_overrides, + greedy_max_real_sequences, validate_mdp_config, ) from megatron.core.mdp.errors import MdpConfigurationError @@ -307,3 +308,28 @@ def test_greedy_packing_is_independent_of_static_packing(): thd_max_packed_sequences=8, ), ) + + +def test_static_packing_reserves_a_sequence_slot_for_the_padding_tail(): + # thd_max_packed_sequences is the FINAL cu_seqlens capacity. Under static + # packing the tail becomes an ordinary dummy sequence, so a bin filled to + # the full cap would need cap + 2 entries and die inside _pad_cu_seqlens. + eager = _options(max_seqlen_per_dp_cp_rank=8192, thd_max_packed_sequences=8) + static = _options( + max_seqlen_per_dp_cp_rank=8192, thd_max_packed_sequences=8, thd_static_packing=True + ) + assert greedy_max_real_sequences(eager) == 8 + assert greedy_max_real_sequences(static) == 7 + assert greedy_max_real_sequences(_options()) is None + + +def test_static_packing_needs_room_for_a_real_sequence_and_the_dummy(): + with pytest.raises(MdpConfigurationError, match="thd_max_packed_sequences >= 2"): + validate_mdp_config( + MdpConfig(enable=True, greedy_packing=True), + _options( + max_seqlen_per_dp_cp_rank=8192, + thd_max_packed_sequences=1, + thd_static_packing=True, + ), + ) diff --git a/tests/unit_tests/mdp/test_packing.py b/tests/unit_tests/mdp/test_packing.py index 1f130965257..7401fa829de 100644 --- a/tests/unit_tests/mdp/test_packing.py +++ b/tests/unit_tests/mdp/test_packing.py @@ -6,11 +6,7 @@ import torch from megatron.core.mdp.errors import MdpConfigurationError, MdpStateError -from megatron.core.mdp.packing import ( - GreedySampleStream, - decoder_sample_length, - greedy_bin_sizes, -) +from megatron.core.mdp.packing import GreedySampleStream, decoder_sample_length def _sample(length, tag=0): @@ -28,48 +24,37 @@ def _bin_lengths(bins): # --------------------------------------------------------------------------- -# greedy_bin_sizes +# GreedySampleStream # --------------------------------------------------------------------------- -def test_bins_respect_the_token_budget(): - lengths = [400, 400, 400, 100, 900] - sizes = greedy_bin_sizes(lengths, token_budget=1000, max_num_seqs=None, num_bins=2) - # 400+400 fits, +400 would be 1200 -> close. Second bin: 400+100+900 = 1400 -> 400+100. - assert sizes == [2, 2] - - -def test_bins_respect_the_sequence_cap(): - lengths = [10] * 20 - sizes = greedy_bin_sizes(lengths, token_budget=1000, max_num_seqs=4, num_bins=3) - assert sizes == [4, 4, 4] +def _stream(lengths, *, mbs, budget, cap=None, align=1): + return GreedySampleStream( + _microbatches(lengths, mbs), + token_budget=budget, + max_num_seqs=cap, + align=align, + length_of=decoder_sample_length, + ) -def test_exactly_num_bins_are_produced(): - lengths = [100] * 50 - sizes = greedy_bin_sizes(lengths, token_budget=250, max_num_seqs=8, num_bins=5) - assert len(sizes) == 5 - assert all(size == 2 for size in sizes) +def test_bins_respect_the_token_budget(): + stream = _stream([400, 400, 400, 100, 900], mbs=5, budget=1000) + # 400+400 fits; +400 would be 1200 -> close. Then 400+100, +900 -> close. + assert _bin_lengths([next(stream), next(stream)]) == [[400, 400], [400, 100]] def test_no_bin_exceeds_the_budget(): lengths = [137, 998, 5, 640, 640, 1, 512, 512, 511] - sizes = greedy_bin_sizes(lengths, token_budget=1024, max_num_seqs=8, num_bins=3) - cursor = 0 - for size in sizes: - assert sum(lengths[cursor : cursor + size]) <= 1024 - cursor += size - assert cursor <= len(lengths) - + stream = _stream(lengths, mbs=3, budget=1024, cap=8) + for _ in range(3): + assert sum(decoder_sample_length(s) for s in next(stream)) <= 1024 -def test_exhausted_stream_is_a_named_error(): - with pytest.raises(MdpStateError, match="non-empty bins per iteration"): - greedy_bin_sizes([10, 10], token_budget=1000, max_num_seqs=None, num_bins=3) - -# --------------------------------------------------------------------------- -# GreedySampleStream -# --------------------------------------------------------------------------- +def test_exactly_num_bins_are_produced(): + stream = _stream([100] * 50, mbs=10, budget=250, cap=8) + bins = [next(stream) for _ in range(5)] + assert [len(b) for b in bins] == [2] * 5 def test_stream_drains_microbatch_lists_sample_by_sample(): From 8a5136198777840d9835f9383dbbfa943ad99064 Mon Sep 17 00:00:00 2001 From: Li Tao Date: Thu, 20 Aug 2026 18:43:48 -0700 Subject: [PATCH 3/5] test(mdp): assert the derived pad_between_seqs, not the old hardcoded True The static-shape test still encoded the first implementation's pad_between_seqs=True. The shipped value is derived from the collator's row alignment, so at TP=CP=1 it is False -- and provably so: cu_seqlens and cu_seqlens_padded coincide, which the test now also asserts. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Li Tao --- examples/multimodal_dev/tests/test_thd_static_packing.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/multimodal_dev/tests/test_thd_static_packing.py b/examples/multimodal_dev/tests/test_thd_static_packing.py index 10326f03c2d..59ae1521ac5 100644 --- a/examples/multimodal_dev/tests/test_thd_static_packing.py +++ b/examples/multimodal_dev/tests/test_thd_static_packing.py @@ -170,7 +170,12 @@ def test_heterogeneous_batch_has_fixed_shapes(self, static_args, mdp_enable, sam assert getattr(psp, name).numel() == MAX_PACKED + 1, name assert psp.max_seqlen_q == MAX_SEQLEN assert psp.max_seqlen_kv == MAX_SEQLEN - assert psp.pad_between_seqs is True + # Derived from the collator's row alignment, not hardcoded: at TP=CP=1 no + # sample is ever padded, so cu_seqlens and cu_seqlens_padded coincide and + # there is provably no gap. Claiming True here is not free -- TE disables + # FlashAttention for THD whenever padding may exist between sequences. + assert psp.pad_between_seqs is False + assert torch.equal(psp.cu_seqlens_q, psp.cu_seqlens_q_padded) def test_shapes_do_not_depend_on_the_batch(self, static_args, mdp_enable, sample_device): static_args(mdp_enable=mdp_enable) From 5200c89a3d549cb8ba5d2977133f60b61f660034 Mon Sep 17 00:00:00 2001 From: Li Tao Date: Thu, 20 Aug 2026 18:49:24 -0700 Subject: [PATCH 4/5] test(mdp): assert FLOPs honesty through the emitted override The static tail policy is append_dummy_seq, so the pad does land in cu_seqlens_q and the collator emits the pre-tail vector as flops_cu_seqlens. The test still asserted the earlier extend_last behavior ('no override emitted'). Assert the property that actually matters instead: the override sums to the real token count, using cu_seqlens_q would count the whole static pad, and accumulate_flops_stats consumes the override. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Li Tao --- .../tests/test_thd_static_packing.py | 34 +++++++++++++++---- 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/examples/multimodal_dev/tests/test_thd_static_packing.py b/examples/multimodal_dev/tests/test_thd_static_packing.py index 59ae1521ac5..04290ade521 100644 --- a/examples/multimodal_dev/tests/test_thd_static_packing.py +++ b/examples/multimodal_dev/tests/test_thd_static_packing.py @@ -243,18 +243,40 @@ def test_static_packing_off_is_byte_identical(static_args): class TestFlopsAccounting: - def test_extend_last_reports_only_real_tokens(self, static_args): + def test_the_emitted_override_reports_only_real_tokens(self, static_args): + """The accumulator must see the real tokens, never the static pad. + + Under static packing the tail is an ``append_dummy_seq`` sequence, so it + lands in ``cu_seqlens_q`` itself; the collator emits the pre-tail vector + separately and ``accumulate_flops_stats`` prefers it. + """ static_args() lengths = [5, 13, 7] batch = [_make_sample(L, base=i * 1000) for i, L in enumerate(lengths)] packed = fs.pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") - # extend_last leaves cu_seqlens_q compact, so no override is emitted and - # the accumulator sees exactly the real tokens. - assert "flops_cu_seqlens" not in packed + override = packed["flops_cu_seqlens"] + assert (override[1:] - override[:-1]).sum().item() == sum(lengths) + + # Without the override the pad would be counted as real work. cu = packed["packed_seq_params"].cu_seqlens_q - seqlens = (cu[1:] - cu[:-1]).tolist() - assert sum(seqlens) == sum(lengths) + assert (cu[1:] - cu[:-1]).clamp(min=0).sum().item() == MAX_SEQLEN + + def test_the_accumulator_prefers_the_override(self, static_args, monkeypatch): + # accumulate_flops_stats imports this lazily from the training module. + import megatron.training.training as training + + seen = [] + monkeypatch.setattr( + training, "update_seqlen_stats_from_cu_seqlens", lambda cu: seen.append(cu) + ) + static_args() + batch = [_make_sample(L, base=i * 1000) for i, L in enumerate([5, 13, 7])] + packed = fs.pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + fs.accumulate_flops_stats( + packed["packed_seq_params"], real_cu_seqlens=packed["flops_cu_seqlens"] + ) + assert seen and torch.equal(seen[0], packed["flops_cu_seqlens"]) def test_dummy_tail_would_pollute_without_the_override(self): """The override exists because the dummy tail lands in cu_seqlens_q.""" From 3c8e72cd5c596a8218d766c254103351c712c24e Mon Sep 17 00:00:00 2001 From: Li Tao Date: Thu, 20 Aug 2026 18:53:12 -0700 Subject: [PATCH 5/5] feat(mdp): run per-layer CUDA graphs on the MDP decoder Three changes make per-layer ('partial') CUDA graphs actually replay on MDP's decoder, and one makes them worth enabling. 1. Widen the predicate. GraphableMegatronModule._is_thd_cuda_graph asked 'is MCore's packing scheduler configured' when it needs to know 'does the incoming THD batch have fixed shapes'. Those coincided only while the scheduler was the sole fixed-shape producer; --thd-static-packing is a third. One shared packed_seq_params.thd_shapes_are_static(config) now backs the four conflated call sites (module.py, transformer_config.py's THD CUDA-graph gate, the same gate in arguments.py, and the MoE dispatcher restriction). _get_thd_varlen_max_num_microbatches deliberately stays dp_balanced-specific. 2. Fix layer discovery. TECudaGraphHelper._discover_layers resolves layers with get_attr_wrapped_model(chunk, 'decoder'), which unwraps only through '.module'; MDP's decoder lives at MultimodalModel.language_model, so the lookup raised, the helper swallowed its own error at DEBUG level, and it captured ZERO layers while reporting success. MultimodalModel now forwards decoder / mtp / rotary_pos_emb / position_embedding_type, and core raises when every chunk fails that lookup. 3. Rewrite the MDP gate. full_iteration stays rejected (it would capture P4 itself); per-layer graphs are accepted only under --thd-static-packing, and still rejected together with --mdp-overlap-window-capture. 4. Stop forcing pad_between_seqs=True in the graph path. It cannot be a graph input and must not be inferred from CUDA tensors during capture, but it does not have to be guessed: under --thd-static-packing it is a function of the collator's row alignment alone. Forcing True is not free -- TE disables FlashAttention for THD whenever padding may exist between sequences, and on this configuration cuDNN fused attention reports 'no backend supports the provided input', so TE fell through to its unfused O(T^2) backend and OOM'd at 32 GiB. With the value derived, attn graphs replay. Measured on 4xGB300, 20 iterations, lognormal lengths, budget 8192: attn graphs 194.6 ms/iter and 320.0 TFLOP/s/GPU vs 215.7 ms and 288.6 for eager with the same static pad -- ~10% faster, loss and grad norm matching. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Li Tao --- examples/multimodal_dev/forward_step.py | 25 +- examples/multimodal_dev/models/base.py | 36 +++ .../tests/test_mdp_cuda_graph.py | 292 ++++++++++++++++++ megatron/core/mdp/README.md | 5 +- megatron/core/mdp/config.py | 102 +++++- megatron/core/mdp/integration.py | 7 +- megatron/core/mdp/knowledge.md | 29 ++ megatron/core/packed_seq_params.py | 78 +++++ megatron/core/transformer/cuda_graphs.py | 25 ++ megatron/core/transformer/module.py | 15 +- .../core/transformer/transformer_config.py | 20 +- .../core/transformer/transformer_layer.py | 17 +- megatron/training/arguments.py | 5 +- tests/unit_tests/mdp/test_config.py | 71 ++++- .../test_thd_static_shape_predicate.py | 53 ++++ 15 files changed, 728 insertions(+), 52 deletions(-) create mode 100644 examples/multimodal_dev/tests/test_mdp_cuda_graph.py create mode 100644 tests/unit_tests/transformer/test_thd_static_shape_predicate.py diff --git a/examples/multimodal_dev/forward_step.py b/examples/multimodal_dev/forward_step.py index 0e0560b74da..a72a384ebe6 100644 --- a/examples/multimodal_dev/forward_step.py +++ b/examples/multimodal_dev/forward_step.py @@ -12,7 +12,11 @@ from examples.multimodal_dev.observability import nvtx_phase from megatron.core import mpu -from megatron.core.packed_seq_params import PackedSeqParams, build_static_thd_metadata +from megatron.core.packed_seq_params import ( + PackedSeqParams, + build_static_thd_metadata, + thd_collate_row_alignment, +) from megatron.core.parallel_state import ( get_tensor_model_parallel_group, get_tensor_model_parallel_rank, @@ -423,10 +427,11 @@ def pack_or_pad_batch( except AssertionError: has_sp = False - if cp_size > 1: - divisible_by = (tp_size * cp_size * 2) if has_sp else (cp_size * 2) - else: - divisible_by = tp_size if has_sp else 1 + divisible_by = thd_collate_row_alignment( + context_parallel_size=cp_size, + tensor_model_parallel_size=tp_size, + sequence_parallel=has_sp, + ) if pad_to_multiple is not None: divisible_by = max(divisible_by, pad_to_multiple) @@ -447,6 +452,16 @@ def pack_or_pad_batch( except AssertionError: static_args = None if static_args is not None and getattr(static_args, "thd_static_packing", False): + # pad_between_seqs below is derived from the row alignment, and the + # CUDA-graph path re-derives the same value from the config alone + # (packed_seq_params.thd_static_pad_between_seqs). An extra + # pad_to_multiple would introduce gaps the config cannot see, so the + # two would silently disagree. + assert pad_to_multiple is None, ( + "thd_static_packing is incompatible with an explicit pad_to_multiple: " + "the CUDA-graph path derives pad_between_seqs from the CP/SP row " + "alignment alone and cannot see it." + ) static_target_T = int(static_args.max_seqlen_per_dp_cp_rank) * cp_size static_max_num_seqs = int(static_args.thd_max_packed_sequences) # append_dummy_seq, matching what --sequence-packing-scheduler diff --git a/examples/multimodal_dev/models/base.py b/examples/multimodal_dev/models/base.py index f4938b5d25d..607b25c677c 100644 --- a/examples/multimodal_dev/models/base.py +++ b/examples/multimodal_dev/models/base.py @@ -158,6 +158,42 @@ def __init__( mtp_block_spec=mtp_block_spec, ) + # ------------------------------------------------------------------ + # Attributes surfaced for per-layer CUDA graph capture. + # + # ``TECudaGraphHelper._discover_layers`` resolves the graphable layers with + # ``get_attr_wrapped_model(chunk, 'decoder')``, which only unwraps through + # ``.module`` (DDP -> Float16Module -> this class). It then reads ``decoder``, + # ``mtp``, ``rotary_pos_emb``, and ``position_embedding_type`` off the object it + # stopped at. Without these forwards the helper raises, catches its own + # RuntimeError, and silently captures zero layers. Same pattern as + # ``shared_embedding_or_output_weight`` below. + # ------------------------------------------------------------------ + + @property + def decoder(self): + """The language model's transformer block.""" + return self.language_model.decoder + + @property + def mtp(self): + """The language model's MTP block. + + Raises ``AttributeError`` when MTP is off, so ``hasattr(chunk, 'mtp')`` + keeps reporting False exactly as it does for a bare ``GPTModel``. + """ + return self.language_model.mtp + + @property + def rotary_pos_emb(self): + """The language model's rotary embedding module.""" + return self.language_model.rotary_pos_emb + + @property + def position_embedding_type(self): + """The language model's position embedding type.""" + return self.language_model.position_embedding_type + def shared_embedding_or_output_weight(self): """Surface the wrapped language model's shared embedding / output weight to the PP grad-finalize step (mirrors the LLaVA wrapper). diff --git a/examples/multimodal_dev/tests/test_mdp_cuda_graph.py b/examples/multimodal_dev/tests/test_mdp_cuda_graph.py new file mode 100644 index 00000000000..0f68036842b --- /dev/null +++ b/examples/multimodal_dev/tests/test_mdp_cuda_graph.py @@ -0,0 +1,292 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Per-layer (partial) CUDA graph support for the MDP decoder. + +MDP's P4 runs the *unmodified* decoder schedule, so a per-layer CUDA graph over a +decoder submodule never interacts with the bridge: MDP only supplies the decoder's +embedding leaves (P3) and consumes their gradients (P5). What did block the feature +is that ``TECudaGraphHelper`` resolves graphable layers with +``get_attr_wrapped_model(chunk, 'decoder')``, which only unwraps through ``.module`` +and therefore stopped at :class:`MultimodalModel` and captured nothing. + +These tests run the SBHD path deliberately: graph replay needs a fixed static input +shape, and the THD side of that contract now comes from ``--thd-static-packing``, +which is an end-to-end property of the data path rather than something a unit test +can construct. What is covered here is everything that is *not* the packing +contract -- layer discovery through the multimodal wrapper, the loud failure when +discovery finds nothing, and per-layer capture/replay numerics on the decoder. The +packing contract itself is covered by the GPU end-to-end run recorded in +``agent_works/mdp-partial-cudagraph-decoder/summary_v2.md``. + +These tests need exactly 1 rank:: + + torchrun --nproc-per-node 1 -m pytest -q \\ + examples/multimodal_dev/tests/test_mdp_cuda_graph.py +""" + +import os +import sys + +import pytest +import torch + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from examples.multimodal_dev.models.base import MultimodalModel +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.num_microbatches_calculator import ( + destroy_num_microbatches_calculator, + init_num_microbatches_calculator, +) +from megatron.core.tensor_parallel.random import ( + initialize_rng_tracker, + model_parallel_cuda_manual_seed, +) +from megatron.core.transformer.cuda_graphs import HAVE_TE_GRAPHS, TECudaGraphHelper +from megatron.core.transformer.enums import CudaGraphModule +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + +if not HAVE_TE_GRAPHS: + pytest.skip( + "Per-layer CUDA graphs require a TransformerEngine build with " + "make_graphed_callables().", + allow_module_level=True, + ) + +NUM_LAYERS = 2 +HIDDEN = 128 +HEADS = 4 +VOCAB = 128 +SEQ = 32 +BATCH = 2 +IMAGE_TOKEN_ID = 7 +IMAGE_POSITIONS = (0, 1, 2, 3) +NUM_VISUAL_TOKENS = BATCH * len(IMAGE_POSITIONS) +DTYPE = torch.bfloat16 + + +class _StubVisionEncoder(MegatronModule): + """Minimal trainable stand-in for the real vision encoder.""" + + def __init__(self, config, hidden_size, dtype=DTYPE): + super().__init__(config=config) + self.proj = torch.nn.Linear(hidden_size, hidden_size, bias=False, dtype=dtype) + + def forward(self, pixel_values, image_grid_thw): + """Project pixel features to decoder-width embeddings.""" + return self.proj(pixel_values) + + +def _make_config(cuda_graph_impl="none", cuda_graph_modules=None): + return TransformerConfig( + num_layers=NUM_LAYERS, + hidden_size=HIDDEN, + ffn_hidden_size=4 * HIDDEN, + num_attention_heads=HEADS, + num_query_groups=HEADS, + bf16=True, + params_dtype=DTYPE, + pipeline_dtype=DTYPE, + hidden_dropout=0.0, + attention_dropout=0.0, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + sequence_parallel=False, + cuda_graph_impl=cuda_graph_impl, + cuda_graph_modules=( + [] if cuda_graph_modules is None else list(cuda_graph_modules) + ), + cuda_graph_warmup_steps=0, + ) + + +class _ModelChunk(torch.nn.Module): + """Stand-in for the DDP wrapper the training loop hands to the helper. + + ``TECudaGraphHelper`` unwraps ``.module`` to reach the decoder and calls + ``zero_grad_buffer()`` on the chunk once capture finishes; DDP provides both. + A real DDP is deliberately not used here: its bucket gradient sync lands + inside the captured backward and invalidates the capture on a 1-rank DP + group, which would test the wrapper rather than the graph. + """ + + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, *args, **kwargs): + """Delegate to the wrapped model.""" + return self.module(*args, **kwargs) + + def zero_grad_buffer(self): + """DDP's gradient-buffer reset; plain ``.grad`` tensors here.""" + self.zero_grad(set_to_none=True) + + +def _build_model(config): + """One model chunk, shaped the way the training loop hands it to the helper.""" + torch.manual_seed(1234) + vision = _StubVisionEncoder(config, HIDDEN) + model = MultimodalModel( + language_config=config, + language_spec=get_gpt_layer_with_transformer_engine_spec(), + vision_encoder=vision, + vocab_size=VOCAB, + max_sequence_length=SEQ, + image_token_id=IMAGE_TOKEN_ID, + position_embedding_type="rope", + parallel_output=False, + pre_process=True, + post_process=True, + ) + return _ModelChunk(model.cuda()) + + +def _grad_norm(model): + """L2 norm of every populated parameter gradient.""" + total = 0.0 + for param in model.parameters(): + if param.grad is not None: + total += param.grad.float().norm(2).item() ** 2 + return total**0.5 + + +def _make_batch(seed=1234): + """Deterministic batch, identical for every model built in this module.""" + g = torch.Generator(device="cuda") + g.manual_seed(seed) + input_ids = torch.randint(0, VOCAB, (BATCH, SEQ), generator=g, device="cuda") + input_ids[input_ids == IMAGE_TOKEN_ID] = (IMAGE_TOKEN_ID + 1) % VOCAB + for pos in IMAGE_POSITIONS: + input_ids[:, pos] = IMAGE_TOKEN_ID + labels = torch.randint(0, VOCAB, (BATCH, SEQ), generator=g, device="cuda") + loss_mask = torch.ones(BATCH, SEQ, device="cuda") + position_ids = torch.arange(SEQ, device="cuda").unsqueeze(0).expand(BATCH, -1).contiguous() + pixel_values = torch.randn( + NUM_VISUAL_TOKENS, HIDDEN, generator=g, device="cuda", dtype=DTYPE + ) + image_grid_thw = torch.tensor([[1, 2, 2]] * BATCH, device="cuda") + return input_ids, labels, loss_mask, position_ids, pixel_values, image_grid_thw + + +def _forward_backward(model): + """One forward + backward on the shared batch; returns (loss, grad norm).""" + input_ids, labels, loss_mask, position_ids, pixel_values, image_grid_thw = _make_batch() + model.zero_grad_buffer() + per_token_loss = model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + labels=labels, + loss_mask=loss_mask, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + ) + flat = per_token_loss.float().view(-1) + mask = loss_mask.float().view(-1) + loss = (flat * mask).sum() / mask.sum().clamp(min=1) + loss.backward() + return loss.item(), _grad_norm(model) + + +@pytest.fixture(scope="module", autouse=True) +def _single_rank_parallel_state(): + """TP=1/PP=1 model-parallel groups plus the microbatch calculator.""" + if torch.distributed.is_initialized() and torch.distributed.get_world_size() != 1: + pytest.skip("needs exactly 1 rank") + initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + model_parallel_cuda_manual_seed(1234) + init_num_microbatches_calculator( + rank=0, + global_batch_size=BATCH, + micro_batch_size=BATCH, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + yield + destroy_num_microbatches_calculator() + Utils.destroy_model_parallel() + + +def test_multimodal_model_surfaces_graph_discovery_attributes(): + """``TECudaGraphHelper`` reads these off the object it unwraps to.""" + model = _build_model(_make_config()) + + model = model.module + + assert model.decoder is model.language_model.decoder + assert model.rotary_pos_emb is model.language_model.rotary_pos_emb + assert model.position_embedding_type == model.language_model.position_embedding_type + # MTP is off here, so `mtp` must stay invisible exactly as on a bare GPTModel. + assert not hasattr(model, "mtp") + assert not hasattr(model.language_model, "mtp") + + +def test_te_helper_discovers_decoder_layers(): + """Regression: discovery used to stop at MultimodalModel and capture nothing.""" + config = _make_config("transformer_engine", [CudaGraphModule.mlp]) + model = _build_model(config) + + helper = TECudaGraphHelper( + model=[model], config=config, seq_length=SEQ, micro_batch_size=BATCH + ) + + assert len(helper.flattened_callables) == NUM_LAYERS + assert helper.flattened_callables == list(model.module.language_model.decoder.layers) + assert helper.flattened_callables_is_mtp == [False] * NUM_LAYERS + + +def test_per_layer_cuda_graph_forward_backward_parity(): + """Loss and grad norm match between graphed and eager decoder MLPs.""" + eager_model = _build_model(_make_config()) + eager_loss, eager_grad_norm = _forward_backward(eager_model) + + config = _make_config("transformer_engine", [CudaGraphModule.mlp]) + graph_model = _build_model(config) + # `_build_model` reseeds, so the two models start from identical weights; assert + # it rather than trusting it, otherwise a numeric mismatch would be ambiguous. + graph_model.load_state_dict(eager_model.state_dict()) + + helper = TECudaGraphHelper( + model=[graph_model], config=config, seq_length=SEQ, micro_batch_size=BATCH + ) + helper.create_cudagraphs() + assert helper.graphs_created(), "no per-layer graph was captured" + assert all( + layer.cuda_graphs for layer in graph_model.module.language_model.decoder.layers + ) + + graph_loss, graph_grad_norm = _forward_backward(graph_model) + + assert graph_loss == pytest.approx(eager_loss, rel=2e-2, abs=2e-2) + assert graph_grad_norm == pytest.approx(eager_grad_norm, rel=5e-2) + + +def test_unreachable_decoder_fails_loudly(): + """A wrapper whose decoder the helper cannot reach must not capture zero layers. + + ``get_attr_wrapped_model`` unwraps only through ``.module``, so before the + forwarding properties above existed the lookup raised, the helper swallowed + its own RuntimeError at DEBUG level, and a run with CUDA graphs explicitly + enabled silently trained in eager mode. + """ + + class _NoDecoder(torch.nn.Module): + def forward(self, *args, **kwargs): + raise AssertionError("not called") + + config = _make_config("transformer_engine", [CudaGraphModule.mlp]) + with pytest.raises(RuntimeError, match="forwarding property"): + TECudaGraphHelper( + model=[_ModelChunk(_NoDecoder())], + config=config, + seq_length=SEQ, + micro_batch_size=BATCH, + ) diff --git a/megatron/core/mdp/README.md b/megatron/core/mdp/README.md index 1e8822831fd..94fffcf2b88 100644 --- a/megatron/core/mdp/README.md +++ b/megatron/core/mdp/README.md @@ -27,7 +27,7 @@ EMPTY`) driving seven phases: | P1 | `begin_iteration` | Capture the iteration window, broadcast fixed-width descriptors from the PP0 endpoint, run deterministic LPT to logical workers, check the plan digest across the group, exchange pixels | | P2 | `begin_iteration` | Grad-enabled chunked encoder forward on encoder THD (`no_grad` for evaluation); outputs retained as a list in the forward handle | | P3 | `begin_iteration` | Exchange detached embeddings; endpoint assembles one detached leaf per vision-bearing microbatch | -| P4 | native schedule | Replay iterators feed the unmodified decoder schedule; the wrapped `finalize_model_grads_func` captures the in-place-reduced global token count | +| P4 | native schedule | Replay iterators feed the unmodified decoder schedule (per-layer CUDA graphs, if enabled, live entirely inside it); the wrapped `finalize_model_grads_func` captures the in-place-reduced global token count | | P5 | `end_iteration` | Exchange leaf gradients back, one multi-tensor backward per producer (native MCore recompute replays here), WORLD sum-reduce with prescale 1, scale by `1/clamp(T_global, 1)` | | P6 | composite optimizer | WORLD MAX overflow union before any scaler update, combined-norm shared clipping, one atomic step for `[decoder_dense, decoder_expert?, encoder]` | @@ -65,8 +65,7 @@ recompute (`None`/`selective`/`full`) via the override channel, text-only microbatches, synchronous global `torch_dist` weight-only checkpoints, `alignment_rows=1` (tests exercise 16). -Rejected at startup: FSDP/HSDP, FP8/MXFP8, full-iteration CUDA graphs, CPU -activation offload, comm overlap (`overlap_grad_reduce`, +Rejected at startup: FSDP/HSDP, FP8/MXFP8, CPU activation offload, comm overlap (`overlap_grad_reduce`, `overlap_param_gather`, delayed reduction), multiple distributed-optimizer instances, `calculate_per_token_loss=False`, non-`torch_dist` checkpoint formats, non-weight-only save/load, invalid rank mappings. diff --git a/megatron/core/mdp/config.py b/megatron/core/mdp/config.py index f3630be9f0b..85870ec49db 100644 --- a/megatron/core/mdp/config.py +++ b/megatron/core/mdp/config.py @@ -22,6 +22,18 @@ # The only checkpoint format supported by the MDP checkpoint facade. SUPPORTED_CHECKPOINT_MODE = "torch_dist" +# CUDA graph implementations MDP accepts. "local" and "transformer_engine" are the +# per-layer ("partial") implementations: they graph individual decoder submodules and +# leave the schedule, the encoder, and every MDP collective in eager mode. MDP's P4 +# runs the unmodified decoder schedule, so the decoder layers a per-layer graph owns +# are never touched by the bridge. +SUPPORTED_CUDA_GRAPH_IMPLS: frozenset = frozenset({"none", "local", "transformer_engine"}) + +# The one implementation that is structurally incompatible with the phase machine: a +# single graph over the whole forward-backward path would swallow P4 together with the +# Python-level control flow P1-P3/P5 depend on. +REJECTED_CUDA_GRAPH_IMPL = "full_iteration" + # Keys that may be overridden on the vision TransformerConfig. Field semantics and # cross-field validation are delegated entirely to MCore's own __post_init__. VISION_CONFIG_OVERRIDE_ALLOWLIST: frozenset = frozenset( @@ -69,7 +81,7 @@ class MdpCompatibilityOptions: bf16: bool fsdp_enabled: bool fp8_enabled: bool - cuda_graph_enabled: bool + cuda_graph_impl: str activation_offload_enabled: bool overlap_grad_reduce: bool overlap_param_gather: bool @@ -91,13 +103,13 @@ def thd_row_alignment(options: "MdpCompatibilityOptions") -> int: per-rank split; SP additionally splits across TP). The greedy token budget must be a multiple of this, or a full bin cannot be partitioned legally. """ - if options.context_parallel_size > 1: - return ( - options.tensor_parallel_size * options.context_parallel_size * 2 - if options.sequence_parallel - else options.context_parallel_size * 2 - ) - return options.tensor_parallel_size if options.sequence_parallel else 1 + from megatron.core.packed_seq_params import thd_collate_row_alignment + + return thd_collate_row_alignment( + context_parallel_size=options.context_parallel_size, + tensor_model_parallel_size=options.tensor_parallel_size, + sequence_parallel=options.sequence_parallel, + ) def _reject(option: str, value: Any, condition: str, why: str, suggestion: str = "") -> None: @@ -270,14 +282,7 @@ def validate_mdp_config(config: MdpConfig, options: MdpCompatibilityOptions) -> "config override channel is reserved for a future FP8 recipe.", "False", ) - if options.cuda_graph_enabled: - _reject( - "cuda_graph_enabled", - options.cuda_graph_enabled, - "full-iteration CUDA graphs disabled", - "MDP buffers are not captured graph-safe in this version.", - "False", - ) + _validate_cuda_graph_options(config, options) if options.activation_offload_enabled: _reject( "activation_offload_enabled", @@ -328,6 +333,71 @@ def validate_mdp_config(config: MdpConfig, options: MdpCompatibilityOptions) -> ) +def _validate_cuda_graph_options( + config: MdpConfig, options: MdpCompatibilityOptions +) -> None: + """Accept per-layer CUDA graphs; keep full-iteration graphs rejected. + + Per-layer graphs own individual decoder submodules. P4 replays the captured + microbatches through the *unmodified* decoder schedule, and the bridge only ever + touches the decoder's embedding leaves (P3) and their gradients (P5) - never the + transformer layers - so a per-layer graph and the phase machine do not interact. + A full-iteration graph would instead capture P4 itself. + """ + impl = options.cuda_graph_impl or "none" + if impl == REJECTED_CUDA_GRAPH_IMPL: + _reject( + "cuda_graph_impl", + impl, + f"cuda_graph_impl != '{REJECTED_CUDA_GRAPH_IMPL}'", + "A full-iteration graph captures the decoder schedule itself, so the " + "Python-level phase machine around it (pixel/embedding/gradient exchange, " + "per-iteration plan, dynamic vision item counts) cannot run. Per-layer " + "graphs are supported instead.", + "transformer_engine", + ) + if impl not in SUPPORTED_CUDA_GRAPH_IMPLS: + _reject( + "cuda_graph_impl", + impl, + f"cuda_graph_impl in {sorted(SUPPORTED_CUDA_GRAPH_IMPLS)}", + "MDP validates against the known CUDA graph implementations only.", + "none", + ) + if impl == "none": + return + + if config.overlap_window_capture: + _reject( + "overlap_window_capture", + config.overlap_window_capture, + "overlap_window_capture == False when per-layer CUDA graphs are enabled", + "Window prefetch runs H2D copies and allocations from a background thread " + "on a side CUDA stream while graph capture is in flight, which " + "cudaStreamCaptureModeGlobal treats as an unsafe concurrent action. The " + "prefetch started in iteration N is still running when the capture step " + "runs, so the conflict is timing-dependent rather than reproducible.", + "False", + ) + + if not options.thd_static_packing: + _reject( + "thd_static_packing", + options.thd_static_packing, + "thd_static_packing == True when per-layer CUDA graphs are enabled", + "MDP requires a THD-packed decoder (window.py asserts qkv_format == 'thd'), " + "and a CUDA graph replays into fixed-size static input buffers: " + "[max_seqlen_per_dp_cp_rank, 1, H] hidden_states plus cu_seqlens of " + "thd_max_packed_sequences + 1 entries. Only --thd-static-packing makes MDP's " + "collator emit that shape; without it every microbatch has a different " + "packed token count and replay fails on the first mismatch. " + "--sequence-packing-scheduler, which produces the same contract for the " + "decoder-only path, is not usable under MDP (see validate_mdp_config).", + "--thd-static-packing --pad-packed-seq-alignment max " + "--max-seqlen-per-dp-cp-rank --thd-max-packed-sequences ", + ) + + def greedy_max_real_sequences(options: "MdpCompatibilityOptions") -> Optional[int]: """Real sequences a greedy bin may hold, or ``None`` for no cap. diff --git a/megatron/core/mdp/integration.py b/megatron/core/mdp/integration.py index b28afd493c8..6e32a573fcb 100644 --- a/megatron/core/mdp/integration.py +++ b/megatron/core/mdp/integration.py @@ -114,7 +114,10 @@ def compatibility_options_from_args(args) -> MdpCompatibilityOptions: or getattr(args, "use_custom_fsdp", False) or getattr(args, "use_megatron_fsdp", False) ) - cuda_graph = getattr(args, "cuda_graph_impl", "none") not in (None, "none") + # Snapshot the implementation itself, not just "any graph enabled": it decides + # whether the graph owns the whole iteration (rejected) or individual decoder + # layers (supported). + cuda_graph_impl = getattr(args, "cuda_graph_impl", "none") or "none" offload = bool( getattr(args, "cpu_offloading", False) or getattr(args, "fine_grained_activation_offloading", False) @@ -149,7 +152,7 @@ def compatibility_options_from_args(args) -> MdpCompatibilityOptions: bf16=bool(args.bf16), fsdp_enabled=fsdp, fp8_enabled=getattr(args, "fp8", None) is not None, - cuda_graph_enabled=cuda_graph, + cuda_graph_impl=cuda_graph_impl, activation_offload_enabled=offload, overlap_grad_reduce=getattr(args, "overlap_grad_reduce", False), overlap_param_gather=getattr(args, "overlap_param_gather", False), diff --git a/megatron/core/mdp/knowledge.md b/megatron/core/mdp/knowledge.md index e3b8700b21d..1659c6dc61f 100644 --- a/megatron/core/mdp/knowledge.md +++ b/megatron/core/mdp/knowledge.md @@ -266,6 +266,35 @@ Primary flags: - `--mdp-greedy-packing` - `--mdp-mock-dataset-config-json` +CUDA graphs: + +- `--cuda-graph-impl local|transformer_engine` -- per-layer ("partial") graphs, + **accepted**. Requires `--thd-static-packing`; rejected together with + `--mdp-overlap-window-capture`. +- `--cuda-graph-impl full_iteration` -- **rejected**: one graph over the whole + forward-backward path swallows P4 along with the Python control flow P1-P3/P5 + depend on. +- Whether the THD CUDA-graph machinery engages is decided by one shared + predicate, `packed_seq_params.thd_shapes_are_static(config)`: true for + `sequence_packing_scheduler`, `dynamic_context_parallel`, **or** + `thd_static_packing`. Before this it asked "is MCore's packing scheduler + configured", which coincided with "are the shapes fixed" only while the + scheduler was the sole fixed-shape producer. Four call sites share it: + `module.py:_is_thd_cuda_graph`, `transformer_config.py`'s + `thd_max_packed_sequences` / `pad_packed_seq_alignment` gate, the same gate in + `arguments.py`, and the MoE dispatcher restriction. + `cuda_graphs.py:_needs_full_local_padding_mask` follows via + `_is_thd_cuda_graph`. `_get_thd_varlen_max_num_microbatches` deliberately does + **not**: it stays `dp_balanced`-specific, and MDP falls through to "runtime" + because greedy variant A keeps `num_microbatches` static. +- `MultimodalModel` forwards `decoder` / `mtp` / `rotary_pos_emb` / + `position_embedding_type` to `self.language_model`. + `TECudaGraphHelper._discover_layers` resolves layers with + `get_attr_wrapped_model(chunk, 'decoder')`, which unwraps only through + `.module`; without the forwards it raised, swallowed its own error at DEBUG, + and captured **zero** layers while reporting success. Core now raises when + *every* chunk fails that lookup. + Packing flags MDP consumes from the core config (all optional, all off by default): diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 2f259e57bf1..5b0b6229dd3 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -293,6 +293,84 @@ def extend_thd_padding_before_cp_slice( return cu_seqlens_padded, max_seqlen, global_target_len +def thd_shapes_are_static(config) -> bool: + """Whether the incoming THD batches have fixed shapes. + + The THD CUDA-graph machinery -- static ``hidden_states``, static + ``cu_seqlens_*``, static ``padding_mask``, and the tensor <-> PackedSeqParams + bridge -- needs exactly one thing from the data path: that ``T`` and the + ``cu_seqlens`` entry count do not vary per microbatch. It does not care + *who* guarantees that. + + Three producers do: + + - ``--sequence-packing-scheduler`` (``dp_balanced`` / ``default_dynamic_cp``); + - ``--dynamic-context-parallel``; + - ``--thd-static-packing``, for collators that pack outside the scheduler + (MDP's greedy packer). + + Deliberately **not** derived from ``pad_packed_seq_alignment is not None``: + that would silently change behavior for existing GPT ``--sft`` runs that set + an alignment without a scheduler. An explicit opt-in cannot. + """ + return bool( + getattr(config, 'sequence_packing_scheduler', None) is not None + or getattr(config, 'dynamic_context_parallel', False) + or getattr(config, 'thd_static_packing', False) + ) + + +def thd_collate_row_alignment( + *, context_parallel_size: int, tensor_model_parallel_size: int, sequence_parallel: bool +) -> int: + """Row alignment a THD collator must pad each packed sample to. + + Zigzag context parallelism needs an even per-rank split, and sequence + parallelism additionally splits the packed rows across TP. Single source of + truth for the rule: the collator pads with it, MDP validates the greedy token + budget against it, and ``thd_static_pad_between_seqs`` derives from it. + """ + if context_parallel_size > 1: + return ( + tensor_model_parallel_size * context_parallel_size * 2 + if sequence_parallel + else context_parallel_size * 2 + ) + return tensor_model_parallel_size if sequence_parallel else 1 + + +def thd_static_pad_between_seqs(config) -> bool: + """Batch-independent ``pad_between_seqs`` for a fixed-shape THD data path. + + ``pad_between_seqs`` cannot be a graph input (it is a capture-time Python + branch) and cannot be inferred from the cu_seqlens tensors during capture (a + device comparison would synchronize), so the CUDA-graph path needs a value + that is correct for *every* replay batch. Answering "True, always" is safe + but expensive: TE disables FlashAttention for THD whenever padding may exist + between sequences, and when cuDNN fused attention does not support the head + configuration either, the fallback is the unfused O(T^2) backend. + + Under ``thd_static_packing`` the answer is knowable without looking at any + batch. A collator setting that flag pads each sample to + ``thd_collate_row_alignment``, so gaps between sequences exist exactly when + that alignment exceeds 1. **That is the contract the flag asserts**; a + collator that leaves gaps at alignment 1 must not set it. + + Without ``thd_static_packing`` (the ``--sequence-packing-scheduler`` path, + which does pad each sub-sample) the conservative ``True`` is retained. + """ + if not getattr(config, 'thd_static_packing', False): + return True + return ( + thd_collate_row_alignment( + context_parallel_size=config.context_parallel_size, + tensor_model_parallel_size=config.tensor_model_parallel_size, + sequence_parallel=config.sequence_parallel, + ) + > 1 + ) + + def build_static_thd_metadata( cu_seqlens: Tensor, cu_seqlens_padded: Tensor, diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 00ec8f6ebc4..ca5f7138d35 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1769,6 +1769,7 @@ def _discover_layers(self): self.callables_per_chunk_is_mtp = [] self.flattened_callables = [] self.flattened_callables_is_mtp = [] + chunks_missing_decoder = 0 for chunk_number, model_chunk in enumerate(self.model): try: chunk_with_decoder = get_attr_wrapped_model( @@ -1776,6 +1777,7 @@ def _discover_layers(self): ) except RuntimeError: num_graphable_layers = 0 + chunks_missing_decoder += 1 log_on_each_pipeline_stage( logger=logger, tp_group=self.tp_group, @@ -1836,6 +1838,29 @@ def _discover_layers(self): f'{len(self.flattened_callables)} graphable layers.', ) + # The per-chunk attribute-lookup failure above is only a DEBUG log, so a + # model whose decoder `get_attr_wrapped_model(chunk, 'decoder')` cannot + # reach captures zero layers and looks like a successful no-op run. When + # *every* chunk failed that lookup it is a wiring bug, not a property of + # the model: `get_attr_wrapped_model` unwraps only through `.module`, so a + # wrapper holding its decoder elsewhere (a multimodal model whose decoder + # is `self.language_model.decoder`, say) needs a forwarding property. + # + # Deliberately keyed on the lookup failing rather than on + # `flattened_callables` being empty: a model whose layers are all + # non-graphable is a legitimate zero-layer case that must stay silent. + if self.model and chunks_missing_decoder == len(self.model): + raise RuntimeError( + f"CUDA graphs are enabled (cuda_graph_impl=" + f"{self.config.cuda_graph_impl!r}) but none of the " + f"{len(self.model)} model chunks exposes a 'decoder' attribute " + "reachable through get_attr_wrapped_model, so capture would be a " + "silent no-op. get_attr_wrapped_model unwraps only through " + "'.module'; if the decoder lives under another attribute, add a " + "forwarding property for 'decoder' (and 'mtp' / 'rotary_pos_emb' / " + "'position_embedding_type') on the wrapper." + ) + def capture_finished(self): """ Returns whether create_cudagraphs() has been called. diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 35c5faab550..014d2d286cd 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -225,11 +225,16 @@ def _te_cuda_graph_backward_dw_graph(self, microbatch_idx): self.cuda_graphs[cg_index].backward_dw() def _is_thd_cuda_graph(self): - """Check if THD format with CUDA Graph is being used.""" - return ( - getattr(self.config, 'sequence_packing_scheduler', None) is not None - and self.config.cuda_graph_impl != "none" - ) + """Check if THD format with CUDA Graph is being used. + + The question is "does the incoming THD batch have fixed shapes", not + "is MCore's packing scheduler configured" -- those coincided only while + the scheduler was the sole fixed-shape producer. See + ``thd_shapes_are_static``. + """ + from megatron.core.packed_seq_params import thd_shapes_are_static + + return thd_shapes_are_static(self.config) and self.config.cuda_graph_impl != "none" def get_layer_static_inputs(self, seq_length, micro_batch_size): """ diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 58261444a02..85336daf54d 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -12,6 +12,7 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.inference.moe import InferenceGroupedGemmBackend +from megatron.core.packed_seq_params import thd_shapes_are_static from megatron.core.quantization.quant_config import RecipeConfig from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, @@ -3447,9 +3448,7 @@ def _scope_to_str(s): self.attention_backend == AttnBackend.flash ), "Batch invariant mode only supports FlashAttention" - if self.cuda_graph_impl != "none" and ( - self.sequence_packing_scheduler is not None or self.dynamic_context_parallel - ): + if self.cuda_graph_impl != "none" and thd_shapes_are_static(self): if self.thd_max_packed_sequences is None: raise ValueError("THD CUDA Graph requires --thd-max-packed-sequences to be set.") assert ( @@ -3464,6 +3463,15 @@ def _scope_to_str(s): f"({self.max_seqlen_per_dp_cp_rank}), got {self.pad_packed_seq_alignment}." ) + # THD packing constrains the MoE dispatcher regardless of *who* produced + # the pack, so this keys off the shape contract rather than on + # sequence_packing_scheduler alone. + if thd_shapes_are_static(self) and self.num_moe_experts is not None: + assert self.moe_token_dispatcher_type in ("alltoall", "flex"), ( + f"sequence packing only supports moe_token_dispatcher_type in " + f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'" + ) + if self.thd_static_packing and self.thd_max_packed_sequences is None: raise ValueError( "thd_static_packing requires --thd-max-packed-sequences: it fixes the " @@ -3504,12 +3512,6 @@ def _scope_to_str(s): # Needed for passing variable sequences between pp stages. self.variable_seq_lengths = True - if self.num_moe_experts is not None: - assert self.moe_token_dispatcher_type in ("alltoall", "flex"), ( - f"sequence_packing only supports moe_token_dispatcher_type in " - f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'" - ) - supported_schedulers = ['dp_balanced', 'default_dynamic_cp'] if ( self.sequence_packing_scheduler is not None diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ea714552464..aaaf3c041dc 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -19,7 +19,10 @@ 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 ( + PackedSeqParams, + thd_static_pad_between_seqs, +) 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 ( @@ -1386,12 +1389,12 @@ def _reconstruct_packed_seq_params_from_kwargs(self, kwargs): cu_seqlens_kv_padded=kwargs.pop('cu_seqlens_kv_padded'), max_seqlen_q=max_seqlen, max_seqlen_kv=max_seqlen, - # CUDA graph inputs do not carry this Python bool. Sequence-packing - # metadata may contain valid/physical gaps, so use the conservative - # graph-static value instead of asking TE to infer it with a CUDA - # tensor comparison during capture. This restricts TE backend - # selection for THD to cuDNN fused attention. - pad_between_seqs=True, + # CUDA graph inputs do not carry this Python bool, and inferring it + # from the cu_seqlens tensors would synchronize during capture, so it + # comes from the config instead - batch-independent by construction. + # Defaults to the conservative True; --thd-static-packing narrows it, + # which matters because True costs FlashAttention eligibility. + pad_between_seqs=thd_static_pad_between_seqs(self.config), ) kwargs['packed_seq_params'] = packed_seq_params diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index a9631c9a7e7..9a73076093c 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -18,6 +18,7 @@ from megatron.core.dist_checkpointing.validation import StrictHandling from megatron.core.fusions.fused_bias_geglu import quick_gelu from megatron.core.model_parallel_config import _parse_pad_packed_seq_alignment +from megatron.core.packed_seq_params import thd_shapes_are_static from megatron.core.msc_utils import MultiStorageClientFeature from megatron.core.quantization.utils import ( kitchen_quantization_recipe_config, @@ -1597,9 +1598,7 @@ def validate_args(args, defaults={}): f'got {args.pad_packed_seq_alignment}.' ) - if args.cuda_graph_impl != "none" and ( - args.sequence_packing_scheduler is not None or args.dynamic_context_parallel - ): + if args.cuda_graph_impl != "none" and thd_shapes_are_static(args): if getattr(args, 'pad_packed_seq_alignment', None) is None: raise ValueError('THD CUDA Graph requires --pad-packed-seq-alignment to be set.') if ( diff --git a/tests/unit_tests/mdp/test_config.py b/tests/unit_tests/mdp/test_config.py index ba510fad057..44fe9ded8fa 100644 --- a/tests/unit_tests/mdp/test_config.py +++ b/tests/unit_tests/mdp/test_config.py @@ -11,6 +11,7 @@ VISION_CONFIG_OVERRIDE_ALLOWLIST, MdpCompatibilityOptions, MdpConfig, + SUPPORTED_CUDA_GRAPH_IMPLS, apply_vision_config_overrides, greedy_max_real_sequences, validate_mdp_config, @@ -34,7 +35,7 @@ def _options(**overrides): bf16=True, fsdp_enabled=False, fp8_enabled=False, - cuda_graph_enabled=False, + cuda_graph_impl="none", activation_offload_enabled=False, overlap_grad_reduce=False, overlap_param_gather=False, @@ -103,7 +104,7 @@ def test_invalid_mdp_config_fields_rejected(config_kwargs, match): (dict(bf16=False), "fp16/bf16"), (dict(fsdp_enabled=True), "fsdp"), (dict(fp8_enabled=True), "fp8"), - (dict(cuda_graph_enabled=True), "cuda_graph"), + (dict(cuda_graph_impl="full_iteration"), "cuda_graph_impl"), (dict(activation_offload_enabled=True), "activation_offload"), (dict(overlap_grad_reduce=True), "overlap_grad_reduce"), (dict(overlap_param_gather=True), "overlap_param_gather"), @@ -333,3 +334,69 @@ def test_static_packing_needs_room_for_a_real_sequence_and_the_dummy(): thd_static_packing=True, ), ) + + +# --------------------------------------------------------------------------- +# Per-layer CUDA graphs +# --------------------------------------------------------------------------- + + +def _graph_options(**overrides): + """Options that satisfy every per-layer CUDA-graph precondition.""" + base = dict( + cuda_graph_impl="transformer_engine", + thd_static_packing=True, + max_seqlen_per_dp_cp_rank=8192, + thd_max_packed_sequences=16, + ) + base.update(overrides) + return _options(**base) + + +@pytest.mark.parametrize("impl", sorted(SUPPORTED_CUDA_GRAPH_IMPLS)) +def test_per_layer_cuda_graphs_are_accepted(impl): + validate_mdp_config(MdpConfig(enable=True), _graph_options(cuda_graph_impl=impl)) + + +def test_full_iteration_graphs_stay_rejected(): + # Regression guard: a full-iteration graph captures the decoder schedule + # itself, so the Python-level phase machine around it cannot run. + with pytest.raises(MdpConfigurationError, match="full_iteration"): + validate_mdp_config( + MdpConfig(enable=True), _graph_options(cuda_graph_impl="full_iteration") + ) + + +def test_unknown_cuda_graph_impl_is_rejected(): + with pytest.raises(MdpConfigurationError, match="cuda_graph_impl"): + validate_mdp_config(MdpConfig(enable=True), _graph_options(cuda_graph_impl="nonsense")) + + +def test_per_layer_graphs_require_static_thd_shapes(): + # Without --thd-static-packing every microbatch has a different packed token + # count and replay fails on the first mismatch. + with pytest.raises(MdpConfigurationError, match="thd_static_packing"): + validate_mdp_config(MdpConfig(enable=True), _graph_options(thd_static_packing=False)) + + +def test_per_layer_graphs_reject_overlap_window_capture(): + with pytest.raises(MdpConfigurationError, match="overlap_window_capture"): + validate_mdp_config( + MdpConfig(enable=True, overlap_window_capture=True), _graph_options() + ) + + +def test_graph_gate_is_inert_without_graphs(): + # thd_static_packing is not required when graphs are off. + validate_mdp_config( + MdpConfig(enable=True, overlap_window_capture=True), _options(cuda_graph_impl="none") + ) + + +def test_mcore_scheduler_still_rejected_with_graphs(): + # The scheduler would satisfy the *shape* contract but not MDP's data + # contract, so the packing rejection must win. + with pytest.raises(MdpConfigurationError, match="sequence_packing_scheduler"): + validate_mdp_config( + MdpConfig(enable=True), _graph_options(sequence_packing_scheduler="dp_balanced") + ) diff --git a/tests/unit_tests/transformer/test_thd_static_shape_predicate.py b/tests/unit_tests/transformer/test_thd_static_shape_predicate.py new file mode 100644 index 00000000000..97d07354184 --- /dev/null +++ b/tests/unit_tests/transformer/test_thd_static_shape_predicate.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""The shared "are THD shapes static" predicate behind the THD CUDA-graph path.""" + +from types import SimpleNamespace + +import pytest + +from megatron.core.packed_seq_params import thd_shapes_are_static + + +def _config(**overrides): + base = dict( + sequence_packing_scheduler=None, + dynamic_context_parallel=False, + thd_static_packing=False, + pad_packed_seq_alignment=None, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +@pytest.mark.parametrize( + "overrides", + [ + dict(sequence_packing_scheduler="dp_balanced"), + dict(sequence_packing_scheduler="default_dynamic_cp"), + dict(dynamic_context_parallel=True), + dict(thd_static_packing=True), + ], +) +def test_every_fixed_shape_producer_is_recognized(overrides): + assert thd_shapes_are_static(_config(**overrides)) + + +def test_no_producer_means_dynamic_shapes(): + assert not thd_shapes_are_static(_config()) + + +def test_alignment_alone_does_not_imply_static_shapes(): + """Existing GPT --sft runs set an alignment without a scheduler. + + Deriving the predicate from ``pad_packed_seq_alignment`` would silently + change their behavior; an explicit opt-in cannot. + """ + assert not thd_shapes_are_static(_config(pad_packed_seq_alignment="max")) + assert not thd_shapes_are_static(_config(pad_packed_seq_alignment=4096)) + + +def test_absent_attributes_are_tolerated(): + """Callers pass both TransformerConfig and the argparse Namespace.""" + assert not thd_shapes_are_static(SimpleNamespace()) + assert thd_shapes_are_static(SimpleNamespace(thd_static_packing=True))