### Basic Setup

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Initializes distributed training using torch.distributed and creates a distributed MoE layer.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch import MoE

# Initialize distributed training
dist.init_process_group(
    backend='nccl',  # GPU: 'nccl', CPU: 'gloo', MPS: 'gloo'
    init_method='env://',
    rank=0,  # Set automatically from environment
    world_size=4,  # Set automatically from environment
)

# Create distributed MoE
moe = MoE(
    dim=512,
    num_experts=16,
    is_distributed=True,  # Enable distributed mode
)

# Clean up
dist.destroy_process_group()
```

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

### Training vs Evaluation Differences Example Setup

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/routing-mechanism.md

Provides a typical setup for initializing an MoE model, specifying different values for training and evaluation parameters like thresholds and capacity factors.

```python
MoE(
    threshold_train=0.2,    # More selective during training
    threshold_eval=0.2,     # Same for evaluation
    capacity_factor_train=1.25,  # Tight during training
    capacity_factor_eval=2.0,    # More generous during evaluation
)
```

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

### Training Loop Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

A comprehensive example of a distributed training loop using ST-MoE, including model initialization, DDP wrapping, data loading, optimizer, and the training/validation process.

```python
import torch
import torch.distributed as dist
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from st_moe_pytorch import MoE, SparseMoEBlock

# Initialize distributed
dist.init_process_group(backend='nccl')
rank = dist.get_rank()
world_size = dist.get_world_size()

# Set device
torch.cuda.set_device(rank)

# Create model
moe = MoE(
    dim=512,
    num_experts=16,
    is_distributed=True,
    allow_var_seq_len=True
)
moe_block = SparseMoEBlock(moe)

# Wrap with DDP for synchronizing gradients
moe_block = DDP(moe_block, device_ids=[rank])

# Create data loader
dataloader = create_dataloader(rank, world_size)

# Optimizer and scheduler
optimizer = optim.AdamW(moe_block.parameters(), lr=1e-4)

# Training loop
for epoch in range(num_epochs):
    moe_block.train() 
    
    for batch_idx, (x, targets) in enumerate(dataloader):
        x = x.cuda(rank)
        targets = targets.cuda(rank)
        
        # Forward pass
        moe_result = moe_block(x, noise_gates=True)
        outputs = moe_result.outputs
        
        # Compute loss
        main_loss = compute_loss(outputs, targets)
        aux_loss = moe_result.total_aux_loss
        total_loss = main_loss + aux_loss
        
        # Backward pass
        optimizer.zero_grad()
        total_loss.backward()
        
        # DDP synchronizes gradients
        optimizer.step()
        
        # Logging (rank 0 only)
        if rank == 0 and batch_idx % 100 == 0:
            print(f"Epoch {epoch}, Step {batch_idx}: "
                  f"main_loss={main_loss:.4f}, "
                  f"aux_loss={aux_loss:.6f}")
    
    # Validation
    if rank == 0:
        moe_block.eval()
        val_loss = evaluate(moe_block, val_dataloader)
        print(f"Validation loss: {val_loss:.4f}")

dist.destroy_process_group()
```

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

### Using torchrun (torch 1.10+)

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Command-line example for launching distributed training using `torchrun`, the recommended tool for distributed launches in PyTorch 1.10 and later.

```bash
torchrun --nproc_per_node=4 train.py
```

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

### Using torch.distributed.launch

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Command-line examples for launching distributed training using `torch.distributed.launch` on a single machine with multiple GPUs or across multiple machines.

```bash
# 4 GPUs on single machine
python -m torch.distributed.launch \
    --nproc_per_node=4 \
    --nnodes=1 \
    train.py

# 2 machines, 4 GPUs each
python -m torch.distributed.launch \
    --nproc_per_node=4 \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr=master_ip \
    --master_port=29500 \
    train.py
```

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

### Expert Offloading

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Example of accessing the experts container for potential offloading to CPU in distributed scenarios.

```python
from st_moe_pytorch.st_moe_pytorch import Experts

experts_container = moe.experts
```

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

### Example: Basic Usage

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/sparse-moe-block.md

Basic usage example of SparseMoEBlock.

```python
import torch
from st_moe_pytorch import MoE, SparseMoEBlock

# Create MoE first
moe = MoE(
    dim=512,
    num_experts=16,
    gating_top_n=2,
)

# Wrap with optional feedforward layers
moe_block = SparseMoEBlock(
    moe,
    add_ff_before=True,
    add_ff_after=True
)

x = torch.randn(4, 1024, 512)
result = moe_block(x)

print(f"Output shape: {result.outputs.shape}")  # (4, 1024, 512)
print(f"Auxiliary loss: {result.total_aux_loss.item():.4f}")
```

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

### Variable Batch Sizes

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Demonstrates how different ranks can process different batch sizes in a distributed setup.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch import MoE

dist.init_process_group(backend='nccl')

moe = MoE(dim=512, num_experts=16, is_distributed=True)

# Each rank with different batch size
rank = dist.get_rank()
batch_size = 4 + rank  # Different: 4, 5, 6, 7

x = torch.randn(batch_size, 1024, 512).cuda(rank)

result = moe(x)
output = result.outputs

dist.destroy_process_group()
```

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

