Loading indexer.py 0 → 100644 +185 −0 Original line number Diff line number Diff line import os import sys import time import torch import torch.distributed as dist from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP from megatron import get_args, get_adlr_autoresume, print_rank_0 from megatron import mpu from megatron.checkpointing import get_checkpoint_tracker_filename, get_checkpoint_name from megatron.data.dataset_utils import get_indexed_dataset_ from megatron.data.realm_dataset import ICTDataset from megatron.data.realm_dataset_utils import BlockSampleData from megatron.data.realm_index import detach, BlockData, FaissMIPSIndex from megatron.data.samplers import DistributedBatchSampler from megatron.initialize import initialize_megatron from megatron.training import get_model from pretrain_bert_ict import get_batch, general_ict_model_provider def pprint(*args): print(*args, flush=True) class IndexBuilder(object): """Object for taking one pass over a dataset and creating a BlockData of its embeddings""" def __init__(self): args = get_args() self.model = None self.dataloader = None self.block_data = None self.load_attributes() self.is_main_builder = args.rank == 0 self.iteration = self.total_processed = 0 def load_attributes(self): """Load the necessary attributes: model, dataloader and empty BlockData""" # TODO: handle from_realm_chkpt correctly self.model = load_ict_checkpoint(only_block_model=True, from_realm_chkpt=False) self.model.eval() self.dataloader = iter(get_one_epoch_dataloader(get_ict_dataset())) self.block_data = BlockData() def track_and_report_progress(self, batch_size): """Utility function for tracking progress""" self.iteration += 1 self.total_processed += batch_size if self.iteration % 10 == 0: print('Batch {:10d} | Total {:10d}'.format(self.iteration, self.total_processed), flush=True) def build_and_save_index(self): """Goes through one epoch of the dataloader and adds all data to this instance's BlockData. The copy of BlockData is saved as a shard, which when run in a distributed setting will be consolidated by the rank 0 process and saved as a final pickled BlockData. """ while True: try: # batch also has query_tokens and query_pad_data _, _, block_tokens, block_pad_mask, block_sample_data = get_batch(self.dataloader) except: break # detach, setup and add to BlockData unwrapped_model = self.model while not hasattr(unwrapped_model, 'embed_block'): unwrapped_model = unwrapped_model.module block_logits = detach(unwrapped_model.embed_block(block_tokens, block_pad_mask)) detached_data = detach(block_sample_data) block_indices = detached_data[:, 3] block_metas = detached_data[:, :3] self.block_data.add_block_data(block_indices, block_logits, block_metas) self.track_and_report_progress(batch_size=block_tokens.shape[0]) # This process signals to finalize its shard and then synchronize with the other processes self.block_data.save_shard() torch.distributed.barrier() del self.model # rank 0 process builds the final copy if self.is_main_builder: self.block_data.merge_shards_and_save() self.block_data.clear() def load_ict_checkpoint(only_query_model=False, only_block_model=False, from_realm_chkpt=False): """load ICT checkpoints for indexing/retrieving. Arguments specify which parts of the state dict to actually use.""" args = get_args() model = get_model(lambda: general_ict_model_provider(only_query_model, only_block_model)) if isinstance(model, torchDDP): model = model.module load_path = args.load if from_realm_chkpt else args.ict_load tracker_filename = get_checkpoint_tracker_filename(load_path) with open(tracker_filename, 'r') as f: iteration = int(f.read().strip()) # assert iteration > 0 checkpoint_name = get_checkpoint_name(load_path, iteration, False) if mpu.get_data_parallel_rank() == 0: print('global rank {} is loading checkpoint {}'.format( torch.distributed.get_rank(), checkpoint_name)) state_dict = torch.load(checkpoint_name, map_location='cpu') ict_state_dict = state_dict['model'] if from_realm_chkpt: print(">>>> Attempting to get ict state dict from realm", flush=True) ict_state_dict = ict_state_dict['retriever']['ict_model'] if only_query_model: ict_state_dict.pop('context_model') if only_block_model: ict_state_dict.pop('question_model') model.load_state_dict(ict_state_dict) torch.distributed.barrier() if mpu.get_data_parallel_rank() == 0: print(' successfully loaded {}'.format(checkpoint_name)) return model def get_ict_dataset(use_titles=True, query_in_block_prob=1): """Get a dataset which uses block samples mappings to get ICT/block indexing data""" args = get_args() block_dataset = get_indexed_dataset_(args.data_path, 'mmap', True) titles_dataset = get_indexed_dataset_(args.titles_data_path, 'mmap', True) kwargs = dict( name='full', block_dataset=block_dataset, title_dataset=titles_dataset, data_prefix=args.data_path, num_epochs=1, max_num_samples=None, max_seq_length=args.seq_length, short_seq_prob=0.0001, # doesn't matter seed=1, query_in_block_prob=query_in_block_prob, use_titles=use_titles, use_one_sent_docs=True ) dataset = ICTDataset(**kwargs) return dataset def get_one_epoch_dataloader(dataset, batch_size=None): """Specifically one epoch to be used in an indexing job.""" args = get_args() world_size = mpu.get_data_parallel_world_size() rank = mpu.get_data_parallel_rank() if batch_size is None: batch_size = args.batch_size global_batch_size = batch_size * world_size num_workers = args.num_workers sampler = torch.utils.data.SequentialSampler(dataset) # importantly, drop_last must be False to get all the data. batch_sampler = DistributedBatchSampler(sampler, batch_size=global_batch_size, drop_last=False, rank=rank, world_size=world_size) return torch.utils.data.DataLoader(dataset, batch_sampler=batch_sampler, num_workers=num_workers, pin_memory=True) if __name__ == "__main__": # This usage is for basic (as opposed to realm async) indexing jobs. initialize_megatron(extra_args_provider=None, args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'}) index_builder = IndexBuilder() index_builder.build_and_save_index() Loading
indexer.py 0 → 100644 +185 −0 Original line number Diff line number Diff line import os import sys import time import torch import torch.distributed as dist from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP from megatron import get_args, get_adlr_autoresume, print_rank_0 from megatron import mpu from megatron.checkpointing import get_checkpoint_tracker_filename, get_checkpoint_name from megatron.data.dataset_utils import get_indexed_dataset_ from megatron.data.realm_dataset import ICTDataset from megatron.data.realm_dataset_utils import BlockSampleData from megatron.data.realm_index import detach, BlockData, FaissMIPSIndex from megatron.data.samplers import DistributedBatchSampler from megatron.initialize import initialize_megatron from megatron.training import get_model from pretrain_bert_ict import get_batch, general_ict_model_provider def pprint(*args): print(*args, flush=True) class IndexBuilder(object): """Object for taking one pass over a dataset and creating a BlockData of its embeddings""" def __init__(self): args = get_args() self.model = None self.dataloader = None self.block_data = None self.load_attributes() self.is_main_builder = args.rank == 0 self.iteration = self.total_processed = 0 def load_attributes(self): """Load the necessary attributes: model, dataloader and empty BlockData""" # TODO: handle from_realm_chkpt correctly self.model = load_ict_checkpoint(only_block_model=True, from_realm_chkpt=False) self.model.eval() self.dataloader = iter(get_one_epoch_dataloader(get_ict_dataset())) self.block_data = BlockData() def track_and_report_progress(self, batch_size): """Utility function for tracking progress""" self.iteration += 1 self.total_processed += batch_size if self.iteration % 10 == 0: print('Batch {:10d} | Total {:10d}'.format(self.iteration, self.total_processed), flush=True) def build_and_save_index(self): """Goes through one epoch of the dataloader and adds all data to this instance's BlockData. The copy of BlockData is saved as a shard, which when run in a distributed setting will be consolidated by the rank 0 process and saved as a final pickled BlockData. """ while True: try: # batch also has query_tokens and query_pad_data _, _, block_tokens, block_pad_mask, block_sample_data = get_batch(self.dataloader) except: break # detach, setup and add to BlockData unwrapped_model = self.model while not hasattr(unwrapped_model, 'embed_block'): unwrapped_model = unwrapped_model.module block_logits = detach(unwrapped_model.embed_block(block_tokens, block_pad_mask)) detached_data = detach(block_sample_data) block_indices = detached_data[:, 3] block_metas = detached_data[:, :3] self.block_data.add_block_data(block_indices, block_logits, block_metas) self.track_and_report_progress(batch_size=block_tokens.shape[0]) # This process signals to finalize its shard and then synchronize with the other processes self.block_data.save_shard() torch.distributed.barrier() del self.model # rank 0 process builds the final copy if self.is_main_builder: self.block_data.merge_shards_and_save() self.block_data.clear() def load_ict_checkpoint(only_query_model=False, only_block_model=False, from_realm_chkpt=False): """load ICT checkpoints for indexing/retrieving. Arguments specify which parts of the state dict to actually use.""" args = get_args() model = get_model(lambda: general_ict_model_provider(only_query_model, only_block_model)) if isinstance(model, torchDDP): model = model.module load_path = args.load if from_realm_chkpt else args.ict_load tracker_filename = get_checkpoint_tracker_filename(load_path) with open(tracker_filename, 'r') as f: iteration = int(f.read().strip()) # assert iteration > 0 checkpoint_name = get_checkpoint_name(load_path, iteration, False) if mpu.get_data_parallel_rank() == 0: print('global rank {} is loading checkpoint {}'.format( torch.distributed.get_rank(), checkpoint_name)) state_dict = torch.load(checkpoint_name, map_location='cpu') ict_state_dict = state_dict['model'] if from_realm_chkpt: print(">>>> Attempting to get ict state dict from realm", flush=True) ict_state_dict = ict_state_dict['retriever']['ict_model'] if only_query_model: ict_state_dict.pop('context_model') if only_block_model: ict_state_dict.pop('question_model') model.load_state_dict(ict_state_dict) torch.distributed.barrier() if mpu.get_data_parallel_rank() == 0: print(' successfully loaded {}'.format(checkpoint_name)) return model def get_ict_dataset(use_titles=True, query_in_block_prob=1): """Get a dataset which uses block samples mappings to get ICT/block indexing data""" args = get_args() block_dataset = get_indexed_dataset_(args.data_path, 'mmap', True) titles_dataset = get_indexed_dataset_(args.titles_data_path, 'mmap', True) kwargs = dict( name='full', block_dataset=block_dataset, title_dataset=titles_dataset, data_prefix=args.data_path, num_epochs=1, max_num_samples=None, max_seq_length=args.seq_length, short_seq_prob=0.0001, # doesn't matter seed=1, query_in_block_prob=query_in_block_prob, use_titles=use_titles, use_one_sent_docs=True ) dataset = ICTDataset(**kwargs) return dataset def get_one_epoch_dataloader(dataset, batch_size=None): """Specifically one epoch to be used in an indexing job.""" args = get_args() world_size = mpu.get_data_parallel_world_size() rank = mpu.get_data_parallel_rank() if batch_size is None: batch_size = args.batch_size global_batch_size = batch_size * world_size num_workers = args.num_workers sampler = torch.utils.data.SequentialSampler(dataset) # importantly, drop_last must be False to get all the data. batch_sampler = DistributedBatchSampler(sampler, batch_size=global_batch_size, drop_last=False, rank=rank, world_size=world_size) return torch.utils.data.DataLoader(dataset, batch_sampler=batch_sampler, num_workers=num_workers, pin_memory=True) if __name__ == "__main__": # This usage is for basic (as opposed to realm async) indexing jobs. initialize_megatron(extra_args_provider=None, args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'}) index_builder = IndexBuilder() index_builder.build_and_save_index()
