تنزيل JAX Toolbox - تنزيل كود مصدر JAX Toolbox

صندوق أدوات جاكس

شفرة المصدر الأخرى

تنزيل

صندوق أدوات جاكس

يوفر JAX Toolbox صور CI عامة وDocker لمكتبات JAX الشهيرة وأمثلة JAX محسنة لتبسيط وتعزيز تجربة تطوير JAX على وحدات معالجة الرسومات NVIDIA. وهو يدعم مكتبات JAX مثل MaxText وPaxml وPallas.

الأطر والنماذج المدعومة

نحن ندعم ونختبر أطر عمل JAX وبنيات النماذج التالية. يمكن العثور على مزيد من التفاصيل حول كل نموذج والحاويات المتاحة في ملفات README الخاصة بكل منها.

نطاق	نماذج	حالات الاستخدام	حاوية
com.maxtext	جي بي تي، لاما، جيما، ميسترال، ميكسترال	التدريب المسبق	`ghcr.io/nvidia/jax:maxtext`
com.paxml	GPT، LLaMA، وزارة التعليم	التدريب المسبق، الضبط الدقيق، LoRA	`ghcr.io/nvidia/jax:pax`
t5x	T5، فيت	التدريب المسبق والضبط الدقيق	`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 الأخرى التي يمكن للمستخدمين تعيينها لتحسين الأداء. للحصول على شرح تفصيلي لهذه العلامات، يرجى الرجوع إلى مستند أداء وحدة معالجة الرسومات. يمكن ضبط علامات XLA حسب سير العمل. على سبيل المثال، يقوم كل برنامج نصي في contrib/gpu/scripts_gpu بتعيين علامات XLA الخاصة به.

للحصول على قائمة بعلامات XLA المستخدمة سابقًا والتي لم تعد هناك حاجة إليها، يرجى أيضًا الرجوع إلى صفحة أداء وحدة معالجة الرسومات.

الإصدارات

أول ليلة مع حاوية قاعدة جديدة	الحاوية الأساسية
2024-11-06	نفيديا/كودا:12.6.2-devel-ubuntu22.04
2024-09-25	نفيديا/كودا:12.6.1-devel-ubuntu22.04
2024-07-24	نفيديا/كودا:12.5.0-devel-ubuntu22.04

التنميط

راجع هذه الصفحة لمزيد من المعلومات حول كيفية تعريف برامج JAX على وحدة معالجة الرسومات.

الأسئلة المتداولة (الأسئلة الشائعة)

``خطأ في الناقل`` عند تشغيل JAX في حاوية عامل إرساء

حل:

 تشغيل عامل الميناء -it --shm-size=1g ...

توضيح: قد يحدث bus error بسبب محدودية حجم /dev/shm . يمكنك معالجة هذه المشكلة عن طريق زيادة حجم الذاكرة المشتركة باستخدام خيار --shm-size عند تشغيل الحاوية الخاصة بك.

يُبلغ enroot/pyxis عن رمز الخطأ 404 عند استيراد صور متعددة الأقواس

وصف المشكلة:

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 لقوائم البيانات متعددة الأقواس ولكنه تحول إلى استخدام تنسيق مبادرة الحاوية المفتوحة (OCI) منذ 20.10. أضاف Enroot دعمًا لتنسيق OCI في الإصدار 3.4.0.