Skip to content

Commit fbd63d2

Browse files
committed
docs(rl): update tutorial with AgenticGRPOLearner for async RL training
1 parent 008b02a commit fbd63d2

1 file changed

Lines changed: 20 additions & 6 deletions

File tree

  • docs/tutorials/posttraining

docs/tutorials/posttraining/rl.md

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -121,11 +121,25 @@ The overview of what this run will do is as follows:
121121

122122
1. We load a policy model and a reference model. Both are copies of the model
123123
checkpoint you specified (e.g., `Llama3.1-8b-Instruct`).
124-
2. Evaluate the policy model's performance on GSM8K math reasoning benchmark.
125-
3. Train the policy model using GRPO.
126-
4. Evaluate the policy model's performance on GSM8K math reasoning benchmark
124+
1. Evaluate the policy model's performance on GSM8K math reasoning benchmark.
125+
1. Train the policy model using GRPO.
126+
1. Evaluate the policy model's performance on GSM8K math reasoning benchmark
127127
after the post-training with GRPO.
128128

129+
By default, the above command will train the model using GRPOLearner from Tunix. To enable
130+
asynchronous RL training with AgenticGRPOLearner, we can set `rl.use_agentic_rollout` to
131+
true. An example command will be:
132+
133+
```
134+
python3 -m maxtext.trainers.post_train.rl.train_rl \
135+
model_name=${MODEL?} \
136+
load_parameters_path=${MAXTEXT_CKPT_PATH?} \
137+
run_name=${RUN_NAME?} \
138+
base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
139+
chips_per_vm=${CHIPS_PER_VM?} \
140+
rl.use_agentic_rollout=True
141+
```
142+
129143
## Run GSPO
130144

131145
Run the following command for GSPO:
@@ -144,7 +158,7 @@ The overview of what this run will do is as follows:
144158

145159
1. We load a policy model and a reference model. Both are copies of the model
146160
checkpoint you specified (e.g., `Llama3.1-8b-Instruct`).
147-
2. Evaluate the policy model's performance on GSM8K math reasoning benchmark.
148-
3. Train the policy model using GSPO.
149-
4. Evaluate the policy model's performance on GSM8K math reasoning benchmark
161+
1. Evaluate the policy model's performance on GSM8K math reasoning benchmark.
162+
1. Train the policy model using GSPO.
163+
1. Evaluate the policy model's performance on GSM8K math reasoning benchmark
150164
after the post-training with GSPO.

0 commit comments

Comments
 (0)