feat(train): split artifact into model/state (#128)

src/dalle_mini/text.py

@@ -116,7 +116,7 @@ def remove_comma_numbers(t):
 
 
 def pre_process_dot_numbers(t):
-    return re.sub("(\w)\.(\w)", fr"\1{temp_token}dot{temp_token}\2", t)
+    return re.sub("(\w)\.(\w)", rf"\1{temp_token}dot{temp_token}\2", t)
 
 
 def post_process_dot_numbers(t):
@@ -126,7 +126,7 @@ def post_process_dot_numbers(t):
 def pre_process_quotes(t):
     # allows quotes only for 's, 't, 'd, 'm, 'll, 're, 've
     return re.sub(
-        r"'(?=([stdm]|(ll)|(re)|(ve)|(ll))\b)", fr"{temp_token}quote{temp_token}", t
+        r"'(?=([stdm]|(ll)|(re)|(ve)|(ll))\b)", rf"{temp_token}quote{temp_token}", t
     )
 
 
@@ -135,7 +135,7 @@ def post_process_quotes(t):
 
 
 def pre_process_dates(t):
-    return re.sub("(\d)/(\d)", fr"\1{temp_token}slash{temp_token}\2", t)
+    return re.sub("(\d)/(\d)", rf"\1{temp_token}slash{temp_token}\2", t)
 
 
 def post_process_dates(t):

tools/train/train.py

@@ -88,6 +88,23 @@ class ModelArguments:
             "help": "Floating-point format in which the computations will be performed (not the model weights). Choose one of `[float32, float16, bfloat16]`."
         },
     )
+    restore_state: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "Restore optimizer and training state associated with a wandb checkpoint."
+        },
+    )
+
+    state_artifact: str = field(init=False)
+
+    def __post_init__(self):
+        if self.restore_state:
+            assert (
+                "/model-" in self.model_name_or_path
+            ), "Restoring state only available with W&B artifact reference"
+            self.state_artifact = self.model_name_or_path.replace(
+                "/model-", "/state-", 1
+            )
 
 
 @dataclass
@@ -319,11 +336,6 @@ class TrainingArguments:
         },
     )
 
-    resume_from_checkpoint: Optional[str] = field(
-        default=None,
-        metadata={"help": "Reference to a wandb artifact for resuming training."},
-    )
-
     wandb_entity: Optional[str] = field(
         default=None,
         metadata={"help": "The wandb entity to use (for teams)."},
@@ -349,6 +361,8 @@ class TrainingArguments:
         },
     )
 
+    dp_devices: int = field(init=False)
+
     def __post_init__(self):
         assert self.optim in [
             "distributed_shampoo",
@@ -470,62 +484,40 @@ def main():
             config=parser.parse_args(),
         )
 
-    if training_args.resume_from_checkpoint is not None:
-        if jax.process_index() == 0:
-            artifact = wandb.run.use_artifact(training_args.resume_from_checkpoint)
-        else:
-            artifact = wandb.Api().artifact(training_args.resume_from_checkpoint)
-        artifact_dir = artifact.download()
+    # Set up our new model config
+    if model_args.config_name:
+        config = DalleBartConfig.from_pretrained(model_args.config_name)
+    else:
+        config = None
 
-        # load model
+    # Load or create new model
+    if model_args.model_name_or_path:
         model = DalleBart.from_pretrained(
-            artifact_dir,
+            model_args.model_name_or_path,
+            config=config,
+            seed=training_args.seed_model,
             dtype=getattr(jnp, model_args.dtype),
             abstract_init=True,
             load_on_cpu=True,
         )
+    else:
+        model = DalleBart(
+            config,
+            seed=training_args.seed_model,
+            dtype=getattr(jnp, model_args.dtype),
+            load_on_cpu=True,
+        )
 
-        # load tokenizer
+    # Load tokenizer
+    if model_args.tokenizer_name is not None:
         tokenizer = DalleBartTokenizer.from_pretrained(
-            artifact_dir,
-            use_fast=True,
+            model_args.tokenizer_name, use_fast=True
         )
-
     else:
-        # Set up our new model config
-        if model_args.config_name:
-            config = DalleBartConfig.from_pretrained(model_args.config_name)
-        else:
-            config = None
-
-        # Load or create new model
-        if model_args.model_name_or_path:
-            model = DalleBart.from_pretrained(
-                model_args.model_name_or_path,
-                config=config,
-                seed=training_args.seed_model,
-                dtype=getattr(jnp, model_args.dtype),
-                abstract_init=True,
-                load_on_cpu=True,
-            )
-        else:
-            model = DalleBart(
-                config,
-                seed=training_args.seed_model,
-                dtype=getattr(jnp, model_args.dtype),
-                load_on_cpu=True,
-            )
-
-        # Load tokenizer
-        if model_args.tokenizer_name is not None:
-            tokenizer = DalleBartTokenizer.from_pretrained(
-                model_args.tokenizer_name, use_fast=True
-            )
-        else:
-            tokenizer = DalleBartTokenizer.from_pretrained(
-                model_args.model_name_or_path,
-                use_fast=True,
-            )
+        tokenizer = DalleBartTokenizer.from_pretrained(
+            model_args.model_name_or_path,
+            use_fast=True,
+        )
 
     # get PartitionSpec for model params (required to be a dict)
     param_spec = set_partitions(model.params)
@@ -655,30 +647,29 @@ def main():
 
     # get PartitionSpec for optimizer state
     def get_opt_state_spec_and_shape(param_spec):
