Test MIX algorithm.
(self)
| 1012 | |
| 1013 | class TestTrainerMIX(BaseTrainerCase): |
| 1014 | def test_trainer(self): |
| 1015 | """Test MIX algorithm.""" |
| 1016 | # gsm8k has 16 tasks, sft_for_gsm8k has 8 tasks |
| 1017 | # total 4 steps, each step: read 4 tasks from gsm8k, 16 tasks from sft_for_gsm8k |
| 1018 | self.config.algorithm.algorithm_type = "mix" |
| 1019 | self.config.algorithm.repeat_times = 4 |
| 1020 | self.config.algorithm.sample_strategy = "mix" |
| 1021 | self.config.algorithm.advantage_fn = "grpo" |
| 1022 | self.config.algorithm.sample_strategy_args = {"expert_data_ratio": 0.5} # rft=4*4 : sft=16 |
| 1023 | self.config.algorithm.policy_loss_fn = "mix" |
| 1024 | self.config.buffer.batch_size = 4 |
| 1025 | self.config.buffer.train_batch_size = 32 |
| 1026 | self.config.buffer.total_steps = 2 |
| 1027 | self.config.buffer.explorer_input.taskset = get_unittest_dataset_config("gsm8k") |
| 1028 | self.config.synchronizer.sync_interval = 1 |
| 1029 | self.config.trainer.save_interval = 1 |
| 1030 | self.config.buffer.trainer_input.auxiliary_buffers[ |
| 1031 | "sft_dataset" |
| 1032 | ] = get_unittest_dataset_config("sft_for_gsm8k") |
| 1033 | self.config.buffer.trainer_input.auxiliary_buffers[ |
| 1034 | "sft_dataset" |
| 1035 | ].total_epochs = 8 # test this works |
| 1036 | self.config.check_and_update() |
| 1037 | self.config.buffer.trainer_input.experience_buffer.max_read_timeout = 20 |
| 1038 | self.config.trainer.max_checkpoints_to_keep = 2 |
| 1039 | both(self.config) |
| 1040 | ray.shutdown(_exiting_interpreter=True) |
| 1041 | |
| 1042 | # check trainer resume metadata |
| 1043 | trainer_meta_file = os.path.join(self.config.checkpoint_job_dir, "trainer_meta.json") |
| 1044 | with open(trainer_meta_file) as f: |
| 1045 | trainer_meta = json.load(f) |
| 1046 | self.assertEqual(trainer_meta["latest_iteration"], 2) |
| 1047 | self.assertEqual( |
| 1048 | trainer_meta["sample_strategy_state"]["expert_buffer"]["current_index"], 32 |
| 1049 | ) |
| 1050 | |
| 1051 | self.config.buffer.total_steps = None |
| 1052 | self.config.buffer.total_epochs = 1 |
| 1053 | self.config.check_and_update() |
| 1054 | ray.init(ignore_reinit_error=True, namespace=self.config.ray_namespace) |
| 1055 | both(self.config) |
| 1056 | |
| 1057 | # check trainer resume metadata |
| 1058 | with open(trainer_meta_file) as f: |
| 1059 | trainer_meta = json.load(f) |
| 1060 | self.assertEqual(trainer_meta["latest_iteration"], 4) |
| 1061 | self.assertEqual( |
| 1062 | trainer_meta["sample_strategy_state"]["expert_buffer"]["current_index"], 64 |
| 1063 | ) |
| 1064 | |
| 1065 | parser = TensorBoardParser(os.path.join(self.config.monitor.cache_dir, "tensorboard")) |
| 1066 | |
| 1067 | # test rollout metrics |
| 1068 | rollout_metrics = parser.metric_list("rollout") |
| 1069 | self.assertGreater(len(rollout_metrics), 0) |
| 1070 | self.assertEqual(parser.metric_max_step(rollout_metrics[0]), 4) |
| 1071 | self.assertEqual( |
nothing calls this directly
no test coverage detected