Add warmup peak memory estimation via --max-num-batched-tokens

[vrdn-23]

Summary

• Adds an opt-in --max-num-batched-tokens flag that estimates the peak activation memory during a single warmup forward pass for Safetensors models — analogous to what vLLM measures in profile_run and TEI measures during its warmup.
• Reports warmup_peak as a new memory component alongside the existing weights and KV cache, and includes it in total_memory.
• Supports causal LMs, conditional generation (VLMs), masked LMs, sequence/token classifiers, base encoders, and Sentence Transformers. Independent of --experimental — works on its own.

Why

hf-mem already reports weights and (under --experimental) KV cache. A user trying to decide whether a model fits on a given GPU still has to guess the activation overhead, which can be a large fraction of total reserved memory — particularly for models with large vocab (LM head logits at the end of the forward pass) or wide FFNs. Inference engines (vLLM, TEI) measure this directly at startup; this PR produces a closed-form approximation from config.json metadata alone, no GPU required.

What it computes

For each forward-pass phase, the formula models PyTorch's caching allocator (a persistent residual stream + the largest transient at any one point):

• residual_stream = T × hidden_size × d
• attn_transient = T × (num_q_heads + 2·num_kv_heads) × head_dim × d
• ffn_transient = 2 × T × intermediate_size × d (× num_experts_per_tok for MoE)
• lm_head_spike: model-class-dependent
  • causal LM / VLM: S × vocab_size × d
  • masked LM: T × vocab_size × d (LM head applied to every token, not just the last)
  • seq classifier: S × num_labels × d
  • token classifier: T × num_labels × d
  • base encoder / Sentence Transformers: S × hidden_size × d
• peak = residual_stream + max(attn_transient, ffn_transient, lm_head_spike)

