feat: vmap optimizer (borisdayma#166)

Author: borisdayma

src/dalle_mini/model/modeling.py

@@ -946,15 +946,6 @@ def __call__(
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-        # postln is already applied in every layer
-        if self.config.use_final_ln_encoder and self.config.ln_positions != "postln":
-            hidden_states = norm(
-                self.config.ln_type,
-                dtype=self.dtype,
-                epsilon=1e-05,
-                use_scale=self.config.force_ln_scale,
-            )(hidden_states)
-
         outputs = [
             hidden_states,
             all_hidden_states,
@@ -1034,7 +1025,7 @@ def __call__(
                 self.config,
                 dtype=self.dtype,
                 add_norm=self.config.ln_positions == "postln",
-                name="FlaxBartEncoderLayers",
+                name="FlaxBartDecoderLayers",
             )(
                 hidden_states,
                 attention_mask,
@@ -1086,15 +1077,6 @@ def __call__(
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-        # postln is already applied in every layer
-        if self.config.use_final_ln_decoder and self.config.ln_positions != "postln":
-            hidden_states = norm(
-                self.config.ln_type,
-                dtype=self.dtype,
-                epsilon=1e-05,
-                use_scale=self.config.force_ln_scale,
-            )(hidden_states)
-
         outputs = [
             hidden_states,
             all_hidden_states,
@@ -1146,6 +1128,17 @@ def setup(self):
             self.config.ln_type, dtype=self.dtype, epsilon=1e-05
         )
 
+        # postln is already applied in every layer
+        if self.config.use_final_ln_encoder and self.config.ln_positions != "postln":
+            self.final_ln = norm(
+                self.config.ln_type,
+                dtype=self.dtype,
+                epsilon=1e-05,
+                use_scale=self.config.force_ln_scale,
+            )
+        else:
+            self.final_ln = None
+
     def __call__(
         self,
         input_ids,
@@ -1177,11 +1170,16 @@ def __call__(
             return_dict=return_dict,
         )
 
+        if self.final_ln is None:
+            final_output = outputs[0]
+        else:
+            final_output = self.final_ln(outputs[0])
+
         if not return_dict:
-            return outputs
+            return (final_output,) + outputs[1:]
 
         return FlaxBaseModelOutput(
-            last_hidden_state=outputs.last_hidden_state,
+            last_hidden_state=final_output,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
@@ -1223,6 +1221,15 @@ def setup(self):
             self.config.ln_type, dtype=self.dtype, epsilon=1e-05
         )
 
+        # postln is already applied in every layer
+        if self.config.use_final_ln_decoder and self.config.ln_positions != "postln":
+            self.final_ln = norm(
+                self.config.ln_type,
+                dtype=self.dtype,
+                epsilon=1e-05,
+                use_scale=self.config.force_ln_scale,
+            )
+
     def __call__(
         self,
         input_ids,
@@ -1260,11 +1267,16 @@ def __call__(
             return_dict=return_dict,
         )
 
+        if self.final_ln is None:
+            final_output = outputs[0]
+        else:
+            final_output = self.final_ln(outputs[0])
+
         if not return_dict:
-            return outputs
+            return (final_output,) + outputs[1:]
 
         return FlaxBaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=outputs.last_hidden_state,
+            last_hidden_state=final_output,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
             cross_attentions=outputs.cross_attentions,

src/dalle_mini/model/partitions.py

@@ -65,7 +65,7 @@ def set_partitions(in_dict, use_scan):
             print(f"Unmatched -> {k}")
     l = list(result.keys())
     if use_scan:
-        # add None dimension to scanned layers
+        # add None dimension to layers
         result = {
             k: (P(*(None,) + v) if v is not None else None)
             if any(x in k for x in ["FlaxBartEncoderLayers", "FlaxBartDecoderLayers"])

tools/train/config/mega/config.json

@@ -7,14 +7,14 @@
   "decoder_attention_heads": 32,
   "decoder_ffn_dim": 4096,
   "decoder_layerdrop": 0.0,
-  "decoder_layers": 26,
+  "decoder_layers": 24,
   "decoder_start_token_id": 16384,
   "do_sample": true,
   "dropout": 0.0,
   "encoder_attention_heads": 32,
   "encoder_ffn_dim": 4096,
   "encoder_layerdrop": 0.0,
-  "encoder_layers": 26,
+  "encoder_layers": 24,
   "encoder_vocab_size": 50272,
   "eos_token_id": 16385,
   "force_ln_scale": false,