### Python Doctest Example with Setup

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Illustrates writing a Python function with a docstring that includes a `.. testsetup::` block for fixtures and a `.. testcode::` block for the executable example.

```python
import torch
from typing import Optional

def my_function(x: Optional[torch.Tensor]) -> torch.Tensor:
    """blah function

    Args:
        input (torch.Tensor): Your guess.

    Returns:
        torch.Tensor: How good your input is.

    Raises:
        ValueError: If your input is negative.

    Example:
        .. testsetup::

            # optional setup section, not shown in docs
            import torch
            x = torch.randn(42)


        .. testcode::

            # shown in docs; runs after testsetup
            my_function(x)
    """
    ...
```

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

### Install and Configure Pre-Commit Hooks

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Install development dependencies and set up pre-commit hooks to enforce style checks before each commit.

```bash
pip install '.[dev]'
pre-commit install
```

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

### Install development dependencies

Source: https://github.com/databricks/megablocks/blob/main/CONTRIBUTING.md

Installs the necessary packages for testing and linting the MegaBlocks codebase.

```bash
pip install -e '.[all]'
```

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

### Complete Training Loop with MegaBlocks MoE

Source: https://context7.com/databricks/megablocks/llms.txt

Example of a complete training loop integrating MegaBlocks MoE layers. Includes model definition, optimizer setup, and basic forward pass.

```python
import torch
import torch.nn as nn
from functools import partial
from megablocks import dMoE, Arguments
from megablocks.layers.moe import batched_load_balancing_loss, clear_load_balancing_loss

class TransformerBlockWithMoE(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.attention = nn.MultiheadAttention(
            args.hidden_size,
            num_heads=8,
            batch_first=False,
            dtype=torch.bfloat16
        )
        self.moe = dMoE(args)
        self.norm1 = nn.LayerNorm(args.hidden_size, dtype=torch.bfloat16)
        self.norm2 = nn.LayerNorm(args.hidden_size, dtype=torch.bfloat16)

    def forward(self, x):
        # Self-attention with residual
        attn_out, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_out)

        # MoE with residual
        moe_out = self.moe(x)
        if isinstance(moe_out, tuple):
            moe_out, bias = moe_out
            if bias is not None:
                moe_out = moe_out + bias
        x = self.norm2(x + moe_out)
        return x

# Configuration
args = Arguments(
    hidden_size=512,
    ffn_hidden_size=2048,
    moe_num_experts=8,
    moe_top_k=2,
    moe_loss_weight=0.1,
    mlp_impl='grouped',
    memory_optimized_mlp=True,
    bf16=True,
    init_method=partial(torch.nn.init.normal_, mean=0.0, std=0.02),
)

# Create model
model = TransformerBlockWithMoE(args).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

```

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

### Clone and setup repository

Source: https://github.com/databricks/megablocks/blob/main/CONTRIBUTING.md

Clones the fork and adds the original MegaBlocks repository as an upstream remote.

```bash
git clone git@github.com:<github_id>/megablocks.git
cd megablocks
git remote add upstream https://github.com/databricks/megablocks.git
```

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

### Python Function Docstring Examples

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Demonstrates various ways to document function arguments, default values, and return types in Python docstrings.

```python
from typing import Optional, Union

def foo(bar: int):
    """Foo.

    Args:
        bar (int): Required bar.
    """
    ...
```

```python
def foo2(bar: int = 42):
    """Foo2.

    Args:
        bar (int, optional): The first Argument. Default: ``42``.
    """
    ...
```

```python
def foo3(bar: Optional[int] = None):
    """Foo3.

    Args:
        bar (int, optional): The first Argument. Default: ``None``.
    """
    ...
```

```python
def foo4(bar: Union[int, str] = 42):
    """Foo4.

    Args:
        bar (int | str, optional): The first Argument. Default: ``42``.
    """
    ...
```

```python
def foo5(bar: int) -> int:
    """Foo5.

    Args:
        bar (int): Required bar.

    Returns:
        int: Description of return statement.
    """
    ...
```

```python
def foo6(bar: int) -> tuple[int, str]:
    """Foo6.

    Args:
        bar (int): Required bar.

    Returns:
        a (int): Returned value.
        b (str): Returned value.
    """
    ...
```

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

### Train MegaBlocks dMoE Model with Megatron-LM

Source: https://context7.com/databricks/megablocks/llms.txt