### Manual Usage Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/top-n-gating.md

Example of manually using the TopNGating class.

```python
import torch
from st_moe_pytorch.st_moe_pytorch import TopNGating

# Create gating layer (typically done internally by MoE)
gating = TopNGating(
    dim=512,
    num_gates=16,
    top_n=2,
    threshold_train=0.2,
    threshold_eval=0.2,
    capacity_factor_train=1.25,
    capacity_factor_eval=2.0
)

x = torch.randn(4, 1024, 512)

# Get routing tensors
dispatch, combine, balance_loss, router_z_loss = gating(x)

print(f"Dispatch shape: {dispatch.shape}")   # (4, 1024, 16, capacity)
print(f"Combine shape: {combine.shape}")     # (4, 1024, 16, capacity)
print(f"Balance loss: {balance_loss.item():.6f}")
print(f"Router z-loss: {router_z_loss.item():.6f}")
```

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

### Example: Manual Usage

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/geglu.md

Demonstrates manual usage of the GEGLU layer.

```python
import torch
from st_moe_pytorch.st_moe_pytorch import GEGLU

# Create GEGLU layer
geglu = GEGLU(dim=256, mult_bias=True)

# Input should be 2×dim (typically from a linear projection)
x = torch.randn(4, 64, 512)  # 512 = 2 * 256
output = geglu(x)

print(f"Input shape: {x.shape}")   # (4, 64, 512)
print(f"Output shape: {output.shape}")  # (4, 64, 256)
```

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

### Example: Training with Loss Accumulation

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/sparse-moe-block.md

Example demonstrating training with SparseMoEBlock, including loss accumulation.

```python
import torch
from st_moe_pytorch import MoE, SparseMoEBlock

moe = MoE(dim=512, num_experts=16, balance_loss_coef=1e-2, router_z_loss_coef=1e-3)
moe_block = SparseMoEBlock(moe, add_ff_before=False, add_ff_after=True)
optimizer = torch.optim.Adam(moe_block.parameters())

x = torch.randn(4, 1024, 512)
targets = torch.randn(4, 1024, 512)

moe_block.train()

# Forward pass
result = moe_block(x, noise_gates=True)

# Main loss
main_loss = torch.nn.functional.mse_loss(result.outputs, targets)

# Total loss including auxiliary losses
total_loss = main_loss + result.total_aux_loss

# Backprop and optimization
total_loss.backward()
optimizer.step()
optimizer.zero_grad()

# Logging
print(f"Main loss: {main_loss.item():.4f}")
print(f"Balance loss: {result.balance_loss.item():.6f}")
print(f"Router z-loss: {result.router_z_loss.item():.6f}")
```

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

### Installation

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

Install the st-moe-pytorch package using pip.

```bash
$ pip install st-moe-pytorch
```

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

### Installation

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/README.md

Command to install the st-moe-pytorch library.

```bash
pip install st-moe-pytorch
```

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

### Example with Differentiable Top-K

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/top-n-gating.md

Demonstrates how to use TopNGating with differentiable top-k selection for gradient flow.

```python
import torch
from st_moe_pytorch.st_moe_pytorch import TopNGating

# Use differentiable top-k for gradient flow through top-k selection
gating = TopNGating(
    dim=512,
    num_gates=16,
    top_n=2,
    differentiable_topk=True,      # Enable differentiable selection
    differentiable_topk_fused=True  # Use Triton optimization
)

x = torch.randn(4, 1024, 512, requires_grad=True)
dispatch, combine, balance_loss, router_z_loss = gating(x)

# Now can backprop through top-k selection
loss = balance_loss + router_z_loss
loss.backward()
```

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

### Example: Using Experts Directly (Advanced)

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/experts.md

An example demonstrating the direct usage of the Experts module.

```python
import torch
import torch.nn as nn
from st_moe_pytorch.st_moe_pytorch import Experts, Expert

# Create expert modules
expert_list = [Expert(dim=512) for _ in range(16)]

# Create container
experts = Experts(expert_list, is_distributed=False)

# Inputs should be (batch, num_experts, expert_capacity, dim)
batch_size, num_experts_num, capacity, dim = 4, 16, 32, 512
x = torch.randn(batch_size, num_experts_num, capacity, dim)
```

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

### Example: Basic Expert

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/expert.md

Creates a basic expert and processes a tensor.

```python
import torch
from st_moe_pytorch.st_moe_pytorch import Expert

# Create an expert for 512-dim embeddings
expert = Expert(dim=512, hidden_mult=4, mult_bias=True)

# Process batch of sequences
x = torch.randn(4, 1024, 512)
output = expert(x)

print(f"Input shape: {x.shape}")   # (4, 1024, 512)
print(f"Output shape: {output.shape}")  # (4, 1024, 512)
```

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

### Example: Basic Usage

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/rms-norm.md

Basic usage example of RMSNorm.

