diff --git a/nightly/_modules/index.html b/nightly/_modules/index.html index 8c3fc97f5..232dbc8e8 100644 --- a/nightly/_modules/index.html +++ b/nightly/_modules/index.html @@ -265,7 +265,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/_modules/torch/_functorch/aot_autograd.html b/nightly/_modules/torch/_functorch/aot_autograd.html index 17e32539a..6dd8ecd60 100644 --- a/nightly/_modules/torch/_functorch/aot_autograd.html +++ b/nightly/_modules/torch/_functorch/aot_autograd.html @@ -265,7 +265,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/_modules/torch/_functorch/compilers.html b/nightly/_modules/torch/_functorch/compilers.html index 285778c94..ca0533d85 100644 --- a/nightly/_modules/torch/_functorch/compilers.html +++ b/nightly/_modules/torch/_functorch/compilers.html @@ -265,7 +265,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/_modules/torch/_functorch/deprecated.html b/nightly/_modules/torch/_functorch/deprecated.html index d4f8c6bf9..5ef04df32 100644 --- a/nightly/_modules/torch/_functorch/deprecated.html +++ b/nightly/_modules/torch/_functorch/deprecated.html @@ -265,7 +265,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/_modules/torch/_functorch/partitioners.html b/nightly/_modules/torch/_functorch/partitioners.html index 7a6c2f618..1172a6d82 100644 --- a/nightly/_modules/torch/_functorch/partitioners.html +++ b/nightly/_modules/torch/_functorch/partitioners.html @@ -265,7 +265,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
@@ -414,6 +414,11 @@

Source code for torch._functorch.partitioners

from torch.utils.checkpoint import CheckpointPolicy from . import config +from ._activation_checkpointing.knapsack import ( + dp_knapsack, + greedy_knapsack, + ilp_knapsack, +) from ._aot_autograd.logging_utils import get_aot_graph_name from ._aot_autograd.utils import is_with_effects from .compile_utils import fx_graph_cse, get_aten_target @@ -1770,126 +1775,6 @@

Source code for torch._functorch.partitioners

