feat: better multi-node support (borisdayma#158)

* reproducible data loader
* custom sharding
* model parallel across multiple nodes

Author: borisdayma

src/dalle_mini/data.py

@@ -43,6 +43,8 @@ def __post_init__(self):
         if self.seed_dataset is None:
             # create a random seed
             self.seed_dataset = random.randint(0, 2**32 - 1)
+        # set numpy rng
+        self.np_rng = np.random.default_rng(self.seed_dataset)
         self.multi_hosts = jax.process_count() > 1
         # feed blank captions only in streaming mode for now
         # otherwise dataset could be cached with same blanked captions
@@ -173,14 +175,17 @@ def preprocess(self, tokenizer, config):
                 blank_caption_function,
                 text_column=self.text_column,
                 blank_caption_prob=self.blank_caption_prob,
+                rng=self.np_rng,
             )
             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,
+                        num_proc=None
+                        if self.seed_dataset
+                        else self.preprocessing_num_workers,
                         load_from_cache_file=False,
                         desc="Blanking some captions",
                     )
@@ -316,8 +321,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:
+def blank_caption_function(example, text_column, blank_caption_prob, rng=None):
+    if (
+        blank_caption_prob
+        and (rng.random() if rng is not None else np.random.random())
+        < blank_caption_prob
+    ):
         example[text_column] = ""
     return example
 

tools/train/config/mega/config.json

@@ -1,30 +1,49 @@
 {
   "activation_dropout": 0.0,
-  "activation_function": "gelu",
+  "activation_function": "swish",
   "attention_dropout": 0.0,
   "bos_token_id": 16385,
   "d_model": 2048,
   "decoder_attention_heads": 32,
-  "decoder_ffn_dim": 8192,
+  "decoder_ffn_dim": 4096,
   "decoder_layerdrop": 0.0,
-  "decoder_layers": 24,
+  "decoder_layers": 25,
   "decoder_start_token_id": 16384,
+  "do_sample": true,
   "dropout": 0.0,
   "encoder_attention_heads": 32,
-  "encoder_ffn_dim": 8192,
+  "encoder_ffn_dim": 4096,
   "encoder_layerdrop": 0.0,
-  "encoder_layers": 24,
-  "encoder_vocab_size": 50264,
+  "encoder_layers": 25,
+  "encoder_vocab_size": 50272,
   "eos_token_id": 16385,
+  "force_ln_scale": false,
+  "gradient_checkpointing": false,
   "image_length": 256,
-  "image_vocab_size": 16391,
+  "image_vocab_size": 16415,
   "init_std": 0.01,
   "is_encoder_decoder": true,
+  "ln_positions": "normformer",
+  "ln_type": "layernorm",
+  "max_length": 257,
   "max_text_length": 64,
+  "min_length": 257,
   "model_type": "dallebart",
   "normalize_text": true,
   "pad_token_id": 16385,
   "scale_embedding": false,
+  "sinkhorn_iters": 1,
+  "tau_init": 0.05,
   "tie_word_embeddings": false,
-  "use_cache": true
+  "use_absolute_position_embeddings": true,
+  "use_alibi": false,
+  "use_bias": false,
+  "use_cache": true,
+  "use_cosine_attention": false,
+  "use_deepnet_scaling": false,
+  "use_final_ln_decoder": true,
+  "use_final_ln_encoder": true,
+  "use_glu": true,
+  "use_head_scale": false,
+  "use_swin_position_embeddings": false
 }

tools/train/config/mini/config.json

@@ -16,7 +16,7 @@
   "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,

tools/train/train.py

@@ -368,6 +368,12 @@ class TrainingArguments:
             "help": "Whether to quantize optimizer (only supported with Distributed Shampoo)."
         },
     )
+    shard_shampoo_across: str = field(
+        default="dp",
+        metadata={
+            "help": "Whether to shard the optimizer across data devices (dp), model devices (mp) or both (2d)."
+        },
+    )
 
     num_train_epochs: int = field(
         default=3, metadata={"help": "Total number of training epochs to perform."}
@@ -450,6 +456,11 @@ class TrainingArguments:
         metadata={"help": "Verify that TPU is not in use."},
     )
 
