add HuggingFace streaming support in data input pipeline

Author: aireenmei

requirements.txt

@@ -1,5 +1,6 @@
 jax>=0.4.30
 jaxlib>=0.4.30
+grain-nightly
 google-cloud-storage==2.17.0
 absl-py
 datasets

src/maxdiffusion/configs/base14.yml

@@ -125,11 +125,20 @@ ici_tensor_parallelism: 1
 # Dataset
 # Replace with dataset path or train_data_dir. One has to be set.
 dataset_name: 'diffusers/pokemon-gpt4-captions'
-# saves transformed dataset of dataset_name.
+dataset_type: 'tf'
+cache_latents_text_encoder_outputs: True
+# cache_latents_text_encoder_outputs only apply to dataset_type="tf",
+# only apply to small dataset that fits in memory
+# prepare image latents and text encoder outputs
+# Reduce memory consumption and reduce step time during training
+# transformed dataset is saved at dataset_save_location
 dataset_save_location: '/tmp/pokemon-gpt4-captions_sd15'
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
+hf_data_dir: ''
+hf_train_files: ''
+hf_access_token: ''
 image_column: 'image'
 caption_column: 'text'
 resolution: 512
@@ -145,11 +154,6 @@ enable_data_shuffling: True
 # checkpoint every number of samples, -1 means don't checkpoint.
 checkpoint_every: -1
 
-# Prepare image latents and text encoder outputs
-# during dataset creation to reduce memory consumption.
-cache_latents_text_encoder_outputs: True
-
-
 # Training loop
 learning_rate: 1.e-7
 scale_lr: False
@@ -205,4 +209,4 @@ class_prompt: ''
 prior_loss_weight: 1.0
 num_class_images: 100
 # If true, set dataset_save_location.
-cache_dreambooth_dataset: False
+cache_dreambooth_dataset: False

src/maxdiffusion/configs/base21.yml

@@ -127,11 +127,20 @@ ici_tensor_parallelism: 1
 # Dataset
 # Replace with dataset path or train_data_dir. One has to be set.
 dataset_name: 'diffusers/pokemon-gpt4-captions'
-# saves transformed dataset of dataset_name.
+dataset_type: 'tf'
+cache_latents_text_encoder_outputs: True
+# cache_latents_text_encoder_outputs only apply to dataset_type="tf",
+# only apply to small dataset that fits in memory
+# prepare image latents and text encoder outputs
+# Reduce memory consumption and reduce step time during training
+# transformed dataset is saved at dataset_save_location
 dataset_save_location: '/tmp/pokemon-gpt4-captions_sd21'
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
+hf_data_dir: ''
+hf_train_files: ''
+hf_access_token: ''
 image_column: 'image'
 caption_column: 'text'
 resolution: 768
@@ -201,4 +210,4 @@ class_prompt: ''
 prior_loss_weight: 1.0
 num_class_images: 100
 # If true, set dataset_save_location.
-cache_dreambooth_dataset: False
+cache_dreambooth_dataset: False

src/maxdiffusion/configs/base_2_base.yml

@@ -140,11 +140,20 @@ ici_tensor_parallelism: 1
 # Dataset
 # Replace with dataset path or train_data_dir. One has to be set.
 dataset_name: 'diffusers/pokemon-gpt4-captions'
-# saves transformed dataset of dataset_name.
+dataset_type: 'tf'
+cache_latents_text_encoder_outputs: True
+# cache_latents_text_encoder_outputs only apply to dataset_type="tf",
+# only apply to small dataset that fits in memory
+# prepare image latents and text encoder outputs
+# Reduce memory consumption and reduce step time during training
+# transformed dataset is saved at dataset_save_location
 dataset_save_location: '/tmp/pokemon-gpt4-captions'
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
+hf_data_dir: ''
+hf_train_files: ''
+hf_access_token: ''
 image_column: 'image'
 caption_column: 'text'
 resolution: 512
@@ -160,11 +169,6 @@ enable_data_shuffling: True
 # checkpoint every number of samples, -1 means don't checkpoint.
 checkpoint_every: -1
 
-# Prepare image latents and text encoder outputs
-# during dataset creation to reduce memory consumption.
-cache_latents_text_encoder_outputs: True
-
-
 # Training loop
 learning_rate: 1.e-7
 scale_lr: False
@@ -218,4 +222,4 @@ class_prompt: ''
 prior_loss_weight: 1.0
 num_class_images: 100
 # If true, set dataset_save_location.
-cache_dreambooth_dataset: False
+cache_dreambooth_dataset: False

src/maxdiffusion/configs/base_xl.yml

@@ -128,11 +128,20 @@ ici_tensor_parallelism: 1
 # Dataset
 # Replace with dataset path or train_data_dir. One has to be set.
 dataset_name: 'diffusers/pokemon-gpt4-captions'
-# saves transformed dataset of dataset_name.
+dataset_type: 'tf'
+cache_latents_text_encoder_outputs: True
+# cache_latents_text_encoder_outputs only apply to dataset_type="tf",
+# only apply to small dataset that fits in memory
+# prepare image latents and text encoder outputs
+# Reduce memory consumption and reduce step time during training
+# transformed dataset is saved at dataset_save_location
 dataset_save_location: '/tmp/pokemon-gpt4-captions_xl'
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
+hf_data_dir: ''
+hf_train_files: ''
+hf_access_token: ''
 image_column: 'image'
 caption_column: 'text'
 resolution: 1024
