Skip to content

fix(verl): restore buffered async inputs from checkpoints - #705

Open
JiahaoTanXX (JiahaoTanXX) wants to merge 1 commit into
microsoft:mainfrom
JiahaoTanXX:fix/restore-async-input-buffer
Open

JiahaoTanXX (JiahaoTanXX) wants to merge 1 commit into
microsoft:mainfrom
JiahaoTanXX:fix/restore-async-input-buffer

Conversation

@JiahaoTanXX

Copy link
Copy Markdown
Contributor

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.pt file 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 when fit() 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

  • 21 CPU regression cases exercise actual StatefulDataLoader and 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.
  • Pyright: 0 errors, 0 warnings. Ruff lint/format, headers, pre-commit hooks, and diff whitespace checks passed.
  • mkdocs build --strict and uv build --no-sources passed.

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.

Copilot AI balanced review requested due to automatic review settings November 12, 2026 02:01

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants