Skip to content

Memory Offload and Multi GPU

Jump to bottom
Qubitium edited this page Oct 28, 2026 · 4 revisions

Memory, Offload, and Multi-GPU

Documentation home · Quickstart · Lifecycle

On this page

Overview

GPT-QModel has several independent memory and placement controls: quantization calculation pools, calibration-cache placement, optional forward data parallelism, lazy checkpoint materialization, and completed-module disk offload. These controls solve different problems. Merely exposing more GPUs does not enable parallel quantization, and inference device_map="auto" is not the quantization pool scheduler.

This page describes commit c1976918cfe3703fbcf6d1ce2440a82b26394375. Device availability, model-definition strategy support, quantization method, Python threading mode, host RAM, and disk bandwidth affect the actual execution path.

Workflow at a glance

Separate memory and device controls. Memory controls govern separate parts of the quantization workload.

Reading the diagram: The lanes are independent concerns, not consecutive stages. Cache storage is calibration-only; lazy loading and disk offload are conditional. Pools do not enable tensor parallelism.

Text version of the workflow
  • Source weights: lazy path: Index checkpoint shards → Read requested tensors → Materialize active work.
  • Calibration cache: Capture input batches → Choose cache storage → Replay cached batches.
  • Compute placement: Find eligible devices → Resolve device pools → Place module groups.
  • Completed-module offload: Finalize module → Write eligible payload → Release tensor storage.

Open full-size SVG

Concepts: what occupies memory

Resource Where it comes from Relevant controls
Original model weights Loading and materializing the source checkpoint LazyTurtle, meta shell loading, device placement.
Active quantization work Module weights, statistics, algorithm workspaces, and intermediate tensors Dense/MoE device pools, strategy, method-specific workspace configuration.
Calibration activations Cached captured inputs and forwarded data calibration_data_device, calibration shape/count, asynchronous activation transfers.
Forward replicas Optional module clones for data-parallel forwarding auto_forward_data_parallel.
Completed quantized modules Packed weights and metadata pending final save offload_to_disk and its storage path.
Native runtime workspace Kernel repacking, scratch, execution plans, and current-call decode Selected inference kernel; see Kernel Backends Reference.
Allocator-reserved memory Backend memory allocator state, separate from live tensor bytes Cache cleanup after safe synchronization; not equivalent to freeing referenced tensors.

Disk offload can lower resident model memory but does not eliminate peak working memory. A module being quantized or serialized, its activations, and its current workspaces can coexist. Similarly, balancing work across devices does not promise equal byte usage or shard an individual matrix calculation across GPUs.

Configuration reference

The following fields are declared on BaseQuantizeConfig in gptqmodel/quantization/config.py and exposed through QuantizeConfig and its method-specific configurations. Defaults below are declaration defaults; temporary path creation and model/method restrictions can modify the effective runtime state.

Field Type Default Meaning
offload_to_disk bool True Offload completed module memory during the quantization loop.
offload_to_disk_path str None Directory for offloaded state; when offload is enabled without a supplied path, configuration creates a managed temporary directory.
compute_device_filter Optional[callable] None Callback fn(devices: List) -> List filtering compute-intensive quantization/forward devices.
calibration_data_device Optional[Union[str, torch.device]] None Captured-calibration placement; accepts an explicit device or the special string "balanced".
auto_forward_data_parallel bool True Allows automatic module replication for multi-device forwarding where the stage permits. Replication can increase VRAM and preparation time.
dense_vram_strategy VramStrategy VramStrategy.EXCLUSIVE Placement strategy for the dense/serial module pool.
dense_vram_strategy_devices Optional[List[Union[str, torch.device]]] None Explicit dense pool, resolved as a subset of visible/eligible compute devices.
moe_vram_strategy VramStrategy VramStrategy.EXCLUSIVE Placement strategy for expert families.
moe_vram_strategy_devices Optional[List[Union[str, torch.device]]] None Explicit expert pool; gate/up/down of one expert stay together.
pack_impl str "auto" Packing implementation preference. The lower-level hf_pack_layer API separately treats a missing pack_impl as original packing.

