הסקת מסקנות של MaxDiffusion במעבדי TPU מדגם v6e

במדריך הזה מוסבר איך להפעיל מודלים של MaxDiffusion ב-TPU v6e. במדריך הזה תלמדו איך ליצור תמונות באמצעות מודל Stable Diffusion XL.

לפני שמתחילים

הכנה להקצאת TPU v6e עם 4 שבבים:

כדי להגדיר Google Cloud פרויקט, להגדיר את Google Cloud CLI, להפעיל את Cloud TPU API ולוודא שיש לכם גישה לשימוש ב-Cloud TPU, פועלים לפי ההוראות במדריך הגדרת סביבת Cloud TPU.
מאמתים את Google Cloud הפרויקט ואת האזור שמוגדרים כברירת מחדל ב-Google Cloud CLI.
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

קיבולת מאובטחת

כשמוכנים לאבטח את קיבולת ה-TPU, אפשר לעיין במאמר מכסות Cloud TPU כדי לקבל מידע נוסף על מכסות Cloud TPU. אם יש לכם שאלות נוספות לגבי אבטחת קיבולת, תוכלו לפנות לצוות המכירות או לצוות ניהול החשבון של Cloud TPU.

הקצאת סביבת Cloud TPU

אפשר להקצות מכונות וירטואליות של TPU באמצעות GKE, באמצעות GKE ו-XPK, או כמשאבים בתור.

דרישות מוקדמות

מוודאים שיש לפרויקט מספיק TPUS_PER_TPU_FAMILY מכסה, שמציינת את המספר המקסימלי של שבבים שאפשר לגשת אליהם בGoogle Cloud פרויקט.
מוודאים שיש בפרויקט מספיק מכסת TPU ל:
- מכסת TPU VM
- מכסת כתובות IP
- מכסת Hyperdisk Balanced
הרשאות משתמש בפרויקט
- אם אתם משתמשים ב-GKE עם XPK, תוכלו לעיין בהרשאות במסוף Cloud בחשבון המשתמש או בחשבון השירות כדי לראות אילו הרשאות נדרשות להפעלת XPK.

הקצאת TPU v6e

   gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
        --node-id TPU_NAME \
        --project PROJECT_ID \
        --zone ZONE \
        --accelerator-type v6e-4 \
        --runtime-version v2-alpha-tpuv6e \
        --service-account SERVICE_ACCOUNT

משתמשים בפקודות list או describe כדי לשאול על הסטטוס של המשאב שנמצא בתור.

   gcloud alpha compute tpus queued-resources describe QUEUED_RESOURCE_ID  \
      --project=PROJECT_ID --zone=ZONE

רשימה מלאה של סטטוסים של בקשות למשאבים בתור זמינה במאמר בנושא משאבים בתור.

התחברות ל-TPU באמצעות SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

יצירת סביבת Conda

יוצרים ספרייה בשביל Miniconda:
```
mkdir -p ~/miniconda3
```

מורידים את סקריפט ההתקנה של Miniconda:

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh

מתקינים את Miniconda:

bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

מסירים את סקריפט ההתקנה של Miniconda:
```
rm -rf ~/miniconda3/miniconda.sh
```
מוסיפים את Miniconda למשתנה PATH:
```
export PATH="$HOME/miniconda3/bin:$PATH"
```
טוענים מחדש ~/.bashrc כדי להחיל את השינויים על המשתנה PATH:
```
source ~/.bashrc
```
יצירת סביבת Conda חדשה:
```
conda create -n tpu python=3.10
```
מפעילים את סביבת Conda:
```
source activate tpu
```

הגדרת MaxDiffusion

משכפלים את מאגר MaxDiffusion GitHub ועוברים לספרייה MaxDiffusion:
```
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
```
עוברים לענף mlperf-4.1:
```
git checkout mlperf4.1
```
מתקינים את MaxDiffusion:
```
pip install -e .
```
יחסי תלות של התקנות:
```
pip install -r requirements.txt
```

מתקינים את JAX:

pip install jax[tpu]==0.4.34 jaxlib==0.4.34 ml-dtypes==0.2.0 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

התקנה של יחסי תלות נוספים:

 pip install huggingface_hub==0.25 absl-py flax tensorboardX google-cloud-storage torch tensorflow transformers

יצירת תמונות

מגדירים משתני סביבה כדי להגדיר את זמן הריצה של TPU:

LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"

יצירת תמונות באמצעות ההנחיה וההגדרות שמוגדרות ב-src/maxdiffusion/configs/base_xl.yml:
```
python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"
```
אחרי שהתמונות נוצרות, חשוב לנקות את משאבי ה-TPU.

הסרת המשאבים

מוחקים את ה-TPU:

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async