Configuration¶
Configuration classes and loading utilities for FlexTensor offloading behavior.
Bases: BaseModel
Configuration for offload manager.
class-attribute
instance-attribute
¶
Whether to enable offload.
class-attribute
instance-attribute
¶
GPU device index to use.
class-attribute
instance-attribute
¶
Whether to use pinned (page-locked) memory for CPU tensors, enabling non-blocking GPU transfers.
class-attribute
instance-attribute
¶
pinned_memory_mode: PinnedMemoryMode = Field(default='torch')
How to pin CPU tensors when pinned_memory=True.
"torch"(default) — :meth:torch.Tensor.pin_memory: copies the tensor into a fresh pinned allocation."host_register"—cudaHostRegister: pins the existing allocation in place (no copy, lower peak host memory). Requires a loadable CUDA runtime.
Has no effect when pinned_memory=False.
Can also be set via the FT_PINNED_MEMORY_MODE env var.
.. note:: Only controls the non-SHM allocator path. When
shm_enabled=True and the SHM segment is built with
pinned_memory=True, that segment is registered in place via
cudaHostRegister regardless of this mode.
.. note:: Records the requested value and never mutates. Two construction-time outcomes can diverge from a naive read of this field:
- **CUDA available, cudart binding broken/missing** —
:class:`~flextensor.tensor_manager.TensorManager` falls back from
``"host_register"`` to ``"torch"`` and emits a ``WARNING``.
- **CUDA unavailable** — constructing
:class:`~flextensor.tensor_manager.TensorManager` with
``pinned_memory=True`` raises ``RuntimeError``. Set
``pinned_memory=False`` on CPU-only hosts.
class-attribute
instance-attribute
¶
Enable cross-process weight sharing via POSIX shared memory.
When true, model weights are stored in shared memory so multiple processes
can share a single copy in CPU RAM. This selects SHM-backed allocation and
requires a non-empty :attr:shm_namespace at configuration time. A
multi-process application must also coordinate one creator with its
followers. Enabling this flag alone does not elect a creator or build
follower models on the meta device.
Can also be set via the FT_SHM_ENABLED env var.
class-attribute
instance-attribute
¶
Base namespace for shared memory blocks.
Must be set explicitly when :attr:shm_enabled is true. Applications that
have a stable model path may derive a value with
:func:flextensor.shm.compute_shm_namespace before initializing a manager.
Can also be set via FT_SHM_NAMESPACE env var.
class-attribute
instance-attribute
¶
Timeout (seconds) for followers waiting on creator.
0 means no overall limit; followers still use creator-heartbeat liveness detection. A positive value also bounds coordination-lock acquisition, though scheduler delays can make observed elapsed time slightly longer. Can also be set via FT_SHM_WAIT_TIMEOUT env var.
class-attribute
instance-attribute
¶
Number of discovery-phase iterations before profiling begins.
Can also be set via the FT_DISCOVERY_ITERS env var.
class-attribute
instance-attribute
¶
Skip the discovery phase and jump directly to profiling.
When True, tensor-to-layer mappings are discovered statically from
modules patched via include_patterns instead of being inferred from
discovery-phase iterations. Falls back to the discovery phase if no
modules are patched; a WARNING is logged and
OffloadManager.skip_discovery_honored flips to False so the
fallback is detectable without log scraping.
Leave at the default False for manual offload_block() blocks.
offload_block() raises when the skip is honored; the no-patched-
modules fallback still allows manual blocks. Set True only on the
auto-trap (include_patterns) path to cut startup time.
.. note:: Minimum-iteration contract
Because ``OffloadManager.update_state()`` runs as a post-forward hook,
the PROFILING→INFERENCE transition can only fire *after* the first
profile forward completes. Any transition — including
``profiling_iters=0`` — therefore always sees at least one profile
forward's worth of collector data. Setting ``profiling_iters=0`` and
``profiling_iters=1`` produce the same runtime behavior; both are
low-quality strategy inputs (``report_profiling_quality`` warns).
The only way to reach INFERENCE without any profile forward is to
restore a saved profile via ``load_profile`` / ``offload_from_profile``.
The value is read once when the underlying TensorManager is first
constructed (during the first offload() call). Calling set_config
with a different skip_discovery afterwards raises RuntimeError;
call release() and re-offload() to pick up a new value.
class-attribute
instance-attribute
¶
Number of profiling-phase iterations to collect timing statistics.
Can also be set via the FT_PROFILING_ITERS env var.
class-attribute
instance-attribute
¶
Prepare offloaded layers for external per-unit torch.compile after INFERENCE.
Does not compile the model itself. It arms the custom-op / rolling-loader
path so a caller (e.g. vLLM or a script) can torch.compile each offloaded
unit after the INFERENCE transition. Pass compile_fn to
:func:~flextensor.offload for the integrated path instead.
Can also be set via FT_EXTERNAL_COMPILE.
Requires a block transfer_mode (allocation_block_transfer or
raw_block_transfer); strategy is incompatible with the rolling-block
custom ops used by compiled offload.
With compile_fn and profile_mode='view' (default), FlexTensor
auto-runs compiled view-profile so the strategy is built from compiled
timings (no post-INFERENCE replan). Use
:meth:~flextensor.OffloadManager.request_strategy_replan after a non-view
profile or after external compile with external_compile=True.
class-attribute
instance-attribute
¶
Multiplier on the time budget for weight transfers.
Values below 1.0 add a safety margin to absorb profiling measurement noise
(e.g., 0.95 = 5% margin). Most relevant in latency mode
(max_gpu_mem_fraction=None); in memory mode the strategy may override
this value to meet the GPU memory target.
class-attribute
instance-attribute
¶
Tensor loading mode. Block transfer loaders (allocation_block_transfer,
raw_block_transfer) do not support shared tensors between layers;
use strategy for models with shared weights.
class-attribute
instance-attribute
¶
Number of memory blocks for block transfer loaders.
More blocks reserve more GPU memory but give the pipelining scheduler more room to overlap transfers with compute. Fewer blocks (e.g., 2) reduce GPU memory overhead and can work well for models whose layers have long compute times (such as diffusers), since the transfer easily finishes within a single layer's compute window.
class-attribute
instance-attribute
¶
Minimum number of blocks for assignment optimization search range.
The optimizer tries block counts from min_blocks to num_blocks and
picks the best. Lower values give more freedom and reduce GPU memory, but
may hurt pipelining throughput. Must be >= 2 (pipelined execution requires
at least two blocks).
Can also be set via the FT_MIN_BLOCKS env var.
class-attribute
instance-attribute
¶
Target maximum GPU memory as a fraction of total device memory (e.g. 0.9 = 90%).
When set, the strategy switches to memory mode and keeps peak GPU usage within this
budget. At runtime, the fractional budget is capped by actual available GPU memory so
the strategy never targets more than what is free. When None, the strategy operates
in latency mode with no explicit memory constraint. Defaults to None.
Can also be set via the FT_MAX_GPU_MEM_FRACTION env var.
class-attribute
instance-attribute
¶
Directory for profile persistence.
class-attribute
instance-attribute
¶
Only load profiles, saving is disabled.
class-attribute
instance-attribute
¶
load_strategy: Annotated[Strategy, WithJsonSchema({type: object})] | None = Field(default=None)
Override automatic strategy selection.
class-attribute
instance-attribute
¶
Whether to capture component init args for debugging.
class-attribute
instance-attribute
¶
Instrumentation output directory.
class-attribute
instance-attribute
¶
Inference offload-timing instrumentation mode.
"off"— disabled (default)."eager"— record per-trap transfer / compute / wait with internal CUDA timing events (module forwards / collect via :meth:OffloadManager.collect_offload_timing)."cuda_graph"— same measurement usingtorch.cuda.Event(..., external=True)soelapsed_time()is fresh after CUDA-graph replay (PyTorch >= 2.8). Required for :meth:OffloadManager.request_strategy_replanwithmanual_update_state=True.
Periodic log cadence and durable measure retention are internal defaults
(not independently tunable). Can also be set via FT_OFFLOAD_TIMING.
Requires a block transfer_mode (allocation_block_transfer or
raw_block_transfer). transfer_mode='strategy' has no
PreallocatedLoader enter/exit timing hooks, so enabling timing with
strategy is rejected at config construction.
One-shot: consumed when the TensorManager is first constructed (first
offload()). A later set_config with a different value warns; call
release() and re-offload() to apply it.
class-attribute
instance-attribute
¶
Policy when a PIECEWISE join forces outstanding H2D onto the critical path.
Under PIECEWISE capture, join_after_forward drains transfer_stream at
the piece boundary. If schedule and wait straddle pieces (rearrange
early-slot, or nested enter/exit such as parent 1 with graphs on
1.1 / 1.2), the transfer still completes but async overlap is lost.
"off"— disable the check."warn"— log a warning (default; lost overlap is not silent)."error"— raise :class:~flextensor.piecewise_prefetch_policy.PiecewisePrefetchPolicyError.
Integrations must call the block loader's join_after_forward() before
each piece's capture_end. There is no in-tree automatic piece hook yet;
without that call only the last-trap join runs, so mid-piece boundaries are
neither joined nor checked.
Can also be set via FT_PIECEWISE_PREFETCH.
One-shot: consumed when the TensorManager is first constructed (first
offload()). A later set_config with a different value warns; call
release() and re-offload() to apply it.
class-attribute
instance-attribute
¶
Module or parameter patterns to include for offloading (supports * and ? wildcards).
Each entry is one of three forms:
<glob>— name-based (module or parameter path), existing behaviour.name:<glob>— name-based, explicit form. Equivalent to<glob>.class:<glob>— class-based; matches modules whose class name matches the glob. Each pattern is tested against both the short class name (type(m).__name__) and the fully-qualified class name (f"{cls.__module__}.{cls.__qualname__}"); a match on either wins. Useful for hybrid architectures where a single Python class (e.g.SharedExpertMLP) appears at varying paths and renames across upstream versions. FQCN globs such asclass:torch.nn.*.Linearprovide an escape hatch when the same short name is used by multiple classes.
Name patterns are matched against both module paths (from
model.named_modules()) and parameter paths (from
model.named_parameters()).
Module-level patterns (e.g., layers.*) patch matching modules and offload
all their parameters.
Parameter-level patterns (e.g., layers.*.weight) automatically derive the
module-level pattern (layers.*) for patching, but only offload parameters that
match the full pattern. Non-matching parameters stay on GPU permanently.
Class patterns patch every module of a matching class; all of that
module's parameters are offloaded. Class patterns are module-level only —
class:*Expert* covers every descendant parameter of any matching module.
Can also be set via the FT_INCLUDE_PATTERNS env var as a comma-separated list.
.. note::
Standalone * has different semantics depending on context:
* **Module selection** — ``*`` matches exactly one path segment.
``layers.*`` matches ``layers.0`` but not ``layers.0.attn``.
* **Parameter filtering** — ``*`` matches one or more segments.
``*.weight`` matches both ``linear.weight`` and ``layers.0.attn.weight``.
* **Class matching** — ``*`` is a single-segment glob spanning the
whole class-name string. Short class names have no dots, so
``class:*Expert*`` behaves as a simple substring glob. Patterns
containing dots only match against the fully-qualified class name
(e.g., ``class:*.SharedExpertMLP``).
Validated at construction; assignment / in-place mutation bypasses validation.
class-attribute
instance-attribute
¶
Module/parameter patterns to exclude from offloading (supports wildcards).
Accepts the same three forms as :attr:include_patterns: bare <glob>,
name:<glob>, or class:<glob>. Class patterns match against both
the short class name and the fully-qualified class name (see
:attr:include_patterns).
Applied after include_patterns: targets matching both include and exclude are NOT offloaded.
Standalone * matches one or more segments in both module selection and parameter
filtering, so foo.* excludes all descendants of foo. A class: match
cascades to every descendant parameter of the matched module.
Can also be set via the FT_EXCLUDE_PATTERNS env var as a comma-separated list.
Example
Keeping Nemotron-H MoE routers and shared-expert MLPs on GPU while offloading the rest of each layer::
exclude_patterns = ["class:SharedExpertMLP", "class:MoELayer", "*.norm"]
Validated at construction; assignment / in-place mutation bypasses validation.
class-attribute
instance-attribute
¶
Whether to log diagnostic information (per-trap duration statistics, block assignment table) after strategy computation.
class-attribute
instance-attribute
¶
Profile-phase mechanism (how the profile model is patched).
"view"— default. Pre-allocates substantial GPU memory during profile; pick"getter"if tight on memory."getter"— property getters route parameter access through the per-trap loader. Lower GPU footprint, attribute-getter overhead in per-trap durations."torch_function"— fallback for models that reject the patching used by"view"/"getter". Heavy per-op overhead; only valid withtransfer_mode='strategy'.
"view" and "getter" are the two profiling variants of the
model-patching ("direct") runtime and affect only the profile phase;
"torch_function" selects the indirect runtime, which also changes
warmup and inference.
See profile_mode in docs/explanation/configuration.md for the
timing model, memory formula, and tradeoffs. Can also be set via the
FT_PROFILE_MODE env var.
Return a validated copy while preserving explicit-field tracking.
Configuration Loading¶
load_config(config_path: str | Path | None = None, env_prefix: str = 'FT_', use_env: bool = True, config_file_env_var: str = 'FT_CONFIG_FILE', **kwargs: Any) -> OffloadConfig
Load OffloadConfig with precedence: file < env vars < kwargs.
This is the unified configuration loading function that combines file and environment variable loading with proper precedence.
Precedence order (highest wins):
- kwargs (explicit arguments)
- Environment variables (if use_env=True)
- Config file values
- OffloadConfig defaults
When loading from environment only (no file), enabled defaults to False.
You must explicitly set FT_ENABLED=1 to enable offloading.
Supported file formats (auto-detected by extension):
- INI (.conf, .ini) — requires [flextensor] section
- JSON (.json)
- YAML (.yaml, .yml)
| PARAMETER | DESCRIPTION |
|---|---|
config_path
|
Path to configuration file. If None, checks config_file_env_var.
TYPE:
|
env_prefix
|
Prefix for environment variable names (default: "FT_")
TYPE:
|
use_env
|
Whether to load and override with environment variables (default: True)
TYPE:
|
config_file_env_var
|
Environment variable for config file path (default: "FT_CONFIG_FILE")
TYPE:
|
**kwargs
|
Additional keyword arguments to override all other values
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
OffloadConfig
|
OffloadConfig instance |
Examples:
load_config_from_file(config_path: str | Path | None = None, env_var: str = 'FT_CONFIG_FILE', **kwargs: Any) -> OffloadConfig
Load OffloadConfig from a configuration file (no env var override).
This is a convenience wrapper around load_config() that loads from file only, without environment variable overrides.
Supported file formats (auto-detected by extension):
- INI (.conf, .ini) — requires [flextensor] section
- JSON (.json)
- YAML (.yaml, .yml)
| PARAMETER | DESCRIPTION |
|---|---|
config_path
|
Path to configuration file. If None, uses env_var.
TYPE:
|
env_var
|
Environment variable name for config file path (default: "FT_CONFIG_FILE")
TYPE:
|
**kwargs
|
Additional keyword arguments to override file values
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
OffloadConfig
|
OffloadConfig instance with values from file and kwargs |
| RAISES | DESCRIPTION |
|---|---|
FileNotFoundError
|
If the config file does not exist |
ValueError
|
If the file format is not supported or file is invalid |
Examples:
load_config_from_env(prefix: str = 'FT_', **kwargs: Any) -> OffloadConfig
Load OffloadConfig from environment variables (no file).
This is a convenience wrapper around load_config() that loads from environment variables only, without file loading.
Environment variables should be named as {prefix}{FIELD_NAME} in uppercase. For example, with the default prefix "FT_":
- FT_ENABLED=1 (required to enable offloading)
- FT_GPU_DEVICE=1
- FT_DISCOVERY_ITERS=5
Note: When loading from environment, enabled defaults to False.
You must explicitly set FT_ENABLED=1 to enable offloading.
Type conversions:
- bool: Case-insensitive parsing of "true"/"false", "1"/"0", "yes"/"no", etc.
- int: Direct integer conversion
- float: Direct float conversion
- str: Used as-is
- None: Case-insensitive "none" or "null"
| PARAMETER | DESCRIPTION |
|---|---|
prefix
|
Prefix for environment variable names (default: "FT_")
TYPE:
|
**kwargs
|
Additional keyword arguments to override environment variables
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
OffloadConfig
|
OffloadConfig instance with values from environment variables and kwargs |
Examples:
Types¶
dataclass
¶
GPU memory usage by FlexTensor in inference mode.
This dataclass provides a breakdown of GPU memory allocated by FlexTensor for tensor offloading operations. It includes memory used for transfer blocks and unmapped tensors that are moved to GPU.
| ATTRIBUTE | DESCRIPTION |
|---|---|
blocks_bytes |
Memory in bytes used by GPU transfer blocks for CPU-to-GPU tensor transfers.
TYPE:
|
unmapped_tensors_bytes |
Memory in bytes used by tensors that are not managed by the offload strategy and were moved directly to GPU.
TYPE:
|
total_bytes |
Total GPU memory in bytes used by FlexTensor.
TYPE:
|
Examples:
>>> import flextensor
>>> from flextensor import OffloadConfig
>>> om = flextensor.get_offload_manager()
>>> config = OffloadConfig(include_patterns=["layers.*"])
>>> model = om.offload(model, config=config)
>>> # Run discovery and profiling iterations (path-aware count)...
>>> for _ in range(om.iters_before_inference):
... model(input)
>>> # Now in inference mode, get memory usage
>>> usage = om.get_gpu_memory_usage()
>>> print(f"FlexTensor using {usage.total_mb:.2f} MB GPU memory")