Have datasets account for tokenizers which incorrectly define PAD

megatron/core/datasets/bert_dataset.py

@@ -139,18 +139,14 @@ def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]:
         assert length_pads >= 0
 
         tokens = numpy.array(tokens, dtype=numpy.int64)
-        tokens = numpy.pad(tokens, (0, length_pads), constant_values=self.config.tokenizer.pad)
+        tokens = numpy.pad(tokens, (0, length_pads), constant_values=self._pad_token_id)
 
         assignments = numpy.array(assignments, dtype=numpy.int64)
-        assignments = numpy.pad(
-            assignments, (0, length_pads), constant_values=self.config.tokenizer.pad
-        )
+        assignments = numpy.pad(assignments, (0, length_pads), constant_values=self._pad_token_id)
 
         # Get the padding mask
-        mask_pads = numpy.ones(length_toks, dtype=numpy.int64)
-        mask_pads = numpy.pad(
-            mask_pads, (0, length_pads), constant_values=self.config.tokenizer.pad
-        )
+        mask_pads = numpy.ones(self.config.sequence_length, dtype=numpy.int64)
+        mask_pads[tokens == self._pad_token_id] = self._pad_token_id
 
         # Mask the labels
         labels = numpy.zeros(self.config.sequence_length, dtype=numpy.int64) - 1
@@ -160,6 +156,10 @@ def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]:
         mask_loss = numpy.zeros(self.config.sequence_length, dtype=numpy.int64)
         mask_loss[masked_positions] = 1
 
