1 files changed, 16 insertions, 7 deletions

diff --git a/train_lora.py b/train_lora.py
index 0d8ee23..29e40b2 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -919,6 +919,8 @@ def main():
             args.num_train_steps / len(lora_datamodule.train_dataset)
         ) * args.gradient_accumulation_steps
         lora_sample_frequency = math.ceil(num_train_epochs * (lora_sample_frequency / args.num_train_steps))
+    num_training_steps_per_epoch = math.ceil(len(lora_datamodule.train_dataset) / args.gradient_accumulation_steps)
+    num_train_steps = num_training_steps_per_epoch * num_train_epochs
     if args.sample_num is not None:
         lora_sample_frequency = math.ceil(num_train_epochs / args.sample_num)
 
@@ -956,15 +958,19 @@ def main():
 
     training_iter = 0
 
+    lora_project = "lora"
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers(lora_project)
+
     while True:
-        training_iter += 1
-        if training_iter > args.auto_cycles:
+        if training_iter >= args.auto_cycles:
             response = input("Run another cycle? [y/n] ")
             if response.lower().strip() == "n":
                 break
 
         print("")
-        print(f"============ LoRA cycle {training_iter} ============")
+        print(f"============ LoRA cycle {training_iter + 1} ============")
         print("")
 
         lora_optimizer = create_optimizer(params_to_optimize)
@@ -976,19 +982,18 @@ def main():
             train_epochs=num_train_epochs,
         )
 
-        lora_project = f"lora_{training_iter}"
-        lora_checkpoint_output_dir = output_dir / lora_project / "model"
-        lora_sample_output_dir = output_dir / lora_project / "samples"
+        lora_checkpoint_output_dir = output_dir / lora_project / f"{training_iter + 1}" / "model"
+        lora_sample_output_dir = output_dir / lora_project / f"{training_iter + 1}" / "samples"
 
         trainer(
             strategy=lora_strategy,
-            project=lora_project,
             train_dataloader=lora_datamodule.train_dataloader,
             val_dataloader=lora_datamodule.val_dataloader,
             optimizer=lora_optimizer,
             lr_scheduler=lora_lr_scheduler,
             num_train_epochs=num_train_epochs,
             gradient_accumulation_steps=args.gradient_accumulation_steps,
+            global_step_offset=training_iter * num_train_steps,
             # --
             group_labels=group_labels,
             sample_output_dir=lora_sample_output_dir,
@@ -996,6 +1001,10 @@ def main():
             sample_frequency=lora_sample_frequency,
         )
 
+        training_iter += 1
+
+    accelerator.end_training()
+
 
 if __name__ == "__main__":
     main()