Skip to content

Fix SFT Pipeline when TP>1 #3268

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jaredcasper merged 8 commits into from
Feb 11, 2026
Merged

Fix SFT Pipeline when TP>1 #3268

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
f58f582
fix maxseqlen broadcast
asolergi-nv Feb 4, 2026
2a880aa
add sft tokenizer to core tokenizers
asolergi-nv Feb 4, 2026
649b353
Guard SFTTokenizer
asolergi-nv Feb 5, 2026
03af6e3
Merge branch 'main' into sft_fix
asolergi-nv Feb 6, 2026
d853d45
lint
asolergi-nv Feb 6, 2026
7478d7f
move tokenizer utils
asolergi-nv Feb 9, 2026
3e6b5aa
Merge branch 'main' into sft_fix
asolergi-nv Feb 9, 2026
7180ccb
Merge branch 'main' into sft_fix
asolergi-nv Feb 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 2 additions & 45 deletions megatron/training/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
Loading