Acelere su ciclo de vida de aprendizaje automático con el nuevo y mejorado SDK de Python de Amazon SageMaker – Parte 1: ModelTrainer

Amazon SageMaker ha rediseñado su SDK de Python para proporcionar una interfaz unificada orientada a objetos que facilita la interacción con los servicios de SageMaker. El nuevo SDK está diseñado teniendo en mente una experiencia de usuario por niveles, donde el nuevo SDK de nivel inferior (Núcleo de SageMaker) proporciona acceso a toda la gama de funciones y configuraciones de SageMaker, lo que permite una mayor flexibilidad y control para los ingenieros de ML. La capa abstracta de nivel superior está diseñada para científicos de datos con experiencia limitada en AWS y ofrece una interfaz simplificada que oculta detalles complejos de la infraestructura.

En esta serie de dos partes, presentamos la capa abstracta del SDK de Python de SageMaker que le permite entrenar e implementar modelos de aprendizaje automático (ML) mediante el uso del nuevo ModelTrainer y las clases mejoradas de ModelBuilder.

En esta publicación, nos centramos en la clase ModelTrainer para simplificar la experiencia de formación. La clase ModelTrainer proporciona mejoras significativas con respecto a la actual Estimador clase, que se analizan en detalle en esta publicación. Le mostramos cómo utilizar la clase ModelTrainer para entrenar sus modelos de ML, lo que incluye la ejecución de entrenamiento distribuido utilizando un script o contenedor personalizado. En parte 2le mostramos cómo crear un modelo e implementarlo en un punto final de SageMaker utilizando la clase ModelBuilder mejorada.

Beneficios de la clase ModelTrainer

La nueva clase ModelTrainer ha sido diseñada para abordar los desafíos de usabilidad asociados con la clase Estimator. En el futuro, ModelTrainer será el enfoque preferido para el entrenamiento de modelos, aportando mejoras significativas que mejoran enormemente la experiencia del usuario. Esta evolución marca un paso hacia el logro de la mejor experiencia de desarrollador para la capacitación de modelos. Los siguientes son los beneficios clave:

• Intuición mejorada – La clase ModelTrainer reduce la complejidad al consolidar las configuraciones en solo unos pocos parámetros principales. Esta optimización minimiza la sobrecarga cognitiva, lo que permite a los usuarios centrarse en el entrenamiento del modelo en lugar de en las complejidades de la configuración. Además, emplea clases de configuración intuitivas para interacciones sencillas con la plataforma.
• Modo script simplificado y BYOC – La transición del desarrollo local a la capacitación en la nube ahora es perfecta. ModelTrainer asigna automáticamente el código fuente, las rutas de datos y las especificaciones de parámetros al entorno de ejecución remota, eliminando la necesidad de protocolos de enlace especiales o procesos de configuración complejos.
• Entrenamiento distribuido simplificado – La clase ModelTrainer proporciona flexibilidad mejorada para que los usuarios especifiquen comandos personalizados y estrategias de entrenamiento distribuidas, lo que le permite proporcionar directamente el comando exacto que desea ejecutar en su contenedor a través del command parámetro en el SourceCode Este enfoque desacopla las estrategias de capacitación distribuida del conjunto de herramientas de capacitación y los estimadores específicos del marco.
• Contratos de hiperparámetros mejorados – La clase ModelTrainer pasa los hiperparámetros del trabajo de entrenamiento como una única variable de entorno, lo que le permite cargar los hiperparámetros usando una única SM_HPSvariable.

Para explicar con más detalle cada uno de estos beneficios, lo demostramos con ejemplos en las siguientes secciones y, finalmente, le mostramos cómo configurar y ejecutar capacitación distribuida para el modelo Meta Llama 3.1 8B utilizando el nuevo ModelTrainer clase.

Lanzar un trabajo de entrenamiento usando la clase ModelTrainer

