Megatron-LM 详细代码架构分析

# 模型架构参数 num_layers: int = 0 hidden_size: int = 0 num_attention_heads: int = 0 ffn_hidden_size: Optional[int] = None # 注意力机制配置 attention_backend: AttnBackend = AttnBackend.auto num_query_groups: Optional[int] = None softmax_scale: Optional[float] = None # MoE (Mixture of Experts) 配置 num_moe_experts: Optional[int] = None moe_router_load_balancing: str = "aux_loss" # 高级特性 multi_latent_attention: bool = False rotary_interleaved: bool = False window_size: Optional[Tuple[int, int]] = None # 并行化配置 sequence_parallel: bool = False expert_model_parallel_size: int = 1

from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_layer import BaseTransformerLayer class TransformerBlock(MegatronModule): def __init__( self, config: TransformerConfig, spec: ModuleSpec, num_layers: int, post_layer_norm: bool = True, pre_process: bool = True, post_process: bool = True, ): # 计算当前pipeline stage需要构建的层数 self.num_layers_per_pipeline_rank = get_num_layers_to_build(config) # 构建transformer层 self.layers = torch.nn.ModuleList() for i in range(self.num_layers_per_pipeline_rank): layer = build_module(spec, config=i+1) self.layers.append(layer)

def forward( self, hidden_states: Tensor, attention_mask: Tensor, context: Optional[Tensor] = None, context_mask: Optional[Tensor] = None, ) -> Tensor: # 应用前处理层 if self.pre_process: hidden_states = self.pre_process_block(hidden_states) # 逐层处理 for index, layer in enumerate(self.layers): hidden_states = layer( hidden_states, attention_mask, context=context, context_mask=context_mask, ) # 应用后处理层 if self.post_process: hidden_states = self.post_process_block(hidden_states) return hidden_states def get_num_layers_to_build(config: TransformerConfig) -> int: # 支持不均匀的pipeline stage层数分配 if config.num_layers_in_first_pipeline_stage is not None: # 第一stage的特殊处理 if parallel_state.is_pipeline_first_stage(ignore_virtual=True): return config.num_layers_in_first_pipeline_stage # 默认均匀分配 return config.num_layers // parallel_state.get_pipeline_model_parallel_world_size()

class SelfAttention(MegatronModule): def __init__( self, config: TransformerConfig, submodules: SelfAttentionSubmodules, layer_number: int, attn_mask_type=AttnMaskType.padding, ): # QKV投影层 self.linear_qkv = submodules.linear_qkv # 核心注意力层 self.core_attention = submodules.core_attention # 输出投影层 self.linear_proj = submodules.linear_proj # 查询和键的LayerNorm if config.qk_layernorm: self.q_layernorm = submodules.q_layernorm self.k_layernorm = submodules.k_layernorm

def backward_dw(self) -> NoReturn: """Execute weight update operations""" try: self._backward_qkv_proj() self._backward_output_proj() except Exception as e: raise RuntimeError(f"Error in SelfAttention backward_dw: {str(e)}") def _backward_qkv_proj(self): """Update weights for QKV projection layer""" self.linear_qkv.backward_dw() def _backward_output_proj(self): """Update weights for output projection layer""" self.linear_proj.backward_dw()

def forward( self, hidden_states: Tensor, attention_mask: Tensor, encoder_output: Optional[Tensor] = None, inference_params: Optional[InferenceParams] = None, ) -> Tensor: # QKV投影 query, key, value = self._get_qkv(hidden_states, encoder_output) # 查询和键的LayerNorm if self.config.qk_layernorm: query = self.q_layernorm(query) key = self.k_layernorm(key) # 核心注意力计算 context = self.core_attention( query, key, value, attention_mask, inference_params ) # 输出投影 output = self.linear_proj(context) return output

class MLP(MegatronModule): def __init__( self, config: TransformerConfig, submodules: MLPSubmodules, is_expert: bool = False, input_size: int = None, ): # 第一个线性层 self.linear_fc1 = submodules.linear_fc1 # 激活函数 self.activation_func = config.activation_func # 门控线性单元支持 if config.gated_linear_unit: self.linear_fc1_gated = submodules.linear_fc1_gated # 第二个线性层 self.linear_fc2 = submodules.linear_fc2