T = max_num_batched_tokens, S = min(T, batch_size), d = activation dtype bytes. The activation dtype is resolved from config.json (torch_dtype / dtype, falling back to F32 for older configs like BERT that don't set it; bf16 for quantized models since vLLM/TEI dequantize for compute).

Formula was derived by reading the actual vLLM V1 worker (vllm/v1/worker/gpu_worker.py:388-416 for the profile_run accounting, gpu_model_runner.py:5408-5431 for the token splitting that gives S = min(T, max_num_seqs)), TEI's warmup pass, and llama.cpp (src/llama-context.cpp:551-612) to confirm GGUF doesn't have a closed-form equivalent.

Sample output

~/dev/os-help/hf-mem vrdn-23/add-warmup-peak-memory = ?3 ❯ uv run hf-mem --model-id HuggingFaceTB/SmolLM2-135M --experimental --max-num-batched-tokens 8192                                                                                                                                                                                                                                                                                                                                                                                                                      10:21:12
/Users/vidamoda/dev/os-help/hf-mem/src/hf_mem/cli.py:135: UserWarning: `--experimental` is set, which means that models with an architecture as `...ForCausalLM` and `...ForConditionalGeneration` will include estimations for the KV Cache as well. You can also provide the args `--max-model-len` and `--batch-size` as part of the estimation. Note that enabling `--experimental` means that the output will be different both when displayed and when dumped as JSON with `--json-output`, so bear that in mind. Warmup peak activation memory is controlled separately via `--max-num-batched-tokens`.
  warnings.warn(
/Users/vidamoda/dev/os-help/hf-mem/src/hf_mem/safetensors/print.py:133: UserWarning: Given that the provided `--model-id HuggingFaceTB/SmolLM2-135M` (with `--revision main`) is longer than 64 characters, the table width will be expanded to fit the provided values within their row, but it might lead to unexpected table views.
  warnings.warn(
┌┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬┬ hf-mem v0.5.5 ┐
├┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┴┤
│                        INFERENCE MEMORY ESTIMATE FOR                        │
│               https://hf.co/HuggingFaceTB/SmolLM2-135M @ main               │
│                     w/ max-model-len=8192, batch-size=1                     │
│                w/ max-num-batched-tokens=8192, batch-size=1                 │
├────────────────┬────────────────────────────────────────────────────────────┤
│ TOTAL MEMORY   │ 0.48 GiB (134.52M PARAMS + KV CACHE + WARMUP PEAK)         │
│ REQUIREMENTS   │ ██████████████████████████████████████████████████████████ │
├────────────────┴────────────────────────────────────────────────────────────┤
│                      MODEL (134.52M PARAMS, 0.25 GiB)                       │
├────────────────┬────────────────────────────────────────────────────────────┤
│ BF16           │ 0.25 / 0.48 GiB                                            │
│ 134.52M PARAMS │ ██████████████████████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░ │
├────────────────┴────────────────────────────────────────────────────────────┤
│                      KV CACHE (8192 TOKENS, 0.18 GiB)                       │
├────────────────┬────────────────────────────────────────────────────────────┤
│ BF16           │ 0.18 / 0.48 GiB                                            │
│ 8192 TOKENS    │ █████████████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ │
├────────────────┴────────────────────────────────────────────────────────────┤
│                     WARMUP PEAK (8192 TOKENS, 0.06 GiB)                     │
├────────────────┬────────────────────────────────────────────────────────────┤
│ BF16           │ 0.06 / 0.48 GiB                                            │
│ 8192 TOKENS    │ ███████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ │
└────────────────┴────────────────────────────────────────────────────────────┘
~/d/o/hf-mem vrdn-23/add-warmup-peak-memory = ?3 ❯ 

JSON shape: {"memory": ..., "kv_cache": ..., "warmup_peak": ..., "total_memory": ...} (with warmup_peak omitted when the flag isn't set, so default behavior is unchanged).

Out of scope

• GGUF: llama.cpp's compute buffer size is determined at runtime by ggml_backend_sched_get_buffer_size after building the actual ggml graph; there is no closed-form formula derivable from GGUF metadata. The flag is silently a no-op for GGUF files (emits a clear warning).
• non_torch_increase: NCCL buffers, FlashAttention kernel workspaces, CUDA staging — only knowable at runtime. The estimate models torch_peak_increase only; real-world peak is typically a bit higher, especially in multi-GPU deployments.
• Tensor parallelism: out of scope (would divide most terms by world size).
• Per-component breakdown (residual / attn / ffn / lm_head) in --details JSON — deferred; the internal WarmupPeak dataclass is designed so adding it later doesn't break the public API.

Implementation notes

• New module: src/hf_mem/safetensors/warmup_peak.py (pure functions; no I/O).
• New dataclass: WarmupPeak in _types.py (mirrors KvCache style).
• Result gains warmup_peak: int | None (parallel to kv_cache) and warmup_peak_metadata: WarmupPeak | None (parallel to kv_cache_metadata).
• run.py's experimental block was refactored to fetch config.json once whether one or both estimates are requested. The ForConditionalGeneration text_config swap was lifted out of the experimental-only branch so warmup peak benefits from it too.
• --batch-size is reused as max_num_seqs (the same physical concurrency that drives both KV cache sizing in vLLM and the LM head spike during warmup); no new flag for that.

Test plan

• Causal LM (HuggingFaceTB/SmolLM2-135M): --max-num-batched-tokens 8192 produces WARMUP PEAK block + positive warmup_peak in JSON.
• Combined --experimental + --max-num-batched-tokens renders both KV CACHE and WARMUP PEAK blocks; total_memory = weights + kv_cache + warmup_peak.
• VLM (Qwen/Qwen2-VL-2B-Instruct) with --experimental correctly reports both KV cache and warmup peak (text_config swap correctly preserved).
• Masked LM (bert-base-uncased): T × vocab_size term dominates (~978 MiB at T=8192 in F32 fallback).
• Sentence Transformers (sentence-transformers/all-MiniLM-L6-v2) and cross-encoder (cross-encoder/ms-marco-MiniLM-L-6-v2) detected as encoders; FFN transient dominates.
• MoE (Qwen/Qwen2.5-7B-Instruct, etc.) — num_experts_per_tok correctly scales the FFN transient.
• GGUF (TheBloke/deepseek-llm-7B-chat-GGUF) emits a no-op warning and skips warmup peak; existing GGUF behavior unchanged.
• Negative --max-num-batched-tokens raises a clear RuntimeError.
• No flag passed → no warmup_peak field in output (default behavior unchanged for existing users).
• (Optional, GPU-only) Comparison against an actual vLLM profile_run log on a reference model — left for follow-up; not blocking.

---

Draft for early review. Will mark ready once the test plan's optional GPU comparison is done (or omitted, depending on feedback).

🤖 Generated with Claude Code

[vrdn-23]

Still need to clean this up a bit more, but this is the general idea of the PR

[alvarobartt]

Thanks @vrdn-23, super useful! LGTM just left some nits, but I'll need to still review it in detail (might as well add Copilot for a quick review), thanks again 🤗

[Copilot]

Pull request overview

This PR adds an opt-in warmup forward-pass peak activation memory estimate for Safetensors models via --max-num-batched-tokens, reports it as a new warmup_peak component, and includes it in total_memory (independent of the existing --experimental KV-cache estimate).

Changes:

• Introduces a new warmup-peak estimator (compute_safetensors_warmup_peak) based on config.json metadata only.
• Extends the result/reporting pipeline (CLI args, JSON output, table rendering) to include warmup_peak and fold it into total_memory.
• Refactors run.py config fetching so the same config.json load supports KV cache and/or warmup peak estimation (including text_config swapping for ForConditionalGeneration).

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/hf_mem/safetensors/warmup_peak.py	Implements architecture classification + closed-form warmup peak activation estimate.
src/hf_mem/run.py	Adds max_num_batched_tokens plumbing, computes/stores warmup_peak, updates total_memory.
src/hf_mem/safetensors/print.py	Prints the WARMUP PEAK block and includes it in the combined totals.
src/hf_mem/cli.py	Adds --max-num-batched-tokens CLI flag + warnings and passes it through to arun.
src/hf_mem/_types.py	Adds WarmupPeak dataclass and threads it through Result.
README.md	Documents the new warmup peak feature and usage.
skills/hf-mem/SKILL.md	Updates skill docs to mention activation/warmup memory support.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[vrdn-23]

@alvarobartt Addressed the feedback and the Copilot comments seemed worth addressing. Let me know if there's anything else I can take care of.

[vrdn-23]

One thing I would like input on is whether you think we should have a separate input called --max-num-seqs which is different from --batch-size?

[vrdn-23]

@alvarobartt just wanted to bump this in case you had time to give it a review

[vrdn-23]

@alvarobartt Just wanted to bump this and see if there was anything more you want me to test