From f58f5827596c4020d7a9979a1ecc4f76e642327f Mon Sep 17 00:00:00 2001 From: asolergi-nv Date: Wed, 4 Feb 2026 17:36:28 +0000 Subject: [PATCH 1/5] fix maxseqlen broadcast --- megatron/training/utils.py | 47 ++------------------------------------ 1 file changed, 2 insertions(+), 45 deletions(-) diff --git a/megatron/training/utils.py b/megatron/training/utils.py index d9681728467..8af1a44f1fb 100644 --- a/megatron/training/utils.py +++ b/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,7 +638,7 @@ 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, @@ -665,13 +646,7 @@ def _broadcast_cu_seqlens(cu_seqlens): ) 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, From 2a880aabeae6c2244bbff9feed119a8c96d4a7e0 Mon Sep 17 00:00:00 2001 From: asolergi-nv Date: Wed, 4 Feb 2026 17:49:00 +0000 Subject: [PATCH 2/5] add sft tokenizer to core tokenizers --- megatron/core/tokenizers/text/utils/build_tokenizer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/megatron/core/tokenizers/text/utils/build_tokenizer.py b/megatron/core/tokenizers/text/utils/build_tokenizer.py index 2017945204d..afb4174bad4 100644 --- a/megatron/core/tokenizers/text/utils/build_tokenizer.py +++ b/megatron/core/tokenizers/text/utils/build_tokenizer.py @@ -60,6 +60,14 @@ def build_tokenizer(args): tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, **kwargs) return tokenizer + elif args.tokenizer_type == 'SFTTokenizer': + # SFTTokenizer uses the legacy tokenizer system + from megatron.training.tokenizer.sft_tokenizer import SFTTokenizer + tokenizer = SFTTokenizer( + args.tokenizer_model, + args.sft_tokenizer_prompt_format, + ) + return tokenizer if args.tokenizer_metadata: metadata = args.tokenizer_metadata From 649b353a113aeeef2529ef99acaf82695edc5140 Mon Sep 17 00:00:00 2001 From: Antoni-Joan Solergibert Date: Thu, 5 Feb 2026 07:24:33 -0800 Subject: [PATCH 3/5] Guard SFTTokenizer --- megatron/core/tokenizers/text/utils/build_tokenizer.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/megatron/core/tokenizers/text/utils/build_tokenizer.py b/megatron/core/tokenizers/text/utils/build_tokenizer.py index afb4174bad4..614621e17fb 100644 --- a/megatron/core/tokenizers/text/utils/build_tokenizer.py +++ b/megatron/core/tokenizers/text/utils/build_tokenizer.py @@ -1,6 +1,11 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. from megatron.core.tokenizers import MegatronTokenizer +try: + from megatron.training.tokenizer.sft_tokenizer import SFTTokenizer + HAVE_SFT_TOKENIZER = True +except ImportError: + HAVE_SFT_TOKENIZER = False MEGATRON_TOKENIZERS = ['BertWordPieceLowerCase', 'BertWordPieceCase', 'GPT2BPETokenizer'] @@ -62,7 +67,10 @@ def build_tokenizer(args): return tokenizer elif args.tokenizer_type == 'SFTTokenizer': # SFTTokenizer uses the legacy tokenizer system - from megatron.training.tokenizer.sft_tokenizer import SFTTokenizer + if not HAVE_SFT_TOKENIZER: + raise ImportError( + "SFTTokenizer currently requires megatron.training module" + ) tokenizer = SFTTokenizer( args.tokenizer_model, args.sft_tokenizer_prompt_format, From d853d453a9e15b1124f8b6a605edce8cecc2a22c Mon Sep 17 00:00:00 2001 From: asolergi-nv Date: Fri, 6 Feb 2026 14:56:46 +0000 Subject: [PATCH 4/5] lint --- .../core/tokenizers/text/utils/build_tokenizer.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/megatron/core/tokenizers/text/utils/build_tokenizer.py b/megatron/core/tokenizers/text/utils/build_tokenizer.py index 614621e17fb..fe655ab5ce0 100644 --- a/megatron/core/tokenizers/text/utils/build_tokenizer.py +++ b/megatron/core/tokenizers/text/utils/build_tokenizer.py @@ -1,8 +1,10 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. from megatron.core.tokenizers import MegatronTokenizer + try: from megatron.training.tokenizer.sft_tokenizer import SFTTokenizer + HAVE_SFT_TOKENIZER = True except ImportError: HAVE_SFT_TOKENIZER = False @@ -68,13 +70,8 @@ def build_tokenizer(args): elif args.tokenizer_type == 'SFTTokenizer': # SFTTokenizer uses the legacy tokenizer system if not HAVE_SFT_TOKENIZER: - raise ImportError( - "SFTTokenizer currently requires megatron.training module" - ) - tokenizer = SFTTokenizer( - args.tokenizer_model, - args.sft_tokenizer_prompt_format, - ) + raise ImportError("SFTTokenizer currently requires megatron.training module") + tokenizer = SFTTokenizer(args.tokenizer_model, args.sft_tokenizer_prompt_format) return tokenizer if args.tokenizer_metadata: From 7478d7f8cf5e20e4fa6c2fe879522678ff1af6f8 Mon Sep 17 00:00:00 2001 From: Antoni-Joan Solergibert Date: Mon, 9 Feb 2026 05:59:57 -0800 Subject: [PATCH 5/5] move tokenizer utils --- megatron/core/tokenizers/{text => }/utils/build_tokenizer.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename megatron/core/tokenizers/{text => }/utils/build_tokenizer.py (100%) diff --git a/megatron/core/tokenizers/text/utils/build_tokenizer.py b/megatron/core/tokenizers/utils/build_tokenizer.py similarity index 100% rename from megatron/core/tokenizers/text/utils/build_tokenizer.py rename to megatron/core/tokenizers/utils/build_tokenizer.py