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

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

AWS Trainium에서 TorchNeuron의 RMSNorm 세 구현을 프로파일링한 결과, 융합 연산과 torch.compile은 개별 연산 방식의 6회 디스패치를 1회로 줄이고 해당 실행의 Neuron 장치 시간을 약 6분의 1로 낮췄다.

📌 핵심 요약

• TorchNeuron은 AWS Trainium용 오픈 소스 네이티브 PyTorch 백엔드로, PrivateUse1 확장 지점을 이용한다. 글은 torch.profiler의 key_averages()와 Neuron Explorer를 함께 사용해 CPU 시간, Neuron 장치 시간, 프레임워크부터 하드웨어까지의 실행 경로를 분석한다.
• RMSNorm 비교에는 동일한 입력과 프로파일러 설정·스케줄을 사용했다. 입력은 bfloat16 형식의 4×512 텐서이며, 개별 PyTorch 연산 구현, torch.rms_norm, 원래 구현에 torch.compile을 적용한 방식의 세 가지를 측정했다.
• 개별 연산 구현은 연산 호출 6회에 Neuron 시간 2,355 μs, Self CPU 시간 약 125 μs를 기록했다. aten::mul 3회가 장치 시간의 49.5%를 차지했고, epsilon을 더하는 aten::add에서는 CPU에 있던 스칼라 때문에 호스트에서 장치로 복사가 발생했다.
• torch.rms_norm은 단일 융합 커널로 실행되어 Neuron 시간 403.3 μs와 Self CPU 시간 47.6 μs를 기록했다. torch.compile 적용 구현도 단일 장치 실행으로 합쳐졌으며 Neuron 시간 375.5 μs, Self CPU 시간 98.4 μs를 기록했다.
• 이 실행들에서 두 융합 방식의 장치 시간은 개별 연산 방식의 약 6분의 1이었다. 융합은 하드웨어 실행 사이의 공백과 중간값의 별도 구체화를 줄였으며, 정확한 관측을 위해 correlate_device_time=True와 프로파일러 루프 안의 torch.neuron.synchronize()가 중요하다.

🧩 주요 포인트

1. 계측 조건 → correlate_device_time=True는 key_averages()에 장치 시간을 연결하고, torch.neuron.synchronize()는 장치 완료 전에 추적이 내보내져 Neuron 시간 열이 비는 문제를 방지한다.
2. 수치와 실행 경로의 연결 → key_averages()에서 비용이 큰 연산을 찾고 Neuron Explorer에서 프레임워크 디스패치·런타임 스케줄링·하드웨어 실행을 따라가면, 개별 연산 사이의 공백과 융합 효과를 확인할 수 있다.
3. 동일한 단일 실행, 서로 다른 비용 → torch.rms_norm과 torch.compile은 모두 1회 장치 실행을 만들지만, 이번 측정에서는 torch.compile의 Neuron 시간이 조금 짧고 Self CPU 시간은 더 길어 두 시간 지표를 함께 비교해야 한다.

🧠 상세 정리

1. 분석 목표와 TorchNeuron의 PyTorch 통합

이 글은 모델을 빠르게 만드는 방법에 대한 직관에서 출발해, 추론이 느린 이유와 시간이 소비되는 위치를 체계적으로 분석하는 방법을 설명한다. 사전 프로파일링 경험을 가정하지 않으며, 딥러닝 학습과 추론을 위해 설계된 AWS Trainium과 오픈 소스 네이티브 PyTorch 백엔드인 TorchNeuron을 실습 대상으로 삼는다. TorchNeuron은 최소한의 코드 변경으로 eager 실행, torch.compile, 표준 PyTorch API를 지원하고, Neuron Kernel Interface(NKI)와 통합해 사용자 정의 연산 및 하드웨어 성능 최적화를 가능하게 한다. Neuron은 PyTorch의 PrivateUse1 확장 지점을 사용하므로, cuda처럼 전용 디스패치 키를 새로 추가하지 않고 프레임워크에 통합된다. torch_neuronx 패키지는 업스트림 PyTorch와 함께 Neuron을 네이티브 백엔드로 등록하며, 이후 분석은 이 구조 위에서 연산의 실행 경로를 추적한다.

2. 프로파일러 구성과 측정 구간 설정

프로파일링에는 수집할 Neuron 데이터 스트림과 NeuronCore, 출력 위치를 정하는 NeuronConfig와 추적 내보내기를 담당하는 NeuronProfiler가 사용된다. 예제는 modes에 ProfileMode.DEVICE와 ProfileMode.RUNTIME을 지정하고, capture_enabled_for_nc='0'으로 추적할 코어를 선택하며, correlate_device_time=True로 장치 시간을 key_averages() 표에 연결한다. NeuronConfig는 torch.profiler.profile의 experimental_config로 전달하고, NeuronProfiler의 export_trace 콜백은 프레임워크 추적인 trace.json을 Neuron 시스템 추적과 함께 저장한다. 활동 목록에는 CPU와 ProfilerActivity.PrivateUse1을 포함하며, Neuron 전용 Kineto 기능을 활성화하려면 experimental_config 전달도 필요하다. 측정 스케줄은 wait=2, warmup=1, active=1로 구성하고 네 번의 prof.step()을 호출해 초기 설정 비용을 제외한 준비된 반복 한 번을 기록한다. 또한 루프 안에서 torch.neuron.synchronize()를 호출해야 장치가 끝나기 전에 추적이 내보내져 Neuron 시간 열이 비는 문제를 피할 수 있다.

