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..a9314c1a787 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,56 @@ 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, scenarios): + """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) + 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)) + 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) + # 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), + scenarios=scenarios, ) + scale = _greedy_sample_scale(args, scenarios) 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..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 +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, @@ -146,7 +150,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 +173,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 +189,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 +238,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 +253,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: @@ -410,16 +427,58 @@ 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) 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): + # 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 + # 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 +539,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 +604,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 +673,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 +711,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 +863,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 +921,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/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_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_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/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..04290ade521 --- /dev/null +++ b/examples/multimodal_dev/tests/test_thd_static_packing.py @@ -0,0 +1,368 @@ +# 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 + # 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) + 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_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") + + 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 + 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.""" + 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..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]` | @@ -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) | @@ -64,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 f08e23d05d6..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( @@ -48,6 +60,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) @@ -68,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 @@ -76,6 +89,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. + """ + 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: @@ -150,6 +184,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: @@ -247,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", @@ -305,6 +333,145 @@ 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. + + ``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. + + ``--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.", + ) + 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, + 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", + ) + + 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..6e32a573fcb 100644 --- a/megatron/core/mdp/integration.py +++ b/megatron/core/mdp/integration.py @@ -29,6 +29,8 @@ SUPPORTED_RANK_ORDER, MdpCompatibilityOptions, MdpConfig, + greedy_max_real_sequences, + thd_row_alignment, validate_mdp_config, ) from megatron.core.mdp.encoder import ( @@ -101,6 +103,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), ) @@ -111,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) @@ -146,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), @@ -154,6 +160,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 +237,13 @@ 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 + ) + greedy_max_num_seqs = greedy_max_real_sequences(compat) _RUNTIME = MdpRuntime( config=mdp_config, rank_map=rank_map, @@ -245,6 +263,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=greedy_max_num_seqs, + 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..1659c6dc61f 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,75 @@ Primary flags: - `--mdp-plan-check-interval` - `--mdp-overlap-window-capture` - `--mdp-debug-plan-payload-check` +- `--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): + +- `--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 +347,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..d0b2bcb9036 --- /dev/null +++ b/megatron/core/mdp/packing.py @@ -0,0 +1,181 @@ +# 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 + + +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..1e56108abe2 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,78 @@ 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) + 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)] = entry + return entry[0] + + def consumed_samples(self) -> Optional[int]: + """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, 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: + 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..5b0b6229dd3 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -293,6 +293,162 @@ 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, + *, + 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/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 4091b647103..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,22 @@ 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 " + "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. @@ -3497,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/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..44fe9ded8fa 100644 --- a/tests/unit_tests/mdp/test_config.py +++ b/tests/unit_tests/mdp/test_config.py @@ -11,7 +11,9 @@ VISION_CONFIG_OVERRIDE_ALLOWLIST, MdpCompatibilityOptions, MdpConfig, + SUPPORTED_CUDA_GRAPH_IMPLS, apply_vision_config_overrides, + greedy_max_real_sequences, validate_mdp_config, ) from megatron.core.mdp.errors import MdpConfigurationError @@ -33,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, @@ -102,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"), @@ -244,3 +246,157 @@ 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, + ), + ) + + +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, + ), + ) + + +# --------------------------------------------------------------------------- +# 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/mdp/test_packing.py b/tests/unit_tests/mdp/test_packing.py new file mode 100644 index 00000000000..7401fa829de --- /dev/null +++ b/tests/unit_tests/mdp/test_packing.py @@ -0,0 +1,142 @@ +# 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 + + +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] + + +# --------------------------------------------------------------------------- +# GreedySampleStream +# --------------------------------------------------------------------------- + + +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_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] + 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_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(): + 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 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))