finetune_utils.py 11.3 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
#
# 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."""

Jared Casper's avatar
Jared Casper committed
18
19
from functools import partial

20
21
import torch

Neel Kant's avatar
Neel Kant committed
22
23
from megatron import get_args
from megatron import print_rank_0
Mohammad's avatar
Mohammad committed
24
from megatron import get_timers
25
from megatron import mpu
Neel Kant's avatar
Neel Kant committed
26
from megatron.checkpointing import load_checkpoint
Mohammad's avatar
Mohammad committed
27
from megatron.checkpointing import save_checkpoint
28
29
30
31
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
32
from megatron.utils import average_losses_across_data_parallel_group
mohammad's avatar
mohammad committed
33
34
from megatron.utils import calc_params_l2_norm
from megatron.utils import check_adlr_autoresume_termination
35
36


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

    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


Jared Casper's avatar
Jared Casper committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def cross_entropy_loss_func(labels, output_tensor):
    logits = output_tensor

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

    # Reduce loss for logging.
    averaged_loss = average_losses_across_data_parallel_group([loss])

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


def _cross_entropy_forward_step(batch, model):
65
    """Simple forward step with cross-entropy loss."""
Mohammad's avatar
Mohammad committed
66
    timers = get_timers()
67
68

    # Get the batch.
mohammad's avatar
mohammad committed
69
    timers('batch-generator').start()
70
71
    try:
        batch_ = next(batch)
Neel Kant's avatar
Neel Kant committed
72
    except BaseException:
73
        batch_ = batch
Mohammad's avatar
Mohammad committed
74
    tokens, types, labels, attention_mask = process_batch(batch_)
mohammad's avatar
mohammad committed
75
    timers('batch-generator').stop()
76
77

    # Forward model.
Jared Casper's avatar
Jared Casper committed
78
    output_tensor = model(tokens, attention_mask, tokentype_ids=types)
79

Jared Casper's avatar
Jared Casper committed
80
    return output_tensor, partial(cross_entropy_loss_func, labels)
81
82