La clase ModelTrainer simplifica la experiencia al permitirle personalizar el trabajo de capacitación, lo que incluye proporcionar un script personalizado, proporcionar directamente un comando para ejecutar el trabajo de capacitación, admitir el modo local y mucho más. Sin embargo, puede iniciar un trabajo de capacitación de SageMaker en modo script proporcionando parámetros mínimos: SourceCode y el URI de la imagen de entrenamiento.

El siguiente ejemplo ilustra cómo puede iniciar un trabajo de entrenamiento con su propio script personalizado proporcionando solo el script y el URI de la imagen de entrenamiento (en este caso, PyTorch) y un archivo de requisitos opcional. El SDK establece automáticamente parámetros adicionales, como el tipo y el tamaño de la instancia, para preestablecer los valores predeterminados, y parámetros como el Gestión de acceso e identidad de AWS (IAM) y la sesión de SageMaker se detectan automáticamente a partir de la sesión actual y las credenciales del usuario. Los administradores y usuarios también pueden sobrescribir los valores predeterminados utilizando el Valores predeterminados del SDK archivo de configuración. Para obtener una lista detallada de valores preestablecidos, consulte la Documentación del SDK.

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

# image URI for the training job
pytorch_image = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310"
# you can find all available images here
# https://docs.aws.amazon.com/sagemaker/latest/dg-ecr-paths/sagemaker-algo-docker-registry-paths.html

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

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

# pass the input data
input_data = InputData(
    channel_name="train",
    data_source=training_input_path,  #s3 path where training data is stored
)

# start the training job
model_trainer.train(input_data_config=[input_data], wait=False)

Con configuraciones diseñadas específicamente, ahora puede reutilizar estos objetos para crear múltiples trabajos de entrenamiento con diferentes hiperparámetros, por ejemplo, sin tener que redefinir todos los parámetros.

Ejecute el trabajo localmente para experimentar.

Para ejecutar el trabajo de entrenamiento anterior localmente, simplemente puede configurar el training_mode parámetro como se muestra en el siguiente código:

from sagemaker.modules.train.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.train()

El trabajo de capacitación se realiza de forma remota porque training_mode está configurado para Mode.LOCAL_CONTAINER. Si no se configura explícitamente, ModelTrainer ejecuta un trabajo de entrenamiento remoto de SageMaker de forma predeterminada. Este comportamiento también se puede aplicar cambiando el valor a Mode.SAGEMAKER_TRAINING_JOB. Para obtener una lista completa de las configuraciones disponibles, incluidas la informática y las redes, consulte la Documentación del SDK.

Lea los hiperparámetros en su script personalizado

ModelTrainer admite varias formas de leer los hiperparámetros que se pasan a un trabajo de entrenamiento. Además del soporte existente para leer los hiperparámetros como argumentos de línea de comando en su script personalizado, ModelTrainer también admite la lectura de los hiperparámetros como variables de entorno individuales, con el prefijo SM_HPS_<hyperparameter-key>o como un diccionario de variables de entorno único, SM_HPS.

Supongamos que se pasan los siguientes hiperparámetros al trabajo de entrenamiento:

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

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

Tienes las siguientes opciones:

• Opción 1 – Cargue los hiperparámetros en un único diccionario JSON usando el SM_HPS variable de entorno en su script personalizado:

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

• Opción 2 – Leer los hiperparámetros como variables de entorno individuales, con el prefijo SM_HP como se muestra en el siguiente código (debe especificar explícitamente el tipo de entrada correcto para estas variables):

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

• Opción 3 – Lea los hiperparámetros como argumentos de AWS CLI utilizando parse.args:

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

Ejecute trabajos de capacitación distribuidos

SageMaker admite la capacitación distribuida para respaldar la capacitación para tareas de aprendizaje profundo, como el procesamiento del lenguaje natural y la visión por computadora, para ejecutar datos seguros y escalables en paralelo y modelar trabajos paralelos. Esto generalmente se logra proporcionando el conjunto correcto de parámetros cuando se utiliza un Estimador. Por ejemplo, para utilizar torchrundefinirías el distribution parámetro en el Estimador PyTorch y configúrelo en "torch_distributed": {"enabled": True}.

