---
title: Working-Set-Aware Batch Size Control
url: https://www.emergentmind.com/topics/working-set-aware-batch-size-control
type: topic
---

# Working-Set-Aware Batch Size Control

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 [2510.08522].

## 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 [2303.05919].

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 [2510.08522].

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 [2303.05919].

## 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 [2510.08522].

| Component | Definition in DYNAMIX | Role |
|---|---|---|
| State $s_t^i$ | $T p_t^i$, $Rtx_t^i$, CPU time-to-wall-clock ratio, memory utilization, $\bar{A}_t^i$, $\sigma_{\text{batch},t}^i$, $\Delta A_t^i$, $T_{\text{iter},t}^i$, $\sigma_{\text{norm},t}^i$, $\sigma_{\text{norm},t}^{2,i}$ | Senses communication, resource pressure, and training dynamics |
| Action $a_t^i$ | $a_t^i \in \{-100,-25,0,+25,+100\}$ with $B_{t+1}^i=\operatorname{clip}(B_t^i+a_t^i,B_{\min},B_{\max})$ | Per-node discrete batch adjustment |
| Bounds | $B_{\min}=32$, $B_{\max}=1024$ | 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
$$
r_t^{\mathrm{SGD}} = \bar{A}_t + \alpha \cdot \max\{0,\Delta A_t\} - \beta\, T_{\mathrm{iter},t} - \delta\left(\log_2(B_t)-5\right).
$$
For adaptive optimizers such as Adam and LAMB, it adds penalties on gradient-normalization statistics:
$$
r_t^{\mathrm{opt}} = \bar{A}_t + \alpha \cdot \max\{0,\Delta A_t\} - \beta\, T_{\mathrm{iter},t} - \eta\left(\sigma_{\mathrm{norm},t}^2 + \sigma_{\mathrm{norm},t}\right) - \delta\left(\log_2(B_t)-5\right).
$$
The discounted objective is
$$
J(\pi) = \mathbb{E}_{\pi}\Big[\sum_{t=0}^{\infty} \gamma^t r_t\Big].
$$
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 [2510.08522].

PPO is used in standard clipped-surrogate form:
$$
L^{\mathrm{CLIP}}(\theta) = \mathbb{E}_t\Big[\min\big(r_t(\theta) A_t,\; \operatorname{clip}(r_t(\theta),\,1-\epsilon,\,1+\epsilon) A_t\big)\Big],
$$
with
$$
r_t(\theta) = \frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)}.
$$
DYNAMIX also states the standard value and entropy terms,
$$
L^{\mathrm{VF}}(\theta) = \big(V_\theta(s_t)-V_t^{\mathrm{target}}\big)^2,\qquad
S[\pi_\theta](s_t) = -\sum_a \pi_\theta(a\mid s_t)\log\pi_\theta(a\mid s_t),
$$
and the combined objective
$$
J(\theta) = \mathbb{E}[L^{\mathrm{CLIP}}] - c_1\,\mathbb{E}[L^{\mathrm{VF}}] + c_2\,\mathbb{E}[S].
$$
Generalized Advantage Estimation is given, if used, by
$$
A_t = \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l},\qquad
\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t).
$$
The implementation reports simplified PPO updates in its low-variance setting [2510.08522].

## 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 [2510.08522].

The paper makes this proxy explicit through a simple footprint model:
$$
W(B) \approx W_{\mathrm{fixed}} + B\,w_{\mathrm{sample}},
$$
where $W_{\mathrm{fixed}}$ includes parameters, optimizer states, and persistent buffers, and $w_{\mathrm{sample}}$ covers per-sample activations and gradients. The corresponding safe-batch constraint is
$$
B_t \le \left\lfloor \frac{\mathrm{Mem}_{\mathrm{free}} - W_{\mathrm{fixed}}}{w_{\mathrm{sample}}} \right\rfloor,
$$
and, with a headroom margin $h$,
$$
B_t \le \left\lfloor \frac{(\mathrm{Mem}_{\mathrm{free}} - h) - W_{\mathrm{fixed}}}{w_{\mathrm{sample}}} \right\rfloor.
$$
This provides a direct interpretation of “working-set-aware” control: batch size is admitted only when the estimated footprint fits comfortably inside available memory [2510.08522].

