Support for GPT-OSS pretraining script in Megatron-Bridge

[fengxy-03]

Is your feature request related to a problem? Please describe.
I am interested in pretraining a model from scratch using the GPT-OSS architecture via Megatron-Bridge.
However, I couldn't find a corresponding configuration or script for GPT-OSS in the current examples/recipes/ directory. Could you please provide a pretrain_gpt_oss.py script (or a similar recipe) to support training this architecture from scratch?

Describe the solution you'd like
Provide a pretrain_gpt_oss.py script (or a similar recipe) to support training GPT-OSS architecture.

Describe alternatives you've considered

Additional context

[fengxy-03]

I've written the pretrain_gpt_oss.py script myself and it runs, but the performance is very poor.

The TFLOPS is stuck around 10 regardless of whether FlashAttention is toggled on or off. It seems FlashAttention is not being triggered correctly in my implementation.

Here is my code:

#!/usr/bin/env python3
# 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.

import argparse
import logging
import os
import sys
from typing import Tuple

import torch
from omegaconf import OmegaConf

from megatron.bridge.recipes.gpt_oss import gpt_oss_20b_pretrain_config as pretrain_config
from megatron.bridge.training.config import ConfigContainer
from megatron.bridge.training.gpt_step import forward_step
from megatron.bridge.training.pretrain import pretrain
from megatron.bridge.training.utils.omegaconf_utils import (
    apply_overrides,
    create_omegaconf_dict_config,
    parse_hydra_overrides,
)


logger: logging.Logger = logging.getLogger(__name__)


def parse_cli_args() -> Tuple[argparse.Namespace, list[str]]:
    """Parse command line arguments, separating known script args from OmegaConf overrides."""
    parser = argparse.ArgumentParser(
        description="Pretrain Llama3 8B model using Megatron-Bridge with YAML and CLI overrides",
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument(
        "--config-file",
        type=str,
        help="Path to the YAML OmegaConf override file. Default: conf/llama3_8b_pretrain_override_example.yaml",
    )
    # Parse known args for the script, remaining will be treated as overrides
    args, cli_dotlist_overrides = parser.parse_known_args()
    return args, cli_dotlist_overrides


def main() -> None:
    """
    Entry point for the Mamba 8B pretraining script.
    """
    args, cli_overrides = parse_cli_args()

    cfg: ConfigContainer = pretrain_config(
        hf_path="",
        train_iters=None,
        lr_warmup_iters=0,
    )

    # Convert the initial Python dataclass to an OmegaConf DictConfig for merging
    merged_omega_conf, excluded_fields = create_omegaconf_dict_config(cfg)

    # Load and merge YAML overrides if a config file is provided
    if args.config_file:
        logger.debug(f"Loading YAML overrides from: {args.config_file}")
        if not os.path.exists(args.config_file):
            logger.error(f"Override YAML file not found: {args.config_file}")
            sys.exit(1)
        yaml_overrides_omega = OmegaConf.load(args.config_file)
        merged_omega_conf = OmegaConf.merge(merged_omega_conf, yaml_overrides_omega)
        logger.debug("YAML overrides merged successfully.")

    # Apply command-line overrides using Hydra-style parsing
    if cli_overrides:
        logger.debug(f"Applying Hydra-style command-line overrides: {cli_overrides}")
        merged_omega_conf = parse_hydra_overrides(merged_omega_conf, cli_overrides)
        logger.debug("Hydra-style command-line overrides applied successfully.")

    # Apply the final merged OmegaConf configuration back to the original ConfigContainer
    logger.debug("Applying final merged configuration back to Python ConfigContainer...")
    final_overrides_as_dict = OmegaConf.to_container(merged_omega_conf, resolve=True)
    # Apply overrides while preserving excluded fields
    apply_overrides(cfg, final_overrides_as_dict, excluded_fields)

    # Start training
    logger.debug("Starting pretraining...")
    pretrain(config=cfg, forward_step_func=forward_step)

    if torch.distributed.is_initialized():
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()

[yaoyu-33]

https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/scripts/performance/configs/gpt_oss can you check this script. we have perf numbers published here https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html

[malay-nagda]

@fengxy-03 let me know if you need help running the script noted by @yaoyu-33

Can you tell what GPU and precision you are using?

[fengxy-03]

@yaoyu-33
https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/scripts/performance/configs/gpt_oss .This repo only has the GPT-OSS configs but doesn't include the pre-training code like pretrain_gpt.py.
Also, I noticed the GPT-OSS-120B is supposed to hit 153 TFLOPS on H100, but when I tested a small 0.15B model on an H100, I only got about 15 TFLOPS. Plus, the VRAM is almost full. The performance seems way off for a model this size on an H100.

[fengxy-03]

@malay-nagda
I tested a small 0.15B GPT-OSS model on an H100, I only got about 15 TFLOPS. Plus, the VRAM is almost full. The precision is bf16 and the train_seq length is 8192.

Also, based on the discussions in this pr: NVIDIA/Megatron-LM#2383, I used Megatron-LM to train a GPT-OSS model with the same parameters. I found that enabling flash attention actually dropped the performance to about 5 TFLOPS. However, when I set --attention-backend fused, the performance improved to around 110 TFLOPS, which seems much more normal.

[malay-nagda]

This repo only has the GPT-OSS configs but doesn't include the pre-training code like pretrain_gpt.py.

@fengxy-03 you can launch pre-training with this config using scripts/performance/setup_experiment.py (this is from login node). This script launches scripts/performance/run_script.py on compute node and run_script.py executes pre-train code.

More info about setup here- https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/scripts/performance.
Can you try this and update your results here?

Also, I noticed the GPT-OSS-120B is supposed to hit 153 TFLOPS on H100, but when I tested a small 0.15B model on an H100, I only got about 15 TFLOPS.

This was for configs in https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/scripts/performance/configs/gpt_oss and 4096 seq length. If you are using 8192 seq length, I think the perf will be lower.

[malay-nagda]

The following steps might make it easier for you to launch pre-training using scripts/performance/setup_experiment.py

python3 -m venv bridge_venv
source bridge_venv/bin/activate
pip3 install git+https://github.com/NVIDIA-NeMo/Run.git

git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
cd Megatron-Bridge
git checkout r0.2.0

export NEMORUN_HOME=<results_dir>

python3 scripts/performance/setup_experiment.py -a <slurm_account> -p <slurm_partition> -g h100 -i nvcr.io/nvidia/nemo:25.11.01 -l <results_dir> -vb 1 -m gpt_oss -s 120b -ng 64 -hf <HF_ACCESS_TOKEN>

You can generate your HuggingFace Access Token from https://huggingface.co/settings/tokens/new?tokenType=fineGrained