@@ -148,10 +157,6 @@ enable_data_shuffling: True
 # checkpoint every number of samples, -1 means don't checkpoint.
 checkpoint_every: -1
 
-# Prepare image latents and text encoder outputs
-# during dataset creation to reduce memory consumption.
-cache_latents_text_encoder_outputs: True
-
 # Training loop
 learning_rate: 4.e-7
 scale_lr: False
@@ -204,4 +209,4 @@ enable_mllog: False
 controlnet_model_name_or_path: 'diffusers/controlnet-canny-sdxl-1.0'
 controlnet_from_pt: True
 controlnet_conditioning_scale: 0.5
-controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png'
+controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png'

src/maxdiffusion/input_pipeline/_hf_data_processing.py

@@ -0,0 +1,142 @@
+"""
+ Copyright 2024 Google LLC
+
+ 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
+
+      https://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.
+ """
+
+import warnings
+import datasets
+from datasets import load_dataset
+from datasets.distributed import split_dataset_by_node
+import grain.python as grain
+
+from maxdiffusion import max_logging
+from maxdiffusion import multihost_dataloading
+
+
+def make_hf_streaming_iterator(
+    config,
+    split,
+    dataloading_host_index,
+    dataloading_host_count,
+    mesh,
+    global_batch_size,
+    tokenize_fn=None,
+    image_transforms_fn=None,
+):
+  """Streaming data from HF Hub or GCS buckect.
+  No download regardless of config.cache_latents_text_encoder_outputs"""
+  ds = load_dataset(
+      config.dataset_name,
+      split=split,
+      data_dir=config.hf_data_dir,
+      data_files=config.hf_train_files,
+      streaming=True,
+      token=config.hf_access_token,
+  )
+
+  ds = ds.shuffle(seed=config.seed)
+
+  if tokenize_fn:
+    ds = ds.map(
+        function=tokenize_fn,
+        batched=True,
+        remove_columns=[config.caption_column],
+    )
+
+  if image_transforms_fn:
+    ds = ds.map(
+        function=image_transforms_fn,
+        batched=True,
+        remove_columns=[config.image_column],
+    )
+
+  ds = HFDataSource(
+      ds,
+      dataloading_host_index,
+      dataloading_host_count,
+  )
+  dummy_index_sampler = grain.IndexSampler(
+      num_records=len(ds),
+      num_epochs=1,
+      shard_options=grain.ShardOptions(
+          shard_index=dataloading_host_index, shard_count=dataloading_host_count, drop_remainder=False
+      ),
+      shuffle=False,
+      seed=0,
+  )
+  operations = [grain.Batch(batch_size=global_batch_size // dataloading_host_count, drop_remainder=True)]
+  dataloader = grain.DataLoader(
+      data_source=ds,
+      operations=operations,
+      sampler=dummy_index_sampler,
+      worker_count=1,  # only supports one worker for now, more workers results in duplicated data
+      worker_buffer_size=1,
+      read_options=grain.ReadOptions(num_threads=1, prefetch_buffer_size=64),
+  )
+  train_iter = multihost_dataloading.MultiHostDataLoadIterator(dataloader, mesh)
+  return train_iter
+
+
+class HFDataSource(grain.RandomAccessDataSource):
+  """A class that makes HuggingFace IterableDataset a grain datasource without random access support"""
+
+  def __init__(
+      self,
+      dataset: datasets.IterableDataset,
+      dataloading_host_index: int,
+      dataloading_host_count: int,
+  ):
+    self.dataset = dataset
+    self.dataloading_host_count = dataloading_host_count
+    self.dataloading_host_index = dataloading_host_index
+    self.n_shards = dataset.n_shards
+    self._check_shard_count()
+    self.current_shard = dataloading_host_index
+    self.dataset_shard = split_dataset_by_node(dataset, world_size=self.n_shards, rank=self.current_shard)
+    self.data_iter = None
+
+  def _check_shard_count(self):
+    if self.n_shards < self.dataloading_host_count:
+      warnings.warn(
+          f"WARNING: Inefficient dataloading. Your train or eval dataset contains {self.n_shards} shards, "
+          "smaller than number of host loading data. This is known to lead to inefficient dataloading. "
+          "see https://github.com/google/maxtext/blob/main/getting_started/Data_Input_Pipeline.md#multihost-dataloading-best-practice"
+      )
+      self.n_shards = self.dataloading_host_count
+
+  def _update_shard(self):
+    new_shard = (self.current_shard + self.dataloading_host_count) % self.n_shards
+    max_logging.log(f"Updating host {self.dataloading_host_index} dataset, was on shard {self.current_shard}")
+    max_logging.log(f"New shard is {new_shard}")
+    self.current_shard = new_shard
+    self.dataset_shard = split_dataset_by_node(self.dataset, world_size=self.n_shards, rank=self.current_shard)
+    self.data_iter = iter(self.dataset_shard)
+
+  def __len__(self):
+    """Return length of the HF dataset. Since HuggingFace IterableDataset does not have length,
+    a fake length bigger than the dataset is returned"""
+    return 10_000_000_000
+
+  def __getitem__(self, index):
+    """Since HuggingFace IterableDataset does not support random access by index.
+    The next item in the iterator is returned."""
+    if not self.data_iter:
+      self.data_iter = iter(self.dataset_shard)
+
+    while True:
+      try:
+        data = next(self.data_iter)
+        return data
+      except StopIteration:
+        self._update_shard()