From 42bf4ca431939b633243d02ea30a4eb853bade59 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 11 Nov 2026 12:18:03 +0000 Subject: [PATCH] feat(rl): reuse rollout and training debug phases for disaggregated runs Mechanical extraction from XTuner #2257 fixed head 43d62b6. No new implementation. --- xtuner/v1/train/rl_trainer.py | 121 +++++++++++++++++++++------------- 1 file changed, 76 insertions(+), 45 deletions(-) diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 32c8dede03..9d4e5dd0dc 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -1114,6 +1114,31 @@ def _save_debug_rollout_batch(self, train_batch: list[list[RolloutState]], train torch.save(serializable_batch, save_path) self.logger.info(f"Debug rollout batch for step {train_step} saved to {save_path}") + def _fit_debug_train(self) -> None: + self._benchmark_start_time_s = time.perf_counter() + self._benchmark_training_samples = 0 + self._benchmark_training_tokens = 0 + + init_train_step = self._cur_step + 1 + for train_step in range(init_train_step, self._total_train_steps + 1): + self.logger.info(f"Debug train step {train_step}/{self._total_train_steps} start") + step_timer_dict: dict[str, float] = {} + with timer("step", step_timer_dict): + with timer("load_debug_rollout", step_timer_dict): + train_batch = self._load_debug_rollout_batch(train_step) + train_log_info = self._train_one_batch( + train_batch, + train_step, + step_timer_dict, + offload_rollout_before_train=False, + resume_train_nccl_before_train=False, + ) + eval_log_info: dict[str, float] = {} + produce_result = ProduceBatchResult(rollout_states=train_batch) + + self._log_step(train_step, step_timer_dict, produce_result, train_log_info, eval_log_info) + self._cur_step = train_step + def _list_debug_rollout_files(self, debug_rollout_dir: Path) -> dict[int, Path]: debug_files = { _parse_debug_rollout_step(path): path @@ -1734,31 +1759,6 @@ def _fit(self): self._log_step(train_step, step_timer_dict, produce_result, train_log_info, eval_log_info) self._cur_step = train_step - def _fit_debug_train(self) -> None: - self._benchmark_start_time_s = time.perf_counter() - self._benchmark_training_samples = 0 - self._benchmark_training_tokens = 0 - - init_train_step = self._cur_step + 1 - for train_step in range(init_train_step, self._total_train_steps + 1): - self.logger.info(f"Debug train step {train_step}/{self._total_train_steps} start") - step_timer_dict: dict[str, float] = {} - with timer("step", step_timer_dict): - with timer("load_debug_rollout", step_timer_dict): - train_batch = self._load_debug_rollout_batch(train_step) - train_log_info = self._train_one_batch( - train_batch, - train_step, - step_timer_dict, - offload_rollout_before_train=False, - resume_train_nccl_before_train=False, - ) - eval_log_info: dict[str, float] = {} - produce_result = ProduceBatchResult(rollout_states=train_batch) - - self._log_step(train_step, step_timer_dict, produce_result, train_log_info, eval_log_info) - self._cur_step = train_step - def _get_colocate_rollout_model_step(self, train_step: int) -> int: previous_step = train_step - 1 return previous_step - (previous_step % self._sync_weights_interval) @@ -1837,14 +1837,23 @@ def __init__(self, cfg: RLDisaggregatedTrainerConfig): train_resources=cfg.train_resources, rollout_resources=cfg.rollout_resources, ) - self._cpu_resource_manager = CPUResourceManager([self._train_pg, self._rollout_pg]) + self._cpu_resource_manager = CPUResourceManager( + [pg for pg in (self._train_pg, self._rollout_pg) if pg is not None] + ) self._cpu_resource_manager.log_initial_snapshot() set_cpu_resource_manager(self._cpu_resource_manager) - self.train_controller = self._train_worker_cfg.build( - self._train_pg, - advantage_estimator=self._advantage_estimator, - distillation=self._distillation, - ) + if not self._debug_rollout: + self.train_controller = self._train_worker_cfg.build( + self._train_pg, advantage_estimator=self._advantage_estimator, distillation=self._distillation + ) + if self._debug_train: + self.tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_path, trust_remote_code=True) + self._debug_train_files = self._list_debug_rollout_files(self._debug_rollout_dir) + if cfg.total_train_steps is None: + self._total_train_steps = max(self._debug_train_files) + return + if self._debug_rollout: + self._rollout_config.skip_load_weights = False self.rollout_controller = self._rollout_config.build(self._rollout_pg) if self._rollout_config.weight_transport_type != "nccl": @@ -1862,6 +1871,9 @@ def __init__(self, cfg: RLDisaggregatedTrainerConfig): "because it does not allow early stopping in production." ) + if self._debug_rollout: + return + bind_train_rollout( train_controller=self.train_controller, rollout_controller=self.rollout_controller, @@ -1882,9 +1894,17 @@ def _build_disaggregated_placement_groups( train_pg_name = f"{pg_name_prefix}_train" rollout_pg_name = f"{pg_name_prefix}_rollout" - train_pg = AutoAcceleratorWorkers.build_placement_group(train_resources, name=train_pg_name) - rollout_pg = AutoAcceleratorWorkers.build_placement_group(rollout_resources, name=rollout_pg_name) - if train_pg.id == rollout_pg.id: + train_pg = ( + None + if self._debug_rollout + else AutoAcceleratorWorkers.build_placement_group(train_resources, name=train_pg_name) + ) + rollout_pg = ( + None + if self._debug_train + else AutoAcceleratorWorkers.build_placement_group(rollout_resources, name=rollout_pg_name) + ) + if train_pg is not None and rollout_pg is not None and train_pg.id == rollout_pg.id: raise RuntimeError( "RLDisaggregatedTrainer requires distinct placement groups for train and rollout, " f"but both resolved to the same placement group id={train_pg.id}. " @@ -1893,8 +1913,8 @@ def _build_disaggregated_placement_groups( self.logger.info( "Created disaggregated placement groups: " - f"train={train_pg_name}(id={train_pg.id}), " - f"rollout={rollout_pg_name}(id={rollout_pg.id})" + f"train={train_pg_name}(id={getattr(train_pg, 'id', None)}), " + f"rollout={rollout_pg_name}(id={getattr(rollout_pg, 'id', None)})" ) return train_pg, rollout_pg @@ -1957,6 +1977,12 @@ async def _get_batch_or_raise_producer_failure( async def _fit(self): self.logger.info("Start RL disaggregated training") + if self._debug_train: + await asyncio.to_thread(self._fit_debug_train) + # Diagnostic actor/optimizer checkpoint. Full manager/replay checkpoint + # acceptance uses the online quiescent save path below. + await asyncio.to_thread(self.train_controller.save, str(self.exp_dir / "debug-train-dcp"), False) + return if self._cur_step >= self._total_train_steps: self.logger.info(f"Train steps {self._total_train_steps} reached, stop training") return @@ -2014,15 +2040,20 @@ async def _fit(self): "unless status is empty EXPIRED_BATCH." ) # 训练路径放到线程里执行，避免阻塞事件循环。 - train_log_info = await asyncio.to_thread( - self._train_one_batch, - train_batch, - train_step, - step_timer_dict, - raw_rewards_sum=produce_result.raw_rewards_sum, - raw_rewards_count=produce_result.raw_rewards_count, - produce_result=produce_result, - ) + if self._debug_rollout: + self._save_debug_rollout_batch(train_batch, train_step) + self._release_train_batch_resources(train_batch) + train_log_info = {} + else: + train_log_info = await asyncio.to_thread( + self._train_one_batch, + train_batch, + train_step, + step_timer_dict, + raw_rewards_sum=produce_result.raw_rewards_sum, + raw_rewards_count=produce_result.raw_rewards_count, + produce_result=produce_result, + ) sync_model_step = train_step # 保存、同步、评测、恢复 producer 都以已完成 model_step 为口径。