Amazon은 Trainium AI 칩 제품군을 위한 오픈 소스 네이티브 PyTorch 백엔드인 TorchNeuron을 발표했습니다. 이는 기계 학습 워크로드를 위한 맞춤형 실리콘을 채택하는 데 있어 가장 큰 실질적인 장벽 중 하나를 제거하는 조치입니다. 10월 20일 북미 PyTorch 컨퍼런스에 앞서 게시된 PyTorch 블로그의 공식 게시물에 따르면, 개발자는 이제 독점 도구에 대한 워크플로를 다시 작성하는 대신 한 줄만 변경하면 Trainium에서 기존 PyTorch 모델 코드를 실행할 수 있습니다.

이번 발표는 분석가들이 오랫동안 Nvidia의 가장 내구성 있는 장점이라고 설명했던 원시 칩 성능이 아니라 개발자를 GPU에 연결하는 CUDA 소프트웨어 생태계를 대상으로 하기 때문에 즉각적인 관심을 끌었습니다. Reuters는 Amazon이 자사의 AI 칩이 AI 분야에서 가장 널리 사용되는 개발 프레임워크와 작동할 것이라고 밝혔으며 TechCrunch는 TorchNeuron이 Nvidia에서 회사의 가장 큰 기회라고 말했습니다. 가속기 시장이 어떻게 변화하고 있는지 더 폭넓게 살펴보기 위해 진행 중인 AI 하드웨어 뉴스 보도에서 이러한 발전 상황을 추적합니다.

TorchNeuron이 실제로 변경하는 것

TorchNeuron의 핵심 약속은 단순성입니다. PyTorch 블로그에 게시된 예에서는 '.to('neuron')'을 사용하여 모델을 새 장치로 이동한 후 Trainium에서 기존 훈련 루프가 실행됩니다. 게시물에 따르면 PyTorch가 권장하는 장치 독립적 API인 'torch.accelerator'에 대해 작성된 코드는 수정되지 않은 채 실행됩니다. 컴파일된 워크로드의 경우 개발자는 `torch.compiler.set_default_backend("neuron")`를 사용하여 기본 백엔드를 한 번 설정한 다음 추가 변경 없이 `model.compile()`을 호출합니다.

내부적으로 TorchNeuron은 맞춤형 하드웨어 백엔드를 위한 PyTorch의 공식 확장 포인트인 PrivateUse1을 통해 Trainium을 등록합니다. 이는 칩이 프레임워크에 고정된 부속 장치가 아닌 표준 PyTorch 장치처럼 작동한다는 것을 의미합니다. AWS는 통합이 여러 실행 모드에 걸쳐 있다고 설명합니다.

  • 실행 오버헤드를 줄이는 실험적인 "적응형 Eager 모드"를 통해 연구 및 디버깅을 위한 Eager 실행
  • 생산 성능을 위한 torch.compile, FlexAttention에 대한 기본 지원을 통해 Neuron 컴파일러가 최적화하는 그래프로 모델 캡처
  • 맞춤형 저수준 커널이 필요한 연구자를 위한 뉴런 커널 인터페이스(NKI)

발표에 따르면 백엔드는 Trainium2와 최신 Trainium3 칩 모두에서 작동합니다.

Nvidia의 소프트웨어 해자에 대한 직접적인 공격

AI 컴퓨팅에 대한 Nvidia의 지배력은 결코 더 빠른 하드웨어에만 의존한 적이 없습니다. 회사의 20년 된 소프트웨어 플랫폼인 CUDA는 대부분의 AI 연구 코드가 작성, 조정 및 배포되는 곳이며 AMD의 Instinct 라인부터 Google의 TPU 및 Amazon의 Trainium에 이르기까지 모든 도전자 칩은 개발자에게 이를 떠나도록 요구하는 비용을 감당해야 했습니다. GeekWire, Benzinga 및 GuruFocus에서 TorchNeuron 출시에 대한 보도는 해당 발표를 해당 해자에 대한 공격으로 균일하게 구성했습니다.

Amazon은 개발자를 PyTorch에서 끌어내려고 하지 않습니다. 대신 Trainium을 PyTorch에 도입하여 프레임워크의 막대한 사용자 기반이 습관을 바꾸지 않고도 칩을 채택할 수 있게 되었습니다. PyTorch 블로그 게시물에는 프레임워크를 유지 관리하는 Meta의 보증이 포함되어 있습니다. "AWS Trainium에 대한 기본 지원은 선택의 폭을 넓혀 연구원과 엔지니어가 이미 알고 있는 Eager Execution 및 분산 API를 사용하여 실험하고 반복할 수 있게 해줍니다."라고 Meta의 PyTorch 엔지니어링 이사인 Jana van Greunen이 게시물에서 말했습니다.

