Simplified API¶

Module-level convenience functions — the recommended way to use FlexTensor. These wrap the OffloadManager singleton for common use cases.

DEFAULT_MANAGER_NAME `module-attribute` ¶

DEFAULT_MANAGER_NAME: str = 'default'

Name used by :func:get_offload_manager when no explicit name is provided.

init ¶

init(config: OffloadConfig | None = None, name: str = DEFAULT_MANAGER_NAME) -> None

offload ¶

offload(model: Module, config: OffloadConfig | None = None, name: str = DEFAULT_MANAGER_NAME, compile_fn: Callable[[Module], Module] | None = None) -> Module

Offload modules using the specified offload manager.

PARAMETER	DESCRIPTION
`model`	PyTorch model to patch. TYPE: `Module`
`config`	OffloadConfig to use for offload. TYPE: `OffloadConfig \| None` DEFAULT: `None`
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`
`compile_fn`	Optional per-unit compiler. Passing it activates the compiled-offload path without any `FT_EXTERNAL_COMPILE` env var: FlexTensor applies `compile_fn` to each offloaded unit (one compiled graph per unit -- slot-alias safe by construction). With default `profile_mode='view'`, compiled view-profile builds the strategy from compiled timings so no post-INFERENCE replan is needed. Use :meth:`OffloadManager.request_strategy_replan` only for getter-profile or external-compile paths. `None` (default) is exactly today's eager offload. See :meth:`OffloadManager.offload` for the full contract. TYPE: `Callable[[Module], Module] \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`Module`	The patched model.

get_offload_manager ¶

get_offload_manager(name: str = DEFAULT_MANAGER_NAME) -> OffloadManager

Get or create an OffloadManager singleton.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

RETURNS	DESCRIPTION
`OffloadManager`	OffloadManager instance

RAISES	DESCRIPTION
`RuntimeError`	If the named manager was created by a different thread.

Examples:

>>> om = get_offload_manager()
>>> om.set_config(OffloadConfig(gpu_device=0))

set_config ¶

set_config(config: OffloadConfig, name: str = DEFAULT_MANAGER_NAME) -> None

Set config for the specified offload manager.

PARAMETER	DESCRIPTION
`config`	OffloadConfig to apply. TYPE: `OffloadConfig`
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

offload_block ¶

offload_block(block_name: str, name: str = DEFAULT_MANAGER_NAME) -> Any

Create offload block using the specified offload manager.

PARAMETER	DESCRIPTION
`block_name`	Name for the offload block. TYPE: `str`
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

get_gpu_memory_usage ¶

get_gpu_memory_usage(name: str = DEFAULT_MANAGER_NAME) -> GPUMemoryUsage

Get GPU memory usage by the specified offload manager.

Returns the memory used by GPU transfer blocks and unmapped tensors. Must be called after the manager has transitioned to inference mode.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

RETURNS	DESCRIPTION
`GPUMemoryUsage`	GPUMemoryUsage dataclass with memory breakdown in bytes and megabytes.

RAISES	DESCRIPTION
`RuntimeError`	If called before inference mode.

save_profile ¶

save_profile(profile_directory: str | None = None, name: str = DEFAULT_MANAGER_NAME) -> None

Save profile using the specified offload manager.

PARAMETER	DESCRIPTION
`profile_directory`	Directory to save profile to. If None, uses config.profile_storage_dir. TYPE: `str \| None` DEFAULT: `None`
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

load_profile ¶

load_profile(profile_directory: str | None = None, model: Module | None = None, name: str = DEFAULT_MANAGER_NAME) -> None

Load profile using the specified offload manager.

PARAMETER	DESCRIPTION
`profile_directory`	Directory containing the profile. If None, uses config.profile_storage_dir. TYPE: `str \| None` DEFAULT: `None`
`model`	Model to restore state to. If None, uses the current model. TYPE: `Module \| None` DEFAULT: `None`
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

offload_from_profile ¶

