Speed up your ML lifecycle utilizing the brand new and improved Amazon SageMaker Python SDK – Half 1: ModelTrainer

Amazon SageMaker has redesigned its Python SDK to supply a unified object-oriented interface that makes it easy to work together with SageMaker companies. The brand new SDK is designed with a tiered consumer expertise in thoughts, the place the brand new lower-level SDK (SageMaker Core) offers entry to full breadth of SageMaker options and configurations, permitting for larger flexibility and management for ML engineers. The upper-level abstracted layer is designed for knowledge scientists with restricted AWS experience, providing a simplified interface that hides complicated infrastructure particulars.

On this two-part collection, we introduce the abstracted layer of the SageMaker Python SDK that lets you prepare and deploy machine studying (ML) fashions by utilizing the brand new ModelTrainer and the improved ModelBuilder courses.

On this publish, we deal with the ModelTrainer class for simplifying the coaching expertise. The ModelTrainer class offers vital enhancements over the present Estimator class, that are mentioned intimately on this publish. We present you use the ModelTrainer class to coach your ML fashions, which incorporates executing distributed coaching utilizing a customized script or container. In Half 2, we present you construct a mannequin and deploy to a SageMaker endpoint utilizing the improved ModelBuilder class.

Advantages of the ModelTrainer class

The brand new ModelTrainer class has been designed to handle usability challenges related to Estimator class. Transferring ahead, ModelTrainer would be the most well-liked method for mannequin coaching, bringing vital enhancements that significantly enhance the consumer expertise. This evolution marks a step in direction of attaining a best-in-class developer expertise for mannequin coaching. The next are the important thing advantages:

• Improved intuitiveness – The ModelTrainer class reduces complexity by consolidating configurations into simply few core parameters. This streamlining minimizes cognitive overload, permitting customers to deal with mannequin coaching moderately than configuration intricacies. Moreover, it employs intuitive config courses for easy platform interactions.
• Simplified script mode and BYOC – Transitioning from native improvement to cloud coaching is now seamless. The ModelTrainer mechanically maps supply code, knowledge paths, and parameter specs to the distant execution surroundings, eliminating the necessity for particular handshakes or complicated setup processes.
• Simplified distributed coaching – The ModelTrainer class offers enhanced flexibility for customers to specify customized instructions and distributed coaching methods, permitting you to straight present the precise command you need to run in your container by way of the command parameter within the SourceCode This method decouples distributed coaching methods from the coaching toolkit and framework-specific estimators.
• Improved hyperparameter contracts – The ModelTrainer class passes the coaching job’s hyperparameters as a single surroundings variable, permitting the you to load the hyperparameters utilizing a single SM_HPSvariable.

To additional clarify every of those advantages, we reveal with examples within the following sections, and eventually present you arrange and run distributed coaching for the Meta Llama 3.1 8B mannequin utilizing the brand new ModelTrainer class.

Launch a coaching job utilizing the ModelTrainer class

The ModelTrainer class simplifies the expertise by letting you customise the coaching job, together with offering a customized script, straight offering a command to run the coaching job, supporting native mode, and rather more. Nonetheless, you may spin up a SageMaker coaching job in script mode by offering minimal parameters—the SourceCode and the coaching picture URI.

The next instance illustrates how one can launch a coaching job with your personal customized script by offering simply the script and the coaching picture URI (on this case, PyTorch), and an non-compulsory necessities file. Extra parameters such because the occasion sort and occasion measurement are mechanically set by the SDK to preset defaults, and parameters such because the AWS Identification and Entry Administration (IAM) function and SageMaker session are mechanically detected from the present session and consumer’s credentials. Admins and customers can even overwrite the defaults utilizing the SDK defaults configuration file. For the detailed listing of pre-set values, seek advice from the SDK documentation.

from sagemaker.modules.prepare import ModelTrainer
from sagemaker.modules.configs import SourceCode, InputData

