finetune_utils.py 9.65 KB
Newer Older
1
# coding=utf-8
Mohammad's avatar
Mohammad committed
2
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Finetune utilities."""

import torch

Neel Kant's avatar
Neel Kant committed
20
21
from megatron import get_args
from megatron import print_rank_0
Mohammad's avatar
Mohammad committed
22
from megatron import get_timers
23
from megatron import mpu
Neel Kant's avatar
Neel Kant committed
24
from megatron.checkpointing import load_checkpoint
Mohammad's avatar
Mohammad committed
25
from megatron.checkpointing import save_checkpoint
26
27
28
29
30
31
32
33
from megatron.training import evaluate_and_print_results
from megatron.training import setup_model_and_optimizer
from megatron.training import train_step
from megatron.training import training_log
from megatron.utils import check_adlr_autoresume_termination
from megatron.utils import reduce_losses


Mohammad's avatar
Mohammad committed
34
def process_batch(batch):
35
    """Process batch and produce inputs for the model."""
Mohammad's avatar
Mohammad committed
36
    args = get_args()
37
38
39
40
41
42
43
44
45
46
47

    tokens = batch['text'].long().cuda().contiguous()
    types = batch['types'].long().cuda().contiguous()
    labels = batch['label'].long().cuda().contiguous()
    attention_mask = batch['padding_mask'].float().cuda().contiguous()
    if args.fp16:
        attention_mask = attention_mask.half()

    return tokens, types, labels, attention_mask


Mohammad's avatar
Mohammad committed
48
def _cross_entropy_forward_step(batch, model):
49
    """Simple forward step with cross-entropy loss."""
Mohammad's avatar
Mohammad committed
50
    timers = get_timers()
51
52
53
54
55

    # Get the batch.
    timers('batch generator').start()
    try:
        batch_ = next(batch)
Neel Kant's avatar
Neel Kant committed
56
    except BaseException:
57
        batch_ = batch
Mohammad's avatar
Mohammad committed
58
    tokens, types, labels, attention_mask = process_batch(batch_)
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
    timers('batch generator').stop()

    # Forward model.
    logits = model(tokens, attention_mask, types)

    # Cross-entropy loss.
    loss_func = torch.nn.CrossEntropyLoss()
    loss = loss_func(logits.contiguous().float(), labels)

    # Reduce loss for logging.
    reduced_loss = reduce_losses([loss])

    return loss, {'lm loss': reduced_loss[0]}


def build_data_loader(dataset, batch_size, num_workers, drop_last):
    """Data loader. Note that batch-size is the local (per GPU) batch-size."""

    # Sampler.
    world_size = mpu.get_data_parallel_world_size()
    rank = mpu.get_data_parallel_rank()
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset, num_replicas=world_size, rank=rank)

    # Data loader. Note that batch size is the per GPU batch size.
    data_loader = torch.utils.data.DataLoader(dataset,
                                              batch_size=batch_size,
                                              sampler=sampler,
                                              shuffle=False,
                                              num_workers=num_workers,
                                              drop_last=drop_last,
                                              pin_memory=True)

    return data_loader


def _build_infinite_size_dataloader(dataloader):
    """Build a looped dataloader with infinite size."""

    iterator = dataloader.__iter__()
    while True:
        try:
            yield iterator.__next__()
        except StopIteration:
            iterator = dataloader.__iter__()


Mohammad's avatar
Mohammad committed
106
def _build_train_valid_dataloaders(train_dataset, valid_dataset):
107
    """Traing and validation dataloaders."""
Mohammad's avatar
Mohammad committed
108
    args = get_args()
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

    print_rank_0('building train and validation dataloaders ...')
    # Training dataset.
    train_dataloader = build_data_loader(train_dataset, args.batch_size,
                                         args.num_workers, not args.keep_last)
    # Set the training iterations.
    args.train_iters_per_epoch = len(train_dataloader)
    args.train_iters = args.epochs * args.train_iters_per_epoch
    # Validation dataset. For this dataset, we do not need to set up
    # shuffling so we can just use a simple infinite loop.
    valid_dataloader_ = build_data_loader(valid_dataset, args.batch_size,
                                          args.num_workers, not args.keep_last)
    valid_dataloader = _build_infinite_size_dataloader(valid_dataloader_)

    return train_dataloader, valid_dataloader


def _train(model, optimizer, lr_scheduler, forward_step,
Mohammad's avatar
Mohammad committed
127
           train_dataloader, valid_dataloader, end_of_epoch_callback):
128
    """Train the model."""
Mohammad's avatar
Mohammad committed
129
130
    args = get_args()
    timers = get_timers()
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

    # Turn on training mode which enables dropout.
    model.train()

    # Tracking loss.
    losses_dict_sum = {}

    # Starting epoch and iteration
    start_epoch = args.iteration // args.train_iters_per_epoch
    start_iteration = args.iteration % args.train_iters_per_epoch
    iteration = args.iteration

    # Memory reporting flag.
    report_memory_flag = True

    # For each remaining epoch
    timers('interval time').start()
    for epoch in range(start_epoch, args.epochs):
Neel Kant's avatar
Neel Kant committed
149
        print_rank_0('working on epoch {} ...'.format(epoch + 1))
150
151
152
153
154
155
156
157
158
159
160
161
162
163

        # Set the data loader epoch to shuffle the index iterator.
        train_dataloader.sampler.set_epoch(args.seed + epoch)

        # For all the batches in the dataset.
        for iteration_, batch in enumerate(train_dataloader):

            # Ignore the iterations before starting value
            if iteration_ < start_iteration:
                continue
            # Set to zero so the next epoch does not skip any batches.
            start_iteration = 0

            # Train for one step.
164
165
            losses_dict, skipped_iter = train_step(forward_step, batch, model,
                                                   optimizer, lr_scheduler)
166
167
168
169
170
171
            iteration += 1

            # Logging.
            report_memory_flag = training_log(losses_dict, losses_dict_sum,
                                              optimizer.param_groups[0]['lr'],
                                              iteration, optimizer.loss_scale,
172
                                              report_memory_flag, skipped_iter)
173
174

            # Autoresume
Neel Kant's avatar
Neel Kant committed
175
            if args.adlr_autoresume and \
176
               (iteration % args.adlr_autoresume_interval == 0):
Mohammad's avatar
Mohammad committed
177
178
                check_adlr_autoresume_termination(iteration, model,
                                                  optimizer, lr_scheduler)
179
180
181
182

            # Checkpointing
            if args.save and args.save_interval and \
               iteration % args.save_interval == 0:
Mohammad's avatar
Mohammad committed
183
                save_checkpoint(iteration, model, optimizer, lr_scheduler)
184
185
186
187
188

            # Evaluation
            if args.eval_interval and iteration % args.eval_interval == 0:
                prefix = 'iteration {}'.format(iteration)
                evaluate_and_print_results(prefix, forward_step,
Mohammad's avatar
Mohammad committed
189
190
                                           valid_dataloader, model,
                                           iteration, False)
191
192
193

        # Checkpointing at the end of each epoch.
        if args.save:
Mohammad's avatar
Mohammad committed
194
            save_checkpoint(iteration, model, optimizer, lr_scheduler)
195
196
197

        # Callback at the end of each epoch.
        if end_of_epoch_callback is not None:
Mohammad's avatar
Mohammad committed
198
            end_of_epoch_callback(model, epoch)
199
200


Mohammad's avatar
Mohammad committed
201
def finetune(train_valid_datasets_provider, model_provider,
202
203
204
             forward_step=_cross_entropy_forward_step,
             end_of_epoch_callback_provider=None):
    """Main finetune function used across all tasks."""
Mohammad's avatar
Mohammad committed
205
206
    args = get_args()
    timers = get_timers()
207
208

    # Train and validation data loaders.
Mohammad's avatar
Mohammad committed
209
    timers('train/valid/test dataset/dataloder').start()
210
    if args.epochs > 0:
Mohammad's avatar
Mohammad committed
211
        train_dataset, valid_dataset = train_valid_datasets_provider()
212
        train_dataloader, valid_dataloader = _build_train_valid_dataloaders(
Mohammad's avatar
Mohammad committed
213
214
            train_dataset, valid_dataset)
    timers('train/valid/test dataset/dataloder').stop()
215
216

    # Build calback function.
Mohammad's avatar
Mohammad committed
217
    timers('callback function').start()
218
219
    end_of_epoch_callback = None
    if end_of_epoch_callback_provider is not None:
Mohammad's avatar
Mohammad committed
220
221
        end_of_epoch_callback = end_of_epoch_callback_provider()
    timers('callback function').stop()
222
223

    # Build model, optimizer and learning rate scheduler.
Mohammad's avatar
Mohammad committed
224
225
226
    timers('model and optimizer').start()
    model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider)
    timers('model and optimizer').stop()
227
228
229
230

    # If pretrained checkpoint is provided and we have not trained for
    # any iteration (i.e., iteration is zero), then load the pretrained
    # checkpoint.
Mohammad's avatar
Mohammad committed
231
    timers('pretrained checkpoint').start()
232
233
234
    if args.iteration == 0 and args.pretrained_checkpoint is not None:
        original_load = args.load
        args.load = args.pretrained_checkpoint
Mohammad's avatar
Mohammad committed
235
        _ = load_checkpoint(model, None, None)
236
237
238
239
240
        args.load = original_load
        # This is critical when only model is loaded. We should make sure
        # master parameters are also updated.
        if args.fp16:
            optimizer._model_params_to_master_params()
Mohammad's avatar
Mohammad committed
241
    timers('pretrained checkpoint').stop()
242

Mohammad's avatar
Mohammad committed
243
244
245
246
247
    # Print setup timing.
    print_rank_0('done with setups ...')
    timers.log(['train/valid/test dataset/dataloder', 'callback function',
                'model and optimizer', 'pretrained checkpoint'])
    print_rank_0('training ...')
248
249
250
251

    # Finetune the model.
    if args.epochs > 0:
        _train(model, optimizer, lr_scheduler, forward_step,
Mohammad's avatar
Mohammad committed
252
               train_dataloader, valid_dataloader, end_of_epoch_callback)
253
254
255
256
    # Or just evaluate.
    else:
        if end_of_epoch_callback is not None:
            print_rank_0('evaluation only mode, setting epoch to -1')
Mohammad's avatar
Mohammad committed
257
            end_of_epoch_callback(model, epoch=-1, output_predictions=True)
258
259

    print_rank_0('done :-)')