3. Neuron Explorer와 추적 파일의 역할

Neuron Explorer는 시스템 추적, 장치 타임라인, 엔진별 실행 내역을 확인하는 브라우저 기반 도구이며, SSH 터널이나 Visual Studio Code 확장으로 접근할 수 있다. 이 글에서는 key_averages()가 수치로 보여 주는 결과를 시각적으로 확인하고, 디스패치 사이에 어떤 일이 발생하는지 살펴보는 데 사용한다. 프로파일링 출력에는 시스템 추적이 저장되며, 장치 데이터를 수집한 경우에는 서로 대응하는 .neff와 .ntff 파일도 포함된다. NEFF에는 컴파일된 Neuron 실행 파일이, NTFF에는 그 실행 파일의 하드웨어 명령 추적이 들어 있고, 두 파일은 같은 숫자 해시로 짝을 식별한다. 시스템 프로파일이나 시스템·장치 결합 프로파일은 전체 출력 디렉터리를 업로드하고, 장치 전용 프로파일은 대응하는 .neff와 .ntff 파일을 각각 업로드한다. 따라서 분석하려는 추적 종류에 맞춰 파일을 열어야 프레임워크 이벤트와 실제 장치 실행을 연결해 볼 수 있다.

4. 개별 PyTorch 연산으로 구현한 RMSNorm

첫 구현은 bfloat16 형식의 4×512 입력과 길이 512의 가중치를 사용해, 제곱·마지막 차원 평균·epsilon 1e-6 추가·역제곱근·정규화 및 가중치 곱셈을 개별 PyTorch 연산으로 수행한다. 세 구현은 같은 입력과 프로파일러 설정·스케줄을 사용하며, 재현용 스크립트는 torch-neuronx 예제 디렉터리에 제공된다. key_averages()를 self_privateuse1_time_total로 정렬한 결과, aten::mul 3회는 Self CPU 60.9 μs와 Self NEURON 1,164.5 μs로 장치 시간의 49.5%를 차지했다. 나머지는 aten::add가 26.6 μs와 434.6 μs·18.5%, aten::rsqrt가 15.6 μs와 415.0 μs·17.6%, aten::mean이 21.1 μs와 340.8 μs·14.5%를 기록했으며, 각 쌍의 앞 수치는 Self CPU, 뒤 수치는 Self NEURON이다. 전체 Self NEURON 시간은 2.355 ms이고, aten::mul의 높은 합계는 입력 제곱, 정규화, 가중치 적용에 걸친 세 번의 곱셈에 대응한다. epsilon을 더하는 aten::add는 스칼라가 CPU에서 시작하기 때문에 호스트에서 장치로 복사를 유발하며, 전체 표에는 이처럼 CPU 측 할당과 전송 관련 이벤트도 포함된다.

5. 디스패치에서 하드웨어까지 실행 경로 읽기

Explorer에서는 시간이 왼쪽에서 오른쪽으로 흐르고, 한 이벤트의 생명주기는 위쪽 프레임워크 계층에서 아래쪽 하드웨어 계층으로 이어진다. 각 ATen 연산은 프레임워크에서 생성된 뒤 런타임의 nrta_execute_schedule을 거쳐 하드웨어의 nc_exec_running으로 실행된다. 기본 그룹화에서는 레인 이름이 instance/PID/subsystem/identifier 형태를 따르며, framework/TID:<id>는 스레드별 연산 디스패치, framework/Stream:<n>은 스트림별 장치 상호작용을 보여 준다. neuron_rt/NC:<n>은 논리적 NeuronCore별 런타임 API 호출을, neuron_hw/NC<n>:<id>는 물리적 NeuronCore별 하드웨어 실행을 나타낸다. eager 모드에서는 연산마다 이 경로를 독립적으로 통과하므로, 의존성 체인 뷰어로 aten::mul.Tensor 하나를 선택하면 프레임워크·런타임·하드웨어 사이의 연결과 개별 실행 사이의 공백을 확인할 수 있다.

6. torch.rms_norm의 명시적 융합

두 번째 구현은 개별 산술 연산을 torch.rms_norm(t, [512], weight) 호출로 바꾸며, Neuron은 이 연산을 하나의 융합 커널로 낮춘다. key_averages()에는 aten::_fused_rms_norm이 한 번 나타나고, Self CPU 시간은 47.6 μs, Self NEURON 시간은 403.3 μs로 측정된다. 표에 표시된 RMSNorm 장치 시간은 이 단일 연산이 100%를 차지하며, 전체 Self NEURON 시간도 403.3 μs다. 개별 구현이 각 산술 연산마다 디스패치 경계를 넘어 프레임워크·런타임·하드웨어 경로를 따르는 데 비해, 이 구현에서는 하나의 Python 호출이 전체 계산을 포괄한다. 따라서 추적에는 디스패치 한 번과 하드웨어 이벤트 한 번이 나타나고, 구성 연산을 각각 실행할 때 생기던 사이사이의 공백이 사라진다. 이 비교는 연산 호출 횟수 감소가 표의 장치 시간뿐 아니라 실제 실행 경로에도 어떻게 드러나는지 보여 준다.