```python
import torch
from st_moe_pytorch.st_moe_pytorch import RMSNorm

# Create normalization layer
norm = RMSNorm(dim=512)

# Apply to batch of sequences
x = torch.randn(4, 1024, 512)
output = norm(x)

print(f"Input shape: {x.shape}")   # (4, 1024, 512)
print(f"Output shape: {output.shape}")  # (4, 1024, 512)

# Check that activations are roughly normalized
print(f"Input mean: {x.mean().item():.4f}, std: {x.std().item():.4f}")
print(f"Output mean: {output.mean().item():.4f}, std: {output.std().item():.4f}")
```

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

### Usage Example

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

Demonstrates how to instantiate and use the MoE module for sparse mixture of experts.

```python
import torch
from st_moe_pytorch import MoE

moe = MoE(
    dim = 512,
    num_experts = 16,               # increase the experts (# parameters) of your model without increasing computation
    gating_top_n = 2,               # default to top 2 gating, but can also be more (3 was tested in the paper with a lower threshold)
    threshold_train = 0.2,          # at what threshold to accept a token to be routed to second expert and beyond - 0.2 was optimal for 2 expert routing, and apparently should be lower for 3
    threshold_eval = 0.2,
    capacity_factor_train = 1.25,   # experts have fixed capacity per batch. we need some extra capacity in case gating is not perfectly balanced.
    capacity_factor_eval = 2.,      # capacity_factor_* should be set to a value >=1
    balance_loss_coef = 1e-2,       # multiplier on the auxiliary expert balancing auxiliary loss
    router_z_loss_coef = 1e-3,      # loss weight for router z-loss
)

inputs = torch.randn(4, 1024, 512)
out, total_aux_loss, balance_loss, router_z_loss = moe(inputs) # (4, 1024, 512), (1,), (1,), (1,)
```

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

### Gradient Accumulation

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Example of implementing gradient accumulation to simulate larger batch sizes with limited GPU memory.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch import MoE

moe = MoE(dim=512, num_experts=16, is_distributed=True)
optimizer = torch.optim.AdamW(moe.parameters())

accumulation_steps = 4

for i, (x, targets) in enumerate(dataloader):
    result = moe(x)
    loss = compute_loss(result.outputs, targets)
    
    # Scale loss for gradient accumulation
    loss = loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
```

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

### Example: Comparison with LayerNorm

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/rms-norm.md

Example comparing RMSNorm with standard LayerNorm.

```python
import torch
import torch.nn as nn
from st_moe_pytorch.st_moe_pytorch import RMSNorm

# RMSNorm version
rms_norm = RMSNorm(dim=512)

# Standard LayerNorm for comparison
layer_norm = nn.LayerNorm(512)

x = torch.randn(4, 1024, 512)

rms_output = rms_norm(x)
layer_output = layer_norm(x)

print(f"RMSNorm output shape: {rms_output.shape}")
print(f"LayerNorm output shape: {layer_output.shape}")

# Both have similar shapes but RMSNorm is more efficient
# (no mean subtraction, just RMS normalization)
```

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

### Basic Usage

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/moe.md

Example of basic usage for the MoE module.

```python
import torch
from st_moe_pytorch import MoE

# Create mixture of experts layer
moe = MoE(
    dim=512,
    num_experts=16,
    gating_top_n=2,
    threshold_train=0.2,
    threshold_eval=0.2,
    capacity_factor_train=1.25,
    capacity_factor_eval=2.0,
)

# Process input tokens
batch_size, seq_len, dim = 4, 1024, 512
x = torch.randn(batch_size, seq_len, dim)

# Forward pass returns all components for loss calculation
result = moe(x)
outputs = result.outputs
total_aux_loss = result.total_aux_loss

print(f"Outputs shape: {outputs.shape}")  # (4, 1024, 512)
print(f"Auxiliary loss: {total_aux_loss.item():.4f}")
```

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

### All-Gather Operations (Start of Expert Processing)

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Gathers all batch data from all ranks at the start of expert processing.

```python
# Gather all batch data from all ranks
gathered_x, batch_sizes = all_gather(x, dim=0)

# batch_sizes = [size_rank0, size_rank1, ..., size_rank_n]
total_batch = sum(batch_sizes)
```

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

### Example: Custom Feedforward with GEGLU

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/geglu.md

Illustrates creating a custom feedforward network using GEGLU.

```python
import torch
import torch.nn as nn
from st_moe_pytorch.st_moe_pytorch import GEGLU

class CustomFeedforward(nn.Module):
    def __init__(self, dim, hidden_mult=4):
        super().__init__()
        hidden_dim = int(dim * hidden_mult)
        
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim * 2),
            GEGLU(hidden_dim, mult_bias=True),
            nn.Linear(hidden_dim, dim)
        )
    
    def forward(self, x):
        return self.net(x)

# Test the custom feedforward
ff = CustomFeedforward(dim=512, hidden_mult=4)
x = torch.randn(4, 1024, 512)
output = ff(x)

