Piper is a PyTorch distributed training system that separates strategy specification from runtime execution, letting users express composed strategies such as PP + DP/EP + ZeRO without hard-coding each combination into the framework. Its core abstraction is a global training DAG that explicitly represents compute, communication, placement, streams, and ordering constraints, enabling joint scheduling across parallelism dimensions. The main result is flexibility without giving up baseline performance: Piper matches common strategies such as ZeRO-1 and 1F1B while enabling DualPipe-like schedules with 6-30% throughput gains, plus PP x ZeRO combinations that other evaluated systems either cannot run or do not memory-optimize correctly.

Abstract

1. Distributed training needs composed strategies Large-scale training increasingly combines data, pipeline, expert parallelism, and memory optimizations such as ZeRO, but deployed foundation-model systems often depend on experts who manually design both high-level sharding and low-level execution.

2. Piper decouples strategy from runtime Piper exposes model annotations and scheduling directives that transform a unified IR, described as “a unified global training DAG that represents all computation and communication”. The runtime then executes per-device plans without being specialized to a fixed strategy.

3. Performance parity plus composed-strategy gains Piper maintains parity on widely used strategies such as ZeRO and enables further gains through joint scheduling of communication and compute in strategies like DeepSeek-V3’s DualPipe.

1 Introduction

1. High-level and low-level strategy jointly determine throughput A training strategy decomposes into a high-level parallelism plan and a per-device low-level execution plan. The high-level plan sets lower bounds on per-device memory, compute, and communication load; the low-level plan determines how close execution gets to that bound.

2. Existing systems either require experts or expose limited strategy spaces Human-engineered systems such as DeepSeek-V3’s DualPipe require custom codesign of PP, EP, and intra-GPU resource usage. General frameworks such as Megatron, DeepSpeed, and TorchTitan expose knobs, but they tend to dispatch each parallelism dimension independently, which makes composed scheduling hard.

3. DualPipe motivates local microbatch overlap DualPipe shares a GPU between forward and backward microbatches to hide EP communication, but that conflicts with framework assumptions that each microbatch owns the full GPU. Compiler systems like JAX/XLA offer generic tensor placement but lack arbitrary PP scheduling and user control over streams/resources.

4. Piper’s goal is extensibility The target is “to build a system that minimizes the effort needed to specify and implement an arbitrary distributed training strategy”. Piper does this with a scheduling API, a global DAG IR, and a strategy-agnostic runtime.

5. Reported contributions are API, IR, runtime, and evaluation Piper contributes a user scheduling interface, a unified global training DAG for joint compute/communication scheduling, an efficient distributed runtime, and an evaluation showing parity on common strategies plus better performance and memory efficiency on composed ones.

2 Background

1. DP and ZeRO reduce redundant state differently In DP, every worker stores a full model replica, computes local gradients, and averages gradients with allreduce or allgather/reduce-scatter. ZeRO reduces redundant state by sharding optimizer state, gradients, and/or weights, rematerializing state with allgather and resharding after use.

2. TP, EP, and CP put communication on the critical path Tensor, expert, and context parallelism shard weights or activations and require collectives according to the sharding plan. Unlike DP, their collectives execute directly on the batch’s critical path.

3. PP needs microbatch schedules Pipeline parallelism shards layers across workers, then uses microbatches to overlap execution across devices. Its performance depends on bubbles, communication overhead, and schedule design; PP is naturally MPMD because ranks run different operations in different orders.

4. Composed strategies must handle heterogeneous submodules The DualPipe-like example combines PP across layers, EP for expert layers, and DP for non-expert attention layers. This matters more as models become heterogeneous, such as Qwen3-Next’s diverse attention layers and multimodal models with modality- specific components.

3 Challenges

1. High-level strategy must include intra-device parallelism As communication overhead grows, scheduling must include communication/compute overlap within a GPU. DualPipeV shows that overlapping forward and backward microbatches helps, but overlapping forward-forward or backward-backward can introduce bubbles that erase gains.

2. Low-level scheduling faces resource contention Multiple parallelism dimensions introduce operations that compete for GPU memory, network bandwidth, streams, and communicators. In the DualPipe-style example, a DP allreduce can interfere with EP all-to-all; measured EP communication slowed by 1.46x under background DP allreduces.

3. Communication partitioning is a tradeoff Using separate streams can avoid sequential delay but create bandwidth interference; putting DP and EP communication on the same stream can delay critical all-to-alls; partitioning DP allreduces into smaller pieces may reduce interference but can lower communication efficiency.

