Skip to content

Commit 9f57b66

Browse files
raymondzouuraymondzouu
committed
Add MaxText Llama 3.1 70B training with GCS recipe
1 parent af2a7cd commit 9f57b66

File tree

5 files changed

+295
-0
lines changed

5 files changed

+295
-0
lines changed

‎training/trillium/Llama3.1-70B-MaxText-with-Storage/README.md

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
# Instructions for training Llama3.1-70B-MaxText on TPU trillium (v6e-256) with Google Cloud Storage (GCS)
2+
3+
## GCS Bucket setup
4+
1. Create two buckets: one to hold the dataset and one to use for checkpoints. To create regional HNS buckets use the following commands:
5+
```
6+
# Set variables
7+
export DATASET_BUCKET="dataloading-bucket-name"
8+
export CHECKPOINT_BUCKET="checkpoint-bucket-name"
9+
export DATASET_STORAGE_NAME="dataset-bucket"
10+
export CHECKPOINT_STORAGE_NAME="checkpoint-bucket"
11+
export REGION="us-central1"
12+
13+
# Create dataset bucket
14+
gcloud storage buckets create gs://${DATASET_BUCKET} --location=${REGION} --default-storage-class=Standard --enable-hierarchical-namespace --uniform-bucket-level-access
15+
16+
# Create checkpoint bucket
17+
gcloud storage buckets create gs://${CHECKPOINT_BUCKET} --location=${REGION} --default-storage-class=Standard --enable-hierarchical-namespace --uniform-bucket-level-access
18+
```
19+
Replace the following values:
20+
- `<DATASET_BUCKET>`:the name of your Cloud Storage bucket with training dataset. Do not include the gs:// prefix
21+
- `<CHECKPOINT_BUCKET>`: the name of your Cloud Storage bucket where checkpoints will written. Do not include the gs:// prefix
22+
- `<REGION>`: the region where your cluster is located ([available locations](https://cloud.google.com/storage/docs/locations#location-r))
23+
24+
2. Follow these [instructions](https://github.com/AI-Hypercomputer/maxtext/blob/b93beba652db6b3f4e6c82dc48a83b03229f5d3a/getting_started/Data_Input_Pipeline.md#tfds-pipeline) to download the Allenai c4 dataset to the dataset bucket.
25+
Then follow these [instructions](https://github.com/google/array_record/tree/main/beam) to convert the dataset into ArrayRecord.
26+
27+
## XPK setup
28+
1. Please follow this [link](https://github.com/AI-Hypercomputer/tpu-recipes/blob/main/training/trillium/XPK_README.md) to create your GKE cluster with XPK.
29+
2. GCSFuse lets you mount and access Cloud Storage buckets as local file systems, so applications can read and write objects in your bucket using standard file system semantics. You'll need to use the below commands to create [XPK storage resources](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#storage) for both the dataset and checkpoint buckets in order to mount them to the MaxText workload using GCSFuse. For the dataset bucket and checkpoint bucket use separate manifest files `checkpoint_pvc.yaml` and `dataset_pvc.yaml` from this repo.
30+
Be sure to update `volumeHandle` in the yamls with your correct bucket names. Creating a bucket and xpk storage is a one time setup.
31+
```
32+
33+
export RECIPE_REPO="path-to-this-recipe-repo" # Update
34+
35+
cd ~/xpk
36+
37+
python3 xpk.py storage attach $DATASET_STORAGE_NAME type=gcsfuse project=$PROJECT cluster=$CLUSTER zone=$ZONE mountpoint=/tmp/dataset readonly=false bucket=$DATASET_BUCKET size=64 automount=false manifest=$RECIPE_REPO/tpu-recipes/training/trillium/Llama3.1-70B-MaxText-with-Storage/dataset_pvc.yaml
38+
39+
python3 xpk.py storage attach $CHECKPOINT_STORAGE_NAME type=gcsfuse project=$PROJECT cluster=$CLUSTER zone=$ZONE mountpoint=/tmp/ckpt readonly=false bucket=$CHECKPOINT_BUCKET size=64 automount=false manifest=$RECIPE_REPO/tpu-recipes/training/trillium/Llama3.1-70B-MaxText-with-Storage/checkpoint_pvc.yaml
40+
```
41+
42+
43+
## Prep for MaxText
44+
45+
### Install MaxText and Build Docker Image
46+
Please follow this [link](https://github.com/AI-Hypercomputer/tpu-recipes/blob/main/training/trillium/MAXTEXT_README.md) to install maxtext and build the docker image.
47+
48+
In step 2, use the jax-stable-stack image containing JAX 0.5.2:
49+
```
50+
BASE_IMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.5.2-rev1
51+
bash docker_build_dependency_image.sh DEVICE=tpu MODE=stable_stack BASEIMAGE=${BASE_IMAGE}
52+
```
53+
54+
## Run MaxText Llama3.1-70B workloads on GKE
55+
56+
### Starting workload
57+
58+
From the MaxText root directory, start your Llama3.1-70B workload.
59+
60+
Run MaxText Llama 3.1 70B with synthetic data and no checkpointing:
61+
```
62+
python3 benchmarks/benchmark_runner.py xpk \
63+
project=$PROJECT \
64+
zone=$ZONE \
65+
device_type=v6e-256 \
66+
num_slices=1 \
67+
cluster_name=$CLUSTER \
68+
base_output_directory=$OUTPUT_DIR \
69+
model_name="llama3_1_70b_8192_synthetic" \
70+
num_steps=100 \
71+
base_docker_image=maxtext_base_image
72+
```
73+
74+
Run MaxText Llama 3.1 70B with checkpointing and loading real data from GCS:
75+
```
76+
python3 benchmarks/benchmark_runner.py xpk \
77+
project=$PROJECT \
78+
zone=$ZONE \
79+
device_type=v6e-256 \
80+
num_slices=1 \
81+
cluster_name=${CLUSTER} \
82+
base_output_directory=/tmp/ckpt \
83+
model_name="llama3_1_70b_8192_rd_ckpt_grain" \
84+
num_steps=100 \
85+
base_docker_image=maxtext_base_image \
86+
xpk_storage=$DATASET_STORAGE_NAME xpk_storage=$CHECKPOINT_STORAGE_NAME
87+
```
88+
89+
If you would like to run on multiple slices of v6e-256, you may modify the `--num_slices` flag.
90+
91+
### Workload Details
92+
93+
For reference, here are the `llama3_1_70b_8192_synthetic` and `llama3_1_70b_8192_rd_ckpt_grain` workload details:
94+
95+
```
96+
MaxTextModel(
97+
model_name="llama3_1-70b-8192",
98+
model_type="llama3.1-70b",
99+
tuning_params={
100+
"per_device_batch_size": 4,
101+
"ici_fsdp_parallelism": -1,
102+
"remat_policy": "custom",
103+
"decoder_layer_input": "offload",
104+
"query_proj": "offload",
105+
"key_proj": "offload",
106+
"value_proj": "offload",
107+
"max_target_length": 8192,
108+
"attention": "flash",
109+
"use_iota_embed": True,
110+
"dataset_path": "gs://max-datasets-rogue",
111+
"dataset_type": "synthetic",
112+
"enable_checkpointing": False,
113+
"sa_block_q": 2048,
114+
"sa_block_kv": 2048,
115+
"sa_block_kv_compute": 2048,
116+
"sa_block_q_dkv": 2048,
117+
"sa_block_kv_dkv": 2048,
118+
"sa_block_kv_dkv_compute": 2048,
119+
"sa_block_q_dq": 2048,
120+
"sa_block_kv_dq": 2048,
121+
"sa_use_fused_bwd_kernel": True,
122+
"profiler": "xplane",
123+
"skip_first_n_steps_for_profiler": 10,
124+
"profiler_steps": 5,
125+
},
126+
xla_flags=(
127+
xla_flags_library.DENSE_VMEM_LIMIT_FLAG
128+
+ xla_flags_library.LAYOUT_FOR_ALL_REDUCE_SCATTER
129+
+ xla_flags_library.DATA_PARALLEL_OVERLAP
130+
+ xla_flags_library.CF_FOR_ALL_GATHER
131+
+ xla_flags_library.HOST_OFFLOAD_FLAGS
132+
),
133+
)
134+
135+
136+
MaxTextModel(
137+
model_name="llama3_1_70b_8192_rd_ckpt_grain",
138+
model_type="llama3.1-70b",
139+
tuning_params={
140+
"per_device_batch_size": 2,
141+
"ici_fsdp_parallelism": -1,
142+
"remat_policy": "custom",
143+
"decoder_layer_input": "offload",
144+
"query_proj": "offload",
145+
"key_proj": "offload",
146+
"value_proj": "offload",
147+
"max_target_length": 8192,
148+
"attention": "flash",
149+
"use_iota_embed": True,
150+
"dataset_path": "/tmp/dataset",
151+
"dataset_type": "grain",
152+
"grain_train_files": "/tmp/dataset/array-record/c4/en/3.0.1/c4-train.array_record*",
153+
"grain_worker_count": 24,
154+
"enable_checkpointing": True,
155+
"async_checkpointing": True,
156+
"checkpoint_period": 20,
157+
"sa_block_q": 2048,
158+
"sa_block_kv": 2048,
159+
"sa_block_kv_compute": 2048,
160+
"sa_block_q_dkv": 2048,
161+
"sa_block_kv_dkv": 2048,
162+
"sa_block_kv_dkv_compute": 2048,
163+
"sa_block_q_dq": 2048,
164+
"sa_block_kv_dq": 2048,
165+
"sa_use_fused_bwd_kernel": True,
166+
},
167+
xla_flags=(
168+
xla_flags_library.DENSE_VMEM_LIMIT_FLAG
169+
+ xla_flags_library.LAYOUT_FOR_ALL_REDUCE_SCATTER
170+
+ xla_flags_library.DATA_PARALLEL_OVERLAP
171+
+ xla_flags_library.CF_FOR_ALL_GATHER
172+
+ xla_flags_library.HOST_OFFLOAD_FLAGS
173+
+ xla_flags_library.ENABLE_SPARSECORE_OFFLOADING_FOR_ALL_REDUCE
174+
+ " --xla_tpu_iova_dma_chunk_size_bytes=104857"
175+
),
176+
)
177+
```
178+
179+
This equivalent workload code can be found in the [maxtext_trillium_model_configs.py](https://github.com/AI-Hypercomputer/maxtext/blob/1e4d513ad70dd4074d975a9f7936295008d4b900/benchmarks/maxtext_trillium_model_configs.py#L1103-L1146) file within the MaxText repository.
180+
181+
## Clean-up
182+
Detach storage
183+
```
184+
# Detach dataset storage
185+
python3 xpk.py storage detach $DATASET_STORAGE_NAME \
186+
--project=$PROJECT --cluster=$CLUSTER --zone=$ZONE
187+
188+
# Detach checkpoint storage
189+
python3 xpk.py storage detach $CHECKPOINT_STORAGE_NAME \
190+
--project=$PROJECT --cluster=$CLUSTER --zone=$ZONE
191+
```
192+

‎training/trillium/Llama3.1-70B-MaxText-with-Storage/checkpoint_pvc.yaml

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
apiVersion: v1
2+
kind: PersistentVolume
3+
metadata:
4+
name: checkpoint-bucket-pv
5+
spec:
6+
accessModes:
7+
- ReadWriteMany
8+
capacity:
9+
storage: 64Gi
10+
persistentVolumeReclaimPolicy: Retain
11+
storageClassName: gcsfuse-sc # dummy storage class
12+
claimRef:
13+
namespace: default
14+
name: checkpoint-bucket-pvc
15+
mountOptions:
16+
- metadata-cache:ttl-secs:-1
17+
- metadata-cache:negative-ttl-secs:0
18+
- metadata-cache:stat-cache-max-size-mb:-1
19+
- metadata-cache:type-cache-max-size-mb:-1
20+
- file-cache:enable-parallel-downloads:false
21+
- file-system:kernel-list-cache-ttl-secs:0
22+
- write:enable-streaming-writes:true
23+
- file-system:precondition-errors:false
24+
csi:
25+
driver: gcsfuse.csi.storage.gke.io
26+
volumeHandle: checkpoint-bucket-name # Update with your checkpoint bucket name
27+
volumeAttributes:
28+
gcsfuseMetadataPrefetchOnMount: "true"
29+
---
30+
apiVersion: v1
31+
kind: PersistentVolumeClaim
32+
metadata:
33+
name: checkpoint-bucket-pvc
34+
namespace: defaultls
35+
spec:
36+
accessModes:
37+
- ReadWriteMany
38+
resources:
39+
requests:
40+
storage: 64Gi
41+
volumeName: checkpoint-bucket-pv
42+
storageClassName: gcsfuse-sc # dummy storage class

‎training/trillium/Llama3.1-70B-MaxText-with-Storage/dataset_pvc.yaml

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
apiVersion: v1
2+
kind: PersistentVolume
3+
metadata:
4+
name: dataset-bucket-pv
5+
spec:
6+
accessModes:
7+
- ReadWriteMany
8+
capacity:
9+
storage: 64Gi
10+
persistentVolumeReclaimPolicy: Retain
11+
storageClassName: gcsfuse-sc # dummy storage class
12+
claimRef:
13+
namespace: default
14+
name: dataset-bucket-pvc
15+
mountOptions:
16+
- metadata-cache:ttl-secs:-1
17+
- metadata-cache:stat-cache-max-size-mb:-1
18+
- metadata-cache:type-cache-max-size-mb:-1
19+
- file-cache:enable-parallel-downloads:false
20+
- file-system:kernel-list-cache-ttl-secs:-1
21+
- write:enable-streaming-writes:true
22+
csi:
23+
driver: gcsfuse.csi.storage.gke.io
24+
volumeHandle: dataloading-bucket-name # Update with your bucket name
25+
volumeAttributes:
26+
gcsfuseMetadataPrefetchOnMount: "true"
27+
---
28+
apiVersion: v1
29+
kind: PersistentVolumeClaim
30+
metadata:
31+
name: dataset-bucket-pvc
32+
namespace: default
33+
spec:
34+
accessModes:
35+
- ReadWriteMany
36+
resources:
37+
requests:
38+
storage: 64Gi
39+
volumeName: dataset-bucket-pv
40+
storageClassName: gcsfuse-sc # dummy storage class

‎training/trillium/Llama3.1-70B-MaxText-with-Storage/llama3-1-70B-real-data-and-ckpt-1xv6e-256.sh

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
python3 benchmarks/benchmark_runner.py xpk \
2+
project=$PROJECT \
3+
zone=$ZONE \
4+
device_type=v6e-256 \
5+
num_slices=1 \
6+
cluster_name=${CLUSTER} \
7+
base_output_directory=/tmp/ckpt \
8+
model_name="llama3_1_70b_8192_rd_ckpt_grain" \
9+
num_steps=100 \
10+
base_docker_image=maxtext_base_image \
11+
xpk_storage=$DATASET_STORAGE_NAME xpk_storage=$CHECKPOINT_STORAGE_NAME

‎training/trillium/Llama3.1-70B-MaxText-with-Storage/llama3-1-70B-synthetic-data-1xv6e-256.sh

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
python3 benchmarks/benchmark_runner.py xpk \
2+
project=$PROJECT \
3+
zone=$ZONE \
4+
device_type=v6e-256 \
5+
num_slices=1 \
6+
cluster_name=$CLUSTER \
7+
base_output_directory=$OUTPUT_DIR \
8+
model_name="llama3_1_70b_8192_synthetic" \
9+
num_steps=100 \
10+
base_docker_image=maxtext_base_image

0 commit comments

Comments
 (0)