Amazon SageMaker AI

Accelerate generative AI development with Amazon SageMaker AI and MLflow

Efficiently manage the machine learning and generative AI lifecycle at scale using MLflow 3.0

Why use Amazon SageMaker with MLflow?

Amazon SageMaker offers a managed MLflow capability for machine learning (ML) and generative AI experimentation. This capability makes it easy for data scientists to use MLflow on SageMaker for model training, registration, and deployment. Admins can quickly set up secure and scalable MLflow environments on AWS. Data scientists and ML developers can efficiently track ML experiments and find the right model for a business problem.

Benefits of Amazon SageMaker AI with MLflow 3.0

Efficiently experiment with foundation models (FMs)

Data Scientists can use MLflow to keep track of all the metrics generated during fine-tuning of a foundation model, evaluate the model, test the model with sample data, compare the outputs of each model side by side on the MLflow UI, and register the right model for their use case. Once they register the model, ML engineers can deploy the model to SageMaker inference.

Zero infrastructure maintenance

You do not need to manage any infrastructure required to host MLflow. Data Scientists can use all the MLflow open source capabilities without admins worrying about the infrastructure overhead. This saves time and cost when setting up data science environments. MLflow is integrated with Amazon Identity and Access Management (IAM), allowing you to set up Role Based Access Control (RBAC) for MLflow Tracking Servers.

Seamlessly transition from model experimentation to production

Models registered in MLflow will automatically be registered to the Amazon SageMaker Model Registry with an associated Amazon SageMaker Model Card. This enables data scientists to transition their models to ML engineers for production deployment without switching context. ML Engineers can deploy models from MLflow to SageMaker endpoints without building custom containers or repackaging the MLflow model artifacts.

Innovate with the open-source community

As the MLflow project evolves, SageMaker AI customers will benefit from the open-source innovation from the MLflow community while enjoying the infrastructure management provided by AWS.

Advanced tracing capabilities

Tracing capabilities in fully managed MLflow 3.0 enable customers to record the inputs, outputs, and metadata at every step of gen AI development to help teams quickly identify the source of bugs or unexpected behaviors. By maintaining records of each model and application version, fully managed MLflow 3.0 offers traceability to connect AI responses to their source components, allowing developers to quickly trace an issue directly to the specific code, data, or parameters that generated it.

Track experiments from anywhere

ML experiments are performed in diverse environments, including local notebooks, IDEs, cloud-based training code, or managed IDEs in Amazon SageMaker Studio. With SageMaker AI and MLflow, you can use your preferred environment to train models, track your experiments in MLflow, and launch the MLflow UI directly or through SageMaker Studio for analysis.

Screenshot of the Amazon SageMaker Studio interface showing the MLflow experiment logging feature, with application options such as JupyterLab, RStudio, and MLflow, and Python code for machine learning experiment tracking and autologging within the Studio environment.

Accelerate generative AI development with MLflow 3.0

Building foundation models is an iterative process, involving hundreds of training iterations to find the best algorithm, architecture, and parameters for optimal model accuracy. Fully-managed MLflow 3.0 enables you to track gen AI experiments, evaluate model performance, and gain deeper insights into the behavior of models and AI applications from experimentation to production. With a single interface, you can visualize-progress training jobs, collaborate with colleagues during experimentation, and maintain version control for each model and application. MLflow 3.0 also offers advanced tracing capabilities that record the inputs, outputs, and metadata at every step of gen AI development, enabling you to quickly identify the source of bugs or unexpected behaviors.

Accelerate generative AI development with MLflow

Screenshot of the MLflow 2.13.2 experiment metadata management dashboard displaying experiment runs, evaluation metrics, and a comparison chart of model accuracy for different runs.

Evaluate experiments

Identifying the best model from multiple iterations requires analysis and comparison of model performance. MLflow offers visualizations such as scatter plots, bar charts, and histograms to compare training iterations. Additionally, MLflow enables the evaluation of models for bias and fairness.

Screenshot of an MLflow interface comparing 3 runs from 1 experiment, displaying a contour plot visualization with accuracy on the X-axis, F1 score on the Y-axis, and precision on the Z-axis.

Centrally manage MLflow models

Multiple teams often use MLflow to manage their experiments, with only some models becoming candidates for production. Organizations need an easy way to keep track of all candidate models to make informed decisions about which models proceed to production. MLflow integrates seamlessly with SageMaker Model Registry, allowing organizations to see their models registered in MLflow automatically appear in SageMaker Model Registry, complete with a SageMaker Model Card for governance. This integration enables data scientists and ML engineers to use distinct tools for their respective tasks: MLflow for experimentation and SageMaker Model Registry for managing the production lifecycle with comprehensive model lineage.

Screenshot of the MLflow model overview page in Amazon SageMaker showing details for the model 'flawless-jay-491', including a Finished status, creation date, experiment ID, source, logged models, and other metadata.

Deploy MLflow Models to SageMaker endpoints

Deploying models from MLflow to SageMaker Endpoints is seamless, eliminating the need to build custom containers for model storage. This integration allows customers to leverage SageMaker’s optimized inference containers while retaining the user-friendly experience of MLflow for logging and registering models.