Train Qwen3 on AWS Trainium with torchtitan

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

AWS Trainium의 단일 trn2.48xlarge에서 torchtitan으로 Qwen3를 학습하는 세 가지 경로와 Hugging Face 체크포인트 호환, NKI 커널 연동 방법을 설명한다.

📌 핵심 요약

• TorchNeuron은 Trainium을 PyTorch의 neuron 장치로 등록하며, 세 구성은 동일한 torchrun 실행 방식과 설정 구조를 사용한다. 실습 환경은 trn2.48xlarge의 Trainium2 가속기 16개, 물리 NeuronCore 128개이며, LNC=2에서 논리 NeuronCore 64개에 학습 프로세스 64개를 실행한다.
• Transformers v5 백엔드는 AutoModelForCausalLM으로 Qwen/Qwen3-4B-Instruct-2507을 구성하고 Hugging Face safetensors를 읽고 저장한다. 검증 실행의 첫 단계 손실은 공개 가중치에서 6.54, 무작위 초기화에서 12.43이었으며, 내보낸 가중치는 토큰 임베딩과 공유되는 lm_head.weight를 제외하고 누락되거나 예상하지 못한 키 없이 다시 로드됐다.
• torchtitan의 네이티브 Qwen3Model 경로는 qwen3_32b()를 시작점으로 삼되 config.model_spec을 4B와 SDPA 설정으로 교체한다. 테스트된 빌드에는 범용 qwen3_4b() 설정 생성기가 없으며, 예제는 10단계 학습, 로컬 배치 1, 글로벌 배치 64와 64개 rank에 걸친 데이터 병렬 샤딩을 사용한다.
• NKI MoE 예제는 Qwen3-30B-A3B를 단일 trn2.48xlarge에 맞추기 위해 4개 레이어와 전문가 64개로 축소하고 토큰마다 전문가 4개를 선택한다. 데이터 병렬 샤딩은 64개 rank, 텐서 병렬 차수는 1이며 전문가 병렬화는 사용하지 않는다. 예제는 eager 모드로 실행되지만 등록된 NKI 연산자는 torch.compile에서도 지원된다.
• 사용자 정의 NKI 커널 예제는 nkilib의 rope_hf를 PyTorch 연산자로 감싸고 텐서 배치 변환과 역전파를 연결한다. 이는 성능 향상을 입증하는 예제가 아니며, 제공된 본문은 register_autograd 호출 도중에 끝나므로 최종 연결 코드 전체는 확인할 수 없다.

🧩 주요 포인트

1. 공통 torchrun 실행 구조와 config.model_spec 교체 → 실행 틀을 유지하면서 Transformers v5 구현과 torchtitan 네이티브 구현을 선택할 수 있다.
2. Hugging Face safetensors 입출력과 재로딩 검증 → 공개 가중치 기반 미세조정 및 학습 결과의 체크포인트 호환성을 확인하는 근거가 된다.
3. 개 레이어·전문가 64개로 축소한 NKI MoE와 성능 검증 없는 rope_hf 예제 → 제시된 결과의 적용 범위는 축소 구성의 실행과 커널 연동에 한정된다.

🧠 상세 정리

1. Trainium을 표준 PyTorch 장치로 사용하는 세 가지 경로

이 글은 AWS Trainium에서 torchtitan을 기존 PyTorch 사용 방식으로 실행하면서 Qwen3를 학습하는 방법을 설명한다. TorchNeuron이 Trainium을 네이티브 PyTorch 장치로 등록하므로 torchtitan은 표준 PyTorch 인터페이스를 통해 해당 장치를 사용하며, 무작위 초기화 사전학습과 공개 가중치 미세조정을 모두 구성할 수 있다. 첫 번째 경로는 Transformers v5의 모델 구현을 사용하고, 두 번째는 config.model_spec을 바꿔 torchtitan의 네이티브 Qwen3 구현을 선택하며, 세 번째는 torchtitan-neuron의 NKI MoE 커널을 적용한다. 세 구성은 단일 trn2.48xlarge에서 같은 torchrun 실행 패턴과 설정 구조를 공유하고, 글은 이어서 사용자 정의 커널을 연결하는 방법도 다룬다.

2. 실습 환경과 버전 고정

