[Task][Megatron] 3.3: Extend DeviceMesh beyond 2D

[lvyufeng]

Parent Issue

Part of #206 — Megatron Roadmap, Phase 3 (Megatron Core Execution MVP): Task 3.3

Description

Extend DeviceMesh to support >2 dimensions, needed for realistic TP × PP × DP × EP × CP layouts.

Requirements

• 3D DeviceMesh (TP × PP × DP)
• 4D DeviceMesh (TP × PP × DP × EP)
• 5D DeviceMesh (TP × PP × DP × EP × CP)
• Submesh extraction for each parallel dimension
• Process group creation from mesh dimensions
• Compatible with Megatron's initialize_model_parallel

Key Files

• src/candle/distributed/ (DeviceMesh)

Blocked By

• [Task][Megatron] 1.1: NCCL/HCCL backend routing strategy for Megatron #252 (Backend routing)

Blocks

• Multi-dimensional parallelism