megatron_patch/model/baichuan2/transformer.py [1225:1556]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if self.apply_residual_connection_post_norm: residual = norm_output else: residual = hidden_states if self.drop_path is None: # jit scripting for a nn.module (with dropout) is not # trigerring the fusion kernel. For now, we use two # different nn.functional routines to account for varying # dropout semantics during training and inference phases. if self.bias_dropout_fusion: if self.training: bias_dropout_add_func = bias_dropout_add_fused_train else: bias_dropout_add_func = bias_dropout_add_fused_inference else: bias_dropout_add_func = get_bias_dropout_add(self.training) if attention_bias is not None: attention_bias = attention_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( attention_output, attention_bias, residual, self.hidden_dropout) else: out = torch.nn.functional.dropout(attention_output + attention_bias, p=self.hidden_dropout, training=self.training) norm_input = residual + self.drop_path(out) # Layer norm post the self attention. norm_output = self.post_attention_norm(norm_input) # Cross attention. if self.layer_type == LayerType.encoder: pass elif self.layer_type == LayerType.decoder: norm_input, norm_output = \ self.default_decoder_cross_attention( encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type == LayerType.retro_encoder: norm_input, norm_output = \ self.retro_encoder_cross_attention( retriever_output, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type in (LayerType.retro_decoder, LayerType.retro_decoder_with_retriever): retriever_output, norm_input, norm_output = \ self.retro_decoder_cross_attention( retriever_input, retriever_output, retriever_attn_mask, norm_input, norm_output, inference_params, bias_dropout_add_func) else: raise Exception("Unsupported layer type, '%s'." % self.layer_type.name) # MLP. mlp_output, mlp_bias = self.mlp(norm_output) # Second residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = norm_input if self.drop_path is None: if mlp_bias is not None: mlp_bias = mlp_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): output = bias_dropout_add_func( mlp_output, mlp_bias, residual, self.hidden_dropout) # Jit compiled function creates 'view' tensor. This tensor # potentially gets saved in the MPU checkpoint function context, # which rejects view tensors. While making a viewless tensor here # won't result in memory savings (like the data loader, or # p2p_communication), it serves to document the origin of this # 'view' tensor. output = core.utils.make_viewless_tensor(inp = output, requires_grad = output.requires_grad, keep_graph = True) else: if mlp_bias is not None: mlp_output = mlp_output + mlp_bias out = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training) output = residual + self.drop_path(out) if self.layer_type == LayerType.retro_decoder_with_retriever: return output, retriever_output else: return output class NoopTransformerLayer(MegatronModule): """A single 'no-op' transformer layer. The sole purpose of this layer is for when a standalone embedding layer is used (i.e., args.standalone_embedding_stage == True). In this case, zero transformer layers are assigned when pipeline rank == 0. Additionally, when virtual pipeline rank >= 1, zero total model parameters are created (virtual rank 0 contains the input embedding). This results in the model's input and output tensors being the same, which causes an error when performing certain memory optimiations on the output tensor (e.g., deallocating it). Thus, this layer disconnects the input from the output via a clone. Since ranks containing a no-op layer are generally under- utilized (both compute and memory), there's no worry of any performance degredation. """ def __init__(self, layer_number): super().__init__() self.layer_number = layer_number def forward(self, hidden_states, attention_mask, encoder_output=None, enc_dec_attn_mask=None, inference_params=None): return hidden_states.clone() def _get_num_layers(args, model_type, is_decoder=False): """Compute the number of transformer layers resident on the current rank.""" is_encoder_and_decoder_model = (model_type == ModelType.encoder_and_decoder) if model_type == ModelType.retro_encoder: num_layers = args.retro_encoder_layers elif mpu.get_pipeline_model_parallel_world_size() > 1: if is_encoder_and_decoder_model: assert args.pipeline_model_parallel_split_rank is not None # When a standalone embedding stage is used, a rank is taken from # the encoder's ranks, to be used for the encoder's embedding # layer. This way, the rank referenced by the 'split rank' remains # the same whether or not a standalone embedding stage is used. num_ranks_in_encoder = ( args.pipeline_model_parallel_split_rank - 1 if args.standalone_embedding_stage else args.pipeline_model_parallel_split_rank ) num_ranks_in_decoder = args.transformer_pipeline_model_parallel_size - num_ranks_in_encoder assert args.encoder_num_layers % num_ranks_in_encoder == 0, \ 'encoder_num_layers (%d) must be divisible by number of ranks given to encoder (%d)' % (args.encoder_num_layers, num_ranks_in_encoder) assert args.decoder_num_layers % num_ranks_in_decoder == 0, \ 'decoder_num_layers (%d) must be divisible by number of ranks given to decoder (%d)' % (args.decoder_num_layers, num_ranks_in_decoder) if mpu.is_pipeline_stage_before_split(): num_layers = ( 0 if args.standalone_embedding_stage and mpu.get_pipeline_model_parallel_rank() == 0 else args.encoder_num_layers // num_ranks_in_encoder ) else: num_layers = args.decoder_num_layers // num_ranks_in_decoder else: assert args.num_layers == args.encoder_num_layers assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \ 'num_layers must be divisible by transformer_pipeline_model_parallel_size' # When a standalone embedding stage is used, all transformer layers # are divided among pipeline rank >= 1, while on pipeline rank 0, # ranks either contain the input embedding layer (virtual pp rank 0), # or no layers at all (virtual pp rank >= 1). num_layers = ( 0 if args.standalone_embedding_stage and mpu.get_pipeline_model_parallel_rank() == 0 else args.num_layers // args.transformer_pipeline_model_parallel_size ) else: if not is_decoder: num_layers = args.encoder_num_layers else: num_layers = args.decoder_num_layers return num_layers def _get_layer_type(model_type, default_layer_type, retro_layer_numbers, layer_number): args = get_args() if args.retro_add_retriever and layer_number in retro_layer_numbers: if model_type == ModelType.retro_decoder: return LayerType.retro_decoder_with_retriever \ if layer_number == retro_layer_numbers[0] \ else LayerType.retro_decoder elif model_type == ModelType.retro_encoder: return LayerType.retro_encoder else: raise Exception("Unsupported model type, '%s'." % model_type) else: return default_layer_type class ParallelTransformer(MegatronModule): """Transformer class.""" def __init__(self, config, model_type, layer_type=LayerType.encoder, self_attn_mask_type=AttnMaskType.padding, post_norm=True, pre_process=True, post_process=True, drop_path_rate=0.0): super(ParallelTransformer, self).__init__() args = get_args() self.layer_type = layer_type self.model_type = model_type self.bf16 = config.bf16 self.fp32_residual_connection = config.fp32_residual_connection self.post_norm = post_norm self.pre_process = pre_process self.post_process = post_process self.input_tensor = None self.drop_path_rate = drop_path_rate self.transformer_impl = args.transformer_impl self.retro_add_retriever = args.retro_add_retriever # Store activation checkpoiting flag. self.recompute_granularity = config.recompute_granularity self.recompute_method = config.recompute_method self.recompute_num_layers = config.recompute_num_layers self.distribute_saved_activations = \ config.distribute_saved_activations and not config.sequence_parallel self.sequence_parallel = config.sequence_parallel # Transformer Engine Init. self.transformer_engine_v_0_10 = False self.transformer_engine_v_0_11 = False self.transformer_engine_v_0_8 = False if self.transformer_impl == 'transformer_engine': global transformer_engine import transformer_engine from importlib.metadata import version from pkg_resources import packaging te_version = packaging.version.Version(version("transformer-engine")) if te_version >= packaging.version.Version("0.8.0"): self.transformer_engine_v_0_8 = True if te_version >= packaging.version.Version("0.10.0"): self.transformer_engine_v_0_10 = True if te_version >= packaging.version.Version("0.11.0"): self.transformer_engine_v_0_11 = True del version, packaging assert not args.squared_relu, "TransformerEngine does not support squared relu activation." self.use_fp8 = args.fp8 is not None self.fp8_recipe = None self.fp8_group = None if self.use_fp8: assert args.transformer_impl == 'transformer_engine', \ 'transformer-engine required for fp8 training and inference' self.fp8_group = mpu.get_amax_reduction_group() if args.fp8 == "e4m3": fp8_format = transformer_engine.common.recipe.Format.E4M3 elif args.fp8 == "hybrid": fp8_format = transformer_engine.common.recipe.Format.HYBRID else: raise ValueError("The DelayedScaling recipe only supports E4M3 and HYBRID formats.") self.fp8_recipe = transformer_engine.common.recipe.DelayedScaling( margin=args.fp8_margin, interval=args.fp8_interval, fp8_format=fp8_format, amax_history_len=args.fp8_amax_history_len, amax_compute_algo=args.fp8_amax_compute_algo, override_linear_precision=(False, False, not args.fp8_wgrad), ) self.num_microbatches_in_previous_step = -1 self.microbatch_count = 0 self.checkpoint_core_attention = config.recompute_granularity == 'selective' # Number of layers. self.num_layers = _get_num_layers(args, model_type, layer_type==LayerType.decoder) self.drop_path_rates = [ rate.item() for rate in torch.linspace(0, self.drop_path_rate, config.num_layers)] self.retro_layer_numbers = None if model_type == ModelType.retro_decoder: retro_layer_start = 6 if config.num_layers <= 15 else 9 self.retro_layer_numbers = \ np.arange(retro_layer_start, args.num_layers + 1, 3).tolist() if model_type == ModelType.retro_encoder: self.retro_layer_numbers = [1] # Transformer layers. if args.retro_add_retriever: assert self.recompute_granularity != 'full', \ "Full recompute not supported for Retro." assert args.transformer_impl == 'local', \ "Transformer engine does not support Retro layers." def build_layer(layer_number): if args.transformer_impl == 'local': current_layer_type = _get_layer_type( model_type, layer_type, self.retro_layer_numbers, layer_number) return ParallelTransformerLayer( config, layer_number, layer_type=current_layer_type, self_attn_mask_type=self_attn_mask_type, drop_path_rate=self.drop_path_rates[layer_number - 1]) else: # This argument is only available from TE v0.10 onwards. extra_transformer_engine_kwargs = {} if self.transformer_engine_v_0_8: extra_transformer_engine_kwargs["bias"] = args.add_bias_linear if self.transformer_engine_v_0_10: extra_transformer_engine_kwargs["activation"] = "swiglu" if args.swiglu else "gelu" if self.transformer_engine_v_0_11: extra_transformer_engine_kwargs["normalization"] = args.normalization - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - megatron_patch/model/llama3/transformer_legacy.py [1183:1514]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if self.apply_residual_connection_post_norm: residual = norm_output else: residual = hidden_states if self.drop_path is None: # jit scripting for a nn.module (with dropout) is not # trigerring the fusion kernel. For now, we use two # different nn.functional routines to account for varying # dropout semantics during training and inference phases. if self.bias_dropout_fusion: if self.training: bias_dropout_add_func = bias_dropout_add_fused_train else: bias_dropout_add_func = bias_dropout_add_fused_inference else: bias_dropout_add_func = get_bias_dropout_add(self.training) if attention_bias is not None: attention_bias = attention_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( attention_output, attention_bias, residual, self.hidden_dropout) else: out = torch.nn.functional.dropout(attention_output + attention_bias, p=self.hidden_dropout, training=self.training) norm_input = residual + self.drop_path(out) # Layer norm post the self attention. norm_output = self.post_attention_norm(norm_input) # Cross attention. if self.layer_type == LayerType.encoder: pass elif self.layer_type == LayerType.decoder: norm_input, norm_output = \ self.default_decoder_cross_attention( encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type == LayerType.retro_encoder: norm_input, norm_output = \ self.retro_encoder_cross_attention( retriever_output, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type in (LayerType.retro_decoder, LayerType.retro_decoder_with_retriever): retriever_output, norm_input, norm_output = \ self.retro_decoder_cross_attention( retriever_input, retriever_output, retriever_attn_mask, norm_input, norm_output, inference_params, bias_dropout_add_func) else: raise Exception("Unsupported layer type, '%s'." % self.layer_type.name) # MLP. mlp_output, mlp_bias = self.mlp(norm_output) # Second residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = norm_input if self.drop_path is None: if mlp_bias is not None: mlp_bias = mlp_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): output = bias_dropout_add_func( mlp_output, mlp_bias, residual, self.hidden_dropout) # Jit compiled function creates 'view' tensor. This tensor # potentially gets saved in the MPU checkpoint function context, # which rejects view tensors. While making a viewless tensor here # won't result in memory savings (like the data loader, or # p2p_communication), it serves to document the origin of this # 'view' tensor. output = core.utils.make_viewless_tensor(inp = output, requires_grad = output.requires_grad, keep_graph = True) else: if mlp_bias is not None: mlp_output = mlp_output + mlp_bias out = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training) output = residual + self.drop_path(out) if self.layer_type == LayerType.retro_decoder_with_retriever: return output, retriever_output else: return output class NoopTransformerLayer(MegatronModule): """A single 'no-op' transformer layer. The sole purpose of this layer is for when a standalone embedding layer is used (i.e., args.standalone_embedding_stage == True). In this case, zero transformer layers are assigned when pipeline rank == 0. Additionally, when virtual pipeline rank >= 1, zero total model parameters are created (virtual rank 0 contains the input embedding). This results in the model's input and output tensors being the same, which causes an error when performing certain memory optimiations on the output tensor (e.g., deallocating it). Thus, this layer disconnects the input from the output via a clone. Since ranks containing a no-op layer are generally under- utilized (both compute and memory), there's no worry of any performance degredation. """ def __init__(self, layer_number): super().__init__() self.layer_number = layer_number def forward(self, hidden_states, attention_mask, encoder_output=None, enc_dec_attn_mask=None, inference_params=None): return hidden_states.clone() def _get_num_layers(args, model_type, is_decoder=False): """Compute the number of transformer layers resident on the current rank.""" is_encoder_and_decoder_model = (model_type == ModelType.encoder_and_decoder) if model_type == ModelType.retro_encoder: num_layers = args.retro_encoder_layers elif mpu.get_pipeline_model_parallel_world_size() > 1: if is_encoder_and_decoder_model: assert args.pipeline_model_parallel_split_rank is not None # When a standalone embedding stage is used, a rank is taken from # the encoder's ranks, to be used for the encoder's embedding # layer. This way, the rank referenced by the 'split rank' remains # the same whether or not a standalone embedding stage is used. num_ranks_in_encoder = ( args.pipeline_model_parallel_split_rank - 1 if args.standalone_embedding_stage else args.pipeline_model_parallel_split_rank ) num_ranks_in_decoder = args.transformer_pipeline_model_parallel_size - num_ranks_in_encoder assert args.encoder_num_layers % num_ranks_in_encoder == 0, \ 'encoder_num_layers (%d) must be divisible by number of ranks given to encoder (%d)' % (args.encoder_num_layers, num_ranks_in_encoder) assert args.decoder_num_layers % num_ranks_in_decoder == 0, \ 'decoder_num_layers (%d) must be divisible by number of ranks given to decoder (%d)' % (args.decoder_num_layers, num_ranks_in_decoder) if mpu.is_pipeline_stage_before_split(): num_layers = ( 0 if args.standalone_embedding_stage and mpu.get_pipeline_model_parallel_rank() == 0 else args.encoder_num_layers // num_ranks_in_encoder ) else: num_layers = args.decoder_num_layers // num_ranks_in_decoder else: assert args.num_layers == args.encoder_num_layers assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \ 'num_layers must be divisible by transformer_pipeline_model_parallel_size' # When a standalone embedding stage is used, all transformer layers # are divided among pipeline rank >= 1, while on pipeline rank 0, # ranks either contain the input embedding layer (virtual pp rank 0), # or no layers at all (virtual pp rank >= 1). num_layers = ( 0 if args.standalone_embedding_stage and mpu.get_pipeline_model_parallel_rank() == 0 else args.num_layers // args.transformer_pipeline_model_parallel_size ) else: if not is_decoder: num_layers = args.encoder_num_layers else: num_layers = args.decoder_num_layers return num_layers def _get_layer_type(model_type, default_layer_type, retro_layer_numbers, layer_number): args = get_args() if args.retro_add_retriever and layer_number in retro_layer_numbers: if model_type == ModelType.retro_decoder: return LayerType.retro_decoder_with_retriever \ if layer_number == retro_layer_numbers[0] \ else LayerType.retro_decoder elif model_type == ModelType.retro_encoder: return LayerType.retro_encoder else: raise Exception("Unsupported model type, '%s'." % model_type) else: return default_layer_type class ParallelTransformer(MegatronModule): """Transformer class.""" def __init__(self, config, model_type, layer_type=LayerType.encoder, self_attn_mask_type=AttnMaskType.padding, post_norm=True, pre_process=True, post_process=True, drop_path_rate=0.0): super(ParallelTransformer, self).__init__() args = get_args() self.layer_type = layer_type self.model_type = model_type self.bf16 = config.bf16 self.fp32_residual_connection = config.fp32_residual_connection self.post_norm = post_norm self.pre_process = pre_process self.post_process = post_process self.input_tensor = None self.drop_path_rate = drop_path_rate self.transformer_impl = args.transformer_impl self.retro_add_retriever = args.retro_add_retriever # Store activation checkpoiting flag. self.recompute_granularity = config.recompute_granularity self.recompute_method = config.recompute_method self.recompute_num_layers = config.recompute_num_layers self.distribute_saved_activations = \ config.distribute_saved_activations and not config.sequence_parallel self.sequence_parallel = config.sequence_parallel # Transformer Engine Init. self.transformer_engine_v_0_10 = False self.transformer_engine_v_0_11 = False self.transformer_engine_v_0_8 = False if self.transformer_impl == 'transformer_engine': global transformer_engine import transformer_engine from importlib.metadata import version from pkg_resources import packaging te_version = packaging.version.Version(version("transformer-engine")) if te_version >= packaging.version.Version("0.8.0"): self.transformer_engine_v_0_8 = True if te_version >= packaging.version.Version("0.10.0"): self.transformer_engine_v_0_10 = True if te_version >= packaging.version.Version("0.11.0"): self.transformer_engine_v_0_11 = True del version, packaging assert not args.squared_relu, "TransformerEngine does not support squared relu activation." self.use_fp8 = args.fp8 is not None self.fp8_recipe = None self.fp8_group = None if self.use_fp8: assert args.transformer_impl == 'transformer_engine', \ 'transformer-engine required for fp8 training and inference' self.fp8_group = mpu.get_amax_reduction_group() if args.fp8 == "e4m3": fp8_format = transformer_engine.common.recipe.Format.E4M3 elif args.fp8 == "hybrid": fp8_format = transformer_engine.common.recipe.Format.HYBRID else: raise ValueError("The DelayedScaling recipe only supports E4M3 and HYBRID formats.") self.fp8_recipe = transformer_engine.common.recipe.DelayedScaling( margin=args.fp8_margin, interval=args.fp8_interval, fp8_format=fp8_format, amax_history_len=args.fp8_amax_history_len, amax_compute_algo=args.fp8_amax_compute_algo, override_linear_precision=(False, False, not args.fp8_wgrad), ) self.num_microbatches_in_previous_step = -1 self.microbatch_count = 0 self.checkpoint_core_attention = config.recompute_granularity == 'selective' # Number of layers. self.num_layers = _get_num_layers(args, model_type, layer_type==LayerType.decoder) self.drop_path_rates = [ rate.item() for rate in torch.linspace(0, self.drop_path_rate, config.num_layers)] self.retro_layer_numbers = None if model_type == ModelType.retro_decoder: retro_layer_start = 6 if config.num_layers <= 15 else 9 self.retro_layer_numbers = \ np.arange(retro_layer_start, args.num_layers + 1, 3).tolist() if model_type == ModelType.retro_encoder: self.retro_layer_numbers = [1] # Transformer layers. if args.retro_add_retriever: assert self.recompute_granularity != 'full', \ "Full recompute not supported for Retro." assert args.transformer_impl == 'local', \ "Transformer engine does not support Retro layers." def build_layer(layer_number): if args.transformer_impl == 'local': current_layer_type = _get_layer_type( model_type, layer_type, self.retro_layer_numbers, layer_number) return ParallelTransformerLayer( config, layer_number, layer_type=current_layer_type, self_attn_mask_type=self_attn_mask_type, drop_path_rate=self.drop_path_rates[layer_number - 1]) else: # This argument is only available from TE v0.10 onwards. extra_transformer_engine_kwargs = {} if self.transformer_engine_v_0_8: extra_transformer_engine_kwargs["bias"] = args.add_bias_linear if self.transformer_engine_v_0_10: extra_transformer_engine_kwargs["activation"] = "swiglu" if args.swiglu else "gelu" if self.transformer_engine_v_0_11: extra_transformer_engine_kwargs["normalization"] = args.normalization - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -