megatron-checkpoint-layout

Diagnose and convert Megatron checkpoint shard layouts across TP, PP, and EP dimensions.

1.2k|78|Updated Sep 16, 2025
One-click install
npx skills add https://github.com/EvolvingLMMs-Lab/LLaVA-OneVision-2 --skill megatron-checkpoint-layout
Or copy as Structured Prompt for Agent
Please help me install this Agent Skill.
Skill: megatron-checkpoint-layout
Source: https://github.com/EvolvingLMMs-Lab/LLaVA-OneVision-2/tree/main/.opencode/skills/megatron-checkpoint-layout
Command: npx skills add https://github.com/EvolvingLMMs-Lab/LLaVA-OneVision-2 --skill megatron-checkpoint-layout

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve?

Diagnose and convert Megatron checkpoint shard layouts across TP, PP, and EP dimensions to ensure correct loading and performance across distributed training setups.

Core Features & Use Cases

  • Layout discovery: Identify whether a checkpoint is 1D, 2D, or 3D based on the presence of EP and the sizes of TP/PP/EP.
  • Naming validation: Enforce recommended patterns mp_rank_{tp}, mp_rank_{tp}{pp}, mp_rank{tp}{pp}{ep}.
  • Migration guidance: Provide actionable steps to convert a layout when upgrading from non-EP to EP-enabled configurations or adjusting PP/EP sizes.

Quick Start

Run the layout validator on a Megatron checkpoint root to determine its dimensionality and the expected shard directory names.

Frequently Asked Questions about megatron-checkpoint-layout

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I identify the TP, PP, and EP dimensions of a Megatron checkpoint?

Convert non-EP Megatron checkpoints to EP-enabled layouts by following specific migration guidance. This process adjusts pipeline and expert parallelism sizes while enforcing the correct 3D mp_rank naming conventions for dense and MoE models.

When do I need to use a 3D-EP loader instead of a non-EP loader for Megatron checkpoints?

You need a 3D-EP loader when your Megatron checkpoint uses expert parallelism, indicated by a 3D layout. The diagnostic tool validates runtime arguments against on-disk shard layouts to determine whether a 3D-EP or non-EP loader is required.

What are the correct mp_rank naming conventions for Megatron-Core checkpoint shards?

The correct mp_rank naming conventions for Megatron-Core shards are mp_rank_{tp} for 1D, mp_rank_{tp}_{pp} for 2D, and mp_rank_{tp}_{pp}_{ep} for 3D layouts. The validator enforces these patterns to ensure proper distributed loading.

Does this layout validator work with Mixture of Experts (MoE) models in Megatron?

Yes, the layout validator works with Mixture of Experts (MoE) models in Megatron. It diagnoses and converts checkpoint shard layouts across TP, PP, and EP dimensions for both dense and MoE configurations.