JAX Toolbox มอบ CI สาธารณะ, อิมเมจ Docker สำหรับไลบรารี JAX ยอดนิยม และตัวอย่าง JAX ที่ปรับให้เหมาะสมเพื่อลดความซับซ้อนและปรับปรุงประสบการณ์การพัฒนา JAX ของคุณบน NVIDIA GPU รองรับไลบรารี JAX เช่น MaxText, Paxml และ Pallas
เราสนับสนุนและทดสอบเฟรมเวิร์ก JAX และสถาปัตยกรรมโมเดลต่อไปนี้ รายละเอียดเพิ่มเติมเกี่ยวกับแต่ละรุ่นและคอนเทนเนอร์ที่มีอยู่สามารถพบได้ใน README ที่เกี่ยวข้อง
| กรอบ | โมเดล | กรณีการใช้งาน | คอนเทนเนอร์ |
|---|---|---|---|
| แม็กซ์เท็กซ์ | GPT, LLaMA, เจมม่า, มิสทรัล, มิกซ์ทรัล | การฝึกอบรมล่วงหน้า | ghcr.io/nvidia/jax:maxtext |
| paxml | GPT, LLaMA, สธ | การฝึกล่วงหน้า การปรับแต่ง LoRA | ghcr.io/nvidia/jax:pax |
| t5x | T5, ViT | การฝึกอบรมล่วงหน้า การปรับแต่ง | ghcr.io/nvidia/jax:t5x |
| t5x | อิมเมจ | ก่อนการฝึกอบรม | ghcr.io/nvidia/t5x:imagen-2023-10-02.v3 |
| วิสัยทัศน์ที่ยิ่งใหญ่ | ปาลีเจมม่า | การปรับแต่ง การประเมินผล | ghcr.io/nvidia/jax:gemma |
| เลแวนเตอร์ | GPT, LLaMA, MPT, เป้สะพายหลัง | การฝึกล่วงหน้า, การปรับแต่งอย่างละเอียด | ghcr.io/nvidia/jax:levanter |
| ส่วนประกอบ | คอนเทนเนอร์ | สร้าง | ทดสอบ |
|---|---|---|---|
 | ghcr.io/nvidia/jax:base |   | [ไม่มีการทดสอบ] |
 | ghcr.io/nvidia/jax:jax |   |        |
 | ghcr.io/nvidia/jax:levanter |   |   |
 | ghcr.io/nvidia/jax:equinox |   | [ปิดการทดสอบ] |
 | ghcr.io/nvidia/jax:triton |  |   |
 | ghcr.io/nvidia/jax:upstream-t5x |   |  |
 | ghcr.io/nvidia/jax:t5x |   |  |
 | ghcr.io/nvidia/jax:upstream-pax |   |  |
 | ghcr.io/nvidia/jax:pax |   |  |
 | ghcr.io/nvidia/jax:maxtext |   |  |
 | ghcr.io/nvidia/jax:gemma |  |   |
