Megatron-Core 大模型训练框架

from megatron.core import transformer

config = transformer.TransformerConfig(
    num_layers=32,                  # Transformer 层数
    hidden_size=4096,               # 隐藏层维度
    num_attention_heads=32,         # 注意力头数
    num_query_groups=32,            # GQA 中的 KV 组数（等于 num_attention_heads 为标准 MHA）
    kv_channels=128,                # 每个注意力头的维度 (= hidden_size / num_attention_heads)
    ffn_hidden_size=11008,          # FFN 中间层维度（通常为 hidden_size 的 2.7~4 倍）
    normalization="LayerNorm",      # 归一化方式："LayerNorm" 或 "RMSNorm"
    norm_epsilon=1e-5,              # 归一化的 epsilon
    # --- 并行策略 ---
    tensor_model_parallel_size=1,   # 张量并行度（TP）
    pipeline_model_parallel_size=1, # 流水线并行度（PP）
    context_parallel_size=1,        # 上下文并行度（CP）
    # --- 混合精度 ---
    bf16=True,                      # 是否使用 BF16
    fp8=False,                      # 是否启用 FP8 训练（需要 Hopper GPU）
    fp8_format="hybrid",            # FP8 格式："hybrid" 或 "e4m3"
    # --- 序列并行 ---
    sequence_parallel=True,         # 是否启用序列并行
    # --- 其他 ---
    apply_residual_connection_post_norm=False,  # 残差连接是否在归一化之后
    hidden_dropout=0.1,             # 隐藏层 dropout
    attention_dropout=0.1,          # 注意力 dropout
    kv_channels=128,                # KV 通道数
    rotary_percent=1.0,             # RoPE 应用的比例
    add_bias_linear=False,          # 线性层是否添加偏置
    gated_linear_unit=True,         # 是否使用门控线性单元（SwiGLU）
    activation_func="silu",         # 激活函数
)

from megatron.core import mpu

# 初始化并行组
mpu.initialize_model_parallel(
    tensor_model_parallel_size=4,    # TP 大小
    pipeline_model_parallel_size=2,  # PP 大小
    virtual_pipeline_model_parallel_size=None,  # 虚拟流水线并行（ interleaved PP）
    context_parallel_size=1,         # CP 大小
)

# 获取各并行维度的 world size 和 rank
tp_rank = mpu.get_tensor_model_parallel_rank()       # 当前 TP rank
tp_world = mpu.get_tensor_model_parallel_world_size() # TP world size

pp_rank = mpu.get_pipeline_model_parallel_rank()     # 当前 PP rank
pp_world = mpu.get_pipeline_model_parallel_world_size()

dp_rank = mpu.get_data_parallel_rank()               # 当前 DP rank
dp_world = mpu.get_data_parallel_world_size()

cp_rank = mpu.get_context_parallel_rank()            # 当前 CP rank
cp_world = mpu.get_context_parallel_world_size()

from megatron.core.tensor_parallel import ColumnParallelLinear

col_linear = ColumnParallelLinear(
    input_size=4096,        # 输入维度
    output_size=11008,      # 输出维度（完整值，会自动按 tp 切分）
    config=config,          # TransformerConfig 实例
    init_method=None,       # 权重初始化方法
    bias=False,             # 是否有偏置
    gather_output=False,    # 是否在输出后 all-gather 汇聚（False=保持切分状态）
    skip_bias_add=False,    # 是否跳过偏置加法（返回偏置供融合操作使用）
)

from megatron.core.tensor_parallel import RowParallelLinear

row_linear = RowParallelLinear(
    input_size=11008,          # 输入维度（完整值，会自动按 tp 切分）
    output_size=4096,          # 输出维度
    config=config,             # TransformerConfig 实例
    init_method=None,          # 权重初始化方法
    bias=False,                # 是否有偏置
    input_is_parallel=True,    # 输入是否已经是按 tp 切分的状态
    skip_bias_add=False,       # 是否跳过偏置加法
)

from megatron.core.tensor_parallel.mappings import TensorParallelMode
from megatron.core.transformer.mlp import ParallelMLP

mlp = ParallelMLP(
    config=config,  # TransformerConfig 实例
)

from megatron.core.transformer.attention import ParallelSelfAttention

attn = ParallelSelfAttention(
    config=config,  # TransformerConfig 实例
)

from megatron.core.transformer.transformer_layer import TransformerLayer

layer = TransformerLayer(
    config=config,             # TransformerConfig 实例
    layer_number=1,            # 层编号（用于 RoPE 等位置相关计算）
    self_attn=None,            # 可选：自定义注意力模块
    mlp=None,                  # 可选：自定义 MLP 模块
)

from megatron.core.transformer.transformer_block import TransformerBlock

block = TransformerBlock(
    config=config,               # TransformerConfig 实例
    transformer_layer_list=None, # 可选：自定义层列表
    post_layer_norm=True,        # 是否在最后添加 LayerNorm
)

