diff --git a/recipes/configs/dev/qwen3B_async_grpo.yaml b/recipes/configs/dev/qwen3B_async_grpo.yaml
index 7f9f409f0d..c39dd101c0 100644
--- a/recipes/configs/dev/qwen3B_async_grpo.yaml
+++ b/recipes/configs/dev/qwen3B_async_grpo.yaml
@@ -18,6 +18,7 @@
 name: grpo_async_qwen3b
 output_dir: /tmp/checkpoints/${name}
 base_model_path: /tmp/Qwen2.5-3B
+device: cuda
 
 # Model architecture (Qwen2.5-3B)
 model:
@@ -39,7 +40,7 @@ orchestration:
   num_inference_workers: 4
   num_postprocessing_workers: 1
   num_training_workers: 2
-  replay_buffer_size: ${inference.batch_size}  # TODO: Right now this can't be bigger, or else we'll get padding issues
+  replay_buffer_size: ${inference.batch_size}  # Number of trajectory batches to keep
   num_steps: 250
 
 # All inference args
diff --git a/recipes/dev/async_grpo_full_finetune_distributed.py b/recipes/dev/async_grpo_full_finetune_distributed.py
index 422210ea24..798fc5b7b8 100644
--- a/recipes/dev/async_grpo_full_finetune_distributed.py
+++ b/recipes/dev/async_grpo_full_finetune_distributed.py
@@ -134,7 +134,9 @@ def setup(self, cfg: DictConfig):
             storage=functools.partial(
                 LazyStackStorage, max_size=cfg.orchestration.replay_buffer_size
             ),
-            batch_size=cfg.training.batch_size,
+            # Each item in the buffer is one whole trajectory batch (see
+            # PostProcessingWorker.run), so a sample of size 1 is a single batch.
+            batch_size=1,
             remote_config={"num_cpus": 10, "num_gpus": 0},
         )
 
diff --git a/recipes/dev/grpo_full_finetune_distributed.py b/recipes/dev/grpo_full_finetune_distributed.py
index 591ddd5c36..c829fa8f24 100644
--- a/recipes/dev/grpo_full_finetune_distributed.py
+++ b/recipes/dev/grpo_full_finetune_distributed.py
@@ -21,7 +21,7 @@
 from torchtune.config._utils import _get_component_from_path
 from torchtune.datasets import ConcatDataset
 from torchtune.dev.rl.generation import generate
-from torchtune.dev.rl.rewards import batched_rewards
+from torchtune.dev.rl.rewards import batched_rewards, group_normalized_advantages
 from torchtune.dev.rl.types import GRPOStats, GRPOTrajectory
 from torchtune.modules import local_kv_cache
 from torchtune.recipe_interfaces import FTRecipeInterface
