diff --git a/megatron/legacy/model/biencoder_model.py b/megatron/legacy/model/biencoder_model.py deleted file mode 100644 index df787686b41..00000000000 --- a/megatron/legacy/model/biencoder_model.py +++ /dev/null @@ -1,329 +0,0 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -import os -import torch -import sys - -from megatron.training import get_args, print_rank_0, get_tokenizer -from megatron.core import mpu -from megatron.training.checkpointing import fix_query_key_value_ordering -from megatron.training.checkpointing import get_checkpoint_tracker_filename -from megatron.training.checkpointing import get_checkpoint_name -from megatron.legacy.model.bert_model import bert_position_ids -from megatron.legacy.model.enums import AttnMaskType -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 - -def get_model_provider(only_query_model=False, only_context_model=False, - biencoder_shared_query_context_model=False): - - def model_provider(pre_process=True, post_process=True): - """Build the model.""" - - print_rank_0('building Bienoder model ...') - model = biencoder_model_provider(only_query_model=only_query_model, - only_context_model = only_context_model, - biencoder_shared_query_context_model = \ - biencoder_shared_query_context_model, - pre_process=pre_process, post_process=post_process) - - return model - - return model_provider - - -def biencoder_model_provider(only_query_model=False, - only_context_model=False, - biencoder_shared_query_context_model=False, - pre_process=True, - post_process=True): - """Build the model.""" - - assert mpu.get_tensor_model_parallel_world_size() == 1 and \ - mpu.get_pipeline_model_parallel_world_size() == 1, \ - "Model parallel size > 1 not supported for ICT" - - print_rank_0('building BiEncoderModel...') - - # simpler to just keep using 2 tokentypes since - # the LM we initialize with has 2 tokentypes - model = BiEncoderModel( - num_tokentypes=2, - parallel_output=False, - only_query_model=only_query_model, - only_context_model=only_context_model, - biencoder_shared_query_context_model=\ - biencoder_shared_query_context_model, - pre_process=pre_process, - post_process=post_process) - - return model - - -class BiEncoderModel(MegatronModule): - """Bert-based module for Biencoder model.""" - - def __init__(self, - num_tokentypes=1, - parallel_output=True, - only_query_model=False, - only_context_model=False, - biencoder_shared_query_context_model=False, - pre_process=True, - post_process=True): - super(BiEncoderModel, self).__init__() - args = get_args() - - bert_kwargs = dict( - num_tokentypes=num_tokentypes, - parallel_output=parallel_output, - pre_process=pre_process, - post_process=post_process) - - self.biencoder_shared_query_context_model = \ - biencoder_shared_query_context_model - assert not (only_context_model and only_query_model) - self.use_context_model = not only_query_model - self.use_query_model = not only_context_model - self.biencoder_projection_dim = args.biencoder_projection_dim - - if self.biencoder_shared_query_context_model: - self.model = PretrainedBertModel(**bert_kwargs) - self._model_key = 'shared_model' - self.query_model, self.context_model = self.model, self.model - else: - if self.use_query_model: - # this model embeds (pseudo-)queries - Embed_input in the paper - self.query_model = PretrainedBertModel(**bert_kwargs) - self._query_key = 'query_model' - - if self.use_context_model: - # this model embeds evidence blocks - Embed_doc in the paper - self.context_model = PretrainedBertModel(**bert_kwargs) - self._context_key = 'context_model' - - def set_input_tensor(self, input_tensor): - """See megatron.legacy.model.transformer.set_input_tensor()""" - # this is just a placeholder and will be needed when model - # parallelism will be used - # self.language_model.set_input_tensor(input_tensor) - return - - def forward(self, query_tokens, query_attention_mask, query_types, - context_tokens, context_attention_mask, context_types): - """Run a forward pass for each of the models and - return the respective embeddings.""" - - if self.use_query_model: - query_logits = self.embed_text(self.query_model, - query_tokens, - query_attention_mask, - query_types) - else: - raise ValueError("Cannot embed query without the query model.") - if self.use_context_model: - context_logits = self.embed_text(self.context_model, - context_tokens, - context_attention_mask, - context_types) - else: - raise ValueError("Cannot embed block without the block model.") - return query_logits, context_logits - - @staticmethod - def embed_text(model, tokens, attention_mask, token_types): - """Embed a batch of tokens using the model""" - logits = model(tokens, - attention_mask, - token_types) - return logits - - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): - """Save dict with state dicts of each of the models.""" - state_dict_ = {} - if self.biencoder_shared_query_context_model: - state_dict_[self._model_key] = \ - self.model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) - else: - if self.use_query_model: - state_dict_[self._query_key] = \ - self.query_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) - - if self.use_context_model: - state_dict_[self._context_key] = \ - self.context_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) - - return state_dict_ - - def load_state_dict(self, state_dict, strict=True): - """Load the state dicts of each of the models""" - if self.biencoder_shared_query_context_model: - print_rank_0("Loading shared query-context model") - self.model.load_state_dict(state_dict[self._model_key], \ - strict=strict) - else: - if self.use_query_model: - print_rank_0("Loading query model") - self.query_model.load_state_dict( \ - state_dict[self._query_key], strict=strict) - - if self.use_context_model: - print_rank_0("Loading context model") - self.context_model.load_state_dict( \ - state_dict[self._context_key], strict=strict) - - def init_state_dict_from_bert(self): - """Initialize the state from a pretrained BERT model - on iteration zero of ICT pretraining""" - args = get_args() - - if args.bert_load is None: - print_rank_0("bert-load argument is None") - return - - tracker_filename = get_checkpoint_tracker_filename(args.bert_load) - if not os.path.isfile(tracker_filename): - raise FileNotFoundError("Could not find BERT checkpoint") - with open(tracker_filename, 'r') as f: - iteration = int(f.read().strip()) - assert iteration > 0 - - checkpoint_name = get_checkpoint_name(args.bert_load, iteration, False) - if mpu.get_data_parallel_rank() == 0: - print('global rank {} is loading BERT checkpoint {}'.format( - torch.distributed.get_rank(), checkpoint_name)) - - # Load the checkpoint. - try: - state_dict = torch.load(checkpoint_name, map_location='cpu') - except ModuleNotFoundError: - from megatron.legacy.fp16_deprecated import loss_scaler - # For backward compatibility. - print_rank_0(' > deserializing using the old code structure ...') - sys.modules['fp16.loss_scaler'] = sys.modules[ - 'megatron.fp16_deprecated.loss_scaler'] - sys.modules['megatron.fp16.loss_scaler'] = sys.modules[ - 'megatron.fp16_deprecated.loss_scaler'] - state_dict = torch.load(checkpoint_name, map_location='cpu') - sys.modules.pop('fp16.loss_scaler', None) - sys.modules.pop('megatron.fp16.loss_scaler', None) - except Exception: - print_rank_0('could not load the BERT checkpoint') - sys.exit() - - checkpoint_version = state_dict.get('checkpoint_version', 0) - - # load the LM state dict into each model - model_dict = state_dict['model']['language_model'] - - if self.biencoder_shared_query_context_model: - self.model.language_model.load_state_dict(model_dict) - fix_query_key_value_ordering(self.model, checkpoint_version) - else: - if self.use_query_model: - self.query_model.language_model.load_state_dict(model_dict) - # give each model the same ict_head to begin with as well - if self.biencoder_projection_dim > 0: - query_proj_state_dict = \ - self.state_dict_for_save_checkpoint()\ - [self._query_key]['projection_enc'] - fix_query_key_value_ordering(self.query_model, checkpoint_version) - - if self.use_context_model: - self.context_model.language_model.load_state_dict(model_dict) - if self.query_model is not None and \ - self.biencoder_projection_dim > 0: - self.context_model.projection_enc.load_state_dict\ - (query_proj_state_dict) - fix_query_key_value_ordering(self.context_model, checkpoint_version) - - -class PretrainedBertModel(MegatronModule): - """BERT-based encoder for queries or contexts used for - learned information retrieval.""" - - def __init__(self, num_tokentypes=2, - parallel_output=True, pre_process=True, post_process=True): - super(PretrainedBertModel, self).__init__() - - args = get_args() - tokenizer = get_tokenizer() - self.pad_id = tokenizer.pad - self.biencoder_projection_dim = args.biencoder_projection_dim - self.parallel_output = parallel_output - self.pre_process = pre_process - self.post_process = post_process - init_method = init_method_normal(args.init_method_std) - scaled_init_method = scaled_init_method_normal( - args.init_method_std, args.num_layers) - - self.language_model, self._language_model_key = get_language_model( - num_tokentypes=num_tokentypes, - add_pooler=False, - encoder_attn_mask_type=AttnMaskType.padding, - init_method=init_method, - scaled_init_method=scaled_init_method, - pre_process=self.pre_process, - post_process=self.post_process) - - if args.biencoder_projection_dim > 0: - self.projection_enc = get_linear_layer(args.hidden_size, - args.biencoder_projection_dim, - init_method) - self._projection_enc_key = 'projection_enc' - - def forward(self, input_ids, attention_mask, tokentype_ids=None): - extended_attention_mask = attention_mask.unsqueeze(1) - #extended_attention_mask = bert_extended_attention_mask(attention_mask) - position_ids = bert_position_ids(input_ids) - - lm_output = self.language_model(input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids) - # This mask will be used in average-pooling and max-pooling - pool_mask = (input_ids == self.pad_id).unsqueeze(2) - - # Taking the representation of the [CLS] token of BERT - pooled_output = lm_output[0, :, :] - - # Converting to float16 dtype - pooled_output = pooled_output.to(lm_output.dtype) - - # Output. - if self.biencoder_projection_dim: - pooled_output = self.projection_enc(pooled_output) - - return pooled_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.biencoder_projection_dim > 0: - state_dict_[self._projection_enc_key] = \ - self.projection_enc.state_dict(prefix=prefix, - keep_vars=keep_vars) - - return state_dict_ - - def load_state_dict(self, state_dict, strict=True): - """Customized load.""" - print_rank_0("loading pretrained weights") - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) - - if self.biencoder_projection_dim > 0: - print_rank_0("loading projection head weights") - self.projection_enc.load_state_dict( - state_dict[self._projection_enc_key], strict=strict) diff --git a/megatron/legacy/model/realm_model.py b/megatron/legacy/model/realm_model.py deleted file mode 100644 index 1999cdb07cd..00000000000 --- a/megatron/legacy/model/realm_model.py +++ /dev/null @@ -1,205 +0,0 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -import os -import torch - -from megatron.training import get_args, print_rank_0 -from megatron.training.checkpointing import get_checkpoint_tracker_filename, get_checkpoint_name -from megatron.legacy.model import BertModel -from .module import MegatronModule -from megatron.core import mpu -from megatron.legacy.model.enums import AttnMaskType -from megatron.legacy.model.utils import get_linear_layer -from megatron.legacy.model.utils import init_method_normal -from megatron.legacy.model.language_model import get_language_model -from megatron.legacy.model.utils import scaled_init_method_normal -from megatron.legacy.model.bert_model import bert_extended_attention_mask, bert_position_ids - - -def general_ict_model_provider(only_query_model=False, only_block_model=False): - """Build the model.""" - args = get_args() - assert args.ict_head_size is not None, \ - "Need to specify --ict-head-size to provide an ICTBertModel" - assert mpu.get_tensor_model_parallel_world_size() == 1 and mpu.get_pipeline_model_parallel_world_size() == 1, \ - "Model parallel size > 1 not supported for ICT" - - print_rank_0('building ICTBertModel...') - - # simpler to just keep using 2 tokentypes since the LM we initialize with has 2 tokentypes - model = ICTBertModel( - ict_head_size=args.ict_head_size, - num_tokentypes=2, - parallel_output=True, - only_query_model=only_query_model, - only_block_model=only_block_model) - - return model - - -class ICTBertModel(MegatronModule): - """Bert-based module for Inverse Cloze task.""" - def __init__(self, - ict_head_size, - num_tokentypes=1, - parallel_output=True, - only_query_model=False, - only_block_model=False): - super(ICTBertModel, self).__init__() - bert_kwargs = dict( - ict_head_size=ict_head_size, - num_tokentypes=num_tokentypes, - parallel_output=parallel_output - ) - assert not (only_block_model and only_query_model) - self.use_block_model = not only_query_model - self.use_query_model = not only_block_model - - if self.use_query_model: - # this model embeds (pseudo-)queries - Embed_input in the paper - self.query_model = IREncoderBertModel(**bert_kwargs) - self._query_key = 'question_model' - - if self.use_block_model: - # this model embeds evidence blocks - Embed_doc in the paper - self.block_model = IREncoderBertModel(**bert_kwargs) - self._block_key = 'context_model' - - def forward(self, query_tokens, query_attention_mask, block_tokens, block_attention_mask): - """Run a forward pass for each of the models and return the respective embeddings.""" - query_logits = self.embed_query(query_tokens, query_attention_mask) - block_logits = self.embed_block(block_tokens, block_attention_mask) - return query_logits, block_logits - - def embed_query(self, query_tokens, query_attention_mask): - """Embed a batch of tokens using the query model""" - if self.use_query_model: - query_types = torch.cuda.LongTensor(*query_tokens.shape).fill_(0) - query_ict_logits, _ = self.query_model.forward(query_tokens, query_attention_mask, query_types) - return query_ict_logits - else: - raise ValueError("Cannot embed query without query model.") - - def embed_block(self, block_tokens, block_attention_mask): - """Embed a batch of tokens using the block model""" - if self.use_block_model: - block_types = torch.cuda.LongTensor(*block_tokens.shape).fill_(0) - block_ict_logits, _ = self.block_model.forward(block_tokens, block_attention_mask, block_types) - return block_ict_logits - else: - raise ValueError("Cannot embed block without block model.") - - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): - """Save dict with state dicts of each of the models.""" - state_dict_ = {} - if self.use_query_model: - state_dict_[self._query_key] \ - = self.query_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) - - if self.use_block_model: - state_dict_[self._block_key] \ - = self.block_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) - - return state_dict_ - - def load_state_dict(self, state_dict, strict=True): - """Load the state dicts of each of the models""" - if self.use_query_model: - print("Loading ICT query model", flush=True) - self.query_model.load_state_dict( - state_dict[self._query_key], strict=strict) - - if self.use_block_model: - print("Loading ICT block model", flush=True) - self.block_model.load_state_dict( - state_dict[self._block_key], strict=strict) - - def init_state_dict_from_bert(self): - """Initialize the state from a pretrained BERT model on iteration zero of ICT pretraining""" - args = get_args() - tracker_filename = get_checkpoint_tracker_filename(args.bert_load) - if not os.path.isfile(tracker_filename): - raise FileNotFoundError("Could not find BERT load for ICT") - with open(tracker_filename, 'r') as f: - iteration = int(f.read().strip()) - assert iteration > 0 - - checkpoint_name = get_checkpoint_name(args.bert_load, iteration, False) - if mpu.get_data_parallel_rank() == 0: - print('global rank {} is loading checkpoint {}'.format( - torch.distributed.get_rank(), checkpoint_name)) - - try: - state_dict = torch.load(checkpoint_name, map_location='cpu') - except Exception: - raise ValueError("Could not load checkpoint") - - # load the LM state dict into each model - model_dict = state_dict['model']['language_model'] - self.query_model.language_model.load_state_dict(model_dict) - self.block_model.language_model.load_state_dict(model_dict) - - # give each model the same ict_head to begin with as well - query_ict_head_state_dict = self.state_dict_for_save_checkpoint()[self._query_key]['ict_head'] - self.block_model.ict_head.load_state_dict(query_ict_head_state_dict) - - -class IREncoderBertModel(MegatronModule): - """BERT-based encoder for queries or blocks used for learned information retrieval.""" - def __init__(self, ict_head_size, num_tokentypes=2, parallel_output=True): - super(IREncoderBertModel, self).__init__() - args = get_args() - - self.ict_head_size = ict_head_size - self.parallel_output = parallel_output - init_method = init_method_normal(args.init_method_std) - scaled_init_method = scaled_init_method_normal(args.init_method_std, - args.num_layers) - - self.language_model, self._language_model_key = get_language_model( - num_tokentypes=num_tokentypes, - add_pooler=True, - encoder_attn_mask_type=AttnMaskType.padding, - init_method=init_method, - scaled_init_method=scaled_init_method) - - self.ict_head = get_linear_layer(args.hidden_size, ict_head_size, init_method) - self._ict_head_key = 'ict_head' - - def forward(self, input_ids, attention_mask, tokentype_ids=None): - extended_attention_mask = bert_extended_attention_mask( - attention_mask, next(self.language_model.parameters()).dtype) - position_ids = bert_position_ids(input_ids) - - lm_output, pooled_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids) - - # Output. - ict_logits = self.ict_head(pooled_output) - return ict_logits, None - - 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) - state_dict_[self._ict_head_key] \ - = self.ict_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) - self.ict_head.load_state_dict( - state_dict[self._ict_head_key], strict=strict) - -