Perfilado de TorchNeuron (Parte 1): Perfilado a nivel de sistema en AWS Trainium.

Muchos profesionales de ML tienen cierta intuición sobre cómo hacer que los modelos se ejecuten rápido, pero menos saben cómo analizar por qué un modelo es lento. En la Parte 1 de esta serie, aprenderá cómo analizar sistemáticamente el rendimiento de inferencia e identificar dónde se dedica el tiempo.

Esta publicación comienza desde cero y asume que no hay experiencia previa en perfilado. El recorrido usa AWS Trainium, un chip de IA diseñado a propósito para cargas de trabajo de entrenamiento e inferencia de aprendizaje profundo, y TorchNeuron, su backend nativo de PyTorch de código abierto. TorchNeuron admite flujos de trabajo estándar de PyTorch en Trainium, incluida la ejecución eager, torch.compile, y las API estándar de PyTorch, con cambios mínimos de código. Además, TorchNeuron se integra con Neuron Kernel Interface (NKI) para habilitar operadores personalizados y optimización de rendimiento para el hardware Trainium. Al final, usted debería saber:

• Cómo configurar torch.profiler con el backend de Neuron y capturar un trace limpio.
• Cómo leer la tabla key_averages() para identificar qué operaciones consumen tiempo del dispositivo Neuron y comparar su Self CPU time.
• Cómo seguir una operación a través de Neuron Explorer, desde el despacho del framework pasando por la programación del runtime hasta la ejecución en hardware en un NeuronCore.
• Cómo cambian estas señales entre tres implementaciones de la misma operación.

Neuron como backend de PyTorch

Antes de perfilar una carga de trabajo, ayuda entender cómo Neuron se integra con PyTorch.

El backend de Neuron usa el punto de extensión PrivateUse1 de PyTorch. Este mecanismo permite que un acelerador se integre con el framework sin agregar una clave de despacho dedicada como cuda. El paquete torch_neuronx funciona con PyTorch upstream para registrar Neuron como un backend nativo para código PyTorch. Para obtener detalles sobre este mecanismo de extensión, consulte Facilitating New Backend Integration by PrivateUse1.

Configurar el profiler

Usted usa dos objetos para perfilar una carga de trabajo de Neuron:

• NeuronConfig controla qué flujos de datos de Neuron y qué NeuronCores capturar y dónde escribir la salida. Usted lo pasa a torch.profiler.profile como experimental_config.
• NeuronProfiler toma un NeuronConfig y proporciona un callback export_trace . El callback escribe el trace del framework, trace.json, junto con los traces del sistema de Neuron.

nc = NeuronConfig(
    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
    profile_output_dir=OUTPUT_DIR,
    capture_enabled_for_nc='0',
    correlate_device_time=True,
)
exporter = NeuronProfiler(nc)

El modes parámetro selecciona los flujos de datos que se capturan, como las instrucciones del dispositivo y los rastros de ejecución. El capture_enabled_for_nc parámetro selecciona los NeuronCores que se rastrean. Establecer correlate_device_time=True propaga el tiempo del dispositivo desde el rastro a la key_averages() tabla, donde puede comparar el tiempo de CPU y el tiempo del dispositivo Neuron. Para la referencia completa de parámetros y los ProfileMode valores disponibles, consulte la documentación de NeuronConfig Profile Modes.

Herramientas: Neuron Explorer

Neuron Explorer es un visor basado en navegador para el rastro del sistema, la línea de tiempo del dispositivo y el desglose por motor. Puede acceder a él mediante un túnel SSH o la extensión de Visual Studio Code. En esta publicación, usa Explorer para confirmar lo que key_averages() reporta numéricamente e inspeccionar lo que sucede entre despachos.

