Package: git-theta
Git-ThetaGit-Theta is a Git extension for collaborative, continual, and communal development of machine learning models.Version control systems like Git enable large distributed teams to collaborate on shared codebases by tracking changes over time and providing tools for merging changes from multiple sources.
Git-Theta is a Git extension that aims to provide similar functionality for machine learning model checkpoints byefficientlyandmeaningfullytrack a model's version history natively through Git.
Specifically, rather than treating the checkpoint as a blob of data (as done by other systems for tracking models with Git), Git-Thetaatomically tracks each parameter "group" (e.g. a weight matrix or bias vector in a neural network)tracks dense or communication-efficient updates likelow-rankorsparsechanges to parameter groupsallows models to be merged automatically or manuallydisplays meaningful "diffs" by showing which parameter groups have changedsupports checkpoint formats from most popular machine learning frameworksenables easy extension of update types, merging methods, and checkpoint formats through a plugin systemGit-Theta is currently under active development and should be used with caution.
For feature discussions and debugging help, please join the #git-theta stream in theCCCML Zulip community.
If you use Git-Theta as part of a published research project, please citeour paper.Quick StartInstalling Git LFSDownload and install Git LFS using the instructions fromthe Git LFS website.Installing Git-ThetaInstall the git-theta Python package:pipinstallgit-thetaBy default, installinggit-thetawithpipwill not install any of the supported machine learning frameworks (PyTorch, TensorFlow, etc.).
If you want to install the framework you intend to use when installinggit-theta, you can specify it when installing (e.g. by runningpip install git-theta[pytorch]for PyTorch).Configure Git to use Git-Theta when tracking model checkpoints:gitthetainstallTracking a modelSay you have a codebase for training a model along with the model's checkpoint:my_codebase
├──model.pt
└──train.pyGit-Theta allows you to use Git to track the changes to your codeandyour model's parameters in tandem.
To use Git-Theta to track the model checkpoint, first rungitthetatrackmodel.ptThis will create or update the.gitattributesfile that tells Git to use Git-Theta to handle the checkpoint file.
You can then add and commit the.gitattributesfile:gitadd.gitattributes
gitcommitAfter tracking the model, you can regular Git commands (add,commit,push,pull,checkout,status,diff, etc.) as if the checkpoint file was any other file.
To add and commit the initial version of the checkpoint, simply rungitaddmodel.pt
gitcommitStoring updates efficientlyAdditionally,git theta addcan be used instead ofgit addto provide optional extra information, including e.g., the checkpoint format with--checkpoint-type, theUpdateused to update parameters with--update-type, and the location of auxiliary information/data for the update with--update-path.
For example, if the model was updated using usingLoRA, the low-rank factors can be efficiently stored by Git-Theta by running:# After training with LoRA and saving the factors to updates.pt...gitthetaaddmodel.pt--update-typelow-rank--update-pathupdates.pt
gitcommitMerging updatesGit-Theta can also handle merging of models trained with differing updates.
For example, if an existing model is further trained on a new branch calledalternate-training:gitcheckout-balternate-training# After performing training...gitaddmodel.pt
gitcommitand is separately trained on the main branch:gitcheckoutmain# After some other training...gitaddmodel.pt
gitcommitWe then can then merge the updates from thealternate-trainingbranch via a standardgit merge:gitmergealternate-trainingGit-Theta supports various methods for automatically merging models, including parameter averaging. The merge tools shows us each parameter that is different between the two models and asks what merge operation to perform.Efficiently tracking updatesGit-Theta supports various workflows for efficiently tracking updates to a checkpoint.Parameter groupsUnder the hood, Git-Theta tracks changes to a checkpoint at theparameter grouplevel.
A parameter group is a semantically-grouped collection of parameters like a weight matrix or bias vector in a neural network.
Parameter groups are determined based on the structure of the checkpoint file itself as specified in the format-specificCheckpoint class.
In the simplest case where all of the parameters of a model are updated, Git-Theta will effectively store an entirely new copy of the checkpoint.
However, if only a subset of the model's parameter groups are updated, Git-Theta will only store the updates to the changed parameter groups, which saves space and communication costs.
Similarly, if a model is updated by adding new parameter groups, Git-Theta will only store the new parameter groups.Parameter-efficient updatesBeyond updating a subset of a model's parameter groups, Git-Theta also natively supportsparameter-efficientupdates.
Examples of parameter-efficient updates include updating a sparse subset of the model's parameters (as inFISH MaskorDiff Pruning) or applying a low-rank update (as inLoRA).
There are multiple workflows for efficiently tracking parameter-efficient updates with Git-Theta.Saving update information as new parameter groupsA simple way to track parameter-efficient updates is to store the information required to produce the update (e.g., the low-rank factors for LoRA or the indices and values for a sparse update) as new parameter groups in the checkpoint file itself.
In this case, model code handles creating and applying the update and the checkpoint is saved and loaded as usual.Pros:Simple to implement.Original checkpoint and updates are bundled together and saving and loading is done as usual without special logic.Cons:Checkpoint saving may result in unnecessary writes of unchanged parameters.If many subsequent parameter-efficient updates are made, the number of parameter groups stored in the checkpoint file could become onerous.After saving update information in the checkpoint, the new checkpoint can be committed simply usinggit addandgit commitas usual.Applying updates to existing parameter groups before savingA second option is to apply the updates to the parameter groups before saving them.
Git-Theta will treat these updates in the same way it treats updating all parameters in a parameter group, so this approach sacrifices any savings to communication or storage costs that would have been achieved by using a parameter-efficient method.Pros:Similar to saving update information as new parameter groups, this is simple to implement and only involves handling a single checkpoint file.The checkpoint can be used as-is without any special logic for re-applying the update.Cons:Sacrifices any communication/storage savings from using a parameter-efficient update.Checkpoint saving may result in unnecessary writes of unchanged parameters.After folding the updates into the parameter groups, the model can be saved, added, and committed as usual.Saving update information externallyAnother option is to save parameter-efficient update information in a separate file from the original checkpoint.
This maintains storage and communication efficiency at the cost of requiring additional implementation overhead.Pros:Only the parameter updates are saved, reducing storage requirements.Only updated parameters are saved during the training loop, removing wasteful writes.Makes it easy to work with multiple datasets via different file names or branches.Cons:Implementation overhead. Training code needs to be able to segment out and save only the parameters that have changed. Inference code needs to know how to load both the original checkpoint and the update from the new checkpoint as well as how to merge them.The original checkpoint and parameter udpates are decoupled, running the risk that one could be changed without appropriately modifying the other.Assuming we have already committed the original model, the auxiliary information checkpoint needs to be separately added and committed as normal.Using Git-Theta to incorporate external update informationTo streamline the workflow of saving update information externally, Git-Theta has functionality for applying the update as part of the version control process.
This ties together the main model checkpoint and the update checkpoint to prevent them from diverging.
In addition, Git-Theta takes care of applying the update so that the model checkpoint can be used as-is after checkout.
Git-Theta assumes assumes that the update information checkpoint uses the same format as the original checkpoint and that the names of updates are prefixed by the name of the parameter group they are applied to.
For example, if a parameter group called/layer1/weightswas updated with a low-rank update, then Git-Theta would look for parameters named/layer1/weights/Rand/layer1/weights/Cin the update information checkpoint based on the naming conventions in theLowRankUpdateclass.
The low-rank update can then be efficiently tracked and applied with Git-Theta viagitthetaadd/path/to/original/checkpoint.ckpt--update-typelow-rank--update-path/path/to/updates.ckpt
gitcommitNote that using this approach requires usinggit theta addinstead of justgit addto allow for additional command line arguments.
Updates that involve modifying existing parameters (rather than just completely replacing them) are referred to by Git-Theta as "incremental updates" and are handled via a plugin system (describedbelow).Managing model development with Git-ThetaGit-Theta provides principled and rigorous way to keep track of different versions of a model based on the standard version control workflow.Tracking the progression of a modelPre-trained models are increasingly being continually updated to make them applicable to new tasks and domains.
For example, a pre-trained language model might beadapted to a new objective,process text in a new domain, andimprove its instruction-following capabilitiesbefore being fine-tuned on a target task.
Git-Theta allows the provenance of these steps to be straightforwardly tracked using Git's built-in functionality.
Apart from committing each model to keep track of a checkpoint's history, other Git functionality liketaggingcan be used to keep track of notable versions.
When checking out a particular version of a model, Git-Theta will only download what's required to reconstruct it and won't download any files that have already been cached.Tracking different versions of a modelModel development is not always straightforward - often we want to try out different versions of a base model, or we might create different versions that are applicable to different tasks.
Git-Theta supports this mode of development natively simply by using Git's branch feature - simply create a new branch (git checkout -b), modify the model, and add and commit it as usual.
This provides a straightforward workflow for trying out different ways to update a model.
If parameter groups are shared across checkpoints being tracked by Git-Theta (whether they are on the same or different branches), Git-Theta will only store a single copy of each parameter group.
Contributors can also develop their own updated versions of a model by forking the base repository.Merging modelsIf different versions of a model are created on different branches or repositories, Git-Theta will handle merging them.
Whengit mergeis run and there is a merge conflict between two histories of a model, Git-Theta will automatically open its merge tool.
Git-Theta's merge tool currently supports basic resolution patterns like choosing the parameters from one of the models or merging parameter groups via averaging.
For more sophisticated merges, the environment variableGIT_THETA_MANUAL_MERGEcan be set to true when performing the merge operation, i.e.exportGIT_THETA_MANUAL_MERGE=True
gitmerge${other-branch}and the merge tool will write out 3 copies of the model, one for each branch being merged and an additional one that represents the model at the most recent commit in the history of both branches.
The merge tool will also specify where to save the merged model.
After the merged model has been saved to the specified location, a merge commit can be created as usual.Sharp EdgesGit-Theta aims to support all standard Git workflows.
However, there are currently some situations that Git-Theta does not currently support.Git RebaseCurrently,git rebaseis not supported when special update types are used.
Additionally, repeated merge-conflict resolution---often encountered in a rebase---can be onerous for large models.Octopus MergesCurrently, git-theta's merge utilities are optimized for (and only tested for) 3-way merges where two branches with a shared ancestor commit are merged together.
We are working on support for Octopus merges where multiple branches are all combined at once.Under the hoodThis section describes how Git-Theta works in more detail.Git-Theta's filtersGit offers several points of customization where specialized, model-aware Git-Theta versions of various tools are run.
Git has a "working tree" where human-facing files live and a "staging area" where a copies of working tree files live before they are stored in Git.
When a file is moved from the working tree to the staging area, the "clean filter" is run.
When it is moved back the "smudge filter" is run.
Git-theta provides model-aware versions of these filters.When a model checkpoint iscleaned(git add):Git-Theta reads the checkpoint from the working tree using a plug-in system to support different deep-learning frameworks.Git-Theta converts the checkpoint into a tree of parameter group names that map to parameter values.Git-Theta records metadata for each parameter group, including a hash of the parameter values.Git-Theta compares the metadata for the current parameter group with its previous value.  If the metadata doesn't match, the parameter is serialized and then saved using Git LFS. The Git LFS metadata is recorded in the metadata file.The metadata is written to the staging area.Thus, Git itself only tracks the model metadata; actual values are stored efficiently Git LFS. Additionally, by checking for matching metadata, only changed parameters are stored.When a model checkpoint issmudged(git checkout):The Git-Theta metadata file is retrieved from Git.For each parameter, theUpdateplug-in system is used to get actual parameter values.
a. For updates that change all parameter values, the Git LFS metadata is used to get the values directly.
b. For parameter-efficient updates, Git LFS metadata is used to get update values, previous parameter values are retrieved from Git itself, and the update is applied.The parameter values are written into the working tree using the checkpoint plug-in system to handle different deep learning frameworks.When installing Git-Theta withgit theta install, the following lines are added to the global~/.gitconfig:[filter "theta"]clean=git-theta-filter clean %fsmudge=git-theta-filter smudge %frequired=true[merge "theta"]name=Merge Models with Git-Thetadriver=git-theta-merge %O %A %B %P[diff "theta"]command=git-theta-diffThis configuration defines twoGit filter driversfor Git-Theta and registers them under the nametheta.
In addition, it defines merge and diff programs, also namedtheta.
Whengit theta track path/to/modelis run, an entry is added to the.gitattributesfile to configure Git to use Git-Theta. The new entry looks likepath/to/model filter=theta merge=theta diff=thetaThis tells git that anytime a file that matches the patternpath/to/modelis processed, use the filter/merge/diff driver namedtheta.Incremental updatesGit-Theta supports updates that are based on the previous version of the parameter values.
For example, if a few entries of a parameter group are updated, Git-Theta can avoid storing a new copy of the parameter group; instead, it can be computed on the fly during a smudge filter based on the sparse update and the previous value.
Such updates are implemented as subclasses of theIncrementalUpdateclass.IncrementalUpdates include references to the commit that holds the last parameter value in their metadata.
Then, when the new value is needed, theIncrementalUpdateclass will fetch the value of the previous parameterfrom gitand apply the current update.
This yields a massive reduction in storage costs.
Additionaly, this can be done recursively, i.e. Git-Theta will continuous fetch previous values and applyIncrementalUpdates until a self-contained update (such as aDenseupdate that replaces all parameter values with new ones) is hit.Locality-sensitive hashingTo avoid processing parameter groups that have not been changed, Git-Theta needs a way to determine whether a given parameter group's values have changed.
Directly testing for equality or comparing bitwise hashes might be overly strict due to numerical instability and noise that could arise from using incremental updates, different hardware, or different software stacks.
Instead, Git-Theta uses uses locality sensitive hashing (LSH) for parameter hashes.
Specifically, an LSH that approximates Euclidean distance and uses the random-pool approach to hash parameters of variable sizes.
Git-Theta's LSH uses 16 hash functions and is calibrated so that two parameter groups with a Euclidean distance less than $1e^{-8}$ will have the same hash with a probability of at least $0.99$.
Additionally, weights with a distance $\in [1e{-8}, 1e^{-6}]$ are double-checked withnumpy.allclose.Plug-insGit-theta makes heavy use ofpython plug-insto enable users to add support for additional checkpoint formats as well as custom merge patterns and incremental updates.
Specifically, Git-Theta currently support plug-ins for theCheckpoint,Update, andMergeclasses.
Third-party users can register a plug-in by creating a small installable package that defines the plugin and registers it as an entry point under the name scopegit_theta.plugins.(checkpoints|updates|merges).
An example plugin for JSON formatted checkpoints can be foundhere.
Alternatively, plug-ins can be added directly to thegit-thetapackage by adding new subclasses to the appropriate modules, then declaring it in theentry_pointsdict insetup.py.Development SetupThis project usesblackfor code formatting andisortfor import statement ordering. Additionally, it includes CI that checks for compliance.
We include pre-commit hooks that will automatically runblackandisortagainst any python files staged for commit.
These hooks can be installed with:$pipinstall-rrequirements-dev.txt
$pre-commitinstallWhen one of these tools must reformat your file, it will show as the pre-commit hook failing and your commit will be cancelled.
Reformatted source files will appear in your working directory ready to be re-added to staging (git add).
Runninggit commit -m ${msg}again will result in the hooks passing and the commit actually happening.Note:As your initial commit was blocked, you will probably want to use the same message in the commit that actually goes through.CitationIf you use git-theta in your work, please cite:@InProceedings{kandpal-etal-2023-git-thetatitle={Git-Theta:AGitExtensionforCollaborativeDevelopmentofMachineLearningModels},author={Kandpal, Nikhil and Lester, Brian and Muqeeth, Mohammed and Mascarenhas, Anisha and Evans, Monty and Baskaran, Vishal and Huang, Tenghao and Liu, Haokun and Raffel, Colin},journal={International Conference on Machine Learning, {ICML}},year={2023},month={july},url={https://arxiv.org/abs/2306.04529},}