Skip to content

Commit 73185d5

Browse files
yaoyu-33yaoyu-33
authored
Merge branch 'main' into zhiyul/fix-moonshot-logprob-mismatch
2 parents ce63df8 + 8f43506 commit 73185d5

15 files changed

Lines changed: 1595 additions & 49 deletions

File tree

‎examples/models/sarvam/README.md‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
# Sarvam Examples
2+
3+
This directory contains example scripts for Sarvam language models.
4+
5+
Sarvam models use a Mixture of Experts (MoE) architecture with QKV layernorm and Grouped Query Attention (GQA).
6+
7+
| Model | HF ID | Architecture | Params |
8+
|---|---|---|---|
9+
| Sarvam 30B | `sarvamai/sarvam-30b` | MoE (128 experts, top-6) | 30B total, 3B active |
10+
| Sarvam 105B | `sarvamai/sarvam-105b` | MoE (128 experts, top-8) | 105B total, 10.3B active |
11+
12+
## Workspace Configuration
13+
14+
All scripts use a `WORKSPACE` environment variable for the base directory. Default: `/workspace`.
15+
16+
```bash
17+
export WORKSPACE=/your/custom/path
18+
```
19+
20+
Directory structure:
21+
- `${WORKSPACE}/models/` - Converted checkpoints
22+
- `${WORKSPACE}/results/` - Training outputs
23+
24+
## Checkpoint Conversion
25+
26+
See [conversion.sh](conversion.sh) for checkpoint conversion examples.
27+
28+
### Import HF → Megatron
29+
30+
```bash
31+
python examples/conversion/convert_checkpoints.py import \
32+
--hf-model sarvamai/sarvam-30b \
33+
--megatron-path ${WORKSPACE}/models/sarvam-30b \
34+
--trust-remote-code
35+
```
36+
37+
### Export Megatron → HF
38+
39+
```bash
40+
python examples/conversion/convert_checkpoints.py export \
41+
--hf-model sarvamai/sarvam-30b \
42+
--megatron-path ${WORKSPACE}/models/sarvam-30b/iter_0000000 \
43+
--hf-path ${WORKSPACE}/models/sarvam-30b-hf-export
44+
```
45+
46+
### Round-trip Validation
47+
48+
```bash
49+
python -m torch.distributed.run --nproc_per_node=8 \
50+
examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
51+
--hf-model-id sarvamai/sarvam-30b \
52+
--megatron-load-path ${WORKSPACE}/models/sarvam-30b/iter_0000000 \
53+
--tp 2 --pp 2 --ep 2 \
54+
--trust-remote-code
55+
```
56+
57+
## Inference
58+
59+
See [inference.sh](inference.sh) for text generation with:
60+
- Hugging Face checkpoint (`sarvamai/sarvam-30b`)
61+
- Imported Megatron checkpoint (after [conversion.sh](conversion.sh) import)
62+
- Exported HF checkpoint (after conversion export)
63+
64+
The default parallelism for 8 GPUs is `--tp 2 --pp 1 --ep 4 --etp 1`.
65+
TP×PP×EP must equal `--nproc_per_node`.
66+
67+
> **Note**: `sarvamai/sarvam-30b` is a base pretrained model published in float32.
68+
> The Bridge conversion is correct (verified by matching outputs between HF and Megatron checkpoints).
69+
> For coherent text generation, use an instruction-tuned variant if available.

‎examples/models/sarvam/conversion.sh‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#!/usr/bin/env bash
2+
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -xeuo pipefail
17+
18+
# Workspace directory for checkpoints and results
19+
WORKSPACE=${WORKSPACE:-/workspace}
20+
21+
MODEL_NAME=sarvam-30b
22+
HF_MODEL_ID=sarvamai/sarvam-30b
23+
24+
# Import HF → Megatron
25+
uv run python examples/conversion/convert_checkpoints.py import \
26+
--hf-model "$HF_MODEL_ID" \
27+
--megatron-path "${WORKSPACE}/models/${MODEL_NAME}" \
28+
--trust-remote-code
29+
30+
# Export Megatron → HF
31+
uv run python examples/conversion/convert_checkpoints.py export \
32+
--hf-model "$HF_MODEL_ID" \
33+
--megatron-path "${WORKSPACE}/models/${MODEL_NAME}/iter_0000000" \
34+
--hf-path "${WORKSPACE}/models/${MODEL_NAME}-hf-export"
35+
36+
# Round-trip validation
37+
uv run python -m torch.distributed.run --nproc_per_node=8 \
38+
examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
39+
--hf-model-id "$HF_MODEL_ID" \
40+
--megatron-load-path "${WORKSPACE}/models/${MODEL_NAME}/iter_0000000" \
41+
--tp 2 --pp 2 --ep 2 \
42+
--trust-remote-code