4. Runtime must be strategy-agnostic but still efficient The runtime must avoid CPU scheduling overhead, allocate memory/streams/communicators efficiently, and jointly schedule communication from different parallelism dimensions. Near capacity, PyTorch’s memory allocator can stall waiting for in-flight work, so ZeRO can improve throughput as well as memory.

5. Existing frameworks have incomplete PP x ZeRO behavior The paper reports that general frameworks do not fully support all ZeRO levels with PP, likely because ZeRO hooks interact poorly with PP’s repeated layer execution across microbatches. Piper’s unified DAG supports 3-8x larger batch sizes in the PP x ZeRO case study.

4 Design

1. Piper has compiler and runtime components The compiler translates annotated models and user schedules into a distributed execution plan; the runtime executes that plan on distributed workers. The design avoids hard-coding specific combinations of parallelism strategies.

2. The IR is a global training DAG The DAG contains Chunk nodes for compute and Comm nodes for point-to-point or collective communication. Nodes have device placement, stream assignment, exec functions, and edges for data dependencies.

3. Annotations define schedulable regions Users annotate meaningful model regions, such as PP stages or expert MLP blocks. Piper converts these into Chunks that can later be placed, replicated, sharded, split, or ordered.

4. Scheduling directives transform the DAG The main directives are Place, Replicate, Shard, Split, and Order. Filters select Chunks by dimensions such as PP, EP, MB, or PASS, where PASS can distinguish forward, backward, backward-input, and backward-weight phases.

5. Place and Replicate insert communication Place assigns nodes to devices and inserts send/recv at cross-device boundaries. Replicate synchronizes gradients with allreduce by default or reduce-scatter when gradient sharding is enabled, with optional streams and bucket sizes.

6. Shard, Split, and Order express expert parallelism, microbatching, and schedules Shard inserts all-to-all before and after matched Chunks, enabling EP when combined with Replicate. Split duplicates a sub-DAG into microbatches. Order adds temporal dependencies and can express overlapped sub-DAGs through nested filter lists.

7. DualPipe can be specified concisely The simplified DualPipe schedule uses streams for PP, EP, and DP communication; places two PP stages across devices; replicates non-expert chunks; shards expert chunks; splits into two microbatches; and orders microbatches so one PP stage overlaps forward and backward work.

4.2 Piper Compiler

1. Compilation starts from TorchDynamo graph capture Piper extracts a PyTorch fx.Graph, initially treating all tensor operators for a forward-backward pass as one Chunk. Annotation boundaries split this graph into subgraphs, each becoming a Chunk’s forward exec function; PyTorch autograd supplies backward execution.

2. Model-state buckets are tied to Chunks The compiler uses tensor operator dependencies to associate parameters, gradients, and optimizer state with each Chunk. Currently, a state bucket can only be associated with Chunks that share the same placement.

3. Directives become graph rewrites The compiler mechanically applies user scheduling directives, inserting Comm nodes such as all-to-all and allreduce. It then removes unnecessary parameter allgathers or gradient reduce-scatters when consecutive Chunks use the same state bucket.

4. The final DAG includes data and temporal constraints The compiler output is a distributed DAG with explicit communication, placement, streams, model data dependencies, and Order dependencies. Missing stream assignments default to the compute stream.

4.3 Piper Runtime

1. The centralized scheduler creates per-device partial orders The scheduler decomposes the global DAG into one unique sub-DAG per PP rank, with workers sharing a PP rank executing SPMD. Tasks on the same stream are totally ordered; tasks on different streams are ordered only when data or temporal dependencies require it.

2. Independent tasks are scheduled by a simple dependency heuristic For overlapping sub-DAGs, Piper creates one queue per stream, repeatedly chooses a ready task with the most downstream dependencies, and appends it to that task’s stream queue. This works well for symmetric DualPipe-style forward/backward overlap.

3. Workers manage streams, communicators, and memory Each Ray actor loads its model weights and dispatches Chunks/Comms according to the scheduler’s plan. Cross-stream dependencies use CUDA events and stream waits, while independent tasks proceed concurrently.

4. Worker dispatch prioritizes communication carefully Piper prioritizes send communication first, defers receive communication to reduce P2P interference, and among other communication tasks prioritizes critical-path operations over reductions. Deterministic ordering avoids collective deadlocks.

5. Separate P2P streams and communicators reduce PP bubbles Piper uses separate send and receive streams plus separate communicators. This avoids requiring a single global P2P order and only requires that downstream workers process data in the same order upstream workers produce it.

