Speedup megatron_mmlu by ~6x via prefill scoring and global batching

Replace autoregressive generation with a single prefill forward pass per
batch and argmax over the four answer-choice token logits. This matches
the log-likelihood approach used by lm-evaluation-harness and avoids the
autoregressive decode loop entirely.

Additional improvements:
- Load dataset once with the "all" config (2 calls) instead of once per
  subject (114 calls), eliminating the main CPU overhead bottleneck
- Batch globally across all subjects sorted by descending sequence length
  to minimise padding waste and fail-fast on OOM
- Skip dev dataset load when few_shots=0
- Rename percentage -> fraction for clearer semantics
- Fix few-shot answer formatting (was emitting integer index, now letter)
- Fix off-by-one: idx > threshold -> idx >= threshold
- Fix avg_correct reset bug inside subject loop
- Add tqdm progress bar (rank-0 only)
- Explicitly del logits/padded after each batch to avoid tensor lifetime
  overlap that caused OOM on long-sequence runs

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

examples/megatron_bridge/prune_minitron.py

@@ -18,7 +18,7 @@
 while skipping pruning of num_attention_heads using following defaults:
     1024 samples from nemotron-post-training-dataset-v2 for calibration,
     at-most 20% depth (num_layers) and 40% width is pruned per prunable hparam (hidden_size, ffn_hidden_size, ...),
-    top-10 candidates are evaluated for MMLU score (5% sampled data) to select the best model.
+    top-10 candidates are evaluated for MMLU score (10% sampled data) to select the best model.
 
     torchrun --nproc_per_node 2 prune_minitron.py \
         --hf_model_name_or_path Qwen/Qwen3-8B \