VRAM strategy enum and placement units

class VramStrategy(str, Enum):
    EXCLUSIVE = "exclusive"
    BALANCED = "balanced"
Strategy Dense pool MoE pool
EXCLUSIVE Keeps the relevant work on its resolved exclusive placement rather than distributing calculation groups round-robin. Keeps expert work on the resolved exclusive placement rather than balancing expert families.
BALANCED Assigns model-tree calculation groups across the dense pool; related q/k/v projections remain co-located. Assigns whole expert families across the expert pool; one expert's gate/up/down projections remain co-located.

The dense pool owns qkv, z, output projections, norms, router, shared expert, and dense MLP modules. It does not mean "all modules of a model that happens to be dense." A MoE model still has dense-path modules; shared experts and routed expert families have distinct placement roles.

ModuleLooper normalizes the requested strategies and checks the model's supported dense/MoE strategies. Unsupported strategies fall back to EXCLUSIVE. In stage_subset.py, balanced placement uses calculation groups/families as units, not independent round-robin assignment of every leaf linear.

Device pool resolution

The looper starts from devices selected for the configured quantization device/family, applies compute_device_filter when supplied, then applies method restrictions and resolves the dense/expert subset lists. ParoQuant's method restriction forces single-device quantization to avoid multi-GPU synchronization issues; supplying several pool devices cannot override that method restriction.

Public configuration normalizes dense_vram_strategy_devices and moe_vram_strategy_devices with _normalize_strategy_devices(): None keeps automatic selection, but an explicitly empty list raises ValueError. Use None, not [], to request the default pool.

The lower-level _resolve_strategy_device_pool() then implements these rules:

  1. None returns the available compute pool. The helper also accepts an empty list internally, but public configuration rejects it before this point.
  2. Each explicit device is normalized and must match an eligible available device.
  3. Duplicate entries are removed while preserving their first occurrence.
  4. Unsupported values or entries outside the available pool raise ValueError.

An empty result from compute_device_filter is different from a pool validation error: the looper warns and retains the unfiltered device list. Do not use an empty callback result to request "no compute devices." Use explicit device names and return the devices intended for the run.

Usage examples

Balanced MoE quantization

Select devices before Python starts. Device names inside the process use logical ordinals after visibility filtering.

CUDA_VISIBLE_DEVICES=0,1 PYTHON_GIL=0 python quantize_moe.py
from gptqmodel import GPTQModel, QuantizeConfig
from gptqmodel.quantization.config import VramStrategy

quant_config = QuantizeConfig(
    bits=4,
    group_size=128,
    device="cuda:0",
    moe_vram_strategy=VramStrategy.BALANCED,
    moe_vram_strategy_devices=["cuda:0", "cuda:1"],
)

model = GPTQModel.load("MODEL_REPO_OR_LOCAL_PATH", quant_config, dtype="bfloat16")
# Supply calibration_dataset using the normal quantize() dataset contract.
model.quantize(calibration_dataset)
model.save("My-MoE-GPTQ")

This follows docs/multi_gpu_quantization.md: the dense strategy remains at its default EXCLUSIVE setting while routed expert families are balanced. For parallel expert quantization, the Python build must support free threading and its GIL must actually be disabled. PYTHON_GIL=0 cannot make a conventional Python build free-threaded. Startup logging reports dense/MoE pools, routing mode, and GIL/parallel eligibility; visibility alone is not evidence of concurrent quantization.

Balanced dense groups and explicit offload storage

from gptqmodel import QuantizeConfig
from gptqmodel.quantization.config import VramStrategy

quant_config = QuantizeConfig(
    bits=4,
    group_size=128,
    device="cuda:0",
    dense_vram_strategy=VramStrategy.BALANCED,
    dense_vram_strategy_devices=["cuda:0", "cuda:1"],
    offload_to_disk=True,
    offload_to_disk_path="/fast-scratch/gptq-offload",
    calibration_data_device="cpu",
    auto_forward_data_parallel=False,
)

Use a dedicated writable scratch directory with sufficient free space. Disabling automatic forward replication can reduce clone-induced VRAM pressure but is not a general speedup guarantee. Balanced placement still preserves calculation-group co-location.