return name_to_node -def greedy_knapsack( - memory: List[float], runtimes: List[float], max_memory: float -) -> Tuple[float, List[int], List[int]]: - n = len(runtimes) - items = list(range(n)) - - # Sort items based on the ratio of runtime to memory in descending order - items = sorted(items, key=lambda i: runtimes[i] / memory[i], reverse=True) - - total_memory = 0.0 - total_runtime = 0.0 - items_to_save = [] - items_to_allow_recomputing = [] - - for i in items: - if total_memory + memory[i] <= max_memory: - total_memory += memory[i] - total_runtime += runtimes[i] - items_to_save.append(i) - else: - items_to_allow_recomputing.append(i) - return total_runtime, items_to_save, items_to_allow_recomputing - - -def ilp_knapsack( - memory: List[float], runtimes: List[float], max_memory: float -) -> Tuple[float, List[int], List[int]]: - import numpy as np - - try: - from scipy.optimize import Bounds, LinearConstraint, milp - except ImportError: - raise RuntimeError( - "To use the ILP for memory budget checkpointing you need to install scipy" - ) from None - - np_memory = np.array(memory) - np_runtimes = np.array(runtimes) - c = -np_runtimes # type: ignore[operator] - - memory_constraint = LinearConstraint(A=np_memory, ub=np.array(max_memory)) - constraints = [memory_constraint] - - integrality = np.ones_like(c) - res = milp( - c=c, constraints=constraints, integrality=integrality, bounds=Bounds(0, 1) - ) - if not res.success: - raise RuntimeError("Somehow scipy solving failed") - - items_to_save = [] - items_to_allow_recomputing = [] - for idx, i in enumerate(res.x): - if i == 1: - items_to_save.append(idx) - else: - items_to_allow_recomputing.append(idx) - return -res.fun, items_to_save, items_to_allow_recomputing - - -def dp_knapsack( - memory: List[float], runtimes: List[float], max_memory: float -) -> Tuple[float, List[int], List[int]]: - # Scaling factor to convert floating point weights to integers - S = 10000 - - # Quantize the memory weights - quantized_memory = torch.tensor( - [int(round(m * S)) for m in memory], dtype=torch.long, device="cpu" - ) - runtimes = torch.tensor(runtimes, dtype=torch.float32, device="cpu") - - # Quantized pseudopolynomial DP for 0-1 Knapsack - quantized_max_memory = int(round(max_memory * S)) - - n = len(memory) - - # Initialize the DP table - # TODO(chilli): I think if needed, this memory can be optimized with sliding - # window trick + Hirschberg trick: - # https://codeforces.com/blog/entry/47247?#comment-316200 - dp = torch.zeros( - (n + 1, quantized_max_memory + 1), dtype=torch.float32, device="cpu" - ) - - for i in range(1, n + 1): - current_memory = quantized_memory[i - 1] - current_runtime = runtimes[i - 1] - - # Copy the previous row - dp[i, :] = dp[i - 1, :] - - # Update dp[i, j] for all j >= current_memory - if current_memory == 0: - dp[i, :] = dp[i - 1, :] + current_runtime - else: - dp[i, current_memory:] = torch.maximum( - dp[i - 1, current_memory:], - dp[i - 1, :-current_memory] + current_runtime, - ) - - # Backtrack to find the items included in the knapsack - saved_items = [] - recomputable_items = [] - j: int = quantized_max_memory - for i in range(n, 0, -1): - if dp[i][j] != dp[i - 1][j]: - saved_items.append(i - 1) # Include this item (indexing from 0) - j -= int(quantized_memory[i - 1].item()) - else: - recomputable_items.append(i - 1) - - saved_items.reverse() # To get items in the order they were added - - # The maximum runtime that can be achieved within the max_memory constraint - max_runtime = dp[n][quantized_max_memory].item() - - return max_runtime, saved_items, recomputable_items - - def _optimize_runtime_with_given_memory( joint_graph: fx.Graph, memory: List[float], diff --git a/nightly/aot_autograd.html b/nightly/aot_autograd.html index e28e8657f..24eacdfbf 100644 --- a/nightly/aot_autograd.html +++ b/nightly/aot_autograd.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/batch_norm.html b/nightly/batch_norm.html index 9d9aa488a..21acc9333 100644 --- a/nightly/batch_norm.html +++ b/nightly/batch_norm.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/experimental.html b/nightly/experimental.html index 2324ad8db..e124cebad 100644 --- a/nightly/experimental.html +++ b/nightly/experimental.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/functorch.html b/nightly/functorch.html index 75a39e975..0f993ae36 100644 --- a/nightly/functorch.html +++ b/nightly/functorch.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.combine_state_for_ensemble.html b/nightly/generated/functorch.combine_state_for_ensemble.html index ebfa49885..222e57515 100644 --- a/nightly/generated/functorch.combine_state_for_ensemble.html +++ b/nightly/generated/functorch.combine_state_for_ensemble.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.compile.aot_function.html b/nightly/generated/functorch.compile.aot_function.html index 3a266e10e..ad703a822 100644 --- a/nightly/generated/functorch.compile.aot_function.html +++ b/nightly/generated/functorch.compile.aot_function.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.compile.aot_module.html b/nightly/generated/functorch.compile.aot_module.html index 5f62ce316..7d3d21ff1 100644 --- a/nightly/generated/functorch.compile.aot_module.html +++ b/nightly/generated/functorch.compile.aot_module.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.compile.default_partition.html b/nightly/generated/functorch.compile.default_partition.html index 249051c03..05b23a129 100644 --- a/nightly/generated/functorch.compile.default_partition.html +++ b/nightly/generated/functorch.compile.default_partition.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.compile.memory_efficient_fusion.html b/nightly/generated/functorch.compile.memory_efficient_fusion.html index d64926fb8..518a7adaf 100644 --- a/nightly/generated/functorch.compile.memory_efficient_fusion.html +++ b/nightly/generated/functorch.compile.memory_efficient_fusion.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.compile.min_cut_rematerialization_partition.html b/nightly/generated/functorch.compile.min_cut_rematerialization_partition.html index 4e6e2715d..b57ca42d7 100644 --- a/nightly/generated/functorch.compile.min_cut_rematerialization_partition.html +++ b/nightly/generated/functorch.compile.min_cut_rematerialization_partition.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.compile.nop.html b/nightly/generated/functorch.compile.nop.html index 4a6dc2de3..52d7331f0 100644 --- a/nightly/generated/functorch.compile.nop.html +++ b/nightly/generated/functorch.compile.nop.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.compile.ts_compile.html b/nightly/generated/functorch.compile.ts_compile.html index 452f2844f..0cf5ce136 100644 --- a/nightly/generated/functorch.compile.ts_compile.html +++ b/nightly/generated/functorch.compile.ts_compile.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.functionalize.html b/nightly/generated/functorch.functionalize.html index 1445f21cb..2419fef5d 100644 --- a/nightly/generated/functorch.functionalize.html +++ b/nightly/generated/functorch.functionalize.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.grad.html b/nightly/generated/functorch.grad.html index f3e8e4c85..5551cf535 100644 --- a/nightly/generated/functorch.grad.html +++ b/nightly/generated/functorch.grad.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.grad_and_value.html b/nightly/generated/functorch.grad_and_value.html index 925ba130d..35dff4885 100644 --- a/nightly/generated/functorch.grad_and_value.html +++ b/nightly/generated/functorch.grad_and_value.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.hessian.html b/nightly/generated/functorch.hessian.html index ec0a1a87d..7be47f3d8 100644 --- a/nightly/generated/functorch.hessian.html +++ b/nightly/generated/functorch.hessian.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.jacfwd.html b/nightly/generated/functorch.jacfwd.html index 26ba150b4..62140839f 100644 --- a/nightly/generated/functorch.jacfwd.html +++ b/nightly/generated/functorch.jacfwd.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.jacrev.html b/nightly/generated/functorch.jacrev.html index f063b559e..46c86c215 100644 --- a/nightly/generated/functorch.jacrev.html +++ b/nightly/generated/functorch.jacrev.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.jvp.html b/nightly/generated/functorch.jvp.html index 977522ea5..39fd8cd5a 100644 --- a/nightly/generated/functorch.jvp.html +++ b/nightly/generated/functorch.jvp.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.make_functional.html b/nightly/generated/functorch.make_functional.html index a5a735626..fe4c7c3bb 100644 --- a/nightly/generated/functorch.make_functional.html +++ b/nightly/generated/functorch.make_functional.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.make_functional_with_buffers.html b/nightly/generated/functorch.make_functional_with_buffers.html index 84e8a8da1..aec985b56 100644 --- a/nightly/generated/functorch.make_functional_with_buffers.html +++ b/nightly/generated/functorch.make_functional_with_buffers.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.vjp.html b/nightly/generated/functorch.vjp.html index a65d61bf3..1810ebd2f 100644 --- a/nightly/generated/functorch.vjp.html +++ b/nightly/generated/functorch.vjp.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/generated/functorch.vmap.html b/nightly/generated/functorch.vmap.html index 1275c1599..b39201e55 100644 --- a/nightly/generated/functorch.vmap.html +++ b/nightly/generated/functorch.vmap.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/genindex.html b/nightly/genindex.html index f62975f13..8b895cbf5 100644 --- a/nightly/genindex.html +++ b/nightly/genindex.html @@ -265,7 +265,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/index.html b/nightly/index.html index ee44f52e2..dbc138aff 100644 --- a/nightly/index.html +++ b/nightly/index.html @@ -266,7 +266,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/install.html b/nightly/install.html index 299b98444..faaedd05f 100644 --- a/nightly/install.html +++ b/nightly/install.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/notebooks/aot_autograd_optimizations.html b/nightly/notebooks/aot_autograd_optimizations.html index d8f5b4b5a..4f11f4288 100644 --- a/nightly/notebooks/aot_autograd_optimizations.html +++ b/nightly/notebooks/aot_autograd_optimizations.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/notebooks/ensembling.html b/nightly/notebooks/ensembling.html index b9572f524..62ea4bb57 100644 --- a/nightly/notebooks/ensembling.html +++ b/nightly/notebooks/ensembling.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/notebooks/jacobians_hessians.html b/nightly/notebooks/jacobians_hessians.html index 1e22d7950..021651078 100644 --- a/nightly/notebooks/jacobians_hessians.html +++ b/nightly/notebooks/jacobians_hessians.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/notebooks/minifier.html b/nightly/notebooks/minifier.html index e378365b2..fb8887ced 100644 --- a/nightly/notebooks/minifier.html +++ b/nightly/notebooks/minifier.html @@ -266,7 +266,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/notebooks/neural_tangent_kernels.html b/nightly/notebooks/neural_tangent_kernels.html index 703e2cb71..0340b1323 100644 --- a/nightly/notebooks/neural_tangent_kernels.html +++ b/nightly/notebooks/neural_tangent_kernels.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/notebooks/per_sample_grads.html b/nightly/notebooks/per_sample_grads.html index d0e0c625e..1310b3f12 100644 --- a/nightly/notebooks/per_sample_grads.html +++ b/nightly/notebooks/per_sample_grads.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/notebooks/whirlwind_tour.html b/nightly/notebooks/whirlwind_tour.html index e0061d281..affff9154 100644 --- a/nightly/notebooks/whirlwind_tour.html +++ b/nightly/notebooks/whirlwind_tour.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/objects.inv b/nightly/objects.inv index bc9964e0f..d1b67baf7 100644 Binary files a/nightly/objects.inv and b/nightly/objects.inv differ diff --git a/nightly/search.html b/nightly/search.html index 6bba201ac..4d59a443b 100644 --- a/nightly/search.html +++ b/nightly/search.html @@ -265,7 +265,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼
diff --git a/nightly/ux_limitations.html b/nightly/ux_limitations.html index 41e3bccf5..cc57e4348 100644 --- a/nightly/ux_limitations.html +++ b/nightly/ux_limitations.html @@ -267,7 +267,7 @@
- nightly (2.6.0a0+git0052943) ▼ + nightly (2.6.0a0+gitb1bb860) ▼