ModelTrainer - SageMaker PySDK Training Redesign

Introductions

The ModelTrainer class in the SageMaker Python SDK simplifies the process of launching and managing training jobs on AWS SageMaker. It provides an intuitive interface for customizing training jobs, including the ability to specify custom scripts, custom containers, distributed training configurations and executing training locally. In this notebook, we outline how to get started with the ModelTrainer class, its features, and examples to help you effectively leverage its capabilities.

Benefits of the ModelTrainer

The ModelTrainer is designed to address the usability challenges associated with the Estimator class. Moving training with the SageMaker PySDK towards achieving a best-in-class developer experience.

Key Improvements Include:

1. Improved Intuitiveness - The ModelTrainer reduces complexity by leveraging configuration classes and minimizing the interface to only a few core parameters.
2. Simplified Script Mode and BYOC - The ModelTrainer natively supports script mode and removes the coupling to the SageMaker Training Toolkit for running a job in Script mode. By removing this runtime dependency, users can bring their own image to launch a training job without a needing to adapt it for script mode on SageMaker.
3. Simplified Distributed Training - The ModelTrainer provides enhanced flexibility for users to specify custom commands and distributed training configurations by specify the exect commands to execut in thier container using the command parameter in the SourceCode class or by leveraging a distributed training configuration class like Torchrun()

Install SageMaker PySDK

ModelTrainer - Basic Exuction

This case example shows a minimal setup for a ModelTrainer. A user need only to provide a desired training image and the commands they wish to execute in the container using the SourceCode class.

ModelTrainer - Script Mode

This examples show cases an abstracted setup for script mode where a user can provide their training image and a SourceCode config with path to their source_dir, enty_script, and any additional requirements to install in the training container for their job.

ModelTrainer - Local Container Mode

This example show cases how a user can leverage the LOCAL_CONTAINER mode to run their training job in their local enviornment as docker containers for local experimentation and testing.

Distributed Training

In this section, we will walk through how the ModelTrainer can be used for more complex Distributed Training jobs.

Setup Variables

Load Data Set

Prepare Dataset

Save Dataset

ModelTrainer - Distributed Training - Explicit Commands

This example shows how a user could perform a more complex setup for DistributedTraining using torchrun directly through the command parameter in the SourceCode class.

ModelTrainer - Distributed Training - Abstraction

This examples shows how a user could perform distributed training using an abstracted approach provided via the Torchrun distributed training configuration class.

ModelTrainer - SageMaker HyperPod Recipes

This example showcases how a user could leverage the pre-defined recipe fine-tuning/deepseek/hf_deepseek_r1_distilled_llama_70b_seq8k_gpu_fine_tuning for fine-tuning a DeepSeek R1 Model using synthetic data.

For More Available Recipes Visit - https://github.com/aws/sagemaker-hyperpod-recipes/