Filter devices explicitly

from gptqmodel import QuantizeConfig

quant_config = QuantizeConfig(
    device="cuda:0",
    compute_device_filter=lambda devices: [d for d in devices if d.index != 0],
)

The callback receives the discovered compute-device list. If only device zero is available, the example returns an empty list and the documented warning/fallback applies. Explicit dense/MoE lists must refer to the resulting eligible pool, not physical IDs hidden by CUDA_VISIBLE_DEVICES.

Multi-GPU inference is a different operation

from gptqmodel import BACKEND, GPTQModel

model = GPTQModel.load(
    "/models/my-quantized-checkpoint",
    backend=BACKEND.AUTO,
    device_map="auto",
)

Inference device-map placement does not enable the quantization calculation scheduler. Conversely, dense/MoE strategy fields do not configure inference tensor parallelism. Device-map normalization ignores Accelerate "disk" entries for kernel selection and excludes CPU offload entries when accelerator entries exist; an all-CPU map retains CPU as its target. Symbolic placement keywords such as "auto" identify a selector device without enumerating the final placement map. For concrete retained targets, the local selector validates each exact ordinal, rejects mixed device families/different GPU compute capabilities, and requires a backend compatible with those targets. See Inference Backends.

Native Q1_0 GGUF loading has a separate memory-sensitive profile decision. In the from_pretrained() path, when native Q1_0 is detected and no quantization configuration is supplied, AUTO resolves to FAST and uses dense Hugging Face GGUF import; a requested non-AUTO backend is ignored on that dense path. LOW_MEMORY instead permits redirection into quantized loading. A quantized source file therefore does not by itself guarantee packed resident weights. This profile handling is not applied by GPTQModel.load()'s already-quantized from_quantized() branch. See Inference Backends and GGUF.

Implementation internals

DeviceThreadPool and auxiliary workers

DeviceThreadPool lives in gptqmodel/utils/threadx.py, not a utils/threadx/ directory. utils/threads.py supplies other helpers such as AsyncManager and SerialWorker; it is not the pool implementation.

def __init__(
    self,
    devices: Optional[Iterable[DeviceLike]] = None,
    *,
    include_cuda: bool = True,
    include_xpu: bool = True,
    include_npu: bool = True,
    include_mps: bool = True,
    include_cpu: bool = True,
    inference_mode: bool = False,
    warmups: Optional[Dict[str, WarmupTask]] = None,
    empty_cache_every_n: int = 50,     # <=0 disables janitor
    workers: Optional[Dict[str, int]] = None,  # e.g. {'cpu':4, 'cuda:per':1, 'cuda:0':3}
    gc_debounce_seconds: float = 0.02,  # absorb bursty triggers before GC
    gc_min_interval_seconds: float = 0.0,  # throttle janitor passes
    pin_cpu_workers: bool = False,
    pin_accelerator_workers: bool = False,
):

The class discovers accelerator/CPU devices or accepts an explicit device iterable, establishes per-thread device contexts, and schedules callables with futures. Worker policies include cpu, mps, cuda:per, xpu:per, npu:per, and per-ordinal overrides such as cuda:0. Virtual aliases share parent-device locks. Unspecified devices default to one worker.

API / facility Behavior
submit(device, fn, /, *args, cuda_event=None, **kwargs) Queue work and return a future; a CUDA event can order work on CUDA devices, but is rejected for non-CUDA devices.
do(...) Synchronous execution interface over the pool.
submit_serial(...) Serial scheduling path.
wait(scope=None, *, lock=False) Drain all, a family, a device key, or an iterable of devices. With lock=True, returns a context manager that drains then acquires writer locks.
Locks/counters Per-device read/write locks, global/family/all locking, completed counters, and in-flight counters.
Warmups WarmupTask / WarmUpCtx can run device-scoped and thread-scoped initialization.
Janitor Periodic backend cache cleanup with debounce/minimum-interval controls; direct constructor default is 50 completions, and nonpositive values disable it.
shutdown(wait=True) Shut down workers with optional waiting.