83
def build_data_loader(dataset, micro_batch_size, num_workers, drop_last):
84
85
86
87
88
89
90
91
92
93
    """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,
94
                                              batch_size=micro_batch_size,
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
                                              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
115
def _build_train_valid_dataloaders(train_dataset, valid_dataset):
116
    """Traing and validation dataloaders."""
Mohammad's avatar
Mohammad committed
117
    args = get_args()
118
119
120

    print_rank_0('building train and validation dataloaders ...')
    # Training dataset.
121
    train_dataloader = build_data_loader(train_dataset, args.micro_batch_size,
122
123
124
125
126
127
                                         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.
128
    valid_dataloader_ = build_data_loader(valid_dataset, args.micro_batch_size,
129
130
131
                                          args.num_workers, not args.keep_last)
    valid_dataloader = _build_infinite_size_dataloader(valid_dataloader_)

132
133
134
135
136
    # Now that we've built the data loaders, set batch_size arguments
    # to the actual batch size the model will see for this dataset.
    # This is necessary so pipeline transfers know what size they are
    # and the LR schedule, which is based on samples seen, gets set
    # correctly.
Jared Casper's avatar
Jared Casper committed
137
138
    args.orig_micro_batch_size = args.micro_batch_size
    args.orig_global_batch_size = args.global_batch_size
139
    if hasattr(train_dataset, 'sample_multiplier'):
140
141
142
143
144
        # If our dataset as a sample_multiplier attribute that means
        # each "sample" from the dataset actually has multiple samples
        # that will collapse into the batch dimension (for example in
        # the RACE dataset that has several options), we need to
        # account for that when setting the micro batch size.
145
        args.micro_batch_size *= train_dataset.sample_multiplier
146
        args.global_batch_size *= train_dataset.sample_multiplier
147

148
149
150
151
    return train_dataloader, valid_dataloader


def _train(model, optimizer, lr_scheduler, forward_step,
Mohammad's avatar
Mohammad committed
152
           train_dataloader, valid_dataloader, end_of_epoch_callback):
153
    """Train the model."""
Mohammad's avatar
Mohammad committed
154
155
    args = get_args()
    timers = get_timers()
156
157

    # Turn on training mode which enables dropout.
Jared Casper's avatar
Jared Casper committed
158
159
    for m in model:
        m.train()
160
161
162
163
164
165
166
167
168
169
170
171
172

    # 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
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
173
    timers('interval-time').start()
174
    for epoch in range(start_epoch, args.epochs):
Neel Kant's avatar
Neel Kant committed
175
        print_rank_0('working on epoch {} ...'.format(epoch + 1))
176
177
178
179
180
181
182
183
184
185
186
187
188
189

        # 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.
Jared Casper's avatar
Jared Casper committed
190
191
            out = train_step(forward_step, batch, model, optimizer, lr_scheduler)
            losses_dict, skipped_iter, grad_norm, num_zeros_in_grad = out
192
193
194
            iteration += 1

            # Logging.
195
196
197
            params_norm = None
            if args.log_params_norm:
                params_norm = calc_params_l2_norm(model)
198
199
            report_memory_flag = training_log(losses_dict, losses_dict_sum,
                                              optimizer.param_groups[0]['lr'],
200
201
                                              iteration,
                                              optimizer.get_loss_scale().item(),
202
                                              report_memory_flag, skipped_iter,
Jared Casper's avatar
Jared Casper committed
203
                                              grad_norm, params_norm, num_zeros_in_grad)
204
205

            # Autoresume
Neel Kant's avatar
Neel Kant committed
206
            if args.adlr_autoresume and \
207
               (iteration % args.adlr_autoresume_interval == 0):
Mohammad's avatar
Mohammad committed
208
209
                check_adlr_autoresume_termination(iteration, model,
                                                  optimizer, lr_scheduler)
210
211
212
213

            # Checkpointing
            if args.save and args.save_interval and \
               iteration % args.save_interval == 0:
Mohammad's avatar
Mohammad committed
214
                save_checkpoint(iteration, model, optimizer, lr_scheduler)
215
216
217
218
219

            # 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
220
221
                                           valid_dataloader, model,
                                           iteration, False)
222
223
224

        # Checkpointing at the end of each epoch.
        if args.save:
Mohammad's avatar
Mohammad committed
225
            save_checkpoint(iteration, model, optimizer, lr_scheduler)
226
227
228

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


Mohammad's avatar
Mohammad committed
232
def finetune(train_valid_datasets_provider, model_provider,
233
234
235
             forward_step=_cross_entropy_forward_step,
             end_of_epoch_callback_provider=None):
    """Main finetune function used across all tasks."""
Mohammad's avatar
Mohammad committed
236
237
    args = get_args()
    timers = get_timers()
238

239
240
241
    assert args.rampup_batch_size is None, \
        'batch size scaling is not supported for finetuning'

242
    # Train and validation data loaders.
Mohammad's avatar
Mohammad committed
243
    timers('train/valid/test dataset/dataloder').start()
244
    if args.epochs > 0:
Mohammad's avatar
Mohammad committed
245
        train_dataset, valid_dataset = train_valid_datasets_provider()
246
        train_dataloader, valid_dataloader = _build_train_valid_dataloaders(
Mohammad's avatar
Mohammad committed
247
            train_dataset, valid_dataset)
248
249
    else:
        args.train_iters = 0
Mohammad's avatar
Mohammad committed
250
    timers('train/valid/test dataset/dataloder').stop()
251
252

    # Build calback function.
Mohammad's avatar
Mohammad committed
253
    timers('callback function').start()
254
255
    end_of_epoch_callback = None
    if end_of_epoch_callback_provider is not None:
Mohammad's avatar
Mohammad committed
256
257
        end_of_epoch_callback = end_of_epoch_callback_provider()
    timers('callback function').stop()
258
259

    # Build model, optimizer and learning rate scheduler.
Mohammad's avatar
Mohammad committed
260
261
262
    timers('model and optimizer').start()
    model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider)
    timers('model and optimizer').stop()
263
264
265
266

    # 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
267
    timers('pretrained checkpoint').start()
268
269
270
    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
271
        _ = load_checkpoint(model, None, None)
272
273
        args.load = original_load
        # This is critical when only model is loaded. We should make sure
274
        # main parameters are also updated.
275
        optimizer.reload_model_params()
Mohammad's avatar
Mohammad committed
276
    timers('pretrained checkpoint').stop()
277

Mohammad's avatar
Mohammad committed
278
279
280
281
282
    # 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 ...')
283
284
285
286

    # Finetune the model.
    if args.epochs > 0:
        _train(model, optimizer, lr_scheduler, forward_step,
Mohammad's avatar
Mohammad committed
287
               train_dataloader, valid_dataloader, end_of_epoch_callback)
288
289
290
291
    # 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
292
            end_of_epoch_callback(model, epoch=-1, output_predictions=True)
293
    print_rank_0('done :-)')