Support UL2 for decoder-only models

megatron/arguments.py

@@ -7,6 +7,8 @@
 
 import torch
 
+from megatron.model.enums import UL2ModelType
+
 def parse_args(extra_args_provider=None, ignore_unknown_args=False):
     """Parse all arguments."""
     parser = argparse.ArgumentParser(description='Megatron-LM Arguments',
@@ -321,6 +323,17 @@ def validate_args(args, defaults={}):
     if args.sequence_parallel:
         args.async_tensor_model_parallel_allreduce = False
 
+    args.ul2_model_type = UL2ModelType(args.ul2_model_type)
+    if (
+            args.ul2_model_type is not UL2ModelType.encoder_decoder
+            and args.decoder_seq_length is not None
+    ):
+        print(
+            f'WARNING: `--decoder_seq_length` is ignored when '
+            f'`--ul2-model-type` is not '
+            f'"{UL2ModelType.encoder_decoder.value}"!'
+        )
+
 
     if os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1":
         if args.sequence_parallel:
@@ -1072,6 +1085,11 @@ def _add_vision_args(parser):
 def _add_ul2_args(parser):
     group = parser.add_argument_group(title="UL2")
 
+    group.add_argument('--ul2-model-type', type=str, default='ED',
+                       choices=['ED', 'ND', 'CD'],
+                       help='What type of model to use for UL2 pretraining. '
+                       'ED = encoder-decoder; ND = non-causal decoder-only; '
+                       'CD = causal decoder-only')
     group.add_argument('--ul2-denoiser-ratios', nargs='+', type=float,
                        default=None,
                        help='Probability of each denoising objective to be '

megatron/data/dataset_utils.py

@@ -599,6 +599,7 @@ def build_dataset(index, name):
                 args = get_args()
                 dataset = UL2Dataset(
                     indexed_dataset=indexed_dataset,
+                    model_type=args.ul2_model_type,
                     denoiser_ratios=args.ul2_denoiser_ratios,
                     denoisers=args.ul2_denoisers,
                     mean_span_lengths=args.ul2_mean_span_lengths,

megatron/data/ul2_dataset.py

@@ -2,6 +2,8 @@
 
 """UL2-style dataset."""
 
+import math
+
 import numpy as np
 
 from megatron import get_tokenizer
@@ -10,16 +12,34 @@
     get_samples_mapping,
     SamplingStyle
 )
-from megatron.data.t5_dataset import pad_and_convert_to_numpy, T5Dataset
+from megatron.data.t5_dataset import (
+    make_history_mask,
+    merge_subsequent_masks,
+    pad_and_convert_to_numpy,
+    T5Dataset,
+)
+from megatron.model.enums import UL2ModelType
+
+
+def is_decoder_only(ul2_model_type):
+    """Return whether we use a decoder-only model."""
+    assert isinstance(ul2_model_type, UL2ModelType)
+    return ul2_model_type is not UL2ModelType.encoder_decoder
+
+
+def is_prefix_lm(ul2_model_type):
+    """Return whether we use a non-causal decoder-only model."""
+    assert isinstance(ul2_model_type, UL2ModelType)
+    return ul2_model_type is UL2ModelType.non_causal_decoder
 
 
 class UL2Dataset(T5Dataset):
 
     def __init__(self, name, indexed_dataset, data_prefix,
-                 num_epochs, max_num_samples, denoiser_ratios,
-                 denoisers, mean_span_lengths, mask_ratios,
-                 denoiser_tokens, max_seq_length, max_seq_length_dec,
-                 short_seq_prob, seed):
+                 num_epochs, max_num_samples, model_type,
+                 denoiser_ratios, denoisers, mean_span_lengths,
+                 mask_ratios, denoiser_tokens, max_seq_length,
+                 max_seq_length_dec, short_seq_prob, seed):
 
         if denoiser_ratios is None:
             # Uniform distribution by default.
@@ -39,6 +59,7 @@ def __init__(self, name, indexed_dataset, data_prefix,
                          short_seq_prob, seed)
 
         # Params to store.
+        self.model_type = model_type
         self.denoiser_ratios = [
             denoiser_ratio / sum(denoiser_ratios)
             for denoiser_ratio in denoiser_ratios
@@ -84,18 +105,17 @@ def __getitem__(self, idx):
                                      self.vocab_id_to_token_dict,
                                      self.cls_ids, self.sep_id,
                                      self.mask_id, self.pad_id,
-                                     self.denoiser_ratios, self.denoisers,
-                                     self.mean_span_lengths, self.mask_ratios,
-                                     np_rng,
-                                     self.bos_id, self.eos_id,
-                                     self.sentinel_tokens)
+                                     self.model_type, self.denoiser_ratios,
+                                     self.denoisers, self.mean_span_lengths,
+                                     self.mask_ratios, np_rng, self.bos_id,
+                                     self.eos_id, self.sentinel_tokens)
 
 
 def build_training_sample(sample, target_seq_length,
                           max_seq_length, max_seq_length_dec,
                           vocab_id_list, vocab_id_to_token_dict,
                           cls_ids, sep_id, mask_id, pad_id,
-                          denoiser_ratios, denoisers,
+                          model_type, denoiser_ratios, denoisers,
                           mean_span_lengths, mask_ratios,
                           np_rng, bos_id=None,
                           eos_id=None, sentinel_tokens=None):
