Package: yet-another-retnet
yet-another-retnetA simple but robust PyTorch implementation of RetNet fromRetentive Network: A Successor to Transformer for Large Language Models.Also see Microsoft's original implementation:RetNet @ microsoft/torchscaleUltimately, their implementation is the ground truth.  I have tried to make my implementation as readable and well-documented as possible, while still being consistent with the original.  My version also includes full type annotations, and a robust set of unit tests.I'm obviously biased, but I find the untyped, config-driven approach inmicrosoft/torchscaleclunky and difficult to adapt to other use cases.TODOEquivalentparallelandrecurrentretention methods.  See:retention.pyRecurrent position embedding implementation.MultiScaleRetentionmodule.  See:retention.pyMake relative position embeddings forMultiScaleRetentionoptional.The retention layer explicitly includes a position embedding update, which is based onxPos.  It does not necessarily translate well to other domains (e.g. computer vision, heterogeneous graphs).  So, I have made it optional.I'm not 100% sure why the authors did this.  It seems overly specific to the language modeling use case, and it's not clear to me that it was necessary.End-to-endRetNetmodule.  See:retnet.pyRetNetDecoderLayerRetNetDecoderPreconfigured 1.3B, 2.7B, and 6.7B models (untrained).  See:retnet.pyReproduce inference memory and throughput benchmarks from the paper.  See:Inference Benchmarks,benchmark_inference.pyRelease stable version on PyPI.PrereleaseStableEquivalentchunkwiseretention method.Basic training example for language modeling.  See:train_project_gutenberg.pyInstallPyPI:pipinstallyet-another-retnetNOTE: To run theexample training script, you will need to include the[train]extra package:pipinstallyet-another-retnet[train]From source:pipinstall"yet-another-retnet @ git+ssh://git@github.com/fkodom/yet-another-retnet.git"For contributors:# Clone/fork the repositoryghrepoclonefkodom/yet-another-retnetcdyet-another-retnet# Install all dev dependencies (tests etc.) in editable modepipinstall-e.[test]# Setup pre-commit hookspre-commitinstallAboutRetNet is a transformer-like architecture that has equivalentparallelandrecurrentformulations.The benefits of this dual formulation are:Accuracy comparable to Transformer-based modelsParallel: high training throughputRecurrent: high inference throughputThis is the "impossible triangle" of language model design, as described by the authors:UsageRetNetUse one of the configurations described in the paper:retnet_1_3bretnet_2_7bretnet_6_7bfromyet_another_retnet.retnetimportretnet_1_3bretnet=retnet_1_3b(num_tokens=10000,device="cuda")or create your ownRetNetmodel directly:fromyet_another_retnet.retnetimportRetNet# a very small RetNet model :Dretnet=RetNet(num_tokens=1000,# vocab size, usually taken from tokenizerd_model=64,nhead=4,num_layers=2,device="cuda",).eval()# Important for reproducibility!Equivalent parallel, recurrent, and chunkwise usage:importtorch# Set deterministic CUDA opstorch.backends.cudnn.deterministic=Truetorch.backends.cudnn.benchmark=False# input shape: (batch_size, seq_len)# integer range: [0, num_tokens)x=torch.randint(0,1000,(1,16),device="cuda")# Parallel usagey_parallel=retnet.forward_parallel(x)# Recurrent usageoutputs=[]# container for collecting step-wise outputsprev_states=[]# cache layer states after each stepforidxinrange(16):# seq_lenout,prev_states=retnet.forward_recurrent(x[:,idx],idx,prev_states)outputs.append(out)y_recurrent=torch.stack(outputs,dim=1)# Chunkwise usageoutputs=[]# container for collecting chunk-wise outputsprev_states=[]# cache layer states after each stepchunk_size=4# number of tokens in each chunkforidxinrange(0,16,chunk_size):out,prev_states=retnet.forward_chunkwise(x[:,idx:idx+chunk_size],idx,prev_states)outputs.append(out)y_chunkwise=torch.cat(outputs,dim=1)# Check that outputs are equaltorch.testing.assert_close(y_parallel,y_recurrent)torch.testing.assert_close(y_parallel,y_chunkwise)NOTE: There is some floating point error accumulation in the recurrent formulation, which I believe is less pronounced in the parallel formulation. Especially for untrained models (when activations are very large), the two outputs may not matchexactly.  The absolute difference should still be very small -- on the order of 1e-5 or less.MultiScaleRetentionEquivalent parallel, recurrent, and chunkwise usage:importtorchfromyet_another_retnet.retentionimportMultiScaleRetentionmhr=MultiScaleRetention(embed_dim=32,num_heads=4,device="cuda").eval()# input shape: (batch_size, seq_len, embed_dim)q=k=v=torch.randn(1,16,32,device="cuda")# Parallel retentiony_parallel,_=mhr.forward_parallel(q,k,v)# Recurrent retentionoutputs=[]prev_state=Noneforidxinrange(16):out,prev_state=mhr.forward_recurrent(q[:,idx],k[:,idx],v[:,idx],idx,prev_state)outputs.append(out)y_recurrent=torch.stack(outputs,dim=1)# Chunkwise retentionoutputs=[]prev_state=Nonechunk_size=4foridxinrange(0,16,chunk_size):out,prev_state=mhr.forward_chunkwise(q[:,idx:idx+chunk_size],k[:,idx:idx+chunk_size],v[:,idx:idx+chunk_size],idx,prev_state,)outputs.append(out)y_chunkwise=torch.cat(outputs,dim=1)# Check that outputs are equaltorch.testing.assert_close(y_parallel,y_recurrent)torch.testing.assert_close(y_parallel,y_chunkwise)NOTE: TheMultiScaleRetentionthat is described in the paper includes an
explicit position embedding (based on xPos) as part of the retention layer.  This
does not translate perfectly to other domains (e.g. computer vision, heterogeneous
graphs), so I have made it optional.Setrelative_position=Falseto disable the position embedding.  Instead, you will
be responsible for adding positional information to the inputs (if needed).# Disable relative position embeddingmhr=MultiScaleRetention(embed_dim=32,num_heads=4,relative_position=False,device="cuda")# Everything else works the same as above.# Just add your own positional embeddings to the inputs.Retention forward passSimilar to the example above, but head projections and positional updates are not internalized byMultiScaleRetention:importtorchfromyet_another_retnet.retentionimport(retention_chunkwise,retention_parallel,retention_recurrent,)# input shape: (batch_size, num_heads, seq_len, head_dim)q=k=v=torch.randn(1,4,32,8,device="cuda")# Parallel retentiony_parallel,_=retention_parallel(q,k,v)# Recurrent retentionoutputs=[]prev_state=Noneforiinrange(32):out,prev_state=retention_recurrent(q[:,:,i],k[:,:,i],v[:,:,i],prev_state)outputs.append(out)y_recurrent=torch.stack(outputs,dim=2)# Chunkwise retentionoutputs=[]prev_state=Nonechunk_size=4foriinrange(0,32,chunk_size):out,prev_state=retention_chunkwise(q[:,:,i:i+chunk_size],k[:,:,i:i+chunk_size],v[:,:,i:i+chunk_size],prev_state,)outputs.append(out)y_chunkwise=torch.cat(outputs,dim=2)# Check that outputs are equaltorch.testing.assert_close(y_parallel,y_recurrent)torch.testing.assert_close(y_parallel,y_chunkwise)Inference BenchmarksNOTE: The benchmarks aren't exact one-to-one comparisons, because I have a much lower-end GPU.  The authors benchmark RetNet 6.7B using an A100 80GB.  I have a 2080 Ti (11GB), so I chose to benchmark RetNet 1.3B instead.  I expect the batch size is also smaller in my benchmarks, although the authors don't specify what their batch size was.My benchmarks clearly show the same scaling trends, which is still helpful as a rough validation.From the paper:From this repo:Parallel vs. Recurrent vs. ChunkwiseWhen should you choose one formulation over the others?  Here is a general rule of thumb:Parallel -> model trainingRecurrent -> incremental token predictionChunkwise ->model training -- if training inputs are very long, and parallel formulation is too memory intensiveencoding long prompts -- if inferencepromptsare very long, chunkwise encoding is more efficient than the recurrent formulation.  The state returned from chunkwise formulation is the same as in recurrent formulation.  So, once the prompt is chunkwise encoded, use recurrent formulation to generate new tokens.Citations@misc{sun2023retentive,title={Retentive Network: A Successor to Transformer for Large Language Models},author={Yutao Sun and Li Dong and Shaohan Huang and Shuming Ma and Yuqing Xia and Jilong Xue and Jianyong Wang and Furu Wei},year={2023},eprint={2307.08621},archivePrefix={arXiv},primaryClass={cs.CL}}