mrm-trace 0.1.11

Scientific instrumentation for LLM inference memory trace collection and MRM research

mrm-trace

A Python research package for collecting, parsing, labelling, and analysing LLM inference memory access traces. Designed as scientific instrumentation for Managed-Retention Memory (MRM) research - it characterises how model weights, KV cache, activations, and runtime allocations are actually accessed during inference.

Primary metrics: retention duration · write-once ratio · read frequency · working set size

---

Install from PyPI

pip install mrm-trace

Linux (or WSL2) is required for perf mem collection. See Requirements below.

---

Requirements

Requirement	Notes
Linux (WSL2 supported)	memray works everywhere; perf mem requires bare-metal or PMU-capable VM
Python ≥ 3.11	Tested on 3.11 and 3.12
sudo / root	Required for native_traces=True (memray) and perf mem

Collector capability by environment

Environment	Best collector	region_map	Timestamps	Cache level
WSL2 (non-root)	memray	empty	0	n/a
WSL2 (root)	memray --native-traces	empty †	0	n/a
Bare-metal Linux (root)	perf mem (spawn mode)	populated ‡	nanoseconds	L1/L2/L3/DRAM
Cloud VM / Colab	memray	empty	0	n/a
Debian bare-metal (root)	perf mem (spawn mode)	populated ‡	nanoseconds	L1/L2/L3/DRAM

WSL2: The Microsoft WSL2 kernel does not expose hardware PMU counters. perf mem will not produce data. Use memray for development; run on bare-metal for publication data.

perf mem — spawn mode: perf mem record -p PID (attach to running process) fails on Ubuntu 24.04 / kernel 6.8 for Intel PEBS load-latency events even as root. Use spawn mode: perf mem record -- python script.py. Both demo_qwen_hf.py and demo_perf_mem.py use spawn mode — perf launches the script itself (via --headless) so the traced process starts under perf from the first instruction.

native_traces: memray native_traces=True requires root. pip-installed llama-cpp-python strips C symbols, so region_map remains empty (†). For a populated region map via memray, build llama-cpp-python from source: CMAKE_BUILD_TYPE=Debug pip install llama-cpp-python.

region_map with PyTorch (‡): perf mem populates region_map only with llama.cpp (GGUF) models. PyTorch/transformers symbols (at::native::*) do not match the ggml_* labeller rules. For PyTorch runs, region_type = unknown for all records — retention and write-once metrics are still valid but undifferentiated by region type.

Timestamps: memray does not record per-allocation timestamps. timestamp_ns = 0 and retention_p99_s = 0 in all memray runs. Only perf mem provides real nanosecond hardware timestamps.

sudo: sudo -E python resolves to the system Python. Always use:

sudo -E venv/bin/python notebooks/scripts/demo_qwen_hf.py

---

Install

# Clone and set up a virtual environment
git clone https://github.com/DhiSys-AI/MRM-Trace
cd MRM-Trace
python -m venv venv
source venv/bin/activate      # Windows WSL: same command

# Install package + test dependencies
pip install -e ".[test]"

# Optional: install matplotlib/seaborn for figures
pip install -e ".[test,plots]"

---

Quick start

# Day-one machine check: can this box collect? (exit 0 = publication-grade perf mem)
sudo mrm-trace validate --machine

# Validate a config file
mrm-trace validate --config config/default_experiment.yaml

# Preview what a run would do (dry run)
mrm-trace plan --config config/default_experiment.yaml

# Run a full experiment (requires model files + sudo for perf)
mrm-trace run --config config/default_experiment.yaml

---

Live demo scripts

End-to-end scripts that run real inference against small models and write all mrm-trace artifacts to a timestamped results directory. Located in notebooks/scripts/.

Setup

# From the repo root (WSL2 or Linux)
source venv/bin/activate
pip install -e ".[test]"
pip install memray

TinyLlama 1.1B (llama-cpp-python + GGUF)

