파이토치로 짠 AI 모델 코드를 다시 쓰지 않고 자체 설계 AI 가속기 AWS 트레이니움(AWS Trainium)에서 돌릴 수 있는 길이 열렸다. 아마존웹서비스(AWS)는 27일 이를 지원하는 오픈소스 '토치뉴런(TorchNeuron)'을 공개했다. 토치뉴런은 아파치 2.0(Apache 2.0) 라이선스로 배포되며 트레이니움2(Trainium2)와 트레이니움3(Trainium3)를 지원한다.
토치뉴런은 파이토치 코드에서 하드웨어를 지정하는 설정을 바꾸면 트레이니움에서 실행되도록 설계됐다. 파이토치에 내장된 하드웨어 등록 표준 메커니즘과 연결되는 방식으로 작동한다.
기술 사양: 설정 변경으로 실행, 분산 학습까지
AWS가 진행한 비교 시연에서는 동일한 학습 스크립트를 기존 칩과 트레이니움에서 각각 실행했고, 학습 내용과 속도가 같았다. 이는 해당 시연 환경의 결과이며, 모든 모델과 하드웨어 조합에서 같은 성능을 보장하는 것은 아니다.
분산 학습도 지원한다. 파이토치의 네이티브 분산 학습 프레임워크인 토치타이탄(TorchTitan)과 통합돼 있고, 모델 매개변수와 학습 상태를 여러 장치에 나누는 FSDP(완전 샤딩 데이터 병렬 처리)로 여러 트레이니움 인스턴스에 걸친 분산 학습을 할 수 있다.
AWS는 빌드마다 수십만 개 AI 모델로 토치뉴런을 테스트하고 있다고 설명했다. 추론 수요가 적은 시간대의 컴퓨팅 용량을 학습으로 전환해 활용도를 높일 수도 있다는 것이 회사 설명이다.
개발 부담과 생태계: 이전 비용을 줄이는 구조
AWS에 따르면 파이토치를 쓰는 AI 연구자와 개발자는 약 85%다. 새로운 AI 가속기에서 실행하려면 코드 수정과 성능·정확성 재검증이 필요해 엔지니어링 부담이 컸다. 토치뉴런은 이런 이전 부담을 줄이고 기존 파이토치 환경을 유지한 채 트레이니움을 고를 수 있게 하려는 것이다. 하드웨어 전환 때 코드 변경을 최소화해 개발 기간과 엔지니어링 비용을 줄이는 것이 목표다.
트레이니움 관련 활용과 검토 사례도 있다. 음성 AI 기업 카르테시아(Cartesia)는 트레이니움의 파이토치 네이티브 지원을 활용해 기존 파이토치 코드를 최소한만 수정한 채 표준 추론 파이프라인을 실행했다. AI 모델 공유 플랫폼 허깅페이스(Hugging Face)는 자사 AI 개발 도구에 트레이니움 네이티브 지원을 적용했다. 조현병 등 신경정신질환 치료법을 연구하는 리버 뇌발달연구소(Lieber Institute for Brain Development)는 분자동역학 시뮬레이션을 AI 도구로 가속하거나 대체하는 방안을 트레이니움에서 검토하고 있다. AWS는 스타트업과 대기업이 코드 재작성을 최소화하며 트레이니움 기반 개발을 진행하고 있다고 설명했다.
밸류체인 위치: 자체 칩에서 오픈소스 소프트웨어까지
추론 분야에서는 AWS가 LLM 추론·서빙에 쓰이는 오픈소스 프로젝트 vLLM에 기여하고 있으며, 트레이니움 성능 최적화 커널을 오픈소스로 공개했다.
아마존은 반도체 설계 기업 안나푸르나랩스(Annapurna Labs)를 2015년에 인수했고, 그 뒤 클라우드 워크로드에 맞춘 반도체 개발에 나섰다. 지난 10년간 AWS가 서로 다른 클라우드 워크로드용으로 직접 설계한 칩 제품군은 4개다.
AWS는 앞으로 파이토치 커뮤니티와 협력하고 변경 사항을 업스트림에 기여할 계획이다. 특정 하드웨어에 종속되는 코드를 줄여 하드웨어 중립적 프레임워크 발전을 지원한다는 구상이다.