+        # For padded sequences, ensure the embedding layer can map the token ID
+        tokens[tokens == self._pad_token_id] = 0
+        labels[labels == self._pad_token_id] = 0
+
         return {
             "text": tokens,
             "types": assignments,

megatron/core/datasets/blended_megatron_dataset_config.py

@@ -77,6 +77,17 @@ class BlendedMegatronDatasetConfig:
        datasets(s).
     """
 
+    allow_ambiguous_pad_tokens: Optional[bool] = False
+    """Whether to prevent pad tokens already present in the dataset from being masked out
+       when the pad token incorrectly shares the same id with other special tokens.
+       Treating such tokens as pad tokens results in training instability and divergence.
+       Such a scenario is best resolved by fixing the tokenizer, but leaving this option as False
+       provides a workaround.
+       This argument will have no effect if the tokenizer is correct. However, should the user
+       desire to train on a dataset that intentionally contains pad tokens - while also using an
+       incorrect tokenizer - this option may be set to True. This is typically not recommended.
+    """
+
     def __post_init__(self) -> None:
         """Do asserts and set fields post init"""
         if self.blend_per_split is not None and any(self.blend_per_split):

megatron/core/datasets/gpt_dataset.py

@@ -20,9 +20,6 @@
 logger = logging.getLogger(__name__)
 
 
-_PAD_TOKEN_ID = -1
-
-
 @dataclass
 class GPTDatasetConfig(BlendedMegatronDatasetConfig):
     """Configuration object for Megatron Core GPT datasets"""
@@ -105,11 +102,6 @@ def __init__(
         self.cached_loss_mask = None
         self.cached_position_ids = None
 
-        try:
-            self._pad_token_id = self.config.tokenizer.pad
-        except Exception:
-            self._pad_token_id = _PAD_TOKEN_ID
-
         (self.document_index, self.sample_index, self.shuffle_index) = (
             self._build_document_sample_shuffle_indices()
         )

megatron/core/datasets/megatron_dataset.py

@@ -2,6 +2,7 @@
 
 import hashlib
 import json
+import warnings
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from typing import Dict, Iterable, List, Optional, Union
@@ -16,6 +17,9 @@
 LowLevelDataset = Union[IndexedDataset, Iterable]
 
 
+_PAD_TOKEN_ID = -1
+
+
 class MegatronDataset(ABC, torch.utils.data.Dataset):
     """The highest level wrapper class from which all dataset classes should inherit
 
@@ -66,6 +70,49 @@ def __init__(
             self.unique_description.encode("utf-8"), usedforsecurity=False
         ).hexdigest()
 
+        # Handle pad token id provided by the tokenizer
+        try:
+            self._pad_token_id = self.config.tokenizer.pad
+        except Exception:
+            self._pad_token_id = _PAD_TOKEN_ID
+
+        # Check if pad token id collides with any other special tokens
+        try:
+            _special_tokens_list = [
+                v for k, v in self.config.tokenizer.special_tokens_dict.items() if k != "pad_token"
+            ]
+        except (AttributeError, IndexError, ValueError):
+            _special_tokens_list = []
+        # If the tokenizer does not have a special_tokens_dict attribute, at least check eos and eod
+        if not _special_tokens_list:
+            try:
+                _special_tokens_list.append(self.config.tokenizer.eos)
+            except (AttributeError, NotImplementedError):
+                pass
+            try:
+                _special_tokens_list.append(self.config.tokenizer.eod)
+            except (AttributeError, NotImplementedError):
+                pass
+
+        if self._pad_token_id in _special_tokens_list:
+            if self.config.allow_ambiguous_pad_tokens:
+                # This will break training, but users must explicitly opt-in to this behavior.
+                warnings.warn(
+                    "The pad token id in the tokenizer collides with another special token id. "
+                    "This may cause instability and lack of covergence during training. "
+                    "Do not ignore this warning if you do not understand the implications. "
+                )
+            else:
+                # Reset the pad token id to a value which is guaranteed not to be in the dataset.
+                self._pad_token_id = _PAD_TOKEN_ID
+                warnings.warn(
+                    "The pad token id in the tokenizer collides with another special token id. "
+                    "This may cause instability and lack of covergence during training. "
+                    "As such, the training flow will avoid masking out any pad tokens already "
+                    "present in the dataset. If you would like to disable this behavior, "
+                    "please provide a tokenizer with a uniquely-defined pad token id."
+                )
+
     @staticmethod
     def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int:
         """Return the number of elements in the underlying low level dataset for the purpose of

megatron/core/datasets/t5_dataset.py

@@ -286,17 +286,19 @@ def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]:
 
         encoder_input = numpy.array(encoder_input, dtype=numpy.int64)
         encoder_input = numpy.pad(
-            encoder_input, (0, length_pads_encoder), constant_values=self.config.tokenizer.pad
+            encoder_input, (0, length_pads_encoder), constant_values=self._pad_token_id
         )
 
         decoder_input = numpy.array(decoder_input, dtype=numpy.int64)
         decoder_input = numpy.pad(
-            decoder_input, (0, length_pads_decoder), constant_values=self.config.tokenizer.pad
+            decoder_input, (0, length_pads_decoder), constant_values=self._pad_token_id
         )
 
         # Create attention and history masks
-        mask_encoder = numpy.array([1] * length_toks_encoder + [0] * length_pads_encoder)
-        mask_decoder = numpy.array([1] * length_toks_decoder + [0] * length_pads_decoder)
+        mask_encoder = numpy.ones(self.config.sequence_length_encoder, dtype=numpy.int64)
+        mask_encoder[encoder_input == self._pad_token_id] = 0
+        mask_decoder = numpy.ones(self.config.sequence_length_decoder, dtype=numpy.int64)
+        mask_decoder[decoder_input == self._pad_token_id] = 0
         mask_encoder_decoder = None
 
         # Mask the labels
@@ -307,6 +309,11 @@ def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]:
         loss_mask = numpy.zeros(self.config.sequence_length_decoder, dtype=numpy.int64)
         loss_mask[:length_toks_decoder] = 1
 
+        # For padded sequences, ensure the embedding layer can map the token ID
+        encoder_input[encoder_input == self._pad_token_id] = 0
+        decoder_input[decoder_input == self._pad_token_id] = 0
+        labels[labels == self._pad_token_id] = 0
+
         return {
             "text_enc": encoder_input,
             "text_dec": decoder_input,

megatron/training/arguments.py

@@ -2894,6 +2894,20 @@ def _add_data_args(parser):
                        help='Path to cache index files when using s3 or msc dataloader')
     group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005,
                        help='The sample surplus to build for the mid-level datasets(s)')
+    group.add_argument('--allow-ambiguous-pad-tokens', action='store_true',
+                       help='Whether to prevent pad tokens already present in the dataset '
+                       'from being masked out when the pad token incorrectly shares the same id '
+                       'with other special tokens in the tokenizer. Note that this argument has '
+                       'no effect when the tokenizer correctly provides a unique id for the pad. '
+                       'Masking out such ambiguous pad tokens results in training instability. '
+                       'Such a scenario is best resolved by fixing the tokenizer; leaving this '
+                       'option as False provides a workaround. '
+                       'When left to the default of False, any token ids that collide with the '
+                       'pad token id - as provided by the tokenizer - will not be masked out of '
+                       'the loss calculation: it cannot be determined whether they are truly pad. '
+                       'If instead this argument is set, the training flow will treat all tokens '
+                       'that share the same id as the pad token as true pad tokens, potentially '
+                       'causing severe training instability.')
     return parser
 
 

pretrain_bert.py

@@ -172,6 +172,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
         masking_use_geometric_distribution=False,
         classification_head=args.bert_binary_head,
         mid_level_dataset_surplus=args.mid_level_dataset_surplus,
+        allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
     )
 
     print_rank_0('> building train, validation, and test datasets '

pretrain_gpt.py

@@ -190,6 +190,7 @@ def core_gpt_dataset_config_from_args(args):
         create_attention_mask=args.create_attention_mask_in_dataloader,
         object_storage_cache_path=args.object_storage_cache_path,
         mid_level_dataset_surplus=args.mid_level_dataset_surplus,
+        allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
     )
 
 

pretrain_mamba.py

@@ -177,6 +177,7 @@ def core_gpt_dataset_config_from_args(args):
         create_attention_mask=args.create_attention_mask_in_dataloader,
         object_storage_cache_path=args.object_storage_cache_path,
         mid_level_dataset_surplus=args.mid_level_dataset_surplus,
+        allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
     )
 
 

pretrain_retro.py

@@ -210,6 +210,7 @@ def train_valid_test_datasets_provider(train_valid_test_num_samples):
         reset_attention_mask=args.reset_attention_mask,
         eod_mask_loss=args.eod_mask_loss,
         mid_level_dataset_surplus=args.mid_level_dataset_surplus,
+        allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
     )
 
     # GPT datasets.

pretrain_t5.py

@@ -233,6 +233,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples: int):
         masking_use_longer_ngrams=False,
         masking_use_geometric_distribution=True,
         mid_level_dataset_surplus=args.mid_level_dataset_surplus,
+        allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
     )
 
     print_rank_0('> building train, validation, and test datasets for T5 ...')

pretrain_vlm.py

@@ -224,6 +224,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
         image_w=args.img_w,
         preprocess_func=_preprocess_data_for_llava,
         mid_level_dataset_surplus=args.mid_level_dataset_surplus,
+        allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
     )
 
     print_rank_0("> building train, validation, and test datasets for multimodal ...")