Example bash script to train a distributed dMoE model using Megatron-LM. Configures distributed training, model parameters, and MoE-specific settings.

```bash
#!/bin/bash
# Example: Train 125M dMoE model on 8 GPUs

EXP_DIR="./experiments/dmoe_125m"
mkdir -p ${EXP_DIR}

# MoE configuration
NUM_EXPERTS=64
TOP_K=1
LOSS_WEIGHT=0.1

# Distributed configuration
DISTRIBUTED_ARGS="
    --nproc_per_node 8
    --nnodes 1
    --master_addr localhost
    --master_port 6000
"

# Model configuration
MODEL_ARGS="
    --num-layers 12
    --hidden-size 768
    --num-attention-heads 12
    --seq-length 1024
    --max-position-embeddings 1024
"

# MoE-specific arguments
MOE_ARGS="
    --moe-num-experts ${NUM_EXPERTS}
    --moe-loss-weight ${LOSS_WEIGHT}
    --moe-top-k ${TOP_K}
"

# Training configuration
TRAINING_ARGS="
    --micro-batch-size 32
    --global-batch-size 512
    --train-iters 20000
    --lr 0.0006
    --min-lr 0.00006
    --lr-decay-style cosine
    --lr-warmup-fraction 0.01
    --clip-grad 1.0
"

# Compute configuration
COMPUTE_ARGS="
    --bf16
    --moe-expert-model-parallelism
    --use-flash-attn
"

torchrun ${DISTRIBUTED_ARGS} \
    third_party/Megatron-LM/pretrain_gpt.py \
    ${MOE_ARGS} \
    ${MODEL_ARGS} \
    ${TRAINING_ARGS} \
    ${COMPUTE_ARGS} \
    --save ${EXP_DIR} \
    --save-interval 2000
```

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

### Training Loop with MegaBlocks

Source: https://context7.com/databricks/megablocks/llms.txt

This snippet demonstrates a typical training loop using MegaBlocks for mixture-of-experts models. It includes forward and backward passes, loss computation with load balancing, and optimizer steps. Ensure necessary imports and model/optimizer setup are done prior to this loop.

```python
for step in range(100):
    # Generate sample data [seq_len, batch_size, hidden_size]
    x = torch.randn(1024, 16, 512, dtype=torch.bfloat16, device='cuda')

    # Forward pass
    output = model(x)

    # Compute loss with load balancing
    task_loss = output.pow(2).mean()  # Example task loss
    lbl_loss = batched_load_balancing_loss(args)
    total_loss = task_loss + lbl_loss

    # Backward pass
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

    # Clear loss state
    clear_load_balancing_loss()

    if step % 10 == 0:
        print(f"Step {step}: task_loss={task_loss.item():.4f}, lbl_loss={lbl_loss:.4f}")
```

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

### Conditionally Import Optional Dependencies

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Use MissingConditionalImportError to handle optional dependencies that are not installed, preventing runtime errors for minimal installations.

```python
from composer import Callback
from composer.utils import MissingConditionalImportError

class SystemMetricsMonitor(Callback)
    try:
        import pynvml
    except ImportError as e:
        raise MissingConditionalImportError(extra_deps_group="system_metrics_monitor",
                                            conda_package="pynvml",
                                            conda_channel="conda-forge",) from e
```

```python
from composer.utils import MissingConditionalImportError

try:
    import deepspeed
except ImportError as e:
    raise MissingConditionalImportError(extra_deps_group="deepspeed",
                                        conda_package="deepspeed>=0.5.5",
                                        conda_channel=None) from e
```

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

### Configure dMoE for Expert Model Parallelism

Source: https://context7.com/databricks/megablocks/llms.txt

Configure a dMoE layer for expert model parallelism, distributing experts across multiple GPUs. This requires distributed training setup and process group creation. The forward pass automatically handles all-to-all communication.

