Skip to content

Commit 8b1d342

Browse files
eagle705eagle705
authored
Merge branch 'main' into fix-get-rope-index
2 parents ba9000d + 15d758f commit 8b1d342

95 files changed

Lines changed: 9888 additions & 3634 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/cicd-main.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -380,6 +380,12 @@ jobs:
380380
# - script: L2_Launch_quantization_export
381381
- script: L2_Launch_recipes_llama_cuda_graphs
382382
- script: L2_Launch_utils
383+
- script: L2_Launch_ckpts_mbridge_to_mlm_llama32_1b
384+
- script: L2_Launch_ckpts_mlm_to_mbridge_llama32_1b
385+
- script: L2_Launch_ckpts_mbridge_to_mlm_qwen3_4b
386+
- script: L2_Launch_ckpts_mlm_to_mbridge_qwen3_4b
387+
- script: L2_Launch_ckpts_mbridge_to_mlm_nemotronh_4b
388+
- script: L2_Launch_ckpts_mlm_to_mbridge_nemotronh_4b
383389
needs: [pre-flight, cicd-unit-tests]
384390
runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
385391
if: |

‎docs/models/llm/gpt-oss.md‎

Lines changed: 1 addition & 167 deletions
Original file line numberDiff line numberDiff line change
@@ -46,180 +46,15 @@ GPT OSS models are supported via the Bridge system with specialized configuratio
4646
- **Bias in Linear Layers**: Linear layers include bias terms
4747
- **Activation Clamping**: Output activations clamped to [-7.0, 7.0] for numerical stability
4848

49-
## Conversion with 🤗 Hugging Face
50-
51-
### Load HF → Megatron
52-
```python
53-
from megatron.bridge import AutoBridge
54-
55-
# Example: GPT OSS 20B
56-
bridge = AutoBridge.from_hf_pretrained("openai/gpt-oss-20b")
57-
provider = bridge.to_megatron_provider()
58-
59-
# Configure parallelism before instantiating the model
60-
provider.tensor_model_parallel_size = 2
61-
provider.pipeline_model_parallel_size = 4
62-
provider.expert_model_parallel_size = 4
63-
provider.sequence_parallel = True
64-
65-
model = provider.provide_distributed_model(wrap_with_ddp=False)
66-
67-
provider.finalize()
68-
69-
# For GPT OSS 120B:
70-
# bridge = AutoBridge.from_hf_pretrained("openai/gpt-oss-120b")
71-
```
72-
73-
### Export Megatron → HF
74-
```python
75-
# Convert from a Megatron checkpoint directory to HF format
76-
bridge.export_ckpt(
77-
megatron_path="/results/gpt_oss_120b/checkpoints/iter_0500000",
78-
hf_path="./gpt-oss-hf-export",
79-
)
80-
```
81-
8249
## Examples
8350

