PACE / AdaEmbed / FIITED on ML-20M

This directory is a self-contained implementation and comparison harness for a dynamic embedding-growth experiment. It does not copy the ML-20M dataset. The default input remains:

/home/ljh/proj/emb_prune/ml-20m/processed_binary

All three methods share the same timestamp-ordered windows, DLRM-style model, element-wise Adagrad update, memory budget, high/low watermarks, evaluation split, seed, and output schema.

What is implemented

PACE (latest method)

• persistent full-dimensional embedding weights and element-wise Adagrad state in W0/W4/W8/W16 packed storage;
• lookup-time local dequantization to FP32, FP32 forward/backward/update, and immediate quantized writeback—there is no persistent FP32 master embedding;
• active-row lazy EMAs for access exposure f, per-access effective Adagrad update g, and vector magnitude h;
• cost C_i(b) = f_i g_i h_i^2 rho(b), including a small nonzero W16 error;
• Lagrangian W0/W4/W8/W16 assignment with a feasible marginal-upgrade repair;
• high-watermark trigger, low-watermark target, and newly arrived-row protection.

AdaEmbed baseline

The allocation unit is a complete embedding row. A retained row stores its FP32 weight and FP32 element-wise Adagrad state; W0 removes both. Online access/gradient utility ranks rows at each compression event.

FIITED baseline

Each row is divided into --chunks chunks. FP32 weight chunks and their matching optimizer-state chunks are independently retained or removed using online access/gradient utility.

The two baselines are state-aware adaptations for this dynamic-growth protocol, not byte-for-byte copies of the original authors' software. Their allocation granularity and online importance principles are preserved, while the trainer, memory accounting, trigger, and evaluation protocol are shared so the comparison is controlled.

Environment

The verified remote environment is Python 3.12, PyTorch 2.4.0 + CUDA 12.1, NumPy 1.26.4, scikit-learn 1.4.2, and an RTX 3090. To check the code:

cd /home/ljh/proj/pace_repro
/home/ljh/anaconda3/bin/python -m pytest -q

The launchers honor PYTHON, DATA_DIR, and OUTPUT_ROOT. run_all.sh also accepts RUN_SUFFIX so independent protocols do not share run directories.

Run one method

--materialization {legacy,granule} selects the storage layer for PACE; legacy is the checkout behaviour, granule uses the pooled allocator.

Defaults use D=128, 20 temporal windows, a 32 MiB budget, high watermark 0.92, low watermark 0.78, and seed 123.

cd /home/ljh/proj/pace_repro
./scripts/run_pace.sh
./scripts/run_adaembed.sh
./scripts/run_fiited.sh

Every script forwards extra arguments. For example:

./scripts/run_pace.sh --budget-mib 24 --dimension 128 --windows 20
./scripts/run_fiited.sh --budget-mib 24 --chunks 8

For a quick end-to-end check:

./scripts/run_all.sh \
  --dimension 32 --windows 2 --max-samples 100000 \
  --budget-mib 0.4 --high-watermark 0.8 --low-watermark 0.6 \
  --batch-size 8192 --eval-batch-size 16384 --probe-size 20000

Run and compare all methods

Use exactly the same extra arguments for all three methods:

./scripts/run_all.sh --budget-mib 32 --dimension 128 --windows 20

Or compare existing run directories:

/home/ljh/anaconda3/bin/python compare_results.py results \
  --csv results/comparison.csv --markdown results/comparison.md

The comparator requires exactly one PACE, AdaEmbed, and FIITED summary and rejects any mismatch in budget, seed, model dimension, windows, or other shared protocol arguments.

Each run produces:

results/<run-name>/
├── windows.csv       # AUC/logloss, memory, arrivals, and trigger per window
├── events.jsonl      # configuration and immediate quality at each compression
└── summary.json      # AUC, memory/peak/migration, compression events, runtime

