Package: truegrad
TrueGradPyTorch interface for TrueGrad OptimizersGetting StartedInstallationpython3-mpipinstalltruegradExamplesTrueGrad supports various backends, each with their own tradeoffs:NameAdvantagesDisadvantagestruegrad.nn* What you see is what you get - Modules not in truegrad.nn and truegrad.nn.functional are not supported* Custom forward/backward for some fused functions* Optimized backward passes* Limited applicability - custom modules can't be used* Requires code modificationtruegrad.utils.patch_torch* Uses truegrad.nn under the hood* Works for many (off-the-shelf!) torch models* No code modification necessary* Uncertainty if model is compatiblebackpack* Highest stability* Loud warnings and errors* Battle-tested* Simple to extend further* High memory usage* High compute usage* Sparse support for torch operationstruegrad.utils.patch_model* Works with custom models* Fails silently on fused functions* ~50% to 100% slower than truegrad.nnpatch_torch + patch_model* Best compatibility* Reduced overheads compared topatch_model(by falling back to faster pre-patchedpatch_torchwhere available)* Fails silently on fused functions outside of torch.nn* Slower than truegrad.nn when truegrad.nn would've been enoughBelow, you'll find examples for each of these backends, as well as ageneral strategyallowing
partial application of TrueGrad.nnThe preferred method of using TrueGrad is by replacingtorch.nnwith performanttruegrad.nnmodules. While other
methods add compute and memory overheads,truegrad.nnandtruegrad.nn.functionalhave hand-crafted gradients. This
is the most powerful method, although it requires code modifications.importtorchfromtruegradimportnnfromtruegrad.optimimportTGAdamW# define model by mixing truegrad.nn and torch.nnmodel=torch.nn.Sequential(nn.Linear(1,10),nn.LayerNorm(10),torch.nn.ReLU(),nn.Linear(10,1))optim=TGAdamW(model.parameters())# truegrad.optim.TGAdamW instead of torch.optim.AdamW# standard training loopwhileTrue:input=torch.randn((16,1))model(input).mean().backward()optim.step()optim.zero_grad()Patch TorchIn some cases, you can't modify the model's source. For example, when importing models fromtorchvision. If that's the
case, or if you simply want to try out TrueGrad, you can usetruegrad.utils.patch_torch(), to
replacetorch.nn.Module's withtruegrad.nn.Module's where possible. For example, the code below can be used to train
a ResNet-18:importtorchfromtorchvision.modelsimportresnet18fromtruegrad.optimimportTGAdamWfromtruegrad.utilsimportpatch_torchpatch_torch()# call before model creation, otherwise complete freedommodel=resnet18().cuda()optim=TGAdamW(model.parameters(),lr=1e-7,weight_decay=0)# constant input/output to overfitinp=torch.randn((2,3,224,224)).cuda()tgt=torch.randint(0,1000,(2,)).cuda()# standard training loopi=0whileTrue:loss=torch.nn.functional.cross_entropy(model(inp),tgt)loss.backward()optim.step()optim.zero_grad()i+=1ifi%5==0:print(i,loss.item())Similarly, most huggingface transformers work out of the box:importtorchimporttransformersfromtorch.nnimportfunctionalasFfromtruegrad.optimimportTGAdamWfromtruegrad.utilsimportpatch_torchpatch_torch()# only added line to get truegrad statistics for TGAdamWmodel=transformers.BertModel.from_pretrained("google/bert_uncased_L-2_H-128_A-2")# any existing modeltokenizer=transformers.BertTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")optim=TGAdamW(model.parameters())# constant input to overfitinput=tokenizer(["Hello World!"],return_tensors="pt")# training loop as normalwhileTrue:out=model(**input)loss=F.l1_loss(out[0],torch.ones_like(out[0]))loss.backward()optim.step()optim.zero_grad()print(loss.item())Note that this works even though transformers have custom modules, which could cause issues. The key factor is that all
parameters come fromtorch.nn.Module's, which are patched bypatch_torch(). Therefore, truegrad handles all
parameter usages. Therefore, any composition oftorch.nn.Module's makes for a truegrad-compatible model.BackPackThe most stable although also memory hungry method to compute TrueGrad statistics is to useBackPack. BackPack is a third-party library that automatically computes the sum
of gradient squares and works for most models by implementing custom backward rules for manytorch.nn.Module's.importbackpackimporttorchfromtorch.nnimportCrossEntropyLossfromtruegrad.optimimportTGAdamWfromtorchvision.modelsimportalexnetmodel=alexnet()# BatchNorm and in-place ops (like ResNet's residual path) aren't supportedoptim=TGAdamW(model.parameters(),lr=1e-7,weight_decay=0)# replace inplace ops like nn.ReLU(inplace=True) where possibleformodinmodel.modules():ifhasattr(mod,"inplace"):mod.inplace=False# backpack relies on module-level pytorch hooksmodel=backpack.extend(model)lossfunc=backpack.extend(CrossEntropyLoss())# constant input/output to overfitinp=torch.randn((2,3,224,224))tgt=torch.randint(0,1000,(2,))# standard training loopi=0whileTrue:# "SumGradSquared" computes the sum of the squared gradientwithbackpack.backpack(backpack.extensions.SumGradSquared()):loss=lossfunc(model(inp),tgt)loss.backward()optim.step()optim.zero_grad()i+=1ifi%5==0:print(i,loss.item())If you're using custom modules with self-defined parameters, this method will not work. Additionally, note that, if
your model has any layer called.outputor you're using PyTorch >= 1.13, you will need to installBackPack-HFviapython3 -m pip install git+https://github.com/ClashLuke/backpack-hf.Patch Custom ModelsAnother option to integrate TrueGrad into existing models is to patch them usingtruegrad.utils.patch_model().patch_model()will go through alltorch.nn.Module's in PyTorch model and convert theirtorch.nn.Parameter's totruegrad.nn.TrueGradParameter's. ATrueGradParameteracts largely the same as atorch.nn.Parameter, but adds
required operations into the model's backward pass. Note that this doesn't give the most effective computation graph,
but works well for many custom models.Importantly, be aware that this does not work for fused functions, such astorch.nn.LayerNormandtorch.nn.MultiheadAttention. However, unfused functions which directly access a parameter, such as multiplication,
work well. Therefore, torch.nn.Linear and HuggingFace's attention work as expected.importtorchfromtruegrad.optimimportTGAdamWfromtruegrad.utilsimportpatch_modelfromtorchvision.modelsimportalexnetmodel=alexnet()# patch_model can't handle fused ops like VGG's and ResNet's BatchNormoptim=TGAdamW(model.parameters())# replace inplace ops like nn.ReLU(inplace=True) where possibleformodinmodel.modules():ifhasattr(mod,"inplace"):mod.inplace=Falsepatch_model(model)# replace torch.nn.Parameter with truegrad.nn.Parameter# constant input/output to overfitinp=torch.randn((2,3,224,224))tgt=torch.randint(0,1000,(2,))# standard training loopi=0whileTrue:loss=torch.nn.functional.cross_entropy(model(inp),tgt)loss.backward()optim.step()optim.zero_grad()i+=1ifi%5==0:print(i,loss.item())Full PatchingOne way of avoidingtruegrad.utils.patch_model's downsides when working with off-the-shelf
models containing custom parameters, such aslucidrains' ViT'sis to alsopatch_torch. This takes care of many fused functions, such as LayerNorm, while still allowing full flexibility in
model design.importtorchfromvit_pytorch.levitimportLeViTfromtruegrad.utilsimportpatch_torch,patch_modelfromtruegrad.optimimportTGAdamWpatch_torch()# before model instantiationlevit=LeViT(image_size=224,num_classes=1000,stages=3,# number of stagesdim=(256,384,512),# dimensions at each stagedepth=4,# transformer of depth 4 at each stageheads=(4,6,8),# heads at each stagemlp_mult=2,dropout=0.1)opt=TGAdamW(levit.parameters())patch_model(levit)# replace torch.nn.Parameter with truegrad.nn.TrueGradParameter# constant input to overfitimg=torch.randn(1,3,224,224)# standard training loopwhileTrue:loss=levit(img).square().mean()loss.backward()opt.step()opt.zero_grad()print(loss.item())Partial TrueGradUnfortunately, it's not always sensible to apply TrueGrad, as some backward passes are too slow, and sometimes it's
impossible to avoid a fused function.
Therefore, it can be an option to use TGAdamW only on specific subsections of the model. To do so, you can
specifydefault_to_adam=Trueto TGAdamW. Adding this option allows TGAdamW to fall back to AdamW if there is
nosum_grad_squaredattribute available.
For example, the code from#nncould be extended in the following way:importtorchfromtruegradimportnnfromtruegrad.optimimportTGAdamWmodel=torch.nn.Sequential(nn.Linear(1,10),# Weights coming from truegrad.nnnn.LayerNorm(10),torch.nn.ReLU(),torch.nn.Linear(10,1))# Weights coming torch.nnoptim=TGAdamW(model.parameters(),default_to_adam=True)# standard training loopi=0whileTrue:input=torch.randn((16,1))loss=model(input).mean()loss.backward()optim.step()optim.zero_grad()i+=1ifi%5==0:print(i,loss.item())