From bb72789d7b63e26cd02d8f4d0b827341072dee6e Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Mon, 22 Jun 2026 19:09:52 -0700
Subject: [PATCH 01/13] test: add FSDP2 distributed coverage for
 AsyncGRPOTrainer

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.
---
 .../experimental/_async_grpo_fsdp2_worker.py  | 177 ++++++++++++++++++
 .../accelerate_configs/fsdp2_reshard.yaml     |  29 +++
 tests/experimental/test_async_grpo_trainer.py |  51 ++++-
 3 files changed, 256 insertions(+), 1 deletion(-)
 create mode 100644 tests/experimental/_async_grpo_fsdp2_worker.py
 create mode 100644 tests/experimental/data/accelerate_configs/fsdp2_reshard.yaml

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
new file mode 100644
index 00000000000..88dda4bca4b
--- /dev/null
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -0,0 +1,177 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Companion worker launched under ``accelerate launch --config_file <fsdp2>`` by the FSDP2 case in
+``test_async_grpo_trainer.py``.
+
+It runs a couple of :class:`AsyncGRPOTrainer` steps on an FSDP2-sharded model, driven by an in-process
+stub rollout worker (no vLLM server, no NCCL weight transfer), and checks that training actually
+progresses under FSDP2: the loss is finite and the parameters change. It then prints one
+machine-parseable result line (``ASYNC_GRPO_FSDP2_RESULT {json}``) that the pytest side asserts on.
+
+This is a *functional* FSDP2 smoke, not a performance microbenchmark. (An earlier version tried to count
+``lm_head.weight`` all-gathers to answer PR #6077's per-chunk re-gather question, but under FSDP2 those
+gathers are driven by autograd unshard hooks, not by ``DTensor.full_tensor``, and the trainer's own
+weight-sync path calls ``full_tensor`` on every parameter every step — so a ``full_tensor`` counter
+cannot isolate the chunked-logprob path. The #6077 question is instead settled by static analysis:
+``patch_chunked_lm_head`` uses a plain custom autograd Function with no ``torch.utils.checkpoint``
+recompute, so the per-chunk re-gather mechanism that PR #6077 fixed for SFT's ``chunked_nll`` is
+structurally absent here.)
+
+Self-contained on purpose (mirrors ``tests/experimental/_openreward_echo_env.py``): it imports only
+public TRL symbols and carries its own stub, so it never imports pytest-internal classes across the
+subprocess boundary.
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+import queue
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer
+from trl.experimental.async_grpo.async_rollout_worker import RolloutSample
+
+
+MODEL_ID = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+RESULT_PREFIX = "ASYNC_GRPO_FSDP2_RESULT"
+
+
+def dummy_reward_func(completions, **kwargs):
+    # Mirrors tests/experimental/test_async_grpo_trainer.py: the stub pre-computes rewards, so this is
+    # only here to satisfy the trainer's required `reward_funcs` argument.
+    return [float(hash(c[0]["content"]) % 100) / 100.0 for c in completions]
+
+
+class _StubRolloutWorker:
+    """Minimal in-process rollout worker — same shape as the one in test_async_grpo_trainer.py.
+
+    Reproduced here (rather than imported) because this module runs as ``__main__`` under
+    ``accelerate launch``, not as a pytest module, so importing the test class would be fragile. Keeping
+    it self-contained matches the openreward companion-script precedent.
+    """
+
+    def __init__(self, tokenizer, dataset, num_generations: int = 3, samples_per_weight_sync: int = 10):
+        self.rollout_buffer = queue.Queue()
+        self._samples_per_weight_sync = samples_per_weight_sync
+        self._model_version = 0
+        self._sample_iter = self._make_sample_iter(tokenizer, dataset, num_generations)
+
+    def _make_sample_iter(self, tokenizer, dataset, num_generations):
+        for row in itertools.cycle(dataset):
+            completions = [
+                [{"role": "assistant", "content": f"{row['completion'][0]['content']} {idx}"}]
+                for idx in range(num_generations)
+            ]
+            prompt_completions = [row["prompt"] + completion for completion in completions]
+            prompt_ids = tokenizer.apply_chat_template(
+                row["prompt"], tokenize=True, add_generation_prompt=True, return_dict=False
+            )
+            prompt_completion_ids = tokenizer.apply_chat_template(
+                prompt_completions, tokenize=True, add_generation_prompt=False, return_dict=False
+            )
+            rewards = np.array(dummy_reward_func(completions))
+            advantages = (rewards - rewards.mean()) / rewards.std()
+            for idx in range(num_generations):
+                completion_ids = prompt_completion_ids[idx][len(prompt_ids) :]
+                yield RolloutSample(
+                    prompt=row["prompt"],
+                    completion=completions[idx],
+                    input_ids=prompt_ids + completion_ids,
+                    completion_mask=[0] * len(prompt_ids) + [1] * len(completion_ids),
+                    old_log_probs=[0.0] * len(prompt_ids) + [-0.5] * len(completion_ids),
+                    advantage=float(advantages[idx]),
+                    model_version=self._model_version,
+                    metrics={"reward": float(rewards[idx]), "reward_std": float(rewards.std())},
+                )
+
+    def _fill_queue(self):
+        for _ in range(self._samples_per_weight_sync):
+            self.rollout_buffer.put(next(self._sample_iter))
+
+    def start(self):
+        self._fill_queue()
+
+    def update_model_version(self, version):
+        self._model_version = version
+        self._fill_queue()
+
+    def stop(self):
+        pass
+
+    def check_health(self, stale_after_s):
+        pass
+
+
+def main() -> None:
+    dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_completion", split="train")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+    # Same minimal, memory-frugal config as the existing single-process test_train, with 2 steps so we
+    # exercise the optimizer loop more than once under FSDP2.
+    args = AsyncGRPOConfig(
+        output_dir="async_grpo_fsdp2_out",
+        learning_rate=0.1,
+        per_device_train_batch_size=3,
+        num_generations=3,
+        max_completion_length=8,
+        max_steps=2,
+        vllm_server_timeout=5.0,
+        report_to="none",
+    )
+    trainer = AsyncGRPOTrainer(
+        model=MODEL_ID,
+        reward_funcs=dummy_reward_func,
+        args=args,
+        train_dataset=dataset,
+        rollout_worker=_StubRolloutWorker(tokenizer, dataset, num_generations=3),
+    )
+
+    # Snapshot params before training so we can confirm FSDP2 training actually updated them.
+    before = {n: p.detach().clone() for n, p in trainer.model.named_parameters()}
+
+    trainer.train()
+
+    # Did any parameter change? Materialize DTensors (full_tensor) and move both operands to CPU before
+    # comparing: the `before` snapshot is captured at construction (pre-FSDP-wrap, plain tensor) while the
+    # post-train param is an FSDP2 DTensor on CUDA, so a direct torch.equal would raise a device mismatch.
+    def _materialize(t):
+        t = t.full_tensor() if isinstance(t, torch.distributed.tensor.DTensor) else t
+        return t.detach().cpu()
+
+    changed = False
+    for n, p in trainer.model.named_parameters():
+        if not torch.equal(_materialize(before[n]), _materialize(p)):
+            changed = True
+            break
+
+    last = trainer.state.log_history[-1] if trainer.state.log_history else {}
+    train_loss = last.get("train_loss")
+    result = {
+        "steps": trainer.state.global_step,
+        "params_changed": changed,
+        "train_loss_finite": train_loss is not None and bool(np.isfinite(train_loss)),
+    }
+    # Only rank 0 prints the asserted line, so the pytest side parses exactly one result.
+    if trainer.accelerator.is_main_process:
+        print(f"{RESULT_PREFIX} {json.dumps(result)}", flush=True)  # noqa: T201 - result channel for the launcher
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/experimental/data/accelerate_configs/fsdp2_reshard.yaml b/tests/experimental/data/accelerate_configs/fsdp2_reshard.yaml
new file mode 100644
index 00000000000..5bd38d8c30d
--- /dev/null
+++ b/tests/experimental/data/accelerate_configs/fsdp2_reshard.yaml
@@ -0,0 +1,29 @@
+# 2-process FSDP2 config for the async-GRPO FSDP2 functional test (test_train_fsdp2).
+#
+# `fsdp_reshard_after_forward: true` is set explicitly so the test exercises the resharding parameter
+# lifecycle rather than relying on FSDP's default. Mirrors `examples/accelerate_configs/fsdp2.yaml`
+# with `num_processes: 2` for a 2-GPU node.
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_activation_checkpointing: false
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_version: 2
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index b35bde58719..e618012ca26 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -13,7 +13,11 @@
 # limitations under the License.
 
 import itertools