@@ -112,6 +132,7 @@ def build_training_sample(sample, target_seq_length,
         sep_id: Separator id.
         mask_id: Mask token id.
         pad_id: Padding token id.
+        model_type: What type of model is used.
         denoiser_ratios: Probability of each denoising objective to be selected.
         denoisers: What type of UL2 denoising objective the other UL2
               configurations refer to.
@@ -174,22 +195,64 @@ def build_training_sample(sample, target_seq_length,
         sampling_style=sampling_style, prefix_lm=prefix_lm,
     )
 
-    # Padding.
-    tokens_enc, tokens_dec_in, labels, enc_mask, \
-    dec_mask, enc_dec_mask, loss_mask \
-        = pad_and_convert_to_numpy(tokens, masked_positions,
-                                   masked_labels, pad_id, max_seq_length,
-                                   max_seq_length_dec, masked_spans,
-                                   bos_id, eos_id, sentinel_tokens)
-
-    train_sample = {
-        'text_enc': tokens_enc,
-        'text_dec': tokens_dec_in,
-        'labels': labels,
-        'loss_mask': loss_mask,
-        'truncated': int(truncated),
-        'enc_mask': enc_mask,
-        'dec_mask': dec_mask,
-        'enc_dec_mask': enc_dec_mask,
-    }
+    if is_decoder_only(model_type):
+        # Concatenate to one sequence.
+        tokens_enc, tokens_dec_in, labels = merge_subsequent_masks(
+            tokens, masked_spans, bos_id, eos_id, sentinel_tokens)
+
+        # Move EOS tokens to end of sequence.
+        while tokens_enc[-1] == eos_id:
+            del tokens_enc[-1]
+            tokens_dec_in.append(eos_id)
+            labels.append(eos_id)
+
+        num_labels = len(labels)
+
+        # Move BOS token to start of sequence.
+        tokens_dec_in = tokens_dec_in[1:]
+        tokens = np.concatenate([
+            np.array([bos_id], dtype=np.int64),
+            tokens_enc,
+            np.array([sep_id], dtype=np.int64),
+            tokens_dec_in,
+        ])
+        labels = np.concatenate([
+            tokens_enc,
+            np.array([sep_id], dtype=np.int64),
+            labels,
+        ])
+
+        loss_mask = np.zeros(len(tokens), dtype=np.int64)
+        loss_mask[-num_labels:] = 1
+
+        dec_mask = make_history_mask(tokens)
+        if is_prefix_lm(model_type):
+            dec_mask[:-num_labels, :-num_labels] = 1
+
+        train_sample = {
+            'text': tokens,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'truncated': int(truncated),
+            'dec_mask': dec_mask,
+        }
+    else:
+        # Padding.
+        tokens_enc, tokens_dec_in, labels, enc_mask, \
+        dec_mask, enc_dec_mask, loss_mask \
+            = pad_and_convert_to_numpy(tokens, masked_positions,
+                                       masked_labels, pad_id, max_seq_length,
+                                       max_seq_length_dec, masked_spans,
+                                       bos_id, eos_id, sentinel_tokens)
+
+        train_sample = {
+            'text_enc': tokens_enc,
+            'text_dec': tokens_dec_in,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'truncated': int(truncated),
+            'enc_mask': enc_mask,
+            'dec_mask': dec_mask,
+            'enc_dec_mask': enc_dec_mask,
+        }
     return train_sample

megatron/model/enums.py

@@ -18,3 +18,8 @@ class AttnMaskType(enum.Enum):
     padding = 1
     causal = 2
     prefix = 3
+
+class UL2ModelType(enum.Enum):
+    encoder_decoder = 'ED'
+    non_causal_decoder = 'ND'
+    causal_decoder = 'CD'

pretrain_ul2.py

@@ -13,39 +13,71 @@
 )
 from megatron.core import tensor_parallel
 from megatron.data.dataset_utils import build_train_valid_test_datasets
-from megatron.model import T5Model, ModelType
+from megatron.data.ul2_dataset import (
+    is_decoder_only as _is_decoder_only,
+    is_prefix_lm as _is_prefix_lm,
+)
+from megatron.model import GPTModel, ModelType, T5Model
+from megatron.model.t5_model import t5_position_ids
 from megatron.training import pretrain
 from megatron.utils import average_losses_across_data_parallel_group
 
 
 """
-Pipeline parallelism for UL2 with T5
-====================================
+Pipeline parallelism for UL2
+============================
 
-Since UL2 re-uses the T5 model architecture, please see its
+Since UL2 re-uses the T5 model architecture for encoder-decoder models
+and the GPT model architecture for decoder-only models, please see their
 documentation for more information.
 """
 
 