6. Memory management explicitly controls state and activations Piper allocates flat buffers for parameter and gradient buckets, stores persistent sharded state for ZeRO, materializes temporary full buffers when needed, and releases buffers after the last consumer completes. Intermediate activations are freed once their final downstream task is scheduled.

5 Implementation

1. Piper is a TorchDynamo backend It hooks into arbitrary PyTorch code but currently requires fully traceable models so it can partition fx.Graphs at compile time.

2. Annotations are Python context managers Annotations attach metadata during execution; graph capture records that metadata and uses it to segment the graph into Chunks.

3. Runtime execution uses Ray The compiler and centralized scheduler run in a driver process; each worker is a Ray actor.

6 Evaluation

1. Evaluation compares against three general-purpose frameworks Piper is evaluated against Megatron-LM 0.18.0, DeepSpeed 0.18.9, and TorchTitan 0.2.2 on 4 AWS EC2 NVIDIA 8xA100 nodes with NVLink and EFA.

2. Common PP schedules are competitive For Qwen3 1B with PP-8 x DP/EP-4 across 32 A100s and Qwen3 9B with PP-4 x DP/EP-4 across 16 A100s, Piper supports 1F1B and interleaved 1F1B. TorchTitan schedule builders were adapted to Piper in 29 LoC and 38 LoC.

3. TorchTitan loses performance from memory and stream behavior TorchTitan’s larger DP memory footprint causes PyTorch CUDA allocator delays, and its interleaved schedule is 14% worse than its 1F1B schedule because sends and receives share one stream. Piper-interleaved-1F1B is 5% higher throughput than Piper- 1F1B due to lower memory use and separate send/receive streams.

4. Megatron benefits from fused kernels Megatron’s single-device stage forward takes about 30 ms for Qwen3 1B, versus 40 ms in Piper. Piper’s Chunk abstraction is orthogonal to fused kernels, so those optimizations could be integrated.

5. PP x ZeRO support differs sharply DeepSpeed and Megatron support PP x ZeRO-1 but not ZeRO-2/3 with PP. TorchTitan claims ZeRO-2/3 support, but gradient and weight states do not reshard between all microbatches, so memory savings are much smaller. Piper supports all PP x ZeRO combinations.

6. ZeRO-1 throughput is similar across systems On Qwen3 1B with DP-2 across 2 A100s, ZeRO-1 throughput is Piper 8641 tokens/s ±701, TorchTitan 8637 ±977, DeepSpeed 9352 ±52, and Megatron 9942 ±1106.

7. Piper gets the expected ZeRO memory savings For Qwen3 9B with 8-way PP and 4-way DP across 32 A100s, PP x ZeRO-1 variants OOM even at the smallest batch size. TorchTitan OOMs at batch size 8 for ZeRO-2 and 16 for ZeRO-3, while Piper runs up to batch size 32 for ZeRO-2 and 40 for ZeRO-3, corresponding to 8x and 3.3x larger batch sizes.

8. DualPipeV is easier to express and faster in Piper The DualPipeV schedule builder from TorchTitan was adapted to Piper in 63 LoC using Order for microbatch overlap. On Qwen3 1B, Piper-DualPipeV improves 13% over Piper-1F1B, while TorchTitan-DualPipeV improves only 3% over TorchTitan-1F1B.

9. Qwen3 9B DualPipeV shows gains where TorchTitan OOMs TorchTitan could not run the 9B DualPipeV case due to OOM. Piper-DualPipeV improves 10% over Piper interleaved 1F1B and 6% over Megatron interleaved 1F1B, while Megatron still benefits from fused kernels that Piper could potentially incorporate.

10. Scalability is reasonable Piper scales Qwen3 1B across 2-, 4-, and 8-way PP and 2- and 4-way DP, with global batch size scaled linearly. Figure 9 shows throughput increasing close to the theoretical curve, reaching roughly 390k tokens/s for PP=8, DP=4.

7 Related Work

1. General frameworks expose fixed strategy sets Megatron, DeepSpeed, and TorchTitan implement common DP/ZeRO, TP/EP/CP, and PP mechanisms, but each dimension dispatches eagerly with little synchronization across dimensions, making joint scheduling of memory and communication bandwidth difficult.

2. Compiler systems inspire Piper but lack low-level user scheduling JAX/XLA and GSPMD-inspired systems provide tensor sharding annotations and compiler-inserted communication, but are limited in arbitrary PP scheduling and do not expose per-device resources such as GPU streams to users.