이러한 프레임은 AWS의 클라우드 비즈니스에 중요합니다. Trainium 인스턴스는 훈련 및 추론을 위해 Nvidia GPU에 대한 저렴한 대안으로 마케팅되었으며 이 칩은 이미 주요 AI 고객을 끌어 모았습니다. 칩이 PyTorch에 고유하게 느껴지도록 만드는 것은 광범위한 채택을 지연시키는 반대, 즉 포팅 노력을 직접적으로 해결합니다.

기술 범위, 노트북부터 클러스터 규모 교육까지

PyTorch 블로그에 나열된 기능은 Amazon이 취미로 하는 실험 이상의 것을 목표로 하고 있음을 시사합니다. TorchNeuron은 BF16, FP32 및 FP16 정밀도로 텍스트, 다중 모드 및 비전 언어 모델은 물론 이미지 및 비디오 생성, 상태 공간 모델 및 세계 모델 전반에 걸쳐 표시되는 연산자 세트를 다룹니다. TorchAO 및 FP8 지원은 다음 릴리스의 일반 출시에 앞서 비공개 베타 버전으로 제공됩니다.

단일 GPU 코드를 클러스터 규모 실행으로 전환하는 방식인 분산 교육은 다차원 프로세스 그룹을 위한 DTensor 및 DeviceMesh와 함께 표준 PyTorch API(DDP, FSDP2, SimpleFSDP, 텐서 및 파이프라인 병렬 처리, ZeRO 스타일 샤딩)를 통해 지원됩니다. 스트림과 이벤트는 PyTorch의 의미를 Trainium의 하드웨어 대기열에 매핑한다고 게시물은 말했습니다.

추론을 위해 Amazon의 Neuron 스택은 동일한 기반 위에 구축됩니다. vLLM Neuron은 대규모 언어 모델을 제공하고 vLLM Omni Neuron은 이미지 및 비디오 생성을 위한 확산 모델을 제공합니다. Neuron Kernel Interface를 사용하면 연구원은 즉시 실행 및 컴파일된 실행을 위한 사용자 지정 연산자를 작성할 수 있으며, AWS는 주류 LLM 및 다중 모드 시나리오에 대한 FlashAttention 및 MoE 블록 단위 계산을 다루는 Trainium 최적화 커널 세트인 NKI-lib를 오픈 소스로 제공하고 있습니다.

성능 도구도 포함되어 있습니다. Neuron Explorer는 표준 `torch.profiler`를 확장하여 프레임워크에서 하드웨어까지 실행을 추적하고 모든 이벤트를 소스 라인에 다시 연결하며 자동화된 병목 현상 분석을 위해 에이전트가 포스트 노트를 사용할 수 있는 Parquet 형식의 메트릭을 내보낼 수 있습니다.

가용성 및 향후 계획

TorchNeuron은 블로그 게시물에 따르면 10월 20일 북미 PyTorch 컨퍼런스에서 오픈 소스 출시를 앞두고 있으며 Hugging Face Transformers v5까지 호환성이 확장됩니다. AWS는 Trainium을 클라우드 규모의 AI 교육 및 추론을 위해 처음부터 설계된 칩 포트폴리오라고 설명합니다.

전략적 맥락은 모든 주요 클라우드 제공업체가 동시에 Nvidia 가속기의 최대 구매자이자 잠재적 경쟁자가 되는 시장입니다. Microsoft, Google, Meta는 모두 맞춤형 실리콘을 개발하고 있으며 Huawei를 포함한 중국 기업은 CUDA의 중력에서 벗어나기 위해 자체 소프트웨어 스택을 구축하고 있습니다. Amazon은 마이그레이션을 요청하는 대신 PyTorch 내에서 개발자를 만나 해당 콘테스트에서 방해가 가장 적은 경로를 선택했으며, 개발자 경험이 유지된다면 실제 교육 워크로드를 Nvidia 하드웨어에서 이동할 가능성이 가장 높은 경로를 선택했습니다.

TorchNeuron이 실제로 CUDA의 이점을 침식하는지 여부는 규모에 따른 성능과 실제 모델의 롱테일에 달려 있으며 출시 자료에는 답이 없는 질문이 있습니다. 그러나 대부분의 PyTorch 개발자가 재작성된 Trainium을 시도조차 하지 못하게 했던 장벽이 이제는 코드 한 줄로 낮아졌습니다.

---

AI보다 앞서 나가세요

최신 AI 뉴스, 분석, 획기적인 소식을 모두 한 곳에서 받아보세요.

AI 뉴스 자세히 보기 →