‎examples/models/sarvam/inference.sh‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#!/usr/bin/env bash
2+
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
# Workspace directory for checkpoints and results
17+
WORKSPACE=${WORKSPACE:-/workspace}
18+
19+
# Inference with Hugging Face checkpoints
20+
uv run python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_text.py \
21+
--hf_model_path sarvamai/sarvam-30b \
22+
--prompt "Hello, how are you?" \
23+
--max_new_tokens 64 \
24+
--tp 2 --pp 1 --ep 4 --etp 1 \
25+
--trust-remote-code
26+
27+
# Inference with imported Megatron checkpoints
28+
uv run python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_text.py \
29+
--hf_model_path sarvamai/sarvam-30b \
30+
--megatron_model_path ${WORKSPACE}/models/sarvam-30b/iter_0000000 \
31+
--prompt "Hello, how are you?" \
32+
--max_new_tokens 64 \
33+
--tp 2 --pp 1 --ep 4 --etp 1 \
34+
--trust-remote-code
35+
36+
# Inference with exported HF checkpoints
37+
uv run python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_text.py \
38+
--hf_model_path ${WORKSPACE}/models/sarvam-30b-hf-export \
39+
--prompt "Hello, how are you?" \
40+
--max_new_tokens 64 \
41+
--tp 2 --pp 1 --ep 4 --etp 1 \
42+
--trust-remote-code

‎scripts/training/run_recipe.py‎

