arXiv CS open access

RPEG: Training Reasoning Language Models Scalably and Cost-Effectively

Maciej Besta et al. · arxiv_cs

arXiv CS · Papers

cloud distributed-computing parallel-computing

Reinforcement learning (RL) is now central to training Reasoning Language Models (RLMs), but RLM pipelines are difficult to scale and costly to run because each RL step couples multiple large models, heterogeneous stages, cross-step state, and complex asynchronous dependencies. Existing systems typi

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

RPEG: Training Reasoning Language Models Scalably and Cost-Effectively Maciej Besta∗† , Leonard Schmidt† , Pierre Pang† , Patrik Okanovic, Robert Gerstenberger, Torsten Hoefler

arXiv:2611.09207v1 [cs.DC] 6 Nov 2026

Department of Computer Science, ETH Zurich

Abstract—Reinforcement learning (RL) is now central to training Reasoning Language Models (RLMs), but RLM pipelines are difficult to scale and costly to run because each RL step couples multiple large models, heterogeneous stages, cross-step state, and complex asynchronous dependencies. Existing systems typically optimize only parts of this design space. We address this complexity by jointly optimizing the pipeline structure as well as its execution using a novel data-centric graph intermediate representation called RLM Pipeline Execution Graph (RPEG). RPEG exposes structural choices such as bounded asynchrony and actor–critic sharing together with execution optimizations such as stage fusion, fine-grained scheduling, and multi-dimensional model parallelism. It is coupled with a comprehensive analytical performance model for RLM training, jointly capturing compute, communication, memory, and hierarchical rack-, node-, and GPU-level effects. Across systems with up to 256 GH200 GPUs, our distributed runtime engine implementing RPEG improves performance by up to 2× over ReaL and reduces step latency by up to 2.48× over verl, while increasing hardware utilization by up to 30%. Overall, RPEG provides a graph-driven foundation for next-generation training pipelines for frontier reasoning LLMs.

Code: https://github.com/spcl/RPEG I. I NTRODUCTION

EINFORCEMENT learning (RL) has become central to post-training large language models (LLMs) for human preferences, safety constraints, and task-specific objectives [4, 38, 48]. More recently, RL-based post-training has become the foundation of reasoning language models (RLMs, also called reasoning models): models that perform extended reasoning, produce chain-of-thought (CoT) trajectories [52], exploit richer structures [6, 7, 14, 58, 60], self-correct [15], and allocate more computation to harder problems [21, 23, 36]. Examples include DeepSeek-R1 [23], GPT-5 Thinking [47], Kimi K2.5 Thinking [3], and Grok 4.7 [56]. Here, RL becomes a mechanism by which models learn to reason. The computational footprint is enormous and continues to grow: frontierscale models consume millions of GPU-hours, while scaling both train-time RL and inference-time reasoning continues to improve model capabilities [20, 30, 35]. These costs grow further in agentic systems [8, 12, 13, 16, 28, 32, 51, 55, 57, 59]: long-horizon tool use and multi-agent execution increase system overhead [26], with reported token usage reaching 4×– 15× that of ordinary chat [1] and roughly three orders of magnitude higher for agentic coding workloads [2]. Yet, RLM post-training is substantially more complex than a single-model workload. In the canonical Proximal Policy ∗ Corresponding author † Core contributions

Optimization (PPO) pipeline, each step comprises a generation stage that samples responses with a trainable actor (also called policy), an assessment stage that evaluates them using a reward model, critic, and frozen reference model, and a training stage that updates the trainable components [31, 62]. Other widely used online RLM-training paradigms, such as Group Relative Policy Optimization (GRPO) [43], simplify parts of this pipeline (for example by removing the critic) but retain essentially the same generation–assessment–training template and its heterogeneous execution characteristics. In PPO, generation is autoregressive and latency-sensitive, assessment is largely teacher-forced and throughput-oriented, while training adds backward passes, optimizer state, and parameter updates. RLM training is therefore simultaneously multi-model, multistage, heterogeneous, and iterative – and fundamentally a parallel and distributed computing problem [53]. Scaling individual models does not by itself optimize an RLM training pipeline. Data (DP), tensor (TP), pipeline (PP), context (CP), sequence (SP), and expert parallelism (EP) must be coordinated with cross-stage placement, state exchange, and execution order [5, 17]. Existing systems pursue three roughly complementary directions: (1) effective execution planning, including placement and parallel-strategy selection in ReaL [31] and HybridFlow/VERL [45], and workflow compilation in RLinf [61]; (2) fine-grained and asynchronous execution, including sample- and microbatch-level fusion in RLHFuse [63] and decoupled rollout–training pipelines in AReaL [22], Laminar [44], and StaleFlow [29]; and (3) adaptive resource orchestration, which adjusts parallelism or reallocates resources as workloads evolve [18, 49, 50]. These mechanisms increasingly coexist within the same system. The key design question is therefore how broadly these decisions can be optimized together. Aspects such as overlap, placement, parameter freshness, model sharing, parallelism, sharding vs. replication, and memory cannot be optimized independently: improving one may increase communication, resharding, or memory pressure elsewhere. This interaction becomes even more important on hierarchical supercomputers, where communication costs differ substantially within a node, across nodes, and across racks. Finally, maximizing throughput alone may still leave expensive hardware poorly utilized. Taken together, these couplings make it difficult to achieve both high performance and high hardware utilization while respecting communication and memory constraints. We address this challenge by optimizing not only how an RLM pipeline executes, but also selected aspects of the pipeline structure itself (contribution 1). We represent both

Simple Established Baselines Fully Co-Located

Actor

Throughput: 11,372 tok/s Makespan: 195.9 s GPU U�liza�on: 100%

Genera�on Assessment Training

Naive Concurrent broadcast rollouts

Oﬄoad actor

Reward

Cri�c

Communica�on

broadcast weights

Throughput: 11,221 tok/s Makespan: 198.6 s GPU U�liza�on: 98.1%

Relevant Baseline Schedules ReaL

Overlaps only invoca�ons without dependencies: genera�on never overlaps assessment or actor training

adapted from [32]

Only 1 itera�on

Devices Load trainable models

reshard actor & cri�c weights

Time

broadcast training signals

Time

RPEG

broadcast rollouts

Synchronous RPEG individual rollouts

Devices

Itera�on boundary

genera�on length quan�les

actor training synchronous signals rollouts

Throughput: 15,876 tok/s Makespan: 140.4 s GPU U�liza�on: 64.6%

Asynchronous RPEG async execu�on => overlap of genera�on & training

Comm. visualized on receiving devices

Throughput: 17,970 tok/s Makespan: 124.0 s GPU U�liza�on: 81.5%

bounded asynchronous rollouts

adapted from [46]

Only 1 itera�on

bubbles

sub-node spli�ng => stage fusion

Time

preﬁll bubbles

shadow copy of actor

no bubbles during steady-state

Searches execu�on plans for a ﬁxed pipeline: bounded asynchrony and actor–cri�c are never explored

reshard actor weights

Devices

sync weights

Devices

sub-node spli�ng ﬁne-grained enables rollout forwarding assessment

broadcast training signals Time

VERL

...

forward early ﬁnished rollouts

Reference

Resharding

Time

broadcast rollouts

broadcast training signals

Time

Fig. 1: Representative PPO schedules selected by different planners. RPEG exposes stage-level and fine-grained overlap, enabling disaggregated placements and bounded-asynchrony that are not available under fully colocated or coarse-stage baselines. No AI was used to conceive or draw the picture.

jointly as a single optimization object C = (G◦ , P ), exposing structural and execution decisions within one searchable configuration space. Specifically, the compact steady-state graph G◦ captures the recurring pipeline structure and cross-step state semantics, including structural variants such as bounded asynchrony and actor–critic (AC) sharing [17]; the execution plan P specifies per-invocation hardware mapping, multidimensional model parallelism, and microbatching. Thus, the pipeline itself becomes an optimization variable rather than merely a fixed graph to be mapped onto hardware. We realize this abstraction in the RLM Pipeline Execution Graph (RPEG), a data-centric graph intermediate representation and optimization framework (contribution 2). Inspired by property graphs from the graph database domain [9– 11], RPEG is property-rich rather than a bare dependency graph, compactly attaching rich data (e.g., the model, role, staleness, execution semantics) needed for subsequent optimization. A staged graph-transformation pipeline materializes execution refinements such as stage fusion and fine-grained sub-invocation execution, and converts each candidate configuration C into a schedulable hardware-aware graph. This allows the optimizer to evaluate structural and execution choices jointly without directly searching a fully materialized multistep graph. We showcase RPEG execution plans in Figure 1. We further develop a distributed-memory runtime (contribution 3) that executes the selected RPEG plans faithfully on supercomputers, connecting the graph-level optimization decisions to their physical execution. RPEG is coupled with a calibrated comprehensive analytical performance model for end-to-end RLM training (contribution 4). The model jointly captures compute, communication, and memory costs of heterogeneous RLM operators, including resharding, pipeline bubbles, optimizer traffic, co-resident model state, and hierarchical GPU-, node-, rack-, and inter-rack communication. Automated hardware calibration replaces reliance on nominal vendor peaks, enabling rapid evaluation of large numbers of candidate plans.

Finally, we comprehensively evaluate RPEG from analytical modeling to large-scale distributed execution (contribution 5). We validate the performance and memory models, study optimizer stability and ablations, and execute RPEG plans on up to 256 GH200 GPUs. Optimized synchronous plans are 1.250–1.471× faster than their executable fully colocated inputs; against ReaL, RPEG improves end-to-end performance by up to 2×; and in the controlled synchronous comparison with separate actor and critic models, RPEG achieves up to 2.477× lower time per step than VERL (median 1.488×). RPEG also improves hardware utilization by up to 30%, while experiments with AC sharing, stage fusion, and bounded asynchrony expose further performance gains and tradeoffs enabled by the representation, all while preserving training semantics and accuracy (within 0.47% of the VERL baseline), and with negligible optimizer overheads even at 2k GPUs. II. BASICS OF RLM P IPELINES We call one generation–assessment–training cycle an RL step. An invocation is one model’s execution in a particular role, such as actor generation or critic training; it may itself contain many token-generation or forward/backward operations. A rollout is a sampled response to a prompt, and its successive token-generation decisions form a trajectory. In PPO-style RLHF, the trainable actor generates responses; the frozen reward model scores them; the trainable critic estimates token-level future returns; and the frozen reference model supplies log-probabilities for policy regularization. One RL step therefore contains six invocations: actor generation, reward/reference/critic inference, and actor/critic training [31, 62]. PPO compares current-policy probabilities with those of the policy that generated the responses and uses reward/value estimates to construct advantages, which guide the update [38, 42]. GRPO removes the critic and estimates advantages from rewards within groups of responses, but retains the generation–assessment–training template [43].

These invocations have different computational profiles. Generation first processes the prompt (prefill) and then decodes tokens sequentially, making it latency-sensitive and increasingly memory-bound as the KV cache grows. Assessment instead processes complete sequences with highly parallel forward passes and is typically throughput-oriented. Training adds backpropagation, optimizer state, and parameter updates, substantially increasing both computation and memory requirements. Intra-model parallelism distributes one invocation using standard techniques such as data (DP), tensor (TP), pipeline (PP), sequence (SP), context (CP), and expert parallelism (EP) [5]. A microbatch is a smaller batch partition processed within an invocation. A microtask, such as one pipeline stage’s forward or backward computation for one microbatch, is a finer execution unit. Sub-invocation execution exposes units smaller than a whole model invocation, such as microbatches, pipeline microtasks, or groups of responses that finish at similar lengths, to the scheduler. Inter-model strategies organize different model invocations [17]. Co-location places them on overlapping GPU sets; disaggregation uses disjoint sets. Different invocations of one model may use different parallel layouts, requiring resharding, i.e., redistribution of parameter tensors. Stage fusion overlaps work within an RL step: assessment of completed responses can overlap continuing generation, or actor/critic training microtasks can interleave while retaining separate parameters. Bounded asynchrony instead overlaps different RL steps by allowing serving invocations (generation or assessment) to use parameter versions delayed by at most a configured number of RL steps. Actor–critic (AC) sharing uses one backbone with policy & value heads instead of two separate models. Unlike fusion, AC sharing changes the trainable model; asynchrony changes parameter freshness. Both are thus explicit configuration choices rather than universally interchangeable implementations of synchronous, separate-model PPO. III. OVERVIEW OF RPEG RPEG’s key idea is to optimize not only how an RLM training pipeline executes, but also selected aspects of the pipeline structure itself. We represent the complete optimization object as an RPEG configuration C = (G◦ , P ); G◦ captures the pipeline structure while P specifies its execution plan. Symbols used in the paper are listed in Table I. The graph G◦ captures the steady-state algorithmic structure of one RL step and cross-step semantics. Its vertices are operators, i.e., individual model invocations in specific roles such as generation, inference, or training. Its edges encode how data and updated model state flow between them within and across steps. The representation is richer than a bare dependency graph. Operators carry their role, associated model, and declared input/output sets used for consistency checks, while edges distinguish within-step dataflow from cross-step model-state publication and carry typed data containers. The execution plan P specifies how each operator is mapped to hardware through device allocation, multidimensional model parallelism, and microbatching. During staged graph transformation, C is progressively refined. For example, a model invocation may be decomposed

into finer-grained executable phases, while communication and redistribution implied by the selected hardware mapping become explicit graph nodes. Thus, the same graph-centric abstraction evolves from a compact logical description into a detailed executable representation that specifies the operations, data movement, and resource requirements on a given cluster (Figure 2). This transformation consists of the following steps: U NROLL

L OWER

L OCAL E VAL

S CHEDULE +S CORE

C −−−−→ Ḡ −−−−→ Ĝ −−−−−−→ Φ(Ĝ) −−−−−−−−−→ Score(C).

Unrolling 1 expands the cyclic steady-state graph IR over a finite horizon Keval , yielding an acyclic graph Ḡ in which cross-step publication dependencies become explicit. By representing bounded-staleness publication already at the compact IR level, this stage exposes asynchronous inter-step overlap while keeping the search space compact. Lowering 2 then combines Ḡ with the execution plan P and produces a hardware-aware executable DAG Ĝ, inserting explicit C OMM and R ESHARD nodes whenever the selected mapping requires communication or redistribution. Lowering also materializes finer-grained executable units, such as response-length phases, when they enable scheduling below a whole model invocation. Local evaluation 3 assigns analytical cost semantics to Ĝ: it produces a node-level time map T and estimates memory requirements under the selected hardware mapping, without yet resolving global resource contention. We denote the resulting graph, annotated with time/memory estimates, by Φ(Ĝ). Scheduling 4 then constructs a resource-aware execution order under precedence and device-reservation constraints, yielding the predicted makespan MS(Ĝ) and scheduled resource usage. Finally, the optimizer scores 5 the configuration using steady-state throughput or, optionally, a throughput–utilization objective, together with a soft penalty for memory-infeasible candidates. The optimizer treats this score as a search signal and iteratively updates 6 both G◦ and P . In this way, the framework jointly optimizes pipeline structure and execution mapping rather than treating the RLHF graph as fixed. After optimization, the runtime reconstructs the selected configuration and its physical mapping and executes 7 it on persistent distributed workers. Operations start only when their actual predecessors have completed, the required model versions are available, and the necessary resources are free. Runtime traces verify the executed dependencies, model versions, transfers, and resource ownership. IV. RPEG T RANSFORMATION P IPELINE We now detail the RPEG transformation pipeline that transforms a compact high-level graph-driven algorithmic description of the RL post-training into a low-level representation schedulable on a give cluster; see also Figure 2. Extensive mathematical details of all the subsections are in Appendix A. A. Graph-Centric RLM Training Representation G◦ defines the algorithmic structure of RLM training: ◦ ◦ G◦ = (V ◦ , Edata , Epub , A, M, τ ).

Here, V ◦ is the finite set of logical operators in one ◦ steady-state training step; Edata is the set of within-step data

Pipeline structure 𝑮°

It deﬁnes the algorithmic structure of RLM training

Data edges: within step dependencies

Publica�on edges: crossstep state ﬂow

Actor (Policy)

Synchronous v0

Inference

Cri�c (Value)

Training

values

weights

cri�c = 0

Reward

rewards

Cri�c (Value)

Reference

Execu�on Plan P GPUs assigned to a given vertex

Actor (Policy)

It assigns execu�on details to each RPEG vertex

v0 v1

Step k=2 v1

Node-spli�ng within the search space allows the op�mizer to overlap genera�on and assesment on disjoint devices

Step k=1

With quan�le-based sub-node spli�ng

Step k=2

Microbatch Setup

... Op�mizer Propose C' based on Score(C) next C' Score S

current C

Conﬁgura�on C

Execu�on 7 (Sec�on VII):

Genera�on in k=1 can start while actor training in k=0 is ongoing

Throughput

Inner loop: execu�on plan P. Simulated annealing over per-operator placement, 3D parallelism, and microbatch counts.

Scheduler

Score S. Steady-state thoughput (op�onally × u�liza�on) with a so� memory-penatly.

Assigns predicted start and ﬁnish �mes under dependencies and resource constraints. start �me

Run�me Executes the best conﬁg found Copt on persistent distributed GPU workers. Comple�on-driven dispatch: opera�ons start once precedessors ﬁnish, model versions are available, and GPUs are free.

Each vertex has device alloca�on + 3D strategy

(Sec�on VI): 5 Op�mizing Compu�ng & improving Makespan, U�liza�on,

Outer loop: pipeline structure 𝑮°. Enumerates structural variants: staleness budget and actor-cri�c sharing.

end �me

Makespan MS(G) = max (end �me) U�liza�on = busy / MS(G) Throughput = #tokens / MS(G)

Scheduling (Sec�on IV-E): Computes pernode run�mes and memory requirements

Evalua�on (Sec�on IV-E): Assigns cost to executable opera�ons, communica�on, and resharding

Evaluated graph Φ(𝑮) It augments each node of the DAG G with es�ma�on of its run�me and memory requirements.

makespan

Run�me: 20.43 s Memory: 70.34 GB

... ...

Devices

Search budged reached: best found Copt

Without subnode spli�ng

cri�c = 0

Step k=1

Step k=1 v0

Update 𝑃 or 𝐺° itera�vely

Communica�on Resharding

actor = 1

Step k=0

It captures executable opera�ons, communica�on, resharding, and their dependencies.

Synchronous Step k=0

log probabili�es

Mul�-dimensional parallelism

Asynchronous

example, 1-step bounded

Step k=0

Step k=2

actor = 0

weights

Executable DAG 𝑮

It expands 𝑮° across training steps, exposing cross-step dependencies and model versions

Acyclic Unrolled Graph 𝑮

prompt comple�ons

Lowering (Sec�on IV-D): Translates the unrolled 2 graph into GPU opera�ons and explicit data transfers.

Unrolling (Sec�on IV-C): Expands

pipeline into a ﬁnite sequence 1 the of interconnected training steps.

Devices

Genera�on

Pipeline structure 𝑮°

Execu�on Plan P

It captures both the pipeline structure and its execu�on

RPEG conﬁgura�on

Run�me: 0.1 s Volume: 300 MB

run�me

Run�me: 15.3 s Memory: 45 GB

Time

...

Use calibrated analy�cal expressions

Performance model Analy�cal cost model for execu�on �me and memory requirements. Calibrated via hardware proﬁling. 900 GB/s 20 µs 990 TFLOP/s 96 GB

990 TFLOP/s 96 GB

Fig. 2: Overview of the RPEG. A steady-state graph template and execution plan are unrolled, lowered to an executable DAG with explicit communication and resharding, evaluated with local analytical costs, and scheduled to obtain a steady-state score. No AI was used to conceive or draw the picture.

◦ dependencies; Epub is the set of cross-step model-state publication dependencies, which determine how updated model parameters become available to later consumers; A is the set of typed dataflow containers, such as completions, rewards, log-probabilities, or model weights; M is a set of all the used models (here, Mtrain is the set of trainable models); and τ : Mtrain → N0 is a per-trainable-model staleness budget. Intuitively, G◦ specifies what recurring RLM computation is performed and which cross-step state dependencies are semantically admissible, while deliberately leaving how it is mapped to hardware to the execution plan. Each vertex v ∈ V ◦ denotes one model invocation. It carries a role r(v) ∈ {generation, inference, training}, its associated model µ(v) ∈ M, and declared input/output sets I(v), O(v) ⊆ A specifying which typed containers the operator may consume and produce. For instance, a generation vertex consumes prompts and actor weights and produces completions, while a reward-inference vertex consumes prompt–completion pairs and produces rewards. Every edge refers to one typed container a ∈ A, whose size metadata (|a|, s(a)) specifies its element count and bytes per element. Its byte volume is |a|s(a). Containers represent both transient products such as completions, log-probabilities, rewards, and values, and persistent state such as model weights. ◦ Data edges Edata encode within-step dependencies, whereas ◦ model-state publication edges Epub encode cross-step state

flow from training to future–step consumers. This distinction makes the cross-step reuse of updated model state explicit and exposes bounded asynchrony through the staleness map τ . The graph must satisfy four well-formedness conditions. First, input/output consistency requires every container carried by an edge to belong to the source output set and the destination input set. Second, model-state publication edges must originate at training vertices and carry model parameters. Third, boundary consistency requires every container without a producer or consumer to be represented as an explicit graph ◦ input or output. Fourth, the within-step graph (V ◦ , Edata ) must be acyclic; cycles may arise only through publication edges representing cross-step model updates. These conditions ensure acyclicity after valid unrolling; physical feasibility additionally depends on the execution plan, model versions, memory, and resource constraints. We use a cyclic representation instead of searching directly on a flat multi-step DAG. First, it keeps the search space compact. Suppose invocation v has qv possible execution-plan assignments. In the compact RPEG, this choice is made once for Q v and reused by all of its copies after unrolling, giving v∈V ◦ qv combinations. If instead each of the Keval unrolled copies could choose its plan independently, invocation v alone Keval would have and the full plan space would Q qv Kcombinations, grow to v∈V ◦ qv eval . Second, staleness becomes a first-class structural variable: synchronous and bounded-asynchronous

TABLE I: Main notation used throughout the paper. Models & their architecture and workload M; Mtrain B; R Lp ; Lr sb

Set of all models; a trainable subset of models (Mtrain ⊂ M). Batch size (input prompts per step); generated rollouts per prompt. Average prompt length; average rollout length. Number of bytes per scalar in working precision (sb = 2 for BF16).

Hardware & its architecture D Nr ; Nn ; Ng α; β Fpeak BWmem γM , γK MGPU , Meff

Device space; D = [Nr ] × [Nn ] × [Ng ]. Number of racks; nodes/rack; GPUs/node. Latency; inverse bandwidth (α–β model). Workload-effective GPU throughput (FLOP/s per GPU). HBM bandwidth (bytes/s per GPU). GEMM SM-utilization constants (batch dim.; hidden dim.). Per-GPU HBM capacity, effective capacity after runtime overhead.

RPEG C G◦ A, a Ḡ Ĝ Φ(Ĝ) Keval τ