offload_from_profile(model: Module, profile_directory: str, config: OffloadConfig | None = None, name: str = DEFAULT_MANAGER_NAME, compile_fn: Callable[[Module], Module] | None = None) -> Module

Initialize, load a saved profile, and offload the model in one step.

Convenience wrapper that combines :func:init, :func:load_profile, and :func:offload. The profile must be loaded before offloading so that initialize_warmup sees the saved state and skips discovery/profiling, going straight to inference mode.

PARAMETER	DESCRIPTION
`model`	PyTorch model to offload. TYPE: `Module`
`profile_directory`	Directory containing the saved profile. TYPE: `str`
`config`	OffloadConfig to apply. TYPE: `OffloadConfig \| None` DEFAULT: `None`
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`
`compile_fn`	Optional per-unit compiler (same as :func:`offload`). TYPE: `Callable[[Module], Module] \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`Module`	The offloaded model proxy, ready for inference.

offload_from_state ¶

offload_from_state(model: Module, state: TensorManagerState, config: OffloadConfig | None = None, name: str = DEFAULT_MANAGER_NAME, *, allow_strategy_replan: bool = True) -> Module

Adopt an in-memory tensor-manager state and enter inference.

The state must describe model and use the transfer mode selected by config. Storage placement is adopted before loader metadata is restored; the model and its modules keep their identity.

If this function raises, cleanup preserves storage ownership and removes FlexTensor modifications where possible, but does not roll back storage placement. The input model is not guaranteed usable and must be discarded.

PARAMETER	DESCRIPTION
`model`	Eager model matching the saved state. TYPE: `Module`
`state`	Validated state to adopt. TYPE: `TensorManagerState`
`config`	Offload configuration. Uses the manager's current configuration when omitted. TYPE: `OffloadConfig \| None` DEFAULT: `None`
`name`	Offload manager name. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`
`allow_strategy_replan`	Preserve source weights for a later compiled-timing replan. TYPE: `bool` DEFAULT: `True`

RETURNS	DESCRIPTION
`Module`	An inference-ready proxy for `model`.

release ¶

release(name: str = DEFAULT_MANAGER_NAME, *, preserve_model: bool = True) -> None

Release the specified offload manager.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`
`preserve_model`	Restore loader-managed tensor values before teardown. Set to `False` only when the managed model will be discarded. TYPE: `bool` DEFAULT: `True`

Profiling Data Control¶

clear_profiling_durations ¶

clear_profiling_durations(name: str = DEFAULT_MANAGER_NAME) -> None

Clear all accumulated profiling duration measurements.

Wipes every duration sample collected so far. Tensor-to-layer mappings (discovery data) are not affected.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

suspend_profiling ¶

suspend_profiling(name: str = DEFAULT_MANAGER_NAME) -> None

Suppress profiling recording (durations and tensor IDs).

In PROFILING, suppression is total — TensorManager.record_all() is a complete no-op, so a paused warmup pass cannot widen per-layer tensor sets on data-dependent models (MoE / conditional branches / mixed-batch shapes). In DISCOVERY, record_tensors() is not affected (the tensor-to-layer mapping is a hard prerequisite for every later phase). The iteration counter is frozen in PROFILING so suppressed passes don't consume the profiling_iters budget. See :meth:OffloadManager.suspend_profiling for the full per-phase contract.

Suspensions are reference-counted: each :func:suspend_profiling must be matched by a :func:resume_profiling, and recording only resumes once every outstanding suspension has been released. Independent callers can therefore bracket their own sections without accidentally resuming someone else's suspension.

.. note:: With OffloadConfig(enabled=False) the active manager is :class:NoOpTensorManager: these functions are no-ops and the profiling-iteration counter is not frozen.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

Example::

model = flextensor.offload(model, config=config)

flextensor.suspend_profiling()
for size in [1, 8, 32, 128]:
    model(make_dummy_input(size))   # profiling suppressed during warmup
flextensor.resume_profiling()

for batch in dataloader:            # durations are now collected for the real workload
    output = model(batch)

resume_profiling ¶