+import json
+import os
 import queue
+import subprocess
+from pathlib import Path
 
 import numpy as np
 import torch
@@ -23,7 +27,14 @@
 from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer
 from trl.experimental.async_grpo.async_rollout_worker import RolloutSample
 
-from ..testing_utils import TrlTestCase
+from ..testing_utils import TrlTestCase, require_torch_multi_accelerator
+
+
+ROOT = Path(__file__).resolve().parents[2]
+_HERE = Path(__file__).parent
+_FSDP2_WORKER = _HERE / "_async_grpo_fsdp2_worker.py"
+_FSDP2_CONFIG = _HERE / "data" / "accelerate_configs" / "fsdp2_reshard.yaml"
+_FSDP2_RESULT_PREFIX = "ASYNC_GRPO_FSDP2_RESULT"
 
 
 def dummy_reward_func(completions, **kwargs):
@@ -128,3 +139,41 @@ def test_train(self):
         for n, param in previous_trainable_params.items():
             new_param = trainer.model.get_parameter(n)
             assert not torch.equal(param, new_param), f"Parameter {n} has not changed."
+
+    @require_torch_multi_accelerator
+    def test_train_fsdp2(self):
+        # Functional smoke: AsyncGRPOTrainer trains under a 2-process FSDP2 group. This exercises the
+        # `patch_chunked_lm_head` chunked-logprob path on FSDP2-sharded parameters end-to-end and confirms
+        # the optimizer actually updates them. The worker uses an in-process stub rollout worker (no vLLM
+        # server / NCCL weight transfer), so the only distributed surface is the FSDP2 parameter lifecycle.
+        #
+        # (This is NOT a #6077 all-gather microbenchmark: under FSDP2 the per-parameter gathers are driven
+        # by autograd unshard hooks, not by `DTensor.full_tensor`, and the trainer's weight-sync path calls
+        # `full_tensor` on every parameter every step — so counting `full_tensor` cannot isolate the chunk
+        # path. The #6077 question is settled by static analysis instead: `patch_chunked_lm_head` has no
+        # `torch.utils.checkpoint` recompute, so the per-chunk re-gather that PR #6077 fixed for SFT's
+        # `chunked_nll` is structurally absent here.)
+        #
+        # Pin the repo root onto PYTHONPATH for the child: `accelerate launch` re-execs each rank via
+        # torch.distributed.elastic, which sets sys.path[0] to the launched script's directory
+        # (tests/experimental/), not cwd. Without this, a non-editable `trl` already in site-packages
+        # would shadow the working tree and the test would exercise the wrong code.
+        env = os.environ.copy()
+        env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]).rstrip(os.pathsep)
+        result = subprocess.run(
+            ["accelerate", "launch", "--config_file", str(_FSDP2_CONFIG), str(_FSDP2_WORKER)],
+            env=env,
+            cwd=ROOT,
+            capture_output=True,
+            text=True,
+        )
+        assert result.returncode == 0, f"FSDP2 worker failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+
+        result_lines = [ln for ln in result.stdout.splitlines() if ln.startswith(_FSDP2_RESULT_PREFIX)]
+        assert len(result_lines) == 1, f"expected exactly one result line, got {result_lines}\n{result.stdout}"
+        measured = json.loads(result_lines[0][len(_FSDP2_RESULT_PREFIX) :].strip())
+
+        # Training actually ran under FSDP2, produced a finite loss, and updated the parameters.
+        assert measured["steps"] >= 1, f"no training steps ran: {measured}"
+        assert measured["train_loss_finite"], f"train loss not finite under FSDP2: {measured}"
+        assert measured["params_changed"], f"parameters did not change under FSDP2: {measured}"

