From 86e29ea1928375cad42c793f5fe5f441a03666ad Mon Sep 17 00:00:00 2001 From: adtian2 Date: Wed, 20 Aug 2025 11:32:29 -0700 Subject: [PATCH 1/3] changes for gpt_oss jobs support --- src/sagemaker/modules/train/sm_recipes/utils.py | 1 + src/sagemaker/pytorch/estimator.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/sagemaker/modules/train/sm_recipes/utils.py b/src/sagemaker/modules/train/sm_recipes/utils.py index 3b7659016e..b6523e14dd 100644 --- 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 index 633317927b..208239e368 100644 --- 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: From be0343d9e1d328e082425c5e7d9a91dc63b1a511 Mon Sep 17 00:00:00 2001 From: adtian2 Date: Wed, 20 Aug 2025 11:36:55 -0700 Subject: [PATCH 2/3] added unit tests --- .../unit/sagemaker/modules/train/sm_recipes/test_utils.py | 5 +++++ tests/unit/test_pytorch.py | 8 ++++++++ 2 files changed, 13 insertions(+) diff --git a/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py b/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py index a58b1f641e..a8965c303e 100644 --- a/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py +++ b/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py @@ -237,6 +237,11 @@ def test_get_args_from_recipe_with_nova_and_role(mock_get_args_from_nova_recipe, "script": "deepseek_pretrain.py", "model_base_name": "deepseek", }, + { + "model_type": "gpt_oss", + "script": "custom_pretrain.py", + "model_base_name": "gpt_oss", + }, ], ) def test_get_trainining_recipe_gpu_model_name_and_script(test_case): diff --git a/tests/unit/test_pytorch.py b/tests/unit/test_pytorch.py index 34d3c6784b..8352f3090b 100644 --- a/tests/unit/test_pytorch.py +++ b/tests/unit/test_pytorch.py @@ -1087,6 +1087,14 @@ def test_training_recipe_for_trainium(sagemaker_session): }, }, }, + { + "script": "custom_pretrain.py", + "recipe": { + "model": { + "model_type": "gpt_oss", + }, + }, + }, ], ) @patch("shutil.copyfile") From 848de311a8634638c728d576fddbbf89645e8418 Mon Sep 17 00:00:00 2001 From: adtian2 Date: Wed, 20 Aug 2025 14:02:47 -0700 Subject: [PATCH 3/3] fixing unit test --- tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py b/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py index a8965c303e..17cfda55b0 100644 --- a/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py +++ b/tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py @@ -240,7 +240,7 @@ def test_get_args_from_recipe_with_nova_and_role(mock_get_args_from_nova_recipe, { "model_type": "gpt_oss", "script": "custom_pretrain.py", - "model_base_name": "gpt_oss", + "model_base_name": "custom_model", }, ], )