changes for gpt_oss jobs support

adtian2 · adtian2 · commit 86e29ea19283 · 2025-08-20T11:32:29.000-07:00
diff --git a/src/sagemaker/modules/train/sm_recipes/utils.py b/src/sagemaker/modules/train/sm_recipes/utils.py
@@ -136,6 +136,7 @@ def _get_trainining_recipe_gpu_model_name_and_script(model_type: str):
         "mistral": ("mistral", "mistral_pretrain.py"),
         "mixtral": ("mixtral", "mixtral_pretrain.py"),
         "deepseek": ("deepseek", "deepseek_pretrain.py"),
+        "gpt_oss": ("custom_model", "custom_pretrain.py"),
     }
 
     for key in model_type_to_script:
diff --git a/src/sagemaker/pytorch/estimator.py b/src/sagemaker/pytorch/estimator.py
@@ -99,6 +99,7 @@ def _get_training_recipe_gpu_script(code_dir, recipe, source_dir):
         "mistral": ("mistral", "mistral_pretrain.py"),
         "mixtral": ("mixtral", "mixtral_pretrain.py"),
         "deepseek": ("deepseek", "deepseek_pretrain.py"),
+        "gpt_oss": ("custom_model", "custom_pretrain.py"),
     }
 
     if "model" not in recipe:

Original file line number	Diff line number	Diff line change
`@@ -136,6 +136,7 @@ def _get_trainining_recipe_gpu_model_name_and_script(model_type: str):`
`136`	`136`	`"mistral": ("mistral", "mistral_pretrain.py"),`
`137`	`137`	`"mixtral": ("mixtral", "mixtral_pretrain.py"),`
`138`	`138`	`"deepseek": ("deepseek", "deepseek_pretrain.py"),`
	`139`	`+ "gpt_oss": ("custom_model", "custom_pretrain.py"),`
`139`	`140`	`}`
`140`	`141`
`141`	`142`	`for key in model_type_to_script:`
Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,7 @@ def _get_training_recipe_gpu_script(code_dir, recipe, source_dir):`
`99`	`99`	`"mistral": ("mistral", "mistral_pretrain.py"),`
`100`	`100`	`"mixtral": ("mixtral", "mixtral_pretrain.py"),`
`101`	`101`	`"deepseek": ("deepseek", "deepseek_pretrain.py"),`
	`102`	`+ "gpt_oss": ("custom_model", "custom_pretrain.py"),`
`102`	`103`	`}`
`103`	`104`
`104`	`105`	`if "model" not in recipe:`