Upload training metadata (#294)

* first pass at uploading metadata

* second pass

* reduce to 1 flag

* lint

* Update open_instruct/dpo_tune.py

Co-authored-by: Costa Huang <[email protected]>

---------

Co-authored-by: Hamish Ivison <[email protected]>
Co-authored-by: Costa Huang <[email protected]>

Dockerfile

@@ -100,7 +100,7 @@ COPY open_instruct open_instruct
 # install the package in editable mode
 COPY pyproject.toml .
 RUN pip install -e .
-COPY .git .
+COPY .git/ ./.git/
 
 COPY eval eval
 COPY configs configs

Dockerfile.olmo

@@ -103,7 +103,7 @@ COPY open_instruct open_instruct
 # install the package in editable mode
 COPY pyproject.toml .
 RUN pip install -e .
-COPY .git .
+COPY .git/ ./.git/
 
 COPY eval eval
 COPY configs configs

open_instruct/dpo_tune.py

@@ -71,6 +71,7 @@
     maybe_use_ai2_hf_entity,
     maybe_use_ai2_wandb_entity,
     submit_beaker_eval_jobs,
+    upload_metadata_to_hf,
 )
 
 logger = get_logger(__name__)
@@ -350,6 +351,8 @@ class FlatArguments:
     """The url of the saved model in the Hugging Face Hub (will be autoset)"""
     try_launch_beaker_eval_jobs: bool = True
     """Whether to launch beaker evaluation jobs after training"""
+    hf_metadata_dataset: Optional[str] = None
+    """What dataset to upload the metadata to. If unset, don't upload metadata"""
 
     def __post_init__(self):
         if self.reduce_loss not in ["mean", "sum"]:
@@ -847,6 +850,7 @@ def load_model():
             experiment_config,
             init_kwargs={"wandb": {"entity": args.wandb_entity, "tags": [args.exp_name] + get_wandb_tags()}},
         )
+        wandb_tracker = accelerator.get_tracker("wandb")
 
     # Train!
     total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
@@ -1034,6 +1038,26 @@ def load_model():
                 location=args.hf_repo_id,
                 hf_repo_revision=args.hf_repo_revision,
             )
+    if args.hf_metadata_dataset and accelerator.is_main_process and is_beaker_job():
+        # dpo script only supports these two options right now for datasets
+        dataset_name = args.dataset_name if args.dataset_name else args.train_file
+        # mainly just focussing here on what would be useful for the leaderboard.
+        # wandb will have even more useful information.
+        metadata_blob = {
+            "model_name": args.exp_name,
+            "model_type": "sft",
+            "datasets": [dataset_name],
+            "base_model": args.model_name_or_path,
+            "wandb_path": wandb_tracker.run.get_url(),
+            "beaker_experiment": beaker_config.beaker_experiment_url,
+            "beaker_datasets": beaker_config.beaker_dataset_id_urls
+        }
+        upload_metadata_to_hf(
+            metadata_blob,
+            "metadata.json",
+            args.hf_metadata_dataset,
+            'results/' + args.hf_repo_revision,  # to match what the auto-evals name as.
+        )
 
     accelerator.wait_for_everyone()
     if args.with_tracking:

open_instruct/finetune.py

@@ -62,6 +62,7 @@
     maybe_use_ai2_hf_entity,
     maybe_use_ai2_wandb_entity,
     submit_beaker_eval_jobs,
+    upload_metadata_to_hf,
 )
 
 logger = get_logger(__name__)
@@ -306,6 +307,12 @@ class FlatArguments:
         default=3,
         metadata={"help": "How many checkpoints to keep in the output directory. -1 for all."},
     )
+    fused_optimizer: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to use fused AdamW or not.",
+        },
+    )
     push_to_hub: bool = True
     """Whether to upload the saved model to huggingface"""
     hf_entity: Optional[str] = None
@@ -318,12 +325,8 @@ class FlatArguments:
     """The url of the saved model in the Hugging Face Hub (will be autoset)"""
     try_launch_beaker_eval_jobs: bool = True
     """Whether to launch beaker evaluation jobs after training"""
