maxtext
v1.0.0
MaxText عبارة عن LLM عالي الأداء وقابل للتطوير بدرجة كبيرة ومفتوح المصدر ومكتوب بلغة Python/Jax النقية ويستهدف Google Cloud TPUs وGPUs للتدريب والاستدلال . يحقق MaxText وحدات MFU عالية ويتوسع من مضيف واحد إلى مجموعات كبيرة جدًا مع البقاء بسيطًا و"خاليًا من التحسين" بفضل قوة Jax ومترجم XLA.
يهدف MaxText إلى أن يكون نقطة انطلاق لمشاريع LLM الطموحة في كل من البحث والإنتاج. نحن نشجع المستخدمين على البدء بتجربة MaxText خارج الصندوق ثم تفرع MaxText وتعديله لتلبية احتياجاتهم.
لقد استخدمنا MaxText لإظهار التدريب عالي الأداء والمتقارب بشكل جيد في int8 وتوسيع نطاق التدريب إلى 51 ألف شريحة تقريبًا.
الميزات الرئيسية المدعومة:
لأول مرة تقوم فيها بتشغيل MaxText، فإننا نقدم تعليمات محددة.
يدعم MaxText التدريب والاستدلال على العديد من النماذج المفتوحة. اتبع أدلة المستخدم في مجلد البدء لمعرفة المزيد.
بعض الأدلة المفيدة الإضافية:
بالإضافة إلى أدلة البدء، هناك دائمًا إمكانيات MaxText أخرى تتم إضافتها باستمرار! المجموعة الكاملة من الاختبارات الشاملة موجودة في end_to_end. نحن نديرهم بإيقاع ليلي. يمكن أن تكون مصدرًا جيدًا لفهم MaxText. وبدلاً من ذلك، يمكنك رؤية اختبارات الوحدة المستمرة التي يتم إجراؤها بشكل مستمر تقريبًا.
يمكن العثور على مزيد من التفاصيل حول إعادة إنتاج هذه النتائج في MaxText/configs/README.md.
| عدد المعلمات | نوع المسرع | تفلوب/رقاقة/ثانية | استخدام النموذج المتخبط (MFU) |
|---|---|---|---|
| 32 ب | v5p-128 | 3.28ه+02 | 71.47% |
| 64 ب | v5p-128 | 3.23ه+02 | 70.31% |
| 128 ب | v5p-256 | 3.15هـ+02 | 68.68% |
| 128 ب | v5p-512 | 3.15هـ+02 | 68.53% |
| 256 ب | v5p-1024 | 3.16هـ+02 | 68.82% |
| 512 ب | v5p-1024 | 2.94ه+02 | 63.99% |
| 1024 ب | v5p-2048 | 2.49هـ+02 | 64.05% |
| 1024 ب | v5p-4096 | 2.97هـ+02 | 64.80% |
| 1160 ق | v5p-7680 | 2.95ه+02 | 64.27% |
| 1160 ق | v5p-12288 | 3.04هـ+02 | 66.23% |
لموديلات 16B و32B و64B و128B. راجع تكوينات التشغيل الكامل في MaxText/configs/v5e/ كـ 16b.sh , 32b.sh , 64b.sh , 128b.sh .
| الأجهزة | 16B TFLOP/ثانية/رقاقة | 16 ب مفو | 32B TFLOP/ثانية/رقاقة | 32 ب مفو | 64B TFLOP/ثانية/رقاقة | 64 ب مفو | 128B TFLOP/ثانية/رقاقة | 128 ب مفو |
|---|---|---|---|---|---|---|---|---|
| 1x v5e-256 | 120 | 61.10% | 132 | 66.86% | 118 | 59.90% | 110 | 56.06% |
| 2x v5e-256 | 117 | 59.37% | 128 | 64.81% | 112 | 56.66% | 110 | 55.82% |
| 4x v5e-256 | 117 | 59.14% | 126 | 64.10% | 110 | 55.85% | 108 | 54.93% |
| 8x v5e-256 | 115 | 58.27% | 125 | 63.67% | 108 | 54.96% | 104 | 52.93% |
| 16x v5e-256 | 111 | 56.56% | 123 | 62.26% | 105 | 53.29% | 100 | 50.86% |
| 32x v5e-256 | 108 | 54.65% | 119 | 60.40% | 99 | 50.18% | 91 | 46.25% |
MaxText مستوحى بشكل كبير من MinGPT/NanoGPT، وتطبيقات GPT المستقلة الأنيقة المكتوبة في PyTorch والتي تستهدف وحدات معالجة الرسومات Nvidia. يعد MaxText أكثر تعقيدًا، حيث يدعم المزيد من النماذج القياسية الصناعية ويتسع لعشرات الآلاف من الرقائق. في نهاية المطاف، يحتوي MaxText على MFU أكثر من ثلاثة أضعاف الـ 17% التي تم الإبلاغ عنها مؤخرًا باستخدام قاعدة التعليمات البرمجية هذه، وهو قابل للتطوير على نطاق واسع وينفذ ذاكرة تخزين مؤقت ذات قيمة رئيسية لفك تشفير الانحدار التلقائي الفعال.
يشبه MaxText إلى حد كبير Nvidia/Megatron-LM، وهو تطبيق LLM تم ضبطه جيدًا ويستهدف وحدات معالجة الرسومات Nvidia. يحقق التطبيقان MFUs قابلة للمقارنة. يسلط الاختلاف في قواعد التعليمات البرمجية الضوء على استراتيجيات البرمجة المختلفة. MaxText هو لغة Python خالصة، ويعتمد بشكل كبير على مترجم XLA لتحقيق الأداء العالي. على النقيض من ذلك، يعد Megatron-LM مزيجًا من Python وCUDA، ويعتمد على نواة CUDA المحسنة جيدًا لتحقيق الأداء العالي.
MaxText يمكن مقارنته أيضًا بـ Pax. مثل Pax، يوفر MaxText تطبيقات عالية الأداء وقابلة للتطوير لـ LLMs في Jax. يركز Pax على تمكين معلمات التكوين القوية، مما يمكّن المطورين من تغيير النموذج عن طريق تحرير معلمات التكوين. على النقيض من ذلك، يعد MaxText تطبيقًا بسيطًا وملموسًا لمختلف دورات LLM التي تشجع المستخدمين على التوسع عن طريق تفرع كود المصدر وتحريره مباشرة.
عند تشغيل مهمة برنامج واحد وبيانات متعددة (SPMD) على المسرعات، يمكن تعليق العملية بأكملها في حالة وجود أي خطأ أو توقف/تعطل أي جهاز افتراضي لسبب ما. في هذا السيناريو، سيساعد التقاط تتبعات المكدس في تحديد المشكلات واستكشاف مشكلات المهام التي يتم تشغيلها على أجهزة TPU الافتراضية وإصلاحها.
ستساعد التكوينات التالية في تصحيح الأخطاء أو عندما يكون البرنامج عالقًا أو معلقًا في مكان ما عن طريق جمع تتبعات المكدس. قم بتغيير قيم المعلمات وفقًا لذلك في MaxText/configs/base.yml :
collect_stack_trace: True لتمكين مجموعة تتبعات المكدس عند حدوث أخطاء أو عند تعليق البرنامج. سيقوم هذا الإعداد بشكل دوري بتفريغ آثار البرنامج للمساعدة في تصحيح الأخطاء. لتعطيل هذا، قم بتعيين collect_stack_trace: False .stack_trace_to_cloud: False لعرض تتبعات المكدس على وحدة التحكم. stack_trace_to_cloud: True إلى إنشاء ملف مؤقت في /tmp/debugging في وحدات TPU لتخزين آثار المكدس. يوجد وكيل يعمل على أجهزة TPU الافتراضية التي ستقوم بشكل دوري بتحميل التتبعات من الدليل المؤقت إلى التسجيل السحابي في مشروع gcp. يمكنك عرض الآثار في Logs Explorer على Cloud Logging باستخدام الاستعلام التالي: logName="projects/<project_name>/logs/tpu.googleapis.com%2Fruntime_monitor"
jsonPayload.verb="stacktraceanalyzer"
stack_trace_interval_seconds إلى المدة بالثواني بين كل حدث مجموعة تتبع المكدس. إعداد stack_trace_interval_seconds: 600 سوف يجمع آثار المكدس كل 600 ثانية (10 دقائق).إليك حزمة PyPI ذات الصلة: https://pypi.org/project/cloud-tpu-diagnostics.
لتجميع التدريب الخاص بك مسبقًا، نوفر لك أداة train_compile.py . تسمح لك هذه الأداة بتجميع train_step الرئيسي في train.py للأجهزة المستهدفة (على سبيل المثال، عدد كبير من أجهزة v5e) دون استخدام المجموعة الكاملة.
لا يجوز لك سوى استخدام وحدة المعالجة المركزية (CPU) أو جهاز افتراضي واحد من عائلة مختلفة للتجميع المسبق لمجموعة TPU. يساعد هذا التجميع في تحقيق هدفين رئيسيين:
سيتم وضع علامة على أي معلومات حول نفاد الذاكرة (OOM)، على سبيل المثال عندما يتم تعيين per_device_batch_size على مستوى مرتفع جدًا، مع تتبع مكدس OOM متطابق كما لو تم تجميعه على الجهاز المستهدف.
يمكن حفظ التجميع المسبق ثم تحميله لبدء التشغيل السريع وأوقات إعادة التشغيل على الجهاز المستهدف.
ترتبط الأداة train_compile.py ارتباطًا وثيقًا بـ train.py وتستخدم نفس ملف التكوين configs/base.yml . على الرغم من أنك لا تحتاج إلى التشغيل على TPU، فأنت بحاجة إلى تثبيت jax[tpu] بالإضافة إلى التبعيات الأخرى، لذلك نوصي بتشغيل setup.sh لتثبيتها إذا لم تكن قد قمت بذلك بالفعل.
بعد تثبيت التبعيات المذكورة أعلاه، تصبح جاهزًا للتجميع مسبقًا:
# Run the below on a single machine, e.g. a CPU
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256 compile_topology_num_slices=2
global_parameter_scale=16 per_device_batch_size=4
سيؤدي هذا إلى تجميع نموذج MaxText ذو معلمة 16B على 2 v5e pods.
فيما يلي مثال يحفظ ثم يحمّل train_step المترجم، بدءًا من الحفظ:
الخطوة 1: تشغيل AOT وحفظ الوظيفة المترجمة
# Run the below on a single machine, e.g. a CPU
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256
compile_topology_num_slices=2
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16
per_device_batch_size=4 steps=10000 learning_rate=1e-3
الخطوة 2: قم بتشغيل Train.py وقم بتحميل الوظيفة المترجمة
لتحميل Train_step المترجم، تحتاج فقط إلى تمرير compiled_trainstep_file=my_compiled_train.pickle إلى train.py :
# Run the below on each host of the target hardware, e.g. each host on 2 slices of v5e-256
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile
compiled_trainstep_file=my_compiled_train.pickle
global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket
في خطوة الحفظ للمثال 2 أعلاه قمنا بتضمين تصدير علامة المترجم LIBTPU_INIT_ARGS و learning_rate لأن ذلك يؤثر على الكائن المترجم my_compiled_train.pickle. يتم إصلاح أحجام النموذج (على سبيل المثال، global_parameter_scale و max_sequence_length و per_device_batch ) عندما تقوم بالتجميع في البداية عبر compile_train.py ، وسترى خطأ في الحجم إذا حاولت تشغيل الكائن المترجم المحفوظ بأحجام مختلفة عن تلك التي تم تجميعها بها. ومع ذلك، هناك ملاحظة دقيقة تتمثل في أن جدول معدل التعلم يتم إصلاحه أيضًا عند تشغيل compile_train - والذي يتم تحديده من خلال كل من steps و learning_rate . يتم تمرير معلمات المحسن مثل adam_b1 فقط ككائنات ذات شكل إلى المترجم - وبالتالي يتم تحديد قيمها الحقيقية عند تشغيل train.py ، وليس أثناء التحويل البرمجي. إذا قمت بالتمرير بأشكال مختلفة (على سبيل المثال، per_device_batch )، فسوف تحصل على رسالة خطأ واضحة تفيد بأن التوقيع المترجم له أشكال متوقعة مختلفة عما تم إدخاله. إذا حاولت التشغيل على أجهزة مختلفة عن أهداف التجميع المطلوبة عبر compile_topology ، فسوف تحصل على خطأ يفيد بوجود فشل في تعيين الأجهزة من المترجم إلى أجهزتك الحقيقية. من المحتمل أن يعمل استخدام إشارات XLA أو LIBTPU مختلفة عما تم تجميعه بصمت مع البيئة التي قمت بتجميعها دون أخطاء. ولكن لا يوجد سلوك مضمون في هذه الحالة؛ يجب أن تعمل في نفس البيئة التي قمت بتجميعها فيها.
يتم أيضًا دعم التجميع المسبق لوحدات معالجة الرسومات مع بعض الاختلافات عن وحدات معالجة الرسومات:
لا يدعم GPU التجميع عبر الأجهزة: لا يزال مضيف GPU مطلوبًا لتشغيل تجميع AoT، ولكن يمكن لمضيف GPU واحد تجميع برنامج لمجموعة أكبر من نفس الأجهزة.
بالنسبة لوحدات معالجة الرسومات السحابية A3، فإن الحد الأقصى لحجم "الشريحة" هو مضيف واحد، وتمثل المعلمة compile_topology_num_slices عدد أجهزة A3 المطلوب التحويل البرمجي لها مسبقًا.
يوضح هذا المثال العلامات التي سيتم استخدامها لتجميع GPU متعدد المضيف الذي يستهدف مجموعة من 4 مضيفين A3:
الخطوة 1: تشغيل AOT وحفظ الوظيفة المترجمة
# Run the below on a single A3 machine
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=a3
compile_topology_num_slices=4
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16
attention=dot_product per_device_batch_size=4 steps=10000 learning_rate=1e-3
الخطوة 2: قم بتشغيل Train.py وقم بتحميل الوظيفة المترجمة
لتحميل Train_step المترجم، تحتاج فقط إلى تمرير compiled_trainstep_file=my_compiled_train.pickle إلى train.py :
# Run the below on each of the 4 target A3 hosts.
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile
compiled_trainstep_file=my_compiled_train.pickle
attention=dot_product global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket
كما في حالة TPU، لاحظ أن بيئة الترجمة يجب أن تتطابق مع بيئة التنفيذ، في هذه الحالة عن طريق تعيين نفس XLA_FLAGS .
يدعم MaxText التحميل التلقائي للسجلات المجمعة في دليل إلى مثيل Tensorboard في Vertex AI. اتبع دليل المستخدم لمعرفة المزيد.
