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

Method test_trainer

tests/trainer/trainer_test.py:87–189  ·  view source on GitHub ↗

Test the both and bench mode.

(self)

Source from the content-addressed store, hash-verified

85)
86class TestTrainerCountdown(BaseTrainerCase):
87 def test_trainer(self):
88 """Test the both and bench mode."""
89 # test both mode
90 self.config.model.rope_scaling = {
91 "rope_type": "yarn",
92 "factor": 2.0,
93 "original_max_position_embeddings": 16384,
94 }
95 self.config.model.rope_theta = 10000
96 self.config.explorer.rollout_model.engine_type = self.engine_type
97 self.config.buffer.explorer_input.taskset = get_unittest_dataset_config("countdown")
98 self.config.buffer.explorer_input.taskset.data_selector = DataSelectorConfig(
99 selector_type="shuffle", seed=42
100 )
101 eval_tasksets = self.config.buffer.explorer_input.eval_tasksets
102 eval_tasksets.append(get_unittest_dataset_config("countdown", "test"))
103 eval_tasksets.append(get_unittest_dataset_config("copy_countdown", "test"))
104 eval_tasksets[0].repeat_times = 4
105 eval_tasksets[1].repeat_times = 4
106 self.config.trainer.save_interval = 4
107 self.config.trainer.save_hf_checkpoint = "never"
108 self.config.trainer.max_checkpoints_to_keep = 2
109 if self.strategy == "megatron":
110 self.config.trainer.trainer_strategy = "megatron"
111 self.config.check_and_update()
112 both(self.config)
113 parser = TensorBoardParser(os.path.join(self.config.monitor.cache_dir, "tensorboard"))
114 rollout_metrics = parser.metric_list("rollout")
115 self.assertGreater(len(rollout_metrics), 0)
116 self.assertEqual(parser.metric_max_step(rollout_metrics[0]), 8)
117 eval_metrics = parser.metric_list("eval")
118 self.assertGreater(len(eval_metrics), 0)
119 self.assertEqual(parser.metric_max_step(eval_metrics[0]), 8)
120 actor_metrics = parser.metric_list("actor")
121 self.assertGreater(len(actor_metrics), 0)
122 self.assertEqual(parser.metric_max_step(actor_metrics[0]), 8)
123 actor_kl_metrics = parser.metric_list("actor/kl")
124 self.assertGreater(len(actor_kl_metrics), 0)
125 actor_kl_loss = parser.metric_values("actor/kl_loss")
126 self.assertEqual(actor_kl_loss[0], 0.0)
127 advantage_kl_metrics = parser.metric_list("advantage/kl")
128 self.assertGreater(len(advantage_kl_metrics), 0)
129 time_metrics = parser.metric_list("time")
130 self.assertGreater(len(time_metrics), 0)
131 self.assertEqual(parser.metric_max_step(time_metrics[0]), 8)
132 ray.shutdown(_exiting_interpreter=True)
133 # check checkpoint
134 checkpoint_step_4, _ = get_checkpoint_dir_with_step_num(
135 checkpoint_root_path=self.config.checkpoint_job_dir,
136 trainer_type=self.config.trainer.trainer_type,
137 step_num=4,
138 )
139 # check save lastest checkpoint
140 checkpoint_step_8, step_num = get_checkpoint_dir_with_step_num(
141 checkpoint_root_path=self.config.checkpoint_job_dir,
142 trainer_type=self.config.trainer.trainer_type,
143 )
144 self.assertGreater(len(os.listdir(os.path.join(checkpoint_step_4, "actor"))), 0)

Callers

nothing calls this directly

Calls 13

metric_listMethod · 0.95
metric_max_stepMethod · 0.95
metric_valuesMethod · 0.95
metric_stepsMethod · 0.95
DataSelectorConfigClass · 0.90
bothFunction · 0.90
TensorBoardParserClass · 0.90
convert_commandFunction · 0.90
benchFunction · 0.90
check_and_updateMethod · 0.80

Tested by

no test coverage detected