The process-global DEVICE_THREAD_POOL is a lazy proxy defined in gptqmodel/__init__.py; first use calls get_device_thread_pool(). Its configured defaults differ from constructing DeviceThreadPool() directly:

Global pool policy Configured value
CUDA workers per ordinal 4
XPU / NPU workers per ordinal 1 / 1
MPS workers 8
CPU workers Computed by _default_cpu_workers(), not one fixed portable integer.
Model-loader CPU alias 2 workers
Empty-cache interval 512 completions

Pool concurrency does not make model mutation thread-safe by itself. Offload and state-dict changes use module-parent locks and scheduler boundaries. wait() drains scheduled work; it is not interchangeable with synchronizing every GPU operation created outside the pool.

AsyncManager in threads.py accepts callables, returns futures, tracks in-flight tasks, and propagates exceptions through future results. Its join() supports a particular future or all tracked tasks with a total timeout budget. It captures the device-telemetry context when work is submitted.

LazyTurtle checkpoint-backed materialization

gptqmodel/utils/structure.py::LazyTurtle is a source-checkpoint reader for a meta shell model. The loader attempts LazyTurtle.maybe_create() in the offload-enabled load path instead of eagerly constructing a complete CPU "turtle" model when a suitable local source is available.

It stores a checkpoint tensor map and materializes requested submodules from the relevant safetensors shards. Local pickle checkpoints with .bin, .pt, .pth, or .ckpt extensions can be converted once to temporary safetensors. This avoids repeatedly loading the whole dense checkpoint just to populate a small active subtree; conversion can still require substantial host memory and scratch disk.

maybe_create() returns None when the source is not an existing local directory or setup fails; construction failures are logged at debug level. In the offload-enabled source loader, a missing lazy source after shell construction raises RuntimeError rather than triggering an eager-load fallback. The separate direct-CPU fallback applies to recognized failures while building the meta shell; MiMo mixed sources reject that fallback. Lazy loading must therefore be observed, not assumed from the offload flag alone. The object declares supports_reload=False and is_lazy_checkpoint_source=True.

def materialize_submodule(
    self,
    *,
    target_model: torch.nn.Module,
    target_submodule: torch.nn.Module,
    device: torch.device,
    non_blocking: bool = False,
    module_path: Optional[str] = None,
    recurse: bool = True,
    tie_weights: bool = True,
    show_progress: bool = True,
) -> torch.nn.Module:

The materializer resolves runtime/checkpoint names through model-definition/module-tree information and reversed Hugging Face conversion maps, including supported MoE aliases and model placement exclusions. It takes a lock while copying checkpoint tensors, can use an explicit module_path to avoid scanning the whole model per operation, and re-ties weights after successful materialization when requested. Meta tensors are placeholders, not allocated weight storage; an ordinary .to() of a meta shell does not reconstruct source checkpoint data.

Lazy source checkpoints and completed-module offload bundles serve opposite directions of the lifecycle: LazyTurtle supplies original dense tensors into active work, while completed-module offload moves processed tensors out to disk. Neither directory is automatically the final quantized model export.

Completed-module disk offload

gptqmodel/utils/offload.py implements the quantization-side offload API on top of Accelerate:

def offload_to_disk(module: List[str] | nn.Module, model: nn.Module, disk_path: str = ".", force: bool = False):
Argument Default Meaning
module Required A module or list of dotted submodule names relative to model; NamedModule wrappers are unwrapped.
model Required Root used to resolve module names and relationships.
disk_path "." Root for per-module offload directories. Use a dedicated directory rather than relying on this low-level default.
force False Bypass the normal small-payload cutoff, including a payload exactly equal to 4 KiB; meta/direct-state checks still apply.

The implementation resolves the module's full name, takes a parent-module lock, and serializes module state. It skips already-meta modules and modules with neither direct parameters nor direct buffers. Modules with payloads of 4 KiB or less (total_bytes <= 4 * 1024) normally avoid disk writes; force=True bypasses this size cutoff, not the meta/direct-state checks. A caller that depends on an offload bundle as the only durable resume record must not assume the default writes every tiny module.

