feat: support pod (borisdayma#139)

Author: borisdayma

src/dalle_mini/data.py

@@ -27,13 +27,19 @@ class Dataset:
     do_eval: bool = True
     seed_dataset: int = None
     shard_by_host: bool = False
+    blank_caption_prob: float = 0.0
     train_dataset: Dataset = field(init=False)
     eval_dataset: Dataset = field(init=False)
     rng_dataset: jnp.ndarray = field(init=False)
     multi_hosts: bool = field(init=False)
 
     def __post_init__(self):
         self.multi_hosts = jax.process_count() > 1
+        # feed blank captions only in streaming mode for now
+        if self.blank_caption_prob:
+            assert (
+                self.streaming is True
+            ), "blank_caption_prob can only be used in streaming mode"
         # define data_files
         if self.train_file is not None or self.validation_file is not None:
             # accept braceexpand notation
@@ -101,6 +107,25 @@ def preprocess(self, tokenizer, config):
                 self.seed_dataset = np.random.get_state()[1][0]
             self.rng_dataset = jax.random.PRNGKey(self.seed_dataset)
 
+        # blank captions
+        if self.blank_caption_prob:
+            partial_blank_caption_function = partial(
+                blank_caption_function,
+                text_column=self.text_column,
+                blank_caption_prob=self.blank_caption_prob,
+            )
+            if hasattr(self, "train_dataset"):
+                self.train_dataset = (
+                    self.train_dataset.map(partial_blank_caption_function)
+                    if self.streaming
+                    else self.train_dataset.map(
+                        partial_blank_caption_function,
+                        num_proc=self.preprocessing_num_workers,
+                        load_from_cache_file=False,
+                        desc="Blanking some captions",
+                    )
+                )
+
         # normalize text
         if normalize_text:
             text_normalizer = TextNormalizer()
@@ -144,6 +169,10 @@ def preprocess(self, tokenizer, config):
                         getattr(self, ds).map(
                             partial_preprocess_function,
                             batched=True,
+                            remove_columns=[
+                                self.text_column,
+                                self.encoding_column,
+                            ],
                         )
                         if self.streaming
                         else getattr(self, ds).map(
@@ -193,8 +222,8 @@ def _dataloader_datasets_streaming(
             while (self.multi_hosts and split == "train") or first_loop:
                 # in multi-host, we run forever (no epoch) as hosts need to stop
                 # at the same time and training data may not be split equally
-                # For validation data we put the entire set on each host as we could lose
-                # too many samples on pods
+                # For validation data we put the entire batch on each host and then
+                # keep only the one specific to each host (could be improved but not necessary)
                 if epoch is not None:
                     assert split == "train"
                     # reshuffle training data at each epoch
@@ -252,6 +281,12 @@ def shift_tokens_right(input_ids: np.array, decoder_start_token_id: int):
     return shifted_input_ids
 
 
+def blank_caption_function(example, text_column, blank_caption_prob):
+    if blank_caption_prob and np.random.rand() < blank_caption_prob:
+        example[text_column] = ""
+    return example
+
+
 def normalize_function(example, text_column, text_normalizer):
     example[text_column] = text_normalizer(example[text_column])
     return example

src/dalle_mini/model/modeling.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2021 The Fairseq Authors and The Google Flax Team Authors And The HuggingFace Inc. team and the DalleBart team. All rights reserved.
+# Copyright 2021-2022 The Fairseq Authors and The Google Flax Team Authors And The HuggingFace Inc. team and & DALL·E Mini team. 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.
@@ -328,6 +328,7 @@ def __init__(
         dtype: jnp.dtype = jnp.float32,
         abstract_init: bool = False,
         load_on_cpu: bool = False,
+        init_weights: bool = True,
         **kwargs,
     ):
         module = self.module_class(config=config, dtype=dtype, **kwargs)
@@ -347,25 +348,34 @@ def __init__(
         self.key = PRNGKey(seed)
         self.dtype = dtype
 
-        # init weights on CPU
-        if load_on_cpu:
-            # init weights on CPU
-            init_fn = jax.jit(self.init_weights, static_argnums=(1,), backend="cpu")
-        else:
-            init_fn = self.init_weights
+        if init_weights:
+            # get shape of params only
+            random_params = self.init_weights(
+                self.key,
+                input_shape,
+                abstract_init=abstract_init,
+                load_on_cpu=load_on_cpu,
+            )
+
+            # save required_params as set
+            self._required_params = set(flatten_dict(unfreeze(random_params)).keys())
+            self.params = random_params
 
-        # randomly initialized parameters
-        random_params = self.init_weights(self.key, input_shape)
+    def init_weights(
+        self, rng=None, input_shape=(1, 1), abstract_init=False, load_on_cpu=False
+    ):
+        if rng is None:
+            rng = self.key
+        init_fn = super().init_weights
+        if load_on_cpu:
+            init_fn = jax.jit(init_fn, static_argnums=(1,), backend="cpu")
         if abstract_init:
             # only set shape and dtype, load parameters separately
             init_fn = partial(init_fn, input_shape=input_shape)
-            random_params = jax.eval_shape(init_fn, self.key)
+            params = jax.eval_shape(init_fn, rng)
         else:
-            random_params = init_fn(self.key, input_shape)
-
-        # save required_params as set
-        self._required_params = set(flatten_dict(unfreeze(random_params)).keys())
-        self.params = random_params
+            params = init_fn(rng, input_shape)
+        return params
 
     @property
     def num_params(self):

src/dalle_mini/model/utils.py

@@ -23,12 +23,6 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
                 else:
                     artifact = wandb.Api().artifact(pretrained_model_name_or_path)
                 pretrained_model_name_or_path = artifact.download(tmp_dir)
-                if artifact.metadata.get("bucket_path"):
-                    pretrained_model_name_or_path = artifact.metadata["bucket_path"]
-
-            if pretrained_model_name_or_path.startswith("gs://"):
-                copy_blobs(pretrained_model_name_or_path, tmp_dir)
-                pretrained_model_name_or_path = tmp_dir
 
             return super(PretrainedFromWandbMixin, cls).from_pretrained(
                 pretrained_model_name_or_path, *model_args, **kwargs

tools/inference/inference_pipeline.ipynb

@@ -83,7 +83,7 @@
     "VQGAN_COMMIT_ID = \"e93a26e7707683d349bf5d5c41c5b0ef69b677a9\"\n",
     "\n",
     "# CLIP model\n",
-    "CLIP_REPO = \"openai/clip-vit-base-patch16\"\n",
+    "CLIP_REPO = \"openai/clip-vit-large-patch14\"\n",
     "CLIP_COMMIT_ID = None"
    ]
   },
@@ -129,7 +129,6 @@
     "from dalle_mini.model import DalleBart, DalleBartTokenizer\n",
     "from vqgan_jax.modeling_flax_vqgan import VQModel\n",
     "from transformers import CLIPProcessor, FlaxCLIPModel\n",
-    "import wandb\n",
     "\n",
     "# Load dalle-mini\n",
     "model = DalleBart.from_pretrained(\n",
@@ -168,9 +167,9 @@
     "if dtype == jnp.bfloat16:\n",
     "    model.params = model.to_bf16(model.params)\n",
     "\n",
-    "model_params = replicate(model.params)\n",
-    "vqgan_params = replicate(vqgan.params)\n",
-    "clip_params = replicate(clip.params)"
+    "model._params = replicate(model.params)\n",
+    "vqgan._params = replicate(vqgan.params)\n",
+    "clip._params = replicate(clip.params)"
    ]
   },
   {
@@ -292,7 +291,7 @@
    },
    "outputs": [],
    "source": [
-    "prompt = \"a blue table\""
+    "prompt = \"view of the beach during sunset\""
    ]
   },
   {
@@ -414,12 +413,12 @@
     "    key, subkey = jax.random.split(key)\n",
     "    # generate images\n",
     "    encoded_images = p_generate(\n",
-    "        tokenized_prompt, shard_prng_key(subkey), model_params, gen_top_k, gen_top_p\n",
+    "        tokenized_prompt, shard_prng_key(subkey), model.params, gen_top_k, gen_top_p\n",
     "    )\n",
     "    # remove BOS\n",
     "    encoded_images = encoded_images.sequences[..., 1:]\n",
     "    # decode images\n",
-    "    decoded_images = p_decode(encoded_images, vqgan_params)\n",
+    "    decoded_images = p_decode(encoded_images, vqgan.params)\n",
     "    decoded_images = decoded_images.clip(0.0, 1.0).reshape((-1, 256, 256, 3))\n",
     "    for img in decoded_images:\n",
     "        images.append(Image.fromarray(np.asarray(img * 255, dtype=np.uint8)))"
@@ -453,7 +452,7 @@
     "    max_length=77,\n",
     "    truncation=True,\n",
     ").data\n",
-    "logits = p_clip(shard(clip_inputs), clip_params)\n",
+    "logits = p_clip(shard(clip_inputs), clip.params)\n",
     "logits = logits.squeeze().flatten()"
    ]
   },