```python
import torch
import torch.distributed as dist
from megablocks import dMoE, Arguments
from megablocks.layers.moe import clear_load_balancing_loss

# Initialize distributed training (must be called first)
# dist.init_process_group(backend='nccl')

# Create expert parallel process group
# world_size = dist.get_world_size()
# expert_parallel_group = dist.new_group(list(range(world_size)))

# Configure for expert parallelism
args = Arguments(
    hidden_size=768,
    ffn_hidden_size=3072,
    moe_num_experts=64,  # Large number of experts
    moe_top_k=2,

    # Enable expert model parallelism
    moe_expert_model_parallelism=True,
    # expert_parallel_group=expert_parallel_group,  # Uncomment in distributed setting

    # Pipeline parallelism settings
    pipeline_model_parallel_size=1,
    num_layers_per_virtual_pipeline_stage=None,

    mlp_impl='grouped',
    bf16=True,
)

# In distributed setting, each rank will only hold a subset of experts
dmoe_layer = dMoE(args)
dmoe_layer.cuda().to(torch.bfloat16)

# Forward pass handles all-to-all communication automatically
x = torch.randn(1024, 16, 768, dtype=torch.bfloat16, device='cuda')
output, _ = dmoe_layer(x)

clear_load_balancing_loss()
```

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

### Run Doctests Locally

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Commands to activate the virtual environment, navigate to the docs directory, build HTML documentation, and then run doctests.

```bash
source path/to/megablocks_venv/bin/activate  # activate your megablocks virtual env
cd megablocks/docs  # cd to the docs folder insde your megablocks clone
make clean
make html  # the html build must be completed first to ensure all doctests are identified
make doctest 2>/dev/null # For more verbosity, do not direct stderr to /dev/null
```

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

### Serve Local Documentation

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Command to serve the generated HTML documentation locally using Python's http.server module.

```bash
cd megablocks/docs
python3 -m http.server --directory _build/html/
```

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

### Activate Virtual Environment and Build Docs

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Commands to activate a Python virtual environment, navigate to the docs directory, clean previous builds, and generate HTML documentation.

```bash
source path/to/megablocks_venv/bin/activate  # activate your megablocks virtual env
cd megablocks/docs  # cd to the docs folder insde your megablocks clone
make clean
make html
```

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

### Trainer Device Parameter Type Hinting

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Demonstrates how to use `Union` for type hinting to allow string or custom Device types for parameters, simplifying user imports.

```python
from typing import Optional, Union

from composer.devices import Device

class Trainer:
    def __init__(
        self,
        device: Union[str, Device],
    ):
        if isinstance(device, str):
            device = Device(device)
        ...
```

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

### Configure pre-commit hooks

Source: https://github.com/databricks/megablocks/blob/main/CONTRIBUTING.md

Initializes pre-commit to automatically format code before each commit.

```bash
pre-commit install
```

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

### Expose Public APIs in __init__.py

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Add all public classes and functions to the module's __init__.py file to simplify imports for users.

```python
from composer.path.to.module.file import MyClass as MyClass
from composer.path.to.module.file import my_func as my_func
```

```python
from composer.path.to.module import my_file as my_file
```

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

### Configure and Use SparseGLU Layer

Source: https://context7.com/databricks/megablocks/llms.txt

Configure a GLU-based MoE layer with specified arguments and perform a forward pass. Ensure the input tensor and tokens_per_expert are on the correct device and dtype.

```python
args = Arguments(
    hidden_size=512,
    ffn_hidden_size=1024,
    moe_num_experts=4,
    moe_top_k=1,
    mlp_impl='grouped',
    mlp_type='glu',
    bf16=True,
)

# Create SparseGLU layer
glu_layer = SparseGLU(args)
glu_layer.cuda().to(torch.bfloat16)

# Forward pass requires tokens and tokens_per_expert
batch_size, seq_len, hidden = 16, 1024, 512
x = torch.randn(batch_size * seq_len, hidden, dtype=torch.bfloat16, device='cuda')

# tokens_per_expert specifies how many tokens each expert processes
tokens_per_expert = torch.tensor([batch_size * seq_len], device='cuda')

output = glu_layer(x, tokens_per_expert)
print(f"GLU output shape: {output.shape}")  # [batch*seq, hidden]
```

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

### Run Yapf and Isort Formatters Manually

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Manually trigger yapf for general code formatting and isort for import sorting using pre-commit.

```bash
pre-commit run yapf --all-files
pre-commit run isort --all-files
```

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

### Convert Megatron-LM Arguments to MegaBlocks Arguments

Source: https://context7.com/databricks/megablocks/llms.txt

Utilize the `from_megatron` utility to convert Megatron-LM configuration arguments to MegaBlocks Arguments for seamless integration. This is useful when migrating or using both libraries.