_bundle_module_state_dict() writes one module.safetensors bundle and one index.json. Index entries include dtype, shape, absolute bundle path, tensor name, and byte offsets into the safetensors file. State tensors are detached, copied to CPU, and made contiguous before writing. That staging itself needs host RAM; disk offload is not a zero-copy promise.

For replacement, _offload_disk_locked() writes a sibling .tmp directory referencing the final bundle path, renames an existing bundle to .old, renames the staged directory into place, and removes the old directory. This staged rename sequence avoids exposing a partially written bundle to the normal reader path; it is not a blanket transaction guarantee across filesystems or arbitrary process failures. Do not place unrelated files in per-module .tmp/.old directories: the writer removes those staging directories when reusing them.

After serialization, Accelerate disk_offload(..., offload_buffers=True, execution_device=original_device) installs hooks and replaces resident state with offload-backed/meta state. Weight tying is restored where the root supports it. The module's apparent meta parameters therefore need not mean the original checkpoint was never loaded; they can indicate completed offloaded state.

Importing this utility also replaces accelerate.utils.modeling.clear_device_cache with a no-op to avoid its thread-unsafe clearing behavior. The original callable is retained as ACCELERATE_CLEAR_DEVICE_CACHE. GPT-QModel has its own coordinated cleanup paths; do not assume every Accelerate helper still has its unmodified process-global behavior after import.

Restoring ordinary resident tensors

def undo_offload_to_disk(
    module: nn.Module,
    device: torch.device = CPU,
    include_buffers: bool = True,
    delete_offload_folders: bool = False,
    dtype: Optional[torch.dtype] = None,
) -> nn.Module:

The restore API materializes parameters and optionally buffers on the requested device, removes Accelerate hooks from the root and descendants, restores normal tensor ownership, and optionally deletes discovered offload folders. It returns the same module. The default target is CPU, buffer restoration is enabled, deletion is disabled, and dtype=None preserves each tensor's dtype.

Restoration first tries a version-detected hook weights_map; otherwise it uses Accelerate alignment and clones/rebinds the leaves. Cloning detaches tensors from memory-mapped or hook-managed views. Special handling keeps restored parameters out of inference-tensor state where they may later be mutated or re-tied. Unrelated meta leaves without offload hooks cannot be reconstructed by this API alone.

Folder discovery and optional deletion are best-effort, based on available hook metadata. Do not delete backing storage while live hooks still depend on it, and do not confuse undoing offload with exporting a complete checkpoint. Restoring everything to one GPU can itself exceed device memory.

Shared temporary-directory lifetime

_SharedTemporaryDirectory in quantization/config.py wraps tempfile.TemporaryDirectory. _create_temp_offload_dir() uses prefix "gptqmodel_". If offload_to_disk=True and no path is supplied, _ensure_offload_temp_dir() stores this wrapper in _offload_temp_dir and sets offload_to_disk_path to its name.

Both __copy__() and __deepcopy__() return the same wrapper. The deep-copy method also records itself in memo. Copied configuration objects therefore share one directory handle rather than creating independent owners that can prematurely remove each other's backing files. The wrapper exposes .name and .cleanup().

Managed temporary storage is a lifecycle convenience, not a durable checkpoint location. Keep the configuration/handle alive for the run, finish save/restoration before cleanup, and supply an explicit persistent scratch location if the run must retain its working files beyond temporary-directory lifetime.

Disk throughput estimation

def estimate_disk_io_speed(write_size_bytes: int = DEFAULT_WRITE_SIZE) -> float:
Property Value
DEFAULT_WRITE_SIZE 100 * 1024 * 1024 bytes (100 MiB).
BUFFER_SIZE 32 * 1024 * 1024 bytes (32 MiB).
Method Sequential writes to a temporary file; flush and fsync before calculating rate.
Return unit Bytes per second, not MB/s. The looper converts for display.
Invalid input Nonpositive write size raises ValueError.
I/O failure Logs a warning and returns 0.0.
Nonpositive measured elapsed time Returns positive infinity.

