MindSpeed框架理解

if not args.use_legacy_models:
        if args.spec is not None:
            transformer_layer_spec = import_module(args.spec)
        else:
            if use_te:
                transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(args.num_experts, args.moe_grouped_gemm)
            else:
                transformer_layer_spec = get_gpt_layer_local_spec(args.num_experts, args.moe_grouped_gemm)
        mtp_block_spec = None
        if args.mtp_num_layers is not None:
            mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec, use_transformer_engine=use_te)

        model = GPTModel(
            config=config,
            transformer_layer_spec=transformer_layer_spec,
            vocab_size=args.padded_vocab_size,
            max_sequence_length=args.max_position_embeddings,
            pre_process=pre_process,
            post_process=post_process,
            fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
            parallel_output=True,
            share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
            position_embedding_type=args.position_embedding_type,
            rotary_percent=args.rotary_percent,
            seq_len_interpolation_factor=args.rotary_seq_len_interpolation_factor,
            mtp_block_spec=mtp_block_spec,
        )
    else:
        if not args.context_parallel_size == 1:
            raise ValueError("Context parallelism is only supported with Megatron Core!")

        model = megatron.legacy.model.GPTModel(
            config,
            num_tokentypes=0,
            parallel_output=True,
            pre_process=pre_process,
            post_process=post_process
        )

def get_batch(data_iterator):
    """Generate a batch."""

    args = get_args()

    is_middle_stage = not (mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage())
    pretrain_not_tnd_flags = not args.is_instruction_dataset and not args.reset_position_ids
    if pretrain_not_tnd_flags and is_middle_stage:
        return (None,) * 5

    # get batches based on the TP rank you are on
    batch, actual_seq_len = get_batch_on_this_tp_rank(data_iterator)

    if args.return_document_ids and mpu.get_context_parallel_rank() == 0 and mpu.get_tensor_model_parallel_rank() == 0 and mpu.get_pipeline_model_parallel_rank() == 0:
        print("current idx: {}, current rank: {}, data_parallel_rank: {}, document_ids: {}".format(batch['idx'], torch.distributed.get_rank(), mpu.get_data_parallel_rank(), batch['document_ids']))
        batch.pop('document_ids', None)
        batch.pop('idx', None)

    if args.reset_position_ids and not args.reset_attention_mask:
        generate_actual_seq_len(batch, actual_seq_len)
        batch = get_batch_on_this_cp_rank(batch)
    else:
        # slice batch along sequence dimension for context parallelism
        batch = get_batch_on_this_cp_rank(batch)
    return batch.values()

def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
    """Loss function.

    Args:
        loss_mask (torch.Tensor): Used to mask out some portions of the loss
        output_tensor (torch.Tensor): The tensor with the losses
    """    
    args = get_args()

    losses = output_tensor.float()
    loss_mask = loss_mask.view(-1).float()
    if args.context_parallel_size > 1:
        loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), loss_mask.sum().view(1)])
        torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group())
        loss = loss[0] / loss[1]
    else:
        loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()

    # Check individual rank losses are not NaN prior to DP all-reduce.
    if args.check_for_nan_in_loss_and_grad:
        global_rank = torch.distributed.get_rank()
        if loss.isnan():
            raise ValueError(f'Rank {global_rank}: found NaN in local forward loss calculation. '
                             f'Device: {torch.cuda.current_device()}, node: {os.uname()[1]}')

    if args.async_log_allreduce:
        # Reduce loss for logging, which is different from megatron pretrain_gpt.py.
        reporting_loss = loss.clone().detach()

        allreduce_handle = torch.distributed.all_reduce(
            reporting_loss, group=mpu.get_data_parallel_group(), async_op=True
        )

        return loss * args.context_parallel_size, ({"lm loss": (reporting_loss)}, allreduce_handle)
    else:

        # Reduce loss for logging.
        averaged_loss = average_losses_across_data_parallel_group([loss])

        return loss * args.context_parallel_size, {'lm loss': averaged_loss[0]}

if args.context_parallel_size > 1:
    loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), 
                      loss_mask.sum().view(1)])
    torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group())
    loss = loss[0] / loss[1]

else:
    loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()

def forward_step(data_iterator, model: GPTModel):
    """Forward training step.

    Args:
        data_iterator : Input data iterator
        model (GPTModel): The GPT Model
    """
    args = get_args()
    timers = get_timers()

    # Get the batch.
    timers('batch-generator', log_level=2).start()
    tokens, labels, loss_mask, attention_mask, position_ids = get_batch(
        data_iterator)
    timers('batch-generator').stop()

    if args.use_legacy_models:
        output_tensor = model(tokens, position_ids, attention_mask,
                              labels=labels)
    else:
        output_tensor = model(tokens, position_ids, attention_mask,
                              labels=labels, loss_mask=loss_mask)

    return output_tensor, partial(loss_func, loss_mask)

