FedCache — Predictive Federated RL for GenAI Model Caching (AFedPG)

A reinforcement-learning system that learns to prefetch GenAI models into GPU VRAM before requests arrive, trained across a fleet of edge nodes with AFedPG (Asynchronous Federated Policy Gradient). It is evaluated on the Alibaba GenTD26 serverless image-generation trace against reactive and learned caching baselines (LRU, IceBreaker, XGBoost, ACPM).

The core claim: a learned, predictive policy reduces cold-start latency that reactive caching structurally cannot avoid — and federation lets many nodes learn one shared policy without sharing their private request traces.

---

Results (5000-step evaluation, N=5, C=15 GB)

Algorithm	Hit %	Avg latency (ms)	VRAM %	Thrash
LRU	90.48	705.6	81.1	476
XGBoost	90.34	713.8	81.2	483
ACPM (LSTM)	89.68	739.7	80.8	621
IceBreaker	89.10	793.6	82.0	637
FedCache (proposed)	94.78	411.4	83.9	454

FedCache has the highest hit rate (+4.3 pts over LRU), the lowest latency (~42% reduction), and the lowest thrash — it prefetches models ahead of demand, cutting cold starts, while a demand-aware admission gate keeps churn below LRU. Its only cost is marginally higher VRAM occupancy. Latency here is essentially a re-expression of the hit rate (avg latency ≈ 50ms·hit% + ~5,300ms·miss%), since a cold-start miss costs ~100× a hit; the two move together.

All evaluation is deterministic and seed-independent: the trace is fixed, the environment has no stochastic dynamics, and the deployed policy acts greedily (argmax). We report a single deterministic run, not seed averages.

---

Repository layout

File	Role
extract_data.py	One-time: unpack the GenTD26 .tar.gz archives into data/.
generate_meta.py	One-time: create data/model_meta.csv with heterogeneous model sizes (LoRA / ControlNet / base / pipeline tiers).
data_loader.py	AlibabaV2026DataLoader — streams the trace in chunks; model sizes; queue window + per-model demand vector.
env.py	GenAIServerlessEnv — the simulation: VRAM, warming/load-delay, cold-start clock, state, macro actions, demand-aware admission gate, multi-objective reward.
models.py	GenAIPPOArchitecture (3-way actor-critic) + AFedPG primitives (lookahead, pseudo-gradient, apply).
train.py	The asynchronous federated (AFedPG) training loop.
evaluate.py	Runs the trained agent and all baselines; prints the multi-objective table.
icebreaker.py, xgboost_policy.py, lstm.py	Baseline caching policies (IceBreaker, XGBoost, ACPM).
config.py	All hyperparameters and constants.
experiments/	Ablation studies (VRAM sweep, Belady, generalization, scaling, learning curve, gate ablation, FedProx μ, reward, queue window) — see experiment.md.
data/README.md	Documentation for the GenTD26 dataset itself.

---

Setup & usage

Requires Python 3.x with torch, numpy, pandas, xgboost.

# 1. One-time data prep
python extract_data.py     # unpack the GenTD26 archives placed in data/
python generate_meta.py    # synthesize heterogeneous model sizes -> data/model_meta.csv

# 2. Train (1000 asynchronous AFedPG steps from scratch -> checkpoints/)
python train.py

# 3. Evaluate (FedCache vs LRU / IceBreaker / XGBoost / ACPM)
python evaluate.py

Training writes periodic checkpoints (checkpoint_step_*.pth), a final model (checkpoint_final.pth), and the best-by-hit-rate snapshot (checkpoint_best.pth). AFedPG fine-tuning can occasionally drift below the policy it had earlier, so the global model is evaluated periodically and the best snapshot is kept; evaluate.py loads checkpoint_best.pth.

Ablation studies live in experiments/ (e.g. python experiments/exp_vram_sweep.py); see experiment.md.

---

How it works

The environment (env.py)

Each step() serves one trace request. A model is in one of three states: absent, warming (load started, VRAM reserved, not yet usable), or available. A LOAD puts a model into warming; it only becomes usable after a size-dependent warming delay Δ(m) = LOAD_BASE_MS + size_gb × LOAD_MS_PER_GB (≈ 5–15 steps). This is the key constraint — you cannot load-and-serve the same request, so only predictive prefetching (loading ahead using the upcoming-request window) avoids cold starts.

