(self)
| 831 | ) |
| 832 | class TestTrainerCheckpointSave(unittest.TestCase): |
| 833 | def setUp(self): |
| 834 | if multiprocessing.get_start_method(allow_none=True) != "spawn": |
| 835 | multiprocessing.set_start_method("spawn", force=True) |
| 836 | self.config = get_template_config() |
| 837 | self.config.buffer.total_steps = 6 |
| 838 | self.config.buffer.batch_size = 4 |
| 839 | self.config.model.model_path = get_model_path() |
| 840 | self.config.explorer.rollout_model.engine_type = "vllm" |
| 841 | self.config.algorithm.repeat_times = 3 |
| 842 | self.config.project = "Trainer-unittest" |
| 843 | self.config.name = f"trainer-{datetime.now().strftime('%Y%m%d%H%M%S')}" |
| 844 | self.config.monitor.monitor_type = "tensorboard" |
| 845 | self.config.checkpoint_root_dir = get_checkpoint_path() |
| 846 | self.config.synchronizer.sync_interval = 1 |
| 847 | self.config.synchronizer.sync_method = SyncMethod.CHECKPOINT |
| 848 | self.config.explorer.eval_interval = 4 |
| 849 | self.config.buffer.explorer_input.taskset = get_unittest_dataset_config("countdown") |
| 850 | self.config.trainer.save_interval = 2 |
| 851 | self.config.trainer.save_hf_checkpoint = "last" |
| 852 | self.config.trainer.trainer_strategy = self.strategy |
| 853 | if self.strategy == "megatron": |
| 854 | self.config.trainer.megatron.tensor_model_parallel_size = 2 |
| 855 | self.config.trainer.max_checkpoints_to_keep = 2 |
| 856 | self.config.check_and_update() |
| 857 | self.process_list = [] |
| 858 | |
| 859 | def test_trainer(self): # noqa: C901 |
| 860 | """Test the checkpoint saving.""" |
nothing calls this directly
no test coverage detected