From d6d198db7a6c6384e74105f5f51cc90ccf380b74 Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Wed, 24 Jun 2026 10:43:25 -0700
Subject: [PATCH 02/13] style: reflow worker docstrings to max_len 119
 (doc-builder)

---
 .../experimental/_async_grpo_fsdp2_worker.py  | 34 +++++++++----------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 88dda4bca4b..66a89dc6b75 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -15,23 +15,21 @@
 """Companion worker launched under ``accelerate launch --config_file <fsdp2>`` by the FSDP2 case in
 ``test_async_grpo_trainer.py``.
 
-It runs a couple of :class:`AsyncGRPOTrainer` steps on an FSDP2-sharded model, driven by an in-process
-stub rollout worker (no vLLM server, no NCCL weight transfer), and checks that training actually
-progresses under FSDP2: the loss is finite and the parameters change. It then prints one
-machine-parseable result line (``ASYNC_GRPO_FSDP2_RESULT {json}``) that the pytest side asserts on.
+It runs a couple of :class:`AsyncGRPOTrainer` steps on an FSDP2-sharded model, driven by an in-process stub rollout
+worker (no vLLM server, no NCCL weight transfer), and checks that training actually progresses under FSDP2: the loss is
+finite and the parameters change. It then prints one machine-parseable result line (``ASYNC_GRPO_FSDP2_RESULT {json}``)
+that the pytest side asserts on.
 
 This is a *functional* FSDP2 smoke, not a performance microbenchmark. (An earlier version tried to count
-``lm_head.weight`` all-gathers to answer PR #6077's per-chunk re-gather question, but under FSDP2 those
-gathers are driven by autograd unshard hooks, not by ``DTensor.full_tensor``, and the trainer's own
-weight-sync path calls ``full_tensor`` on every parameter every step — so a ``full_tensor`` counter
-cannot isolate the chunked-logprob path. The #6077 question is instead settled by static analysis:
-``patch_chunked_lm_head`` uses a plain custom autograd Function with no ``torch.utils.checkpoint``
-recompute, so the per-chunk re-gather mechanism that PR #6077 fixed for SFT's ``chunked_nll`` is
-structurally absent here.)
-
-Self-contained on purpose (mirrors ``tests/experimental/_openreward_echo_env.py``): it imports only
-public TRL symbols and carries its own stub, so it never imports pytest-internal classes across the
-subprocess boundary.
+``lm_head.weight`` all-gathers to answer PR #6077's per-chunk re-gather question, but under FSDP2 those gathers are
+driven by autograd unshard hooks, not by ``DTensor.full_tensor``, and the trainer's own weight-sync path calls
+``full_tensor`` on every parameter every step — so a ``full_tensor`` counter cannot isolate the chunked-logprob path.
+The #6077 question is instead settled by static analysis: ``patch_chunked_lm_head`` uses a plain custom autograd
+Function with no ``torch.utils.checkpoint`` recompute, so the per-chunk re-gather mechanism that PR #6077 fixed for
+SFT's ``chunked_nll`` is structurally absent here.)
+
+Self-contained on purpose (mirrors ``tests/experimental/_openreward_echo_env.py``): it imports only public TRL symbols
+and carries its own stub, so it never imports pytest-internal classes across the subprocess boundary.
 """
 
 from __future__ import annotations
@@ -62,9 +60,9 @@ def dummy_reward_func(completions, **kwargs):
 class _StubRolloutWorker:
     """Minimal in-process rollout worker — same shape as the one in test_async_grpo_trainer.py.
 
-    Reproduced here (rather than imported) because this module runs as ``__main__`` under
-    ``accelerate launch``, not as a pytest module, so importing the test class would be fragile. Keeping
-    it self-contained matches the openreward companion-script precedent.
+    Reproduced here (rather than imported) because this module runs as ``__main__`` under ``accelerate launch``, not as
+    a pytest module, so importing the test class would be fragile. Keeping it self-contained matches the openreward
+    companion-script precedent.
     """
 
     def __init__(self, tokenizer, dataset, num_generations: int = 3, samples_per_weight_sync: int = 10):

From 2fc1fe06ccdecc3a09562d6bc1166d17deb06dba Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Fri, 3 Jul 2026 08:57:02 -0700
Subject: [PATCH 03/13] test: mark AsyncGRPO FSDP2 smoke as slow so it runs in
 the multi-GPU 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.
---
 tests/experimental/_async_grpo_fsdp2_worker.py |  8 +-------
 tests/experimental/test_async_grpo_trainer.py  | 16 ++++++----------
 2 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 66a89dc6b75..1ea634e146a 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -20,13 +20,7 @@
 finite and the parameters change. It then prints one machine-parseable result line (``ASYNC_GRPO_FSDP2_RESULT {json}``)
 that the pytest side asserts on.
 
