이전 게시물에서 우리는 AWS Trainium을 위한 네이티브 PyTorch 백엔드인 TorchNeuron(패키지 torch-neuronx)을 소개했습니다. 이를 통해 Transformers는 다른 모든 장치에서와 마찬가지로 자체 모델링 코드를 Trainium에서 실행할 수 있으므로, 허브에 있는 모든 Transformers 모델이 실행될 수 있습니다. 실제로 어떤 모델들이 실행될까요?
torch-neuronx 모델 지원 Space는 허브의 가장 인기 있는 모델들에 대해 이 질문에 답합니다. 이 공간은 특정 모델이 수정 없이 Trainium에서 실행되는지 여부를 알려주며, 직접 실행할 수 있는 스크립트도 제공합니다.
Space는 Selected(선택됨) 보기로 열립니다. 여기서는 다운로드 횟수와 트렌딩 점수를 기준으로 허브의 상위 모델들 중에서 저희가 선별한 모델들을 보여줍니다. 각 행은 결과, 모델이 분할된 Trainium 코어 수, 실행 시간(컴파일 포함), 그리고 최대 장치 메모리를 제공합니다.
특정 모델을 확인하려면 Trending(트렌딩) 또는 Downloads(다운로드) 보기로 전환하여 이름의 일부를 입력하세요. 이 보기들은 허브 자체 순위를 따라가며, 모델이 테스트되었는지 여부와 관계없이 한 번에 하나의 작업에 대해 작동합니다. 만약 특정 모델이 의도적으로 제외된 경우, Space는 그 이유를 알려줍니다 (대부분은 우선순위가 낮았기 때문이며, 때로는 정말로 호환되지 않기 때문입니다). 또한 Neuron 컴파일러(neuronx-cc)와 Transformers의 버전을 선택하여 확인할 수도 있습니다.
작성 시점 기준으로, Transformers 모델만 테스트되었으며 저희의 선택은 text-generation 및 image-text-to-text 모델로 제한됩니다. 우리는 diffusers, sentence transformers, 그리고 다른 트랜스포머 작업으로 확장할 계획입니다.
모델을 선택하여 실행 방법을 확인하세요. 이 Space는 해당 모델이 테스트된 Neuron 스택, 위에 설치해야 할 패키지들, 그리고 실행에 사용된 run.py 스크립트를 나열합니다. 이 스크립트는 TorchNeuron 발표의 레시피를 따릅니다: from_pretrained로 모델을 로드하고, 필요한 경우 분할(sharded)한 다음, 정적 캐시와 torch.compile(backend="neuron")을 사용하여 생성합니다.
예를 들어 AWS Neuron Deep Learning Container 내의 Trainium2 (trn2) 인스턴스에서 해당 스택으로 패키지를 설치한 후, 다음 명령어를 실행하세요:
python3 run.py
torchrun은 필요하지 않습니다.
스크립트는 모델이 샤딩(sharding)될 때 자체적으로 torch.distributed 랭크를 실행합니다.
첫 실행은 모델 그래프를 컴파일하는 것부터 시작합니다. 이 단계를 건너뛰려면, 테스트 중에 동일한 컴파일러가 구축한 그래프를 공개 NEFF 캐시 버킷에서 복사해야 합니다. Space는 올바른 컴파일러 버전과 함께 해당 명령을 제공합니다:
hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache
각 테스트는 모델을 trn2 인스턴스에 로드하며, 가중치를 담을 수 있는 가장 적은 수의 Trainium 코어에 걸쳐 샤딩됩니다. 대부분의 모델은 텐서 병렬성(tensor parallelism)으로 샤딩됩니다. 텐서 병렬 계획이 없는 Mixture-of-experts 모델은 전문가 병렬성(expert parallelism)만 사용합니다. 이후 테스트는 32개의 토큰을 생성하고 이를 확인합니다. 단일 코어에서는 해당 체크포인트의 CPU 실행 결과와 토큰을 비교합니다. 모델이 샤딩될 때는 작은 심사 모델(judge model)이 텍스트가 유창한 영어인지 확인합니다.
작성 시점 기준으로, 최대 16개의 코어(Trainium2 장치 4개에 적합한 모델)만 테스트됩니다. PASS는 처리량(throughput)이나 더 긴 생성 길이와 관련된 정보는 제공하지 않습니다.
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"]
데이터셋 카드는 각 결과의 정확한 의미를 알려줍니다.
다음에 모델이 Trainium에서 실행되는지 궁금할 때, 이 정보를 찾아보세요.
