OffloadManager

The OffloadManager class provides full control over the offloading lifecycle. Obtain an instance with get_offload_manager() rather than constructing one directly.

OffloadManager

OffloadManager(name: str)

Simplified offload manager API.

Provides a high-level interface for weight offloading with automatic phase management, profiling, and module patching capabilities.

external_compile_active property

external_compile_active: bool

Whether the active run uses external_compile with no compile_fn.

Describes the effective live run as resolved at offload() — mutating config.external_compile afterwards cannot flip it (that field is re-read only by the next offload()). False for eager offload and for FlexTensor-driven compile_fn runs, which do not need the caller to compile each unit.

strategy_replan_status property

strategy_replan_status: StrategyReplanStatus

Passive snapshot of the current replan lifecycle.

Reading this never arms or advances the tail, so it is safe both as a preflight check before :meth:request_strategy_replan and as an outcome check after the tail finishes. Field semantics are documented on :class:~flextensor.StrategyReplanStatus; for the preflight-then-drive pattern see docs/how-to/torch-compile.md.

can_request is required by both :meth:request_strategy_replan variants but sufficient only for the default one: the CUDA-graph variant (manual_update_state=True) adds requirements — offload_timing= 'cuda_graph' plus external CUDA timing events — not reflected here.

init

init(config: OffloadConfig | None = None, *, _shared_offload_owner_token: object | None = None) -> None

set_config

set_config(config: OffloadConfig) -> None

Set configuration for the offload manager.

Most fields take effect on the next phase transition or on the operation that reads them. The fields in :data:_TENSOR_MANAGER_ONESHOT_FIELDS are baked into the active :class:TensorManager when it is constructed at the first offload() call — _initialize_tensor_manager does not run again — so changing them here cannot affect the live manager. To honor a new value, release() the current manager and re-offload().

skip_discovery is the one field that raises rather than warns: the live TensorManager would keep its captured value while self.config.skip_discovery reported the new one, and offload_block()'s guard reads the live config — so a True→False flip would silently permit manual blocks while the manager stays in skip-mode and never captured their tensor mappings. Every other one-shot field emits a UserWarning naming it.

PARAMETER	DESCRIPTION
config	OffloadConfig instance with settings TYPE: OffloadConfig

RAISES	DESCRIPTION
RuntimeError	If skip_discovery differs from the value the active TensorManager captured at first offload().

WARNS	DESCRIPTION
UserWarning	If any other one-shot field differs from the value the active TensorManager captured.

offload

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

Offload modules in the model based on config patterns.

Patches modules matching config.include_patterns (and applies config.exclude_patterns) to automatically manage tensor transfers. Supports wildcards in include/exclude patterns.