실습에는 Neuron용 AWS Deep Learning AMI 또는 Neuron PyTorch 학습 컨테이너를 사용하는 단일 trn2.48xlarge가 필요하며, 두 환경은 Neuron 드라이버·런타임·컴파일러를 제공한다. 인스턴스에는 Trainium2 가속기 16개와 물리 NeuronCore 128개가 있고, LNC=2에서는 논리 NeuronCore 64개를 사용한다. 재현에 사용된 버전은 torch 2.14.0, torch_neuronx 2.14.3, transformers 5.16.0, torchtitan-neuron release-2.33의 0.2.2+neuron.1이며, 글은 릴리스에 따라 통합 기능과 Trainium 전용 오버라이드 모듈이 달라질 수 있어 버전 고정을 권한다. 설치 전 기존 패키지의 버전과 경로를 확인하고 Neuron용 torch 빌드를 보존하도록 의존성을 설치한 뒤, Hugging Face Hub에서 Qwen/Qwen3-4B-Instruct-2507 자료를 로컬 디렉터리로 내려받아 설정에 전달한다. torch_neuronx가 설치되면 neuron 장치가 등록되고, eager 모드에서는 실행되는 연산을, torch.compile에서는 캡처된 그래프를 Neuron이 컴파일한다.

3. Transformers v5 백엔드로 Qwen3 실행

첫 번째 구성에서는 Transformers v5의 AutoModelForCausalLM이 Qwen3를 생성하고, Hub 모델 정의가 순전파와 역전파에 사용할 구현을 제공한다. 시작점인 transformers_modeling_backend_full()은 Hugging Face 설정 키에 대응하는 필드를 미지정 상태로 두어 AutoConfig.from_pretrained()가 Hub에서 모델 구조를 읽도록 한다. 저장소 루트에 작성하는 neuron_qwen3.py에서 Hub ID와 로컬 모델 자료 경로를 지정하고, c4_test 데이터셋·학습 10단계·로컬 배치 크기 1·데이터 병렬 샤딩 차수 -1을 설정한다. torchrun의 --nproc_per_node=64는 논리 NeuronCore마다 하나씩 프로세스를 실행하며, --module과 --config로 로컬 설정 모듈과 함수를 선택하고 모델은 64개 rank 전체에 걸쳐 샤딩된다.

4. 공개 가중치 로딩과 Hugging Face 체크포인트 왕복 검증

공개 가중치에서 미세조정을 시작하려면 체크포인트 기능을 켜고 initial_load_path에 모델 자료의 절대 경로를 지정하며, Hugging Face 형식의 모델 가중치만 초기 로드하도록 설정한다. 상태 딕셔너리 어댑터는 키 앞의 model. 접두사를 조정하고 공유 임베딩을 처리하며, 마지막 체크포인트는 모델 가중치만 Hugging Face safetensors 형식으로 저장한다. Qwen/Qwen3-4B-Instruct-2507의 검증 실행에서는 첫 학습 단계 손실이 공개 가중치 사용 시 6.54, 무작위 초기화 시 12.43으로 나타나 초기 가중치 로딩을 빠르게 점검하는 근거가 됐다. 학습 후 AutoModelForCausalLM과 load_state_dict로 내보낸 체크포인트를 다시 읽었을 때, 토큰 임베딩과 가중치를 공유하는 lm_head.weight를 제외하고 누락되거나 예상하지 못한 키가 없었으며 글은 이를 표준 Hugging Face 체크포인트 유지의 확인으로 제시한다.

5. torchtitan 네이티브 Qwen3 구현으로 전환

두 번째 구성은 Transformers의 모델 정의 대신 torchtitan의 네이티브 Qwen3Model을 생성한다. 테스트된 torchtitan 빌드에는 범용 qwen3_4b() 설정 생성기가 없어 qwen3_32b()에서 시작하지만, 곧바로 config.model_spec을 model_registry("4B", attn_backend="sdpa")로 교체하므로 실제 선택 대상은 4B 모델과 SDPA다. config.model_spec은 모델 클래스뿐 아니라 torchtitan이 모델을 샤딩하는 방식과 체크포인트 키를 매핑하는 방식도 결정한다. 나머지 설정은 로컬 모델 자료 경로, c4_test 데이터셋, 학습 10단계, 로컬 배치 크기 1, 글로벌 배치 크기 64와 전체 rank에 대한 데이터 병렬 샤딩으로 검증 실행을 재현하며, 실행 명령은 동일한 구조에서 설정 함수 이름을 native_qwen3_4b로 바꾼다.

6. 축소한 Qwen3 MoE에 NKI 커널 적용