@@ -140,11 +140,11 @@ def get_args() -> argparse.Namespace:
     parser.add_argument(
         "--prune_score_func",
         type=str,
-        default="mmlu_5pct",
+        default="mmlu_10pct",
         help=(
             "Score function to use for NAS-based pruning (--prune_target_params). Only supports MMLU at the moment. "
             "Format: mmlu_<N>pct where <N> is the percentage of MMLU data to sample per subject "
-            "(e.g. mmlu_5pct for 5%, mmlu_100pct for full eval)."
+            "(e.g. mmlu_10pct for 10%, mmlu_100pct for full eval)."
         ),
     )
     parser.add_argument(
@@ -300,15 +300,17 @@ def main(args: argparse.Namespace):
         if not match:
             raise ValueError(
                 f"Invalid score function: {args.prune_score_func}. "
-                "Expected format: mmlu_<N>pct (e.g. mmlu_5pct)"
+                "Expected format: mmlu_<N>pct (e.g. mmlu_10pct)"
             )
         mmlu_pct = int(match.group(1))
         if not 0 < mmlu_pct <= 100:
             raise ValueError("--prune_score_func percentage must be in the range [1, 100].")
-        _mmlu_pct = mmlu_pct / 100.0
+        _mmlu_frac = mmlu_pct / 100.0
 
         def score_func(m):
-            return megatron_mmlu(m, tokenizer, percentage=_mmlu_pct)
+            return megatron_mmlu(
+                m, tokenizer, few_shots=0, fraction=_mmlu_frac, batch_size=args.calib_mbs
+            )
 
         pruning_config["score_func"] = score_func
         pruning_config["max_width_pruning"] = args.max_width_pruning

examples/pruning/README.md

@@ -124,7 +124,7 @@ This mode can be useful when you don't know the exact dimensions you want to pru
 from modelopt.torch.utils.plugins.megatron_mmlu import megatron_mmlu
 
 def score_func(m):
-    return megatron_mmlu(m, tokenizer, percentage=0.05)  # 5% sampled data for faster eval
+    return megatron_mmlu(m, tokenizer, fraction=0.1, batch_size=4)  # 10% sampled data for faster eval
 
 # Specify target parameter count and configure the auto pruning algorithm
 # Save minitron scores at checkpoint so we can resume pruning without running the forward loop again
@@ -147,7 +147,7 @@ mtp.prune(...)
 
 1. **Importance Scoring**: Same as manual pruning - computes activation magnitudes for all parameters (takes ~5 minutes for an 8B model)
 2. **Search Space Construction**: Generates a search space of possible architectures based search space config and other configs (`max_width_pruning`, `max_depth_pruning`, `hparams_to_skip`)
-3. **Architecture Search**: Find candidate architectures that meet the parameter constraint and evaluate `top_k` (based on number of parameters) of them using `score_func` e.g. MMLU, negative validation loss, etc. (takes ~10 mins per candidate for an 8B model pruning)
+3. **Architecture Search**: Find candidate architectures that meet the parameter constraint and evaluate `top_k` (based on number of parameters) of them using `score_func` e.g. MMLU, negative validation loss, etc. (takes ~1 min per candidate for an 8B model MMLU score with 10% sampled data)
 4. **Best Architecture Selection**: Returns the architecture (best `export_config`) with the highest actual score from the top-K evaluated architectures
 5. **Weight Slicing**: Slices the model weights according to the best pruned architecture found
 

modelopt/torch/utils/plugins/megatron_mmlu.py

@@ -40,62 +40,53 @@
 
 """A simple MMLU evaluation for Megatron LM models."""
 
-import requests
 import torch
-import transformers
 from datasets import load_dataset
+from tqdm import tqdm
+from transformers import PreTrainedTokenizer
 
-from .megatron_generate import megatron_generate
+from .. import distributed as dist
+from .megatron_generate import megatron_prefill
 
 __all__ = ["megatron_mmlu"]
 
-
-def _get_all_subjects():
-    """All subjects (anatomy, ...) can be acquired from querying all subsets and splits."""
-    response = requests.get(
-        "https://datasets-server.huggingface.co/splits?dataset=cais/mmlu", timeout=10
-    )
-    data = response.json()
-    all_subjects = set()
-    for split in data["splits"]:
-        all_subjects.add(split["config"])
-    for name in ["all", "auxiliary_train"]:
-        all_subjects.discard(name)
-    return sorted(all_subjects)
+_CHOICES = ["A", "B", "C", "D"]
 
 
 def megatron_mmlu(
     model,
-    tokenizer: transformers.PreTrainedTokenizer,
+    tokenizer: PreTrainedTokenizer,
     few_shots: int = 0,
-    percentage: float = 0.05,
-    enable_kv_cache: bool = False,
+    fraction: float = 0.05,
+    batch_size: int = 1,
 ) -> float:
-    """Evaluate the model on MMLU.
+    """Evaluate the model on MMLU using log-likelihood scoring over batched prefill passes.
+
+    Instead of autoregressively generating tokens, a single prefill forward pass is run per
+    batch and the answer is selected as argmax over the four choice token logits at the last
+    prompt position. This is the same approach used by lm-evaluation-harness.
 
     Args:
         model: The model to evaluate.
         tokenizer: The tokenizer to use.
         few_shots: The number of few-shot examples to use.
-        percentage: The percentage of the test set to evaluate on.
-        enable_kv_cache: Whether to disable KV-cache.
+        fraction: The fraction of the test set to evaluate on.
+        batch_size: Number of examples to process in one forward pass.
     """
-    all_correct = {}
-    all_subjects = _get_all_subjects()
+    # Token IDs for " A", " B", " C", " D" — the last subword handles edge cases.
+    choice_ids = [tokenizer.encode(f" {c}", add_special_tokens=False)[-1] for c in _CHOICES]
 
     def _format_example(example, include_answer: bool = True):
-        """Format an example into a multi-choices problem."""
         prompt = example["question"]
-        for choice, answer in zip(["A", "B", "C", "D"], example["choices"]):
+        for choice, answer in zip(_CHOICES, example["choices"]):
             prompt += f"\n{choice}. {answer}"
         if include_answer:
-            prompt += "Answer: {}\n\n".format(example["answer"])
+            prompt += "Answer: {}\n\n".format(_CHOICES[example["answer"]])
         else:
             prompt += "\nAnswer:"
         return prompt
 
     def _generate_prompt(test_example, dev_examples, few_shots=0):
-        """Generating few-shot prompts."""
         prompt = "The following are multiple choice questions (with answers) about {}.\n\n".format(
             " ".join(test_example["subject"].split("_"))
         )
@@ -104,51 +95,97 @@ def _generate_prompt(test_example, dev_examples, few_shots=0):
         prompt += _format_example(test_example, include_answer=False)
         return prompt
 
-    if torch.distributed.get_rank() == 0:
-        print(f"\nMMLU ({percentage * 100}%, {few_shots}-shot) evaluation started...\n", flush=True)
+    # Load all subjects in two dataset calls instead of 2x num_subjects calls.
+    # The "all" config includes a "subject" field for per-subject reporting.
+    test_dataset = load_dataset("cais/mmlu", "all", split="test")
+    dev_dataset = load_dataset("cais/mmlu", "all", split="dev") if few_shots > 0 else None
+
+    # Group dev examples by subject for few-shot prompt construction.
+    dev_by_subject: dict = {}
+    if dev_dataset is not None:
+        for ex in dev_dataset:
+            dev_by_subject.setdefault(ex["subject"], []).append(ex)
+
+    # Collect all examples, tracking subject membership for per-subject reporting.
+    all_subjects_seen: list[str] = []
+    all_prompts: list[str] = []
+    all_labels: list[str] = []
+
+    # Count test examples per subject to apply the fraction cutoff correctly.
+    subject_counts: dict[str, int] = {}
+    for ex in test_dataset:
+        subject_counts[ex["subject"]] = subject_counts.get(ex["subject"], 0) + 1
+
+    subject_idx: dict[str, int] = {}
+    for ex in test_dataset:
+        subj = ex["subject"]
+        idx = subject_idx.get(subj, 0)
+        if idx >= fraction * subject_counts[subj]:
+            continue
+        subject_idx[subj] = idx + 1
+        prompt = _generate_prompt(ex, dev_by_subject.get(subj, []), few_shots=few_shots)
+        all_prompts.append(prompt)
+        all_labels.append(_CHOICES[ex["answer"]])
+        all_subjects_seen.append(subj)
+
+    # Tokenize all prompts and sort by length to minimise padding waste within batches.
+    encoded = [tokenizer(p, return_tensors="pt").input_ids[0] for p in all_prompts]
+    lengths = [e.shape[0] for e in encoded]
+    order = sorted(range(len(encoded)), key=lambda i: lengths[i], reverse=True)
+
+    sorted_encoded = [encoded[i] for i in order]
+    sorted_lengths = [lengths[i] for i in order]
+
+    # Run inference in global batches.
+    predictions: list[str] = [""] * len(encoded)
+    n_batches = (len(sorted_encoded) + batch_size - 1) // batch_size
+    pbar = tqdm(
+        range(0, len(sorted_encoded), batch_size),
+        total=n_batches,
+        desc="MMLU",
+        unit="batch",
+        disable=not dist.is_master(),
+    )
+    for batch_start in pbar:
+        batch_enc = sorted_encoded[batch_start : batch_start + batch_size]
+        batch_len = sorted_lengths[batch_start : batch_start + batch_size]
+        max_len = max(batch_len)
+
+        # Right-pad to max_len; causal mask means the last real token is unaffected by padding.
+        padded = torch.zeros(len(batch_enc), max_len, dtype=torch.long)
+        for i, (e, seq_len) in enumerate(zip(batch_enc, batch_len)):
+            padded[i, :seq_len] = e
+
+        logits = megatron_prefill(model, padded.cuda())  # [B, max_len, vocab]
+
+        for i, seq_len in enumerate(batch_len):
+            answer_logits = logits[i, seq_len - 1, choice_ids]
+            predictions[order[batch_start + i]] = _CHOICES[answer_logits.argmax().item()]
+
+        examples_done = min(batch_start + batch_size, len(sorted_encoded))
+        pbar.set_postfix(examples=f"{examples_done}/{len(sorted_encoded)}")
+
+    # Compute per-subject accuracy and overall average.
+    subject_correct: dict[str, list[bool]] = {}
+    for pred, label, subj in zip(predictions, all_labels, all_subjects_seen):
+        subject_correct.setdefault(subj, []).append(pred == label)
+
+    all_correct = [pred == label for pred, label in zip(predictions, all_labels)]
+    n_total = len(all_correct)
+    avg = sum(all_correct) / n_total
+
+    if dist.is_master():
+        print(f"\nMMLU ({fraction * 100}%, {few_shots}-shot) evaluation started...\n", flush=True)
         print("{:48} | (ACC) | Count/Total".format("Subject"), flush=True)
         print("{:48} | {:5} | {:11}".format("-" * 48, "-" * 5, "-" * 11), flush=True)
-
-    for subject in all_subjects:
-        test_data = load_dataset("cais/mmlu", subject, split="test")
-        dev_data = load_dataset("cais/mmlu", subject, split="dev")
-
-        correct = []
-        for idx, test_example in enumerate(test_data):
-            if idx > percentage * len(test_data):
-                break
-            prompt = _generate_prompt(test_example, dev_data, few_shots=few_shots)
-            label = ["A", "B", "C", "D"][test_example["answer"]]
-            tokens = tokenizer(prompt, return_tensors="pt")
-            generated_ids = megatron_generate(
-                model,
-                tokens.input_ids.cuda(),
-                osl=2,
-                disable_tqdm=True,
-                enable_kv_cache=enable_kv_cache,
-            )
-            predict = tokenizer.batch_decode(generated_ids)[0].strip()
-            correct += [True] if predict.startswith(label) else [False]
-        all_correct[subject] = correct
-
-        if torch.distributed.get_rank() == 0:
-            print(
-                f"{subject:48} | {sum(correct) / len(correct):.3f} | {sum(correct):5}/{len(correct):5}",
-                flush=True,
-            )
-
-        avg_correct = []
-
-    for subject, correct in all_correct.items():
-        avg_correct += correct
-
-    if torch.distributed.get_rank() == 0:
+        for subj in sorted(subject_correct):
+            correct = subject_correct[subj]
+            n = len(correct)
+            print(f"{subj:48} | {sum(correct) / n:.3f} | {sum(correct):5}/{n:5}", flush=True)
         print("{:48} | {:5} | {:11}".format("-" * 48, "-" * 5, "-" * 11), flush=True)
         print(
-            "{:48} | {:.3f} | {:5}/{:5}".format(
-                "average", sum(avg_correct) / len(avg_correct), sum(avg_correct), len(avg_correct)
-            ),
+            "{:48} | {:.3f} | {:5}/{:5}".format("average", avg, sum(all_correct), n_total),
             flush=True,
         )
 
-    return sum(avg_correct) / len(avg_correct)
+    return avg

tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py

@@ -25,9 +25,7 @@
 
 def _test_megatron_generate_and_mmlu(rank, size):
     initialize_for_megatron(tensor_model_parallel_size=size, seed=SEED)
-
     model = get_mcore_qwen3_600m(tensor_model_parallel_size=size).cuda().eval()
-
     tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
 
     messages = [
@@ -42,9 +40,9 @@ def _test_megatron_generate_and_mmlu(rank, size):
     model_inputs = tokenizer([text], return_tensors="pt").to(device="cuda")
     output_ids = megatron_generate(model, model_inputs["input_ids"])
     output_text = tokenizer.batch_decode(output_ids)
-    print(output_text)
+    print(rank, output_text)
 
-    assert megatron_mmlu(model, tokenizer) > 0.24
+    assert 0.37 < megatron_mmlu(model, tokenizer, fraction=0.1, batch_size=16) < 0.38
 
 
 def test_megatron_generate_and_mmlu(dist_workers):