TorchNeuron 프로파일링 1부: AWS Trainium에서 시스템 수준 추론 분석
많은 ML 실무자들이 모델을 가속화하는 방법을 이해하지만, 모델이 왜 느린지 분석하려면 체계적인 프로파일링이 필요합니다. 이 가이드는 TorchNeuron을 사용하여 AWS Trainium에서 추론 성능을 분석하는 방법을 안내하며, 시간 소비 핫스팟을 식별하고 커널 융합의 성능 이점을 입증하는 포괄적인 방법을 제공합니다.
PyTorch 백엔드로서의 Neuron 통합
Neuron 백엔드는 PrivateUse1 확장 지점을 사용하여 PyTorch와 통합됩니다. 이 메커니즘은 cuda와 같은 전용 디스패치 키 없이도 가속기가 프레임워크와 통합할 수 있게 합니다. torch_neuronx 패키지는 PyTorch의 고급 확장 기능을 활용하여 Neuron을 PyTorch 코드의 네이티브 백엔드로 등록합니다.
프로파일러 설정
Neuron 워크로드를 프로파일링하려면 NeuronConfig와 NeuronProfiler 두 객체를 구성해야 합니다. NeuronConfig는 캡처할 Neuron 데이터 스트림과 NeuronCore를 제어하고 출력 위치를 지정합니다. NeuronProfiler는 이 구성을 사용하여 프레임워크 추적(trace.json)을 Neuron 시스템 추적과 함께 작성하는 export_trace 콜백을 제공합니다.
주요 구성 매개변수는 다음과 같습니다:
modes: 장치 명령어 및 런타임 추적과 같은 캡처할 데이터 스트림을 선택합니다.capture_enabled_for_nc: 추적할 특정 NeuronCore를 선택합니다.correlate_device_time=True: 이 설정은 장치 시간을 추적에서key_averages()테이블로 전파하여 CPU와 Neuron 장치 시간을 직접 비교할 수 있게 합니다.
도구: Neuron Explorer
Neuron Explorer는 시스템 추적, 장치 타임라인, 엔진별 분석을 시각화하도록 설계된 브라우저 기반 뷰어입니다. 사용자가 key_averages()의 수치 데이터를 확인하고 작업의 상세 수명 주기를 검사할 수 있습니다.
프로파일링 후 출력 디렉터리에는 시스템 추적과 함께, 장치 데이터가 캡처된 경우 일치하는 .neff(Neuron 실행 파일 형식) 및 .ntff(하드웨어 명령어 추적) 파일이 포함됩니다. Neuron Explorer에서 프로필을 보려면 시스템 또는 결합 프로필의 경우 전체 출력 디렉터리를 업로드하고, 장치 전용 프로필의 경우 일치하는 .neff 및 .ntff 파일을 개별적으로 업로드해야 합니다.
사례 연구: RMSNorm 프로파일링
RMSNorm은 산술 및 축소 작업의 명확한 시퀀스를 가진 작은 작업이므로 이 분석에 사용되었습니다. 성능을 비교하기 위해 RMSNorm의 세 가지 구현을 프로파일링했습니다:
- Naive RMSNorm: 개별 PyTorch 작업(곱셈, 평균, 덧셈, 역제곱근, 곱셈)을 사용하여 구현.
torch.rms_norm(): 단일 융합 Neuron 커널로 낮춰지는 PyTorch 내장 함수.- Naive RMSNorm +
torch.compile: Neuron 백엔드로torch.compile을 사용하여 컴파일된 원래 naive 함수.
모든 스크립트는 결과의 직접 비교 가능성을 보장하기 위해 동일한 입력, 프로파일러 구성, 스케줄링을 사용했습니다.
프로필 1: Naive PyTorch 작업
RMSNorm이 개별 PyTorch 작업으로 구현되면 각 줄은 하나 이상의 ATen 작업에 매핑되고 중간 텐서를 생성할 수 있습니다. 프로파일링 결과 세 개의 aten::mul 호출이 Neuron 장치 시간의 가장 큰 부분을 차지하는 것으로 나타났습니다.
| 이름 | 호출 수 | 자체 CPU | 자체 NEURON | % NEURON |
|---|---|---|---|---|
| aten::mul | 3 | 60.9 μs | 1,164.5 μs | 49.5% |
| aten::add | 1 | 26.6 μs | 434.6 μs | 18.5% |
| aten::rsqrt | 1 | 15.6 μs | 415.0 μs | 17.6% |
| aten::mean | 1 | 21.1 μs | 340.8 μs | 14.5% |
자체 NEURON 총계: 2.355 ms
이 naive 구현에서 각 작업은 디스패치 경계를 독립적으로 교차합니다. 예를 들어 aten::add 작업은 스칼라 엡실론이 CPU에서 시작되므로 호스트-장치 복사를 트리거합니다.
프로필 2: torch.rms_norm()
개별 작업을 PyTorch의 내장 rms_norm 함수로 대체하면 Neuron은 전체 작업을 하나의 융합 커널로 낮춥니다. 키 평균 테이블은 이 작업을 단일 장치 작업으로 보여줍니다.
| 이름 | 호출 수 | 자체 CPU | 자체 NEURON | % NEURON |
|---|---|---|---|---|
| aten::_fused_rms_norm | 1 | 47.6 μs | 403.3 μs | 100% |
자체 NEURON 총계: 403.3 μs
이 융합 구현은 단일 프레임워크-하드웨어 경로를 따르므로 각 산술 작업을 개별적으로 디스패치하는 오버헤드를 피합니다.
프로필 3: torch.compile을 사용한 Naive PyTorch
원래 naive RMSNorm 함수가 @torch.compile(backend='neuron', fullgraph=True)를 사용하여 컴파일되면 TorchDynamo는 Python 바이트코드를 FX 그래프로 캡처합니다. Neuron 백엔드는 이 전체 그래프를 수신하여 융합하고 단일 NEFF 아티팩트로 컴파일합니다. 이 아티팩트는 단일 하드웨어 디스패치로 실행됩니다.
| 이름 | 호출 수 | 자체 CPU | 자체 NEURON | % NEURON |
|---|---|---|---|---|
| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |
자체 NEURON 총계: 375.5 μs
이 컴파일된 그래프는 또한 런타임 스케줄링 및 하드웨어 실행을 통해 단일 경로를 따릅니다. 여기서 달성된 Neuron 시간(375 μs)은 특수 목적의 torch.rms_norm(403 μs)과 비슷합니다.
나란히 비교 성능 요약
| 변형 | NEURON 시간 | 자체 CPU 시간 | 연산자 호출 |
|---|---|---|---|
| Naive 작업 | 2,355 µs | ~125 µs | 6 |
| torch.rms_norm | 403 µs | ~48 µs | 1 |
| Naive + torch.compile | 376 µs | ~98 µs | 1 |
torch.rms_norm과 torch.compile 모두 RMSNorm을 하나의 커널 디스패치로 성공적으로 융합합니다. 이러한 융합 방법은 장치 시간을 naive 구현의 장치 시간의 약 1/6로 줄입니다. 융합은 하드웨어 실행 사이의 간격을 제거하고 개별적으로 디스패치된 작업 사이에 x_sq 및 mean_sq와 같은 중간 값의 구체화를 방지합니다.
결론
key_averages() 테이블은 어떤 작업이 장치 시간을 소비하는지 식별하는 반면, 전체 출력은 CPU 측 설정 및 전송에 대한 컨텍스트를 제공합니다. Neuron Explorer는 해당 시각적 추적을 제공하여 사용자가 프레임워크 디스패치에서 런타임 스케줄링, 하드웨어 실행까지 작업을 따라가고 융합이 제거하는 간격을 관찰할 수 있습니다.
이 워크플로는 더 큰 모델에 적용할 수 있습니다: 테이블에서 성능 핫스팟을 식별하고, 추적에서 시각적으로 확인한 다음, 융합 작업 또는 torch.compile이 비용을 줄일 수 있는지 결정합니다. 향후 분석은 장치 수준 프로파일링으로 이동하여 TensorEngine, VectorEngine 및 DMA 분석을 검사하여 커널이 계산 바운드인지 메모리 바운드인지 확인할 것입니다.