@@ -656,9 +656,7 @@ def generate_trajectory(
         rewards = rewards.sum(dim=-1)  # [B, G]
         successes = successes.sum(dim=-1)  # [B, G]
 
-        advantages = (rewards - rewards.mean(1, keepdim=True)) / (
-            rewards.std(1, keepdim=True) + 1e-4
-        )
+        advantages = group_normalized_advantages(rewards)  # [B, G]
         advantages = advantages.reshape(batch_size * grpo_size)  # flatten
         del responses
         torch.cuda.empty_cache()
diff --git a/tests/torchtune/dev/rl/test_async_data_path.py b/tests/torchtune/dev/rl/test_async_data_path.py
new file mode 100644
index 0000000000..03a53026e4
--- /dev/null
+++ b/tests/torchtune/dev/rl/test_async_data_path.py
@@ -0,0 +1,284 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import itertools
+
+import pytest
+import torch
+from torch import nn
+
+from tensordict import NonTensorStack
+
+from torchtune.dev.rl.datatypes.trajectory import Trajectory
+from torchtune.dev.rl.linear_grpo_loss import LinearGRPOLoss
+from torchtune.dev.rl.rewards import group_normalized_advantages
+from torchrl.data import LazyStackStorage, ReplayBuffer, RoundRobinWriter
+
+
+def _make_trajectory(batch_size, seq_len, resp_len, num_funcs, seed=0, advantages=None):
+    """Build a trajectory in the same layout PostProcessingWorker emits."""
+    torch.manual_seed(seed)
+    rewards = torch.randn(batch_size, num_funcs)
+    successes = torch.rand(batch_size, num_funcs) > 0.5
+    func_names = [f"reward_fn_{i}" for i in range(num_funcs)]
+    if advantages is None:
+        advantages = torch.randn(batch_size)
+    return Trajectory(
+        query_responses=torch.randint(0, 100, (batch_size, seq_len)),
+        responses=torch.randint(0, 100, (batch_size, resp_len)),
+        logprobs=torch.randn(batch_size, resp_len),
+        ref_logprobs=torch.randn(batch_size, resp_len),
+        query_response_padding_masks=torch.ones(batch_size, seq_len, dtype=torch.bool),
+        seq_lens=torch.randint(1, resp_len, (batch_size,)),
+        answers=NonTensorStack(*[f"ans{i}" for i in range(batch_size)]),
+        policy_version=3,
+        advantages=advantages,
+        rewards=rewards,
+        successes=successes,
+        reward_func_names=NonTensorStack(*[func_names for _ in range(batch_size)]),
+        batch_size=[batch_size],
+        sequence_ids=NonTensorStack(
+            *[f"worker0_{i}" for i in range(batch_size)]
+        ),
+    )
+
+
+class TestAdvantageGradientOracle:
+    """Advantages must flow into the optimizer proportionally to their value.
+
+    The GRPO policy loss with ``kl_coeff=0`` reduces to ``-advantages.mean()``
+    per batch, so the gradient with respect to any model parameter is exactly
+    linear in the per-sample advantages. These tests pin that down so that any
+    regression that collapses advantages (e.g. all samples seeing the same
+    value, the #2943 symptom) is caught by construction.
+    """
+
+    def _setup(self, batch_size=4, seq_len=6):
+        torch.manual_seed(0)
+        num_output_chunks = 1
+        loss_fn = LinearGRPOLoss(
+            num_output_chunks=num_output_chunks, kl_coeff=0.0
+        )
+        head = nn.Linear(8, 32)
+        loss_fn.linear_projection = head
+        hidden = torch.randn(batch_size, seq_len, 8)
+        targets = torch.randint(0, 32, (batch_size, seq_len))
+        ref_logprobs = torch.randn(batch_size, seq_len)
+        masks = torch.ones(batch_size, seq_len, dtype=torch.bool)
+        return loss_fn, hidden, targets, ref_logprobs, masks
+
+    def test_loss_equals_neg_advantage_mean(self):
+        loss_fn, hidden, targets, ref_logprobs, masks = self._setup()
+        advantages = torch.tensor([1.0, -2.0, 0.5, 0.25])
+        loss, policy_loss, *_ = loss_fn(
+            hidden, targets, ref_logprobs, advantages, masks
+        )
+        # kl_coeff=0 -> the policy term is exactly -advantages[:, None] per token
+        assert torch.allclose(loss, -advantages.mean(), atol=1e-6)
+        assert torch.allclose(policy_loss, advantages.mean(), atol=1e-6)
+
+    def test_gradient_scales_linearly_with_advantages(self):
+        loss_fn, hidden, targets, ref_logprobs, masks = self._setup()
+
+        def grad_norm(advantages):
+            loss, *_ = loss_fn(hidden, targets, ref_logprobs, advantages, masks)
+            loss.backward(retain_graph=True)
+            norm = torch.cat(
+                [p.grad.flatten() for p in loss_fn.parameters() if p.grad is not None]
+            )
+            loss_fn.zero_grad()
+            return norm
+
+        adv = torch.tensor([1.0, -2.0, 0.5, 0.25])
+        g1 = grad_norm(adv)
+        g2 = grad_norm(2.0 * adv)
+        assert torch.allclose(g2, 2.0 * g1, atol=1e-6)
+
+    def test_gradient_is_additive_over_samples(self):
+        """Per-sample advantages add: grad([2,1]) == 2*grad([1,0]) + grad([0,1])."""
+        loss_fn, hidden, targets, ref_logprobs, masks = self._setup(batch_size=2)
+        ones_hot = [torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])]
+
+        def grad_vec(advantages):
+            loss, *_ = loss_fn(hidden, targets, ref_logprobs, advantages, masks)
+            loss.backward(retain_graph=True)
+            vec = torch.cat(
+                [p.grad.flatten() for p in loss_fn.parameters() if p.grad is not None]
+            )
+            loss_fn.zero_grad()
+            return vec
+
+        g_a, g_b = (grad_vec(v) for v in ones_hot)
+        combined = grad_vec(torch.tensor([2.0, 1.0]))
+        assert torch.allclose(combined, 2.0 * g_a + g_b, atol=1e-6)
+
+
+class TestAsyncDataPathFuzz:
+    """Fuzz the postprocessing -> replay buffer -> sampling round trip."""
+
+    @pytest.mark.parametrize("seed", range(5))
+    @pytest.mark.parametrize(
+        "batch_size,group_size,resp_len,num_funcs,capacity",
+        [
+            (2, 2, 8, 1, 2),  # config-like: buffer capacity in batches
+            (4, 2, 16, 2, 3),
+            (1, 4, 32, 3, 1),  # capacity 1: only the last batch survives
+            (2, 3, 11, 4, 5),
+        ],
+    )
+    def test_round_trip_preserves_everything(
+        self, seed, batch_size, group_size, resp_len, num_funcs, capacity
+    ):
+        grpo_samples = batch_size * group_size
+        buffer = ReplayBuffer(
+            storage=LazyStackStorage(max_size=capacity),
+            batch_size=1,
+            writer=RoundRobinWriter(),
+        )
+        batch = _make_trajectory(
+            grpo_samples, resp_len + 8, resp_len, num_funcs, seed=seed
+        )
+
+        # Storage counts each whole batch as one item (postprocessing contract).
+        buffer.extend(batch.unsqueeze(0))
+        sampled = buffer.sample(1).squeeze(0)
+
+        assert tuple(sampled.batch_size) == (grpo_samples,)
+        # The only batch in the buffer is returned wholesale: per-sample fields
+        # are bit-for-bit the ones that went in.
+        assert torch.equal(sampled.advantages, batch.advantages)
+        assert torch.equal(sampled.rewards, batch.rewards)
+        assert torch.equal(sampled.successes, batch.successes)
+        assert sampled.reward_func_names == batch.reward_func_names
+        assert sampled.sequence_ids == batch.sequence_ids
+        assert sampled.policy_version == batch.policy_version
+        assert torch.equal(sampled.logprobs, batch.logprobs)
+
+    @pytest.mark.parametrize("seed", range(5))
+    def test_advantages_vary_within_groups_after_round_trip(
+        self, seed, group_size=4, batch_size=3, resp_len=10, num_funcs=2
+    ):
+        """Regression for #2943: sampled advantages must NOT be identical."""
+        grpo_samples = batch_size * group_size
+        rewards = torch.randn(grpo_samples, num_funcs)
+        advantages = group_normalized_advantages(
+            rewards.reshape(batch_size, group_size, -1).sum(-1)
+        ).reshape(-1)
+
+        buffer = ReplayBuffer(
+            storage=LazyStackStorage(max_size=2),
+            batch_size=1,
+            writer=RoundRobinWriter(),
+        )
+        traj = _make_trajectory(grpo_samples, resp_len + 8, resp_len, num_funcs, seed)
+        traj = Trajectory(
+            query_responses=traj.query_responses,
+            responses=traj.responses,
+            logprobs=traj.logprobs,
+            ref_logprobs=traj.ref_logprobs,
+            query_response_padding_masks=traj.query_response_padding_masks,
+            seq_lens=traj.seq_lens,
+            answers=traj.answers,
+            policy_version=traj.policy_version,
+            advantages=advantages,
+            rewards=rewards,
+            successes=traj.successes,
+            reward_func_names=traj.reward_func_names,
+            batch_size=grpo_samples,
+            sequence_ids=traj.sequence_ids,
+        )
+        buffer.extend(traj.unsqueeze(0))
+        sampled = buffer.sample(1).squeeze(0)
+
+        adv = sampled.advantages.reshape(batch_size, group_size)
+        # Every group is centered at zero (GRPO normalization survived the trip).
+        assert torch.allclose(adv.mean(-1), torch.zeros(batch_size), atol=1e-5)
+        # The per-sample advantages are distinct within a group (the #2943 bug
+        # produced a constant value across all samples instead).
+        assert adv.std(-1).gt(1e-6).all()
+
+    def test_capacity_counts_batches_across_multiple_writes(self):
+        buffer = ReplayBuffer(
+            storage=LazyStackStorage(max_size=2),
+            batch_size=1,
+            writer=RoundRobinWriter(),
+        )
+        first = _make_trajectory(4, 16, 8, 1, seed=1)
+        second = _make_trajectory(4, 16, 8, 1, seed=2)
+        third = _make_trajectory(4, 16, 8, 1, seed=3)
+        buffer.extend(first.unsqueeze(0))
+        buffer.extend(second.unsqueeze(0))
+        buffer.extend(third.unsqueeze(0))
+
+        # Capacity 2 batches: the first batch was evicted, the last two remain,
+        # and sampling always returns one coherent whole batch.
+        for _ in range(10):
+            sampled = buffer.sample(1).squeeze(0)
+            assert torch.equal(
+                sampled.advantages, second.advantages
+            ) or torch.equal(sampled.advantages, third.advantages)
+
+
+class TestRewardFuncNamesLabeling:
+    def test_per_sample_names_from_reward_outputs(self):
+        """The collector emits trajectories without reward labels; the
+        postprocessing worker labels them from the reward functions it ran."""
+        from torch import tensor
+
+        from torchtune.dev.rl.rewards import RewardOutput
+        from torchtune.dev.rl.workers.postprocessing import (
+            reward_func_names_per_sample,
+        )
+
+        reward_outputs = [
+            RewardOutput(
+                reward_base_name="math_correctness",
+                total_reward=tensor([1.0, 0.0]),
+                successes=tensor([1, 0]),
+            ),
+            RewardOutput(
+                reward_base_name="formatting",
+                total_reward=tensor([1.0, 1.0]),
+                successes=tensor([1, 1]),
+            ),
+        ]
+        names = reward_func_names_per_sample(reward_outputs, num_samples=3)
+        assert len(names) == 3
+        for i in range(3):
+            assert names[i].data == ["math_correctness", "formatting"]
+
+    def test_labels_survive_buffer_round_trip(self):
+        """Per-sample reward names survive sampling from the replay buffer."""
+        buffer = ReplayBuffer(
+            storage=LazyStackStorage(max_size=2),
+            batch_size=1,
+            writer=RoundRobinWriter(),
+        )
+        traj = _make_trajectory(4, 16, 8, 2, seed=0)
+        buffer.extend(traj.unsqueeze(0))
+        sampled = buffer.sample(1).squeeze(0)
+        assert sampled.reward_func_names == traj.reward_func_names
+
+
+class TestGroupAdvantageAggregation:
+    def test_aggregation_matches_reference(self):
+        """Reproduce the exact aggregation used by the async pipeline.
+
+        The reference is computed from first principles in float64, which is
+        independent of the implementation's eps/mean/std choices.
+        """
+        torch.manual_seed(7)
+        for batch_size, group_size, num_funcs in itertools.product(
+            (1, 2, 4), (2, 4, 8), (1, 3)
+        ):
+            rewards = torch.randn(batch_size, group_size, num_funcs)
+            r = rewards.sum(-1).double()
+            expected = (r - r.mean(1, keepdim=True)) / (
+                r.std(1, keepdim=True) + 1e-4
+            )
+            got = group_normalized_advantages(rewards.sum(-1))
+            assert torch.allclose(got, expected.float(), atol=1e-6)
+            assert torch.allclose(got.mean(-1), torch.zeros(batch_size), atol=1e-6)
diff --git a/tests/torchtune/dev/rl/test_replay_buffer_contract.py b/tests/torchtune/dev/rl/test_replay_buffer_contract.py
new file mode 100644
index 0000000000..51e1e8fca5
--- /dev/null
+++ b/tests/torchtune/dev/rl/test_replay_buffer_contract.py
@@ -0,0 +1,203 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import pytest
+import torch
+from tensordict import NonTensorStack
+from torchrl.data import LazyStackStorage, ReplayBuffer
+
+from torchtune.dev.rl.datatypes import Trajectory
+from torchtune.dev.rl.rewards import group_normalized_advantages
+
+
+def _make_trajectory(
+    batch_size: int = 16,
+    seq_len: int = 64,
+    resp_len: int = 20,
+    num_funcs: int = 2,
+    seed: int = 0,
+) -> Trajectory:
+    """Build a trajectory shaped like the one PostProcessingWorker produces.
+
+    The advantages and reward metadata are per-sample and varied, so any loss of
+    per-sample information in the replay buffer is detectable by exact equality.
+    """
+    torch.manual_seed(seed)
+    advantages = torch.randn(batch_size) * 2.0 - 0.5
+    rewards = torch.randn(batch_size, num_funcs)
+    successes = torch.rand(batch_size, num_funcs) > 0.5
+    func_names = [f"reward_fn_{i}" for i in range(num_funcs)]
+    return Trajectory(
+        query_responses=torch.randint(0, 100, (batch_size, seq_len)),
+        responses=torch.randint(0, 100, (batch_size, resp_len)),
+        logprobs=torch.randn(batch_size, resp_len),
+        ref_logprobs=torch.randn(batch_size, resp_len),
+        query_response_padding_masks=torch.ones(batch_size, seq_len, dtype=torch.bool),
+        seq_lens=torch.randint(3, resp_len, (batch_size,)),
+        answers=NonTensorStack(*[f"ans{i}" for i in range(batch_size)]),
+        policy_version=7,
+        advantages=advantages,
+        rewards=rewards,
+        successes=successes,
+        reward_func_names=NonTensorStack(*[func_names for _ in range(batch_size)]),
+        batch_size=[batch_size],
+        sequence_ids=NonTensorStack(
+            *[f"worker0_{i}" for i in range(batch_size)]
+        ),
+    )
+
+
+class TestReplayBufferContract:
+    """The async GRPO recipe stores each whole trajectory batch in the replay
+    buffer as a single item and samples exactly one item per training step.
+
+    This contract prevents two failure modes observed in production:
+
+    1. Extending a batched trajectory directly lets the storage capacity count
+       in *rows*: with ``max_size < batch_size`` all but the last few samples
+       are silently overwritten, and sampling returns repeated copies of the
+       same rows -- i.e. identical advantages for every sample in the batch.
+       See https://github.com/meta-pytorch/torchtune/issues/2943.
+    2. Sampling more than one item returns repeated copies of the whole batch,
+       duplicating every sample num_samples times.
+    """
+
+    def _make_buffer(self, max_size: int = 2, batch_size: int = 1):
+        return ReplayBuffer(
+            storage=LazyStackStorage(max_size=max_size), batch_size=batch_size
+        )
+
+    def test_extending_batched_trajectory_directly_corrupts_advantages(self):
+        """Regression test for the reported bug: the old contract (extend a
+        batched trajectory, sample batch_size items) collapses the batch."""
+        batch_size, max_size = 16, 2
+        traj = _make_trajectory(batch_size=batch_size, seed=1)
+        buf = self._make_buffer(max_size=max_size, batch_size=batch_size)
+        buf.extend(traj)
+
+        sampled = buf.sample()
+        assert sampled.advantages.numel() == batch_size
+        # At most max_size distinct advantages can survive the round trip, and
+        # with max_size < batch_size the sampled advantages are repeated.
+        assert len(torch.unique(sampled.advantages)) <= max_size
+        # The surviving rows are the last max_size rows of the batch.
+        assert torch.allclose(
+            torch.unique(sampled.advantages),
+            torch.unique(traj.advantages[-max_size:]),
+        )
+
+    def test_sampling_more_than_one_item_repeats_the_batch(self):
+        batch_size = 8
+        traj = _make_trajectory(batch_size=batch_size, seed=2)
+        buf = self._make_buffer(max_size=2, batch_size=1)
+        buf.extend(traj.unsqueeze(0))
+
+        sampled = buf.sample(3)
+        assert tuple(sampled.batch_size) == (3, batch_size)
+        for i in range(3):
+            assert torch.equal(sampled.advantages[i], traj.advantages)
+
+    def test_round_trip_preserves_per_sample_fields(self):
+        batch_size = 16
+        traj = _make_trajectory(batch_size=batch_size, seed=3)
+        buf = self._make_buffer(max_size=2, batch_size=1)
+        buf.extend(traj.unsqueeze(0))
+
+        sampled = buf.sample(1).squeeze(0)
+        assert tuple(sampled.batch_size) == (batch_size,)
+
+        # Every per-sample tensor field must survive exactly.
+        for field in (
+            "query_responses",
+            "responses",
+            "logprobs",
+            "ref_logprobs",
+            "query_response_padding_masks",
+            "seq_lens",
+            "advantages",
+            "rewards",
+            "successes",
+        ):
+            assert torch.equal(getattr(sampled, field), getattr(traj, field)), field
+
+        # Non-tensor per-sample fields must be preserved and aligned with rows.
+        assert list(sampled.sequence_ids) == list(traj.sequence_ids)
+        assert list(sampled.answers) == list(traj.answers)
+        assert [list(f) for f in sampled.reward_func_names] == [
+            list(f) for f in traj.reward_func_names
+        ]
+        assert sampled.policy_version == traj.policy_version
+
+    def test_capacity_counts_batches_not_rows(self):
+        batch_size = 16
+        max_size = 3
+        buf = self._make_buffer(max_size=max_size, batch_size=1)
+        for i in range(5):
+            buf.extend(_make_trajectory(batch_size=batch_size, seed=10 + i).unsqueeze(0))
+        assert len(buf) == max_size
+
+        # The buffer holds whole batches: any sample is one coherent batch whose
+        # per-sample fields are internally consistent with its sequence ids.
+        for _ in range(20):
+            sampled = buf.sample(1).squeeze(0)
+            assert tuple(sampled.batch_size) == (batch_size,)
+            assert len(sampled.advantages) == batch_size
+            # Sequence ids within a sampled batch all come from one batch.
+            assert sampled.sequence_ids[0].split("_")[0] == sampled.sequence_ids[
+                -1
+            ].split("_")[0]
+
+    def test_sample_returns_different_batches(self):
+        batch_size, max_size = 8, 3
+        buf = self._make_buffer(max_size=max_size, batch_size=1)
+        for i in range(3):
+            buf.extend(_make_trajectory(batch_size=batch_size, seed=20 + i).unsqueeze(0))
+
+        # With max_size > 1, sampling repeatedly with replacement can return
+        # different batches, each internally varied (no degenerate duplicates).
+        seen = set()
+        for _ in range(100):
+            sampled = buf.sample(1).squeeze(0)
+            seen.add(tuple(sampled.advantages.tolist()))
+        assert len(seen) > 1
+        for adv in seen:
+            assert len(set(adv)) > 1
+
+
+class TestGroupNormalizedAdvantages:
+    def test_mean_and_std_are_normalized_per_group(self):
+        rewards = torch.tensor(
+            [
+                [10.0, 0.0, 5.0],
+                [8.0, 4.0, 6.0],
+            ]
+        )
+        adv = group_normalized_advantages(rewards)
+        # Within each group: mean 0 and unit variance (up to the eps term).
+        assert torch.allclose(adv.mean(-1), torch.zeros(2), atol=1e-6)
+        assert torch.allclose(adv.std(-1), torch.ones(2), atol=1e-3)
+
+    def test_matches_reference_implementation(self):
+        torch.manual_seed(0)
+        rewards = torch.randn(4, 8)
+        mean = rewards.mean(1, keepdim=True)
+        std = rewards.std(1, keepdim=True)
+        expected = (rewards - mean) / (std + 1e-4)
+        assert torch.allclose(group_normalized_advantages(rewards), expected)
+
+    def test_shape_is_preserved(self):
+        rewards = torch.randn(3, 5)
+        assert tuple(group_normalized_advantages(rewards).shape) == (3, 5)
+
+    def test_constant_group_is_finite(self):
+        rewards = torch.tensor([[1.0, 1.0, 1.0]])
+        adv = group_normalized_advantages(rewards)
+        assert torch.isfinite(adv).all()
+        assert torch.allclose(adv, torch.zeros_like(adv), atol=1e-5)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])
diff --git a/tests/torchtune/dev/rl/workers/test_postprocessing.py b/tests/torchtune/dev/rl/workers/test_postprocessing.py
index 7bbf8cb48e..2b40c0e0dc 100644
--- a/tests/torchtune/dev/rl/workers/test_postprocessing.py
+++ b/tests/torchtune/dev/rl/workers/test_postprocessing.py
@@ -141,7 +141,9 @@ def test_run(self, cfg, log_file):
                         sequence_ids=None,
                         policy_version=None,
                         advantages=None,
