Merge pull request #52 from AI-Hypercomputer/llama31_70b

Add Llama 3.1 70B training recipe

Author: raymondzouu

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

@@ -0,0 +1,90 @@
+# Instructions for training Llama3.1-70B-MaxText on TPU trillium (v6e-256)
+
+## XPK setup
+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
+
+## Prep for Maxtext
+
+### Install MaxText and Build Docker Image
+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. The following variables should be set:
+
+In step 1, use the MaxText [tpu-recipes-v0.1.0](https://github.com/AI-Hypercomputer/maxtext/releases/tag/tpu-recipes-v0.1.0) tag to run this recipe:
+```
+git checkout tpu-recipes-v0.1.0
+```
+
+In step 2, use the jax-stable-stack image containing JAX 0.5.2:
+```
+BASE_IMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.5.2-rev1
+bash docker_build_dependency_image.sh DEVICE=tpu MODE=stable_stack BASEIMAGE=${BASE_IMAGE}
+```
+
+## Run MaxText Llama3.1-70B workloads on GKE
+
+### Starting workload
+
+From the MaxText root directory, start your Llama3.1-70B workload
+```
+python3 benchmarks/benchmark_runner.py xpk \
+    --project=$PROJECT \
+    --zone=$ZONE \
+    --device_type=v6e-256 \
+    --num_slices=1  \
+    --cluster_name=${CLUSTER_NAME} \
+    --base_output_directory=${OUTPUT_DIR} \
+    --model_name="llama3_1_70b_8192" \
+    --base_docker_image=maxtext_base_image
+```
+
+From your workload logs, you should start seeing step time logs like the following:
+```
+completed step: 7, seconds: 34.562, TFLOP/s/device: 456.442, Tokens/s/device: 948.086, total_weights: 8388608, loss: 8.946
+```
+If you would like to run on multiple slices of v6e-256, you may modify the `--num_slices` flag.
+
+### Workload Details
+
+For reference, here are the `llama3_1_70b_8192` workload details as found in `[email protected]`:
+
+```
+  MaxTextModel(
+    model_name="llama3_1-70b-8192",
+    model_type="llama3.1-70b",
+    tuning_params={
+        "per_device_batch_size": 4,
+        "ici_fsdp_parallelism": -1,
+        "remat_policy": "custom",
+        "decoder_layer_input": "offload",
+        "query_proj": "offload",
+        "key_proj": "offload",
+        "value_proj": "offload",
+        "max_target_length": 8192,
+        "attention": "flash",
+        "use_iota_embed": True,
+        "dataset_path": "gs://max-datasets-rogue",
+        "dataset_type": "synthetic",
+        "enable_checkpointing": False,
+        "sa_block_q": 2048,
+        "sa_block_kv": 2048,
+        "sa_block_kv_compute": 2048,
+        "sa_block_q_dkv": 2048,
+        "sa_block_kv_dkv": 2048,
+        "sa_block_kv_dkv_compute": 2048,
+        "sa_block_q_dq": 2048,
+        "sa_block_kv_dq": 2048,
+        "sa_use_fused_bwd_kernel": True,
+        "profiler": "xplane",
+        "skip_first_n_steps_for_profiler": 10,
+        "profiler_steps": 5,
+    },
+    xla_flags=(
+        xla_flags_library.DENSE_VMEM_LIMIT_FLAG
+        + xla_flags_library.LAYOUT_FOR_ALL_REDUCE_SCATTER
+        + xla_flags_library.DATA_PARALLEL_OVERLAP
+        + xla_flags_library.CF_FOR_ALL_GATHER
+        + xla_flags_library.HOST_OFFLOAD_FLAGS
+    ),
+  )
+```
+
+This equivalent workload code can be found in the [maxtext_trillium_model_configs.py](https://github.com/AI-Hypercomputer/maxtext/blob/243b25e480f7550a0c389fa95cd3adcc716fe0df/benchmarks/maxtext_trillium_model_configs.py#L932-L972) file within the MaxText repository.

training/trillium/Llama3.1-70B-MaxText/llama3-1-70B-1xv6e-256.sh

@@ -0,0 +1,2 @@
+python3 benchmarks/benchmark_runner.py --project=$PROJECT --zone=$ZONE --device_type=v6e-256 --num_slices=1  --cluster_name=${CLUSTER_NAME} --base_output_directory=${OUTPUT_DIR} \
+--model_name="llama3_1_70b_8192" --base_docker_image maxtext_base_image