print(f"Output shape: {output.shape}")  # (4, 1024, 512)
```

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

### SparseMoEBlock Constructor Options - Full Configuration

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/configuration.md

Example showing a full configuration of MoE and SparseMoEBlock, including setting various parameters and a forward pass.

```python
from st_moe_pytorch import MoE, SparseMoEBlock

moe = MoE(
    dim=512,
    num_experts=16,
    gating_top_n=2,
    threshold_train=0.2,
    threshold_eval=0.2,
    capacity_factor_train=1.25,
    capacity_factor_eval=2.0,
    balance_loss_coef=1e-2,
    router_z_loss_coef=1e-3,
)

moe_block = SparseMoEBlock(
    moe,
    add_ff_before=True,
    add_ff_after=True
)

import torch
x = torch.randn(4, 1024, 512)
outputs, total_aux_loss, balance_loss, router_z_loss = moe_block(x)
```

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

### Distributed Training Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/experts.md

An example demonstrating how to set up and use the Experts module in a distributed training environment.

```python
import torch
import torch.nn as nn
import torch.distributed as dist
from st_moe_pytorch.st_moe_pytorch import Experts, Expert

# Initialize distributed training
dist.init_process_group(backend='nccl')

expert_list = [Expert(dim=512) for _ in range(16)]

# Enable distributed mode
experts = Experts(expert_list, is_distributed=True, allow_var_seq_len=False)

# In distributed mode, each rank processes its assigned experts
x = torch.randn(2, 16, 32, 512).to(dist.get_rank() % torch.cuda.device_count())
outputs = experts(x)

dist.destroy_process_group()
```

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

### Example: Prenormalization Pattern

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/rms-norm.md

Example demonstrating the prenormalization pattern with RMSNorm in a TransformerLayer.

```python
import torch
import torch.nn as nn
from st_moe_pytorch.st_moe_pytorch import RMSNorm

