Repository navigation
fix(verl): restore buffered async inputs from checkpoints - #705
Open
JiahaoTanXX (JiahaoTanXX) wants to merge 1 commit into
Open
JiahaoTanXX (JiahaoTanXX) wants to merge 1 commit into
JiahaoTanXX (JiahaoTanXX) wants to merge 1 commit into
Conversation
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The new test module imports optional training dependencies before guarding them, causing pytest tests collection to fail rather than skip VERL tests when those dependencies are absent.
1 open issue
What changed in this PR
This PR fixes skipped inputs when resuming asynchronous VERL training by saving the unused batch tail alongside the dataloader position and input epoch.
Changes:
- Atomically saves versioned async input state and restores it without relying on optimizer-step epoch inference.
- Adds CPU regression tests for checkpoint round trips, resume modes, shuffled inputs, compatibility, and failed writes.
- Documents the extra checkpoint file and recovery limits, preserving synchronous checkpoint behavior.
| File | Description |
|---|---|
tests/verl/test_async_input_checkpoint.py |
Adds regression coverage for async input checkpoint recovery. |
docs/20-trainer-configuration.md |
Explains checkpoint contents, compatibility, and recovery limits. |
agentlightning/verl/trainer.py |
Saves and restores buffered async inputs and preserves the restored epoch at startup. |
Review effort: Balanced
Comment on lines
+11
to
+13
| import numpy as np | ||
| import pytest | ||
| import torch |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

Async input refills can consume only part of a dataloader batch. For example, requesting two inputs from a four-input batch leaves two inputs in
_train_dataloader_buf, while the dataloader has advanced past all four. VERL's checkpoint saves the advanced loader but not this buffer, so resuming silently skips the two inputs that were never enqueued.Save a versioned
agentlightning_input.ptfile containing the loader state, unused batch tail, and input epoch for async training. Write it atomically before delegating to VERL's model save and latest-checkpoint publication. On resume, restore that exact input position and retain the epoch whenfit()starts. This also avoids VERL's optimizer-step-based epoch-boundary heuristic, which does not describe asynchronous input consumption.Synchronous checkpoints keep their existing format. Older checkpoints remain loadable and emit an async warning because missing buffered inputs cannot be reconstructed. The configuration docs explain the extra file and the recovery boundary: this does not restore in-flight carry-over rollouts or external Server/Controller state, or guarantee identical distributed training after a restart.
Validation
StatefulDataLoaderand VERL checkpoint save/load with only GPU worker operations stubbed. They cover partial/empty buffers, cross-epoch and shuffled input order, automatic and explicit resume into a different output directory, optimizer steps divisible by loader length, legacy/synchronous checkpoints, disabled/no-checkpoint resume,fit()startup, unknown versions, and failed writes leaving the previous checkpoint selected. Tensor, NumPy object-array, and list payloads are checked. Before the fix, 11 of the initial 16 cases failed on unchanged main.pytest -q tests: 262 passed, 4 skipped.mkdocs build --strictanduv build --no-sourcespassed.Validated on macOS CPU with the locked Python 3.12 / VERL 0.8.0 development environment. No GPU rollout or end-to-end training was run.
Related work: #694 addresses epoch-limit admission and explicitly leaves checkpoint-resume bookkeeping unchanged. #683 adds process-credit checkpoint metadata for a synchronous-only path. This PR addresses async buffered inputs independently; the changes touch adjacent trainer code.
AI assistance was used to investigate the bug, implement the change, write tests and documentation, and perform independent automated review.