-                        reward_outputs=None,
+                        rewards=None,
+                        successes=None,
+                        reward_func_names=None,
                     )
                 )
             replay_buffer = []
diff --git a/torchtune/dev/rl/datatypes/trajectory.py b/torchtune/dev/rl/datatypes/trajectory.py
index 7f4e55aae7..3adfc25b48 100644
--- a/torchtune/dev/rl/datatypes/trajectory.py
+++ b/torchtune/dev/rl/datatypes/trajectory.py
@@ -5,8 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
+from tensordict import NonTensorStack
 from tensordict import TensorClass
-from torchtune.dev.rl.rewards import RewardOutput
 
 
 class Trajectory(TensorClass["nocast"]):
@@ -19,5 +19,7 @@ class Trajectory(TensorClass["nocast"]):
     answers: torch.Tensor
     policy_version: int
     advantages: torch.Tensor
-    reward_outputs: list[RewardOutput]
+    rewards: torch.Tensor
+    successes: torch.Tensor
+    reward_func_names: NonTensorStack
     sequence_ids: list[str]
diff --git a/torchtune/dev/rl/rewards.py b/torchtune/dev/rl/rewards.py
index 42dc40a74f..5c217f9dfc 100644
--- a/torchtune/dev/rl/rewards.py
+++ b/torchtune/dev/rl/rewards.py
@@ -267,6 +267,34 @@ def extract_tags(text: str) -> tuple[str, str]:
     return cot, potential_answer
 
 
