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