def forward(self, hidden_states: Tensor) -> Tensor: # 第一个线性变换 intermediate_parallel = self.linear_fc1(hidden_states) # 门控线性单元 if self.config.gated_linear_unit: intermediate_parallel_gated = self.linear_fc1_gated(hidden_states) intermediate_parallel = self.activation_func(intermediate_parallel_gated) * intermediate_parallel else: intermediate_parallel = self.activation_func(intermediate_parallel) # 第二个线性变换 output = self.linear_fc2(intermediate_parallel) return output def backward_dw(self): """Execute MLP weight update operations""" try: self.linear_fc2.backward_dw() self.linear_fc1.backward_dw() if self.config.gated_linear_unit: self.linear_fc1_gated.backward_dw() except Exception as e: raise RuntimeError(f"MLP backward_dw execution failed: {str(e)}")

class ColumnParallelLinear(MegatronModule): def __init__( self, input_size: int, output_size: int, *, config: ModelParallelConfig, init_method: Callable, bias: bool = True, gather_output: bool = True, skip_bias_add: bool = False, skip_weight_param_allocation: bool = False, is_expert: bool = False, ): # 输出维度分片 self.output_size_per_partition = divide(output_size, get_tensor_model_parallel_world_size()) # 权重初始化 if not skip_weight_param_allocation: self.weight = Parameter(torch.empty( self.output_size_per_partition, input_size, device=torch.cuda.current_device(), dtype=config.params_dtype, ))

def forward(self, input_: Tensor) -> Tensor: # 设置输入的并行属性 set_tensor_model_parallel_attributes(input_, False, False, False) # 矩阵乘法 output_parallel = F.linear(input_, self.weight, self.bias) # 收集输出 (如果需要) if self.gather_output: output = gather_from_tensor_model_parallel_region(output_parallel) else: output = output_parallel return output def backward_dw(self): """Custom weight gradient computation for tensor parallel layers""" if hasattr(self, 'split_bw') and self.split_bw: # 使用Transformer Engine的权重梯度计算 super().wgrad_comp()

def set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride): """Sets tp attributes to tensor""" setattr(tensor, 'tensor_model_parallel', is_parallel) setattr(tensor, 'partition_dim', dim) setattr(tensor, 'partition_stride', stride) def param_is_not_tensor_parallel_duplicate(param): """Returns true if param is not a duplicate on another TP rank""" return (hasattr(param, 'tensor_model_parallel') and param.tensor_model_parallel) or ( get_tensor_model_parallel_rank() == 0 )

def pretrain( model_provider: Callable, model_type: ModelType, args, data_loader_provider: Callable = None, process_non_loss_data_func: Callable = None, model_init_fn: Callable = None, ): # 初始化Megatron initialize_megatron(args) # 设置模型、优化器、学习率调度器 model, optimizer, lr_scheduler = setup_model_and_optimizer( model_provider, model_type, args, model_init_fn ) # 设置数据加载器 if data_loader_provider is None: data_loader_provider = build_pretraining_data_loader dataloader = data_loader_provider(args) # 训练循环 if args.train_iters > 0: iteration = train( forward_step_func=model_provider.module_provider.forward_step, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, dataloader=dataloader, args=args, process_non_loss_data_func=process_non_loss_data_func, ) return iteration

