test: add FSDP2 distributed coverage for AsyncGRPOTrainer - #6144
test: add FSDP2 distributed coverage for AsyncGRPOTrainer#6144behroozazarkhalili wants to merge 25 commits into
Conversation
Adds a 2-process FSDP2 functional test for the experimental AsyncGRPOTrainer, which previously had no distributed-test coverage. test_train_fsdp2 launches a companion worker via accelerate launch under a 2-process FSDP2 config and asserts the trainer trains end-to-end on FSDP2-sharded parameters (steps run, loss finite, params update). It uses an in-process stub rollout worker (no vLLM server / NCCL weight transfer), so the only distributed surface exercised is the FSDP2 parameter lifecycle. Guarded by require_torch_multi_accelerator so it skips when fewer than 2 accelerators are available.
…-coverage # Conflicts: # tests/experimental/test_async_grpo_trainer.py
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
albertvillanova
left a comment
albertvillanova
left a comment
There was a problem hiding this comment.
Thanks @behroozazarkhalili — this is careful work and the companion-worker pattern is a clean match for _openreward_echo_env.py.
The blocker is that the test won't actually run anywhere in CI: the experimental lane runs on a single-GPU aws-g4dn-2xlarge, so @require_torch_multi_accelerator skips it (the green "Tests (experimental)" check here never executed it), and the only multi-GPU lane is the slow-tests multi_gpu job, which collects -m slow only — and this test isn't marked slow. Could you add @pytest.mark.slow so it lands in that lane (and confirm the runner really exposes 2 GPUs)? Otherwise it's a permanent skip that will quietly bitrot.
Separately, the #6077 narrative in the module docstring and the test comment is very long for what it ultimately concludes — trimming to a one-line note would keep the test readable. The duplicated stub is fine given the subprocess boundary.
…CI lane test_train_fsdp2 requires @require_torch_multi_accelerator, but the experimental lane is single-GPU and would skip it permanently. The only multi-GPU lane is slow_tests, which collects -m slow. Mark the test slow so it runs there. Also trim the #6077 narrative in the test comment and the worker module docstring to a one-line note per review feedback.
|
Thanks @albertvillanova, both points addressed in 1. 2. Docstring/comment trimmed. The long #6077 narrative is now a one-line note in both the test comment and the worker module docstring. On "confirm the runner really exposes 2 GPUs", being precise about what I could and could not verify from outside your infra:
The nice property is that it is self-verifying and cannot go falsely red: If you would prefer a louder signal than a silent skip, I can add an explicit cc @qgallouedec |
…po_trainer.py (keep new async-GRPO suite + FSDP2 smoke test)
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
… API Two defects in the companion worker, both surfaced by review on this PR and both introduced by merging main into a branch whose stub predated the current API: - `RolloutSample` now requires `group_id`, which the stub never passed, so the first sample raised `TypeError` and training never started. Every completion of one prompt shares a group id. - Without an explicit `weight_transfer`, `AsyncGRPOTrainer` builds the default `WeightTransferClient`, which needs a live vLLM server. The smoke is meant to exercise only the FSDP2 parameter lifecycle, so pass a no-op implementation of the protocol.
…c_grpo_trainer.py import block #6715 added `from collections import defaultdict` in the same stdlib import block this branch added `subprocess` and `pathlib.Path` to. Both sides are needed, so the resolution keeps all four in isort order. No other hunk conflicted.
…ains #6715 widened `RolloutWorkerProtocol` with a required `metrics_queue`, and `AsyncGRPOTrainer.log()` drains it unconditionally on the main process. The FSDP2 worker's stub carried only `rollout_buffer`, so the smoke would have raised AttributeError at the first log step. Line copied from the sibling stub in test_async_grpo_trainer.py so the two stay identical.
`token_budget` was unset, so `get_train_dataloader` defaulted it to the vLLM server's max_model_len and rank 0 called `wait_for_server_ready()` first. This smoke starts no vLLM server, so it died on the 5s timeout before reaching a single FSDP step. The sibling stubbed `test_train` already sets `token_budget=256` for exactly this reason; line and comment copied from it. The params_changed loop broke at the first changed parameter, but `_materialize` calls the collective `DTensor.full_tensor()`. With `fsdp_cpu_ram_efficient_loading` the pre-wrap snapshot holds real weights on rank 0 and placeholders elsewhere, so in the no-update regression this smoke exists to catch, the ranks break at different indices and rank 0 blocks forever on an unpartnered collective: the test hung instead of reporting the failure. Every rank now walks the full parameter list and decides afterwards. The list comprehension is deliberate, since `any()` over a generator short-circuits exactly like the `break` did.
…ock again #5911 added `unittest.mock` to the same stdlib import block, and its `json` / `os` additions were already here. Union of both sides in isort order; no other hunk conflicted.
`TokenBudgetBatcher` emits a micro-batch only when the next sample fits in no
row, and a bounded producer of short samples never triggers that. Running TRL's
own batcher on this worker's inputs gives:
10 samples x 41 tok, budget 256, 2 ranks -> rows [205, 205], 0 micro-batches
24 samples, budget 0, 2 ranks -> 4 micro-batches, no empty row
At 0 micro-batches the rollout consumer blocks forever on an empty queue.
`token_budget=256` came from the sibling single-process `test_train`, where a
single row of 256 overflows at the 7th sample. Two ranks double the capacity, so
the overflow that made the sibling work never happens here.
`token_budget=0` is the only value that both selects the count-based
`FixedCountBatcher` (`> 0` is false) and skips `wait_for_server_ready()`
(`is None` is false). The run therefore needs no vLLM server and no packing
arithmetic.
`samples_per_weight_sync=24` covers `max_steps=2` with no weight-sync refill:
micro-batch = per_device_train_batch_size 3 x 2 ranks = 6
samples needed = 6 x grad accum 1 x max_steps 2 = 12
previous value = 10
The `subprocess.run` timeout bounds a failure mode this fix does not remove.
`RolloutQueueDataset.__iter__` loops on `queue.Empty` and calls only
`check_health`, a no-op in this stub, so a future starvation would hang the
pytest process rather than fail it. The timeout makes it a readable failure.
`AsyncGRPOTrainer` loads every model with `attn_implementation="kernels-community/flash-attn3"` (`async_grpo_trainer.py:811`), and Flash Attention rejects a head_size that is not a multiple of 8. `trl-internal-testing/tiny-Qwen2ForCausalLM-2.5` reports hidden_size 8 over 4 attention heads, so head_size is 2 and the forward pass raises. The three sibling tests in this file already carry the #6837 xfail. `test_train_fsdp2` drives the same model through the same trainer without it. The comment above the test also claimed `@pytest.mark.slow` routes it to the multi-GPU `slow_tests` lane. It does not. `slow_tests` is `pytest -m "slow" tests/`, a recursive collection, and `norecursedirs` excludes `tests/experimental`: pytest -m slow tests/ --collect-only -> 0 matches for test_train_fsdp2 pytest tests/experimental --collect-only -> 1 match `test_experimental` passes an explicit path, so it does collect the test, but its runner is single-GPU and `@require_torch_multi_accelerator` skips it there. The comment now says that instead.
|
@albertvillanova a correction to what I told you on 2026-07-03, since you adjusted your review on the strength of it. It changes the conclusion rather than a detail, so the short version first: I wrote: "With the marker it is now collectable by the The slow lane never collects it. The slow workflow would also never fire for this change. The lane that does collect it cannot run it. That is run 32436517199, from today, and the same line appears in all six completed #6837 would stop the test even on a two-GPU lane. Because that skip is silent, the same trap stays open for whatever is added next to it. Which way to fix that is a repo-layout call. I would rather hand you the measurement than a recommendation I have not tested. Adding |
… smoke `test_train_fsdp2` was marked xfail because the trainer loads the model with Flash Attention, which rejects a `head_size` that is not a multiple of 8, and the worker used `tiny-Qwen2ForCausalLM-2.5` at `head_size=2`. #6853 and #6854 fixed the same class of failure by swapping the model rather than skipping the test, and removed every xfail carrying this reason. This does the same for the one test they did not cover: the worker now loads `small-Qwen2ForCausalLM-2.5` (`head_size=32`), and the xfail is gone. Measured from the two configs rather than assumed: tiny-Qwen2ForCausalLM-2.5 hidden=8 heads=4 head_size=2 rejected small-Qwen2ForCausalLM-2.5 hidden=128 heads=4 head_size=32 accepted Also merges main, which brings in those two PRs and drops the six inherited xfails this branch was still carrying.
The worker set `output_dir` to a relative path while the launcher runs it with `cwd` at the repo root, so every run left trainer artifacts in the working tree with nothing to clean them up. The five sibling tests in this file already pass `self.tmp_dir`; the worker is a subprocess, so it receives that path through the environment dict the launcher already builds for `PYTHONPATH`. No default for the variable: this worker is launched only by `test_train_fsdp2`, so a missing value should fail loudly rather than fall back to the old behavior.
…very configured step, and make stub rewards distinct Three defects from an adversarial review of the FSDP2 test, each reproduced before the change. The timeout used subprocess.run, which kills only the accelerate launcher. torch elastic starts each rank with start_new_session=True, so on a hang the ranks survived the timeout and kept the GPUs. The launcher now runs in its own process group and the whole group is killed on timeout. The parent accepted steps >= 1 although the worker configures max_steps=2 to exercise the optimizer loop more than once, so an early stop after step 1 passed every assertion. The worker reports its configured step count and the parent requires the measured count to match it. The stub derived rewards from hash() of the completion text. Under a randomized PYTHONHASHSEED the three completions of a group can collide modulo 100, the reward standard deviation is then zero and the advantages are NaN, a failure unrelated to FSDP2 (reproduced with PYTHONHASHSEED=5062). Rewards are now a fixed linspace, distinct by construction.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 2a112af. Configure here.
…uncher's process group The previous timeout path put the launcher in its own process group and killed that group. torch elastic starts every rank with start_new_session=True, so the ranks were never in that group: a hung run killed only the launcher, the ranks kept the GPUs and held the stdout pipe open, and the read after the timeout never returned. The descendants are now collected from /proc while the launcher is still their parent, and the launcher and every descendant get SIGKILL before the pipes are drained. Reproduced with a launcher whose child detached into its own session: the child is found, the drain returns at once, and nothing survives.
The worker only reported step count, a finite loss and changed parameters, all of which a replicated single-process run also produces, so the test could not tell an FSDP2 launch from a plain one. The result now carries the process count, the distributed type, the FSDP version and the number of DTensor parameters, and the test pins them to two ranks, FSDP, version 2 and at least one sharded parameter.
3 participants
What
Adds a 2-process FSDP2 functional test for the experimental
AsyncGRPOTrainer, which previously had no distributed-test coverage (tests/experimental/test_async_grpo_trainer.pyonly ran single-process).test_train_fsdp2launches a companion worker viaaccelerate launchunder a 2-process FSDP2 config and asserts the trainer trains end-to-end on FSDP2-sharded parameters: training steps run, the loss is finite, and parameters update. It uses an in-process stub rollout worker (no vLLM server / NCCL weight transfer), so the only distributed surface exercised is the FSDP2 parameter lifecycle. The test is guarded by@require_torch_multi_accelerator, so it skips automatically when fewer than 2 accelerators are available.Why
AsyncGRPOTrainerruns under FSDP2 in practice but had no test confirming it trains under a real 2-process FSDP2 group. This closes that gap with a lightweight functional smoke that mirrors the existing_openreward_echo_env.pycompanion-script pattern.Test plan
Verified on 2x H100:
Skips cleanly on single-GPU / CPU.
ruff checkandruff formatpass.Note
Low Risk
Changes are limited to experimental test harness and Accelerate config; no production trainer or runtime behavior is modified.
Overview
Adds distributed FSDP2 coverage for experimental
AsyncGRPOTrainer, which previously only had single-process training tests.A new
test_train_fsdp2(marked slow, gated by@require_torch_multi_accelerator) spawns_async_grpo_fsdp2_worker.pyviaaccelerate launchand a dedicatedfsdp2_reshard.yaml(2 ranks, FSDP v2,fsdp_reshard_after_forward: true). The worker runs two training steps with an in-process stub rollout queue and a no-op weight transfer so the test exercises FSDP2 parameter/optimizer lifecycle without vLLM or NCCL weight sync. It emits a singleASYNC_GRPO_FSDP2_RESULTJSON line; pytest asserts finite loss, parameter updates, full step count, and launch shape (2 processes, FSDP2, DTensors).The pytest side sets
PYTHONPATHto the repo root, passes output dir via env, and on timeout kills the full/procdescendant tree so detached elastic ranks cannot hang the suite.Reviewed by Cursor Bugbot for commit d03ecc2. Bugbot is set up for automated code reviews on this repo. Configure here.