아마존 웹 서비스(AWS)가 10월 19일(현지시간) PyTorch 공식 블로그에서 자사 AI 칩 Trainium을 표준 PyTorch 장치로 쓰게 해 주는 네이티브 백엔드 'TorchNeuron'을 발표했다. 소스코드는 10월 20일 미국 새너제이에서 열리는 PyTorch Conference North America에 맞춰 Apache 2.0 라이선스로 공개되고 같은 날 공개 베타가 시작된다.
핵심 정리
- TorchNeuron은 Trainium2·Trainium3를 PyTorch의 PrivateUse1 확장 지점에 등록해 model.to('neuron') 한 줄로 기존 학습 코드를 옮길 수 있게 한다.
- AWS는 2026년 9월 기준 Hugging Face 모델 10만 개를 축소 구성으로 돌려 로딩·순전파·역전파·컴파일 여부를 점검했다고 밝혔다.
- 음성 AI 기업 Deepgram은 자체 평가에서 Trainium이 기존 GPU 대비 달러당 약 60% 더 많은 입력 오디오(초)를 처리했다고 밝혔다.
TorchNeuron이란 무엇이고 무엇이 달라지나
TorchNeuron은 Trainium을 CUDA GPU처럼 PyTorch의 일반 장치로 다루게 하는 오픈소스 백엔드다. AWS는 PyTorch 블로그 발표문에서 별도의 Neuron 전용 프로그래밍 모델 없이 기존 모델 코드와 라이브러리, 도구를 그대로 쓰는 것이 목표라고 설명했다. PyTorch가 지난 몇 년간 장치 독립 API와 하드웨어 등록용 공식 확장 지점을 갖춘 덕분에 가능해진 방식이다.
즉시 실행(eager) 모드에서는 모델을 'neuron' 장치로 보내는 것 외에 코드를 바꿀 필요가 없다. 파이썬 제어 흐름과 print, 중단점도 그대로 쓴다. 실험 기능인 적응형 eager 모드는 연산을 묶어 실행 비용을 줄여 컴파일 수준 성능에 다가가는 것을 목표로 한다. 운영 환경에서는 torch.compile에 Neuron 백엔드를 지정하면 TorchDynamo와 Torch-MLIR, StableHLO를 거쳐 칩 전용 실행 파일이 만들어진다.
개발자가 쓸 수 있는 기능과 아직 빠진 기능
분산 학습은 DDP와 FSDP2, 텐서·파이프라인 병렬화, ZeRO 같은 표준 API를 그대로 지원한다. 노드 안에서는 Trainium의 전용 집합 통신 코어를, 노드 사이에서는 EFA 네트워크를 쓴다. 성능이 중요한 구간은 NKI 커널 언어로 직접 짤 수 있고, FlashAttention과 MoE 커널을 담은 NKI-lib도 함께 공개된다. 프로파일링은 torch.profiler에 NeuronConfig를 붙이는 방식이다.
한계도 문서에 적혀 있다. 8월 기준 PyTorch 업스트림 테스트 결과 약 15만4천 건 중 6만9천 건가량이 건너뛰어졌고, 이 가운데 약 3만9천 건이 Neuron 쪽 의도적 제외였다. 가장 큰 비중은 아직 네이티브 변환이 없는 연산자 약 2만2천 건이다. 희소 텐서, NestedTensor, 동적 형상 등은 첫 릴리스에서 빠졌다. TorchAO와 FP8은 비공개 베타로, 다음 릴리스에서 정식 지원할 예정이다.
Hugging Face·Meta는 어떻게 반응했나
생태계 연동은 이미 상당 부분 진행됐다. Hugging Face와 AWS는 Transformers v5와 Diffusers, TRL에 Trainium을 업스트림 지원 장치로 넣었고, device_map="neuron"으로 모델을 바로 올릴 수 있다. 추론용으로는 vLLM Neuron 플러그인이, 대규모 학습용으로는 TorchTitan 연동판이 나왔다. Meta의 PyTorch 엔지니어링 디렉터 Jana van Greunen은 연구자가 이미 아는 eager 실행·분산 API로 실험할 수 있게 됐다고 평가했다.
AWS는 지난달 에이전트 하네스를 공개하며 오픈소스 개발자 생태계를 넓혀 왔다(AWS의 에이전트 하네스 Strands 오픈소스 공개). 이번 발표는 그 흐름을 하드웨어 쪽으로 옮긴 셈이다.
전망: GPU 의존을 줄일 수 있을까
AI 가속기 시장에서 소프트웨어 호환성은 칩 성능만큼 중요한 진입 장벽이다. AWS가 Trainium을 PyTorch 표준 장치로 만든 것은 GPU 클라우드 대신 자사 칩을 고르는 전환 비용을 낮추려는 시도로 읽힌다. 다만 범용 처리량·지연 시간 벤치마크는 공개되지 않았고, 10만 개 모델 점검의 통과율도 밝히지 않았다. 실제 이식 난도는 공개 베타 이후 개발자 사용기로 확인될 전망이다.
자주 묻는 질문
TorchNeuron은 오픈소스인가?
그렇다. Apache 2.0 라이선스로 10월 20일 GitHub의 aws-neuron/torch-neuronx 저장소에 공개되며, pip install torch-neuronx로 설치한다. 같은 날 공개 베타가 시작된다.
TorchNeuron은 어떤 칩에서 쓸 수 있나?
AWS 발표 기준 Trainium2와 Trainium3를 지원한다. 지원 자료형은 BF16, FP32, FP16, 정수, 불리언이며 FP8은 아직 비공개 베타다.
기존 CUDA 코드를 그대로 돌릴 수 있나?
torch.accelerator 같은 장치 독립 API로 짠 코드는 수정 없이 돌아간다고 AWS는 설명했다. 다만 CUDA 그래프나 NCCL 전용 경로처럼 CUDA에 묶인 기능은 지원 대상이 아니다.
