@@ -121,11 +121,25 @@ The overview of what this run will do is as follows:
121121
1221221 . 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
131145Run the following command for GSPO:
@@ -144,7 +158,7 @@ The overview of what this run will do is as follows:
144158
1451591 . 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