{"id":"ebcb3c83-c3cb-482a-a9ec-b5a4b031a97f","arxiv_id":"2412.14374","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":2,"one_line_summary":"JaxPP introduces a user-defined MPMD pipeline schedule API on top of JAX/GSPMD and reports throughput gains up to 1.11x over SPMD training on H100 clusters.","lead":"JaxPP is a new system that lets JAX users add pipeline parallelism to large model training with a few annotations, running different pipeline stages on different groups of GPUs. On GPT-3 175B and Llama2 70B benchmarks it reports up to 1.11x higher GPU utilization than the best single-program (SPMD) JAX configuration and a 44.6% speedup over JAX's SPMD pipeline parallelism.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"No correctness validation for transformed autodiff: the reported throughput numbers could be for a program that computes the wrong updates.","rationale":"The reader's weakest assumption identifies the lack of validation for arbitrary user schedules, which is real and important. My read agrees that correctness is the weakest point, but I see it as broader: the paper provides no evidence that the transformed autodiff program preserves the semantics of the original training step, even for the specific schedules it evaluates. The performance measurements are presented as step times and TFLOPS, so a transform that silently produces incorrect gradients or that deadlocks only on unexercised schedules would not affect any reported number. This makes the headline claim load-bearing on a semantic guarantee that the paper neither proves nor tests. I considered the fairness of the 1.11x comparison against JAX FSDP, which uses different parallelism configurations; that is a meaningful caveat but not decisive because the paper explicitly compares whole systems and reports a carefully qualified \"up to\" improvement. I also considered the lack of formal proof for the Section 4.2 communication-ordering algorithm; that is a concrete gap, but it is subsumed by the absence of any correctness testing. Since the paper is already CONDITIONAL and the main remedy is the same (add correctness validation, ideally with code), my stress-test does not change the reader's verdict. I would make the condition explicit: acceptance requires demonstrating, on at least one small model, that JaxPP reproduces reference training loss and gradients, and that a representative space of user schedules runs without deadlock or data corruption.","tokens_in":20512,"tokens_out":10703,"duration_ms":97599,"concrete_test":"Run a differential test on a small tied-embedding transformer with two pipeline stages: train for 10 steps in JaxPP and in a plain JAX `jax.grad` loop with identical seeds, hyperparameters, and global batch, using per-microbatch gradient-norm logging inside the `accumulate_grads` body. Compare per-step losses, gradients, and final parameters. Separately, enumerate 100 random schedules over a 3-actor, 3-stage, 4-microbatch DAG and assert no deadlock and identical results within BF16 tolerance. Any mismatch invalidates the implicit precondition of the Section 3.4 rewrite or the claimed general schedule correctness.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The evaluation contains no correctness check: no loss curves, no gradient comparison with a reference implementation, no deadlock stress tests, and no released code. The central performance claim depends on JaxPP's Jaxpr transformations preserving the semantics of the user's training step, and this is nontrivial. Section 3.4's loop-commuting rewrite replaces the carried total-gradient accumulator with per-use partial-gradient accumulators plus a final addition; this is only valid if partial gradients are used nowhere except the final sum. The paper does not state that precondition, and Section 3.1 permits arbitrary accumulation of \"additional metrics\" in the loop body. If a user's loop body clips or logs per-microbatch gradients, the rewritten program changes behavior. Similarly, `pipeline_yield` inside a `jax.value_and_grad` (Figure 4) is not defined as a differentiable primitive; without an explicit VJP rule, autodiff treats it as identity, so the backward-pass stage boundaries may not be the ones shown. The paper never specifies how autodiff propagates stage markers. If either transformation is unsound, the reported 457 TFLOPS/device and 1.11x speedup are throughput numbers for a program that is not the intended training computation. The schedule-deadlock issue raised by the reader is a special case of this broader absence of semantic validation.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"JaxPP extends JAX's SPMD programming model with user-controlled MPMD pipeline parallelism. The paper introduces a programming model based on accumulate_grads and pipeline_yield, a driver-side task-graph transformation that unrolls the gradient accumulation loop and infers placement and communication, and a single-controller runtime built on Ray and XLA. The evaluation on GPT-3 175B and Llama2 70B reports 457 TFLOPS/device, 44.6% higher throughput than JAX SPMD pipeline parallelism, 1.11x over JAX FSDP, and 92.87% weak scaling efficiency from 64 to 1024 GPUs. The central claim is that pipeline parallelism implemented in JaxPP improves hardware utilization by up to 1.11x with respect to the best performing SPMD configuration.","tokens_in":20902,"tokens_out":3897,"duration_ms":34005,"significance":"If the semantic transformations are sound, JaxPP is a useful contribution: it gives JAX users a model-agnostic way to express 1F1B and interleaved pipeline schedules with asynchronous point-to-point communication, and the reported performance numbers are internally consistent with Table 1 and Figure 8. The paper is clearly written, the design's main pieces (loop unrolling, placement inference, fused task dispatch) are sensible, and the authors report exact configurations and a performance breakdown. However, the significance is currently bounded by the absence of any correctness validation of the Jaxpr transformations and by evaluation conditions that make the headline improvements difficult to attribute; both are fixable in revision. I also credit the authors for comparing against external systems (JAX FSDP, JAX SPMD PP, NeMo) rather than only self-baselines.","major_comments":[{"comment":"The loop-commuting rewrite in Section 3.4, g = sum_i(g1^{(i)} + g2^{(i)} + ...) ⇝ sum_i g1^{(i)} + sum_i g2^{(i)} + ..., is only valid if each partial gradient is used solely in the final addition. The paper does not state this precondition, and Section 3.1 explicitly allows the loop body to return 'additional metrics' whose reference semantics collect a per-iteration loss list. If a user's loop body clips, normalizes, or logs per-microbatch gradients, the rewritten program computes a different result. The evaluation contains no gradient comparison with a reference implementation, no loss curves, and no convergence check, so the reported TFLOPS and speedups are not shown to correspond to the intended training computation.","section":"Section 3.4 and Section 3.1"},{"comment":"pipeline_yield appears inside loss_fn, which is passed to jax.value_and_grad, but the paper never defines how automatic differentiation treats pipeline_yield: there is no custom VJP/JVP rule or marker-propagation rule specified. If pipeline_yield is treated as an identity by autodiff, the backward-pass stage boundaries need not coincide with the forward-pass stage boundaries shown in Figure 3, and the executed schedule may not be the schedule the user specified. This is load-bearing for the system's core claim of supporting user-defined schedules and must be specified and validated.","section":"Section 3.2 and Figure 4"},{"comment":"The runtime accepts an arbitrary list of Task(i, ty, stage) per actor and infers send/receive pairs in topological order, but the paper states only that 'care has to be taken' and gives one example. There is no formal statement of the schedule validity condition, no proof that the inference preserves the dataflow partial order for all schedules expressible through the API, and no deadlock stress test across a space of schedules. Because arbitrary user-defined schedules are a central selling point, the absence of a correctness argument for schedule inference is a substantive gap.","section":"Section 4.2"},{"comment":"The headline comparisons vary multiple confounding factors: for GPT-3 175B, JAX SPMD PP uses GA=128 and PP=16 while JaxPP uses GA=32 and PP=8, and JAX FSDP uses GA=1 and GBS=128 while JaxPP uses GA=32 and GBS=128. The claimed 44.6% and 1.11x improvements therefore conflate schedule choice, gradient accumulation count, and system implementation. The paper should either compare at matched global batch size and controlled GA/PP configurations, or decompose the gains. Additionally, all headline numbers appear to be single unrepeated runs; at least a few repetitions or variance estimates are needed, and the absence of released code or benchmark scripts prevents independent verification.","section":"Section 5 and Table 1"}],"minor_comments":[{"comment":"The text says 'When training Llama2 70B on 8 DGX H100 nodes (8 GPUs)', but a DGX H100 node has 8 GPUs, so the correct count is 64 GPUs, matching Table 1.","section":"Section 5.2"},{"comment":"The claim of requiring '1K fewer lines of user code' is not backed by any code-size comparison; please provide concrete line counts or remove the claim.","section":"Section 5.2"},{"comment":"The labels 'SPMD' and 'MPMD' in Figure 2 are not explained in the caption; a reader cannot tell which schedule corresponds to which row without inferring it from the text.","section":"Figure 2"},{"comment":"The rewrite rule is typeset inline with a line break; please use display math for readability.","section":"Section 3.4"},{"comment":"The term 'circular repeat' is used without definition before its first use; please define it when introducing Interleaved 1F1B.","section":"Section 2.2.1"}],"recommendation":"major_revision","confidential_remarks":"This is a well-scoped systems paper whose central design is plausible, but the lack of any semantic correctness validation is a serious gap for a claim that the system executes user-defined pipeline schedules. The comparison methodology also needs tightening. I would not accept without the authors adding correctness experiments and at least one controlled comparison; both are achievable within the manuscript's scope."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"You should know: JaxPP is the first thing I've seen that gives JAX users real control over pipeline schedules without forking the compiler. The pipeline_yield and accumulate_grads API is a genuine usability improvement over the SPMD encoding in GSPMD, and the single-controller MPMD runtime with async P2P is a plausible design. The reported numbers are internally consistent—457 TFLOPS/device, 44.6% over SPMD PP, 1.11x over FSDP, 92.87% weak scaling from 64 to 1024 GPUs—and I don't see arithmetic contradictions in Table 1 or Figure 8. I also credit the authors for publishing the NeMo comparison even though JaxPP only reaches 91.4% of NeMo throughput; that is honest.\n\nWhere the paper falls short is validation. There is no code release, no error bars, no loss curves, no gradient check against a reference, and the baselines use different parallelism configurations. For a systems paper in this area, that is not a minor omission; it is the core evidence that the system actually trains the model it claims to train.\n\nThe stress-test note about autodiff semantics points at a real gap. Section 3.4's loop-commuting rewrite is algebraically fine for associative gradient accumulation, but the paper never states the precondition that partial gradients are not consumed elsewhere in the loop body. And pipeline_yield is described as 'auto-differentiable' in Section 3.2, yet the paper neither gives a VJP rule nor explains how stage markers propagate through jax.value_and_grad. The stronger worry—that the throughput numbers could be for a program computing wrong updates—is not supported by anything in the text; there is no indication the transformations are unsound. But the absence of a correctness check means we simply don't know. The deadlock-freedom argument for arbitrary user schedules is likewise an algorithm plus an example, not a proof.\n\nBottom line: this is a paper for the pipeline-parallelism community, and it deserves a serious referee. The core idea is useful, the writing is clear, and the architecture is credible. I would send it to review with the explicit demand that the authors either release code or add a convergence/gradient-correctness section. Without that, the headline speedups are performance measurements in search of a validation.","headline":"A credible MPMD pipeline parallelism design for JAX with internally consistent speedups, but no correctness validation and no code make the headline numbers hard to trust.","tokens_in":21344,"tokens_out":3664,"would_cite":false,"duration_ms":31340,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"JaxPP lets JAX users run arbitrary pipeline schedules, beating SPMD by up to 11%.","keywords":["pipeline parallelism","MPMD runtime","SPMD","JAX","task graph","gradient accumulation","1F1B","distributed training"],"falsifier":"Run a small two-actor experiment where each actor sends two messages to the other in opposite orders across non-adjacent stages (a schedule that is dataflow-valid but whose inferred local send/recv sequence mismatches the remote order); if the run deadlocks or corrupts buffers, the inference rule is not sufficient for arbitrary schedules.","tokens_in":20365,"feed_emoji":"📈","tokens_out":7649,"duration_ms":59222,"temperature":0.7,"pith_summary":"JaxPP is a system that extends JAX's Single-Program-Multiple-Data (SPMD) model with user-driven pipeline parallelism, and this paper claims that doing so recovers performance that the SPMD model leaves on the table. The central claim is that by letting users write arbitrary gradient-accumulation schedules as lists of tasks (iteration, forward/backward, stage) and annotating stage boundaries with pipeline_yield, JaxPP can automatically split a training step into a distributed task graph, infer all cross-stage communication, and execute it with an MPMD runtime. On GPT-3 175B training, the pipeline implementation improves hardware utilization by up to 1.11x over the best SPMD configuration, reaches 457 TFLOPS per device, and maintains 92.87% weak-scaling efficiency from 64 to 1024 GPUs. The consequence for practitioners is that memory-saving and throughput-improving schedules like 1F1B and interleaved 1F1B become available in JAX with only small code changes and no hand-written communication.","feed_headline":"JaxPP runs 1F1B pipelines in JAX, beating SPMD by 11%","feed_subtitle":"JaxPP's MPMD runtime hits 457 TFLOPS/device on GPT-3 175B with 1F1B pipelines.","key_machinery":"The load-bearing object is the task graph produced from the gradient accumulation loop: each task is a triple (iteration, forward/backward, stage) assigned to a specific SPMD actor, and the schedule is just the ordered list of these triples per actor. JaxPP's inference pass turns that assignment into a concrete execution plan by walking tasks in topological order and placing asynchronous send and receive pairs immediately after the task that produces the data, which both avoids the deadlock that naive local send/recv ordering can cause and overlaps communication with computation. A buffer liveness pass then schedules deletions, and a loop-commuting rewrite for weight sharing delays the addition of partial gradients until after the loop so that embedding-sized tensors are not repeatedly transmitted. All local schedules are fused into one fused MPMD program per actor, so a training step costs one RPC per actor.","core_discovery":"JaxPP establishes that the SPMD encoding of pipeline parallelism is unnecessarily limiting, and that an MPMD task-based runtime on top of GSPMD can express and efficiently execute schedules the SPMD partitioner cannot. The system unrolls the gradient accumulation loop into a task graph, assigns tasks to long-lived SPMD actors according to a user-supplied schedule, and automatically infers the send and receive operations needed between (potentially non-adjacent) stages. It then fuses all task dispatches for a step into a single RPC per actor. On the paper's benchmarks this design delivers 457 TFLOPS/device for GPT-3 175B, a 44.6% step-time reduction over SPMD pipeline parallelism, up to 1.11x hardware utilization versus the best SPMD configuration, and 92.87% weak-scaling efficiency from 64 to 1024 GPUs.","pith_inferences":["A natural next test is whether novel schedules from the zero-bubble literature can be dropped into this API unchanged; the runtime's task-graph model suggests they would run without new inference or plumbing.","The 1.11x number is relative to the evaluated SPMD configurations; on models with heterogeneous stages or smaller per-stage compute, the gap could be larger because the SPMD encoding is restricted to homogeneous stages.","The loop-commuting rewrite for tied weights could be adopted by other pipeline runtimes as a general optimization for any parameter used in multiple stages.","If schedule validity were formalized, the compiler could reject an invalid user schedule at trace time instead of risking runtime deadlock."],"forward_implications":["JAX users can adopt 1F1B and interleaved 1F1B pipelines by adding a few annotations, without writing send/receive code or restructuring models as separate functions.","The 1.11x utilization gain over the best SPMD configuration implies that SPMD-only training leaves measurable throughput on the table for large models like GPT-3 175B.","Since weak-scaling efficiency matches a highly optimized FSDP baseline, pipeline parallelism under user schedules is not inherently less scalable.","Fusing dispatches into one RPC per actor per step keeps control-plane overhead small enough for long-running training loops."],"supporting_citations":[{"why":"Supplies the SPMD partitioning model JaxPP extends and the homogeneous-stage pipeline encoding it argues is limiting.","marker":"[Xu et al., 2021]"},{"why":"Defines GPipe, the baseline pipeline schedule whose memory cost motivates 1F1B.","marker":"[Huang et al., 2019]"},{"why":"Introduces the 1F1B schedule that improves memory and throughput, which JaxPP makes available to JAX users.","marker":"[Narayanan et al., 2019]"},{"why":"Introduces interleaved 1F1B and the Megatron-LM configuration used as a throughput benchmark.","marker":"[Narayanan et al., 2021]"},{"why":"Provides single-controller asynchronous dataflow runtime design that JaxPP shares implementation details with.","marker":"[Barham et al., 2022]"},{"why":"Demonstrates MPMD pipeline execution on JAX/XLA; JaxPP differentiates by user-controlled schedules rather than automated partitioning.","marker":"[Zheng et al., 2022]"},{"why":"Supplies the RPC layer and worker processes that the JaxPP runtime uses for orchestration.","marker":"[Moritz et al., 2018]"},{"why":"Activation rematerialization costs appear in the performance breakdown that explains the speedup over SPMD pipeline parallelism.","marker":"[Chen et al., 2016]"}],"fun_headline_variants":["JaxPP: MPMD cuts GPT-3 step time 44.6% over SPMD","JaxPP: 457 TFLOPS/device with MPMD pipelines","JaxPP: 92.87% scaling efficiency to 1024 GPUs","MPMD beats SPMD: JaxPP improves utilization 1.11x","JaxPP: 1F1B pipelines in JAX, 44.6% faster step time"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The runtime stays correct only if every user-supplied schedule, expressed as a list of (iteration, forward/backward, stage) tasks, respects the dataflow ordering that JaxPP's send/receive inference assumes; the paper gives an algorithm and an example but no formal proof that all valid schedules avoid deadlock.","fun_headline_variants_meta":{"raw":{"variants":["JaxPP: MPMD cuts GPT-3 step time 44.6% over SPMD","JaxPP: 457 TFLOPS/device with MPMD pipelines","JaxPP: 92.87% scaling efficiency to 1024 GPUs","MPMD beats SPMD: JaxPP improves utilization 1.11x","JaxPP: 1F1B pipelines in JAX, 44.6% faster step time"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000908,"raw_usage":{"total_tokens":3849,"prompt_tokens":835,"completion_tokens":3014,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":451,"completion_tokens_details":{"reasoning_tokens":2894}},"tokens_in":451,"tokens_out":3014,"duration_ms":19029,"temperature":1.0,"reasoning_tokens":2894,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T12:17:14.119812+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run a small two-actor experiment where each actor sends two messages to the other in opposite orders across non-adjacent stages (a schedule that is dataflow-valid but whose inferred local send/recv sequence mismatches the remote order); if the run deadlocks or corrupts buffers, the inference rule is not sufficient for arbitrary schedules.","supporting_citations":[{"cited_title":"X., Lee, H., Ngiam, J., Le, Q","cited_arxiv_id":null,"evidence_quote":"Defines GPipe, the baseline pipeline schedule whose memory cost motivates 1F1B."},{"cited_title":"E., Thekkath, C","cited_arxiv_id":null,"evidence_quote":"Provides single-controller asynchronous dataflow runtime design that JaxPP shares implementation details with."},{"cited_title":"P., Gonzalez, J","cited_arxiv_id":null,"evidence_quote":"Demonstrates MPMD pipeline execution on JAX/XLA; JaxPP differentiates by user-controlled schedules rather than automated partitioning."},{"cited_title":"I., and Stoica, I","cited_arxiv_id":null,"evidence_quote":"Supplies the RPC layer and worker processes that the JaxPP runtime uses for orchestration."}],"review_version":1}