3. DSLs and scheduling systems overlap with Piper’s goals CoCoNeT, AutoSP, DynaFlow, Slapo, and TVM relate through annotation, compiler support, or schedule languages. Piper extends these ideas to distributed tensor programs with explicit intra-device parallelism and a flexible runtime.

4. Auto-parallelism systems need broader execution support Many auto-parallel systems restrict the search space to remain tractable and may miss strategies such as full communication/compute overlap. Piper aims to serve as a common runtime for such systems by supporting richer composed strategies through a unified IR.

5. nnScaler is closest but lacks intra-device parallelism nnScaler allows generic constraints similar to Piper directives, but does not support intra-device parallelism. Future Piper work could integrate profile-guided and dynamic approaches from systems such as DeepCompile and Tessera.

8 Conclusion

1. Piper’s central contribution is explicit strategy representation Piper decouples distributed execution strategy from model and runtime using a unified global training DAG. By making placement, granularity, and ordering explicit, it can express PP, DP, EP, and ZeRO combinations without runtime specialization.

2. The interface supports both expert schedules and automated search Piper is positioned as a programmable substrate for “both expert-designed schedules and future automated search over a rich space of composed strategies”.

Design

• (1) Functionality specifications, which use concepts like Hoare logic (pre/post-conditions) and invariants to describe the behavior of individual modules.
  • Specification methods
    • pre- and post-conditions, following Hoare Logic, define the contractual obligations for each function by specifying the required state before execution and the guaranteed state upon completion.
    • invariants are properties that must hold true across all state transitions, ensuring the module’s integrity
    • a system algorithm outlines the high-level logic for how a function should achieve its state transition, guiding the LLM’s implementation strategy. A high-level intent is often sufficient. The intent can be regarded as a lightweight system algorithm that, expressed in natural language, guides the LLM in generating the desired implementation.
• (2) Modularity specifications, which decompose the system into distinct components and use a rely-guarantee discipline [21] to ensure they can be composed correctly and developed independently.
  • (1) module implementations must respect their declared dependencies (Rely)
  • (2) provide guarantees about their behavior (Guarantee)
  • (3) compose through logical implication of these contracts
• (3) Concurrency specifications, which explicitly define locking protocols and other concurrency-related behaviors that are notoriously difficult for LLMs to infer on their own.
  • During code generation, our toolchain first directs the LLM to generate a correct sequential version of the code, focusing only on the primary functionality.
  • In a second pass, using the dedicated concurrency specification to instrument the code with the required locking and other concurrent behaviors

The SYSSPEC Toolchain

• SpecCompiler
  • Two-phase prompting - First phase, it generates a correct sequential implementation of the module, focusing only on its core functionality. - Second phase, it uses the dedicated concurrency specification to instrument this sequential code with the necessary locking and concurrent behaviors.
  • Retry-with-feedback loop used within each phase: a CodeGen agent generates the implementation, and a separate, reasoning-focused SpecEval agent reviews the output against the specification. If the SpecEval agent identifies a flaw, it does not simply report failure; instead, it generates specific, actionable feedback (e.g., “The case where function foo() fails is not handled”). This feedback is then appended to the original prompt, and the CodeGen agent retries. This refinement cycle continues until the generated code satisfies the specification or an attempt-limit is reached.

Why $frac{1}{\sqrt(dk)}$

X, Y are two independent random variable. E[X] = 0, E[Y] = 0. Var[X] = 1, Var[Y] = 1.

$Var[XY] = E[X]^2 \dot Var[Y] + E[Y]^2 \dot Var[X] + Var[X] \dot Var[Y]$

Since E[x] and E[Y] = 0, $Var[XY] = Var[X] * Var[Y]$

For $C = Q * K^T$, $C_{ij} = \sum_{k=1}^{N} Q_{ik} * K_{kj}$. Each entry is a sum of N products of random variables.

$Var[QK^T] = dk * 1 * 1$, where N = dk, Var[Q] = 1, Var[k] = 1.

We want to find a value to scale the $QK^T$ so that Var[a \dot QK^T] = 1$

$Var[aX] = a^2 * Var[X]$, $Var[aQK^T] = a^2 * Var[QK^T] = a^2 * dk = 1$. $a = \frac{1}{\sqrt(dk)}$

Motivation

Contexts should function not as concise summaries, but as comprehensive, evolving playbooks—detailed, inclusive, and rich with domain insights.

Component