An RPEG configuration: C = (G◦ , P ). ◦ ◦ An RPEG: G◦ = (V ◦ , Edata , Epub , A, M, τ ). ◦ ◦ ◦ V : RPEG vertices, Edata : data edges, Epub : publication edges. A set of containers and a single container a = (|a|, s(a)) ∈ A. The unrolling of G◦ : Ḡ = (V̄ , Ēdata , Ēpub ). The lowering of Ḡ: Ĝ = (V̂ , Ê). Evaluation of Ĝ (estimating runtimes of all vertices). Unrolling horizon (#steps considered when unrolling G◦ into Ḡ). Staleness budget per trainable model; τ : Mtrain → N.

RPEG’s Execution Plan P σ(v) |σ(v)| Π(v), mb(v)

RPEG’s execution plan: P = (σ, Π, mb). Device allocation of operator v ∈ V ◦ (the set of GPUs reserved by v). The number of GPUs in this allocation. Data/pipeline/tensor parallelism Π(v) = (Pd , Pp , Pt ); microbatch count.

Evaluation Pipeline T Wshard (v) peak(d) Score(C) Feff (v) I(v), η(v)

Node-level time map (runtime estimations for operators), T : V̂ → R≥0 . Per-GPU parameter count (weight shard) for operator v. Per-device peak memory (persistent + worst-case transient). Throughput estimate for configuration C. Achieved per-GPU GEMM throughput of operator v. Arithmetic intensity and calibrated GEMM utilization of operator v.

variants are obtained by modifying τ and re-unrolling rather than by manually rewiring a materialized multi-step graph. Third, the cyclic form expresses steady-state semantics directly, independently of the finite horizon used for evaluation. The execution plan P = (σ, Π, mb) assigns hardwaredependent attributes to each RPEG vertex. Here, σ(v) is a device allocation, Π(v) is a multi-dimensional model-parallel configuration, and mb(v) is the number of microbatches. In the evaluated implementation, Π(v) = (Pd , Pp , Pt ) instantiates nested data, pipeline, and tensor parallelism. We model the cluster as a hierarchical device space D = [Nr ] × [Nn ] × [Ng ], comprising racks, nodes per rack, and GPUs per node. Placement therefore determines not only co-residency and disaggregation but also whether communication remains intra-node, crosses nodes within a rack, or crosses racks.

model staleness map τ . AC sharing changes the model structure and associated dependencies, while changing τ changes the cross-step state dependencies and hence the admissible degree of asynchronous overlap. For every such structural candidate G◦ , the optimizer searches its execution plan variants P . G◦ and P are thus evaluated jointly rather than optimizing only a hardware mapping for one fixed pipeline. Additional execution refinements are materialized during lowering. Stage fusion can realize multiple operations through a composite fine-grained execution structure, while responselength-aware splitting replaces generation and downstream inference invocations by matching executable phases when their selected placements permit overlap. These transformations enlarge the executable design space without treating arbitrary graph topology as an outer-loop search variable.

B. Joint Structural and Execution Configuration

C. Unrolling: Making Asynchrony Explicit

◦

Crucially, G is itself part of the optimization space rather than a fixed workload graph. The current optimizer explicitly searches structural variants defined by AC sharing and the per-

The cyclic RPEG cannot be scheduled directly because model-state publication edges connect training in one step to consumers in later steps. Unrolling expands G◦ over a finite

horizon Keval into an acyclic graph Ḡ = (V̄ , Ēdata , Ēpub ), where each vertex v ∈ V ◦ is replicated once for every step index k ∈ [0, Keval − 1]. Each within-step edge is copied inside the same replicated step, whereas each publication edge is instantiated across steps according to the staleness budget. Version zero is the initial checkpoint; training at step j produces version j + 1, and successive updates of the same optimizer state remain ordered. A serving consumer (generation or inference) at step k requires version max{0, k−τ (m)}: training step k − τ (m) − 1 supplies it when k > τ (m), and otherwise the initial checkpoint supplies it. The synchronous case τ (m) = 0 consumes the immediately preceding update for k > 0, and initialization for k = 0. This construction makes bounded asynchrony explicit in the graph. If k ≤ τ (m), no publication edge is created and the corresponding early steps consume the initial parameter version; because they have no training predecessor, their serving operations may overlap subject to all remaining constraints (e.g., resource or update order). Generally, publication edges always point from smaller to larger step indices, while withinstep edges preserve the step index. Hence, if the within-step dependency graph is acyclic, the unrolled graph is also acyclic. Unrolling thus exposes inter-step overlap without sacrificing the compact steady-state search space at the RPEG level. D. Lowering: Delivering Executable Representation Unrolling makes semantic dependencies explicit, but not the physical consequences of the selected execution plan. Lowering combines Ḡ with the candidate execution plan P and produces an executable DAG Ĝ = (V̂ , Ê) with three kinds of vertices: Op nodes for unrolled operators, Comm nodes for explicit data transfers, and Reshard nodes for parameter redistribution between incompatible layouts. Every Op node carries its placement, model-parallel configuration, and microbatch count; every Comm or Reshard node carries an explicit byte payload and the source and destination GPU allocations. The component σ(v) specifies the set of GPUs assigned to invocation v ∈ V ◦ , with |σ(v)| = Pd (v)Pp (v)Pt (v); its unrolled Op nodes inherit this allocation. For analytical scheduling, σ(x) denotes the GPUs reserved by an executable node x: the assigned GPUs for Op nodes and the destination (and optionally source) GPUs for Comm/Reshard nodes. Lowering also realizes the principal mechanisms behind fine-grained overlap and asynchrony. First, when generation and downstream inference are placed on disjoint devices, generation can be split into quantile-based sub-nodes and inference into matching sub-nodes using a profiled responselength distribution. Inference for early-finishing responses can then overlap with later generation phases. Second, lowering detects shadow pairs: if a trainable model is consumed asynchronously (τ (m) > 0) and its producer and consumer use disjoint allocations, the two sides retain distinct parameter copies and intra-step resharding can be avoided, enabling cross-step overlap. For non-shadow pairs, lowering inserts Reshard nodes when layouts differ on the same devices, or precedence edges and explicit model-state transfers when

mappings use different device sets. Transfers also account for partially overlapping source and destination allocations. E. Scoring and Scheduling A schedule assigns a start time to every node of the executable DAG. Formally, it is a function s : V̂ → R≥0 ; the finish-time function is f (v) = s(v) + T (v), where the analytical model of Section V supplies T . Precedence requires f (u) ≤ s(v) for every (u, v) ∈ Ê. Resource exclusivity requires that nodes with overlapping reserved GPU sets, σ(u) ∩ σ(v) ̸= ∅, do not execute concurrently. The predicted makespan is MS(Ĝ) = max f (v). v∈V̂

We construct this schedule with a greedy scheduler, whose input is the DAG, a map T , and GPU reservations σ(v) defined during lowering, and whose output is the start/finish times of each node. At each step, it considers nodes whose predecessors have already been assigned finish times, computes their earliest possible start times subject to precedence and GPU availability, and schedules the node that can start earliest. Ties first favor Comm/Reshard nodes, reducing unnecessary delays in transfers needed by downstream computation. Remaining ties favor nodes with the longest remaining dependency path, measured by the sum of predicted execution times along a path to a terminal node (one without successors), since delaying such work is more likely to increase the makespan. Generation phases inherit the priority of their first phase, discouraging unrelated work from interrupting consecutive phases that reuse the same KV cache. We score a configuration C = (G◦ , P ) by Score(C) =

Keval BR(Lp + Lr ) MS(Ĝ)

where B is the number of prompts per RL step, R is the number of rollouts per prompt, and Lp , Lr are the average prompt and rollout lengths. The numerator counts logical prompt–response tokens, not repeated processing by every model role. The score amortizes execution over the unrolled horizon. Memory infeasibility is handled by a soft penalty, peak(d) −1 . Scorepen (C) = Score(C) exp −βoom max 0, max d∈D Meff

Here βoom > 0 controls the penalty strength and Meff = (1 − εrt )MGPU reserves part of physical HBM for runtime overhead and fragmentation (εrt = 0.10 in the analytical baseline). A candidate slightly above this usable capacity may be explored, but the returned configuration must satisfy peak(d) ≤ Meff for every GPU. “Nearly feasible” candidates are not accepted as executable plans. For a fixed GPU budget, throughput is the default objective. The optional efficiency objective also uses Util(C) =

tbusy (d)

d∈D

|D| MS(Ĝ)

where tbusy (d) is the time GPU d is occupied by scheduled nodes. For a user-specified throughput tolerance δ, it is ( Scoreeff (C) = Score(C) ×

Util(C), λfloor ,

Score(C) ≥ (1 − δ)Score⋆ , otherwise,

where Score⋆ is the highest raw throughput found so far and λfloor = 10−12 strongly penalizes plans below the throughput floor. The same memory penalty is then applied to Scoreeff . Thus, the floor is a search preference, whereas the final memory bound is mandatory. GPU-hours are reported separately as reserved GPUs multiplied by elapsed hours. V. A NALYTICAL P ERFORMANCE M ODEL FOR RLM S We now describe the performance model for RLM training. Given the executable DAG Ĝ produced by unrolling and lowering, we assign each node an analytical execution-time estimate and check whether the resulting configuration is memory-feasible. The performance model therefore consists of (1) a node-level time map T : V̂ → R≥0 , and (2) a memory-feasibility predicate MemFeasible(C) with C = (G◦ , P ) being a given configuration. Together, these model the runtime consequences of different placement strategies, multi-dimensional parallelism, microbatching, communication effects such as resharding, bounded asynchrony, and coresidency on a hierarchical supercomputing machine. Due to space constraints, we focus on key aspects specific to RL post-training; detailed derivations appear in Appendix D, while calibration, modeling assumptions, and implementation details appear in Appendix C. A. Model Configuration We first discuss several setup aspects of the used model. 1) Modeling Communication: Communication arises both between executable DAG nodes, through explicit C OMM and R ESHARD nodes, and within operator O P nodes (through multi-dimensional parallelism). We use a standard α–β model [25] T (M ) = α + βM , for a message of M bytes. The appropriate (α, β) pair is selected from the tier spanned by the communicating allocations: intra-node, inter-node within a rack, or inter-rack. If an operation spans multiple tiers, we conservatively use the slowest one. For collectives, we use ring-based reduction/gather estimates and a pairwise-exchange estimate for all-to-all, summarized in Table II. The message-size convention depends on the primitive: M is the buffer size for AllReduce/ReduceScatter, the final gathered buffer size for AllGather, and the aggregate exchanged bytes per GPU for AlltoAll. P2P uses the transferred bytes. Collective costs are zero for p = 1. 2) Model Input: The model is parameterized by three input groups. The workload specifies batch size B, rollouts per prompt R, average prompt length Lp , and average rollout length Lr . Then, for each model m ∈ M, the model architecture specification provides layer count, hidden dimension, feed-forward dimension, vocabulary size, head dimension, and working precision sb (i.e., the number of bytes per scalar in the runtime tensor format, e.g., sb = 2 for BF16). The hardware description consists of a hierarchical device space D = [Nr ]×[Nn ]×[Ng ], with Nr racks, Nn nodes per rack, and Ng GPUs per node, plus calibrated performance parameters.

TABLE II: Communication estimates for p participating GPUs; M is the message size; (α, β) are calibrated for the primitive and communication tier. Primitive

Time

Application

AllReduce AllGather ReduceScatter AlltoAll P2P

2(p−1)α + 2 p−1 p M β (p−1)α + p−1 p M β (p−1)α + p−1 p M β (p−1)α + M β α+Mβ

TP partial-result aggregation ZeRO-2 weight gather ZeRO-2 gradient reduce R ESHARD nodes PP send/recv, C OMM nodes

3) Hardware Calibration: Rather than relying on vendor peak numbers, we calibrate effective hardware parameters on the target system. These include workload-effective compute throughput Fpeak , HBM bandwidth BWmem , separate (α, β) parameters for collective and point-to-point communication at all three tiers, GEMM efficiency constants γM , γK , and GPU memory capacity MGPU . This distinction is important because both compute and communication are strongly topology-dependent on modern hierarchical systems, and because NCCL collectives and P2P transfers can differ substantially in achieved bandwidth. B. Intra-Model Runtime For each operator node v ∈ V̂ , we model runtime under the evaluated 3D (DP/PP/TP) strategy (Pd , Pp , Pt ) and microbatch count mb by the decomposition Top (v) = Tgemm (v)+Tbubble (v)+Tcomm (v)+Tnon-gemm (v)+Topt (v),

where the five terms capture GEMM execution, pipeline bubbles, intra-operator communication, non-GEMM overheads, and optimizer traffic. The dominant GEMM compute term is obtained from the operator work W (v), i.e., the total FLOP count induced by the operator role and model type. For an operator with execution strategy (Pd , Pp , Pt ), we define Pd Pp Pt = |σ(v)| as the number of GPUs allocated to that operator. Assuming balanced sharding across this 3D mesh, the GEMM work is divided across these devices, giving Tgemm (v) =

W (v) , Pd Pp Pt Feff (v)

The denominator states that the devices allocated to v execute its data, pipeline, and tensor shards concurrently (other operators may use disjoint allocations at the same time if the schedule permits). Feff (v) is the achieved per-GPU GEMM throughput for v after accounting for arithmetic intensity and calibrated GEMM shape efficiency. Thus, W (v) specifies how much computation the operator performs, while Feff (v) specifies how fast that computation can be sustained in practice. We model Feff (v) with a roofline-style throughput corrected for both bandwidth limits and calibrated GEMM shape under-utilization. This captures inefficient GEMM execution at small local batch sizes or narrow tensor-parallel dimensions. The separate terms Tnon-gemm and Topt account for layer operations whose latency does not decrease proportionally with tensor parallelism (e.g., normalization, residual additions, positional embeddings, attention softmax, activation functions) and memory-bandwidth-bound optimizer traffic. Concretely, Feff (v) = min(Fpeak η(v), I(v)BWmem ); I(v) is the operator’s arithmetic intensity and η(v) is a calibrated GEMM shape efficiency factor. In our implementation, η(v) decreases for

small local batch/sequence dimensions and for narrow tensorparallel hidden dimensions, parameterized by γM and γK . The work term W (v) is itself role- and model-specific: generation, inference, and training have different FLOP counts, and actor/reference, critic/value, and reward operators differ further through their output heads and whether backward and recomputation phases are present; optimizer traffic is accounted for separately. Finally, generation receives an additional RLM-specific refinement by splitting its work into a compute-oriented prefill phase and a memory-bandwidthbound decode phase sensitive to weight and KV-cache traffic, which is essential to capture long-response reasoning workloads. We use work terms for different RLM model roles as presented in a recent extensive work–depth analysis [17]. C. Inter-Model Communication and Resharding Explicit communication and redistribution nodes are computed separately from operator nodes. If a C OMM node transfers container a between operators u and v, its runtime is |a|s(a) β; σ(u) is the device allocation of operator α+ min(|σ(u)|,|σ(v)|) u, |σ(u)| is its size in #GPUs (for α − β, the tier-specific point-to-point values are used). The division by the smaller allocation reflects the limiting side in concurrent send/recv. If a R ESHARD node transforms weights between two incompatible layouts on the same device set, we model it as an all-to-all m sb over that allocation TA2A W , |σ(u)| , where Wm sb is the |σ(u)| full parameter size of a given model in bytes and TA2A (M, p) is the collective all-to-all cost for per-GPU message size M across p GPUs. Together with the operator-node formulas from Section V-B, these define the complete time map T over Ĝ. D. Memory Model Many configurations that look attractive under the time model are infeasible in practice as they exceed HBM capacity. We thus use an analytical memory model both to exclude clearly infeasible candidates during search and to softly penalize memory-violating candidates instead of discovering OOMs only at runtime. We assume mixed-precision training [33] with BF16 parameters and FP32 Adam [27] optimizer state, and activation checkpointing [19]. Optimizer state and gradients are sharded across data-parallel ranks via ZeRO-2 [39]. We target high-throughput RLM training on GPU-rich HPC systems, where 3D parallelism and ZeRO-2 sharding enable GPU-resident execution [34, 39]. CPU offloading primarily addresses insufficient GPU memory [40, 41], but introduces host–device transfers and potentially CPU-side optimizer work that can reduce throughput and disrupt the pipeline overlap targeted by RPEG [46]. 1) Parameter and Persistent-State Accounting: We use the dominant parameter counts derived in a recent work [17] and retain the placement-specific byte accounting required here. Let Wshard (v) denote the maximum per-GPU parameter shard size of invocation v, accounting for its pipeline and tensor partitioning and any replicated embeddings or heads. For device-level accounting, Wshard (m, d) denotes the corresponding parameter shard of model m on GPU d under the relevant execution layout.

The persistent footprint depends on how each GPU uses the model. A training GPU retains FP32 master weights and two Adam moment arrays, sharded across Pd data-parallel ranks under ZeRO-2 (12 bytes per parameter before sharding). Its BF16 weights are instead materialized on demand. A GPU serving a trainable model with τ (m) > 0 on an allocation disjoint from training retains a persistent BF16 shadow copy; frozen models such as reference and reward likewise retain BF16 weights. If training and serving allocations overlap, no additional shadow copy is needed. Synchronous serving-only allocations retain no persistent weights between invocations. Thus, the per-model persistent footprint S(m, d) is   12 Wshard (m, d)/Pd , training, S(m, d) = Wshard (m, d) sb , shadow/frozen,  0, no persistent copy.

Here, Pd is the training layout’s data-parallel degree, and sb is the BF16 scalar size in bytes. 2) Transient-State Accounting: In addition to persistent state, each invocation v requires a peak transient allocation A(v), including activations, generation KV cache, BF16 weights materialized on demand, training gradients, and local working buffers. Under ZeRO-2, the sharded BF16 gradient term is Wshard (v)sb /Pd (v). State already included in S(m, d) is excluded from A(v) to avoid double counting. Since invocations sharing a GPU cannot execute simultaneously, only the largest transient allocation on that device needs to be considered. 3) Co-Residency-Aware Device Peaks: The key inter-model issue is that RL pipelines place multiple models on the same physical devices. A GPU may participate, for example, in actor generation, reference inference, and actor training within one overlapping steady-state schedule. Thus, memory cannot be evaluated operator by operator in isolation. We thus define device-level peak memory as persistent coresident state plus the largest transient term among operators mapped to that device: X

peak(d) =

S(m, d) + max A(v) . v placed on d

m placed on d

persistent

}

(1)

}

worst transient

To account for runtime memory not explicitly captured by the analytical model (e.g., allocator reservations, communication workspaces, fragmentation) we reserve a calibrated safety fraction εrt of GPU memory. The effective capacity is Meff = (1 − εrt ) MGPU . Finally, a given configuration is memory-feasible iff every device respects the effective capacity bound: MemFeasible(C) = true ⇐⇒ ∀d ∈ D : peak(d) ≤ Meff . (2)

Figure 3 illustrates this persistent/transient classification for five representative two-GPU cases: (1) synchronous disjoint execution, where inference receives weights transiently via communication and holds no persistent copy; (2) asynchronous disjoint execution, where a shadow pair is active and both sides keep persistent weight copies; (3) partial overlap where one device is shared between training and inference while the other remains inference-only; (4) partial overlap with the complementary overlap pattern; and (5) full colocation

with resharding, where the same devices execute both roles and a layout transformation bridges inference and training. These cases show how bounded asynchrony changes memory behavior: enabling a shadow pair increases static memory by keeping persistent copies on both sides, but it avoids repeated redistribution and allows cross-step overlap. inference phase communica�on communica�on training phase Case 1: No overlap (sync) within RL itera�on 1 of weights of cri�c values within RL itera�on 0 weights inference@0 inference@1 GPU 0 val training@0 GPU 1

MEM 0 MEM 1

weights optim

Case 2: No overlap (async) inference@0 inference@1 GPU 0 val training@0 GPU 1 MEM 0 MEM 1

Case 3: Par�al overlap (i) inference@0 GPU 0 GPU 1 MEM 0 weights,optim#1 optim#2 MEM 1

val

Case 5: Full overlap GPU 0 inference@0 GPU 1

training@0

inference@2 training@1

weights weights,optim, grad,actv

weights#2

weights#1,optim#1,grad#1,actv#1 weights#2,optim#2,grad#2,actv#2

Case 4: Par�al overlap (ii) GPU 0 inference@0 val GPU 1 MEM 0 MEM 1

weights val

weights weights,optim,grad,actv

weights optim

weights#1,optim weights#2

weights optim

weights, optim, grad,actv

inference@1

A. Pipeline Structure Candidates For a fixed RL algorithm, the explicitly searched structural choices are actor–critic sharing and bounded staleness. We represent them by y = (ac, τ ), where ac ∈ {0, 1} determines whether actor and critic parameters are shared, and τ (m) is the allowed staleness of each trainable model m. AC sharing is considered only when both an actor and critic exist. Each τ (m) is chosen from a small user-defined finite set. We denote the resulting set of legal tuples by Ycand . Each y ∈ Ycand determines one RPEG G◦ (y).

weights,optim#1 optim#2

B. Execution Plan Candidates training@0 weights#2

weights

training@0

inference@1 weights#1,optim weights#2

weights,optim,grad,actv

weights

MEM 0 weights#1',optim#1 weights#1,optim#1,grad#1,actv#1 MEM 1 weights#2',optim#2 weights#2,optim#2,grad#2,actv#2

inference@1

weights#1',optim#1 weights#2',optim#2

Time

Resharding

This search has two levels. First I , we enumerate a small set of pipeline-structure candidates G◦ . Second, for each such graph, we search its much larger space of execution plans P II . Algorithm 1 summarizes the complete procedure. In the following, we detail the key aspects of the algorithm; further details are in Appendix A-J.

Cri�c Workload

Communica�on

Idle

op�m: op�mizer state grad: gradients ac�v: ac�va�ons #1: at a given GPU

Fig. 3: Per-device memory classification for 5 allocation and staleness configurations of one trainable model (critic) on two GPUs. In Case 1 (sync, disjoint), GPU 0 receives weights transiently via a C OMM node. In Case 2 (async, disjoint), both GPUs hold persistent weight copies (shadow pair). Cases 3–4 show partial device overlap. Case 5 shows full colocation where a R ESHARD node transforms the weight layout between inference and training. No AI was used to conceive or draw the picture.

VI. J OINT S TRUCTURAL AND E XECUTION O PTIMIZATION The preceding sections define how to evaluate a configuration C = (G◦ , P ). For each such candidate, the transformation pipeline (Section IV) delivers an executable DAG Ĝ. The analytical performance model (Section V) estimates node execution times and per-GPU memory demands, while scheduling (Section IV-E) determines the predicted makespan and score, including memory penalties. During optimization, this evaluation is repeated for each newly proposed configuration; the resulting scores guide the search, while the best memory-feasible configuration found is retained. The optimization problem is defined by four fixed inputs: (1) the RL training algorithm (e.g., PPO, GRPO, or DPO), (2) workload parameters such as batch and sequence sizes, (3) model specifications, and (4) the available hardware. The RL algorithm determines the initial RPEG structure; for example, PPO contains actor, critic, reward, and reference invocations, whereas GRPO has no critic. Together, these inputs establish the initial compact RPEG and its optimization constraints.

For every model invocation v in a fixed G◦ , we enumerate its legal execution choices as Xcand (v) = {(σ, Π, mb) : execution constraints hold}, where σ is the GPU allocation, Π = (Pd , Pp , Pt ) is the DP/PP/TP configuration, and mb is the microbatch count. “Legal” means that the choice can be mapped to the given workload, model, and machine: for example, pipeline stages must receive valid layer partitions, tensor parallelism is kept within a node, and the allocation must contain the GPUs required by Pd Pp Pt . We enumerate aligned contiguous GPU blocks satisfying these constraints and discard obviously infeasible choices early. An execution plan P assigns one legal choice to each invocation v: P (v) = σ(v), Π(v), mb(v) ∈ Xcand (v). The search starts from an initial plan and repeatedly replaces the assignment of one or two invocations, as described next. C. Stochastic Plan Search For each structural candidate G◦ (y), we search its execution plans using simulated annealing. Starting from the current plan P , one search round proposes up to Sn = 6 nearby plans. Each proposal defines a complete configuration C = (G◦ , P ), which is evaluated using the pipeline from Section III. Thus, unrolling, lowering, analytical modeling, scheduling, and scoring are performed for every newly proposed configuration. At the beginning of each search round, the optimizer chooses one of two proposal mechanisms. With probability ppair = 0.28, it uses the two-invocation mechanism described below; otherwise, with probability 1 − ppair = 0.72, it uses the single-invocation mechanism. For a single-invocation move, the optimizer first chooses the invocation v whose execution-plan assignment will be modified. With probability pcp = 0.6, v is sampled from the current schedule’s critical (bottleneck) path with probability proportional to its predicted execution time; with probability 1 − pcp = 0.4, it is sampled uniformly from all invocations. The optimizer then samples alternative assignments for v from Xcand (v), producing up to Sn neighboring plans. The

subsequent evaluation determines whether any of them is better than the current plan. For a two-invocation move, the optimizer instead changes two allocations jointly. It samples a donor invocation vd with probability proportional to its current GPU count |σ(vd )|, so invocations occupying more GPUs are more likely to donate resources. It then chooses a receiver invocation vr among invocations whose current allocations overlap with the donor’s allocation. The donor is assigned a legal smaller allocation, while the receiver is resampled from Xcand (vr ) so that it may use some of the released GPUs. A receiver can already be feasible and still benefit from additional GPUs: feasibility only means that its assignment is legal, whereas a larger allocation may reduce its execution time and hence the overall makespan. The coordinated move allows such a resource redistribution to be tested directly instead of first passing through a plan that leaves the donor’s released GPUs unused. D. Proposal Evaluation and Acceptance Regardless of which proposal mechanism is used, one search round produces up to Sn complete candidate configurations C1′ , . . . , CS′ n . They are evaluated independently, with uncached evaluations executed in parallel when possible. The optimizer then selects C ′ = arg max Scorepen (Ci′ ) ′ Ci

and only this highest-scoring proposal is considered for updating the current plan. Let s = Scorepen (C) and s′ = Scorepen (C ′ ). If s′ > s, C ′ is accepted. Otherwise it is accepted with probability ′ s −s . pacc = min 1, exp |s|t The temperature t > 0 controls how readily the search accepts a temporarily worse configuration: larger t increases exploration, whereas smaller t makes the search more greedy. The normalization by |s| makes the temperature largely independent of the absolute score scale. A sequence of proposal-and-acceptance rounds forms a search chain. The temperature remains fixed within each chain. The first chain uses temperature tinit . Each subsequent, warm-restarted chain uses twarm > tinit , allowing broader exploration around a previously discovered high-scoring plan. If the selected proposal is rejected in ten consecutive rounds, the chain is considered stuck and a new one is started. The search budget is specified either as wall-clock time or as a maximum number of full configuration evaluations. The budget is allocated equally across structural candidates: in wall-clock mode, the remaining time is divided evenly among the remaining y ∈ Ycand ; in evaluation-budget mode, each candidate receives the same fixed budget of uncached configuration evaluations. E. Caching Evaluating a configuration is considerably more expensive than generating a new proposal, so evaluated configurations

