Changes the DPO + finetune scripts to provide progress updates in the Beaker description. (#1127)

* Now, we get num_attention_heads from the hf config.

* Update code

* Added test that we match manual values

* Updated calculations

* Updated code with check_calculation

* Updated code

* Now, tests pass.

* Updated code to normalize properly

* Added some fixes

* Updated code

* Updated code

* Another fix

* Cleaned up tests.

* Cleaned up PR

* Update MFU/MBU code.

* Now, mbu tests pass.

* Moved to json file

* Added test data

* undid changes and simplified test function.

* An attempt at a fix

* Update code with patches

* now, tests pass

* Added MFU to DPO

* updated script

* uses uv for dpo

* Added a chat template to the DPO script.

* Added trackign

* Updated code to handle tracking when none

* Added description updates

* undid changes

* Check out dpo script

* updated script

* Update code to remove whitespace

* fix finetune timing

* Fixed bugs pointed out by cursor.

Author: finbarrtimbers

open_instruct/dpo_tune_cache.py

@@ -79,6 +79,7 @@
     is_beaker_job,
     launch_ai2_evals_on_weka,
     maybe_get_beaker_config,
+    maybe_update_beaker_description,
     maybe_use_ai2_hf_entity,
     maybe_use_ai2_wandb_entity,
 )
@@ -498,6 +499,7 @@ def main(args: FlatArguments, tc: TokenizerConfig):
             },
         )
         wandb_tracker = accelerator.get_tracker("wandb")
+        maybe_update_beaker_description(wandb_url=wandb_tracker.run.get_url() if args.with_tracking else None)
 
     if accelerator.is_main_process:
         pprint([args, tc])
@@ -813,6 +815,7 @@ def load_model():
     print("=============after cache logprobs; clear cache")
     print_gpu_stats(init_gpu_memory)
     # Only show the progress bar once on each machine.
+    start_time = time.perf_counter()
     progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
     # update the progress_bar if load from checkpoint
     progress_bar.update(completed_steps)
@@ -936,6 +939,12 @@ def load_model():
                     logger.info(logger_str)
                     if args.with_tracking:
                         accelerator.log(metrics_to_log, step=completed_steps)
+                    maybe_update_beaker_description(
+                        current_step=completed_steps,
+                        total_steps=args.max_train_steps,
+                        start_time=start_time,
+                        wandb_url=wandb_tracker.run.get_url() if args.with_tracking else None,
+                    )
                     # Reset the local metrics
                     local_metrics.zero_()
 
@@ -989,7 +998,7 @@ def load_model():
             path=args.output_dir,
             leaderboard_name=args.hf_repo_revision,
             oe_eval_max_length=args.oe_eval_max_length,
-            wandb_url=wandb_tracker.run.get_url(),
+            wandb_url=wandb_tracker.run.get_url() if args.with_tracking else None,
             oe_eval_tasks=args.oe_eval_tasks,
             gs_bucket_path=args.gs_bucket_path,
         )

open_instruct/finetune.py

@@ -62,6 +62,7 @@
     is_beaker_job,
     launch_ai2_evals_on_weka,
     maybe_get_beaker_config,
+    maybe_update_beaker_description,
     maybe_use_ai2_hf_entity,
     maybe_use_ai2_wandb_entity,
 )
@@ -438,6 +439,7 @@ def main(args: FlatArguments, tc: TokenizerConfig):
             },
         )
         wandb_tracker = accelerator.get_tracker("wandb")
+        maybe_update_beaker_description(wandb_url=wandb_tracker.run.get_url())
     else:
         wandb_tracker = None  # for later eval launching
 
@@ -727,7 +729,7 @@ def main(args: FlatArguments, tc: TokenizerConfig):
     local_total_tokens_this_log_period = torch.tensor(0, dtype=torch.int64, device=accelerator.device)
     local_pred_tokens_this_log_period = torch.tensor(0, dtype=torch.int64, device=accelerator.device)
     total_token_including_padding = torch.tensor(0, dtype=torch.int64, device=accelerator.device)
-    start_time = time.time()
+    start_time = time.perf_counter()
     skipped_batches = False
     for epoch in range(starting_epoch, args.num_train_epochs):
         model.train()
