AWS 트레니움 가속기에서 토치타이탄으로 Qwen3 모델 훈련 지원

트레니움 환경 지원 개요

AWS 트레니움 가속기 기반에서 대규모 언어 모델 훈련 라이브러리인 토치타이탄을 활용해 Qwen3 모델을 훈련하는 세 가지 접근법이 정리되었다. 트레니움 환경에서는 토치뉴런 확장을 통해 장치가 파이토치 기본 디바이스로 등록되므로, 개발자는 기존 표준 인터페이스와 분산 실행 도구를 유지하면서 훈련을 수행할 수 있다.

이번에 다뤄진 방식은 허깅페이스 트랜스포머스 버전 5를 백엔드로 활용하는 경로, 토치타이탄 고유의 네이티브 모델 사양을 지정하는 경로, 그리고 뉴런 커널 인터페이스 기반의 전문가 혼합 연산자를 결합하는 경로로 구분된다. 세 구성 모두 단일 대형 인스턴스에서 동일한 분산 실행 패턴과 설정 구조를 공유하도록 설계되었다.

표준 백엔드와 체크포인트 연동

첫 번째 방식인 트랜스포머스 버전 5 백엔드 구성에서는 허깅페이스 허브에 등록된 인과적 언어 모델 정의를 직접 불러와 순전파와 역전파 연산을 수행한다. 이 방식은 허브 모델 식별자를 지정하는 방식으로 동작하며, 공개된 안전 텐서 가중치를 직접 읽고 쓸 수 있는 편의성을 제공한다.

사전 공개 가중치를 기반으로 미세 조정을 수행할 때는 상태 딕셔너리 어댑터가 모델 접두사와 임베딩 바인딩 구조를 적절히 보정해 처리한다. 실제 검증 환경에서 40억 파라미터 지시 튜닝 모델 기준 첫 훈련 단계 손실값은 공개 가중치 사용 시 6.54를 기록했으며, 무작위 초기화 상태의 손실값 12.43과 뚜렷한 차이를 보였다. 훈련 종료 후 내보낸 가중치 역시 허깅페이스 표준 형식과 호환되어 추가 누락 키 없이 다시 로드할 수 있다.

네이티브 모델과 분산 병렬화

두 번째 방식에서는 토치타이탄 고유의 Qwen3 모델 구현체를 사용하도록 모델 사양을 교체하여 구동한다. 설정에서 40억 파라미터 사양과 스케일드 닷 프로덕트 어텐션 백엔드를 선택함으로써, 토치타이탄 내부의 모델 분할 규칙과 체크포인트 키 매핑 방식을 일관되게 적용할 수 있다.

단일 트레니움2 인스턴스는 16개의 가속기와 128개의 물리 코어를 포함하며, 논리 코어 배율 설정에 따라 64개의 논리 뉴런 코어로 인식된다. 데이터 병렬 샤딩 수준을 전체 코어로 지정하면 64개 훈련 프로세스가 모델 매개변수와 그래디언트, 옵티마이저 상태를 완전히 나누어 관리한다. 여기에 텐서 병렬화를 결합하면 다차원 장치 메시 상에서 레이어별 가중치를 분할하고 시퀀스 병렬 처리를 동시에 수행할 수 있다.

NKI 커널 및 사용자 정의 연산

세 번째 방식은 전문가 혼합 모델 구조를 위해 뉴런 커널 언어로 작성된 고성능 연산자를 결합하는 기법이다. 토치타이탄 뉴런 레시피에서 제공하는 설정을 통해 그룹화된 전문가 연산자를 배치하고, 토큰별 전문가 라우팅 연산을 단일 인스턴스 자원에 맞춰 분할 실행할 수 있다. 이 연산자는 즉시 실행 모드와 그래프 컴파일 모드 모두에서 등록된 파이토치 연산자로 동작한다.

개발자는 텐서 형태나 메모리 이동을 특정 모델에 맞추기 위해 뉴런 라이브러리의 커널을 가져오거나 직접 구현하여 덮어쓸 수 있다. 파이썬 데코레이터와 오버라이드 레지스트리를 사용하면 특정 레이어 경로의 연산자만을 정밀하게 교체할 수 있으며, 전방 계산 시 회전 캐시를 보존하여 역전파 자동미분 함수와 연결하는 방식으로 일관된 수치 정밀도를 유지한다.

시스템 요구 사양과 제약 사항

해당 훈련 절차를 재현하려면 트레니움2 가속기 인스턴스와 함께 뉴런 전용 딥러닝 머신 이미지 또는 뉴런 파이토치 컨테이너 환경이 필수적으로 요구된다. 드라이버와 런타임, 컴파일러 버전을 고정해야 하며, 트랜스포머스 백엔드 및 오버라이드 모듈은 릴리스 버전에 따라 동작 차이가 발생할 수 있다.

또한 사용자 정의 커널을 적용할 때는 시퀀스 길이가 특정 배수 조건을 충족해야 하는 메모리 정렬 제약이 따르므로 사전 유효성 검증이 필요하다. 전문가 혼합 모델의 경우 단일 인스턴스 메모리 한계에 맞춰 레이어 수와 활성 전문가 수를 조절해야 하며, 모델 규모나 배치 크기 확장에 따른 코어별 고대역폭 메모리 점유율을 면밀히 감시해야 한다.