diff --git a/agentlightning/verl/trainer.py b/agentlightning/verl/trainer.py index 7090881a3..62419c875 100644 --- a/agentlightning/verl/trainer.py +++ b/agentlightning/verl/trainer.py @@ -55,6 +55,10 @@ RolloutManagerT = TypeVar("RolloutManagerT", bound=AglRolloutManagerBase) +class _TrainingDataExhausted(Exception): + """The configured training epochs have no more input batches.""" + + def _batch_dict_len(batch: dict[str, Any] | None) -> int: """Leading-dim length of a dataloader batch dict (0 if absent/empty).""" if batch is None or not batch: @@ -274,10 +278,13 @@ def _next_train_batch_dict_for_rollout(self) -> dict[str, Any]: async_train_batch_size = async_cfg.async_train_batch_size n_carry_over = len({rollout.data_id for rollout in self._carry_over_rollouts}) n_new = int(async_train_batch_size) - n_carry_over - return self._next_train_batch_dict_with_size(n_new) + batch = self._next_train_batch_dict_with_size(n_new) + if not batch and not self._carry_over_rollouts: + raise _TrainingDataExhausted + return batch def _next_train_batch_dict(self) -> dict[str, Any]: - while True: + while self.epoch <= self.config.trainer.total_epochs: if self._train_dataloader_iter is None: self._train_dataloader_iter = iter(self.train_dataloader) @@ -286,6 +293,7 @@ def _next_train_batch_dict(self) -> dict[str, Any]: self.epoch += 1 self._train_dataloader_iter = None + raise _TrainingDataExhausted def _next_train_batch_dict_with_size(self, size: int) -> dict[str, Any]: if size <= 0: @@ -340,15 +348,12 @@ def concat_batches(batches: list[dict[str, Any]]) -> dict[str, Any]: remaining = 0 while remaining > 0: - if self._train_dataloader_iter is None: - self._train_dataloader_iter = iter(self.train_dataloader) - try: - batch_dict = next(self._train_dataloader_iter) - except StopIteration: - self.epoch += 1 - self._train_dataloader_iter = None - continue + batch_dict = self._next_train_batch_dict() + except _TrainingDataExhausted: + # The last refill may be partial. An empty refill still lets + # the async runner finish groups admitted by an earlier step. + break batch_size = _batch_dict_len(batch_dict) if batch_size <= remaining: @@ -360,7 +365,7 @@ def concat_batches(batches: list[dict[str, Any]]) -> dict[str, Any]: self._train_dataloader_buf = tail remaining = 0 - return concat_batches(collected) + return concat_batches(collected) if collected else {} def _rollout(self, gen_batch: DataProto, is_train: bool) -> tuple[DataProto, dict[str, Any]]: """Run Agent Lightning rollouts and return the resulting DataProto plus metrics.""" @@ -381,11 +386,14 @@ def _rollout(self, gen_batch: DataProto, is_train: bool) -> tuple[DataProto, dic rollout_manager.delete_model() rollout_manager.register_model(server_addresses) previous_carry_over_rollouts = list(self._carry_over_rollouts) + available_groups = _batch_dict_len(data_dict) + len( + {rollout.data_id for rollout in previous_carry_over_rollouts} + ) completed_rollouts, new_carry_over_rollouts = rollout_manager.enqueue_and_wait_until_group_completed( data_dict, previous_carry_over_rollouts, is_train=True, - target_finished_group_num=self.config.data.train_batch_size, + target_finished_group_num=min(self.config.data.train_batch_size, available_groups), ) self._carry_over_rollouts = new_carry_over_rollouts async_rollout_metrics = self._compute_async_rollout_metrics( @@ -680,6 +688,7 @@ def fit(self): self.global_steps = 0 self.epoch = 0 self._train_dataloader_iter = None + self._train_dataloader_buf = None self._carry_over_rollouts = [] self._load_checkpoint() # Push loaded weights before the first rollout or validation. @@ -698,9 +707,10 @@ def fit(self): self.global_steps += 1 self.epoch += 1 last_val_metrics = None + last_update_step: int | None = None while True: - if self.global_steps > self.total_training_steps or self.epoch > self.config.trainer.total_epochs: + if self.global_steps > self.total_training_steps: progress_bar.close() return @@ -712,8 +722,22 @@ def fit(self): else False ) - with marked_timer("step", timing_raw): - result = self._train_step(timing_raw, curr_step_profile) + try: + with marked_timer("step", timing_raw): + result = self._train_step(timing_raw, curr_step_profile) + except _TrainingDataExhausted: + # Exhaustion is discovered when requesting the next input, so + # finalize the last update without counting an extra step. + if last_update_step is not None: + self.global_steps = last_update_step + if self.config.trainer.test_freq > 0 and self.global_steps % self.config.trainer.test_freq != 0: + last_val_metrics = self._validate() + logger.log(data=last_val_metrics, step=self.global_steps) + if self.config.trainer.save_freq > 0 and self.global_steps % self.config.trainer.save_freq != 0: + self._save_checkpoint() + pprint(f"Final validation metrics: {last_val_metrics}") + progress_bar.close() + return if result is None: print("AgentLightningRayPPOTrainer: train step returned no batch; skipping update.") # SC experiment budgets count optimizer updates, not empty collections. @@ -721,6 +745,7 @@ def fit(self): self.global_steps += 1 continue metrics, step_batch = result + last_update_step = self.global_steps # Compute timing and throughput after the step timer records its total duration. metrics.update(compute_timing_metrics(batch=step_batch, timing_raw=timing_raw)) @@ -735,8 +760,7 @@ def fit(self): ): with marked_timer("validate", timing_raw, color="green"): val_metrics = self._validate() - if is_last_step: - last_val_metrics = val_metrics + last_val_metrics = val_metrics metrics.update(val_metrics) if self.config.trainer.save_freq > 0 and ( diff --git a/docs/20-trainer-configuration.md b/docs/20-trainer-configuration.md index acb491fde..4ea2d8464 100644 --- a/docs/20-trainer-configuration.md +++ b/docs/20-trainer-configuration.md @@ -193,6 +193,19 @@ Transition mode does not use `trajectory_max_prompt_length` or `trajectory_max_r Use transition mode when every request-response call should remain a separate training sample. +## Training limits + +`trainer.total_epochs` limits how many times the trainer reads the training dataloader. +After the final epoch, synchronous collection stops before starting another rollout batch. +Asynchronous collection stops admitting new inputs and finishes the groups already in flight; +its last refill may contain fewer inputs than `async_train_batch_size` requests. The configured +`trainer.total_training_steps` remains an independent upper bound and can stop training earlier. + +When the epoch limit ends training, enabled validation and checkpointing run at the last +successful training step, unless they already ran there on their regular schedule. No extra +optimizer step is taken to discover that the data is exhausted. Final partial collections still +follow the PPO mini-batch alignment rules described below. + ## Algorithm correctness The following settings control how rollout data contributes to optimization: diff --git a/tests/verl/test_training_epoch_limit.py b/tests/verl/test_training_epoch_limit.py new file mode 100644 index 000000000..47b02f9bf --- /dev/null +++ b/tests/verl/test_training_epoch_limit.py @@ -0,0 +1,206 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Exercise epoch termination through fit without launching GPU workers.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf +from verl import DataProto + +from agentlightning.verl import trainer as trainer_module + + +def _trainer( + monkeypatch: pytest.MonkeyPatch, + *, + rows: int = 4, + epochs: int = 1, + async_size: int | None = None, + total_steps: int = 100, + frequency: int = 10, + skip_updates: bool = False, +) -> Any: + trainer = cast(Any, object.__new__(trainer_module.AgentLightningRayPPOTrainer)) + trainer.config = OmegaConf.create( + { + "trainer": { + "project_name": "test", + "experiment_name": "epochs", + "logger": [], + "val_before_train": False, + "total_epochs": epochs, + "test_freq": frequency, + "save_freq": frequency, + }, + "global_profiler": {"steps": None}, + "actor_rollout_ref": {"actor": {"policy_loss": {"loss_mode": "score_centering"}}, "rollout": {"n": 1}}, + "agentlightning": { + "async_rollout": {"async_train_batch_size": async_size}, + "trace_aggregator": {}, + "reward_fillna_value": 0.0, + }, + "data": {"train_batch_size": 2, "max_prompt_length": 1, "max_response_length": 1}, + } + ) + trainer.is_async = async_size is not None + trainer.total_training_steps = total_steps + trainer.train_dataloader = [ + {"item": np.arange(start, min(start + 2, rows)), "fake_ids": torch.ones(min(2, rows - start), 1)} + for start in range(0, rows, 2) + ] + trainer._load_checkpoint = lambda: None + trainer.checkpoint_manager = SimpleNamespace(update_weights=lambda _step: None) + trainer.resource_pool_manager = SimpleNamespace(get_n_gpus=lambda: 0) + trainer.tokenizer = SimpleNamespace(pad_token_id=0) + trainer.async_rollout_manager = SimpleNamespace(server_addresses=["test"]) + trainer._resume_all_rollout_generation = lambda: None + trainer._resume_gateway = lambda: None + trainer._pause_and_drain_gateway = lambda **_kwargs: {} + trainer._compute_async_rollout_metrics = lambda **_kwargs: {} + trainer._rollout_lifecycle_metrics = lambda _rollouts: {} + trainer.fetched = [] + trainer.completed = [] + trainer.updates = [] + trainer.validations = [] + trainer.checkpoints = [] + trainer.targets = [] + trainer._save_checkpoint = lambda: trainer.checkpoints.append(trainer.global_steps) + + def validate(): + trainer.validations.append(trainer.global_steps) + return {"validation/reward": 1.0} + + trainer._validate = validate + monkeypatch.setattr(trainer_module, "Tracking", lambda **_kwargs: SimpleNamespace(log=lambda **_kwargs: None)) + monkeypatch.setattr(trainer_module, "compute_timing_metrics", lambda **_kwargs: {}) + + class Manager: + next_id = 0 + + def delete_model(self): + pass + + def register_model(self, _addresses): + pass + + def enqueue_and_wait_until_group_completed(self, data, carry, *, is_train, target_finished_group_num): + assert is_train + active = list(carry) + for item in data.get("item", []): + active.append(SimpleNamespace(data_id=str(self.next_id), item=int(item))) + self.next_id += 1 + # A target above the available groups would hang the real manager. + assert 0 < target_finished_group_num <= len(active) + trainer.targets.append(target_finished_group_num) + return active[:target_finished_group_num], active[target_finished_group_num:] + + class Adapter: + def __init__(self, **_kwargs): + pass + + def get_train_data_batch(self, completed, **_kwargs): + trainer.completed.extend(row.item for row in completed) + return DataProto.from_dict(tensors={"fake_ids": torch.ones(len(completed), 1)}), {} + + manager = Manager() + trainer._make_rollout_manager = lambda _cls: manager + monkeypatch.setattr(trainer_module, "RolloutAdapter", Adapter) + + def train_step(_timing, _profile): + batch_dict = trainer._next_train_batch_dict_for_rollout() + trainer.fetched.extend(batch_dict.get("item", [])) + batch = DataProto.from_single_dict(batch_dict) + if trainer.is_async: + batch, _ = trainer._rollout(batch, is_train=True) + if skip_updates: + return None + trainer.updates.append(trainer.global_steps) + return {}, batch + + trainer._train_step = train_step + return trainer + + +@pytest.mark.parametrize("epochs", [1, 2]) +def test_fit_stops_at_epoch_limit_before_an_extra_update(monkeypatch, epochs): + trainer = _trainer(monkeypatch, epochs=epochs) + + trainer.fit() + + assert trainer.fetched == [0, 1, 2, 3] * epochs + assert trainer.updates == list(range(1, 2 * epochs + 1)) + assert trainer.validations == [2 * epochs] + assert trainer.checkpoints == [2 * epochs] + + +@pytest.mark.parametrize("rows,async_size", [(4, 4), (6, 5), (3, 5)]) +@pytest.mark.parametrize("epochs", [1, 2]) +def test_async_fit_drains_admitted_groups_without_refilling_past_epoch_limit(monkeypatch, rows, async_size, epochs): + trainer = _trainer(monkeypatch, rows=rows, epochs=epochs, async_size=async_size) + + trainer.fit() + + assert trainer.fetched == list(range(rows)) * epochs + assert trainer.completed == trainer.fetched + assert trainer._carry_over_rollouts == [] + assert sum(trainer.targets) == rows * epochs + assert len(trainer.updates) == (rows * epochs + 1) // 2 + assert trainer.validations == [trainer.updates[-1]] + assert trainer.checkpoints == [trainer.updates[-1]] + + +def test_epoch_completion_does_not_repeat_scheduled_validation_or_checkpoint(monkeypatch): + trainer = _trainer(monkeypatch, frequency=2) + + trainer.fit() + + assert trainer.validations == [2] + assert trainer.checkpoints == [2] + + +def test_explicit_step_limit_still_stops_training_first(monkeypatch): + trainer = _trainer(monkeypatch, epochs=2, total_steps=1) + + trainer.fit() + + assert trainer.fetched == [0, 1] + assert trainer.updates == trainer.validations == trainer.checkpoints == [1] + + +def test_epoch_limit_stops_empty_collections_without_advancing_optimizer_steps(monkeypatch): + trainer = _trainer(monkeypatch, skip_updates=True) + + trainer.fit() + + assert trainer.fetched == [0, 1, 2, 3] + assert trainer.updates == trainer.validations == trainer.checkpoints == [] + + +def test_actual_train_step_accepts_an_empty_refill_for_carry_over(monkeypatch): + trainer = _trainer(monkeypatch, async_size=4) + trainer.epoch = 2 + trainer.global_steps = 2 + trainer._train_dataloader_iter = None + trainer._carry_over_rollouts = [SimpleNamespace(data_id="pending")] + + class ReachedRollout(Exception): + pass + + def rollout(gen_batch, *, is_train): + assert is_train + assert gen_batch.non_tensor_batch == {} + assert len(gen_batch) == 0 + assert gen_batch.meta_info["global_steps"] == 2 + raise ReachedRollout + + trainer._rollout = rollout + # Exercise DataProto construction and the inherited _get_gen_batch, not the + # simplified train-step stub used to observe the fit loop above. + with pytest.raises(ReachedRollout): + trainer_module.AgentLightningRayPPOTrainer._train_step(trainer, {}, False)