• Generator: produces reasoning trajectories.
• Reflector: distills concrete insights from successes and errors
  • separates evaluation and insight extraction from curation, improving context quality and downstream performance.
• Curator: integrates these insights into structured context updates.

Incremental Delta Updates

• context as a collection of structured, itemized bullets, rather than a single monolithic prompt.
  • (1) metadata: including a unique identifier and counters tracking how often it was marked helpful or harmful.
  • (2) content: capturing a small unit such as a reusable strategy, domain concept, or common failure mode.
• Properties:
  • (1) localization, so only the relevant bullets are updated.
  • (2) fine-grained retrieval, so the Generator can focus on the most pertinent knowledge.
  • (3) incremental adaptation, allowing efficient merging, pruning, and de-duplication during inference.
• Incrementally produces compact delta contexts: small sets of candidate bullets distilled by the Reflector and integrated by the Curator.

Grow-and-Refine

1. bullets with new identifiers are appended, while existing bullets are updated in place (e.g., incrementing counters).
2. A de-duplication step then prunes redundancy by comparing bullets via semantic embeddings. This refinement can be performed proactively (after each delta) or lazily (only when the context window is exceeded), depending on application requirements for latency and accuracy.

Observation

LMT function (Python functions, GPU/CPU kernel functions, memory operations, etc.) executions exhibit two significant characteristics.

• Within a single worker, most low-level functions are executed repeatedly. This is because training involves many identical iterations, and models are typically composed of repetitive submodules (e.g., transformer blocks).
• Runtime behaviors of functions are highly identical across workers, because modern parallelisms (e.g., at data, pipeline, tensor, and expert levels) distribute workloads evenly with frequent synchronization operations.
• Most performance issues can be diagnosed by observing abnormal function runtime behaviors in comparison to all other function executions

Insight

• (1) Performance issues can be observed by profiling the behavior of function executions
• (2) We can troubleshoot performance issues using differential observability that localizes the offending function executions with abnormal behavior (e.g., low average GPU-NIC throughput without fluctuation)
• (3) We do not need to analyze fine-grained raw observability data of all the functions; instead, we only need to summarize their runtime behavior patterns.

Design

• (1) detecting performance degradation of LMT to trigger online profiling (per worker).
• (2) summarizing runtime behavior patterns of each function from raw profiling data (per worker).
• (3) a centralized localization algorithm that pinpoints the root-cause function based on the behavior patterns (global).

Detecting Performance Degradation

• Indicators of iteration time: A PyTorch training iteration always involves several dataloader.next() calls, followed by several optimizer.step() calls (the number depends on training parameters like pipeline parallelism). The duration from the first dataloader.next() to the last optimizer.step() is regarded as the duration of a complete training iteration.
• After detecting 𝑀 (=10 in practice) identical sequences starting with dataloader.next() and ending with optimizer.step(), this sequence is defined as the training iteration sequence.
• Performance degradation detection
  • (1) The average duration of the recent 𝑁 (=50 in practice) iterations exceeds the recent shortest iteration time by more than 5%.
  • (2) The current training iteration sequence has not yet been fully matched, but the time elapsed since the last event received is at least 5× the average iteration duration (indicating the training is blocked)
    • If PerfTracker fails to match a training iteration after 𝐾 (=200 in practice) consecutive event receptions, it goes back to the previous iteration detection phase to redetect the training iteration sequence.
• Profile generation (default 20s)
  • Torch Profiler: function execution events for Python functions, CPU operations, memory operations, and CUDA kernels
  • nsys to sample hardware metrics at 10 kHz, such as GPU, DRAM, NVLink, PCIe, and the network.

Summarization

• (1) finding the function execution events on the critical path, including GPU computation kernels, collective communication functions, memory operations, Python functions, and all other functions executed in LMT
• (2) clustering all execution events of each function (for Python functions, the entire call stack must be identical to be considered the same function), then defining several patterns to summarize the behavior of each function

Localization

Performance issues

• (1) a common problem of all LMT workers, like hardware misconfigurations and low-efficiency code implementation
• (2) a special problem on only a part of the workers, like hardware issues

Motivation

How can we extract maximal learning signal from every expensive rollout to enable effective adaptation of complex, modular AI systems in low-data or budgetconstrained settings?

Overview

Genetic prompt evolution