-        if training_args.optim in ["adam", "adafactor"]:
-            # get opt_state shape without actual init
-            opt_state_shape = jax.eval_shape(optimizer.init, model.params)
-
-            if training_args.optim == "adam":
-
-                def _opt_state_spec_per_leaf(x):
-                    if isinstance(x, FrozenDict):
-                        # variables with same structure as params
-                        return param_spec
-                    else:
-                        # other variables such as count
-                        return None
-
-                opt_state_spec = jax.tree_map(
-                    _opt_state_spec_per_leaf,
-                    opt_state_shape,
-                    # return None spec for empty elements
-                    is_leaf=lambda x: isinstance(x, (FrozenDict, optax.EmptyState)),
-                )
+        # get opt_state shape without actual init
+        opt_state_shape = jax.eval_shape(optimizer.init, model.params)
+
+        if training_args.optim == "adam":
+
+            def _opt_state_spec_per_leaf(x):
+                if isinstance(x, FrozenDict):
+                    # variables with same structure as params
+                    return param_spec
+                else:
+                    # other variables such as count
+                    return None
+
+            opt_state_spec = jax.tree_map(
+                _opt_state_spec_per_leaf,
+                opt_state_shape,
+                # return None spec for empty elements
+                is_leaf=lambda x: isinstance(x, (FrozenDict, optax.EmptyState)),
+            )
 
-            elif training_args.optim == "adafactor":
-                # factorized state must be replicated (rank different than params)
-                opt_state_spec = None
+        elif training_args.optim == "adafactor":
+            # factorized state must be replicated (rank different than params)
+            opt_state_spec = None
 
         elif training_args.optim == "distributed_shampoo":
             opt_state_spec = opt_fn.pspec_fn(
@@ -686,7 +677,6 @@ def main():
                 params_partition_spec=param_spec,
                 partition_spec_for_statistics=PartitionSpec(None, "batch", None),
             )
-            opt_state_shape = opt_fn.shape_and_dtype_fn(model.params)
         else:
             raise NotImplementedError
         return opt_state_spec, opt_state_shape
@@ -698,7 +688,7 @@ def main():
     devices = np.asarray(jax.devices()).reshape(*mesh_shape)
     mesh = maps.Mesh(devices, ("batch", "mp"))
 
-    # Create state spec
+    # define state spec
     state_spec = TrainState(
         params=param_spec,
         opt_state=opt_state_spec,
@@ -713,7 +703,7 @@ def main():
 
     # create training state
     with maps.mesh(mesh.devices, mesh.axis_names):
-        if training_args.resume_from_checkpoint is None:
+        if not model_args.restore_state:
 
             def init_state(params):
                 return TrainState.create(
@@ -731,6 +721,13 @@ def main():
             )(model.params)
 
         else:
+            # get state files from artifact
+            if jax.process_index() == 0:
+                artifact = wandb.run.use_artifact(model_args.state_artifact)
+            else:
+                artifact = wandb.Api().artifact(model_args.state_artifact)
+            artifact_dir = artifact.download()
+
             # restore opt_state
             with (Path(artifact_dir) / "opt_state.msgpack").open("rb") as f:
                 opt_state = from_bytes(opt_state_shape, f.read())
@@ -760,7 +757,7 @@ def main():
             del opt_state
 
     # free memory
-    del model._params
+    del model._params, opt_state_spec, opt_state_shape
 
     # define batch specs
     keys = ["attention_mask", "decoder_input_ids", "input_ids", "labels"]
@@ -998,51 +995,46 @@ def main():
                     f,
                 )
 
-            if jax.process_index() == 0:
-                # save to W&B
-                if training_args.log_model:
-                    # save some space
-                    c = wandb.wandb_sdk.wandb_artifacts.get_artifacts_cache()
-                    c.cleanup(wandb.util.from_human_size("10GB"))
-
-                    metadata = dict(state_dict)
-                    metadata["num_params"] = num_params
-                    if eval_metrics is not None:
-                        metadata["eval"] = eval_metrics
-                    artifact = wandb.Artifact(
-                        name=f"model-{wandb.run.id}",
-                        type="bart_model",
-                        metadata=metadata,
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "flax_model.msgpack")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "config.json")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "tokenizer.json")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "tokenizer_config.json")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "vocab.json")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "merges.txt")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "special_tokens_map.json")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "opt_state.msgpack")
-                    )
-                    artifact.add_file(
-                        str(Path(training_args.output_dir) / "training_state.json")
+            # save to W&B
+            if training_args.log_model:
+                # save some space
+                c = wandb.wandb_sdk.wandb_artifacts.get_artifacts_cache()
+                c.cleanup(wandb.util.from_human_size("10GB"))
+
+                metadata = dict(state_dict)
+                metadata["num_params"] = num_params
+                if eval_metrics is not None:
+                    metadata["eval"] = eval_metrics
+
+                # create model artifact
+                artifact = wandb.Artifact(
+                    name=f"model-{wandb.run.id}",
+                    type="DalleBart_model",
+                    metadata=metadata,
+                )
+                for filename in [
+                    "config.json",
+                    "flax_model.msgpack",
+                    "merges.txt",
+                    "special_tokens_map.json",
+                    "tokenizer.json",
+                    "tokenizer_config.json",
+                    "vocab.json",
+                ]:
+                    artifact.add_file(f"{Path(training_args.output_dir) / filename}")
+                wandb.run.log_artifact(artifact)
+
+                # create state artifact
+                artifact_state = wandb.Artifact(
+                    name=f"state-{wandb.run.id}",
+                    type="DalleBart_state",
+                    metadata=metadata,
+                )
+                for filename in ["opt_state.msgpack", "training_state.json"]:
+                    artifact_state.add_file(
+                        f"{Path(training_args.output_dir) / filename}"
                     )
-
-                    wandb.run.log_artifact(artifact)
+                wandb.run.log_artifact(artifact_state)
 
     # init variables
     last_time = time.perf_counter()