https://github.com/kir-gadjello/zipslicer Skip to content Toggle navigation Sign up * Product + Actions Automate any workflow + Packages Host and manage packages + Security Find and fix vulnerabilities + Codespaces Instant dev environments + Copilot Write better code with AI + Code review Manage code changes + Issues Plan and track work + Discussions Collaborate outside of code + Explore + All features + Documentation + GitHub Skills + Blog * Solutions + For + Enterprise + Teams + Startups + Education + By Solution + CI/CD & Automation + DevOps + DevSecOps + Case Studies + Customer Stories + Resources * Open Source + GitHub Sponsors Fund open source developers + The ReadME Project GitHub community articles + Repositories + Topics + Trending + Collections * Pricing [ ] * # In this repository All GitHub | Jump to | * No suggested jump to results * # In this repository All GitHub | Jump to | * # In this user All GitHub | Jump to | * # In this repository All GitHub | Jump to | Sign in Sign up {{ message }} kir-gadjello / zipslicer Public * Notifications * Fork 1 * Star 19 A library for incremental loading of large PyTorch checkpoints License View license 19 stars 1 fork Star Notifications * Code * Issues 0 * Pull requests 0 * Actions * Projects 0 * Security * Insights More * Code * Issues * Pull requests * Actions * Projects * Security * Insights kir-gadjello/zipslicer This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository. main Switch branches/tags [ ] Branches Tags Could not load branches Nothing to show {{ refName }} default View all branches Could not load tags Nothing to show {{ refName }} default View all tags Name already in use A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch? Cancel Create 1 branch 2 tags Code * Local * Codespaces * Clone HTTPS GitHub CLI [https://github.com/k] Use Git or checkout with SVN using the web URL. [gh repo clone kir-ga] Work fast with our official CLI. Learn more. * Open with GitHub Desktop * Download ZIP Sign In Required Please sign in to use Codespaces. Launching GitHub Desktop If nothing happens, download GitHub Desktop and try again. Launching GitHub Desktop If nothing happens, download GitHub Desktop and try again. Launching Xcode If nothing happens, download Xcode and try again. Launching Visual Studio Code Your codespace will open once ready. There was a problem preparing your codespace, please try again. Latest commit @kir-gadjello kir-gadjello fix setup.py ... 3d32941 Mar 3, 2023 fix setup.py 3d32941 Git stats * 19 commits Files Permalink Failed to load latest commit information. Type Name Latest commit message Commit time .github/workflows CI update February 6, 2023 06:44 examples update CI January 30, 2023 16:45 tests add support for _extra_state in state_dict January 31, 2023 00:26 zipslicer more careful usage of pytorch._utils February 5, 2023 05:43 .gitignore enable handling for tensors of size 1, add example, add readme, cleanup January 30, 2023 03:41 LICENSE Initial commit January 29, 2023 22:07 MANIFEST.in preparing for first pypi release February 6, 2023 06:01 README.md preparing for first pypi release February 6, 2023 06:01 requirements.txt clarify README, update requirements January 30, 2023 19:31 setup.py fix setup.py March 4, 2023 00:42 View code ZIPSLICER [?] Synopsis Prerequisites Install Notes Prospective features we are considering README.md ZIPSLICER [?] Lint and Test Python package Published to PyPI A library for incremental loading of large PyTorch checkpoints Read a blogpost introduction by yours truly Synopsis import torch import zipslicer # Could be a private custom recurrent sentient transformer # instead of a garden variety resnet my_complicated_network = torch.hub.load( "pytorch/vision:v0.10.0", "resnet18", pretrained=True ) s_dict = my_complicated_network.state_dict() torch.save(s_dict, "my_network_checkpoint_v123.pth") del my_complicated_network # Later, on a smaller unrelated machine you load a "LazyStateDict" # Which is just like a regular state dict, but it loads tensors only when it has to lazy_s_dict = zipslicer.load("my_network_checkpoint_v123.pth") layer3_tensors = {} for k in lazy_s_dict.keys(): if k.startswith("layer3"): layer3_tensors[k] = lazy_s_dict[k] # Now you have layer3's tensors and you can analyze them without breaking your RAM. # Or you can instantiate the layers' classes in sequence and compute the whole # network's output for a given input by threading the activations through them. # But we will just print the tensors instead: print(layer3_tensors) Run this example and unit-tests: python examples/example_resnet18.py pytest -o log_cli=true --capture=tee-sys -p no:asyncio Test your checkpoint for compatibility: python tests/test_checkpoint_readonly.py your_magnificent_checkpoint.pth If it's all green, it will work. Prerequisites * Supported python and torch versions: python-3.10 + torch- (1.11,1.12,stable) python-3.11 + torch:stable * Generally, zipslicer should work with modern enough install of PyTorch - use included safe test to check for compatibility of zipslicer with your PyTorch and your checkpoint. This is a pure Python library, so specific CPU architecture shouldn't matter. * A checkpoint produced by saving your model's state_dict via vanilla torch.save(...) - default settings should suffice, as Torch doesn't use ZIP compression. * An application that can take advantage of incrementally-loaded checkpoint - i.e. if your app just loads all state_dict.items() in a loop right away it doesn't make much sense to use this library. Make sure your code reads state_dict.keys() (and state_dict.get_meta(k) if necessary) and uses these intelligently to work on a subset of state_dict[k] tensors at a time. For general inspiration you might read this (HF) and this (arxiv). With some additional engineering it should be possible to run Large Language Models like BLOOM-176B or FLAN-T5-XXL on a single mid-range GPU at home - if you are willing to wait for a night's worth of time. In the large batch regime this might even make some practical sense, for example to process a set of documents into embeddings. Install Generally, copying the zipslicer/zipslicer directory into your project's source tree is enough. If you are a fan of official ceremony-driven install processes for executable modules of dubious provenance, soon there will be a possibility of installing this boutique software module via pip: pip install zipslicer Notes * This library is only for reading pytorch tensors from checkpoints. We leave writing for future work. * Writing to loaded state_dict is frowned upon, but it will work - though you should avoid doing this while iterating over keys for now and expecting the keys to reflect this update. * Perhaps more importantly, general-purpose pickles are not supported - the design of this library doesn't allow you to load whole neural network class instances. Usually this isn't necessary, and pytorch official documentation recommends you to use state_dict for model serialization. We support state_dict's. * Some rare tensor types (i.e: pytorch quantized tensors - not to be confused with integer tensors which work fine) are not yet supported. If this bothers you, share your experience in issues. * We say "Hi" to HF safetensors project, but note that in comparison to theirs, our approach doesn't require checkpoint conversion which takes significant time and storage. In fact, both approaches could be complementary, as you will have to load tensors from the pytorch checkpoint somehow to convert it to safetensors - and the default loading mechanism is constrained by available RAM. Prospective features we are considering If you are interested in some of these features, consider creating an issue: * Effective loading of tensor slices - to implement tensor parallelism in sharded deployments * Accessing the source checkpoint over a network * Writing to a checkpoint in-place * Incremental conversion to other checkpoint formats About A library for incremental loading of large PyTorch checkpoints Topics deep-learning zip torch pytorch checkpoint Resources Readme License View license Stars 19 stars Watchers 1 watching Forks 1 fork Releases 2 0.8.1-alpha Latest Mar 3, 2023 + 1 release Packages 0 No packages published Languages * Python 100.0% Footer (c) 2023 GitHub, Inc. Footer navigation * Terms * Privacy * Security * Status * Docs * Contact GitHub * Pricing * API * Training * Blog * About You can't perform that action at this time. You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session.