Prepare e execute um trabalho de treinamento com o Profiler do SageMaker

A configuração para executar um trabalho de treinamento com o Profiler do SageMaker consiste em duas etapas: adaptar o script de treinamento e configurar o iniciador de trabalhos de treinamento do SageMaker.

Etapa 1: adapte o script de treinamento usando os módulos Python do Profiler do SageMaker

Para começar a capturar as execuções do kernel em GPUs enquanto o trabalho de treinamento está em execução, modifique o script de treinamento usando os módulos Python do Profiler do SageMaker. Importe a biblioteca e adicione os métodos start_profiling() e stop_profiling() para definir o início e o fim da criação de perfil. Você também pode usar anotações personalizadas opcionais para adicionar marcadores no script de treinamento para visualizar as atividades do hardware durante operações específicas em cada etapa.

Observe que os anotadores extraem operações das GPUs. Para operações de criação de perfil em CPUs, você não precisa adicionar nenhuma anotação adicional. A criação de perfil de CPU também é ativada quando você especifica a configuração de criação de perfil, a qual você praticará em Etapa 2: criar um estimador de framework do SageMaker AI e ativar o SageMaker Profiler.

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

Abordagem 1. Use o gerenciador de contexto smprof.annotate para anotar funções completas

Você pode agrupar funções completas com o gerenciador de contexto smprof.annotate(). Esse wrapper é recomendado se você quiser criar perfis por funções em vez de linhas de código. O script de exemplo a seguir mostra como implementar o gerenciador de contexto para encapsular o ciclo de treinamento e as funções completas em cada iteração.

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()

Abordagem 2. Use smprof.annotation_begin() e smprof.annotation_end() para anotar uma linha de código específica nas funções

Você também pode definir anotações para traçar o perfil de linhas de código específicas. Você pode definir o ponto inicial e final exatos da criação de perfil no nível das linhas de código individuais, não pelas funções. Por exemplo, no script a seguir, o step_annotator é definido no início de cada iteração e termina no final da iteração. Enquanto isso, outros anotadores detalhados para cada operação são definidos e envolvem as operações de destino em cada iteração.

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()

Depois de anotar e configurar os módulos de iniciação do Profiler, salve o script para enviar usando um iniciador de tarefas de treinamento do SageMaker na próxima Etapa 2. O iniciador de exemplo presume que o script de treinamento seja chamado de train_with_profiler_demo.py.

Etapa 2: criar um estimador de framework do SageMaker AI e ativar o SageMaker Profiler

O procedimento a seguir mostra como preparar um estimador de framework do SageMaker AI para treinamento usando o SageMaker AI Python SDK.

Durante a execução do trabalho de treinamento ou após a conclusão do trabalho, você pode ir para o próximo tópico em Abra a aplicação Profiler UI do SageMaker e começar a explorar e visualizar os perfis salvos.

Se você quiser acessar diretamente os dados do perfil salvos no bucket do Amazon S3, use o script a seguir para recuperar o URI do 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) Instalar o pacote SageMaker Profiler Python

Para usar o SageMaker Profiler em imagens da estrutura PyTorch ou TensorFlow não listadas em Imagens do framework do SageMaker AI pré-instaladas com o SageMaker Profiler, ou em seu próprio contêiner do Docker personalizado para treinamento, você pode instalar o SageMaker Profiler usando um dos Arquivos binários do pacote SageMaker Profiler Python.

Opção 1: instalar o pacote SageMaker Profiler ao iniciar um trabalho de treinamento

Se você quiser usar o SageMaker Profiler para treinar tarefas usando imagens do PyTorch ou do TensorFlow não listadas em Imagens do framework do SageMaker AI pré-instaladas com o SageMaker Profiler, crie um arquivo requirements.txt e localize-o no caminho especificado para o parâmetro source_dir do estimador de framework do SageMaker AI na Etapa 2. Para obter mais informações sobre como configurar um arquivo requirements.txt em geral, consulte Uso de bibliotecas de terceiros na documentação do SDK do SageMaker Python. No arquivo requirements.txt, adicione um dos caminhos de bucket do S3 para o Arquivos binários do pacote SageMaker Profiler Python.

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

Opção 2: instalar o pacote SageMaker Profiler em seus contêineres do Docker personalizados

Se você usa um contêiner do Docker personalizado para treinamento, adicione um dos Arquivos binários do pacote SageMaker Profiler Python ao seu 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 obter orientações sobre como executar um contêiner do Docker personalizado para treinamento no SageMaker AI em geral, consulte Adaptar seu próprio contêiner de treinamento.