+    use_vmap_trick: bool = field(
+        default=True,
+        metadata={"help": "Verify that TPU is not in use."},
+    )
+
     mp_devices: Optional[int] = field(
         default=1,
         metadata={
@@ -500,6 +511,11 @@ def __post_init__(self):
                 f"Output directory ({self.output_dir}) already exists and is not empty."
                 "Use --overwrite_output_dir to overcome."
             )
+        assert self.shard_shampoo_across in [
+            "dp",
+            "mp",
+            "2d",
+        ], f"Shard shampoo across {self.shard_shampoo_across} not supported."
         assert (
             self.mp_devices > 0
         ), f"Number of devices for model parallelism must be > 0"
@@ -530,6 +546,12 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
+    # check arguments
+    if training_args.mp_devices > jax.local_device_count():
+        assert (
+            data_args.seed_dataset is not None
+        ), "Seed dataset must be provided when model is split over multiple hosts"
+
     # Make one log on every process with the configuration for debugging.
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -748,8 +770,20 @@ def create_learning_rate_fn() -> Callable[[int], jnp.array]:
             graft_type=graft_type,
             nesterov=False,
             exponent_override=0,
-            statistics_partition_spec=PartitionSpec(None, "dp", None),
-            preconditioner_partition_spec=PartitionSpec("dp", None, None),
+            statistics_partition_spec=PartitionSpec(
+                None, training_args.shard_shampoo_across, None
+            )
+            if training_args.shard_shampoo_across != "2d"
+            else PartitionSpec(None, "dp", "mp"),
+            preconditioner_partition_spec=PartitionSpec(
+                training_args.shard_shampoo_across, None, None
+            )
+            if training_args.shard_shampoo_across != "2d"
+            else PartitionSpec(
+                "mp" if training_args.mp_devices > training_args.dp_devices else "dp",
+                None,
+                None,
+            ),
             num_devices_for_pjit=training_args.dp_devices,
             shard_optimizer_states=True,
             inverse_failure_threshold=0.1,
@@ -917,7 +951,7 @@ def loss_fn(logits, labels):
 
     # "vmap trick" avoids a crash when mp_devices > 1 (not sure why it happens)
     # lead to better perf: see https://wandb.ai/dalle-mini/dalle-mini/reports/JAX-pmap-vs-pjit--VmlldzoxNDg1ODA2
-    use_vmap_trick = True
+    use_vmap_trick = training_args.use_vmap_trick
 
     # make grad_param_spec for vmap
     if use_vmap_trick:
@@ -1145,7 +1179,8 @@ def update_state_metrics(self, state):
                 self.log_time("train_per_log", delta_time, offset=False)
 
         def log_time(self, key, duration, offset=True):
-            wandb.log({f"time/{key}": duration, **self.state_dict})
+            if jax.process_index() == 0:
+                wandb.log({f"time/{key}": duration, **self.state_dict})
             if offset:
                 self.offset_time += duration
 
@@ -1191,7 +1226,11 @@ def run_evaluation():
         # ======================== Evaluating ==============================
         if training_args.do_eval:
             start_eval_time = time.perf_counter()
-            eval_loader = dataset.dataloader("eval", eval_batch_size_per_step)
+            eval_loader = dataset.dataloader(
+                "eval",
+                eval_batch_size_per_step
+                * max(1, training_args.mp_devices // jax.local_device_count()),
+            )
             eval_steps = (
                 len_eval_dataset // eval_batch_size_per_step
                 if len_eval_dataset is not None
@@ -1353,10 +1392,12 @@ def run_save_model(state, eval_metrics=None):
             metrics_logger.update_state_metrics(local_state)
             metrics_logger.log({})
 
-            # Generate an epoch by shuffling sampling indices from the train dataset
+            # load data - may be replicated on multiple nodes
+            node_groups = max(1, training_args.mp_devices // jax.local_device_count())
+            loader_bs = batch_size_per_node * node_groups
             train_loader = dataset.dataloader(
                 "train",
-                batch_size_per_node,
+                loader_bs,
                 epoch,
             )
             # train
@@ -1373,12 +1414,12 @@ def run_save_model(state, eval_metrics=None):
 
                 # set correct shape to batch
                 # - add grad_step dim if gradient_accumulation_steps > 1
-                # - split per dp device if not multi-host for vmap trick (does not work in multi-host)
                 bs_shape = (
-                    (batch_size_per_node_per_grad_step,)
+                    (batch_size_per_node_per_grad_step * node_groups,)
                     if not use_vmap_trick
                     else (
                         jax.local_device_count()
+                        * node_groups
                         // training_args.mp_devices,  # local dp devices
                         training_args.per_device_train_batch_size,
                     )