diff --git a/megatron/core/datasets/bert_dataset.py b/megatron/core/datasets/bert_dataset.py index 314efb46cd6..d36f618349d 100644 --- a/megatron/core/datasets/bert_dataset.py +++ b/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, diff --git a/megatron/core/datasets/blended_megatron_dataset_config.py b/megatron/core/datasets/blended_megatron_dataset_config.py index 3222ece836f..fd7132acc0f 100644 --- a/megatron/core/datasets/blended_megatron_dataset_config.py +++ b/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): diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 55270faa0e7..710a4c684ff 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/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() ) diff --git a/megatron/core/datasets/megatron_dataset.py b/megatron/core/datasets/megatron_dataset.py index 0980ef92d36..00b249de4c1 100644 --- a/megatron/core/datasets/megatron_dataset.py +++ b/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 diff --git a/megatron/core/datasets/t5_dataset.py b/megatron/core/datasets/t5_dataset.py index 85da1480e10..edcba299005 100644 --- a/megatron/core/datasets/t5_dataset.py +++ b/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, diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 5d1bc2e40a3..aa523983bd9 100644 --- a/megatron/training/arguments.py +++ b/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 diff --git a/pretrain_bert.py b/pretrain_bert.py index a5e2728db89..401c32b4cb9 100644 --- a/pretrain_bert.py +++ b/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 ' diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 7f1a5e1c6b4..6316aef03bf 100644 --- a/pretrain_gpt.py +++ b/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, ) diff --git a/pretrain_mamba.py b/pretrain_mamba.py index ae14b472df2..ba084c73478 100644 --- a/pretrain_mamba.py +++ b/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, ) diff --git a/pretrain_retro.py b/pretrain_retro.py index 100cf605657..63abbac5e39 100644 --- a/pretrain_retro.py +++ b/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. diff --git a/pretrain_t5.py b/pretrain_t5.py index 6e6d9ad2c06..e74e7d8809e 100644 --- a/pretrain_t5.py +++ b/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 ...') diff --git a/pretrain_vlm.py b/pretrain_vlm.py index ce1a5102444..524931d2727 100644 --- a/pretrain_vlm.py +++ b/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 ...")