허깅페이스, AWS 트레니움 네이티브 백엔드 '토치뉴런' 지원 발표

트레니움 지원의 대대적 개편

허깅페이스는 AWS 트레니움 환경을 위한 파이토치 신규 네이티브 백엔드인 토치뉴런을 공개했다. 이번 개편으로 기존 파이토치 및 허깅페이스 핵심 라이브러리인 트랜스포머스, 디퓨저스 등이 트레니움을 일반 가속기처럼 직접 지원하게 된다. 토치뉴런은 2026년 10월 20일부터 패키지 설치 도구를 통해 공개 베타로 배포된다.

새로운 백엔드가 도입됨에 따라 기존에 쓰이던 전용 래퍼인 옵티멈 뉴런 대신 파이토치 기본 생태계를 그대로 활용하는 방식으로 전환된다. 개발자는 모델 실행을 위해 하드웨어 전용 라이브러리를 거칠 필요가 없으며, 즉시 실행 모드와 토치 컴파일, 텐서 병렬화 및 완전 분산 데이터 병렬 처리를 네이티브 인터페이스로 다룰 수 있다.

컴파일러 중심 구조의 한계 극복

지금까지 AWS 트레니움과 인퍼런시아 같은 인공지능 전용 칩은 파이토치 XLA에 기반한 패키지와 옵티멈 뉴런에 의존해 왔다. 이 방식은 지연 텐서 연산과 정적 형상 제약, 그래프 중단 시마다 발생하는 잦은 재컴파일 등 컴파일러 구조에 맞춘 복잡한 프로그래밍 사고를 강제하는 구조였다.

허깅페이스는 이러한 복잡도를 감추기 위해 전용 모델링 클래스와 트레이너 코드를 독자적으로 유지 관리해야 했다. 토치뉴런은 뉴런 칩을 파이토치의 네이티브 디바이스로 직접 노출하는 플러그인 방식을 택해, 이러한 중간 추상화 계층 없이도 가속기를 범용 하드웨어처럼 다룰 수 있게 만들었다.

생성 모델 지원율과 추론 최적화

트랜스포머스 라이브러리에서 텍스트 생성 모델의 네이티브 지원 범위가 크게 늘어났다. 기존 옵티멈 뉴런 0.4.5 버전에서 상위 100대 텍스트 생성 모델 아키텍처 중 46퍼센트만 지원되던 것과 비교해, 토치뉴런 환경에서는 지원율이 89퍼센트까지 향상됐다. 장치 설정 인자만 지정하면 모델이 뉴런 장치로 자동 배치된다.

추론 단계에서는 빈번한 재컴파일을 방지하기 위해 정적 캐시 사용이 권장된다. 텐서 병렬화 기능도 기본적으로 결합할 수 있어 여러 뉴런 코어로 모델을 나누어 처리하는 작업이 수월해졌다. 디퓨저스 라이브러리 역시 플럭스 1과 2, 큐원 이미지 계열 등 다양한 모델 파이프라인에서 즉시 실행과 컴파일, 텐서 병렬화를 지원한다.

단일 스크립트 기반 분산 훈련

모델 훈련 환경에서도 트레니움 전용 트레이너나 하드웨어별 모델 코드 분기가 불필요해졌다. 트랜스포머스, 액셀러레이트, PEFT, TRL 같은 표준 도구를 그래픽 처리 장치 환경과 동일한 형태로 사용할 수 있다. 분산 설정에 텐서 병렬화와 완전 분산 크기를 지정하면 가중치를 적재하는 과정에서 모델이 텐서 병렬과 완전 분산 그리드에 맞추어 자동으로 분할된다.

실제 검증에서는 32개 뉴런 코어로 구성된 트레니움2 칩 8개 환경에서 텐서 병렬 크기 4와 분산 크기 8의 그리드 분산 훈련이 정상 완주됐다. 훈련이 끝난 뒤 저장되는 체크포인트는 하드웨어에 종속되지 않는 표준 트랜스포머스 형식이어서 허깅페이스 허브에 그대로 등록하거나 다른 장치에서 바로 불러올 수 있다.

적용 조건과 주요 제약 사항

개선된 구조에도 불구하고 트레니움 아키텍처 특성상 준수해야 할 제약 사항이 존재한다. 우선 토치 컴파일에서 제공하는 자동 동적 형상 컴파일 기능은 아직 지원되지 않으므로 비활성화해야 한다. 입력 배치 길이를 특정 단위 배수로 맞추고 시퀀스를 묶는 패킹 방식을 사용해 입력 형상의 가짓수를 제한해야 컴파일 그래프를 재사용할 수 있다.

또한 트랜스포머스에 내장된 비동기 가중치 적재 기능은 현재 뉴런 백엔드에서 지원되지 않아 환경 설정을 통해 명시적으로 꺼두어야 한다. 더 세밀한 연산 최적화가 필요한 경우 뉴런 커널 인터페이스를 이용해 저수준 제어가 가능한 커스텀 커널을 직접 작성할 수 있으며, 비전 트랜스포머 작업에서 최대 2.4배의 가속을 달성한 사례가 보고됐다.