class TransformerLayer(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = RMSNorm(dim)
        self.feed_forward = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
    
    def forward(self, x):
        # Prenorm pattern: normalize before processing
        normalized = self.norm(x)
        output = self.feed_forward(normalized)
        return output + x  # Residual connection

layer = TransformerLayer(dim=512)
x = torch.randn(4, 1024, 512)
output = layer(x)
```

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

### MoE Constructor Options - Different Thresholds

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/configuration.md

Example demonstrating how to set different probability thresholds for routing to non-primary experts during training and evaluation.

```python
from st_moe_pytorch import MoE

# Different thresholds for 2nd and 3rd experts
moe = MoE(
    dim=512,
    num_experts=16,
    gating_top_n=3,
    threshold_train=(0.2, 0.1),  # 2nd expert: 0.2, 3rd expert: 0.1
    threshold_eval=(0.2, 0.1),
)
```

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

### Example: Custom MoE with Different Experts

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/expert.md

Demonstrates creating custom experts with varying configurations and using them within a MoE module.

```python
import torch
import torch.nn as nn
from st_moe_pytorch import MoE
from st_moe_pytorch.st_moe_pytorch import Expert

# Create custom experts with different configurations
custom_experts = [
    Expert(dim=512, hidden_mult=4, mult_bias=True),
    Expert(dim=512, hidden_mult=6, mult_bias=True),  # Larger
    Expert(dim=512, hidden_mult=2, mult_bias=True),  # Smaller
]

# Replicate to get 16 total experts
custom_experts = custom_experts * 5 + custom_experts[:1]

# Use with MoE
moe = MoE(
    dim=512,
    num_experts=16,
    experts=nn.ModuleList(custom_experts)
)

x = torch.randn(4, 1024, 512)
result = moe(x)
```

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

### Example: In SparseMoEBlock

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/rms-norm.md

Example of RMSNorm usage within SparseMoEBlock.

```python
import torch
from st_moe_pytorch import MoE, SparseMoEBlock

moe = MoE(dim=512, num_experts=16)
moe_block = SparseMoEBlock(moe, add_ff_before=False, add_ff_after=False)

# Inside SparseMoEBlock, RMSNorm is applied before MoE
x = torch.randn(4, 1024, 512)
result = moe_block(x)

# The actual computation inside is:
# output = moe(norm(x)) + x
```

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

### Example: In Expert Architecture

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/geglu.md

Shows how GEGLU is used internally within an Expert module.

```python
import torch
import torch.nn as nn
from st_moe_pytorch.st_moe_pytorch import Expert, GEGLU

# Expert internally uses GEGLU
expert = Expert(dim=512)

# Equivalent to what happens inside Expert.forward():
# 1. Linear(512 → hidden * 2)
# 2. GEGLU(hidden)
# 3. Linear(hidden → 512)

x = torch.randn(4, 1024, 512)
output = expert(x)

print(output.shape)  # (4, 1024, 512)
```

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

### With Noise and Training Loop

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/moe.md

Example of using MoE with noise gates and a training loop.

```python
import torch
from st_moe_pytorch import MoE

moe = MoE(dim=512, num_experts=16)
optimizer = torch.optim.Adam(moe.parameters(), lr=1e-4)

x = torch.randn(4, 1024, 512)
targets = torch.randn(4, 1024, 512)

moe.train()

# Forward pass with noise gates for exploration
result = moe(x, noise_gates=True, noise_mult=1.0)
outputs = result.outputs

# Compute main task loss
main_loss = torch.nn.functional.mse_loss(outputs, targets)

# Add auxiliary losses
total_loss = main_loss + result.total_aux_loss

total_loss.backward()
optimizer.step()
optimizer.zero_grad()

```

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

### Example: In a Transformer

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/sparse-moe-block.md

Example of using SparseMoEBlock within a Transformer layer.

```python
import torch
import torch.nn as nn
from st_moe_pytorch import MoE, SparseMoEBlock

class TransformerLayer(nn.Module):
    def __init__(self, dim, num_heads, num_experts):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        
        moe = MoE(dim=dim, num_experts=num_experts)
        self.moe_block = SparseMoEBlock(moe, add_ff_before=False, add_ff_after=True)
        
        self.norm_attn = nn.LayerNorm(dim)
        self.norm_moe = nn.LayerNorm(dim)
    
    def forward(self, x):
        # Attention
        attn_out, _ = self.self_attn(x, x, x)
        x = x + attn_out
        x = self.norm_attn(x)
        
        # MoE block
        moe_result = self.moe_block(x)
        x = moe_result.outputs
        x = self.norm_moe(x)
        
        return x, moe_result.total_aux_loss

# Create a small transformer layer
layer = TransformerLayer(dim=512, num_heads=8, num_experts=16)
x = torch.randn(4, 1024, 512)
output, aux_loss = layer(x)

print(f"Output shape: {output.shape}")
print(f"Auxiliary loss: {aux_loss.item():.4f}")
```

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

### Basic All-Gather Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/distributed.md

Demonstrates a basic all-gather operation where each rank has a different batch size.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch.distributed import AllGather

# Initialize distributed training
dist.init_process_group(backend='nccl')

all_gather = AllGather(dim=0)

# Each rank has its own batch
if dist.get_rank() == 0:
    x = torch.randn(4, 512)  # Batch size 4
elif dist.get_rank() == 1:
    x = torch.randn(6, 512)  # Batch size 6
else:
    x = torch.randn(5, 512)  # Batch size 5

# Gather all batches
gathered, sizes = all_gather(x)

print(f"Local shape: {x.shape}")      # e.g., (4, 512)
print(f"Gathered shape: {gathered.shape}")  # (15, 512) = 4+6+5
print(f"Sizes per rank: {sizes}")     # tensor([4, 6, 5])

dist.destroy_process_group()
```

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

### Environment Variables

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Typical environment variables set by a job launcher for distributed training.

```bash
# Set by launcher
export MASTER_ADDR=localhost
export MASTER_PORT=29500
export RANK=0
export LOCAL_RANK=0
export WORLD_SIZE=4
```

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

### Example: With Prenormalization

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/expert.md

Creates an expert with prenormalization enabled, suitable for residual blocks.

```python
import torch
from st_moe_pytorch.st_moe_pytorch import Expert

# Expert with prenorm, useful for residual blocks
expert = Expert(dim=512, prenorm=True)

x = torch.randn(2, 256, 512)

# Prenorm is applied inside forward
output = expert(x)

# Can be used in residual: x + expert(x)
output_with_residual = x + expert(x)
```

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

### Parameter Effects on Routing: threshold_train Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/routing-mechanism.md

Illustrates how the `threshold_train` parameter controls the probability of routing to secondary experts, showing examples of probability calculation and clamping.

```python
For expert k with gate weight g:
    P(route) = min(1, g / 0.2)    
If g = 0.4: P = 2.0 → clamped to 1.0 (always route)
If g = 0.1: P = 0.5 (route 50% of time)
If g = 0.05: P = 0.25 (route 25% of time)
```

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

### Basic Usage

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/README.md

Basic usage example demonstrating how to create and use the MoE layer.

```python
import torch
from st_moe_pytorch import MoE

# Create mixture of experts layer
moe = MoE(
    dim=512,
    num_experts=16,
    gating_top_n=2,
    balance_loss_coef=1e-2,
    router_z_loss_coef=1e-3,
)

# Forward pass
x = torch.randn(4, 1024, 512)  # (batch, seq_len, dim)
result = moe(x)

# Unpack results
outputs = result.outputs
auxiliary_loss = result.total_aux_loss

# Training loop
main_loss = compute_loss(outputs, targets)
total_loss = main_loss + auxiliary_loss
total_loss.backward()
```

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

### Example: Keep Only One Expert on GPU

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/experts.md

Demonstrates how to keep only a specific expert on the GPU.

```python
experts.all_experts_to_cpu_besides(expert_index)
# Only experts[expert_index] remains on GPU
# All others moved to CPU
```

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

### All-Gather Overhead Calculation

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Formula and typical values for calculating all-gather overhead.

```python
# Time = gather_size × 2 / (bandwidth × num_hops)

# Typical (16 experts, batch=64, seq=1024, dim=512):
# - Single machine (8 GPUs): <1ms
# - Multi-machine: 10-50ms
# - Wide area network: 100-500ms
```

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

### Capacity Overflow Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/errors.md

Demonstrates a capacity overflow scenario in TopNGating by setting a very low capacity factor during training.

```python
import torch
from st_moe_pytorch import MoE

moe = MoE(
    dim=512,
    num_experts=16,
    capacity_factor_train=0.5  # Very low capacity
)

# Many tokens, few experts → capacity overflow
x = torch.randn(8, 1024, 512)  # Large batch, long sequence

result = moe(x)

# Some tokens may have been dropped due to capacity limits
# This can cause training instability
```

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

### Example: Evaluation Mode

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/moe.md

This snippet demonstrates how to use the MoE module in evaluation mode, highlighting the increased capacity factor and the absence of gating noise.

```python
import torch
from st_moe_pytorch import MoE

moe = MoE(
    dim=512,
    num_experts=16,
    capacity_factor_train=1.25,
    capacity_factor_eval=2.0,  # More capacity during evaluation
)

x = torch.randn(4, 1024, 512)

moe.eval()

# During eval, no gating noise and different thresholds/capacity are used
with torch.no_grad():
    result = moe(x)  # noise_gates defaults to False in eval mode
    outputs = result.outputs
```

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

### Pattern 1: Standard All-Gather-Route-Gather

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Illustrates a common distributed training pattern involving all-gathering inputs, routing through local experts, and then all-gathering outputs.

```python
# 1. All-gather input
all_x, batch_sizes = all_gather(x)

# 2. Route through experts (local on this rank)
expert_outputs = experts(routed_x)

# 3. All-gather outputs
all_outputs, _ = all_gather(expert_outputs, sizes=batch_sizes)

# 4. Split back to rank
rank_outputs = all_outputs.split(batch_sizes)[rank]
```

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

### Tune Capacity Factors

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Configures capacity factors for training and evaluation in the MoE layer.

```python
moe = MoE(
    dim=512,
    num_experts=16,
    capacity_factor_train=1.5,   # Balance between capacity and efficiency
    capacity_factor_eval=2.0,    # More generous during evaluation
)
```

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

### Auto-Detection

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

MoE auto-detects distributed training if torch.distributed is already initialized.

```python
from st_moe_pytorch import MoE

# MoE auto-detects distributed training if torch.distributed is initialized
moe = MoE(
    dim=512,
    num_experts=16,
    # is_distributed=None (default) — auto-detects
)

# Equivalent to setting is_distributed=True if dist.is_initialized()
```

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

### Integration with MoE

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/distributed.md

Example showing how the distributed utilities are used internally by MoE, particularly when `is_distributed` and `allow_var_seq_len` are enabled.

```python
moe = MoE(
    dim=512,
    num_experts=16,
    is_distributed=True,
    allow_var_seq_len=True
)

# MoE internally uses AllGather and related utilities
x = torch.randn(4, 1024, 512)  # Variable batch size per rank
result = moe(x)
```

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

### gather_sizes example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/distributed.md

Demonstrates gathering the sizes of a tensor along a specific dimension across different ranks.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch.distributed import gather_sizes

dist.init_process_group(backend='nccl')

x = torch.randn(dist.get_rank() + 2, 512)  # Variable sizes
sizes = gather_sizes(x, dim=0)

print(f"Sizes across ranks: {sizes}")

dist.destroy_process_group()
```

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

### MixtureOfExpertsReturn Usage Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/types.md

Demonstrates how to instantiate and use the MoE module, and how to access the outputs and auxiliary losses from the MixtureOfExpertsReturn object.

```python
import torch
from st_moe_pytorch import MoE

moe = MoE(
    dim=512,
    num_experts=16,
    balance_loss_coef=1e-2,
    router_z_loss_coef=1e-3,
)

x = torch.randn(4, 1024, 512)
result = moe(x)

output = result.outputs  # Shape: (4, 1024, 512)
total_loss = result.total_aux_loss  # Scalar tensor

# Add auxiliary loss to main loss during training
loss = compute_main_loss(output, targets) + result.total_aux_loss
```

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

### Synchronize All Ranks

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Demonstrates using `dist.barrier()` to block execution until all distributed processes reach a certain point, useful for debugging and synchronization.

```python
import torch.distributed as dist

# Block until all ranks reach this point
dist.barrier()

# Useful for debugging one rank at a time:
if dist.get_rank() == 0:
    print("Rank 0 debug info")
dist.barrier()
```

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

### All-Gather Operations (After Expert Processing)

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Gathers expert outputs back and splits them to per-rank data.

```python
# Gather expert outputs back
gathered_outputs, _ = all_gather(outputs, sizes=batch_sizes)

# Split gathered outputs back to per-rank data
rank_outputs = gathered_outputs.split(batch_sizes)[current_rank]
```

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

### Initial Model Broadcast and Output Gathering

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Broadcasts the initial model state dictionary to all ranks and gathers outputs after a distributed forward pass.

```python
dist.broadcast_object_list([model.state_dict()], src=0)

x = torch.randn(4, 512, 512).cuda()
out = model(x)

gathered_outs = [None] * dist.get_world_size()
dist.all_gather_object(gathered_outs, out.cpu())

if dist.get_rank() == 0:
    for i, out_i in enumerate(gathered_outs):
        print(f"Rank {i} output: {out_i.shape}")

dist.destroy_process_group()
```

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

### Minimal MoE Block (Expert FFN only)

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/sparse-moe-block.md

Example demonstrating the minimal configuration of a SparseMoEBlock, utilizing only the MoE routing without additional feedforward layers.

```python
import torch
from st_moe_pytorch import MoE, SparseMoEBlock

moe = MoE(dim=512, num_experts=8)

# Only MoE routing, no additional feedforward
moe_block = SparseMoEBlock(moe, add_ff_before=False, add_ff_after=False)

x = torch.randn(2, 256, 512)
result = moe_block(x)
```

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

### Validate Distributed Correctness

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

A placeholder snippet showing the initialization of a distributed process group and creating model copies, likely intended for validation checks.

```python
import torch
import torch.distributed as dist
from copy import deepcopy

dist.init_process_group(backend='nccl')

# Create model and gather all weights
model = create_model()
model_copy = deepcopy(model)
```

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

### Distributed Training with MoE

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/README.md

Example of initializing distributed training and using MoE with distributed support, including variable sequence lengths.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch import MoE

# Initialize distributed training
dist.init_process_group(backend='nccl')

# Create MoE with distributed support
moe = MoE(
    dim=512,
    num_experts=16,
    is_distributed=True,
    allow_var_seq_len=True
)

# Forward on different batch sizes per rank
local_batch_size = 4 + dist.get_rank()  # Variable
x = torch.randn(local_batch_size, 1024, 512)

result = moe(x)
output = result.outputs

dist.destroy_process_group()
```

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

### Device Mismatch Error Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/errors.md

Demonstrates a scenario that causes a RuntimeError due to tensors being on different devices (CPU vs. GPU) in distributed mode.

```python
import torch
from st_moe_pytorch import MoE

moe = MoE(dim=512).cuda()

# Trying to forward data on CPU to GPU model
x = torch.randn(4, 1024, 512)  # On CPU

try:
    result = moe(x)  # Model is on GPU, input is on CPU
except RuntimeError as e:
    print(f"Device error: {e}")
```

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

### Constructor

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/top-n-gating.md

Constructor for the TopNGating class.

```python
TopNGating(
    dim: int,
    num_gates: int,
    eps: float = 1e-9,
    top_n: int = 2,
    threshold_train: float | Tuple[float, ...] = 0.2,
    threshold_eval: float | Tuple[float, ...] = 0.2,
    capacity_factor_train: float = 1.25,
    capacity_factor_eval: float = 2.,
    straight_through_dispatch_tensor: bool = True,
    differentiable_topk: bool = False,
    differentiable_topk_fused: bool = True
)
```

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

### Memory-Efficient Attention

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/integration-patterns.md

Example of integrating memory-efficient attention variants from xformers.

```python
# Use memory-efficient attention variants
from xformers.ops import memory_efficient_attention

class EfficientMoEBlock(nn.Module):
    def forward(self, x):
        # Use memory-efficient attention
        attn = memory_efficient_attention(q, k, v)
        x = x + attn
        
        # Standard MoE
        result = self.moe(x)
        return result
```

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

### All-Gather with Pre-computed Sizes Example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/distributed.md

Demonstrates using pre-computed sizes for efficiency when performing all-gather operations, especially along a non-batch dimension.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch.distributed import AllGather, gather_sizes

dist.init_process_group(backend='nccl')

all_gather = AllGather(dim=-2)  # Gather along sequence dimension

x = torch.randn(2, 256, 512)  # (batch, seq, dim)

# Pre-compute sizes if gathering multiple times
sizes = gather_sizes(x, dim=-2)

# Now use pre-computed sizes for efficiency
gathered, _ = all_gather(x, sizes=sizes)

print(f"Gathered shape: {gathered.shape}")

dist.destroy_process_group()
```

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

### has_only_one_value example

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/distributed.md

Demonstrates checking if all elements in a tensor are identical.

```python
import torch
from st_moe_pytorch.distributed import has_only_one_value

sizes = torch.tensor([64, 64, 64])  # All same
print(has_only_one_value(sizes))  # True

sizes = torch.tensor([64, 65, 64])  # Different
print(has_only_one_value(sizes))  # False
```

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

### Pattern 1: Variable Batch Sizes

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/distributed.md

Example demonstrating how to handle variable batch sizes across distributed ranks using the AllGather utility.

```python
import torch
import torch.distributed as dist
from st_moe_pytorch.distributed import AllGather

dist.init_process_group(backend='nccl')

all_gather = AllGather(dim=0)

# Each rank has different batch size
batch_sizes = [4, 6, 5, 7]  # for 4 ranks
local_batch = batch_sizes[dist.get_rank()]
x = torch.randn(local_batch, 512)

gathered, sizes = all_gather(x)

# Now have all batches together
print(f"Total batch size: {gathered.shape[0]}")

dist.destroy_process_group()
```

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

### Constructor

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/api-reference/experts.md

Initializes the Experts module with a list of expert modules.

```python
Experts(
    experts: List[Module],
    is_distributed: bool | None = None,
    allow_var_seq_len: bool = False
)
```

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

### Transformer Layer with MoE

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/README.md

Example of integrating SparseMoEBlock into a standard Transformer layer.

```python
import torch
import torch.nn as nn
from st_moe_pytorch import MoE, SparseMoEBlock

class TransformerLayer(nn.Module):
    def __init__(self, dim, num_experts):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(dim, num_heads=8, batch_first=True)
        
        moe = MoE(dim=dim, num_experts=num_experts)
        self.moe_block = SparseMoEBlock(moe, add_ff_before=False, add_ff_after=True)
        
        self.norm_attn = nn.LayerNorm(dim)
        self.norm_moe = nn.LayerNorm(dim)
    
    def forward(self, x):
        # Attention branch
        attn_out, _ = self.self_attn(x, x, x)
        x = x + attn_out
        x = self.norm_attn(x)
        
        # MoE branch
        moe_result = self.moe_block(x)
        x = moe_result.outputs
        x = self.norm_moe(x)
        
        return x, moe_result.total_aux_loss

# Usage
layer = TransformerLayer(dim=512, num_experts=16)
x = torch.randn(4, 1024, 512)
output, aux_loss = layer(x)
```

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

### Hang During All-Gather Fix

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Ensures all ranks perform the same operations to prevent hangs during all-gather.

```python
# Ensure all ranks do the same operations
dist.barrier()
result = all_gather(x)
dist.barrier()
```

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

### Training with Gradient Accumulation

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/integration-patterns.md

Illustrates how to use gradient accumulation with SparseMoEBlock to effectively increase the batch size for training.

```python
from st_moe_pytorch import MoE, SparseMoEBlock

moe = MoE(dim=512, num_experts=16)
moe_block = SparseMoEBlock(moe)
optimizer = torch.optim.AdamW(moe_block.parameters())

accumulation_steps = 4

for epoch in range(num_epochs):
    for i, (x, targets) in enumerate(dataloader):
        result = moe_block(x)
        loss = compute_loss(result.outputs, targets)
        
        # Scale loss for accumulation
        loss = loss / accumulation_steps
        loss.backward()
        
        # Update weights every N steps
        if (i + 1) % accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(moe_block.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()
```

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

### Communication-Compute Overlap

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Enables overlap of computation and communication, and uses gradient checkpointing to reduce memory.

```python
# Enable overlap of computation and communication
torch.cuda.enable_peer_access(allow_tf32=True)

# Use gradient checkpointing to reduce memory
from torch.utils.checkpoint import checkpoint

output = checkpoint(moe_block, x, use_reentrant=False)
```

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

### Training with Learning Rate Scheduling

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/integration-patterns.md

Shows how to integrate a learning rate scheduler, such as CosineAnnealingLR, with SparseMoEBlock for stable training.

```python
from torch.optim.lr_scheduler import CosineAnnealingLR

moe_block = SparseMoEBlock(MoE(dim=512, num_experts=16))
optimizer = torch.optim.AdamW(moe_block.parameters(), lr=1e-4)

# Cosine annealing for stable training
scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)

for epoch in range(num_epochs):
    train_loss = 0
    
    for x, targets in dataloader:
        result = moe_block(x)
        loss = compute_loss(result.outputs, targets) + result.total_aux_loss
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item()
    
    scheduler.step()
    print(f"Epoch {epoch}: loss={train_loss/len(dataloader):.4f}")
```

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

### Custom Expert Networks

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/README.md

Example of creating and using custom expert networks within MoE.

```python
import torch.nn as nn
from st_moe_pytorch import MoE
from st_moe_pytorch.st_moe_pytorch import Expert

# Create custom experts
custom_experts = [
    Expert(dim=512, hidden_mult=4),
    Expert(dim=512, hidden_mult=6),  # Larger expert
    Expert(dim=512, hidden_mult=2),  # Smaller expert
]

moe = MoE(
    dim=512,
    num_experts=16,
    experts=nn.ModuleList(custom_experts * 5 + custom_experts[:1])
)
```

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

### Reduce All-Gather Frequency

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Increases gradient accumulation steps to reduce the frequency of expert routing all-gathers.

```python
# Increase gradient accumulation steps
accumulation_steps = 4

# This reduces frequency of expert routing all-gathers
```

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

### Variable Sequence Lengths

Source: https://github.com/lucidrains/st-moe-pytorch/blob/main/_autodocs/distributed-training.md

Enables variable sequence lengths per rank, with internal padding and restoration.

```python
# Each rank can have different sequence length
seq_lens = [256, 256, 512, 512]  # Different per rank

# Internally padded to max
padded_seq_len = 512

# After processing, padded back to original lengths

moe = MoE(
    dim=512,
    num_experts=16,
    is_distributed=True,
    allow_var_seq_len=True  # Enable variable seq len
)
```