resume_profiling(name: str = DEFAULT_MANAGER_NAME) -> None

Resume profiling duration recording after :func:suspend_profiling.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

pause_profiling ¶

pause_profiling(name: str = DEFAULT_MANAGER_NAME) -> Any

Context-manager form of :func:suspend_profiling / :func:resume_profiling.

Shares the refcount with the raw functions, so it nests freely with them.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

Example::

model = flextensor.offload(model, config=config)

with flextensor.pause_profiling():
    for size in [1, 8, 32, 128]:
        model(make_dummy_input(size))   # profiling suppressed during warmup

for batch in dataloader:                # durations are now collected for the real workload
    output = model(batch)

Offload timing and strategy replan¶

Measure H2D / compute / wait during inference (OffloadConfig.offload_timing) and rebuild the offload strategy after compile or CUDA-graph capture. After INFERENCE, offload_units() / resident_units() list the modules to compile per unit. OffloadManager.external_compile_active and OffloadManager.strategy_replan_status report replan availability and outcome without touching the model. See Measure transfer overlap and torch.compile.

collect_offload_timing ¶

collect_offload_timing(name: str = DEFAULT_MANAGER_NAME) -> OffloadTimingReport | None

Collect aggregate offload timing from the durable measure store.

Convenience wrapper around :meth:OffloadManager.collect_offload_timing. After CUDA-graph replay, call :func:update_offload_timing first so passes are published; this drains the store and does not re-finalize.

PARAMETER	DESCRIPTION
`name`	Name for the offload manager. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`

RETURNS	DESCRIPTION
`OffloadTimingReport \| None`	class:`~flextensor.offload_timing.OffloadTimingReport`, or `None`
`OffloadTimingReport \| None`	when timing is `"off"`, the manager has no `TensorManager` yet, or
`OffloadTimingReport \| None`	the durable measure store is empty after flush.

reset_offload_timing ¶

reset_offload_timing(name: str = DEFAULT_MANAGER_NAME) -> None

Start a fresh durable offload-timing measure window.

Convenience wrapper around :meth:OffloadManager.reset_offload_timing.

update_offload_timing ¶

update_offload_timing(name: str = DEFAULT_MANAGER_NAME, *, replay_generation: int = -1) -> bool

Store the current offload-timing pass into the collector.

Convenience wrapper around :meth:OffloadManager.update_offload_timing. Measurement only. Call after each graph.replay() you want to keep (or :func:update_state does it when replan used manual_update_state=True).

RETURNS	DESCRIPTION
`bool`	`True` on success / timing inactive; `False` if finalize failed.

request_strategy_replan ¶

request_strategy_replan(name: str = DEFAULT_MANAGER_NAME, *, manual_update_state: bool = False) -> int

Remeasure under the current runtime and rebuild the offload strategy.

Convenience wrapper around :meth:OffloadManager.request_strategy_replan.

PARAMETER	DESCRIPTION
`name`	Offload manager name. Defaults to DEFAULT_MANAGER_NAME. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`
`manual_update_state`	Caller-driven measure path — arm post-capture timing and require :func:`update_state` after each opaque step (today: `graph.replay()`). After rebuild, recapture CUDA graphs before serving. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`int`	Forwards / replays to drive, or `0` when no replan was armed.

StrategyReplanStatus `dataclass` ¶

StrategyReplanStatus(state: StrategyReplanState, can_request: bool, applied: bool | None)

Passive snapshot of :meth:~flextensor.OffloadManager.request_strategy_replan state.

Read from :attr:~flextensor.OffloadManager.strategy_replan_status; reading never arms or advances the replan tail.

ATTRIBUTE	DESCRIPTION
`state`	`"idle"` before any replan request, `"in_progress"` while the warm/measure tail rides the caller's forwards, `"done"` once it finished (applied or refused), `"failed"` when the rebuild raised — inference is unsafe after `"failed"`. TYPE: `StrategyReplanState`
`can_request`	Whether calling `request_strategy_replan()` now is supported. A `"done"` replan that refused to rebuild can still be retried while FlexTensor holds the source weights; `"in_progress"` and `"failed"` cannot. TYPE: `bool`
`applied`	`None` before and during a request, `True` after a successful rebuild, `False` after a completed soft refusal (strategy kept) or a failure. TYPE: `bool \| None`