from megatron.core.tensor_parallel import (
    TensorParallelMode,
    get_tensor_model_parallel_group,
)

# TensorParallelMode 枚举值：
# - Column: 列并行模式（用于 QKV 投影、MLP 第一个线性层）
# - Row: 行并行模式（用于输出投影、MLP 第二个线性层）

from megatron.core.tensor_parallel import (
    copy_to_tensor_model_parallel_region,   # 恒等操作（前向），all-reduce（反向）
    gather_from_tensor_model_parallel_region, # all-gather（前向），恒等（反向）
    reduce_from_tensor_model_parallel_region, # reduce-scatter（前向），恒等（反向）
    scatter_to_tensor_model_parallel_region,  # 恒等（前向），all-reduce（反向）
)

import torch
from megatron.core import mpu
from megatron.core.transformer import TransformerConfig, TransformerBlock

# Step 1: 初始化分布式环境
torch.distributed.init_process_group(backend="nccl")

# Step 2: 初始化模型并行组
mpu.initialize_model_parallel(
    tensor_model_parallel_size=4,     # 4路张量并行
    pipeline_model_parallel_size=1,   # 不使用流水线并行
    context_parallel_size=1,          # 不使用上下文并行
)

# Step 3: 创建模型配置
config = TransformerConfig(
    num_layers=32,
    hidden_size=4096,
    num_attention_heads=32,
    num_query_groups=8,               # GQA: 8组KV头
    ffn_hidden_size=11008,
    seq_length=4096,
    normalization="RMSNorm",
    bf16=True,
    tensor_model_parallel_size=4,
    sequence_parallel=True,           # 启用序列并行
    gated_linear_unit=True,           # SwiGLU FFN
    add_bias_linear=False,
)

# Step 4: 构建模型
model = TransformerBlock(config=config)

# Step 5: 前向传播
# 输入: [seq_len, batch_size, hidden_size]
input_ids = torch.randint(0, 32000, (4096, 2)).cuda()
hidden_states = torch.randn(4096, 2, 4096, dtype=torch.bfloat16).cuda()
output = model(hidden_states)

import torch
import torch.nn as nn
from megatron.core import mpu
from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
from megatron.core.transformer import TransformerConfig

torch.distributed.init_process_group(backend="nccl")
mpu.initialize_model_parallel(tensor_model_parallel_size=4)

config = TransformerConfig(
    num_layers=1,
    hidden_size=4096,
    num_attention_heads=32,
    tensor_model_parallel_size=4,
    bf16=True,
)

# 手动构建一个张量并行的两层 MLP
class TensorParallelMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense_h_to_4h = ColumnParallelLinear(
            input_size=config.hidden_size,          # 4096
            output_size=config.ffn_hidden_size,     # 11008，自动按4切分为2752
            config=config,
            gather_output=False,   # 不汇聚，保持切分状态传给下一层
            bias=False,
        )
        self.dense_4h_to_h = RowParallelLinear(
            input_size=config.ffn_hidden_size,      # 11008，输入已按4切分
            output_size=config.hidden_size,         # 4096
            config=config,
            input_is_parallel=True,  # 输入已经是切分状态
            bias=False,
        )
        self.activation = nn.SiLU()

    def forward(self, x):
        # x: [seq, batch, hidden] - 每个 GPU 持有完整副本
        intermediate = self.dense_h_to_4h(x)   # 每GPU: [seq, batch, ffn/tp]
        intermediate = self.activation(intermediate)
        output = self.dense_4h_to_h(intermediate)  # All-Reduce → [seq, batch, hidden]
        return output

model = TensorParallelMLP(config).cuda()
x = torch.randn(512, 2, 4096, dtype=torch.bfloat16).cuda()
y = model(x)  # 输出形状: [512, 2, 4096]，所有GPU结果一致

import torch
from megatron.core import mpu
from megatron.core.transformer import TransformerConfig, TransformerBlock

torch.distributed.init_process_group(backend="nccl")
mpu.initialize_model_parallel(
    tensor_model_parallel_size=2,
    pipeline_model_parallel_size=4,  # 4 级流水线
)

config = TransformerConfig(
    num_layers=32,
    hidden_size=4096,
    num_attention_heads=32,
    tensor_model_parallel_size=2,
    pipeline_model_parallel_size=4,
    bf16=True,
)

# TransformerBlock 自动根据 PP rank 获取对应的层
# PP rank 0: layer 0-7, PP rank 1: layer 8-15, ...
model = TransformerBlock(config=config).cuda()

# 在实际训练中，使用 Megatron 训练框架处理流水线调度
# 此处展示概念性前向传播
pp_rank = mpu.get_pipeline_model_parallel_rank()
num_layers_per_stage = config.num_layers // mpu.get_pipeline_model_parallel_world_size()

print(f"Pipeline stage {pp_rank}: handles layers "
      f"{pp_rank * num_layers_per_stage} - {(pp_rank + 1) * num_layers_per_stage - 1}")