```python
from megablocks.layers.arguments import Arguments, from_megatron

# Example Megatron-LM args (typically from argparse)
class MegatronArgs:
    hidden_size = 768
    ffn_hidden_size = 3072
    num_layers = 12
    moe_num_experts = 8
    moe_top_k = 2
    moe_loss_weight = 0.1
    fp16 = False
    bf16 = True

megatron_args = MegatronArgs()

# Convert to MegaBlocks Arguments
megablocks_args = from_megatron(megatron_args)
print(megablocks_args)
```

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

### Run Pyright Type Checker Manually

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Manually execute the pyright type checker on all files in the repository using pre-commit.

```bash
pre-commit run pyright --all-files
```

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

### Configure MoE Arguments

Source: https://context7.com/databricks/megablocks/llms.txt

Configure MoE layer behavior including model dimensions, expert settings, parallelism, and compute options using the Arguments dataclass. Set moe_capacity_factor to 0 for dropless behavior.

```python
from functools import partial
import torch
from megablocks import Arguments

# Basic MoE configuration
args = Arguments(
    # Model dimensions
    hidden_size=768,
    ffn_hidden_size=3072,
    num_layers=12,
    bias=True,

    # MoE configuration
    moe_num_experts=8,
    moe_top_k=2,
    moe_capacity_factor=1,  # Set to 0 for dropless behavior
    moe_normalize_expert_weights=1,  # L1 normalization
    moe_loss_weight=0.1,  # Load balancing loss weight
    moe_jitter_eps=0.01,  # Router jitter for training stability

    # Compute settings
    mlp_impl='grouped',  # 'grouped' (recommended) or 'sparse'
    mlp_type='mlp',  # 'mlp' or 'glu'
    memory_optimized_mlp=True,

    # Precision
    fp16=False,
    bf16=True,

    # Initialization
    init_method=partial(torch.nn.init.normal_, mean=0.0, std=0.02),
    device=torch.cuda.current_device(),
)
```

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

### Run Pre-Commit Hooks Manually

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Execute all pre-commit hooks on changed or all files in the repository.

```bash
pre-commit run
pre-commit run --all-files
```

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

### Create a new branch

Source: https://github.com/databricks/megablocks/blob/main/CONTRIBUTING.md

Creates a new local branch for implementing changes.

```bash
git checkout -b cool-new-feature
```

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

### Data Validation with ValueError

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Illustrates how to perform data validation using `ValueError` exceptions, as `assert` statements can be disabled.

```python
if parameter is None:
    raise ValueError("parameter must be specified and cannot be None")
```

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

### Configure dMoE with Shared Expert

Source: https://context7.com/databricks/megablocks/llms.txt

Configure a dMoE layer to include a shared expert that processes all tokens alongside routed experts. The shared expert's output is automatically combined with the routed expert output. Ensure to clear load balancing loss after use.

```python
import torch
from megablocks import dMoE, Arguments
from megablocks.layers.moe import clear_load_balancing_loss

# Configure dMoE with shared expert
args = Arguments(
    hidden_size=512,
    ffn_hidden_size=2048,
    moe_num_experts=8,
    moe_top_k=2,

    # Shared expert configuration
    shared_expert=True,
    shared_expert_hidden_size=1024,  # Can differ from ffn_hidden_size
    shared_expert_weighted_sum=True,  # Weight by number of experts

    mlp_impl='grouped',
    bf16=True,
)

dmoe_layer = dMoE(args)
dmoe_layer.cuda().to(torch.bfloat16)

# The shared expert output is automatically combined with routed expert output
x = torch.randn(1024, 16, 512, dtype=torch.bfloat16, device='cuda')
output, _ = dmoe_layer(x)
print(f"Output with shared expert: {output.shape}")

clear_load_balancing_loss()
```

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

### Configure and Use dMoE with Router Z-Loss

Source: https://context7.com/databricks/megablocks/llms.txt

Configure a dMoE layer with router z-loss enabled for training stability. Ensure losses are combined and gradients are computed correctly. Remember to clear both load balancing and z-loss accumulators after backward pass.