offload_units ¶

offload_units(name: str = DEFAULT_MANAGER_NAME) -> tuple[tuple[str, Module], ...]

Return outermost offloaded modules as (qualified_name, module) pairs.

Convenience wrapper around :meth:OffloadManager.offload_units.

resident_units ¶

resident_units(name: str = DEFAULT_MANAGER_NAME) -> tuple[tuple[str, Module], ...]

Return outermost GPU-resident modules as (qualified_name, module) pairs.

Convenience wrapper around :meth:OffloadManager.resident_units.

update_state ¶

update_state(name: str = DEFAULT_MANAGER_NAME, *, replay_generation: int = -1) -> None

Advance offload phase / armed replan by one step.

Convenience wrapper around :meth:OffloadManager.update_state. Needed explicitly after graph.replay() when replan was armed with manual_update_state=True (replay does not run module forward hooks).

PARAMETER	DESCRIPTION
`name`	Offload manager name. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`
`replay_generation`	Optional CUDA-graph replay generation for the manual measure path (dedupes with a gen-aware serving hook). TYPE: `int` DEFAULT: `-1`

format_offload_timing_table ¶

format_offload_timing_table(data: OffloadTimingSnapshot | OffloadTimingReport) -> str

Format a timing snapshot or report as a human-readable table.

Accepts either a single-pass :class:OffloadTimingSnapshot or a multi-pass :class:OffloadTimingReport.

OffloadTimingReport `dataclass` ¶

OffloadTimingReport(per_trap: Sequence[TrapTimingStats] = (), passes: Sequence[OffloadTimingSnapshot] = ())

Aggregate timing report across multiple forward passes.

Produced by :meth:~flextensor.OffloadManager.collect_offload_timing (drains TensorManager's durable measure store). The collector's rolling _snapshots window is for periodic logging only.

Canonical pass data is :attr:passes. num_passes and pass-total sums/averages are derived from it. :attr:per_trap holds precomputed min/max/avg/median/std stats built once when the report is assembled.

num_passes `property` ¶

num_passes: int

total_compute_sum `property` ¶

total_compute_sum: float

total_transfer_sum `property` ¶

total_transfer_sum: float

total_wait_sum `property` ¶

total_wait_sum: float

total_compute_avg `property` ¶

total_compute_avg: float

total_transfer_avg `property` ¶

total_transfer_avg: float

total_wait_avg `property` ¶

total_wait_avg: float

compute_budgets_by_profile_label ¶

compute_budgets_by_profile_label(profile_labels: Sequence[str], *, conservative: bool = False) -> dict[str, float]

Map this report's compute budgets onto profile trap labels by index.

Timing row i pairs with profile_labels[i] (report labels may be shortened relative to TensorManager stats / LayerStatistics.label).

Replan feeds only these compute budgets into :meth:~flextensor.tensor_manager.TensorManager.replan_from_compiled_durations (rewriting per-trap duration). Measured transfer_ms and wait_ms are not strategy inputs:

Transfer cost stays on the profiling size→time curve (memory_transfer_stats). H2D volume and host↔device bandwidth are assumed unchanged under CUDA-graph replay, so the size-based model remains valid; runtime transfer_ms is a diagnostic check that the curve still matches.
Wait is an outcome of the current schedule (exposed stall when H2D was not hidden), not a planner cost. wait_ms > 0 means the prior plan lost overlap; replan adjusts the compute/hiding window so the knapsack can re-pack. Feeding wait back as a cost would be circular.

OffloadTimingSnapshot `dataclass` ¶

OffloadTimingSnapshot(per_trap: Sequence[TrapTimingRecord] = ())

Timing data from a single forward pass.

Canonical data is :attr:per_trap. Aggregate total_*_ms values are derived from those records so they cannot drift.

total_wait_ms `property` ¶

total_wait_ms: float

Cumulative stall across all traps in this pass.

total_transfer_ms `property` ¶

total_transfer_ms: float

Cumulative H2D transfer time across all traps in this pass.

total_compute_ms `property` ¶

total_compute_ms: float

Cumulative compute / budget window across all traps in this pass.

TrapTimingRecord `dataclass` ¶

TrapTimingRecord(label: str, transfer_ms: float = 0.0, compute_ms: float = 0.0, wait_ms: float = 0.0)

Per-trap timing measurements from a single forward pass.

All durations are in milliseconds. A wait_ms value of zero means the transfer finished before compute needed the data (fully hidden).

label is the stable trap name string (same concept as LayerStatistics.label until that legacy rename lands).

TrapTimingStats `dataclass` ¶

TrapTimingStats(label: str, compute_min: float = 0.0, compute_max: float = 0.0, compute_avg: float = 0.0, compute_median: float = 0.0, compute_std: float = 0.0, transfer_min: float = 0.0, transfer_max: float = 0.0, transfer_avg: float = 0.0, transfer_median: float = 0.0, transfer_std: float = 0.0, wait_min: float = 0.0, wait_max: float = 0.0, wait_avg: float = 0.0, wait_median: float = 0.0, wait_std: float = 0.0)

Aggregate timing statistics for one trap across multiple passes.

compute_budget_ms ¶

compute_budget_ms(*, conservative: bool = False) -> float

Positive compute budget (ms) for offload knapsack / graph re-plan.

Under CUDA-graph replay the compute column is the per-trap hiding window. Prefer median for a typical budget; conservative=True prefers min (tightest window across noisy probes).

Only this column is consumed by replan. Measured transfer_ms / wait_ms stay diagnostic (see :meth:OffloadTimingReport.compute_budgets_by_profile_label).

Lazy Model Initialization¶

load_model_from_profile ¶

load_model_from_profile(tensor_manager: TensorManager, profile_directory: str, model_class: type[Module], model_args: tuple = (), safetensor_path: str | None = None, model_kwargs: dict | None = None, buffer_device: str | None = None, diagnose: bool = False, strict: bool = False) -> Module

Load a model from a saved profile state with optimized weight loading.

This function creates a model with parameters on meta device and buffers on GPU, then loads only the necessary weights from safetensors based on the profile state. This is more memory-efficient than loading the full model.

PARAMETER	DESCRIPTION
`tensor_manager`	The TensorManager instance to use for state restoration TYPE: `TensorManager`
`profile_directory`	Directory containing the profile (profile.json) TYPE: `str`
`model_class`	The model class to instantiate (e.g., Transformer) TYPE: `type[Module]`
`model_args`	Positional arguments for the model constructor TYPE: `tuple` DEFAULT: `()`
`safetensor_path`	Path to the safetensors file. Required. TYPE: `str \| None` DEFAULT: `None`
`model_kwargs`	Keyword arguments for the model constructor TYPE: `dict \| None` DEFAULT: `None`
`buffer_device`	Device for model buffers (default: uses tensor manager's device) TYPE: `str \| None` DEFAULT: `None`
`diagnose`	If True, run diagnostics on nested parameter tracking TYPE: `bool` DEFAULT: `False`
`strict`	If True, raise KeyError when any weight key from the profile is missing in the safetensors file. Default False. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`Module`	The loaded model in eval mode, ready for use with tensor_manager

RAISES	DESCRIPTION
`KeyError`	If strict=True and any profile weight keys are missing from safetensors

Example

from flextensor import TensorManager from flextensor.lazy_model_init import load_model_from_profile

tensor_manager = TensorManager(device_gpu, ...) model = load_model_from_profile( ... tensor_manager, ... "./profile", ... Transformer, ... model_args=(model_args,), ... safetensor_path="/path/to/model.safetensors", ... ) tensor_manager.set_model(model) model.set_tensor_manager(tensor_manager) model = tensor_manager.initialize_warmup()