Algorithm 1: RPEG optimization Input: Fixed RL algorithm, workload, model, and hardware inputs; legal structural choices Ycand = {(ac, τ )}; wall-clock or uncached-evaluation search budget; temperatures tinit , twarm Output: C ⋆ : highest-scoring memory-feasible configuration found within the search budget ⋆ 1 Initialize cache and C ← ⊥. Define Eval(C) to return (penalized score Scorepen (C), memory feasibility, schedule) via cached unrolling/lowering (Sec. IV), performance modeling (Sec. V), and scheduling/scoring (Sec. IV-E) 2 foreach y ∈ Ycand I do 3 Construct G◦ (y) (Sec. VI-A); For each v ∈ V ◦ (y), enumerate/reuse Xcand (v) by filtering (σ, Π, mb) combinations for legal parallelism, device packing, and HBM feasibility (Sec. VI-B, App. A-J2a) 4 first ← true 5 while budget remains for plan-search on G◦ (y) II do 6 if first then 7 Initialize P using the concurrent baseline: group independent invocations, partition GPUs, and choose P (v) ∈ Xcand (v) to minimize estimated step time (Sec. VI-B, App. A-J2h) 8 t ← tinit ; first ← false 9 else 10 Initialize P from a cached high-scoring feasible plan with a compatible vertex set (Sec. VI-E, App. A-J2h) 11 t ← twarm 12 end 13 (s, ok, Q) ← Eval((G◦ (y), P )) 14 If ok and the score improves C ⋆ , set C ⋆ ← (G◦ (y), P ) 15 rejected ← 0 16 while budget remains and rejected < 10 do 17 if pair kernel selected with ppair = 0.28 then 18 Sample donor u proportional to its GPU count and receiver v from overlapping allocations; for i ≤ Sn = 6, copy Pi′ ← P , shrink Pi′ (u) using Xcand (u), and resample Pi′ (v) from Xcand (v) (Sec. VI-C, App. A-J2f) 19 else 20 Select v using critical-path weights from schedule Q with probability pcp = 0.6, otherwise uniformly; for i ≤ Sn = 6, copy Pi′ ← P and sample Pi′ (v) ∼ Unif(Xcand (v)) (Sec. VI-C, App. A-J2f) 21 end 22 For each Pi′ , evaluate (s′i , ok′i , Q′i ) ← Eval((G◦ (y), Pi′ )) 23 i⋆ ← arg maxi s′i ; s′ ← s′i⋆ 24 Update C ⋆ from all evaluated memory-feasible proposals with higher scores 25 if accept the best proposal with probability pacc (s, s′ , t) (Sec. VI-D, App. A-J2e) then 26 P ← Pi′⋆ ; s ← s′ ; Q ← Q′i⋆ ; ok ← ok′i⋆ ; rejected ← 0 27 else 28 rejected ← rejected + 1 29 end 30 If ok, retain the current P as a possible restart seed 31 end 32 end 33 end ⋆ 34 Polish C by repeatedly testing every x ∈ Xcand (v) for each invocation v, replacing only P (v), and greedily accepting feasible score improvements until none remain (Sec. VI-F, App. A-J2j) III ⋆ 35 return C

are cached using their complete graph and execution-plan description. Repeated proposals therefore reuse their previous score. Structural candidates with the same invocation set also reuse the corresponding Xcand (v) sets. The implementation additionally retains high-scoring feasible plans with compatible invocation sets as seeds for subsequent search chains.

F. Stopping and Final Refinement The search stops when its wall-clock or evaluation budget is exhausted and always retains the the highest-scoring memoryfeasible configuration found so far. Finally III , we deterministically refine that plan. We visit every invocation in a fixed order and try its legal execution choices while keeping all other assignments unchanged; if any choice improves the score, we retain the best one. These passes repeat until an entire pass produces no improvement. Thus, the final “greedy coordinate descent” simply removes any remaining improvement that can be obtained by changing one invocation at a time. VII. C OMPLETION -D RIVEN D ISTRIBUTED RUNTIME The RPEG runtime realizes a configuration C ⋆ selected by the optimized on physical GPUs, preserving its execution dependencies, placement, parallelism, and model-version semantics. Crucially, while predicted execution times guide plan selection, actual node completion events govern execution, allowing the runtime to adapt to varying computation and communication durations without changing the selected plan. Execution proceeds in three stages, see Alg. 2. First A , graph realization unrolls and lowers C ⋆ into an executable DAG Ĝ = (V̂ , Ê) (as in Sec. IV), and assigns concrete GPU/worker bindings and resource requirements to every executable node x ∈ V̂ . Then B , persistent worker groups instantiate the selected DP/PP/TP configurations. Each logical RPEG invocation v ∈ V ◦ reuses its assigned workers across unrolled training steps, avoiding repeated model initialization and distributed-group setup. Third C , completion-driven dispatch executes nodes as their predecessors finish and required physical resources become available. The controller operates directly on the executable DAG, including O P, C OMM, and R ESHARD nodes. A node x becomes ready when all predecessors in Ĝ have completed, and may execute when its exclusive resource set R(x) is available. These resources include assigned workers and GPUs and, for direct transfers, the participating endpoints. Among ready nodes, priority is given to earlier RL steps, then communication/resharding, and finally a fixed topological order. The graph also determines the exact parameter version consumed by each invocation and when intermediate data or obsolete parameter copies can be released, thereby preserving boundedstaleness semantics while enabling safe overlap. Model invocations use PyTorch/Megatron with the selected DP/PP/TP groups. NCCL supports tensor-native parameter transfers and redistribution, while the principal experimental campaign uses Gloo for semantic data transfers. On Clariden, worker processes share a pooled Slurm launch with NCCL communication over OFI/libfabric/CXI. These backend choices realize the graph’s operations and transfers without changing its logical dependencies. Physical realization. Table III summarizes how RPEG’s optimization decisions are translated into physical execution, connecting the selected graph and execution plan to concrete Megatron workers, communication operations, and parameter versions. Execution verification. Since RPEG jointly optimizes graph structure, placement, parallelism, and scheduling, its

Algorithm 2: Completion-driven RPEG execution Input: Optimized configuration C ⋆ = (G◦ , P ), execution horizon Krun Output: Completed execution and runtime trace ⋆ 1 (Ĝ, b, R) ← RealizeAndValidate(C , Krun ) A 2 W ← LaunchPersistentWorkers(b) B 3 remaining(x) ← |pred (x)| for every x ∈ V̂ Ĝ 4 Busy, Launched, Done ← ∅ 5 while Done ̸= V̂ C do 6 Ready ← {x ∈ V̂ \ Launched : remaining(x) = 0} 7 Order Ready by RL step, C OMM/R ESHARD before O P, then topological order 8 foreach x ∈ Ready in this order do 9 if R(x) ∩ Busy = ∅ then 10 Reserve R(x) in Busy 11 Bind inputs and required parameter versions 12 Record dispatch; submit x asynchronously to W 13 Launched ← Launched ∪ {x} 14 end 15 end 16 Wait for newly completed nodes 17 foreach newly completed node x do 18 Record completion; publish outputs and model versions 19 Release R(x) from Busy 20 Done ← Done ∪ {x} 21 foreach successor z of x do 22 remaining(z) ← remaining(z) − 1 23 end 24 Release artifacts and parameter copies with no remaining consumers 25 end 26 end 27 Stop workers; return execution trace

performance gains are meaningful only if the runtime executes the configuration actually selected by the optimizer. Otherwise, unintended changes in execution could invalidate the predicted-to-measured performance correspondence or alter training semantics. RPEG therefore identifies each optimized plan (and fused microtask schedule, if applicable) by a hash and verifies its physical GPU bindings, parallelism, dependencies, parameter versions, and execution order against the runtime trace. VIII. E VALUATION We show the advantages of RPEG along four complementary dimensions: accuracy of the analytical time/memory model, schedule quality, improvements over prior systems, and optimizer efficiency & stability. Experiments use two HPC platforms: an A100 system with eight 40-GB GPUs, intra-node NVSwitch, and four 200-Gb/s InfiniBand interfaces per node; and a GH200 system with four GPUs and four Slingshot-11 interfaces per node. Due to space constraints, further setup details are in Appendix B. A. Model Validation We first validate both execution-time and peak-memory predictions against measurements on NVIDIA GH200 GPUs. The study covers 123 configurations spanning three LLaMA model sizes (3B, 7B, 13B), batch sizes B ∈ {32, . . . , 2048}, GPU counts from 1 to 128, and 25 distinct (Pd , Pp , Pt ) strategies, including DP-only, TP-only, PP-only, and 3D-hybrid configurations. We fix Lp = 512, Lr = 128, and R = 1. All runs use BF16, FlashAttention-2, activation checkpointing

Optimized RPEG component Runtime action G◦ and its typed edges Placement σ(v) Parallelism Π(v) Microbatch count mb(v) Publication edges and τ (m) Cross-layout dependencies Resource set R(x)

Instantiate step-indexed model operations, data dependencies, and explicit transfer nodes. Bind the invocation’s worker ranks to its selected physical GPUs. Form DP replica groups, PP layer stages, and TP tensor shards/collectives in Megatron. Split each DP replica’s local batch into the selected number of microbatches and issue the corresponding model operations. Resolve each consumer’s required parameter version and install its corresponding shards. Route tensor fragments between producer and consumer layouts through C OMM/R ESHARD operations. Reserve the GPUs and worker groups needed by x until its completion; transfers reserve their participating endpoints.

TABLE III: Example concrete realization of RPEG components within its runtime.

Training

Generation

MAPE=19% ρ=0.913 n=38

Predicted (GB)

104

Inference MAPE=15% ρ=0.974 n=51

103 103

104

102

Measured (ms) 1-GPU

103

104

Measured (ms) DP

103

102

101

102

101

Measured (GB)

1-GPU

102

Measured (GB) PP

Fig. 5: Predicted vs. measured peak GPU memory. The model achieves 15.1% MAPE with Spearman ρ = 0.972, correctly ordering memory use for 93.7% of configuration pairs. The remaining underprediction in inference and generation reflects unmodeled CUDA context overhead (∼2–3 GB per GPU). 1e4

Lr = 128, B = 512

1e4

+23%

2.4

OOM

1.8

OOM

Lr = 384, B = 256 +44%

2.0

+25%

1e4

+44%

Lr = 896, B = 128 +50%

+50%

1.2

1.5

0.9

OOM

+33%

1.2

1.0

0.6

0.5

0.0

OOM

16G 13B/7B

32G 34B/7B

64G 70B/7B

0.0 8G 7B/7B 1e4

+5%

2.4

16G 13B/7B

32G 34B/7B

64G 70B/7B

8G 7B/7B 1e4

+22% OOM

OOM

+44%

2.0

1.6

OOM

1.2

1.0

0.8

0.6

0.5

32G 7B/34B

64G 7B/70B

1e4 4

16G 7B/13B

32G 7B/34B

64G 7B/70B

+39%

OOM OOM

1.6

+42%

1.5 +52%

1.0

+39%

+48%

0.8

0.5

0.0

32G 64G 128G 13B/13B 34B/34B 70B/70B

Greedy heuristic

64G 7B/70B

2.0

16G 7B/7B

32G 7B/34B

16G 32G 64G 128G 7B/7B 13B/13B 34B/34B 70B/70B Model sizes shown as Actor / Critic Concurrent heuristic RPEG (throughput)

+44%

16G 7B/7B

Scale Both

OOM

16G 7B/13B

+60%

2.4

8G 7B/7B 1e4

1e4 +48%

+22%

+39%

0.0 8G 7B/7B

3.2

+25%

64G 70B/7B

+47%

0.4

0.0 16G 7B/13B

32G 34B/7B

+53%

OOM

1.5

0.0

16G 13B/7B +49%

+24%

1.8

8G 7B/7B

+23%

0.3

0.0 8G 7B/7B 1e4

0.6

Scale Critic

We also compare the execution plans selected by RPEG against heuristic baselines under the analytical model. This isolates the effect of joint optimization of placement, parallelism, overlap, and bounded asynchrony from runtime-system overheads. We compare to two baselines: the greedy baseline chooses a feasible strategy for each invocation on the full GPU set and executes invocations sequentially; the concurrent baseline assigns independent invocations disjoint allocations. 1) Weak Scaling: The weak scaling results are in Figures 6 and 7; we increase actor size, critic/reward size, or both as

MAPE=14% ρ=0.941 n=38

Scale Actor

B. Predicted Schedule Advantages

100

Measured (GB)

for training, and ZeRO-2 gradient sharding for Pd > 1. Hardware parameters come from the automated calibration suite described in Section V-A3. Execution Time. Figure 4 compares predicted and measured operator latency for generation, inference, and training. The model attains 15.5% mean absolute percentage error (MAPE) and Spearman rank correlation ρ = 0.990, correctly ranking 95.4% of all configuration pairs. NCCL microbenchmarks directly calibrate the communication parameters (α, β); the operator measurements then test their combined effect with compute and other costs including TP AllReduce, PP pointto-point transfers, and ZeRO-2 gradient synchronization. Peak Memory. Figure 5 compares predicted and measured per-device peak memory. The memory model achieves 15.1% MAPE with Spearman ρ = 0.972, correctly identifying which configuration uses more memory in 93.7% of pairs. The systematic underprediction for inference and generation reflects unmodeled CUDA context overhead of roughly 2–3 GB per GPU, for which the εrt margin in Meff reserves capacity.

Training

MAPE=14% ρ=0.987 n=51

101

104

Fig. 4: Predicted vs. measured operator latency. Each point is one (workload, strategy) configuration. Marker shape indicates the parallelism strategy. The solid line denotes exact agreement and the shaded band marks ±25% error. Overall MAPE is 15.5% with Spearman ρ = 0.990 (n = 123), correctly ranking 95.4% of all configuration pairs.

Inference

MAPE=19% ρ=0.991 n=34

Measured (ms)

Throughput (tokens/s)

Predicted (ms)

Generation MAPE=13% ρ=0.945 n=34

32G 64G 128G 13B/13B 34B/34B 70B/70B

RPEG (util)

Fig. 6: Weak scaling: predicted throughput (tokens/s) across three scaling experiments (rows) and three workload regimes (columns). OOM marks infeasible baselines. Percentage labels indicate the optimizer’s gain over the concurrent heuristic. Models: LLaMA-2.

GPU count grows from 8 to 128. The (B, Lr ) settings vary workload characteristics. RPEG improves throughput by 5–60% over the concurrent heuristic where that baseline is feasible, and finds feasible plans in all tested cases, including those with OOM baselines. Gains are largest for long responses, which expose more generation–assessment overlap. Next, RPEG does not optimize for raw device occupancy alone. Instead, it assigns generation, assessment, and training (and their subunits) to different GPU subsets with different 3D strategies. The resulting GPU utilization-aware plans attain 78–100% scheduled

Lr = 128, B = 512 +1pp

+7pp

1.00

+13pp

Lr = 896, B = 128

OOM

+2pp

+19pp

+11pp

1.00

0.75

0.50

0.25

0.00

0.00 16G 13B/7B

+2pp

1.00

32G 34B/7B

+1pp

64G 70B/7B OOM

OOM

+6pp

+1pp

16G 13B/7B

32G 34B/7B OOM

+2pp

64G 70B/7B OOM +9pp

+10pp

8G 7B/7B

1.00

0.75

0.50

0.25

16G 7B/13B

+6pp

32G 7B/34B

+3pp

1.00

OOM

+4pp

64G 7B/70B

+8pp

1.00

16G 7B/13B

32G 7B/34B

64G 7B/70B

+19pp

+8pp

+7pp

1.00

64G 70B/7B

+7pp

+20pp

0.75

0.50

0.25

0.00

0.00 16G 32G 64G 128G 7B/7B 13B/13B 34B/34B 70B/70B Model sizes shown as Actor / Critic Concurrent heuristic RPEG (throughput)

32G 7B/34B

64G 7B/70B

+12pp

+6pp

+15pp

+25pp

16G 7B/7B

32G 64G 128G 13B/13B 34B/34B 70B/70B

GPU utilization. Relative to throughput-only optimization, they trade 1–3% throughput for up to 25 percentage points higher utilization. This is scheduled resource occupancy, not a measurement of GPU arithmetic utilization. 2) Strong Scaling: In the strong-scaling regime (Figure 8) the model is fixed and GPU count increases. With 13B models, B = 1024, Lp = 128, and Lr = 384, RPEG finds a feasible 8-GPU plan at 8,966 tokens/s while both baselines are OOM. From 16 to 128 GPUs, predicted throughput is 20–49% higher than the greedy baseline, reaching 104,930 tokens/s at 128 GPUs. Improvement over greedy (%)

Throughput (tokens/s)

(a) Throughput

80k 60k 40k 20k 0k

OOM

128

Number of GPUs Greedy heuristic

Concurrent heuristic

(b) Optimizer advantage

+49% +49%

40 30

+27%

+29% +29%

+25% +25%

20 +20%

10 8

16 GPU, 13B (PFLOP/s)

32 GPU, 34B (PFLOP/s)

0.233

0.656

64 GPU, 70B (PFLOP/s)

Avg.

1.339

–

0.203 (-12.8%) 0.569 (-13.2%) 0.189 (-18.8%) 0.497 (-24.3%) 0.238 (+2.1%) 0.672 (+2.4%)

1.318 (-1.5%) -9.2% 1.270 (-5.2%) -16.1% 1.105 (-17.5%) -4.3%

A7: polish phase 0.233 (+0.0%) 0.647 (-1.5%) A9: geometric cooling 0.233 (+0.0%) 0.607 (-7.5%)

1.269 (-5.2%) -2.2% 1.327 (-0.9%) -2.8%

A1: gen-inf overlap A2: async execution A3: mb tuning

A4/A5/A6/A8

< 1% average impact across all tested configurations < 1%

RPEG (util)

Fig. 7: Weak scaling: predicted GPU utilization for the same configurations as Figure 6. The throughput×utilization objective (red) consistently improves utilization over the throughput-only objective (green) by up to +25 pp (absolute difference between utilization percentages). Models: LLaMA-2.

100k

TABLE IV: Ablation study. Search space features (A1–A3) dominate. The last row (A4–A8) summarises pair moves, warm restarts, critical-path vertex selection, and soft OOM penalty, which show < 1% impact across 9 configurations spanning memory-constrained and wide-search-space regimes. We describe the details of each ablation configuration in Appendix B. Ablation

16G 7B/13B

Scale Both

0.75

Greedy heuristic

32G 34B/7B

mb = 1 for Pp = 1; otherwise, it fixes mb = Pp for generation/inference and mb = 2Pp for training. All A1–A9 modifications are identified in the table.

Full system (τ =1)

+9pp

0.50

32G 64G 128G 13B/13B 34B/34B 70B/70B

+21pp

+13pp

8G 7B/7B

0.75

16G 7B/7B

16G 13B/7B

0.00 8G 7B/7B

OOM OOM

+21pp

0.25

0.00 8G 7B/7B

+17pp

Scale Critic

0.75

0.00

+6pp

0.00 8G 7B/7B

1.00

+11pp

Scale Actor

+1pp

0.75

8G 7B/7B

GPU Utilization

Lr = 384, B = 256

OOM

1.00

128

Number of GPUs RPEG (throughput)

RPEG (util)

Fig. 8: Strong scaling. (a) Absolute throughput (tokens/s). At 8 GPUs, both baselines are OOM. (b) Percentage improvement over the greedy heuristic. The optimizer’s advantage grows with GPU count. Model: LLaMA-2 13B.

3) Ablations: Table IV varies one feature at a time (out of nine) across four scenarios: memory-constrained execution (16 GPUs, 50-GB HBM limit), a wide search space (128 GPUs), heterogeneous model sizes (34B actor and 7B critic/reward), and deep pipelines (70B critic/reward across two racks). Each setting uses five seeds, for 200 runs in total. The largest average losses in predicted performance come from disabling bounded asynchrony (A2, 16.1%), generation–inference overlap (A1, 9.2%), and microbatch tuning (A3, 4.3%). A3 fixes

The effects are workload-dependent, e.g., A3 improves performance by 2.1–2.4% in the two smaller configurations but reduces it by 17.5% in the 70B case. Search-algorithm changes have smaller average effects: disabling final polishing (A7) and using geometric cooling (A9) reduce performance by 2.2% and 2.8%, respectively, while A4/A5/A6/A8 have less than 1% average impact. Thus, most of the gain is due to the available execution choices, not to any single search heuristic. A representative case study illustrates the trade-off between minimizing individual operator times and the overall schedule makespan (Figure P 9). We quantify effective operator parallelism as Peff = v∈V̂Op Top (v)/MS(Ĝ), i.e., the sum of individual operator durations divided by the total schedule makespan. The greedy heuristic minimizes individual operator times under a fully colocated allocation, but executes them sequentially (Peff = 1). The concurrent heuristic overlaps operators (Peff = 1.6–1.7), but splitting GPUs among them increases their individual execution times, largely negating the benefit of overlap. RPEG also makes locally suboptimal choices: most operators execute more slowly than under greedy, yet their overlap (Peff = 1.8–1.9) reduces overall makespan by 17–22%. Notably, actor generation is an exception: RPEG accelerates it despite allocating fewer GPUs, because the larger local batch improves utilization during memory-bound decoding. Thus, minimizing individual operator times or maximizing overlap alone is insufficient; the objective must be the global schedule makespan. C. Measured Performance Improvements We next assess end-to-end performance. 1) Comparison with Existing Systems: We first evaluate RPEG against two existing systems under synchronous PPO, without bounded asynchrony. On A100 GPUs, we compare RPEG against ReaL (Figure 10), using models from 7B to 70B. RPEG is faster in 11 of the 13 jointly executable configurations and comparable in two (7B/16 GPUs and 32B/32 GPUs). Throughput improves by 1.4×–2×, with up to 30% higher effective GPU utilization. For 32B/16 GPUs, RPEG also finds and executes a memory-feasible plan where ReaL finds none. On GH200 GPUs, we compare RPEG against

Time (s) ↓ Llama 3.1 8B / 8B

The greedy heuris�c minimizes individual operator �mes under a fully colocated alloca�on, but executes them sequen�ally

Throughput (tokens/s) ↑ Llama 3.3 70B / 70B

verl RPEG

Llama 3.1 8B / 8B

Llama 3.3 70B / 70B

10000

5000

128 256

2000

Number of GPUs

Llama 3.1 8B / 8B

100

Llama 3.3 70B / 70B

The concurrent heuris�c overlaps operators, but spli�ng GPUs among them increase their individual execu�on �mes, largely nega�n the beneﬁt of overlap

GPU occupancy (%) ↑

128 256

Number of GPUs

128 256

Number of GPUs

Fig. 11: RPEG versus VERL. System: GH200s.

their respective baselines by 1.250×–1.471× (mean 1.365×), reducing graph execution time by approximately 20–32%. This establishes the benefit of optimization within each structural variant, independently of switching variants.

Minimizing individual operator �mes or maximizing overlap alone is insuﬃcient; the objec�ve must be the global schedule makespan, as in RPEG

TABLE V: Benefits of RPEG optimization within fixed actor– critic architectures. Model: Llama 3.1 8B. System: GH200s.

Fig. 9: Predicted per-operator wall-time breakdown for LLaMA-2 13B PPO on 16, 32, and 64 GPUs; B = 1024, Lp = 128, and Lr = 384. The greedy heuristic (left) executes all operators sequentially. The concurrent heuristic (center) partitions GPUs but does not reduce makespan significantly. The optimizer (right) reduces makespan by 17–22% via GPU partitioning, gen-inf overlap, and asynchronous execution.

Time (s) ↓

LLaMA-2 7B / 7B

1000

LLaMA-2 13B / 13B

Qwen2.5 32B / 3B

LLaMA-2 70B / 7B

64 128

ReaL RPEG

500

Throughput (tokens/s) ↑

500

GPU occupancy (%) ↑

100

64 128

Number of GPUs

Fig. 10: RPEG versus ReaL. System: A100s. VERL (Figure 11), using Llama 3.1 8B and Llama 3.3 70B with separate actor/critic models. RPEG achieves 1.024×– 2.477× speedup (median 1.488×). 2) Joint Pipeline-Structure and Execution Optimization: We next investigate deeper the benefits of jointly optimizing pipeline structure G◦ and execution plan P . a) Optimization Within A Fixed Training Structure: We compare feasible, fully colocated baseline configurations against RPEG-optimized configurations executed using the same runtime on GH200 GPUs, eliminating cross-framework differences. For each baseline, we preserve the AC architecture (shared or separate), synchronous execution, model size, workload, and GPU budget, while optimizing placement, DP/PP/TP parallelism, microbatching, and applicable execution refinements. Seven Llama 3.1 8B baselines yield 20 measured comparisons (Table V). All optimized configurations outperform

Nodes

AC structure

2 4 4 8 8 16 16

Shared Shared Separate Shared Separate Shared Separate

Baseline time RPEG time Speedup [seconds/step] [seconds/step] 11.678 14.311 16.752 12.822 16.770 13.189 16.852

8.026–8.158 1.432–1.455× 10.726–10.960 1.306–1.334× 11.392–12.436 1.347–1.471× 9.129–9.167 1.399–1.404× 11.837–12.780 1.312–1.417× 10.479–10.552 1.250–1.259× 12.244–13.031 1.293–1.376×

b) Optimization Across Structural Variants: We summarize the effects of structural optimizations using Llama 3.1 8B and Llama 3.3 70B on GH200 GPUs. Bounded asynchrony (τ = 1) improves every available same-architecture comparison against synchronous RPEG, accelerating separatemodel execution by 1.346×–1.429× at 8B and 1.217×– 1.662× at 70B. AC sharing generally reduces execution time by eliminating duplicated model computation and state, by 12– 34% (1.14×–1.51× speedup). Stage fusion preserves separate actor/critic parameters while interleaving their execution: a topology-matched 8B test reduces runtime by 7%, although expensive resharding can reverse the benefit. Importantly, for a given fixed PPO configuration, RPEG preserves outcome accuracy: it optimizes the scheduling of PPO operations without changing the underlying generation, reward computation, advantage estimation, or loss functions; bounded asynchrony only changes parameter freshness. Following the evaluation methodology of AsyncFlow, we compare RPEG’s asynchronous execution (τ = 1) against synchronous VERL in terms of mean rollout reward and response length, using comparable PPO hyperparameters (including update rate, batch size, learning rate, and maximum sequence lengths), generation settings, and old-policy log-probability computation. RPEG achieves a mean reward on the final rollout within 0.47% relative difference of the VERL baseline, providing empirical evidence of comparable training outcomes under the evaluated settings. D. Optimizer Efficiency and Stability A richer optimization space is only useful if it can still be searched efficiently. RPEG remains practical because candidate evaluation is analytical and heavily cached. Search Cost. A convergence study spanning 256–2048 GPUs (Figure 12, left) shows that the optimizer reduces predicted makespan by approximately 14–19% relative to the

TABLE VI: Design dimensions of representative RLM-training systems. “ ” denotes explicit support, “ ” partial or restricted support, and “é” no first-class support in the cited design. Brief annotations indicate the mechanism or scope of support. Cross-iteration asynchrony

