Basil Wong 7b2af25f80 [1/n] Support Dynamic Memory Budget in Auto AC (#143539) ... # Summary: Full Context: https://docs.google.com/document/d/1-j5KSbfGFJQcH4sYh7BIeJXso3zYzl5G5yFQqXdKx_o/edit?usp=sharing tl;dr This change introduces classes which help determine a dynamic memory budget. This will mostly be helpful for models with many implicit graph breaks. --- New Classes: *GraphInfoProvider* * Takes the joint_graph as well as the input memories and runtimes and parses the graph + values into usable forms for the SolverEvaluator. *KnapsackEvaluator* * Provides a function: Given all of the four inputs (solver function as a callable, max_dynamic_memory_budget, min_dynamic_memory_budget, dynamic_memory_budget_pareto_granularity) it returns an approximation of the knee point of the pareto distribution. # Test Plan: ### LintRunner LintRunner Output: P1700445547 ### Unit Tests ``` $ buck test @mode/opt //caffe2/test/functorch:test_ac_knapsack `@mode/opt` was specified, but not found. Using file at `//mode/opt`. This behavior is being deprecated. Please use `"@//mode/opt"` instead File changed: fbcode//caffe2/.ruff_cache/0.7.4/.tmpB6PmDS File changed: fbsource//xplat/caffe2/test/functorch/test_ac_knapsack.py File changed: fbcode//caffe2/.ruff_cache/0.7.4/.tmpyjCiPn 20 additional file change events Buck UI: https://www.internalfb.com/buck2/414ead46-9ede-4192-8e1a-5d3c52bdb9cc Test UI: https://www.internalfb.com/intern/testinfra/testrun/6473924710342830 Network: Up: 0B Down: 0B (reSessionID-159794b9-9d61-477e-8e63-9bdeaa537dca) Analyzing targets. Remaining 0/214 Executing actions. Remaining 0/6933 0.1s exec time total Command: test. Finished 1 local Time elapsed: 18.5s Tests finished: Pass 15. Fail 0. Fatal 0. Skip 0. Build failure 0 ``` ### Test Run Updated the config: ``` activation_memory_budget_solver: DYNAMIC_MEMORY_BUDGET_DP ``` Confirming proper execution via: [aps-fb_fm_v4_768_01_dynamic-2a792ba8af](https://www.internalfb.com/mlhub/pipelines/runs/mast/aps-fb_fm_v4_768_01_dynamic-2a792ba8af?job_attempt=0&version=0&env=PRODUCTION) Pull Request resolved: https://github.com/pytorch/pytorch/pull/143539 Approved by: https://github.com/jansel		2024-12-21 07:38:52 +00:00
..
_activation_checkpointing	[1/n] Support Dynamic Memory Budget in Auto AC (#143539)	2024-12-21 07:38:52 +00:00
_aot_autograd	Remove unused Python variables in torch/[_-a]* (#133492)	2024-12-12 17:39:14 +00:00
__init__.py
aot_autograd.py	[easy] Set feature use for aot autograd remote cache (#143674)	2024-12-21 01:40:18 +00:00
apis.py
autograd_function.py
batch_norm_replacement.py	remove allow-untyped-defs for torch/_functorch/batch_norm_replacement.py (#143438)	2024-12-18 09:01:06 +00:00
benchmark_utils.py
compile_utils.py
compilers.py	Remove unused Python variables in torch/[_-a]* (#133492)	2024-12-12 17:39:14 +00:00
config.py	Turn on AOTAutogradCache by default on open source (#141981)	2024-12-12 04:21:11 +00:00
deprecated.py
eager_transforms.py	Remove unused Python variables in torch/[_-a]* (#133492)	2024-12-12 17:39:14 +00:00
functional_call.py	[torch.func] expand stack_module_state's typing (#142169)	2024-12-06 01:08:53 +00:00
fx_minifier.py
make_functional.py
partitioners.py	[1/n] Support Dynamic Memory Budget in Auto AC (#143539)	2024-12-21 07:38:52 +00:00
pyfunctorch.py
python_key.py
pytree_hacks.py
top_operators_github_usage.py
utils.py
vmap.py