@@ -479,6 +478,13 @@
     "    display(images[idx])\n",
     "    print(f\"Score: {logits[idx]:.2f}\\n\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

tools/train/config/medium/config.json

@@ -28,6 +28,5 @@
   "pad_token_id": 16385,
   "scale_embedding": false,
   "tie_word_embeddings": false,
-  "transformers_version": "4.13.0.dev0",
   "use_cache": true
 }

tools/train/config/mega/config.json

@@ -5,21 +5,20 @@
   "bos_token_id": 16385,
   "classifier_dropout": 0.0,
   "d_model": 2048,
-  "decoder_attention_heads": 16,
-  "decoder_ffn_dim": 4096,
+  "decoder_attention_heads": 32,
+  "decoder_ffn_dim": 8192,
   "decoder_layerdrop": 0.0,
-  "decoder_layers": 31,
+  "decoder_layers": 24,
   "decoder_start_token_id": 16384,
-  "dropout": 0.1,
-  "encoder_attention_heads": 16,
-  "encoder_ffn_dim": 4096,
+  "dropout": 0.0,
+  "encoder_attention_heads": 32,
+  "encoder_ffn_dim": 8192,
   "encoder_layerdrop": 0.0,
-  "encoder_layers": 31,
+  "encoder_layers": 24,
   "encoder_vocab_size": 50264,
   "eos_token_id": 16385,
