What does this PR do?

Problem:

In /src/llama_recipes/utils/train_utils.py the training loss is correctly divided by the # of gradient accumulation steps to scale down the gradient:

loss = loss / gradient_accumulation_steps

The training loss is then accumulated

total_loss += loss.detach().float()

and used in the following to calculate the average loss across all samples in the epoch:

train_epoch_loss = total_loss / len(train_dataloader)

As the accumulated loss is scaled down by gradient_accumulation_steps and len(train_dataloader) includes all steps (even the gradient accumulation ones), train_epoch_loss is gradient_accumulation_steps times lower than it should be.

Solution:

Accumulate the loss
total_loss += loss.detach().float()

before scaling it down

loss = loss / gradient_accumulation_steps

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• [ X ] Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Thanks for contributing 🎉!