dingo-actual / infini-transformer Goto Github PK

The MLPs in both transformer modules currently have ReLU hard-coded as the activation function. It would help to have options for nonlinear activations commonly used in recent LLMs (GeLU, SwiGLU, GeGLU, etc.)

Exactly where to use causal masking?

Build a module that uses infini-attention in a mixture-of-depths layer (https://arxiv.org/abs/2404.02258).

Isn’t is unnecessary to perform q, k, v projections inside the loop? Due to data copying during the calculation process this makes the whole attention operation slower. Ideally this whole thing would be done with a fused kernel … future personal project of mine.

After changing the way positional embeddings are incorporated into the transformer classes, the docs and example need to be updated.

Currently, in MoDInfiniTransformer, sampling_factor must evenly divide segment_len. This restriction is unnecessary and logic needs to be added to handle sampling_factors in the interval (1, segment_len].

The example script isn't complicated, but it should have comments all the same.

While running some tests, I noticed several memory inefficiencies that nullify the advantages of the Infini-Transformer. In particular:

• The loop over segments needs to be moved out of the memory module and into the transformer module.
• The final projection step should be performed on a per-segment basis, rather than for the full sequence.

Currently there's no documentation around the activations functions. Also need citations for the activations, notably "GLU Variants Improve Transformer" (https://arxiv.org/abs/2002.05202) and "Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks" (https://arxiv.org/abs/2310.02244).

Currently, the only positional embeddings that can be handled are those that take place before passing embeddings to the transformer blocks. It would be nice to have positional embeddings that can be utilized within the transformer blocks.

Need to look through positional embedding papers to see which ones are compatible with Infini-Transformer.

Currently, the token sampling for MoD Infini-Former at inference time can result in different length sequences for each observation in the batch. The current workaround is to force the batch size to one and loop through the observations in the batch, which is highly inefficient.

There are two main options for handling this efficiently:

1. Pad the sampled sequences to the longest sequence length in such a way that the additional tokens contribute nothing to downstream calculations.
2. Wait for PyTorch to implement a ragged tensor type

I'm likely to pursue the first because there's no telling how long it'll be before the PyTorch devs add ragged tensors.

Great implementation! I have a couple of questions:
Does you implementation have O(N) complexity like the papers? What kind of memory savings should i expect compared classic sdpa attention? It would be great if you can give me some numbers on the VRAM usage during your training test compared to classic SDPA.
Thanks!

There are a couple of changes that need to be made to the training loop when using MoD Infini-Former, as well as a minor change to the .forward() method of the Module that contains MoD Infini-Former layers.

Need to put example usage in the readme.

thoughts about this exception during inference?

  File "/workspace/voltronformers/src/voltronformer/model.py", line 303, in forward                               
    output = residual + self.attn(h, position_ids=position_ids)[0]                                                                                                                                                                  
modules/module.py", line 1520, in _call_impl    tor h/nn/modules/module.py",eline 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)                                                                       
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl             
    return forward_call(*args, **kwargs)                                                                                                                                                                                            
  File "/workspace/voltronformers/src/voltronformer/infini_attention.py", line 109, in forward                                                                                                                                      
    return torch.concat(out, dim=1)                                                                                                                                                                                                 
RuntimeError: torch.cat(): expected a non-empty list of Tensors