Profiling TorchNeuron (Parte 1): System-level profiling on AWS Trainium.

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

Esta publicación comienza desde cero y no asume experiencia previa en profiling. La guía usa AWS Trainium, un chip de IA diseñado específicamente para cargas de trabajo de entrenamiento e inferencia de deep learning, y TorchNeuron, su backend nativo de PyTorch de código abierto. TorchNeuron es compatible con los flujos de trabajo estándar de PyTorch en Trainium, incluida la ejecución eager, torch.compile, y las APIs 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 de 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 key_averages() tabla 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 hasta la programación en tiempo de ejecución y 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 hacer profiling de una carga de trabajo, ayuda entender cómo Neuron se integra con PyTorch.

El backend de Neuron utiliza el punto de extensión PrivateUse1 de PyTorch. Este mecanismo permite que un acelerador se integre con el framework sin agregar una dispatch key 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 hacer profiling de 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 tabla key_averages() , donde puede comparar el tiempo de CPU y el tiempo del dispositivo Neuron. Para la referencia completa de parámetros y los valores ProfileMode 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 a través de un túnel SSH o de 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 los envíos.

Después del perfilado, 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 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 para un perfil de sistema o combinado. Para un perfil solo de dispositivo, cargue los archivos .neff y .ntff correspondientes individualmente. Vea Step 5: Upload your profile para las instrucciones de la interfaz de usuario y de la línea de comandos.

Perfilado de RMSNorm

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

Perfila 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 Neuron fusionado.
3. RMSNorm ingenuo con torch.compile, usando las mismas operaciones que la primera implementación.

Los scripts utilizan las mismas entradas, la misma configuración del profiler y el mismo calendario, por lo 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 Naive 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 específicas de Kineto para Neuron.

El calendario utiliza 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 la columna de tiempo de Neuron vacía.

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

Después de que finalice el profiling, 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 en el dispositivo porque la función realiza tres multiplicaciones.

Total Self NEURON: 2.355 ms

Cada fila corresponde directamente a 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 del framework pasando por la programación del runtime hasta la ejecución en los lanes del hardware Neuron.

Lea 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 de ATen sigue el mismo patrón: creación en la lane del framework, planificación en la lane del runtime como nrta_execute_scheduley ejecución en el dispositivo en la lane del hardware como nc_exec_running.

La vista de system trace 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 planificació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 las lanes del framework, el runtime y el hardware.

Perfil 2: torch.rms_norm

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

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

El RMSNorm fusionado 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. El trace 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 sola 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 programación del 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 sigue 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 original de Python. En esta ejecución, el tiempo de Neuron es comparable al de torch.rms_norm, con 375 μs frente a 403 μs.

Resumen lado a lado

Tanto torch.rms_norm como torch.compile integran RMSNorm en un único 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 en el lado de la CPU. Explorer proporciona el rastro visual correspondiente. Puedes seguir una operación desde el despacho del framework, pasando por la planificación del runtime, hasta la ejecución en hardware, y ver los huecos que la fusión elimina.

Puedes usar el mismo flujo de trabajo en modelos más grandes: encuentra un punto crítico en la tabla, verifícalo en el rastro 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 de TensorEngine, 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!