Pengantar Trillium (v6e)
v6e digunakan untuk merujuk ke Trillium dalam dokumentasi, TPU API, dan log ini. v6e mewakili TPU generasi ke-6 Google.
Dengan 256 chip per Pod, arsitektur v6e memiliki banyak kesamaan dengan v5e. Sistem ini dioptimalkan untuk pelatihan, penyesuaian, dan penayangan transformer, teks ke gambar, dan jaringan saraf konvolusi (CNN).
Untuk mengetahui informasi selengkapnya tentang konfigurasi dan arsitektur sistem v6e, lihat TPU v6e.
Dokumen pengantar ini berfokus pada proses pelatihan dan penayangan model menggunakan framework JAX, PyTorch, atau TensorFlow. Dengan setiap framework, Anda dapat menyediakan TPU menggunakan resource dalam antrean atau GKE. Penyiapan GKE dapat dilakukan menggunakan perintah XPK atau GKE.
Prosedur umum untuk melatih atau menayangkan model menggunakan v6e
- Menyiapkan Google Cloud project
- Kapasitas aman
- Menyediakan lingkungan Cloud TPU
- Menjalankan workload pelatihan atau inferensi model
Menyiapkan project Google Cloud
Sebelum dapat menggunakan Cloud TPU, Anda harus:
- Membuat akun dan project Google Cloud dengan penagihan diaktifkan
- Menginstal komponen alfa Google Cloud CLI
- Mengaktifkan Cloud TPU API
- Membuat agen layanan Cloud TPU
- Membuat akun layanan Cloud TPU dan memberikan izin
Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU.
Kapasitas aman
Hubungi Google Cloud dukungan untuk meminta kuota Cloud TPU v6e dan untuk menjawab pertanyaan apa pun tentang kapasitas.
Menyediakan lingkungan Cloud TPU
Cloud TPU v6e dapat disediakan dan dikelola dengan GKE, dengan GKE dan XPK (alat CLI wrapper melalui GKE), atau sebagai resource dalam antrean.
Prasyarat
- Pastikan project Anda memiliki kuota
TPUS_PER_TPU_FAMILY
yang cukup, yang menentukan jumlah maksimum chip yang dapat Anda akses dalam project Google Cloud. - v6e telah diuji dengan konfigurasi berikut:
- Python
3.10
atau yang lebih baru - Versi software harian:
- JAX Harian
0.4.32.dev20240912
- LibTPU Harian
0.1.dev20240912+nightly
- JAX Harian
- Versi software stabil:
- JAX + JAX Lib v0.4.37
- Python
Pastikan project Anda memiliki cukup kuota untuk:
- Kuota VM Cloud TPU
- Kuota alamat IP
Kuota Hyperdisk Balanced
Jika Anda menggunakan GKE dengan XPK, lihat Izin Konsol Cloud di akun pengguna atau layanan untuk mengetahui izin yang diperlukan untuk menjalankan XPK.
Membuat variabel lingkungan
Di Cloud Shell, buat variabel lingkungan berikut:
export NODE_ID=your-tpu-name export PROJECT_ID=your-project-id export ACCELERATOR_TYPE=v6e-16 export ZONE=us-east1-d export RUNTIME_VERSION=v2-alpha-tpuv6e export SERVICE_ACCOUNT=your-service-account export QUEUED_RESOURCE_ID=your-queued-resource-id export VALID_DURATION=your-duration # Additional environment variable needed for Multislice: export NUM_SLICES=number-of-slices # Use a custom network for better performance as well as to avoid having the default network becoming overloaded. export NETWORK_NAME=${PROJECT_ID}-mtu9k export NETWORK_FW_NAME=${NETWORK_NAME}-fw
Deskripsi flag perintah
Variabel | Deskripsi |
NODE_ID | ID Cloud TPU yang ditetapkan pengguna yang dibuat saat permintaan resource dalam antrean dialokasikan. |
PROJECT_ID | Google Cloud nama project. Gunakan project yang sudah ada atau buat project baru. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan Google Cloud project. |
ZONA | Lihat dokumen Region dan zona Cloud TPU untuk mengetahui zona yang didukung. |
ACCELERATOR_TYPE | Lihat Jenis Akselerator. |
RUNTIME_VERSION | v2-alpha-tpuv6e
|
SERVICE_ACCOUNT | Ini adalah alamat email untuk akun layanan yang dapat Anda temukan di
Google Cloud Konsol -> IAM -> Akun Layanan
Contoh: |
NUM_SLICES | Jumlah slice yang akan dibuat (hanya diperlukan untuk Multislice). |
QUEUED_RESOURCE_ID | ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan. |
VALID_DURATION | Durasi validitas permintaan resource yang diantrekan. |
NETWORK_NAME | Nama jaringan sekunder yang akan digunakan. |
NETWORK_FW_NAME | Nama firewall jaringan sekunder yang akan digunakan. |
Mengoptimalkan performa jaringan
Untuk performa terbaik,gunakan jaringan dengan MTU 8.896 (unit transmisi maksimum).
Secara default, Virtual Private Cloud (VPC) hanya menyediakan MTU sebesar 1.460 byte yang akan memberikan performa jaringan yang kurang optimal. Anda dapat menetapkan MTU jaringan VPC ke nilai apa pun antara 1.300 byte dan 8.896 byte (inklusif). Ukuran MTU kustom umum adalah 1.500 byte (Ethernet standar) atau 8.896 byte (maksimum yang memungkinkan). Untuk mengetahui informasi selengkapnya, lihat Ukuran MTU jaringan VPC yang valid.
Untuk informasi selengkapnya tentang cara mengubah setelan MTU untuk jaringan yang ada atau default, lihat Mengubah setelan MTU jaringan VPC.
Contoh berikut membuat jaringan dengan 8.896 MTU.
export RESOURCE_NAME=your-resource-name export NETWORK_NAME=${RESOURCE_NAME}-privatenetwork export NETWORK_FW_NAME=${RESOURCE_NAME}-privatefirewall gcloud compute networks create ${NETWORK_NAME} --mtu=8896 --project=${PROJECT_ID} \ --subnet-mode=auto --bgp-routing-mode=regional gcloud compute firewall-rules create ${NETWORK_FW_NAME} --network=${NETWORK_NAME} \ --allow tcp,icmp,udp --project=${PROJECT_ID}
Menggunakan multi-NIC (opsi untuk Multislice)
Variabel lingkungan berikut diperlukan untuk subnet sekunder saat Anda menggunakan lingkungan Multislice.
export NETWORK_NAME_2=${RESOURCE_NAME} export SUBNET_NAME_2=${RESOURCE_NAME} export FIREWALL_RULE_NAME=${RESOURCE_NAME} export ROUTER_NAME=${RESOURCE_NAME}-network-2 export NAT_CONFIG=${RESOURCE_NAME}-natconfig-2 export REGION=your-region
Gunakan perintah berikut untuk membuat perutean IP kustom untuk jaringan dan subnet.
gcloud compute networks create ${NETWORK_NAME_2} --mtu=8896 \
--bgp-routing-mode=regional --subnet-mode=custom --project=${PROJECT_ID}
gcloud compute networks subnets create ${SUBNET_NAME_2} \
--network=${NETWORK_NAME_2} \
--range=10.10.0.0/18 --region=${REGION} \
--project=${PROJECT_ID}
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
--network=${NETWORK_NAME_2} --allow tcp,icmp,udp \
--source-ranges 10.10.0.0/18 --project=${PROJECT_ID}
gcloud compute routers create ${ROUTER_NAME} \
--project=${PROJECT_ID} \
--network=${NETWORK_NAME_2} \
--region=${REGION}
gcloud compute routers nats create ${NAT_CONFIG} \
--router=${ROUTER_NAME} \
--region=${REGION} \
--auto-allocate-nat-external-ips \
--nat-all-subnet-ip-ranges \
--project=${PROJECT_ID} \
--enable-logging
Setelah membuat slice multi-jaringan, Anda dapat memvalidasi bahwa kedua kartu antarmuka jaringan (NIC) digunakan dengan menyiapkan cluster XPK dan menambahkan flag --command ifconfig
ke perintah pembuatan beban kerja XPK.
Gunakan perintah xpk workload
berikut untuk menampilkan output perintah ifconfig
di log konsol Google Cloud dan pastikan eth0 dan eth1 memiliki mtu=8896.
python3 xpk.py workload create \ --cluster CLUSTER_NAME \ {--base-docker-image maxtext_base_image | --docker-image your-cloud-image-name} \ --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \ --tpu-type=${ACCELERATOR_TYPE} \ --num-slices=${NUM_SLICES} \ --on-demand \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --command "ifconfig"
Jika Anda ingin mengaktifkan log debug atau menggunakan Vertex AI TensorBoard, tambahkan argumen opsional berikut ke perintah:
--enable-debug-logs \ --use-vertex-tensorboard
Verifikasi bahwa eth0 dan eth1 memiliki mtu=8.896. Anda dapat memverifikasi bahwa multi-NIC
berjalan dengan menambahkan flag --command ifconfig
ke perintah pembuatan
workload XPK. Periksa output workload xpk tersebut di log konsol Google Cloud dan verifikasi bahwa eth0 dan eth1 memiliki mtu=8896.
Meningkatkan setelan TCP
Jika membuat Cloud TPU menggunakan antarmuka resource yang diantrekan, Anda dapat menjalankan perintah berikut untuk meningkatkan performa jaringan dengan meningkatkan batas buffering penerimaan TCP.
gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \ --project "${PROJECT}" \ --zone "${ZONE}" \ --node=all \ --command='sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"' \ --worker=all
Penyediaan dengan resource dalam antrean
Anda dapat membuat Cloud TPU v6e menggunakan resource yang diantrekan. Resource dalam antrean memungkinkan Anda menerima kapasitas setelah tersedia. Anda dapat menentukan waktu mulai dan waktu berakhir opsional untuk saat permintaan harus diisi. Untuk mengetahui informasi selengkapnya, lihat Mengelola resource dalam antrean.
Menyediakan Cloud TPU v6e dengan GKE atau XPK
Jika menggunakan perintah GKE dengan v6e, Anda dapat menggunakan perintah Kubernetes atau XPK untuk menyediakan Cloud TPU dan melatih atau menayangkan model. Lihat Merencanakan Cloud TPU di GKE untuk mempelajari cara merencanakan konfigurasi Cloud TPU di cluster GKE. Bagian berikut memberikan perintah untuk membuat cluster XPK dengan dukungan NIC tunggal dan dukungan multi-NIC.
Membuat cluster XPK dengan dukungan NIC tunggal
export CLUSTER_NAME=xpk-cluster-name export ZONE=us-east1-d export PROJECT_ID=your-project-id export TPU_TYPE=v6e-256 export NUM_SLICES=2 export NETWORK_NAME=${CLUSTER_NAME}-mtu9k export NETWORK_FW_NAME=${NETWORK_NAME}-fw
gcloud compute networks create ${NETWORK_NAME} \ --mtu=8896 \ --project=${PROJECT_ID} \ --subnet-mode=auto \ --bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} \ --network=${NETWORK_NAME} \ --allow tcp,icmp,udp \ --project=${PROJECT_ID}
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
python3 xpk.py cluster create --cluster=${CLUSTER_NAME} \ --cluster-cpu-machine-type=n1-standard-8 \ --num-slices=${NUM_SLICES} \ --tpu-type=${TPU_TYPE} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --on-demand \ --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \ --create-vertex-tensorboard
Deskripsi flag perintah
Variabel | Deskripsi |
CLUSTER_NAME | Nama yang ditetapkan pengguna untuk cluster XPK. |
PROJECT_ID | Google Cloud nama project. Gunakan project yang sudah ada atau buat project baru. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan Google Cloud project. |
ZONA | Lihat dokumen Region dan zona Cloud TPU untuk mengetahui zona yang didukung. |
TPU_TYPE | Lihat Jenis Akselerator. |
NUM_SLICES | Jumlah slice yang ingin Anda buat |
CLUSTER_ARGUMENTS | Jaringan dan subnetwork yang akan digunakan.
Contoh: |
NUM_SLICES | Jumlah slice yang akan dibuat. |
NETWORK_NAME | Nama jaringan sekunder yang akan digunakan. |
NETWORK_FW_NAME | Nama firewall jaringan sekunder yang akan digunakan. |
Membuat cluster XPK dengan dukungan multi-NIC
export CLUSTER_NAME=xpk-cluster-name export ZONE=us-east1-d export PROJECT_ID=your-project-id export TPU_TYPE=v6e-256 export NUM_SLICES=2 export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE} export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE} export NETWORK_FW_NAME_1=${NETWORK_NAME_1}-fw-1-${ZONE} export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-1-${ZONE} export ROUTER_NAME=${CLUSTER_NAME}-network-1-${ZONE} export NAT_CONFIG=${CLUSTER_NAME}-natconfig-1-${ZONE}
gcloud compute networks create ${NETWORK_NAME_1} \ --mtu=8896 \ --bgp-routing-mode=regional \ --subnet-mode=custom \ --project=${PROJECT_ID}
gcloud compute networks subnets create ${SUBNET_NAME_1} \ --network=${NETWORK_NAME_1} \ --range=10.11.0.0/18 \ --region=${REGION} \ --project=${PROJECT_ID}
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \ --network=${NETWORK_NAME_1} \ --allow tcp,icmp,udp \ --project=${PROJECT_ID}
gcloud compute routers create ${ROUTER_NAME} \ --project=${PROJECT_ID} \ --network=${NETWORK_NAME_1} \ --region=${REGION}
gcloud compute routers nats create ${NAT_CONFIG} \ --router=${ROUTER_NAME} \ --region=${REGION} \ --auto-allocate-nat-external-ips \ --nat-all-subnet-ip-ranges \ --project=${PROJECT_ID} \ --enable-logging
# Secondary subnet for multi-nic experience.
# Need custom IP routing to be different from the first network's subnet.
export NETWORK_NAME_2=${CLUSTER_NAME}-privatenetwork-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-2-${ZONE}
gcloud compute networks create ${NETWORK_NAME_2} \ --mtu=8896 \ --bgp-routing-mode=regional \ --subnet-mode=custom \ --project=${PROJECT_ID}
gcloud compute networks subnets create ${SUBNET_NAME_2} \ --network=${NETWORK_NAME_2} \ --range=10.10.0.0/18 \ --region=${REGION} \ --project=${PROJECT_ID}
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \ --network=${NETWORK_NAME_2} \ --allow tcp,icmp,udp \ --project=${PROJECT_ID}
gcloud compute routers create ${ROUTER_NAME} \ --project=${PROJECT_ID} \ --network=${NETWORK_NAME_2} \ --region=${REGION}
gcloud compute routers nats create ${NAT_CONFIG} \ --router=${ROUTER_NAME} \ --region=${REGION} \ --auto-allocate-nat-external-ips \ --nat-all-subnet-ip-ranges \ --project=${PROJECT_ID} \ --enable-logging
export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking
--network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network
network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"
python3 ~/xpk/xpk.py cluster create \
--cluster=${CLUSTER_NAME} \
--num-slices=${NUM_SLICES} \
--tpu-type=${TPU_TYPE} \
--zone=${ZONE} \
--project=${PROJECT_ID} \
--on-demand \
--custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
--custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
--create-vertex-tensorboard
Deskripsi flag perintah
Variabel | Deskripsi |
CLUSTER_NAME | Nama yang ditetapkan pengguna untuk cluster XPK. |
PROJECT_ID | Google Cloud nama project. Gunakan project yang sudah ada atau buat project baru. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan Google Cloud project. |
ZONA | Lihat dokumen Region dan zona Cloud TPU untuk mengetahui zona yang didukung. |
TPU_TYPE | Lihat Jenis Akselerator. |
NUM_SLICES | Jumlah slice yang ingin Anda buat |
CLUSTER_ARGUMENTS | Jaringan dan subnetwork yang akan digunakan.
Contoh: |
NODE_POOL_ARGUMENTS | Jaringan node tambahan yang akan digunakan.
Contoh: |
NUM_SLICES | Jumlah slice yang akan dibuat (hanya diperlukan untuk Multislice). |
NETWORK_NAME | Nama jaringan sekunder yang akan digunakan. |
NETWORK_FW_NAME | Nama firewall jaringan sekunder yang akan digunakan. |
Penyiapan framework
Bagian ini menjelaskan proses penyiapan umum untuk pelatihan model ML menggunakan framework JAX, PyTorch, atau TensorFlow. Jika menggunakan GKE, Anda dapat menggunakan perintah XPK atau Kubernetes untuk penyiapan framework.
Penyiapan untuk JAX
Bagian ini memberikan petunjuk penyiapan untuk menjalankan workload JAX di GKE, dengan atau tanpa XPK, serta menggunakan resource yang diantrekan.
Menyiapkan JAX menggunakan GKE
Satu slice di satu host
Contoh berikut menyiapkan node pool satu host 2x2 menggunakan file YAML Kubernetes.
apiVersion: v1
kind: Pod
metadata:
name: tpu-pod-jax-v6e-a
spec:
restartPolicy: Never
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 2x2
containers:
- name: tpu-job
image: python:3.10
securityContext:
privileged: true
command:
- bash
- -c
- |
pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true python3 -c 'import jax; print("Total TPU chips:", jax.device_count())'
resources:
requests:
google.com/tpu: 4
limits:
google.com/tpu: 4
Setelah berhasil diselesaikan, Anda akan melihat pesan berikut di log GKE:
Total TPU chips: 4
Satu slice di multi-host
Contoh berikut menyiapkan node pool multi-host 4x4 menggunakan file YAML Kubernetes.
apiVersion: v1
kind: Service
metadata:
name: headless-svc
spec:
clusterIP: None
selector:
job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
name: tpu-available-chips
spec:
backoffLimit: 0
completions: 4
parallelism: 4
completionMode: Indexed
template:
spec:
subdomain: headless-svc
restartPolicy: Never
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 4x4
containers:
- name: tpu-job
image: python:3.10
ports:
- containerPort: 8471 # Default port using which TPU VMs communicate
- containerPort: 8431 # Port to export TPU runtime metrics, if supported.
securityContext:
privileged: true
command:
- bash
- -c
- |
pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true python -c 'import jax; print("Total TPU chips:", jax.device_count())'
resources:
requests:
google.com/tpu: 4
limits:
google.com/tpu: 4
Setelah berhasil diselesaikan, Anda akan melihat pesan berikut di log GKE:
Total TPU chips: 16
Multi-slice di multi-host
Contoh berikut menyiapkan dua node pool multi-host 4x4 menggunakan file YAML Kubernetes.
Sebagai prasyarat, Anda harus menginstal JobSet v0.2.3 atau yang lebih baru.
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: multislice-job
annotations:
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
failurePolicy:
maxRestarts: 4
replicatedJobs:
- name: slice
replicas: 2
template:
spec:
parallelism: 4
completions: 4
backoffLimit: 0
template:
spec:
hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 4x4
hostNetwork: true
containers:
- name: jax-tpu
image: python:3.10
ports:
- containerPort: 8471
- containerPort: 8080
- containerPort: 8431
securityContext:
privileged: true
command:
- bash
- -c
- |
pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true python -c 'import jax; print("Total TPU chips:", jax.device_count())'
resources:
limits:
google.com/tpu: 4
requests:
google.com/tpu: 4
Setelah berhasil diselesaikan, Anda akan melihat pesan berikut di log GKE:
Total TPU chips: 32
Untuk mengetahui informasi selengkapnya, lihat Menjalankan workload Multislice dalam dokumentasi GKE.
Untuk performa yang lebih baik, Aktifkan hostNetwork.
Multi-NIC
Untuk memanfaatkan multi-NIC di GKE, manifes Pod Kubernetes harus memiliki anotasi tambahan. Berikut adalah contoh manifes workload multi-NIC non-TPU.
apiVersion: v1
kind: Pod
metadata:
name: sample-netdevice-pod-1
annotations:
networking.gke.io/default-interface: 'eth0'
networking.gke.io/interfaces: |
[
{"interfaceName":"eth0","network":"default"},
{"interfaceName":"eth1","network":"netdevice-network"}
]
spec:
containers:
- name: sample-netdevice-pod
image: busybox
command: ["sleep", "infinity"]
ports:
- containerPort: 80
restartPolicy: Always
tolerations:
- key: "google.com/tpu"
operator: "Exists"
effect: "NoSchedule"
Jika menggunakan perintah exec
untuk terhubung ke Pod Kubernetes, Anda akan melihat NIC tambahan menggunakan kode berikut.
$ k exec --stdin --tty sample-netdevice-pod-1 -- /bin/sh
/ # ip a
1: lo: <LOOPBACK,UP,LOWER_UP> mtu 65536 qdisc noqueue qlen 1000
link/loopback 00:00:00:00:00:00 brd 00:00:00:00:00:00
inet 127.0.0.1/8 scope host lo
valid_lft forever preferred_lft forever
2: eth0@if11: <BROADCAST,MULTICAST,UP,LOWER_UP,M-DOWN> mtu 1460 qdisc noqueue
link/ether da:be:12:67:d2:25 brd ff:ff:ff:ff:ff:ff
inet 10.124.2.6/24 brd 10.124.2.255 scope global eth0
valid_lft forever preferred_lft forever
3: eth1: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1460 qdisc mq qlen 1000
link/ether 42:01:ac:18:00:04 brd ff:ff:ff:ff:ff:ff
inet 172.24.0.4/32 scope global eth1
valid_lft forever preferred_lft forever
Menyiapkan JAX menggunakan GKE dengan XPK
Untuk menyiapkan JAX menggunakan GKE dan XPK, lihat README xpk.
Untuk menyiapkan dan menjalankan XPK dengan MaxText, lihat Cara menjalankan MaxText.
Menyiapkan JAX menggunakan resource yang diantrekan
Instal JAX di semua VM Cloud TPU di slice atau beberapa slice secara bersamaan menggunakan perintah gcloud alpha compute tpus tpu-vm ssh
. Untuk Multislice, tambahkan flag --node=all
.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \
--zone ${ZONE} --worker=all \
--command='pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Anda dapat menjalankan perintah berikut untuk memeriksa jumlah core Cloud TPU yang tersedia di slice dan menguji apakah semuanya diinstal dengan benar:
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \
--zone ${ZONE} --worker=all \
--command='python3 -c "import jax; print(jax.device_count(), jax.local_device_count())"'
Outputnya mirip dengan berikut ini saat berjalan di slice v6e-16:
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16 4
16 4
16 4
16 4
jax.device_count()
menampilkan jumlah total chip dalam slice yang diberikan.
jax.local_device_count()
menunjukkan jumlah chip yang dapat diakses oleh satu VM di slice ini.
gcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} --zone=${ZONE} --worker=all \
--command='git clone -b mlperf4.1 https://github.com/google/maxdiffusion.git &&
cd maxdiffusion && git checkout 975fdb7dbddaa9a53ad72a421cdb487dcdc491a3 &&
&& pip install -r requirements.txt && pip install . '
Memecahkan masalah penyiapan JAX
Tips umum adalah mengaktifkan logging panjang dalam manifes workload GKE Anda. Kemudian, berikan log ke dukungan GKE.
TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0
Pesan error
no endpoints available for service 'jobset-webhook-service'
Error ini berarti set tugas tidak diinstal dengan benar. Periksa untuk melihat apakah Pod Kubernetes deployment jobset-controller-manager sedang berjalan. Untuk mengetahui informasi selengkapnya, lihat dokumentasi pemecahan masalah JobSet untuk mengetahui detailnya.
TPU initialization failed: Failed to connect
Pastikan versi node GKE Anda adalah 1.30.4-gke.1348000 atau yang lebih baru (GKE 1.31 tidak didukung).
Penyiapan untuk PyTorch
Bagian ini menjelaskan cara mulai menggunakan PJRT di v6e dengan PyTorch/XLA. Python 3.10 adalah versi Python yang direkomendasikan.
Menyiapkan PyTorch menggunakan GKE dengan XPK
Anda dapat menggunakan penampung Docker berikut dengan XPK yang telah menginstal dependensi PyTorch:
us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_20241028
Untuk membuat beban kerja XPK, gunakan perintah berikut:
python3 xpk.py workload create \
--cluster ${CLUSTER_NAME} \
{--docker-image | --base-docker-image} us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_20241028 \
--workload ${USER} -xpk-${ACCELERATOR_TYPE} -${NUM_SLICES} \
--tpu-type=${ACCELERATOR_TYPE} \
--num-slices=${NUM_SLICES} \
--on-demand \
--zone ${ZONE} \
--project ${PROJECT_ID} \
--enable-debug-logs \
--command 'python3 -c "import torch; import torch_xla; import torch_xla.runtime as xr; print(xr.global_runtime_device_count())"'
Menggunakan --base-docker-image
akan membuat image Docker baru dengan direktori kerja
saat ini yang di-build ke dalam Docker baru.
Menyiapkan PyTorch menggunakan resource dalam antrean
Ikuti langkah-langkah berikut untuk menginstal PyTorch menggunakan resource yang diantrekan dan menjalankan skrip kecil di v6e.
Menginstal dependensi menggunakan SSH untuk mengakses VM
Gunakan perintah berikut untuk menginstal dependensi di semua VM Cloud TPU. Untuk
Multislice, tambahkan tanda --worker=all
:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='sudo apt install -y libopenblas-base pip3 && \
pip install torch~=2.6.0 "torch_xla[tpu]~=2.6.0" -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html'
Meningkatkan performa model dengan alokasi yang cukup besar dan sering
Untuk model yang memiliki alokasi yang cukup besar dan sering, penggunaan fungsi tcmalloc
akan meningkatkan performa secara signifikan dibandingkan dengan implementasi fungsi malloc
default, sehingga fungsi malloc
default yang digunakan di VM Cloud TPU adalah
tcmalloc
. Namun, bergantung pada beban kerja Anda (misalnya, dengan DLRM yang
memiliki alokasi yang sangat besar untuk tabel penyematan), fungsi tcmalloc
dapat
menyebabkan pelambatan. Dalam hal ini, Anda dapat mencoba membatalkan penetapan variabel berikut
menggunakan fungsi malloc
default:
unset LD_PRELOAD
Menggunakan skrip Python untuk melakukan penghitungan di VM v6e
Gunakan perintah berikut untuk menjalankan skrip yang membuat dua tensor, menambahkannya bersama-sama, dan mencetak hasilnya.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} --worker all --command='
unset LD_PRELOAD
python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"
'
Tindakan ini akan menghasilkan output yang mirip dengan berikut ini:
SSH: Attempting to connect to worker 0...
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
xla:0
tensor([[ 0.3355, -1.4628, -3.2610],
[-1.4656, 0.3196, -2.8766],
[ 0.8668, -1.5060, 0.7125]], device='xla:0')
Penyiapan untuk TensorFlow
Anda dapat mereset runtime Cloud TPU dengan versi TensorFlow yang kompatibel dengan v6e dengan menjalankan perintah berikut:
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \
--zone ${ZONE} --worker=all --command="sudo sed -i 's/TF_DOCKER_URL=.*/TF_DOCKER_URL=gcr.io\/cloud-tpu-v2-images\/grpc_tpu_worker:v6e\"/' /etc/systemd/system/tpu-runtime.service"
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \
--zone ${ZONE} --worker=all --command='sudo systemctl daemon-reload && sudo systemctl restart tpu-runtime'
Gunakan SSH untuk mengakses worker-0:
$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \
--zone ${ZONE}
Instal TensorFlow di pekerja-0:
sudo apt install -y libopenblas-base
pip install cloud-tpu-client
pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310
pip install cloud-tpu-client
pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl \
-f https://storage.googleapis.com/libtpu-tf-releases/index.html --force
Ekspor variabel lingkungan TPU_NAME
:
export TPU_NAME=v6e-16
Anda dapat menjalankan skrip Python berikut untuk memeriksa jumlah core Cloud TPU yang tersedia di slice dan untuk menguji apakah semuanya diinstal dengan benar:
import TensorFlow as tf
print("TensorFlow version " + tf.__version__)
@tf.function
def add_fn(x,y):
z = x + y
return z
cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.TPUStrategy(cluster_resolver)
x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)
Outputnya mirip dengan berikut ini saat berjalan di slice v6e-16:
PerReplica:{
0: tf.Tensor(2.0, shape=(), dtype=float32),
1: tf.Tensor(2.0, shape=(), dtype=float32),
2: tf.Tensor(2.0, shape=(), dtype=float32),
3: tf.Tensor(2.0, shape=(), dtype=float32),
4: tf.Tensor(2.0, shape=(), dtype=float32),
5: tf.Tensor(2.0, shape=(), dtype=float32),
6: tf.Tensor(2.0, shape=(), dtype=float32),
7: tf.Tensor(2.0, shape=(), dtype=float32)
}
v6e dengan SkyPilot
Anda dapat menggunakan Cloud TPU v6e dengan SkyPilot. Gunakan langkah-langkah berikut untuk menambahkan informasi lokasi dan harga terkait v6e ke SkyPilot.
Tambahkan kode berikut ke bagian akhir file
~/.sky/catalogs/v5/gcp/vms.csv
:,,,tpu-v6e-1,1,tpu-v6e-1,us-south1,us-south1-a,0,0 ,,,tpu-v6e-1,1,tpu-v6e-1,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-1,1,tpu-v6e-1,us-east5,us-east5-b,0,0 ,,,tpu-v6e-4,1,tpu-v6e-4,us-south1,us-south1-a,0,0 ,,,tpu-v6e-4,1,tpu-v6e-4,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-4,1,tpu-v6e-4,us-east5,us-east5-b,0,0 ,,,tpu-v6e-8,1,tpu-v6e-8,us-south1,us-south1-a,0,0 ,,,tpu-v6e-8,1,tpu-v6e-8,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-8,1,tpu-v6e-8,us-east5,us-east5-b,0,0 ,,,tpu-v6e-16,1,tpu-v6e-16,us-south1,us-south1-a,0,0 ,,,tpu-v6e-16,1,tpu-v6e-16,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-16,1,tpu-v6e-16,us-east5,us-east5-b,0,0 ,,,tpu-v6e-32,1,tpu-v6e-32,us-south1,us-south1-a,0,0 ,,,tpu-v6e-32,1,tpu-v6e-32,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-32,1,tpu-v6e-32,us-east5,us-east5-b,0,0 ,,,tpu-v6e-64,1,tpu-v6e-64,us-south1,us-south1-a,0,0 ,,,tpu-v6e-64,1,tpu-v6e-64,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-64,1,tpu-v6e-64,us-east5,us-east5-b,0,0 ,,,tpu-v6e-128,1,tpu-v6e-128,us-south1,us-south1-a,0,0 ,,,tpu-v6e-128,1,tpu-v6e-128,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-128,1,tpu-v6e-128,us-east5,us-east5-b,0,0 ,,,tpu-v6e-256,1,tpu-v6e-256,us-south1,us-south1-a,0,0 ,,,tpu-v6e-256,1,tpu-v6e-256,europe-west4,europe-west4-a,0,0 ,,,tpu-v6e-256,1,tpu-v6e-256,us-east5,us-east5-b,0,0
Tentukan resource berikut dalam file YAML:
# tpu_v6.yaml resources: accelerators: tpu-v6e-16 # Fill in the accelerator type you want to use accelerator_args: runtime_version: v2-alpha-tpuv6e # Official suggested runtime
Luncurkan cluster dengan Cloud TPU v6e:
sky launch tpu_v6.yaml -c tpu_v6
Hubungkan ke Cloud TPU v6e menggunakan SSH:
ssh tpu_v6
Tutorial inferensi
Tutorial berikut menunjukkan cara menjalankan inferensi di Cloud TPU v6e:
Contoh pelatihan
Bagian berikut memberikan contoh untuk melatih model MaxText, MaxDiffusion, dan PyTorch di Cloud TPU v6e.
Pelatihan MaxText dan MaxDiffusion di VM Cloud TPU v6e
Bagian berikut membahas siklus proses pelatihan model MaxText dan MaxDiffusion.
Secara umum, langkah-langkah tingkat tingginya adalah:
- Build image dasar workload.
- Jalankan workload Anda menggunakan XPK.
- Buat perintah pelatihan untuk beban kerja.
- Deploy workload.
- Ikuti beban kerja dan lihat metrik.
- Hapus workload XPK jika tidak diperlukan.
- Hapus cluster XPK jika tidak diperlukan lagi.
Mem-build image dasar
Instal MaxText atau MaxDiffusion dan build image Docker:
Clone repositori yang ingin Anda gunakan dan ubah ke direktori untuk repositori tersebut:
MaxText:
git clone https://github.com/google/maxtext.git && cd maxtext
MaxDiffusion:
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
Konfigurasikan Docker untuk menggunakan Google Cloud CLI:
gcloud auth configure-docker
Build image Docker menggunakan perintah berikut atau menggunakan JAX Stable Stack. Untuk mengetahui informasi selengkapnya tentang JAX Stable Stack, lihat Mem-build image Docker dengan JAX Stable Stack.
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35
Jika Anda meluncurkan beban kerja dari mesin yang tidak memiliki image yang di-build secara lokal, upload image:
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${USER}_runner
Menjalankan workload menggunakan XPK
Tetapkan variabel lingkungan berikut jika Anda tidak menggunakan nilai default yang ditetapkan oleh MaxText atau MaxDiffusion:
export BASE_OUTPUT_DIR=gs://YOUR_BUCKET export PER_DEVICE_BATCH_SIZE=2 export NUM_STEPS=30 export MAX_TARGET_LENGTH=8192
Build skrip model Anda. Skrip ini akan disalin sebagai perintah pelatihan pada langkah berikutnya.
Jangan jalankan skrip model terlebih dahulu.
MaxText
MaxText adalah LLM open source berperforma tinggi dan sangat skalabel yang ditulis dalam Python dan JAX murni serta menargetkan Google Cloud TPU dan GPU untuk pelatihan dan inferensi.
JAX_PLATFORMS=tpu,cpu \ ENABLE_PJRT_COMPATIBILITY=true \ TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \ TPU_SLICE_BUILDER_DUMP_ICI=true && \ python /deps/MaxText/train.py /deps/MaxText/configs/base.yml \ base_output_directory=${BASE_OUTPUT_DIR} \ dataset_type=synthetic \ per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ enable_checkpointing=false \ gcs_metrics=true \ profiler=xplane \ skip_first_n_steps_for_profiler=5 \ steps=${NUM_STEPS} # attention='dot_product'"
Gemma2
Gemma adalah serangkaian LLM bobot terbuka yang dikembangkan oleh Google DeepMind, berdasarkan riset dan teknologi Gemini.
python3 MaxText/train.py MaxText/configs/base.yml \ model_name=gemma2-27b \ run_name=gemma2-27b-run \ base_output_directory=${BASE_OUTPUT_DIR} \ max_target_length=${MAX_TARGET_LENGTH} \ per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ steps=${NUM_STEPS} \ enable_checkpointing=false \ use_iota_embed=true \ gcs_metrics=true \ dataset_type=synthetic \ profiler=xplane \ attention=flash
Mixtral 8x7b
Mixtral adalah model AI canggih yang dikembangkan oleh Mistral AI, yang menggunakan arsitektur campuran pakar (MoE) yang jarang.
python3 MaxText/train.py MaxText/configs/base.yml \ base_output_directory=${BASE_OUTPUT_DIR} \ per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ model_name=mixtral-8x7b \ steps=${NUM_STEPS} \ max_target_length=${MAX_TARGET_LENGTH} \ tokenizer_path=assets/tokenizer.mistral-v1 \ attention=flash \ dtype=bfloat16 \ dataset_type=synthetic \ profiler=xplane
Llama3-8b
Llama adalah serangkaian LLM bobot terbuka yang dikembangkan oleh Meta.
python3 MaxText/train.py MaxText/configs/base.yml \ model_name=llama3-8b \ base_output_directory=${BASE_OUTPUT_DIR} \ dataset_type=synthetic \ tokenizer_path=assets/tokenizer_llama3.tiktoken \ per_device_batch_size=${PER_DEVICE_BATCH_SIZE} # set to 4 \ gcs_metrics=true \ profiler=xplane \ skip_first_n_steps_for_profiler=5 \ steps=${NUM_STEPS} \ max_target_length=${MAX_TARGET_LENGTH} \ attention=flash
MaxDiffusion
MaxDiffusion adalah kumpulan implementasi referensi dari berbagai model difusi laten yang ditulis dalam Python murni dan JAX yang berjalan di perangkat XLA, termasuk Cloud TPU dan GPU. Stable Diffusion adalah model text-to-image laten yang menghasilkan gambar fotorealistik dari input teks apa pun.
Anda perlu menginstal cabang Git tertentu untuk menjalankan MaxDiffusion seperti yang ditunjukkan dalam perintah
git checkout
berikut.git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout e712c9fc4cca764b0930067b6e33daae2433abf0 && pip install -r requirements.txt && pip install .
Skrip pelatihan:
cd maxdiffusion && OUT_DIR=${BASE_OUTPUT_DIR} \ python src/maxdiffusion/train_sdxl.py \ src/maxdiffusion/configs/base_xl.yml \ revision=refs/pr/95 \ activations_dtype=bfloat16 \ weights_dtype=bfloat16 \ resolution=1024 \ per_device_batch_size=1 \ output_dir=${OUT_DIR} \ jax_cache_dir=${OUT_DIR}/cache_dir/ \ max_train_steps=200 \ attention=flash run_name=sdxl-ddp-v6e
Ekspor variabel berikut:
export CLUSTER_NAME=CLUSTER_NAME export ACCELERATOR_TYPE=ACCELERATOR_TYPE export NUM_SLICES=NUM_SLICES export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT
Deskripsi variabel lingkungan
Variabel Deskripsi CLUSTER_NAME
Nama cluster XPK Anda. ACCELERATOR_TYPE
Lihat Jenis Akselerator. NUM_SLICES
Jumlah slice TPU. YOUR_MODEL_SCRIPT
Skrip model yang akan dieksekusi sebagai perintah pelatihan. Jalankan model menggunakan skrip yang Anda buat di langkah sebelumnya. Anda harus menentukan flag
--base-docker-image
untuk menggunakan gambar dasar MaxText atau menentukan flag--docker-image
dan gambar yang ingin Anda gunakan.Opsional: Anda dapat mengaktifkan logging debug dengan menyertakan flag
--enable-debug-logs
. Untuk informasi selengkapnya, lihat Men-debug JAX di MaxText.Opsional: Anda dapat membuat Eksperimen Vertex AI untuk mengupload data ke Vertex AI TensorBoard dengan menyertakan tanda
--use-vertex-tensorboard
. Untuk mengetahui informasi selengkapnya, lihat Memantau JAX di MaxText menggunakan Vertex AI.python3 xpk.py workload create \ --cluster ${CLUSTER_NAME} \ {--base-docker-image maxtext_base_image|--docker-image ${CLOUD_IMAGE_NAME}} \ --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \ --tpu-type=${ACCELERATOR_TYPE} \ --num-slices=${NUM_SLICES} \ --on-demand \ --zone=${ZONE} \ --project=${PROJECT_ID} \ [--enable-debug-logs] \ [--use-vertex-tensorboard] \ --command=${YOUR_MODEL_SCRIPT}
Outputnya menyertakan link untuk mengikuti beban kerja Anda. Buka link dan klik tab Logs untuk melacak beban kerja Anda secara real time.
Men-debug JAX di MaxText
Gunakan perintah XPK tambahan untuk mendiagnosis alasan cluster atau beban kerja tidak berjalan.
- Daftar workload XPK
- XPK inspector
- Aktifkan logging panjang di log beban kerja menggunakan flag
--enable-debug-logs
saat Anda membuat beban kerja XPK.
Memantau JAX di MaxText menggunakan Vertex AI
Lihat data skalar dan profil melalui TensorBoard terkelola Vertex AI.
- Tingkatkan permintaan pengelolaan resource (CRUD) untuk zona yang Anda gunakan dari 600 menjadi 5.000. Hal ini mungkin bukan masalah untuk beban kerja kecil yang menggunakan kurang dari 16 VM.
Instal dependensi seperti
cloud-accelerator-diagnostics
untuk Vertex AI:# xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI cd ~/xpk pip install .
Buat cluster XPK menggunakan flag
--create-vertex-tensorboard
, seperti yang didokumentasikan dalam Membuat Vertex AI TensorBoard. Anda juga dapat menjalankan perintah ini di cluster yang ada.Buat eksperimen Vertex AI saat menjalankan beban kerja XPK menggunakan tanda
--use-vertex-tensorboard
dan tanda--experiment-name
opsional. Untuk mengetahui daftar lengkap langkah-langkahnya, lihat Membuat Vertex AI Experiment untuk mengupload data ke Vertex AI TensorBoard.
Log menyertakan link ke Vertex AI TensorBoard, mirip dengan hal berikut:
View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name
Anda juga dapat menemukan link Vertex AI TensorBoard di Konsol Google Cloud. Buka Vertex AI Experiments di konsol Google Cloud. Pilih wilayah yang sesuai dari drop-down.
Direktori TensorBoard juga ditulis ke bucket Cloud Storage yang Anda tentukan dengan ${BASE_OUTPUT_DIR}
.
Menghapus workload XPK
Gunakan perintah xpk workload delete
untuk menghapus satu atau beberapa beban kerja berdasarkan awalan tugas atau status tugas. Perintah
ini mungkin berguna jika Anda mengirim workload XPK yang tidak perlu lagi dijalankan,
atau jika Anda memiliki tugas yang macet dalam antrean.
Menghapus cluster XPK
Gunakan perintah xpk cluster delete
untuk menghapus cluster:
python3 xpk.py cluster delete --cluster ${CLUSTER_NAME} \ --zone=${ZONE} --project=${PROJECT_ID}
Pelatihan Llama dan PyTorch/XLA di VM Cloud TPU v6e
Tutorial ini menjelaskan cara melatih model Llama menggunakan PyTorch/XLA di Cloud TPU v6e menggunakan set data WikiText.
Dapatkan akses ke Hugging Face dan model Llama 3
Anda memerlukan token akses pengguna Hugging Face untuk menjalankan tutorial ini. Untuk informasi tentang cara membuat dan menggunakan token akses pengguna, lihat dokumentasi Hugging Face tentang token akses pengguna.
Anda juga memerlukan izin untuk mengakses model Llama 3 8B di Hugging Face. Untuk mendapatkan akses, buka model Meta-Llama-3-8B di HuggingFace dan minta akses.
Membuat VM Cloud TPU
Buat Cloud TPU v6e dengan 8 chip untuk menjalankan tutorial.
Siapkan variabel lingkungan:
export ACCELERATOR_TYPE=v6e-8 export VERSION=v2-alpha-tpuv6e export TPU_NAME=$USER-$ACCELERATOR_TYPE export PROJECT_ID=your-project-id export ZONE=us-east1-d
Buat VM Cloud TPU:
gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${VERSION} \ --accelerator-type=${ACCELERATOR_TYPE} \ --zone=${ZONE} \ --project=${PROJECT_ID}
Penginstalan
Instal fork pytorch-tpu/transformers
transformer dan dependensi Hugging Face. Tutorial ini telah diuji dengan
versi dependensi berikut yang digunakan dalam contoh ini:
torch
: kompatibel dengan 2.5.0torch_xla[tpu]
: kompatibel dengan 2.5.0jax
: 0.4.33jaxlib
: 0.4.33
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} --zone ${ZONE} \ --worker=all --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git cd transformers sudo pip3 install -e . pip3 install datasets pip3 install evaluate pip3 install scikit-learn pip3 install accelerate pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html pip install jax==0.4.38 jaxlib==0.4.38 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html'
Menyiapkan konfigurasi model
Perintah pelatihan di bagian berikutnya, Jalankan model, menggunakan dua file konfigurasi JSON untuk menentukan parameter model dan konfigurasi FSDP (Fully Sharded Data Parallel). Sharding FSDP digunakan agar bobot model sesuai dengan ukuran batch yang lebih besar saat pelatihan. Saat berlatih dengan model yang lebih kecil, mungkin cukup menggunakan paralelisme data dan mereplikasi bobot di setiap perangkat. Untuk mengetahui informasi selengkapnya tentang cara melakukan shard tensor di seluruh perangkat di PyTorch/XLA, lihat Panduan Pengguna SPMD PyTorch/XLA.
Buat file konfigurasi parameter model. Berikut adalah konfigurasi parameter model untuk Llama3-8B. Untuk model lainnya, temukan konfigurasi di Hugging Face. Misalnya, lihat konfigurasi Llama2-7B.
cat > llama-config.json << EOF { "architectures": [ "LlamaForCausalLM" ], "attention_bias": false, "attention_dropout": 0.0, "bos_token_id": 128000, "eos_token_id": 128001, "hidden_act": "silu", "hidden_size": 4096, "initializer_range": 0.02, "intermediate_size": 14336, "max_position_embeddings": 8192, "model_type": "llama", "num_attention_heads": 32, "num_hidden_layers": 32, "num_key_value_heads": 8, "pretraining_tp": 1, "rms_norm_eps": 1e-05, "rope_scaling": null, "rope_theta": 500000.0, "tie_word_embeddings": false, "torch_dtype": "bfloat16", "transformers_version": "4.40.0.dev0", "use_cache": false, "vocab_size": 128256 } EOF
Buat file konfigurasi FSDP:
cat > fsdp-config.json << EOF { "fsdp_transformer_layer_cls_to_wrap": [ "LlamaDecoderLayer" ], "xla": true, "xla_fsdp_v2": true, "xla_fsdp_grad_ckpt": true } EOF
Untuk informasi selengkapnya tentang FSDP, lihat FSDPv2.
Upload file konfigurasi ke VM Cloud TPU Anda menggunakan perintah berikut:
gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \ --worker=all \ --project=${PROJECT_ID} \ --zone=${ZONE}
Menjalankan model
Dengan menggunakan file konfigurasi yang Anda buat di bagian sebelumnya, jalankan skrip run_clm.py
untuk melatih model Llama 3 8B pada set data WikiText. Skrip pelatihan
memerlukan waktu sekitar 10 menit untuk dijalankan di Cloud TPU v6e-8.
Login ke Hugging Face di Cloud TPU Anda menggunakan perintah berikut:
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT} \ --zone ${ZONE} \ --worker=all \ --command=' pip3 install "huggingface_hub[cli]" huggingface-cli login --token HUGGING_FACE_TOKEN'
Jalankan pelatihan model:
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT} \ --zone ${ZONE} \ --worker=all \ --command=' export PJRT_DEVICE=TPU export XLA_USE_SPMD=1 export ENABLE_PJRT_COMPATIBILITY=true # Optional variables for debugging: export XLA_IR_DEBUG=1 export XLA_HLO_DEBUG=1 export PROFILE_EPOCH=0 export PROFILE_STEP=3 export PROFILE_DURATION_MS=100000 # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path export PROFILE_LOGDIR=PROFILE_PATH python3 transformers/examples/pytorch/language-modeling/run_clm.py \ --dataset_name wikitext \ --dataset_config_name wikitext-2-raw-v1 \ --per_device_train_batch_size 16 \ --do_train \ --output_dir /home/$USER/tmp/test-clm \ --overwrite_output_dir \ --config_name /home/$USER/llama-config.json \ --cache_dir /home/$USER/cache \ --tokenizer_name meta-llama/Meta-Llama-3-8B \ --block_size 8192 \ --optim adafactor \ --save_strategy no \ --logging_strategy no \ --fsdp "full_shard" \ --fsdp_config /home/$USER/fsdp-config.json \ --torch_dtype bfloat16 \ --dataloader_drop_last yes \ --flash_attention \ --max_steps 20'
Memecahkan masalah PyTorch/XLA
Jika Anda menetapkan variabel opsional untuk proses debug di bagian sebelumnya,
profil untuk model akan disimpan di lokasi yang ditentukan oleh
variabel PROFILE_LOGDIR
. Anda dapat mengekstrak file xplane.pb
yang disimpan
di lokasi ini dan menggunakan tensorboard
untuk melihat profil di browser
menggunakan petunjuk TensorBoard
Jika PyTorch/XLA tidak berperforma seperti yang diharapkan, lihat panduan pemecahan masalah,
yang berisi saran untuk men-debug, membuat profil, dan mengoptimalkan model Anda.
Pelatihan DLRM DCN v2 di v6e
Tutorial ini menunjukkan cara melatih model DLRM DCN v2 di Cloud TPU v6e. Anda perlu menyediakan TPU v6e dengan 64, 128, atau 256 chip.
Jika Anda menjalankan di TPU multi-host, reset tpu-runtime
dengan versi TensorFlow yang sesuai dengan menjalankan perintah berikut. Jika menjalankannya di TPU host tunggal, Anda tidak perlu menjalankan dua perintah berikut.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command="sudo sed -i 's/TF_DOCKER_URL=.*/TF_DOCKER_URL=gcr.io\/cloud-tpu-v2-images\/grpc_tpu_worker:v6e\"/' /etc/systemd/system/tpu-runtime.service"
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='sudo systemctl daemon-reload && sudo systemctl restart tpu-runtime'
Menghubungkan ke worker-0 menggunakan SSH
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} --project ${PROJECT_ID}
Menetapkan nama Cloud TPU
export TPU_NAME=your-tpu-name
Menjalankan DLRM v2
Salin cuplikan kode berikut dalam file bernama script.sh
:
pip install --user setuptools==65.5.0
pip install cloud-tpu-client
pip install gin-config && pip install tensorflow-datasets && pip install tf-keras-nightly --no-deps
pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl -f https://storage.googleapis.com/libtpu-tf-releases/index.html --force
git clone https://github.com/tensorflow/recommenders.git
git clone https://github.com/tensorflow/models.git
export PYTHONPATH=~/recommenders/:~/models/
export TF_XLA_FLAGS='--tf_mlir_enable_mlir_bridge=true --tf_xla_sparse_core_disable_table_stacking=true --tf_mlir_enable_convert_control_to_data_outputs_pass=true --tf_mlir_enable_merge_control_flow_pass=true'
TF_USE_LEGACY_KERAS=1 TPU_LOAD_LIBRARY=0 python3 ./models/official/recommendation/ranking/train.py --mode=train --model_dir=gs://ptxla-debug/tf/sc/dlrm/runs/2/ --params_override="
runtime:
distribution_strategy: tpu
mixed_precision_dtype: 'mixed_bfloat16'
task:
use_synthetic_data: false
use_tf_record_reader: true
train_data:
input_path: 'gs://trillium-datasets/criteo/train/day_*/*'
global_batch_size: 16384
use_cached_data: true
validation_data:
input_path: 'gs://trillium-datasets/criteo/eval/day_*/*'
global_batch_size: 16384
use_cached_data: true
model:
num_dense_features: 13
bottom_mlp: [512, 256, 128]
embedding_dim: 128
interaction: 'multi_layer_dcn'
dcn_num_layers: 3
dcn_low_rank_dim: 512
size_threshold: 8000
top_mlp: [1024, 1024, 512, 256, 1]
use_multi_hot: true
concat_dense: false
dcn_use_bias: true
vocab_sizes: [40000000,39060,17295,7424,20265,3,7122,1543,63,40000000,3067956,405282,10,2209,11938,155,4,976,14,40000000,40000000,40000000,590152,12973,108,36]
multi_hot_sizes: [3,2,1,2,6,1,1,1,1,7,3,8,1,6,9,5,1,1,1,12,100,27,10,3,1,1]
max_ids_per_chip_per_sample: 128
max_ids_per_table: [280, 128, 64, 272, 432, 624, 64, 104, 368, 352, 288, 328, 304, 576, 336, 368, 312, 392, 408, 552, 2880, 1248, 720, 112, 320, 256]
max_unique_ids_per_table: [104, 56, 40, 32, 72, 32, 40, 32, 32, 144, 64, 192, 32, 40, 136, 32, 32, 32, 32, 240, 1352, 432, 120, 80, 32, 32]
use_partial_tpu_embedding: false
size_threshold: 0
initialize_tables_on_host: true
trainer:
train_steps: 10000
validation_interval: 1000
validation_steps: 660
summary_interval: 1000
steps_per_loop: 1000
checkpoint_interval: 0
optimizer_config:
embedding_optimizer: 'Adagrad'
dense_optimizer: 'Adagrad'
lr_config:
decay_exp: 2
decay_start_steps: 70000
decay_steps: 30000
learning_rate: 0.025
warmup_steps: 0
dense_sgd_config:
decay_exp: 2
decay_start_steps: 70000
decay_steps: 30000
learning_rate: 0.00025
warmup_steps: 8000
train_tf_function: true
train_tf_while_loop: true
eval_tf_while_loop: true
use_orbit: true
pipeline_sparse_and_dense_execution: true"
Jika Anda menjalankan TensorFlow di GKE, instal wheel TensorFlow Cloud TPU dan libtpu menggunakan perintah berikut:
pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl \
-f https://storage.googleapis.com/libtpu-tf-releases/index.html --force
Tetapkan flag berikut, yang diperlukan untuk menjalankan beban kerja rekomendasi (seperti DLRM DCN):
ENV TF_XLA_FLAGS='--tf_mlir_enable_mlir_bridge=true \
--tf_mlir_enable_tpu_variable_runtime_reformatting_pass=false \
--tf_mlir_enable_convert_control_to_data_outputs_pass=true \
--tf_mlir_enable_merge_control_flow_pass=true --tf_xla_disable_full_embedding_pipelining=true' \
ENV LIBTPU_INIT_ARGS="--xla_sc_splitting_along_feature_dimension=auto \
--copy_with_dynamic_shape_op_output_pjrt_buffer=true"
Jalankan script.sh
:
chmod +x script.sh
./script.sh
Hasil benchmark
Bagian berikut berisi hasil benchmark untuk DLRM DCN v2 dan MaxDiffusion di v6e.
DLRM DCN v2
Skrip pelatihan DLRM DCN v2 dijalankan pada skala yang berbeda. Lihat throughput dalam tabel berikut.
v6e-64 | v6e-128 | v6e-256 | |
---|---|---|---|
Langkah-langkah pelatihan | 7.000 | 7.000 | 7.000 |
Ukuran batch global | 131072 | 262144 | 524288 |
Throughput (contoh/dtk) | 2975334 | 5111808 | 10066329 |
MaxDiffusion
Kami menjalankan skrip pelatihan untuk MaxDiffusion di v6e-4, v6e-16, dan dua v6e-16. Lihat throughput dalam tabel berikut.
v6e-4 | v6e-16 | Dua v6e-16 | |
---|---|---|---|
Langkah-langkah pelatihan | 0,069 | 0,073 | 0,13 |
Ukuran batch global | 8 | 32 | 64 |
Throughput (contoh/dtk) | 115,9 | 438,4 | 492,3 |