### Install MEGABYTE-pytorch

Source: https://github.com/lucidrains/megabyte-pytorch/blob/main/README.md

Install the MEGABYTE-pytorch library using pip.

```bash
pip install MEGABYTE-pytorch
```

--------------------------------

### Complete Training Loop Example

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Demonstrates a full training loop using the enwik8 dataset. Includes data loading, optimization, validation, and periodic text generation for monitoring training progress. Requires the enwik8 dataset to be downloaded.

```python
import torch
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import gzip
import numpy as np
from MEGABYTE_pytorch import MEGABYTE

# Hyperparameters
NUM_BATCHES = 100000
BATCH_SIZE = 4
GRADIENT_ACCUMULATE_EVERY = 4
LEARNING_RATE = 2e-4
VALIDATE_EVERY = 100
GENERATE_EVERY = 500
SEQ_LEN = 8192

# Initialize model
model = MEGABYTE(
    num_tokens=256,            # byte-level
    dim=(768, 512, 256),       # three-level hierarchy
    depth=(6, 4, 2),
    max_seq_len=(512, 4, 4),   # 512 * 4 * 4 = 8192 total
    flash_attn=True
).cuda()

# Dataset class for sampling sequences
class TextSamplerDataset(Dataset):
    def __init__(self, data, seq_len):
        self.data = data
        self.seq_len = seq_len

    def __getitem__(self, index):
        rand_start = torch.randint(0, self.data.size(0) - self.seq_len, (1,)) 
        return self.data[rand_start:rand_start + self.seq_len].long().cuda()

    def __len__(self):
        return self.data.size(0) // self.seq_len

# Load enwik8 data
with gzip.open('./data/enwik8.gz') as file:
    x = np.frombuffer(file.read(int(95e6)), dtype=np.uint8).copy()
    train_x, valid_x = np.split(x, [int(90e6)])
    data_train, data_val = map(torch.from_numpy, (train_x, valid_x))

# Create data loaders
def cycle(loader):
    while True:
        for data in loader:
            yield data

train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
val_dataset = TextSamplerDataset(data_val, SEQ_LEN)
train_loader = cycle(DataLoader(train_dataset, batch_size=BATCH_SIZE))
val_loader = cycle(DataLoader(val_dataset, batch_size=BATCH_SIZE))

# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Training loop
for i in range(NUM_BATCHES):
    model.train()

    # Gradient accumulation
    for _ in range(GRADIENT_ACCUMULATE_EVERY):
        loss = model(next(train_loader), return_loss=True)
        loss.backward()

    # Gradient clipping and optimization step
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
    optimizer.step()
    optimizer.zero_grad()

    if i % VALIDATE_EVERY == 0:
        model.eval()
        with torch.no_grad():
            val_loss = model(next(val_loader), return_loss=True)
            print(f"Step {i} | Train loss: {loss.item():.4f} | Val loss: {val_loss.item():.4f}")

    if i != 0 and i % GENERATE_EVERY == 0:
        model.eval()
        with torch.no_grad():
            sample = model.generate(temperature=0.9, filter_thres=0.9)
            text = ''.join([chr(max(32, t)) for t in sample.flatten().tolist()[:200]])
            print(f"Sample: {text}")
```

--------------------------------

### Initialize and Use MEGABYTE Model

Source: https://github.com/lucidrains/megabyte-pytorch/blob/main/README.md

Initialize the MEGABYTE model with specified parameters for token count, dimensions, sequence lengths, and transformer depth. This example demonstrates forward pass for loss calculation and generating logits, as well as sampling.

```python
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
)

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

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

# then after much training

logits = model(x)

# and sample from the logits accordingly
# or you can use the generate function

sampled = model.generate(temperature = 0.9, filter_thres = 0.9) # (1, 1024, 4)
```

--------------------------------

### Configure and Initialize MEGABYTE Model

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Initialize the MEGABYTE model with hierarchical parameters and verify parameter counts.

