Skip to content

Clarify RL checkpoint environment variables #4212

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
huytransformer wants to merge 1 commit into
base: main
Choose a base branch
from
+5 −6
Open

Clarify RL checkpoint environment variables #4212

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 5 additions & 6 deletions docs/tutorials/posttraining/rl.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ placeholders with your actual values.
```bash
# -- Model configuration --
# The MaxText model name. See `src/maxtext/configs/types.py` for `ModelName` for a
# full list of supported models.
# full list of supported models. This value is case-sensitive.
export MODEL=<MODEL_NAME> # e.g. 'llama3.1-8b-Instruct'

# -- MaxText configuration --
Expand All @@ -78,15 +78,14 @@ export BASE_OUTPUT_DIRECTORY=<GCS_BUCKET> # e.g., gs://my-bucket/maxtext-runs
# Note: Kubernetes requires workload names to be valid DNS labels (lowercase, no underscores or periods).
export RUN_NAME=<RUN_NAME>

# Number of accelerator chips per VM.
# - TPU v5e (single host): 8
# - TPU v5p (single host): 4
# - TPU v6e (single host): 8
# Number of accelerator chips per VM. For example, use 8 for a v5p-8 VM.
export CHIPS_PER_VM=<CHIPS_PER_VM>
```

## Get your model checkpoint

This section explains how to prepare your model checkpoint for use with MaxText. It sets `MAXTEXT_CKPT_PATH`, which is passed to `load_parameters_path` in the training command below.

### Option 1: Using an existing MaxText checkpoint

If you already have a MaxText-compatible model checkpoint, simply set the
Expand All @@ -98,7 +97,7 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/

### Option 2: Converting from a Hugging Face checkpoint

Refer the steps in [Hugging Face to MaxText](hf-to-maxtext) to convert a hugging face checkpoint to MaxText. Make sure you have correct checkpoint files converted and saved. Similar to Option 1, set the following environment variable and move on to the next section.
Refer to [Hugging Face to MaxText](hf-to-maxtext) to convert a Hugging Face checkpoint to MaxText format. After conversion finishes, set `MAXTEXT_CKPT_PATH` to the converted MaxText checkpoint path.

```bash
export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
Expand Down
Loading