Después de la generación de perfiles, el directorio de salida contiene el rastro del sistema y, cuando captura datos del dispositivo, los archivos .neff y .ntff correspondientes. El Neuron Executable File Format (NEFF) contiene el ejecutable de Neuron compilado, mientras que el NTFF correspondiente contiene su rastro de instrucciones de hardware. Los archivos comparten un hash numérico que los identifica como un par. Para abrir el perfil en Neuron Explorer, cargue el directorio de salida completo de un perfil de sistema o combinado. Para un perfil solo del dispositivo, cargue individualmente los archivos .neff y .ntff correspondientes. Vea Step 5: Upload your profile para las instrucciones de la interfaz de usuario y de la línea de comandos.

Generar el perfil de RMSNorm

Para este primer análisis, genera el perfil de RMSNorm, una operación pequeña con una secuencia clara de operaciones aritméticas y de reducción.

Genera el perfil de tres implementaciones de la misma operación:

1. RMSNorm ingenuo con operaciones individuales de PyTorch: multiplicación, media, suma, raíz cuadrada recíproca y multiplicación.
2. torch.rms_norm(), la función integrada de PyTorch reducida a un único kernel fusionado de Neuron.
3. RMSNorm ingenuo con torch.compile, usando las mismas operaciones que la primera implementación.

Los scripts usan las mismas entradas, la misma configuración del profiler y el mismo cronograma, de modo que los resultados son directamente comparables.

Todos los scripts para reproducir este trabajo están disponibles en el torch-neuronx directorio de ejemplos.

Perfil 1: Operaciones ingenuas de PyTorch

Comience escribiendo RMSNorm con operaciones individuales de PyTorch. Cada línea se corresponde con una o más operaciones de ATen y puede crear tensores intermedios.

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t                                    # square
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # mean of squares
    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)
    return t * rms * weight                         # normalize + scale

Llame al profiler con dos argumentos requeridos para el backend de Neuron:

1. Incluya ProfilerActivity.PrivateUse1 para recopilar actividad del dispositivo Neuron.
2. Pase NeuronConfig como experimental_config para habilitar las características de Kineto específicas de Neuron.

El cronograma usa wait=2, warmup=1, y active=1 a lo largo de cuatro prof.step() llamadas. Esto registra una iteración de calentamiento sin incluir la sobrecarga de la configuración inicial.