7. torch.compile이 발견하는 그래프 융합

세 번째 구현은 처음의 개별 연산 RMSNorm 함수에 @torch.compile(backend='neuron', fullgraph=True)을 적용한다. TorchDynamo는 Python 바이트코드를 FX 그래프로 포착하고 AOTAutograd는 이를 ATen 연산으로 낮추며, Neuron 백엔드는 연산을 하나씩 전달받는 대신 전체 그래프를 받는다. 백엔드는 연산을 융합하고 메모리 접근을 최적화한 뒤 그래프를 단일 NEFF 산출물로 컴파일하며, 이 산출물은 한 번의 하드웨어 디스패치로 실행된다. key_averages()에는 Torch-Compiled Region: 0/0이 한 번 나타나고, Self CPU 시간 98.4 μs와 Self NEURON 시간 375.5 μs를 기록한다. torch.rms_norm과 마찬가지로 장치 실행은 한 번이지만, 미리 마련된 융합 연산을 호출하는 대신 컴파일러가 원래 Python 함수에서 융합을 찾아낸다는 점에서 경로가 다르다. 원문은 이 실행의 장치 시간을 torch.rms_norm의 약 403 μs와 비교해 약 375 μs로 설명하며, 두 결과를 비슷한 수준으로 평가한다.

8. 세 구현의 비교와 분석 결과의 범위

나란히 제시한 비교표에서 개별 연산 구현은 Neuron 시간 2,355 μs·Self CPU 약 125 μs·호출 6회, torch.rms_norm은 403 μs·약 48 μs·호출 1회, torch.compile 적용 구현은 376 μs·약 98 μs·호출 1회로 정리된다. 두 융합 방식 모두 이번 실행들에서 개별 연산 방식의 약 6분의 1에 해당하는 장치 시간을 사용했으며, 이 수치는 해당 입력과 측정 조건에서 얻은 결과다. 융합은 하드웨어 실행 사이의 공백을 없애고, 별도로 디스패치되는 연산 사이에서 x_sq와 mean_sq 같은 중간값을 구체화할 필요를 줄인다. 다만 torch.compile은 장치 시간이 조금 더 짧은 반면 Self CPU 시간은 torch.rms_norm보다 길어, 두 지표는 서로 다른 비용을 보여 준다. 결론은 key_averages()에서 장치 시간을 많이 쓰는 연산과 CPU 측 준비·전송을 확인하고, Explorer에서 대응하는 실행 경로와 공백을 살펴보는 분석 흐름을 강조한다. 제공된 source_body는 이 흐름을 더 큰 모델에 적용한다는 설명 도중 문장이 끊겨 있으므로, 그 뒤의 구체적인 결론은 확인할 수 없다.

🧾 핵심 주장 / 시사점

• aten::mul의 장치 시간 합계가 가장 크다는 사실은 세 번의 호출을 함께 반영한다. 연산 이름별 시간뿐 아니라 호출 횟수와 전체 디스패치 구조를 함께 읽는 것이 융합 효과를 이해하는 데 중요하다.
• torch.rms_norm과 torch.compile은 서로 다른 경로로 단일 커널 실행에 도달한다. 하나는 내장 융합 연산을 사용하고, 다른 하나는 원래 함수의 전체 그래프에서 컴파일러가 융합을 발견한다.
• 장치 시간 감소와 CPU 비용 감소는 같은 비율로 움직이지 않았다. 이번 결과에서 torch.compile의 Neuron 시간은 torch.rms_norm보다 짧았지만 Self CPU 시간은 더 길었다.

✅ 액션 아이템

• correlate_device_time=True와 프로파일러 루프 안의 torch.neuron.synchronize()를 적용해 Neuron 장치 시간이 기록되는 측정 조건 확보.
• key_averages()에서 비용이 큰 연산과 호출 횟수를 확인하고, Neuron Explorer에서 프레임워크 디스패치부터 하드웨어 실행까지의 경로와 공백 확인.
• 동일한 입력과 프로파일러 설정·스케줄에서 개별 PyTorch 연산, torch.rms_norm, torch.compile의 Neuron 시간과 Self CPU 시간 비교.

❓ 열린 질문

• 동일한 입력과 프로파일러 설정·스케줄로 측정을 반복하면 약 6분의 1이라는 Neuron 장치 시간 비율이 얼마나 일관되게 나타나는가?
• torch.rms_norm과 torch.compile 중 Neuron 시간과 Self CPU 시간을 함께 고려할 때 어느 방식이 더 적합한가?
• 개별 PyTorch 연산의 Neuron 시간 2,355 μs에서 디스패치 사이의 공백과 중간값 구체화가 각각 차지하는 비용은 얼마인가?