Add zl

Author: daviswer

fms_fsdp/utils/train_utils.py

@@ -89,6 +89,7 @@ def train(
         output = output.logits if hasattr(output, "logits") else output
         ce_loss = torch.nn.CrossEntropyLoss()
         loss = ce_loss(output.view(-1, output.size(-1)), label.view(-1).long())
+        loss = loss + .0001 * torch.logsumexp(output, dim=-1).pow(2).mean()
 
         loss.backward()
         ddp_stats[1] += model.clip_grad_norm_(cfg.grad_clip_thresh).item()