module_looper.py::io_write_performance() caches the write-speed result for the process when disk writes are enabled. The benchmark uses the system temporary-file location, so it may measure a different filesystem from an explicitly configured offload path. Treat it as a diagnostic estimate, not a benchmark of every destination or a guarantee that offload will be bandwidth-bound.

Calibration-cache placement and forwarding

StageInputsCapture in looper/stage_inputs_capture.py materializes the first decoder layer on CPU unless a deferred floatx-source path intentionally leaves it lazy. In the ordinary capture path, cur_layer_device is CPU. calibration_data_device=None uses that current layer device; it does not mean "automatically distribute calibration tensors over all GPUs."

Setting Capture behavior
None Store captured inputs on the current capture-layer device, CPU in the ordinary path.
"cpu" or an explicit accelerator Move captured data to that device. This is cache storage placement, not a request to run every quantization stage there.
"balanced" Select devices of the quantization device family, apply compute_device_filter, and assign captured batches round-robin. An empty filtered result falls back to the discovered list; if no candidates exist, the current layer device is used.

Captured hidden states, attention masks, position IDs, and supported nested keyword tensors follow the selected batch-storage device. Balanced activation placement is independent from balanced dense/MoE weight/work placement. Selecting both can increase cross-device traffic even when it lowers per-device storage.

auto_forward_data_parallel=True permits, but does not force, forward replicas. stage_subset.py combines it with other serial constraints; an explicit forward placement map, model placement overrides, or stage requirements can make a subset serial. Setting it to False forces that subset toward the serial forwarding path without removing the separately configured calculation pools. Model-specific placement exclusions are applied before preferred quantization-device state is recorded.

CUDA event activation buffers

utils/cuda_activation_buffer.py provides an asynchronous device-to-host transport separate from ordinary calibration-cache configuration.

def __init__(
    self,
    device: torch.device | str | int,
    stream: Optional[torch.cuda.Stream] = None,
    pin_memory: bool = True,
    host_allocator: Optional[Callable[[torch.Size, torch.dtype, torch.layout], torch.Tensor]] = None,
    host_reclaimer: Optional[Callable[[torch.Tensor], None]] = None,
) -> None:
def capture_async(
    self,
    activation: torch.Tensor,
    *,
    meta: Any = None,
    enqueue: bool = True,
) -> ActivationPacket:
Component / method Contract
ActivationPacket Dataclass with event, host_tensor, optional meta=None, and created_at from time.perf_counter.
Device CUDA must be available; constructor rejects non-CUDA devices and resolves an unspecified CUDA ordinal to the current device. Captured tensors must match the buffer device exactly.
Copy Detaches the activation, makes it contiguous if needed, allocates CPU storage, makes the copy stream wait for the producer stream, enqueues a nonblocking copy, and records an event.
pin_memory Defaults to True for internal host allocation. A custom allocator must return pinned CPU tensors.
enqueue Defaults to True; False returns a packet without putting it on the drain queue. Caller then owns its synchronization/lifetime.
drain(wait=True, max_items=None) Collect queued packets FIFO, synchronizing each event by default; with wait=False, consumers must synchronize before reading.
recycle(packet) Invokes an optional host reclaimer; it is not automatic deallocation of every packet.
pending_count() / len(buffer) Lock-protected approximate queued count, not a count of every live host or GPU tensor.
Context exit Drains with wait=True. 
from gptqmodel.utils.cuda_activation_buffer import CudaEventActivationBuffer

# activation must be a CUDA tensor on cuda:0.
with CudaEventActivationBuffer("cuda:0") as buffer:
    buffer.capture_async(activation, meta={"layer": 0})
    for packet in buffer.drain():
        # drain() has synchronized the copy event.
        process_cpu_activation(packet.host_tensor, packet.meta)
        buffer.recycle(packet)

Do not read or recycle host storage before the event completes. Retain source-storage lifetime until its asynchronous copy is safe; stream/event ordering is not permission to overwrite an in-flight operand. The internal queue is not a memory-budget mechanism: consumers must drain promptly and manage how many activations they retain. Pinned host RAM remains a limited resource even when GPU pressure drops.

