Megablocks compatibility#
2025-09-05
3 min read time
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.
ROCm support for Megablocks is hosted in the official ROCm/megablocks repository.
Due to independent compatibility considerations, this location differs from the stanford-futuredata/Megatron-LM upstream repository.
Use the prebuilt Docker image with ROCm, PyTorch, and Megablocks preinstalled.
See the ROCm Megablocks installation guide to install and get started.
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 | 24.04 |