Skip to content

Feat: Optimize() validations across TRT, VLLM, Neuron container optimizations #4927

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 3 commits into from
Nov 19, 2024
Merged

Feat: Optimize() validations across TRT, VLLM, Neuron container optimizations #4927

merged 3 commits into from
Nov 19, 2024
+1,166 −21

Conversation

gwang111
Copy link
Collaborator

Issue #, if available:

Description of changes:

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.

@gwang111 gwang111 requested a review from Lokiiiiii November 13, 2024 21:30
@gwang111 gwang111 requested a review from a team as a code owner November 13, 2024 21:30
@gwang111 gwang111 removed the request for review from mufaddal-rohawala November 13, 2024 21:30
gwang111
gwang111 commented Nov 13, 2024
View reviewed changes
src/sagemaker/serve/builder/model_builder.py Outdated
@@ -1157,6 +1161,15 @@ def optimize(
Model: A deployable ``Model`` object.
"""

# TODO: ideally these dictionaries need to be sagemaker_core shapes
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: Ideally validations should happen at the start of the flow, not after everything is set up in the wrapper. Lets fast follow this and gradually shift the validations to helper functions in the outer scope here

Lokiiiiii
Lokiiiiii suggested changes Nov 14, 2024
View reviewed changes
src/sagemaker/serve/validations/optimization.py Outdated
logger = logging.getLogger(__name__)


class OptimizationContainer(Enum):
Copy link
Contributor

@Lokiiiiii Lokiiiiii Nov 14, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can use these for internal organization. But we should not burden customers with knowing how our optimizations work, which containers are used under which scenarios, etc.

Consider the following error messages -

  • Optimizations that use Compilation and Speculative Decoding are not currently supported on GPU & Neuron instances.
  • Optimizations that use Compilation and Quantization are not currently supported on Neuron instances.
  • Optimizations that use Quantization:SmoothQuant are only supported with Compilation on GPU instances.
  • Sharding is mutually exclusive and only supported on GPU instances.

The general theme we want to maintain is

Optimizations {A,B,C} are {not|only} supported with optimizations {X,Y,Z} on instances {L,M,N}

Lokiiiiii
Lokiiiiii suggested changes Nov 18, 2024
View reviewed changes
@cj-zhang
Copy link
Collaborator

serve UTs are failing, I see TypeError: unsupported operand type(s) for |: 'type' and 'NoneType'. Can you look into this please?

Lokiiiiii
Lokiiiiii suggested changes Nov 18, 2024
View reviewed changes
Lokiiiiii
Lokiiiiii suggested changes Nov 18, 2024
View reviewed changes
Ashish Gupta and others added 2 commits November 19, 2024 05:48
@gwang111
changes for blackbird - model sharding
f5b568a 
changes for blackbird - model sharding

add more tests

fix sharded model flag

add optimization validations

fix formatting and msging

fixing validation bugs

add UTs

simplify logic

update messaging

formatting

fix UTs

add more UTs

fix validations

update ruleset

update formatting

update validation logic

update bug fixes

Disable network isolation if using sharded models.

check sharding + network iso pre optimization

add more UTs for sharding

add more UTs
@gwang111
fix rebase issues
0707798
mufaddal-rohawala
mufaddal-rohawala reviewed Nov 19, 2024
View reviewed changes
src/sagemaker/serve/validations/optimization.py
}


def _validate_optimization_configuration(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is such a big method, we should look to optimise this?

Copy link
Collaborator Author

@gwang111 gwang111 Nov 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should be able to view the UTs here: https://github.com/aws/sagemaker-python-sdk/pull/4927/files#diff-5f58a16d071696e26bf4081f3f01294336f6c669407a9474ed7c0747af0e629fR3479. and Ack, we will look to optimize post reinvent

@mufaddal-rohawala mufaddal-rohawala merged commit efd6c80 into aws:master Nov 19, 2024
13 of 14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mufaddal-rohawala mufaddal-rohawala mufaddal-rohawala approved these changes

+2 more reviewers

@cj-zhang cj-zhang cj-zhang approved these changes

@Lokiiiiii Lokiiiiii Lokiiiiii approved these changes

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@gwang111 @cj-zhang @Lokiiiiii @mufaddal-rohawala