diff --git a/examples/post_training/modelopt/finetune.py b/examples/post_training/modelopt/finetune.py index d31fbcbed0c..20777ba14c8 100755 --- a/examples/post_training/modelopt/finetune.py +++ b/examples/post_training/modelopt/finetune.py @@ -17,7 +17,6 @@ from megatron.core import mpu, tensor_parallel from megatron.core.enums import ModelType from megatron.core.models.gpt import GPTModel -from megatron.core.tokenizers.text.libraries.huggingface_tokenizer import HuggingFaceTokenizer from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.loss_func import loss_func from megatron.post_training.model_builder import modelopt_gpt_mamba_builder @@ -122,7 +121,7 @@ def __init__( self, num_packed_samples: int, hf_dataset: str, - tokenizer: HuggingFaceTokenizer, + tokenizer: transformers.PreTrainedTokenizerBase, seq_length: int, num_shards: int = 1, shard_index: int = 0, @@ -143,8 +142,8 @@ def __init__( num_shards: number of shards for distributed training shard_index: shard index for distributed training """ - if not isinstance(tokenizer, HuggingFaceTokenizer): - raise ValueError("SFTDataset only supports HuggingFaceTokenizer!") + if not isinstance(tokenizer, transformers.PreTrainedTokenizerBase): + raise ValueError("SFTDataset only supports transformers.PreTrainedTokenizerBase!") self.num_packed_samples = num_packed_samples self.hf_dataset = hf_dataset @@ -284,7 +283,7 @@ def _process_example(self, example: Dict[str, Any]): return None # We always add eos between samples for training purpose. - input_ids = self.tokenizer.apply_chat_template(example, self.tokenizer.chat_template)["input_ids"] + input_ids = self.tokenizer.apply_chat_template(example) current_loss_mask = [1] * len(input_ids) input_ids = input_ids + [get_eos_id()] current_loss_mask += [0] @@ -344,8 +343,8 @@ def train_valid_test_sft_datasets_provider(train_val_test_num_samples): args = get_args() tokenizer = get_tokenizer() - if not isinstance(tokenizer._tokenizer, HuggingFaceTokenizer): - raise ValueError("SFTDataset only supports HuggingFaceTokenizer!") + if not isinstance(tokenizer._tokenizer, transformers.PreTrainedTokenizerBase): + raise ValueError("SFTDataset only supports transformers.PreTrainedTokenizerBase!") if args.micro_batch_size > 1: raise ValueError("SFTDataloader only supports micro_batch_size=1.") diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index cfe7f9a8eb9..bf65ae0430f 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -11,12 +11,16 @@ import sys import warnings -import modelopt.torch.quantization as mtq import torch import torch.distributed +from tqdm import tqdm + +# NOTE: Needs to be before modelopt imports in case megatron.core is not installed. +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))) + +import modelopt.torch.quantization as mtq from modelopt.torch.export import import_mcore_gpt_from_hf from modelopt.torch.utils.dataset_utils import get_dataset_dataloader -from tqdm import tqdm try: import modelopt.torch.quantization.plugins.psx_formats as mtq_psx @@ -31,8 +35,6 @@ mtq_luts = None warnings.warn("luts is not installed. LUTs quantization configs will not be available.") -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))) - from megatron.core.utils import get_batch_on_this_cp_rank from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index ac677550cde..322f12a4122 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -49,6 +49,13 @@ def __post_init__(self): cu_seqlens_with_max = torch.cat([cu_seqlens, total_tokens_tensor]) # Example: [0, 5, 7, 11, 16] -> [5, 2, 4, 5] seq_lengths = cu_seqlens_with_max[1:] - cu_seqlens_with_max[:-1] + # Clamp to non-negative: cu_seqlens_q_padded may not be strictly + # monotonic when context parallelism slices sequences across ranks, + # or when padded cumulative lengths exceed total_tokens (e.g. the + # appended total_tokens sentinel is smaller than cu_seqlens[-1] + # due to padding). In either case the diff can go negative, which + # causes torch.repeat_interleave to fail. + seq_lengths = seq_lengths.clamp(min=0) # Example: [5, 2, 4, 5] -> [0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3] self.seq_idx = ( torch.repeat_interleave( diff --git a/megatron/training/training.py b/megatron/training/training.py index 0e8dac37c59..46801cc4a24 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3256,16 +3256,14 @@ def evaluate( # Reduce across processes. for key in loss_dicts[0].keys(): if key not in total_loss_dict: - total_loss_dict[key] = torch.tensor( - [0.0, 0.0], dtype=torch.float - ).cuda() + total_loss_dict[key] = torch.tensor([0.0, 0.0], dtype=torch.float, device='cuda') val = [x[key].view(-1) for x in loss_dicts] if val[0].numel() == 2: if args.sft: # normalize over micro batch instead of global val = torch.vstack(val) - val = val[:, 0] / val[:, 1] + val = val[:, 0] / val[:, 1].clamp(min=1) val = val.mean() torch.distributed.all_reduce( val, diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 81768944623..31eee0f4dc6 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -26,6 +26,7 @@ from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset from megatron.core.enums import ModelType +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.models.gpt import GPTModel from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer @@ -62,11 +63,60 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): - """Generate a batch.""" + """Generate a batch. + + Packed sequence support (SFT / ``--sft`` flag): + When ``args.sft`` is True, the dataset emits THD-format batches where + multiple sequences are concatenated into a single flat token tensor. + The batch includes ``cu_seqlens`` (cumulative sequence lengths, shape + ``[1, S+1]``) and ``max_seqlen`` (shape ``[1]``) that describe the + individual sequence boundaries. + + This function validates and squeezes those fields: + - ``cu_seqlens``: asserted to have shape ``[1, S+1]`` (micro-batch + size must be 1 for packing), then squeezed to ``[S+1]``. + - ``max_seqlen``: asserted to be 1-D; kept as a tensor and passed + to ``get_thd_batch_on_this_cp_rank`` which performs the final + scalar conversion internally. + + Pipeline stage handling: + - First/last PP stages: fetch the full batch (tokens + labels) and + route through ``get_thd_batch_on_this_cp_rank`` to produce a + ``PackedSeqParams`` object that carries ``cu_seqlens`` and + ``max_seqlen`` to the attention kernel. + - Middle PP stages: only ``cu_seqlens`` and ``max_seqlen`` are + needed for attention masking; all other fields are returned as + ``None`` with a ``PackedSeqParams`` built directly here. + - MTP ranks (``mtp_on_this_rank``) also receive the full batch, + regardless of pipeline stage. + + Difference from ``pretrain_mamba.py``: + - Return format: GPT returns a 6-tuple + ``(tokens, labels, loss_mask, attention_mask, position_ids, + packed_seq_params)`` where ``packed_seq_params`` is a + ``PackedSeqParams`` dataclass. Mamba returns 7 values via + ``batch.values()`` with ``cu_seqlens`` and ``max_seqlen`` as + separate dict entries (no ``PackedSeqParams`` wrapper). + - Middle-stage return: GPT returns ``(None×5, PackedSeqParams)``; + Mamba returns an ``empty_batch`` dict with ``cu_seqlens`` and + ``max_seqlen`` set. + - CP with packed sequences: GPT delegates to + ``get_thd_batch_on_this_cp_rank`` (MCore utility); Mamba + implements the ``tex.thd_get_partitioned_indices`` CP slicing + inline and does not call that helper. + - MTP: GPT passes ``mtp_on_this_rank`` to ``get_batch_on_this_tp_rank`` + and uses it to gate the early-return; Mamba has no MTP support. + - ``max_seqlen`` conversion: Mamba converts to a Python int scalar + before returning (``int(max_seqlen[0].item())``); GPT keeps it as + a tensor and lets ``get_thd_batch_on_this_cp_rank`` convert it, + except for the middle-stage ``PackedSeqParams`` where conversion + is done inline. + """ args = get_args() config = core_transformer_config_from_args(args) # TODO: this is pretty hacky, find a better way - if not is_first_or_last_pipeline_stage(vp_stage) and ( + is_packed_sequence = get_args().sft # SFT always uses packed sequence + if not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence and ( (not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage))): return None, None, None, None, None, None @@ -83,16 +133,33 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): if local_cp_size is not None: local_cp_size = int(local_cp_size.item()) + if cu_seqlens is not None: + assert ( + cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1 + ), "micro-batch-size must be 1 for packing" + cu_seqlens = cu_seqlens[0] + assert max_seqlen.dim() == 1 + + # For middle pipeline stages with packed sequences, only cu_seqlens and + # max_seqlen are needed (for attention masking); skip the full batch. + if not is_first_or_last_pipeline_stage(vp_stage) and is_packed_sequence: + return None, None, None, None, None, PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=int(max_seqlen[0].item()), + max_seqlen_kv=int(max_seqlen[0].item()), + qkv_format='thd', + ) + if cu_seqlens is None and local_cp_size is None: # slice batch along sequence dimension for context parallelism batch = get_batch_on_this_cp_rank(batch) # The implementation of this function is in MCore packed_seq_params = None elif local_cp_size is None: # Packed THD format - assert max_seqlen.dim() == 1 batch, packed_seq_params = get_thd_batch_on_this_cp_rank(batch, cu_seqlens, cu_seqlens_padded, max_seqlen) else: # Hybrid CP format batch, packed_seq_params = get_batch_on_this_hybrid_cp_rank(batch, local_cp_size) - + return (*batch.values(), packed_seq_params) @@ -201,13 +268,17 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa return output_tensor, partial(loss_func, loss_mask, model=model) -def is_dataset_built_on_rank(vp_stage=None): +def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): args = get_args() config = core_transformer_config_from_args(args) + if parallel_state.get_tensor_model_parallel_rank() != 0: + return False + elif is_packed_sequence: + return True return ( is_first_or_last_pipeline_stage(vp_stage) or mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage) - ) and parallel_state.get_tensor_model_parallel_rank() == 0 + ) def core_gpt_dataset_config_from_args(args): @@ -285,8 +356,11 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None config = core_gpt_dataset_config_from_args(args) + + is_packed_sequence = False if args.sft: dataset_type = SFTDataset + is_packed_sequence = True # SFT always uses packed sequence else: if args.mock_data: dataset_type = MockGPTDataset @@ -297,9 +371,9 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None print_rank_0("> building train, validation, and test datasets for GPT ...") - is_dataset_built = partial(is_dataset_built_on_rank, vp_stage=vp_stage) + is_dataset_built = partial(is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence) train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( - dataset_type, train_val_test_num_samples, partial(is_dataset_built_on_rank, vp_stage=vp_stage), config + dataset_type, train_val_test_num_samples, is_dataset_built, config ).build() print_rank_0("> finished creating GPT datasets ...")