-This is a *functional* FSDP2 smoke, not a performance microbenchmark. (An earlier version tried to count
-``lm_head.weight`` all-gathers to answer PR #6077's per-chunk re-gather question, but under FSDP2 those gathers are
-driven by autograd unshard hooks, not by ``DTensor.full_tensor``, and the trainer's own weight-sync path calls
-``full_tensor`` on every parameter every step — so a ``full_tensor`` counter cannot isolate the chunked-logprob path.
-The #6077 question is instead settled by static analysis: ``patch_chunked_lm_head`` uses a plain custom autograd
-Function with no ``torch.utils.checkpoint`` recompute, so the per-chunk re-gather mechanism that PR #6077 fixed for
-SFT's ``chunked_nll`` is structurally absent here.)
+This is a *functional* FSDP2 smoke, not a #6077 all-gather microbenchmark.
 
 Self-contained on purpose (mirrors ``tests/experimental/_openreward_echo_env.py``): it imports only public TRL symbols
 and carries its own stub, so it never imports pytest-internal classes across the subprocess boundary.
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index 1e7c0f25b3f..6c1a230f71a 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -146,19 +146,15 @@ def test_train(self):
             new_param = trainer.model.get_parameter(n)
             assert not torch.equal(param, new_param), f"Parameter {n} has not changed."
 
+    @pytest.mark.slow
     @require_torch_multi_accelerator
     def test_train_fsdp2(self):
-        # Functional smoke: AsyncGRPOTrainer trains under a 2-process FSDP2 group. This exercises the
-        # `patch_chunked_lm_head` chunked-logprob path on FSDP2-sharded parameters end-to-end and confirms
-        # the optimizer actually updates them. The worker uses an in-process stub rollout worker (no vLLM
-        # server / NCCL weight transfer), so the only distributed surface is the FSDP2 parameter lifecycle.
+        # Functional smoke: AsyncGRPOTrainer trains under a 2-process FSDP2 group, confirming the optimizer
+        # updates the FSDP2-sharded parameters. Uses an in-process stub rollout worker (no vLLM server /
+        # NCCL weight transfer), so the only distributed surface is the FSDP2 parameter lifecycle.
         #
-        # (This is NOT a #6077 all-gather microbenchmark: under FSDP2 the per-parameter gathers are driven
-        # by autograd unshard hooks, not by `DTensor.full_tensor`, and the trainer's weight-sync path calls
-        # `full_tensor` on every parameter every step — so counting `full_tensor` cannot isolate the chunk
-        # path. The #6077 question is settled by static analysis instead: `patch_chunked_lm_head` has no
-        # `torch.utils.checkpoint` recompute, so the per-chunk re-gather that PR #6077 fixed for SFT's
-        # `chunked_nll` is structurally absent here.)
+        # `@pytest.mark.slow` so it runs in the multi-GPU `slow_tests` lane; the experimental lane is
+        # single-GPU, where `@require_torch_multi_accelerator` would otherwise skip it permanently.
         #
         # Pin the repo root onto PYTHONPATH for the child: `accelerate launch` re-execs each rank via
         # torch.distributed.elastic, which sets sys.path[0] to the launched script's directory

From 002095dfef0c4dac022eb78a7d37972d3ce4eb0b Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Wed, 5 Aug 2026 17:06:40 -0700
Subject: [PATCH 04/13] test: fix the AsyncGRPO FSDP2 worker stub against the
 current rollout 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.
---
 .../experimental/_async_grpo_fsdp2_worker.py  | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 1ea634e146a..a3db5e2d17c 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -66,7 +66,7 @@ def __init__(self, tokenizer, dataset, num_generations: int = 3, samples_per_wei
         self._sample_iter = self._make_sample_iter(tokenizer, dataset, num_generations)
 
     def _make_sample_iter(self, tokenizer, dataset, num_generations):
-        for row in itertools.cycle(dataset):
+        for group_id, row in enumerate(itertools.cycle(dataset)):
             completions = [
                 [{"role": "assistant", "content": f"{row['completion'][0]['content']} {idx}"}]
                 for idx in range(num_generations)
@@ -90,6 +90,7 @@ def _make_sample_iter(self, tokenizer, dataset, num_generations):
                     old_log_probs=[0.0] * len(prompt_ids) + [-0.5] * len(completion_ids),
                     advantage=float(advantages[idx]),
                     model_version=self._model_version,
+                    group_id=group_id,  # every completion of one prompt belongs to the same group
                     metrics={"reward": float(rewards[idx]), "reward_std": float(rewards.std())},
                 )
 
@@ -111,6 +112,27 @@ def check_health(self, stale_after_s):
         pass
 
 
+class _NoOpWeightTransfer:
+    """No-op `WeightTransferProtocol`, so the smoke exercises only the FSDP2 parameter lifecycle.
+
+    Without it `AsyncGRPOTrainer` builds the default `WeightTransferClient`, which needs a live vLLM server to stream
+    weights into over NCCL.
+    """
+
+    def init_weight_transfer(self) -> None: ...
+
+    def pause(self) -> None: ...
+
+    def send_weights(self, iterator) -> None:
+        # Drain the iterator so the trainer's weight-gathering path still runs end to end.
+        for _ in iterator:
+            pass
+
+    def resume(self) -> None: ...
+
+    def destroy(self) -> None: ...
+
+
 def main() -> None:
     dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_completion", split="train")
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -133,6 +155,7 @@ def main() -> None:
         args=args,
         train_dataset=dataset,
         rollout_worker=_StubRolloutWorker(tokenizer, dataset, num_generations=3),
+        weight_transfer=_NoOpWeightTransfer(),
     )
 
     # Snapshot params before training so we can confirm FSDP2 training actually updated them.

From cfb6e30adf492909c5576764f9264971c8b6ec3f Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Wed, 19 Aug 2026 22:58:53 -0700
Subject: [PATCH 05/13] test: give the FSDP2 stub worker the metrics_queue the
 trainer now drains

#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.
---
 tests/experimental/_async_grpo_fsdp2_worker.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index a3db5e2d17c..a1a7edab433 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -61,6 +61,7 @@ class _StubRolloutWorker:
 
     def __init__(self, tokenizer, dataset, num_generations: int = 3, samples_per_weight_sync: int = 10):
         self.rollout_buffer = queue.Queue()
+        self.metrics_queue = queue.Queue()  # drained by the trainer in `log()`; this stub measures nothing
         self._samples_per_weight_sync = samples_per_weight_sync
         self._model_version = 0
         self._sample_iter = self._make_sample_iter(tokenizer, dataset, num_generations)

From 9d68df7424dce7b767ee1e24d6f8988eea0c71a5 Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Thu, 20 Aug 2026 04:37:52 -0700
Subject: [PATCH 06/13] test: fix the two defects that kept the FSDP2 smoke
 from ever reporting

`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.
---
 tests/experimental/_async_grpo_fsdp2_worker.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index a1a7edab433..5a6991af7ee 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -146,6 +146,7 @@ def main() -> None:
         per_device_train_batch_size=3,
         num_generations=3,
         max_completion_length=8,
+        token_budget=256,  # set explicitly; the stub worker has no real vLLM server to query for max_model_len
         max_steps=2,
         vllm_server_timeout=5.0,
         report_to="none",
@@ -171,11 +172,14 @@ def _materialize(t):
         t = t.full_tensor() if isinstance(t, torch.distributed.tensor.DTensor) else t
         return t.detach().cpu()
 
-    changed = False
-    for n, p in trainer.model.named_parameters():
-        if not torch.equal(_materialize(before[n]), _materialize(p)):
-            changed = True
-            break
+    # Compare every parameter on every rank before deciding: `_materialize` calls the collective
+    # `full_tensor()`, so breaking early would leave the ranks issuing different numbers of collectives and
+    # rank 0 would hang instead of reporting `params_changed: false`. A list comprehension is deliberate,
+    # since `any()` over a generator short-circuits the same way `break` does. Only rank 0's verdict is
+    # asserted: with `fsdp_cpu_ram_efficient_loading`, the pre-wrap snapshot on other ranks holds
+    # placeholders rather than the loaded weights.
+    diffs = [not torch.equal(_materialize(before[n]), _materialize(p)) for n, p in trainer.model.named_parameters()]
+    changed = any(diffs)
 
     last = trainer.state.log_history[-1] if trainer.state.log_history else {}
     train_loss = last.get("train_loss")

From 6c431cd91d2fee4adaf6d916273c932e5dbbfc46 Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Thu, 20 Aug 2026 18:28:45 -0700
Subject: [PATCH 07/13] test: stop the FSDP2 smoke test starving on
 token-budgeted batching

`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.
---
 .../experimental/_async_grpo_fsdp2_worker.py  |  9 ++++++--
 tests/experimental/test_async_grpo_trainer.py | 23 +++++++++++++------
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 5a6991af7ee..3ee80d7414e 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -146,7 +146,10 @@ def main() -> None:
         per_device_train_batch_size=3,
         num_generations=3,
         max_completion_length=8,
-        token_budget=256,  # set explicitly; the stub worker has no real vLLM server to query for max_model_len
+        # 0 selects the count-based FixedCountBatcher and, being non-None, still skips the vLLM
+        # max_model_len lookup. A positive budget starves here: the stub's samples are short enough
+        # that 2 rank-rows never fill, so TokenBudgetBatcher would never emit a micro-batch.
+        token_budget=0,
         max_steps=2,
         vllm_server_timeout=5.0,
         report_to="none",
@@ -156,7 +159,9 @@ def main() -> None:
         reward_funcs=dummy_reward_func,
         args=args,
         train_dataset=dataset,
-        rollout_worker=_StubRolloutWorker(tokenizer, dataset, num_generations=3),
+        # 24 = 4 x microbatch_size (per_device_train_batch_size 3 x 2 ranks); max_steps=2 needs 2,
+        # so the initial fill covers the whole run without relying on a weight-sync refill.
+        rollout_worker=_StubRolloutWorker(tokenizer, dataset, num_generations=3, samples_per_weight_sync=24),
         weight_transfer=_NoOpWeightTransfer(),
     )
 
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index fb88c208c45..8c4bb2648fd 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -334,13 +334,22 @@ def test_train_fsdp2(self):
         # would shadow the working tree and the test would exercise the wrong code.
         env = os.environ.copy()
         env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]).rstrip(os.pathsep)