-  "gradient_checkpointing": false,
   "image_length": 256,
-  "image_vocab_size": 16384,
+  "image_vocab_size": 16391,
   "init_std": 0.01,
   "is_encoder_decoder": true,
   "max_text_length": 64,
@@ -28,6 +27,5 @@
   "pad_token_id": 16385,
   "scale_embedding": false,
   "tie_word_embeddings": false,
-  "transformers_version": "4.13.0.dev0",
   "use_cache": true
 }

tools/train/config/micro/config.json

@@ -4,22 +4,21 @@
   "attention_dropout": 0.0,
   "bos_token_id": 16385,
   "classifier_dropout": 0.0,
-  "d_model": 1024,
-  "decoder_attention_heads": 16,
-  "decoder_ffn_dim": 2048,
+  "d_model": 256,
+  "decoder_attention_heads": 2,
+  "decoder_ffn_dim": 256,
   "decoder_layerdrop": 0.0,
   "decoder_layers": 2,
   "decoder_start_token_id": 16384,
   "dropout": 0.0,
-  "encoder_attention_heads": 16,
-  "encoder_ffn_dim": 2048,
+  "encoder_attention_heads": 2,
+  "encoder_ffn_dim": 256,
   "encoder_layerdrop": 0.0,
   "encoder_layers": 2,
   "encoder_vocab_size": 50264,
   "eos_token_id": 16385,
-  "gradient_checkpointing": false,
   "image_length": 256,
-  "image_vocab_size": 16384,
+  "image_vocab_size": 16391,
   "init_std": 0.02,
   "is_encoder_decoder": true,
   "max_text_length": 64,
@@ -28,6 +27,5 @@
   "pad_token_id": 16385,
   "scale_embedding": false,
   "tie_word_embeddings": false,
-  "transformers_version": "4.13.0.dev0",
   "use_cache": true
 }

tools/train/config/mini/config.json

@@ -28,6 +28,5 @@
   "pad_token_id": 16385,
   "scale_embedding": false,
   "tie_word_embeddings": false,
-  "transformers_version": "4.13.0.dev0",
   "use_cache": true
 }

tools/train/scalable_shampoo/README.md

@@ -0,0 +1,7 @@
+# Notes
+
+Files copied from [google-research/scalable_shampoo/optax](https://github.com/google-research/google-research/tree/master/scalable_shampoo/optax).
+
+Imports have been modified to be relative.
+
+This will be replaced with `optax-shampoo` package eventually.