# picture URI for the coaching job
pytorch_image = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310"
# you could find all obtainable photos right here
# https://docs.aws.amazon.com/sagemaker/newest/dg-ecr-paths/sagemaker-algo-docker-registry-paths.html

# outline the script to be run
source_code = SourceCode(
    source_dir="basic-script-mode",
    necessities="necessities.txt",
    entry_script="custom_script.py",
)

# outline the ModelTrainer
model_trainer = ModelTrainer(
    training_image=pytorch_image,
    source_code=source_code,
    base_job_name="script-mode",
)

# cross the enter knowledge
input_data = InputData(
    channel_name="prepare",
    data_source=training_input_path,  #s3 path the place coaching knowledge is saved
)

# begin the coaching job
model_trainer.prepare(input_data_config=[input_data], wait=False)

With purpose-built configurations, now you can reuse these objects to create a number of coaching jobs with totally different hyperparameters, for instance, with out having to re-define all of the parameters.

Run the job regionally for experimentation

To run the previous coaching job regionally, you may merely set the training_mode parameter as proven within the following code:

from sagemaker.modules.prepare.model_trainer import Mode

...
model_trainer = ModelTrainer(
    training_image=pytorch_image,
    source_code=source_code,
    base_job_name="script-mode-local",
    training_mode=Mode.LOCAL_CONTAINER,
)
model_trainer.prepare()

The coaching job runs remotely as a result of training_mode is ready to Mode.LOCAL_CONTAINER. If not explicitly set, the ModelTrainer runs a distant SageMaker coaching job by default. This habits can be enforced by altering the worth to Mode.SAGEMAKER_TRAINING_JOB. For a full listing of the obtainable configs, together with compute and networking, seek advice from the SDK documentation.

Learn hyperparameters in your customized script

The ModelTrainer helps a number of methods to learn the hyperparameters which might be handed to a coaching job. Along with the present assist to learn the hyperparameters as command line arguments in your customized script, ModelTrainer additionally helps studying the hyperparameters as particular person surroundings variables, prefixed with SM_HPS_, or as a single surroundings variable dictionary, SM_HPS.

Suppose the next hyperparameters are handed to the coaching job:

hyperparams = {
    "learning_rate": 1e-5,
    "epochs": 2,
}

model_trainer = ModelTrainer(
    ...
    hyperparameters=hyperparams,
    ...
)

You’ve the next choices:

• Choice 1 – Load the hyperparameters right into a single JSON dictionary utilizing the SM_HPS surroundings variable in your customized script:

def important():
    hyperparams = json.hundreds(os.environ["SM_HPS"])
    learning_rate = hyperparams.get("learning_rate")
    epochs = hyperparams.get("epochs", 1)
    ...

• Choice 2 – Learn the hyperparameters as particular person surroundings variables, prefixed by SM_HP as proven within the following code (you could explicitly specify the right enter sort for these variables):

