[data] Add openmathinstruct and gsm8k datasets (#2264)

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

scripts/training/run_recipe.py

@@ -18,42 +18,72 @@
 
 This script works with any model family that uses GPT-style training
 (Llama, Gemma, Qwen, GPT, etc.). It dynamically loads recipes and supports
-CLI overrides.
+CLI overrides. The --dataset flag selects the dataset type and automatically
+infers pretrain vs finetune mode.
 
 Usage:
-    Pretrain (single-GPU):
-        uv run python run_recipe.py \
-            --recipe llama32_1b_pretrain_config
-
-    Pretrain (multi-GPU):
-        uv run torchrun --nproc_per_node=8 run_recipe.py \
-            --recipe llama32_1b_pretrain_config
-
-    SFT (full finetuning):
-        uv run torchrun --nproc_per_node=8 run_recipe.py \
-            --recipe llama32_1b_sft_config
-
-    PEFT (LoRA/DoRA):
-        uv run torchrun --nproc_per_node=8 run_recipe.py \
-            --recipe llama32_1b_peft_config
-
-    With CLI overrides:
-        uv run torchrun --nproc_per_node=8 run_recipe.py \
-            --recipe llama32_1b_pretrain_config \
-            train.train_iters=5000 \
+    Pretrain (mock data):
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe llama32_1b_pretrain_config \\
+            --dataset llm-pretrain-mock
+
+    Pretrain (real data):
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe llama32_1b_pretrain_config \\
+            --dataset llm-pretrain \\
+            'dataset.blend=[[/data/my_dataset_text_document],null]'
+
+    Finetune (SQuAD, default):
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe llama32_1b_sft_config \\
+            --dataset llm-finetune
+
+    Finetune (GSM8K):
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe llama32_1b_sft_config \\
+            --dataset llm-finetune \\
+            dataset.dataset_name=gsm8k
+
+    Finetune (user-supplied JSONL):
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe llama32_1b_sft_config \\
+            --dataset llm-finetune-preloaded \\
+            dataset.dataset_root=/data/my_finetune_data
+
+    VLM with HF dataset:
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe qwen3_vl_8b_peft_config \\
+            --dataset vlm-hf \\
+            --step_func qwen3_vl_step \\
+            dataset.maker_name=cord_v2 \\
+            dataset.hf_processor_path=Qwen/Qwen3-VL-8B-Instruct \\
+            checkpoint.pretrained_checkpoint=/path/to/checkpoint
+
+    VLM with Energon dataset:
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe qwen3_vl_8b_peft_energon_config \\
+            --dataset vlm-energon \\
+            --step_func qwen3_vl_step \\
+            dataset.path=/data/energon \\
+            checkpoint.pretrained_checkpoint=/path/to/checkpoint
+
+    VLM with preloaded JSON:
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe qwen3_vl_8b_peft_config \\
+            --dataset vlm-preloaded \\
+            --step_func qwen3_vl_step \\
+            dataset.train_data_path=/data/vlm_train.json \\
+            dataset.image_folder=/data/vlm_images \\
+            dataset.hf_processor_path=Qwen/Qwen3-VL-8B-Instruct \\
+            checkpoint.pretrained_checkpoint=/path/to/checkpoint
+
+    With CLI overrides (Hydra-style, works for any config field):
+        uv run torchrun --nproc_per_node=8 run_recipe.py \\
+            --recipe llama32_1b_pretrain_config \\
+            --dataset llm-pretrain-mock \\
+            train.train_iters=5000 \\
             optimizer.lr=0.0003
 
-    With VLM step function:
-        uv run torchrun --nproc_per_node=8 run_recipe.py \
-            --recipe qwen25_vl_finetune_config \
-            --step_func vlm_step
-
-    With packed sequences and custom sequence length:
-        uv run torchrun --nproc_per_node=8 run_recipe.py \
-            --recipe llama32_1b_pretrain_config \
-            --packed_sequence \
-            --seq_length 2048
-
 Recipe Arguments:
     Generic scripts call recipes with no arguments: recipe().
 
@@ -67,6 +97,11 @@
 
 import megatron.bridge.recipes as recipes
 from megatron.bridge.models.qwen_vl.qwen3_vl_step import forward_step as qwen3_vl_forward_step
+from megatron.bridge.recipes.utils.dataset_utils import (
+    DATASET_TYPES,
+    apply_dataset_override,
+    infer_mode_from_dataset,
+)
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.finetune import finetune
 from megatron.bridge.training.gpt_step import forward_step as gpt_forward_step
@@ -83,16 +118,16 @@
     "llava_step": llava_forward_step,
 }
 
