Working-Set-Aware Batch Size Control

• The paper presents an adaptive algorithm that uses memory footprint signals to balance throughput, stability, and convergence.
• It integrates real-time telemetry from memory, CPU, and network metrics to adjust per-node batch sizes in distributed training.
• Empirical results show improved model accuracy and training speed, with scalability across heterogeneous clusters and policy transfer.

Working-set-aware batch size control is the adaptive selection of batch size from signals that represent the memory footprint that must remain resident, together with compute, communication, and training-dynamics indicators. In distributed machine learning, this problem is explicit in DYNAMIX, which formulates batch size optimization as a sequential decision-making problem solved with Proximal Policy Optimization (PPO): the controller observes memory utilization as a working-set proxy, CPU utilization, throughput, retransmissions, accuracy-derived statistics, and iteration time, then emits per-node batch size adjustments under safety constraints to balance throughput, convergence, and stability (Dai et al., 9 Oct 2025).

1. Definition and operating context

The formal notion of a working set originates in Denning’s definition of the working set at time $t$ with window size $\Delta$ as the set of distinct pages referenced in the interval $[t-\Delta,t]$:

$$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$

In systems terms, WSS is the memory footprint that must be resident to avoid thrashing; if residency falls short of the current working set, page fault activity, reclamation pressure, and latency increase. This differs from RSS, which counts currently mapped physical pages regardless of whether they are actively used, and from total physical memory usage, which is not attributable to a single workload (Lian et al., 2023).

In distributed training, working-set-aware control treats batch size as a systems-and-optimization variable rather than a fixed hyperparameter. The “right” batch size depends on model phase, optimizer, and the hardware and network state of each node. Static or heuristic batch sizing often fails on heterogeneous clusters because nodes differ in GPU memory, CPU speed, and NIC bandwidth; contention varies over time; the optimal batch size shifts across training phases; and the memory footprint associated with activations, gradients, parameters, and optimizer state changes with micro-batching, mixed precision, and checkpointing. Communication regimes also change: when network congestion rises, large batches reduce synchronization frequency; when GPUs are underutilized, increasing batch size can improve throughput; when memory is tight, smaller batches avoid OOM and paging or thrashing (Dai et al., 9 Oct 2025).

A common misconception is that working-set awareness is equivalent to tracking raw memory occupancy. The systems literature instead distinguishes current usage from temporal locality: WSS measures the set of pages actually touched in a time window, whereas RSS can be inflated by cold pages or deflated by sharing and reclamation. This is why WSS is described as the right signal for dynamic batch size control when the objective is to sustain throughput without triggering page faults and OS reclaim (Lian et al., 2023).

2. Control formulation as a sequential decision problem

DYNAMIX formalizes batch size control as a Markov decision process in which observations are aggregated every $k$ iterations to stabilize decisions. The local state for worker $i$ at decision step $t$ concatenates network-level metrics, system-level utilization, and training statistical efficiency signals; a shared global state includes signals such as global loss trajectory and validation accuracy under BSP (Dai et al., 9 Oct 2025).

Component	Definition in DYNAMIX	Role
State $s_t^i$	$T p_t^i$, $Rtx_t^i$, CPU time-to-wall-clock ratio, memory utilization, $\Delta$0, $\Delta$1, $\Delta$2, $\Delta$3, $\Delta$4, $\Delta$5	Senses communication, resource pressure, and training dynamics
Action $\Delta$6	$\Delta$7 with $\Delta$8	Per-node discrete batch adjustment
Bounds	$\Delta$9, $[t-\Delta,t]$0	Enforces safe operating range
Reward	Accuracy terms, iteration-time penalty, optimizer-specific stability penalties	Balances convergence, throughput, and stability

The reward explicitly couples statistical progress to systems efficiency. For SGD, DYNAMIX uses

$[t-\Delta,t]$1

For adaptive optimizers such as Adam and LAMB, it adds penalties on gradient-normalization statistics:

$[t-\Delta,t]$2

The discounted objective is

$[t-\Delta,t]$3

The stated rationale is that the reward encourages accuracy and positive learning trajectory, penalizes slow steps and unstable gradient statistics, and discourages extreme batches, thereby implicitly rewarding working-set-aware decisions that keep memory and compute within safe, efficient regimes (Dai et al., 9 Oct 2025).

PPO is used in standard clipped-surrogate form:

$[t-\Delta,t]$4

with

$[t-\Delta,t]$5

DYNAMIX also states the standard value and entropy terms,

$[t-\Delta,t]$6

and the combined objective

$[t-\Delta,t]$7

Generalized Advantage Estimation is given, if used, by

$[t-\Delta,t]$8

The implementation reports simplified PPO updates in its low-variance setting (Dai et al., 9 Oct 2025).

3. Working-set modeling, memory footprint, and telemetry

