Add Loss Averaging and no_backward_sync() in Gradient Accumulation Logic (#357)

Co-authored-by: Carlos Mocholí <[email protected]>

Author: waitzkin

finetune/adapter.py

@@ -131,9 +131,9 @@ def train(
         t0 = time.time()
 
         input_ids, targets = get_batch(fabric, train_data)
-        logits = model(input_ids)
-        loss = loss_fn(logits, targets)
         with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_iters != 0)):
+            logits = model(input_ids)
+            loss = loss_fn(logits, targets)
             fabric.backward(loss / gradient_accumulation_iters)
 
         if (iter_num + 1) % gradient_accumulation_iters == 0:

finetune/adapter_v2.py

@@ -137,9 +137,9 @@ def train(
         t0 = time.time()
 
         input_ids, targets = get_batch(fabric, train_data)
-        logits = model(input_ids)
-        loss = loss_fn(logits, targets)
         with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_iters != 0)):
+            logits = model(input_ids)
+            loss = loss_fn(logits, targets)
             fabric.backward(loss / gradient_accumulation_iters)
 
         if (iter_num + 1) % gradient_accumulation_iters == 0:

finetune/full.py

@@ -114,12 +114,12 @@ def train(
                 param_group['lr'] = lr
 
         t0 = time.time()
-
+        
+        input_ids, targets = get_batch(fabric, train_data)
         with fabric.no_backward_sync(model, enabled=is_accumulating):
-            input_ids, targets = get_batch(fabric, train_data)
             logits = model(input_ids)
             loss = loss_fn(logits, targets)
-            fabric.backward(loss)
+            fabric.backward(loss / gradient_accumulation_iters)
 
         if not is_accumulating:
             optimizer.step()

finetune/lora.py

@@ -108,9 +108,10 @@ def train(
         t0 = time.time()
 
         input_ids, targets = get_batch(fabric, train_data)
-        logits = model(input_ids)
-        loss = loss_fn(logits, targets)
-        fabric.backward(loss)
+        with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_iters != 0)):
+            logits = model(input_ids)
+            loss = loss_fn(logits, targets)
+            fabric.backward(loss / gradient_accumulation_iters)
 
         if (iter_num + 1) % gradient_accumulation_iters == 0:
             optimizer.step()