AWS Trainium 모델 지원: Hugging Face, torch-neuronx 호환성 조회 도구 출시

요약 — Hugging Face와 AWS가 torch-neuronx 모델 지원 Space를 출시했습니다. 이 공개 조회 도구는 Hugging Face Hub의 특정 Transformers 모델이 AWS Trainium에서 수정 없이 실행되는지 여부를 알려주고, 테스트에 사용된 정확한 스크립트와 스택을 제공합니다. "이 모델이 Trainium에서 작동하나?"라는 시행착오 질문을 쿼리 가능한 데이터셋으로 바꿔줍니다.

이 Space는 torch-neuronx 패키지에 포함된 AWS Trainium용 네이티브 PyTorch 백엔드인 TorchNeuron을 기반으로 합니다. TorchNeuron은 이전 Hugging Face 게시물에서 소개되었습니다. TorchNeuron을 사용하면 Transformers는 다른 장치에서와 마찬가지로 Trainium에서 자체 모델링 코드를 실행하므로, 원칙적으로 Hub의 모든 Transformers 모델이 실행될 수 있습니다. 새 Space는 Hub에서 가장 인기 있는 모델 중 실제로 실행되는 모델이 무엇인지 답변합니다.

발표 내용

aws-neuron/torch-neuronx-model-support Space는 AWS Trainium용 공개 브라우징 호환성 테이블입니다. 각 모델에 대해 다음을 보고합니다:

모델이 Trainium에서 수정 없이 실행되는지 여부 (PASS 결과)
모델이 분할된 Trainium 코어 수
컴파일을 포함한 실행 벽시계 시간
최대 장치 메모리 사용량

각 모델에 대해 Space는 실행된 스크립트도 공개하므로 결과를 신뢰하는 대신 재현할 수 있습니다.

모델 조회 방법

Space는 선택됨 보기로 열리며, 작성자가 Hub의 다운로드 수와 트렌드 점수 기준 상위 모델 중에서 직접 선택한 모델로 설명됩니다.

트렌드 및 다운로드 보기는 Hub의 자체 순위를 따르며, 한 번에 하나의 작업을 처리합니다. 이 보기에는 테스트 여부와 관계없이 모델이 나열되므로 누락된 항목은 결코 모호하지 않습니다. 모델이 의도적으로 제외된 경우 Space는 그 이유를 명시합니다.

작성 시점 기준으로 Transformers 모델만 테스트되었으며, 선택 범위는 text-generation 및 image-text-to-text 모델로 제한됩니다. diffusers, sentence transformers 및 기타 transformers 작업으로 확장할 계획입니다.

Space에서는 Neuron 컴파일러(neuronx-cc) 및 Transformers 버전을 선택하여 결과를 단일 날짜 없는 스냅샷이 아닌 스택 버전별로 필터링할 수 있습니다.

특정 모델을 찾으려면 트렌드 또는 다운로드 보기로 전환하고 이름의 일부를 입력하세요. Space는 모델이 제외된 이유를 보고합니다. 대부분 우선순위가 낮아서이며, 때로는 Trainium과 실제로 호환되지 않기 때문입니다.

직접 모델 실행 방법

Space에서 모델을 선택하면 실행 방법이 표시됩니다. Space는 모델이 테스트된 Neuron 스택, 해당 스택 위에 설치할 패키지, 실제로 실행된 run.py 스크립트를 나열합니다.

스크립트는 TorchNeuron 발표의 레시피를 따릅니다. from_pretrained로 모델을 로드하고, 모델이 필요할 때 분할한 다음, 정적 캐시와 torch.compile(backend="neuron")로 생성합니다.

해당 스택이 포함된 Trainium2(trn2) 인스턴스(예: AWS Neuron Deep Learning Container 내부)에서 나열된 패키지를 설치한 후 다음을 실행합니다:

python3 run.py

torchrun 호출은 필요하지 않습니다. 모델이 분할된 경우 스크립트는 자체 torch.distributed 랭크를 시작합니다.

NEFF 캐시로 컴파일 건너뛰기

