Strategies¶
Offloading strategies determine which model weights to move from GPU to CPU memory.
Pass a strategy instance to OffloadConfig(load_strategy=...).
Strategy Protocol¶
Bases: Protocol
Protocol for offload strategies.
All offload strategies must implement this protocol to be usable with the OffloadConfig and TensorManager.
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Compute offload strategy.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
Per-trap statistics containing tensor info and durations.
TYPE:
|
memory_stats
|
Optional dict mapping tensor size (bytes) to transfer time (ms).
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit in bytes. When set (memory mode),
the strategy must respect this constraint. When
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and optional block_data. |
Result Types¶
dataclass
¶
StrategyResult(strategy_map: dict[str, list[TensorStatistics]], block_data: BlockStrategyData | None = None)
Result from strategy.compute().
All strategies return this structure for consistency. TensorManager uses the appropriate fields based on loader_type.
instance-attribute
¶
Maps trap label to list of weights to offload.
class-attribute
instance-attribute
¶
block_data: BlockStrategyData | None = None
Block allocation data. None if n_blocks=0 or not computed.
dataclass
¶
BlockStrategyData(label_to_size_map: dict[str, int], allocation_ordered: dict[int, list[str]], block_sizes: dict[int, int] | list[int], label_to_block_id: dict[str, int], transfer_to_compute_map: dict[str, str])
Block allocation data computed by strategies.
This contains all data needed for block-based loaders to execute pipelined CPU→GPU transfers.
instance-attribute
¶
Maps trap label to total size of weights to offload.
instance-attribute
¶
Maps block_id to ordered list of trap labels assigned to it.
instance-attribute
¶
Size of each block in bytes (max trap size per block).
instance-attribute
¶
Maps trap label to assigned block index.
instance-attribute
¶
Maps transfer trap label to compute trap label for pipelining.
Knapsack Strategies¶
KnapsackStrategy(scale: float = 1.0, cyclic: bool = True, group_size: int = 1, threshold_mb: float = 0.1, n_blocks: int = 4)
Knapsack-based offloading strategy using duration-based optimization.
Initialize KnapsackStrategy for weight offloading.
| PARAMETER | DESCRIPTION |
|---|---|
scale
|
Duration multiplier (baseline).
TYPE:
|
cyclic
|
Whether to apply cyclic strategy extension.
TYPE:
|
group_size
|
Number of layers to merge into groups.
TYPE:
|
threshold_mb
|
Minimum tensor size threshold in MB.
TYPE:
|
n_blocks
|
Number of blocks for block-based loaders. Set to 0 to skip block data computation.
TYPE:
|
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Compute offload strategy using knapsack optimization.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics containing tensor info and durations.
TYPE:
|
memory_stats
|
Optional dict mapping tensor size (bytes) to transfer time (ms). If not provided, extracts from layer stats.
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit in bytes (never exceed).
When
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and optional block_data. |
KnapsackBlockStrategy(scale: float = 1.0, group_size: int = 1, threshold_mb: float = 0.1, n_blocks: int = 4)
Block-based knapsack strategy that uses externally provided memory transfer statistics.
This strategy expects memory_stats to be provided by TensorManager, which handles the live GPU benchmarking. This separates concerns: TensorManager manages profiling, strategies focus on optimization logic.
Initialize KnapsackBlockStrategy for weight offloading.
| PARAMETER | DESCRIPTION |
|---|---|
scale
|
Duration multiplier (baseline).
TYPE:
|
group_size
|
Number of layers to merge into groups.
TYPE:
|
threshold_mb
|
Minimum tensor size threshold in MB.
TYPE:
|
n_blocks
|
Number of blocks for block-based loaders.
TYPE:
|
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Compute offload strategy using block-based knapsack optimization.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics containing tensor info and durations.
TYPE:
|
memory_stats
|
Dict mapping tensor size (bytes) to transfer time (ms). Should be provided by TensorManager via live GPU benchmarking. If not provided, falls back to extracting from layer stats.
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit in bytes (never exceed).
When
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and block_data. |
AdaptiveKnapsackStrategy(scale: float = 1.0, loader_type: str = 'strategy', cyclic: bool = True, group_size: int = 1, threshold_mb: float = 0.1)
Adaptive strategy that selects between KnapsackStrategy and KnapsackBlockStrategy based on the loader type configuration.
This provides a unified interface that automatically selects the appropriate strategy implementation based on how tensors will be loaded during inference.
Initialize AdaptiveKnapsackStrategy.
| PARAMETER | DESCRIPTION |
|---|---|
scale
|
Duration multiplier (baseline). See
TYPE:
|
loader_type
|
Type of tensor loader to use. Options: - "strategy": Uses KnapsackStrategy (profiled tensor data) - "raw_block_transfer": Uses KnapsackBlockStrategy - "allocation_block_transfer": Uses KnapsackBlockStrategy
TYPE:
|
cyclic
|
Whether to apply cyclic strategy extension.
TYPE:
|
group_size
|
Number of layers to merge into groups.
TYPE:
|
threshold_mb
|
Minimum tensor size threshold in MB.
TYPE:
|
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Compute offload strategy using the appropriate strategy implementation.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics containing tensor info and durations.
TYPE:
|
memory_stats
|
Optional dict mapping tensor size (bytes) to transfer time (ms).
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit in bytes (never exceed).
When
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and block_data. |
Budget-First Strategy¶
BudgetFillStrategy(n_blocks: int = 0, *, threshold_mb: float = 1.0, scale: float = 1.0, assignment_strategy: AssignmentStrategy | None = None, enable_layer_de: bool = True, enable_tensor_de: bool = True, pop_size: int = 20, epoch: int = 60, max_early_stop: int = 20, seed: int | None = None)
BudgetFill facade: greedy, then optional layer / tensor DE solvers.
Always runs :class:BudgetFillGreedyStrategy. Optionally runs
:class:BudgetFillLayerDEStrategy (enable_layer_de, default True) and
:class:BudgetFillTensorDEStrategy (enable_tensor_de, default True),
keeping the best feasible plan with the fewest offloaded bytes.
Pinching the budget does not prove min-offload, so DE is not skipped solely
for a tight peak — Adaptive constructs BudgetFill with
enable_tensor_de=False and uses
:class:~flextensor.strategy.global_strategy.GlobalTensorSelectionStrategy
under extra_optimization instead.
Each solver remains independently usable for A/B comparison.
| PARAMETER | DESCRIPTION |
|---|---|
n_blocks
|
Pipeline block count (
TYPE:
|
threshold_mb
|
Ignore tensors at or below this size (MiB).
TYPE:
|
scale
|
Multiplier on previous-layer duration for transfer-fit ranking.
TYPE:
|
assignment_strategy
|
Trap-to-block assignment (default Strict).
TYPE:
|
enable_layer_de
|
Run per-layer DE after greedy.
TYPE:
|
enable_tensor_de
|
Run per-tensor DE after greedy/layer DE (default True; Adaptive keeps this False).
TYPE:
|
pop_size / epoch / max_early_stop
|
Caps forwarded to DE solvers.
|
seed
|
RNG seed for DE solvers.
TYPE:
|
property
¶
Name of the solver that produced the last :meth:compute result.
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Run greedy, then optional scipy solvers; keep the best feasible plan.
BudgetFillGreedyStrategy(n_blocks: int = 0, *, threshold_mb: float = 1.0, scale: float = 1.0, assignment_strategy: AssignmentStrategy | None = None)
Ranked-prefix BudgetFill: scan fit-then-size prefixes under a hard peak budget.
| PARAMETER | DESCRIPTION |
|---|---|
n_blocks
|
Pipeline block count.
TYPE:
|
threshold_mb
|
Ignore tensors at or below this size (MiB).
TYPE:
|
scale
|
Multiplier on previous-layer duration for transfer-fit ranking.
TYPE:
|
assignment_strategy
|
Trap-to-block assignment. Defaults to
:class:
TYPE:
|
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Select a minimal ranked offload prefix that meets max_gpu_mem_bytes.
BudgetFillLayerDEStrategy(n_blocks: int = 0, *, threshold_mb: float = 1.0, scale: float = 1.0, assignment_strategy: AssignmentStrategy | None = None, pop_size: int = 20, epoch: int = 60, max_early_stop: int = 20, seed: int | None = None)
Per-layer differential-evolution BudgetFill (one offload fraction per layer).
Same hard/soft goals as :class:BudgetFillGreedyStrategy, with search space
~n_layers. Within each layer, tensors are taken in fit-then-size order up
to the chosen byte budget. Used by :class:BudgetFillStrategy when
enable_layer_de=True, or alone for comparison.
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Optimize per-layer offload fractions under max_gpu_mem_bytes via DE.
BudgetFillTensorDEStrategy(n_blocks: int = 0, *, threshold_mb: float = 1.0, scale: float = 1.0, assignment_strategy: AssignmentStrategy | None = None, pop_size: int = 30, epoch: int = 80, max_early_stop: int = 25, seed: int | None = None)
Per-tensor differential-evolution BudgetFill (slow; opt-in fallback).
Same hard/soft goals as :class:BudgetFillGreedyStrategy, but searches
arbitrary subsets via differential evolution
(scipy.optimize.differential_evolution).
Needed when the best feasible set is not a within-layer fit→size prefix
(e.g. take a mid-ranked tensor without its higher-ranked neighbors). Prefer
:class:BudgetFillLayerDEStrategy for speed; enable via
:class:BudgetFillStrategy enable_tensor_de=True or run alone.
Adaptive leaves this off and uses :class:GlobalTensorSelectionStrategy
for expensive tensor-level search under extra_optimization.
Seeds the population with the greedy ranked-prefix plan.
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Optimize a binary offload mask under max_gpu_mem_bytes via DE.
Simple Strategies¶
Greedy offloading strategy based on cumulative compute time heuristic.
Initialize GreedyStrategy.
| PARAMETER | DESCRIPTION |
|---|---|
threshold_mb
|
Minimum tensor size threshold in MB.
TYPE:
|
n_blocks
|
Number of blocks for block-based loaders.
TYPE:
|
scale
|
Multiplier on the cumulative compute budget (< 1 adds safety margin, > 1 allows more transfers).
TYPE:
|
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Compute greedy offload strategy.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics.
TYPE:
|
memory_stats
|
Unused, accepted for interface consistency.
TYPE:
|
max_gpu_mem_bytes
|
Unused, accepted for interface consistency.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and optional block_data. |
Offloading strategy that offloads every Nth layer.
Initialize NthLayerStrategy.
| PARAMETER | DESCRIPTION |
|---|---|
nth_layer
|
Interval between offloaded layers.
TYPE:
|
threshold_mb
|
Minimum tensor size threshold in MB.
TYPE:
|
n_blocks
|
Number of blocks for block-based loaders.
TYPE:
|
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Compute Nth layer offload strategy.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics.
TYPE:
|
memory_stats
|
Unused, accepted for interface consistency.
TYPE:
|
max_gpu_mem_bytes
|
Unused, accepted for interface consistency.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and optional block_data. |
Global Strategies¶
GlobalOffloadStrategy(n_blocks: int = 4, scale: float = 1.0, threshold_mb: float = 0.1, min_blocks: int | None = None, max_blocks: int | None = None, assignment_strategy: AssignmentStrategy | None = None)
Global offload strategy using assignment-based layer-to-block mapping for pipelined execution.
Pipeline Execution Model: - Layer N computes using Block X (data transferred during Layer N-1) - While computing Layer N, we transfer Block Y for Layer N+1 - Constraint: Block Y ≠ Block X (can't transfer and compute same block simultaneously) - Block X becomes free after Layer N finishes → can be reused for Layer N+2
This means: - Minimum 2 blocks required for pipelining - With 2 blocks: alternating pattern (A, B, A, B, ...) - With 3+ blocks: more flexibility for scheduling
The strategy uses an AssignmentStrategy (default: OptimizedRoundRobinAssignment) to determine layer-to-block assignments, then calculates block sizes based on the maximum weight size assigned to each block.
Initialize GlobalOffloadStrategy.
| PARAMETER | DESCRIPTION |
|---|---|
n_blocks
|
Number of reusable memory blocks (minimum 2 for pipelining).
TYPE:
|
scale
|
Duration multiplier (baseline). Controls how much of the previous
layer's compute time is available for transfers.
TYPE:
|
threshold_mb
|
Minimum weight size threshold in MB. Weights smaller than this are kept on GPU and not offloaded.
TYPE:
|
min_blocks
|
Minimum number of blocks to use (default: 2). - 2 blocks: With alternating pattern (A,B,A,B,...)
TYPE:
|
max_blocks
|
Maximum number of blocks to use (default: n_blocks).
TYPE:
|
assignment_strategy
|
Strategy for layer-to-block assignment. Default: OptimizedRoundRobinAssignment (best GPU utilization). Available strategies: - OptimizedRoundRobinAssignment: Optimized cyclic pattern (recommended) - StrictRoundRobinAssignment: Strict i % n_blocks pattern
TYPE:
|
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
If n_blocks < 2 (minimum required for pipelining). |
ValueError
|
If min_blocks > max_blocks. |
ValueError
|
If min_blocks < 2 or max_blocks > n_blocks. |
ValueError
|
If scale <= 0. |
staticmethod
¶
Calculate total size of tensors in a layer.
staticmethod
¶
Calculate peak GPU memory for pipelined execution.
Peak memory includes: - All blocks (must be pre-allocated on GPU for pipelining) - Non-offloaded tensors (always stay on GPU)
| PARAMETER | DESCRIPTION |
|---|---|
block_sizes
|
Size of each block in bytes (integers).
TYPE:
|
non_offloaded_memory
|
Memory used by tensors that are not offloaded.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
int
|
Total peak GPU memory in bytes (integer). |
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Compute offload strategy using pipelined block optimization.
When max_gpu_mem_bytes is set (memory mode), the strategy will
automatically increase scale beyond its initial value if the
initial offloading is insufficient to meet the memory constraint.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics containing tensor info and durations.
TYPE:
|
memory_stats
|
Optional dict mapping tensor size (bytes) to transfer time (ms). If provided, uses this data. Otherwise falls back to extraction from layer stats.
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit in bytes (never exceed).
When
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and block_data. |
GlobalTensorSelectionStrategy(scale: float = 1.0, scale_ub: float | None = None, n_blocks: int = 4, min_blocks: int | None = None, max_blocks: int | None = None, threshold_mb: float = 0.1, pop_size: int = 50, epoch: int = 200, max_early_stop: int = 25, optimizer: str = 'DE', seed_ratio: float = 0.25, seed: int | None = None, assignment_strategy: AssignmentStrategy | None = None, transfer_window: TransferWindowCalculator | None = None)
Weight-level metaheuristic optimization strategy for GPU memory offloading.
This strategy optimizes the binary decision of whether to offload each weight, maximizing GPU memory utilization while respecting memory limits and ensuring transfers fit within pipelined execution windows.
Unlike GlobalOffloadStrategy which optimizes block sizes and layer assignments, this strategy performs fine-grained weight selection to find the optimal subset of weights to offload.
Decision variables: - Binary x[i] for each weight: 1 = offload, 0 = keep on GPU - Continuous scale: multiplier for compute time (transfer window)
Constraints: - Hard: Peak GPU memory <= max_gpu_mem_bytes (when set) - Hard: Transfer time for layer N <= (layer N-1's compute time) * scale
Objective (maximize): - GPU utilization relative to max_gpu_mem (when set) - Offload ratio (secondary) - Prefer lower scale (less latency overhead)
Memory Model: - Tensors NOT offloaded stay on GPU permanently (preloaded) - Peak GPU = sum(all kept tensors) + sum(block sizes) - Blocks are pre-allocated for pipelined transfers
Supported Optimizers (via scipy): - DE: Differential Evolution (scipy.optimize.differential_evolution) - RECOMMENDED default. Natively handles binary variables via integrality constraints. Best for high-dimensional problems with many tensors. - SA: Simulated Annealing (scipy.optimize.dual_annealing) - Treats variables as continuous; less efficient for binary decisions. May be useful for small problems with few tensors.
Initialize GlobalTensorSelectionStrategy.
| PARAMETER | DESCRIPTION |
|---|---|
scale
|
Duration multiplier (baseline / lower bound for the optimizer).
TYPE:
|
scale_ub
|
Upper bound for scale the optimizer may explore.
When
TYPE:
|
n_blocks
|
Number of memory blocks (max). Used if assignment_strategy is None.
TYPE:
|
min_blocks
|
Minimum number of blocks to try (default: 2).
TYPE:
|
max_blocks
|
Maximum number of blocks to try (default: n_blocks).
TYPE:
|
threshold_mb
|
Minimum tensor size threshold in MB (smaller tensors always kept on GPU).
TYPE:
|
pop_size
|
Population size for optimization.
TYPE:
|
epoch
|
Maximum number of optimization epochs.
TYPE:
|
max_early_stop
|
Stop if no improvement for this many epochs.
TYPE:
|
optimizer
|
Optimizer to use (default: "DE"). Available options: - "DE": Differential Evolution (scipy.optimize.differential_evolution). RECOMMENDED. Natively handles binary variables via integrality constraints. Best for high-dimensional problems with many tensors. - "SA": Simulated Annealing (scipy.optimize.dual_annealing). Treats variables as continuous; less efficient for binary decisions. May be useful for small problems with few tensors.
TYPE:
|
seed_ratio
|
Ratio of population to seed with feasible solutions (0.0-1.0). Lower ratio = more random exploration.
TYPE:
|
seed
|
Random seed for reproducibility. If None, uses non-deterministic seed.
TYPE:
|
assignment_strategy
|
Strategy for assigning layers to blocks. Default: OptimizedRoundRobinAssignment (tries min_blocks to max_blocks). Options: StrictRoundRobinAssignment, OptimizedRoundRobinAssignment.
TYPE:
|
transfer_window
|
Calculator for per-layer transfer windows. Default: SingleLayerWindow (previous layer's duration only). Use GapAwareWindow to exploit empty layers for larger transfer budgets.
TYPE:
|
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Run the core optimizer.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics containing tensor info and durations.
TYPE:
|
memory_stats
|
Optional dict mapping tensor size (bytes) to transfer time (ms). If provided, uses this data. Otherwise falls back to extraction from layer stats.
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit in bytes (never exceed).
When
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
StrategyResult containing strategy_map and block_data. |
Optimizer parameter values
GlobalTensorSelectionStrategy accepts an optimizer parameter with two accepted values: "DE" (differential evolution via scipy.optimize.differential_evolution) and "SA" (dual annealing via scipy.optimize.dual_annealing). "DE" is recommended for most cases as it natively handles binary variables. "SA" treats variables as continuous and is less efficient for binary decision problems; it may be useful for small problems with few tensors.
Adaptive Strategy¶
AdaptiveStrategy(scale: float = 1.0, loader_type: str = 'allocation_block_transfer', cyclic: bool = False, group_size: int = 1, threshold_mb: float = 0.1, min_blocks: int = 2, n_blocks: int = 4, extra_optimization: bool = False)
Run all available strategies and return the best result.
For block-based loaders (the common path), the following candidates are evaluated:
- :class:
KnapsackBlockStrategy - :class:
GlobalOffloadStrategy(Optimized assignment) - :class:
GlobalOffloadStrategy(Strict assignment) - :class:
BudgetFillStrategy(only whenmax_gpu_mem_bytesis set; layer DE on, tensor DE off — useenable_tensor_de=Trueonly when calling BudgetFill directly as a non-prefix fallback)
With extra_optimization=True (slower):
- :class:
GlobalTensorSelectionStrategy(Optimized assignment) - :class:
GlobalTensorSelectionStrategy(Strict assignment)
For non-block loaders only :class:KnapsackStrategy is used.
After all candidates run, the one with the lowest overhead that also
satisfies the memory constraint (if any) is selected. The winning
strategy's name is available via :attr:selected_strategy_name.
Example
strategy = AdaptiveStrategy(scale=1.0, n_blocks=4) result = strategy.compute(layer_stats, memory_stats, max_gpu_mem_bytes=50 * 1024**3) print(strategy.selected_strategy_name) 'KnapsackBlock'
Initialize AdaptiveStrategy.
| PARAMETER | DESCRIPTION |
|---|---|
scale
|
Duration multiplier (baseline).
TYPE:
|
loader_type
|
Type of tensor loader. Determines which strategy candidates are instantiated.
TYPE:
|
cyclic
|
Whether to apply cyclic strategy extension (KnapsackStrategy only).
TYPE:
|
group_size
|
Number of layers to merge into groups (Knapsack strategies).
TYPE:
|
threshold_mb
|
Minimum tensor size threshold in MB.
TYPE:
|
min_blocks
|
Minimum number of blocks for OptimizedRoundRobinAssignment search range. Does not affect Strict or Knapsack candidates.
TYPE:
|
n_blocks
|
Number of memory blocks.
TYPE:
|
extra_optimization
|
If True, include slower
:class:
TYPE:
|
property
¶
Name of the strategy that produced the best result.
| RAISES | DESCRIPTION |
|---|---|
RuntimeError
|
If accessed before :meth: |
property
¶
all_scores: list[StrategyScore]
Scores from all candidates evaluated in the last :meth:compute call.
compute(layer_stats: list[LayerStatistics], memory_stats: dict[int, float] | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyResult
Run all strategy candidates and return the best result.
| PARAMETER | DESCRIPTION |
|---|---|
layer_stats
|
List of layer statistics containing tensor info and durations.
TYPE:
|
memory_stats
|
Optional dict mapping tensor size (bytes) to transfer time (ms).
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit in bytes (never exceed).
When
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
StrategyResult
|
class: |
Assignment Strategies¶
Bases: Protocol
Protocol for trap-to-block assignment strategies.
Compute layer-to-block assignment.
| PARAMETER | DESCRIPTION |
|---|---|
layer_sizes
|
Size of each layer in bytes.
TYPE:
|
n_blocks
|
Number of available blocks.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
list[int]
|
List of block indices, one per layer. |
Strict round-robin assignment.
Assigns traps to blocks using strict pattern: trap_i -> block (i % n_blocks). Always uses all available blocks in a fixed cyclic pattern (0,1,2,3,0,1,2,3,...). Does not optimize for memory — all blocks sized to fit the largest trap.
Compute simple cyclic assignment.
OptimizedRoundRobinAssignment(min_blocks: int = 3, max_blocks: int = 4, distribution_weight: float = 0.5)
Optimized round-robin assignment that minimizes memory while satisfying constraints.
Features: - Tries different numbers of blocks (min_blocks to max_blocks) - Uses exhaustive search for small problems, greedy for larger ones - Balances memory efficiency with reuse distance (distribution_weight)
| PARAMETER | DESCRIPTION |
|---|---|
min_blocks
|
Minimum number of blocks to use.
TYPE:
|
max_blocks
|
Maximum number of blocks to use.
TYPE:
|
distribution_weight
|
Weight for reuse distance vs memory (0.0-1.0). 0.0 = pure memory optimization, 1.0 = pure distribution optimization.
TYPE:
|
Compute optimized assignment within block count range.
Transfer Windows¶
Bases: Protocol
Protocol for computing effective transfer windows per trap.
In the pipelined offload model, the transfer for trap i starts during an earlier trap's compute. How much earlier depends on the strategy:
SingleLayerWindowassumes transfer happens during layer i-1 only.GapAwareWindowextends the window backward through consecutive layers that carry no transfer of their own.
Compute the effective transfer window for a given layer.
| PARAMETER | DESCRIPTION |
|---|---|
layer_idx
|
Index of the layer whose transfer window to compute.
TYPE:
|
layer_offload_sizes
|
Per-layer offload sizes in bytes. For pre-optimisation checks this reflects permanent (known) offload; during optimisation it reflects the candidate solution's offload pattern.
TYPE:
|
layer_durations
|
Per-layer compute durations in milliseconds.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
float
|
Effective transfer duration in milliseconds. |
Transfer window limited to the immediately preceding layer's compute duration.
This is the conservative (original) behaviour: the transfer for layer i must complete within layer i-1's compute time.
Return layer_durations[layer_idx - 1].
Transfer window that exploits gap layers for a larger budget.
Sums backward from layer i-1 through consecutive layers whose next layer has no offloaded tensors (i.e. no competing transfer). This captures both permanent gaps (layers with no offloadable tensors at all) and dynamic gaps (layers the optimiser chose not to offload in the current candidate solution).
A layer j is considered "busy" (stops the backward scan) when
layer_offload_sizes[j + 1] > 0 — meaning that during layer j's
compute, the transfer stream is occupied moving layer j+1's data.
Return the sum of durations of consecutive free layers before layer_idx.
Evaluation Utilities¶
Bases: Exception
Raised when a strategy cannot produce a valid result for the given inputs.
This covers recoverable, expected failures such as insufficient layers for
pipelining or missing memory-transfer data. :class:AdaptiveStrategy
catches this exception to skip a failing candidate and try the next one.
Programming errors (TypeError, AttributeError, etc.) should not
be wrapped in this exception — they must propagate so that bugs surface
during development.
dataclass
¶
StrategyScore(strategy_name: str, peak_memory_bytes: int, estimated_overhead: float, consecutive_violations: int, is_valid: bool)
Evaluation metrics for a single strategy result.
| ATTRIBUTE | DESCRIPTION |
|---|---|
strategy_name |
Human-readable name of the strategy that produced the result.
TYPE:
|
peak_memory_bytes |
Estimated peak GPU memory in bytes (blocks + non-offloaded tensors).
TYPE:
|
estimated_overhead |
Fractional overhead from synchronous transfers.
TYPE:
|
consecutive_violations |
Number of adjacent layer pairs assigned to the same block.
TYPE:
|
is_valid |
TYPE:
|
evaluate_strategy_result(result: StrategyResult, layer_stats: list[LayerStatistics], strategy_name: str, interpolator: MemoryTransferInterpolator | None = None, max_gpu_mem_bytes: int | None = None) -> StrategyScore
Score a strategy result for comparison.
| PARAMETER | DESCRIPTION |
|---|---|
result
|
The :class:
TYPE:
|
layer_stats
|
Layer statistics from profiling (used for durations and tensor sizes).
TYPE:
|
strategy_name
|
Display name for this strategy candidate.
TYPE:
|
interpolator
|
Transfer-time interpolator. Required for overhead calculation.
When
TYPE:
|
max_gpu_mem_bytes
|
Hard GPU memory limit. When set, results exceeding this limit are marked invalid.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
A
|
class:
TYPE:
|
strategy_has_transfer_gaps(strategy_map: dict[str, list[TensorStatistics]], layer_stats: list[LayerStatistics]) -> bool
Check if a strategy has unused layer slots where transfer rearrangement would help.
In the pipelined model, each of the first N-1 layers can serve as a
transfer slot (transferring the next layer's data during its compute).
If some of those slots are unused (because the corresponding next layer
has no offloaded tensors), rearrange_transfers can move other
transfers into those gaps for better amortisation.
| PARAMETER | DESCRIPTION |
|---|---|
strategy_map
|
Mapping of layer labels to tensors being transferred during that layer's compute.
TYPE:
|
layer_stats
|
Layer statistics for all layers in the model.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
bool
|
|