세 번째 구성은 qwen3_nki_moe_config()를 사용해 MoE 블록을 NKI의 grouped-experts 구현으로 대체하며, 테스트된 빌드에서는 SDPA와 gather-free rotary embedding도 함께 선택한다. 예제는 Qwen3-30B-A3B를 단일 trn2.48xlarge에 맞추기 위해 레이어 수를 4개, 전문가 수를 64개로 줄이고 top_k=4로 각 토큰을 전문가 4개에 전달한다. 데이터 병렬 샤딩은 64개 rank에 걸쳐 수행하고 텐서 병렬 차수는 1이며 전문가 병렬화는 사용하지 않으므로, num_experts=64를 전문가 병렬 차수로 해석해서는 안 된다. CompileConfig(enable=False)는 eager 실행을 선택하지만 TorchNeuron은 등록된 NKI 연산자를 eager와 torch.compile 양쪽에서 지원하며, torch.compile을 켤 때 기울기 노름 클리핑도 컴파일하려면 CompileConfig.components에 grad_clip을 추가해야 한다.

7. NKI 커널의 PyTorch 연산자 등록과 자동 미분

NKI 커널은 PyTorch 연산자로 등록되어 학습 과정에 들어가며, wrap_nki는 커널을 PyTorch에서 호출할 수 있도록 감싸고 nki_op는 해당 연산자를 등록한다. register_autograd는 순전파 중 저장한 라우팅 메타데이터 등을 이용하는 역전파를 연결하므로, 커널 호출과 학습에 필요한 미분 처리를 함께 구성한다. 글은 모델의 특정 텐서 형태에 맞춰 타일링과 데이터 이동을 조정하거나 성능상 중요한 PyTorch 연산 시퀀스를 하나의 최적화 연산자로 대체하려는 경우 사용자 정의 커널을 고려할 수 있다고 설명한다. 커널은 직접 작성하거나 nkilib에서 가져올 수 있으며, 이어지는 예제는 자체 역전파 모드를 지원하는 nkilib.core.embeddings의 rope_hf로 등록과 자동 미분 연동을 보여주지만 성능 향상을 입증하지는 않는다.

8. rope_hf의 텐서 배치 변환과 역전파 연결

사용자 정의 예제는 저장소 루트의 my_neuron_kernels.py에 rope_hf를 감싸는 연산자를 정의하며, torchtitan의 query와 key 배치인 [batch, seq, heads, head_dim]을 커널이 요구하는 [batch, heads, seq, head_dim]으로 변환한다. 순전파 연산자 rope_fwd는 변환한 텐서를 연속 메모리 형태로 만들고 torch.empty_like()로 출력 텐서를 할당하며, rotary cache를 재구성해 커널에 전달한 뒤 결과를 원래 배치로 되돌린다. 역전파 연산자 rope_bwd는 같은 배치 변환을 기울기에 적용하고 rope_hf를 backward=True로 호출하며, 순전파에서 저장해 둔 rotary cache를 재사용한다. 캐시는 미분 대상이 아니므로 해당 기울기는 None을 반환하도록 구성하지만, 제공된 source_body는 rope_fwd.register_autograd 호출 도중에 끝나므로 최종 등록 코드와 그 이후의 연결 과정 전체는 확인할 수 없다.

🧾 핵심 주장 / 시사점

• 실행 구조를 공유하면서 모델 구현을 바꾸는 방식은 Trainium 학습 환경에서 Transformers v5와 torchtitan 네이티브 모델을 선택하는 변경 지점을 분명하게 보여준다.
• 첫 단계 손실 비교는 공개 가중치 로딩의 빠른 점검 근거이고, 내보낸 가중치의 재로딩은 체크포인트 호환성의 근거이므로 두 검증은 서로 다른 사항을 확인한다.
• 축소된 MoE 구성의 실행과 사용자 정의 커널의 자동 미분 연결만으로 전체 Qwen3-30B-A3B의 실행 가능성이나 커널 성능 향상을 판단할 수는 없다.

✅ 액션 아이템

• 공통 torchrun 실행 구조와 config.model_spec을 기준으로 Transformers v5 및 torchtitan 네이티브 Qwen3Model의 적용 경로 검토.
• Qwen/Qwen3-4B-Instruct-2507의 첫 단계 손실 6.54와 12.43, Hugging Face safetensors 재로딩 결과를 기준으로 가중치 로딩 및 호환성 확인.
• NKI MoE의 4개 레이어·전문가 64개 축소 조건과 rope_hf 예제의 성능 검증 부재, register_autograd 코드 절단을 적용 범위 판단에 반영.

❓ 열린 질문

• Qwen/Qwen3-4B-Instruct-2507의 첫 단계 손실 6.54와 12.43은 다른 실행 조건에서도 공개 가중치 로딩을 구분하는 근거가 되는가?
• 4개 레이어·전문가 64개로 축소한 NKI MoE를 넘어 전체 Qwen3-30B-A3B를 학습하려면 단일 trn2.48xlarge의 수용 조건이 어떻게 달라지는가?
• rope_hf의 최종 register_autograd 연결을 완성했을 때 학습 동작과 성능은 어떻게 검증되는가?