indexer.py 8.44 KB
Newer Older
1
2
3
import os
import time

Neel Kant's avatar
Neel Kant committed
4
5
6
7
8
9
10
import torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP

from megatron import get_args
from megatron import mpu
from megatron.checkpointing import get_checkpoint_tracker_filename, get_checkpoint_name
from megatron.data.bert_dataset import get_indexed_dataset_
11
from megatron.data.realm_dataset import ICTDataset
12
from megatron.data.realm_index import detach, BlockData, RandProjectionLSHIndex
Neel Kant's avatar
Neel Kant committed
13
14
from megatron.data.samplers import DistributedBatchSampler
from megatron.initialize import initialize_megatron
Neel Kant's avatar
Neel Kant committed
15
from megatron.model import REALMRetriever
Neel Kant's avatar
Neel Kant committed
16
17
18
19
from megatron.training import get_model
from pretrain_bert_ict import get_batch, model_provider


Neel Kant's avatar
Neel Kant committed
20
def test_retriever():
21
    # TODO: Update this because it's outdated and definitely won't run.
Neel Kant's avatar
Neel Kant committed
22
23
    initialize_megatron(extra_args_provider=None,
                        args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})
Neel Kant's avatar
Neel Kant committed
24
    args = get_args()
Neel Kant's avatar
Neel Kant committed
25
    model = load_ict_checkpoint(only_block_model=True)
Neel Kant's avatar
Neel Kant committed
26
    model.eval()
Neel Kant's avatar
Neel Kant committed
27
    dataset = get_ict_dataset()
Neel Kant's avatar
Neel Kant committed
28
    hashed_index = HashedIndex.load_from_file(args.hash_data_path)
Neel Kant's avatar
Neel Kant committed
29
    retriever = REALMRetriever(model, dataset, hashed_index)
Neel Kant's avatar
Neel Kant committed
30
31
32
33
34
35
36
37
38
39

    strs = [
        "The last monarch from the house of windsor",
        "married to Elvis Presley",
        "tallest building in the world today",
        "who makes graphics cards"
    ]

    for s in strs:
        retriever.retrieve_evidence_blocks_text(s)
Neel Kant's avatar
Neel Kant committed
40
41


Neel Kant's avatar
Neel Kant committed
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def main():

    # TODO
    # consider broadcasting/all-reducing all in memory rather than using the filesystem
    # create a different process group in the same nccl world - don't have to use chkpts on disc or transfer things on disc
    # torch distributed new group, constains a list of rank, gives back a group which I can hand to the collective operations
    # create a training process group, indexing process group
    # pass the training group to the distributed DDP, instead of the large world process group
    # use indexing process group for the shard-combining
    # communication group between process "8" and process "0" which tells training group that there's a new index
    # also, process 0 sends process 8 the new model

    # if i want to launch a separate process for indexing, may have to work with environment variables to
    # allocate the resources well. Have to subsequently assign the correct gpus to the indexing job
    # consider initializing everything in a single group and break off processes based on the ranks

Neel Kant's avatar
Neel Kant committed
58
59
60
    # for debugging purposes, make it so that the training process group checks every some number of intervals
    # and if it isn't ready, then wait so that it's consistent. Start with using the filesystem

Neel Kant's avatar
Neel Kant committed
61
62
63
    initialize_megatron(extra_args_provider=None,
                        args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})
    args = get_args()
64
    ran_once = False
Neel Kant's avatar
Neel Kant committed
65
    while True:
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
        model = load_ict_checkpoint(only_block_model=True, no_grad=True, from_realm_chkpt=ran_once)
        model.eval()
        dataset = get_ict_dataset()
        data_iter = iter(get_one_epoch_dataloader(dataset))
        all_block_data = BlockData()
        hashed_index = RandProjectionLSHIndex(embed_size=128, num_buckets=32, whiten=True)

        i = 1
        total = 0
        while True:
            with torch.no_grad():
                try:
                    query_tokens, query_pad_mask, \
                    block_tokens, block_pad_mask, block_index_data = get_batch(data_iter)
                except:
81
82
                    break

83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
                block_index_data = detach(block_index_data)
                block_indices = block_index_data[:, 3]
                block_meta = block_index_data[:, :3]

                block_logits = detach(model(None, None, block_tokens, block_pad_mask, only_block=True))
                all_block_data.add_block_data(block_indices, block_logits, block_meta)

                total += block_indices.size
                i += 1
                if i % 20 == 0:
                    print('Batch {:10d} | Total {:10d}'.format(i, total), flush=True)
                    if args.debug:
                        break

        all_block_data.save_shard(args.rank)
        torch.distributed.barrier()
        del model

        if args.rank == 0:
            all_block_data.consolidate_shards_and_save()
            hashed_index.hash_whitened_block_embeds(all_block_data)
            hashed_index.save_to_file()
        else:
            all_block_data.clear()

        ran_once = True
        set_index_com_file_ready()
        torch.distributed.barrier()
        while not check_model_com_file_ready():
            time.sleep(5)

        set_model_com_file_not_ready()


