diff --git a/megatron/legacy/model/__init__.py b/megatron/legacy/model/__init__.py index 1482c11f475..bcf70a423e6 100644 --- a/megatron/legacy/model/__init__.py +++ b/megatron/legacy/model/__init__.py @@ -3,7 +3,6 @@ from .fused_layer_norm import MixedFusedLayerNorm as LayerNorm from .rms_norm import RMSNorm -from .bert_model import BertModel from .gpt_model import GPTModel from .t5_model import T5Model from .language_model import get_language_model diff --git a/megatron/legacy/model/bert_model.py b/megatron/legacy/model/bert_model.py deleted file mode 100644 index 451667dd2fa..00000000000 --- a/megatron/legacy/model/bert_model.py +++ /dev/null @@ -1,257 +0,0 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. - -"""BERT model.""" - -import torch - -from megatron.training import get_args -from megatron.core import tensor_parallel -from megatron.legacy.model.enums import AttnMaskType -from megatron.legacy.model.language_model import parallel_lm_logits -from megatron.legacy.model.language_model import get_language_model -from megatron.legacy.model.utils import get_norm -from megatron.legacy.model.utils import openai_gelu, erf_gelu -from megatron.legacy.model.utils import get_linear_layer -from megatron.legacy.model.utils import init_method_normal -from megatron.legacy.model.utils import scaled_init_method_normal -from .module import MegatronModule - - -def bert_extended_attention_mask(attention_mask): - # We create a 3D attention mask from a 2D tensor mask. - # [b, 1, s] - attention_mask_b1s = attention_mask.unsqueeze(1) - # [b, s, 1] - attention_mask_bs1 = attention_mask.unsqueeze(2) - # [b, s, s] - attention_mask_bss = attention_mask_b1s * attention_mask_bs1 - # [b, 1, s, s] - extended_attention_mask = attention_mask_bss.unsqueeze(1) - - # Convert attention mask to binary: - extended_attention_mask = (extended_attention_mask < 0.5) - - return extended_attention_mask - -def bert_position_ids(token_ids): - # Create position ids - seq_length = token_ids.size(1) - position_ids = torch.arange(seq_length, dtype=torch.long, - device=token_ids.device) - position_ids = position_ids.unsqueeze(0).expand_as(token_ids) - - return position_ids - - -class BertLMHead(MegatronModule): - """Masked LM head for Bert - - Args: - config: TransformerConfig object - mpu_vocab_size: model parallel size of vocabulary. - parallel_output: whether output logits being distributed or not. - """ - - def __init__(self, mpu_vocab_size, config, parallel_output): - super().__init__(config=config) - - args = get_args() - self.bias = torch.nn.Parameter(torch.zeros(mpu_vocab_size)) - tensor_parallel.set_tensor_model_parallel_attributes(self.bias, True, 0, 1) - self.parallel_output = parallel_output - - self.dense = get_linear_layer(config.hidden_size, config.hidden_size, config.init_method) - setattr(self.dense.weight, 'sequence_parallel', config.sequence_parallel) - setattr(self.dense.bias, 'sequence_parallel', config.sequence_parallel) - - self.norm = get_norm(config) - self.gelu = torch.nn.functional.gelu - if args.openai_gelu: - self.gelu = openai_gelu - elif args.onnx_safe: - self.gelu = erf_gelu - - def forward(self, hidden_states, word_embeddings_weight): - hidden_states = self.dense(hidden_states) - hidden_states = self.gelu(hidden_states) - hidden_states = self.norm(hidden_states) - output = parallel_lm_logits(hidden_states, - word_embeddings_weight, - self.parallel_output, - bias=self.bias) - return output - - def load_state_dict(self, state_dict, strict=True): - """Customize load.""" - - # Handle renaming layernorm -> norm in component names - state_dict_ = {} - for key in state_dict.keys(): - newkey = key.replace("layernorm", "norm") - state_dict_[newkey] = state_dict[key] - - super().load_state_dict(state_dict_, strict) - - -def post_language_model_processing(lm_output, pooled_output, - lm_head, binary_head, - lm_labels, - logit_weights, - fp16_lm_cross_entropy): - # Output. - lm_logits = lm_head( - lm_output, logit_weights) - - binary_logits = None - if binary_head is not None: - binary_logits = binary_head(pooled_output) - - if lm_labels is None: - # [s b h] => [b s h] - return lm_logits.transpose(0,1).contiguous(), binary_logits - else: - # [b s] => [s b] - lm_labels = lm_labels.transpose(0,1).contiguous() - # lm_logits : [s, b, h] and lm_labels: [s, b] - if fp16_lm_cross_entropy: - assert lm_logits.dtype == torch.half - lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits, lm_labels) - else: - lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits.float(), - lm_labels) - # [s, b] => [b s] - lm_loss = lm_loss.transpose(0,1).contiguous() - return lm_loss, binary_logits - - -class BertModel(MegatronModule): - """Bert Language model.""" - - def __init__(self, - config, - num_tokentypes=2, - add_binary_head=True, - parallel_output=True, - pre_process=True, - post_process=True): - super().__init__(config=config) - args = get_args() - - # TODO this option is not yet implemented in BERT - assert args.untie_embeddings_and_output_weights is False - - self.fp16_lm_cross_entropy = args.fp16_lm_cross_entropy - self.add_binary_head = add_binary_head - self.parallel_output = parallel_output - self.pre_process = pre_process - self.post_process = post_process - - self.return_embeddings = args.output_bert_embeddings - if self.return_embeddings: - assert self.post_process and self.add_binary_head - - self.language_model, self._language_model_key = get_language_model( - config=config, - num_tokentypes=num_tokentypes, - add_pooler=self.add_binary_head, - encoder_attn_mask_type=AttnMaskType.padding, - pre_process=self.pre_process, - post_process=self.post_process) - - self.initialize_word_embeddings() - if self.post_process: - self.lm_head = BertLMHead(self.shared_embedding_or_output_weight().size(0), config, parallel_output) - self._lm_head_key = 'lm_head' - self.binary_head = None - if self.add_binary_head: - self.binary_head = get_linear_layer(config.hidden_size, 2, - config.init_method) - self._binary_head_key = 'binary_head' - - def set_input_tensor(self, input_tensor): - """See megatron.legacy.model.transformer.set_input_tensor()""" - self.language_model.set_input_tensor(input_tensor) - - def forward(self, bert_model_input, attention_mask, - tokentype_ids=None, lm_labels=None, inference_context=None): - - extended_attention_mask = bert_extended_attention_mask(attention_mask) - input_ids = bert_model_input - position_ids = bert_position_ids(input_ids) - - lm_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids - ) - - if self.post_process and self.add_binary_head: - lm_output, pooled_output = lm_output - - # Return pooled output (e.g., when computing Bert embeddings). - if self.return_embeddings: - - # Sum attention mask. - embeddings = torch.transpose(lm_output, 0, 1) - masks = torch.sum(attention_mask, dim=1) - - # Collect masked embeddings. - output = torch.zeros( - size=(embeddings.shape[0], embeddings.shape[2]), - dtype=torch.float32, - device=torch.cuda.current_device()) - for i, (embedding, mask) in enumerate(zip(embeddings, masks)): - output[i, :] = torch.mean(embedding[1: mask - 1], dim=0) - - return output - - else: - pooled_output = None - - if self.post_process: - return post_language_model_processing(lm_output, pooled_output, - self.lm_head, self.binary_head, - lm_labels, - self.shared_embedding_or_output_weight(), - self.fp16_lm_cross_entropy) - else: - return lm_output - - - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): - """For easy load when model is combined with other heads, - add an extra key.""" - - state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) - if self.post_process: - state_dict_[self._lm_head_key] \ - = self.lm_head.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) - if self.post_process and self.add_binary_head: - state_dict_[self._binary_head_key] \ - = self.binary_head.state_dict(prefix=prefix, keep_vars=keep_vars) - # Save word_embeddings. - if self.post_process and not self.pre_process: - state_dict_[self._word_embeddings_for_head_key] \ - = self.word_embeddings.state_dict(prefix=prefix, keep_vars=keep_vars) - return state_dict_ - - def load_state_dict(self, state_dict, strict=True): - """Customized load.""" - - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) - if self.post_process: - self.lm_head.load_state_dict( - state_dict[self._lm_head_key], strict=strict) - if self.post_process and self.add_binary_head: - self.binary_head.load_state_dict( - state_dict[self._binary_head_key], strict=strict) - # Load word_embeddings. - if self.post_process and not self.pre_process: - self.word_embeddings.load_state_dict( - state_dict[self._word_embeddings_for_head_key], strict=strict) diff --git a/megatron/legacy/model/classification.py b/megatron/legacy/model/classification.py deleted file mode 100644 index c9fe165280e..00000000000 --- a/megatron/legacy/model/classification.py +++ /dev/null @@ -1,101 +0,0 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. - -"""Classification model.""" - -import torch - -from megatron.training import get_args, print_rank_last -from megatron.legacy.model.enums import AttnMaskType -from megatron.legacy.model.bert_model import bert_extended_attention_mask, bert_position_ids -from megatron.legacy.model.language_model import get_language_model -from megatron.legacy.model.utils import get_linear_layer -from megatron.legacy.model.utils import init_method_normal -from megatron.legacy.model.utils import scaled_init_method_normal -from .module import MegatronModule - - -class Classification(MegatronModule): - - def __init__(self, - config, - num_classes, - num_tokentypes=2, - pre_process=True, - post_process=True): - super().__init__(config=config, share_embeddings_and_output_weights=False) - args = get_args() - - self.num_classes = num_classes - self.pre_process = pre_process - self.post_process = post_process - - self.language_model, self._language_model_key = get_language_model( - config=config, - num_tokentypes=num_tokentypes, - add_pooler=True, - encoder_attn_mask_type=AttnMaskType.padding, - pre_process=self.pre_process, - post_process=self.post_process) - - # Multi-choice head. - if self.post_process: - self.classification_dropout = torch.nn.Dropout(args.hidden_dropout) - self.classification_head = get_linear_layer(args.hidden_size, - self.num_classes, - config.init_method) - self._classification_head_key = 'classification_head' - - def set_input_tensor(self, input_tensor): - """See megatron.legacy.model.transformer.set_input_tensor()""" - self.language_model.set_input_tensor(input_tensor) - - def forward(self, model_input, attention_mask, tokentype_ids=None): - - extended_attention_mask = bert_extended_attention_mask(attention_mask) - input_ids = model_input - position_ids = bert_position_ids(input_ids) - - lm_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids - ) - - if self.post_process: - _, pooled_output = lm_output - classification_output = self.classification_dropout(pooled_output) - classification_logits = self.classification_head(classification_output) - - # Reshape back to separate choices. - classification_logits = classification_logits.view(-1, self.num_classes) - - return classification_logits - return lm_output - - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): - """For easy load when model is combined with other heads, - add an extra key.""" - - state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) - if self.post_process: - state_dict_[self._classification_head_key] \ - = self.classification_head.state_dict(prefix=prefix, keep_vars=keep_vars) - return state_dict_ - - def load_state_dict(self, state_dict, strict=True): - """Customized load.""" - - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) - if self.post_process: - if self._classification_head_key in state_dict: - self.classification_head.load_state_dict( - state_dict[self._classification_head_key], strict=strict) - else: - print_rank_last('***WARNING*** could not find {} in the checkpoint, ' - 'initializing to random'.format( - self._classification_head_key)) diff --git a/megatron/legacy/model/multiple_choice.py b/megatron/legacy/model/multiple_choice.py deleted file mode 100644 index bec0548c405..00000000000 --- a/megatron/legacy/model/multiple_choice.py +++ /dev/null @@ -1,112 +0,0 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. - -"""Multiple choice model.""" - -import torch - -from megatron.training import get_args, print_rank_last -from megatron.legacy.model.enums import AttnMaskType -from megatron.legacy.model.bert_model import bert_extended_attention_mask, bert_position_ids -from megatron.legacy.model.language_model import get_language_model -from megatron.legacy.model.utils import get_linear_layer -from megatron.legacy.model.utils import init_method_normal -from megatron.legacy.model.utils import scaled_init_method_normal -from .module import MegatronModule - - -class MultipleChoice(MegatronModule): - - def __init__(self, - config, - num_tokentypes=2, - pre_process=True, - post_process=True): - super(MultipleChoice, self).__init__(share_embeddings_and_output_weights=False) - args = get_args() - - self.pre_process = pre_process - self.post_process = post_process - - self.language_model, self._language_model_key = get_language_model( - config=config, - num_tokentypes=num_tokentypes, - add_pooler=True, - encoder_attn_mask_type=AttnMaskType.padding, - pre_process=self.pre_process, - post_process=self.post_process) - - # Multi-choice head. - if self.post_process: - self.multichoice_dropout = torch.nn.Dropout(args.hidden_dropout) - self.multichoice_head = get_linear_layer(args.hidden_size, 1, - init_method) - self._multichoice_head_key = 'multichoice_head' - - def set_input_tensor(self, input_tensor): - """See megatron.legacy.model.transformer.set_input_tensor()""" - self.language_model.set_input_tensor(input_tensor) - - def forward(self, model_input, attention_mask, tokentype_ids=None): - - # [batch, choices, sequence] --> [batch * choices, sequence] --> - # transformer --> [batch, choices] --> softmax - - # Ensure the shape is [batch-size, choices, sequence] - assert len(attention_mask.shape) == 3 - num_choices = attention_mask.shape[1] - - # Reshape and treat choice dimension the same as batch. - attention_mask = attention_mask.view(-1, attention_mask.size(-1)) - extended_attention_mask = bert_extended_attention_mask(attention_mask) - - input_ids = model_input - # Do the same as attention_mask for input_ids, tokentype_ids - assert len(input_ids.shape) == 3 - assert len(tokentype_ids.shape) == 3 - input_ids = input_ids.view(-1, input_ids.size(-1)) - tokentype_ids = tokentype_ids.view(-1, tokentype_ids.size(-1)) - position_ids = bert_position_ids(input_ids) - - lm_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids - ) - if self.post_process: - _, pooled_output = lm_output - multichoice_output = self.multichoice_dropout(pooled_output) - multichoice_logits = self.multichoice_head(multichoice_output) - - # Reshape back to separate choices. - multichoice_logits = multichoice_logits.view(-1, num_choices) - - return multichoice_logits - return lm_output - - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): - """For easy load when model is combined with other heads, - add an extra key.""" - - state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) - if self.post_process: - state_dict_[self._multichoice_head_key] \ - = self.multichoice_head.state_dict(prefix=prefix, keep_vars=keep_vars) - return state_dict_ - - def load_state_dict(self, state_dict, strict=True): - """Customized load.""" - - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) - if self.post_process: - if self._multichoice_head_key in state_dict: - self.multichoice_head.load_state_dict( - state_dict[self._multichoice_head_key], strict=strict) - else: - print_rank_last('***WARNING*** could not find {} in the checkpoint, ' - 'initializing to random'.format( - self._multichoice_head_key)) diff --git a/pretrain_bert.py b/pretrain_bert.py index 9b11908811b..3a66a124b6e 100644 --- a/pretrain_bert.py +++ b/pretrain_bert.py @@ -12,7 +12,6 @@ from megatron.training import get_timers from megatron.core import tensor_parallel from megatron.core.enums import ModelType -import megatron.legacy.model from megatron.core.models.bert.bert_model import BertModel from megatron.training import pretrain from megatron.training.utils import average_losses_across_data_parallel_group @@ -36,35 +35,26 @@ def model_provider(pre_process=True, post_process=True, vp_stage=None, config=No config = core_transformer_config_from_args(args) num_tokentypes = 2 if args.bert_binary_head else 0 - if args.use_legacy_models: - model = megatron.legacy.model.BertModel( - config=config, - num_tokentypes=num_tokentypes, - add_binary_head=args.bert_binary_head, - parallel_output=True, - pre_process=pre_process, - post_process=post_process) - else: - if args.spec is None: - transformer_layer_spec = bert_layer_with_transformer_engine_spec #default spec - elif args.spec[0] == 'local': - print_rank_0('Using Local spec for transformer layers') - transformer_layer_spec = bert_layer_local_spec - else : - transformer_layer_spec = import_module(args.spec) - - model = BertModel( - config=config, - transformer_layer_spec=transformer_layer_spec, - vocab_size=args.padded_vocab_size, - max_sequence_length=args.max_position_embeddings, - num_tokentypes=num_tokentypes, - add_binary_head=args.bert_binary_head, - share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, - parallel_output=True, - pre_process=pre_process, - post_process=post_process, - vp_stage=vp_stage) + if args.spec is None: + transformer_layer_spec = bert_layer_with_transformer_engine_spec #default spec + elif args.spec[0] == 'local': + print_rank_0('Using Local spec for transformer layers') + transformer_layer_spec = bert_layer_local_spec + else : + transformer_layer_spec = import_module(args.spec) + + model = BertModel( + config=config, + transformer_layer_spec=transformer_layer_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + num_tokentypes=num_tokentypes, + add_binary_head=args.bert_binary_head, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + parallel_output=True, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage) return model diff --git a/tools/bert_embedding/embed.py b/tools/bert_embedding/embed.py index effc6f6d91e..84b6a55480e 100644 --- a/tools/bert_embedding/embed.py +++ b/tools/bert_embedding/embed.py @@ -17,7 +17,6 @@ from megatron.core import parallel_state from megatron.core.enums import ModelType from megatron.core.pipeline_parallel import get_forward_backward_func -from megatron.legacy.model import BertModel from megatron.training.training import setup_model_and_optimizer from pretrain_bert import model_provider, get_batch, loss_func, forward_step