Fix SFT Pipeline when TP>1 (NVIDIA#3268)

Author: asolergi-nv

megatron/training/utils.py

@@ -600,25 +600,6 @@ def _broadcast_cu_seqlens(cu_seqlens):
             _broadcast(batch['loss_mask'])
             _broadcast(batch['attention_mask'])
 
-        def _broadcast_cu_seqlens(cu_seqlens):
-            dev = torch.cuda.current_device()
-
-            n = 0 if cu_seqlens is None else int(cu_seqlens.numel())
-            n_tensor = torch.tensor(n, dtype=torch.int64, device=dev)
-            _broadcast(n_tensor)
-
-            if n == 0:
-                buf = torch.empty(0, dtype=torch.int32, device=dev)
-            else:
-                assert isinstance(cu_seqlens, torch.Tensor)
-                assert cu_seqlens.dtype == torch.int32
-                assert cu_seqlens.shape[0] == 1, "micro-batch-size must be 1 for packing"
-                buf = cu_seqlens.to(device=dev, non_blocking=True).contiguous()
-            _broadcast(buf)
-
-        _broadcast_cu_seqlens(batch['cu_seqlens'])
-        _broadcast(batch['max_seqlen'])
-
     else:
         if args.hybrid_context_parallel:
             seq_len = torch.tensor(0, dtype=torch.int32, device=torch.cuda.current_device())
@@ -657,21 +638,15 @@ def _broadcast_cu_seqlens(cu_seqlens):
             device=torch.cuda.current_device(),
         )
         cu_seqlens = None
-        if args.sft:
+        if args.hybrid_context_parallel or args.sft:
             max_seqlen = torch.empty(
                 1,
                 dtype=torch.int32,
                 device=torch.cuda.current_device(),
             )
         else:
             max_seqlen = None
-
-        cu_seqlens = None
-        max_seqlen = torch.empty(
-            1,
-            dtype=torch.int32,
-            device=torch.cuda.current_device(),
-        ) if args.hybrid_context_parallel else None
+        
         local_cp_size = torch.empty(
             1,
             dtype=torch.int32,
@@ -726,24 +701,6 @@ def _broadcast_cu_seqlens():
             _broadcast(loss_mask)
             _broadcast(attention_mask)
 
-        def _broadcast_cu_seqlens():
-            dev = torch.cuda.current_device()
-
-            n = torch.empty((), dtype=torch.int64, device=dev)
-            _broadcast(n)
-            n = int(n.item())
-
-            if n == 0:
-                cu_seqlens = torch.empty(0, dtype=torch.int32, device=dev)
-            else:
-                cu_seqlens = torch.empty((args.micro_batch_size, n), dtype=torch.int32, device=dev)
-            _broadcast(cu_seqlens)
-
-            return cu_seqlens if n > 0 else None
-
-        cu_seqlens = _broadcast_cu_seqlens()
-        _broadcast(max_seqlen)
-
         batch = {
             'tokens': tokens,
             'labels': labels,