![\"\"](\"https://jedrzej.maczan.pl/knapsack.png\")
TL;DR: If you use a default\ntorch._functorch.config.activation_memory_budget_solver=dp, you can change it to dp_knapsack_sliding_hirschberg\nand decrease your peak RAM usage during memory planning by a factor of 20.
\n
When PyTorch builds a computation graph of your ML model, it tries many different things to make it fast and\nmemory efficient. One technique is to take a joint forward and backward computation graph and choose which\noperations should be saved (stored in a memory and reused) and which operations should be recomputed.\nEach operation takes some time to run and takes up some memory if stored and we only can use a limited amount\nof memory.
\nImagine you are a thief and you just broke into a place with many valuable things inside. You have a knapsack\nthat can carry up to 10 kilograms. You want to maximize the value of things you steal. Ideally, you'd just take\nall of them. But if you put more than 10kg to your knaspack, it will break and you will get away with nothing.\nAlternatively, you can use some heuristics to pick the items (like take only small items or only green ones),\nbut then your solution\nmight not be optimal - you\nwill steal some\nvalue, but perhaps less than what you could gain in an optimal choice. For your rescue, there come 0/1\nknapsack solvers. There are different algorithms that compute optimal items choice within a given capacity (max\nweight of things in a knapsack) and given a list of items, where each of items has a value and a cost (weight).\nIn PyTorch, this is what we actually do, which is: within a RAM budget (knapsack capacity), we choose which\noperations (items) should be saved to\nthe memory to maximize the time savings in runtime (value).
\n![\"\"](\"https://jedrzej.maczan.pl/thief.png\")
Currently, the default implementation of knapsack solver in PyTorch is dp_knapsack, which is a dynamic\nprogramming approach to\nsolving 0/1 knapsack. The algoritmh is actually simple once you understand it. It helps a lot if you have had a\nprior experience with dynamic\nprogramming (I do not). Look at this nice GIF from Wikipedia:
\n![\"\"](\"https://jedrzej.maczan.pl/knapsack_dp.gif\")
By RDSEED - Own work,\nCC\nBY-SA 4.0, Link
\nOn the left, there is a list of 4 items. i is an index, v is a value and w is a weight. Units irrelevant.\nCapacity is 6. That's why on the right side there are 6 columns (each column is a part of total capacity)\nand 4 rows (one per item). Ok, actually there are 7 and 5, because of padding with zeros, but that's only\nfor a computation convenience (every row and column has a parent, index arithmetic gets much easier). You could\nactually do it without padding, but I tried and fallbacked to padding anyways, because index arithmetic was a\nbit annoying. The table on the\nright is called a DP table. It's where the computation happens.
\nDP 0/1 Knapsack algorithm goes like this: get an item. In a DP table, go through a corresponding row, column\nby\ncolumn and compare your item's value to value in current DP table cell (row, column; current column index is\ncurrent max weight\n(capacity) of your knapsack). If your item's weight is bigger than current column's index, move on to the\nnext column. If your\nitem weight is equal to a column index and your item value is higher than current cell, replace the cell\nwith your item value. If your item weight is smaller than column index and it fits in a spare capacity,\nadd\nit to the cell's value (a free capacity is current capacity minus capacity of picked items). Repeat until you\niterate\nthrough all the items. At this point, you learn what is the maximum value you can get within the weight\nbudget. If our goal would be to measure a max possible value, then we would stop at this point. Unfortunately,\nwe\nwant to know more.
\n![\"\"](\"https://jedrzej.maczan.pl/knapsack_no_backtrack.png\")
We need to know which items we should take to get this max possible value. So in case of PyTorch, which\noperations which we should save to memory and which operations we should recompute. To find out which items\ncontribute to the max value (in other words, which items were chosen during building DP table), we need to go\nbackwards through the DP table. We go from the last row and last column upward. We compare a cell with a cell\nabove (same column, previous row). If they are different, then the current item\ncontributes to the max value - because each row corresponds to a single item. If they are equal, this\nitem\nwill not be stored in a memory and becomes recomputable. Now, we traverse the DP table upward. If the item was\nchosen, we move upward by a single row and we\nmove the position of\nthe current cell to the left by the\nweight of the chosen item. If item wasn't chosen, we move upward by a single row and we start on the\nsame position as previously. Repeat until we get to the top of the DP table.
\n![\"\"](\"https://jedrzej.maczan.pl/knapsack_backtrack.png\")
This solver works correctly. It provides an exact solution and for most of the time you don't need anything\nbetter.\nBy better I mean 1. faster and 2. less memory hungry. Because look again at the algorithm and the GIF above.\nThere are\na few issues with it. To get a result, we allocate a full 2D table of shape=(number of items, max_weight), which\nis\nnot great. Horace He (Chillee) who\noriginally built a great part of (all of?) the memory planning in PyTorch left a note in the DP knapsack solver code:
\n# TODO(chilli): I think if needed, this memory can be optimized with sliding\n# window trick + Hirschberg trick:\n# https://codeforces.com/blog/entry/47247?#comment-316200What he proposes are two optimizations.
\nWindow trick means that instead of building the a DP table, we slide\nover a table and use only a previous row and a current row. When moving to the next row, next becomes a new\ncurrent and old current becomes a previous row. This improvement alone would be enough, if we wanted to just\ncompute the max value of items out\nof\nour list, within a given capacity. But again, we need to know which items contribute to this max value. That's\nwhy Horace suggested using a sliding window together with Hirschberg trick.
\n![\"\"](\"https://jedrzej.maczan.pl/sliding_window.png\")
It's a divide and conquer approach, kinda similar to quicksort, because\nyou split problem into smaller\nones and solve them and continue splitting until you solve the whole problem. The main benefit of Hirschberg\ntrick for usis that it gets rid of backtracking.
\nFrom an implementation perspective, I sort of do a recursion here, but without putting new frames on the stack.\nI have a Python list that I use as an explicit stack (last in, first out). The list stores a range of item\nindexes and a currently available capacity. We start with all the items and all the capacity. Then the algo goes\nlike this: First, we pop an item from the stack. If\nit's a single item (item index range end - start = 1), check if item's weight fits in the\navailable capacity. If\nyes, include it in saved items, if\nnot, include it in recomputable items. This part is imo the biggest strength of Hirschberg algorithm for us -\nbecause you can see now that we won't need a backtracking, since we build saved and recomputable lists inside\nthe algorithm. No need to build it afterwards, unlike in original 0/1 DP knapsack algorithm. If the items range\nis\nbigger than 1, split the items into a two parts and for\neach of\nthem compute a DP table. We don't want to build a full DP table, so we do the sliding window trick here and\nreuse a current and previous DP row to compute the final DP row.
\nOnce we have both left and right DP rows, add\nthem together. Find the position of the highest value. It will be the new splitting point for our memory. Now we\nsplit the\nmemory\ninto two, using this newly computed position and we run the algorithm again, using both memory and items split.\nRepeat until done.\nBtw. It might\nhappen that one of the splitted parts will be empty. It might happen when the best split is either the first or\nlast item. That's completely legit, we just skip\nsuch an stack frame and move on.
\n![\"\"](\"https://jedrzej.maczan.pl/hirschberg_step_by_step.png\")
\nHere is a great explanation with drawings\nstep by step of Hirschberg algo by KokiYmgch on CodeForces
No. Look - with these two optimizations combined, we can decrease the DP table size from shape=(number of items,\nmax weight) to shape=(2, max weight),\neffectively\nreducing a peak memory by a factor of 20x - on my machine with 64GB RAM dp_knapsack crashes on problems with 100\nitems, while dp_knapsack_sliding_hirschberg is able to handle problems with 2000 items\nwithin this RAM budget. We also gain some\nruntime speedup, roughly 37%, but I didn't measure it very rigorously, though.\nThat being said - if you don't mind\nusing SciPy, then you probably should use ilp_knapsack instead. It's MUCH faster than any existing DP\nknapsack implementation in PyTorch. It's a bit slower than a greedy algorithm, but still very fast, yet unlike\ngreedy and similarly to both DP solvers, it provides an exact solution.
\nIf you care about speed above correctness, use greedy_knapsack.
\nElse if you're ok with pulling SciPy into project, use ilp_knapsack.
\nElse if you install PyTorch from a package manager, noop - PyTorch will use dp_knapsack, which is a\ndefault solver and you don't need to do anything.
\nElse use dp_knapsack_sliding_hirschberg.
\n![\"\"](\"https://jedrzej.maczan.pl/pick_solver.png\")
import torch\nimport torch._functorch.config as fconfig\nfconfig.activation_memory_budget_solver = fconfig.dp_knapsack_sliding_hirschbergIf you spot any bugs when using dp_knapsack_sliding_hirschberg, pls register an issue and tag me @jmaczan and I'll hop on to solve it.
\nThanks to Edward Yang for a code review and a merge and\nthanks to you for reading
","headings":[]}}