pretrain_dynamic_dataloader.py

# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
#
# 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.

"""Dataset loader for the pre-training with dynamic sequence length."""
from typing import Optional, Tuple

import dataclasses
import tensorflow as tf, tf_keras

from official.core import config_definitions as cfg
from official.core import input_reader
from official.nlp.data import data_loader_factory
from official.nlp.data import pretrain_dataloader


@dataclasses.dataclass
class BertPretrainDataConfig(cfg.DataConfig):
  """Data config for BERT pretraining task (tasks/masked_lm)."""
  input_path: str = ''
  global_batch_size: int = 512
  is_training: bool = True
  seq_bucket_lengths: Tuple[int, ...] = (128, 256, 384, 512,)
  # TODO(rxsang): `seq_bucket_window_scale` is only useful when round robin
  # tf.data service is disabled. Deprecate this flag once we always enable round
  # robin tf.data service.
  seq_bucket_window_scale: int = 8
  use_next_sentence_label: bool = True
  use_position_id: bool = False
  deterministic: bool = False
  enable_tf_data_service: bool = False
  enable_round_robin_tf_data_service: bool = False
  tf_data_service_job_name: str = 'bert_pretrain'
  use_v2_feature_names: bool = False


@data_loader_factory.register_data_loader_cls(BertPretrainDataConfig)
class PretrainingDynamicDataLoader(pretrain_dataloader.BertPretrainDataLoader):
  """Dataset loader for bert-style pretraining with dynamic sequenece length.

  Bucketizes the input id features by the seq_bucket_lengths and features are
  padded to the bucket boundaries. The mask features are usually short than
  input id features and can also be dynamic. We require the mask feature lengths
  within a bucket must be the same. For example, with [128, 256] buckets,
  the mask features for bucket 128 should always have the length as X and
  features for bucket 256 should always have the length as Y.

  The dataloader does not filter out empty masks. Make sure to handle this
  in the model.
  """

  def __init__(self, params):
    self._params = params
    if len(params.seq_bucket_lengths) < 1:
      raise ValueError('The seq_bucket_lengths cannot be empty.')
    self._seq_bucket_lengths = params.seq_bucket_lengths
    self._seq_bucket_window_scale = params.seq_bucket_window_scale
    self._global_batch_size = params.global_batch_size
    self._use_next_sentence_label = params.use_next_sentence_label
    self._use_position_id = params.use_position_id
    self._drop_remainder = params.drop_remainder
    self._enable_tf_data_service = params.enable_tf_data_service
    self._enable_round_robin_tf_data_service = (
        params.enable_round_robin_tf_data_service)
    self._mask_keys = [
        'masked_lm_positions', 'masked_lm_ids', 'masked_lm_weights'
    ]

  def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_mask': tf.io.VarLenFeature(tf.int64),
        'masked_lm_positions': tf.io.VarLenFeature(tf.int64),
        'masked_lm_ids': tf.io.VarLenFeature(tf.int64),
        'masked_lm_weights': tf.io.VarLenFeature(tf.float32),
    }
    if self._params.use_v2_feature_names:
      input_ids_key = 'input_word_ids'
      segment_key = 'input_type_ids'
      name_to_features.update({
          input_ids_key: tf.io.VarLenFeature(tf.int64),
          segment_key: tf.io.VarLenFeature(tf.int64),
      })
    else:
      input_ids_key = 'input_ids'
      segment_key = 'segment_ids'
      name_to_features.update({
          input_ids_key: tf.io.VarLenFeature(tf.int64),
          segment_key: tf.io.VarLenFeature(tf.int64),
      })
    if self._use_next_sentence_label:
      name_to_features['next_sentence_labels'] = tf.io.FixedLenFeature([1],
                                                                       tf.int64)
    dynamic_keys = [input_ids_key, 'input_mask', segment_key]
    if self._use_position_id:
      name_to_features['position_ids'] = tf.io.VarLenFeature(tf.int64)
      dynamic_keys.append('position_ids')

    example = tf.io.parse_single_example(record, name_to_features)
    for key in dynamic_keys + self._mask_keys:
      example[key] = tf.sparse.to_dense(example[key])

    # Truncate padded data after the first non pad in the
    # sequence length dimension.
    # Pad before the first non pad from the back should not be removed.
    mask = tf.math.greater(
        tf.math.cumsum(example[input_ids_key], reverse=True), 0)
    for key in dynamic_keys:
      example[key] = tf.boolean_mask(example[key], mask)

    # masked_lm_ids should be 0 padded.
    # Change mask features to -1 padding so that we can differentiate
    # padding from data or from bucketizing.
    mask = tf.math.not_equal(example['masked_lm_ids'], 0)
    example['masked_lm_ids'] = tf.where(
        mask, example['masked_lm_ids'],
        -tf.ones(tf.shape(example['masked_lm_ids']), dtype=example[key].dtype))

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    # tf.data service uses dataset graph fingerprint to distinguish input
    # pipeline jobs, thus we sort the keys here to make sure they are generated
    # in a deterministic order each time the dataset function is traced.
    for name in sorted(list(example.keys())):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def _bucketize_and_batch(
      self,
      dataset,
      input_context: Optional[tf.distribute.InputContext] = None):
    """Bucketize by sequence length and batch the datasets."""
    per_replica_batch_size = input_context.get_per_replica_batch_size(
        self._global_batch_size) if input_context else self._global_batch_size

    def element_length_func(example, seq_len_dim):
      return tf.shape(example['input_word_ids'])[seq_len_dim]

    bucket_boundaries = [length + 1 for length in self._seq_bucket_lengths]
    bucket_batch_sizes = [per_replica_batch_size] * (len(bucket_boundaries) + 1)

    # Bucketize and batch the dataset with per replica batch size first.
    dataset = dataset.apply(
        tf.data.experimental.bucket_by_sequence_length(
            lambda example: tf.cast(element_length_func(example, 0), tf.int32),
            bucket_boundaries,
            bucket_batch_sizes,
            pad_to_bucket_boundary=True,
            drop_remainder=self._drop_remainder))
    if input_context:
      window_size = input_context.num_replicas_in_sync
      if self._enable_tf_data_service and (
          not self._enable_round_robin_tf_data_service):
        # If tf.data service is enabled but round-robin behavior is not enabled,
        # different TPU workers may fetch data from one tf.data service worker
        # in different speed. We set the window size to be
        # `seq_bucket_window_scale` larger to leave buffer if some workers are
        # fetching data faster than others, so all the data within the same
        # global batch can still have more chances to be in the same bucket.
        window_size *= self._seq_bucket_window_scale

      # Group `num_replicas_in_sync` batches from same bucket together, so all
      # replicas can get the same sequence length for one global step.
      dataset = dataset.apply(
          tf.data.experimental.group_by_window(
              key_func=lambda example: tf.cast(  # pylint: disable=g-long-lambda
                  element_length_func(example, 1), tf.int64),
              reduce_func=lambda _, x: tf.data.Dataset.from_tensors(x),
              window_size=window_size))
      dataset = dataset.flat_map(lambda x: x)

    def _remove_pads_from_bucketize(features):
      # All mask features must have the same effective length.
      # The real masked ids padding token is -1 and 0 comes from
      # bucket_by_sequence_length.
      mask = tf.math.not_equal(features['masked_lm_ids'], 0)

      mask_per_example = tf.math.reduce_sum(tf.cast(mask, tf.int32), axis=1)
      normalized = tf.cast(
          mask_per_example / tf.math.reduce_max(mask_per_example), tf.int32)
      assert_op = tf.debugging.assert_equal(
          tf.math.reduce_sum(normalized), per_replica_batch_size,
          'Number of non padded mask tokens is not the same for each example '
          'in the same sequence length.')
      with tf.control_dependencies([assert_op]):
        for key in self._mask_keys:
          features[key] = tf.reshape(
              tf.boolean_mask(
                  features[key], mask), [per_replica_batch_size, -1])
      # Revert masked_lm_ids to be 0-padded.
      mask = tf.math.not_equal(features['masked_lm_ids'], -1)
      features['masked_lm_ids'] = tf.where(
          mask, features['masked_lm_ids'],
          tf.zeros(
              tf.shape(features['masked_lm_ids']),
              dtype=features['masked_lm_ids'].dtype))
      return features

    dataset = dataset.map(_remove_pads_from_bucketize)
    return dataset

  def load(self, input_context: Optional[tf.distribute.InputContext] = None):
    """Returns a tf.dataset.Dataset."""
    reader = input_reader.InputReader(
        params=self._params,
        decoder_fn=self._decode,
        parser_fn=self._parse,
        transform_and_batch_fn=self._bucketize_and_batch)
    return reader.read(input_context)