Preparar y ejecutar un trabajo de entrenamiento con Generador de perfiles de SageMaker

La configuración para ejecutar una tarea de entrenamiento con Generador de perfiles de SageMaker consta de dos pasos: adaptar el script de entrenamiento y configurar el lanzador de tareas de entrenamiento de SageMaker.

Paso 1: Adaptar el script de entrenamiento con los módulos de Python del Generador de perfiles de SageMaker

Para empezar a capturar las ejecuciones del kernel en las GPU mientras se ejecuta el trabajo de entrenamiento, modifique el script de entrenamiento mediante los módulos de Python del Generador de perfiles de SageMaker. Importe la biblioteca y añada los métodos start_profiling() y stop_profiling() para definir el principio y el final de la creación de perfiles. También puede utilizar anotaciones personalizadas opcionales para añadir marcadores en el script de entrenamiento a fin de visualizar las actividades del hardware durante determinadas operaciones de cada paso.

Ten en cuenta que los anotadores extraen las operaciones de las GPU. Para las operaciones de creación de un perfil en las CPU, no es necesario añadir ninguna anotación adicional. La creación de perfiles de CPU también se activa cuando especificas la configuración de creación de perfiles, con la que practicará en Paso 2: creación de un estimador de marco de SageMaker AI y activación del Generador de perfiles de SageMaker.

try:
    import smprof 
except ImportError:
    # backward-compatability for TF 2.11 and PT 1.13.1 images
    import smppy as smprof

Método 1. Utilice el administrador de contexto smprof.annotate para anotar todas las funciones

El administrador de contextos de smprof.annotate() puede encapsular todas las funciones. Se recomienda utilizar este contenedor si desea crear perfiles por funciones en lugar de por líneas de código. El siguiente script de ejemplo muestra cómo implementar el administrador de contexto para encapsular el ciclo de entrenamiento y todas las funciones en cada iteración.

import smprof

SMProf = smprof.SMProfiler.instance()
config = smprof.Config()
config.profiler = {
    "EnableCuda": "1",
}
SMProf.configure(config)
SMProf.start_profiling()

for epoch in range(args.epochs):
    if world_size > 1:
        sampler.set_epoch(epoch)
    tstart = time.perf_counter()
    for i, data in enumerate(trainloader, 0):
        with smprof.annotate("step_"+str(i)):
            inputs, labels = data
            inputs = inputs.to("cuda", non_blocking=True)
            labels = labels.to("cuda", non_blocking=True)
    
            optimizer.zero_grad()
    
            with smprof.annotate("Forward"):
                outputs = net(inputs)
            with smprof.annotate("Loss"):
                loss = criterion(outputs, labels)
            with smprof.annotate("Backward"):
                loss.backward()
            with smprof.annotate("Optimizer"):
                optimizer.step()

SMProf.stop_profiling()

Método 1. Utilice smprof.annotation_begin() y smprof.annotation_end() anote una línea de código específica en funciones

También puede definir anotaciones para perfilar líneas de código específicas. Puede establecer el punto de inicio y el punto de conexión exactos de la creación de perfiles a nivel de líneas de código individuales, no mediante las funciones. Por ejemplo, en el siguiente script, el step_annotator se define el al principio de cada iteración y termina al final de la iteración. Mientras tanto, se definen otros anotadores detallados para cada operación y encapsulan las operaciones objetivo a lo largo de cada iteración.

import smprof

SMProf = smprof.SMProfiler.instance()
config = smprof.Config()
config.profiler = {
    "EnableCuda": "1",
}
SMProf.configure(config)
SMProf.start_profiling()

for epoch in range(args.epochs):
    if world_size > 1:
        sampler.set_epoch(epoch)
    tstart = time.perf_counter()
    for i, data in enumerate(trainloader, 0):
        step_annotator = smprof.annotation_begin("step_" + str(i))

        inputs, labels = data
        inputs = inputs.to("cuda", non_blocking=True)
        labels = labels.to("cuda", non_blocking=True)
        optimizer.zero_grad()

        forward_annotator = smprof.annotation_begin("Forward")
        outputs = net(inputs)
        smprof.annotation_end(forward_annotator)

        loss_annotator = smprof.annotation_begin("Loss")
        loss = criterion(outputs, labels)
        smprof.annotation_end(loss_annotator)

        backward_annotator = smprof.annotation_begin("Backward")
        loss.backward()
        smprof.annotation_end(backward_annotator)

        optimizer_annotator = smprof.annotation_begin("Optimizer")
        optimizer.step()
        smprof.annotation_end(optimizer_annotator)

        smprof.annotation_end(step_annotator)

