4 ms·

PyTorch saves tensor data and metadata separately (metadata in Pickle format a.k.a. pickle VM, which is famously not mmap friendly because it builds up Python o

by liuliu 4y ago

PyTorch saves tensor data and metadata separately (metadata in Pickle format a.k.a. pickle VM, which is famously not mmap friendly because it builds up Python object through the VM execution). These then zipped together in a zip file. Whether you can mmap depends on two things: if the zip file has compression enabled, you are out of luck. If not, CUDA's unified memory requires I think 64-byte alignment, and that may at odds with the zip archive, otherwise CPU depending on optimization may require some other alignment too (typically 16 bytes or 4 bytes).

jart 4y ago
The same was true with llama.cpp. The trick I used was basically to identify which part of the program is performing the transformation, where the unwieldy file from disk is being turned into the aligned data structures that actually get used. In the llama.cpp codebase, that function was called llama_model_load(). So I basically just wrapped that code in a magic transaction: main() { magic_init(); // replaces malloc() with ./magic.dat file if (magic->was_committed) { model = magic->model; // use the heap built from last run } else { // create a new heap model = new model; llama_model_load(model); // pre-existing model loading code magic->model = model; magic_commit(); // mark memory transaction complete } // rest of pre-existing code } You could write a Python extension for this code https://github.com/ggerganov/llama.cpp/commit/5b8023d935401072b73b63ea995aaae040d57b87 https://github.com/ggerganov/llama.cpp/commit/5b8023d9354010... and see if it works for your use case.
- beagle3 4y ago
  That would work well only if there are no pointers in the load result, which sometimes (if rarely) happens without planning in C or C++ code, but almost never -- or actually never -- happens with Python code. Or am I missing something?
  - jart 4y ago
    You can have all the pointers you want. The way my code works is it uses MAP_FIXED so that memory is always allocated to the same addresses. Even if you have ASLR it will be the same every time. The issue is mainly stack-allocated memory. For example, our `model` object here needed to be allocated on the heap, rather than being declared as an automatic variable. Since automatic objects is really an optimization unique to low-level languages like C++, I doubt that's going to be an issue for something like Python, where everything is probably a heap object. Oh! Also be careful about .data objects (i.e. static variables). Those could prove problematic too. Although my code could easily be updated to memcpy() the memory between _etext and _end into the file.
    - beagle3 4y ago
      Ah! That’s brilliant. Thanks!

by liuliu 4y ago

jart 4y ago

The same was true with llama.cpp. The trick I used was basically to identify which part of the program is performing the transformation, where the unwieldy file from disk is being turned into the aligned data structures that actually get used. In the llama.cpp codebase, that function was called llama_model_load(). So I basically just wrapped that code in a magic transaction: main() { magic_init(); // replaces malloc() with ./magic.dat file if (magic->was_committed) { model = magic->model; // use the heap built from last run } else { // create a new heap model = new model; llama_model_load(model); // pre-existing model loading code magic->model = model; magic_commit(); // mark memory transaction complete } // rest of pre-existing code } You could write a Python extension for this code https://github.com/ggerganov/llama.cpp/commit/5b8023d935401072b73b63ea995aaae040d57b87 https://github.com/ggerganov/llama.cpp/commit/5b8023d9354010... and see if it works for your use case.

beagle3 4y ago
That would work well only if there are no pointers in the load result, which sometimes (if rarely) happens without planning in C or C++ code, but almost never -- or actually never -- happens with Python code. Or am I missing something?
- jart 4y ago
  You can have all the pointers you want. The way my code works is it uses MAP_FIXED so that memory is always allocated to the same addresses. Even if you have ASLR it will be the same every time. The issue is mainly stack-allocated memory. For example, our `model` object here needed to be allocated on the heap, rather than being declared as an automatic variable. Since automatic objects is really an optimization unique to low-level languages like C++, I doubt that's going to be an issue for something like Python, where everything is probably a heap object. Oh! Also be careful about .data objects (i.e. static variables). Those could prove problematic too. Although my code could easily be updated to memcpy() the memory between _etext and _end into the file.
  - beagle3 4y ago
    Ah! That’s brilliant. Thanks!