Device inspection and usage measurements

utils/device.py contains representative-device queries and DeviceSMI-backed memory readings.

Helper Behavior / limitation
get_device(obj) A tensor returns its own device; a module returns the first parameter device, otherwise first buffer, otherwise CPU. This is not a uniformity assertion.
get_device_new(obj, recursive=False, assert_mode=False, expected=None, check_index=False) Optional recursive scan and placement assertions. Type equality is checked by default; check_index=True also requires matching ordinals.
get_gpu_usage_memory() DeviceSMI memory usage for CUDA_0, divided by 1024 cubed: GiB, not an aggregate over all GPUs.
get_cpu_usage_memory() DeviceSMI CPU memory usage divided by 1024 cubed: GiB.

A model sharded over several devices can have a representative first device without all of its tensors living there. Use explicit assertions and per-device instrumentation when checking placement. DeviceSMI usage, Torch live allocations, Torch allocator reservations, and process memory are different metrics.

Structured device telemetry

utils/device_telemetry.py uses a ContextVar with default False and a lock-protected process record list. TelemetryConfig.device also defaults to False; enable it with the config's telemetry object or mapping.

from gptqmodel import QuantizeConfig
from gptqmodel.utils.device_telemetry import (
    clear_device_telemetry_records,
    get_device_telemetry_records,
)

quant_config = QuantizeConfig(telemetry={"device": True})
clear_device_telemetry_records()
# Load and quantize using quant_config, then inspect emitted placement events.
records = get_device_telemetry_records()
API Behavior
device_telemetry_scope(enabled) Strict boolean validation; set/reset the current context flag.
with_quantization_device_telemetry Derive the invocation setting from model.quantize_config.telemetry.device.
capture_device_telemetry Capture the setting for queued work and reset reused worker context afterward; the thread pool and async helper use it.
emit_device_telemetry(event, **fields) When enabled, normalize fields, add event and rounded wall-clock ts, append under lock, and log the record.
get_device_telemetry_records() Return copied record dictionaries.
clear_device_telemetry_records() Clear stored records under lock.

Telemetry reports instrumented placement/scheduling events; it is not continuous allocator tracing. Records are process-local and accumulate until cleared. Context propagation prevents one job's setting from leaking into later work on the same thread, but applications still need to distinguish their own runs when sharing a process.

MemTracker: accounting and cache cleanup

def __init__(self, auto_gc_bytes: int | str | None = "auto") -> None:

MemTracker accounts for explicit tensor/module objects, not every allocation automatically. allocate() adds measured bytes; free() subtracts measured bytes, accumulates freed counters, and may trigger backend cache cleanup. Neither method destroys the supplied Python object. Drop actual owners/references separately.

API / behavior Detail
Accepted objects A tensor, module, or tuple of tensors/modules; unsupported tuple members raise TypeError.
Module scan Parameters and buffers recursively, including packed state.
Storage deduplication Within an object measurement, shared dense storage is counted once using storage pointer and byte size; meta tensors are skipped. Sparse components are measured separately.
Multiple top-level inputs Each tuple item is measured separately before aggregation; shared storage across separate top-level items can therefore be counted more than once.
allocated(device=None) / freed(device=None) Return (byte_count, formatted_string); omitted device aggregates counters. Family-only queries can sum ordinals.
reset() Clear allocation/freed/GC counters, not free model storage.
set_auto_gc(...) Accept nonnegative integer thresholds, "auto", or None. Here None requests auto detection, not unconditional disabling; use 0 to disable threshold triggering.

Auto detection uses one third of the smallest visible CUDA device's total memory. It is disabled when CUDA is unavailable or detection fails. Thresholds track accounted freed bytes per device; after successful backend empty-cache cleanup, that device's freed counter resets and GC counters increase. The helper supports CUDA/XPU/MPS/NPU cache calls where available; exceptions or unsupported backends leave cleanup unsuccessful.

This is accounting-triggered allocator cleanup, not a proof that all references were released or that a live model fits. Repeatedly allocating/accounting the same object without matching bookkeeping can inflate counters. Shared-storage deduplication within one measurement is not a global ownership database.

