Skip to content

Commit bdff25e

Browse files
SurbhiJainUSCSurbhiJainUSC
committed
Update post-training docs to point to single source of truth for installation
1 parent 487bb6f commit bdff25e

3 files changed

Lines changed: 14 additions & 65 deletions

File tree

‎docs/tutorials/posttraining/rl.md‎

Lines changed: 2 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -42,47 +42,9 @@ rely on the vLLM library.
4242

4343
Let's get started!
4444

45-
## Create virtual environment and Install MaxText dependencies
45+
## Install MaxText and post-training dependencies
4646

47-
```bash
48-
# Create a virtual environment
49-
export VENV_NAME=<your virtual env name> # e.g., maxtext_venv
50-
pip install uv
51-
uv venv --python 3.12 --seed ${VENV_NAME?}
52-
source ${VENV_NAME?}/bin/activate
53-
```
54-
55-
### Option 1: From PyPI releases (Recommended)
56-
57-
Run the following commands to get all the necessary installations.
58-
59-
```bash
60-
uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
61-
install_maxtext_tpu_post_train_extra_deps
62-
```
63-
64-
It installs MaxText and then for post-training, it installs primarily the following:
65-
66-
a. [Tunix](https://github.com/google/tunix) as the LLM Post-Training Library, and
67-
68-
b. `vllm-tpu` which is
69-
[vllm](https://github.com/vllm-project/vllm) and
70-
[tpu-inference](https://github.com/vllm-project/tpu-inference) and thereby
71-
providing TPU inference for vLLM, with unified JAX and PyTorch support.
72-
73-
### Option 2: From Github
74-
75-
For using a version newer than the latest PyPI release, you could also install the latest vetted versions of the dependencies from MaxText in the following way:
76-
77-
```bash
78-
# 1. Clone the repository
79-
git clone https://github.com/AI-Hypercomputer/maxtext.git
80-
cd maxtext
81-
82-
# 2. Install dependencies in editable mode
83-
uv pip install -e .[tpu-post-train] --resolution=lowest
84-
install_maxtext_tpu_post_train_extra_deps
85-
```
47+
For instructions on installing MaxText with post-training dependencies on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/latest/install_maxtext.html) and use the `maxtext[tpu-post-train]` installation path to include all necessary post-training dependencies.
8648

8749
## Setup environment variables
8850

‎docs/tutorials/posttraining/sft.md‎

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -26,20 +26,7 @@ In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get st
2626

2727
## Install MaxText and Post-Training dependencies
2828

29-
```bash
30-
# Create a virtual environment
31-
export VENV_NAME=<your virtual env name> # e.g., maxtext_venv
32-
pip install uv
33-
uv venv --python 3.12 --seed ${VENV_NAME?}
34-
source ${VENV_NAME?}/bin/activate
35-
```
36-
37-
Run the following commands to get all the necessary installations.
38-
39-
```bash
40-
uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
41-
install_maxtext_tpu_post_train_extra_deps
42-
```
29+
For instructions on installing MaxText with post-training dependencies on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/latest/install_maxtext.html) and use the `maxtext[tpu-post-train]` installation path to include all necessary post-training dependencies.
4330

4431
## Setup environment variables
4532

@@ -53,7 +40,7 @@ Set the following environment variables before running SFT.
5340

5441
```sh
5542
# -- Model configuration --
56-
export PRE_TRAINED_MODEL=<model name> # e.g., 'llama3.1-8b-Instruct'
43+
export MODEL=<MaxText Model> # e.g., 'llama3.1-8b-Instruct'
5744

5845
# -- MaxText configuration --
5946
export BASE_OUTPUT_DIRECTORY=<output directory to store run logs> # e.g., gs://my-bucket/my-output-directory
@@ -76,15 +63,15 @@ This section explains how to prepare your model checkpoint for use with MaxText.
7663
If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section.
7764

7865
```sh
79-
export PRE_TRAINED_MODEL_CKPT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
66+
export MAXTEXT_CKPT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
8067
```
8168

8269
### Option 2: Converting a Hugging Face checkpoint
8370

8471
Refer the steps in [Hugging Face to MaxText](../../guides/checkpointing_solutions/convert_checkpoint.md#hugging-face-to-maxtext) to convert a hugging face checkpoint to MaxText. Make sure you have correct checkpoint files converted and saved. Similar as Option 1, you can set the following environment and move on.
8572

8673
```sh
87-
export PRE_TRAINED_MODEL_CKPT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
74+
export MAXTEXT_CKPT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
8875
```
8976

9077
## Run SFT on Hugging Face Dataset
@@ -95,8 +82,8 @@ Now you are ready to run SFT using the following command:
9582
python3 -m maxtext.trainers.post_train.sft.train_sft \
9683
run_name=${RUN_NAME?} \
9784
base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
98-
model_name=${PRE_TRAINED_MODEL?} \
99-
load_parameters_path=${PRE_TRAINED_MODEL_CKPT_PATH?} \
85+
model_name=${MODEL?} \
86+
load_parameters_path=${MAXTEXT_CKPT_PATH?} \
10087
per_device_batch_size=${PER_DEVICE_BATCH_SIZE?} \
10188
steps=${STEPS?} \
10289
hf_path=${DATASET_NAME?} \

‎docs/tutorials/posttraining/sft_on_multi_host.md‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ export STEPS=<Fine-Tuning Steps> # e.g., 1000
6262
export HF_TOKEN=<Hugging Face Access Token>
6363

6464
# -- Model Configuration --
65-
export MODEL_NAME=<Model Name> # e.g., deepseek3-671b
65+
export MODEL=<MaxText Model> # e.g., deepseek3-671b
6666

6767
# -- Dataset configuration --
6868
export DATASET_NAME=<Hugging Face Dataset Name> # e.g., HuggingFaceH4/ultrachat_200k
@@ -79,10 +79,10 @@ This section explains how to prepare your model checkpoint for use with MaxText.
7979
If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section.
8080

8181
```bash
82-
export MODEL_CHECKPOINT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
82+
export MAXTEXT_CKPT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
8383
```
8484

85-
**Note:** Make sure that `MODEL_CHECKPOINT_PATH` has the checkpoints created using the correct storage flags:
85+
**Note:** Make sure that `MAXTEXT_CKPT_PATH` has the checkpoints created using the correct storage flags:
8686

8787
```
8888
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX.
@@ -95,7 +95,7 @@ checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS))
9595
Refer the steps in [Hugging Face to MaxText](../../guides/checkpointing_solutions/convert_checkpoint.md#hugging-face-to-maxtext) to convert a hugging face checkpoint to MaxText. Make sure you have correct checkpoint files converted and saved. Similar as Option 1, you can set the following environment and move on.
9696

9797
```bash
98-
export MODEL_CHECKPOINT_PATH=<gcs path for MaxText checkpoint> # gs://my-bucket/my-checkpoint-directory/0/items
98+
export MAXTEXT_CKPT_PATH=<gcs path for MaxText checkpoint> # gs://my-bucket/my-checkpoint-directory/0/items
9999
```
100100

101101
## Submit workload on GKE cluster
@@ -113,7 +113,7 @@ xpk workload create \
113113
--workload=${WORKLOAD_NAME?} \
114114
--tpu-type=${TPU_TYPE?} \
115115
--num-slices=${TPU_SLICE?} \
116-
--command "python3 -m maxtext.trainers.post_train.sft.train_sft run_name=${WORKLOAD_NAME?} base_output_directory=${OUTPUT_PATH?} model_name=${MODEL_NAME?} load_parameters_path=${MODEL_CHECKPOINT_PATH?} hf_access_token=${HF_TOKEN?} per_device_batch_size=1 steps=${STEPS?} profiler=xplane hf_path=${DATASET_NAME?} train_split=${TRAIN_SPLIT?} train_data_columns=${TRAIN_DATA_COLUMNS?}"
116+
--command "python3 -m maxtext.trainers.post_train.sft.train_sft run_name=${WORKLOAD_NAME?} base_output_directory=${OUTPUT_PATH?} model_name=${MODEL?} load_parameters_path=${MAXTEXT_CKPT_PATH?} hf_access_token=${HF_TOKEN?} per_device_batch_size=1 steps=${STEPS?} profiler=xplane hf_path=${DATASET_NAME?} train_split=${TRAIN_SPLIT?} train_data_columns=${TRAIN_DATA_COLUMNS?}"
117117
```
118118

119119
Once the fine-tuning is completed, you can access your model checkpoints at `$OUTPUT_PATH/$WORKLOAD_NAME/checkpoints`.
@@ -131,7 +131,7 @@ xpk workload create-pathways \
131131
--workload=${WORKLOAD_NAME?} \
132132
--tpu-type=${TPU_TYPE?} \
133133
--num-slices=${TPU_SLICE?} \
134-
--command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 python3 -m maxtext.trainers.post_train.sft.train_sft run_name=${WORKLOAD_NAME?} base_output_directory=${OUTPUT_PATH?} model_name=${MODEL_NAME?} load_parameters_path=${MODEL_CHECKPOINT_PATH?} hf_access_token=${HF_TOKEN?} per_device_batch_size=1 steps=${STEPS?} profiler=xplane checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) enable_single_controller=True"
134+
--command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 python3 -m maxtext.trainers.post_train.sft.train_sft run_name=${WORKLOAD_NAME?} base_output_directory=${OUTPUT_PATH?} model_name=${MODEL?} load_parameters_path=${MAXTEXT_CKPT_PATH?} hf_access_token=${HF_TOKEN?} per_device_batch_size=1 steps=${STEPS?} profiler=xplane checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) enable_single_controller=True"
135135
```
136136

137137
Once the fine-tuning is completed, you can access your model checkpoints at `$OUTPUT_PATH/$WORKLOAD_NAME/checkpoints`.

0 commit comments

Comments
 (0)