[rl] Fetch generator weights from the StorageVolume into GPU by d4l3k · Pull Request #5652 · pytorch/torchtitan

d4l3k · 2026-10-28T18:01:01Z

Root issue: RL generators prefetch each policy from the StorageVolume into a full pinned CPU copy of the weights before loading it into vLLM. GB300 generator hosts have little CPU memory and OOM on this copy.

Fix: one load path, VLLMModelWrapper.load_model_state_dict(fetch), walks the FSDP groups one at a time (reshard, restore sharded params, fetch the group's subset in place, unshard and free; buffers last). Every parameter must be filled. Generator.Config.model_state_dict_load_mode picks the fetch:

`model_state_dict_load_mode`	prefetch	fetch during load	generator host memory
`"cpu_prefetch"` (default, previous behavior)	StorageVolume -> pinned CPU copy, overlaps generation	`copy_` from the pinned copy	+1 full weight copy
`"gpu_fetch"`	no-op	TorchStore read straight into the GPU weights	none

A gloo all-reduce agrees on load failure across generator ranks before rank 0 replies; on a fetch failure every rank still unshards every group (avoids an all-gather hang) and then stops the engine loop. gpu_fetch tradeoff: the transfer runs inside the load while the engine is paused rather than overlapping generation. gpu_fetch rejects quantized parameters at init. Parameters whose dtype differs from the model dtype (float32 GDN A_log/dt_bias, #5691) are fetched into model-dtype staging tensors and cast, since the trainer sends the model dtype and torchstore requires matching dtypes. The trainer-side duplicate copy is tracked in #5654.

Requires meta-pytorch/torchstore#252 (merged).

Test plan:

CPU: pytest tests/rl/unit_tests/cpu with no GPU: all pass (gpu_fetch load path covered by test_vllm_wrapper/test_generator, incl. float32 param staging).
GPU (GB200): torchrun --nproc-per-node=2 -m pytest tests/rl/unit_tests/gpu/test_bitwise_parity.py::TestBitwiseParityQwen35DebugVarlen: 5 passed on both ranks. Trainer vs vLLM prefill and decode are bitwise identical, including a reload through load_model_state_dict after CUDA graph replay.
E2E async RL on GB200 through TorchStore (before merging main):
- Qwen3 debugmodel, gen TP2, 4 steps: logprob_diff/max 0.0099-0.012 in both modes.
- Qwen3 8B, gen TP1, 5 syncs: logprob_diff/max 0.21-0.31 in both modes (same as before this PR; untrained weights, not batch invariant).

Generator memory, Qwen3 8B, gen TP1:

mode	host RSS	GPU growth in sync	sync time
`cpu_prefetch`	24.7 GiB (16 GiB pinned copy)	+1 GiB on first sync only	0.48 s, overlapped with generation
`gpu_fetch`	3.8-3.9 GiB	0	1.1-1.2 s, engine paused

Lines changed: impl +306/-88, tests +436/-100, CI 0.

tianyu-l · 2026-10-29T23:54:20Z

@claude review

claude · 2026-10-29T23:54:41Z

Claude finished @tianyu-l's task in 4m 36s —— View job

Needs discussion: the design is sound, but the flag only works once meta-pytorch/torchstore#252 lands (torchstore isn't pinned), and #5484/#5585 change the same weight-fetch path. Full review: #5652 (review)

claude

Verdict: Needs discussion -- The problem (GB300 generator hosts OOM on the pinned CPU copy) is real and the design is sound. But the flag is broken until torchstore#252 lands, torchstore is unpinned, and two other open PRs add knobs to the same weight-fetch path. Maintainers should settle the merge order and the overall design.

What this PR does. It adds Generator.Config.fetch_model_state_dict_to_gpu. When it is on, the generator skips the pinned CPU copy, and prefetch_model_state_dict does nothing. The load then calls ts.get_state_dict(strict=False) once per FSDP group, writing straight into that group's reallocated sharded GPU params, then unshards the group. After the load, a gloo all-reduce agrees on failure across ranks before rank 0 replies.

Example. 8B, TP1, numbers from the PR description:

staged (default):  prefetch -> 16 GiB pinned CPU copy (RSS 24.7 GiB); load = H2D copy, 0.52 s/sync
direct (new flag): prefetch no-op (RSS ~3.8 GiB); load = per-group RDMA into GPU params, 1.2-2.3 s/sync
                   with hot_swap=True, in-flight requests stall for those 1.2-2.3 s

Must fix

Merge order. CI installs torchstore from git+...torchstore.git@main (.github/scripts/validate_release_gpu.sh:119); there is no pin. With torchstore before #252, the direct load raises on the second group, and CI cannot detect that (per the description). Hold this PR until meta-pytorch/torchstore#252 has merged.
torchtitan/rl/generator/vllm/model_wrapper.py:279 -- a new private FSDP _to_sharded() call, without a TODO.
torchtitan/rl/generator/generator.py:498 -- torch.cuda.current_stream() in engine-agnostic code; use torch.accelerator.
torchtitan/rl/generator/vllm/model_wrapper.py:407, engine_loop.py:143, registry.py:135 -- internal defaults duplicate the config default; make the arguments required.
tests/rl/unit_tests/cpu/test_generator.py:638 (and L547, L614, L835, test_vllm_wrapper.py:483) -- these tests check plumbing or repeat other tests; delete them (about 120 lines).

Should consider

torchtitan/rl/generator/vllm/model_wrapper.py:562 -- express the staged load_state_dict as load_state_dict_direct with a copy_ fetch, so there is one per-group load sequence instead of two.
torchtitan/rl/generator/engine.py:97-109 -- make load_model_state_dict_direct abstract instead of raising at the first sync.
tests/rl/unit_tests/gpu/test_bitwise_parity.py:932 -- skip the inherited batch-invariance tests in the Direct subclass to save GPU CI time.

Questions

Overlap. #5484 (felipemello1) rewrites the same pinned-buffer block (model_wrapper.py:489-504). On the staged path it keeps attention scales and makes unpublished keys fail. #5585 (wwwjn) adds Generator.Config.fetch_weights_per_module, also per-module gets on the same fetch. Merged together, the three PRs give three independent knobs on one transfer, and system.py:133 already has TODO: move to a model state dict transport. Should these become one transport config, chosen by config type, where staged and direct are siblings instead of a bool checked in three places (Generator prefetch, EngineLoop, VLLMModelWrapper.__init__)? And which PR lands first?
Is the staged path meant to remain long term? Once trainer-side staging goes away (#5659 / #5654), is there a plan to overlap the direct fetch with generation, so that hot_swap doesn't stall?

Details

Failure handling: an exception in _load_model_state_dict propagates out of the main loop's try (engine_loop.py:383-389) on every rank after the all-reduce, so all ranks fail together and rank 0 fails the pending load replies. A failed fetch still runs _reshard/_restore/_unshard for each later group, so the other ranks' all-gathers complete.
Tied weights: the trainer's state_dict() publishes both tok_embeddings.weight and lm_head.weight, so the per-group missing check passes. The cost is a double fetch, as the description says.
Smaller alternative for the flag plumbing: the engine already knows the mode, since prefetched_model_state_dict is None. Generator.prefetch_model_state_dict could check that instead of re-reading the config. A transport config (question above) would remove the bool entirely.

Not verified: the GPU parity and end-to-end memory numbers (needs GB200 and torchstore#252); copy_ into _ShardedFSDPTensor for the unification suggestion.

claude · 2026-10-29T23:58:32Z

+    for param_group in module._get_fsdp_state()._fsdp_param_groups:
+        param_group._to_sharded()


This adds a new call to the private FSDP API FSDPParamGroup._to_sharded(), and correctness depends on it: it's what refreshes the tied group left in FORWARD. The neighboring helpers mark their private calls with TODO: replace this with FSDPModule._restore_sharded_params(). This one has no TODO, so an upstream change to _to_sharded (or to the training-state check in reshard()) would break the direct load without any marker in the code. Add a TODO, with an upstream issue link if one exists, for a public "reshard unconditionally" API:

Suggested change

for param_group in module._get_fsdp_state()._fsdp_param_groups:

param_group._to_sharded()

# TODO: replace with a public FSDP API that reshards regardless of training state.

for param_group in module._get_fsdp_state()._fsdp_param_groups:

param_group._to_sharded()

claude · 2026-10-29T23:58:35Z

+        """
+        # The read writes outside the current stream; wait for queued work that
+        # still uses state_dict's storage, e.g. its reallocation.
+        torch.cuda.current_stream().synchronize()


generator.py is the engine-agnostic layer, and it already uses torch.accelerator (L264-265). torch.cuda.current_stream() is the only CUDA-specific call in it, against the hardware-agnostic rule. Suggested fix:

Suggested change

torch.cuda.current_stream().synchronize()

torch.accelerator.current_stream().synchronize()

The CPU tests that patch torch.cuda.current_stream would then patch torch.accelerator.current_stream.

claude · 2026-10-29T23:58:38Z

        vllm_config: VllmConfig,
        prefix: str = "",
        dist_moe_runtime: DistMoeRuntime.Config | None = None,
+        fetch_model_state_dict_to_gpu: bool = False,


This default duplicates Generator.Config.fetch_model_state_dict_to_gpu's default, which the rules forbid for internal functions. The only caller is registry.py, which always passes the value, and the signature is keyword-only, so the field can be required. The same applies to EngineLoop(fetch_model_state_dict=... | None = None) (engine_loop.py:142-143): its one production caller and _bare_loop both pass it explicitly. register_to_vllm (registry.py:135) also has callers that rely on the default (generate.py, several GPU tests); make it required there too and update those callers.

Suggested change

fetch_model_state_dict_to_gpu: bool = False,

fetch_model_state_dict_to_gpu: bool,

claude · 2026-10-29T23:58:41Z

+    def load_model_state_dict_direct(
+        self, fetch: Callable[[dict[str, torch.Tensor]], Iterable[str]]
+    ) -> None:
+        """Load a model state dict by having ``fetch`` fill the model's weights in place, and
+        reset the KV cache when configured. Used with ``fetch_model_state_dict_to_gpu``.
+
+        Args:
+            fetch: Fills each tensor of the given subset of the model state dict in place
+                with the new weights and returns the keys it filled.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support fetch_model_state_dict_to_gpu"
+        )


Should consider: VLLMEngine is the only GenerationEngine subclass. With a default that raises NotImplementedError, an engine that doesn't support this fails at the first weight sync, after the engine is built and CUDA graphs are captured. That is a runtime failure, not a config-time one. Make it @abc.abstractmethod like load_prefetched_model_state_dict. If a future engine can't support it, it should reject fetch_model_state_dict_to_gpu=True at construction.

claude · 2026-10-29T23:58:46Z

+
+
+@pytest.mark.parametrize("fetch_to_gpu", [False, True])
+def test_fetch_model_state_dict_to_gpu_reaches_engine_and_engine_loop(


Several of the new tests only check that a bool reaches the next layer, or repeat assertions made elsewhere. If the plumbing broke, TestBitwiseParityQwen35DebugVarlenDirect::test_no_pinned_cpu_copy and the direct-load engine-loop tests would already fail. Suggest deleting (about 120 lines):

this test: heavily mocked Generator.__init__, only checks that kwargs are forwarded;

test_fetch_model_state_dict_to_gpu_reaches_model_registration (L835): kwarg forwarding;

test_load_model_state_dict_direct_fetches_into_the_model (L547): a one-line forward plus _maybe_reset_kv_cache, which the existing reset test already covers;

test_prefetch_model_state_dict_is_noop_with_fetch_to_gpu (L614): test_direct_load_reads_torchstore_on_the_actor_loop_during_the_load already asserts reads == [] after prefetch;

test_direct_load_groups_rejects_unsharded_params (test_vllm_wrapper.py:483): tests an internal assert.

claude · 2026-10-29T23:58:49Z

        self._check_vllm_prefill_and_decode_batch_invariance()


+class TestBitwiseParityQwen35DebugVarlenDirect(TestBitwiseParityQwen35DebugVarlen):


Should consider: this subclass inherits all four parent tests, so CI runs a second setUpClass (trainer build, vLLM build, graph capture) plus test_batch_invariance and test_vllm_prefill_and_decode_batch_invariance a second time. Neither batch-invariance test depends on how the weights were loaded. The tests that exercise the direct path are test_trainer_vs_vllm_prefill, test_vllm_decode_vs_prefill, test_reload_after_generate and test_no_pinned_cpu_copy. Skipping the two batch-invariance tests here (override them with unittest.skip) cuts this job's added GPU time.

claude · 2026-10-29T23:58:54Z

+    def load_state_dict_direct(
+        self, fetch: Callable[[dict[str, torch.Tensor]], Iterable[str]]
+    ) -> None:
+        """Fetch new weights straight into `self.model`, one FSDP group at a time.
+
+        Follows the same allocate -> copy -> unshard and free sequence per group as
+        `load_state_dict`, but `fetch` writes the group's sharded parameters in place
+        (e.g. over RDMA), so no full CPU copy of the weights is needed. Requires
+        `fetch_model_state_dict_to_gpu=True` at construction.
+
+        Args:
+            fetch: Fills each tensor of the given subset of `self.model.state_dict()`
+                in place with the new weights and returns the keys it filled. Each
+                group's sharded storage is reallocated just before its fetch, so a
+                fetch writing outside the current stream (e.g. RDMA) must first wait
+                for queued device work, and must finish writing before it returns.
+        """
+        error = None
+        for module, keyed_params in self._direct_load_groups:


Should consider: there are now two implementations of the same per-FSDP-group allocate -> fill -> unshard/free sequence. One is hook-based (load_state_dict, L512-560); this one is an explicit loop that also does _reshard. The staged path could become this loop with a copy fetch, e.g. self.load_state_dict_direct(lambda sub: [t.copy_(state_dict[k]) or k for k, t in sub.items()]), which is what _sync_trainer_weights_to_vllm already does in the GPU test. That deletes the hook registration (about 25 lines), computes _direct_load_groups unconditionally, and gives both paths one load order and one tied-group handling. The quantized rejection would then move from _direct_load_groups into the TorchStore fetch, since a copy_ fetch can write a _ShardedFSDPTensor the way nn.Module.load_state_dict does (unverified). This could be a follow-up.

If possible, I wonder if we could unify the two paths more

As claude suggests, can we share more code so that it's easier to read and maintain? We can let agent to try its best, as long as constraints mentioned in [rl] FSDP generator: sync weights one FSDP module at a time #5268 are not violated.

For the code paths that definitely couldn't be shared, can we make the calling convention more consistent. I'll comment inline where I think it could be improved.

brandondfeng

loosk good for default-off flag

tianyu-l · 2026-11-01T05:36:23Z

+    def load_state_dict_direct(
+        self, fetch: Callable[[dict[str, torch.Tensor]], Iterable[str]]
+    ) -> None:
+        """Fetch new weights straight into `self.model`, one FSDP group at a time.
+
+        Follows the same allocate -> copy -> unshard and free sequence per group as
+        `load_state_dict`, but `fetch` writes the group's sharded parameters in place
+        (e.g. over RDMA), so no full CPU copy of the weights is needed. Requires
+        `fetch_model_state_dict_to_gpu=True` at construction.
+
+        Args:
+            fetch: Fills each tensor of the given subset of `self.model.state_dict()`
+                in place with the new weights and returns the keys it filled. Each
+                group's sharded storage is reallocated just before its fetch, so a
+                fetch writing outside the current stream (e.g. RDMA) must first wait
+                for queued device work, and must finish writing before it returns.
+        """
+        error = None
+        for module, keyed_params in self._direct_load_groups:


If possible, I wonder if we could unify the two paths more

As claude suggests, can we share more code so that it's easier to read and maintain? We can let agent to try its best, as long as constraints mentioned in [rl] FSDP generator: sync weights one FSDP module at a time #5268 are not violated.

For the code paths that definitely couldn't be shared, can we make the calling convention more consistent. I'll comment inline where I think it could be improved.

tianyu-l · 2026-11-01T05:37:44Z

            handle.remove()
        return incompatible_keys

+    def load_state_dict_direct(


some times we use "load ... direct"; some times we use "fetch"; some times we use "fetch .. to_gpu"

tianyu-l · 2026-11-01T05:38:48Z

+        group = sharded_param_to_group.get(id(value))
+        if group is None:
+            # A parameter here would be loaded outside the reshard/unshard sequence.
+            assert id(value) in buffers, f"{key} is neither FSDP-sharded nor a buffer"


since other_keys is always buffers, can we rename to reflect that for better readability?

tianyu-l · 2026-11-01T05:40:31Z

+        if self._fetch_model_state_dict is None:
+            self._engine.load_prefetched_model_state_dict()
+        else:
+            await self._load_model_state_dict_direct(self._fetch_model_state_dict)


one is load via engine API, the other is self._load_model_state_dict_direct which calls engine API; but yeah it seems this one is more legit because we need to do gloo collective and error out

tianyu-l · 2026-11-01T05:42:45Z

+        # TODO: unify with the pinned CPU prefetch into one load path (e.g.
+        # direct fetch overlapped with generation) and drop this flag.


curious what is this unification about -- introducing chunking in CPU prefetching?

tianyu-l · 2026-11-01T05:44:52Z


+        # TODO: unify with the pinned CPU prefetch into one load path (e.g.
+        # direct fetch overlapped with generation) and drop this flag.
+        fetch_model_state_dict_to_gpu: bool = False


if this flag is going to be dropped, then it's fine to not over optimize the readability; o/w I would prefer treating both CPU prefetch + load and GPU direct load as first class citizens and have symmetric code paths than hardcoding "fetch...to_gpu=False" in the config. Instead it should be "model_state_dict_load_mode" being CPU prefetch+load vs. GPU pull

tianyu-l · 2026-11-01T05:47:43Z

@claude review latest code and review my review

claude · 2026-11-01T05:48:03Z

Claude finished @tianyu-l's task in 3m 16s —— View job

Request changes: the earlier automated must-fix items are resolved, but tianyu-l's Nov 1 comments are still open. Most of them (unifying the two load paths, consistent naming, a first-class config) are correct, and one change resolves them all together. Full review: #5652 (review)

claude

Verdict: Request changes -- The must-fix items from the earlier automated review are resolved at HEAD. tianyu-l's Nov 1 comments are still open, because the last commit (9f8a102) came before them. The main one, unifying the two load paths, changes the engine API and the config, so the PR needs another look after it.

What this PR does. With Generator.Config.fetch_model_state_dict_to_gpu=True, the generator skips the pinned CPU copy. Each load reads the StorageVolume snapshot over TorchStore straight into the engine's GPU weights, one FSDP group at a time. A gloo all-reduce makes every generator rank agree on failure before rank 0 replies.

Example. 8B, TP1, numbers from the description:

staged: prefetch overlaps generation, 16 GiB pinned copy (RSS 24.7 GiB); load 0.52 s
direct: prefetch is a no-op (RSS ~3.8 GiB); load 1.2-2.3 s while the engine is paused

Earlier automated review, resolved: torchstore#252 merged; _reshard TODO (model_wrapper.py:314); torch.accelerator stream (generator.py:500); required args (registry.py:199, model_wrapper.py:444); abstract load_model_state_dict_direct; tests trimmed; batch-invariance tests skipped (test_bitwise_parity.py:937-944).

tianyu-l's review

model_wrapper.py:617, unify the paths -- Correct, and it fits within #5268's constraints. The staged load is the direct loop with a fetch that copies from the pinned buffers. That is the same copy_ that nn.Module.load_state_dict does, and it keeps the per-group allocate -> fill -> unshard/free order and the memory peak (see Details). The hook-based load_state_dict goes away. The engine then needs only one load_model_state_dict(fetch) method, and EngineLoop no longer branches.
model_wrapper.py:599, naming -- Correct. One mode has four names: fetch_model_state_dict_to_gpu, load_*_direct, _direct_load_groups, and fetch_model_state_dict (which is really the transport callable). After (1), only the config value names the mode. fetch stays the callable's name, and _direct_load_groups becomes e.g. _fsdp_load_groups.
model_wrapper.py:401, other_keys -- Correct. The assert on that line guarantees they are buffers. Rename to buffer_keys / _direct_load_buffer_keys, and update the docstring.
engine_loop.py:505 -- Partly correct. Only the direct path needs the all-reduce before replying, since other ranks may still be reading the snapshot. The staged path has a related gap, though: if one rank's load raises, that rank stops alone, rank 0 replies success, and the next decision broadcast hangs. After (1), run the all-reduce on both paths (it reduces a single integer over gloo) and inline _load_model_state_dict_direct, which has one caller.
generator.py:162, what is the unification? -- It isn't chunked CPU prefetch, and the TODO gives no concrete plan. To overlap a direct fetch with generation, the new weights need somewhere to live while CUDA graphs still read the old ones. That means keeping the sharded buffers allocated: 1/FSDP-degree of the weights in GPU memory, taken from the KV cache, and a full second copy at FSDP 1 (#5268's OOM case). Chunked CPU prefetch would save host memory, but the transfer would still happen during the pause. The author should say what the TODO means, or delete it.
generator.py:163, first-class paths -- Correct. Neither path is better on both axes: 0.52 s overlapped vs 1.2-2.3 s paused, and 24.7 vs 3.8 GiB RSS. Both will stay, so drop "and drop this flag". Use a Literal, e.g. model_state_dict_transfer: Literal["cpu_prefetch", "gpu_direct"] = "cpu_prefetch". If a variant later gets its own fields (e.g. #5585's per-module requests), switch to sibling configs selected by type.

Questions

Merge order with #5484 and #5585. Both are still open and touch the same pinned buffers and fetch.

Details: unification sketch

# VLLMModelWrapper.__init__: always compute the groups; pinned buffers only for cpu_prefetch
self._fsdp_load_groups, self._buffer_keys = _fsdp_load_groups(self.model)

def load_state_dict(self, fetch): ...  # today's load_state_dict_direct body

# Generator, cpu_prefetch: the fetch reads the pinned copy
def _copy_from_prefetched(state_dict):
    with torch.no_grad():
        for key, tensor in state_dict.items():
            tensor.copy_(prefetched[key])
    return state_dict.keys()

_reshard replaces FSDP's load pre-hook. The tied group is filled in one fetch, so its buffers are no longer allocated across the whole load.
test_bitwise_parity.py:302-310 already loads with a copy_ fetch, so the staged parity class and the direct one would differ only in the fetch.
Move the quantized rejection out of the grouping and into the TorchStore path: a copy_ into a quantized sharded param should go through the tensor subclass as it does today (unverified; test_generator_dist_moe_mxfp8.py covers it).
Deletes about 25 lines of hook code, load_prefetched_model_state_dict (abstract and vLLM implementation), and the EngineLoop branch. The two CPU wrapper tests (test_load_state_dict_allocates_one_fsdp_group_at_a_time and the direct ones) can merge.

Not verified: the GPU parity and E2E numbers. CI checks the direct path only with a copy_ fetch, so the real TorchStore read into GPU tensors is exercised only by the author's E2E run.

tianyu-l · 2026-11-01T06:19:55Z

@claude review again

claude · 2026-11-01T06:20:17Z

Claude finished @tianyu-l's task in 3m 36s —— View job

Approve -- please address the items below before merging: all open review comments are resolved by the unification, but the reload-after-generate GPU test should cover the default path, and the E2E numbers need a re-run at HEAD. Full review: #5652 (review)

claude

Verdict: Approve -- please address the items below before merging. The unification in 926cd74/795aa03 resolves every open comment from tianyu-l and from the earlier automated review. What remains is a test move and refreshed evidence for the default path, which this PR now also changes. A maintainer still needs to review, because this changes how the default RL weight sync loads weights.

What this PR does. Generator.Config.model_state_dict_load_mode (Literal["cpu_prefetch", "gpu_fetch"]) chooses the fetch callable that EngineLoop passes to the single engine method, load_model_state_dict(fetch). VLLMModelWrapper.load_model_state_dict then goes through the FSDP groups one at a time: reshard, reallocate, fetch, unshard and free, with buffers last. A gloo all-reduce makes every rank agree on failure before rank 0 replies, in both modes.

Example. The default path, before vs after:

before: model.load_state_dict(pinned, strict=True) + FSDP pre/post hooks per group
after:  for each group: _reshard -> _restore -> tensor.copy_(pinned[key]) -> unshard/free
        then buffers; all-reduce(failed) over gloo; bump version
gpu_fetch: same loop, but fetch = ts.get_state_dict(user_state_dict=group, strict=False)

Resolved since the last review:

Paths unified: one load_model_state_dict(fetch). load_prefetched_model_state_dict, the hook code, and the EngineLoop branch are gone.
Naming: the mode lives only in model_state_dict_load_mode, and _direct_load_groups became _fsdp_load_groups.
other_keys became buffer_keys.
The all-reduce runs on both paths, and _load_model_state_dict_direct is inlined.
The "drop this flag" TODO is replaced by a Literal with documented trade-offs.
The quantized rejection moved to the gpu_fetch init. The cpu_prefetch copy_ into a _ShardedFSDPTensor is covered by test_generator_dist_moe_mxfp8.py.

Must fix

tests/rl/unit_tests/gpu/test_bitwise_parity.py:953: move test_reload_after_generate into the default-path class, then drop the GPUFetch class and its CI step (see the inline comment).
PR description: the E2E and memory numbers come from the pre-unification revision. 926cd74 changes the default cpu_prefetch load (a new loop, _reshard, an all-reduce per load). Before merging, re-run the Qwen3.5 async RL E2E at HEAD with cpu_prefetch, and with gpu_fetch if possible. Report the logprob_diff and the s/sync for each.

Should consider

Literal["cpu_prefetch", "gpu_fetch"] is spelled out in five production signatures (generator.py, vllm/engine.py x2, model_wrapper.py, registry.py). One alias next to GenerationEngine would keep them in sync.

Questions

#5484 and #5585 are still open and touch the same fetch. If #5585's per-module request options apply to only one mode, should the Literal become sibling configs selected by type, as in the earlier review?

Details

Default-path equivalence: nn.Module.load_state_dict did param.copy_(input) under no_grad into the registered param. After _reshard, that param is fsdp_param.sharded_param, which is the tensor the new fetch writes. Buffers come from the same model.state_dict() key set as the pinned dict, so the result should be bitwise identical. The tied group is now filled in one fetch instead of staying allocated for the whole load, so peak memory is the same or lower.
_fsdp_load_groups now runs at init in both modes. Its "neither FSDP-sharded nor a buffer" assert can't fire, because the root is asserted to be an FSDPModule (model_wrapper.py:478).
Bypassing nn.Module.load_state_dict skips the NVFP4 _load_from_state_dict overrides (quantization/nvfp4/linear.py:455, experts.py:57). That is harmless: they only refresh a tuple built from the non-persistent, constant _rht_sign_vector.
The new all-reduce adds one integer gloo collective per load, with .item() on a CPU tensor and no device sync.

Not verified: GPU parity and E2E numbers at HEAD.

claude · 2026-11-01T06:23:34Z

+        wrapper = self.engine.model_executor.driver_worker.get_model()
+        self.assertIsNone(wrapper._prefetched_model_state_dict)
+
+    def test_reload_after_generate(self):


After 926cd74, both modes run the same VLLMModelWrapper.load_model_state_dict loop, including _reshard for the tied group that compute_logits leaves in FORWARD. The only GPU test that reloads after generation is this one, and it only runs for gpu_fetch, so the default cpu_prefetch path has no GPU coverage for that case.

At the same time, this class's fetch is a test-only copy_ from the trainer, so it doesn't exercise the real TorchStore read. Apart from this test, it differs from the parent only in test_no_pinned_cpu_copy, and that just re-checks the if model_state_dict_load_mode == "cpu_prefetch" branch. The cost is a second setUpClass in GPU CI (trainer build, vLLM build, graph capture).

Suggested fix: move test_reload_after_generate into TestBitwiseParityQwen35DebugVarlen so it covers the default path. Then delete TestBitwiseParityQwen35DebugVarlenGPUFetch, its CI step in unit_test_gpu_rl.yaml, and the model_state_dict_load_mode class attribute and prefetched is None branch in _sync_trainer_weights_to_vllm.

…generator GPU

…ert direct-load keys

…to the load

…re all_reduce; add direct-load GPU parity test and 2-rank fetch-error test

…wiring and 2-rank direct load values

…ter all collectives

… vLLM attention scales)

…ad, trim tests

…plicate direct-load tests, unification TODO

…lias

…t-gpu # Conflicts: # torchtitan/rl/generator/vllm/model_wrapper.py

pytorch-bot Bot added the ciflow/rl label Oct 28, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 28, 2026

d4l3k mentioned this pull request Oct 28, 2026

Routed per-subset state dict gets into NIXL CUDA destinations meta-pytorch/torchstore#252

Merged

d4l3k changed the title ~~[rl] Direct RDMA weight load from trainer CPU into generator GPU~~ [rl] Fetch generator weights from the StorageVolume into GPU Oct 28, 2026

d4l3k marked this pull request as ready for review October 29, 2026 00:31

d4l3k requested a review from felipemello1 as a code owner October 29, 2026 00:31

pytorch-bot Bot added the ciflow/fake-pg Run 1-GPU Fake PG integration tests label Oct 29, 2026

d4l3k requested review from fegin, pzhan9, tianyu-l, wconstab and wwwjn as code owners October 29, 2026 00:31

d4l3k requested review from acisseJZhong and shuhuayu as code owners October 29, 2026 17:14

claude Bot reviewed Oct 29, 2026

View reviewed changes

d4l3k force-pushed the d4l3k/titan-gen-direct-gpu branch from 9112f70 to 0f77270 Compare October 30, 2026 00:18

pytorch-bot Bot added the ciflow/b200 Trigger B200 CI label Oct 30, 2026

brandondfeng approved these changes Oct 30, 2026

View reviewed changes

tianyu-l reviewed Nov 1, 2026

View reviewed changes

claude Bot reviewed Nov 1, 2026

View reviewed changes

claude Bot approved these changes Nov 1, 2026

View reviewed changes

tianyu-l approved these changes Nov 1, 2026

View reviewed changes

d4l3k added 2 commits October 31, 2026 23:36

[rl] Add direct_rdma_model_state_dict: read trainer CPU weights into …

2823775

…generator GPU

Address review: docs, RuntimeError for missing torchstore client, ass…

7d1a402

…ert direct-load keys

d4l3k added 18 commits October 31, 2026 23:36

Reshard tied FSDP group unconditionally before direct load

91fe4b1

Reject quantized params with direct RDMA loads

c6c500b

Test direct RDMA load wiring

d986a50

Wait for every generator rank to finish direct reads before replying …

20309c2

…to the load

Keep unsharding every FSDP group after a direct fetch fails

f845598

Wait for direct reads on cancel; cancelled ranks still join the failu…

8146de9

…re all_reduce; add direct-load GPU parity test and 2-rank fetch-error test

Keep waiting for the direct read through repeated cancels; test flag …

35c2a9c

…wiring and 2-rank direct load values

Test a direct reload after generation replays CUDA graphs

9d7855d

Reload test: reset prefix cache, check decode after reload, assert af…

237574e

…ter all collectives

Fetch generator weights from the StorageVolume into GPU

e93d655

Direct loads: require every parameter, keep unpublished buffers (e.g.…

ec60b5e

… vLLM attention scales)

Direct loads: test subset fetch, document hot_swap stall

616fe22

Address review: required args, accelerator stream, abstract direct lo…

2988557

…ad, trim tests

Address review gate: accelerator stream in engine loop test, merge du…

574ac44

…plicate direct-load tests, unification TODO

Pass direct=False to mxfp8 test weight syncs

b74d652

Unify generator model state dict loads; model_state_dict_load_mode

e042bf3

Address review gate: narrow load failure docs

cd2b3a7

Address review: reload test on default path; ModelStateDictLoadMode a…

03d4521

…lias

d4l3k force-pushed the d4l3k/titan-gen-direct-gpu branch from 795aa03 to 03d4521 Compare November 1, 2026 06:42

d4l3k added 2 commits November 2, 2026 11:19

Merge remote-tracking branch 'origin/main' into d4l3k/titan-gen-direc…

c5023b0

…t-gpu # Conflicts: # torchtitan/rl/generator/vllm/model_wrapper.py

gpu_fetch: fetch float32 GDN params in the model dtype, then cast

fdb352d

		for param_group in module._get_fsdp_state()._fsdp_param_groups:
		param_group._to_sharded()

	torch.cuda.current_stream().synchronize()
	torch.accelerator.current_stream().synchronize()

	fetch_model_state_dict_to_gpu: bool = False,
	fetch_model_state_dict_to_gpu: bool,



		@pytest.mark.parametrize("fetch_to_gpu", [False, True])
		def test_fetch_model_state_dict_to_gpu_reaches_engine_and_engine_loop(

		self._check_vllm_prefill_and_decode_batch_invariance()


		class TestBitwiseParityQwen35DebugVarlenDirect(TestBitwiseParityQwen35DebugVarlen):

		# TODO: unify with the pinned CPU prefetch into one load path (e.g.
		# direct fetch overlapped with generation) and drop this flag.

Conversation

d4l3k commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tianyu-l commented Oct 29, 2026

Uh oh!

claude Bot commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Must fix

Should consider

Questions

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Oct 29, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Nov 1, 2026

Choose a reason for hiding this comment

Uh oh!

brandondfeng left a comment

Choose a reason for hiding this comment

Uh oh!

tianyu-l Nov 1, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Nov 1, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Nov 1, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Nov 1, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Nov 1, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l Nov 1, 2026

Choose a reason for hiding this comment

Uh oh!

tianyu-l commented Nov 1, 2026

Uh oh!

claude Bot commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

tianyu-l's review

Questions

Uh oh!

tianyu-l commented Nov 1, 2026

Uh oh!

claude Bot commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Must fix

Should consider

Questions

d4l3k commented Oct 28, 2026 •

edited

Loading

claude Bot commented Oct 29, 2026 •

edited

Loading

claude Bot commented Nov 1, 2026 •

edited

Loading

claude Bot commented Nov 1, 2026 •

edited

Loading