Framework

Intra-iteration overlap

Sub-invocation execution

Execution plan search

ReaL [31] é (indep. ops) é HybridFlow/VERL [45] (VERL runtime) é RLHFuse [63] é (fusion) (sample+mb) AReaL [22] (rollout/train overlap) é (samples) AsyncFlow [24] (streaming pipeline) (pipeline) Laminar [44] é (samples) StaleFlow [29] (bounded) (disaggregation) (samples) RLinf [61] (pipelining) (elastic pipeline) (microflows) OrchestrRL [49] (1-step overlap) (disaggregation) (generation) Libra [18] (rollout/train overlap) (disaggregation) (requests) RPEG [This work]

(bounded)

(place+3D+mb) é (place+3D) é (fusion plans) (fusion) (resource split) é (load balance) é (repacking) é (rollout coord.) é (place+pipeline) (flow rewrite) (generation only) é (stage-level) é

(phase+ microtask)

(fusion)

(place+3D+mb)

7B / 256 GPUs 13B / 512 GPUs 34B / 1024 GPUs 70B / 2048 GPUs

0.975 0.950

Within ≈300 seconds, all configurations achieve makespans within 5% of their respective results after one hour of searching.

0.925 0.900 0.875 0.850 0.825 0.800

10−1

100

101

102

Search time (s)

103

104

100%

Runs within gap (%)

Makespan relative to initial plan

initial plan. Importantly, within just 300 seconds, even at 2048 GPUs, the best discovered plan is already within 5% of the makespan achieved after a full hour of searching. Thus, most optimization benefits are obtained within minutes. For perspective, amortizing a one-time 300-second (600-second) search over 1,000 PPO updates at our measured 70B/64-node rate of 23.7 seconds/update yields an estimated overhead of only 1.3% (2.5%) of total runtime. Run-to-Run Stability. To quantify robustness, we evaluate 200 independent runs per configuration over scales from 64 to 1024 GPUs. Figure 12 (right) shows score gaps relative to the best result found across seeds for each configuration. The coefficient of variation ranges from 0.27% to 2.31%. All runs at 64–512 GPUs are within 6% of the best observed score; at 1024 GPUs they are within 11%. 1.000

Pipeline structure search

(AC+fusion+ staleness)

Performance + memory model

Plan verification

Resource-eff. objective

é é é é (version control) é é (staleness checks) é é é é é (staleness checks) é (profiled) é é (compute+net sim.) é (network cost) (analytical) (state sync.) (resource split) (profiled) (simulation) (simulation)

(calibrated, analytical)

é é é

(hash+versioning)

(throughput× utilization)

objectives, which explicitly consider resource use beyond raw runtime or throughput. Existing systems expose different subsets of these capabilities [18, 22, 24, 29, 31, 44, 45, 49, 61, 63]. Our work extends this landscape along three main axes. First, RPEG exposes a richer graph configuration that enables searching over AC-sharing and bounded-staleness settings together with per-invocation placement, parallelism, and microbatching, all while supporting fusion and sub-invocation response-length refinements. Second, RPEG couples this space to a comprehensive analytical model that captures heterogeneous computation, communication and resharding, hierarchical hardware, co-resident and versioned state, and explicit throughput–resource-efficiency tradeoffs. Third, RPEG provides an execution-faithful runtime: optimized graph structure, placement, model versions, transfers, and resource ownership are reconstructed and verified during distributed execution. Table VI provides a more fine-grained comparison.

80%

All seeds fall within 6% of the best at 64–512 GPUs and 11% at 1024 GPUs

60%

64 GPUs 128 GPUs 256 GPUs 512 GPUs 1024 GPUs

40% 20% 0%

Gap to best observed score (%)

Fig. 12: (Left: optimizer search effectiveness study) Best discovered predicted makespan relative to the initial plan (lower is better) versus search time, across 256–2048 modeled GPUs. (Right: optimizer stability study) Cumulative distribution of gaps to the best found outcome across 200 independent seeds per configuration (600-second searches).

IX. R ELATED W ORK Recent RLM-training systems can be compared along several complementary architectural features: cross-iteration asynchrony, which overlaps different RL iterations; intraiteration overlap, which overlaps or fuses generation, assessment, and training within an iteration; sub-invocation execution, which schedules units smaller than a whole model invocation; execution plan search, which automatically selects placement, parallelism, or microbatching; pipeline structure search, which automates changing the pipeline structure or execution semantics themselves; performance and memory modeling, which predicts candidate execution costs; plan verification, which checks that the selected logical and physical plan is faithfully realized at runtime; and resource-efficiency

X. C ONCLUSION RLM training is no longer a single-model scaling problem: it is a heterogeneous, multi-model, asynchronous pipeline optimization problem. RPEG addresses this by making pipeline structure, overlap, placement, and parallelization explicit within one graph-centric representation, and by coupling that representation to calibrated HOC-focused analytical models of time and memory. This enables the optimizer to search jointly over structural and execution decisions, yielding better schedules than prior baselines and translating those gains into measured improvements in throughput, makespan, and utilization on real hardware. More broadly, RPEG points toward a general systems layer beyond PPO-style RLM training, jointly optimizing heterogeneous operators, bounded asynchrony, communication hierarchy, and cost. Such a layer could support multi-agent and tool-augmented LLM systems, training with richer verification or search, and inference-time reasoning. Thus, RPEG is not merely a training optimizer, but a step toward graph-native orchestration of future large-scale LLM systems. ACKNOWLEDGMENTS We thank Hussein Harake, Colin McMurtrie, Mark Klein, Angelo Mangili, and the whole CSCS team granting access to the Ault and Alps machines, and for their excellent technical

support. We gratefully acknowledge Polish high-performance computing infrastructure PLGrid (HPC Center: ACK Cyfronet AGH) for providing compute facilities within computational grant no. PLG/2026/019437; we also thank Łukasz Flis and the whole Cyfronet team for their excellent technical support. We thank Timo Schneider for help with infrastructure at SPCL. This project received funding from the European Research Council (Project PSAP, No. 101002047). This project was supported by the ETH Future Computing Laboratory (EFCL), financed by a donation from Huawei Technologies. ChatGPT [37] served as an editorial tool for this manuscript, while ideas and content are original work of the authors. R EFERENCES [1] Anthropic, “How We Built Our Multi-Agent Research System,” https://www.anthropic.com/engineering/multi-a gent-research-system, Jun. 2025, accessed 2026-10-18. [2] L. Bai, Z. Huang, X. Wang, J. Sun, R. Mihalcea, E. Brynjolfsson, A. Pentland, and J. Pei, “How do ai agents spend your money? analyzing and predicting token consumption in agentic coding tasks,” Oct. 2026, arXiv:2604.22750. [Online]. Available: https: //arxiv.org/abs/2604.22750 [3] T. Bai, Y. Bai, Y. Bao, S. H. Cai, Y. Cao, Z. Chai, Y. Charles, H. S. Che, C. Chen, G. Chen, et al., “Kimi K2.5: Visual Agentic Intelligence,” Aug. 2026, arXiv:2602.02276. [Online]. Available: https://arxiv.org/abs/2602.02276 [4] Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan et al., “Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback,” Apr. 2022, arXiv:2204.05862. [Online]. Available: https://arxiv.org/abs/2204.05862 [5] T. Ben-Nun and T. Hoefler, “Demystifying Parallel and Distributed Deep Learning: An In-depth Concurrency Analysis,” ACM Comput. Surv., vol. 52, no. 4, pp. 65:1–65:43, Aug. 2019. [Online]. Available: https: //doi.org/10.1145/3320060 [6] M. Besta, J. Barth, E. Schreiber, A. Kubicek, A. Catarino, R. Gerstenberger, P. Nyczyk, P. Iff, Y. Li, S. Houliston et al., “Reasoning Language Models: A Blueprint,” Jun. 2025, arXiv:2501.11223. [Online]. Available: https: //arxiv.org/abs/2501.11223 [7] M. Besta, N. Blach, A. Kubicek, R. Gerstenberger, M. Podstawski, L. Gianinazzi, J. Gajda, T. Lehmann, H. Niewiadomski, P. Nyczyk, and T. Hoefler, “Graph of Thoughts: Solving Elaborate Problems with Large Language Models,” Proceedings of the AAAI Conference on Artificial Intelligence, vol. 38, no. 16, pp. 17 682– 17 690, Mar. 2024. [Online]. Available: https://ojs.aaai.o rg/index.php/AAAI/article/view/29720 [8] M. Besta, S. Chandran, R. Gerstenberger, M. Lindner, M. Chrapek, S. H. Martschat, T. Ghandi, P. Iff, H. Niewiadomski, P. Nyczyk et al., “Psychologically Enhanced AI Agents,” Sep. 2025, arXiv:2509.04343. [Online]. Available: https://arxiv.org/abs/2509.04343

[9] M. Besta, R. Gerstenberger, M. Fischer, M. Podstawski, N. Blach, B. Egeli, G. Mitenkov, W. Chlapek, M. Michalewicz, H. Niewiadomski, J. Mueller, and T. Hoefler, “The graph database interface: Scaling online transactional and analytical graph workloads to hundreds of thousands of cores,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, ser. SC ’23. Denver, CO, USA: Association for Computing Machinery, Nov. 2023, pp. 22:1–22:18. [Online]. Available: https://doi.org/10.1145/3581784.3607068 [10] M. Besta, R. Gerstenberger, E. Peter, M. Fischer, M. Podstawski, C. Barthels, G. Alonso, and T. Hoefler, “Demystifying Graph Databases: Analysis and Taxonomy of Data Organization, System Designs, and Graph Queries,” ACM Comput. Surv., vol. 56, no. 2, pp. 31:1–31:40, Sep. 2023. [Online]. Available: https://doi.org/10.1145/3604932 [11] M. Besta, P. Iff, F. Scheidl, K. Osawa, N. Dryden, M. Podstawski, T. Chen, and T. Hoefler, “Neural Graph Databases,” in Proceedings of the First Learning on Graphs Conference (LOG ’22), ser. Proceedings of Machine Learning Research, B. Rieck and R. Pascanu, Eds., vol. 198. Virtual Event: PMLR, Dec. 2022, pp. 31:1–31:38. [Online]. Available: https://proceedings.mlr. press/v198/besta22a.html [12] M. Besta, Ł. Jarmocik, O. Hrycyna, S. Klaiman, K. Maczka, R. Gerstenberger, J. Müller, P. Nyczyk, H. Niewiadomski, and T. Hoefler, “GraphSeek: NextGeneration Graph Analytics with LLMs,” Mar. 2026, arXiv:2602.11052. [Online]. Available: https://arxiv.org/ abs/2602.11052 [13] M. Besta, A. Kubicek, R. Gerstenberger, M. Chrapek, R. Niggli, P. Okanovic, Y. Zhu, P. Iff, M. Podstawski, L. Weitzendorf et al., “Multi-Head RAG: Solving Multi-Aspect Problems with LLMs,” Feb. 2026, arXiv:2406.05085. [Online]. Available: https://arxiv.org/abs/2406.05085 [14] M. Besta, F. Memedi, Z. Zhang, R. Gerstenberger, G. Piao, N. Blach, P. Nyczyk, M. Copik, G. Kwaśniewski, J. Müller et al., “Demystifying Chains, Trees, and Graphs of Thoughts,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 47, no. 12, pp. 10 967–10 989, Dec. 2025. [Online]. Available: https://ieeexplore.ieee.org/document/11123142 [15] M. Besta, L. Paleari, M. Copik, R. Gerstenberger, A. Kubicek, P. Nyczyk, P. Iff, E. Schreiber, T. Srindran, T. Lehmann et al., “CheckEmbed: Effective Verification of LLM Solutions to Open-Ended Tasks,” Jul. 2025, arXiv:2406.02524. [Online]. Available: https://arxiv.org/ abs/2406.02524 [16] M. Besta, L. Paleari, J. H. A. Jiang, R. Gerstenberger, Y. Wu, J. G. Hannesson, P. Iff, A. Kubicek, P. Nyczyk, D. Khimey et al., “Affordable AI Assistants with Knowledge Graph of Thoughts,” Oct. 2025, arXiv:2504.02670. [Online]. Available: https://arxiv.org/abs/2504.02670 [17] M. Besta, L. Schmidt, L. Nonino, R. Gerstenberger,

P. Pang, P. Okanovic, A. Kubicek, T. Chen, B. Lipshitz, and T. Hoefler, “Performance foundations of parallel & distributed reasoning language models,” Aug. 2026, arXiv:2608.27046. [Online]. Available: https://arxiv.org/ abs/2608.27046 [18] K. Chen, X. Tan, J. Li, Z. Zhou, C. Li, J. Liu, M. Jie, J. Jiang, and H. Xu, “Libra: Efficient Resource Management for Agentic RL Post-Training,” Sep. 2026, arXiv:2606.03077. [Online]. Available: https://arxiv.org/abs/2606.03077 [19] T. Chen, B. Xu, C. Zhang, and C. Guestrin, “Training Deep Nets with Sublinear Memory Cost,” Apr. 2016, arXiv:1604.06174. [Online]. Available: https://arxiv.org/ abs/1604.06174 [20] B. Cottier, R. Rahman, L. Fattorini, N. Maslej, T. Besiroglu, and D. Owen, “The Rising Costs of Training Frontier AI Models,” Feb. 2025, arXiv:2405.21015. [Online]. Available: https://arxiv.org/abs/2405.21015 [21] A. Du, B. Gao, B. Xing, C. Jiang, C. Chen, C. Li, C. Xiao, C. Du, C. Liao, C. Tang et al., “Kimi k1.5: Scaling Reinforcement Learning with LLMs,” Jun. 2025, arXiv:2501.12599. [Online]. Available: https: //arxiv.org/abs/2501.12599 [22] W. Fu, J. Gao, X. Shen, C. Zhu, Z. Mei, C. He, S. Xu, G. Wei, J. Mei, W. Jiashu, T. Yang, B. Yuan, and Y. Wu, “AREAL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning,” in Proceedings of the Thirty-Ninth Annual Conference on Neural Information Processing Systems (NeurIPS ’25), ser. Advances in Neural Information Processing Systems, D. Belgrave, C. Zhang, H. Lin, R. Pascanu, P. Koniusz, M. Ghassemi, and N. Chen, Eds., vol. 38. San Diego, CA, USA: Curran Associates, Dec. 2025, pp. 36 256–36 282. [Online]. Available: https://proceedings.neurips.cc/paper files/pap er/2025/hash/33c00862bfa29ac72ecf630a41e19352-Abs tract-Conference.html [23] D. Guo, D. Yang, H. Zhang, J. Song, P. Wang, Q. Zhu, R. Xu, R. Zhang, S. Ma, X. Bi et al., “DeepSeek-R1 Incentivizes Reasoning in LLMs through Reinforcement Learning,” Nature, vol. 645, no. 8081, pp. 633–638, Sep. 2025. [Online]. Available: https: //www.nature.com/articles/s41586-025-09422-z [24] Z. Han, A. You, H. Wang, K. Luo, G. Yang, W. Shi, M. Chen, S. Zhang, Z. Lan, C. Deng et al., “AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training,” Sep. 2026, arXiv:2507.01663. [Online]. Available: https://arxiv.org/abs/2507.01663 [25] R. W. Hockney, “The Communication Challenge for MPP: Intel Paragon and Meiko CS-2,” Parallel Computing, vol. 20, no. 3, pp. 389–398, Mar. 1994. [Online]. Available: https://www.sciencedirect.com/scie nce/article/abs/pii/S0167819106800219 [26] J. Kim, B. Shin, J. Chung, and M. Rhu, “The Cost of Dynamic Reasoning: Demystifying AI Agents and TestTime Scaling from an AI Infrastructure Perspective,” in Proceedings of the IEEE International Symposium on High-Performance Computer Architecture, ser. HPCA

’26. Sydney, Australia: IEEE Press, Feb. 2026, pp. 1–16. [Online]. Available: https://ieeexplore.ieee.org/do cument/11408569 [27] D. P. Kingma and J. Ba, “Adam: A Method for Stochastic Optimization,” in Proceedings of the 3rd International Conference on Learning Representations, ser. ICLR ’15, San Diego, CA, USA, May 2015, pp. 1–15. [Online]. Available: https://arxiv.org/abs/1412.6980 [28] G. Li, H. Hammoud, H. Itani, D. Khizbullin, and B. Ghanem, “CAMEL: Communicative Agents for “Mind” Exploration of Large Language Model Society,” in Proceedings of the Thirty-Seventh Annual Conference on Neural Information Processing Systems (NeurIPS ’23), ser. Advances in Neural Information Processing Systems, A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, Eds., vol. 36. New Orleans, LA, USA: Curran Associates, Dec. 2023, pp. 51 991– 52 008. [Online]. Available: https://proceedings.neurips. cc/paper files/paper/2023/hash/a3621ee907def47c1b952 ade25c67698-Abstract-Conference.html [29] H. Li, S. Lin, F. Fu, Y. Zhou, X. Ji, Y. Zhao, L. Wang, J. Jiang, and B. Cui, “StaleFlow: Staleness-Aware Data Management for Mitigating Data Skewness in Fully Disaggregated RL Post-Training,” Proc. ACM Manag. Data, vol. 4, no. 4, pp. 287:1–287:29, Sep. 2026. [Online]. Available: https://doi.org/10.1145/3837125 [30] A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan et al., “DeepSeekV3 Technical Report,” Feb. 2025, arXiv:2412.19437. [Online]. Available: https://arxiv.org/abs/2412.19437 [31] Z. Mei, W. Fu, K. Li, G. Wang, H. Zhang, and Y. Wu, “ReaL: Efficient RLHF Training for Large Language Models through Parameter Reallocation,” in Proceedings of the Eighth Annual Conference on Machine Learning and Systems (MLSys ’25), ser. Proceedings of Machine Learning and Systems, M. Zaharia, G. Joshi, and Y. Lin, Eds., vol. 7, Santa Clara, CA, USA, May 2025, pp. 1–20. [Online]. Available: https://proceedings.mlsys.org/ paper files/paper/2025/hash/3b3889d313ba9476c12c2d7 7ea66b24f-Abstract-Conference.html [32] G. Mialon, C. Fourrier, T. Wolf, Y. LeCun, and T. Scialom, “GAIA: A Benchmark for General AI Assistants,” in Proceedings of the Twelfth International Conference on Learning Representations, ser. ICLR ’24, B. Kim, Y. Yue, S. Chaudhuri, K. Fragkiadaki, M. Khan, and Y. Sun, Eds., Vienna, Austria, May 2024, pp. 9025–9049. [Online]. Available: https://proceedings. iclr.cc/paper files/paper/2024/hash/25ae35b5b1738d80f 1f03a8713e405ec-Abstract-Conference.html [33] P. Micikevicius, S. Narang, J. Alben, G. Diamos, E. Elsen, D. Garcia, B. Ginsburg, M. Houston, O. Kuchaiev, G. Venkatesh, and H. Wu, “Mixed Precision Training,” in Proceedings of the Sixth International Conference on Learning Representations, ser. ICLR ’18. Vancouver, Canada: OpenReview, Apr. 2018, pp. 1–12. [Online]. Available: https: //openreview.net/forum?id=r1gs9JgRZ [34] D. Narayanan, M. Shoeybi, J. Casper, P. LeGresley,

M. Patwary, V. Korthikanti, D. Vainbrand, P. Kashinkunti, J. Bernauer, B. Catanzaro, A. Phanishayee, and M. Zaharia, “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, ser. SC ’21. St. Louis, MO, USA: Association for Computing Machinery, Nov. 2021, pp. 58:1–58:15. [Online]. Available: https://doi.org/10.1145/3458817.34 76209 [35] OpenAI, “OpenAI o3,” https://openai.com/index/introdu cing-o3-and-o4-mini/, Apr. 2025, accessed 2026-08-06. [36] ——, “OpenAI o3 and o4-mini System Card,” https: //openai.com/index/o3-o4-mini-system-card/, Apr. 2025, accessed 2026-08-06. [37] ——, “ChatGPT,” https://chatgpt.com/, Apr. 2026, accessed 2026-08-06. [38] L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray et al., “Training Language Models to Follow Instructions with Human Feedback,” in Proceedings of the ThirtySixth Annual Conference on Neural Information Processing Systems (NeurIPS ’22), ser. Advances in Neural Information Processing Systems, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, Eds., vol. 35. New Orleans, LA, USA: Curran Associates, Dec. 2022, pp. 27 730–27 744. [Online]. Available: https://proceedings.neurips.cc/paper files/pap er/2022/hash/b1efde53be364a73914f58805a001731-Abs tract-Conference.html [39] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He, “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, ser. SC ’20. Atlanta, GA, USA: IEEE Press, Nov. 2020, pp. 20:1–20:16. [Online]. Available: https://ieeexplore.ieee.org/document/9355301 [40] S. Rajbhandari, O. Ruwase, J. Rasley, S. Smith, and Y. He, “ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, ser. SC ’21. St. Louis, MO, USA: Association for Computing Machinery, Nov. 2021, pp. 59:1–59:14. [Online]. Available: https://doi.org/10.1145/3458817.3476205 [41] J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y. He, “ZeROOffload: Democratizing Billion-Scale Model Training,” in Proceedings of the USENIX Annual Technical Conference (USENIX ATC), ser. ATC ’21. Virtual Event: USENIX Association, Jul. 2021, pp. 551–564. [Online]. Available: https://www.usenix.org/conference/ atc21/presentation/ren-jie [42] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal Policy Optimization Algorithms,” Aug. 2017, arXiv:1707.06347. [Online]. Available: https://arxiv.org/abs/1707.06347

[43] Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, X. Bi, H. Zhang, M. Zhang, Y. Li, Y. Wu, and D. Guo, “DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models,” Apr. 2024, arXiv:2402.03300. [Online]. Available: https://arxiv.org/abs/2402.03300 [44] G. Sheng, Y. Tong, B. Wan, W. Zhang, C. Jia, X. Wu, Y. Wu, X. Li, C. Zhang, Y. Peng, H. Lin, X. Liu, and C. Wu, “Laminar: A scalable asynchronous rl posttraining framework,” in Proceedings of the 21st European Conference on Computer Systems, ser. EuroSys ’26. Edinburgh, Scotland, UK: Association for Computing Machinery, Apr. 2026, pp. 400–422. [Online]. Available: https://doi.org/10.1145/3767295.3803580 [45] G. Sheng, C. Zhang, Z. Ye, X. Wu, W. Zhang, R. Zhang, Y. Peng, H. Lin, and C. Wu, “HybridFlow: A Flexible and Efficient RLHF Framework,” in Proceedings of the Twentieth European Conference on Computer Systems, ser. EuroSys ’25. Rotterdam, Netherlands: Association for Computing Machinery, Mar. 2025, pp. 1279–1297. [Online]. Available: https: //doi.org/10.1145/3689031.3696075 [46] Y. Sheng, L. Zheng, B. Yuan, Z. Li, M. Ryabinin, B. Chen, P. Liang, C. Re, I. Stoica, and C. Zhang, “FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU,” in Proceedings of the 40th International Conference on Machine Learning (ICML ’23), ser. Proceedings of Machine Learning Research, A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, Eds., vol. 202. Honolulu, HI, USA: PMLR, Jul. 2023, pp. 31 094–31 116. [Online]. Available: https: //proceedings.mlr.press/v202/sheng23a.html [47] A. Singh, A. Fry, A. Perelman, A. Tart, A. Ganesh, A. El-Kishky, A. McLaughlin, A. Low, A. Ostrow, A. Ananthram et al., “OpenAI GPT-5 System Card,” May 2026, arXiv:2601.03267. [Online]. Available: https://arxiv.org/abs/2601.03267 [48] N. Stiennon, L. Ouyang, J. Wu, D. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei, and P. F. Christiano, “Learning to Summarize with Human Feedback,” in Proceedings of the Thirty-Fourth Annual Conference on Neural Information Processing Systems (NeurIPS ’20), ser. Advances in Neural Information Processing Systems, H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, Eds., vol. 33. Virtual Event: Curran Associates, Dec. 2020, pp. 3008–3021. [Online]. Available: https: //proceedings.neurips.cc/paper files/paper/2020/hash/1f8 9885d556929e98d3ef9b86448f951-Abstract.html [49] X. Tan, Y. Feng, Y. Zhou, Y. Jiang, Y. Zhu, and H. Xu, “OrchestrRL: Dynamic Compute and Network Orchestration for Disaggregated RL,” Mar. 2026, arXiv:2601.01209. [Online]. Available: https: //arxiv.org/abs/2601.01209 [50] Z. Tan, M. Wang, S. Wang, Y. Yin, Q. Wang, X. Chu, and S. Shi, “Bidirectional Resource Scheduling for Disaggregated and Asynchronous RL Post-Training,” Aug. 2026, arXiv:2607.09207. [Online]. Available:

