Merge pull request #118 from borisdayma/feat-optim

src/dalle_mini/data.py

@@ -161,7 +161,7 @@ class Dataset:
         ):
             """
             Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
-            Shuffle batches if `shuffle` is `True`.
+            Shuffle batches if rng is set.
             """
             steps_per_epoch = len(dataset) // batch_size
 
@@ -182,19 +182,20 @@ class Dataset:
                 yield batch
 
         def _dataloader_datasets_streaming(
-            dataset: Dataset, batch_size: int, epoch: int
+            dataset: Dataset, split: str, batch_size: int, epoch: int
         ):
-            # epoch is only use for multi-host
             keys = ["input_ids", "attention_mask", "labels", "decoder_input_ids"]
             batch = {k: [] for k in keys}
-            first_loop = True
-            while self.multi_hosts or first_loop:
+            first_loop = True  # stop after one loop in some cases
+            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 we don't know how much data is on each host
-                if not first_loop:
-                    # multi-host setting, we reshuffle shards
-                    epoch += 1
+                # 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
+                if epoch is not None:
+                    # reshuffle training data at each epoch (not applicable with validation set)
                     dataset.set_epoch(epoch)
+                    epoch += 1
                 for item in dataset:
                     for k, v in item.items():
                         batch[k].append(v)
@@ -213,9 +214,7 @@ class Dataset:
             raise ValueError(f'split must be "train" or "eval", got {split}')
 
         if self.streaming:
-            if split == "train":
-                ds.set_epoch(epoch)
-            return _dataloader_datasets_streaming(ds, batch_size, epoch)
+            return _dataloader_datasets_streaming(ds, split, batch_size, epoch)
         else:
             if split == "train":
                 self.rng_dataset, input_rng = jax.random.split(self.rng_dataset)

tools/inference/inference_pipeline.ipynb

@@ -70,15 +70,15 @@
     "# Model references\n",
     "\n",
     "# dalle-mini\n",
-    "DALLE_MODEL = 'dalle-mini/dalle-mini/model-3bqwu04f:latest'  # can be wandb artifact or 🤗 Hub or local folder\n",
+    "DALLE_MODEL = \"dalle-mini/dalle-mini/model-3bqwu04f:latest\"  # can be wandb artifact or 🤗 Hub or local folder\n",
     "DALLE_COMMIT_ID = None  # used only with 🤗 hub\n",
     "\n",
     "# VQGAN model\n",
-    "VQGAN_REPO = 'dalle-mini/vqgan_imagenet_f16_16384'\n",
-    "VQGAN_COMMIT_ID = 'e93a26e7707683d349bf5d5c41c5b0ef69b677a9'\n",
+    "VQGAN_REPO = \"dalle-mini/vqgan_imagenet_f16_16384\"\n",
+    "VQGAN_COMMIT_ID = \"e93a26e7707683d349bf5d5c41c5b0ef69b677a9\"\n",
     "\n",
     "# CLIP model\n",
-    "CLIP_REPO = 'openai/clip-vit-base-patch16'\n",
+    "CLIP_REPO = \"openai/clip-vit-base-patch16\"\n",
     "CLIP_COMMIT_ID = None"
    ]
   },
@@ -121,18 +121,28 @@
     "import wandb\n",
     "\n",
     "# Load dalle-mini\n",
-    "if ':' in DALLE_MODEL:\n",
+    "if \":\" in DALLE_MODEL:\n",
     "    # wandb artifact\n",
     "    artifact = wandb.Api().artifact(DALLE_MODEL)\n",
     "    # we only download required files (no need for opt_state which is large)\n",
-    "    model_files = ['config.json', 'flax_model.msgpack', 'merges.txt', 'special_tokens_map.json', 'tokenizer.json', 'tokenizer_config.json', 'vocab.json']\n",
+    "    model_files = [\n",
+    "        \"config.json\",\n",
+    "        \"flax_model.msgpack\",\n",
+    "        \"merges.txt\",\n",
+    "        \"special_tokens_map.json\",\n",
+    "        \"tokenizer.json\",\n",
+    "        \"tokenizer_config.json\",\n",
+    "        \"vocab.json\",\n",
+    "    ]\n",
     "    for f in model_files:\n",
-    "        artifact.get_path(f).download('model')\n",
-    "    model = DalleBart.from_pretrained('model', dtype=dtype, abstract_init=True)\n",
-    "    tokenizer = AutoTokenizer.from_pretrained('model')\n",
+    "        artifact.get_path(f).download(\"model\")\n",
+    "    model = DalleBart.from_pretrained(\"model\", dtype=dtype, abstract_init=True)\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(\"model\")\n",
     "else:\n",
     "    # local folder or 🤗 Hub\n",