Several standard memory-saving techniques change $W(B)$ and therefore change the admissible control region. Mixed precision reduces both $w_{\mathrm{sample}}$ and $W_{\mathrm{fixed}}$; activation checkpointing reduces $w_{\mathrm{sample}}$ 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 [2510.08522].

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 $\Delta t$ 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 $0.0744$, average update latency $0.0581$ s per sample window, a 65$\times$ reduction in time overhead relative to Brendan Gregg’s vanilla WSS tool, and a 98.5% overhead reduction relative to traditional VM-based methods [2303.05919]. 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 $s_t^i$ plus shared global state $s_t^{\mathrm{global}}$, 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 [2510.08522].

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 $T p$ and retransmissions $Rtx$ in the state to sense compute-bound versus communication-bound regimes and adapt batch size accordingly [2510.08522].

Operationally, DYNAMIX aggregates observations over $k$ iterations before each decision, uses discrete increments $\pm 25$ and $\pm 100$ 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 $k$ iterations at current $B_i$; collect $T p$, $Rtx$, CPU ratio, memory utilization, $\bar{A}$, $\sigma_{\text{batch}}$, $\Delta A$, $T_{\text{iter}}$, $\sigma_{\text{norm}}$, and $\sigma_{\text{norm}}^2$; build local and global state; compute action; enforce constraints; apply the new batch; run the next $k$ iterations; compute reward; and update PPO. Optional guardrails include hysteresis, cooldowns, and per-node clipping based on recent OOM or timeout history [2510.08522].

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 [2510.08522].

## 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
$$
g = \epsilon \left(\frac{N}{B}-1\right) \approx \frac{\epsilon N}{B},
$$
and, with momentum $m$,
$$
g \approx \frac{\epsilon N}{B(1-m)}.
$$
This yields the operational scaling rules $B \propto \epsilon$ and $B \propto 1/(1-m)$, together with the recommendation to keep $B \ll N$ and often below roughly $N/10$ before reverting to learning-rate decay [1711.00489]. Stagewise Enlargement of Batch Size (SEBS) provides a complementary theoretical account: under bounded gradient variance, $L$-smoothness, $\alpha$-weak quasi-convexity, and the $\mu$-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 $\mathcal{O}\!\left(\sigma^2/(\alpha^2\mu\epsilon)\right)$ and reducing iteration complexity to $\mathcal{O}\!\left((L/(\alpha^2\mu))\log(1/\epsilon)\right)$ [2002.11601].

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
$$
\tau^* \in \arg\min_{\tau\in\{1,\dots,n\}} \max\Big\{\tau\,\mathcal{L}(\tau),\;\frac{2}{\epsilon\mu}\tau\,\sigma(x^*,\tau)\Big\},
$$
where $\mathcal{L}(\tau)$ is an expected-smoothness term induced by the sampling rule and $\sigma(x^*,\tau)$ is the gradient noise at the optimum under batch size $\tau$ [2005.01097]. For Transformers, another dynamic criterion monitors the angle between successive accumulated gradients,
$$
a(g_0^{k-1}, g_0^k)=\arccos\Big(\frac{g_0^{k-1}\cdot g_0^k}{\|g_0^{k-1}\|\|g_0^k\|}\Big),
$$
and triggers an optimizer step when the newest angle exceeds $\alpha a_{\min}$, with reported default $\alpha=1.1$ [2005.02008]. 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 [2510.08522]. 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 [2310.03882]. 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 [2510.08522].

| Setting | Static baseline | DYNAMIX |
|---|---|---|
| BytePS heterogeneous cluster, 8 GPUs (RTX 3090/T4 mix) | $B=64$, $\sim 20{,}000$ s, 71.4% | $\sim 16{,}000$ s, 80% |
| OSC, 8 nodes, VGG16/SGD | $B=128$, 85.3% in 853 s | 91.3% in 652 s |
| OSC, 16 nodes, VGG16/SGD | $B=128$, 83.4% in 543 s | 91.5% in 479 s |
| OSC, 32 nodes, VGG16/SGD | $B=64$, 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 [2510.08522].

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 $w_{\mathrm{sample}}$, iteration time $T_{\mathrm{iter}}$, samples per second, batch and validation accuracy, loss slope, gradient norm and variance proxies $\sigma_{\mathrm{norm}}$ and $\sigma_{\mathrm{norm}}^2$, 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 $h$, for example 5–10% of device memory, rejecting actions that violate the footprint constraint, capping iteration time with rollback if $T_{\mathrm{iter}}$ exceeds a threshold, and aligning decisions with DDP or BSP barriers [2510.08522].

The stated limitations are characteristic of RL-based runtime control. Non-stationarity and delayed rewards challenge credit assignment; longer aggregation windows $k$ 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 $W(B)$ 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 [2510.08522].

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 [2503.05248]. 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.

Source: https://www.emergentmind.com/topics/working-set-aware-batch-size-control