# Install backend
pip install llama-cpp-python \
  --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu

# Download model (~670 MB, one-time)
mkdir -p models
wget -P models/ https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf

# Run (non-root — Python-level trace only)
python notebooks/scripts/demo_tinyllama.py \
  --model models/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf

# Run as root — enables native_traces=True (C-level symbols, populated region_map)
sudo -E python notebooks/scripts/demo_tinyllama.py \
  --model models/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf \
  --native-traces

Qwen2.5-0.5B-Instruct (transformers, no GGUF needed)

# Install backend (model auto-downloads from HuggingFace, ~1 GB)
pip install transformers torch accelerate

# Auto-detects collector: memray on WSL2/non-root, perf mem on bare-metal Linux as root
python notebooks/scripts/demo_qwen_hf.py

# Force memray (any environment):
python notebooks/scripts/demo_qwen_hf.py --collector memray

# Force perf mem (bare-metal Linux, run as root with full venv path):
sudo -E venv/bin/python notebooks/scripts/demo_qwen_hf.py --collector perf-mem

# Larger variant
python notebooks/scripts/demo_qwen_hf.py --model Qwen/Qwen2.5-1.5B-Instruct

Note on region_map: The Qwen demo uses PyTorch/transformers. region_map will be empty and all records will have region_type = unknown regardless of collector. Retention and write-once metrics are still measured correctly but cannot be separated by region type. For classified regions, use demo_perf_mem.py with a GGUF model.

perf mem + TinyLlama GGUF (bare-metal Linux, classified regions)

# Requires root and hardware PMU — does NOT work on WSL2
# Use full venv path: sudo -E python resolves to system Python, not venv
sudo -E venv/bin/python notebooks/scripts/demo_perf_mem.py \
  --model models/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf

Script options

Flag	Default	Description
--model PATH	varies	GGUF path or HuggingFace model ID
--ctx N	2048	Context length
--tokens N	128/256	Max output tokens per prompt
--batch-size N	1	Number of prompts in the batch
--rep N	0	Repetition index (0 = warmup; 1+ = data runs)
--out DIR	results/	Output base directory
--collector auto|memray|perf-mem	auto	Override collector selection (Qwen demo only)
--native-traces	auto (root check)	Force native_traces=True for memray
--no-native-traces	—	Force native_traces=False

All scripts run 5 real prompts and write trace.parquet, region_map.parquet, kv_block_lifecycle.parquet, metrics.csv, metadata.json, and manifest.json to results/<model_id>/<run_id>/.

---

Notebooks

Notebook	Description	Run on Colab
001 - Getting Started	Install, synthetic trace, label, analyse, export, schema versioning, validity
002 - YAML Config & Experiment Planning	Write & validate configs, sweep expansion, multi-model runs, collector tuning
003 - Real Collection Walkthrough	Real memray capture, parse raw trace, understand symbols, real-model guide

All three notebooks run without root or model files (001 and 002 use synthetic data; 003 uses memray on a simulated workload). They are a good first stop for new contributors and researchers.

---

Running tests

# Every commit - fast, no I/O
pytest -m unit

# Pre-merge - includes integration tests
pytest -m "unit or integration"

# Before dataset release - scientific correctness checks
pytest -m validity

# Property-based invariant tests (Hypothesis)
pytest tests/property/

# Performance benchmarks (excluded from default run)
pytest -m performance --benchmark-only

# Full suite (excludes slow + performance)
pytest

The test suite has three tiers:

Tier	Marker	Purpose
1	unit	Individual functions behave correctly
2	integration	Components work together
3	validity	Measurements are scientifically sound

Tier-3 validity tests are the most important: they verify that known synthetic inputs produce known metric outputs (e.g. a 30s weight retention window must yield retention_p99_s ≈ 30.0).

---

Output layout

Each run writes to results/<model_id>/<run_id>/:

