delete loss before bwd

Author: Felipe Mello

recipes/full_finetune_distributed.py

@@ -523,6 +523,8 @@ def train(self) -> None:
                 logits = logits.transpose(1, 2)
                 # Compute loss
                 loss = self._loss_fn(logits, labels)
+                # free logits otherwise it peaks backward memory
+                del logits
 
                 loss = loss / self._gradient_accumulation_steps
                 running_loss += loss

recipes/full_finetune_single_device.py

@@ -452,6 +452,8 @@ def train(self) -> None:
                 logits = logits.transpose(1, 2)
                 # Compute loss
                 loss = self._loss_fn(logits, labels)
+                # free logits otherwise it peaks backward memory
+                del logits
 
                 loss = loss / self._gradient_accumulation_steps
                 running_loss += loss

recipes/lora_finetune_single_device.py

@@ -573,6 +573,9 @@ def train(self) -> None:
                     logits = logits.transpose(1, 2)
                     # Compute loss
                     loss = self._loss_fn(logits, labels)
+                    # free logits otherwise it peaks backward memory
+                    del logits
+
                     loss = loss / self._gradient_accumulation_steps
                     running_loss += loss
                     loss.backward()

recipes/qat_distributed.py

@@ -572,6 +572,8 @@ def train(self) -> None:
                 logits = logits.transpose(1, 2)
                 # Compute loss
                 loss = self._loss_fn(logits, labels)
+                # free logits otherwise it peaks backward memory
+                del logits
 
                 loss = loss / self._gradient_accumulation_steps
                 running_loss += loss