feat: Add support for deepseek recipes #5011

benieric · 2025-01-30T20:19:34Z

Issue #, if available:

Description of changes:

Adding support for deepseek recipes in Pytorch Estimator and ModelTrainer

Testing done:

Merge Checklist

Put an x in the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

I have read the CONTRIBUTING doc
I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the Python SDK team
I used the commit message format described in CONTRIBUTING
I have passed the region in to all S3 and STS clients that I've initialized as part of this change.
I have updated any necessary documentation, including READMEs and API docs (if appropriate)

Tests

I have added tests that prove my fix is effective or that my feature works (if appropriate)
I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes
I have checked that my tests are not configured for a specific region or account (if appropriate)
I have used unique_name_from_base to create resource names in integ tests (if appropriate)
If adding any dependency in requirements.txt files, I have spell checked and ensured they exist in PyPi

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

rohithn1 · 2025-01-30T20:27:23Z

tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py

+                "model_base_name": "mistral",
+            },
+            {
+                "model_type": "deepseek_llamav3",


to be extra safe can we align on deepseek_r1_llama_v3 and deepseek_r1_qwen_v2?

Would these have different training scripts? Or you mean just add a test case for

no im just saying instead deepseek_llamav3 lets do deepseek_r1_llama_v3

I think it will be easier to extend support to the v3 and native r1 model from deepseek if we use deepseek_r1_llama_v3 over deepseek_llamav3

schinmayee · 2025-01-30T21:03:56Z

tests/unit/sagemaker/modules/train/sm_recipes/test_utils.py

@@ -178,3 +179,37 @@ def test_get_args_from_recipe_compute(
            assert mock_gpu_args.call_count == 0
            assert mock_trainium_args.call_count == 0
            assert args is None
+
+    @pytest.mark.parametrize(


Thanks Erick! Can you copy this test case/change to tests/unit/test_pytorch.py as well?

Added some for pytorch as well

schinmayee

Thanks for the quick changes!

nargokul · 2025-01-31T00:07:53Z

src/sagemaker/modules/train/sm_recipes/utils.py

+    model_type_to_script = {
+        "llama_v3": ("llama", "llama_pretrain.py"),
+        "mistral": ("mistral", "mistral_pretrain.py"),
+        "mixtral": ("mixtral", "mixtral_pretrain.py"),
+        "deepseek": ("deepseek", "deepseek_pretrain.py"),
+    }


Can we remove this redundancy and hardcode it somewhere ?

* feat: Add support for deeepseek recipes * pylint * add unit test

…urated Hub Phase 2 (#5070) * change: update image_uri_configs 01-27-2025 06:18:13 PST * fix: skip TF tests for unsupported versions (#5007) * fix: skip TF tests for unsupported versions * flake8 * change: update image_uri_configs 01-29-2025 06:18:08 PST * chore: add new images for HF TGI (#5005) * feat: add pytorch-tgi-inference 2.4.0 * add tgi 3.0.1 image * skip faulty test * formatting * formatting * add hf pytorch training 4.46 * update version alias * add py311 to training version * update tests with pyversion 311 * formatting --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: use jumpstart deployment config image as default optimization image (#4992) Co-authored-by: Erick Benitez-Ramos <[email protected]> * prepare release v2.238.0 * update development version to v2.238.1.dev0 * Fix ssh host policy (#4966) * Fix ssh host policy * Filter policy by algo- * Add docstring * Fix pylint * Fix docstyle summary * Unit test * Fix unit test * Change to unit test * Fix unit tests * Test comment out flaky tests * Readd the flaky tests * Remove flaky asserts * Remove flaky asserts --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * change: Allow telemetry only in supported regions (#5009) * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions --------- Co-authored-by: Roja Reddy Sareddy <[email protected]> * mpirun protocol - distributed training with @Remote decorator (#4998) * implemented multi-node distribution with @Remote function * completed unit tests * added distributed training with CPU and torchrun * backwards compatibility nproc_per_node * fixing code: permissions for non-root users, integration tests * fixed docstyle * refactor nproc_per_node for backwards compatibility * refactor nproc_per_node for backwards compatibility * pylint fix, newlines * added unit tests for bootstrap_environment remote * added mpirun protocol for distributed training with @Remote decorator * aligned mpi_utils_remote.py to mpi_utils.py for estimator * updated docstring for sagemaker sdk doc --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: Add support for deepseek recipes (#5011) * feat: Add support for deeepseek recipes * pylint * add unit test * feat: [JumpStart] Add access configs and training instance type variants artifact uri handling for Curated Hub Phase 2 training integration (#1653) * Add access config to training input for Curated Hub Training Integration * Add support to retrieve instance specific training artifact keys * Fix some typos and naming issues * Fix more typos * fix formatting issues with black * modify access config logic so accept_eula is passed into fit * update black formatting * Add more unit tests for passing access configs * fix style errors * fix for failing integ test * fix styles and integ test error * skip blocking integ test * fix formatting * remove env vars when access configs are being used * fix docstyle issue * update usage of access configs, remove conversion of training artifact key to uri * fix styling issues * fix styling issues * fix unit tests * fix adding hubaccessconfig only if hubcontentarn exists * move logic to JumpStartEstimator from Job * Fix styling issues * Remove unused code * fix styling issues * fix unit test failure * fix some formatting, add comments * remove typing for estimator in get_access_configs function * fix circular import dependency * fix styling issues --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * Always add code channel, regardless of network isolation (#1657) * fix formatting issue * fix formatting issue * fix formatting issue * fix tensorflow file --------- Co-authored-by: sagemaker-bot <[email protected]> Co-authored-by: Erick Benitez-Ramos <[email protected]> Co-authored-by: varunmoris <[email protected]> Co-authored-by: Gary Wang <[email protected]> Co-authored-by: ci <ci> Co-authored-by: parknate@ <[email protected]> Co-authored-by: rsareddy0329 <[email protected]> Co-authored-by: Roja Reddy Sareddy <[email protected]> Co-authored-by: Bruno Pistone <[email protected]>

…urated Hub Phase 2 (aws#5070) * change: update image_uri_configs 01-27-2025 06:18:13 PST * fix: skip TF tests for unsupported versions (aws#5007) * fix: skip TF tests for unsupported versions * flake8 * change: update image_uri_configs 01-29-2025 06:18:08 PST * chore: add new images for HF TGI (aws#5005) * feat: add pytorch-tgi-inference 2.4.0 * add tgi 3.0.1 image * skip faulty test * formatting * formatting * add hf pytorch training 4.46 * update version alias * add py311 to training version * update tests with pyversion 311 * formatting --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: use jumpstart deployment config image as default optimization image (aws#4992) Co-authored-by: Erick Benitez-Ramos <[email protected]> * prepare release v2.238.0 * update development version to v2.238.1.dev0 * Fix ssh host policy (aws#4966) * Fix ssh host policy * Filter policy by algo- * Add docstring * Fix pylint * Fix docstyle summary * Unit test * Fix unit test * Change to unit test * Fix unit tests * Test comment out flaky tests * Readd the flaky tests * Remove flaky asserts * Remove flaky asserts --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * change: Allow telemetry only in supported regions (aws#5009) * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions --------- Co-authored-by: Roja Reddy Sareddy <[email protected]> * mpirun protocol - distributed training with @Remote decorator (aws#4998) * implemented multi-node distribution with @Remote function * completed unit tests * added distributed training with CPU and torchrun * backwards compatibility nproc_per_node * fixing code: permissions for non-root users, integration tests * fixed docstyle * refactor nproc_per_node for backwards compatibility * refactor nproc_per_node for backwards compatibility * pylint fix, newlines * added unit tests for bootstrap_environment remote * added mpirun protocol for distributed training with @Remote decorator * aligned mpi_utils_remote.py to mpi_utils.py for estimator * updated docstring for sagemaker sdk doc --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: Add support for deepseek recipes (aws#5011) * feat: Add support for deeepseek recipes * pylint * add unit test * feat: [JumpStart] Add access configs and training instance type variants artifact uri handling for Curated Hub Phase 2 training integration (aws#1653) * Add access config to training input for Curated Hub Training Integration * Add support to retrieve instance specific training artifact keys * Fix some typos and naming issues * Fix more typos * fix formatting issues with black * modify access config logic so accept_eula is passed into fit * update black formatting * Add more unit tests for passing access configs * fix style errors * fix for failing integ test * fix styles and integ test error * skip blocking integ test * fix formatting * remove env vars when access configs are being used * fix docstyle issue * update usage of access configs, remove conversion of training artifact key to uri * fix styling issues * fix styling issues * fix unit tests * fix adding hubaccessconfig only if hubcontentarn exists * move logic to JumpStartEstimator from Job * Fix styling issues * Remove unused code * fix styling issues * fix unit test failure * fix some formatting, add comments * remove typing for estimator in get_access_configs function * fix circular import dependency * fix styling issues --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * Always add code channel, regardless of network isolation (aws#1657) * fix formatting issue * fix formatting issue * fix formatting issue * fix tensorflow file --------- Co-authored-by: sagemaker-bot <[email protected]> Co-authored-by: Erick Benitez-Ramos <[email protected]> Co-authored-by: varunmoris <[email protected]> Co-authored-by: Gary Wang <[email protected]> Co-authored-by: ci <ci> Co-authored-by: parknate@ <[email protected]> Co-authored-by: rsareddy0329 <[email protected]> Co-authored-by: Roja Reddy Sareddy <[email protected]> Co-authored-by: Bruno Pistone <[email protected]>

* feat: Add support for deeepseek recipes * pylint * add unit test

…urated Hub Phase 2 (aws#5070) * change: update image_uri_configs 01-27-2025 06:18:13 PST * fix: skip TF tests for unsupported versions (aws#5007) * fix: skip TF tests for unsupported versions * flake8 * change: update image_uri_configs 01-29-2025 06:18:08 PST * chore: add new images for HF TGI (aws#5005) * feat: add pytorch-tgi-inference 2.4.0 * add tgi 3.0.1 image * skip faulty test * formatting * formatting * add hf pytorch training 4.46 * update version alias * add py311 to training version * update tests with pyversion 311 * formatting --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: use jumpstart deployment config image as default optimization image (aws#4992) Co-authored-by: Erick Benitez-Ramos <[email protected]> * prepare release v2.238.0 * update development version to v2.238.1.dev0 * Fix ssh host policy (aws#4966) * Fix ssh host policy * Filter policy by algo- * Add docstring * Fix pylint * Fix docstyle summary * Unit test * Fix unit test * Change to unit test * Fix unit tests * Test comment out flaky tests * Readd the flaky tests * Remove flaky asserts * Remove flaky asserts --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * change: Allow telemetry only in supported regions (aws#5009) * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions --------- Co-authored-by: Roja Reddy Sareddy <[email protected]> * mpirun protocol - distributed training with @Remote decorator (aws#4998) * implemented multi-node distribution with @Remote function * completed unit tests * added distributed training with CPU and torchrun * backwards compatibility nproc_per_node * fixing code: permissions for non-root users, integration tests * fixed docstyle * refactor nproc_per_node for backwards compatibility * refactor nproc_per_node for backwards compatibility * pylint fix, newlines * added unit tests for bootstrap_environment remote * added mpirun protocol for distributed training with @Remote decorator * aligned mpi_utils_remote.py to mpi_utils.py for estimator * updated docstring for sagemaker sdk doc --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: Add support for deepseek recipes (aws#5011) * feat: Add support for deeepseek recipes * pylint * add unit test * feat: [JumpStart] Add access configs and training instance type variants artifact uri handling for Curated Hub Phase 2 training integration (aws#1653) * Add access config to training input for Curated Hub Training Integration * Add support to retrieve instance specific training artifact keys * Fix some typos and naming issues * Fix more typos * fix formatting issues with black * modify access config logic so accept_eula is passed into fit * update black formatting * Add more unit tests for passing access configs * fix style errors * fix for failing integ test * fix styles and integ test error * skip blocking integ test * fix formatting * remove env vars when access configs are being used * fix docstyle issue * update usage of access configs, remove conversion of training artifact key to uri * fix styling issues * fix styling issues * fix unit tests * fix adding hubaccessconfig only if hubcontentarn exists * move logic to JumpStartEstimator from Job * Fix styling issues * Remove unused code * fix styling issues * fix unit test failure * fix some formatting, add comments * remove typing for estimator in get_access_configs function * fix circular import dependency * fix styling issues --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * Always add code channel, regardless of network isolation (aws#1657) * fix formatting issue * fix formatting issue * fix formatting issue * fix tensorflow file --------- Co-authored-by: sagemaker-bot <[email protected]> Co-authored-by: Erick Benitez-Ramos <[email protected]> Co-authored-by: varunmoris <[email protected]> Co-authored-by: Gary Wang <[email protected]> Co-authored-by: ci <ci> Co-authored-by: parknate@ <[email protected]> Co-authored-by: rsareddy0329 <[email protected]> Co-authored-by: Roja Reddy Sareddy <[email protected]> Co-authored-by: Bruno Pistone <[email protected]>

…urated Hub Phase 2 (#5070) * change: update image_uri_configs 01-27-2025 06:18:13 PST * fix: skip TF tests for unsupported versions (#5007) * fix: skip TF tests for unsupported versions * flake8 * change: update image_uri_configs 01-29-2025 06:18:08 PST * chore: add new images for HF TGI (#5005) * feat: add pytorch-tgi-inference 2.4.0 * add tgi 3.0.1 image * skip faulty test * formatting * formatting * add hf pytorch training 4.46 * update version alias * add py311 to training version * update tests with pyversion 311 * formatting --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: use jumpstart deployment config image as default optimization image (#4992) Co-authored-by: Erick Benitez-Ramos <[email protected]> * prepare release v2.238.0 * update development version to v2.238.1.dev0 * Fix ssh host policy (#4966) * Fix ssh host policy * Filter policy by algo- * Add docstring * Fix pylint * Fix docstyle summary * Unit test * Fix unit test * Change to unit test * Fix unit tests * Test comment out flaky tests * Readd the flaky tests * Remove flaky asserts * Remove flaky asserts --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * change: Allow telemetry only in supported regions (#5009) * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions --------- Co-authored-by: Roja Reddy Sareddy <[email protected]> * mpirun protocol - distributed training with @Remote decorator (#4998) * implemented multi-node distribution with @Remote function * completed unit tests * added distributed training with CPU and torchrun * backwards compatibility nproc_per_node * fixing code: permissions for non-root users, integration tests * fixed docstyle * refactor nproc_per_node for backwards compatibility * refactor nproc_per_node for backwards compatibility * pylint fix, newlines * added unit tests for bootstrap_environment remote * added mpirun protocol for distributed training with @Remote decorator * aligned mpi_utils_remote.py to mpi_utils.py for estimator * updated docstring for sagemaker sdk doc --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * feat: Add support for deepseek recipes (#5011) * feat: Add support for deeepseek recipes * pylint * add unit test * feat: [JumpStart] Add access configs and training instance type variants artifact uri handling for Curated Hub Phase 2 training integration (#1653) * Add access config to training input for Curated Hub Training Integration * Add support to retrieve instance specific training artifact keys * Fix some typos and naming issues * Fix more typos * fix formatting issues with black * modify access config logic so accept_eula is passed into fit * update black formatting * Add more unit tests for passing access configs * fix style errors * fix for failing integ test * fix styles and integ test error * skip blocking integ test * fix formatting * remove env vars when access configs are being used * fix docstyle issue * update usage of access configs, remove conversion of training artifact key to uri * fix styling issues * fix styling issues * fix unit tests * fix adding hubaccessconfig only if hubcontentarn exists * move logic to JumpStartEstimator from Job * Fix styling issues * Remove unused code * fix styling issues * fix unit test failure * fix some formatting, add comments * remove typing for estimator in get_access_configs function * fix circular import dependency * fix styling issues --------- Co-authored-by: Erick Benitez-Ramos <[email protected]> * Always add code channel, regardless of network isolation (#1657) * fix formatting issue * fix formatting issue * fix formatting issue * fix tensorflow file --------- Co-authored-by: sagemaker-bot <[email protected]> Co-authored-by: Erick Benitez-Ramos <[email protected]> Co-authored-by: varunmoris <[email protected]> Co-authored-by: Gary Wang <[email protected]> Co-authored-by: ci <ci> Co-authored-by: parknate@ <[email protected]> Co-authored-by: rsareddy0329 <[email protected]> Co-authored-by: Roja Reddy Sareddy <[email protected]> Co-authored-by: Bruno Pistone <[email protected]>

feat: Add support for deeepseek recipes

b695743

benieric requested a review from a team as a code owner January 30, 2025 20:19

benieric requested a review from mollyheamazon January 30, 2025 20:19

benieric temporarily deployed to auto-approve January 30, 2025 20:19 — with GitHub Actions Inactive

rohithn1 suggested changes Jan 30, 2025

View reviewed changes

pylint

0bd5104

benieric temporarily deployed to auto-approve January 30, 2025 20:38 — with GitHub Actions Inactive

schinmayee reviewed Jan 30, 2025

View reviewed changes

add unit test

0eb8359

benieric temporarily deployed to auto-approve January 30, 2025 21:30 — with GitHub Actions Inactive

schinmayee approved these changes Jan 30, 2025

View reviewed changes

haohanchen-aws approved these changes Jan 30, 2025

View reviewed changes

nargokul reviewed Jan 31, 2025

View reviewed changes

nargokul approved these changes Jan 31, 2025

View reviewed changes

benieric changed the title ~~feat: Add support for deeepseek recipes~~ feat: Add support for deepseek recipes Jan 31, 2025

benieric merged commit ae1146c into aws:master Jan 31, 2025
14 checks passed

Narrohag pushed a commit to Narrohag/sagemaker-python-sdk that referenced this pull request Mar 3, 2025

feat: Add support for deepseek recipes (aws#5011)

4533790

* feat: Add support for deeepseek recipes * pylint * add unit test

evakravi pushed a commit to evakravi/sagemaker-python-sdk that referenced this pull request Mar 20, 2025

feat: Add support for deepseek recipes (aws#5011)

c9124f1

* feat: Add support for deeepseek recipes * pylint * add unit test

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add support for deepseek recipes #5011

feat: Add support for deepseek recipes #5011

benieric commented Jan 30, 2025

rohithn1 Jan 30, 2025

benieric Jan 30, 2025

rohithn1 Jan 31, 2025 •

edited

Loading

schinmayee Jan 30, 2025

benieric Jan 30, 2025

schinmayee left a comment

nargokul Jan 31, 2025

feat: Add support for deepseek recipes #5011

feat: Add support for deepseek recipes #5011

Conversation

benieric commented Jan 30, 2025

Merge Checklist

General

Tests

rohithn1 Jan 30, 2025

Choose a reason for hiding this comment

benieric Jan 30, 2025

Choose a reason for hiding this comment

rohithn1 Jan 31, 2025 • edited Loading

Choose a reason for hiding this comment

schinmayee Jan 30, 2025

Choose a reason for hiding this comment

benieric Jan 30, 2025

Choose a reason for hiding this comment

schinmayee left a comment

Choose a reason for hiding this comment

nargokul Jan 31, 2025

Choose a reason for hiding this comment

rohithn1 Jan 31, 2025 •

edited

Loading