[graph_trainer] Overlap chunk-owned paged stashing with HybridEP

.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml

@@ -109,6 +109,8 @@ jobs:
         # Exercise real HybridEP transfers and full-CUDA overlap contracts.
         CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestHybridEPContracts -v
         CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_eager_overlap -v
+        CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_eager_paging -v
+        CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_eager_no_remat_paging -v
 
         # Run precompile integration tests (DSv3 with EP; Llama3 runs in default workflow)
         NCCL_NVLS_ENABLE=0 python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests $RUNNER_TEMP/artifacts-to-be-uploaded/precompile --ngpu 8 --test_name aot_fx_trace_deepseek_v3_precompile_fsdp_tp_ep

torchtitan/experiments/graph_trainer/README.md

@@ -153,7 +153,7 @@ Router arguments and the full-call auxiliary-loss denominator are shared.
 Dense feed-forward modules are not chunked.
 
 The workflow is eager chunking -> tracing -> unshard normalization -> memory
-policy -> collective bucketing -> EP scheduling -> regional
+policy and paging -> collective bucketing -> EP scheduling -> regional
 attention compilation -> whole forward/backward CUDA capture and replay.
 Shared parametrized weights are read once outside chunk ownership, so both
 chunks accumulate through one parameter-gradient reduction edge.
@@ -179,11 +179,17 @@ accept compiler layouts and normalize native contiguous inputs internally;
 waits preserve storage rather than making independent layout copies. The
 scheduler issues the second chunk's launch
 before the first chunk's wait, then prefers ready launches, independent work
-and finally waits. It preserves data and mutation dependencies;
+and finally waits. It preserves data, mutation and paging queue dependencies;
 it does not infer scheduling roles from node names or tensor shapes.
 Among ready computations, launch ancestors go first so input gradients expose
 their exchange before independent weight gradients. The latter can then cover
 the final backward combine rather than being drained before it launches.
+With eager paging, parameter-gradient annotations identify wgrad-only branches
+outside the exchanges' input dependencies. The routed chunk0 branch waits to
+become eligible until its terminal backward combine is issued. Independent
+shared-expert and chunk1 work remain available earlier; reload and EP waits
+stay late. Unchunked PP retains complementary F/B work rather than reserving
+every backward wgrad for its own combine.
 
 The whole-call output row budget is
 `B = num_max_tokens_per_rank * ep_size * min(num_local_experts, top_k) * factor`.
@@ -203,13 +209,50 @@ differences must not substitute for checking completion and storage ownership.
 Native capacity validation observes the real eager MoE wrapper at factor 4
 and an odd, aligned budget; it does not hand-bind the expected capacity.
 
+Set `config.compile.memory_policy = "sac_and_paged_stash"` to page saved expert
+activations. Counts and stash commits belong to each layer/chunk, not just the
+layer. Annotated stash, reload, prefetch and wait operations share the existing
+paging stream and allocator. A true stash overflow invalidates backward and
+uses the existing device-side optimizer-skip protocol; it is not a native
+capacity drop. Initialization measures live pages once, retaining compact live
+prefixes rather than every capacity-padded activation until backward. Reload
+restores the capacity shape with the same live-prefix contract as real paging;
+padding is not meaningful. Full-count inputs and unpaged reruns keep the original
+tensor. This sizing step changes neither model computation nor recomputation.
+Steady-state exchanges and paging do not read device counts on the host. The
+CUDA page-buffer factor is headroom over that measured live-page peak,
+independent of HybridEP's
+static output capacity factor. Routing can change after calibration; choose
+headroom for that variability or the deferred overflow contract stops the run.
+
+Page copies and reloads use one high-priority stream, which serializes freelist
+updates. Set `TORCHTITAN_PAGED_STASH_MAX_CTAS` before import to bound copy/pop
+CTAs (2048 by default), without changing live-row counts or buffer sizing.
+Copy/pop CTAs stride over flat multi-token tiles and copy payloads without
+arithmetic. Page allocation, records, FIFO recycling and overflow semantics
+are unchanged; lowering the grid cap does not shorten the copied live prefix.
+Source annotations preserve each source's issue order and storage lifetimes;
+independent multiplexed F/B sources may interleave on that one stream without
+an additional whole-backward-before-forward paging chain. Per-slot records
+remain FIFO and backward generations stay stage-owned.
+
+For paging without activation recomputation, select
+`config.compile.memory_policy = "none_and_paged_stash"` and
+`config.activation_remat = None`. This saves forward activations and uses the
+same annotated page operations and runner; the existing `sac_and_paged_stash`
+mode retains its selective-recompute policy.
+
 The focused integration recipes use full forward/backward CUDA graphs for