Llame a torch.neuron.synchronize() dentro del bucle del profiler. Sin esta sincronización, el profiler puede exportar antes de que el dispositivo termine y dejar vacía la columna de tiempo de Neuron.

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
    for step in range(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

Una vez que finalice la creación de perfiles, inspeccione la tabla key_averages() . Ordene por self_privateuse1_time_total para colocar primero las operaciones con el mayor costo en el dispositivo. La tabla que se muestra aquí está filtrada a las operaciones de RMSNorm que reportan tiempo del dispositivo Neuron. La salida completa de key_averages() también incluye eventos del lado de la CPU, como asignaciones y transferencias de host a dispositivo.

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

La tabla de promedios clave muestra que tres llamadas a aten::mul representan la mayor parte del tiempo del dispositivo Neuron.

La tabla muestra qué operaciones de ATen se ejecutaron, con qué frecuencia se ejecutaron y cuánto tiempo pasaron en la CPU y en el dispositivo Neuron. aten::mul tiene el mayor tiempo total del dispositivo porque la función realiza tres multiplicaciones.

Self NEURON total: 2.355 ms

Cada fila se corresponde directamente con la función. aten::mul se ejecuta durante x * x, x * rms, y la multiplicación final por weight. aten::mean reduce la última dimensión. aten::add añade el valor de epsilon y desencadena una copia de host a dispositivo porque el escalar comienza en la CPU. aten::rsqrt calcula la raíz cuadrada recíproca.

Neuron Explorer sigue una operación aten::mul desde el despacho en el framework, pasando por la programación en tiempo de ejecución, hasta la ejecución en los lanes del hardware Neuron.

Leer el trace en Explorer

En Explorer, el tiempo fluye de izquierda a derecha y el ciclo de vida de los eventos fluye de arriba abajo. Cada operación ATen sigue el mismo patrón: creación en el lane del framework, programación en el lane del runtime como nrta_execute_schedule, y ejecución en el dispositivo en el lane del hardware como nc_exec_running.

La vista de trace del sistema organiza los eventos en lanes. En el modo de agrupación Default, las etiquetas siguen el patrón instancia/PID/subsistema/identificador:

En modo eager, cada operación se ejecuta de forma independiente. Use el visor de cadena de dependencias para seguir aten::mul.Tensor desde la capa del framework, pasando por la programación del runtime, hasta la ejecución en el dispositivo.

El visor de cadena de dependencias resalta el ciclo de vida de una operación a través de los lanes del framework, el runtime y el hardware.

Perfil 2: torch.rms_norm

Reemplace las operaciones individuales por la función rms_norm integrada de PyTorch. Neuron reduce la operación a un único kernel fusionado:

def rms_norm_fn(t):
    return torch.rms_norm(t, [512], weight)

La RMSNorm fusionada aparece como una sola operación de dispositivo en la tabla de promedios clave.

Self NEURON total: 403.3 μs

El RMSNorm fusionado sigue una única ruta de framework a hardware en lugar de despachar cada operación aritmética por separado.

Comparar la ruta de lowering

En la implementación ingenua, cada operación cruza la frontera de despacho de forma independiente y sigue su propia ruta de framework a runtime a hardware. En la implementación fusionada, una sola llamada de Python cubre el cálculo completo. La traza contiene un despacho y un evento de hardware, sin huecos entre las operaciones componentes.

Perfil 3: PyTorch ingenuo con torch.compile

Para el ejemplo final, compile la función RMSNorm ingenua original con torch.compile.

Con @torch.compile(backend="neuron"), TorchDynamo captura el bytecode de Python como un grafo FX y AOTAutograd lo reduce a operaciones ATen. El backend de Neuron recibe el grafo completo en lugar de despachar cada operación por separado. Fusiona las operaciones, optimiza el acceso a memoria y compila el grafo en un único artefacto NEFF. El artefacto se ejecuta como un solo despacho de hardware.

@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
    x_sq = t * t
    mean_sq = x_sq.mean(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

La función compilada aparece como una Torch-Compiled Region en la tabla de promedios clave.

Self NEURON total: 375.5 μs

El grafo compilado sigue una única ruta a través de la planificación en runtime y la ejecución en hardware.

La tabla muestra una sola ejecución en el dispositivo, como en el Perfil 2, pero el código toma una ruta diferente. En lugar de llamar a una operación fusionada construida a propósito, el compilador descubre la fusión a partir de la función Python original. En esta ejecución, el tiempo de Neuron es comparable a torch.rms_norm, con 375 μs frente a 403 μs.

Resumen lado a lado

Ambos torch.rms_norm y torch.compile fusionan RMSNorm en un solo despacho de kernel. En estas ejecuciones, cualquiera de los dos enfoques toma aproximadamente un sexto del tiempo de dispositivo de la implementación ingenua. La fusión elimina los huecos entre ejecuciones de hardware y evita materializar valores intermedios como x_sq y mean_sq entre operaciones despachadas por separado.

Conclusión

key_averages() muestra qué operaciones consumen tiempo de dispositivo, mientras que su salida completa también revela la configuración y las transferencias del lado de la CPU. Explorer proporciona la traza visual correspondiente. Puedes seguir una operación desde el despacho del framework pasando por la programación en tiempo de ejecución hasta la ejecución en hardware y ver los huecos que elimina la fusión.

Puedes usar el mismo flujo de trabajo en modelos más grandes: encuentra un punto crítico en la tabla, verifícalo en la traza y luego determina si una operación fusionada o torch.compile puede reducir su costo.

En la Parte 2, pasaremos al perfil a nivel de dispositivo. Abres el NEFF, inspeccionas el desglose del TensorEngine, el VectorEngine y el acceso directo a memoria (DMA), y colocas el kernel en el roofline para determinar si está limitado por cómputo o por memoria.

Un agradecimiento especial a Aritra Roy Gosthipaty por la inspiración para comenzar esta serie, y a mis colegas Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia y al resto por la revisión!