-    "    model = DalleBart.from_pretrained(DALLE_MODEL, revision=DALLE_COMMIT_ID, dtype=dtype, abstract_init=True)\n",
+    "    model = DalleBart.from_pretrained(\n",
+    "        DALLE_MODEL, revision=DALLE_COMMIT_ID, dtype=dtype, abstract_init=True\n",
+    "    )\n",
     "    tokenizer = AutoTokenizer.from_pretrained(DALLE_MODEL, revision=DALLE_COMMIT_ID)\n",
     "\n",
     "# Load VQGAN\n",
@@ -191,7 +201,7 @@
     "from functools import partial\n",
     "\n",
     "# model inference\n",
-    "@partial(jax.pmap, axis_name=\"batch\", static_broadcasted_argnums=(3,4))\n",
+    "@partial(jax.pmap, axis_name=\"batch\", static_broadcasted_argnums=(3, 4))\n",
     "def p_generate(tokenized_prompt, key, params, top_k, top_p):\n",
     "    return model.generate(\n",
     "        **tokenized_prompt,\n",
@@ -203,11 +213,13 @@
     "        top_p=top_p\n",
     "    )\n",
     "\n",
+    "\n",
     "# decode images\n",
     "@partial(jax.pmap, axis_name=\"batch\")\n",
     "def p_decode(indices, params):\n",
     "    return vqgan.decode_code(indices, params=params)\n",
     "\n",
+    "\n",
     "# score images\n",
     "@partial(jax.pmap, axis_name=\"batch\")\n",
     "def p_clip(inputs, params):\n",
@@ -235,7 +247,7 @@
     "import random\n",
     "\n",
     "# create a random key\n",
-    "seed = random.randint(0, 2**32-1)\n",
+    "seed = random.randint(0, 2 ** 32 - 1)\n",
     "key = jax.random.PRNGKey(seed)"
    ]
   },
@@ -287,7 +299,7 @@
    },
    "outputs": [],
    "source": [
-    "prompt = 'a red T-shirt'"
+    "prompt = \"a red T-shirt\""
    ]
   },
   {
@@ -323,7 +335,13 @@
     "repeated_prompts = [processed_prompt] * jax.device_count()\n",
     "\n",
     "# tokenize\n",
-    "tokenized_prompt = tokenizer(repeated_prompts, return_tensors='jax', padding='max_length', truncation=True, max_length=128).data\n",
+    "tokenized_prompt = tokenizer(\n",
+    "    repeated_prompts,\n",
+    "    return_tensors=\"jax\",\n",
+    "    padding=\"max_length\",\n",
+    "    truncation=True,\n",
+    "    max_length=128,\n",
+    ").data\n",
     "tokenized_prompt"
    ]
   },
@@ -408,12 +426,14 @@
     "    # get a new key\n",
     "    key, subkey = jax.random.split(key)\n",
     "    # generate images\n",
-    "    encoded_images = p_generate(tokenized_prompt, shard_prng_key(subkey), model_params, gen_top_k, gen_top_p)\n",
+    "    encoded_images = p_generate(\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 = decoded_images.clip(0., 1.).reshape((-1, 256, 256, 3))\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)))"
    ]
@@ -436,7 +456,14 @@
    "outputs": [],
    "source": [
     "# get clip scores\n",
-    "clip_inputs = processor(text=[prompt] * jax.device_count(), images=images, return_tensors='np', padding='max_length', max_length=77, truncation=True).data\n",
+    "clip_inputs = processor(\n",
+    "    text=[prompt] * jax.device_count(),\n",
+    "    images=images,\n",
+    "    return_tensors=\"np\",\n",
+    "    padding=\"max_length\",\n",
+    "    max_length=77,\n",
+    "    truncation=True,\n",
+    ").data\n",
     "logits = p_clip(shard(clip_inputs), clip_params)\n",
     "logits = logits.squeeze().flatten()"
    ]
@@ -458,10 +485,10 @@
    },
    "outputs": [],
    "source": [
-    "print(f'Prompt: {prompt}\\n')\n",
+    "print(f\"Prompt: {prompt}\\n\")\n",
     "for idx in logits.argsort()[::-1]:\n",
     "    display(images[idx])\n",
-    "    print(f'Score: {logits[idx]:.2f}\\n')"
+    "    print(f\"Score: {logits[idx]:.2f}\\n\")"
    ]
   }
  ],