-the two-chunk control and overlap. They select native
+the two-chunk control, overlap, and overlap with paging. They select native
 varlen attention, which uses the existing FA4 provider on Blackwell, without
 compiling attention through Inductor. They check actual
 capture and compare ten training steps of full-precision loss and gradient
 norm, including repeated replay; chunked-versus-unchunked equality is not the
-comparison contract.
+comparison contract. Paging covers both last-layer residency and
+`config.compile.paged_stash.skip_immediate_backward = False`, which requires
+a terminal stash join before backward reload. Keepalive activations belong to
+the join, not to a functional reload that can return one of them during
+calibration or residency.
 
 Varlen backend selection is explicit: use
 `VarlenInnerAttention.Config(sdpa_backend=SDPBackend.FLASH_ATTENTION)` to select
@@ -221,7 +264,7 @@ Current limitations:
 
 - `expert_parallel_degree > 1` and `pipeline_parallel_degree == 1` are required.
 - DeepEP and GraphPP overlap are outside this feature's scope.
-- Full-graph Inductor is outside this feature's scope; expert ops
+- Full-graph Inductor is outside this feature's scope; expert and paging ops
   execute natively. CUDA graphs capture the whole forward/backward callable.
 - The token count must be even.
 - Additional overlap or sharding modes require their own loss and trace checks.

torchtitan/experiments/graph_trainer/common_utils.py

@@ -35,6 +35,7 @@
 
 
 BOXED_CODEGEN_META = "graph_trainer_boxed_codegen"
+_GRAPH_PP_DIRECTION = "graph_pp_direction"
 LossResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, dict[str, Any]]
 AnnotatedLossFn: TypeAlias = Callable[..., LossResult]
 

torchtitan/experiments/graph_trainer/configs.py

@@ -195,6 +195,7 @@ class GraphTrainerCompileConfig:
         "min_cut",
         "sac_and_offload",
         "sac_and_paged_stash",
+        "none_and_paged_stash",
     ] = "default"
     """
     Memory optimization policy for activation management (SAC, offload, stash).
@@ -216,6 +217,8 @@ class GraphTrainerCompileConfig:
             rows stay resident. Requires a token dispatcher with a static
             capacity factor (e.g. HybridEP with non_blocking_capacity_factor
             set). Tuned by the paged_stash settings.
+        none_and_paged_stash: save forward activations without recomputation,
+            paging capacity-padded expert activations with the same stash.
     """
 
     full_recompute_save_ops: str = ""
@@ -257,8 +260,8 @@ class GraphTrainerCompileConfig:
     Tensors are selected largest-first until the budget is exhausted."""
 
     paged_stash: PagedStashConfig = field(default_factory=PagedStashConfig)
-    """Settings for MoE paged stashing, which
-    ``memory_policy = "sac_and_paged_stash"`` enables."""
+    """MoE paged-stash settings for the sac_and_paged_stash and
+    none_and_paged_stash memory policies."""
 
     enable_fsdp_ag_rs_overlap: bool = False
     """When True, run ``overlap_fsdp_ag_rs_pass``. The pass moves backward

torchtitan/experiments/graph_trainer/cpu_offload.py

