많은 ML 실무자들이 모델을 빠르게 실행하는 방법에 대한 직관은 가지고 있지만, 어떤 이유로 모델이 느린지 분석하는 방법을 아는 사람은 적습니다. 이 시리즈의 Part 1에서는 추론 성능을 체계적으로 분석하고 시간이 어디에 사용되는지 식별하는 방법을 배울 것입니다.
본 게시물은 아무런 사전 프로파일링 경험이 없다고 가정하고 처음부터 시작합니다. 워크스루에는 딥러닝 학습 및 추론 워크로드를 위해 특별히 설계된 AI 칩인 AWS Trainium과 그 오픈 소스 네이티브 PyTorch 백엔드인 TorchNeuron을 사용합니다. TorchNeuron은 최소한의 코드 변경만으로 eager execution, torch.compile, 그리고 표준 PyTorch API를 포함하여 Trainium에서 표준 PyTorch 워크플로우를 지원합니다. 또한, TorchNeuron은 Neuron Kernel Interface (NKI)와 통합되어 사용자 정의 연산자(custom operators) 및 Trainium 하드웨어에 대한 성능 최적화를 가능하게 합니다. 끝날 무렵에는 다음을 알게 될 것입니다:
- Neuron 백엔드를 사용하여
torch.profiler를 구성하고 깨끗한 트레이스(trace)를 캡처하는 방법.
- 어떤 연산이 Neuron 장치 시간을 소비하는지 식별하고 Self CPU 시간과 비교하기 위해
key_averages() 테이블을 읽는 방법.
- 프레임워크 디스패치부터 런타임 스케줄링, 그리고 NeuronCore에서의 하드웨어 실행에 이르기까지 하나의 연산이 Neuron Explorer를 통해 어떻게 추적되는지 확인하는 방법.
- 동일한 연산의 세 가지 구현에서 이러한 신호가 어떻게 변화하는지 확인하는 방법.
워크로드를 프로파일링하기 전에, Neuron이 PyTorch와 어떻게 통합되는지 이해하는 것이 도움이 됩니다.
Neuron 백엔드는 PyTorch의 PrivateUse1 확장 지점(extension point)을 사용합니다. 이 메커니즘은 가속기(accelerator)가 cuda와 같은 전용 디스패치 키를 추가하지 않고도 프레임워크와 통합할 수 있게 합니다. torch_neuronx 패키지는 상위 PyTorch와 협력하여 Neuron을 PyTorch 코드의 네이티브 백엔드로 등록합니다. 이 확장 메커니즘에 대한 자세한 내용은 Facilitating New Backend Integration by PrivateUse1을 참조하십시오.
Neuron 워크로드를 프로파일링하려면 두 개의 객체를 사용합니다:
NeuronConfig는 어떤 Neuron 데이터 스트림과 NeuronCore를 캡처하고 출력을 어디에 쓸지 제어합니다. 이를 torch.profiler.profile에 전달합니다.
asexperimental_config
.NeuronProfiler
은 NeuronConfig를 받아 export_trace 콜백을 제공합니다. 이 콜백은 프레임워크 추적(framework trace)인 trace.json과 함께 Neuron 시스템 추적을 기록합니다.
nc = NeuronConfig(modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
profile_output_dir=OUTPUT_DIR,
...
modes 매개변수는 장치 명령어 및 런타임 추적과 같이 캡처할 데이터 스트림을 선택합니다. capture_enabled_for_nc 매개변수는 추적할 NeuronCore를 선택합니다. correlate_device_time=True로 설정하면, 장치 시간을 추적하여 key_averages() 테이블에 전파하며, 여기서 CPU 시간과 Neuron 장치 시간을 비교할 수 있습니다. 전체 매개변수 참조 및 사용 가능한 ProfileMode 값은 NeuronConfig Profile Modes 문서를 참조하십시오.
Neuron Explorer는 시스템 추적, 장치 타임라인, 엔진별 상세 분석을 위한 브라우저 기반 뷰어입니다. SSH 터널 또는 Visual Studio Code 확장을 통해 접근할 수 있습니다. 본 게시물에서는 Explorer를 사용하여 key_averages()가 수치적으로 보고하는 내용을 확인하고 디스패치(dispatches) 사이에 발생하는 일을 검사합니다.
프로파일링 후, 출력 디렉터리에는 시스템 추적과 장치 데이터를 캡처한 경우 매칭되는 .neff 및 .ntff 파일이 포함됩니다. Neuron Executable File Format (NEFF)은 컴파일된 Neuron 실행 파일을 담고 있으며, 이에 대응하는 NTFF는 해당 하드웨어 명령어 추적을 담고 있습니다. 이 파일들은 쌍으로 식별되는 숫자 해시를 공유합니다. Neuron Explorer에서 프로파일을 열려면 시스템 또는 결합(combined) 프로파일의 전체 출력 디렉터리를 업로드하십시오. 장치 전용 프로파일의 경우, 매칭되는 .neff 및 .ntff 파일을 개별적으로 업로드하십시오. UI 및 명령줄 지침은 5단계: 프로파일 업로드를 참조하십시오.
첫 번째 분석을 위해, 산술 연산과 축소(reduction) 연산이 명확한 작은 연산인 RMSNorm을 프로파일링합니다.
동일한 연산의 세 가지 구현체를 프로파일링합니다:
- 개별 PyTorch 연산(multiply, mean, add, reciprocal square root, multiply)을 사용한 순수 RMSNorm:
torch.rms_norm()
이는 PyTorch의 내장 함수가 단일 융합된 Neuron 커널로 낮춰진 형태입니다.
- 첫 번째 구현과 동일한 연산을 사용하는
torch.compile을 이용한 순수 RMSNorm.
스크립트는 동일한 입력, 프로파일러 설정 및 스케줄을 사용하므로 결과는 직접 비교가 가능합니다.
이 작업을 재현하는 모든 스크립트는 torch-neuronx examples 디렉터리에서 확인할 수 있습니다.
먼저 개별 PyTorch 연산으로 RMSNorm을 작성해 보세요. 각 줄은 하나 이상의 ATen 연산에 매핑되며 중간 텐서를 생성할 수 있습니다.
device = torch.device('neuron')
x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)
...
Neuron 백엔드에 필요한 두 개의 인수를 사용하여 프로파일러를 호출합니다:
- Neuron 장치에서 활동을 수집하려면
ProfilerActivity.PrivateUse1을 포함합니다.
- Neuron 전용 Kineto 기능을 활성화하려면
experimental_config로 NeuronConfig를 전달합니다.
스케줄은 네 번의 prof.step() 호출에 걸쳐 wait=2, warmup=1, active=1을 사용합니다. 이는 초기 설정 오버헤드를 포함하지 않고 하나의 웜업(warmup) 반복을 기록합니다.
프로파일러 루프 내부에서는 torch.neuron.synchronize()를 호출해야 합니다. 이 동기화가 없으면 프로파일러가 장치가 완료되기 전에 내보내져 Neuron 시간 열이 비어 있게 됩니다.
with torch.profiler.profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
experimental_config=nc,
...
프로파일링을 마친 후에는 key_averages() 테이블을 검사합니다. self_privateuse1_time_total로 정렬하여 장치 비용이 가장 높은 연산을 먼저 배치합니다. 여기에 표시된 표는 Neuron 장치 시간을 보고하는 RMSNorm 연산으로 필터링되었습니다. 전체 key_averages() 출력에는 할당(allocations) 및 호스트-투-장치 전송과 같은 CPU 측 이벤트도 포함됩니다.
print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))
키 평균값 테이블은 세 개의 aten::mul 호출이 Neuron 장치 시간의 가장 큰 비중을 차지함을 보여줍니다.
이 테이블은 어떤 ATen 연산이 실행되었는지, 얼마나 자주 실행되었는지, 그리고 CPU와 Neuron 장치에서 얼마만큼의 시간을 사용했는지를 보여줍니다. aten::mul은 함수가 세 번의 곱셈을 수행하기 때문에 가장 높은 총 장치 시간(total device time)을 가집니다.
| Name | # Calls | Self CPU | Self NEURON | % NEURON |
|---|
| aten::mul | 3 | 60.9 μs | 1,164.5 μs | 49.5% |
| ... | | | | |
| Self NEURON 총합: 2.355 ms | | | | |
각 행은 함수에 직접 매핑됩니다. aten::mul은 x * x, x * rms, 그리고 최종적으로 weight를 곱하는 연산을 수행합니다. aten::mean은 마지막 차원을 축소(reduce)합니다. aten::add는 엡실론 값(epsilon value)을 더하고, 스칼라가 CPU에서 시작하기 때문에 호스트-투-디바이스 복사(host-to-device copy)를 유발합니다. aten::rsqrt는 역제곱근(reciprocal square root)을 계산합니다.
Neuron Explorer는 프레임워크 디스패치부터 런타임 스케줄링, 그리고 Neuron 하드웨어 레인에서의 실행에 이르는 aten::mul 연산의 흐름을 추적합니다.
Explorer에서 시간은 왼쪽에서 오른쪽으로 흐르고 이벤트 수명 주기(event lifecycle)는 위에서 아래로 흐릅니다. 각 ATen 연산은 동일한 패턴을 따릅니다: 프레임워크 레인에서의 생성(creation), nrta_execute_schedule로서의 런타임 레인에서의 스케줄링(scheduling), 그리고 nc_exec_running으로서의 하드웨어 레인에서의 장치 실행(device execution)입니다.
시스템 트레이스 뷰는 이벤트를 레인별로 구성합니다. 기본 그룹화 모드(Default grouping mode)에서는 레이블이 instance/PID/subsystem/identifier 패턴을 따릅니다:
| Lane | Label suffix | What it shows |
|---|
| Framework dispatch | framework/TID:<id> | Operator dispatch events (PrivateUse1 runtime), grouped by thread ID. |
| Framework device | framework/Stream:<n> | Device interaction events (PrivateUse1 driver), grouped by stream number. |
| Neuron Runtime | neuron_rt/NC:<n> | Neuron Runtime API calls, grouped by logical NeuronCore. |
| Neuron Hardware | neuron_hw/NC<n>:<id> | Hardware execution events, grouped by physical NeuronCore. |
eager 모드에서는 각 연산이 독립적으로 실행됩니다. 의존성 체인 뷰어(dependency chain viewer)를 사용하여 프레임워크 계층에서 런타임 스케줄링을 거쳐 장치 실행까지 aten::mul.Tensor의 흐름을 추적할 수 있습니다.
의존성 체인 뷰어는 하나의 연산이 프레임워크, 런타임, 하드웨어 레인을 가로지르는 라이프사이클을 강조합니다.
개별 연산을 PyTorch의 내장 rms_norm 함수로 대체해 보세요. Neuron은 이 연산을 단일 융합 커널(fused kernel)로 낮춥니다:
def rms_norm_fn(t):
return torch.rms_norm(t, [512], weight)
융합된 RMSNorm은 주요 평균값 표에서 하나의 장치 연산으로 나타납니다.
| Name | # Calls | Self CPU | Self NEURON |
|---|
| aten::_fused_rms_norm | 1 | 47.6 μs | 403.3 μs |
Self NEURON 총합: 403.3 μs
융합된 RMSNorm은 개별 산술 연산을 분산(dispatch)하는 대신 하나의 프레임워크-하드웨어 경로를 따릅니다.
순수 구현(naive implementation)에서는 각 연산이 독립적으로 디스패치 경계를 넘고 자체적인 프레임워크-런타임-하드웨어 경로를 따릅니다. 융합된 구현에서는 하나의 Python 호출이 전체 계산을 포괄합니다. 트레이스는 컴포넌트 연산 사이에 간극 없이 단일 디스패치와 단일 하드웨어 이벤트를 포함합니다.
마지막 예시를 위해, 원래의 순수 RMSNorm 함수를 torch.compile로 컴파일해 보세요.
@torch.compile(backend=
, TorchDynamo는 Python 바이트코드를 FX 그래프로 포착하고 AOTAutograd가 이를 ATen 연산으로 낮춥니다. Neuron 백엔드는 각 연산을 개별적으로 디스패치하는 대신 전체 그래프를 받습니다. 이는 연산들을 융합(fuses)하고, 메모리 접근을 최적화하며, 그래프를 단일 NEFF 아티팩트로 컴파일합니다. 이 아티팩트는 하나의 하드웨어 디스패치로 실행됩니다.
@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
x_sq = t * t
...
컴파일된 함수는 주요 평균값 테이블에서 하나의 Torch-Compiled Region으로 나타납니다.
| Name | # Calls | Self CPU | Self NEURON | % NEURON |
|---|
| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |
Self NEURON 총합: 375.5 μs
컴파일된 그래프는 런타임 스케줄링과 하드웨어 실행을 통해 하나의 경로를 따릅니다.
이 표는 Profile 2와 같이 단일 장치(device) 실행을 보여주지만, 코드는 다른 경로를 따릅니다. 컴파일러가 목적에 맞게 설계된 융합 연산을 호출하는 대신, 원래 Python 함수로부터 융합을 발견합니다. 이 실행에서 Neuron 시간은 torch.rms_norm의 375 μs와 비교하여 403 μs인 것과 유사합니다.
| Variant | NEURON time | Self CPU time | Operator calls |
|---|
| Naive ops | 2,355 µs | ~125 µs | 6 |
| ... | | | |
둘 다 torch.rms_norm과 torch.compile은 RMSNorm을 하나의 커널 디스패치로 융합합니다. 이러한 실행에서 두 접근 방식 모두 기본 구현의 장치 시간 중 약 6분의 1을 소요합니다. 융합은 하드웨어 실행 사이의 간격을 제거하고, 개별적으로 디스패치된 연산들 사이에 x_sq와 mean_sq 같은 중간 값을 메모리화하는 것을 방지합니다. | | | |
key_averages()는 어떤 연산이 장치 시간을 소비하는지 보여주며, 전체 출력은 CPU 측 설정 및 전송도 밝혀줍니다. Explorer는 이에 상응하는 시각적 추적을 제공합니다. 프레임워크 디스패치부터 런타임 스케줄링을 거쳐 하드웨어 실행까지 하나의 연산을 따라가면서 융합이 제거하는 간격들을 확인할 수 있습니다.
더 큰 모델에서도 동일한 워크플로우를 사용할 수 있습니다. 표에서 핫스팟을 찾고, 트레이스에서 이를 검증한 다음, 결합된 연산(fused operation)이나 torch.compile이 비용을 줄일 수 있는지 판단할 수 있습니다.
Part 2에서는 장치 레벨 프로파일링으로 넘어갈 것입니다. NEFF를 열어 TensorEngine, VectorEngine, 그리고 직접 메모리 접근(DMA)의 세부 내역을 검사하고, 커널을 루프라인(roofline)에 배치하여 계산 바운드인지 아니면 메모리 바운드인지를 판단합니다.
이 시리즈를 시작할 영감을 준 Aritra Roy Gosthipaty님께 특별한 감사를 표하며, 리뷰해 주신 동료 Pinak Panigrahi님, Samir Araujo님, Akhil Azhikodan님, Esha Lakhotia님과 모든 분들께도 감사드립니다!