tools/train/train.py

@@ -65,7 +65,7 @@ class ModelArguments:
     config_name: Optional[str] = field(
         default=None,
         metadata={
-            "help": "Pretrained config name or path if not the same as model_name"
+            "help": "Pretrained config name or path if not the same as model_name_or_path"
         },
     )
     tokenizer_name: Optional[str] = field(
@@ -77,7 +77,7 @@ class ModelArguments:
     dtype: Optional[str] = field(
         default="float32",
         metadata={
-            "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
+            "help": "Floating-point format in which the computations will be performed (not the model weights). Choose one of `[float32, float16, bfloat16]`."
         },
     )
 
@@ -106,11 +106,15 @@ class DataTrainingArguments:
     )
     train_file: Optional[str] = field(
         default=None,
-        metadata={"help": "The input training data file (glob acceptable)."},
+        metadata={
+            "help": "The input training data file (glob & braceexpand acceptable)."
+        },
     )
     validation_file: Optional[str] = field(
         default=None,
-        metadata={"help": "An optional input evaluation data file (glob acceptable)."},
+        metadata={
+            "help": "An optional input evaluation data file (glob & braceexpand acceptable)."
+        },
     )
     # data loading should not be a bottleneck so we use "streaming" mode by default
     streaming: Optional[bool] = field(
@@ -132,15 +136,13 @@ class DataTrainingArguments:
     max_train_samples: Optional[int] = field(
         default=None,
         metadata={
-            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
-            "value if set."
+            "help": "For debugging purposes or quicker training, truncate the number of training examples."
         },
     )
     max_eval_samples: Optional[int] = field(
         default=None,
         metadata={
-            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
-            "value if set."
+            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples."
         },
     )
     preprocessing_num_workers: Optional[int] = field(
@@ -191,42 +193,40 @@ class TrainingArguments:
 
     do_train: bool = field(default=False, metadata={"help": "Whether to run training."})
     do_eval: bool = field(
-        default=False, metadata={"help": "Whether to run eval on the dev set."}
+        default=False, metadata={"help": "Whether to run eval on the validation set."}
     )
 
     per_device_train_batch_size: int = field(
-        default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for training."}
+        default=8, metadata={"help": "Batch size per GPU/TPU/CPU for training."}
     )
     per_device_eval_batch_size: int = field(
-        default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."}
+        default=8, metadata={"help": "Batch size per GPU/TPU/CPU for evaluation."}
     )
 
     gradient_accumulation_steps: int = field(
         default=1,
         metadata={
-            "help": "Number of updates steps to accumulate before performing a backward/update pass."
+            "help": "Number of updates steps to accumulate before performing an update pass."
         },
     )
 
     learning_rate: float = field(
         default=5e-5, metadata={"help": "The initial learning rate."}
     )
-    adafactor: bool = field(
-        default=False,
-        metadata={"help": "Use Adafactor instead of AdamW."},
-    )
-    distributed_shampoo: bool = field(
-        default=False,
-        metadata={"help": "Use Distributed Shampoo optimizer instead of AdamW."},
-    )
-    weight_decay: float = field(
-        default=None, metadata={"help": "Weight decay if we apply some."}
+    optim: str = field(
+        default="distributed_shampoo",
+        metadata={
+            "help": 'The optimizer to use. Can be "distributed_shampoo" (default), "adam" or "adafactor"'
+        },
     )
-    adam_beta1: float = field(
-        default=0.9, metadata={"help": "Beta1 for AdamW optimizer"}
+    weight_decay: float = field(default=None, metadata={"help": "Weight decay."})
+    beta1: float = field(
+        default=0.9,
+        metadata={"help": "Beta1 for Adam & Distributed Shampoo."},
     )
-    adam_beta2: float = field(
-        default=0.999, metadata={"help": "Beta2 for AdamW optimizer"}
+    beta2: float = field(
+        default=0.999,
+        metadata={"help": "Beta2 for for Adam & Distributed Shampoo."},
     )
     adam_epsilon: float = field(
         default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."}
@@ -234,9 +234,47 @@ class TrainingArguments:
     max_grad_norm: float = field(
         default=1.0, metadata={"help": "Max gradient norm for Adafactor."}
     )
-    use_decay: bool = field(
+    block_size: int = field(
+        default=1024,
+        metadata={"help": "Chunked size for large layers with Distributed Shampoo."},
+    )
+    preconditioning_compute_steps: int = field(
+        default=10, metadata={"help": "Number of steps to update preconditioner."}
+    )
+    skip_preconditioning_dim_size_gt: int = field(
+        default=4096,
+        metadata={"help": "Max size for preconditioning with Distributed Shampoo."},
+    )
+    optim_quantized: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to quantize optimizer (only supported with Distributed Shampoo)."
+        },
+    )
+
+    lr_decay: str = field(
+        default=None,
+        metadata={
+            "help": "Decay to be used in the learning rate scheduler. Can be None (default), linear or exponential."
+        },
+    )
+    lr_transition_steps: int = field(
+        default=None,
+        metadata={
+            "help": "Number of transition steps associated with learning rate decay when using exponential decay."
+        },
+    )
+    lr_decay_rate: float = field(
+        default=None,
+        metadata={
+            "help": "Decay rate associated with learning rate when using exponential decay."
+        },
+    )
+    lr_staircase: bool = field(
         default=False,
-        metadata={"help": "Whether to use decay in the learning rate scheduler."},
+        metadata={
+            "help": "Whether to use staircase or continuous learning rate when using exponential decay."
+        },
     )
 
     num_train_epochs: float = field(
@@ -267,18 +305,18 @@ class TrainingArguments:
         },
     )
 
-    push_to_hub: bool = field(
-        default=False,
-        metadata={
-            "help": "Whether or not to upload the trained model to the model hub after training."
-        },
-    )
-
     resume_from_checkpoint: Optional[str] = field(
         default=None,
         metadata={"help": "Reference to a wandb artifact for resuming training."},
     )
 
+    def __post_init__(self):
+        assert self.optim in [
+            "distributed_shampoo",
+            "adam",
+            "adafactor",
+        ], f"Selected optimizer not supported: {self.optim}"
+
 
 class TrainState(train_state.TrainState):
     dropout_rng: jnp.ndarray = None
@@ -309,33 +347,6 @@ class TrainState(train_state.TrainState):
         )
 
 
-def create_learning_rate_fn(
-    num_warmup_steps: int,
-    learning_rate: float,
-    use_decay: bool,
-    num_train_steps: int = None,  # used only with `use_decay`, typically train_size // batch_size * num_epochs
-) -> Callable[[int], jnp.array]:
-    """Returns a linear warmup, linear_decay learning rate function."""
-    if use_decay:
-        assert (
-            num_train_steps is not None
-        ), "Learning rate with decay requires number of training steps"
-    warmup_fn = optax.linear_schedule(
-        init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps
-    )
-    if not use_decay:
-        return warmup_fn
-    decay_fn = optax.linear_schedule(
-        init_value=learning_rate,
-        end_value=0,
-        transition_steps=num_train_steps - num_warmup_steps,
-    )
-    schedule_fn = optax.join_schedules(
-        schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps]
-    )
-    return schedule_fn
-
-
 class MetricsLogger:
     def __init__(self, state):
         self.step = state.step
@@ -529,12 +540,37 @@ def main():
     num_params = model.num_params
 
     # Create learning rate schedule
-    learning_rate_fn = create_learning_rate_fn(
-        training_args.warmup_steps,
-        training_args.learning_rate,
-        training_args.use_decay,
-        num_train_steps,
-    )
+    def create_learning_rate_fn() -> Callable[[int], jnp.array]:
+        """Create the learning rate function."""
+        warmup_fn = optax.linear_schedule(
+            init_value=0.0,
+            end_value=training_args.learning_rate,
+            transition_steps=training_args.warmup_steps,
+        )
+        if training_args.lr_decay is None:
+            return warmup_fn
+        elif training_args.lr_decay == "linear":
+            assert (
+                num_train_steps is not None
+            ), "linear decay requires knowing the dataset length"
+            decay_fn = optax.linear_schedule(
+                init_value=training_args.learning_rate,
+                end_value=0,
+                transition_steps=num_train_steps - training_args.warmup_steps,
+            )
+        elif training_args.lr_decay == "exponential":
+            decay_fn = optax.exponential_decay(
+                init_value=training_args.learning_rate,
+                transition_steps=training_args.lr_transition_steps,
+                decay_rate=training_args.lr_decay_rate,
+                staircase=training_args.lr_staircase,
+            )
+        schedule_fn = optax.join_schedules(
+            schedules=[warmup_fn, decay_fn], boundaries=[training_args.warmup_steps]
+        )
+        return schedule_fn
+
+    learning_rate_fn = create_learning_rate_fn()
 
     # We use Optax's "masking" functionality to not apply weight decay
     # to bias and LayerNorm scale parameters. decay_mask_fn returns a
@@ -558,29 +594,22 @@ def main():
         return traverse_util.unflatten_dict(flat_mask)
 
     # create adam optimizer
-    if training_args.adafactor:
-        # We use the default parameters here to initialize adafactor,
-        # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
-        optimizer = optax.adafactor(
-            learning_rate=learning_rate_fn,
-            weight_decay_rate=training_args.weight_decay,
-            weight_decay_mask=decay_mask_fn,
-            clipping_threshold=training_args.max_grad_norm,
-        )
-    elif training_args.distributed_shampoo:
+    if training_args.optim == "distributed_shampoo":
         # parameters from https://github.com/tensorflow/lingvo/blob/03ee9d7cd50764b0424c7c863733c91fc0b053ec/lingvo/jax/optimizers.py#L729
         # Notes:
-        # - mask for weight decay is not implemented but we don't use it anyway
+        # - mask for weight decay is not implemented
         optimizer = distributed_shampoo(
             learning_rate_fn,
-            block_size=1024,  # recommended default for large LM is 1536
-            beta1=0.9,
-            beta2=0.999,
+            block_size=training_args.block_size,
+            beta1=training_args.beta1,
+            beta2=training_args.beta2,
             diagonal_epsilon=1e-10,
             matrix_epsilon=1e-8,
-            weight_decay=0.0,
-            start_preconditioning_step=1001,
-            preconditioning_compute_steps=10,
+            weight_decay=training_args.weight_decay
+            if training_args.weight_decay is not None
+            else 0.0,
+            start_preconditioning_step=training_args.warmup_steps,
+            preconditioning_compute_steps=training_args.preconditioning_compute_steps,
             statistics_compute_steps=1,
             best_effort_shape_interpretation=True,
             graft_type=GraftingType.RMSPROP_NORMALIZED,
@@ -589,23 +618,32 @@ def main():
             batch_axis_name="batch",
             inverse_failure_threshold=0.1,
             moving_average_for_momentum=True,
-            skip_preconditioning_dim_size_gt=4096,
+            skip_preconditioning_dim_size_gt=training_args.skip_preconditioning_dim_size_gt,
             clip_by_scaled_gradient_norm=None,
             precision=jax.lax.Precision.HIGHEST,
-            best_effort_memory_usage_reduction=False,
+            best_effort_memory_usage_reduction=training_args.optim_quantized,
         )
 
-    else:
+    elif training_args.optim == "adam":
         optimizer = optax.adamw(
             learning_rate=learning_rate_fn,
-            b1=training_args.adam_beta1,
-            b2=training_args.adam_beta2,
+            b1=training_args.beta1,
+            b2=training_args.beta2,
             eps=training_args.adam_epsilon,
             weight_decay=training_args.weight_decay
             if training_args.weight_decay is not None
             else 0.0,
             mask=decay_mask_fn,
         )
+    elif training_args.optim == "adafactor":
+        # We use the default parameters here to initialize adafactor,
+        # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
+        optimizer = optax.adafactor(
+            learning_rate=learning_rate_fn,
+            weight_decay_rate=training_args.weight_decay,
+            weight_decay_mask=decay_mask_fn,
+            clipping_threshold=training_args.max_grad_norm,
+        )
 
     # add gradient accumulation
     if training_args.gradient_accumulation_steps > 1:
@@ -821,16 +859,6 @@ def main():
 
                     wandb.run.log_artifact(artifact)
 
-                # save to the hub
-                if training_args.push_to_hub:
-                    model.save_pretrained(
-                        training_args.output_dir,
-                        params=params,
-                        push_to_hub=training_args.push_to_hub,
-                        commit_message=f"Saving weights and logs at step {unreplicate(state.step)+1}",
-                        temp_dir=True,  # avoid issues with being in a repository
-                    )
-
     # init variables
     last_time = time.perf_counter()
     train_metrics = None
@@ -841,7 +869,7 @@ def main():
         metrics_logger.log({"train/epoch": epoch}, step=unreplicate(state.step))
 
         # Generate an epoch by shuffling sampling indices from the train dataset
-        train_loader = dataset.dataloader("train", train_batch_size)
+        train_loader = dataset.dataloader("train", train_batch_size, epoch)
         # train
         for batch in tqdm(
             train_loader,