def important():
    learning_rate = float(os.environ.get("SM_HP_LEARNING_RATE", 3e-5))
    epochs = int(os.environ.get("SM_HP_EPOCHS", 1)
    ...

• Choice 3 – Learn the hyperparameters as AWS CLI arguments utilizing parse.args:

def important():
    parser = argparse.ArgumentParser()
    parser.add_argument("--learning_rate", sort=float, default=3e-5)
    parser.add_argument("--epochs", sort=int, default=1)
    
    args = parse_args()
    
    learning_rate = args.learning_rate
    epochs = args.epochs

Run distributed coaching jobs

SageMaker helps distributed coaching to assist coaching for deep studying duties similar to pure language processing and pc imaginative and prescient, to run safe and scalable knowledge parallel and mannequin parallel jobs. That is often achieved by offering the precise set of parameters when utilizing an Estimator. For instance, to make use of torchrun, you’ll outline the distribution parameter within the PyTorch Estimator and set it to "torch_distributed": {"enabled": True}.

The ModelTrainer class offers enhanced flexibility for customers to specify customized instructions straight by way of the command parameter within the SourceCode class, and helps torchrun, torchrun smp, and the MPI methods. This functionality is especially helpful when you could launch a job with a customized launcher command that isn’t supported by the coaching toolkit.

Within the following instance, we present fine-tune the newest Meta Llama 3.1 8B mannequin utilizing the default launcher script utilizing Torchrun on a customized dataset that’s preprocessed and saved in an Amazon Easy Storage Service (Amazon S3) location:

from sagemaker.modules.prepare import ModelTrainer
from sagemaker.modules.distributed import Torchrun
from sagemaker.modules.configs import Compute, SourceCode, InputData

# present  picture URI - replace the URI if you happen to're in a unique area
pytorch_image = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.2.0-gpu-py310"

# Outline the supply code configuration for the distributed coaching job
source_code = SourceCode(
    source_dir="distributed-training-scripts",    
    necessities="necessities.txt",  
    entry_point="fine_tune.py",
)

torchrun = Torchrun()

hyperparameters = {
    ...
}

# Compute configuration for the coaching job
compute = Compute(
    instance_count=1,
    instance_type="ml.g5.12xlarge",
    volume_size_in_gb=96,
    keep_alive_period_in_seconds=3600,
)


# Initialize the ModelTrainer with the desired configurations
model_trainer = ModelTrainer(
    training_image=pytorch_image,  
    source_code=source_code,
    compute=compute,
    distributed_runner=torchrun,
    hyperparameters=hyperparameters,
)

# cross the enter knowledge
input_data = InputData(
    channel_name="dataset",
    data_source="s3://your-bucket/your-prefix",  # that is the s3 path the place processed knowledge is saved
)

# Begin the coaching job
model_trainer.prepare(input_data_config=[input_data], wait=False)

In case you needed to customise your torchrun launcher script, you may as well straight present the instructions utilizing the command parameter:

# Outline the supply code configuration for the distributed coaching job
source_code = SourceCode(
    source_dir="distributed-training-scripts",    
    necessities="necessities.txt",    
    # Customized command for distributed coaching launcher script
    command="torchrun --nnodes 1 
            --nproc_per_node 4 
            --master_addr algo-1 
            --master_port 7777 
            fine_tune_llama.py"
)


# Initialize the ModelTrainer with the desired configurations
model_trainer = ModelTrainer(
    training_image=pytorch_image,  
    source_code=source_code,
    compute=compute,
)

# Begin the coaching job
model_trainer.prepare(..)

For extra examples and end-to-end ML workflows utilizing the SageMaker ModelTrainer, seek advice from the GitHub repo.

Conclusion

The newly launched SageMaker ModelTrainer class simplifies the consumer expertise by decreasing the variety of parameters, introducing intuitive configurations, and supporting complicated setups like bringing your personal container and working distributed coaching. Knowledge scientists can even seamlessly transition from native coaching to distant coaching and coaching on a number of nodes utilizing the ModelTrainer.

We encourage you to check out the ModelTrainer class by referring to the SDK documentation and pattern notebooks on the GitHub repo. The ModelTrainer class is accessible from the SageMaker SDK v2.x onwards, at no extra cost. In Half 2 of this collection, we present you construct a mannequin and deploy to a SageMaker endpoint utilizing the improved ModelBuilder class.

In regards to the Authors

Durga Sury is a Senior Options Architect on the Amazon SageMaker staff. Over the previous 5 years, she has labored with a number of enterprise clients to arrange a safe, scalable AI/ML platform constructed on SageMaker.

Shweta Singh is a Senior Product Supervisor within the Amazon SageMaker Machine Studying (ML) platform staff at AWS, main SageMaker Python SDK. She has labored in a number of product roles in Amazon for over 5 years. She has a Bachelor of Science diploma in Laptop Engineering and a Masters of Science in Monetary Engineering, each from New York College.