ตรวจสอบโมเดล Transformers บน AWS Trainium ได้ง่ายๆ ผ่าน Space ใหม่

ใน โพสต์ก่อนหน้า เราได้แนะนำ TorchNeuron ซึ่งเป็น backend ของ PyTorch แบบ native สำหรับ AWS Trainium (แพ็กเกจ torch-neuronx) โดย Transformers สามารถรันโค้ดการสร้างโมเดลบน Trainium ได้เช่นเดียวกับอุปกรณ์อื่น ส่งผลให้โมเดล Transformers บน Hub มีโอกาสรันได้ทั้งหมด แต่คำถามสำคัญคือโมเดลไหนบ้างที่ใช้งานได้จริง

Space torch-neuronx model support ถูกสร้างขึ้นเพื่อตอบคำถามนี้สำหรับโมเดลยอดนิยมบน Hub โดยจะระบุว่าโมเดลนั้นสามารถรันบน Trainium ได้ทันทีโดยไม่ต้องแก้ไขหรือไม่ พร้อมทั้งจัดเตรียมสคริปต์สำหรับนำไปทดสอบด้วยตัวคุณเอง

วิธีการค้นหาโมเดล

เมื่อเปิด Space จะพบกับมุมมอง Selected ซึ่งเป็นการคัดเลือกโมเดลระดับท็อปจาก Hub ตามยอดดาวน์โหลดและแนวโน้มความนิยม แต่ละแถวจะแสดงผลลัพธ์การทดสอบ จำนวนคอร์ของ Trainium ที่ใช้ทำ sharding ระยะเวลาที่ใช้รันจริง (wall time) ซึ่งรวมขั้นตอนการ compile และหน่วยความจำสูงสุดที่ใช้บนอุปกรณ์

หากต้องการค้นหาโมเดลเฉพาะเจาะจง สามารถสลับไปที่มุมมอง Trending หรือ Downloads แล้วพิมพ์ชื่อโมเดลได้ทันที โดยระบบจะอ้างอิงตามการจัดอันดับของ Hub ในแต่ละ task ซึ่งหากโมเดลใดถูกคัดออก ระบบจะระบุเหตุผลไว้ เช่น ไม่ได้อยู่ในลำดับความสำคัญหรือพบปัญหาความไม่เข้ากัน นอกจากนี้ยังเลือกเวอร์ชันของ Neuron compiler (neuronx-cc) และเวอร์ชันของ Transformers เพื่อตรวจสอบย้อนหลังได้

ข้อมูล ณ ขณะที่เขียนบทความนี้ มีเฉพาะโมเดล Transformers ประเภท text-generation และ image-text-to-text เท่านั้นที่ถูกทดสอบ โดยมีแผนจะขยายไปยัง diffusers, sentence transformers และ task อื่นๆ ในอนาคต

การรันโมเดลด้วยตัวคุณเอง

เมื่อเลือกโมเดลที่ต้องการ Space จะแสดงรายละเอียด Neuron stack ที่ใช้ทดสอบ แพ็กเกจที่ต้องติดตั้ง และสคริปต์ run.py ซึ่งใช้หลักการเดียวกับการประกาศเปิดตัว TorchNeuron โดยจะโหลดโมเดลด้วย from_pretrained ทำ sharding เมื่อจำเป็น และสร้างข้อความด้วย static cache ร่วมกับ torch.compile(backend="neuron")

สำหรับการรันบนอินสแตนซ์ Trainium2 (trn2) ที่ใช้ AWS Neuron Deep Learning Container ให้ติดตั้งแพ็กเกจที่กำหนดแล้วรันคำสั่งด้านล่าง ซึ่งไม่จำเป็นต้องใช้ torchrun เพราะสคริปต์จะจัดการ torch.distributed ranks ให้เองอัตโนมัติ:

python3 run.py
``` เนื่องจากการรันครั้งแรกต้องมีการ compile กราฟโมเดล คุณสามารถประหยัดเวลาได้โดยคัดลอกกราฟที่เคย compile ไว้แล้วจาก bucket สาธารณะ [NEFF cache](https://huggingface.co/buckets/aws-neuron/torch-neuronx-neff-cache) ผ่านคำสั่งที่ Space เตรียมไว้ให้ดังนี้:
 
```bash
hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

นิยามของคำว่า "รันได้" (runs)

การทดสอบจะโหลดโมเดลบนอินสแตนซ์ trn2 โดยใช้จำนวนคอร์น้อยที่สุดที่รองรับค่าน้ำหนัก (weights) ได้ ส่วนใหญ่จะทำ sharding ด้วย tensor parallelism ยกเว้นโมเดล Mixture-of-experts ที่จะใช้ expert parallelism แทน จากนั้นจะทำการสร้าง 32 tokens เพื่อตรวจสอบความถูกต้อง หากเป็นคอร์เดี่ยวจะเทียบกับผลลัพธ์จาก CPU แต่หากมีการทำ sharding จะใช้โมเดลตัดสินขนาดเล็กตรวจสอบความเป็นธรรมชาติของภาษา

ปัจจุบันการทดสอบจำกัดเฉพาะโมเดลที่ใช้งานบนอุปกรณ์ Trainium2 ไม่เกิน 4 เครื่อง (สูงสุด 16 คอร์) โดยผลลัพธ์ที่เป็น PASS จะยืนยันเพียงว่าโมเดลสามารถรันได้สำเร็จ แต่ไม่ได้การันตีถึงประสิทธิภาพ throughput หรือความสามารถในการสร้างข้อความที่ยาวกว่าที่ทดสอบ

ข้อมูลในรูปแบบชุดข้อมูลสาธารณะ

ข้อมูลทั้งหมดที่ปรากฏบน Space มาจากชุดข้อมูลสาธารณะ aws-neuron/torch-neuronx-model-support ในรูปแบบไฟล์ Parquet ซึ่งแบ่งตามเวอร์ชันของ neuronx-cc และไลบรารีต่างๆ

ชุดข้อมูลนี้ประกอบด้วยผลลัพธ์การรัน ข้อผิดพลาดที่พบ เวอร์ชันของ stack และสคริปต์ 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"]]

สำหรับรายละเอียดความหมายของผลลัพธ์แต่ละแบบ สามารถศึกษาเพิ่มเติมได้ที่ dataset card และหากต้องการทราบว่าโมเดลของคุณรองรับหรือไม่ ลองไปตรวจสอบได้ทันที