In DYNAMIX, memory utilization is the core working-set signal. It functions as a proxy for when activations, gradients, parameters, and optimizer state approach device capacity. Combined with iteration time and CPU or GPU utilization, this enables the policy to back off when memory utilization is high or fluctuating, exploit headroom when utilization is stable and below thresholds, and distinguish communication-bound from compute-bound phases (Dai et al., 9 Oct 2025).

The paper makes this proxy explicit through a simple footprint model:

$[t-\Delta,t]$9

where $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$0 includes parameters, optimizer states, and persistent buffers, and $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$1 covers per-sample activations and gradients. The corresponding safe-batch constraint is

$$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$2

and, with a headroom margin $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$3,

$$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$4

This provides a direct interpretation of “working-set-aware” control: batch size is admitted only when the estimated footprint fits comfortably inside available memory (Dai et al., 9 Oct 2025).

Several standard memory-saving techniques change $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$5 and therefore change the admissible control region. Mixed precision reduces both $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$6 and $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$7; activation checkpointing reduces $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$8 at a compute cost; gradient accumulation and micro-batching decouple per-device memory from global batch; and per-buffer instrumentation can sharpen pre-checks before an increase. DYNAMIX notes that users can extend the state with per-GPU memory used and free, allocation and free headroom, per-buffer footprints, cache and NUMA counters, PCIe or NVLink utilization, disk and I/O queue depth, and NIC bandwidth (Dai et al., 9 Oct 2025).

A stricter notion of working-set telemetry comes from OS-level WSS estimation. An eBPF-based method attaches a kprobe to __handle_mm_fault, records per-pid page fault counts in BPF_HASH maps, emits samples through BPF_PERF_OUTPUT(events), and uses page fault count together with inter-sample interval $$\mathrm{WSS}(t, \Delta) = \left|\{p \in \mathcal{P} : p\text{ is referenced at least once in }[t-\Delta,\, t]\}\right|.$$9 as features for a LightGBM 3.3.2 regression model. The reported implementation uses Linux 5.13.0 and BCC v0.23.0, reports average RMSE $k$0, average update latency $k$1 s per sample window, a 65$k$2 reduction in time overhead relative to Brendan Gregg’s vanilla WSS tool, and a 98.5% overhead reduction relative to traditional VM-based methods (Lian et al., 2023). DYNAMIX itself does not require such explicit WSS estimation; it uses memory utilization as the working-set proxy. A plausible implication is that eBPF-style WSS estimators could refine the state when page-fault behavior is a dominant failure mode.

4. Distributed operation and per-node adaptation

Working-set-aware batch size control is particularly motivated by BSP-style distributed training, where per-rank imbalance produces stragglers and communication stalls. DYNAMIX uses a centralized agent that emits per-node actions from local state $k$3 plus shared global state $k$4, allowing fast nodes to increase batch size while slower or memory-constrained nodes decrease it. The stated effect is mitigation of stragglers under BSP and improved alignment between per-rank compute and communication capacity (Dai et al., 9 Oct 2025).

The communication logic is bidirectional rather than monotone. Larger batches reduce synchronization frequency, amortize all-reduce latency, and lower sensitivity to retransmissions or low throughput; however, they also increase per-synchronization payload, so further increases can hurt throughput if NIC, PCIe, or NVLink utilization saturates. DYNAMIX therefore includes throughput $k$5 and retransmissions $k$6 in the state to sense compute-bound versus communication-bound regimes and adapt batch size accordingly (Dai et al., 9 Oct 2025).

Operationally, DYNAMIX aggregates observations over $k$7 iterations before each decision, uses discrete increments $k$8 and $k$9 to smooth updates, and applies guardrails such as bounds, headroom checks, and rejection of increases after OOM or timeout events. The reported online loop is: run $i$0 iterations at current $i$1; collect $i$2, $i$3, CPU ratio, memory utilization, $i$4, $i$5, $i$6, $i$7, $i$8, and $i$9; build local and global state; compute action; enforce constraints; apply the new batch; run the next $t$0 iterations; compute reward; and update PPO. Optional guardrails include hysteresis, cooldowns, and per-node clipping based on recent OOM or timeout history (Dai et al., 9 Oct 2025).

The implementation claims low systems overhead: workers collect metrics with low overhead using eBPF for system and network telemetry and the training loop for accuracy and timing signals, and the decision-making overhead is reported as less than 0.1% of iteration time across tested settings. This is significant because explicit memory and communication models are avoided; instead, the controller learns from observed signatures of the working set (Dai et al., 9 Oct 2025).

5. Relation to batch-size theory and adjacent adaptive methods

Batch size has long been understood as a control on optimization noise. One influential result showed that learning-rate decay can often be replaced by increasing batch size while preserving training and test curves, with the noise scale for SGD approximated by

$t$1

and, with momentum $t$2,

$t$3

This yields the operational scaling rules $t$4 and $t$5, together with the recommendation to keep $t$6 and often below roughly $t$7 before reverting to learning-rate decay (Smith et al., 2017). Stagewise Enlargement of Batch Size (SEBS) provides a complementary theoretical account: under bounded gradient variance, $t$8-smoothness, $t$9-weak quasi-convexity, and the $s_t^i$0-PL condition, it links proper batch size to the gap between initialization and optimum and enlarges batch size geometrically by stage while maintaining computation complexity $s_t^i$1 and reducing iteration complexity to $s_t^i$2 (Zhao et al., 2020).