@@ -824,10 +826,12 @@ def main(args: FlatArguments, tc: TokenizerConfig):
                         "avg_tokens_per_batch": avg_tokens_per_batch,
                         "avg_tokens_per_batch_including_padding": avg_tokens_per_batch_including_padding,
                         "avg_pred_tokens_per_batch": avg_pred_tokens_per_batch,
-                        "per_device_tps": total_tokens / accelerator.num_processes / (time.time() - start_time),
+                        "per_device_tps": total_tokens
+                        / accelerator.num_processes
+                        / (time.perf_counter() - start_time),
                         "per_device_tps_including_padding": total_tokens_including_padding
                         / accelerator.num_processes
-                        / (time.time() - start_time),
+                        / (time.perf_counter() - start_time),
                         "reserved_mem_GiB": torch.cuda.max_memory_reserved(device=torch.cuda.current_device()) / 2**30,
                         "allocated_mem_GiB": torch.cuda.max_memory_allocated(device=torch.cuda.current_device())
                         / 2**30,
@@ -855,7 +859,7 @@ def main(args: FlatArguments, tc: TokenizerConfig):
                     avg_loss = sum_loss / total_fwd_passes
                     metrics_to_log["train_loss"] = avg_loss
                     if args.verbose:
-                        sec_per_step = (time.time() - start_time) / (completed_steps - resume_step)
+                        sec_per_step = (time.perf_counter() - start_time) / (completed_steps - resume_step)
                         steps_remaining = args.max_train_steps - completed_steps
                         secs_remaining = steps_remaining * sec_per_step
                         accelerator.print(
@@ -869,17 +873,23 @@ def main(args: FlatArguments, tc: TokenizerConfig):
                             / args.logging_steps
                         )
                         logger.info(
-                            f"  Step: {completed_steps}, LR: {lr_scheduler.get_last_lr()[0]}, Loss: {avg_loss}, Aux Loss: {avg_aux_loss}, TPS: {total_tokens / (time.time() - start_time)}"
+                            f"  Step: {completed_steps}, LR: {lr_scheduler.get_last_lr()[0]}, Loss: {avg_loss}, Aux Loss: {avg_aux_loss}, TPS: {total_tokens / (time.perf_counter() - start_time)}"
                         )
                         metrics_to_log["aux_loss"] = avg_aux_loss
                     else:
                         logger.info(
-                            f"  Step: {completed_steps}, LR: {lr_scheduler.get_last_lr()[0]}, Loss: {avg_loss}, TPS: {total_tokens / (time.time() - start_time)}"
+                            f"  Step: {completed_steps}, LR: {lr_scheduler.get_last_lr()[0]}, Loss: {avg_loss}, TPS: {total_tokens / (time.perf_counter() - start_time)}"
                         )
                     if args.verbose:
                         accelerator.print(f"{metrics_to_log=}")
                     if args.with_tracking:
                         accelerator.log(metrics_to_log, step=completed_steps)
+                    maybe_update_beaker_description(
+                        current_step=completed_steps,
+                        total_steps=args.max_train_steps,
+                        start_time=start_time,
+                        wandb_url=wandb_tracker.run.get_url() if wandb_tracker is not None else None,
+                    )
                     total_loss = 0
                     total_aux_loss = 0
 

scripts/train/debug/dpo.sh

@@ -1,8 +1,16 @@
-python mason.py \
+#!/bin/bash
+BEAKER_IMAGE="${1:-nathanl/open_instruct_auto}"
+
+uv run python mason.py \
     --cluster ai2/neptune \
+    --cluster ai2/saturn \
+    --cluster ai2/jupiter \
+    --cluster ai2/prior \
+    --description "Single GPU DPO run, for debugging purposes." \
     --workspace ai2/tulu-thinker \
     --priority high \
-    --image nathanl/open_instruct_auto --pure_docker_mode \
+    --image "$BEAKER_IMAGE" \
+    --pure_docker_mode \
     --preemptible \
     --num_nodes 1 \
     --budget ai2/oe-adapt \
@@ -26,5 +34,6 @@ python mason.py \
     --logging_steps 1 \
     --dataset_mixer_list allenai/tulu-3-wildchat-reused-on-policy-8b 100 \
     --add_bos \
+    --chat_template_name olmo \
     --seed 123
-    # --with_tracking
+    # --with_tracking