ในทุกกรณี ghcr.io/nvidia/jax:XXX ชี้ไปที่บิลด์ล่าสุดของคอนเทนเนอร์สำหรับ XXX สำหรับการอ้างอิงที่เสถียร ให้ใช้ ghcr.io/nvidia/jax:XXX-YYYY-MM-DD
นอกจาก CI สาธารณะแล้ว เรายังทำการทดสอบ CI ภายในกับ H100 SXM 80GB และ A100 SXM 80GB อีกด้วย
อิมเมจ JAX ถูกฝังด้วยแฟล็กและตัวแปรสภาพแวดล้อมต่อไปนี้สำหรับการปรับแต่งประสิทธิภาพของ XLA และ NCCL:
| ธง XLA | ค่า | คำอธิบาย |
|---|---|---|
--xla_gpu_enable_latency_hiding_scheduler | true | อนุญาตให้ XLA ย้ายกลุ่มการสื่อสารเพื่อเพิ่มการทับซ้อนกับเคอร์เนลประมวลผล |
--xla_gpu_enable_triton_gemm | false | ใช้ cuBLAS แทนเคอร์เนล Trition GeMM |
| ตัวแปรสภาพแวดล้อม | ค่า | คำอธิบาย |
|---|---|---|
CUDA_DEVICE_MAX_CONNECTIONS | 1 | ใช้คิวเดียวสำหรับงาน GPU เพื่อลดเวลาแฝงของการดำเนินการสตรีม ตกลงเนื่องจาก XLA สั่งให้เปิดตัวแล้ว |
NCCL_NVLS_ENABLE | 0 | ปิดใช้งาน NVLink SHARP (1) รุ่นต่อๆ ไปจะเปิดใช้งานฟีเจอร์นี้อีกครั้ง |
มีการตั้งค่าสถานะ XLA อื่นๆ อีกมากมายที่ผู้ใช้สามารถตั้งค่าเพื่อปรับปรุงประสิทธิภาพได้ สำหรับคำอธิบายโดยละเอียดเกี่ยวกับแฟล็กเหล่านี้ โปรดดูเอกสารประสิทธิภาพ GPU สามารถปรับค่าสถานะ XLA ตามเวิร์กโฟลว์ได้ ตัวอย่างเช่น แต่ละสคริปต์ใน contrib/gpu/scripts_gpu จะตั้งค่าสถานะ XLA ของตัวเอง
สำหรับรายการแฟล็ก XLA ที่ใช้ก่อนหน้านี้ซึ่งไม่จำเป็นอีกต่อไป โปรดดูที่หน้าประสิทธิภาพของ GPU
| คืนแรกด้วยคอนเทนเนอร์ฐานใหม่ | ภาชนะฐาน |
|---|---|
| 06-11-2024 | NVIDIA/cuda:12.6.2-devel-ubuntu22.04 |
| 25-09-2024 | NVIDIA/cuda:12.6.1-devel-ubuntu22.04 |
| 24-07-2024 | NVIDIA/cuda:12.5.0-devel-ubuntu22.04 |
ดูหน้านี้สำหรับข้อมูลเพิ่มเติมเกี่ยวกับวิธีการโปรไฟล์โปรแกรม JAX บน GPU
สารละลาย:
นักเทียบท่าวิ่ง -it --shm-size=1g ...
คำอธิบาย: bus error อาจเกิดขึ้นเนื่องจากข้อจำกัดด้านขนาดของ /dev/shm คุณสามารถแก้ไขปัญหานี้ได้โดยการเพิ่มขนาดหน่วยความจำที่ใช้ร่วมกันโดยใช้ตัวเลือก --shm-size เมื่อเปิดใช้คอนเทนเนอร์ของคุณ
คำอธิบายปัญหา:
slurmstepd: error: pyxis: [INFO] Authentication succeeded slurmstepd: error: pyxis: [INFO] Fetching image manifest list slurmstepd: error: pyxis: [INFO] Fetching image manifest slurmstepd: error: pyxis: [ERROR] URL https://ghcr.io/v2/nvidia/jax/manifests/returned error code: 404 Not Found
วิธีแก้ไข: อัปเกรด enroot หรือใช้แพตช์ไฟล์เดียวตามที่กล่าวไว้ในบันทึกประจำรุ่น enroot v3.4.0
คำอธิบาย: Docker เดิมใช้ Docker Schema V2.2 สำหรับรายการ multi-arch manifest แต่ได้เปลี่ยนไปใช้รูปแบบ Open Container Initiative (OCI) ตั้งแต่ 20.10 น. Enroot เพิ่มการรองรับรูปแบบ OCI ในเวอร์ชัน 3.4.0
AWS
เพิ่มการบูรณาการ EFA
ตัวอย่างโค้ด SageMaker
จีซีพี
เริ่มต้นใช้งานแอปพลิเคชันหลายโหนด JAX พร้อม NVIDIA GPU บน Google Kubernetes Engine
สีฟ้า
เร่งแอปพลิเคชัน AI โดยใช้เฟรมเวิร์ก JAX บนเครื่องเสมือน NDm A100 v4 ของ Azure
โอซีไอ
การรันเวิร์กโหลดการเรียนรู้เชิงลึกด้วย JAX บนคลัสเตอร์ multi-GPU แบบหลายโหนดบน OCI
แจ๊กซ์ | คอนเทนเนอร์ NVIDIA NGC
การรวมการกำหนดค่า Slurm และ OpenMPI เป็นศูนย์
การเพิ่ม Ops GPU ที่กำหนดเอง
การถดถอยแบบ Triaging
Equinox สำหรับ JAX: รากฐานของระบบนิเวศสำหรับวิทยาศาสตร์และการเรียนรู้ของเครื่องจักร
ปรับขนาด Grok ด้วย JAX และ H100
JAX Supercharged บน GPU: LLM ประสิทธิภาพสูงพร้อม JAX และ OpenXLA
มีอะไรใหม่ใน JAX | GTC ฤดูใบไม้ผลิ 2024
มีอะไรใหม่ใน JAX | GTC ฤดูใบไม้ผลิ 2023
