del logits=(bs, seq_len, vocab_size) to save 3.9G memory (#391)

logits=(bs, seq_len, vocab_size). call `del logits` to free it before
backward

<img width="1607" alt="Screenshot 2024-06-12 at 11 10 36 AM"
src="https://github.com/pytorch/torchtitan/assets/134637289/82db2792-59a3-40c4-9591-842be3dd9284">

Author: weifengpy

train.py

@@ -351,6 +351,9 @@ def loss_fn(pred, labels):
                 with loss_parallel_ctx():
                     pred = model(input_ids)
                     loss = loss_fn(pred, labels)
+                    # pred.shape=(bs, seq_len, vocab_size)
+                    # need to free to before bwd to avoid peaking memory
+                    del pred
                     loss.backward()
 
             # clip gradients