diff --git a/.github/scripts/validate_release_gpu.sh b/.github/scripts/validate_release_gpu.sh index 6c288fd3ed1..e6ee823fb4a 100644 --- a/.github/scripts/validate_release_gpu.sh +++ b/.github/scripts/validate_release_gpu.sh @@ -299,7 +299,7 @@ run_graph_trainer_h100_tests() { "${ARTIFACTS}/graph_trainer_h100" \ --ngpu 8 python -m pytest "${tests_root}/test_numerics.py::TestGraphTrainerNumerics" \ - -v -k "moe or hybridep_eager" + -v -k "moe or hybridep_eager or hybridep_combined_pp or hybridep_dual_pipe_v" python -m pytest "${tests_root}/test_numerics.py::TestHybridEPContracts" -v python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests \ "${ARTIFACTS}/graph_trainer_h100_precompile" \ diff --git a/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml b/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml index 3ec0a613885..13c3248adf5 100644 --- a/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml +++ b/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml @@ -111,6 +111,9 @@ jobs: 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 + CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_combined_pp_overlap_and_paging -v + CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_dual_pipe_v_overlap_and_paging -v + CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_combined_pp_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 diff --git a/torchtitan/experiments/graph_trainer/README.md b/torchtitan/experiments/graph_trainer/README.md index 47d04e7f82f..2badb539dae 100644 --- a/torchtitan/experiments/graph_trainer/README.md +++ b/torchtitan/experiments/graph_trainer/README.md @@ -105,8 +105,7 @@ live `param.grad`. Internal values remain flat because they never leave GraphPP graph execution: saved-for-backward tensors, unsharded FSDP params, raw grad leaves, reduce-grad inputs, and multiplexed intermediate outputs. -Current limitations: GraphPP does not load precompile artifacts yet, and -EP-overlap annotations will be composed with GraphPP in a later PR. +Current limitation: GraphPP does not load precompile artifacts yet. ### Compiler Optimizations @@ -269,6 +268,58 @@ Current limitations: - The token count must be even. - Additional overlap or sharding modes require their own loss and trace checks. +### Pipeline Communication Overlap + +Pipeline overlap uses `Interleaved1F1BCombined` or `DualPipeV` and graph +multiplexing, not EP token chunking. `Interleaved1F1BCombined` adds one warmup +forward, then pairs an independent +forward and backward as an explicit `OVERLAP_F_B` action. Both actions' receives +and parameter unshards precede the pair; normal send, gradient accumulation and +reduction ownership remain in the existing pipeline schedule. + +`DualPipeV` uses the upstream V-shaped mapping with exactly two stages per +rank. Its combined actions use full backward; separate input-gradient and +weight-gradient actions fill the bubbles. Their saved activations remain live +until weight backward completes. The focused HybridEP recipes use native +varlen attention for both schedules, including actual dI/dW graph splitting. + +```python +config.parallelism.pipeline_parallel_schedule = "Interleaved1F1BCombined" +config.compile.enable_graph_pp_comm_compute_overlap = True +config.compile.ep_overlap.enabled = False +config.compile.inductor_compilation = "regional" +config.training.disable_cuda_graphs = False +``` + +Use non-blocking `HybridEPTokenDispatcher` capacity. Enabling the GraphPP pass +selects explicit asynchronous native launches/waits without splitting tokens +or dividing each microbatch's capacity. Configure `sac_and_paged_stash` to +include page operations. Source graphs retain explicit forward/backward +ownership when multiplexed. The pass validates native request/input lifetimes, +preserves tensor/mutation and paging queue dependencies, launches ready +communication early and postpones waits behind independent compute. +Each EP or paging launch gives ready work in the opposite source graph priority, +including its ready completion waits, instead of draining one source first. +Full CUDA capture covers the entire pipeline forward/backward step, including every +microbatch's distinct native routing state. + +For long multistream captures, PyTorch's native allocator can reclaim storage +after graph dependencies prove its asynchronous uses have joined: + +```bash +PYTORCH_ALLOC_CONF=graph_capture_record_stream_reuse:True +``` + +This upstream option is experimental. Without it, cross-stream frees can wait +until capture ends, growing the graph pool with the microbatch count. Keep the +same allocator setting in matched comparisons; it does not replace request +ownership, stream joins, or the explicit overlap annotations. + +The pass can be disabled with `compile.disable_passes = +["graph_pp_comm_compute_overlap_pass"]` for a same-schedule comparison. Without +combined actions, the requested overlap emits a warning. Full-graph Inductor +and DeepEP are outside this feature's scope. + ### Experimental AutoParallel Sharding GraphTrainer can use AutoParallel to solve SPMD placement for supported models, diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 13b4c5661c5..32c0fd51ddb 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -283,6 +283,11 @@ class GraphTrainerCompileConfig: ep_overlap: EpOverlapConfig = field(default_factory=EpOverlapConfig) """Configuration for EP-overlap chunking and scheduling.""" + enable_graph_pp_comm_compute_overlap: bool = False + """Overlap HybridEP and paged-stash operations across multiplexed PP F/B + graphs. Requires a combined-action schedule and non-blocking HybridEP. + This does not chunk tokens or divide per-microbatch capacity budgets.""" + precompile_artifact_dir: str = "" """ Directory for precompiled artifacts. Setting this enables precompile: diff --git a/torchtitan/experiments/graph_trainer/ep_overlap_pass.py b/torchtitan/experiments/graph_trainer/ep_overlap_pass.py index 248ab4fddfb..bb1ff83dee2 100644 --- a/torchtitan/experiments/graph_trainer/ep_overlap_pass.py +++ b/torchtitan/experiments/graph_trainer/ep_overlap_pass.py @@ -214,12 +214,14 @@ def _prioritized_order( launches: set[fx.Node], waits: set[fx.Node], regions: list[tuple[set[fx.Node], set[fx.Node]]], + yield_to: dict[fx.Node, set[fx.Node]] | None = None, ) -> list[fx.Node]: """Topologically schedule ordered scopes with early launches and late waits. Each region contains its eligible scope and the owned body whose completion - advances to the next region. - + advances to the next region. A multiplexed pair uses one whole-graph scope. + Optional launch boundaries prefer complementary work among equally ranked + ready nodes, preventing one source from draining all its completion waits. 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. For eager paging, @@ -284,6 +286,7 @@ def _prioritized_order( emitted = set() result = [] active_region = 0 + preferred = set() while ready: while active_region < len(regions) and regions[active_region][1] <= emitted: active_region += 1 @@ -306,6 +309,7 @@ def _prioritized_order( if candidate in launch_ancestors else 2 ), + candidate not in preferred, order[candidate], ), ) @@ -314,6 +318,8 @@ def _prioritized_order( ready.remove(node) emitted.add(node) result.append(node) + if yield_to is not None and node in yield_to: + preferred = yield_to[node] for successor in successors[node]: remaining[successor] -= 1 if remaining[successor] == 0: diff --git a/torchtitan/experiments/graph_trainer/graph_builder.py b/torchtitan/experiments/graph_trainer/graph_builder.py index e6a04c6fa6f..fd87e02eee1 100644 --- a/torchtitan/experiments/graph_trainer/graph_builder.py +++ b/torchtitan/experiments/graph_trainer/graph_builder.py @@ -299,6 +299,11 @@ def prepare_graphs( required_overlap_pairs = stage_builder._required_multiplex_pairs(schedule) if not required_overlap_pairs: + if self.config.compile.enable_graph_pp_comm_compute_overlap: + logger.warning( + "GraphPP communication overlap requested, but the schedule " + "contains no combined forward/backward actions." + ) self._overlap_graphs = {} overlap_graphs: dict[tuple[int, int], OverlapStageGraphs] = {} elif self._overlap_graphs is None: diff --git a/torchtitan/experiments/graph_trainer/graph_pp/graph_multiplex.py b/torchtitan/experiments/graph_trainer/graph_pp/graph_multiplex.py index c94a253c0fd..442bee2c5d1 100644 --- a/torchtitan/experiments/graph_trainer/graph_pp/graph_multiplex.py +++ b/torchtitan/experiments/graph_trainer/graph_pp/graph_multiplex.py @@ -8,8 +8,10 @@ from itertools import dropwhile import torch.fx as fx +from torch._dynamo.graph_deduplication import _stable_topological_sort from torch.fx.graph_module import _assign_attr, _get_attr +from torchtitan.experiments.graph_trainer.common_utils import _GRAPH_PP_DIRECTION from torchtitan.experiments.graph_trainer.debug_utils import tlparse_log_graph_pass from torchtitan.experiments.graph_trainer.graph_pp.utils import ( _find_fake_mode, @@ -17,12 +19,86 @@ ) +def graph_pp_comm_compute_overlap_pass( + gm: fx.GraphModule, example_inputs: tuple | None = None +) -> fx.GraphModule: + """Overlap explicitly owned F/B communications in a multiplexed graph. + + Native exchanges retain their exact inputs through completion. Tensor and + mutation edges, plus the paging queue order, define the legal schedule. + Both source graphs form one scope: launches first, ready compute next, + completion waits last. Each communication launch yields to ready work in + the opposite source graph, so neither source drains all its waits first. + No tokens or computation are split or duplicated. + """ + del example_inputs + from torchtitan.experiments.graph_trainer.ep_overlap_pass import ( + _collect_native_exchanges, + _graph_dependencies, + _prioritized_order, + ) + from torchtitan.experiments.graph_trainer.paged_stash_memory_policy import ( + _paging_dependencies, + ) + + exchanges = _collect_native_exchanges(gm.graph.nodes, context="GraphPP F/B pair") + deps = _graph_dependencies(gm) + launches, waits = _paging_dependencies(gm, deps) + source_nodes = { + direction: { + node + for node in gm.graph.nodes + if node.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION) == direction + } + for direction in ("forward", "backward") + } + yield_to = {} + for exchange in exchanges: + direction = exchange.launch.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION) + if direction not in ("forward", "backward") or direction != ( + exchange.wait.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION) + ): + raise ValueError( + f"GraphPP exchange {exchange.launch.name} lost its source F/B ownership" + ) + launches.add(exchange.launch) + waits.add(exchange.wait) + for launch in launches: + direction = launch.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION) + if direction not in source_nodes: + raise ValueError( + f"GraphPP launch {launch.name} lost its source F/B ownership" + ) + yield_to[launch] = source_nodes[ + "backward" if direction == "forward" else "forward" + ] + if not launches and not waits: + return gm + nodes = set(gm.graph.nodes) + planned = _prioritized_order( + deps, + order={node: index for index, node in enumerate(gm.graph.nodes)}, + launches=launches, + waits=waits, + regions=[(nodes, nodes)], + yield_to=yield_to, + ) + sequencing = {node: {previous} for previous, node in zip(planned, planned[1:])} + _stable_topological_sort(gm.graph, sequencing) + gm.graph.lint() + gm.recompile() + if list(gm.graph.nodes) != planned: + raise ValueError("GraphPP communication schedule did not materialize its order") + return gm + + def _copy_prefixed_get_attrs( dst: fx.GraphModule, src: fx.GraphModule, *, prefix: str, ) -> dict[str, str]: + """Copy referenced attributes under unique names and return the remapping.""" remap: dict[str, str] = {} for node in src.graph.find_nodes(op="get_attr"): attr_name = str(node.target) @@ -63,13 +139,10 @@ def multiplex_fw_bw_graph( replace the output tuple with backward outputs followed by forward outputs The forward graph remains the destination module because its ShapeEnv owns - the dynamic collective-size constraints needed by full Inductor for MoE - all-to-all outputs. The backward graph is inserted in topological order - before the existing forward compute, with disjoint placeholders and - prefixed attributes, so this helper does not need dependency analysis or a - scheduling policy. EP-overlap annotations are intentionally not applied - inside this graph; pass ordering keeps EP-overlap on the standalone no-FSDP - forward/backward graphs. + the traced collective-size constraints. The backward graph is inserted + before forward compute, with disjoint placeholders and prefixed attributes. + Explicit source ownership allows the optional communication/compute pass + to reorder the pair; multiplexing itself preserves each source's order. Args: fw_gm (fx.GraphModule): Forward graph whose ShapeEnv and module @@ -90,6 +163,11 @@ def multiplex_fw_bw_graph( # Preserve the forward ShapeEnv exactly as traced. Reconstructing it from # copied metadata can lose collective-size hints that full Inductor needs. multiplexed_gm = copy.deepcopy(fw_gm) + for node in multiplexed_gm.graph.nodes: + node.meta["custom"] = { + **node.meta.get("custom", {}), + _GRAPH_PP_DIRECTION: "forward", + } dst_fake_mode = _find_fake_mode(multiplexed_gm) meta_transfer = _MetaShapeEnvTransfer(dst_fake_mode) for node in bw_gm.graph.nodes: @@ -141,6 +219,10 @@ def multiplex_fw_bw_graph( node, lambda arg: old_to_new[arg] ) new_node.meta = meta_transfer.copy_meta(node.meta) + new_node.meta["custom"] = { + **new_node.meta.get("custom", {}), + _GRAPH_PP_DIRECTION: "backward", + } if new_node.op == "get_attr": attr_name = str(new_node.target) if attr_name not in bw_get_attr_remap: diff --git a/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py b/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py index f0611148205..c4d8494822e 100644 --- a/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py +++ b/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py @@ -56,6 +56,8 @@ from torchtitan.experiments.graph_trainer.graph_pp.stage import GraphPipelineStage from torchtitan.experiments.graph_trainer.registry import PASS_PIPELINE_REGISTRY from torchtitan.models.base.model import BaseModel +from torchtitan.models.common.moe import MoE +from torchtitan.models.common.moe.token_dispatcher import HybridEPTokenDispatcher if TYPE_CHECKING: @@ -196,6 +198,19 @@ def _validate_graph_pp_config( compile_config: GraphTrainerCompileConfig, parallelism: ParallelismConfig, ) -> None: + if compile_config.enable_graph_pp_comm_compute_overlap: + if compile_config.ep_overlap.enabled: + raise ValueError( + "GraphPP communication overlap does not use EP token chunking" + ) + if compile_config.inductor_compilation != "regional": + raise ValueError( + "GraphPP communication overlap requires regional compilation" + ) + if parallelism.expert_parallel_degree <= 1: + raise ValueError( + "GraphPP HybridEP overlap requires expert_parallel_degree > 1" + ) if compile_config.precompile_artifact_dir: raise ValueError( "GraphPP does not support compile.precompile_artifact_dir yet. " @@ -261,6 +276,8 @@ def resolve_graph_execution_plan( pp_enabled: Whether pipeline parallelism is enabled. fsdp_enabled: Whether FSDP is enabled. """ + if compile_config.enable_graph_pp_comm_compute_overlap and not pp_enabled: + raise ValueError("GraphPP communication overlap requires pipeline parallelism") if pp_enabled: _validate_graph_pp_config( compile_config=compile_config, @@ -850,6 +867,18 @@ def graph_pipeline_llm( activation_remat=activation_remat, dump_folder=dump_folder, ) + if compile_config.enable_graph_pp_comm_compute_overlap: + for module in model_part.modules(): + if not isinstance(module, MoE): + continue + dispatcher = module.routed_experts.token_dispatcher + if not isinstance(dispatcher, HybridEPTokenDispatcher): + raise ValueError("GraphPP communication overlap requires HybridEP") + if dispatcher.non_blocking_capacity_factor is None: + raise ValueError( + "GraphPP overlap requires non-blocking HybridEP capacity" + ) + dispatcher.async_op = True logger.info( "PP rank %s is building GraphPP stage_idx %s with modules %s", pp_mesh.get_local_rank(), diff --git a/torchtitan/experiments/graph_trainer/graph_pp/stage_builder.py b/torchtitan/experiments/graph_trainer/graph_pp/stage_builder.py index 3dca18b5e33..2ec2c1d8b1d 100644 --- a/torchtitan/experiments/graph_trainer/graph_pp/stage_builder.py +++ b/torchtitan/experiments/graph_trainer/graph_pp/stage_builder.py @@ -19,6 +19,7 @@ from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig from torchtitan.experiments.graph_trainer.graph_pp.graph_multiplex import ( + graph_pp_comm_compute_overlap_pass, multiplex_fw_bw_graph, ) from torchtitan.experiments.graph_trainer.graph_pp.runner import ( @@ -35,8 +36,10 @@ SplitStageGraphs, ) from torchtitan.experiments.graph_trainer.graph_pp.utils import ( + example_inputs_from_placeholders, overlap_fw_bw_sub_actions, ) +from torchtitan.experiments.graph_trainer.passes import apply_graph_passes class _StageGraphModules(Protocol): @@ -246,6 +249,13 @@ def _build_graph_pp_overlap_graphs( fw_graphs.modules.fw, bw_graphs.modules.full_bw, ) + if compile_config.enable_graph_pp_comm_compute_overlap: + multiplexed_graph = apply_graph_passes( + multiplexed_graph, + example_inputs_from_placeholders(multiplexed_graph), + [graph_pp_comm_compute_overlap_pass], + compile_config=compile_config, + ) compiled_graph = compile_graph_module( multiplexed_graph, compile_config=compile_config, diff --git a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py index 09e5b7cab11..19413246cfe 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py +++ b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py @@ -1051,6 +1051,46 @@ def test_runtime_schedule_validation_accepts_interleaved(self) -> None: parallelism=ParallelismConfig(pipeline_parallel_schedule="Interleaved1F1B"), ) + def test_graph_pp_overlap_configuration(self) -> None: + """Exercise the real execution-plan configuration contract.""" + for invalid in (None, "no_pp", "no_ep", "ep_chunking", "full_compile"): + with self.subTest(invalid=invalid): + compile_config = GraphTrainerCompileConfig( + enable_graph_pp_comm_compute_overlap=True, + inductor_compilation="regional", + ) + parallelism = ParallelismConfig( + pipeline_parallel_schedule="Interleaved1F1BCombined", + pipeline_parallel_degree=2, + expert_parallel_degree=2, + ) + if invalid == "no_ep": + parallelism.expert_parallel_degree = 1 + elif invalid == "ep_chunking": + compile_config.ep_overlap.enabled = True + elif invalid == "full_compile": + compile_config.inductor_compilation = "full" + kwargs = dict( + num_microbatches=4, + parallelism=parallelism, + pp_enabled=invalid != "no_pp", + fsdp_enabled=True, + ) + if invalid is None: + plan = resolve_graph_execution_plan(compile_config, **kwargs) + self.assertEqual( + (plan.unshard, plan.reduce_grad), ("schedule", "schedule") + ) + else: + message = { + "no_pp": "requires pipeline parallelism", + "no_ep": "expert_parallel_degree > 1", + "ep_chunking": "does not use EP token chunking", + "full_compile": "requires regional compilation", + }[invalid] + with self.assertRaisesRegex(ValueError, message): + resolve_graph_execution_plan(compile_config, **kwargs) + def test_split_fsdp_boundaries_reject_dense_region_overlap(self) -> None: compile_config = GraphTrainerCompileConfig( enable_fsdp_dense_region_overlap=True diff --git a/torchtitan/experiments/graph_trainer/tests/test_numerics.py b/torchtitan/experiments/graph_trainer/tests/test_numerics.py index cba7b04479a..07261102062 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_numerics.py +++ b/torchtitan/experiments/graph_trainer/tests/test_numerics.py @@ -490,6 +490,26 @@ def test_hybridep_eager_paging(self): "deepseek_v3_hybridep_chunk_control", target ) + def test_hybridep_combined_pp_overlap_and_paging(self): + """Compare real PP2/EP2 training with full pipeline CUDA replay.""" + for target in ( + "deepseek_v3_hybridep_combined_pp_overlap", + "deepseek_v3_hybridep_combined_pp_overlap_paged", + ): + self._assert_hybridep_training_matches( + "deepseek_v3_hybridep_combined_pp_control", target + ) + + def test_hybridep_dual_pipe_v_overlap_and_paging(self): + """Exercise DualPipeV full and dI/dW backward with native varlen attention.""" + for target in ( + "deepseek_v3_hybridep_dual_pipe_v_overlap", + "deepseek_v3_hybridep_dual_pipe_v_overlap_paged", + ): + self._assert_hybridep_training_matches( + "deepseek_v3_hybridep_dual_pipe_v_control", target + ) + def test_dense_llama3_hsdp_gradient_accumulation_aot_fx_trace_vs_eager(self): self.assertTrue(_run_hsdp_llama3_gradient_accumulation_loss_compare()) @@ -501,6 +521,14 @@ def test_hybridep_eager_no_remat_paging(self): no_remat=True, ) + def test_hybridep_combined_pp_no_remat_paging(self): + """Compare no-remat combined PP with and without paged overlap.""" + self._assert_hybridep_training_matches( + "deepseek_v3_hybridep_combined_pp_no_remat_control", + "deepseek_v3_hybridep_combined_pp_no_remat_paged", + no_remat=True, + ) + def test_dense_llama3_hsdp_graph_pp_vs_eager(self): self.assertTrue(_run_hsdp_llama3_pp_loss_compare()) diff --git a/torchtitan_recipes/tests/graph_trainer/numerics.py b/torchtitan_recipes/tests/graph_trainer/numerics.py index c49989f5c62..3779e72af09 100644 --- a/torchtitan_recipes/tests/graph_trainer/numerics.py +++ b/torchtitan_recipes/tests/graph_trainer/numerics.py @@ -121,6 +121,105 @@ def deepseek_v3_hybridep_no_remat_paged(): return _deepseek_v3_hybridep_chunks(overlap=True, paging=True, no_remat=True) +def _deepseek_v3_hybridep_pp( + *, + schedule: Literal["Interleaved1F1BCombined", "DualPipeV"], + overlap: bool, + paging: bool, + no_remat: bool = False, +): + """Build matched native-varlen PP2/EP2 recipes without EP token chunking.""" + config = graph_trainer_deepseek_v3_debugmodel( + attn_backend="varlen", + parallelism=ParallelismConfig( + data_parallel_shard_degree=2, + expert_parallel_degree=2, + pipeline_parallel_degree=2, + pipeline_parallel_schedule=schedule, + pipeline_parallel_layers_per_stage=2, + num_pp_microbatches=4, + ), + ) + config.training.num_tokens_per_microbatch_per_dp_rank = 512 + config.training.num_tokens_per_train_step = 4096 + config.dataloader.max_num_documents = c4_document_capacity( + config.training.num_tokens_per_microbatch_per_dp_rank + ) + config.training.disable_cuda_graphs = False + config.compile.inductor_compilation = "regional" + config.compile.enable_graph_pp_comm_compute_overlap = overlap + if no_remat: + config.activation_remat = None + config.compile.memory_policy = "none" if no_remat else "default" + if paging: + config.compile.memory_policy = ( + "none_and_paged_stash" if no_remat else "sac_and_paged_stash" + ) + config.compile.paged_stash.buffer_size_factor_cuda = 2.1 + config.debug.seed = 42 + config.debug.deterministic = True + return apply_transforms( + config, + [ + TokenDispatcherTransform( + dispatcher=HybridEPTokenDispatcher, + num_max_tokens_per_rank=512, + kwargs={"non_blocking_capacity_factor": 4.0, "async_op": True}, + ) + ], + ) + + +def deepseek_v3_hybridep_combined_pp_control(): + """Run the combined PP schedule with sequential multiplexed graphs.""" + return _deepseek_v3_hybridep_pp( + schedule="Interleaved1F1BCombined", overlap=False, paging=False + ) + + +def deepseek_v3_hybridep_combined_pp_overlap(): + """Overlap native EP across combined PP F/B graphs without token chunks.""" + return _deepseek_v3_hybridep_pp( + schedule="Interleaved1F1BCombined", overlap=True, paging=False + ) + + +def deepseek_v3_hybridep_combined_pp_overlap_paged(): + """Overlap native EP and paged stashing across combined PP F/B graphs.""" + return _deepseek_v3_hybridep_pp( + schedule="Interleaved1F1BCombined", overlap=True, paging=True + ) + + +def deepseek_v3_hybridep_dual_pipe_v_control(): + """Run native DualPipeV full/split backward without EP/page scheduling.""" + return _deepseek_v3_hybridep_pp(schedule="DualPipeV", overlap=False, paging=False) + + +def deepseek_v3_hybridep_dual_pipe_v_overlap(): + """Overlap HybridEP across DualPipeV's combined forward/backward graphs.""" + return _deepseek_v3_hybridep_pp(schedule="DualPipeV", overlap=True, paging=False) + + +def deepseek_v3_hybridep_dual_pipe_v_overlap_paged(): + """Overlap native HybridEP and paged stashing under DualPipeV.""" + return _deepseek_v3_hybridep_pp(schedule="DualPipeV", overlap=True, paging=True) + + +def deepseek_v3_hybridep_combined_pp_no_remat_control(): + """Run the combined schedule without activation recomputation or overlap.""" + return _deepseek_v3_hybridep_pp( + schedule="Interleaved1F1BCombined", overlap=False, paging=False, no_remat=True + ) + + +def deepseek_v3_hybridep_combined_pp_no_remat_paged(): + """Overlap native EP and pages under the combined schedule without remat.""" + return _deepseek_v3_hybridep_pp( + schedule="Interleaved1F1BCombined", overlap=True, paging=True, no_remat=True + ) + + def _without_inplace_wgrad_accum(config): """Turn off in-place WGRAD accumulation in an eager baseline.