#megatron 训练初始化。 主要是分布式环境。
def initialize_megatron(
    extra_args_provider=None,
    args_defaults={},
    ignore_unknown_args=False,
    allow_no_cuda=False,
    skip_mpu_initialization=False,
    get_embedding_ranks=None,
    get_position_embedding_ranks=None
):
    """Set global variables, initialize distributed, and
    set autoresume and random seeds.
    `allow_no_cuda` should not be set unless using megatron for cpu only
    data processing. In general this arg should not be set unless you know
    what you are doing.
    Returns a function to finalize distributed env initialization
    (optionally, only when args.lazy_mpu_init == True)
    """
    if not allow_no_cuda:
        # Make sure cuda is available.
        assert torch.cuda.is_available(), "Megatron requires CUDA."

    # Parse arguments 解析参数
    args = parse_args(extra_args_provider, ignore_unknown_args)

    if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
        assert args.load is not None, "--use-checkpoints-args requires --load argument"
        load_args_from_checkpoint(args)

    if args.yaml_cfg is not None:
        args = validate_yaml(args, args_defaults)
    else:
        validate_args(args, args_defaults)


    # set global args, build tokenizer, and set adlr-autoresume,
    # tensorboard-writer, and timers.
    set_global_variables(args) #解析的参数全局化。 基于 global 函数。

    # set logging level
    setup_logging() #配置日志

    # torch.distributed initialization
    def finish_mpu_init():
        args = get_args()
        # Pytorch distributed.  分布式初始化
        _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks)

        # Random seeds for reprod-ucibility.
        if args.rank == 0:
            print("> setting random seeds to {} ...".format(args.seed))
        _set_random_seed(args.seed, args.data_parallel_random_init)

    if skip_mpu_initialization:
        return None

    args = get_args()
    if args.lazy_mpu_init:
        # TODO is this still a necessary option?
        args.use_cpu_initialization = True
        # delayed initialization of DDP-related stuff
        # We only set basic DDP globals
        mpu.set_tensor_model_parallel_world_size(args.tensor_model_parallel_size)
        # and return function for external DDP manager
        # to call when it has DDP initialized
        mpu.set_tensor_model_parallel_rank(args.rank)
        return finish_mpu_init
    else:
        # Megatron's MPU is the master. Complete initialization right away.
        finish_mpu_init()

        # Autoresume.
        _init_autoresume()

        # Compile dependencies.
        _compile_dependencies()

        if args.tp_comm_overlap:
           _initialize_tp_communicators()

        # No continuation function
        return None

