feat(mdp): run per-layer CUDA graphs on the MDP decoder

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",

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)
     )

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,21 +137,23 @@ 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))
 
     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