• State (463-dim): VRAM usage, per-model cache state (available/warming), recency, a recency-decayed demand vector over the upcoming window, an unmet-demand vector (demand for models not yet handled), and the raw queue window.
• Macro actions (3): do-nothing, prefetch (the env loads the soonest catchable upcoming model), evict (the env removes the lowest-demand resident model). The concrete model id is resolved by the environment, so the policy learns timing, not a 100-way model selection.
• Demand-aware admission gate: a prefetch is admitted only if there is free VRAM, or the target's demand exceeds that of the eviction victim (_prefetch_is_beneficial). This is what keeps prediction low-thrash.
• Reward: +1 per cache hit (plus a small upcoming-completion bonus); a 5–15 s (size-scaled) penalty per cold-start miss; small penalties for loading, evicting, and VRAM occupancy.

The learner: federated PPO + AFedPG (train.py, models.py)

Each of 5 edge nodes runs PPO from scratch on its own slice of traffic; AFedPG aggregates their learning:

1. Local PPO — each node extrapolates its received weights (lookahead), rolls out, runs clipped PPO, and returns a pseudo-gradient (the net weight displacement), not its data or full model.
2. Asynchronous aggregation — the server applies each node's pseudo-gradient the moment it arrives (θ ← θ − η · g), no waiting for stragglers.
3. Delay-adaptive lookahead — before sampling, each node extrapolates θ̃ = θ + ((1−α)/α)(θ − θ_prev) to compensate for the staleness it will incur while computing.
4. Staleness damping — updates are scaled by 1/√(1+staleness), so stale stragglers can't destabilize the global policy.
5. Targeted broadcast — the refreshed policy goes back only to the node that just reported in.

With the macro-action reformulation, RL reaches the ~98% hit-rate ceiling from a random initialization within ~10–20 global updates — no behavior cloning or warm start is used or needed. The final global policy is the deployed agent that evaluate.py benchmarks.

---

Key configuration (config.py)

Group	Keys
System	NUM_MODELS=100, MAX_VRAM_GB=15, QUEUE_WINDOW=20
Federation	NUM_CLIENTS=5, TOTAL_GLOBAL_STEPS=1000, SERVER_LR=0.3 (≈ 1/N), STALENESS_ENABLED, CLIENT_MIN/MAX_DELAY
Federation (optional)	FEDPROX_MU=0.0 (proximal anchor; 0 = plain aggregation), NON_IID=False (shard the trace across clients)
Lookahead	LOOKAHEAD_ALPHA=0.8
PPO	ROLLOUT_STEPS=512, LOCAL_EPOCHS=4, LEARNING_RATE=3e-4, GAMMA, GAE_LAMBDA, PPO_CLIP, ENTROPY_COEF
Environment	LOAD_BASE_MS=5000, LOAD_MS_PER_GB=770, COLD_START_MS_MIN/MAX, BASE_SERVICE_MS, reward coefficients (CACHE_HIT_REWARD, VRAM_HOARD_COEF, LOAD_PENALTY, EVICT_PENALTY)

Notable knobs for ablations:

• QUEUE_WINDOW (queue-window length H): must exceed the load delay. The predictive advantage appears only once H clears Δ (≈5–15 steps); H≈10–20 is the sweet spot.
• SERVER_LR: must stay near 1/NUM_CLIENTS; full updates from N concurrent clients otherwise sum to an N× overshoot.
• VRAM_HOARD_COEF / LOAD_PENALTY: dial the agent toward leaner (fewer/cheaper prefetches) vs more aggressive caching.

---

Notes

• The system trains from scratch — caching-from-scratch is tamed by the macro-action reformulation (the target model is resolved deterministically, so the agent only learns timing), not by any imitation or warm start. AFedPG then maintains the converged policy.
• Hit rate and average latency are not independent here (latency ≈ 50 + miss% × ~5,300ms); they are two views of the same cold-start-avoidance win.
• Real-hardware validation (H100 MIG, k3s) lives in validation/ — see validation/EXPERIMENT_WRITEUP.md: the hit-rate advantage strengthens on real GPUs, the latency advantage depends on disk speed, and a demand-priority IO scheduler helps only where there is spare IO bandwidth.
• See data/README.md for full details on the GenTD26 dataset and its fields.