-TRAIN_MODES = {
+TRAIN_FUNCTIONS = {
     "pretrain": pretrain,
     "finetune": finetune,
 }
 
-# Error message constants
 ERR_UNKNOWN_STEP = "Unknown step type: {step_type}. Choose from: {choices}"
 ERR_INFER_MODE_FAILED = (
-    "Unable to infer training mode from recipe name. "
-    "Please include 'pretrain', 'sft', 'peft', or 'finetune' in the recipe name or pass --mode explicitly."
+    "Unable to infer training mode. "
+    "Pass --dataset to specify the dataset type, or include 'pretrain' or 'finetune' "
+    "(or 'sft'/'peft') in the recipe name."
 )
 
 
@@ -102,19 +137,30 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
         description="Generic training script for GPT-based models",
         formatter_class=argparse.RawTextHelpFormatter,
     )
-    parser.add_argument(
-        "--mode",
-        type=str,
-        default=None,
-        choices=sorted(TRAIN_MODES.keys()),
-        help="Training mode (optional). If omitted, inferred from recipe name.",
-    )
     parser.add_argument(
         "--recipe",
         type=str,
         required=True,
         help="Recipe function name (e.g., llama32_1b_pretrain_config, gemma3_1b_sft_config, gemma3_1b_peft_config)",
     )
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        choices=DATASET_TYPES,
+        help=(
+            "Dataset type. Training mode (pretrain/finetune) is inferred from this.\n"
+            "LLM datasets:\n"
+            "  llm-pretrain           GPT pretrain data (set dataset.blend=<path>)\n"
+            "  llm-pretrain-mock      Mock pretrain data for testing\n"
+            "  llm-finetune           HF finetune dataset (set dataset.dataset_name=squad|gsm8k|openmathinstruct2)\n"
+            "  llm-finetune-preloaded User-supplied JSONL (set dataset.dataset_root=<path>)\n"
+            "VLM datasets:\n"
+            "  vlm-energon            Energon multimodal (set dataset.path=<path>)\n"
+            "  vlm-hf                 HF VLM dataset (set dataset.maker_name=<name>)\n"
+            "  vlm-preloaded          User-supplied VLM JSON (set dataset.train_data_path=<path>)"
+        ),
+    )
     parser.add_argument(
         "--step_func",
         type=str,
@@ -227,7 +273,7 @@ def load_forward_step(step_type: str) -> Callable:
 
 
 def infer_train_mode(recipe_name: str) -> str:
-    """Infer training mode from the recipe name."""
+    """Infer training mode from the recipe name (fallback when --dataset is not passed)."""
     lowered = recipe_name.lower()
     has_pretrain = "pretrain" in lowered
     has_finetune = "finetune" in lowered or "sft" in lowered or "peft" in lowered
@@ -248,15 +294,35 @@ def main() -> None:
         args.hf_path,
     )
 
+    if args.dataset is not None:
+        mode = infer_mode_from_dataset(args.dataset)
+        config = apply_dataset_override(
+            config,
+            dataset_type=args.dataset,
+            packed_sequence=args.packed_sequence,
+            seq_length=args.seq_length,
+            cli_overrides=cli_overrides,
+        )
+    else:
+        mode = infer_train_mode(args.recipe)
+
     config = process_config_with_overrides(
         config,
         cli_overrides=cli_overrides or None,
     )
 
-    mode = args.mode or infer_train_mode(args.recipe)
+    # Ensure dataset.seq_length and model.seq_length stay in sync after CLI overrides
+    if (
+        hasattr(config, "model")
+        and config.model is not None
+        and hasattr(config, "dataset")
+        and config.dataset is not None
+    ):
+        if hasattr(config.dataset, "seq_length") and config.model.seq_length != config.dataset.seq_length:
+            config.model.seq_length = config.dataset.seq_length
 
     forward_step = load_forward_step(args.step_func)
-    train_func = TRAIN_MODES[mode]
+    train_func = TRAIN_FUNCTIONS[mode]
     train_func(config=config, forward_step_func=forward_step)
 
 

src/megatron/bridge/data/builders/hf_dataset.py

@@ -358,5 +358,18 @@ def _load_dataset(self) -> DatasetDict:
             )
         else:
             raise ValueError("Expected `dataset_name` to be str, got " + str(type(self.dataset_name)))
+        # If a specific split was requested, load_dataset returns a Dataset, not DatasetDict.
+        # We need to wrap it in a DatasetDict for downstream processing.
+        if isinstance(dataset, Dataset):
+            normalized_split = (self.split or "train").split("[", 1)[0]
+            if normalized_split.startswith("train"):
+                split_key = "train"
+            elif normalized_split.startswith(("validation", "valid", "val", "eval")):
+                split_key = "validation"
+            elif normalized_split.startswith("test"):
+                split_key = "test"
+            else:
+                split_key = "train"
+            dataset = DatasetDict({split_key: dataset})
 
         return dataset