+def is_decoder_only():
+    """Return whether we use a decoder-only model."""
+    args = get_args()
+    return _is_decoder_only(args.ul2_model_type)
+
+
+def is_prefix_lm():
+    """Return whether we use a non-causal decoder-only model."""
+    args = get_args()
+    return _is_prefix_lm(args.ul2_model_type)
+
+
 def model_provider(pre_process=True, post_process=True,
                    add_encoder=True, add_decoder=True):
     """Build the model."""
 
     print_rank_0('building UL2 model ...')
-    model = T5Model(num_tokentypes=0,
-                    parallel_output=True,
-                    pre_process=pre_process,
-                    post_process=post_process,
-                    add_encoder=add_encoder,
-                    add_decoder=add_decoder)
+    if is_decoder_only():
+        print_rank_0('Using decoder-only UL2 model.')
+        model = GPTModel(
+            num_tokentypes=0,
+            parallel_output=True,
+            pre_process=pre_process,
+            post_process=post_process,
+            prefix_lm=True
+        )
+    else:
+        print_rank_0('Using encoder-decoder UL2 model.')
+        model = T5Model(num_tokentypes=0,
+                        parallel_output=True,
+                        pre_process=pre_process,
+                        post_process=post_process,
+                        add_encoder=add_encoder,
+                        add_decoder=add_decoder)
     return model
 
 
 def get_batch(data_iterator):
     """Build the batch."""
 
-    keys = ['text_enc', 'text_dec', 'labels', 'loss_mask',
-            'enc_mask', 'dec_mask', 'enc_dec_mask']
+    if is_decoder_only():
+        keys = ['text', 'labels', 'loss_mask', 'dec_mask']
+    else:
+        keys = ['text_enc', 'text_dec', 'labels', 'loss_mask',
+                'enc_mask', 'dec_mask', 'enc_dec_mask']
     datatype = torch.int64
 
     # Broadcast data.
@@ -56,17 +88,25 @@ def get_batch(data_iterator):
     data_b = tensor_parallel.broadcast_data(keys, data, datatype)
 
     # Unpack.
-    tokens_enc = data_b['text_enc'].long()
-    tokens_dec = data_b['text_dec'].long()
-    labels = data_b['labels'].long()
-    loss_mask = data_b['loss_mask'].float()
+    if is_decoder_only():
+        tokens = data_b['text'].long()
+        labels = data_b['labels'].long()
+        loss_mask = data_b['loss_mask'].float()
+
+        dec_mask = (data_b['dec_mask'] < 0.5)
+        return tokens, loss_mask, labels, dec_mask
+    else:
+        tokens_enc = data_b['text_enc'].long()
+        tokens_dec = data_b['text_dec'].long()
+        labels = data_b['labels'].long()
+        loss_mask = data_b['loss_mask'].float()
 
-    enc_mask = (data_b['enc_mask'] < 0.5)
-    dec_mask = (data_b['dec_mask'] < 0.5)
-    enc_dec_mask = (data_b['enc_dec_mask'] < 0.5)
+        enc_mask = (data_b['enc_mask'] < 0.5)
+        dec_mask = (data_b['dec_mask'] < 0.5)
+        enc_dec_mask = (data_b['enc_dec_mask'] < 0.5)
 
-    return tokens_enc, tokens_dec, loss_mask, labels, \
-           enc_mask, dec_mask, enc_dec_mask
+        return tokens_enc, tokens_dec, loss_mask, labels, \
+               enc_mask, dec_mask, enc_dec_mask
 
 
 def loss_func(loss_mask, output_tensor):
@@ -87,18 +127,28 @@ def forward_step(data_iterator, model):
 
     # Get the batch.
     timers('batch generator', log_level=2).start()
-    tokens_enc, tokens_dec, loss_mask, lm_labels, enc_mask, dec_mask, enc_dec_mask \
-        = get_batch(data_iterator)
+    if is_decoder_only():
+        (tokens, loss_mask, lm_labels, dec_mask) = get_batch(data_iterator)
+    else:
+        (
+            tokens_enc, tokens_dec, loss_mask, lm_labels,
+            enc_mask, dec_mask, enc_dec_mask,
+        ) = get_batch(data_iterator)
     timers('batch generator').stop()
 
     # Forward model lm_labels
-    output_tensor = model(tokens_enc,
-                          tokens_dec,
-                          enc_mask,
-                          dec_mask,
-                          enc_dec_mask,
-                          tokentype_ids=None,
-                          lm_labels=lm_labels)
+    if is_decoder_only():
+        position_ids = t5_position_ids(tokens)
+        output_tensor = model(tokens, position_ids, dec_mask,
+                              labels=lm_labels)
+    else:
+        output_tensor = model(tokens_enc,
+                              tokens_dec,
+                              enc_mask,
+                              dec_mask,
+                              enc_dec_mask,
+                              tokentype_ids=None,
+                              lm_labels=lm_labels)
 
     return output_tensor, partial(loss_func, loss_mask)