← Back to reference

Megatron Integration Guide

Install

Megatron integration ships in denpex_sdk.megatron. The denpex package detects Megatron via import megatron and auto-attaches.

# At the top of your pretrain_*.sh script:
import denpex.megatron  # noqa

Or:

pip install denpex[megatron]

Capture

Per TransformerLayer we register forward hooks that record:

  • tensor_parallel_size
  • pipeline_parallel_size
  • context_parallel_size
  • sequence_parallel_size
  • virtual_pipeline_model_parallel_size
  • microbatch_id
  • step_id
  • pipeline_bubble_seconds (idle time per rank)

What We Detect

  • TP all-reduce hang (rank-isolated).
  • PP bubble (bubble > 30s).
  • Sequence parallel mismatch (size drift).
  • Virtual pipeline imbalance.

Compatibility Tested

  • Megatron-LM core_r0.6, core_r0.7.
  • Transformer Engine 1.13+.

Limitations

  • Megatron-Core v0.7+ uses a different layer API; pending.

Debug Recipes

docs/user-guides/megatron-recipes.md:

  • "all-reduce on cross-node hangs" → check NCCL_IB_HCA.
  • "bubble 100 % on rank 0" → TP configuration skew.

Integration with mpu

When mpu.get_tensor_model_parallel_world_size() is available we use its values; otherwise we degrade to env vars (TP_SIZE, PP_SIZE, etc.).