```python
model = MEGABYTE(
    # Required parameters
    num_tokens=256,                # vocabulary size
    dim=(512, 256),                # model dimensions per level
    depth=(6, 4),                  # transformer depth per level
    max_seq_len=(1024, 4),         # sequence length per level

    # Attention configuration
    dim_head=64,                   # dimension per attention head
    heads=8,                       # number of attention heads
    flash_attn=True,               # use Flash Attention (requires PyTorch 2.0+)

    # Positional embeddings
    rel_pos=False,                 # use rotary positional embeddings (RoPE)
    pos_emb=False,                 # use learned absolute positional embeddings

    # Regularization
    attn_dropout=0.1,              # dropout in attention layers
    ff_dropout=0.1,                # dropout in feedforward layers
    ff_mult=4,                     # feedforward hidden dimension multiplier

    # Special tokens
    pad_id=0                       # padding token ID
)

# Check model parameters
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")

# Verify sequence length computation
total_seq_len = 1024 * 4
print(f"Maximum sequence length: {total_seq_len} tokens")
```

--------------------------------

### Model Configuration Options

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Illustrates the import statement for the MEGABYTE model, indicating its availability for use. Further configuration options are detailed in the surrounding documentation.

```python
import torch
from MEGABYTE_pytorch import MEGABYTE
```

--------------------------------

### Initialize MEGABYTE Models

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Configure hierarchical transformer models by defining dimensions, depths, and sequence lengths for each level.

```python
import torch
from MEGABYTE_pytorch import MEGABYTE

# Basic two-level hierarchical model (global + local)
model = MEGABYTE(
    num_tokens=16000,          # vocabulary size (e.g., 256 for bytes, 16000 for subwords)
    dim=(512, 256),            # transformer dimensions: 512 for global, 256 for local
    max_seq_len=(1024, 4),     # sequence lengths: 1024 patches of 4 tokens each = 4096 total
    depth=(6, 4),              # transformer layers: 6 global layers, 4 local layers
    dim_head=64,               # dimension per attention head
    heads=8,                   # number of attention heads
    flash_attn=True            # enable Flash Attention for efficiency
)

# Three-level hierarchical model for even longer sequences
model_3level = MEGABYTE(
    num_tokens=256,            # byte-level vocabulary
    dim=(768, 512, 256),       # dimensions decrease at each level
    depth=(6, 4, 2),           # layers at each hierarchical level
    max_seq_len=(512, 4, 4),   # 512 * 4 * 4 = 8192 total sequence length
    flash_attn=True
)

print(f"Total sequence length: {512 * 4 * 4}")  # Output: Total sequence length: 8192
```

--------------------------------

### Train MEGABYTE Model

Source: https://github.com/lucidrains/megabyte-pytorch/blob/main/README.md

Train the MEGABYTE model on character-level enwik8 with patches of size 4 and a length of 8192.

```bash
python train.py
```

--------------------------------

### Use Attend Class for Efficient Attention

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Utilize the Attend module for memory-efficient attention with optional Flash Attention and masking support.

```python
import torch
from MEGABYTE_pytorch.attend import Attend

# Create attention module with flash attention
attend = Attend(
    causal=True,          # use causal (autoregressive) attention mask
    dropout=0.1,          # attention dropout rate
    flash=True            # enable Flash Attention (requires PyTorch 2.0+)
)

# Prepare query, key, value tensors
batch_size, heads, seq_len, dim_head = 2, 8, 512, 64
q = torch.randn(batch_size, heads, seq_len, dim_head).cuda()
k = torch.randn(batch_size, heads, seq_len, dim_head).cuda()
v = torch.randn(batch_size, heads, seq_len, dim_head).cuda()

# Compute attention
with torch.no_grad():
    output = attend(q, k, v)
    print(f"Attention output shape: {output.shape}")  # (2, 8, 512, 64)

# With optional mask
mask = torch.ones(batch_size, seq_len, dtype=torch.bool).cuda()
mask[:, 256:] = False  # mask out second half

output_masked = attend(q, k, v, mask=mask)
print(f"Masked attention output shape: {output_masked.shape}")
```