Lines changed: 112 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -18,42 +18,72 @@
1818
1919
This script works with any model family that uses GPT-style training
2020
(Llama, Gemma, Qwen, GPT, etc.). It dynamically loads recipes and supports
21-
CLI overrides.
21+
CLI overrides. The --dataset flag selects the dataset type and automatically
22+
infers pretrain vs finetune mode.
2223
2324
Usage:
24-
Pretrain (single-GPU):
25-
uv run python run_recipe.py \
26-
--recipe llama32_1b_pretrain_config
27-
28-
Pretrain (multi-GPU):
29-
uv run torchrun --nproc_per_node=8 run_recipe.py \
30-
--recipe llama32_1b_pretrain_config
31-
32-
SFT (full finetuning):
33-
uv run torchrun --nproc_per_node=8 run_recipe.py \
34-
--recipe llama32_1b_sft_config
35-
36-
PEFT (LoRA/DoRA):
37-
uv run torchrun --nproc_per_node=8 run_recipe.py \
38-
--recipe llama32_1b_peft_config
39-
40-
With CLI overrides:
41-
uv run torchrun --nproc_per_node=8 run_recipe.py \
42-
--recipe llama32_1b_pretrain_config \
43-
train.train_iters=5000 \
25+
Pretrain (mock data):
26+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
27+
--recipe llama32_1b_pretrain_config \\
28+
--dataset llm-pretrain-mock
29+
30+
Pretrain (real data):
31+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
32+
--recipe llama32_1b_pretrain_config \\
33+
--dataset llm-pretrain \\
34+
'dataset.blend=[[/data/my_dataset_text_document],null]'
35+
36+
Finetune (SQuAD, default):
37+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
38+
--recipe llama32_1b_sft_config \\
39+
--dataset llm-finetune
40+
41+
Finetune (GSM8K):
42+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
43+
--recipe llama32_1b_sft_config \\
44+
--dataset llm-finetune \\
45+
dataset.dataset_name=gsm8k
46+
47+
Finetune (user-supplied JSONL):
48+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
49+
--recipe llama32_1b_sft_config \\
50+
--dataset llm-finetune-preloaded \\
51+
dataset.dataset_root=/data/my_finetune_data
52+
53+
VLM with HF dataset:
54+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
55+
--recipe qwen3_vl_8b_peft_config \\
56+
--dataset vlm-hf \\
57+
--step_func qwen3_vl_step \\
58+
dataset.maker_name=cord_v2 \\
59+
dataset.hf_processor_path=Qwen/Qwen3-VL-8B-Instruct \\
60+
checkpoint.pretrained_checkpoint=/path/to/checkpoint
61+
62+
VLM with Energon dataset:
63+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
64+
--recipe qwen3_vl_8b_peft_energon_config \\
65+
--dataset vlm-energon \\
66+
--step_func qwen3_vl_step \\
67+
dataset.path=/data/energon \\
68+
checkpoint.pretrained_checkpoint=/path/to/checkpoint
69+
70+
VLM with preloaded JSON:
71+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
72+
--recipe qwen3_vl_8b_peft_config \\
73+
--dataset vlm-preloaded \\
74+
--step_func qwen3_vl_step \\
75+
dataset.train_data_path=/data/vlm_train.json \\
76+
dataset.image_folder=/data/vlm_images \\
77+
dataset.hf_processor_path=Qwen/Qwen3-VL-8B-Instruct \\
78+
checkpoint.pretrained_checkpoint=/path/to/checkpoint
79+
80+
With CLI overrides (Hydra-style, works for any config field):
81+
uv run torchrun --nproc_per_node=8 run_recipe.py \\
82+
--recipe llama32_1b_pretrain_config \\
83+
--dataset llm-pretrain-mock \\
84+
train.train_iters=5000 \\
4485
optimizer.lr=0.0003
4586
46-
With VLM step function:
47-
uv run torchrun --nproc_per_node=8 run_recipe.py \
48-
--recipe qwen25_vl_finetune_config \
49-
--step_func vlm_step
50-
51-
With packed sequences and custom sequence length:
52-
uv run torchrun --nproc_per_node=8 run_recipe.py \
53-
--recipe llama32_1b_pretrain_config \
54-
--packed_sequence \
55-
--seq_length 2048
56-
5787
Recipe Arguments:
5888
Generic scripts call recipes with no arguments: recipe().
5989
@@ -67,6 +97,11 @@
6797

6898
import megatron.bridge.recipes as recipes
6999
from megatron.bridge.models.qwen_vl.qwen3_vl_step import forward_step as qwen3_vl_forward_step
100+
from megatron.bridge.recipes.utils.dataset_utils import (
101+
DATASET_TYPES,
102+
apply_dataset_override,
103+
infer_mode_from_dataset,
104+
)
70105
from megatron.bridge.training.config import ConfigContainer
71106
from megatron.bridge.training.finetune import finetune
72107
from megatron.bridge.training.gpt_step import forward_step as gpt_forward_step
@@ -83,16 +118,16 @@
83118
"llava_step": llava_forward_step,
84119
}
85120