1. Onitializing a candidate pool P, where a candidate is a concrete instantiation of the learnable parameters of the compound system, ⟨Π, Θ⟩Φ. Initially, the candidate pool consists only of the base system’s parameters as the sole candidate.
2. Proposes increasingly effective candidates by modifying existing ones through mutation or crossover, informed by learning signals from newly gathered rollouts and while tracking each new candidates’ ancestry. Each new candidate inherits learning signals from its parents, as well as signals from the current rollout.
3. During each iteration, GEPA identifies promising candidates from the candidate pool (candidate selection), proposes a new candidate—possibly by mutating prompts in a module based on reflective feedback or by performing crossover between two candidates—and evaluates this new variant on a minibatch of tasks. If the newly proposed candidate demonstrates improved performance relative to its parent(s) on the local minibatch, then GEPA adds the new candidate to the candidate pool P. This involves tracking internal data structures including tracking the ancestry of the new candidate, along with the full evaluation of the new candidate on a Dpareto, a validation set used for candidate selection.
4. After the budget is depleted, GEPA returns the candidate with the best aggregate performance on Dpareto.

Reflection using natural language feedback

1. Given a selected candidate to mutate in the current iteration of the optimization loop, GEPA updates the system with the candidate parameters, selects a target module within the system to improve (via round robin to ensure all modules receive updates), and generates a few rollouts over a minibatch sampled from the training dataset, recording their outcomes (success/failure).
2. By examining the execution traces of the system, GEPA identifies the target module’s inputs, outputs, and reasoning. With this, GEPA uses an LLM to reflectively examine this information, attributing successes or failures to elements of the module’s prompt (or omission thereof), and propose new instructions for the target module.
3. A new candidate is then proposed as a copy of the current candidate, with the target module’s prompt updated to the new proposed prompt.

System Aware Merge

Pareto-based candidate selection

1. Identifies the highest score achieved for each individual training instance across all candidates in the pool, creating a “Pareto frontier” of scores achieved by the optimization process so far.
2. Compiles a list of candidates that achieve the best score on at least one training task. This filters the pool down to candidates that incorporate “winning” strategies, preserving every valuable insight discovered in any reflective mutation.
3. Prunes candidates that are strictly dominated: for instance, if Candidate 2 has the best score on Task 1 only, but Candidate 3 achieves that same best score on Task 1 and the best on Task 2, Candidate 2 is removed.
4. Stochastically samples a candidate from this pruned list, assigning higher selection probability to candidates that achieved the best score across more training instances.

• State: describes the agent’s status with respect to the environment.
• State space: The set of all the states is called the state space, denoted as S = {s1, . . . , $s_n$}.
• Action:
• Policy: tells the agent what action to take in a given state.
• Reward: After executing an action at a state, the agent obtains a reward, denoted as r, as feedback from the environment. The reward is a function of the state s and action a. Hence, it is also denoted as r(s, a).
• Trajectory: state-action-reward chain.
• Return of the the trajectory: the sum of all rewards received along the trajectory.
• Discounted return with discount rate $\gamma$: $\sum_{t=0}^{\infty} \gamma^t r(s_t, a_t)$
• Episode: When interacting with the environment by following a policy, the agent may stop at some terminal states. The resulting trajectory is called an episode (or a trial ).

Markov Decision Process

• Sets:
  • State space: the set of all states, $S$.
  • Action space: a set of actions, $A(s)$, associated with each state $s \in S$.
  • Reward set: a set of rewards, denoted as $R(s, a)$, associated with each state-action pair (s, a)
• Model:
  • State transition probability: In state $s$, when taking action $a$, the probability of transitioning to state s′ is $p(s’|s,a)$.

  \(\sum_{s′ \in S} p(s′|s, a) = 1\)

  • Reward probability: In state $s$, when taking action $a$, the probability of obtaining reward $r$ is $p(r|s, a)$. For any $(s, a)$,
  \[\sum_{r \in R(s,a)} p(r|s, a) = 1\]

• Policy: In state s, the probability of choosing action $a$ is $\pi(a|s)$. For any $s \in S$,

  \[\sum_{a \in A(s)} \pi(a|s) = 1\]

• Markov property: memoryless property of a stochastic process.

  \[p(s_{t+1}|s_t, a_t, s_{t-1}, a_{t-1}, \ldots, s_0, a_0) = p(s_{t+1}|s_t, a_t)\] \[p(r_{t+1}|s_t, a_t, s_{t-1}, a_{t-1}, \ldots, s_0, a_0) = p(r_{t+1}|s_t, a_t)\]

State value:

• Discounted return along the trajectory: $G_t = \sum_{k=0}^{\infty} \gamma^k R_{t+k+1}$, $\gamma \in [0, 1)$ is the discount rate.

