Update kd_loss.py

gather tensors in kd loss to enable TP > 1 with KD

Author: Separius

nemo_automodel/components/loss/kd_loss.py

@@ -15,6 +15,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.distributed.tensor import DTensor
 
 
 class KDLoss(nn.Module):
@@ -68,6 +69,14 @@ def forward(
             teacher_logits = teacher_logits.view(-1, teacher_logits.shape[-1])
         if labels.ndim > 1:
             labels = labels.view(-1)
+        
+        if isinstance(teacher_logits, DTensor):
+            teacher_logits = teacher_logits.full_tensor()
+        if isinstance(student_logits, DTensor):
+            student_logits = student_logits.full_tensor()
+        if isinstance(labels, DTensor):
+            labels = labels.full_tensor()
+        
         t_logits = teacher_logits[valid_mask]
         s_logits = student_logits[valid_mask]
         labels = labels[valid_mask]