The primary quality metric is time_weighted_auc. Also report immediate AUC change at compression, final budget compliance, payload/metadata bytes, and runtime. PACE counts packed weights, packed Adagrad state, quantization scales, and active VHPI/statistic metadata. The Python prototype's maximum-ID lookup arrays and stop-the-world migration workspace are reported as implementation artifacts rather than charged as active VHPI payload.

Measured cost of the storage layer

docs/memory_management.md describes the pooled storage layer: how payload and metadata are separated, how the buddy pool is laid out and shared, how a precision change is grouped and scheduled, and what the invariants are.

docs/measurements.md records the measured four-way comparison (AdaEmbed, FIITED, PACE with the legacy materialization, PACE with pooled storage) for training throughput, reconfiguration workspace, run-wide allocator peak and embedding lookup cost, together with the environment, the dataset provenance and how to reproduce each table. Headline numbers from the full ML-20M run:

materialization	time-weighted AUC	migration MiB	event workspace MiB	admission workspace MiB	peak allocated MiB
legacy	0.740382	412.5	360.35	360.30	420.1
granule	0.740345	120.9	13.08	0.45	239.7

Training throughput is unchanged within run-to-run variation; the pooled storage costs about 0.7 s of extra interpreter time per run because reconfiguration is per row in Python while the legacy rebuild is vectorized.

Verified full-run result

The checked-in code was run on all 20,000,263 timestamp-ordered ratings with D=128, 20 windows, a 32 MiB budget, and seed 123. The generated artifacts are under results/{pace,adaembed,fiited}_seed123/; results/comparison.md and results/comparison.csv contain the complete table.

Method	Time-weighted AUC	Final AUC	Final persistent MiB	Peak pre-compression MiB	Events
PACE (legacy storage, RTX 3090)	0.740436	0.731810	24.9599	33.5852	8
PACE (pooled storage, RTX 4060)	0.740400	0.731486	24.9600	33.5851	8
AdaEmbed	0.711960	0.678173	24.9594	36.6453	17
FIITED	0.717327	0.685277	24.9600	38.6204	17

These values are produced by the controlled implementation in this repository; they should not be presented as the original papers' reported numbers.

migration_volume_bytes is the total payload read plus write volume actually moved by the reconfiguration. peak_migration_workspace_bytes is the analytic workspace of one event: old plus new payload plus, for the packed PACE implementation, the FP32 rebuild buffer. Because that formula omits the per-precision-class temporaries raised inside the dequantizer it understates the real footprint; measured_peak_migration_workspace_bytes reports the same quantity as observed with torch.cuda.max_memory_allocated, and max_admission_migration_workspace_bytes reports the peak during admission reconfiguration, which the analytic value does not cover at all. These metrics are intentionally not conflated.

Code map

run_experiment.py             unified CLI
compare_results.py            summary table/CSV generator
common/config.py              shared experiment configuration
common/trainer.py             temporal trainer and AdaEmbed/FIITED implementations
common/results.py             result loading and comparison
pace/granule_pool.py          flat buddy allocator with a relocator hook
pace/row_store.py             row payload, hoisted scales, chain migration
pace/online_cost.py           PACE-2 lazy active-row EMAs
pace/allocator.py             cost table and Lagrangian allocator
pace/packed_tensor.py         packed row tensor primitive
pace/packed_embedding.py      heterogeneous weight/state packed table
baselines/README.md           exact baseline semantics used here
tests/                        unit and CLI regression tests
scripts/prepare_ml20m.py      raw MovieLens-20M csv to processed_binary
scripts/run_materialization_ab.sh   four-way storage comparison
scripts/summarize_perf.py     table/CSV for that comparison
scripts/compare_materialization.py  legacy against pooled A/B with drift count
scripts/bench_lookup.py       embedding lookup microbenchmark
docs/memory_management.md     the pooled storage layer design
docs/measurements.md          measured results and how to reproduce them