https://arxiv.org/abs/2607.09207 [51] L. Wang, C. Ma, X. Feng, Z. Zhang, H. Yang, J. Zhang, Z. Chen, J. Tang, X. Chen, Y. Lin, W. X. Zhao, Z. Wei, and J. Wen, “A Survey on Large Language Model Based Autonomous Agents,” Frontiers of Computer Science, vol. 18, no. 6, pp. 186 345:1– 186 345:26, Mar. 2024. [Online]. Available: https: //link.springer.com/article/10.1007/s11704-024-40231-1 [52] J. Wei, X. Wang, D. Schuurmans, M. Bosma, B. Ichter, F. Xia, E. Chi, Q. V. Le, and D. Zhou, “Chainof-Thought Prompting Elicits Reasoning in Large Language Models,” in Proceedings of the ThirtySixth Annual Conference on Neural Information Processing Systems (NeurIPS ’22), ser. Advances in Neural Information Processing Systems, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, Eds., vol. 35. New Orleans, LA, USA: Curran Associates, Dec. 2022, pp. 24 824–24 837. [Online]. Available: https://proceedings.neurips.cc/paper files/pap er/2022/hash/9d5609613524ecf4f15af0f7b31abca4-Abs tract-Conference.html [53] A. Weingram, D. Zhang, Z. Chen, H. Qi, and X. Lu, “HPC-R1: Characterizing R1-Like Large Reasoning Models on HPC,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, ser. SC ’25. St. Louis, MO, USA: Association for Computing Machinery, Nov. 2025, pp. 1368–1380. [Online]. Available: https://doi.org/10.1145/3712285.3759827 [54] S. J. Wright, “Coordinate Descent Algorithms,” Mathematical Programming, vol. 151, no. 1, pp. 3–34, Mar. 2015. [Online]. Available: https: //link.springer.com/article/10.1007/s10107-015-0892-3 [55] Q. Wu, G. Bansal, J. Zhang, Y. Wu, B. Li, E. Zhu, L. Jiang, X. Zhang, S. Zhang, J. Liu, A. H. Awadallah, R. W. White, D. Burger, and C. Wang, “AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation,” in Proceedings of the First Conference on Language Modeling, ser. COLM ’24. Philadelphia, PA, USA: OpenReview, Oct. 2024, pp. 1–46. [Online]. Available: https://openreview.net/forum?id=BAakY1hN KS [56] xAI, “Introducing Grok 4.7,” https://x.ai/news/grok-4-7, Sep. 2026, accessed 2026-10-18. [57] J. Yang, C. E. Jimenez, A. Wettig, K. Lieret, S. Yao, K. Narasimhan, and O. Press, “SWEagent: Agent-Computer Interfaces Enable Automated Software Engineering,” in Proceedings of the ThirtyEighth Annual Conference on Neural Information Processing Systems (NeurIPS ’24), ser. Advances in Neural Information Processing Systems, A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, Eds., vol. 37. Vancouver, Canada: Curran Associates, Dec. 2024, pp. 50 528–50 652. [Online]. Available: https://proceedings.neurips.cc/paper files/paper/2024/hash/5a7c947568c1b1328ccc5230172 e1e7c-Abstract-Conference.html [58] S. Yao, D. Yu, J. Zhao, I. Shafran, T. Griffiths, Y. Cao,

and K. Narasimhan, “Tree of Thoughts: Deliberate Problem Solving with Large Language Models,” in Proceedings of the Thirty-Seventh Annual Conference on Neural Information Processing Systems (NeurIPS ’23), ser. Advances in Neural Information Processing Systems, A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, Eds., vol. 36. New Orleans, LA, USA: Curran Associates, Dec. 2023, pp. 11 809– 11 822. [Online]. Available: https://proceedings.neurips. cc/paper files/paper/2023/hash/271db9922b8d1f4dd7aae f84ed5ac703-Abstract-Conference.html [59] S. Yao, J. Zhao, D. Yu, N. Du, I. Shafran, K. R. Narasimhan, and Y. Cao, “ReAct: Synergizing Reasoning and Acting in Language Models,” in Proceedings of the Eleventh International Conference on Learning Representations, ser. ICLR ’23. Kigali, Rwanda: OpenReview, May 2023, pp. 1–24. [Online]. Available: https://openreview.net/forum?id=1PL1NIMMrw [60] Y. Yao, Z. Li, and H. Zhao, “GoT: Effective Graph-ofThought Reasoning in Language Models,” in Findings of the Association for Computational Linguistics: NAACL 2024, K. Duh, H. Gomez, and S. Bethard, Eds. Mexico City, Mexico: Association for Computational Linguistics, Jun. 2024, pp. 2901–2921. [Online]. Available: https: //aclanthology.org/2024.findings-naacl.183/ [61] C. Yu, Y. Wang, Z. Guo, H. Lin, S. Xu, H. Zang, Q. Zhang, Y. Wu, C. Zhu, J. Hu, Z. Huang, M. Wei, Y. Xie, K. Yang, B. Dai, Z. Xu, J. Du, X. Wang, X. Fu, L. Shi, Z. Liu, K. Chen, W. Liu, G. Liu, B. Li, J. Yang, Z. Yang, G. Dai, and Y. Wang, “RLinf: Flexible and Efficient Large-Scale Reinforcement Learning via Macro-to-Micro Flow Transformation,” in Proceedings of the 20th USENIX Symposium on Operating Systems Design and Implementation, ser. OSDI ’26. Seattle, WA, USA: USENIX Association, Jul. 2026, pp. 829– 846. [Online]. Available: https://www.usenix.org/confere nce/osdi26/presentation/yu-chao [62] R. Zheng, S. Dou, S. Gao, Y. Hua, W. Shen, B. Wang, Y. Liu, S. Jin, Q. Liu, Y. Zhou et al., “Secrets of RLHF in Large Language Models Part I: PPO,” Jul. 2023, arXiv:2307.04964. [Online]. Available: https://arxiv.org/abs/2307.04964 [63] Y. Zhong, Z. Zhang, B. Wu, S. Liu, Y. Chen, C. Wan, H. Hu, L. Xia, R. Ming, Y. Zhu, and X. Jin, “Optimizing RLHF Training for Large Language Models with Stage Fusion,” in Proceedings of the 22nd USENIX Symposium on Networked Systems Design and Implementation, ser. NSDI ’25. Philadelphia, PA, USA: USENIX Association, Apr. 2025, pp. 489–503. [Online]. Available: https://www.usenix.org/conference/nsdi25/pre sentation/zhong

A PPENDIX A RPEG P IPELINE : F ORMULATION & D ESIGN D ETAILS

Actor [GEN]

We now provide extended mathematical details of RPEG. Critic [INF]

Reference [INF] Frozen

Reward [INF] Frozen

A. RPEG Pipeline Execution Graph Definition A.1 (RLM Pipeline Execution Graph): An RLM Pipeline Execution Graph (RPEG) is a tuple

Critic [TRAIN]

Actor [TRAIN]

◦ ◦ G◦ = V ◦ , Edata , Epub , A, M, τ , where: V ◦ is a finite set of RPEG vertices (representing model invocations); ◦ ◦ • Edata and Epub are finite sets of data and publication edges, where each edge e = (u → v, a) connects a source vertex u to a destination vertex v and carries a single container a ∈ A; • A is the finite set of containers referenced by the RPEG; the boundary subsets A◦in , A◦out ⊆ A (external inputs and outputs) are characterized by well-formedness condition 3; • M is a set of all models represented in the RPEG pipeline (e.g., actor, critic, reference, reward for PPObased training); • Mtrain ⊂ M is the set of trainable models, where M denotes the full set of models; • τ : Mtrain → N assigns a bounded-staleness budget per trainable model. •

Figure 13 illustrates the PPO RPEG. A complete formal instantiation – including all container definitions, operator input and output interfaces, and edge annotations – is given in Section A-D. We now further detail each component of the RPEG. 1) Vertices and Operator Types: Each RPEG vertex v ∈ V ◦ represents a model invocation. A vertex is characterized by three attributes: a role r(v) ∈ {generation, inference, training}, a model µ(v) ∈ M , and • input/output container interfaces I(v), O(v) ⊆ A declaring which containers the operator may read and write. • •

We write κ(v) = (r(v), µ(v), I(v), O(v)) for the operator kind of v. The vertex set therefore determines which models participate in the algorithm and how they appear within one iteration. Different policy optimization algorithms correspond to differ◦ ◦ ent choices of V ◦ and dependency structure (Edata , Epub ). For example, the base PPO RPEG comprises six vertices – generation, three inference passes, and two training passes – while GRPO eliminates the critic and reduces to four. DPO removes online generation entirely, collapsing the graph to three.

Fig. 13: RPEG for a PPO-based RLM training pipeline. The trainable model set is Mtrain = {πθ , Vψ } (actor and critic), both with staleness budget τ = 0 (synchronous execution; depicted like this only in this example). External inputs are A◦in = {X, Wref , WR }, and no external outputs are exposed (A◦out = ∅). Containers: X = prompts, Y = completions, v = per-token value predictions, r = percompletion reward, ℓref = reference log-probs, ℓπ = actor log-probs, and Wm = model m’s weights. Edges encode scheduling dependencies, not all data accesses: boundary inputs A◦in are accessible to all operators that declare them in their interface, and containers such as Y and ℓπ reach training transitively through the assessment stage without requiring direct edges. The structure is shown for a fixed workload.

2) Containers: A container has an identity and size metadata (#a, s(a)), where #a is the element count and s(a) is the element size in bytes. The total byte volume of a container is #a · s(a). Containers make communication volume explicit in the IR: the byte volume #a · s(a) of a container on a connecting edge determines the data transfer cost when source and destination are placed on different devices. The set A collects all containers in the RPEG, covering both within-step data (e.g., completions, log-probabilities, or rewards) and cross-step model weights. 3) Data and Publication Edges: Each RPEG edge carries exactly one container and represents the algorithmic dependence of the destination on the data produced by the source. The RPEG distinguishes two classes of edges: ◦ Data edges Edata represent within-step dependencies. Generation produces completions consumed by inference; inference produces log-probabilities, rewards, and values consumed by training. ◦ • Publication edges Epub represent cross-step weight updates. After training updates a model’s parameters, a publication edge carries the new weights to the operators that consume them in subsequent iterations. These are the edges whose unrolling semantics are governed by the staleness budget τ . •

An important distinction separates edges from operator interfaces. Each operator declares input and output interfaces (I and O) that specify which containers it may access, but not every declared input requires a corresponding incoming edge, since edges impose both a data-transfer requirement and a scheduling dependency. Two cases lack edges: boundary inputs A◦in (e.g., prompts, frozen weights) are accessible to all operators that declare them since they are external to the

graph; and containers already resident from a prior stage can be accessed transitively without a direct edge. For example, training operators read Y and ℓπ from generation, but the RPEG by default omits direct generation-to-training edges – unless explicitly determined by the user.

GPU Device Allocation

B. Well-Formedness Not every tuple of vertices, edges, and containers constitutes a valid RPEG. We therefore impose a set of static well-formedness conditions on G◦ that are fundamental for lowering, scheduling, and costing the RPEG. Definition A.2 (Well-Formedness): An RPEG G◦ is wellformed if the following conditions hold: 1) Interface consistency. For every edge e = (u → v, a) ∈ ◦ ◦ Edata ∪ Epub carrying container a, the container must satisfy a ∈ O(u) ∧ a ∈ I(v). That is, edges may only transfer containers that the source declares as output and the destination declares as input. 2) Publication edges carry model parameters. For every ◦ publication edge e = (u → v, a) ∈ Epub , the source vertex u must be a training operator and the container a must represent the weight tensor of the trained model. 3) Boundary consistency. If a container a ∈ A is read by some vertex but has no producer, then a ∈ A◦in . If a container a is written by some vertex but has no consumer, then a ∈ A◦out . This condition ensures that all data entering or leaving the RPEG is accounted for explicitly. 4) Acyclicity of within-step dependencies. The directed ◦ ) induced by data edges is acyclic. graph (V ◦ , Edata Cyclic dependencies may only arise through publication edges, which represent cross-iteration weight updates. These conditions are checked statically on the RPEG and ensure that G◦ represents a well-defined steady-state computation before unrolling, lowering, scheduling, and costing. C. Execution Plan and Device Space While the RPEG specifies the algorithmic structure of a steady-state RLM training iteration, execution-level decisions determine how this structure is mapped onto a concrete hardware topology. We formalize both the device space and the execution plan in this section. 1) Device Space: Definition A.3 (Device Space): We model the physical hardware as a hierarchical device space D = [Nr ] × [Nn ] × [Ng ]1 , where Nr denotes the number of racks, Nn the number of nodes per rack, and Ng the number of GPUs per node. The total number of devices is therefore Nr · Nn · Ng . 1 Here we interpret [N ] := {0, . . . , N − 1} (analogously for [N ], r r n [Ng ]), and we use half-open integer intervals [a, b) := {a, . . . , b − 1} when

specifying contiguous blocks.

Rack

Node

Fig. 14: Device space and allocation. The device space D = [Nr ]×[Nn ]×[Ng ] organizes GPUs into a 3D hierarchy. An allocation D(o, δ) selects a contiguous rectangular region, ensuring that all devices assigned to an operator are physically co-located.

Definition A.4 (Device Allocation): A device allocation is a contiguous hyper-rectangle D(o, δ) ⊆ D, defined by an origin o = (or , on , og ) and an extent δ = (δr , δn , δg ): D(o, δ) = {(r, n, g) ∈ D | r ∈ [or , or + δr ), n ∈ [on , on + δn ), g ∈ [og , og + δg )}. We write |D(o, δ)| = δr · δn · δg for the number of devices in an allocation. Figure 14 illustrates the device space and a sample allocation. Example A.1 (Device Allocation): Given a cluster with two racks, eight nodes per rack and four GPUs per node, the device space is D = [2] × [8] × [4]. Allocating the first two nodes in the first rack corresponds to o = (0, 0, 0) and δ = (1, 2, 4). 2) Execution Plan: The execution plan is a set of functions over V ◦ that assign hardware-dependent attributes – placement, 3D parallelization strategy, and microbatch count – to each RPEG vertex. It does not alter the structure of G◦ . Definition A.5 (Execution Plan): Given an RPEG G◦ , an execution plan is a tuple P = (σ, Π, mb), where for every RPEG vertex v ∈ V ◦ : • σ(v) ⊆ D is a device allocation; 3 • Π(v) = (Pd , Pp , Pt ) ∈ N>0 specifies data, pipeline, and tensor parallel degrees; • mb(v) ∈ N>0 is the number of pipeline microbatches. The parallelization tuple Π(v) determines how operator v’s computation is partitioned into shards, while σ(v) specifies where these shards are placed in the device space. We require for every vertex v that |σ(v)| = Pd · Pp · Pt , ensuring that the number of logical shards induced by the parallel degrees matches the number of allocated devices.

Furthermore, if Pp = 1, then mb(v) = 1, since microbatching is meaningful only in the presence of pipeline parallelism.

Vertex

I(v)

O(v)

This section provides the complete formal instantiation of the PPO RPEG.

gπ iref iR iV tπ tV

{Wπ , X} {Wref , Y } {WR , Y } {WV , Y } {Wπ , Y, ℓπ , ℓref , r, v} {WV , Y, ℓref , r, v}

{Y, ℓπ } {ℓref } {r} {v} {Wπ } {WV }

1) Models: The PPO pipeline uses four models M = {πθ , πref , Vψ , Rφ }, of which Mtrain = {πθ , Vψ } are trainable. The RPEG has six vertices (Table VII):

TABLE IX: PPO operator input/output interfaces. Training operators declare containers in their read set that are not delivered by incoming edges (e.g., tπ reads Y and ℓπ from device memory rather than via a direct edge from gπ ).

D. PPO RPEG Instantiation

Vertex

Role

Model

Description

gπ

generation

πθ

Autoregressive rollout: produces completions Y and log-probabilities ℓπ . Reference forward pass: produces baseline log-probabilities ℓref . Reward scoring: produces per-sequence rewards r. Value estimation: produces per-token values v. Actor policy-gradient update. Critic value-loss update.

iref

inference

πref

inference

Rφ

inference

Vψ

tπ tV

training training

πθ Vψ

TABLE VII: PPO RPEG vertices with operator kinds. Each vertex is characterized by a role r(v) and a model µ(v) ∈ M .

2) Containers: Table VIII lists all containers. Let B denote the batch size, R the number of rollouts per prompt, Lp the average prompt length, Lr the average rollout length, sb the working precision (2 bytes for BF16), and #Wm the parameter counts. Container

s(a)

Description

Weight containers (carried by publication edges) Wπ WV Wref

#Wπθ #WVψ #Wπref

sb sb sb

#WRφ

Actor parameters. Critic parameters. Reference parameters (frozen). Reward parameters (frozen).

Data containers (carried by data edges) X Y

BLp BR(Lp +Lr )

sb sb

ℓπ ℓref r v

BRLr BRLr BR BRLr

sb sb sb sb

Prompt tokens. Completions (prompt + response). Actor log-probabilities. Reference log-probabilities. Rewards (per sequence). Value estimates (per token).

TABLE VIII: PPO RPEG containers. Boundary inputs: Ain = {X, Wref , WR }; no external outputs (Aout = ∅).

3) Operator Interfaces: Each vertex declares which containers it reads and writes (Table IX). These interfaces determine interface consistency (Section A-B) and are used by the lowering pass to derive communication payloads.

4) Edges: We also formulate PPO edges. Data edges.  Y Y Y  g −→ iref , gπ −→ iR , gπ −→ iV ,   π ◦ ℓref ℓref Edata = iref −− → tπ , iref −− → tV ,    r r v v iR − → tπ , iR − → tV , iV − → tπ , iV − → tV

      

No direct edges connect generation to training. Publication edges. WV Wπ ◦ Epub = tπ −−→ gπ , tV −−→ iV . Staleness. Each publication edge is governed by the staleness budget of its source model: τ (πθ ) for the actor publication tπ → gπ , and τ (Vψ ) for the critic publication tV → iV . Setting both to zero recovers synchronous execution; setting τ (πθ ) = 1 allows one step of actor overlap (see Section A-F). E. Other RLM Training Algorithms The RPEG formalism applies to any RLM training algorithm by adjusting the vertex and edge sets. GRPO removes the critic vertices (iV , tV ) and derives training signals from relative rewards within groups of sampled responses, reducing the graph to four vertices. DPO removes both the reward model and the online generation loop, collapsing the graph to an offline preference-optimization pipeline. The unrolling, lowering, cost modeling, and optimization procedures from the preceding sections apply without modification to all three variants. Fig. 15 shows the synchronous PPO case. Fig. 16 illustrates how varying τ changes the structure of the unrolled graph. F. Unrolling Unrolling expands the cyclic RPEG into a finite acyclic graph Ḡ over Keval steps, resolving cross-step publication edges according to the staleness budget τ . We use PPO as the running example throughout this section, but the same procedures apply to GRPO and DPO graphs. Before the cyclic RPEG can be unrolled into a finite DAG, we must resolve an ambiguity: the staleness budget τ (m) bounds how old the consumed weights may be, but does not specify which training step provides them. We first formalize this publisher selection, then define the unrolled graph that results from it.

Reward [INF] Frozen

Reward [INF] Frozen Actor [TRAIN]

Actor [GEN]

Actor [TRAIN]

Reference [INF] Frozen

Actor [GEN]

Reference [INF] Frozen

Critic [TRAIN]

Critic [INF]

Fig. 15: Unrolled RPEG for a PPO-based RLM training pipeline. The cyclic RPEG from Fig. 13 is unrolled over Keval = 2 synchronous steps (τ = 0). Each RPEG vertex is replicated per step, data edges are copied within each step, and publication edges are wired across step boundaries according to the staleness budget τ .

a) Synchronous Reward [INF]

Reward [INF]

...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Reward [INF]

...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Reward [INF]

b) 1-Step Asynchrony Reward [INF] ...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Reward [INF]

c) 2-Step Asynchrony Reward [INF] ...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Actor [TRAIN]

Actor [GEN]

Ref. [INF]

...

Critic [TRAIN]

Critic [INF]

...

Fig. 16: Unrolled RPEGs under different staleness budgets (simplified; operator attributes and edge containers omitted for clarity). (a) Synchronous execution (τ = 0): publication edges connect consecutive steps, enforcing strict ordering. (b) One-step bounded staleness (τ (πθ ) = 1): actor publication edges span two steps, enabling concurrent generation and training. (c) Two-step bounded staleness (τ (πθ ) = 2, τ (Vψ ) = 1): deeper overlap with longer publication arcs. Red arrows denote publication edges carrying weight containers Wπ and WV ; black arrows denote within-step data edges.

1) Publication and Bounded Staleness: When τ (m) > 0 for some model m ∈ Mtrain , generation at step k+1 may overlap with training at step k, using slightly stale policy weights for rollout generation. As noted above, τ (m) does not determine which specific step provides the weights. To resolve this during unrolling, we introduce a publisher-selection function sm (k) that maps each step k to the concrete training step whose output it consumes. Definition A.6 (Publisher Selection): Fix Keval > 0. For each trainable model m ∈ Mtrain , define sm (k) = k − τ (m) − 1,

τ (m) < k < Keval .

This choice selects the maximally stale admissible publisher and is deterministic. It satisfies 0 ≤ sm (k) < k,

(k − 1) − sm (k) ≤ τ (m).

The synchronous case is recovered by enforcing τ (m) = 0 for all trainable models, yielding sm (k) = k − 1. 2) Unrolled Graph: Given the deterministic publisher selection function, we can now formally define the unrolling process. Definition A.7 (Unrolled Graph): Given a well-formed RPEG G◦ and an evaluation horizon Keval > 0, the unrolled graph is Ḡ = (V̄ , Ēdata , Ēpub ), with the following components: • Vertices. V̄ = V ◦ × [Keval ], •

(v, k) ∈ V̄ .

Each RPEG vertex is replicated once per step index. Within-step edges. For each data edge (u → v, a) ∈ ◦ Edata and each k ∈ [Keval ], a

(u, k) − → (v, k) ∈ Ēdata , carrying the same container a as the RPEG edge.2 • Publication edges. For each publication edge (u → ◦ v, a) ∈ Epub and each k ∈ {1, . . . , Keval − 1} with k > τ (µ(u)), where µ(u) denotes the model invoked by vertex u, a

(u, sµ(u) (k)) − → (v, k) ∈ Ēpub , carrying the same weight container a as the RPEG edge. For early steps k ≤ τ (µ(u)), no publication edge is created: these steps use the initial parameter version and have no training predecessor. Because no dependency connects them to prior training, the early serving invocations may overlap when the remaining data, ordered-update, and physical-resource constraints permit it; they read the same initial version but process distinct data batches.

Algorithm 3: U NROLL: Unroll cyclic RPEG ◦ ◦ , . . . , τ ), evaluation horizon Input: RPEG G◦ = (V ◦ , Edata , Epub Keval Output: Unrolled DAG Ḡ = (V̄ , Ēdata , Ēpub )

V̄ ← ∅, Ēdata ← ∅, Ēpub ← ∅; /* Step 1: replicate vertices */ 2 for k ← 0 to Keval − 1 do ◦ 3 foreach v ∈ V do 4 add (v, k) to V̄ ; 5 end 6 end /* Step 2: replicate within-step data edges */ 7 for k ← 0 to Keval − 1 do ◦ 8 foreach (u → v) ∈ Edata do 9 add (u, k) → (v, k) to Ēdata ; 10 end 11 end /* Step 3: instantiate publication edges with bounded staleness */ 12 for k ← 1 to Keval − 1 do ◦ 13 foreach (u → v) ∈ Epub do 14 m ← µ(u); 15 if k ≤ τ (m) then 16 continue ; /* step k uses initial weights of model m */ 17 end 18 k′ ← max{0, k − (τ (m) + 1)}; 19 add (u, k′ ) → (v, k) to Ēpub ; 20 end 21 end 22 return Ḡ; 1

◦ Lemma A.1 (Acyclicity): If (V ◦ , Edata ) is acyclic, then the unrolled graph (V̄ , Ēdata ∪ Ēpub ) is a DAG.

Proof. Within-step edges preserve k and are acyclic by the well-formedness assumption on G◦ . Publication edges strictly increase k because sm (k) < k. Together, no directed cycle can exist. □ 3) Ensuring Ordering of Training Updates Under Asynchrony: Successive training updates of the same model are ordered independently of the staleness budget. For each trainable model m with training operator tm , unrolling additionally introduces W

(tm , k − 1) −−m → (tm , k),

1 ≤ k < Keval .

These dependencies ensure that consecutive updates of the same optimizer state remain ordered. 4) Unrolling Algorithm: Algorithm 3 gives the complete unrolling procedure. The runtime is ◦ ◦ O(Keval · (|V ◦ | + |Edata | + |Epub |)),

as each of the three steps in the pseudocode iterates over one component of the RPEG, replicated Keval times. Each unrolled edge inherits the container from its RPEG edge. The algorithm omits this for brevity. G. Lowering

The key structural property of the unrolled graph is acyclicity, which we verify next. 2 Recall that [K eval ] = {0, . . . , Keval −1}.

Lowering incorporates the execution plan P to produce the executable graph Ĝ, inserting explicit communication and resharding nodes wherever placement or strategy differences require data movement. Thus, lowering is the stage at which

Actor [RESHARD] Payload:

Reward [INF] Frozen

[COMM] Payload: Actor [TRAIN]

[COMM] Payload:

Actor [GEN] [COMM] Payload:

Reference [INF] Frozen

[COMM] Payload:

Critic [TRAIN] [COMM] Payload: Critic [INF]

[COMM] Payload:

Critic [RESHARD] Payload:

Fig. 17: Executable graph for one step of a PPO-based RLM training pipeline. Derived from the RPEG in Fig. 13 via unrolling and lowering. C OMM and R ESHARD nodes are inserted wherever the execution plan requires data movement between different device allocations or parallel layouts.

the unrolled graph Ḡ and the execution plan P are combined into a single hardware-aware DAG. The resulting executable graph Ĝ serves as input to the cost model and scheduler defined further. Both transformations are defined generically over any RPEG. We use PPO as the running example throughout this section, but the same procedures apply to GRPO and DPO graphs. Concretely, lowering takes the pair (Ḡ, P ) and constructs the executable graph Ĝ by traversing the nodes and edges of Ḡ and inserting C OMM and R ESHARD nodes wherever the execution plan requires data movement or weight redistribution. Definition A.8 (Executable Graph): The executable graph is a DAG Ĝ = (V̂ , Ê), whose nodes are of three kinds: • O P nodes correspond to unrolled RPEG vertices, each attributed with a device allocation, parallel strategy, and microbatch count from the execution plan. • C OMM nodes represent data transfer between O P nodes placed on different device allocations (on both data and publication edges). • R ESHARD nodes represent weight redistribution when successive invocations of the same model use different parallel layouts on the same devices. Fig. 17 shows the PPO-based executable graph for one step. Before describing the lowering procedure itself, we specify the attributes that distinguish the three node kinds.