-        result = subprocess.run(
-            ["accelerate", "launch", "--config_file", str(_FSDP2_CONFIG), str(_FSDP2_WORKER)],
-            env=env,
-            cwd=ROOT,
-            capture_output=True,
-            text=True,
-        )
+        # Bound the child: the trainer's rollout consumer blocks indefinitely on an empty queue (it only
+        # calls `check_health`, which this stub implements as a no-op), so a starved run would hang the
+        # pytest process rather than fail it. The timeout turns that into a readable failure.
+        try:
+            result = subprocess.run(
+                ["accelerate", "launch", "--config_file", str(_FSDP2_CONFIG), str(_FSDP2_WORKER)],
+                env=env,
+                cwd=ROOT,
+                capture_output=True,
+                text=True,
+                timeout=900,
+            )
+        except subprocess.TimeoutExpired as exc:
+            stdout = exc.stdout.decode() if isinstance(exc.stdout, bytes) else (exc.stdout or "")
+            stderr = exc.stderr.decode() if isinstance(exc.stderr, bytes) else (exc.stderr or "")
+            pytest.fail(f"FSDP2 worker timed out after {exc.timeout}s:\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}")
         assert result.returncode == 0, f"FSDP2 worker failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
 
         result_lines = [ln for ln in result.stdout.splitlines() if ln.startswith(_FSDP2_RESULT_PREFIX)]

