Skip to content

[Feature] Support rollout-only and train-only debug in disaggregated RL #2298

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
xs1997zju wants to merge 1 commit into
base: xs1997zju/agent-rl-stack-20261111/04-agent-core
Choose a base branch
from
Draft

[Feature] Support rollout-only and train-only debug in disaggregated RL #2298

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 76 additions & 45 deletions xtuner/v1/train/rl_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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":
Expand All @@ -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,
Expand All @@ -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}. "
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 为口径。
Expand Down
Loading