84-
- Checkpoint conversion: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py)
85-
86-
## Pretrain recipes
87-
88-
### Example usage (GPT OSS 20B)
89-
```python
90-
from megatron.bridge.recipes.gpt_oss import gpt_oss_20b_pretrain_config
91-
92-
cfg = gpt_oss_20b_pretrain_config(
93-
name="gpt_oss_20b_pretrain",
94-
data_paths=["/path/to/dataset.nvjsonl"],
95-
dir="/results/gpt_oss_20b",
96-
train_iters=500_000,
97-
global_batch_size=512,
98-
seq_length=4096,
99-
)
100-
```
101-
102-
### Example usage (GPT OSS 120B)
103-
```python
104-
from megatron.bridge.recipes.gpt_oss import gpt_oss_120b_pretrain_config
105-
106-
cfg = gpt_oss_120b_pretrain_config(
107-
name="gpt_oss_120b_pretrain",
108-
data_paths=["/path/to/dataset.nvjsonl"],
109-
dir="/results/gpt_oss_120b",
110-
train_iters=500_000,
111-
global_batch_size=512,
112-
seq_length=4096,
113-
)
114-
```
115-
116-
### Key configuration options
117-
- **Parallelism (20B)**: Default TP=2, PP=4, EP=4 for efficient MoE training
118-
- **Parallelism (120B)**: Default TP=2, PP=4, EP=16 for large-scale training
119-
- **Sequence parallel**: Enabled by default for memory efficiency
120-
- **Context parallel**: Supports CP for long sequences (131K tokens)
121-
- **Manual GC**: Aggressive garbage collection (interval=100) for stable memory usage
122-
- **MoE optimizations**: Grouped GEMM and permute fusion enabled by default
123-
124-
## Finetuning recipes
125-
126-
### Example usage (GPT OSS 20B - LoRA finetuning)
127-
```python
128-
from megatron.bridge.recipes.gpt_oss import gpt_oss_20b_finetune_config
129-
130-
cfg = gpt_oss_20b_finetune_config(
131-
hf_path="openai/gpt-oss-20b",
132-
name="gpt_oss_20b_lora_finetune",
133-
pretrained_checkpoint="path/to/gpt_oss/checkpoint",
134-
peft="lora", # or "dora" for DoRA
135-
train_iters=1000,
136-
global_batch_size=128,
137-
finetune_lr=1e-4,
138-
)
139-
```
140-
141-
### Example usage (GPT OSS 20B - Full SFT)
142-
```python
143-
cfg = gpt_oss_20b_finetune_config(
144-
hf_path="openai/gpt-oss-20b",
145-
name="gpt_oss_20b_full_sft",
146-
pretrained_checkpoint="path/to/gpt_oss/checkpoint",
147-
peft=None, # Full supervised finetuning
148-
train_iters=1000,
149-
global_batch_size=128,
150-
finetune_lr=5e-6, # Lower LR for full SFT
151-
)
152-
```
153-
154-
### Example usage (GPT OSS 120B - LoRA finetuning)
155-
```python
156-
from megatron.bridge.recipes.gpt_oss import gpt_oss_120b_finetune_config
157-
158-
cfg = gpt_oss_120b_finetune_config(
159-
hf_path="openai/gpt-oss-120b",
160-
name="gpt_oss_120b_lora_finetune",
161-
pretrained_checkpoint="path/to/gpt_oss/checkpoint",
162-
peft="lora",
163-
train_iters=1000,
164-
global_batch_size=128,
165-
finetune_lr=1e-4,
166-
)
167-
```
168-
169-
### Default configurations
170-
171-
#### GPT OSS 20B
172-
173-
**LoRA/DoRA (1 node, 8 GPUs)**
174-
- TP=1, PP=1, EP=1, LR=1e-4
175-
- Optimized for parameter-efficient training
176-
- Lower memory footprint
177-
178-
**Full SFT (1 node, 8 GPUs)**
179-
- TP=1, PP=1, EP=8, LR=5e-6
180-
- Full model training with expert parallelism
181-
- Higher throughput with distributed experts
182-
183-
#### GPT OSS 120B
184-
185-
**LoRA/DoRA (2 nodes, 16 GPUs)**
186-
- TP=1, PP=4, EP=8, LR=1e-4
187-
- Optimized for parameter-efficient training
188-
- Pipeline parallelism for memory efficiency
189-
190-
**Full SFT (2 nodes, 16 GPUs)**
191-
- TP=1, PP=1, EP=8, LR=5e-6
192-
- Full model training with expert parallelism
193-
- Requires more memory per GPU
51+
For checkpoint conversion, inference, finetuning recipes, and step-by-step training guides, see the [GPT-OSS Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/gpt_oss/README.md).
19452

19553
## API reference
19654

19755
- GPT OSS recipes: [bridge.recipes.gpt_oss](../../apidocs/bridge/bridge.recipes.gpt_oss.md)
19856
- GPT OSS model provider: [bridge.models.gpt_oss.GPTOSSProvider](../../apidocs/bridge/bridge.models.gpt_oss.md)
19957

200-
## Performance optimizations
201-
202-
### Memory efficiency
203-
- **Manual GC**: Aggressive garbage collection (interval=100) for stable memory usage
204-
- **Precision-aware optimizer**: BF16 gradients and optimizer states
205-
- **Expert parallelism**: Distributes experts across GPUs (EP=4 for 20B, EP=16 for 120B)
206-
- **Sequence parallel**: Reduces activation memory across tensor parallel ranks
207-
- **Context parallel**: Splits long sequences across multiple GPUs
208-
209-
### Compute efficiency
210-
- **MoE permute fusion**: Fuses expert permutation operations
211-
- **Grouped GEMM**: Optimized expert computation with grouped matrix multiplications
212-
- **AllToAll dispatcher**: Efficient token routing across expert parallel ranks
213-
- **Bias activation fusion**: Fuses bias addition with activation functions
214-
- **Gradient overlapping**: Overlaps gradient all-reduce with backward computation
215-
216-
### Attention optimizations
217-
- **Sliding window attention**: Reduces attention computation with 128-token windows
218-
- **Window attention skip frequency**: Alternates between windowed and full attention every 2 layers
219-
- **Sink attention**: Learnable softmax offsets for improved attention stability
220-
- **Flash Attention**: FlashAttention-2 support via Transformer Engine
221-
- **Activation clamping**: Prevents numerical instability with [-7.0, 7.0] clamping
222-
22358
## Hugging Face model cards
22459

