Fix step iteration bug in finetuning scripts (#1794)

litgpt/finetune/adapter.py

@@ -268,10 +268,13 @@ def fit(
  total_lengths = 0
  total_t0 = time.perf_counter()
 
- while step_count < max_steps and train_iterator.epoch < train.epochs:
+ while step_count < max_steps:
  iter_num += 1
  iter_t0 = time.perf_counter()
  batch = next(train_iterator)
+ if train_iterator.epoch >= train.epochs:
+ break
+
  input_ids, targets = batch["input_ids"], batch["labels"]
 
  is_accumulating = iter_num % train.gradient_accumulation_iters(devices) != 0

litgpt/finetune/adapter_v2.py

@@ -268,10 +268,13 @@ def fit(
  total_lengths = 0
  total_t0 = time.perf_counter()
 
- while step_count < max_steps and train_iterator.epoch < train.epochs:
+ while step_count < max_steps:
  iter_num += 1
  iter_t0 = time.perf_counter()
  batch = next(train_iterator)
+ if train_iterator.epoch >= train.epochs:
+ break
+
  input_ids, targets = batch["input_ids"], batch["labels"]
 
  is_accumulating = iter_num % train.gradient_accumulation_iters(devices) != 0

litgpt/finetune/full.py

@@ -244,10 +244,12 @@ def fit(
  )
  fabric.barrier()
 
- while state["step_count"] < max_steps and train_iterator.epoch < train.epochs:
+ while state["step_count"] < max_steps:
  state["iter_num"] += 1
  iter_t0 = time.perf_counter()
  batch = next(train_iterator)
+ if train_iterator.epoch >= train.epochs:
+ break
  input_ids, targets = batch["input_ids"], batch["labels"]
 
  is_accumulating = state["iter_num"] % train.gradient_accumulation_iters(devices) != 0

litgpt/finetune/lora.py

@@ -299,10 +299,12 @@ def fit(
  total_lengths = 0
  total_t0 = time.perf_counter()
 
- while step_count < max_steps and train_iterator.epoch < train.epochs:
+ while step_count < max_steps:
  iter_num += 1
  iter_t0 = time.perf_counter()
  batch = next(train_iterator)
+ if train_iterator.epoch >= train.epochs:
+ break
  input_ids, targets = batch["input_ids"], batch["labels"]
 
  is_accumulating = iter_num % train.gradient_accumulation_iters(devices) != 0