• State value:

  \[\begin{aligned} v_{\pi}(s) &= \mathbb{E}[G_t | S_t = s] \\\\ &= \mathbb{E}[R_{t+1} + \gamma G_{t+1} | S_t = s] \\\\ &= \mathbb{E}[R_{t+1}|S_t=s] + \gamma \mathbb{E}[G_{t+1}|S_t=s] \end{aligned}\]

Summary: Mycroft, a novel, lightweight distributed tracing and root cause analysis system designed to enhance reliability in Large Language Model (LLM) training by addressing issues within collective communication libraries (CCLs). These libraries often act as “black boxes,” making troubleshooting difficult and leading to wasted resources due to failures and slowdowns in large-scale training jobs. Mycroft functions by implementing Coll-level tracing, which captures fine-grained control and data dependencies—specifically Flow-level and Chunk-level traces—to expose internal communication states with minimal performance overhead. The system uses a real-time trigger mechanism to detect anomalies quickly and a dependency-driven analysis to efficiently pinpoint the root cause of both gray failures and performance degradation in distributed LLM training environments, showing effectiveness when deployed at scale, such as at ByteDance.

Instrumentation

Trace point format

• Metadata: IP, comm_id, Gid, GPU_id, Channel_id, QP_id
• Operation: timestamps, op_name, op_seq, msg_size
• Chunk: stuck_time, total_chunks, GPU_ready, RDMA_transmitted, RDMA_done

Tracepoints

• completion log: when a CollOp finishes. Include metadata (start and end timestamps, bytes transmitted, local and remote NIC information, and other CollOp metadata).
• realtime state log: every 100ms while a CollOp is in progress.
  • Includes data transmission progress across all devices, covering Stream Multiprocessor (SM) copies and RDMA writes.
  • Generated in each window until the CollOp completes or the NCCL proxy thread exits or crashes.

Real time trigger

• sampling a subset of ranks across the cluster and monitoring all CollOps.
• At least one rank per DP group to ensure coverage in common LLM parallelism architectures.
• Limit sampling to at most 10 ranks.

// LogData represents the aggregated trace data for a specific IP within a time window.
// In the source, this includes "completion logs" and "real-time state logs" [4].
type LogData struct {
	IP             string
	CompletedOps   int     // Count of completed Coll Ops
	HasRealTimeLog bool    // True if "real-time state log" exists (indicates active running state) [3]
	Throughput     float64 // Current throughput measurement
	OpInterval     float64 // Current interval between operations
}

// TriggerMechanism maintains the state for Mycroft's anomaly detection.
type TriggerMechanism struct {
	NormalThroughput float64 // Baseline for healthy throughput [2]
	NormalInterval   float64 // Baseline for healthy operation interval [2]
	Delta            float64 // Time window size (e.g., seconds) [1]
}

// NewTriggerMechanism initializes the trigger system.
func NewTriggerMechanism(delta float64) *TriggerMechanism { return &TriggerMechanism{ Delta: delta}}

// Trigger implements Algorithm 1: Trigger Mechanism [1].
// Input: Sampled IP list (S_ips), time (t)
// Output: Trigger type (string) and Abnormal IP (string), or empty if normal.
func (tm *TriggerMechanism) Trigger(sampledIPs []string, t float64) (string, string) {
	// Line 2: logs <- Acquire(S_ips, t - Delta, t) [1]
	logs := tm.Acquire(sampledIPs, t-tm.Delta, t)

	for ip, log := range logs {
		// Line 3: "if no Coll Ops completed in logs" [1]
		// Context from Section 4.3: Mycroft checks if the rank "stalls mid-operation
		// with real-time state log but without producing completion log" [3].
		if log.HasRealTimeLog && log.CompletedOps == 0 {
			// Line 4: return failure trigger, abnormal IP [1]
			return "failure_trigger", ip
		}

		// Line 6: "if throughput drops by half or Coll Op interval doubles" [2]
		// These heuristics are configurable but based on practical experience [3].
		isThroughputDrop := tm.NormalThroughput > 0 && log.Throughput < (tm.NormalThroughput*0.5)
		isIntervalSpike := tm.NormalInterval > 0 && log.OpInterval > (tm.NormalInterval*2.0)

		if isThroughputDrop || isIntervalSpike {
			// Line 7: return straggler trigger, abnormal IP [2]
			return "straggler_trigger", ip
		}

		// Line 8: update normal throughput and Coll Op interval [2]
		// Updates the baseline for the next window if no anomalies are found.
		tm.NormalThroughput = log.Throughput
		tm.NormalInterval = log.OpInterval
	}

	// Return empty strings if no active abnormal trigger is found.
	return "", ""
}