첫 실행은 모델의 그래프 컴파일로 시작하며, 보고된 벽시계 시간의 대부분이 여기에 소요됩니다. 이 단계를 건너뛰려면 테스트 중 동일한 컴파일러가 빌드한 그래프를 공개 NEFF 캐시 버킷에서 복사할 수 있습니다. Space는 올바른 컴파일러 버전이 채워진 명령을 제공합니다:

hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

캐시를 사용하면 원래 테스트와 동일한 neuronx-cc 버전을 사용하는 한 재컴파일을 건너뛸 수 있습니다.

"실행"의 의미

호환성 주장은 좁고 구체적이며 Space는 이를 정확히 문서화합니다.

각 테스트는 trn2 인스턴스에서 모델을 로드하고, 가중치를 담을 수 있는 가장 적은 수의 Trainium 코어로 분할합니다. 대부분의 모델은 텐서 병렬 처리로 분할되며, 텐서 병렬 계획이 없는 mixture-of-experts 모델은 전문가 병렬 처리만 사용합니다. 그런 다음 테스트는 32개의 토큰을 생성하고 확인합니다.

정확성 검사는 분할 구성에 따라 다릅니다:

단일 코어에서 생성된 토큰은 동일한 체크포인트의 CPU 실행 결과와 비교됩니다.
모델이 분할된 경우 작은 판별 모델이 생성된 텍스트가 유창한 영어인지 확인합니다.

PASS의 범위는 의도적으로 제한됩니다. 작성 시점 기준으로 Trainium2 장치 4개에 맞는 모델만 테스트되며, 이는 최대 16코어를 의미합니다. PASS는 처리량이나 더 긴 생성에 대해 아무것도 말하지 않습니다.

결과는 데이터셋입니다

Space에 표시되는 모든 것은 공개 aws-neuron/torch-neuronx-model-support 데이터셋에서 비롯되며, 각 neuronx-cc 및 라이브러리 버전 쌍에 대해 하나의 Parquet 파일을 저장합니다. 각 행에는 모델의 결과, 실패 시 오류, 전체 스택 버전, 실행된 run.py가 포함됩니다.

렌더링된 테이블이 아닌 데이터셋이므로 결과를 직접 쿼리할 수 있습니다:

import pandas as pd

df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]

데이터셋 카드는 각 결과 값의 정확한 의미를 정의하므로 PASS 또는 실패를 추측이 아닌 정확한 방법론에 따라 해석할 수 있습니다.

적용 범위 제한 및 향후 계획

현재 Space가 증명하는 범위를 제한하는 네 가지 제약이 있습니다:

Transformers 모델만 테스트됩니다.
text-generation 및 image-text-to-text 작업만 다룹니다.
Trainium2 장치 4개(16코어)에 맞는 모델만 테스트됩니다.
PASS는 32개 생성 토큰에 대한 정확성을 보고하며, 처리량이나 장기 컨텍스트 동작은 보고하지 않습니다.

명시된 계획은 diffusers, sentence transformers 및 기타 Transformers 작업으로 적용 범위를 확장하는 것입니다.

이것이 중요한 이유

하드웨어 이식성 질문은 일반적으로 암묵적 지식, 임시 이슈 또는 비용이 많이 드는 시험 실행으로 답변됩니다. Trainium 호환성을 버전이 지정된 쿼리 가능한 데이터셋으로 게시하고 재현 스크립트를 첨부하면 이 질문의 경제성이 바뀝니다. 답은 조회이고 증거는 검사 가능합니다.

두 가지 설계 선택이 이 리소스를 리더보드보다 더 유용하게 만듭니다. 첫째, 결과는 특정 neuronx-cc 및 라이브러리 버전에 연결되므로 실제 배포할 스택에 대해 호환성을 평가할 수 있습니다. 둘째, 데이터셋은 성공과 함께 실패 및 이를 생성한 오류를 기록하므로 이전에 문서화되지 않은 비호환성을 검색할 수 있습니다.

실용적인 요점: AWS Trainium에서 모델을 실행하기 전에 aws-neuron/torch-neuronx-model-support Space를 확인하거나 Parquet 데이터셋을 직접 쿼리하고 게시된 run.py 및 NEFF 캐시 항목을 재사용하여 trn2 인스턴스에서 결과를 재현하세요.