22560
### GPT OSS 20B
@@ -234,4 +69,3 @@ cfg = gpt_oss_120b_finetune_config(
23469
- Training configuration: [Configuration overview](../../training/config-container-overview.md)
23570
- Training entry points: [Entry points](../../training/entry-points.md)
23671
- Attention optimizations: [Attention optimizations](../../training/attention-optimizations.md)
237-

‎examples/models/gpt_oss/README.md‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
# GPT-OSS Examples
2+
3+
This directory contains example scripts for GPT-OSS 20B language models.
4+
5+
For model introduction and architecture details, see the GPT-OSS documentation.
6+
7+
## Workspace Configuration
8+
9+
All scripts use a `WORKSPACE` environment variable to define the base directory for checkpoints and results. By default, this is set to `/workspace`. You can override it:
10+
11+
```bash
12+
export WORKSPACE=/your/custom/path
13+
```
14+
15+
Directory structure:
16+
- `${WORKSPACE}/models/` - Converted checkpoints
17+
- `${WORKSPACE}/results/` - Training outputs and experiment results
18+
19+
## Checkpoint Conversion
20+
21+
See the [conversion.sh](conversion.sh) script for checkpoint conversion examples.
22+
23+
- **Import**: Use `openai/gpt-oss-20b` as the source Hugging Face model.
24+
- **Export**: Use `unsloth/gpt-oss-20b-BF16` as the reference HF model for export because the exported Megatron checkpoint is unquantized (bf16), which matches that repo's format.
25+
26+
### Import HF → Megatron
27+
28+
To import the HF model to your desired Megatron path:
29+
30+
```bash
31+
python examples/conversion/convert_checkpoints.py import \
32+
--hf-model openai/gpt-oss-20b \
33+
--megatron-path ${WORKSPACE}/models/gpt-oss-20b \
34+
--trust-remote-code
35+
```
36+
37+
### Export Megatron → HF
38+
39+
The export uses `unsloth/gpt-oss-20b-BF16` as the reference so the saved HF checkpoint matches that unquantized format:
40+
41+
```bash
42+
python examples/conversion/convert_checkpoints.py export \
43+
--hf-model unsloth/gpt-oss-20b-BF16 \
44+
--megatron-path ${WORKSPACE}/models/gpt-oss-20b/iter_0000000 \
45+
--hf-path ${WORKSPACE}/models/gpt-oss-20b-hf-export
46+
```
47+
48+
### Round-trip Validation
49+
50+
Multi-GPU round-trip validation between formats:
51+
52+
```bash
53+
python -m torch.distributed.run --nproc_per_node=8 \
54+
examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
55+
--hf-model-id unsloth/gpt-oss-20b-BF16 \
56+
--megatron-load-path ${WORKSPACE}/models/gpt-oss-20b/iter_0000000 \
57+
--tp 2 --pp 2 \
58+
--trust-remote-code
59+
```
60+
61+
## Training Recipes
62+
63+
- See: [bridge.recipes.gpt_oss](../../../src/megatron/bridge/recipes/gpt_oss/gpt_oss.py)
64+
- Available recipes:
65+
- `gpt_oss_20b_pretrain_config`: Pretraining configuration for 20B
66+
- `gpt_oss_20b_finetune_config`: Finetuning configuration for 20B with PEFT support
67+
- `gpt_oss_120b_pretrain_config`: Pretraining configuration for 120B
68+
- `gpt_oss_120b_finetune_config`: Finetuning configuration for 120B with PEFT support
69+
70+
Before training, ensure the following are configured:
71+
1. **Container Image**: Set `CONTAINER_IMAGE` in the SLURM scripts to your container path
72+
2. **Container Mounts**: (optional) Set `CONTAINER_MOUNTS` for data and workspace directories
73+
3. **Environment Variables**:
74+
- `HF_TOKEN`: to download models from HF Hub (if required)
75+
- `HF_HOME`: (optional) to avoid re-downloading models and datasets
76+
- `WANDB_API_KEY`: (optional) to enable WandB logging
77+
78+
All training scripts use SLURM for containerized multi-node training.
79+
80+
### Pretrain
81+
82+
Pretrain uses the **DCLM** dataset by default when `DCLM_DATA_DIR` and `DCLM_CACHE` are set (see [slurm_pretrain.sh](slurm_pretrain.sh)).
83+
84+
To use your own preprocessed DCLM data, set the dataset config as follows (e.g. in the recipe or via overrides):
85+
86+
```python
87+
cfg.dataset.blend = [
88+
[f"/path/to/dclm/preprocessed/dclm_{i:02d}_text_document" for i in range(1, 11)],
89+
None,
90+
]
91+
cfg.dataset.split = "9999,8,2"
92+
cfg.dataset.path_to_cache = "/path/to/cache"
93+
```
94+
95+
Preprocess your data using the [DCLM data preprocessing tutorial](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/tutorials/data/dclm).
96+
97+
W&B report coming soon.
98+
99+
### Supervised Fine-Tuning (SFT)
100+
101+
See the [slurm_sft.sh](slurm_sft.sh) script for full parameter fine-tuning. The recipe uses sequence packing by default.
102+
103+
W&B report coming soon.
104+
105+
### Parameter-Efficient Fine-Tuning (PEFT) with LoRA
106+
107+
See the [slurm_peft.sh](slurm_peft.sh) script for LoRA fine-tuning. The recipe uses sequence packing by default.
108+
109+
W&B report coming soon.
110+
111+
## Inference
112+
113+
See [inference.sh](inference.sh) for text generation with:
114+
- Hugging Face checkpoint (`unsloth/gpt-oss-20b-BF16`)
115+
- Imported Megatron checkpoint (after [conversion.sh](conversion.sh) import)
116+
- Exported HF checkpoint (after conversion export)
117+
- **SFT (finetuned) checkpoint**: set `SFT_CHECKPOINT` to your [slurm_sft.sh](slurm_sft.sh) result dir and run:
118+
119+
```bash
120+
uv run python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_text.py \
121+
--hf_model_path unsloth/gpt-oss-20b-BF16 \
122+
--megatron_model_path ${WORKSPACE}/results/gpt_oss_20b_finetune_tp2_pp2_ep4_spTrue_cp1 \
123+
--prompt "Hello, how are you?" \
124+
--max_new_tokens 64 \
125+
--tp 2 --pp 2 --ep 2 --etp 1 \
126+
--trust-remote-code
127+
```
128+
129+
TP×PP×EP must equal `--nproc_per_node`. Adjust parallelism to match your SFT run.
130+
131+
## Evaluation
132+
133+
Coming soon.

‎examples/models/gpt_oss/conversion.sh‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
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=gpt-oss-20b
22+
# Import: use openai/gpt-oss-20b as the source HF model
23+
HF_MODEL_ID_IMPORT=openai/gpt-oss-20b
24+
# Export: use unsloth/gpt-oss-20b-BF16 so exported checkpoint matches that repo's unquantized (bf16) format
25+
HF_MODEL_ID_EXPORT=unsloth/gpt-oss-20b-BF16
26+
27+
# Import HF → Megatron
28+
uv run python examples/conversion/convert_checkpoints.py import \
29+
--hf-model "$HF_MODEL_ID_IMPORT" \
30+
--megatron-path "${WORKSPACE}/models/${MODEL_NAME}" \
31+
--trust-remote-code
32+
33+
# Export Megatron → HF
34+
uv run python examples/conversion/convert_checkpoints.py export \
35+
--hf-model "$HF_MODEL_ID_EXPORT" \
36+
--megatron-path "${WORKSPACE}/models/${MODEL_NAME}/iter_0000000" \
37+
--hf-path "${WORKSPACE}/models/${MODEL_NAME}-hf-export"
38+
39+
# Round-trip validation
40+
uv run python -m torch.distributed.run --nproc_per_node=8 \
41+
examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
42+
--hf-model-id "$HF_MODEL_ID_EXPORT" \
43+
--megatron-load-path "${WORKSPACE}/models/${MODEL_NAME}/iter_0000000" \
44+
--tp 2 --pp 2 \
45+
--trust-remote-code

0 commit comments

Comments
 (0)