results/llama-7b/run_20240101_120000/
├── trace.parquet                  ← labelled memory access trace
├── region_map.parquet             ← one row per region (weight, kv_cache, …)
├── kv_block_lifecycle.parquet     ← per-block write / read / eviction timestamps
├── metrics.csv                    ← per-region-type summary (human-readable)
├── metadata.json                  ← hardware, software, observer effect, run validity
├── manifest.json                  ← SHA-256 checksums for all files
└── raw/
    ├── perf.data
    ├── perf_script.txt
    └── memray.bin

---

Run validity classification

Every run is automatically classified based on observer overhead:

Class	Criteria
clean	observer CPU < 10 %, observer mem < 5 % of target RSS, no throttle, baseline CPU < 15 %
marginal	observer CPU < 20 %, observer mem < 15 % of target RSS, ≤ 2 throttle events
contaminated	anything worse than marginal

Contaminated runs are archived but excluded from aggregated metrics and paper figures.

---

Architecture

mrm_trace/
├── cli.py              CLI (typer)
├── api.py              Python API (Experiment class)
├── schema_version.py   Schema version registry and compatibility checking
├── engines/            llama.cpp / vLLM wrappers
├── collector/          perf mem / memray / process_monitor
├── parser/             perf script + memray parsers → trace.parquet
├── labeller/           symbol + address-range region classification
├── analyser/           retention / write-once / read-freq / working-set / IAI / suitability
├── telemetry/          baseline capture / thermal / observer effect / validity classifier
├── reporter/           CSV + Parquet export / figures / manifest / RunExporter
└── utils/              logging / IDs / file helpers

Key design decisions:

• Streaming parser - generators throughout; never loads full trace into RAM (ADR-2)
• Phase-aware tracing - weight_load / generation / teardown phases distinguish weight from KV (ADR-3)
• Observer effect as mandatory output - every run records overhead and validity class (ADR-4)
• Parquet + zstd - column-oriented, ~3× better compression than gzip (ADR-8)

---

MRM suitability labels

Label	Criteria
high_mrm	write-once ratio ≥ 0.8 and retention p99 ≥ 10 s
medium_mrm	write-once ratio ≥ 0.5 and retention p50 ≥ 1 s
low_mrm	everything else

In practice: model weights → high_mrm, short-lived KV blocks → low_mrm.

---

Schema versioning

All output files carry a mrm_trace_schema_version in their Parquet metadata. The version registry is in mrm_trace/schema_version.py. Readers validate major-version compatibility on load; a major bump is a breaking change.

from mrm_trace.schema_version import check_parquet_schema
check_parquet_schema("results/.../trace.parquet", "trace")  # raises on incompatibility

---

Python API

from mrm_trace.labeller import TraceLabeller
from mrm_trace.analyser import compute_all
from mrm_trace.reporter import RunExporter

# Label a stream of raw trace rows
labeller = TraceLabeller()
labelled = list(labeller.label(raw_rows))
region_map   = labeller.region_map()    # call after consuming label()
kv_lifecycle = labeller.kv_lifecycle()

# Analyse
import pandas as pd
trace = pd.DataFrame(labelled)
results = compute_all(trace)
# results keys: retention_per_region, retention_summary, write_once,
#               read_freq, working_set_per_region, working_set_summary,
#               locality_per_region, locality_summary, iai, suitability

# Export a publication-ready run directory
exporter = RunExporter("results/llama-7b/run_001")
exporter.export(trace, region_map, kv_lifecycle, results,
                metadata={"run_id": "run_001"}, run_id="run_001")

---

Collector hierarchy

1. perf mem — primary; requires Linux PMU + root; bare-metal or PMU-capable VM only; does not work on WSL2
2. memray — fallback; Python-level allocations (no root) or C-level (root + native_traces=True); works everywhere
3. process_monitor — always runs in parallel as coarse RSS/CPU baseline (psutil)

See Collector capability by environment for a full comparison.

---

Reporting issues and contact

• Bug reports / feature requests: GitHub Issues
• Email: [e-mail]