feat: Add GPT-OSS support via mcore (#1452)

Signed-off-by: ashors1 <[email protected]>

Author: ashors1

examples/configs/recipes/llm/grpo-gptoss-20b-8n8g-megatron.yaml

@@ -0,0 +1,26 @@
+defaults: ../../grpo_math_1B.yaml
+grpo:
+  num_prompts_per_step: 64
+  num_generations_per_prompt: 32
+loss_fn:
+  use_importance_sampling_correction: true
+policy:
+  model_name: openai/gpt-oss-20b
+  train_micro_batch_size: 1
+  max_total_sequence_length: 4096
+  megatron_cfg:
+    enabled: true
+    expert_model_parallel_size: 8
+    tensor_model_parallel_size: 4
+    sequence_parallel: true
+    moe_permute_fusion: true
+  dtensor_cfg:
+    enabled: false
+  sequence_packing:
+    enabled: false
+  generation:
+    vllm_cfg:
+      tensor_parallel_size: 2
+cluster:
+  num_nodes: 8
+  gpus_per_node: 8

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -20,6 +20,7 @@
 
 import ray
 import torch
+from transformers import AutoConfig
 
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches
@@ -305,6 +306,17 @@ def _patch_vllm_init_workers_ray():
             self.cfg["vllm_cfg"].get("hf_overrides", {}) or {}
         )
 
+        # Override HF config for gpt-oss models to ensure compatibility with megatron
+        # The megatron --> hf export is done in bf16, so we disable quantization
+        hf_config = AutoConfig.from_pretrained(self.model_name, trust_remote_code=True)
+        if "GptOssForCausalLM" in getattr(hf_config, "architectures", []):
+            if "quantization_config" in hf_config:
+                assert load_format == "dummy", (
+                    "Loading quantized GPT-OSS models is currently only supported with load_format='dummy'."
+                )
+                # disable quantization
+                vllm_kwargs["hf_overrides"]["quantization_config"] = {}
+
         llm_kwargs = dict(
             model=self.model_name,
             served_model_name=self.model_name,

tests/test_suites/llm/grpo-gptoss-20b-8n8g-megatron.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=8
+STEPS_PER_RUN=60
+MAX_STEPS=60
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=240
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_grpo_math.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    "$@" \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'mean(data["train/gen_kl_error"]) < 0.002' \
+        'data["train/reward"]["60"] > 0.60' \
+        'mean(data["timing/train/total_step_time"], -6, -1) < 210'
+fi

tests/test_suites/release.txt

@@ -23,6 +23,9 @@ tests/test_suites/llm/dapo-qwen2.5-7b.sh
 # Deepseek-V3 on DAPO dataset
 tests/test_suites/llm/grpo-dapomath17k-dsv3-megatron.sh
 
+# GPT-OSS
+tests/test_suites/llm/grpo-gptoss-20b-8n8g-megatron.sh
+
 #######
 # SFT #
 #######