Allocation. Each O P node occupies exactly one device allocation (as specified by the execution plan). C OMM and R ESHARD nodes are attributed with both the source and destination allocations. Payload. Every node in Ĝ carries a payload in bytes: • each C OMM node carries a payload equal to the byte volume of the container on the corresponding data edge, • each R ESHARD node carries a payload equal to the full parameter size of the associated model, • O P nodes carry zero payload. We now describe how the executable graph is constructed by lowering the unrolled graph with the execution plan. 1) Stage 1: Operator Instantiation: For every unrolled vertex (v, k) ∈ V̄ , we create an O P node xv,k ∈ V̂ with device allocation σ(v), parallel strategy Π(v), and microbatch count mb(v). When generation quantile splitting is enabled (see below), generation and inference vertices may be replaced by multiple sub-nodes during this stage. Response-length–aware sub-node splitting. Autoregressive generation produces responses of highly variable length. In reasoning workloads, response-length distributions are heavily skewed: many sequences finish quickly while a small number of long-tail rollouts dominate wall-clock time [17, 44, 63]. In a monolithic generation operator, the entire batch must finish before any downstream inference can begin – the longest response delays all subsequent work. When generation and

a) Executable Graph

b) Schedule [COMM]

Actor [GEN]

[COMM]

Reward [INF]

Reference [INF]

Critic [INF]

Makespan

GPU 0

Actor [GEN]

GPU 1

idle

Reward [Inf]

GPU 2

idle

Reference [INF]

GPU 3

idle

Critic [Inf] Time

Fig. 18: Without sub-node splitting. (a) Executable graph: generation on GPU 0 feeds three inference operators on GPUs 1–3 via C OMM nodes. (b) Schedule: inference GPUs are idle for the entire duration of generation, since no partial results are available until the full batch completes. For clarity, only device allocations (origin o, extent δ) are shown on each node; parallel strategies, microbatch counts, and read/write interfaces are omitted.

a) Executable Graph Actor@0 [GEN]

Actor@1 [GEN]

[COMM] [COMM]

Actor@2 [GEN]

[COMM]

Reward@0 [INF]

[COMM]

Reward@1 [INF]

[COMM]

Reward@2 [INF]

[COMM]

Reference@0 [INF]

Reference@1 [INF]

Reference@2 [INF]

Critic@0 [INF]

Critic@1 [INF]

Critic@2 [INF]

b) Schedule

Makespan GPU 0

Actor [GEN]

GPU 1

idle

Reward [Inf]

Rew. [Inf]

GPU 2

idle

Reference [INF]

Ref. [INF]

GPU 3

idle

Critic [Inf]

Crit. [Inf]

Time

Fig. 19: With quantile-based sub-node splitting (Q = 3). (a) Executable graph: generation and each inference operator are split into three sub-nodes chained sequentially. Each generation sub-node feeds the corresponding inference sub-nodes via C OMM nodes. (b) Schedule: inference sub-nodes from early-finishing sequences overlap with later generation phases, substantially reducing idle time and makespan. For clarity, only device allocations (origin o, extent δ) are shown on each node; parallel strategies, microbatch counts, and read/write interfaces are omitted.

inference are placed on disjoint devices, the lowering pass can exploit this skewness by splitting generation into phases based on the response-length distribution, allowing inference to start processing early-finishing sequences while longer ones continue decoding. This is related to the inter-stage fusion mechanism proposed by RLHFuse [63], but expressed here as a structural graph transformation within the lowering pass rather than a runtime scheduling decision. Definition A.9 (Response-Length Quantiles): The responselength distribution is specified by Q quantile breakpoints (f1 , t1 ), (f2 , t2 ), . . . , (fQ , tQ ), where each fq ∈ (0, 1] is a cumulative fraction of sequences and tq is the number of generated tokens at which that fraction completes. We require fQ = 1 (all sequences finish by tQ = Lr ) and f1 < f2 < · · · < fQ . Example A.2 (Response-Length Quantiles): With Q = 4 and breakpoints (0.25, 512), (0.50, 1024), (0.85, 1536), (1.0, 2048) , 25% of sequences finish within 512 tokens, 50% within 1024, 85% within 1536, and all by 2048. Obtaining the response-length distribution. The quantile breakpoints (fq , tq ) are obtained empirically by profiling generation lengths on a representative sample of prompts. Before optimization, a small subset of the training prompts (e.g., a few hundred) is processed through the current model checkpoint, and the empirical response-length distribution is recorded. The quantile breakpoints are then fitted directly from this sample. The number of quantiles Q controls the granularity of the sub-node splitting: Q = 1 recovers the monolithic generation operator, while larger Q enables finer-grained overlap at the cost of additional nodes in the executable graph. In practice, Q ∈ {3, 4, 5} captures the relevant structure of the distribution without introducing excessive graph complexity. Splitting mechanism. The lowering pass splits a generation vertex into sub-nodes only when the execution plan places generation and at least one downstream inference operator on disjoint devices – otherwise no overlap is possible and the operator remains an indivisible O P node. Concretely, for each generation vertex (v, k) with a disjoint downstream inference vertex: 1) The single O P node xv,k is replaced by Q generation subgQ 1 nodes xgv,k , . . . , xv,k , chained sequentially on the same gq devices. Sub-node xv,k represents the decode phase during which the fraction fq − fq−1 of sequences complete (with f0 = 0). The first sub-node additionally includes the prefill phase. 2) Each downstream inference vertex w on disjoint devices is similarly split into Q inference sub-nodes iQ 1 xiw,k , . . . , xw,k , also chained sequentially on their own devices. 3) Generation sub-node q is connected to inference sub-node q via a C OMM node with payload #a · s(a) · (fq − fq−1 )

bytes – the container volume scaled by the fraction of sequences completing in that phase. All sub-nodes inherit the device allocation, parallel strategy, and microbatch count from the original RPEG vertex. Figs. 18 and 19 illustrate the effect of sub-node splitting on both the executable graph and the resulting schedule for a 4-GPU example. Crucially, this wiring makes gen-inf overlap (inter-stage fusion) possible but does not enforce it. The optimizer controls whether overlap occurs by choosing the device allocations: if it places generation and inference on disjoint devices, the scheduler can execute inference sub-nodes concurrently with later generation sub-nodes. If the optimizer places them on the same devices, the scheduler serializes them and no overlap occurs – in this case, no splitting is performed, since subnodes on shared devices would add graph complexity without enabling any concurrency. 2) Stage 2: Communication Node Insertion: When subnode splitting has been applied, the edges between generation and inference sub-nodes created in Stage 1 are processed by the following stages in the same way as edges between unsplit O P nodes – C OMM and R ESHARD insertion does not distinguish between them. a We traverse the data edges of Ḡ. For each edge (u, k) − → (v, k) ∈ Ēdata with container a: • Same allocation (σ(u) = σ(v)): no transfer is needed. We add a precedence edge xu,k → xv,k . • Different allocation (σ(u) ̸= σ(v)): data must move between device allocations. We insert a C OMM node c with payload #a · s(a) bytes and add edges xu,k → c → xv,k . After this stage, all cross-allocation data movement is represented explicitly. 3) Stage 3: Intra-Iteration Weight Redistribution: When two operators of the same trainable model appear within one iteration, the lowering pass must ensure that the second operator receives weights in a compatible layout. To identify such pairs, we compute a topological ordering Topo(G◦ ) of ◦ (V ◦ , Edata ), collect all vertices with µ(v) = m for each trainable model m (frozen models have fixed weights and are never resharded), and define Pairs(m) as the consecutive pairs in this ordering. Not every pair requires resharding. Depending on the staleness budget and device placement, a pair falls into one of two categories – shadow pairs and non-shadow pairs – which the lowering pass handles differently. Shadow pairs. In synchronous RLM training, generation and training of the same model execute sequentially and may share a single weight copy that is resharded between them as needed [31]. When τ (m) > 0 and the optimizer places the two operators on disjoint device allocations, a more efficient pattern becomes possible: each device allocation maintains its own persistent weight copy, avoiding intra-step resharding between those copies while retaining cross-step publication and enabling concurrent execution. We call this a shadow pair.

Definition A.10 (Shadow Pair): A consecutive pair (u, v) ∈ Pairs(m) of a trainable model m ∈ Mtrain is a shadow pair under execution plan P if: 1) τ (m) > 0 (bounded asynchrony), and 2) σ(u) ∩ σ(v) = ∅ (device disjointness). These conditions identify the shadow-pair optimization. Synchronous persistent serving workers can still retain a copy to avoid reconstruction; therefore a single-copy policy is not universally faster. Overlapping allocations restrict concurrency, but physical copy lifetime is determined by actual workers and required versions, not solely by the shadow-pair predicate. When a shadow pair is active, two weight instances coexist in the device space: • Training instance on σ(tm ): holds up-to-date weights, optimizer state (FP32 master weights, first and second moments), and gradients. This is the source of parameter updates. • Shadow instance on the consumer’s devices (e.g., σ(gm ) for generation or σ(im ) for inference): holds a read-only BF16 weight snapshot. This copy is periodically refreshed by a C OMM node that transfers updated weights from the training instance (Stage 4). Because no intra-step reshard or precedence edge connects the two sides, they are free to execute concurrently. The idea of maintaining stale parameter replicas for concurrent execution appears in recent asynchronous RLM training systems [22]; our framework exposes it as a structural property of the execution plan that the optimizer can discover automatically. For shadow pairs, the lowering pass skips the intra-step reshard entirely: no R ESHARD node and no precedence edge are inserted between xu,k and xv,k within any step k. Non-shadow pairs. For each non-shadow pair (u, v) ∈ Pairs(m) and each k ∈ [Keval ]: • Same allocation and strategy (σ(u) = σ(v), Π(u) = Π(v)): the weight layout is already compatible; no action. • Same allocation, different strategy (σ(u) = σ(v), Π(u) ̸= Π(v)): the weights reside on the same devices but in an incompatible parallel layout. We create a R ESHARD node r ∈ V̂ with payload #Wm · sb bytes (the full parameter count of model m) and add edges xu,k → r → xv,k to Ê. • Different allocation (σ(u) ̸= σ(v)): the two operators run on different (possibly overlapping) device allocations – for example, inference on 8 GPUs and training on 16 GPUs with partial overlap. Since this is not a shadow pair, τ (m) = 0 or the allocations are not fully disjoint. Each side holds its own weight copy: training devices have weights from the optimizer state, while inference devices load them transiently via a prior publication transfer and free them after the operator completes. No redistribution is needed, but a bare precedence edge xu,k → xv,k enforces that the first operator completes before the second begins. 4) Stage 4: Cross-Iteration Publication Edges: Stages 1–3 handle within-step structure. The final stage wires the crossstep publication edges that carry updated weights from training

at an earlier step to consumers at a later step. For each a publication edge (u, k ′ ) − → (v, k) ∈ Ēpub with k ′ < k, the lowering pass selects one of three actions based on the relative placement and strategy of source and destination: • Same allocation and strategy (σ(u) = σ(v), Π(u) = Π(v)): a bare precedence edge xu,k′ → xv,k suffices. • Same allocation, different strategy (σ(u) = σ(v), Π(u) ̸= Π(v)): we insert a R ESHARD node as in Stage 3. • Different allocation (σ(u) ̸= σ(v)): the updated weights must be sent from the training devices to the consumer’s devices. We insert a C OMM node c with edges xu,k′ → c → xv,k ; the source retains its copy. When the source and destination allocations partially overlap, devices in the intersection can produce their weight shard locally, reducing the transfer to #Wm ·sb ·(nmax −noverlap )/nmax , where nmax = max(|σ(u)|, |σ(v)|) and noverlap = |σ(u) ∩ σ(v)|. 5) Lowering Algorithm: Algorithm 4 gives the complete lowering procedure. The runtime is ◦ ◦ O Keval · (|V ◦ | + |Edata | + |Epub |) , since the number of intra-step model pairs is bounded by |V ◦ |. H. Scheduling Given the executable graph Ĝ and the node-level time map T (delivered by our performance model), this section details the derivation of the execution time of the entire graph and derives a throughput-based score that serves as the optimization objective. The key challenge is that the graph’s makespan cannot be obtained by simply summing node execution times: nodes may execute concurrently when their precedence constraints are satisfied and their device allocations do not conflict, but must serialize when they compete for the same devices. Computing an optimal schedule under these constraints is NP-hard; we therefore use a greedy list-scheduling heuristic. Definition A.11 (Schedule): A schedule (s(·), f (·)) for Ĝ = (V̂ , Ê) assigns each node v ∈ V̂ a start time s(v) ∈ R≥0 and finish time f (v) = s(v) + T (v) such that: 1) Precedence. ∀(u → v) ∈ Ê : f (u) ≤ s(v). That is, a node may only start after all its predecessors have finished. 2) Device reservation. ∀u ̸= v ∈ V̂ : σ(u) ∩ σ(v) ̸= ∅ =⇒ [s(u), f (u)) ∩ [s(v), f (v)) = ∅. That is, nodes that reserve overlapping devices may not execute concurrently. Here σ(v) denotes the reserved device allocation: the full allocation for O P nodes, and the destination allocation for C OMM/R ESHARD nodes (source devices are free after send initiation).

Algorithm 4: L OWER: lower unrolled DAG to executable graph Input: Unrolled DAG Ḡ, execution plan P , RPEG G◦ , quantiles (fq , tq )Q q=1 (optional) Output: Executable DAG Ĝ = (V̂ , Ê) 1

V̂ ← ∅, Ê ← ∅;

/* Stage 1: Op nodes (optional quantile splitting) foreach (v, k) ∈ V̄ do 3 if r(v) = generation and quantiles given and ∃ disjoint downstream inf then gQ 1 4 replace xv,k by sub-nodes xgv,k , . . . , xv,k and chain g Q 1 → · · · → xv,k ; xgv,k

rank(v) = T (v) +

end

rank(w),

with max ∅ = 0. The rank approximates the remaining work on the longest path from v to any sink.

create O P node xv,k with σ(v), Π(v), mb(v);

max (v→w)∈Ê

else

Definition A.13 (Downstream Rank):

iQ 1 split each disjoint downstream inf into xiw,k , . . . , xw,k ; iq gq connect xv,k → C OMM → xw,k with scaled payload;

1) Scheduling Heuristic: The scheduler constructs a feasible schedule by processing nodes greedily, assigning each node the earliest start time that satisfies both constraints from the schedule definition above. When multiple nodes are ready simultaneously, the scheduler must decide which to schedule first. We prioritize nodes that lie on or near the critical path, since delaying them is most likely to increase the makespan. This is captured by the downstream rank:

end

/* Stage 2: Comm insertion on data edges */ /* For refined operators, resolve the corresponding input/output interfaces; do not duplicate phase edges already created in Stage 1. */ a 11 foreach (u, k) − → (v, k) ∈ Ēdata do 12 if σ(u) = σ(v) then 13 add xu,k → xv,k to Ê; 14 else 15 create C OMM node c with payload #a · s(a); 16 add xu,k → c → xv,k to Ê; 17 end 18 end

Definition A.14 (Earliest Feasible Start Time): For a ready node v (all predecessors scheduled), the earliest feasible start time is n o EST(v) = max max f (u), max avail(d) , (3)

/* Stage 3: intra-step weight redistribution */ Compute Pairs(m) from Topo(G◦ ) for each trainable m; 20 foreach (u, v, m) ∈ Pairs and k ∈ [Keval ] do 21 if τ (m) > 0 and σ(u) ∩ σ(v) = ∅ then 22 skip ; /* shadow pair */ 23 end 24 if σ(u) = σ(v) and Π(u) = Π(v) then 25 skip; 26 end 27 if σ(u) = σ(v) then 28 create R ESHARD r with payload #Wm · sb ; 29 add xu,k → r → xv,k ; 30 else 31 add xu,k → xv,k ; /* precedence only */ 32 end 33 end

Selection and tie-breaking. At each step, the scheduler selects the ready node with the smallest EST. Ties are broken by: 1) C OMM/R ESHARD nodes before O P nodes (to avoid artificial serialization of transfers), 2) larger downstream rank rank(v). After scheduling v at s(v) = EST(v), the availability of all devices in σ(v) is updated to f (v). Two additional tie-breakers (sub-node rank inheritance and RPEG-order consistency) are described after the algorithm.

/* Stage 4: cross-step publication edges a ′ → 34 foreach (u, k ) − (v, k) ∈ Ēpub do 35 if the exact required version already has an equivalent installation dependency then 36 reuse that dependency; 37 end 38 if σ(u) = σ(v) and Π(u) = Π(v) then 39 add xu,k′ → xv,k ; 40 end 41 else if σ(u) = σ(v) then 42 create R ESHARD r with payload #Wm · sb ; 43 add xu,k′ → r → xv,k ; 44 end 45 else 46 create C OMM c with payload #Wm · sb · (nmax − noverlap )/nmax ; 47 add xu,k′ → c → xv,k ; 48 end 49 end 50 return Ĝ;

Definition A.12 (Makespan): The makespan of a schedule (s(·), f (·)) is the latest finish time across all nodes: MS(Ĝ) = max f (v). v∈V̂

(u→v)∈Ê

d∈σ(v)

where avail(d) tracks when device d becomes free. The first term enforces precedence; the second ensures all blocking devices are available.

2) Scheduling Algorithm: Algorithm 5 returns MS(Ĝ). Since |V̂ | = O(Keval · |V ◦ |) and |Ê| = O(Keval · |E ◦ |), the runtime is O(Keval · |V ◦ | log(Keval · |V ◦ |) + Keval · |E ◦ |). The implementation uses a min-heap with lazy deletion, which accounts for the log factor. Sub-node rank inheritance. When generation quantile 1 splitting (Section A-G) produces a chain of sub-nodes xgv,k → gQ · · · → xv,k on the same devices, rank naturally decreases along the chain (later sub-nodes have shorter remaining paths). Without correction, the scheduler may interleave unrelated operators between generation phases, which disrupts KV-cache residency (the cache built during earlier phases would be evicted before later phases use it). To prevent this, all subnodes in a chain inherit the rank of the chain head (the first sub-node). RPEG-order tie-breaking. Beyond rank, the scheduler uses the RPEG vertex index as an additional tie-breaker: operators from the same RPEG position across different unrolled

penalty steepness. When rpeak ≤ 1 (feasible), the penalty is 1 and the score is unchanged. For moderately infeasible configurations (1–3% overflow), the penalty is small enough that the optimizer can temporarily accept high-throughput configurations that slightly exceed memory, then refine them toward feasibility through subsequent moves. At 10% overflow, the factor is e−4 ≈ 0.018; it decreases further for larger violations. The optimizer never returns an infeasible configuration as its final result; the penalty only guides intermediate search decisions.

Algorithm 5: G REEDY S CHEDULE(Ĝ, T ) Input: Executable DAG Ĝ = (V̂ , Ê), time map T , device space D. Output: Schedule (s, f ) and makespan. Compute rank(v) for all v in reverse topological order; Initialize avail(d) ← 0 for all d ∈ D; 3 R ← {v ∈ V̂ : v has no predecessors}; 4 while R ̸= ∅ do 5 Select v ⋆ ∈ arg minv∈R EST(v); /* Tie-break: Comm/Reshard before Op, then larger rank */ 6 s(v ⋆ ) ← EST(v ⋆ ), f (v ⋆ ) ← s(v ⋆ ) + T (v ⋆ ); ⋆ 7 foreach d ∈ σ(v ) do 8 avail(d) ← f (v ⋆ ); 9 end 10 Update R: remove v ⋆ , add newly-ready successors; 11 end 12 return max v∈V̂ f (v); 1 2

steps are scheduled in the same relative order. This produces steady-state regularity in the schedule, which is important for asynchronous execution where overlapping iterations should exhibit consistent device-usage patterns. I. Optimization Objective The executable graph Ĝ spans Keval unrolled iterations. Under bounded staleness, work from different iterations may overlap in time, so a per-iteration latency metric is not applicable: it cannot capture how consecutive iterations overlap. We instead use global throughput over the full evaluation horizon as the optimization objective. 1) Throughput Objective: Since Wstep is constant for a fixed workload, maximizing throughput is equivalent to minimizing the makespan. The optimizer uses Scorepenalized (C) (Equation 5) as its search objective; when the utilization objective is enabled, the penalty is applied to Scoreeff (C) (Equation 6) instead. The score captures the end-to-end effect of all optimization decisions: the RPEG structure G◦ , the staleness budget τ , the execution plan P = (σ, Π, mb), and structural choices such as stage fusion and placement strategies. Definition A.15 (Throughput Score): The throughput score of a configuration C = (G◦ , P ) is Score(C) =

Keval · Wstep MS(Ĝ)

(4)

where Wstep = BR(Lp + Lr ) is the number of tokens processed per iteration and MS(Ĝ) is the makespan from Algorithm 5. Memory penalty. Not all configurations are memoryfeasible. Rather than rejecting infeasible configurations outright (which would create hard boundaries in the search landscape), the optimizer applies a soft exponential penalty to the raw throughput score: Scorepenalized (C) = Score(C) · e

−βoom ·max(0, rpeak −1)

, (5)

where rpeak = maxd peak(d)/Meff is the worst-case peakto-capacity ratio across all devices and βoom = 40 controls the

2) Cost-Efficiency (Throughput×Utilization) Objective: Plans that partition devices across operators may achieve high throughput but leave a fraction of devices idle between operators. For example, two plans on the same 32-GPU cluster may achieve similar throughput, but one keeps all devices busy while the other leaves 8 devices idle whenever generation is not active. To steer the optimizer toward plans that are both fast and resource-efficient, we define an optional efficiencyaware score. Let tbusy (d) denote the total time device d spends executing vertices of Ĝ, and define the mean device utilization as 1 X tbusy (d) . Util(C) = |D| MS(Ĝ) d∈D Definition A.16 (Throughput×Utilization Score): Given the mean device utilization of configuration C, we define the efficiency score (throughput×utilization) as Scoreeff (C) = Score(C)·

( Util(C), λfloor ,

if Score(C) ≥ (1 − ζ) · Score⋆ , otherwise, (6)

where Score⋆ is the best raw throughput found so far during the search, ζ is a user-specified maximum throughput loss, and λfloor = 10−12 penalizes plans that fall below the throughput floor. J. Search and Optimization The preceding sections defined a deterministic evaluation pipeline that compiles any configuration C = (G◦ , P ) into a quantitative score. This section defines how the optimizer searches over the configuration space to maximize that score. The search operates on two levels. An outer loop enumerates structural variants of the RPEG (staleness budgets, actor-critic sharing), while an inner loop optimizes the execution plan P = (σ, Π, mb) for each fixed RPEG using simulated annealing with multi-vertex moves and elite-seeded warm restarts. The procedure is an anytime heuristic: the best feasible configuration found so far can be returned at any point, making it practical for users who simply set a time or evaluation budget. 1) Optimization Problem: Throughout this section, we fix the workload W, model specifications S, hardware H, and evaluation horizon Keval . The optimization variable is the configuration C = (G◦ , P ).

Definition A.17 (Optimization Objective): Given the scoring function from Definition A.15 (or Definition A.16), the optimizer seeks C ⋆ = argmax Score(C),

Ω = {C : VALID(C)}.

(7)

pipeline depth or the available GPUs after TP). A single data-parallel replica spans nrep = Pp /s nodes (divisibility guaranteed by P8 below), so the operator’s total allocation requires nneed = Pd · nrep nodes.

C∈Ω

(P6) Feasible layout. Either nneed ≤ Nn (all required nodes fit within a single rack) or Nn | nneed with nneed /Nn ≤ Nr A configuration C = (G◦ , P ) is valid if: (the allocation spans one or more complete racks). ◦ (V1) Well-formedness. G is well-formed (Section A-B). (P7) Aligned origin. The allocation origin is a multiple of the ◦ (V2) Plan feasibility. For every v ∈ V : extent in each dimension, so allocations tile the device (σ(v), Π(v), mb(v)) ∈ Xcand (v). space without overlap. ◦ (V3) Staleness feasibility. τ ∈ Ycand (G ) (Section A-J2b). (P8) Exact packing. s | Pp , so each allocated GPU hosts (V4) Memory feasibility. The peak memory on every device exactly one logical shard of the model: |σ(v)| = Pd · Pp · fits within GPU capacity. The final returned configuration P . t must satisfy this condition. During search, however, the Combining all of the above, the per-vertex plan candidate hard constraint is relaxed via the soft penalty from set collects every admissible combination of placement, parEquation 5, which allows the optimizer to temporarily allelization strategy, and microbatch count: explore near the memory boundary and refine toward feasibility. Definition A.18 (Plan Candidate Set): 2) Search Space: The configuration space is a union of structure-specific plan spaces: sharing or fusion can change the invocation set and therefore the available plan assignments. It has two conceptual components: the structural space over RPEGs G◦ and the plan space over execution plans P = (σ, Π, mb). a) Plan-Level Search Space: Using the nested 3D parallelization model, the plan for each vertex v must satisfy several constraints to ensure that the logical parallel structure maps cleanly onto the physical device topology. Let B be the batch size, Lm the layer count of the model invoked at v, and Blocal = BR/Pd (v) the per-replica sequence count. Parallelization degrees. (P1) Tensor parallelism. Pt (v) | Ng : TP is restricted to intranode placement, so the group size must divide the GPU count per node. (P2) Pipeline parallelism. Pp (v) | Lm and Pp (v) is 1 or even: each stage receives an equal share of layers. Odd depths > 1 cause asymmetric packing and are excluded. (P3) Data parallelism. Pd (v) | B and Pd (v) ≤ Nr · Nn · Ng (the total GPU count). Every replica processes the same number of samples. When the pipeline fits on a single node, intra-node DP is permitted. Microbatch count. (P4) No pipeline, no microbatching. If Pp = 1, then mb = 1. (P5) Microbatch alignment. If Pp > 1: mb(v) | Blocal ,

