MCPcopy Create free account
hub / github.com/agentscope-ai/Trinity-RFT / test_trainer

Method test_trainer

tests/trainer/trainer_test.py:197–241  ·  view source on GitHub ↗

Test the explore step ahead trainer.

(self)

Source from the content-addressed store, hash-verified

195
196class TestStepAheadAsyncRL(BaseTrainerCase):
197 def test_trainer(self):
198 """Test the explore step ahead trainer."""
199 # train 4 step, sync_offset=1, sync_interval=2
200 # Explorer:
201 # | 1 | 2 | 3 |sync| 4 |
202 # |---|---|---|sync|---|
203 # Trainer:
204 # | 1 | 2 |sync| 3 | 4 |
205 # |---|---|sync|---|---|
206 self.config.buffer.batch_size = 6
207 self.config.buffer.total_steps = 4
208 # use 3 GPU in a 2 x 2 cluster, the trainer only have 1 GPU
209 self.config.explorer.rollout_model.engine_num = 3
210 self.config.explorer.rollout_model.tensor_parallel_size = 1
211 self.config.buffer.explorer_input.taskset = get_unittest_dataset_config("countdown")
212 self.config.trainer.save_interval = 4
213 self.config.trainer.max_checkpoints_to_keep = 1
214 self.config.synchronizer.sync_interval = 2
215 self.config.synchronizer.sync_offset = 1
216 self.config.check_and_update()
217
218 both(self.config)
219 parser = TensorBoardParser(os.path.join(self.config.monitor.cache_dir, "tensorboard"))
220 rollout_metrics = parser.metric_list("rollout")
221 self.assertGreater(len(rollout_metrics), 0)
222 self.assertEqual(parser.metric_max_step(rollout_metrics[0]), 4)
223 actor_metrics = parser.metric_list("actor")
224 self.assertGreater(len(actor_metrics), 0)
225 self.assertEqual(parser.metric_max_step(actor_metrics[0]), 4)
226 actor_kl_metrics = parser.metric_list("actor/kl")
227 self.assertGreater(len(actor_kl_metrics), 0)
228 advantage_kl_metrics = parser.metric_list("advantage/kl")
229 self.assertGreater(len(advantage_kl_metrics), 0)
230 time_metrics = parser.metric_list("time")
231 self.assertGreater(len(time_metrics), 0)
232 self.assertEqual(parser.metric_max_step(time_metrics[0]), 4)
233 ray.shutdown(_exiting_interpreter=True)
234 # check checkpoint
235
236 checkpoint_step_4, step_num = get_checkpoint_dir_with_step_num(
237 checkpoint_root_path=self.config.checkpoint_job_dir,
238 trainer_type=self.config.trainer.trainer_type,
239 )
240 self.assertEqual(step_num, 4)
241 self.assertTrue(os.path.exists(checkpoint_step_4))
242
243 def tearDown(self):
244 # remove dir only when the test passed

Callers

nothing calls this directly

Calls 8

metric_listMethod · 0.95
metric_max_stepMethod · 0.95
bothFunction · 0.90
TensorBoardParserClass · 0.90
check_and_updateMethod · 0.80
shutdownMethod · 0.45

Tested by

no test coverage detected