허깅페이스의 언어 모델 학습 도구가 아마존 전용 학습 칩을 지원하고, 비동기 강화학습에서 반복되는 입력의 중복 계산을 줄였다. 학습 환경을 넓히려는 개발자와 같은 질문에서 여러 답변을 생성해 학습하는 팀이 각각 살펴볼 변화다.

1. Trainium에서 기존 학습 코드를 쓴다

2026년 10월 15일 공개된 TRL v1.16.0은 AWS Trainium 지원과 사용 안내를 추가했다. 릴리스 설명에 따르면 앞선 버전의 통합 언어 모델 출력 계층을 바탕으로 남은 지원 공백을 보완했다.

사용 안내는 torch-neuronx를 설치한 뒤 같은 환경에 TRL을 설치하면 기존 트레이너 코드를 바꾸지 않고 실행할 수 있다고 설명한다. 아래는 지도 파인튜닝, 즉 예시 답변을 따라 모델을 학습시키는 문서의 빠른 시작 코드다.

from datasets import load_dataset
from trl import SFTTrainer

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

2. 길이를 맞춰 재컴파일을 줄인다

Neuron은 입력 모양이 달라질 때마다 계산 그래프를 컴파일하며, 한 번에 약 1분이 걸린다. 이번 버전은 pad_to_multiple_of의 기본값을 자동 선택으로 바꿨다. Neuron에서는 입력 길이를 128의 배수로 채워 서로 다른 모양을 줄이고, GPU에서는 기존 동작을 유지한다.

Trainium2에서 Qwen3-0.6B를 뉴런 코어 하나, 배치 크기 4로 지도 파인튜닝한 결과, 20스텝에 기본값은 130초, 길이 맞춤을 끄면 260초가 걸렸다. 이는 해당 예제의 비교이지 다른 모델의 학습 속도까지 보장하는 수치는 아니다.

모델 컴파일은 선택 사항이다. torch_compile=True를 쓰려면 Dynamo의 자동 동적 모양 기능을 꺼야 한다. 문서의 같은 모델·배치 크기·코어 수에서 첫 컴파일은 약 2분이며, 이후 스텝당 시간은 컴파일하지 않을 때 1.71초에서 1.15초로 줄었다.

3. 같은 앞부분은 한 번만 계산한다

AsyncGRPO의 트리 어텐션은 작은 학습 묶음의 입력을 접두사 트리로 묶는다. 같은 질문에서 갈라진 답변을 나뭇가지처럼 연결해, 공유하는 앞부분을 답변마다 되풀이하지 않고 한 번만 계산하는 방식이다.

내부적으로 갈림길 없는 부분은 구간으로 합치고, 각 구간은 자신에게 이어지는 경로의 키와 값을 모아 참조한다. 어텐션 외 계층은 고유 토큰마다 한 번 실행하고, 공유 경로에서 키와 값만 중복된다.

릴리스의 Qwen3-1.7B 측정은 생성 결과당 32K 토큰의 순방향·역방향 계산을 합친 것이다. 공유 프롬프트, 패킹 비율 3.37 조건에서 기존 시퀀스 패킹은 8179ms, 트리 패킹은 2755ms로, 제시된 속도 배수는 2.97배다. 최대 메모리는 35.2GiB에서 13.6GiB로 줄었다. 측정 하드웨어는 제시되지 않았다.

별도의 패킹 선택 옵션은 없다. 공유되는 부분이 없으면 기존 시퀀스 패킹과 같은 형태가 되며, 릴리스는 이때 추가 비용을 약 1%로 설명한다.

4. 온라인 생성은 캐시와 자료형을 확인한다

이 환경에서 학습 중 답변을 생성하는 온라인 트레이너는 vLLM을 쓸 수 없어 Transformers로 생성한다. 기본 동적 캐시는 토큰마다 입력 모양을 바꾸므로, 문서는 생성 중 모양을 고정하는 정적 캐시 사용을 안내한다.

현재 정적 캐시는 bf16 모델을 요구한다. 기본 fp32 모델에서는 키와 값의 자료형이 달라지는 문제가 보고됐다. 생성 전 프롬프트 길이도 맞추지만, 새로운 캐시 크기가 나타나면 다시 컴파일해야 한다.

데이터 병렬 실행에도 코어 배치를 확인해야 한다. Trainium2 칩 하나를 담은 trn2.3xlarge의 기본 설정은 물리 코어 8개를 논리 코어 4개로 묶는다. 문서는 이 구성에서 Accelerate로 4개 프로세스를 실행한 지도 파인튜닝이 코어 하나보다 3.6배 빠르다고 설명한다. 일부 코어만 고를 때는 문서가 허용한 조합을 따라야 한다.

5. 절약되는 계산과 준비 비용을 나눠 본다

두 변화는 서로 다른 비용을 줄인다. 트리 방식은 공유 입력의 중복 연산과 메모리를 줄이고, Trainium의 길이 맞춤은 여분의 채움 토큰을 처리하는 대신 재컴파일할 모양을 줄인다. 모델 컴파일을 선택하면 시작 대기 시간을 감수하고 이후 반복 실행 시간을 줄이는 구조다.

트리 어텐션 성능 측정과 Trainium 사용 예제는 별도로 제시됐다. 두 기능을 결합한 성능 수치는 제시되지 않았으므로 각각의 개선 폭을 합쳐 해석해서는 안 된다.

이번 업데이트는 학습 하드웨어 선택지와 공유 입력 처리 효율을 넓히지만, 도입 판단은 입력의 공유 정도와 컴파일 준비 비용을 따로 따져야 한다.