mb(v) ≥ Pp (v),

Blocal ≥ 8. mb(v)

The first ensures equal-sized microbatches. The second provides enough microbatches to fill the pipeline. The third maintains a minimum per-device microbatch size for SM utilization. Device packing. Given (Pd , Pp , Pt ), each node can host s = min(Pp , ⌊Ng /Pt ⌋) pipeline stages (limited by either the

Xcand (v) = {(σ, Π, mb) | P1–P8 hold}. Together, these constraints guarantee a unique deterministic mapping of logical shards to physical GPUs. b) Structure-Level Search Space: Plan-level modifications optimize intra-model parallelism and placement for a fixed RPEG. Inter-model modifications, in contrast, change the RPEG G◦ itself. The optimizer re-runs the plan-level search for each structural variant. The base RPEG (e.g., PPO) is specified chosen by the user and fixed for a given optimization run. Within that graph, the optimizer searches over two structural dimensions: (i) whether actor and critic parameters are shared, and (ii) the staleness budgets τ that control asynchronous execution. A third structural transformation, response-length quantile splitting, is not represented as an explicit structural move but applied automatically during lowering whenever the optimizer places generation and inference operators on disjoint device allocations. We represent inter-model choices by an inter-model configuration y = (ac, τ ),

ac ∈ {0, 1},

τ : Mtrain → N,

where ac = 1 indicates that actor–critic parameter sharing is enabled and τ is the staleness budget. Structural and inter-model constraints. We enforce the following constraints on y: (S1) Applicability of AC-sharing. The actor–critic shared variant (ac = 1) may only be selected if the base RPEG contains both an actor and a critic model. (S2) Staleness candidate set. Since τ : Mtrain → N is a function, the search enumerates its image componentwise: for each trainable model m ∈ Mtrain , τ (m) ∈ Tcand (m), where Tcand (m) is a finite set of allowable budgets (typically {0, 1}).

h ← H ASH(C); if h ∈ Cache then 3 return Cache[h]; 4 end ◦ 5 Ḡ ← U NROLL (G , Keval ) ; // Alg. 3 6 Ĝ ← L OWER (Ḡ, P ) ; // Alg. 4 7 (s, f ) ← G REEDY S CHEDULE (Ĝ) ; // Alg. 5 8 Score(C) ← Keval · Wstep / MS(Ĝ); 9 Apply memory penalty (Equation 5) to obtain Scorepen (C); 10 Cache[h] ← Scorepen (C); 11 return Scorepen (C);

2) Inner loop (plan search). For a fixed RPEG G◦ , optimize the execution plan P via multi-start simulated annealing (SA). Each iteration selects one or two RPEG vertices, proposes alternative plan assignments, and accepts or rejects the best proposal according to the Metropolis criterion. Multiple SA chains (restarts) are run per structural variant, seeded from an elite set of previously discovered plans. The overall procedure is anytime: it maintains the best feasible configuration across all structural variants and can return it whenever a wall-clock time limit or evaluation budget Emax is reached.

Definition A.19 (Inter-Model Candidate Set): The intermodel candidate set collects all admissible combinations of sharing flag and staleness budget: o n constraints (S1)–(S2) hold . Ycand = (ac, τ )

e) Score-Normalized Acceptance.: We use a stochastic acceptance rule to balance exploration and exploitation. Given the current configuration C with penalized score S > 0 and a proposed configuration C ′ with score S ′ > 0, define the relative score change

Selecting (ac, τ ) ∈ Ycand determines the structural RPEG variant and the staleness budget, which controls how publication edges are instantiated during unrolling (Section A-F). c) Evaluation Pipeline: Having defined what the optimizer searches over, we now describe how it scores each candidate configuration. Since local search frequently proposes configurations that have already been evaluated (e.g., a candidate that was explored in an earlier restart), results are memoized in a hash-based cache. The cache key is a canonical hash H ASH(C) derived from the RPEG structure (vertices, edges, containers, staleness budget τ ) and the execution plan P = (σ, Π, mb). The evaluation budget E counts only cache misses (full pipeline executions); cache hits are free. The algorithm first computes the raw score via the standard pipeline (unroll, lower, schedule), then applies the soft memory penalty from Equation 5. Runtime. On a cache miss, E VALUATE C ACHED executes the full pipeline. The dominant step is scheduling (Section A-H1), so the total cost per evaluation is O Keval · |V ◦ | log(Keval · |V ◦ |) + Keval · |E ◦ | .

S′ − S . S The proposal is accepted with probability ∆(C, C ′ ) pacc (C → C ′ ) = min 1, exp . t

Algorithm 6: E VALUATE C ACHED: configuration scoring Input: Configuration C = (G◦ , P ), evaluation horizon Keval , cache. Output: Penalized score Scorepen (C). 1 2

Cache hits are O(1) (hash lookup). d) Optimization Strategy: The inter-model candidate set Ycand is small and can be enumerated exhaustively. The combinatorial difficulty lies in the plan space: each vertex v admits Q |Xcand (v)| valid choices, and the joint plan space grows as v∈V ◦ |Xcand (v)|, which is exponential in the number of RPEG vertices. Similar to ReaL [31], we associate each configuration with a cost defined by its (penalized) score and formulate the search as sampling from an energy-based distribution over configurations. We augment this approach with multi-vertex proposals, best-of-Kx candidate selection, and elite-seeded warm restarts. We therefore adopt a two-level strategy: 1) Outer loop (structural enumeration). Enumerate all y ∈ Ycand . For each structural variant, run an independent plan-level search and retain the best feasible configuration found.

∆(C, C ′ ) =

(8)

(9)

The temperature t > 0 controls exploration: improving proposals are always accepted, while worsening proposals may also be accepted. Normalization by S makes the acceptance rule largely independent of the absolute score scale. This is a stochastic search heuristic rather than exact sampling from a prescribed target distribution. Since Score(C) is a product of a positive throughput and a positive penalty factor, S > 0 always holds. The normalization by |S| makes the acceptance threshold scale-invariant: the same temperature t produces comparable acceptance rates regardless of the absolute score magnitude. This differs from the standard Metropolis ratio exp((S ′ − S)/t), which would require re-tuning t for every workload. Improving proposals are always accepted. Crucially, worsening proposals are accepted with non-zero probability, which allows the chain to escape local optima by temporarily traversing lower-scoring regions of the landscape. This is the central advantage over greedy search: a purely greedy optimizer can become permanently trapped in a suboptimal basin, whereas the stochastic acceptance mechanism provides a principled way to trade short-term score regression for long-term exploration. As detailed in Section A-J2g, t is held constant throughout each chain rather than decreased as in classical SA. The exploration–exploitation tradeoff is instead controlled through the multi-start mechanism described in Section A-J2h. f) Move Proposals: Evaluating a full configuration through the pipeline is expensive relative to enumerating valid strategy tuples. To avoid wasting evaluations on invalid or obviously poor candidates, the optimizer pre-computes a candidate directory for each RPEG vertex before the SA search begins. The directory contains all valid strategy tuples satisfying constraints P1–P8, filtered by per-vertex memory

feasibility (strategies that exceed GPU capacity even in isolation are excluded; the soft penalty from Equation 5 handles joint infeasibility across co-resident operators during search) and sorted by estimated execution time. Candidates are sampled uniformly from the full directory. Each SA iteration proposes up to Kx = 6 neighbor configurations by modifying the plan of one or two RPEG vertices while keeping G◦ fixed. The proposal mechanism uses two move kernels. Single-vertex kernel. A target vertex is selected with a bias toward scheduling bottlenecks: with probability 0.6, the vertex is sampled proportionally to its execution time on the critical path (the longest path through the current schedule). Otherwise it is chosen uniformly at random. For the selected vertex, Kx candidate strategies are sampled uniformly from the directory and each is paired with a randomly chosen aligned device origin. Donor-receiver pair kernel. Single-vertex moves cannot redistribute devices between operators in a single step: shrinking one vertex’s allocation without simultaneously expanding another leaves devices idle. To address this, a donor-receiver kernel proposes coordinated two-vertex moves. A donor is selected with probability proportional to its device count, and a receiver is chosen among vertices whose allocation overlaps with the donor’s. The donor is shrunk to a smaller contiguous region, and the receiver is re-sampled to potentially expand into the freed devices. For example, if generation occupies two full racks while inference is constrained to a single node, a pair move can transfer one rack from generation to inference in one step. No sequence of single-vertex moves can reach this configuration without passing through intermediate states where devices are left unassigned. The pair kernel is selected with probability ppair = 0.28. Best-of-Kx selection. All Kx proposals are evaluated via E VALUATE C ACHED (Section A-J2c), with cache misses dispatched to a thread pool for parallel evaluation. Only the highest-scoring proposal Cbest is subjected to the acceptance step (Equation 9). This best-of-Kx strategy amortizes evaluation cost across multiple local directions before committing to a single chain update. g) Temperature and Chain Termination: Unlike classical simulated annealing, which decreases t over the course of a single run, our optimizer holds t constant within each chain. The first chain runs at t = tinit ; warm-restarted chains run at t = twarm > tinit , which allows broader exploration around the elite seed. The higher restart temperature is safe because the chain starts from a known-good plan rather than an arbitrary initial point. The exploration–exploitation tradeoff is therefore managed across chains (via the two temperature levels and the multi-start mechanism in Section A-J2h) rather than within a single chain. The implementation also supports geometric cooling (t ← max{tmin , γ ·t}) as an alternative; our ablation finds that it degrades throughput by 1–5% compared to constant temperature.

Chains are not run for a fixed number of iterations. Instead, the optimizer monitors consecutive rejections: if wstag = 10 proposals are rejected in a row, the chain is terminated and a new restart begins. This ensures that productive chains continue while stuck chains are recycled quickly. h) Restarts and Elite Seeding: Random restarts are a standard technique to escape local optima in SA, but they waste budget rediscovering plans that earlier chains already found. We instead use warm restarts: the first chain (r = 0) is seeded from a concurrent baseline heuristic that groups operators into concurrent groups – sets of operators with no data dependencies that can run in parallel – and partitions the device space so that each operator within a group receives a disjoint allocation, selecting the partition that minimizes the estimated step time. All subsequent chains begin from the best feasible plan discovered so far (the elite seed) and run at twarm > tinit , allowing broader exploration around the knowngood starting point. The optimizer maintains an elite set of the top-Kelite = 3 feasible plans per vertex set, shared across staleness variants that share the same vertices (since a good plan under one τ is a structurally sound seed for another). i) Outer Loop and Budget: The outer loop iterates over all inter-model candidates y = (ac, τ ) ∈ Ycand . Termination is governed by either a wall-clock time limit or an evaluation budget Emax (cache misses), whichever is reached first. In time-limited mode, the remaining wall-clock time is split evenly among the remaining structural variants; in budgetlimited mode, each variant receives a fixed budget Estruct . The candidate directory (Section A-J2f) is rebuilt only when the vertex set changes (e.g., when transitioning from the base RPEG to the AC-shared variant); staleness variants that share the same vertices reuse the same directory. The best feasible configuration C ⋆ is maintained across all structural variants. In practice, a 10-minute time budget suffices for workloads up to 1024 GPUs, producing 75K–220K evaluations depending on cluster and model size. After the SA search completes, a deterministic polish phase (Section A-J2j) performs greedy coordinate descent to extract any remaining local improvements. Algorithm 7 gives the complete procedure. j) Polish Phase: The SA search is stochastic: even when the chain has converged near a good basin, it may not have evaluated every neighbor of the final plan. The polish phase addresses this by performing deterministic greedy coordinate descent [54] on the best plan returned by O PTIMIZE. It visits each RPEG vertex in a fixed order and, for each vertex, evaluates all feasible strategies in the candidate directory. If any strategy strictly improves the score, the best one is applied before moving to the next vertex. This process repeats in passes over all vertices until a full pass yields no improvement. Two properties make the polish phase efficient. First, isolated local operator costs can be reused across identical strategy/extent choices on homogeneous hardware. Second, many candidates will already reside in the evaluation cache from the preceding SA search, so the polish phase primarily

Algorithm 7: O PTIMIZE: multi-start SA with structural enumeration Input: Baseline C0 , candidates Ycand , time limit or budgets Emax /Estruct , temperature t, stagnation window wstag , proposal count Kx . Output: Best feasible configuration C ⋆ . foreach (ac, τ ) ∈ Ycand do for each restart r until budget exhausted do 3 Seed C from baseline (r=0) or elite plan (r≥1); 4 repeat 5 Propose Kx neighbors (single-vertex or pair kernel) ; // Section A-J2f 6 Evaluate proposals via E VALUATE C ACHED ; // Alg. 6 7 Cbest ← argmax among proposals; 8 Accept Cbest with prob. pacc (C → Cbest ) ; // Eq. 9 9 Update elite set and global best C ⋆ ; 10 until wstag consecutive rejections or budget exhausted; 11 end 12 end ⋆ ⋆ 13 C ← P OLISH (C ) ; /* Section A-J2j */ ⋆ 14 return C ; 1

performs hash lookups rather than full pipeline evaluations. Cache misses are dispatched to the thread pool in parallel, as in the SA inner loop. The polish phase is guaranteed to terminate (the score increases strictly at each accepted move, and the candidate set is finite) and is deterministic (no random choices). It cannot discover plans that require coordinated multi-vertex changes, but it reliably extracts any single-vertex improvements that the stochastic search may have missed.

A PPENDIX B E VALUATION : S UPPLEMENTARY D ETAILS This appendix specifies different evaluation configurations. Platforms. GH200 execution and hardware calibration use Clariden at CSCS, with four NVIDIA GH200 superchips per node (each with a 72-core Grace CPU and 120 GB HBM3), intra-node NVLink connectivity (≈160 GB/s per GPU pair), and four HPE Slingshot-11 interfaces per node. The runtime uses persistent PyTorch/Megatron workers launched through Slurm, NCCL over OFI/libfabric/CXI for GPU collectives and weight transfers, and Gloo for semantic data. The A100 experiments use Athena at AGH Cyfronet, with dual 64-core AMD EPYC 7742X CPUs, 1 TB RAM, and eight NVIDIA A100 40-GB GPUs per node, interconnected via NVLink/NVSwitch and four 200-Gb/s HDR InfiniBand interfaces. The A100 comparison with ReaL uses a ReaL-derived runtime (Python 3.10, CUDA 12.1; one worker per GPU) and matched synchronous PPO workloads, without response-length splitting or bounded asynchrony. Analytical optimizer searches run on Piora CPU nodes at CSCS, each with dual 64-core AMD EPYC 7773X processors and 1 TB RAM, using six CPU cores per concurrent search task. Workload hyperparameters. We use one PPO epoch, actor/critic learning rates 10−6 , zero weight decay, gradient clipping at 1.0, PPO clip 0.2, γ = 1, GAE λ = 0.95, and KL coefficient 0.01. Analytical model validation. The 123 configurations comprise 34 generation, 51 inference, and 38 training cases. These cases span 3B, 7B, and 13B models, 1–128 GPUs, batch sizes of 32–2048, and 25 DP/PP/TP strategies. Measurements use BF16, FlashAttention-2, Lp = 512, Lr = 128, and one rollout per prompt; training uses activation checkpointing and ZeRO-2 for Pd > 1. The memory comparison accounts for checkpointing and ZeRO-2 profiling conventions; Pd = 1 traces recorded before Adam-state materialization are excluded. Predicted-schedule configurations. Schedules unroll Keval = 8 steps with one-step bounded staleness and three response-length phases: completion fractions (0.50, 0.85, 1.0) at lengths (Lr /4, Lr /2, Lr ). The throughput–utilization objective uses δ = 0.10. Three scaling sweeps cover actoronly, critic/reward, and joint model growth. The first two use 8/16/32/64 GPUs: the scaled models grow as 7/13/34/70B, while unscaled models remain at 7B. The joint sweep uses 16/32/64/128 GPUs, with actor and critic sizes 7/13/34/70B and reward sizes 7/7/13/13B. Each configuration is evaluated at (B, Lr ) ∈ {(512, 128), (256, 384), (128, 896)}, yielding 36 analytical cases. The strong-scaling study fixes 13B models and the workload specified in the main evaluation. Optimizer stability. Table X specifies five modeled configurations used in the 200-seed stability study. Each CPUside analytical search runs for 600 seconds; its throughput is the number of evaluated plans, not the number of physical training runs. A separate convergence study evaluates 7B/256, 13B/512, 34B/1024, and 70B/2048 model/GPU configurations over 3,600 seconds, tracking the best predicted makespan relative to the initial plan. Ablations. Each experiment disables one mechanism, with five seeds per scenario. The 50 GB HBM limit in M EM -

TABLE X: Optimizer-stability workloads. A/C/R denotes actor/critic/reward model size (billions of parameters); topology denotes racks × nodes/rack × GPUs/node. GPUs Topology

A/C/R (B)

64 2×8×4 13/13/7 128 2×16×4 34/13/13 256 4×16×4 70/34/13 512 8×16×4 70/70/34 1024 16×16×4 175/70/34

B Lp /Lr Evals/600 s 1024 128/384 1024 128/384 1024 128/384 2048 128/384 2048 128/384

∼220K ∼200K ∼160K ∼105K ∼75K

TIGHT is imposed analytically to evaluate memory-constrained planning. Table XI identifies the four workloads; results are reported in Table IV.

TABLE XI: Analytical ablation workloads. A/C/R denotes actor/critic/reward model size (billions of parameters); topology denotes racks × nodes/rack × GPUs/node. Scenario M EM - TIGHT W IDE H ETERO D EEP

GPUs Topology HBM/GPU A/C/R (B) B; Lp /Lr 16 1×4×4 50 GB 128 1×32×4 102 GB 32 1×8×4 102 GB 64 2×8×4 102 GB

13/13/13 7/7/7 34/7/7 7/70/70

128;128/896 1024;128/384 512;128/128 512;128/128

A1 disables response-length splitting; A2 disables bounded asynchrony; A3 fixes microbatch counts to Pp for generation/inference and 2Pp for training; A4 removes coordinated donor–receiver moves; A5 disables warm restarts; A6 uses uniform rather than critical-path invocation selection; A7 removes final polishing; A8 replaces soft memory penalties with hard rejection; and A9 uses geometric cooling instead of constanttemperature chains.

A PPENDIX C A NALYTICAL P ERFORMANCE AND M EMORY M ODEL : A DDITIONAL D ETAILS This appendix complements Section V with the assumptions, calibration procedure, and implementation choices behind the analytical performance model. We follow the same order as the main paper: model configuration and communication, intra-model execution time, inter-model communication, and device-level memory feasibility. The role-specific work– depth–memory foundations are given in [17]. The RPEGspecific execution-time, communication, and physical devicememory equations are derived in Appendix D; we do not duplicate them here. The analytical model assigns node times T (x) and checks memory feasibility MemFeasible(C) to rank execution plans, not to prescribe physical start times. A. Model Configuration and Calibration 1) Inputs and Modeling Assumptions: The analytical model receives three groups of inputs. The workload W = (B, R, Lp , Lr ) specifies input prompts per step, rollouts per prompt, average prompt length, and average rollout length. Each invoked model m ∈ M supplies its transformer architecture, model role, and working scalar precision sb . The hardware description supplies the hierarchical device topology, GPU memory capacity, and calibrated compute and communication parameters. The candidate execution plan P = (σ, Π, mb) assigns each invocation v its GPU allocation σ(v), DP/PP/TP degrees Π(v) = (Pd , Pp , Pt ), and microbatch count mb(v). We use mixed-precision BF16 computation, FP32 Adam master weights and moment arrays, activation checkpointing, and ZeRO-2 sharding for optimizer state and reduced gradients. The estimator targets GPU-resident training rather than CPU offloading. These choices determine both the HBM accounting and the communication primitives costed within a training invocation. Parameter formulas retain the dominant attention projections, FFN weights, embeddings, and task heads; small bias and normalization parameters are omitted in the baseline. The nominal hardware topology is given as D = [Nr ] × [Nn ] × [Ng ]. Device placements are mapped to physical nodes and racks, rather than assuming a flat, uniform interconnect. The time model distinguishes three communication tiers: intranode, inter-node within a rack, and inter-rack. When a transfer spans more than one tier, it uses the slowest applicable tier as a conservative proxy. For tensor parallelism, the legal plan space constrains each TP group to one node. These assumptions are specific to the evaluated modeling implementation; they are not architectural requirements of RPEG. In particular, architectures with grouped-query attention, gated FFNs, vocabulary-sharded heads, or nonstandard pipeline partitions require architecture- and layout-specific dimensions in the cited work terms and our physical residency estimator (Appendix D-H). 2) Automated Hardware Calibration: Rather than substituting advertised GPU FLOP/s or network line rates, RPEG measures effective parameters with an automated suite on the

target system. Its calibration campaign, typically on the order of minutes, comprises the following measurements: NCCL AllReduce message-size sweeps; NCCL send/receive sweeps; actual transformer-forward timing; contiguous HBM transfers; GEMM-shape sweeps; and non-GEMM layer timings at a reference batch size. The resulting configuration is stored in a hardware-parameter file and supplied to the optimizer together with the workload, model specifications, and available devices. The measured quantities have distinct meanings. Transformer-forward benchmarks provide an effective compute ceiling Fpeak , while memory-copy tests provide an effective BWmem . Shape sweeps fit the saturation constants γM and γK ; non-GEMM measurements fit cng at the reference batch size Bcal . Communication measurements fit separate (α, β) pairs for collective and P2P transfers at each hardware tier; α is latency and β is inverse effective bandwidth, not bandwidth itself. The AllReduce benchmark calibrates a collective-family proxy used by the ring-style expressions for AllReduce, ReduceScatter, and AllGather. NCCL may select different algorithms as message sizes, rank counts, and network topology change, so a fitted (α, β) pair should be understood as an empirical approximation over the measured range. Similarly, P2P and collective parameters are not interchangeable: a high aggregate collective bandwidth does not imply equivalent point-to- point transfer performance. Note that the model can evaluate new plans cheaply, but its parameters need not transfer accurately across different kernel implementations, model architectures, message-size regimes, or cold versus warmed process states. 3) Hierarchical Communication Model: Communication appears in two places: the intra-operator term Tcomm (v) for TP, PP, and DP transfers embedded in O P nodes, and the explicit transfer costs Txfer and Treshard for C OMM/R ESHARD nodes inserted by lowering. Both use the α–β model and the primitive estimates in Table II of the main paper. For those estimates, M is the per-GPU input buffer for AllReduce and ReduceScatter, the final gathered buffer for AllGather, the aggregate per-GPU exchange size for AlltoAll, and the transferred P2P payload. Collective costs vanish when the participant count is one. The table uses ring-based reduction/gather expressions and a pairwise- exchange upper-bound estimate for AlltoAll. The all-to-all formula counts M β as an aggregate bandwidth term, including any locally retained share; a perfectly balanced pairwise exchange instead sends (p − 1)M/p remote bytes when M includes the self-share. This and tier selection make the model a cost proxy rather than a packet-exact simulation of the network. Per-message control overhead, dynamic communicator setup, and transport-specific tensor packing are outside this simple expression. B. Intra-Model Runtime Details 1) Operator Work and Effective Throughput: For every O P node v, the five-term operator-time model Top (v) in Section V-B separates GEMM compute Tgemm (v), pipeline bubbles Tbubble (v), intra-operator communication Tcomm (v),

non-GEMM overhead Tnon-gemm (v), and Adam update traffic Topt (v). The role-specific FLOP expressions for generation, teacher-forced inference, and training are adopted from [17]. They distinguish policy, reference, value, reward, and sharedmodel heads. The sum defining Top (v) is an additive prediction. Actual computation, collectives, and pipeline transfers can partially overlap, so it is not an exact critical-path reconstruction of the kernels executed inside a model invocation. The roofline-limited GEMM throughput Feff (v) depends on the GEMM-shape efficiency η(v) and arithmetic intensity I(v), not just nominal work W (v). In particular, increasing Pd reduces the local batch, while increasing Pt narrows local projection matrices. Either can lower achieved per-GPU GEMM throughput enough to offset the benefit of additional devices. Arithmetic intensity I(v) = Wdev (v)/(ΩW + ΩA ) uses estimated HBM bytes read/written per microbatch ΩW + ΩA , whereas peak(d) measures simultaneous GPU-memory occupancy: a resident tensor can be accessed repeatedly. For training, I(v) uses forward-equivalent FLOPs and byte traffic for forward, recomputation, and backward. Generation uses prefill decode separate effective throughputs Feff (v) and Feff (v) in gen Tgemm (v) because the two phases have different memoryaccess patterns. In particular, decode repeatedly loads or accesses parameters and historical KV state while producing relatively little new arithmetic work. The evaluated estimator uses representative lengths Lp , Lr rather than tracking timevarying active-sequence counts or per-rank KV-head layouts during decoding. 2) Pipeline Bubbles and Intra-Operator Transfers: The 1F1B bubble multiplier (mb(v) + Pp (v) − 1)/mb(v) applied to Tgemm (v) models fill and drain under balanced execution. It decreases as the microbatch count grows but ignores unbalanced stage durations and communication contention. The search therefore optimizes microbatch count mb(v) jointly with the parallel degrees Π(v) = (Pd , Pp , Pt ) and placement σ(v), rather than fixing it solely from Pp (v). A larger number of microbatches can reduce bubbles while increasing latencysensitive collective calls. Tensor-parallel AllReduce contributes to Ttp (v) by aggregating partial layer results; TP groups are restricted to one node because the repeated per-layer reductions favor lowlatency, high-bandwidth local links. Pipeline P2P operations move activations between consecutive stages and contribute to Tpp (v) according to the boundaries crossed by the selected placement. Data-parallel ReduceScatter aggregates and partitions gradients for ZeRO-2, while the associated AllGather restores the updated parameter view across replicas; both contribute to Tdp (v). The latter is a parameter-publication cost and need not be adjacent to gradient reduction in the actual training schedule. Table XII records the primitive counts used by the simplified estimator. TP assumes two AllReduces per transformer layer and forward pass (attention and FFN), or three forwardequivalent passes for training. PP counts approximate criticalpath fill/drain steps: mb + Pp − 2 = (mb − 1) + (Pp − 1) for inference and mb + 2Pp − 3 = (mb − 1) + 2(Pp − 1) for 1F1B

