Pith. sign in

REVIEW 4 major objections 5 minor 33 references

Scaling Deep Learning Training with MPMD Pipeline Parallelism

T0 review · 4 major / 5 minor · reviewed 2026-08-11 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read JaxPP lets JAX users run arbitrary pipeline schedules, beating SPMD by up to 11%.

desk verdict 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. read the letter →

arxiv 2412.14374 v1 pith:ZWRUDBOA submitted 2024-12-18 cs.DC cs.LGcs.PL

Anxhelo Xhebraj , Sean Lee , Hanfeng Chen , Vinod Grover This is my paper · ORCID
classification cs.DCcs.LGcs.PL
keywords pipelineparallelismMPMDruntimeSPMDJAXtaskgraphgradientaccumulation1F1Bdistributedtraining
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

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.

What carries the argument

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.

What would settle it

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.

Watch

Extended reading notes

Core claim

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.

Load-bearing premise

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.

Editorial extensions

If this is right

  • 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.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • 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.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 5 minor

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.

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 (4)
  1. [Section 3.4 and Section 3.1] 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.
  2. [Section 3.2 and Figure 4] 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.
  3. [Section 4.2] 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.
  4. [Section 5 and Table 1] 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.
minor comments (5)
  1. [Section 5.2] 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.
  2. [Section 5.2] 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.
  3. [Figure 2] 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.
  4. [Section 3.4] The rewrite rule is typeset inline with a line break; please use display math for readability.
  5. [Section 2.2.1] The term 'circular repeat' is used without definition before its first use; please define it when introducing Interleaved 1F1B.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity found: performance claims are external measurements; the loop-commuting rewrite is a semantic-preserving transformation, and the missing correctness validation is a correctness risk, not a circular step.

full rationale

This is a systems and measurement paper; its central claims are empirical comparisons against independently implemented baselines (JAX FSDP, JAX SPMD PP, NeMo) on standard models (GPT-3 175B, Llama2 70B) with reported TFLOPS/device and step times. I examined the derivation chain for circular reductions. The only program transformation with an equation-like statement is the loop-commuting rewrite in Section 3.4, which replaces a sum of per-microbatch partial gradients aggregated in loop-carried state with per-use partial-gradient accumulators plus a final addition; this is a compiler rewrite justified by associativity and commutativity of the accumulation operators, not a fitted parameter and not a prediction derived from the result. The scheduling API (pipeline_yield, accumulate_grads) defines what JaxPP does, but the evaluation does not use the system to predict its own schedule; the 1F1B and interleaved schedules are standard external schedules, and the SPMD baseline is an external system. There are no self-citations by the present authors and no uniqueness theorem imported from prior work. The absence of a correctness or loss-curve validation, noted in the skeptic's attack, is a substantive correctness risk, but it is not circularity: even if the transformed program computed the wrong updates, the throughput numbers would still be measurements of that program rather than quantities forced by the paper's inputs. The selection of JaxPP's own configuration, such as circular repeat size and microbatch size, is a tuning choice that could affect fairness, not a case of fitting a parameter and then predicting the same data. I therefore find no circular step.

Assumptions & free parameters 2 free parameters · 4 assumptions · 0 invented entities

The central performance claims depend on benchmark configuration choices, on the correctness of XLA/GSPMD, and on the semantic preservation of unproven compiler rewrites. No code, proof, or artifact is provided.

free parameters (2)
  • Circular repeat size = 6
    Hand-chosen for JaxPP experiments in Figures 6-8; throughput varies with this value and the choice is not derived or justified beyond a tradeoff sweep.
  • Microbatch count (GA) and microbatch size = GPT-3 175B: GA=32, MBS=2; Llama2 70B: GA=16, MBS=2
    Selected per benchmark in Table 1; the reported TFLOPS figures depend on these settings, with sensitivity only partially explored in Figures 6 and 7.
assumptions (4)
  • domain assumption GSPMD/XLA SPMD partitioning correctly handles sharding annotations and inserts all necessary collectives.
    JaxPP builds on GSPMD; if XLA's SPMD partitioner is incorrect, all JaxPP tasks inherit that error (Sections 2.1 and 3).
  • ad hoc to paper A user-specified pipeline schedule is a valid dataflow partial order over stages and iterations, and JaxPP's topological send/receive inference prevents deadlock for such schedules.
    Core feature (Section 4.2) is assumed correct across all schedules; only examples are shown, no proof or validation over schedule space.
  • domain assumption Backward computation for a layer is scheduled on the same actor as its forward computation.
    Placement inference in Section 3.3 states: 'We assume that the loop schedule maps backward computations to the same actor of the corresponding forward computation.'
  • ad hoc to paper The gradient accumulation API restriction (no dependencies from later stages of the current iteration to earlier stages of the previous iteration) is sufficient to make the unrolled task graph schedule-safe.
    Section 3.1 states this API restriction is intentional, but correctness is not formally proven.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Scaling Deep Learning Training with MPMD Pipeline Parallelism." pith.science (2026). https://pith.science/paper/ZWRUDBOA