def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,         # 张量模型并行尺寸，默认1（不分割）
    pipeline_model_parallel_size: int = 1,       # 管道模型并行尺寸，默认1（不分割）
    virtual_pipeline_model_parallel_size: Optional[int] = None,  # 虚拟管道并行阶段数，用于层交错
    pipeline_model_parallel_split_rank: Optional[int] = None,    # 弃用参数，分割编码器/解码器的rank位置
    use_sharp: bool = False,                     # 是否启用SHARP通信优化
    context_parallel_size: int = 1,              # 上下文并行尺寸，分割输入序列长度
    expert_model_parallel_size: int = 1,         # 专家模型并行尺寸（MoE相关）
    nccl_communicator_config_path: Optional[str] = None,  # NCCL通信配置文件路径
    distributed_timeout_minutes: int = 30,       # 分布式操作超时时间（分钟）
    order: str = "tp-cp-ep-dp-pp",               # 并行维度初始化顺序
    encoder_pipeline_model_parallel_size: Optional[int] = None,  # 编码器专用管道并行尺寸
    get_embedding_ranks: Optional[Callable[[List[int], Optional[int]], List[int]]] = None,  # 自定义嵌入层rank分配函数
    get_position_embedding_ranks: Optional[Callable[[List[int], Optional[int]], List[int]]] = None,  # 自定义位置嵌入rank分配函数
) -> None:
    """初始化模型数据并行组，构建多种并行策略的进程组"""

    # 设置默认的嵌入层和位置嵌入层rank分配函数
    if get_embedding_ranks is None:
        get_embedding_ranks = partial(default_embedding_ranks, split_rank=pipeline_model_parallel_split_rank)
    
    if get_position_embedding_ranks is None:
        get_position_embedding_ranks = partial(default_position_embedding_ranks, split_rank=pipeline_model_parallel_split_rank)

    # 处理编码器专用管道并行设置
    if encoder_pipeline_model_parallel_size is not None:
        global _PIPELINE_MODEL_PARALLEL_DECODER_START
        _PIPELINE_MODEL_PARALLEL_DECODER_START = encoder_pipeline_model_parallel_size

    # 检查分布式环境是否已初始化
    assert torch.distributed.is_initialized()
    world_size: int = torch.distributed.get_world_size()

    # 验证总GPU数能被各并行尺寸整除
    if (world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size) != 0):
        raise RuntimeError(f"world_size必须能被各并行尺寸乘积整除")

    # 计算数据并行尺寸
    data_parallel_size: int = world_size // (
        tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size
    )

    # 验证专家并行尺寸有效性
    if data_parallel_size % expert_model_parallel_size != 0:
        raise RuntimeError(f"data_parallel_size必须能被expert_model_parallel_size整除")

    # 处理虚拟管道并行相关设置
    if virtual_pipeline_model_parallel_size is not None:
        if not pipeline_model_parallel_size > 1:
            raise RuntimeError("使用虚拟管道并行时管道尺寸必须大于1")
        # 设置全局虚拟管道并行状态
        global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK, _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
        _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = 0
        _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = virtual_pipeline_model_parallel_size

    # 处理弃用的管道分割参数
    if pipeline_model_parallel_split_rank is not None:
        global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
        _PIPELINE_MODEL_PARALLEL_SPLIT_RANK = pipeline_model_parallel_split_rank

    # 获取当前进程rank
    rank = torch.distributed.get_rank()

    # 加载NCCL通信配置
    nccl_comm_cfgs = {}
    if nccl_communicator_config_path is not None:
        with open(nccl_communicator_config_path, "r") as stream:
            nccl_comm_cfgs = yaml.safe_load(stream)

    # 初始化rank生成器（核心逻辑）
    rank_generator = RankGenerator(
        tp=tensor_model_parallel_size,
        ep=expert_model_parallel_size,
        dp=data_parallel_size,
        pp=pipeline_model_parallel_size,
        cp=context_parallel_size,
        order=order,  # 指定并行维度组合顺序
    )
    timeout = timedelta(minutes=distributed_timeout_minutes)

    # 构建数据并行组 ======================================================
    global _DATA_PARALLEL_GROUP, _DATA_PARALLEL_GROUP_GLOO, _DATA_PARALLEL_GLOBAL_RANKS
    assert _DATA_PARALLEL_GROUP is None, '数据并行组已初始化'
    
    # 生成所有数据并行组的rank列表
    for ranks in rank_generator.get_ranks('dp'):
        # 创建NCCL通信组和备用Gloo组
        group = torch.distributed.new_group(ranks, timeout=timeout, pg_options=get_nccl_options('dp', nccl_comm_cfgs))
        group_gloo = torch.distributed.new_group(ranks, timeout=timeout, backend="gloo")
        if rank in ranks:
            _DATA_PARALLEL_GROUP = group
            _DATA_PARALLEL_GROUP_GLOO = group_gloo
            _DATA_PARALLEL_GLOBAL_RANKS = ranks

    # 构建包含上下文并行的数据并行组
    for ranks_with_cp in rank_generator.get_ranks('dp-cp'):
        group_with_cp = torch.distributed.new_group(
            ranks_with_cp, timeout=timeout, pg_options=get_nccl_options('dp_cp', nccl_comm_cfgs)
        )
        if rank in ranks_with_cp:
            _DATA_PARALLEL_GROUP_WITH_CP = group_with_cp

    # 应用SHARP通信优化
    if use_sharp:
        # 使用SHARP需要同步数据并行组并设置环境变量
        torch.distributed.barrier(group=get_data_parallel_group(with_context_parallel=True))
        os.environ["NCCL_COLLNET_ENABLE"] = "0"  # 限制SHARP仅用于DP组

    # 构建上下文并行组 ====================================================
    global _CONTEXT_PARALLEL_GROUP
    for ranks in rank_generator.get_ranks('cp'):
        group = torch.distributed.new_group(ranks, timeout=timeout, pg_options=get_nccl_options('cp', nccl_comm_cfgs))
        if rank in ranks:
            _CONTEXT_PARALLEL_GROUP = group

    # 构建模型并行组（张量+管道）==========================================
    global _MODEL_PARALLEL_GROUP
    for ranks in rank_generator.get_ranks('tp-pp'):
        group = torch.distributed.new_group(ranks, timeout=timeout, pg_options=get_nccl_options('mp', nccl_comm_cfgs))
        if rank in ranks:
            _MODEL_PARALLEL_GROUP = group

    # 构建专家模型并行相关组 ==============================================
    global _MODEL_AND_EXPERT_PARALLEL_GROUP
    for ranks in rank_generator.get_ranks('tp-ep-pp', independent_ep=True):
        group = torch.distributed.new_group(ranks, pg_options=get_nccl_options('mp_exp', nccl_comm_cfgs))
        if rank in ranks:
            _MODEL_AND_EXPERT_PARALLEL_GROUP = group

    # 构建张量模型并行组 ==================================================
    global _TENSOR_MODEL_PARALLEL_GROUP
    for ranks in rank_generator.get_ranks('tp'):
        group = torch.distributed.new_group(ranks, timeout=timeout, pg_options=get_nccl_options('tp', nccl_comm_cfgs))
        if rank in ranks:
            _TENSOR_MODEL_PARALLEL_GROUP = group

    # 构建管道模型并行组及嵌入组 ==========================================
    global _PIPELINE_MODEL_PARALLEL_GROUP, _EMBEDDING_GROUP
    for ranks in rank_generator.get_ranks('pp'):
        # 管道组
        group = torch.distributed.new_group(ranks, timeout=timeout, pg_options=get_nccl_options('pp', nccl_comm_cfgs))
        if rank in ranks:
            _PIPELINE_MODEL_PARALLEL_GROUP = group
        
        # 嵌入组（自定义或默认）
        embedding_ranks = get_embedding_ranks(ranks)
        group = torch.distributed.new_group(embedding_ranks, pg_options=get_nccl_options('embd', nccl_comm_cfgs))
        if rank in embedding_ranks:
            _EMBEDDING_GROUP = group

        # 位置嵌入组
        position_embedding_ranks = get_position_embedding_ranks(ranks)
        group = torch.distributed.new_group(position_embedding_ranks, pg_options=get_nccl_options('embd', nccl_comm_cfgs))
        if rank in position_embedding_ranks:
            _POSITION_EMBEDDING_GROUP = group

    # 构建混合并行组（张量+数据）==========================================
    global _TENSOR_AND_DATA_PARALLEL_GROUP
    for ranks in rank_generator.get_ranks('tp-dp'):
        group = torch.distributed.new_group(ranks, pg_options=get_nccl_options('tp_dp', nccl_comm_cfgs))
        if rank in ranks:
            _TENSOR_AND_DATA_PARALLEL_GROUP = group

    # 初始化全局内存缓冲区 ================================================
    _set_global_memory_buffer()

