Amazon, PyTorch 코드 수정 없이 AWS Trainium에서 실행하는 TorchNeuron 발표

TorchNeuron는 PyTorch가 갖춘 새 하드웨어 등록 표준 메커니즘에 연결되는 방식으로 구현되었다. 이에 따라 테스트, 디버깅, 스케일링에 쓰는 기존 도구를 Trainium에서도 그대로 쓸 수 있다고 한다. Amazon 엔지니어들은 같은 학습 스크립트를 기존 칩과 Trainium에서 실행하는 비교 시연을 진행했으며, 설정 차이만으로 모델이 같은 속도로 같은 방식으로 학습했다고 설명했다.

Amazon은 AI 개발 비용이 칩 가격뿐 아니라 엔지니어의 작업 시간에도 있다고 보고 있다. 새 칩으로 옮길 때 코드 재작성과 재검증이 필요 없어지면, 특히 스타트업은 제품 개선에 더 많은 시간을 쓸 수 있다. 대기업의 경우에는 비피크 시간대에 추론용 용량을 학습 작업에 돌리는 운용 방식이 제시되었다.

초기 사용자로는 음성 AI 기업 Cartesia와 신약 연구를 수행하는 Lieber Institute for Brain Development가 소개되었다. Cartesia는 표준 추론 파이프라인을 최소한의 변경으로 Trainium에서 구동했다. Lieber Institute는 분자 동역학 시뮬레이션을 가속하거나 대체하는 AI 도구를 검토하는 데 TorchNeuron을 사용하고 있다.

생태계 측면에서는 특정 칩에 종속된 코드를 없애기 위해 Amazon이 PyTorch 본체에 변경 사항을 upstream으로 제공했다. Hugging Face는 자사 도구에서 Trainium의 네이티브 지원을 활성화했으며, TorchNeuron은 분산 학습 프레임워크 TorchTitan과 FSDP도 지원한다. 아울러 vLLM에 대한 대규모 기여와 성능 커널 세트의 오픈소스 공개도 이루어졌다. 빌드마다 수십만 개의 AI 모델로 테스트하고 있다는 설명도 있었다.