TorchNeuron, PyTorch 코드 한 줄만 바꾸면 될까?

장치 선택은 간단해졌지만 호환 확인은 남는다

TorchNeuron은 기존 PyTorch 모델을 AWS Trainium에서 실행하는 네이티브 백엔드다. AWS의 10월 20일 발표는 .to('neuron')으로 장치를 바꾸는 간편함을 앞세웠다. 공개 베타의 대상은 Amazon EC2 Trainium2·Trainium3 인스턴스다.

개발자는 익숙한 PyTorch 모델 정의와 도구를 이어 쓸 수 있다. 다만 모델을 장치로 옮겼다고 해서 그 모델의 모든 연산이 지원되는 것은 아니다. 입력 텐서도 해당 장치로 옮겨야 하고, 기존 XLA 기반 Neuron 코드라면 새 neuron 장치에 맞춰 수정해야 한다.

디버깅은 eager, 그래프 최적화는 torch.compile이다

PyTorch에 공개된 개발진 설명은 두 실행 방식을 제시한다. eager는 Python의 조건문·출력·중단점을 활용해 연산을 따라가기 좋다. torch.compile은 모델 전체나 일부를 그래프로 잡아 최적화한다. 성능이 중요한 연산에는 NKI로 작성한 맞춤 커널을 연결할 수 있다.

개발진은 모델 검증에서도 로딩 성공만 보지 않았다. 축소한 대표 구성에서 순전파·역전파·컴파일 여부와 CPU로 넘어간 연산을 조사했다고 밝혔다. 모델이 한 번 실행됐다는 사실만으로 실제 학습·추론 조건의 검증이 끝나는 것은 아니다.

컴파일 경로도 GPU와 다르다. Neuron 백엔드는 TorchInductor 대신 StableHLO와 Neuron 컴파일러를 사용한다. 기존 CUDA 전용 컴파일 옵션을 그대로 가져오는 방식으로 생각해서는 안 된다.

공개 베타는 버전 조합과 미지원 API가 정해져 있다

Neuron 2.33.0 릴리스 문서에 명시된 조건은 다음과 같다.

CPU 대체 실행은 전체 프로그램을 계속 돌리는 데 도움이 되지만, 그 연산이 Trainium에서 처리됐다는 뜻은 아니다. torch_neuronx.get_fallback_ops()로 기록된 대체 연산을 확인할 수 있다.

입력 크기와 첫 실행 시간을 따로 확인한다

현재 torch.compile은 정적 입력 크기를 사용하며 dynamic=False가 필요하다. 입력 크기가 달라지면 새 그래프를 컴파일할 수 있다. 공식 문서는 패딩이나 버킷 구분으로 크기의 종류를 줄여 캐시 재사용을 늘리는 방법을 제시한다.

컴파일 실패 처리도 구분해야 한다. eager의 개별 연산 CPU 대체와 달리, 컴파일·실행 실패 때 원래의 비컴파일 그래프로 돌아가는 기능은 기본적으로 꺼져 있다. 기본 동작은 오류를 드러내는 것이다.

이전 시험에서는 대표 입력 크기별로 결과와 실행 시간을 보고, 캐시가 없는 첫 호출과 캐시가 준비된 호출을 나눠 기록하는 편이 유용하다. CPU 대체 연산까지 함께 보면 단순 실행 성공과 원하는 성능을 구분할 수 있다. 한 줄 변경은 이전의 시작이다. 완료 판단에는 실제 모델의 연산과 입력 조건이 필요하다.

자주 묻는 질문