src/megatron/bridge/data/hf_processors/__init__.py

@@ -19,9 +19,13 @@
 and are designed to work with the HFDatasetConfig and HFDatasetBuilder classes.
 """
 
+from .gsm8k import process_gsm8k_example
+from .openmathinstruct2 import process_openmathinstruct2_example
 from .squad import process_squad_example
 
 
 __all__ = [
+    "process_gsm8k_example",
+    "process_openmathinstruct2_example",
     "process_squad_example",
 ]

src/megatron/bridge/data/hf_processors/gsm8k.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Processing functions for GSM8K (Grade School Math 8K) dataset.
+
+Dataset: https://huggingface.co/datasets/openai/gsm8k
+
+GSM8K contains 8.5K grade school math word problems. Each example has a
+``question`` and an ``answer`` field where the answer contains chain-of-thought
+reasoning followed by ``####`` and the final numerical answer.
+"""
+
+from typing import Any, Optional
+
+from megatron.bridge.data.builders.hf_dataset import ProcessExampleOutput
+from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer
+
+
+def _extract_final_answer(answer: str) -> str:
+    """Extract the final numerical answer after the ``####`` delimiter."""
+    if "####" in answer:
+        return answer.split("####")[-1].strip()
+    return answer.strip()
+
+
+def process_gsm8k_example(
+    example: dict[str, Any], _tokenizer: Optional[MegatronTokenizer] = None
+) -> ProcessExampleOutput:
+    """Process a single GSM8K example into the required format.
+
+    Transforms a raw GSM8K dataset example into the standard format expected by
+    the HFDatasetBuilder for fine-tuning.
+
+    Args:
+        example: Raw GSM8K example containing 'question' and 'answer'
+        tokenizer: Optional tokenizer (not used in this processor)
+
+    Returns:
+        ProcessExampleOutput with formatted input/output and original answers
+
+    Example:
+        >>> example = {
+        ...     "question": "Janet has 3 apples. She buys 2 more. How many does she have?",
+        ...     "answer": "Janet starts with 3 apples and buys 2 more. 3 + 2 = <<3+2=5>>5.\\n#### 5",
+        ... }
+        >>> result = process_gsm8k_example(example)
+        >>> print(result["input"])
+        Question: Janet has 3 apples. She buys 2 more. How many does she have? Answer:
+    """
+    _input = f"Question: {example['question']} Answer:"
+    _output = example["answer"]
+    final_answer = _extract_final_answer(example["answer"])
+
+    return ProcessExampleOutput(input=_input, output=_output, original_answers=[final_answer])

src/megatron/bridge/data/hf_processors/openmathinstruct2.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Processing functions for OpenMathInstruct-2 dataset.
+
+Dataset: https://huggingface.co/datasets/nvidia/OpenMathInstruct-2
+
+OpenMathInstruct-2 contains math problems with generated solutions. Each example
+has ``problem``, ``generated_solution``, and ``expected_answer`` fields.
+"""
+
+from typing import Any, Optional
+
+from megatron.bridge.data.builders.hf_dataset import ProcessExampleOutput
+from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer
+
+
+def process_openmathinstruct2_example(
+    example: dict[str, Any], _tokenizer: Optional[MegatronTokenizer] = None
+) -> ProcessExampleOutput:
+    """Process a single OpenMathInstruct-2 example into the required format.
+
+    Transforms a raw OpenMathInstruct-2 dataset example into the standard format
+    expected by the HFDatasetBuilder for fine-tuning.
+
+    Args:
+        example: Raw example containing 'problem', 'generated_solution', and 'expected_answer'
+        tokenizer: Optional tokenizer (not used in this processor)
+
+    Returns:
+        ProcessExampleOutput with formatted input/output and original answers
+
+    Example:
+        >>> example = {
+        ...     "problem": "What is 2 + 3?",
+        ...     "generated_solution": "We add 2 and 3 to get 5.",
+        ...     "expected_answer": "5",
+        ... }
+        >>> result = process_openmathinstruct2_example(example)
+        >>> print(result["input"])
+        Problem: What is 2 + 3? Solution:
+    """
+    _input = f"Problem: {example['problem']} Solution:"
+    _output = example["generated_solution"]
+    expected_answer = example["expected_answer"]
+
+    return ProcessExampleOutput(input=_input, output=_output, original_answers=[expected_answer])