import torch
from megatron.core import mpu
from megatron.core.transformer import TransformerConfig, TransformerLayer

torch.distributed.init_process_group(backend="nccl")
mpu.initialize_model_parallel(tensor_model_parallel_size=4)

config = TransformerConfig(
    num_layers=32,
    hidden_size=4096,
    num_attention_heads=32,
    tensor_model_parallel_size=4,
    sequence_parallel=True,   # 关键：启用序列并行
    bf16=True,
)

# 当 sequence_parallel=True 时:
# - LayerNorm/Dropout 的输入/输出沿序列维度切分
# - 每个 GPU 只持有 1/4 的序列长度
# - 注意力计算前 all-gather 恢复完整序列
# - 注意力计算后 reduce-scatter 重新切分

layer = TransformerLayer(config=config, layer_number=1).cuda()

# SP 模式下的输入: [seq_len/tp, batch, hidden]
# 例如完整序列 4096, tp=4 → 每个 GPU 输入 seq_len=1024
sp_seq_len = 4096 // 4  # = 1024
hidden_states = torch.randn(sp_seq_len, 2, 4096, dtype=torch.bfloat16).cuda()
output = layer(hidden_states)  # 输出同样为 [seq_len/tp, batch, hidden]

import torch
from megatron.core import mpu
from megatron.core.transformer import TransformerConfig, TransformerBlock

torch.distributed.init_process_group(backend="ncll")
mpu.initialize_model_parallel(
    tensor_model_parallel_size=2,
    context_parallel_size=4,  # 4路上下文并行
)

config = TransformerConfig(
    num_layers=32,
    hidden_size=4096,
    num_attention_heads=32,
    tensor_model_parallel_size=2,
    context_parallel_size=4,
    seq_length=131072,         # 128K 超长序列
    bf16=True,
)

model = TransformerBlock(config=config).cuda()

# CP 模式下：每个 GPU 处理 seq_len/cp 的序列片段
# 通过 Ring Attention 机制在 CP 组内通信
# 实现超长序列的高效训练
cp_seq_len = 131072 // 4  # = 32768
hidden_states = torch.randn(cp_seq_len, 1, 4096, dtype=torch.bfloat16).cuda()
output = model(hidden_states)

import torch
from megatron.core import mpu
from megatron.core.transformer import TransformerConfig, TransformerBlock
from megatron.core.optimizer import get_megatron_optimizer, OptimizerConfig

# --- 初始化 ---
torch.distributed.init_process_group(backend="nccl")
mpu.initialize_model_parallel(
    tensor_model_parallel_size=4,
    pipeline_model_parallel_size=2,
)

# --- 配置 ---
model_config = TransformerConfig(
    num_layers=32,
    hidden_size=4096,
    num_attention_heads=32,
    ffn_hidden_size=11008,
    tensor_model_parallel_size=4,
    pipeline_model_parallel_size=2,
    sequence_parallel=True,
    bf16=True,
)

optim_config = OptimizerConfig(
    optimizer="adam",           # 优化器类型
    lr=1e-4,                    # 学习率
    weight_decay=0.01,          # 权重衰减
    adam_beta1=0.9,             # Adam beta1
    adam_beta2=0.95,            # Adam beta2
    adam_eps=1e-8,              # Adam epsilon
    clip_grad=1.0,              # 梯度裁剪阈值
)

# --- 构建模型 ---
model = TransformerBlock(config=model_config).cuda()

# --- 构建优化器 ---
optimizer = get_megatron_optimizer(optim_config, [model])

# --- 训练循环 ---
for step in range(1000):
    # 生成模拟数据
    batch_size = 4
    seq_len = 2048
    input_ids = torch.randint(0, 32000, (seq_len, batch_size)).cuda()
    labels = torch.randint(0, 32000, (seq_len, batch_size)).cuda()

    # 前向传播
    optimizer.zero_grad()
    hidden_states = model(input_ids)  # 简化示例，实际需要embedding层

    # 计算损失（简化）
    loss = hidden_states.float().mean()

    # 反向传播
    loss.backward()

    # 梯度裁剪与参数更新
    optimizer.step()

    if step % 100 == 0:
        print(f"Step {step}, Loss: {loss.item():.4f}")

# 选择性激活重计算配置
config = TransformerConfig(
    # ...
    recompute_granularity="selective",  # "full" 或 "selective"
    recompute_method="block",           # "uniform" 或 "block"
    recompute_num_layers=8,             # 重计算的层数
)

# 在自定义框架中使用 Megatron-Core 的并行层
from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
from megatron.core.transformer import TransformerConfig

# 只使用并行层，不依赖完整训练框架
config = TransformerConfig(
    num_layers=1, hidden_size=4096, num_attention_heads=32,
    tensor_model_parallel_size=4, bf16=True,
)
linear = ColumnParallelLinear(4096, 11008, config=config, gather_output=False)