DOCSEARCHbyAlgolia
Submit Docs

MEGABYTE-pytorch

repository·main·Indexed 20 days ago

https://github.com/lucidrains/megabyte-pytorch

A PyTorch implementation of the MEGABYTE architecture for predicting million-byte sequences using multiscale transformers. It supports multiple local models and provides functionality for training with loss calculation and sequence generation via the .generate() method.

Tokens
692
Snippets
4
Records
4
Agent score
21%

What's inside megabyte-pytorch

  2. Train the MEGABYTE model

    main

    To train the model, pass a tensor of token indices to the model. Setting return_loss = True will return the calculated loss for backpropagation.

    import torch
from MEGABYTE_pytorch import MEGABYTE

# ... model initialization ...

x = torch.randint(0, 16000, (1, 1024, 4))

loss = model(x, return_loss = True)
loss.backward()

  3. Initialize the MEGABYTE model

    main

    The MEGABYTE class implements a multiscale transformer capable of handling long sequences by using multiple local models. You can specify dimensions, sequence lengths, and depths for both global and local scales using tuples.

    import torch
from MEGABYTE_pytorch import MEGABYTE

model = MEGABYTE(
    num_tokens = 16000,             # number of tokens
    dim = (512, 256),               # transformer model dimension (512 for coarsest, 256 for fine in this example)
    max_seq_len = (1024, 4),        # sequence length for global and then local. this can be more than 2
    depth = (6, 4),                 # number of layers for global and then local. this can be more than 2, but length must match the max_seq_len's
    dim_head = 64,                  # dimension per head
    heads = 8,                      # number of attention heads
    flash_attn = True               # use flash attention
)

  4. Generate sequences with MEGABYTE

    main

    After training, you can use the .generate() method to sample new sequences from the model. You can control the output using temperature and filter_thres.

    # ... model initialization and training ...

# Generate new tokens
sampled = model.generate(temperature = 0.9, filter_thres = 0.9) # returns shape (1, 1024, 4)