A different line of work seeks to learn or estimate the optimal batch size online. For strongly convex and smooth objectives, adaptive SGD can choose

$s_t^i$3

where $s_t^i$4 is an expected-smoothness term induced by the sampling rule and $s_t^i$5 is the gradient noise at the optimum under batch size $s_t^i$6 (Alfarra et al., 2020). For Transformers, another dynamic criterion monitors the angle between successive accumulated gradients,

$s_t^i$7

and triggers an optimizer step when the newest angle exceeds $s_t^i$8, with reported default $s_t^i$9 (Xu et al., 2020). DYNAMIX adopts yet another control surface: it does not estimate an optimal batch size from optimization statistics alone, but jointly conditions on training signals, resource utilization, and communication state (Dai et al., 9 Oct 2025). This suggests a broader formulation in which classical schedule-based enlargement and systems-aware control occupy different points on the same design spectrum.

The importance of this distinction is reinforced by domains in which smaller, not larger, batches improve learning. In replay-based deep RL, reducing batch size from 32 to 8 in QR-DQN produced significant gains on 38 of 60 ALE games, an average performance improvement of 98.25%, and a 29% wall-time speedup, indicating that optimal batch behavior is highly objective- and regime-dependent rather than universally increasing with system capacity (Obando-Ceron et al., 2023). Working-set-aware control therefore does not imply monotone enlargement; it implies conditioning batch size on the active statistical and systems regime.

6. Empirical results, deployment practice, and limitations

DYNAMIX reports up to 6.3% improvement in final model accuracy and up to 46% reduction in total training time relative to static batch baselines, with scalability to 32 nodes and policy transfer across related architectures (Dai et al., 9 Oct 2025).

Setting	Static baseline	DYNAMIX
BytePS heterogeneous cluster, 8 GPUs (RTX 3090/T4 mix)	$T p_t^i$0, $T p_t^i$1 s, 71.4%	$T p_t^i$2 s, 80%
OSC, 8 nodes, VGG16/SGD	$T p_t^i$3, 85.3% in 853 s	91.3% in 652 s
OSC, 16 nodes, VGG16/SGD	$T p_t^i$4, 83.4% in 543 s	91.5% in 479 s
OSC, 32 nodes, VGG16/SGD	$T p_t^i$5, 81.3% in 734 s	92.6% in 421 s

The paper also reports policy transfer: policies trained on VGG16 transfer to VGG19, and policies trained on ResNet34 transfer to ResNet50, maintaining benefits without retraining. The evaluated baselines are static batch sizes chosen via standard practice; heuristic controllers and linear scaling rules are discussed in related work but are not directly benchmarked. No ablations are reported that isolate individual state features or reward components (Dai et al., 9 Oct 2025).

Deployment guidance centers on instrumentation and guardrails. The recommended telemetry includes GPU memory used and free, allocator headroom, OOM counters, footprint proxies for parameters and optimizer states, activation and gradient buffers, per-sample memory estimate $T p_t^i$6, iteration time $T p_t^i$7, samples per second, batch and validation accuracy, loss slope, gradient norm and variance proxies $T p_t^i$8 and $T p_t^i$9, CPU and GPU utilization, NIC throughput, retransmissions, PCIe or NVLink bandwidth, queue backlogs, and optional cache or NUMA metrics. The paper recommends maintaining memory headroom $Rtx_t^i$0, for example 5–10% of device memory, rejecting actions that violate the footprint constraint, capping iteration time with rollback if $Rtx_t^i$1 exceeds a threshold, and aligning decisions with DDP or BSP barriers (Dai et al., 9 Oct 2025).

The stated limitations are characteristic of RL-based runtime control. Non-stationarity and delayed rewards challenge credit assignment; longer aggregation windows $Rtx_t^i$2 improve stability but slow responsiveness. The paper identifies constrained RL as a way to encode explicit OOM or time caps and communication budgets, model-based predictors of $Rtx_t^i$3 or communication time as a way to tighten guardrails and reduce exploration costs, and deeper working-set awareness via per-layer activation footprint, cache and NUMA locality, and allocator fragmentation models as directions for future work (Dai et al., 9 Oct 2025).

The same general principle appears outside training. In LLM inference, memory-aware and SLA-constrained dynamic batching monitors residual memory budget and latency, treating KV cache as the dominant dynamic component of the working set; the reported outcome is throughput gains of 8% to 28% and capacity improvements of 22% relative to static batching in vLLM-compatible settings (Pang et al., 7 Mar 2025). This broader pattern supports the view that working-set-aware batch size control is not a single algorithm but a systems-optimization interface: batch size becomes a runtime control variable governed by resident footprint, communication regime, and task-level efficiency criteria.