--------------------------------

### Autoregressive Text Generation

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Performs autoregressive sampling to produce complete sequences. Supports temperature-based sampling and top-k filtering for controlling output diversity. Use for generating text from scratch or continuing existing text.

```python
import torch
from MEGABYTE_pytorch import MEGABYTE

model = MEGABYTE(
    num_tokens=256,            # byte-level (ASCII)
    dim=(512, 256),
    max_seq_len=(128, 4),      # generates 128 * 4 = 512 bytes
    depth=(6, 4),
    flash_attn=True
).cuda()

# Generate from scratch (no prime)
with torch.no_grad():
    generated = model.generate(
        temperature=0.9,       # higher = more random, lower = more deterministic
        filter_thres=0.9       # top-k filtering threshold (keeps top 10% of logits)
    )
    print(f"Generated shape: {generated.shape}")  # (1, 128, 4)

    # Flatten and decode to text
    flat_output = generated.flatten(1)  # (1, 512)
    text = ''.join([chr(max(32, t)) for t in flat_output[0].tolist()])
    print(f"Generated text: {text[:100]}...")

# Generate with a prime (continuation)
prime_text = "The quick brown fox"
prime_tokens = torch.tensor([[ord(c) for c in prime_text]]).cuda()

with torch.no_grad():
    continued = model.generate(
        prime=prime_tokens,
        temperature=0.8,
        filter_thres=0.9
    )
    print(f"Continued shape: {continued.shape}")
```

--------------------------------

### Process Flat Sequences

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Pass 1D sequences directly to the model; it will automatically handle padding and reshaping.

```python
import torch
from MEGABYTE_pytorch import MEGABYTE

model = MEGABYTE(
    num_tokens=256,
    dim=(512, 256),
    max_seq_len=(1024, 4),
    depth=(6, 4),
    pad_id=0,                  # padding token ID for auto-padding
    flash_attn=True
).cuda()

# Flat sequence input (will be auto-padded to nearest multiple of patch size)
flat_sequence = torch.randint(0, 256, (1, 3000)).cuda()  # 3000 tokens

# Model automatically pads to 3000 -> 3004 (next multiple of 4)
# Then reshapes to (1, 751, 4) for hierarchical processing
loss = model(flat_sequence, return_loss=True)
print(f"Loss with flat input: {loss.item():.4f}")

# Inference returns logits matching original sequence length
with torch.no_grad():
    logits = model(flat_sequence, return_loss=False)
    print(f"Output logits shape: {logits.shape}")  # (1, 3000, 256)
```

--------------------------------

### Perform Forward Pass and Loss Computation

Source: https://context7.com/lucidrains/megabyte-pytorch/llms.txt

Execute training or inference using hierarchical input tensors.

```python
import torch
from MEGABYTE_pytorch import MEGABYTE

model = MEGABYTE(
    num_tokens=16000,
    dim=(512, 256),
    max_seq_len=(1024, 4),
    depth=(6, 4),
    flash_attn=True
).cuda()

# Create input tensor matching the hierarchical shape: (batch, global_seq, local_seq)
x = torch.randint(0, 16000, (2, 1024, 4)).cuda()  # batch_size=2

# Training mode: compute cross-entropy loss
loss = model(x, return_loss=True)
print(f"Training loss: {loss.item():.4f}")

# Backpropagation
loss.backward()

# Inference mode: get logits for next token prediction
with torch.no_grad():
    logits = model(x, return_loss=False)
    print(f"Logits shape: {logits.shape}")  # (2, 1024, 4, 16000)

    # Get predictions for the last position
    next_token_logits = logits[:, -1, -1, :]
    predicted_tokens = next_token_logits.argmax(dim=-1)
    print(f"Predicted next tokens: {predicted_tokens}")
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.