DOCSEARCHbyAlgolia
Submit Docs

ring-attention-pytorch

repository·main·Indexed 19 days ago

https://github.com/lucidrains/ring-attention-pytorch

A PyTorch implementation of Ring Attention designed for extremely long sequences by splitting data across the sequence dimension. It includes the RingAttention class for ring reduce attention matrix tiles, Striped Attention for workload balancing, and the tree_attn_decode function for Tree Attention Decoding.

Tokens
684
Snippets
4
Records
4
Agent score
19%

What's inside ring-attention-pytorch

  2. Run project tests

    main

    To test the implementation, first install the requirements, then run the assertion scripts for specific attention types.

    Testing autoregressive striped ring attention on CUDA:

    $ python assert.py --use-cuda --causal --striped-ring-attn

    Testing tree attention:

    $ python assert_tree_attn.py --use-cuda --seq-len 8192
    $ pip install -r requirements.txt

$ python assert.py --use-cuda --causal --striped-ring-attn

$ python assert_tree_attn.py --use-cuda --seq-len 8192

  3. Use the RingAttention class

    main

    The RingAttention class implements Ring Attention by splitting data across the sequence dimension and applying ring reduce to attention matrix tiles. It supports features like causal masking, automatic sequence sharding, and grouped query attention.

    Key arguments for RingAttention:

    • dim: Dimension of the input tokens.
    • dim_head: Dimension of each attention head.
    • heads: Number of attention heads.
    • causal: Boolean, whether to use causal masking.
    • auto_shard_seq: Boolean, whether to automatically shard the sequence.
    • ring_attn: Boolean, enables ring attention logic.
    • ring_seq_size: The size of the sequence chunks used in the ring reduction.
    import torch
from ring_attention_pytorch import RingAttention

attn = RingAttention(
    dim = 512,
    dim_head = 64,
    heads = 8,
    causal = True,
    auto_shard_seq = True,
    ring_attn = True,
    ring_seq_size = 512
)

tokens = torch.randn(1, 1024, 512)
attended = attn(tokens)

assert attended.shape == tokens.shape

  4. Use tree_attn_decode for Tree Attention Decoding

    main

    The tree_attn_decode function implements Tree Attention Decoding. This function expects the query (q), key (k), and value (v) tensors to already exist across all machines in a distributed environment.

    from ring_attention_pytorch import tree_attn_decode

out = tree_attn_decode(q, k, v) # where q, k, v exists across all machines