```python
import torch
from megablocks import dMoE, Arguments
from megablocks.layers.moe import batched_load_balancing_loss, clear_load_balancing_loss
from megablocks.layers.router import batched_router_zloss, clear_router_zloss

# Configure with z-loss enabled
args = Arguments(
    hidden_size=512,
    ffn_hidden_size=2048,
    moe_num_experts=8,
    moe_top_k=2,
    moe_loss_weight=0.1,
    moe_zloss_weight=1e-3,  # Enable router z-loss
    moe_zloss_in_fp32=True,  # Compute z-loss in FP32 for stability
    mlp_impl='grouped',
    bf16=True,
)

dmoe_layer = dMoE(args)
dmoe_layer.cuda().to(torch.bfloat16)

# Training with both load balancing and z-loss
x = torch.randn(1024, 16, 512, dtype=torch.bfloat16, device='cuda')
x.requires_grad_(True)

output, _ = dmoe_layer(x)

# Combine all losses
lbl_loss = batched_load_balancing_loss(args)
zloss = batched_router_zloss(args)
total_loss = output.sum() + lbl_loss + zloss.sum()
total_loss.backward()

# Cleanup both loss accumulators
clear_load_balancing_loss()
clear_router_zloss()
```

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

### Avoid No-Op Functions with Falsy Values

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Demonstrates the anti-pattern of checking for falsy values within a function, leading to a no-op, and the preferred pattern where the caller handles such checks.

```python
from typing import Optional

def configure_deepspeed(deepspeed_config: Optional[dict]):
    if deepspeed_config is None:
        # Don't do this check in the callee, which results in a no-op
        return
    ...
```

```python
from typing import Optional

def configure_deepspeed(deepspeed_config: dict):
    ...

def trainer(deepspeed_config: Optional[dict]):
    if deepspeed_config is not None:
        # Do this check in the caller function
        configure_deepspeed(deepspeed_config)
    ...
```

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

### Optional Sequence Type Hinting with ensure_tuple

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Shows how to use `Optional[Union[Tensor, Sequence[Tensor]]]` for parameters that accept single tensors, sequences of tensors, or None, and how `ensure_tuple` normalizes the input.

```python
from torch import Tensor
from typing import Optional, Sequence, Union
from composer.utils import ensure_tuple

def foo(x: Optional[Union[Tensor, Sequence[Tensor]]]) -> tuple[Tensor, ...]:
    return ensure_tuple(x)  # ensures that the result is always a (potentially empty) tuple of tensors
```

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

### SparseGLU Layer

Source: https://context7.com/databricks/megablocks/llms.txt

Utilize the SparseGLU class for MoE experts, which provides a Gated Linear Unit (GLU) variant with an additional gating mechanism for enhanced model expressiveness.

```python
import torch
from megablocks import Arguments, SparseGLU
```

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

### Handle Union Types with Type Checks

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Illustrates how to safely perform operations on variables that can be one of multiple types (e.g., int or None) by adding explicit checks.

```python
from typing import Union

def foo(x: Union[int, None]):
    if x is None:
        raise TypeError("x must be an integer, not None!")
    return x + 5  # valid
```

```python
from typing import Union

def foo(x: Union[int, None]):
    assert x is not None, "x should never be None"
    return x + 5  # valid
```

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

### Define Public Exports with __all__

Source: https://github.com/databricks/megablocks/blob/main/STYLE_GUIDE.md

Explicitly define __all__ in public modules to control re-exports and ensure documentation tools only include intended members.

```python
"""Log memory usage during training."""
import logging
from typing import Union

import torch.cuda

from composer.core import State
from composer.loggers import Logger
from composer.core.callback import Callback

log = logging.getLogger(__name__)

__all__ = ["MemoryMonitor"]  # export only the MemoryMonitor, not other imports like `Logger`, `State`, or `Callback`


class MemoryMonitor(Callback):
    ...
```

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

### Convert Megatron Arguments to MegaBlocks Arguments

Source: https://context7.com/databricks/megablocks/llms.txt

Converts arguments from Megatron-LM format to MegaBlocks format for use in MegaBlocks models. Ensure Megatron arguments are properly loaded before conversion.

```python
args = from_megatron(megatron_args)
print(f"Converted args: hidden_size={args.hidden_size}, experts={args.moe_num_experts}")
```

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

### Standard MoE Layer Forward Pass

Source: https://context7.com/databricks/megablocks/llms.txt

Implement a standard mixture-of-experts layer with configurable capacity factor and token dropping. This layer combines a learned router with parallel expert computation. Ensure to compute and accumulate the load balancing loss during training.