86-
TRAIN_MODES = {
121+
TRAIN_FUNCTIONS = {
87122
"pretrain": pretrain,
88123
"finetune": finetune,
89124
}
90125

91-
# Error message constants
92126
ERR_UNKNOWN_STEP = "Unknown step type: {step_type}. Choose from: {choices}"
93127
ERR_INFER_MODE_FAILED = (
94-
"Unable to infer training mode from recipe name. "
95-
"Please include 'pretrain', 'sft', 'peft', or 'finetune' in the recipe name or pass --mode explicitly."
128+
"Unable to infer training mode. "
129+
"Pass --dataset to specify the dataset type, or include 'pretrain' or 'finetune' "
130+
"(or 'sft'/'peft') in the recipe name."
96131
)
97132

98133

@@ -102,19 +137,30 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
102137
description="Generic training script for GPT-based models",
103138
formatter_class=argparse.RawTextHelpFormatter,
104139
)
105-
parser.add_argument(
106-
"--mode",
107-
type=str,
108-
default=None,
109-
choices=sorted(TRAIN_MODES.keys()),
110-
help="Training mode (optional). If omitted, inferred from recipe name.",
111-
)
112140
parser.add_argument(
113141
"--recipe",
114142
type=str,
115143
required=True,
116144
help="Recipe function name (e.g., llama32_1b_pretrain_config, gemma3_1b_sft_config, gemma3_1b_peft_config)",
117145
)
146+
parser.add_argument(
147+
"--dataset",
148+
type=str,
149+
default=None,
150+
choices=DATASET_TYPES,
151+
help=(
152+
"Dataset type. Training mode (pretrain/finetune) is inferred from this.\n"
153+
"LLM datasets:\n"
154+
" llm-pretrain GPT pretrain data (set dataset.blend=<path>)\n"
155+
" llm-pretrain-mock Mock pretrain data for testing\n"
156+
" llm-finetune HF finetune dataset (set dataset.dataset_name=squad|gsm8k|openmathinstruct2)\n"
157+
" llm-finetune-preloaded User-supplied JSONL (set dataset.dataset_root=<path>)\n"
158+
"VLM datasets:\n"
159+
" vlm-energon Energon multimodal (set dataset.path=<path>)\n"
160+
" vlm-hf HF VLM dataset (set dataset.maker_name=<name>)\n"
161+
" vlm-preloaded User-supplied VLM JSON (set dataset.train_data_path=<path>)"
162+
),
163+
)
118164
parser.add_argument(
119165
"--step_func",
120166
type=str,
@@ -227,7 +273,7 @@ def load_forward_step(step_type: str) -> Callable:
227273

228274

229275
def infer_train_mode(recipe_name: str) -> str:
230-
"""Infer training mode from the recipe name."""
276+
"""Infer training mode from the recipe name (fallback when --dataset is not passed)."""
231277
lowered = recipe_name.lower()
232278
has_pretrain = "pretrain" in lowered
233279
has_finetune = "finetune" in lowered or "sft" in lowered or "peft" in lowered
@@ -248,15 +294,35 @@ def main() -> None:
248294
args.hf_path,
249295
)
250296

297+
if args.dataset is not None:
298+
mode = infer_mode_from_dataset(args.dataset)
299+
config = apply_dataset_override(
300+
config,
301+
dataset_type=args.dataset,
302+
packed_sequence=args.packed_sequence,
303+
seq_length=args.seq_length,
304+
cli_overrides=cli_overrides,
305+
)
306+
else:
307+
mode = infer_train_mode(args.recipe)
308+
251309
config = process_config_with_overrides(
252310
config,
253311
cli_overrides=cli_overrides or None,
254312
)
255313

256-
mode = args.mode or infer_train_mode(args.recipe)
314+
# Ensure dataset.seq_length and model.seq_length stay in sync after CLI overrides
315+
if (
316+
hasattr(config, "model")
317+
and config.model is not None
318+
and hasattr(config, "dataset")
319+
and config.dataset is not None
320+
):
321+
if hasattr(config.dataset, "seq_length") and config.model.seq_length != config.dataset.seq_length:
322+
config.model.seq_length = config.dataset.seq_length
257323

258324
forward_step = load_forward_step(args.step_func)
259-
train_func = TRAIN_MODES[mode]
325+
train_func = TRAIN_FUNCTIONS[mode]
260326
train_func(config=config, forward_step_func=forward_step)
261327

262328

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -358,5 +358,18 @@ def _load_dataset(self) -> DatasetDict:
358358
)
359359
else:
360360
raise ValueError("Expected `dataset_name` to be str, got " + str(type(self.dataset_name)))
361+
# If a specific split was requested, load_dataset returns a Dataset, not DatasetDict.
362+
# We need to wrap it in a DatasetDict for downstream processing.
363+
if isinstance(dataset, Dataset):
364+
normalized_split = (self.split or "train").split("[", 1)[0]
365+
if normalized_split.startswith("train"):
366+
split_key = "train"
367+
elif normalized_split.startswith(("validation", "valid", "val", "eval")):
368+
split_key = "validation"
369+
elif normalized_split.startswith("test"):
370+
split_key = "test"
371+
else:
372+
split_key = "train"
373+
dataset = DatasetDict({split_key: dataset})
361374

362375
return dataset

0 commit comments

Comments
 (0)