def train( forward_step_func: Callable, model: MegatronModule, optimizer: Optimizer, lr_scheduler: OptimizerParamScheduler, dataloader: torch.utils.data.DataLoader, args, process_non_loss_data_func: Callable = None, ): # 获取前向后向函数 forward_backward_func = get_forward_backward_func( args.virtual_pipeline_model_parallel_size, args.pipeline_model_parallel_size, args.num_layers_per_virtual_pipeline_stage, ) # 训练循环 for iteration in range(args.iteration, args.train_iters): # 前向后向传播 losses_reduced = forward_backward_func( forward_step_func=forward_step_func, data_iterator=data_iterator, model=model, num_microbatches=get_num_microbatches(), seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, decoder_seq_length=args.decoder_seq_length, layernorm_epsilon=args.layernorm_epsilon, hidden_size=args.hidden_size, params_dtype=args.params_dtype, fp16=args.fp16, bf16=args.bf16, fp32_residual_connection=args.fp32_residual_connection, async_comm=args.async_comm, ) # 梯度同步 if args.DDP_impl == 'local': grads = [param.grad for param in model.parameters() if param.grad is not None] _allreduce_gradients(grads, args) # 优化器步骤 optimizer.step() lr_scheduler.step() # 检查点保存 if args.save and args.save_interval and iteration % args.save_interval == 0: save_checkpoint(iteration, model, optimizer, lr_scheduler, args) # 日志记录 training_log(losses_reduced, optimizer, lr_scheduler, iteration, args)

def add_megatron_arguments(parser, include_defaults=True): # 模型架构参数 group = parser.add_argument_group('model architecture') group.add_argument('--num-layers', type=int, default=None, help='Number of transformer layers') group.add_argument('--hidden-size', type=int, default=None, help='Transformer hidden size') group.add_argument('--num-attention-heads', type=int, default=None, help='Number of transformer attention heads') # 并行化参数 group = parser.add_argument_group('parallelism configuration') group.add_argument('--tensor-model-parallel-size', type=int, default=1, help='Degree of tensor model parallelism') group.add_argument('--pipeline-model-parallel-size', type=int, default=1, help='Degree of pipeline model parallelism') group.add_argument('--expert-model-parallel-size', type=int, default=1, help='Degree of expert model parallelism') # 训练参数 group = parser.add_argument_group('training parameters') group.add_argument('--micro-batch-size', type=int, default=None, help='Micro batch size per GPU') group.add_argument('--global-batch-size', type=int, default=None, help='Global batch size') group.add_argument('--lr', type=float, default=None, help='Learning rate')

def parse_args(extra_args_provider=None, defaults={}, ignore_unknown_args=False, allow_no_cuda=False): # 创建参数解析器 parser = argparse.ArgumentParser(description='Megatron-LM Arguments', allow_abbrev=False) # 添加Megatron参数 add_megatron_arguments(parser, include_defaults=False) # 添加额外参数 if extra_args_provider is not None: extra_args_provider(parser) # 解析参数 args, unknown = parser.parse_known_args() # 从环境变量读取参数 args = _read_args_from_environment(args) # 验证参数 args = validate_args(args) return args

class TELinear(te.pytorch.Linear): def __init__( self, input_size: int, output_size: int, *, parallel_mode: str = None, bias: bool = True, skip_bias_add: bool = False, is_expert: bool = False, tp_size: int = 1, config: ModelParallelConfig = None, init_method: Callable = None, **kwargs ): # 调用父类初始化 super().__init__( input_size, output_size, parallel_mode=parallel_mode, bias=bias, skip_bias_add=skip_bias_add, **kwargs ) # 设置张量并行属性 self.tp_size = tp_size self.is_expert = is_expert self.split_bw = config.split_bw if config else False

def backward_dw(self): """Custom weight gradient computation for TE layers""" if self.split_bw: # 使用Transformer Engine的权重梯度计算 super().wgrad_comp() class TENorm: """Transformer Engine LayerNorm/RMSNorm wrapper""" def __new__(cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5): if config.normalization == "LayerNorm": instance = te.pytorch.LayerNorm( hidden_size=hidden_size, eps=eps, sequence_parallel=config.sequence_parallel, zero_centered_gamma=config.layernorm_zero_centered_gamma, **_get_extra_te_kwargs(config), ) elif config.normalization == "RMSNorm": instance = te.pytorch.RMSNorm( hidden_size=hidden_size, eps=eps, sequence_parallel=config.sequence_parallel, zero_centered_gamma=config.layernorm_zero_centered_gamma, **_get_extra_te_kwargs(config), ) return instance