La clase ModelTrainer proporciona flexibilidad mejorada para que los usuarios especifiquen comandos personalizados directamente a través del command parámetro en el SourceCode clase y apoyos torchrun, torchrun smpy las estrategias del IPM. Esta capacidad es particularmente útil cuando necesita iniciar un trabajo con un comando de inicio personalizado que no es compatible con el kit de herramientas de capacitación.

En el siguiente ejemplo, mostramos cómo ajustar la última versión MetaLlama 3.1 8B modelo usando el script de inicio predeterminado usando Torchrun en un conjunto de datos personalizado que se preprocesa y guarda en un Servicio de almacenamiento simple de Amazon (Amazon S3) ubicación:

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

# provide  image URI - update the URI if you're in a different region
pytorch_image = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.2.0-gpu-py310"

# Define the source code configuration for the distributed training job
source_code = SourceCode(
    source_dir="distributed-training-scripts",    
    requirements="requirements.txt",  
    entry_point="fine_tune.py",
)

torchrun = Torchrun()

hyperparameters = {
    ...
}

# Compute configuration for the training 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 specified configurations
model_trainer = ModelTrainer(
    training_image=pytorch_image,  
    source_code=source_code,
    compute=compute,
    distributed_runner=torchrun,
    hyperparameters=hyperparameters,
)

# pass the input data
input_data = InputData(
    channel_name="dataset",
    data_source="s3://your-bucket/your-prefix",  # this is the s3 path where processed data is stored
)

# Start the training job
model_trainer.train(input_data_config=[input_data], wait=False)

Si quisieras personalizar tu torchrun script de inicio, también puede proporcionar directamente los comandos usando el command parámetro:

# Define the source code configuration for the distributed training job
source_code = SourceCode(
    source_dir="distributed-training-scripts",    
    requirements="requirements.txt",    
    # Custom command for distributed training 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 specified configurations
model_trainer = ModelTrainer(
    training_image=pytorch_image,  
    source_code=source_code,
    compute=compute,
)

# Start the training job
model_trainer.train(..)

Para obtener más ejemplos y flujos de trabajo de aprendizaje automático de un extremo a otro utilizando SageMaker ModelTrainer, consulte la repositorio de GitHub.

Conclusión

La clase SageMaker ModelTrainer recientemente lanzada simplifica la experiencia del usuario al reducir la cantidad de parámetros, introducir configuraciones intuitivas y admitir configuraciones complejas, como traer su propio contenedor y ejecutar capacitación distribuida. Los científicos de datos también pueden realizar una transición sin problemas de la capacitación local a la capacitación remota y a la capacitación en múltiples nodos utilizando ModelTrainer.

Le recomendamos que pruebe la clase ModelTrainer consultando el Documentación del SDK y cuadernos de muestra sobre el repositorio de GitHub. La clase ModelTrainer está disponible desde SageMaker SDK v2.x en adelante, sin costo adicional. En parte 2 En esta serie, le mostramos cómo crear un modelo e implementarlo en un punto final de SageMaker utilizando la clase ModelBuilder mejorada.

Acerca de los autores

Durga Sury es arquitecto senior de soluciones en el equipo de Amazon SageMaker. Durante los últimos cinco años, ha trabajado con varios clientes empresariales para configurar una plataforma AI/ML segura y escalable construida en SageMaker.

Shweta Singh es gerente senior de productos en el equipo de la plataforma Amazon SageMaker Machine Learning (ML) en AWS, y lidera el SDK de SageMaker Python. Ha trabajado en varios puestos de productos en Amazon durante más de 5 años. Tiene una licenciatura en Ingeniería Informática y una Maestría en Ciencias en Ingeniería Financiera, ambas de la Universidad de Nueva York.