VRAM estimates and layer sizes

utils/vram.py supplies static size estimates:

def compute_module_sizes(
    model: nn.Module,
    dtype: Optional[Union[str, torch.device]] = None,
    special_dtypes: Optional[Dict[str, Union[str, torch.device]]] = None,
    buffers_only: bool = False,
):
Helper Meaning
dtype_byte_size(dtype) Logical per-element size; special aliases include "int2" → 0.25, "int4" → 0.5, "fp8" → 1. torch.bool is modeled as 1/8 byte, so this helper is not a literal allocator measurement.
named_module_tensors(...) Parameter/buffer enumeration, with nonpersistent-buffer handling.
compute_module_sizes(...) Aggregate tensor sizes through every dotted parent path, optionally using dtype estimates, per-name special dtypes, or buffers only. Integer/bool storage keeps its original modeled size under a floating dtype override.
get_all_layer_size(modules, module_sizes, no_split_module_classes) Identify leaf/no-split layer units and format their sizes.
get_vram(model) Return (formatted_total_size, layer_size_list), respecting the model's _no_split_modules list when available.

Static sizes do not include every runtime scratch allocation, activation cache, Hessian/workspace, allocator reservation, module replica, temporary dequantization, or external-engine overhead. Logical dtype overrides do not simulate a backend's packing metadata or alignment. For capacity planning, combine static packed-state estimates with actual stage/device telemetry and measured peak memory.

The dtype helper is intentionally limited, not a general Tensor.element_size() replacement. In addition to its logical aliases and modeled bool size, it handles FP16/BF16, FP32/INT32, and the byte-sized floatx dtypes listed in the source when present in Torch. Other integer storage dtypes, including torch.int8, torch.int16, torch.int64, and torch.uint8, raise ValueError. Consequently compute_module_sizes()/get_vram() can fail on a packed model whose buffers use those types; the integer-preservation rule does not imply that every integer dtype has a size entry. Measure actual tensor storage or provide a suitable per-name special-dtype estimate when this limitation applies.

Pitfalls and operational checks

Symptom / mistake Check
Several GPUs are visible but work stays serial Verify the relevant dense versus MoE strategy, effective device pool, method restriction, model support, GIL state, and stage's serial-forward rules.
Invalid pool device despite the physical GPU existing Use process-logical ordinals after CUDA_VISIBLE_DEVICES; ensure explicit pools are subsets after filtering.
OOM after enabling forward parallelism Module replicas and their state can increase memory; try auto_forward_data_parallel=False independently of pool strategy.
Calibration cache still occupies CPU RAM None follows capture placement; choose an explicit device or "balanced" if appropriate, and account for masks/kwargs too.
More GPU memory available but host memory rises Pinned activation queues, serialization staging, pickle conversion, and retained calibration tensors still consume RAM.
Temporary offload directory disappeared Keep shared config/handle lifetime valid; use an explicit durable scratch path when required and do not manually clean live backing files.
Offload is unexpectedly slow Check the actual destination filesystem/free space, not only the temporary-directory write-speed estimate.
Meta tensor cannot be moved normally Distinguish lazy original-checkpoint placeholders from Accelerate offload-backed tensors; use the corresponding materialization/restore path.
free()/empty_cache() did not reduce live memory Accounting and allocator cleanup do not remove Python references or safely cancel in-flight operations.
Inference spans incompatible GPUs The kernel selector rejects mixed families/capabilities; quantization device-pool availability is not a cross-architecture inference guarantee.

Do not clean module working directories until finalizers, offload hooks, queued copies, and saving have finished. Do not treat a successful disk write as proof that the full quantized model has been exported. Use the normal model.save() lifecycle for the final checkpoint and keep source checkpoints/offload work areas separate.

Source files

See also

Documentation home · Quickstart · Troubleshooting · Repository · PyPI

Source baseline: GPT-QModel 7.7.0 / c197691. Maintained by ModelCloud.ai.

Home

Getting started

Quantization methods

Complementary configs

Using a model

Platforms

Internals

Clone this wiki locally