@@ -127,13 +127,17 @@ def _is_view(node: Node) -> bool:
 
 
 def _get_storage_chain(
-    node: Node, *, is_view: Callable[[Node], bool] = _is_view
+    node: Node,
+    *,
+    is_view: Callable[[Node], bool] = _is_view,
+    is_backward: Callable[[Node], bool] = _is_backward_node,
 ) -> tuple[set[Node], bool]:
     """Walk the view chain to find all nodes sharing this tensor's storage.
 
     Views alias the base tensor's storage, so any node reachable through view
     edges (and its non-view consumers) depends on the storage being alive.
-    ``is_view`` decides which nodes extend the chain.
+    ``is_view`` decides which nodes extend the chain; ``is_backward`` decides
+    which consumers read the saved storage rather than retain a copy keepalive.
 
     Returns:
         (chain_nodes, has_bwd) where chain_nodes is every forward
@@ -149,7 +153,7 @@ def _walk(n: Node) -> None:
         for user in n.users:
             if user.op != "call_function":
                 continue
-            if _is_backward_node(user):
+            if is_backward(user):
                 has_bwd = True
                 continue
             chain_nodes.add(user)

torchtitan/experiments/graph_trainer/ep_overlap_pass.py

@@ -222,15 +222,59 @@ def _prioritized_order(
 
     Prefer compute that exposes another launch over unrelated compute. This
     lets input gradients issue their exchange before independent weight grads,
-    while waits still follow every available computation.
+    while waits still follow every available computation. For eager paging,
+    retain chunk0's wgrad-only branch until its final combine is issued. Shared
+    experts, chunk1 and unchunked PP remain eligible to cover earlier exchanges.
     """
+    from torchtitan.experiments.graph_trainer.common_utils import (
+        _is_backward_node,
+        PARAMETER_GRADIENT_FQNS_META,
+    )
+    from torchtitan.experiments.graph_trainer.ep_pass_utils import _chunk_owner
+    from torchtitan.experiments.graph_trainer.paged_stash_memory_policy import (
+        PAGED_STASH_REGION,
+    )
+
     launch_ancestors = set()
     pending = list(launches)
     while pending:
         node = pending.pop()
         for predecessor in deps[node] - launch_ancestors:
             launch_ancestors.add(predecessor)
             pending.append(predecessor)
+    terminal_combines = {}
+    for launch in launches:
+        owner = _chunk_owner(launch)
+        if (
+            owner is not None
+            and owner.is_backward
+            and owner.chunk_id == 0
+            and launch.meta.get("custom", {}).get(_EP_TOKEN_EXCHANGE) == "combine"
+        ):
+            terminal_combines[owner] = launch
+    deferred_wgrads = {}
+    if terminal_combines:
+        # Shared data feeding an exchange is not a wgrad-only branch.
+        pending = [
+            node
+            for node in deps
+            if node.meta.get("custom", {}).get(PARAMETER_GRADIENT_FQNS_META)
+        ]
+        visited = set()
+        while pending:
+            node = pending.pop()
+            if node in visited or node in launch_ancestors:
+                continue
+            visited.add(node)
+            custom = node.meta.get("custom", {})
+            if not (
+                _is_backward_node(node) or custom.get(PARAMETER_GRADIENT_FQNS_META)
+            ):
+                continue
+            owner = _chunk_owner(node)
+            if custom.get(PAGED_STASH_REGION) and owner in terminal_combines:
+                deferred_wgrads[node] = terminal_combines[owner]
+            pending.extend(node.all_input_nodes)
     successors: dict[fx.Node, list[fx.Node]] = {node: [] for node in deps}
     for node, predecessors in deps.items():
         for predecessor in predecessors:
@@ -255,6 +299,9 @@ def _prioritized_order(
                         if candidate in launches
                         else 3
                         if candidate in waits
+                        else 4
+                        if candidate in deferred_wgrads
+                        and deferred_wgrads[candidate] not in emitted
                         else 1
                         if candidate in launch_ancestors
                         else 2
@@ -285,10 +332,13 @@ def _plan_order(
     exchanges: list[dict[int, tuple[_TokenExchange, ...]]],
 ) -> list[fx.Node]:
     """Choose the earliest legal launches and postpone waits behind ready work."""
+    from torchtitan.experiments.graph_trainer.paged_stash_memory_policy import (
+        _paging_dependencies,
+    )
 
     order = ordered_nodes(gm)
     deps = _graph_dependencies(gm)
-    launches, waits = set(), set()
+    launches, waits = _paging_dependencies(gm, deps)
     for region, chunk_exchanges in zip(regions, exchanges, strict=True):
         chunk_order = (1, 0) if region.is_backward else (0, 1)
         for first, second in zip(