def __init__(self, tp: int, ep: int, dp: int, pp: int, cp: int, order: str) -> None:
    # 参数说明：
    - tp: 张量并行 (Tensor Parallelism) 的组数
    - ep: 专家并行 (Expert Parallelism) 的组数（用于 MoE 模型）
    - dp
    : 数据并行 (Data Parallelism) 的组数
    - pp: 流水线并行 (Pipeline Parallelism) 的组数
    - cp: 上下文并行 (Context Parallelism) 的组数
    - order: 并行维度的优先级顺序（如 "tp-cp-ep-dp-pp"）

if 'ep' in order:
    if 'ep-dp' not in order and 'dp-ep' not in order:
        raise RuntimeError("EP 和 DP 必须在 order 中相邻")

for name in self.name_to_size.keys():
    if name not in order and self.name_to_size[name] != 1:
        raise RuntimeError(f"并行维度 {name} 的 size 不为 1，但未在 order 中指定")
    elif name not in order:
        order = order + '-' + name

self.order_w_ep = order  # 包含 EP 的原始 order
self.order_wo_ep = '-'.join([token for token in order.split('-') if token != 'ep'])  # 去除 EP 的 order

# 根据是否包含 EP 调整 DP 的 size
for token in order.split('-'):
    if token == 'dp':
        self.ordered_size_w_ep.append(self.dp // self.ep)  # DP 的 size 调整为 dp//ep
        self.ordered_size_wo_ep.append(self.dp)
    elif token == 'ep':
        self.ordered_size_w_ep.append(self.ep)
    else:
        self.ordered_size_w_ep.append(self.name_to_size[token])
        self.ordered_size_wo_ep.append(self.name_to_size[token])

if independent_ep:
    parallel_size = self.ordered_size_w_ep  # 包含 EP 的调整后 size
    order = self.order_w_ep                 # 包含 EP 的原始 order
else:
    parallel_size = self.ordered_size_wo_ep # 不包含 EP 的 size
    order = self.order_wo_ep                # 去除 EP 的 order

mask = self.get_mask(order, token)

ranks = generate_masked_orthogonal_rank_groups(self.world_size, parallel_size, mask)