Root cause analysis

Principle

local and remote root cause

Deployment

Data volume: a training job utilizing 10,000 GPUs generates approximately 3 TB of data per day. This data is retained for one day before being discarded.

Debugging toolset

• Dump Python stacks of each rank using py-spy to identify the data loader or checkpoint stuck
  • Stack traces of all relevant Python programs are dumped automatically when Mycroft detects a trigger. Stacks are then grouped by process commands across all GPUs and mapped into a topology grid, where each block represents a GPU rank with its current Python call stack. In this grid, identical call stacks will be marked in the same color.
    • This is particularly useful as stuck threads typically have different call stacks from the rest, making them stand out on the grid for ease of troubleshooting.
• Dump Python CollOps collected by Flight Recorder to identify the synchronization issues
  • Stores the traces of the latest N CollOps in a ring buffer.
  • Each CollOp trace includes the CollOp ID, the sizes of input and output tensors, execution state, and communication process group ID.
  • Extract CollOp stacks to analyze the rank synchronization issue by figuring out the device with the last operation on each CUDA stream in a process group, and aggregating all the trace stacks to visualize and identify abnormal devices.

Control plane

1. Robust controller
   1. Robust controller react to events from proactive realtime check:
      1. High confidence event for a specific machine: force drain node, evict the machine, skip stop time diagnostics
         1. GPU unavailability
         2. Disk fault
      2. Network issue: tolerate several alerts before evict the machine. Policy: twice within 5 mins empirically
         1. NIC and Network switch flapping can auto recover
      3. If the restarted job fails again after machine eviction, it goes to a stop-time check procedure.
   2. Robust controller analyze the logs
      1. User space error: traceable to specific code modules from logs and exit codes, python exception -> code rollback
      2. Training crashes/abnormal metrics, e.g., NaN losses, arise without a clear culprit -> stop time check
      3. Performance anomalies: 0 RDMA traffic within 10 mins, low TensorCore utilization -> aggregation analysis
2. Runtime analyzer
   1. Aggregation analysis
      1. Parses process trees in each training pod to identify training related processes (torchrun, dataloader, checkpoint process)
      2. Stack traces from these identified processes are aggregated into multiple groups via string matching to differentiate abnormal sources
         1. The dominant groups are deemed healthy.
         2. Remaining groups are classified as outliers.
      3. Find shared parallel groups for those outliers and isolate the corresponding machines
   2. Fail slow (MFU decline):
      1. Repeat aggregation every 10s, flag the parallel group with the most outliers at each round.
      2. Parallel group with the highest cumulative flag count across 5 rounds is marked as the degrader for over-eviction.

Data plane (Robust agent in each training pod)

1. Monitor
   1. System inspection (Proactive realtime check):
      1. Network: NIC down or jitter, packet loss rate, switches down
      2. GPU: status of DCGM service, PCIe bandwidth, memory row remapping, and GPU temperature, etc.
      3. Host: OS kernel event (Xid in dmesg)
   2. Metics collection
      1. Workload specific training metrics from wandb: loss, gradient norm, MFR, etc
         1. Fatal signal: 5× increase in loss/gradient norms, NaN values
      2. stdout/stderr logs and process exit codes, which serve as hints for diagnostics
      3. Events: Significant declines serve as a signal for potential job hangs and MFU declines.
         1. CUDA
         2. RDMA: traffic
         3. Host
         4. Storage
2. Diagnoser
   1. NaN loss diagnosis
      1. Standard GPU and network tests first (EUD and NCCL tests).
         1. Intra machine all-to-all test to verify bandwidth
         2. Inter machine all-gather NCCL test to verify connectivity and integrity of data transfer with “neighboring” machines.
      2. Bitwise alignment test
         1. Each machine initiates a reference model whose structure matches that of the target training job (dense models or MoE models).
         2. Load predefined weights, employs a specific parallelism configuration (e.g., TP=2, PP=2, DP=2 or EP=2, PP=2, DP=2), and executes one training step on fixed input to ensure reproducibility.
         3. The outputs from all machines are collected and analyzed to verify bit-wise accuracy.
         4. Machines that yield incorrect results are promptly isolated and removed.
         5. If this test does not identify any defective machines, reattempt and rollback are sequentially employed to settle potential transient failures and human errors.
   2. Job hang and MFU declines
3. On demand tracer
4. CKPT manager