GPT-OSS 20B Verification Card (#5397)

Signed-off-by: Maanu Grover <maanug@nvidia.com>

Author: maanug-nv

examples/model_verification_cards/gpt-oss-20b/card.yaml

@@ -60,7 +60,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider
         provider.normalization = "RMSNorm"
         provider.gated_linear_unit = True
         provider.add_bias_linear = True
-        provider.add_qkv_bias = False
+        provider.add_qkv_bias = getattr(hf_pretrained.config, "attention_bias", True)
         provider.share_embeddings_and_output_weights = False
         provider.position_embedding_type = "yarn"
 
@@ -101,6 +101,13 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider
 
         return provider
 
+    @classmethod
+    def megatron_to_hf_config(cls, provider) -> dict:
+        """Convert Megatron provider config to a GPT-OSS Hugging Face config."""
+        hf_config = super().megatron_to_hf_config(provider)
+        hf_config["attention_bias"] = True
+        return hf_config
+
     def maybe_modify_loaded_hf_weight(
         self, hf_param: str | dict[str, str], hf_state_dict: Mapping[str, torch.Tensor]
     ) -> torch.Tensor:

src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py

@@ -19,6 +19,7 @@
     gpt_oss_20b_pretrain_config,
     gpt_oss_20b_pretrain_fp8_current_scaling_config,
     gpt_oss_20b_pretrain_mxfp8_config,
+    gpt_oss_20b_sft_32k_config,
     gpt_oss_20b_sft_config,
     gpt_oss_20b_sft_fp8_current_scaling_config,
     gpt_oss_20b_sft_mxfp8_config,
@@ -34,6 +35,7 @@
     "gpt_oss_20b_pretrain_fp8_current_scaling_config",
     "gpt_oss_120b_pretrain_config",
     "gpt_oss_20b_sft_config",
+    "gpt_oss_20b_sft_32k_config",
     "gpt_oss_20b_sft_openmathinstruct2_thinking_packed_config",
     "gpt_oss_20b_sft_fp8_current_scaling_config",
     "gpt_oss_120b_sft_config",

src/megatron/bridge/recipes/gpt_oss/__init__.py

@@ -38,6 +38,9 @@
 from megatron.bridge.recipes.gpt_oss.h100.gpt_oss import (
     gpt_oss_20b_pretrain_16gpu_h100_fp8mx_config as gpt_oss_20b_pretrain_mxfp8_config,
 )
+from megatron.bridge.recipes.gpt_oss.h100.gpt_oss import (
+    gpt_oss_20b_sft_8gpu_h100_bf16_32k_config as gpt_oss_20b_sft_32k_config,
+)
 from megatron.bridge.recipes.gpt_oss.h100.gpt_oss import (
     gpt_oss_20b_sft_8gpu_h100_bf16_config as gpt_oss_20b_sft_config,
 )
@@ -72,6 +75,7 @@
     "gpt_oss_20b_pretrain_fp8_current_scaling_config",
     "gpt_oss_20b_pretrain_mxfp8_config",
     "gpt_oss_20b_sft_config",
+    "gpt_oss_20b_sft_32k_config",
     "gpt_oss_20b_sft_fp8_current_scaling_config",
     "gpt_oss_20b_sft_mxfp8_config",
     "gpt_oss_20b_sft_openmathinstruct2_thinking_packed_config",

src/megatron/bridge/recipes/gpt_oss/gpt_oss.py

@@ -26,6 +26,7 @@
     "gpt_oss_20b_pretrain_16gpu_h100_fp8cs_config",
     "gpt_oss_20b_pretrain_16gpu_h100_fp8mx_config",
     "gpt_oss_20b_sft_8gpu_h100_bf16_config",
+    "gpt_oss_20b_sft_8gpu_h100_bf16_32k_config",
     "gpt_oss_20b_sft_8gpu_h100_bf16_openmathinstruct2_thinking_packed_config",
     "gpt_oss_20b_sft_8gpu_h100_fp8cs_config",
     "gpt_oss_20b_sft_8gpu_h100_fp8mx_config",

src/megatron/bridge/recipes/gpt_oss/h100/__init__.py

@@ -421,6 +421,23 @@ def gpt_oss_20b_sft_8gpu_h100_bf16_config() -> ConfigContainer:
     return cfg
 
 
+def gpt_oss_20b_sft_8gpu_h100_bf16_32k_config() -> ConfigContainer:
+    """Return the 32K-context full SFT config for GPT-OSS 20B."""
+    cfg = gpt_oss_20b_sft_8gpu_h100_bf16_config()
+
+    seq_length = 32768
+    cfg.model.seq_length = seq_length
+    cfg.dataset.seq_length = seq_length
+    cfg.dataset.offline_packing_specs.packed_sequence_size = seq_length
+    cfg.train.global_batch_size = 32
+
+    # Keep the complete process environment visible on the recipe.
+    cfg.env_vars = {
+        **COMMON_RECIPE_ENV_VARS,
+    }
+    return cfg
+
+
 def gpt_oss_120b_sft_32gpu_h100_bf16_config() -> ConfigContainer:
     """Return a full SFT config for GPT-OSS 120B.
 
@@ -936,6 +953,7 @@ def gpt_oss_20b_sft_8gpu_h100_bf16_openmathinstruct2_thinking_packed_config() ->
     "gpt_oss_20b_pretrain_16gpu_h100_fp8cs_config",
     "gpt_oss_20b_pretrain_16gpu_h100_fp8mx_config",
     "gpt_oss_20b_sft_8gpu_h100_bf16_config",
+    "gpt_oss_20b_sft_8gpu_h100_bf16_32k_config",
     "gpt_oss_20b_sft_8gpu_h100_bf16_openmathinstruct2_thinking_packed_config",
     "gpt_oss_20b_sft_8gpu_h100_fp8cs_config",
     "gpt_oss_20b_sft_8gpu_h100_fp8mx_config",

src/megatron/bridge/recipes/gpt_oss/h100/gpt_oss.py

@@ -40,6 +40,7 @@ def gpt_oss_cfg(self):
             "vocab_size": 201088,
             "hidden_act": "silu",
             "sliding_window": 4096,
+            "attention_bias": True,
         }
 
     @pytest.fixture
@@ -64,6 +65,15 @@ def test_provider_bridge_maps_config(self, mock_pretrained):
         # Key fields mapped from HF config
         assert provider.num_layers == mock_pretrained.config.num_hidden_layers
         assert provider.num_moe_experts == mock_pretrained.config.num_local_experts
+        assert provider.add_qkv_bias == mock_pretrained.config.attention_bias
         # dtype mapping
         assert provider.bf16 is True
         assert provider.params_dtype == torch.bfloat16
+
+    def test_megatron_to_hf_config_preserves_attention_bias(self, mock_pretrained):
+        bridge = GPTOSSBridge()
+        provider = bridge.provider_bridge(mock_pretrained)
+
+        hf_config = bridge.megatron_to_hf_config(provider)
+
+        assert hf_config["attention_bias"] is True

tests/unit_tests/models/gpt_oss/test_gpt_oss_bridges.py

@@ -152,6 +152,22 @@ def test_fp8_sft_inherits_base_sft_config():
     assert base.model.moe_router_padding_for_fp8 is False
 
 
+def test_gpt_oss_20b_sft_32k_overrides_batch_and_sequence_length():
+    """gpt_oss_20b_sft_32k_config should inherit SFT defaults with a 32K context and GBS32."""
+    from megatron.bridge.recipes.gpt_oss import gpt_oss_20b_sft_32k_config
+    from megatron.bridge.recipes.gpt_oss.h100.gpt_oss import gpt_oss_20b_sft_8gpu_h100_bf16_32k_config
+
+    assert gpt_oss_20b_sft_32k_config is gpt_oss_20b_sft_8gpu_h100_bf16_32k_config
+
+    cfg = gpt_oss_20b_sft_32k_config()
+    _assert_basic_config(cfg)
+    assert cfg.model.seq_length == 32768
+    assert cfg.dataset.seq_length == 32768
+    assert cfg.dataset.offline_packing_specs.packed_sequence_size == 32768
+    assert cfg.train.global_batch_size == 32
+    assert cfg.train.micro_batch_size == 1
+
+
 def test_fp8_peft_inherits_base_peft_config():
     """gpt_oss_20b_peft_fp8_current_scaling_config should inherit base PEFT settings."""
     from megatron.bridge.recipes.gpt_oss import (