+def group_normalized_advantages(
+    group_rewards: torch.Tensor, eps: float = 1e-4
+) -> torch.Tensor:
+    """Normalize per-group rewards into GRPO advantages.
+
+    Rewards are normalized within each group (second dimension), so that each
+    group has zero mean and unit variance. This is the advantage baseline used
+    by both the synchronous and asynchronous GRPO recipes.
+
+    Args:
+        group_rewards (torch.Tensor): rewards of shape ``[B, G]``, where ``G``
+            is the group size and ``B`` the batch size.
+        eps (float): small constant added to the standard deviation for numerical
+            stability.
+
+    Returns:
+        torch.Tensor: advantages of shape ``[B, G]``.
+
+    Example:
+        >>> group_rewards = torch.tensor([[10.0, 0.0, 5.0]])
+        >>> group_normalized_advantages(group_rewards)
+        tensor([[1.3363, -1.0690, -0.2673]])
+    """
+    mean = group_rewards.mean(1, keepdim=True)
+    std = group_rewards.std(1, keepdim=True)
+    return (group_rewards - mean) / (std + eps)
+
+
 def batched_rewards(
     tokenizer: Union[ModelTokenizer, HuggingFaceModelTokenizer],
     completions: torch.Tensor,
diff --git a/torchtune/dev/rl/workers/datacollectors/sync.py b/torchtune/dev/rl/workers/datacollectors/sync.py
index b9c73d7771..4142ca93f8 100644
--- a/torchtune/dev/rl/workers/datacollectors/sync.py
+++ b/torchtune/dev/rl/workers/datacollectors/sync.py
@@ -153,9 +153,9 @@ def _postprocess_for_queue(self, data):
         seq_lens = training.get_unmasked_sequence_lengths(response_padding_masks)
         del response_padding_masks
 
-        # Generate unique sequence IDs for the batch
-        # FIXME: it outputs a list[list[str]] when sampling from replay buffer, with shape num_samples X 16.
-        # It should have shape num_samplesX1, so we can log a single sequence_id per sequence.
+        # Generate unique sequence IDs for the batch. Each sequence in the batch
+        # gets a single ID; the replay buffer stores whole batches as one item,
+        # so IDs come back flat (one string per sample) when sampled out.
         batch_size = query_responses.shape[0]
         sequence_ids = NonTensorStack(
             *[
@@ -174,8 +174,10 @@ def _postprocess_for_queue(self, data):
             seq_lens=seq_lens,
             answers=answers,
             policy_version=policy_version,
-            reward_outputs=None,
             advantages=None,
+            rewards=None,
+            successes=None,
+            reward_func_names=None,
             sequence_ids=sequence_ids,
         )
 
diff --git a/torchtune/dev/rl/workers/postprocessing.py b/torchtune/dev/rl/workers/postprocessing.py
index 332094446b..2e467c1ad0 100644
--- a/torchtune/dev/rl/workers/postprocessing.py
+++ b/torchtune/dev/rl/workers/postprocessing.py
@@ -10,14 +10,33 @@
 import torch
 import torchtune.training as training
 from omegaconf import DictConfig
+from tensordict import NonTensorStack
 
 from torchtune import config, generation, rlhf, utils
 from torchtune.dev.rl.datatypes import Trajectory
-from torchtune.dev.rl.rewards import Reward, RewardOutput
+from torchtune.dev.rl.rewards import Reward, RewardOutput, group_normalized_advantages
 
 log = utils.get_logger("DEBUG")
 
 
+def reward_func_names_per_sample(
+    reward_outputs: list[RewardOutput], num_samples: int
+) -> NonTensorStack:
+    """Label every sample with the names of the reward functions that scored it.
+
+    Args:
+        reward_outputs (list[RewardOutput]): one ``RewardOutput`` per reward
+            function, in the same order as the columns of the per-sample rewards
+            tensor built from them.
+        num_samples (int): the number of samples in the batch.
+
+    Returns:
+        NonTensorStack: one list of reward function names per sample, shape ``[num_samples]``.
+    """
+    func_names = [reward_output.reward_base_name for reward_output in reward_outputs]
+    return NonTensorStack(*[func_names for _ in range(num_samples)])
+
+
 @ray.remote(num_cpus=8, num_gpus=1)
 class PostProcessingWorker:
     def __init__(self, *args, **kwargs):
@@ -31,7 +50,7 @@ def __init__(self, *args, **kwargs):
         self.cfg = kwargs.pop("cfg")
         self.rollout_queue = kwargs.pop("rollout_queue")
         self.replay_buffer = kwargs.pop("replay_buffer")
-        device_type = "cuda"
+        device_type = cfg.get("device", "cuda")
         self._device = utils.get_device(device=device_type)
         self._tokenizer = config.instantiate(self.cfg.tokenizer)
         self._dtype = training.get_dtype("bf16", device=self._device)
@@ -257,16 +276,19 @@ def run(self):
             for reward_fn in self.reward_functions:
                 reward_outputs.append(reward_fn(response_ids, responses_str, answers))
 
-            group_rewards = torch.stack(
+            # Per-sample reward metadata for the training worker, shape (B * G, num_funcs).
+            # Rewards are also aggregated below to compute advantages.
+            rewards_all = torch.stack(
                 [reward_output.total_reward for reward_output in reward_outputs], dim=-1
             )  # (B * G, num_funcs)
-            group_rewards = group_rewards.reshape(batch_size, group_size, -1)
+            successes_all = torch.stack(
+                [reward_output.successes for reward_output in reward_outputs], dim=-1
+            )  # (B * G, num_funcs)
+            group_rewards = rewards_all.reshape(batch_size, group_size, -1)
             # Compute advantages: B, G, num_funcs -> B, G
             group_rewards = group_rewards.sum(-1)
             # To compute advantage, subtract the mean of the group rewards from each group reward
-            group_advantages = (group_rewards - group_rewards.mean(1, keepdim=True)) / (
-                group_rewards.std(1, keepdim=True) + 1e-4
-            )  # (B, G)
+            group_advantages = group_normalized_advantages(group_rewards)  # (B, G)
             # Repack trajectory with policy_version
 
             trajectory = Trajectory(
@@ -279,17 +301,26 @@ def run(self):
                 answers=trajectory.answers,
                 policy_version=trajectory.policy_version,
                 advantages=group_advantages.reshape(batch_size * group_size),  # (B, G)
+                rewards=rewards_all,
+                successes=successes_all,
+                # The collector does not know which reward functions will score
+                # the rollout, so label the per-sample rewards here.
+                reward_func_names=reward_func_names_per_sample(
+                    reward_outputs, batch_size * group_size
+                ),
                 batch_size=batch_size * group_size,
                 sequence_ids=trajectory.sequence_ids,
-                reward_outputs=reward_outputs,
             )
 
             log.info(f"Constructed trajectory: {trajectory}")
             # Move tensors to CPU before putting into the queue
             trajectory = trajectory.cpu()
 
-            # Update circular queue
-            self.replay_buffer.extend(trajectory)
+            # Store the whole batch as a single replay buffer item: the storage
+            # counts each stored element as one item, so extending a batched
+            # trajectory directly would let its capacity in *rows* silently
+            # overwrite all but the last few samples of the batch.
+            self.replay_buffer.extend(trajectory.unsqueeze(0))
 
             # End of step timing
             time_total_ref_step = time.perf_counter() - time_step_start
diff --git a/torchtune/dev/rl/workers/trainers/training.py b/torchtune/dev/rl/workers/trainers/training.py
index 557ad38d6b..85da3c1c1b 100644
--- a/torchtune/dev/rl/workers/trainers/training.py
+++ b/torchtune/dev/rl/workers/trainers/training.py
@@ -58,7 +58,7 @@ def __init__(
         self.replay_buffer = replay_buffer
 
         # Device and dtype setup
-        device_type = "cuda"  # Harcoded for now
+        device_type = cfg.get("device", "cuda")
         self._device = utils.get_device(device=device_type)
         self._dtype = training.get_dtype("bf16", device=self._device)
 
@@ -91,6 +91,7 @@ def __init__(
 
         # Training configuration
         self._clip_grad_norm = cfg.training.get("clip_grad_norm", None)
+        self._cfg_batch_size = cfg.training.batch_size
 
         # Activation checkpointing and offloading
         self._enable_activation_checkpointing = cfg.training.get(
@@ -199,7 +200,7 @@ def init_model_update_group(self, master_address, master_port, rank, world_size)
             master_port,
             rank,
             world_size,
-            torch.device("cuda:0"),  # FIXME: Hardcoded device
+            torch.device(self._device),  # was hardcoded to cuda:0
         )
 
     def load_checkpoint(self, cfg_checkpointer: DictConfig) -> dict[str, Any]:
@@ -319,9 +320,9 @@ def grpo_step(
             GRPOStats: Instance of :class:`~torchtune.rlhf.GRPOStats`
         """
         # Create an output mask to avoid computing model.output on tokens we won't train
-        # FIXME: when bsz>1, don't we have multiple context_length?
-        # FIXME: because of chunked CE, the outout of pi_logits is a chunked list, so masking after the fact is
-        # more annoying. Masking before the chunking is easier, but we have to figure out masking for bsz>1
+        # context_length is a single value per batch: the collector pads every
+        # response to the same length (pad_output=True), so the context is
+        # uniform across all samples in the batch.
         output_mask = torch.zeros_like(
             trajectory.query_responses, dtype=torch.bool, device=self._device
         )
@@ -543,7 +544,6 @@ def _log_table(data: list, table_name: str) -> None:
 
         # Iterate over each sample
         for idx in range(num_samples):
-            func_names = metadata["reward_metadata"][idx]["func_names"]
             sequence_id = metadata["sequence_ids"][idx]
             seq_len = grpo_trajectory.seq_lens[idx].item()
 
@@ -570,10 +570,23 @@ def _log_table(data: list, table_name: str) -> None:
             per_sample_dict["prompt"] = prompt
             per_sample_dict["response"] = response
             per_sample_dict["answers"] = grpo_trajectory.answers[idx]
-            per_sample_dict["policy_version"] = metadata["policy_version"][idx]
+            # policy_version is a single int shared by every sample in the batch
+            policy_version = metadata["policy_version"]
+            per_sample_dict["policy_version"] = (
+                policy_version[idx]
+                if isinstance(policy_version, list)
+                else policy_version
+            )
 
-            for reward_output in metadata["reward_outputs"][idx]:
-                per_sample_dict.update(reward_output.log(prefix="rewards"))
+            for reward_func_idx, reward_func_name in enumerate(
+                metadata["reward_func_names"][idx]
+            ):
+                per_sample_dict[f"rewards/{reward_func_name}"] = metadata["rewards"][
+                    idx, reward_func_idx
+                ].item()
+                per_sample_dict[
+                    f"rewards/{reward_func_name}/successes"
+                ] = metadata["successes"][idx, reward_func_idx].item()
 
             # Add GRPO statistics, handling per-sample vs. scalar cases
             # TODO: currently has one scalar per batch. We should enable a scalar per sentence.
@@ -694,7 +707,16 @@ def train(self):
                 log.info("waiting for replay buffer")
                 time.sleep(1.0)
 
-            trajectory = self.replay_buffer.sample().to(self._device)
+            # The replay buffer stores each trajectory batch as a single item
+            # (see PostProcessingWorker.run), so sampling one item yields one
+            # [batch_size, T] batch. Sampling more than one item would return
+            # repeated copies of the whole batch.
+            trajectory = self.replay_buffer.sample(1).squeeze(0).to(self._device)
+            batch_size = trajectory.batch_size[0]
+            if batch_size != self._cfg_batch_size:
+                log.warning(
+                    f"sampled batch has {batch_size} samples, expected {self._cfg_batch_size}"
+                )
             time_waiting_buffer = time.perf_counter() - time_waiting_buffer_start
             if self._is_rank_zero:
                 log.info(f"{self.rank=} got from queue traj {trajectory}")
@@ -913,8 +935,10 @@ def _prepare_trajectory(
             "avg_policy_age": avg_policy_age,
             "sequence_ids": raw_trajectory.sequence_ids,
             "policy_version": raw_trajectory.policy_version,
-            "reward_outputs": raw_trajectory.reward_outputs,
             "query_response_padding_masks": raw_trajectory.query_response_padding_masks,
+            "rewards": raw_trajectory.rewards,
+            "successes": raw_trajectory.successes,
+            "reward_func_names": raw_trajectory.reward_func_names,
         }
 
         return prepared_trajectory, context_length, metadata