SMProf.stop_profiling()

Tras anotar y configurar los módulos de iniciación del generador de perfiles, guarde el script para enviarlo mediante un lanzador de tareas de entrenamiento de SageMaker en el siguiente paso 2. En el ejemplo del lanzador se presupone que el script de entrenamiento se llama train_with_profiler_demo.py.

Paso 2: creación de un estimador de marco de SageMaker AI y activación del Generador de perfiles de SageMaker

En el siguiente procedimiento, se muestra cómo preparar un estimador de marco de SageMaker AI para su entrenamiento mediante el SageMaker Python SDK.

Mientras ejecuta el trabajo de entrenamiento o una vez finalizado el trabajo, puede pasar al tema siguiente en Abra la aplicación de interfaz de usuario del Generador de perfiles de SageMaker y empezar a explorar y visualizar los perfiles guardados.

Si desea acceder directamente a los datos del perfil guardados en el bucket de Amazon S3, utilice el siguiente script para recuperar el URI de S3.

import os
# This is an ad-hoc function to get the S3 URI
# to where the profile output data is saved
def get_detailed_profiler_output_uri(estimator):
    config_name = None
    for processing in estimator.profiler_rule_configs:
        params = processing.get("RuleParameters", dict())
        rule = config_name = params.get("rule_to_invoke", "")
        if rule == "DetailedProfilerProcessing":
            config_name = processing.get("RuleConfigurationName")
            break
    return os.path.join(
        estimator.output_path, 
        estimator.latest_training_job.name, 
        "rule-output",
        config_name,
    )

print(
    f"Profiler output S3 bucket: ", 
    get_detailed_profiler_output_uri(estimator)
)

(Opcional) Instale el paquete de Python del Generador de perfiles de SageMaker

Para usar el Generador de perfiles de SageMaker en PyTorch o imágenes de marco de TensorFlow que no aparecen en Imágenes de marcos de SageMaker AI preinstaladas con el Generador de perfiles de SageMaker, o en su propio contenedor de Docker personalizado para entrenamiento, puede instalar el Generador de perfiles de SageMaker mediante uno de los Archivos binarios del paquete de Python del generador de perfiles de SageMaker.

Opción 1: instalación del paquete del Generador de perfiles de SageMaker al iniciar un trabajo de entrenamiento

Si desea utilizar el Generador de perfiles de SageMaker para trabajos de entrenamiento con imágenes de PyTorch o TensorFlow que no aparecen enumeradas en Imágenes de marcos de SageMaker AI preinstaladas con el Generador de perfiles de SageMaker, cree un archivo requirements.txt y ubíquelo en la ruta que especifique para el parámetro source_dir del estimador de marco de SageMaker AI en el Paso 2. Para obtener más información sobre la configuración de un archivo requirements.txt en general, consulte Using third-party libraries en la documentación del SageMaker Python SDK. En el archivo requirements.txt, añada una de las rutas de bucket de S3 para los Archivos binarios del paquete de Python del generador de perfiles de SageMaker.

# requirements.txt
https://smppy.s3.amazonaws.com/tensorflow/cu112/smprof-0.3.332-cp39-cp39-linux_x86_64.whl

Opción 2: instalación del paquete del Generador de perfiles de SageMaker en sus contenedores de Docker personalizados

Si usa un contenedor de Docker personalizado para entrenamiento, agregue uno de los Archivos binarios del paquete de Python del generador de perfiles de SageMaker a su Dockerfile.

# Install the smprof package version compatible with your CUDA version
RUN pip install https://smppy.s3.amazonaws.com/tensorflow/cu112/smprof-0.3.332-cp39-cp39-linux_x86_64.whl

Para obtener orientación sobre cómo ejecutar un contenedor de Docker personalizado para el entrenamiento en SageMaker AI en general, consulte Adapting your own training container.