PARAMETER	DESCRIPTION
model	PyTorch model to patch TYPE: Module
config	OffloadConfig to use for offload. Set config.external_compile=True (or FT_EXTERNAL_COMPILE=1 via :func:~flextensor.config.load_config``) to enable the compile-transparent ``pre_compute/post_compute`` custom-op forwards and auto-register the rolling loader at INFERENCE so the caller can apply external per-unit ``torch.compile`` after the eager forwards counted by :attr:iters_before_inference` (one graph per offloaded unit — not whole-model compile). TYPE: OffloadConfig | None DEFAULT: None
compile_fn	Optional per-unit compiler. When given, this activates the compiled-offload path (no FT_EXTERNAL_COMPILE env var needed): FlexTensor patches the offloaded units so the offload barriers are visible to the compiler, and applies compile_fn to each offloaded unit -- one compiled graph per unit (lambda m: torch.compile(m, fullgraph=True), a Torch-TensorRT compile, an eager passthrough lambda m: m, ...). Compiling at the offload granularity is slot-alias safe by construction: each graph reads exactly one rolling slot, so no graph can ever alias two same-slot units (the monolithic-graph miscompile). The callable takes one module and returns a callable module with the same forward signature; if it raises on a unit, that unit is left eager. With the default profile_mode='view', FlexTensor also runs compiled view-profile under compile_fn, so the offload strategy is built from compiled timings and no post-INFERENCE replan is needed. Call :meth:request_strategy_replan only when timings were not collected under compile (profile_mode='getter', or external torch.compile after external_compile=True). None (default) is exactly today's eager offload -- no behavior change. TYPE: Callable[[Module], Module] | None DEFAULT: None

RETURNS	DESCRIPTION
Module	A wrapper around the model that tracks state transitions

Examples:

Basic eager offload (see docs/quick-start.md for a runnable end-to-end example). With the default skip_discovery=False, discovery_iters + profiling_iters forwards drive the state machine into INFERENCE (the post-forward hook transitions on the Nth call); the extra iteration below runs the first real inference forward::

import torch
import flextensor as ft

model = MyModel().to("cpu").eval()
config = ft.OffloadConfig(include_patterns=["embed", "layers.*", "head"])
proxy = ft.offload(model, config)
om = ft.get_offload_manager()

x = torch.randn(batch, seq, dim, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
    for _ in range(om.iters_before_inference + 1):
        proxy(x)   # drives discovery -> profile -> inference

Per-unit torch.compile via compile_fn (preferred; see docs/how-to/torch-compile.md). Whole-model torch.compile(proxy) is not supported across phase transitions::

def compile_fn(module: torch.nn.Module) -> torch.nn.Module:
    return torch.compile(module, fullgraph=True)

proxy = ft.offload(model, config, compile_fn=compile_fn)
om = ft.get_offload_manager()
with torch.no_grad():
    for _ in range(om.iters_before_inference):
        proxy(x)  # discovery + compiled view-profile -> INFERENCE

For external per-unit compile after INFERENCE, set OffloadConfig(external_compile=True), drive om.iters_before_inference eager forwards, torch.compile each offloaded unit, then :meth:request_strategy_replan.

offload_from_state

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

Adopt a saved state and return the offloaded model proxy.

If this method 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.

offload_block

offload_block(name: str) -> Any

Create an offload block context manager.

PARAMETER	DESCRIPTION
name	Name of the offload block (used for tracking) TYPE: str

RETURNS	DESCRIPTION
Any	OffloadBlock context manager (Trap, WarmupTrap, TrapInfer, …)

RAISES	DESCRIPTION
RuntimeError	If offload() has not been called yet, or if skip_discovery was honored. Manual blocks need discovery; leave skip_discovery=False (the default). Allowed on the no-patched-modules fallback.

Examples:

>>> with om.offload_block("encoder"):
...     output = encoder(hidden_states)

Note

Compiled offload is driven by auto-patched forwards that call pre_compute/post_compute directly. Manual offload_block stays on the eager trap path (:class:~flextensor.trap_tensor_mode.TrapInferDirect / :class:~flextensor.trap_tensor_mode.TrapInfer); do not wrap an auto-patched compiled unit or transfers will double-schedule.

offload_units

offload_units() -> tuple[tuple[str, Module], ...]

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

Same set compile_fn wraps: nested patched modules are omitted so each graph would read one rolling slot. Order matches patch order; shared modules are identity-deduplicated (first path wins). Available after :meth:offload; the set is stable once the manager is in INFERENCE.

To torch.compile these units, enable external_compile=True (or pass compile_fn) and compile only after INFERENCE. Prefer module.compile(...) so the object in the tree stays the same; torch.compile(module) returns a new wrapper and is not written back until a later replace helper exists.

RETURNS	DESCRIPTION
tuple[str, Module]	Immutable (name, module) pairs. Empty when no model is bound
...	or nothing was patched.

resident_units

resident_units() -> tuple[tuple[str, Module], ...]

Return outermost modules that are not offload units and contain none.

These subtrees stay GPU-resident at serving time and may be compiled without compiled-offload custom ops. The root, containers that hold offload units, registry containers without an executable forward (ModuleList, ModuleDict), and children inside an offload unit are omitted.

Compile only after INFERENCE: earlier phases can swap the live module tree. Resident-only compile does not require external_compile or :meth:request_strategy_replan.

Boundaries are the complement of :meth:offload_units live identities, so a compile_fn replacement that dropped patch markers is still omitted, and a second parent of a shared offload unit is not listed as resident.

RETURNS	DESCRIPTION
tuple[str, Module]	Immutable (name, module) pairs in named_children walk
...	order, identity-deduplicated. Empty when no model is bound.

get_gpu_memory_usage

get_gpu_memory_usage() -> GPUMemoryUsage

Get GPU memory usage by FlexTensor in inference mode.

Returns the memory used by GPU transfer blocks and unmapped tensors that were moved to GPU. This method should be called after the manager has transitioned to inference mode.

RETURNS	DESCRIPTION
GPUMemoryUsage	Memory breakdown with per-component bytes and MB values. See GPUMemoryUsage for field details. TYPE: GPUMemoryUsage

RAISES	DESCRIPTION
RuntimeError	If called before the manager has transitioned to inference mode.

Examples:

>>> om = 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 GPU memory: {usage.total_mb:.2f} MB")
>>> print(f"  Transfer blocks: {usage.blocks_mb:.2f} MB")
>>> print(f"  Unmapped tensors: {usage.unmapped_tensors_mb:.2f} MB")

save_profile

save_profile(profile_directory: str | None = None) -> None

Save current offload profile to directory.

PARAMETER	DESCRIPTION
profile_directory	Directory to save profile to. If None, uses config.profile_storage_dir. TYPE: str | None DEFAULT: None

RAISES	DESCRIPTION
RuntimeError	If tensor manager is not initialized or no profile data exists.
ValueError	If no directory is provided and config.profile_storage_dir is not set, or if config.profile_read_only is True.

load_profile

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

Load offload profile from directory and restore state to model.

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

RAISES	DESCRIPTION
RuntimeError	If no model is available to restore state to.
FileNotFoundError	If profile directory or file doesn't exist.
ValueError	If no directory is provided and config.profile_storage_dir is not set.

clear_profiling_durations

clear_profiling_durations() -> None

Clear all accumulated duration measurements.

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

RAISES	DESCRIPTION
RuntimeError	if called before :func:flextensor.offload — no active TensorManager to clear.

.. note:: With OffloadConfig(enabled=False) the active manager is :class:NoOpTensorManager: this method is a no-op because no durations are ever recorded.

suspend_profiling

suspend_profiling() -> None

Suppress profiling recording and (in PROFILING phase) freeze the iteration counter.

While suspended:

• record_all() (the PROFILING recorder) is a complete no-op — neither tensor IDs nor duration samples are stored, so a paused warmup pass cannot widen per-layer tensor sets on data-dependent models (MoE / conditional branches / mixed-batch shapes).
• record_duration() skips duration recording.
• record_tensors() (the DISCOVERY recorder) is not affected — discovery's tensor-to-layer mapping is a hard prerequisite for every later phase.
• In DISCOVERY the iteration counter still advances (discovery durations are unused).
• In PROFILING the counter is paused so suppressed passes don't consume the profiling_iters budget.

Suspensions are reference-counted: each call must be matched by a :meth: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.

RAISES	DESCRIPTION
RuntimeError	if called before :func:flextensor.offload — no active TensorManager to suspend.

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

resume_profiling

resume_profiling() -> None

Release one outstanding :meth:suspend_profiling call.

Recording only resumes once every outstanding suspension has been released.

RAISES	DESCRIPTION
RuntimeError	if called before :func:flextensor.offload (no active TensorManager), or if called while not suspended (unbalanced suspend / resume).

pause_profiling

pause_profiling() -> Any

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

RAISES	DESCRIPTION
RuntimeError	if called before :func:flextensor.offload — no active TensorManager, so the with block cannot actually run with profiling paused.

collect_offload_timing

collect_offload_timing() -> OffloadTimingReport | None

Collect aggregate offload timing from the durable measure store.

Flushes any pending eager pass before draining (see :meth:TensorManager.collect_offload_timing). Under CUDA-graph replay, call :meth:update_offload_timing (or :meth:update_state during a manual replan) after each replay so passes are published before collect — this method does not re-finalize the last replay (that would duplicate it).

Must be called after all forward passes have completed (e.g. after pipe() returns or at shutdown). Calling this drains the durable measure store, so the next call returns data only for passes since :meth:reset_offload_timing / this collect. Retention before drain is capped by :data:~flextensor.offload_timing.OFFLOAD_TIMING_MEASURE_MAX_PASSES.

RETURNS	DESCRIPTION
OffloadTimingReport | None	class:~flextensor.offload_timing.OffloadTimingReport with
OffloadTimingReport | None	per-trap min/max/avg/std for transfer, compute, and wait, or
OffloadTimingReport | None	None when:
OffloadTimingReport | None	:attr:OffloadConfig.offload_timing is "off",
OffloadTimingReport | None	the :class:TensorManager is not initialized yet, or
OffloadTimingReport | None	the durable measure store is empty after flush (no passes in the current window).

reset_offload_timing

reset_offload_timing() -> None

Start a fresh durable offload-timing measure window.

Clears accumulated measure passes and the collector's rolling log window. Drops any pending eager pass on the collector so it cannot be published into this new window by a later on_pass_start. Call before a measure segment you care about (serving → shutdown :meth:collect_offload_timing).

update_offload_timing

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

Store the current offload-timing pass into the collector.

Measurement only — does not advance phase or rebuild strategy. Needed under CUDA graphs: each graph.replay() overwrites the same captured event handles, so call this after every replay you want to keep (or :meth:update_state does it when replan used manual_update_state=True).

RETURNS	DESCRIPTION
bool	True when a pass was published or same-generation deduped.
bool	False when finalize raised or unexpectedly no-op'd (inactive
bool	pass / capturing) — callers must not advance a measure replan
bool	counter after False.

request_strategy_replan

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

Remeasure under the current runtime and rebuild the offload strategy.

• manual_update_state=True — CUDA graphs: run graph.replay() then :meth:update_state for the returned count (recapture graphs after).
• manual_update_state=False — external compile: run that many model forwards; the forward hook calls :meth:update_state for you.

Needed after external torch.compile (external_compile=True), after compile_fn with profile_mode='getter', or after CUDA-graph capture. Not required for default compile_fn + profile_mode='view'.

PARAMETER	DESCRIPTION
manual_update_state	When True, caller drives each measure step (e.g. graph.replay() + :meth:update_state). TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
int	Measure (and warmup) forwards / replays to drive, or 0 when
int	nothing was armed. Both paths arm only while
int	attr:strategy_replan_status reports can_request; otherwise
int	they warn and leave the current tail untouched.

update_state

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

Advance the phase state machine by one iteration.

In INFERENCE, phase transitions are finished, but when a replan tail is armed this method still advances it:

• CUDA-graph (request_strategy_replan(manual_update_state=True)) — :meth:update_offload_timing then advance the measure counter; the last call rebuilds from graph budgets.
• Compiled / external-compile — advance the passive warm→measure tail via :meth:CompiledOffload.on_forward (usually from the forward hook; no explicit call needed).

PARAMETER	DESCRIPTION
replay_generation	Forwarded to :meth:update_offload_timing on the CUDA-graph measure path. Pass the same generation used by a gen-aware serving hook so -1 is not required and double-publish is avoided. TYPE: int DEFAULT: -1

No-op in NOT_INITIALIZED.

When profiling is suspended:

• During DISCOVERY the iteration counter still advances (discovery durations are unused and cleared at the transition anyway).
• During PROFILING the counter is paused so that suppressed passes don't consume the profiling_iters budget.

Compiled-profile compile-warmup model forwards (while profile_compile_warm_remaining > 0) likewise do not advance the PROFILING counter, so the logged profiling_iters measured samples are collected in full after warmup.

release

release(*, preserve_model: bool = True, _shared_offload_owner_token: object | None = None) -> None

Release this manager and its TensorManager resources.

PARAMETER	DESCRIPTION
preserve_model	Restore loader-managed tensor values before teardown. Set to False only when the managed model will be discarded; this avoids private copies of SHM-backed weights and leaves that model unusable. TYPE: bool DEFAULT: True