```python
import torch
from megablocks import MoE, Arguments, get_load_balancing_loss
from megablocks.layers.moe import clear_load_balancing_loss, batched_load_balancing_loss

# Configure MoE layer
args = Arguments(
    hidden_size=512,
    ffn_hidden_size=2048,
    moe_num_experts=8,
    moe_top_k=2,
    moe_capacity_factor=1.25,  # Allow 25% overflow
    moe_loss_weight=0.1,
    mlp_impl='grouped',
    bf16=True,
)

# Create MoE layer
moe_layer = MoE(args)
moe_layer.cuda().to(torch.bfloat16)

# Forward pass: input shape is [sequence_length, batch_size, hidden_size]
x = torch.randn(1024, 16, 512, dtype=torch.bfloat16, device='cuda')
output = moe_layer(x)

# Output maintains same shape as input
# output: tuple of (tensor, bias) if return_bias=True, else tensor
out_tensor, bias = output
print(f"Output shape: {out_tensor.shape}")  # [1024, 16, 512]

# Compute load balancing loss for training
loss = out_tensor.sum() + batched_load_balancing_loss(args)
loss.backward()

# Clear accumulated loss state after backward pass
clear_load_balancing_loss()
```

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

### Dropless MoE (dMoE) Layer Forward Pass

Source: https://context7.com/databricks/megablocks/llms.txt

Implement a dropless mixture-of-experts layer using block-sparse operations to avoid token dropping. This layer maintains computational efficiency through optimized sparse matrix operations. The capacity_factor argument is ignored for dMoE.

```python
import torch
from megablocks import dMoE, Arguments
from megablocks.layers.moe import batched_load_balancing_loss, clear_load_balancing_loss

# Configure dMoE layer (capacity_factor is ignored)
args = Arguments(
    hidden_size=512,
    ffn_hidden_size=2048,
    moe_num_experts=16,
    moe_top_k=2,
    moe_loss_weight=0.1,
    mlp_impl='grouped',  # Use grouped GEMM for Hopper GPUs
    memory_optimized_mlp=True,  # Enable memory optimization
    bf16=True,
)

# Create dMoE layer
dmoe_layer = dMoE(args)
dmoe_layer.cuda().to(torch.bfloat16)

# Training forward pass
x = torch.randn(1024, 16, 512, dtype=torch.bfloat16, device='cuda')
x.requires_grad_(True)

output = dmoe_layer(x)
out_tensor, bias = output

# Training with load balancing loss
loss = out_tensor.sum() + batched_load_balancing_loss(args)
loss.backward()

# Access gradients
print(f"Input gradient shape: {x.grad.shape}")

# Cleanup
dmoe_layer.zero_grad(set_to_none=True)
clear_load_balancing_loss()
```

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

### MegaBlocks Low-Level Operations API

Source: https://context7.com/databricks/megablocks/llms.txt

Provides low-level sparse operations for advanced use cases. Requires PyTorch and MegaBlocks.Ops. Ensure CUDA is available for GPU operations.

```python
import torch
import megablocks.ops as ops

# Histogram: count tokens per expert
top_experts = torch.tensor([0, 1, 0, 2, 1, 0], dtype=torch.int32, device='cuda')
num_experts = 4
tokens_per_expert = ops.histogram(top_experts, num_experts)
print(f"Tokens per expert: {tokens_per_expert}")  # tensor([3, 2, 1, 0])
```

```python
# Sort: get indices for expert-grouped permutation
sorted_experts, indices = ops.sort(top_experts, end_bit=2)
print(f"Sorted: {sorted_experts}, Indices: {indices}")
```

```python
# Cumulative sum for bin boundaries
bins = ops.inclusive_cumsum(tokens_per_expert, dim=0)
print(f"Bin boundaries: {bins}")  # tensor([3, 5, 6, 6])
```

```python
# Round up to blocking factor
blocking = 128
padded = ops.round_up(tokens_per_expert, blocking)
print(f"Padded token counts: {padded}")
```

```python
# Gather and scatter for token routing
x = torch.randn(6, 512, device='cuda')
gathered = ops.gather(x, indices, sorted_experts, bins, top_k=1)
```

```python
# Scatter back with expert weights
expert_weights = torch.ones(6, device='cuda')
scattered = ops.scatter(gathered, indices, sorted_experts, expert_weights, bins, top_k=1)
```

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