심층 해설 · 약 5분 읽기
Trainium에서 Qwen3 학습하기: 가중치와 커널을 검증하는 torchtitan의 세 경로
AWS의 10월 21일 가이드는 Trainium2 인스턴스 한 대에서 Transformers, 네이티브 Qwen3, NKI 커널을 연결한다. 속도나 비용 우위를 입증한 자료는 아니다.
- 사이트 게시일
- 원문 게시일
- 출처 확인일
AWS가 10월 21일 공개한 기술 가이드는 Trainium에서 torchtitan으로 Qwen3를 학습하는 세 가지 구성을 설명한다. Transformers v5 백엔드, torchtitan 네이티브 모델, NKI 커널을 쓰는 축소 혼합 전문가 예제다. 유용한 진전은 체크포인트와 수치 검증을 포함한 구체적인 통합 절차다. 새로운 Qwen 모델 출시나 더 저렴한 학습을 증명한 벤치마크는 아니다.
뉴스 듣기
이 기사를 읽어 드리거나 뉴스를 계속 재생합니다.
읽기를 준비하는 중…
기기 음성
읽기는 브라우저나 기기에서 처리합니다. AISubHub는 내용을 오디오 API로 보내지 않습니다. 사용 가능한 음성과 잠금 화면 재생은 기기에 따라 다릅니다.
계속 재생할 때 짧은 구절이 반복될 수 있습니다.
핵심 요약
- 예제는 trn2.48xlarge 한 대를 사용한다. Trainium2 가속기 16개, 물리 NeuronCore 128개이며 LNC=2에서 논리 NeuronCore 64개를 제공한다. 실행 시 학습 프로세스 64개를 시작한다.
- 밀집 모델 예제는 Qwen3-4B-Instruct-2507 파일과 c4_test 학습 10단계를 사용한다. MoE 예제는 4개 레이어와 전문가 64개로 축소하고 토큰마다 전문가 4개를 고른다. 수정하지 않은 전체 30B-A3B 모델을 학습하는 예제가 아니다.
- AWS가 보고한 첫 단계 손실은 공개 가중치에서 6.54, 무작위 초기화에서 12.43이다. 사용자 정의 회전 위치 임베딩의 float32 테스트에서는 최대 절대 차이가 0.0이었다. 두 수치는 서로 다른 문제를 확인한다.
하나의 장치 통합 위에서 모델 구현을 세 가지로 선택한다
TorchNeuron은 Trainium을 PyTorch의 네이티브 neuron 장치로 등록해 torchtitan이 일반 인터페이스를 사용하게 한다. 첫 번째 경로에서는 Transformers v5가 AutoModelForCausalLM으로 모델을 만든다. 두 번째는 config.model_spec으로 torchtitan 자체 Qwen3 구현을 선택한다. 세 번째는 혼합 전문가 레이어에 Trainium 전용 NKI 연산자를 연결하고 사용자 정의 대체도 보여준다. 공통 학습 진입점 뒤의 구현 선택이지 새로 출시된 Qwen 모델 세 개가 아니다.
가이드는 torch 2.14.0, torch_neuronx 2.14.3, Transformers 5.16.0, torchtitan-neuron release-2.33의 0.2.2+neuron.1 버전을 고정한다. 드라이버, 런타임, 컴파일러가 포함된 Neuron Deep Learning AMI 또는 학습 컨테이너가 필요하다. 익숙한 PyTorch 코드가 통합 범위를 좁혀준다고 해석할 수는 있지만 시험한 의존성 조합이 선택 사항이 되는 것은 아니다. AWS도 Transformers 통합과 대체 모듈은 torchtitan 릴리스마다 바뀔 수 있다고 명시한다.
4B 예제는 초기화와 체크포인트 왕복을 확인한다
첫 구성은 Qwen/Qwen3-4B-Instruct-2507을 지정하고 c4_test, 10단계, 로컬 배치 크기 1을 설정한다. 데이터 병렬 상태는 시작한 64개 랭크에 분할한다. 아키텍처를 읽는 것과 공개 가중치를 읽는 것은 별개다. 미세 조정 구성은 내려받은 파일에서 체크포인트를 읽도록 명시적으로 활성화하며 최종 모델을 Hugging Face safetensors로 내보낸다. 보고된 첫 단계 손실은 공개 가중치에서 6.54, 무작위 초기화에서 12.43이다.
이 차이는 초기화 검증에 유용하지만 10단계 후 더 좋은 비서가 만들어지거나 최종 수렴했다는 증거는 아니다. AWS는 내보낸 체크포인트를 Transformers로 다시 읽어 키를 검사하고 임베딩과 묶인 lm_head.weight를 고려한다. 네이티브 경로가 qwen3_32b 빌더에서 시작하는 이유는 시험한 버전에 일반적인 4B 빌더가 없기 때문이다. 이후 모델 명세를 4B로 바꾸고 전역 배치 크기를 64로 설정한다. 빌더 이름만으로 32B 모델을 학습했다고 판단해서는 안 된다.
MoE 구성은 모델 계열 이름보다 의도적으로 작다
Qwen3-30B-A3B 예제는 인스턴스 한 대에 맞추기 위해 레이어 4개와 전문가 64개로 축소된다. top_k=4는 토큰마다 전문가 4개로 보내는 설정이다. 원문은 64개 랭크의 데이터 병렬 분할, 텐서 병렬도 1을 사용하며 전문가 병렬은 사용하지 않는다. 전문가 수는 모델 내부 설정이지 전문가 병렬 장치 수가 아니다. 이를 구분해야 작은 통합 예제를 원래 전체 모델의 학습 결과로 확대하지 않을 수 있다.
이 예제는 구성에서 컴파일을 끈 eager 모드로 실행된다. 등록한 NKI 연산자는 torch.compile에서도 지원하지만 실행 모드 변경은 별도로 결정할 문제다. 시험한 버전의 구성은 MoE 블록을 그룹 전문가 커널로 바꾸고 스케일 점곱 어텐션과 gather 없는 회전 위치 임베딩을 선택한다. AWS는 순전파 중 저장한 정보를 이용해 연산자에 autograd를 연결하는 방법을 설명한다. 이 특정 구성이 더 빠르다고 입증하는 처리량 비교는 제시하지 않는다.
사용자 정의 커널은 배치 순서, 기울기, 적용 대상을 검증해야 한다
회전 위치 임베딩 예제는 커널 등록이 단순한 함수 호출 이상임을 보여준다. torchtitan의 query와 key 순서는 배치, 시퀀스, 헤드, 헤드 차원이지만 rope_hf는 배치, 헤드, 시퀀스, 헤드 차원을 기대한다. 래퍼는 순서를 바꾸고 출력 공간을 할당한 뒤 원래대로 변환한다. 별도 역전파 연산자는 저장한 회전 캐시를 재사용하며 캐시 자체에는 기울기를 계산하지 않는다. 시퀀스 길이는 128 곱하기 LNC의 배수여야 하며 LNC가 2이면 256의 배수다.
대체 레지스트리는 이후 의도한 모듈 구성을 교체한다. 시작 로그에는 Qwen3 레이어마다 nki_rope 대체가 하나씩 보여야 하며 더 좁은 이름 패턴으로 특정 레이어만 지정할 수 있다. AWS는 PyTorch 구현과 float32로 비교했을 때 query와 key의 최대 절대 차이가 모두 0.0이었다고 보고한다. 이는 그 수치 테스트를 검증할 뿐 모든 자료형, 형태, 학습 실행의 동등성을 보장하지 않는다. 원문도 성능 향상이 아니라 커널 및 자동 미분 통합 예제임을 명확히 한다.
가이드를 측정 가능한 파일럿으로 바꾸기
첫 단계에는 컴파일이 포함돼 이후 단계보다 오래 걸린다. 로그는 손실, 기울기 노름, 고대역폭 메모리 사용량, 처리량, 모델 FLOPs 활용률을 보여주지만 가이드는 통제된 전체 비용 비교를 제공하지 않는다. 모델 크기, 시퀀스 길이, 로컬 배치가 커질수록 랭크별 메모리가 실질적인 제약이다. 다차원 DeviceMesh에서 텐서 병렬과 완전 분할 데이터 병렬을 결합할 수 있지만 제시된 MoE 구성을 그런 대규모 실험과 동일시해서는 안 된다.
AISubHub가 제안하는 순서는 시험한 버전을 유지하고 모델과 랭크 64개를 확인하며 체크포인트 로딩과 내보내기를 검증한 뒤, 긴 학습 전에 사용자 정의 커널 출력을 비교하는 것이다. 이후 준비 구간과 지속 실행을 나눠 측정하고 목표 데이터에서 수렴을 평가해야 한다. 임대한 하드웨어를 중지하거나 종료하기 전에 필요한 체크포인트를 저장해야 한다. 이는 가이드에서 도출한 평가 절차이지 우리가 재현한 결과가 아니며, 구독 권리나 무료 학습량 또는 Trainium의 보편적 우위를 뜻하지 않는다.
중요한 이유
AISubHub 분석: 팀은 본격적인 학습에 앞서 모델 호환성, 체크포인트 이식성, 커널 정확성을 작은 검증 과제로 나눌 수 있다. 공통 실행 구조는 구성 변경을 줄이지만 버전, 메모리, 컴파일, 수치 조건까지 없애지는 않는다.
제한 및 이용 가능성
AWS의 Josh Longenecker와 Alex Hoffman이 작성한 기술 가이드이며 AISubHub의 하드웨어 시험이 아니다. 비교 가능한 전체 학습 비용, 최종 수렴 결과, 사용자 정의 커널의 입증된 가속 효과는 제공하지 않는다. 인스턴스 가용성과 요금, 정확한 소프트웨어 버전, 작업 부하 검증은 운영자가 확인해야 한다. 소비자 구독 혜택도 발표하지 않았다.
출처와 근거
원문 게시일: · 출처 확인일:
AI 보조 편집 기사를 명시한 일차 출처와 대조합니다. 업체 주장은 출처를 밝히며 게재는 추천을 뜻하지 않습니다.