Fix z loss (PrimeIntellect-ai#76)

* maybe fix oom

* Revert "maybe fix oom"

This reverts commit 0c09f83.

* maybe fix oom 2

* maybe fix oom 3

* Revert "maybe fix oom 3"

This reverts commit 4341086.

* maybe fix oom 4

* maybe fix oom 5

* Revert "maybe fix oom 5"

This reverts commit 55b17ec.

* maybe fix oom 6

* maybe fix oom 6

* maybe fix oom 9

Author: samsja

src/zeroband/loss.py

@@ -1,7 +1,9 @@
 from torch import Tensor
+import torch
 import torch.nn.functional as F
 
 
+@torch.compile
 def cross_entropy_max_z_loss(
     logits: Tensor,
     targets: Tensor,

src/zeroband/train.py

@@ -351,24 +351,31 @@ def train(config: Config):
                 flatten_logits = rearrange(logits, "b seq vocab -> (b seq) vocab")
                 flatten_labels = rearrange(labels, "b seq -> (b seq)")
 
-                if config.optim.z_loss is not None:
+                if config.optim.z_loss:
                     ce_loss, z_loss = cross_entropy_max_z_loss(
                         flatten_logits, flatten_labels, config.optim.z_loss_weight
                     )
-
                     ce_loss /= gradient_accumulation_steps
                     z_loss /= gradient_accumulation_steps
 
-                    loss_batch += ce_loss.detach()
-                    z_loss_batch += z_loss.detach()
-
+                    del logits
                     loss = ce_loss + z_loss
+                    loss.backward()
 
                 else:
                     loss = F.cross_entropy(flatten_logits, flatten_labels) / gradient_accumulation_steps
-                    loss_batch += loss.detach()
+                    del logits
+                    loss.backward()
+
+                if config.optim.z_loss:
+                    loss_batch += ce_loss.clone().detach()
+                    z_loss_batch += z_loss.clone().detach()
+                else:
+                    loss_batch += loss.clone().detach()
 
-                loss.backward()
+            dist.all_reduce(tensor=loss_batch, op=dist.ReduceOp.AVG, group=elastic_device_mesh.local_pg)
+            if config.optim.z_loss:
+                dist.all_reduce(tensor=z_loss_batch, op=dist.ReduceOp.AVG, group=elastic_device_mesh.local_pg)
 
             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
             inner_optimizer.step()
@@ -379,9 +386,6 @@ def train(config: Config):
             training_progress.step += 1
             inner_lr = [group["lr"] for group in inner_optimizer.param_groups][0]
 
-            dist.all_reduce(tensor=loss_batch, op=dist.ReduceOp.AVG, group=elastic_device_mesh.local_pg)
-            dist.all_reduce(tensor=z_loss_batch, op=dist.ReduceOp.AVG, group=elastic_device_mesh.local_pg)
-
             # syncing loss across all data parallel rank within a nodes
 
             new_tokens = config.data.seq_length * config.optim.batch_size