שימוש ב-ML Diagnostics עם MaxText

MaxText היא ספרייה של מודלי שפה גדולים (LLM) בקוד פתוח, שנועדה לפעול ב-TPU וב-GPU לצורך אימון מודלים. ‫ML Diagnostics SDK משולב מראש עם MaxText. אפשר להפעיל את אבחון ה-ML באמצעות MaxText באמצעות הדגל managed-mldiagnostics.

אחרי ההפעלה, תוכלו לבצע את הפעולות הבאות:

  • יוצרים הפעלה מנוהלת של MachineLearning עם כל ההגדרות של MaxText.
  • מעלים עקבות של פרופילים, אם הפרופילים מופעלים באמצעות הדגל profiler="xplane".
  • מעלים מדדי אימון במרווח מוגדר באמצעות הדגל log_period.

כשמשתמשים ב-ML Diagnostics עם MaxText, צריך להשתמש בדגלים הבאים:

managed_mldiagnostics: True  # Enable the managed diagnostics
managed_mldiagnostics_run_group: GROUP_NAME  # Used to group multiple runs. (Optional)

כדי להפעיל את ML Diagnostics ב-MaxText, אפשר לשנות את קובץ ההגדרות של ההרצה או להעביר את הדגלים משורת הפקודה. לדוגמה, מעבירים את הדגלים הבאים עם הפקודה MaxText.train:

python3 -m MaxText.train \
  src/MaxText/configs/base.yml \
  run_name="demo-mldiagnostics-run-2" \
  model_name="<your_chosen_model>" \
  base_output_directory=gs://<your_gcs_folder>/ \
  dataset_type=synthetic \
  steps=100 \
  log_period=10 \
  profiler=xplane \
  upload_all_profiler_results=True \
  managed_mldiagnostics=True \
  managed_mldiagnostics_run_group="demo-mldiagnostics-group"

בדוגמה הזו, הגדרת הדגל upload_all_profiler_results=True מאפשרת ללכוד פרופילים של כמה מארחים מכל המארחים.