training. ZeRO-2 contributes one gradient ReduceScatter and one updated-weight AllGather per update (Appendix D-D). Generation TP is split into prompt prefill and decode; the decode count assumes that the full local DP batch is processed together. A decoder that independently pipelines distinct microbatches needs the corresponding alternate count, as explained in Appendix D-D. The pipeline-transfer byte count Mpp (ℓ) presumes TPsharded stage-boundary activations; frameworks that send replicated hidden states need a different transfer volume. All PP costs vanish for Pp = 1, and TP/DP collectives vanish for a single participating TP/DP rank. 3) Non-GEMM Kernels and Optimizer Updates: The GEMM work formulas define W (v) using dominant matrix multiplications and exclude normalization, residual operations, nonlinearities, positional operations, and other nonGEMM kernels, whose residual latency is represented by Tnon-gemm (v) rather than added to the GEMM FLOP count. At Pt = 1, the full-transformer-calibrated compute ceiling Fpeak already absorbs part of this overhead. The separate calibrated non-GEMM term Tnon-gemm (v) therefore adds only a residual correction as tensor parallelism Pt increases. Its batch scaling and four-pass-equivalent training multiplier are heuristics; launch latency or fused-kernel behavior can depend on microbatch count and sequence length in ways not captured gen by one constant. The baseline sets Tnon -gemm (v) = 0, which is an approximation rather than an assertion that prefill and decode contain no such operations. The bandwidth-bound Adam estimate is Topt (v) = 18Nm /(Pd Pp Pt BWmem ). Here Nm /(Pd Pp Pt ) approximates the optimizer-owned parameters per GPU, and 18 bytes/parameter is an effective coefficient, not the physical read/write total: an uncached update requires 24 bytes for FP32 master/moment reads and writes, 2 for the BF16 gradient read, and 2 for the BF16 weight write (28 bytes per parameter). Replicated boundary embeddings or heads can make the worst-stage count larger (Appendix D-F). C. Inter-Model Communication and Resharding Details A C OMM node has predicted time Txfer (u, v, a) for container a between allocations σ(u), σ(v). The estimator divides its payload |a|s(a) by the assumed number of concurrent P2P pairs min(|σ(u)|, |σ(v)|) and applies the communication tier’s (αp2p , βp2p ). This assumes balanced transfers and ideal overlap among those pairs. A R ESHARD node redistributes parameters when the same model is invoked with incompatible parallel layouts on an identical GPU allocation. The analytical proxy treats the transformation as a symmetric all-toall TA2A (Nm sb /|σ(u)|, |σ(u)|) over that allocation. Actual Treshard (u, v, m) depends on the layout and backend: multiple collectives, nonuniform slices, or host-mediated copies can be much slower than this idealized NCCL all-to-all, even at equal payload. These costs are especially important for bounded asynchrony τ (m) > 0 and training/serving handoffs. The runtime may hold two parameter versions, reserve both source and

TABLE XII: Simplified intra-operator communication costs per GPU. Here ℓ = Lp + Lr for teacher-forced inference/training; prefill has ℓ = Lp . PP counts approximate critical-path steps. The generation decode row assumes one decode step over the full local DP batch. All payloads are per participating GPU. Role

Parallelism

Primitive

Count

Message bytes

Generation prefill Generation decode Inference Training Inference Training Training

TP TP TP TP PP PP DP

AllReduce AllReduce AllReduce AllReduce P2P P2P ReduceScatter + AllGather

2(Lm /Pp ) mb 2(Lm /Pp )(Lr − 1) 2(Lm /Pp ) mb 6(Lm /Pp ) mb mb + Pp − 2 mb + 2Pp − 3 One each

BRLp dm sb /(Pd mb) BRdm sb /Pd BRℓdm sb /(Pd mb) BRℓdm sb /(Pd mb) BRℓdm sb /(Pd Pt mb) BRℓdm sb /(Pd Pt mb) Nm sb /(Pp Pt )

destination worker allocations, and require redistribution before an invocation can begin. Their exact physical constraints must be checked during lowering and execution; the analytical model supplies a plan-ranking approximation, not a guarantee that a transfer follows its idealized route. The formula-level details appear in Appendix D-G. D. Memory Model Details 1) Parameter Shards and Persistent State: The memory model distinguishes the logical model parameter count Nm from the worst-stage per-GPU parameter count Wshard (v). Transformer layers are partitioned across pipeline stages and tensor-parallel ranks, whereas boundary embeddings and task heads may remain replicated in the simplified layout. The worst-stage parameter shard Wshard (v) is therefore not obtained by dividing the model’s total parameter count uniformly by all GPUs. Tied input/output embeddings may share storage when both are on one stage; separate boundary stages can retain physically separate copies of one logically tied matrix. The full stage-level count is derived in Appendix D-H. For a training device, the persistent analytical footprint S(m, d) comprises ZeRO-2-sharded FP32 master weights and Adam moments. BF16 working weights are counted when materialized, along with any live gradient buffers. A frozenmodel or disjoint asynchronous serving device keeps the appropriate BF16 parameter shard persistently. A synchronous serving allocation that only receives parameters transiently does not, under the simplified lifetime policy, retain a separate full persistent copy between invocations. Distinct models colocated on one GPU contribute their persistent footprints S(m, d) additively. For a shared actor– critic backbone, shared parameters must be counted once, with separate policy/value heads added according to their actual residence. For partial overlap or multiple versions of one model, residence and copy lifetimes must be determined by exact allocations and version dependencies, not by the number of logical invocation vertices alone. 2) Transient Activations, Gradients, and KV Cache: T Training activation memory Mact (v) depends on up to min(mb(v), Pp (v)) live pipeline microbatches, sequence length Lp + Lr , activation-checkpoint interval c, and intermediate tensor widths. The square-root checkpoint interval is a practical heuristic rather than the exact optimum for arbitrary activation-buffer coefficients. The baseline activation expression conservatively reserves a recomputation buffer for each live warm-up microbatch; a tighter expression is also

derived in Appendix D-I, but was not substituted into the reported baseline calculations. Teacher-forced inference retains only local forward intermeI diates Mact (v), whereas autoregressive generation additionally retains a KV cache Mkv (v) whose peak scales with the number of active sequences, their admitted context lengths, and the physical KV-head width on each GPU. Average lengths alone need not bound worst-case cache residence. In groupedquery attention, KV storage cannot generally be estimated by dividing the full hidden width by tensor-parallel degree: ranks may own different numbers of KV heads, including replicated groups. The role-specific head also matters for transient buffers: a critic or reward scalar head does not generally require a vocabulary-wide output allocation. The transient footprint A(v, d) additionally contains materialized BF16 weights, live gradients and buckets, transfer/resharding workspaces, and kernel scratch buffers when they are not already classified as persistent. Double counting must be avoided when BF16 serving weights are already resident. The simple activation/KV expressions are only components of the complete transient bound, not its entirety. 3) Co-Residency and Device Feasibility: As described in Section V-D3, the analytical device peak peak(d) adds all co-resident persistent model state S(m, d) and the largest transient allocation A(v) assigned to GPU d. The maximum is appropriate when operations assigned to one GPU execute exclusively and their transient buffers have the assumed lifetime. When several worker processes persist on one GPU, allocator reservations, retained model versions, and transfer buffers can overlap beyond that bound. The device-feasibility predicate MemFeasible(C) checks every GPU against Meff = (1 − εrt )MGPU , with a calibrated safety fraction for fragmentation and other runtime overheads. The nominal 10% reserve used in the analytical setup is a modeling parameter, not a universal physical guarantee; actual allocated and reserved HBM and the effective per-run capacity govern runtime admission. Search may assign a soft penalty to memory-violating proposals, but only feasible final plans are selected for execution. The five placement cases in Figure 3 of the main paper illustrate the distinction: synchronous disjoint serving uses transient weights; asynchronous disjoint serving maintains shadow copies; partial overlaps mix the two; and fully colocated execution may require layout resharding. The equations for S(m, d), A(v, d), per-device peaks, and the memoryfeasibility predicate are given in Appendix D-J.

A PPENDIX D A NALYTICAL D ERIVATIONS

b) Arithmetic intensity.: For inference and training operators processing sequences of length

This appendix gives the RPEG-specific performance-model derivations omitted from Section V. We reuse the work–depth– memory foundations of [17] instead of repeating that paper’s per-role FLOP and logical-storage derivations. Our focus is the calibrated execution-time estimator and its physical perGPU memory accounting. We denote node runtime by T . For an operator node v, let

ℓ = Lp + Lr ,

Π(v) = (Pd , Pp , Pt ),

P = P d Pp Pt ,

and m = µ(v) be the invoked model. The model architecture is (Lm , dm , dff,m , Vm , dh,m ), where Lm is the number of transformer layers, dm is the hidden dimension, dff,m is the FFN dimension, Vm is the vocabulary size, and dh,m is the attention head dimension. The runtime tensor precision uses sb bytes per scalar. We use B prompts per step, R rollouts per prompt, and representative average prompt and rollout lengths Lp and Lr .

A. Operator-Node Runtime For an operator node v ∈ V̂ , the runtime model decomposes execution time into five terms: T (v) = Tgemm (v) + Tbubble (v) + Tcomm (v) + Tnon-gemm (v) + Topt (v). The terms capture GEMM compute, pipeline bubbles, intraoperator communication, non-GEMM overheads, and optimizer traffic, respectively. The decomposition is necessary because a single peak-throughput estimate T (v) ≈

W (v) Pd Pp Pt Fpeak

does not capture small-batch GEMM under-utilization, nonGEMM operations that do not shrink with tensor parallelism, or memory-bandwidth-bound optimizer traffic. a) GEMM time.: Under 3D parallelism, the total GEMM work W (v) is distributed over the P devices. Assuming balanced work across the 3D mesh, the raw GEMM time is Tgemm (v) =

W (v) , P Feff (v)

where Feff (v) is the achieved per-device throughput. We model it by a roofline-style bound corrected for calibrated GEMM utilization: Feff (v) = min(Fpeak η(v), I(v)BWmem ) . Here, Fpeak is workload-effective peak throughput, BWmem is HBM bandwidth, I(v) is arithmetic intensity, and η(v) is GEMM utilization.

the per-device arithmetic intensity is I(v) =

Wdev (v) , ΩW + Ω A

where Wdev (v) is per-device work for one microbatch. For inference, Wdev = WI (v)/(P mb); for training, Wdev = WT (v)/(4P mb), representing one forward-equivalent pass. The HBM-traffic terms ΩW and ΩA count weight and activation bytes, respectively, per device and microbatch: Lm 4d2m + 2dm dff,m sb , ΩW = Pp Pt Lm ΩA = 2 bloc ℓ dm sb , Pp . with local microbatch size bloc = PBR d mb The factor four above converts training FLOPs to a single forward-equivalent pass. ΩW assumes one weight read, and ΩA one activation read and write per local layer. For generation, prefill uses ℓ = Lp and decode uses ℓ = 1 per token. c) GEMM utilization and generation.: The calibrated GEMM utilization is min(dm /Pt , dff,m /Pt ) bloc ℓ η(v) = · . bloc ℓ + γM min(dm /Pt , dff,m /Pt ) + γK The first factor captures under-utilization from small row dimension bloc ℓ; the second captures under-utilization from tensor-parallel narrowing of the hidden or FFN dimension. The constants γM , γK are obtained from calibration. For generation, we evaluate the two phases separately: gen Tgemm (v) =

Wprefill (v) prefill P Feff (v)

Wdecode (v) . decode (v) P Feff

B. Role-Specific Work Terms The work–depth–memory analysis of generation, inference, and training (including the actor, reference, value, reward, and shared-model heads) is already derived in [17]; we do not reproduce it here. We use the corresponding leadingorder role-specific FLOP terms WG (m), WI (m), and WT (m) to define W (v) for each executable invocation, expressed in the workload notation (B, R, Lp , Lr ). The following sections instead derive the RPEG-specific translation of this work into predicted GPU time, communication costs, and physical memory feasibility. C. Pipeline Bubble When Pp > 1, the Pp pipeline stages process mb microbatches. Under a 1F1B schedule, the fill/drain bubble contributes Pp − 1 idle slots out of mb + Pp − 1 total slots. Equivalently, the combined GEMM and bubble time is mb + Pp − 1 Tgemm (v) + Tbubble (v) = Tgemm (v) . mb When Pp = 1, the factor is one and Tbubble (v) = 0.

D. Intra-Operator Communication The intra-operator communication term is Tcomm (v) = Ttp (v) + Tpp (v) + Tdp (v). The three terms correspond to tensor-parallel AllReduce, pipeline-parallel point-to-point activation transfer, and dataparallel gradient synchronization. For generation and inference, Tdp (v) = 0. a) Tensor-parallel AllReduce: Each transformer layer requires two tensor-parallel AllReduces per forward pass: one after the attention block and one after the FFN block. For inference over length ℓ = Lp + Lr , BR ℓ dm sb Lm I mb TAllReduce , Pt . Ttp (v) = 2 Pp Pd mb For training, the simplified forward/recompute/backward count uses three forward-equivalent phases, giving 6Lm /Pp collectives per microbatch: BR(Lp + Lr )dm sb Lm T , Pt . mb TAllReduce Ttp (v) = 6 Pp Pd mb For generation, prefill and decode are evaluated separately: Lm G Ttp (v) = 2 · Pp BR Lp dm sb , Pt + mb TAllReduce P mb d BR dm sb , Pt . (Lr − 1)TAllReduce Pd The above decode term represents one batched decode step over all local DP sequences, while prefill is microbatched. If decode is also executed as mb distinct microbatches, consistency instead requires Lm BRdm sb G mb(Lr −1) TAllReduce , Pt . Ttp,decode (v) = 2 Pp Pd mb The appropriate expression depends on the implementation’s decode batching; the two must not be combined. b) Pipeline-parallel P2P: Pipeline stages exchange TPsharded activations. Let BR ℓ dm sb Mpp (ℓ) = Pd Pt mb be the per-microbatch activation message size at sequence length ℓ. The division by Pt assumes that the communicated hidden tensor is physically TP-sharded. Replicated-output pipeline layouts require the corresponding actual message size. A pipeline with Pp stages has Pp − 1 stage boundaries. Let nintra and ninter be the number of intra-node and inter-node boundaries, respectively, with nintra + ninter = Pp − 1. Under nested placement, if a node contains Ng GPUs and each TP group has size Pt , then approximately Ng /Pt pipeline stages fit per node, so Pp ninter = max 0, −1 , Ng /Pt nintra = Pp − 1 − ninter .

For inference, I Tpp (v) = (mb + Pp − 2)· nintra p2p p2p + αintra + Mpp (Lp + Lr )βintra Pp − 1 ninter p2p p2p . αtier + Mpp (Lp + Lr )βtier Pp − 1

For training under 1F1B, the number of pipeline communication steps is mb + 2Pp − 3: T Tpp (v) = (mb + 2Pp − 3)· nintra p2p p2p + αintra + Mpp (Lp + Lr )βintra Pp − 1 ninter p2p p2p . αtier + Mpp (Lp + Lr )βtier Pp − 1

When Pp = 1, we set Tpp (v) = 0. c) Data-parallel synchronization: For training operators, the simplified ZeRO-2 communication estimate comprises a ReduceScatter and an AllGather. More precisely, ReduceScatter aggregates/partitions the gradients; AllGather communicates the subsequently updated parameter partitions so that data-parallel replicas agree. These collectives can occur at different phases and need not form an adjacent back-to-back pair. Let Nm be the total number of scalar parameters of model m. The BF16 shard size per DP group is Mdp =

Nm sb . Pp Pt

Thus, Tdp (v) = TReduceScatter (Mdp , Pd ) + TAllGather (Mdp , Pd ). The collective tier is selected from the full data-parallel allocation. The common payload Nm sb /(Pp Pt ) treats parameters as uniformly PP/TP-sharded. Replicated boundary embeddings or heads and mixed-precision gradient buckets require the appropriate stage- specific byte count in a layout-exact implementation. E. Non-GEMM Overhead The work terms W (v) count matrix multiplication FLOPs. Transformer layers also execute normalization, residual additions, positional embeddings, FlashAttention score computation, and activation functions. These operations have low arithmetic intensity and are therefore assumed to be memorybandwidth-bound. At Pt = 1, their cost is assumed to be absorbed into the calibrated Fpeak , because calibration uses full transformer operations. At Pt > 1, GEMM work shrinks with tensor parallelism, whereas some non-GEMM operations, especially normalization and residual additions, do not shrink proportionally. We model the residual TP-invariant overhead with a calibrated per-layer constant cng : 1 bloc Lm cng 1 − , Tnon-gemm (v) = Pp Pt Bcal where Bcal is the per-device batch size used when calibrating cng . The expression assumes that the fitted cng has the same

aggregation scope as this operator-level correction. For training, we multiply this term by 4, corresponding to forward, recompute, and backward phases. For generation, we set Tnon-gemm (v) = 0, because decode processes one token per step and its runtime is assumed to be dominated by weight and KV-cache traffic. F. Optimizer Step For training, we approximate the bandwidth-bound Adam update using an effective traffic coefficient of 18 bytes per parameter: Nm 18 Topt (v) = · , Pd Pp Pt BWmem where Nm /(Pd Pp Pt ) is the average logical number of parameters per device. For generation and inference operators, Topt (v) = 0. G. Explicit Communication and Reshard Nodes The executable DAG also contains explicit C OMM and R ESHARD nodes. a) Comm nodes: A C OMM node is inserted when the source and destination operators (u, v) are placed on different allocations. If container a has element count |a| and element size s(a), the byte volume is |a|s(a). Since concurrent send/recv pairs are limited by the side with fewer devices, the transfer time is |a|s(a) βp2p , Txfer (u, v, a) = αp2p + min(|σ(u)|, |σ(v)|) where (αp2p , βp2p ) are selected at the communication tier spanned by σ(u) and σ(v). This is a balanced, fully parallel transfer approximation. It omits per-pair setup, serialization, packing/repartitioning, and asymmetric source/destination routes; the physical controller may also reserve both endpoint allocations while the analytical scheduler is more optimistic. b) Reshard nodes: A R ESHARD node is inserted when two consecutive invocations of the same model use the same allocation but incompatible parallel layouts. All devices in the allocation participate symmetrically in the layout transformation. We model this as an all-to-all: Nm sb , |σ(u)| , Treshard (u, v, m) = TA2A |σ(u)| where Nm sb is the full parameter size of model m in Bytes, and TA2A (M, p) is the all-to-all cost for per-GPU message size M over p GPUs. H. Parameter Accounting and Wshard We now derive Wshard (v), the per-GPU parameter footprint of operator v on its worst pipeline stage. The memory model counts dominant weight matrices: attention projections, FFN projections, embeddings, and task heads. Biases and normalization parameters are omitted.

The logical parameter counts of the policy, reference, critic, reward, and shared actor–critic are derived in [17]. Here we require only the dominant per-layer count to determine physical residency on a particular pipeline stage and GPU: layer Wm = 4d2m + 2dm dff,m .

This leading-order expression assumes standard attention and a two-projection FFN. The remaining calculation is placementspecific: pipeline stages hold different layers and boundary embeddings/heads, which is not captured by a logical permodel parameter total. Under pipeline parallelism, transformer layers are split across Pp stages, so each stage receives Lm /Pp layers (assuming Pp | Lm ). However, embeddings and task heads are pinned to boundary stages rather than spread uniformly. Let Em = Vm dm be the embedding matrix size. Define two indicators: ( 1, m ∈ {πθ , πref , πAC }, LM χm = 0, otherwise, ( 1, m ∈ {Vψ , Rϕ , πAC }, χscal m = 0, otherwise. Then the non-layer resident component on pipeline stage p ∈ {1, . . . , Pp } is nonsplit Wm (p) = 1[p = 1] Em

scal + 1[p = Pp ] χLM m Em + χm dm . The transformer-layer component on stage p is split Wm (p) =

Lm layer W . Pp m

Tensor parallelism shards transformer-layer weights across Pt ranks, but embeddings and heads are replicated on each tensorparallel rank. Hence the per-GPU resident parameter footprint on stage p is split (p) Wm nonsplit + Wm (p). Pt The worst pipeline stage determines the per-GPU parameter footprint: gpu Wm (p) =

Wshard (v) =

max

p∈{1,...,Pp }

gpu Wm (p).

Equivalently,

1 Lm (4d2m + 2dm dff,m ) Pt Pp + 1[p = 1]Vm dm

Wshard (v) = max p

+ 1[p = Pp ]

scal χLM m Vm dm + χm dm

If tied input/output embeddings are physically shared on the same pipeline stage, the duplicate embedding term can be stored once. The formula above is the conservative residentfootprint model used for boundary-stage memory pressure. From Wshard (v), the model-state components follow: Mw (v) = Wshard (v)sb

(v) for one BF16 weight copy, Mopt (v) = 12WPshard for ZeROd (v) (v)sb 2-sharded FP32 Adam state, and M∇ (v) = Wshard for Pd (v) ZeRO-2-sharded BF16 gradients.

This expression assumes ordinary multi-head attention and an evenly sharded KV dimension. With grouped-query attention, the physical per-TP-rank KV width dkv,local replaces dm /Pt , giving

I. Transient Memory

GQA Mkv (v) = 2sb

Transient memory is allocated when an operator starts and released when it finishes. We denote the transient peak of operator v by A(v). Let bloc =

BR Pd mb

be the per-device microbatch size. a) Training activations: Training retains activations for the backward pass. With activation checkpointing, only every c-th layer output is retained, and the missing activations are recomputed during backward. We use %) ( $s Lm . c = max 1, Pp The checkpoint memory stores Lm /Pp c hidden states of width dm . The recomputation buffer must hold the largest single-layer intermediate. We approximate the per-token peak by 6dm 2dff,m layer Mbuf = max 2dm + , 2dm + . Pt Pt Under a 1F1B schedule, the warmup retains activations for min(mb, Pp ) microbatches. Thus T Mact (v) = sb min(mb, Pp ) bloc (Lp + Lr )· Lm /Pp layer dm + c Mbuf . c

The displayed expression conservatively reserves both the checkpoint and recomputation-buffer contributions for up to min(mb, Pp ) live microbatches. b) Inference activations: For inference and generation, only one layer’s intermediates are live at a time. The peak is the maximum of the attention intermediate, FFN intermediate, and LM-head intermediate: The baseline uses the LM-head width as a conservative upper bound. For critic/reward-only scalar heads, a role-specific estimator can replace dm + Vm with dm + 1. 2dff,m 6dm I Mact (v) = sb bloc ℓ · max 2dm + , 2dm + , d m + Vm , Pt Pt

where ℓ = Lp + Lr for teacher-forced inference over the full prompt–response sequence and ℓ = Lp for generation prefill. c) KV cache for generation: Autoregressive generation maintains a key-value cache. At the last decode step, the cache covers a representative Lp + Lr tokens. Each layer stores one key and one value vector per token, sharded across pipeline and tensor parallelism: Lm dm B · R · (Lp + Lr ) · · . Mkv (v) = 2sb · Pd P p Pt The factor 2sb accounts for both key and value caches.

Lm BR (Lp + Lr ) dkv,local . Pd Pp

Replicated KV groups must be counted as replicated rather than divided again by Pt . Training additionally includes transient BF16 weights and sharded gradients. Overall, we have:  T Mact (v)    , r(v) = training,  +Mw (v) + M∇ (v) A(v) = M I (v), r(v) = inference,  act     I Mact (v) + Mkv (v), r(v) = generation. J. Co-Residency-Aware Device Peaks The formulas above describe memory for an isolated operator. RLM training pipelines colocate multiple models and roles on the same physical devices, so memory feasibility must account for persistent state that remains resident across the schedule and transient state that appears only during operator execution. For a model m on device d, define S(m, d) as the persistent footprint, excluding transient activations and KV cache: For a device holding neither training, shadow, nor frozen state for model m, set S(m, d) = 0. A training device stores the modeled sharded master/moment arrays; any BF16 copy not persistently retained must appear in the transient term.   12 Wshard (m, d) if d serves m as a   ,  Pd training device, S(m, d) = if d serves m as a shadow    Wshard (m, d)sb , or frozen device. Here, Wshard (m, d) denotes the per-GPU resident parameter shard of model m on device d, derived from the corresponding operator-level Wshard (v). A training device stores persistent FP32 Adam state; BF16 weights are materialized from the FP32 master weights on demand and released afterward. A shadow device stores a persistent BF16 copy for asynchronous generation or inference. A frozen device stores persistent BF16 weights for a nontrainable model such as the reference or reward model. The per-device peak memory is persistent co-resident state plus the worst transient operator mapped to that device: X peak(d) = S(m, d) + max A(v). m placed on d

v placed on d

A configuration is memory-feasible iff ∀d ∈ D :

peak(d) ≤ Meff ,

where Meff = (1 − εrt )MGPU reserves a safety margin for runtime overhead and fragmentation.

RPEG: Training Reasoning Language Models Scalably and Cost-Effectively

Analyze & explore related research in Conceptio

Related documents