-    fused_optimizer: bool = field(
-        default=True,
-        metadata={
-            "help": "Whether to use fused AdamW or not.",
-        },
-    )
+    hf_metadata_dataset: Optional[str] = None
+    """What dataset to upload the metadata to. If unset, don't upload metadata"""
 
     def __post_init__(self):
         if self.reduce_loss not in ["mean", "sum"]:
@@ -852,6 +855,7 @@ def main(args: FlatArguments):
             experiment_config,
             init_kwargs={"wandb": {"entity": args.wandb_entity, "tags": [args.exp_name] + get_wandb_tags()}},
         )
+        wandb_tracker = accelerator.get_tracker("wandb")
 
     # Train!
     total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
@@ -1013,6 +1017,34 @@ def main(args: FlatArguments):
                 location=args.hf_repo_id,
                 hf_repo_revision=args.hf_repo_revision,
             )
+    if args.hf_metadata_dataset:
+        # dpo script only supports these two options right now for datasets
+        if args.dataset_mixer:
+            dataset_list = args.dataset_mixer.keys()
+        elif args.dataset_mixer_list:
+            dataset_list = args.dataset_mixer_list[::2]  # even indices
+        elif args.dataset_name:
+            dataset_list = [args.dataset_name]
+        else:
+            dataset_list = [args.train_file]
+        beaker_config = maybe_get_beaker_config()
+        # mainly just focussing here on what would be useful for the leaderboard.
+        # wandb will have even more useful information.
+        metadata_blob = {
+            "model_name": args.exp_name,
+            "model_type": "sft",
+            "datasets": dataset_list,
+            "base_model": args.model_name_or_path,
+            "wandb_path": wandb_tracker.run.get_url(),
+            "beaker_experiment": beaker_config.beaker_experiment_url,
+            "beaker_datasets": beaker_config.beaker_dataset_id_urls
+        }
+        upload_metadata_to_hf(
+            metadata_blob,
+            "metadata.json",
+            args.hf_metadata_dataset,
+            'results/' + args.hf_repo_revision,  # to match what the auto-evals name as.
+        )
 
     accelerator.wait_for_everyone()
     if args.with_tracking:

open_instruct/utils.py

@@ -625,7 +625,7 @@ def get_beaker_whoami() -> Optional[str]:
 
 
 def maybe_get_beaker_config():
-    beaker_dataset_ids = [get_beaker_dataset_ids(os.environ["BEAKER_WORKLOAD_ID"])]
+    beaker_dataset_ids = get_beaker_dataset_ids(os.environ["BEAKER_WORKLOAD_ID"])
     beaker_dataset_id_urls = [f"https://beaker.org/ds/{dataset_id}" for dataset_id in beaker_dataset_ids]
     return BeakerRuntimeConfig(
         beaker_workload_id=os.environ["BEAKER_WORKLOAD_ID"],
@@ -690,3 +690,23 @@ def submit_beaker_eval_jobs(
     logger.info(f"Beaker evaluation jobs: Stdout:\n{stdout.decode()}")
     logger.error(f"Beaker evaluation jobs: Stderr:\n{stderr.decode()}")
     logger.info(f"Beaker evaluation jobs: process return code: {process.returncode}")
+
+
+def upload_metadata_to_hf(
+        metadata_dict,
+        filename,
+        hf_dataset_name,
+        hf_dataset_save_dir,
+):
+    # upload a random dict to HF. Originally for uploading metadata to HF
+    # about a model for leaderboard displays.
+    with open("tmp.json", "w") as f:
+        json.dump(metadata_dict, f)
+    api = HfApi(token=os.getenv("HF_TOKEN", None))
+    api.upload_file(
+        path_or_fileobj="tmp.json",
+        path_in_repo=f"{hf_dataset_save_dir}/{filename}",
+        repo_id=hf_dataset_name,
+        repo_type="dataset",
+    )
+    os.remove("tmp.json")