INDEX_COM_FILE = 'ready.index'
MODEL_COM_FILE = 'ready.model'


def setup_index_com_file():
    set_index_com_file_not_ready()
123

Neel Kant's avatar
Neel Kant committed
124

125
126
127
def set_index_com_file_not_ready():
    with open(INDEX_COM_FILE, 'w') as com_file:
        com_file.write('0')
Neel Kant's avatar
Neel Kant committed
128

129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

def set_index_com_file_ready():
    with open(INDEX_COM_FILE, 'w') as com_file:
        com_file.write('1')


def check_index_com_file_ready():
    if os.path.exists(INDEX_COM_FILE):
        with open(INDEX_COM_FILE, 'r') as com_file:
            return bool(com_file.readline())

    return False


def setup_model_com_file():
    set_model_com_file_not_ready()


def set_model_com_file_not_ready():
    with open(MODEL_COM_FILE, 'w') as com_file:
        com_file.write('0')


def set_model_com_file_ready():
    with open(MODEL_COM_FILE, 'w') as com_file:
        com_file.write('1')


def check_model_com_file_ready():
    if os.path.exists(MODEL_COM_FILE):
        with open(MODEL_COM_FILE, 'r') as com_file:
            return bool(com_file.readline())

    return False


def load_ict_checkpoint(only_query_model=False, only_block_model=False, no_grad=False, from_realm_chkpt=False):
Neel Kant's avatar
Neel Kant committed
166
    args = get_args()
Neel Kant's avatar
Neel Kant committed
167
    model = get_model(lambda: model_provider(only_query_model, only_block_model))
Neel Kant's avatar
Neel Kant committed
168

169
170
    load_path = args.load if from_realm_chkpt else args.ict_load

Neel Kant's avatar
Neel Kant committed
171
172
    if isinstance(model, torchDDP):
        model = model.module
173
    tracker_filename = get_checkpoint_tracker_filename(load_path)
Neel Kant's avatar
Neel Kant committed
174
175
176
177
    with open(tracker_filename, 'r') as f:
        iteration = int(f.read().strip())

    assert iteration > 0
178
    checkpoint_name = get_checkpoint_name(load_path, iteration, False)
Neel Kant's avatar
Neel Kant committed
179
180
181
182
183
    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')
184
185
186
187
    ict_state_dict = state_dict['model']
    if from_realm_chkpt:
        ict_state_dict = ict_state_dict['retriever']['ict_model']

Neel Kant's avatar
Neel Kant committed
188
    if only_query_model:
189
        ict_state_dict.pop('context_model')
Neel Kant's avatar
Neel Kant committed
190
    if only_block_model:
191
        ict_state_dict.pop('question_model')
Neel Kant's avatar
Neel Kant committed
192
193
    if no_grad:
        with torch.no_grad():
194
            model.load_state_dict(ict_state_dict)
Neel Kant's avatar
Neel Kant committed
195
    else:
196
        model.load_state_dict(ict_state_dict)
Neel Kant's avatar
Neel Kant committed
197
198
199
200
201
202
203
204
    torch.distributed.barrier()

    if mpu.get_data_parallel_rank() == 0:
        print(' successfully loaded {}'.format(checkpoint_name))

    return model


205
def get_ict_dataset(use_titles=True):
Neel Kant's avatar
Neel Kant committed
206
    args = get_args()
Neel Kant's avatar
Neel Kant committed
207
    block_dataset = get_indexed_dataset_(args.data_path, 'mmap', True)
Neel Kant's avatar
Neel Kant committed
208
    titles_dataset = get_indexed_dataset_(args.titles_data_path, 'mmap', True)
Neel Kant's avatar
Neel Kant committed
209
210
211

    kwargs = dict(
        name='full',
Neel Kant's avatar
Neel Kant committed
212
213
        block_dataset=block_dataset,
        title_dataset=titles_dataset,
Neel Kant's avatar
Neel Kant committed
214
        data_prefix=args.data_path,
Neel Kant's avatar
Neel Kant committed
215
216
        num_epochs=1,
        max_num_samples=None,
Neel Kant's avatar
Neel Kant committed
217
218
        max_seq_length=288,  # doesn't matter
        short_seq_prob=0.0001,  # doesn't matter
219
220
        seed=1,
        use_titles=use_titles
Neel Kant's avatar
Neel Kant committed
221
    )
222
    dataset = ICTDataset(**kwargs)
Neel Kant's avatar
Neel Kant committed
223
224
225
    return dataset


Neel Kant's avatar
Neel Kant committed
226
def get_one_epoch_dataloader(dataset):
Neel Kant's avatar
Neel Kant committed
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
    args = get_args()

    world_size = mpu.get_data_parallel_world_size()
    rank = mpu.get_data_parallel_rank()
    global_batch_size = args.batch_size * world_size
    num_workers = args.num_workers

    sampler = torch.utils.data.SequentialSampler(dataset)
    batch_sampler = DistributedBatchSampler(sampler,
                                            batch_size=global_batch_size,
                                            drop_last=True,
                                            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__":
Neel Kant's avatar
Neel Kant committed
248
    main()