@misc{pith2026241214374,
  author       = {Pith},
  title        = {Pith review of: Scaling Deep Learning Training with MPMD Pipeline Parallelism},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/ZWRUDBOA}},
  note         = {Machine review of arXiv:2412.14374}
}
abstract

We present JaxPP, a system for efficiently scaling the training of large deep learning models with flexible pipeline parallelism. We introduce a seamless programming model that allows implementing user-defined pipeline schedules for gradient accumulation. JaxPP automatically distributes tasks, corresponding to pipeline stages, over a cluster of nodes and automatically infers the communication among them. We implement a MPMD runtime for asynchronous execution of SPMD tasks. The pipeline parallelism implementation of JaxPP improves hardware utilization by up to $1.11\times$ with respect to the best performing SPMD configuration.

Figures

Figures reproduced from arXiv: 2412.14374 by the authors.

Figure 1
Figure 1. Configurable Parallelism Through Named Axes in JAX (Bradbury et al., 2018) Top left (1a): Model implementation where array axes are annotated with logical names. Bottom left (1b): Partitioning specification mapping logical axis names to mesh axes. Right (1c): Two parallel instantiation, data-parallel on the top with mesh shape [("data", 2) ("model", 1)] while a tensor-parallel implementation at the bottom when the m… view at source ↗
Figure 2
Figure 2. Comparison between GPipe and 1F1B. In GPipe, at any time, all pipeline-parallel groups perform the same compu￾tation. Bubbles are implemented as redundant discarded com￾putation (gray Z blocks). In 1F1B, all groups perform different computations. degree of circular repeat increases, stages become smaller, enabling finer-grained scheduling. This approach improves throughput, but introduces additional communication ov… view at source ↗
Figure 3
Figure 3. System Overview. The left box shows the code in the driver process describing the computation and annotating pipeline stage boundaries. Auto-differentiation produces additional stages corresponding to the “backward” computations for the gradients. The user specifies a mapping of stages to SPMD actors and a schedule for the loop. Each call to the step_fn function schedules tasks 3 JAXPP OVERVIEW We now describe JaxPP… view at source ↗
Figures from the paper (6 more)
Figure 4
Figure 4. Figure 4: Training loop in JaxPP 3.2 Stage Marking The user specifies the start and end of “logical stages” through pipeline_yield. Any computation arising before the first call to pipeline_yield is implicitly scheduled on the first stage, with each call “opening” a new stage. I…
Figure 5
Figure 5. Figure 5: Inference of send and receive operations based on uses and definitions in the task graph. Instead, JaxPP iterates over the tasks in their topological order and schedules asynchronous send and receive pairs im￾mediately after the corresponding task has produced the data…
Figure 6
Figure 6. Figure 6: Performance of GPT-3 175B training on 64 GPUs with global batch size of 128 instances across various configurations for interleaving/circular repeat and microbatch size. be exacerbated when using configurations that try to reduce pipeline bubbles, such as: (1) slicing …
Figure 8
Figure 8. Figure 8: JaxPP’s weak scaling in comparison to a highly opti￾mized JAX FSDP implementation. 1.11× over JAX’s FSDP. JaxPP achieves 91.4% throughput of NeMo’s pipeline parallelism while being entirely model￾agnostic. When training Llama2 70B on 8 DGX H100 nodes (8 GPUs), JaxPP de…
Figure 9
Figure 9. Figure 9: Performance comparison between SPMD pipeline paral￾lelism, JaxPP, and NeMo on GPT-3 175B and Llama2 70B. 5.3 Performance Breakdown To understand the sources of performance gains achieved by JaxPP over SPMD pipeline parallelism on GPT-3 175B, we present [PITH_FULL_IMAG…
Figure 10
Figure 10. Figure 10: Overhead of JAX SPMD PP compared to JaxPP. Re￾materialization cost and asynchronous point-to-point send and receive operations account for the majority of the performance differences. 6 RELATED WORK There are numerous works to facilitate scaling the training of large …

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

33 extracted references · 14 canonical work pages

  1. [1]

    URL https://www.top500.org/system/180239

    EOS NVIDIA DGX SuperPod - NVIDIA DGX H100 , 2024. URL https://www.top500.org/system/180239

    work page 2024

  2. [2]

    PartIR: Composing SPMD Partitioning Strategies for Machine Learning

    Alabed, S., Chrzaszcz, B., Franco, J., Grewe, D., Maclaurin, D., Molloy, J., Natan, T., Norman, T., Pan, X., Paszke, A., Rink, N. A., Schaarschmidt, M., Sitdikov, T., Swietlik, A., Vytiniotis, D., and Wee, J. PartIR : Composing SPMD Partitioning Strategies for Machine Learning , January 2024. URL http://arxiv.org/abs/2401.11202

    work page Pith review arXiv 2024

  3. [3]

    E., Thekkath, C

    Barham, P., Chowdhery, A., Dean, J., Ghemawat, S., Hand, S., Hurt, D., Isard, M., Lim, H., Pang, R., Roy, S., Saeta, B., Schuh, P., Sepassi, R., Shafey, L. E., Thekkath, C. A., and Wu, Y. Pathways: Asynchronous Distributed Dataflow for ML . In Proceedings of Machine Learning and Systems 2022, MLSys 2022, Santa Clara , CA , USA , August 29 - September 1, 2...

    work page 2022

  4. [4]

    J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne , S., and Zhang, Q

    Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne , S., and Zhang, Q. JAX : Composable transformations of Python + NumPy programs, 2018. URL http://github.com/google/jax

    work page 2018

  5. [5]

    Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss , A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A....

    work page 1901

  6. [6]

    Training deep nets with sublinear memory cost

    Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. CoRR, abs/1604.06174, 2016. URL http://arxiv.org/abs/1604.06174

    arXiv 2016

  7. [7]

    cudnn: Efficient primitives for deep learning, 2014

    Chetlur, S., Woolley, C., Vandermersch, P., Cohen, J., Tran, J., Catanzaro, B., and Shelhamer, E. cudnn: Efficient primitives for deep learning, 2014. URL https://arxiv.org/abs/1410.0759

    arXiv 2014

  8. [8]

    Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari , G., Yin, P., Duke, T., Levsk...

    arXiv 2022

Show all 33 references
  1. [9]

    An Image is Worth 16x16 Words : Transformers for Image Recognition at Scale

    Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An Image is Worth 16x16 Words : Transformers for Image Recognition at Scale . In 9th International Confere...

    2021

  2. [10]

    Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., Goyal, A., Hartshorn, A., Yang, A., Mitra, A., Sravankumar, A., Korenev, A., Hinsvark, A., Rao, A., Zhang, A., Rodriguez, A., Gregerson, A., Spataru, A., Rozie...

    2024 arXiv

  3. [11]

    Switch Transformers : Scaling to Trillion Parameter Models with Simple and Efficient Sparsity , June 2022

    Fedus, W., Zoph, B., and Shazeer, N. Switch Transformers : Scaling to Trillion Parameter Models with Simple and Efficient Sparsity , June 2022. URL http://arxiv.org/abs/2101.03961

    2022 arXiv

  4. [12]

    NeMo: a toolkit for Conversational AI and Large Language Models

    Harper, E., Majumdar, S., Kuchaiev, O., Jason, L., Zhang, Y., Bakhturina, E., Noroozi, V., Subramanian, S., Nithin, K., Jocelyn, H., Jia, F., Balam, J., Yang, X., Livne, M., Dong, Y., Naren, S., and Ginsburg, B. NeMo: a toolkit for Conversational AI and Large Language Models ....

  5. [13]

    The Hardware Lottery , September 2020

    Hooker, S. The Hardware Lottery , September 2020. URL http://arxiv.org/abs/2009.06489

    2020 arXiv

  6. [14]

    DISTMM : Accelerating distributed multimodal model training

    Huang, J., Zhang, Z., Zheng, S., Qin, F., and Wang, Y. DISTMM : Accelerating distributed multimodal model training. In 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24), pp.\ 1157--1171, Santa Clara, CA, April 2024. USENIX Association. ISBN 978-1-9...

    2024

  7. [15]

    X., Lee, H., Ngiam, J., Le, Q

    Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M. X., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., and Chen, Z. GPipe : Efficient Training of Giant Neural Networks using Pipeline Parallelism . In Wallach, H. M., Larochelle, H., Beygelzimer, A., d'Alch \'e -Buc , F., Fox...

    2019

  8. [16]

    Megascale: Scaling large language model training to more than 10,000 gpus

    Jiang, Z., Lin, H., Zhong, Y., Huang, Q., Chen, Y., Zhang, Z., Peng, Y., Li, X., Xie, C., Nong, S., Jia, Y., He, S., Chen, H., Bai, Z., Hou, Q., Yan, S., Zhou, D., Sheng, Y., Jiang, Z., Xu, H., Wei, H., Zhang, Z., Nie, P., Zou, L., Zhao, S., Xiang, L., Liu, Z., Li, Z., Jia, X....

    2024

  9. [17]

    Breadth- First Pipeline Parallelism , November 2022

    Lamy-Poirier , J. Breadth- First Pipeline Parallelism , November 2022. URL http://arxiv.org/abs/2211.05953

    2022 arXiv

  10. [18]

    Mlir: Scaling compiler infrastructure for domain specific computation

    Lattner, C., Amini, M., Bondhugula, U., Cohen, A., Davis, A., Pienaar, J., Riddle, R., Shpeisman, T., Vasilache, N., and Zinenko, O. Mlir: Scaling compiler infrastructure for domain specific computation. In 2021 IEEE/ACM International Symposium on Code Generation and Optimizat...

    2021

  11. [19]

    GShard : Scaling Giant Models with Conditional Computation and Automatic Sharding , June 2020

    Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. GShard : Scaling Giant Models with Conditional Computation and Automatic Sharding , June 2020. URL http://arxiv.org/abs/2006.16668

    2020 arXiv

  12. [20]

    Torchtitan: One-stop pytorch native solution for production ready llm pre-training, 2024

    Liang, W., Liu, T., Wright, L., Constable, W., Gu, A., Huang, C.-C., Zhang, I., Feng, W., Huang, H., Wang, J., Purandare, S., Nadathur, G., and Idreos, S. Torchtitan: One-stop pytorch native solution for production ready llm pre-training, 2024. URL https://arxiv.org/abs/2410.06511

    2024 arXiv

  13. [21]

    nnScaler : Constraint-Guided parallelization plan generation for deep learning training

    Lin, Z., Miao, Y., Zhang, Q., Yang, F., Zhu, Y., Li, C., Maleki, S., Cao, X., Shang, N., Yang, Y., Xu, W., Yang, M., Zhang, L., and Zhou, L. nnScaler : Constraint-Guided parallelization plan generation for deep learning training. In 18th USENIX Symposium on Operating Systems D...

    2024

  14. [22]

    I., and Stoica, I

    Moritz, P., Nishihara, R., Wang, S., Tumanov, A., Liaw, R., Liang, E., Elibol, M., Yang, Z., Paul, W., Jordan, M. I., and Stoica, I. Ray: A distributed framework for emerging AI applications. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pp...

    2018

  15. [23]

    R., Ganger, G

    Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. PipeDream : Generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles , pp.\ 1--15, Huntsv...

    2019

  16. [24]

    Efficient large-scale language model training on GPU clusters using megatron- LM

    Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Zaharia, M. Efficient large-scale language model training on GPU clusters using megatron- LM . In Proceedings ...

    2021

  17. [25]

    Efficiently Scaling Transformer Inference , November 2022

    Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently Scaling Transformer Inference , November 2022. URL http://arxiv.org/abs/2211.05102

    2022 arXiv

  18. [26]

    Zero bubble (almost) pipeline parallelism

    Qi, P., Wan, X., Huang, G., and Lin, M. Zero bubble (almost) pipeline parallelism. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=tuzTN0eIO5

    2024

  19. [27]

    Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters

    Rasley, J., Rajbhandari, S., Ruwase, O., and He, Y. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD '20, pp.\ 3505–35...

    2020

  20. [28]

    Megatron- LM : Training Multi-Billion Parameter Language Models Using Model Parallelism , March 2020

    Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron- LM : Training Multi-Billion Parameter Language Models Using Model Parallelism , March 2020. URL http://arxiv.org/abs/1909.08053

    2020 arXiv

  21. [29]

    Y., Bernauer, J., Song, X., Shoeybi, M., He, Y., Houston, M., Tiwary, S., and Catanzaro, B

    Smith, S., Patwary, M., Norick, B., LeGresley, P., Rajbhandari, S., Casper, J., Liu, Z., Prabhumoye, S., Zerveas, G., Korthikanti, V., Zhang, E., Child, R., Aminabadi, R. Y., Bernauer, J., Song, X., Shoeybi, M., He, Y., Houston, M., Tiwary, S., and Catanzaro, B. Using DeepSpee...

    2022 arXiv

  22. [30]

    Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N....

    2023 arXiv

  23. [31]

    GSPMD : General and Scalable Parallelization for ML Computation Graphs , December 2021

    Xu, Y., Lee, H., Chen, D., Hechtman, B., Huang, Y., Joshi, R., Krikun, M., Lepikhin, D., Ly, A., Maggioni, M., Pang, R., Shazeer, N., Wang, S., Wang, T., Wu, Y., and Chen, Z. GSPMD : General and Scalable Parallelization for ML Computation Graphs , December 2021. URL http://arx...

    2021 arXiv

  24. [32]

    P., Gonzalez, J

    Zheng, L., Li, Z., Zhang, H., Zhuang, Y., Chen, Z., Huang, Y., Wang, Y., Xu, Y., Zhuo, D., Xing, E. P., Gonzalez, J. E., and Stoica, I. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning . In Aguilera, M. K. and Weatherspoon, H. (eds.), 16th U...

    2022

  25. [33]

    write newline

    " write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 gl...

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.