1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

def pretrain(
    train_valid_test_dataset_provider,
    model_provider,
    forward_step_func,
    valid_forward_step_func=None,
    extra_args_provider=None,
    args_defaults={},
):  
    # 1.初始化分布式环境(源码解读1内容)
    initialize_megatron(
        extra_args_provider=extra_args_provider, args_defaults=args_defaults
    )
    ...
    # 2、模型并行：定义模型架构，并切割模型（本文重点）
    model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider)
    ...
    
    # 3、构造train/val/test数据集（下一篇将讲述）
    ... (
            train_data_iterator,
            valid_data_iterator,
            test_data_iterator,
        ) = build_train_valid_test_data_iterators(train_valid_test_dataset_provider) 
    
    ...
    # 4、训练（下下一篇将讲述）
    iteration = train(
            forward_step_func,
            valid_forward_step_func,
            model,
            optimizer,
            lr_scheduler,
            train_data_iterator,
            valid_data_iterator,
        )
    
    ...

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202

class _VocabParallelCrossEntropy(torch.autograd.Function):
    @staticmethod
    def forward(ctx, vocab_parallel_logits, target):
        """
        前向传播：计算词汇表并行下的交叉熵损失

        输入：
        - vocab_parallel_logits: [batch_size, seq_len, partition_vocab_size]
          当前 GPU 上维护的部分词汇表 logits
        - target: [batch_size, seq_len]
          真实标签（词汇表索引）

        输出：
        - loss: [batch_size, seq_len]
          每个位置的交叉熵损失
        """

        # ============ 第一步：数值稳定性处理 ============
        # 在词汇表维度上找到最大值（只在当前 GPU 的分区上找）
        # logits_max: [batch_size, seq_len]
        # 数学表达：m_i = max(s_i) 其中 s_i 是第 i 个样本的 logits
        logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]

        # 通过 AllReduce 获取全局最大值（跨所有 GPU）
        # 这样确保所有 GPU 使用相同的最大值进行归一化
        # 数学表达：m_global = max(m_1, m_2, ..., m_N) 其中 N 是 GPU 数量
        torch.distributed.all_reduce(
            logits_max,
            op=torch.distributed.ReduceOp.MAX,
            group=get_tensor_model_parallel_group(),
        )

        # 减去最大值，实现数值稳定
        # logits' = logits - max(logits)
        # 这样可以防止 exp(logits) 溢出，同时不改变 softmax 结果
        # 数学表达：s'_ij = s_ij - m_global
        vocab_parallel_logits.sub_(logits_max.unsqueeze(dim=-1))

        # ============ 第二步：确定当前 GPU 的词汇表范围 ============
        # 获取当前 GPU 负责的词汇表索引范围
        # 例如：V=10000, N=4, rank=0 -> [0, 2500), rank=1 -> [2500, 5000), ...
        get_vocab_range = VocabUtility.vocab_range_from_per_partition_vocab_size
        partition_vocab_size = vocab_parallel_logits.size()[-1]  # 当前 GPU 的词汇表分区大小
        rank = get_tensor_model_parallel_rank()  # 当前 GPU 的 rank
        world_size = get_tensor_model_parallel_world_size()  # 总 GPU 数量
        vocab_start_index, vocab_end_index = get_vocab_range(
            partition_vocab_size, rank, world_size
        )

        # ============ 第三步：创建目标掩码 ============
        # 判断哪些目标在当前 GPU 的词汇表范围内
        # target_mask=True 表示目标不在当前 GPU 的词汇表范围内
        # 数学表达：mask_i = 1 if (target_i < start) or (target_i >= end), else 0
        target_mask = (target < vocab_start_index) | (target >= vocab_end_index)

        # 将全局目标索引转换为当前 GPU 的局部索引
        # 数学表达：target'_i = target_i - vocab_start_index
        masked_target = target.clone() - vocab_start_index

        # 将不在当前范围内的目标设为 0（之后会被 mask 掉）
        masked_target[target_mask] = 0

        # ============ 第四步：提取目标位置的 logit ============
        # 将 logits 展平为 2D：[batch_size * seq_len, partition_vocab_size]
        # 将 target 展平为 1D：[batch_size * seq_len]
        # 这样方便使用高级索引
        logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
        masked_target_1d = masked_target.view(-1)

        # 创建行索引 [0, 1, 2, ..., batch_size * seq_len - 1]
        arange_1d = torch.arange(
            start=0, end=logits_2d.size()[0], device=logits_2d.device
        )

        # 提取每个样本对应目标位置的 logit
        # 数学表达：s'_target = logits[arange_1d, masked_target_1d]
        # 即取出每个样本在其目标词汇位置的 logit 值
        predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
        predicted_logits_1d = predicted_logits_1d.clone().contiguous()

        # 恢复原始形状：[batch_size, seq_len]
        predicted_logits = predicted_logits_1d.view_as(target)

        # 将不在当前 GPU 范围内的目标的 logit 设为 0
        # 这样做是因为之后会通过 AllReduce 求和，只需要贡献自己的部分
        predicted_logits[target_mask] = 0.0

        # 通过 AllReduce 从所有 GPU 收集目标的 logit
        # 每个 GPU 只贡献自己分区内的目标 logit
        # 最终得到所有目标位置的完整 logit 值
        # 数学表达：s_y = Σ(s'_y_i) 对所有 GPU 求和
        torch.distributed.all_reduce(
            predicted_logits,
            op=torch.distributed.ReduceOp.SUM,
            group=get_tensor_model_parallel_group(),
        )

        # ============ 第五步：计算归一化项（log-sum-exp）============
        # 计算 exp(logits')
        # 数学表达：e'_ij = exp(s'_ij)
        exp_logits = vocab_parallel_logits
        torch.exp(vocab_parallel_logits, out=exp_logits)

        # 在当前 GPU 的词汇表维度上求和
        # 数学表达：E'_i = Σ(e'_ij) 对当前 GPU 的词汇表维度求和
        sum_exp_logits = exp_logits.sum(dim=-1)

        # 通过 AllReduce 从所有 GPU 收集完整的指数和
        # 数学表达：E_i = Σ(E'_i) 对所有 GPU 求和
        # 这就是 softmax 的分母：Σ(exp(s_j))
        torch.distributed.all_reduce(
            sum_exp_logits,
            op=torch.distributed.ReduceOp.SUM,
            group=get_tensor_model_parallel_group(),
        )

        # ============ 第六步：计算最终损失 ============
        # 交叉熵损失：L = log(sum(exp(logits))) - predicted_logit
        # 数学表达：L = log(E_i) - s_y
        # 这对应：L = log(Σ(exp(s_j))) - s_y
        # 即：L = -log(exp(s_y) / Σ(exp(s_j))) = -softmax_y
        loss = torch.log(sum_exp_logits) - predicted_logits

        # ============ 第七步：保存反向传播所需的中间变量 ============
        # 计算 softmax 并保存，用于反向传播
        # 数学表达：softmax_ij = exp(s'_ij) / E_i
        exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))

        # 保存 softmax、target_mask 和 masked_target_1d 用于反向传播
        ctx.save_for_backward(exp_logits, target_mask, masked_target_1d)

        return loss

    @staticmethod
    def backward(ctx, grad_output):
        """
        反向传播：计算交叉熵损失对 logits 的梯度

        数学推导：
        前向传播：L = log(Σ(exp(s_j))) - s_y

        对 logits 的梯度：
        ∂L/∂s_i = softmax_i - 1_{i=y}

        其中：
        - softmax_i = exp(s_i) / Σ(exp(s_j))
        - 1_{i=y} 是指示函数，当 i=y 时为 1，否则为 0

        在词汇表并行的场景下：
        - 当前 GPU 只计算自己分区内的梯度
        - 对于目标位置：grad = softmax - 1
        - 对于非目标位置：grad = softmax
        """

        # ============ 第一步：恢复前向传播保存的中间变量 ============
        # softmax: [batch_size, seq_len, partition_vocab_size]
        # target_mask: [batch_size, seq_len] (布尔掩码)
        # masked_target_1d: [batch_size * seq_len] (局部索引)
        softmax, target_mask, masked_target_1d = ctx.saved_tensors

        # ============ 第二步：初始化梯度 ============
        # 交叉熵损失的梯度基础是 softmax 值
        # 对于所有位置，梯度从 softmax 开始
        # 数学表达：∂L/∂s_i = softmax_i （基础项）
        grad_input = softmax

        # 为了方便索引，转换为 2D 格式
        # grad_2d: [batch_size * seq_len, partition_vocab_size]
        partition_vocab_size = softmax.size()[-1]
        grad_2d = grad_input.view(-1, partition_vocab_size)

        # ============ 第三步：处理目标位置的梯度 ============
        # 对于目标位置，需要减去 1（指示函数）
        # 数学表达：∂L/∂s_y = softmax_y - 1
        # 其中 1 是真实标签的指示函数 1_{i=y}

        # 创建行索引
        arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)

        # 对目标位置的梯度进行调整：
        # 1. 如果目标在当前 GPU 的词汇表范围内（target_mask=False）：
        #    grad = softmax - 1 （减 1）
        # 2. 如果目标不在当前范围内（target_mask=True）：
        #    grad = softmax - 0 （不减 1，因为目标在其他 GPU 上）
        # 数学表达：
        # ∂L/∂s_y = softmax_y - (1 - mask_y)
        # 其中 mask_y = 1 表示目标不在当前 GPU，mask_y = 0 表示目标在当前 GPU
        grad_2d[arange_1d, masked_target_1d] -= 1.0 - target_mask.view(-1).float()

        # ============ 第四步：乘以上层梯度 ============
        # 根据链式法则，需要乘以损失对输出的梯度
        # 数学表达：∂L_total/∂s = (∂L/∂s) × (∂L_total/∂L)
        # 其中 grad_output 是 ∂L_total/∂L
        grad_input.mul_(grad_output.unsqueeze(dim=-1))

        # 返回梯度（第二个输入 target 不需要梯度，返回 None）
        return grad_input, None


def vocab_parallel_cross_entropy(vocab_parallel_logits, target):
    """Helper function for the cross entropy."""
    return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target)