From 2433b1e3e9a8796d693270ff8748bb07eeb82af9 Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Thu, 20 Aug 2026 18:45:09 -0700
Subject: [PATCH 08/13] test: xfail the FSDP2 smoke on the FA head_size limit
 (#6837)

`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.
---
 tests/experimental/test_async_grpo_trainer.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index 8c4bb2648fd..b0ec22aed3d 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -320,13 +320,20 @@ def reset(self, **kwargs): ...
 
     @pytest.mark.slow
     @require_torch_multi_accelerator
+    @pytest.mark.xfail(
+        reason="Flash Attention rejects a head_size that is not a multiple of 8, and the tiny models are "
+        "hidden_size=8 over 4 attention heads (head_size=2), so the forward pass raises "
+        "(https://github.com/huggingface/trl/issues/6837)",
+    )
     def test_train_fsdp2(self):
         # Functional smoke: AsyncGRPOTrainer trains under a 2-process FSDP2 group, confirming the optimizer
         # updates the FSDP2-sharded parameters. Uses an in-process stub rollout worker (no vLLM server /
         # NCCL weight transfer), so the only distributed surface is the FSDP2 parameter lifecycle.
         #
-        # `@pytest.mark.slow` so it runs in the multi-GPU `slow_tests` lane; the experimental lane is
-        # single-GPU, where `@require_torch_multi_accelerator` would otherwise skip it permanently.
+        # `@pytest.mark.slow` marks the cost, but no lane collects this test today: `slow_tests` is
+        # `pytest -m "slow" tests/`, and `norecursedirs` excludes `tests/experimental` from that recursive
+        # collection. `test_experimental` passes an explicit path so it does collect the test, but its runner
+        # is single-GPU, where `@require_torch_multi_accelerator` skips it.
         #
         # Pin the repo root onto PYTHONPATH for the child: `accelerate launch` re-execs each rank via
         # torch.distributed.elastic, which sets sys.path[0] to the launched script's directory

From 72bf811c05787d1eca3f77c60f079f727101ab9e Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Fri, 21 Aug 2026 10:35:16 -0700
Subject: [PATCH 09/13] test(async-grpo): use a Flash Attention compatible
 model in the FSDP2 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.
---
 tests/experimental/_async_grpo_fsdp2_worker.py | 4 +++-
 tests/experimental/test_async_grpo_trainer.py  | 5 -----
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 3ee80d7414e..787c540c487 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -41,7 +41,9 @@
 from trl.experimental.async_grpo.async_rollout_worker import RolloutSample
 
 
-MODEL_ID = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+# The trainer loads the model with Flash Attention, which requires a `head_size` multiple of 8. Hence the `small-*`
+# model (`head_size=32`) below, rather than the usual `tiny-*` one (`head_size=2`).
+MODEL_ID = "trl-internal-testing/small-Qwen2ForCausalLM-2.5"
 RESULT_PREFIX = "ASYNC_GRPO_FSDP2_RESULT"
 
 
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index 3080037d558..c7dfe3b888a 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -314,11 +314,6 @@ def reset(self, **kwargs): ...
 
     @pytest.mark.slow
     @require_torch_multi_accelerator
-    @pytest.mark.xfail(
-        reason="Flash Attention rejects a head_size that is not a multiple of 8, and the tiny models are "
-        "hidden_size=8 over 4 attention heads (head_size=2), so the forward pass raises "
-        "(https://github.com/huggingface/trl/issues/6837)",
-    )
     def test_train_fsdp2(self):
         # Functional smoke: AsyncGRPOTrainer trains under a 2-process FSDP2 group, confirming the optimizer
         # updates the FSDP2-sharded parameters. Uses an in-process stub rollout worker (no vLLM server /

From c75c6596b48a6981c5dbe598d2f90f5c71c812db Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Fri, 21 Aug 2026 10:57:10 -0700
Subject: [PATCH 10/13] test(async-grpo): write FSDP2 worker artifacts to the
 test's temp dir

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.
---
 tests/experimental/_async_grpo_fsdp2_worker.py | 5 ++++-
 tests/experimental/test_async_grpo_trainer.py  | 2 ++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 787c540c487..9ac0f99af5d 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -30,6 +30,7 @@
 
 import itertools
 import json
+import os
 import queue
 
 import numpy as np
@@ -143,7 +144,9 @@ def main() -> None:
     # Same minimal, memory-frugal config as the existing single-process test_train, with 2 steps so we
     # exercise the optimizer loop more than once under FSDP2.
     args = AsyncGRPOConfig(
-        output_dir="async_grpo_fsdp2_out",
+        # The launcher runs this worker with `cwd` at the repo root, so a relative `output_dir` would drop trainer
+        # artifacts into the working tree. The test owns a temporary directory and passes it in.
+        output_dir=os.environ["ASYNC_GRPO_FSDP2_OUTPUT_DIR"],
         learning_rate=0.1,
         per_device_train_batch_size=3,
         num_generations=3,
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index c7dfe3b888a..f2775486db6 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -330,6 +330,8 @@ def test_train_fsdp2(self):
         # would shadow the working tree and the test would exercise the wrong code.
         env = os.environ.copy()
         env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]).rstrip(os.pathsep)
+        # `cwd` below is the repo root, so hand the worker somewhere else to write trainer artifacts.
+        env["ASYNC_GRPO_FSDP2_OUTPUT_DIR"] = str(self.tmp_dir)
         # Bound the child: the trainer's rollout consumer blocks indefinitely on an empty queue (it only
         # calls `check_health`, which this stub implements as a no-op), so a starved run would hang the
         # pytest process rather than fail it. The timeout turns that into a readable failure.

From 2a112af35ce87d42cb2385db8923ffc45d54a7b7 Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Wed, 2 Sep 2026 22:21:11 -0700
Subject: [PATCH 11/13] test(async-grpo): kill the whole launcher group on
 timeout, require every 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.
---
 .../experimental/_async_grpo_fsdp2_worker.py  | 10 ++++--
 tests/experimental/test_async_grpo_trainer.py | 36 ++++++++++++-------
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 9ac0f99af5d..845444501e2 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -46,6 +46,8 @@
 # model (`head_size=32`) below, rather than the usual `tiny-*` one (`head_size=2`).
 MODEL_ID = "trl-internal-testing/small-Qwen2ForCausalLM-2.5"
 RESULT_PREFIX = "ASYNC_GRPO_FSDP2_RESULT"
+# Reported alongside the measured step count so the launcher can require the whole loop to have run.
+_MAX_STEPS = 2
 
 
 def dummy_reward_func(completions, **kwargs):
@@ -82,7 +84,10 @@ def _make_sample_iter(self, tokenizer, dataset, num_generations):
             prompt_completion_ids = tokenizer.apply_chat_template(
                 prompt_completions, tokenize=True, add_generation_prompt=False, return_dict=False
             )
-            rewards = np.array(dummy_reward_func(completions))
+            # Distinct rewards by construction. Hash-derived rewards can collide within a group under a randomized
+            # PYTHONHASHSEED, which makes `rewards.std()` zero and the advantages NaN, so a run would fail for a reason
+            # unrelated to FSDP2.
+            rewards = np.linspace(0.0, 1.0, num_generations)
             advantages = (rewards - rewards.mean()) / rewards.std()
             for idx in range(num_generations):
                 completion_ids = prompt_completion_ids[idx][len(prompt_ids) :]
@@ -155,7 +160,7 @@ def main() -> None:
         # max_model_len lookup. A positive budget starves here: the stub's samples are short enough
         # that 2 rank-rows never fill, so TokenBudgetBatcher would never emit a micro-batch.
         token_budget=0,
-        max_steps=2,
+        max_steps=_MAX_STEPS,
         vllm_server_timeout=5.0,
         report_to="none",
     )
@@ -195,6 +200,7 @@ def _materialize(t):
     train_loss = last.get("train_loss")
     result = {
         "steps": trainer.state.global_step,
+        "max_steps": _MAX_STEPS,
         "params_changed": changed,
         "train_loss_finite": train_loss is not None and bool(np.isfinite(train_loss)),
     }
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index 656ce4bee19..1ce39f7db85 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -19,6 +19,7 @@
 import multiprocessing as mp
 import os
 import queue
+import signal
 import subprocess
 from collections import defaultdict
 from pathlib import Path
@@ -335,19 +336,25 @@ def test_train_fsdp2(self):
         # Bound the child: the trainer's rollout consumer blocks indefinitely on an empty queue (it only
         # calls `check_health`, which this stub implements as a no-op), so a starved run would hang the
         # pytest process rather than fail it. The timeout turns that into a readable failure.
+        # `accelerate launch` starts each rank through torch elastic, which puts the workers in their own session
+        # (`start_new_session=True`), so a plain `subprocess.run(timeout=...)` would kill only the launcher and leave
+        # the ranks running. Put the launcher in a fresh process group and kill the whole group on timeout.
+        proc = subprocess.Popen(
+            ["accelerate", "launch", "--config_file", str(_FSDP2_CONFIG), str(_FSDP2_WORKER)],
+            env=env,
+            cwd=ROOT,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            start_new_session=True,
+        )
         try:
-            result = subprocess.run(
-                ["accelerate", "launch", "--config_file", str(_FSDP2_CONFIG), str(_FSDP2_WORKER)],
-                env=env,
-                cwd=ROOT,
-                capture_output=True,
-                text=True,
-                timeout=900,
-            )
-        except subprocess.TimeoutExpired as exc:
-            stdout = exc.stdout.decode() if isinstance(exc.stdout, bytes) else (exc.stdout or "")
-            stderr = exc.stderr.decode() if isinstance(exc.stderr, bytes) else (exc.stderr or "")
-            pytest.fail(f"FSDP2 worker timed out after {exc.timeout}s:\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}")
+            stdout, stderr = proc.communicate(timeout=900)
+        except subprocess.TimeoutExpired:
+            os.killpg(proc.pid, signal.SIGKILL)
+            stdout, stderr = proc.communicate()
+            pytest.fail(f"FSDP2 worker timed out after 900s:\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}")
+        result = subprocess.CompletedProcess(proc.args, proc.returncode, stdout, stderr)
         assert result.returncode == 0, f"FSDP2 worker failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
 
         result_lines = [ln for ln in result.stdout.splitlines() if ln.startswith(_FSDP2_RESULT_PREFIX)]
@@ -355,7 +362,10 @@ def test_train_fsdp2(self):
         measured = json.loads(result_lines[0][len(_FSDP2_RESULT_PREFIX) :].strip())
 
         # Training actually ran under FSDP2, produced a finite loss, and updated the parameters.
-        assert measured["steps"] >= 1, f"no training steps ran: {measured}"
+        # The worker configures more than one step so the optimizer loop runs repeatedly under FSDP2; accepting fewer
+        # would let an early stop after step 1 pass.
+        assert measured["max_steps"] > 1, f"worker must configure more than one step: {measured}"
+        assert measured["steps"] == measured["max_steps"], f"not every configured step ran: {measured}"
         assert measured["train_loss_finite"], f"train loss not finite under FSDP2: {measured}"
         assert measured["params_changed"], f"parameters did not change under FSDP2: {measured}"
 

From 059e072dcaf2b1b73d5640fc97d8dc180b0c9bf7 Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Wed, 2 Sep 2026 22:59:11 -0700
Subject: [PATCH 12/13] test(async-grpo): kill the whole launcher tree on
 timeout, not the launcher'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.
---
 tests/experimental/test_async_grpo_trainer.py | 31 ++++++++++++++++---
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index 1ce39f7db85..f7686d28936 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -75,6 +75,22 @@ def dummy_reward_func(completions, **kwargs):
     return [float(hash(c[0]["content"]) % 100) / 100.0 for c in completions]
 
 
+def _descendant_pids(pid: int) -> list[int]:
+    """Every process below `pid` in the /proc tree, so a timeout can kill ranks that torch elastic detached into their
+    own sessions."""
+    found, stack = [], [pid]
+    while stack:
+        parent = stack.pop()
+        try:
+            children = Path(f"/proc/{parent}/task/{parent}/children").read_text().split()
+        except OSError:
+            children = []
+        for child in map(int, children):
+            found.append(child)
+            stack.append(child)
+    return found
+
+
 class _StubRolloutWorker:
     """Minimal rollout worker stub for testing the trainer in isolation."""
 
@@ -336,9 +352,11 @@ def test_train_fsdp2(self):
         # Bound the child: the trainer's rollout consumer blocks indefinitely on an empty queue (it only
         # calls `check_health`, which this stub implements as a no-op), so a starved run would hang the
         # pytest process rather than fail it. The timeout turns that into a readable failure.
-        # `accelerate launch` starts each rank through torch elastic, which puts the workers in their own session
-        # (`start_new_session=True`), so a plain `subprocess.run(timeout=...)` would kill only the launcher and leave
-        # the ranks running. Put the launcher in a fresh process group and kill the whole group on timeout.
+        # `accelerate launch` starts each rank through torch elastic, which puts every worker in its own session
+        # (`start_new_session=True`), so neither `subprocess.run(timeout=...)` nor a kill of the launcher's process
+        # group reaches the ranks: they would keep the GPUs and hold the stdout pipe open, and the read after the
+        # timeout would never return. Collect the descendants from /proc while the launcher is still their parent
+        # and kill the whole tree.
         proc = subprocess.Popen(
             ["accelerate", "launch", "--config_file", str(_FSDP2_CONFIG), str(_FSDP2_WORKER)],
             env=env,
@@ -346,12 +364,15 @@ def test_train_fsdp2(self):
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             text=True,
-            start_new_session=True,
         )
         try:
             stdout, stderr = proc.communicate(timeout=900)
         except subprocess.TimeoutExpired:
-            os.killpg(proc.pid, signal.SIGKILL)
+            for pid in [proc.pid, *_descendant_pids(proc.pid)]:
+                try:
+                    os.kill(pid, signal.SIGKILL)
+                except ProcessLookupError:
+                    pass
             stdout, stderr = proc.communicate()
             pytest.fail(f"FSDP2 worker timed out after 900s:\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}")
         result = subprocess.CompletedProcess(proc.args, proc.returncode, stdout, stderr)

From 0c43a43c38b1413808ac77b39ef793f288cfec1e Mon Sep 17 00:00:00 2001
From: Behrooz Azarkhalili <ermiaazarkhalili@gmail.com>
Date: Thu, 3 Sep 2026 18:15:29 -0700
Subject: [PATCH 13/13] test(async_grpo): make the FSDP2 worker report its
 launch shape

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.
---
 tests/experimental/_async_grpo_fsdp2_worker.py | 7 +++++++
 tests/experimental/test_async_grpo_trainer.py  | 6 ++++++
 2 files changed, 13 insertions(+)

diff --git a/tests/experimental/_async_grpo_fsdp2_worker.py b/tests/experimental/_async_grpo_fsdp2_worker.py
index 845444501e2..bf4763024a0 100644
--- a/tests/experimental/_async_grpo_fsdp2_worker.py
+++ b/tests/experimental/_async_grpo_fsdp2_worker.py
@@ -198,11 +198,18 @@ def _materialize(t):
 
     last = trainer.state.log_history[-1] if trainer.state.log_history else {}
     train_loss = last.get("train_loss")
+    # The pytest side asserts on the launch shape too: a replicated single-process run would also change the
+    # parameters and report a finite loss, so world size, the distributed type and sharded parameters are reported.
+    accelerator = trainer.accelerator
     result = {
         "steps": trainer.state.global_step,
         "max_steps": _MAX_STEPS,
         "params_changed": changed,
         "train_loss_finite": train_loss is not None and bool(np.isfinite(train_loss)),
+        "num_processes": accelerator.num_processes,
+        "distributed_type": accelerator.distributed_type.value,
+        "fsdp_version": accelerator.state.fsdp_plugin.fsdp_version if accelerator.state.fsdp_plugin else None,
+        "sharded_params": sum(isinstance(p, torch.distributed.tensor.DTensor) for p in trainer.model.parameters()),
     }
     # Only rank 0 prints the asserted line, so the pytest side parses exactly one result.
     if trainer.accelerator.is_main_process:
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index f7686d28936..7093b74af22 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -389,6 +389,12 @@ def test_train_fsdp2(self):
         assert measured["steps"] == measured["max_steps"], f"not every configured step ran: {measured}"
         assert measured["train_loss_finite"], f"train loss not finite under FSDP2: {measured}"
         assert measured["params_changed"], f"parameters did not change under FSDP2: {measured}"
+        # A replicated single-process run would pass the checks above too, so pin the launch shape the config asks
+        # for: two ranks, FSDP version 2, and parameters that are DTensors after wrapping.
+        assert measured["num_processes"] == 2, f"expected a two-process launch: {measured}"
+        assert measured["distributed_type"] == "FSDP", f"not launched under FSDP: {measured}"
+        assert measured["fsdp_version"] == 2, f"not FSDP version 2: {measured}"
+        assert measured["sharded_params"] > 0, f"no parameter was sharded as a DTensor: {measured}"
 
 
 def _vision_parameter_names(model) -> set[str]: