megalodon-jax 0.2.1

JAX/Equinox implementation of Megalodon for efficient long-context training and inference

megalodon-jax

A JAX/Equinox implementation of Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length, aligned with the original released PyTorch/CUDA implementation.

Highlights

• Long-context language modeling with chunked attention and ComplexEMA
• JIT-compatible streaming inference with fixed-capacity caches
• Packed-sequence training with document-isolated attention, recurrent state, normalization, and loss
• SafeTensors checkpoints and original-upstream PyTorch checkpoint conversion
• Source-level parity checks and the official released tokenizer bundle

Installation

For CPU:

pip install megalodon-jax

For NVIDIA GPUs using the pip-managed CUDA 13 runtime:

pip install "megalodon-jax[cuda13]"

Python 3.11 or newer is required. See Installation and setup for CUDA requirements, optional features, unreleased main installs, development setup, and troubleshooting.

Quick start

import jax
from megalodon_jax import MegalodonConfig, MegalodonForCausalLM

model_key, input_key = jax.random.split(jax.random.PRNGKey(0))
config = MegalodonConfig(
    vocab_size=32_000,
    model_dim=512,
    num_layers=8,
    num_heads=1,
    chunk_size=256,
    cema_ndim=16,
)
model = MegalodonForCausalLM(config, key=model_key)

input_ids = jax.random.randint(input_key, (1, 128), 0, config.vocab_size)
logits, _ = model(input_ids)
print(logits.shape)  # (1, 128, 32000)

Where to go next

• Installation and setup: CUDA compatibility, optional extras, source installs, and development setup
• Long-context streaming: generation, cache behavior, sliding attention, padding, and packed training
• Dtypes and numerical stability: precision policies, training, memory-bounded loss, and data parallelism
• JAX and PyTorch interoperability: checkpoints, conversion, resume state, and parity gates
• Paper and source differences: named model presets and deliberate compatibility choices
• Development: tests, release process, correctness verification, and benchmarks

The documentation index covers the remaining implementation and compatibility references.

Related

• Original Megalodon - released PyTorch/CUDA source

Citation

@misc{ma2024megalodon,
      title={Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length},
      author={Xuezhe Ma and Xiaomeng Yang and Wenhan Xiong and Beidi Chen and Lili Yu and Hao Zhang and Jonathan May and Luke Zettlemoyer and Omer Levy and Chunting Zhou},
      year={2024},
      eprint={2404.08801},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}