Megablocks compatibility#

2025-07-31

3 min read time

Applies to Linux

Megablocks is a light-weight library for mixture-of-experts (MoE) training. The core of the system is efficient “dropless-MoE” and standard MoE layers. Megablocks is integrated with stanford-futuredata/Megatron-LM, where data and pipeline parallel training of MoEs is supported.

Note

Megablocks is supported on ROCm 6.3.0.

Supported devices#

  • Officially Supported: AMD Instinct MI300X

  • Partially Supported (functionality or performance limitations): AMD Instinct MI250X, MI210X

Supported models and features#

This section summarizes the Megablocks features supported by ROCm.

  • Distributed Pre-training

  • Activation Checkpointing and Recomputation

  • Distributed Optimizer

  • Mixture-of-Experts

  • dropless-Mixture-of-Experts

Use cases and recommendations#

The ROCm Megablocks blog posts guide how to leverage the ROCm platform for pre-training using the Megablocks framework. It features how to pre-process datasets and how to begin pre-training on AMD GPUs through:

  • Single-GPU pre-training

  • Multi-GPU pre-training

Docker image compatibility#

AMD validates and publishes ROCm Megablocks images with ROCm and Pytorch backends on Docker Hub. The following Docker image tags and associated inventories represent the latest Megatron-LM version from the official Docker Hub. The Docker images have been validated for ROCm 6.3.0. Click to view the image on Docker Hub.

Docker image

ROCm

Megablocks

PyTorch

Ubuntu

Python

rocm/megablocks

6.3.0

0.7.0

2.4.0

24.04

3.12.9