docs(training): clarify workflow and checkpoint guidance (#4259)

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

docs/index.md

@@ -59,6 +59,7 @@ models/stepfun/index.md
 
 training/config-container-overview.md
 training/entry-points.md
+training/data-preparation.md
 training/training-loop-settings.md
 training/optimizer-scheduler.md
 training/logging.md

docs/megatron-lm-to-megatron-bridge.md

@@ -55,7 +55,7 @@ uv run python scripts/translate_mlm_to_bridge.py --reverse \
 | `--ffn-hidden-size N` | `model.ffn_hidden_size=N` | |
 | `--num-attention-heads N` | `model.num_attention_heads=N` | |
 | `--num-query-groups N` | `model.num_query_groups=N` | |
-| `--seq-length N` | `model.seq_length=N dataset.sequence_length=N` | Dual mapping |
+| `--seq-length N` | `model.seq_length=N dataset.seq_length=N` | Dual mapping |
 | `--swiglu` | `model.gated_linear_unit=true model.activation_func=silu` | Expanded to two keys |
 | `--squared-relu` | `model.activation_func=squared_relu` | |
 | `--data-path PATH [W PATH...]` | `dataset.data_path=PATH` | Space-separated paths (and optional weights) |
@@ -95,15 +95,17 @@ Flags not present in Bridge (e.g., `--use-mcore-models`, `--use-flash-attn`) are
 
 ## Quick start
 
-Run your example training entrypoint and override config keys directly:
+Run the generic recipe launcher and override config keys directly:
 
 ```bash
-uv run python examples/models/llama/pretrain_llama3_8b.py \
+uv run python scripts/training/run_recipe.py \
+  --recipe llama3_8b_pretrain_config \
+  --dataset llm-pretrain \
   train.micro_batch_size=2 \
   train.global_batch_size=128 \
   model.num_layers=32 model.hidden_size=4096 model.num_attention_heads=32 \
   model.max_position_embeddings=4096 \
-  dataset.sequence_length=4096 \
+  dataset.seq_length=4096 \
   checkpoint.save=/workspace/ckpts checkpoint.save_interval=1000 \
   logger.wandb_project=my_proj logger.wandb_exp_name=exp1
 ```

docs/nemo2-migration-guide.md

@@ -286,12 +286,12 @@ recipe = llm.llama3_8b.pretrain_recipe(name="my_run", num_nodes=2)
 
 **Megatron Bridge**: Recipes in `megatron.bridge.recipes/`  
 ```python
-from megatron.bridge.recipes.llama.llama3_8b import pretrain_config
-from megatron.bridge.training import pretrain
+from megatron.bridge.recipes.llama import llama3_8b_pretrain_config
 from megatron.bridge.training.gpt_step import forward_step
+from megatron.bridge.training.pretrain import pretrain
 
 # Use pre-built recipe
-cfg = pretrain_config()
+cfg = llama3_8b_pretrain_config()
 
 # Customize as needed
 cfg.train.train_iters = 10000
@@ -384,7 +384,7 @@ from megatron.bridge.training.config import (
 )
 from megatron.core.optimizer import OptimizerConfig
 from megatron.bridge.models import GPTModelProvider
-from megatron.bridge.training import pretrain
+from megatron.bridge.training.pretrain import pretrain
 
 def llama3_8b_config(
     # Model/parallelism params
@@ -1221,8 +1221,9 @@ result = llm.finetune(
 In Megatron Bridge, training entry points take a single `ConfigContainer` and a `forward_step_func`:
 
 ```python
-from megatron.bridge.training import pretrain, finetune
 from megatron.bridge.training.config import ConfigContainer
+from megatron.bridge.training.finetune import finetune
+from megatron.bridge.training.pretrain import pretrain
 
 # Create unified configuration
 cfg = ConfigContainer(
@@ -1285,8 +1286,8 @@ For GPT models, use the provided {py:func}`bridge.training.gpt_step.forward_step
 Use `pretrain()` with `GPTDatasetConfig` for training models from scratch:
 
 ```python
-from megatron.bridge.training import pretrain
 from megatron.bridge.training.gpt_step import forward_step
+from megatron.bridge.training.pretrain import pretrain
 
 config = ConfigContainer(
     model=GPTModelProvider(
@@ -1321,8 +1322,8 @@ Use `finetune()` with `FinetuningDatasetConfig` for both full fine-tuning (SFT)
 Full fine-tuning without PEFT - all model parameters are updated:
 
 ```python
-from megatron.bridge.training import finetune
 from megatron.bridge.training.gpt_step import forward_step
+from megatron.bridge.training.finetune import finetune
 
 config = ConfigContainer(
     model=GPTModelProvider(),

docs/recipe-usage.md

@@ -1,6 +1,6 @@
 # Using Recipes
 
-Megatron Bridge provides production-ready training recipes for several popular models. You can find an overview of supported recipes and 🤗 HuggingFace bridges [here](index.md#supported-models).
+Megatron Bridge provides production-ready training recipes for several popular models. You can find an overview of supported recipes and 🤗 Hugging Face bridges [here](index.md#supported-models).
 This guide will cover the next steps to make use of a training recipe, including how to [override configuration](#overriding-configuration) and how to [launch a job](#launch-methods).
 
 ## Overview
@@ -10,23 +10,41 @@ This guide will cover the next steps to make use of a training recipe, including
 - **Integration**: Recipes return a single `ConfigContainer` that plugs directly into our training [entry points](training/entry-points.md) (see the published docs as well: https://docs.nvidia.com/nemo/megatron-bridge/latest/training/entry-points.html).
 - **Customization**: You can override any part of the recipe (Python, YAML, CLI) to adapt to your data, scale, and objectives.
 
+## Choosing a recipe or a new config
+
+Start from an exported recipe when the model family and workflow already exist in `megatron.bridge.recipes`. Recipe functions such as `llama3_8b_pretrain_config`, `llama32_1b_sft_config`, and `qwen3_8b_peft_config` provide model, optimizer, scheduler, precision, dataset, logger, and checkpoint defaults in one `ConfigContainer`. Override those defaults for your dataset, checkpoint paths, run length, parallelism, or precision before creating a new recipe.
+
+Create a new recipe or config when the base model architecture is not represented by an existing model provider, the checkpoint conversion needs a new bridge, the forward step or dataset provider is model-specific, or you need a reusable configuration that will be shared across jobs. If the Hugging Face model is already supported by `AutoBridge`, you usually only need to start from the closest recipe and override the model provider or `hf_path`.
+
+Training mode follows the recipe and dataset type:
+
+| Workflow | Typical config | Entry point | Checkpoint expectation |
+|----------|----------------|-------------|------------------------|
+| LLM pretraining or continued pretraining | `GPTDatasetConfig` | `pretrain()` | No checkpoint for from-scratch runs; use `checkpoint.load` for full resume or `checkpoint.pretrained_checkpoint` for model-weight initialization |
+| Full SFT | `FinetuningDatasetConfig`, `HFDatasetConfig`, or a dataset provider | `finetune()` | Use `checkpoint.pretrained_checkpoint` for the base model, or `checkpoint.load` for a full native Megatron resume |
+| PEFT / LoRA / DoRA | Same as SFT, plus `cfg.peft` | `finetune()` | `checkpoint.pretrained_checkpoint` is required for the frozen base model; `checkpoint.load` resumes adapter training |
+| VLM SFT or PEFT | VLM dataset provider such as Energon, HF, or preloaded JSON provider | `finetune()` with a VLM step function | Use the model-specific checkpoint guidance in the recipe or model docs |
+
+For dataset fields, prefer `seq_length` in Bridge examples. LLM pretraining uses `GPTDatasetConfig` with `data_path`, `blend`, or `blend_per_split`; SFT and PEFT use `dataset_root` for local JSONL data. Do not use `data_path` for SFT/PEFT JSONL roots.
+
 ## Overriding configuration
 
 Recipes are provided through a {py:class}`~bridge.training.config.ConfigContainer` object. This is a dataclass that holds all configuration objects needed for training. You can find a more detailed overview of the `ConfigContainer` [here](training/config-container-overview.md).
 The benefit of providing the full recipe through a pythonic structure is that it is agnostic to any configuration approach that a user may prefer, whether that's YAML, `argparse` or something else. In other words, the user may override the recipe however they see fit.
 
-The following sections detail a few different ways to override the configuration recipe. For a complete training script, please see [this example](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/llama/pretrain_llama3_8b.py).
+The following sections detail a few different ways to override the configuration recipe. For a generic recipe launcher, see [`scripts/training/run_recipe.py`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/scripts/training/run_recipe.py).
 
 
 ### Python
 
 If you prefer to manage configuration in Python, you can directly modify attributes of the `ConfigContainer`:
 
 ```python
-from megatron.bridge.recipes.llama.llama3_8b import pretrain_config
+from megatron.bridge.recipes.llama import llama3_8b_pretrain_config
+from megatron.bridge.training.config import ConfigContainer
 
 # Get the base ConfigContainer from the recipe
-cfg: ConfigContainer = pretrain_config()
+cfg: ConfigContainer = llama3_8b_pretrain_config()
 
 # Apply overrides. Note the hierarchical structure
 cfg.train.train_iters = 20
@@ -38,32 +56,28 @@ cfg.logger.log_interval = 1
 You can also replace entire sub-configs of the `ConfigContainer`:
 
 ```python
-from megatron.bridge.recipes.llama.llama3_8b import pretrain_config
-from megatron.bridge.models.llama import Llama3ModelProvider
+from megatron.bridge.recipes.llama import llama32_1b_pretrain_config, llama3_8b_pretrain_config
+from megatron.bridge.training.config import ConfigContainer
 
-cfg: ConfigContainer = pretrain_config()
+cfg: ConfigContainer = llama3_8b_pretrain_config()
 
-small_llama = Llama3ModelProvider(
-    num_layers=2,
-    hidden_size=768,
-    ffn_hidden_size=2688,
-    num_attention_heads=16,
-)
-cfg.model = small_llama
+small_cfg: ConfigContainer = llama32_1b_pretrain_config()
+cfg.model = small_cfg.model
 ```
 
 ### YAML
 Overriding a configuration recipe with a YAML file can be done using OmegaConf utilities:
 
 ```python
 from omegaconf import OmegaConf
-from megatron.bridge.recipes.llama.llama3_8b import pretrain_config
+from megatron.bridge.recipes.llama import llama3_8b_pretrain_config
+from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.utils.omegaconf_utils import (
     apply_overrides,
     create_omegaconf_dict_config,
 )
 
-cfg: ConfigContainer = pretrain_config()
+cfg: ConfigContainer = llama3_8b_pretrain_config()
 yaml_filepath = "conf/llama3-8b-benchmark-cfg.yaml"
 
 # Convert the initial Python dataclass to an OmegaConf DictConfig for merging
@@ -88,14 +102,15 @@ Megatron Bridge provides some utilities to update the ConfigContainer using Hydr
 ```python
 import sys
 from omegaconf import OmegaConf
-from megatron.bridge.recipes.llama.llama3_8b import pretrain_config
+from megatron.bridge.recipes.llama import llama3_8b_pretrain_config
+from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.utils.omegaconf_utils import (
     apply_overrides,
     create_omegaconf_dict_config,
     parse_hydra_overrides,
 )
 
-cfg: ConfigContainer = pretrain_config()
+cfg: ConfigContainer = llama3_8b_pretrain_config()
 cli_overrides = sys.argv[1:]
 
 # Convert the initial Python dataclass to an OmegaConf DictConfig for merging
@@ -117,6 +132,27 @@ A script containing the above code could be called like so:
 uv run python -m torch.distributed.run <torchrun arguments> pretrain_cli_overrides.py model.tensor_model_parallel_size=4 train.train_iters=100000 ...
 ```
 
+Common dataset overrides:
+
+```python
+from megatron.bridge.recipes.llama import llama32_1b_sft_config, llama3_8b_pretrain_config
+
+pretrain_cfg = llama3_8b_pretrain_config()
+finetune_cfg = llama32_1b_sft_config()
+
+# LLM pretraining data on a pretrain recipe:
+# prefix path without .bin/.idx suffixes
+pretrain_cfg.dataset.data_path = "/data/dclm/preprocessed_text_document"
+pretrain_cfg.dataset.seq_length = 8192
+
+# SFT/PEFT local JSONL data on a finetune recipe:
+# directory containing training.jsonl, validation.jsonl, and optionally test.jsonl
+finetune_cfg.dataset.dataset_root = "/data/sft_jsonl"
+finetune_cfg.dataset.seq_length = 4096
+```
+
+For more detail on accepted dataset layouts, see [Data Preparation](training/data-preparation.md).
+
 ## Launch methods
 
 Megatron Bridge supports launching scripts with both `torchrun` and [NeMo-Run](https://github.com/NVIDIA-NeMo/Run).
@@ -184,7 +220,7 @@ if __name__ == "__main__":
     train_script = run.Script(..., args=args_to_fwd)
 ```
 
-For a complete example of the `run.Script` API, including argument forwarding, please see [this script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/llama/pretrain_llama3_8b_nemo_run_script.py).
+For a complete example of the `run.Script` API, including argument forwarding, see [`scripts/training/launch_with_nemo_run.py`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/scripts/training/launch_with_nemo_run.py).
 
 #### Plugins
 

docs/training/README.md

@@ -7,7 +7,7 @@ This directory contains comprehensive documentation for training and customizing
 ### I want to
 
 **🚀 Get started with training**
-→ Start with [Configuration Container Overview](config-container-overview.md) to understand the training setup
+→ Start with [Configuration Container Overview](config-container-overview.md) and [Data Preparation](data-preparation.md) to understand the training setup
 
 **⚙️ Configure training parameters**
 → See [Training Loop Settings](training-loop-settings.md) and [Optimizer & Scheduler](optimizer-scheduler.md)
@@ -32,6 +32,7 @@ This directory contains comprehensive documentation for training and customizing
 |----------|---------|--------------|
 | **[Configuration Container Overview](config-container-overview.md)** | Central configuration object for all training settings | First time setting up training |
 | **[Entry Points](entry-points.md)** | Training entry points and execution flow | Understanding how training starts |
+| **[Data Preparation](data-preparation.md)** | Dataset formats for pretraining, SFT, PEFT, and VLM fine-tuning | Preparing data or choosing dataset config fields |
 | **[Training Loop Settings](training-loop-settings.md)** | Training loop parameters and configuration | Configuring batch sizes, iterations, validation |
 
 ### Optimization and Performance
@@ -71,7 +72,7 @@ This directory contains comprehensive documentation for training and customizing
 A typical training workflow involves:
 
 1. **Configure Training** - Set up `ConfigContainer` with model, data, and training parameters
-2. **Prepare Data** - Configure dataset loading and preprocessing
+2. **Prepare Data** - Configure dataset loading and preprocessing with the right data format
 3. **Set Optimization** - Configure optimizer, scheduler, and mixed precision
 4. **Enable Monitoring** - Set up logging and profiling
 5. **Configure Checkpointing** - Set up checkpoint saving and resuming
@@ -93,9 +94,10 @@ A typical training workflow involves:
 ### 🆕 First-Time Training Setup
 
 1. [Configuration Container Overview](config-container-overview.md) - Understand the configuration system
-2. [Entry Points](entry-points.md) - Learn how to start training
-3. [Training Loop Settings](training-loop-settings.md) - Configure basic training parameters
-4. [Logging](logging.md) - Set up monitoring
+2. [Data Preparation](data-preparation.md) - Choose the right dataset format and config fields
+3. [Entry Points](entry-points.md) - Learn how to start training
+4. [Training Loop Settings](training-loop-settings.md) - Configure basic training parameters
+5. [Logging](logging.md) - Set up monitoring
 
 ### ⚡ Performance Optimization