From 766988c114983a45c3c6ba817b1c028e5de64518 Mon Sep 17 00:00:00 2001 From: kaiyuan Date: Fri, 5 Jun 2026 13:33:07 +0800 Subject: [PATCH 01/56] ascend: adapt vime for NPU --- docker/Dockerfile.npu | 73 + docker/patch/latest/megatron-bridge.patch | 100996 +++++++++++++++ docker/patch/latest/megatron-npu.patch | 2337 + docker/patch/latest/mindspeed.patch | 70 + docker/patch/latest/vllm-ascend.patch | 78 + docker/patch/latest/vllm.patch | 323 + requirements.txt | 6 +- tools/convert_hf_to_torch_dist.py | 10 +- train.py | 1 + vime/backends/megatron_utils/__init__.py | 11 +- vime/backends/megatron_utils/actor.py | 24 +- vime/backends/megatron_utils/data.py | 6 +- vime/backends/megatron_utils/loss.py | 4 +- .../quantizer_compressed_tensors.py | 2 +- .../megatron_utils/megatron_to_hf/qwen2.py | 5 + .../backends/megatron_utils/model_provider.py | 9 +- .../megatron_utils/update_weight/common.py | 9 +- .../hf_weight_iterator_direct.py | 6 +- .../update_weight_from_distributed.py | 29 +- .../update_weight_from_tensor.py | 4 +- vime/backends/vllm_utils/arguments.py | 11 +- vime/backends/vllm_utils/vllm_engine.py | 45 +- vime/ray/actor_group.py | 58 +- vime/ray/placement_group.py | 9 +- vime/ray/rollout.py | 30 +- vime/ray/train_actor.py | 11 +- vime/ray/utils.py | 4 +- vime/rollout/vllm_rollout.py | 5 +- vime/utils/arguments.py | 10 +- vime/utils/common.py | 9 + vime/utils/external_utils/command_utils.py | 78 + vime/utils/memory_utils.py | 14 +- vime/utils/reloadable_process_group.py | 2 +- vime/utils/routing_replay.py | 4 +- vime/utils/tensor_backper.py | 8 +- 35 files changed, 104186 insertions(+), 115 deletions(-) create mode 100644 docker/Dockerfile.npu create mode 100644 docker/patch/latest/megatron-bridge.patch create mode 100644 docker/patch/latest/megatron-npu.patch create mode 100644 docker/patch/latest/mindspeed.patch create mode 100644 docker/patch/latest/vllm-ascend.patch create mode 100644 vime/utils/common.py diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu new file mode 100644 index 000000000..526a06470 --- /dev/null +++ b/docker/Dockerfile.npu @@ -0,0 +1,73 @@ +ARG BASE_IMAGE=registry-cbu.huawei.com/atelier/verl:verl_0.8.0-vllm_0.17.0-mindspeed_0.16.0-pytorch_2.9.0-cann_8.5.2-py_3.11-hce_2.0.2512-aarch64-snt9b-20260513174224-90d255f +FROM ${BASE_IMAGE} + +ARG PATCH_VERSION=latest +ARG MEGATRON_COMMIT=3bec9aa + +WORKDIR /home/ma-user + +RUN pip install vllm==0.17.0 --no-deps || true + +RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ + cd Megatron-LM && git checkout ${MEGATRON_COMMIT} + +RUN git clone https://github.com/vllm-project/vllm-ascend.git && \ + cd vllm-ascend && git checkout v0.17.0rc1 + +RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@a193d9dd1b877d33c64a41cfb3db9f867df2d926 --no-cache-dir --force-reinstall || true + +COPY docker/patch/${PATCH_VERSION}/megatron.patch /home/ma-user/Megatron-LM/ +COPY docker/patch/${PATCH_VERSION}/megatron-npu.patch /home/ma-user/Megatron-LM/ +COPY docker/patch/${PATCH_VERSION}/megatron-bridge.patch /home/ma-user/Megatron-LM/ +RUN cd Megatron-LM && \ + git update-index --refresh && \ + git apply megatron.patch --3way && \ + git apply megatron-npu.patch --3way && \ + git apply megatron-bridge.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm megatron.patch megatron-npu.patch megatron-bridge.patch && \ + pip install -e . --no-build-isolation + +COPY docker/patch/${PATCH_VERSION}/vllm.patch /home/ma-user/vllm/ +RUN cd vllm && \ + git update-index --refresh && \ + git apply vllm.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm vllm.patch + +COPY docker/patch/${PATCH_VERSION}/vllm-ascend.patch /home/ma-user/vllm-ascend/ +RUN cd vllm-ascend && \ + git update-index --refresh && \ + git apply vllm-ascend.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm vllm-ascend.patch + +COPY docker/patch/${PATCH_VERSION}/mindspeed.patch /home/ma-user/MindSpeed/ +RUN cd MindSpeed && \ + git update-index --refresh && \ + git apply mindspeed.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm mindspeed.patch + +ARG VIME_COMMIT=ascend +RUN git clone -b ${VIME_COMMIT} https://github.com/vllm-project/vime.git /home/ma-user/vime && \ + cd /home/ma-user/vime && \ + pip install -e . --no-deps + +ENV PYTHONPATH=/home/ma-user/Megatron-LM:/home/ma-user/vllm:/home/ma-user/vime:${PYTHONPATH} +ENV CUDA_DEVICE_MAX_CONNECTIONS=1 + +ENTRYPOINT [] +CMD ["/bin/bash"] diff --git a/docker/patch/latest/megatron-bridge.patch b/docker/patch/latest/megatron-bridge.patch new file mode 100644 index 000000000..f0208883f --- /dev/null +++ b/docker/patch/latest/megatron-bridge.patch @@ -0,0 +1,100996 @@ +diff --git a/megatron/bridge/__init__.py b/megatron/bridge/__init__.py +new file mode 100755 +index 0000000..185ac77 +--- /dev/null ++++ b/megatron/bridge/__init__.py +@@ -0,0 +1,37 @@ ++# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++"""Megatron Bridge - A component of the Megatron ecosystem.""" ++ ++from megatron.bridge.models.conversion.auto_bridge import AutoBridge ++from megatron.bridge.package_info import ( ++ __contact_emails__, ++ __contact_names__, ++ __download_url__, ++ __homepage__, ++ __package_name__, ++ __repository_url__, ++ __version__, ++) ++ ++ ++__all__ = [ ++ "__version__", ++ "__package_name__", ++ "__contact_names__", ++ "__contact_emails__", ++ "__homepage__", ++ "__repository_url__", ++ "__download_url__", ++ "AutoBridge", ++] +diff --git a/megatron/bridge/data/__init__.py b/megatron/bridge/data/__init__.py +new file mode 100755 +index 0000000..e69de29 +diff --git a/megatron/bridge/data/builders/__init__.py b/megatron/bridge/data/builders/__init__.py +new file mode 100755 +index 0000000..e69de29 +diff --git a/megatron/bridge/data/builders/finetuning_dataset.py b/megatron/bridge/data/builders/finetuning_dataset.py +new file mode 100755 +index 0000000..0dff48f +--- /dev/null ++++ b/megatron/bridge/data/builders/finetuning_dataset.py +@@ -0,0 +1,342 @@ ++# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++ ++import logging ++from pathlib import Path ++from typing import Any, Optional, Union ++ ++import torch ++from megatron.core.msc_utils import MultiStorageClientFeature ++from megatron.core.tokenizers.text.libraries import HuggingFaceTokenizer ++ ++from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs ++from megatron.bridge.data.datasets.sft import create_sft_dataset ++from megatron.bridge.utils.common_utils import get_rank_safe, print_rank_0 ++ ++ ++logger = logging.getLogger(__name__) ++ ++ ++class FinetuningDatasetBuilder: ++ """Builder class for fine-tuning datasets. ++ ++ This class provides methods to build datasets for fine-tuning large language models. ++ It follows a builder pattern similar to BlendedMegatronDatasetBuilder but adapted for ++ fine-tuning scenarios. ++ ++ Args: ++ dataset_root (Union[str, Path]): The root directory containing training, validation, and test data. ++ tokenizer: The tokenizer to use for preprocessing text. ++ is_built_on_rank (Callable): Function that returns True if the dataset should be built on current rank. ++ seq_length (int, optional): The maximum sequence length. Defaults to 2048. ++ seed (int, optional): Random seed for data shuffling. Defaults to 1234. ++ memmap_workers (int, optional): Number of worker processes for memmap datasets. Defaults to 1. ++ max_train_samples (int, optional): Maximum number of training samples. Defaults to None. ++ packed_sequence_specs (Optional[PackedSequenceSpecs], optional): Specifications for packed sequences. Defaults to None. ++ dataset_kwargs (Optional[dict[str, Any]], optional): Additional dataset creation arguments. Defaults to None. ++ do_validation (bool, optional): Whether to build the validation dataset. Defaults to True. ++ do_test (bool, optional): Whether to build the test dataset. Defaults to True. ++ """ ++ ++ def __init__( ++ self, ++ dataset_root: Union[str, Path], ++ tokenizer, ++ seq_length: int = 2048, ++ seed: int = 1234, ++ memmap_workers: int = 1, ++ max_train_samples: Optional[int] = None, ++ packed_sequence_specs: Optional[PackedSequenceSpecs] = None, ++ dataset_kwargs: Optional[dict[str, Any]] = None, ++ do_validation: bool = True, ++ do_test: bool = True, ++ ): ++ if MultiStorageClientFeature.is_enabled(): ++ msc = MultiStorageClientFeature.import_package() ++ self.dataset_root = msc.Path(dataset_root) ++ else: ++ self.dataset_root = Path(dataset_root) ++ self.tokenizer = tokenizer ++ self.seq_length = seq_length ++ self.seed = seed ++ self.memmap_workers = memmap_workers ++ self.max_train_samples = max_train_samples ++ self.packed_sequence_specs = packed_sequence_specs ++ self.packed_sequence_size = -1 if not packed_sequence_specs else packed_sequence_specs.packed_sequence_size ++ self.dataset_kwargs = dataset_kwargs or {} ++ self._pad_cu_seqlens = False if not packed_sequence_specs else packed_sequence_specs.pad_cu_seqlens ++ self._pad_seq_to_mult = None if not packed_sequence_specs else packed_sequence_specs.pad_seq_to_mult ++ ++ self.do_validation = do_validation ++ self.do_test = do_test ++ ++ print_rank_0(f"Building FinetuningDatasetBuilder with root={self.dataset_root}") ++ ++ if self.packed_sequence_size > 0: ++ print_rank_0(f"Using packed sequences with size {self.packed_sequence_size}") ++ ++ def prepare_data(self) -> None: ++ """Prepare data if needed.""" ++ self.prepare_packed_data() ++ ++ def prepare_packed_data(self) -> None: ++ """Prepare packed sequence data files if configured.""" ++ if self.packed_sequence_size > 0: ++ from megatron.bridge.data.datasets.packed_sequence import prepare_packed_sequence_data ++ ++ if not self.train_path_packed.is_file(): ++ print_rank_0(f"Preparing packed training data at {self.train_path_packed}") ++ prepare_packed_sequence_data( ++ input_path=self.train_path, ++ output_path=self.train_path_packed, ++ packed_sequence_size=self.packed_sequence_size, ++ tokenizer=self.tokenizer, ++ max_seq_length=self.seq_length, ++ seed=self.seed, ++ output_metadata_path=self.pack_metadata, ++ dataset_kwargs=self.dataset_kwargs, ++ pad_seq_to_mult=self._pad_seq_to_mult, ++ ) ++ ++ if self.do_validation and not self.validation_path_packed.is_file(): ++ print_rank_0(f"Preparing packed validation data at {self.validation_path_packed}") ++ prepare_packed_sequence_data( ++ input_path=self.validation_path, ++ output_path=self.validation_path_packed, ++ packed_sequence_size=self.packed_sequence_size, ++ tokenizer=self.tokenizer, ++ max_seq_length=self.seq_length, ++ seed=self.seed, ++ output_metadata_path=self.pack_metadata, ++ dataset_kwargs=self.dataset_kwargs, ++ pad_seq_to_mult=self._pad_seq_to_mult, ++ ) ++ ++ def build(self) -> list[Optional[Any]]: ++ """Build train, validation, and test datasets. ++ ++ This method creates the necessary datasets based on the configuration. ++ It first ensures data preparation (e.g., packing) is done (on rank 0), ++ then builds the datasets potentially using the prepared files. ++ ++ Returns: ++ A list containing the train, validation, and test datasets. ++ Elements can be None if the corresponding data file doesn't exist ++ or if dataset building is skipped for the split. ++ """ ++ # Prepare packed data if needed ++ if get_rank_safe() == 0: ++ self.prepare_data() ++ ++ if torch.distributed.is_initialized(): ++ torch.distributed.barrier() ++ ++ # This needs to be called on all ranks ++ datasets: list[Optional[Any]] = self._build_datasets() ++ return datasets ++ ++ def _build_datasets(self) -> list[Optional[Any]]: ++ """Internal method to build all datasets. ++ ++ Returns: ++ list[Optional[Any]]: The train, validation, and test datasets. ++ """ ++ train_ds = self._create_dataset( ++ self.train_path if self.packed_sequence_size <= 0 else self.train_path_packed, ++ pack_metadata_path=None if self.packed_sequence_size <= 0 else self.pack_metadata, ++ max_num_samples=self.max_train_samples, ++ **self.dataset_kwargs, ++ ) ++ ++ if self.do_validation: ++ valid_ds = self._create_dataset( ++ self.validation_path if self.packed_sequence_size <= 0 else self.validation_path_packed, ++ pack_metadata_path=None if self.packed_sequence_size <= 0 else self.pack_metadata, ++ is_test=True, ++ **self.dataset_kwargs, ++ ) ++ else: ++ valid_ds = None ++ ++ if self.do_test: ++ test_ds = self._create_dataset( ++ self.test_path, ++ is_test=True, ++ **self.dataset_kwargs, ++ ) ++ else: ++ test_ds = None ++ ++ return [train_ds, valid_ds, test_ds] ++ ++ def _create_dataset( ++ self, ++ path: Union[str, Path], ++ pack_metadata_path: Optional[Union[str, Path]] = None, ++ is_test: bool = False, ++ **kwargs: Any, ++ ) -> Optional[Any]: ++ """Create a single dataset instance (train, validation, or test). ++ ++ Args: ++ path: Path to the dataset file ++ pack_metadata_path: Path to the packed sequence metadata ++ is_test: Whether this is a test dataset ++ **kwargs: Additional arguments to pass to the dataset constructor ++ ++ Returns: ++ The created dataset ++ """ ++ if MultiStorageClientFeature.is_enabled(): ++ msc = MultiStorageClientFeature.import_package() ++ path_exists = msc.Path(path).exists() ++ else: ++ path_exists = Path(path).exists() ++ ++ if not path_exists: ++ print_rank_0(f"Warning: Dataset path {path} does not exist") ++ return None ++ ++ is_not_packing = self.packed_sequence_size <= 0 ++ return create_sft_dataset( ++ path, ++ tokenizer=self.tokenizer, ++ seq_length=(self.seq_length if is_not_packing else self.packed_sequence_size), ++ memmap_workers=self.memmap_workers, ++ seed=self.seed, ++ is_test=is_test, ++ pack_metadata_file_path=None if is_not_packing else pack_metadata_path, ++ pad_cu_seqlens=False if is_not_packing else self._pad_cu_seqlens, ++ pad_seq_to_mult=1 if is_not_packing else self._pad_seq_to_mult, ++ **kwargs, ++ ) ++ ++ @property ++ def train_path(self) -> Path: ++ """Path to the training dataset file (training.jsonl).""" ++ return self.dataset_root / "training.jsonl" ++ ++ @property ++ def default_pack_path(self) -> Path: ++ """The default directory path for storing packed sequence files. ++ ++ Constructed based on the dataset root and tokenizer model name. ++ Creates the directory if it doesn't exist. ++ ++ Returns: ++ The Path object for the default packing directory. ++ """ ++ tokenizer_model_name = self._extract_tokenizer_model_name() ++ default_pack_path = ( ++ self.dataset_root / "packed" / f"{tokenizer_model_name}_pad_seq_to_mult{self._pad_seq_to_mult}" ++ ) ++ if not default_pack_path.exists(): ++ default_pack_path.mkdir(parents=True, exist_ok=True) ++ logger.info(f"Using default path for packing files: {str(default_pack_path)}") ++ ++ return default_pack_path ++ ++ @property ++ def pack_metadata(self) -> Path: ++ """Path to the metadata file for packed sequences. ++ ++ Determined by `packed_sequence_specs` or defaults based on the ++ `default_pack_path` and `packed_sequence_size`. ++ ++ Returns: ++ The Path object for the packed sequence metadata file. ++ ++ Raises: ++ ValueError: If packed sequences are not configured. ++ """ ++ if self.packed_sequence_size > 0: ++ if self.packed_sequence_specs.packed_metadata_path is not None: ++ return self.packed_sequence_specs.packed_metadata_path ++ return self.default_pack_path / f"{self.packed_sequence_size}_metadata.jsonl" ++ else: ++ raise ValueError("pack_metadata invalid since packed sequence size is not specified.") ++ ++ @property ++ def train_path_packed(self) -> Path: ++ """Path to the packed training dataset file (.npy). ++ ++ Determined by `packed_sequence_specs` or defaults based on the ++ `default_pack_path` and `packed_sequence_size`. ++ ++ Returns: ++ The Path object for the packed training data file. ++ ++ Raises: ++ ValueError: If packed sequences are not configured. ++ """ ++ if self.packed_sequence_size > 0: ++ if self.packed_sequence_specs.packed_train_data_path is not None: ++ return self.packed_sequence_specs.packed_train_data_path ++ return self.default_pack_path / f"training_{self.packed_sequence_size}.npy" ++ else: ++ raise ValueError("`train_path_packed` invalid since packed sequence size is not specified.") ++ ++ @property ++ def validation_path_packed(self) -> Path: ++ """Path to the packed validation dataset file (.npy). ++ ++ Determined by `packed_sequence_specs` or defaults based on the ++ `default_pack_path` and `packed_sequence_size`. ++ ++ Returns: ++ The Path object for the packed validation data file. ++ ++ Raises: ++ ValueError: If packed sequences are not configured. ++ """ ++ if self.packed_sequence_size > 0: ++ if self.packed_sequence_specs.packed_val_data_path is not None: ++ return self.packed_sequence_specs.packed_val_data_path ++ return self.default_pack_path / f"validation_{self.packed_sequence_size}.npy" ++ else: ++ raise ValueError("`validation_path_packed` invalid since packed sequence size is not specified.") ++ ++ @property ++ def validation_path(self) -> Path: ++ """Path to the validation dataset file (validation.jsonl).""" ++ return self.dataset_root / "validation.jsonl" ++ ++ @property ++ def test_path(self) -> Path: ++ """Path to the test dataset file (test.jsonl).""" ++ return self.dataset_root / "test.jsonl" ++ ++ def _extract_tokenizer_model_name(self) -> str: ++ """Automatically get the model name from model path.""" ++ # Legacy tokenizer compatibility ++ tokenizer_cls = HuggingFaceTokenizer ++ tokenizer_instance = self.tokenizer._tokenizer ++ ++ if self.packed_sequence_specs and self.packed_sequence_specs.tokenizer_model_name is not None: ++ return self.packed_sequence_specs.tokenizer_model_name ++ elif isinstance(tokenizer_instance, tokenizer_cls): ++ name = self.tokenizer.path ++ ++ if name.endswith("context/nemo_tokenizer"): ++ # NEMO_HOME/hf_org/hf_model/context/nemo_tokenizer => hf_org--hf_model ++ tokenizer_model_name = "--".join(name.split("/")[-4:-2]) ++ elif name.endswith("nemo_tokenizer"): ++ # NEMO_HOME/hf_org/hf_model/nemo_tokenizer => hf_org--hf_model ++ tokenizer_model_name = "--".join(name.split("/")[-3:-1]) ++ else: ++ # hf_org/hf_model => hf_org--hf_model ++ tokenizer_model_name = name.replace("/", "--") ++ return tokenizer_model_name ++ else: ++ return f"unknown_tokenizer_{hash(self.tokenizer)}" +diff --git a/megatron/bridge/data/builders/hf_dataset.py b/megatron/bridge/data/builders/hf_dataset.py +new file mode 100755 +index 0000000..6b90bd1 +--- /dev/null ++++ b/megatron/bridge/data/builders/hf_dataset.py +@@ -0,0 +1,362 @@ ++# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++ ++import glob ++import json ++import logging ++import os ++import shutil ++from dataclasses import dataclass ++from pathlib import Path ++from typing import Any, Callable, Optional, Protocol, TypedDict, Union, cast ++ ++from datasets import Dataset, DatasetDict, load_dataset ++from tqdm import tqdm ++ ++from megatron.bridge.data.builders.finetuning_dataset import FinetuningDatasetBuilder ++from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs ++from megatron.bridge.data.datasets.sft import get_dataset_root ++from megatron.bridge.training.config import FinetuningDatasetConfig ++from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer ++from megatron.bridge.utils.common_utils import print_rank_0 ++ ++ ++logger = logging.getLogger(__name__) ++ ++ ++class ProcessExampleOutput(TypedDict): ++ """Expected output structure from a `ProcessExampleFn`.""" ++ ++ input: str ++ output: str ++ original_answers: list[str] ++ ++ ++class ProcessExampleFn(Protocol): ++ """Protocol defining the signature for a function that processes a single dataset example.""" ++ ++ def __call__( ++ self, example: dict[str, Any], tokenizer: Optional[MegatronTokenizer] = None ++ ) -> ProcessExampleOutput: ... ++ ++ ++@dataclass(kw_only=True) ++class HFDatasetConfig(FinetuningDatasetConfig): ++ """Configuration specific to using Hugging Face datasets for finetuning. ++ ++ Inherits from FinetuningDatasetConfig and adds HF-specific options. ++ ++ Attributes: ++ dataset_name: Name of the dataset on the Hugging Face Hub. ++ process_example_fn: A callable conforming to ProcessExampleFn protocol ++ to process raw examples into the desired format. ++ dataset_subset: Optional subset name if the dataset has multiple subsets. ++ dataset_dict: Optional pre-loaded DatasetDict to use instead of downloading. ++ split: Optional specific split to load (e.g., 'train[:10%]'). ++ download_mode: Download mode for load_dataset (e.g., 'force_redownload'). ++ val_proportion: Proportion of the training set to use for validation if ++ no validation set is present. ++ split_val_from_train: If True, creates validation set from training set. ++ If False, uses test set to create validation set. ++ delete_raw: If True, delete the raw downloaded dataset files after processing. ++ rewrite: If True, rewrite existing processed files. ++ hf_kwargs: Additional keyword arguments to pass to `load_dataset`. ++ hf_filter_lambda: Optional function to filter the loaded dataset. ++ hf_filter_lambda_kwargs: Optional keyword arguments for `hf_filter_lambda`. ++ """ ++ ++ dataset_name: str ++ process_example_fn: ProcessExampleFn ++ dataset_subset: Optional[str] = None ++ dataset_dict: Optional[DatasetDict] = None ++ split: Optional[str] = None ++ download_mode: Optional[str] = None ++ val_proportion: Optional[float] = 0.05 ++ split_val_from_train: bool = True ++ delete_raw: bool = False ++ rewrite: bool = True ++ hf_kwargs: Optional[dict[str, Any]] = None ++ hf_filter_lambda: Optional[Callable] = None ++ hf_filter_lambda_kwargs: Optional[dict[str, Any]] = None ++ ++ ++def preprocess_and_split_data( ++ dset: DatasetDict, ++ dataset_name: str, ++ dataset_root: Path, ++ tokenizer: MegatronTokenizer, ++ process_example_fn: ProcessExampleFn, ++ split_val_from_train: bool = True, ++ val_proportion: Optional[float] = None, ++ train_aliases: tuple[str] = ("train", "training"), ++ test_aliases: tuple[str] = ("test", "testing"), ++ val_aliases: tuple[str] = ("val", "validation", "valid", "eval"), ++ delete_raw: bool = False, ++ seed: int = 1234, ++ rewrite: bool = False, ++ do_test: bool = True, ++ do_validation: bool = True, ++): ++ """Download, preprocess, split, and save a Hugging Face dataset to JSONL files. ++ ++ Handles splitting into train/validation/test sets based on available splits ++ and the `val_proportion` parameter. Processes each example using the ++ provided `process_example_fn` and saves the results. ++ ++ Args: ++ dset: The loaded Hugging Face DatasetDict. ++ dataset_name: Name of the dataset (for logging). ++ dataset_root: The root directory to save the processed JSONL files. ++ tokenizer: The tokenizer instance. ++ process_example_fn: Function to process individual examples. ++ split_val_from_train: If True, split validation from train set. ++ Otherwise, split from test set (if available). ++ val_proportion: Proportion of data to use for validation split. ++ train_aliases: Tuple of possible names for the training split. ++ test_aliases: Tuple of possible names for the test split. ++ val_aliases: Tuple of possible names for the validation split. ++ delete_raw: If True, delete raw HF dataset cache after processing. ++ seed: Random seed for splitting. ++ rewrite: If True, overwrite existing processed files. ++ """ ++ logger.info(f"Preprocessing {dataset_name} to jsonl format and splitting...") ++ save_splits = {} ++ train_set: Dataset | None = None ++ val_set: Dataset | None = None ++ test_set: Dataset | None = None ++ ++ for alias in train_aliases: ++ train_set = dset.get(alias) ++ if train_set is not None: ++ break ++ ++ if do_validation: ++ for alias in val_aliases: ++ val_set = dset.get(alias) ++ if val_set is not None: ++ break ++ ++ if do_test: ++ for alias in test_aliases: ++ test_set = dset.get(alias) ++ if test_set is not None: ++ break ++ ++ assert train_set, f"Train set with aliases: {train_aliases} not found in dataset" ++ train_set = cast(Dataset, train_set) ++ ++ if val_proportion: ++ if split_val_from_train: ++ split_dataset = train_set.train_test_split(test_size=val_proportion, seed=seed) ++ save_splits["training"] = split_dataset["train"] ++ save_splits["validation"] = split_dataset["test"] ++ if val_set: ++ save_splits["test"] = val_set ++ else: ++ assert val_set, f"Validation set with aliases: {val_aliases} not found in dataset" ++ val_set = cast(Dataset, val_set) ++ split_dataset = val_set.train_test_split(test_size=val_proportion, seed=seed) ++ save_splits["training"] = train_set ++ save_splits["validation"] = split_dataset["test"] ++ save_splits["test"] = split_dataset["train"] ++ else: ++ save_splits["training"] = train_set ++ if val_set: ++ save_splits["validation"] = val_set ++ if test_set: ++ save_splits["test"] = test_set ++ ++ if test_set: ++ test_set = cast(Dataset, test_set) ++ save_splits["test"] = test_set ++ ++ for split_name, dataset in save_splits.items(): ++ output_file = dataset_root / f"{split_name}.jsonl" ++ ++ if output_file.exists() and output_file.is_file(): ++ if not rewrite: ++ logger.info(f"{output_file} exists, skipping...") ++ continue ++ else: ++ logger.info(f"{output_file} exists, deleting and rewriting...") ++ os.remove(output_file) ++ for p in glob.glob(str(output_file) + "*"): ++ if os.path.exists(p): ++ os.remove(p) ++ ++ with output_file.open("w", encoding="utf-8") as f: ++ for example in tqdm(dataset, desc=f"Processing {split_name} split"): ++ json_line = {} ++ ++ processed_example = process_example_fn(example, tokenizer) ++ # Write each example as a JSON line in the output file ++ json_line["input"] = processed_example["input"] ++ json_line["output"] = processed_example["output"] ++ if split_name == "test": ++ json_line["original_answers"] = processed_example["original_answers"] ++ f.write(json.dumps(json_line) + "\n") ++ ++ logger.info(f"{split_name} split saved to {output_file}") ++ ++ if delete_raw: ++ for p in dataset_root.iterdir(): ++ if p.is_dir(): ++ shutil.rmtree(p) ++ elif ".jsonl" not in str(p.name): ++ p.unlink() ++ ++ ++class HFDatasetBuilder(FinetuningDatasetBuilder): ++ """Builder class for Hugging Face datasets. ++ ++ This class extends FinetuningDatasetBuilder to work with Hugging Face datasets instead of file paths. ++ It provides methods to build datasets from Hugging Face's datasets library. ++ """ ++ ++ def __init__( ++ self, ++ dataset_name: str, ++ tokenizer, ++ process_example_fn: ProcessExampleFn, ++ dataset_dict: Optional[DatasetDict] = None, ++ dataset_subset: Optional[str] = None, ++ dataset_root: Optional[Union[str, Path]] = None, ++ split=None, ++ seq_length=1024, ++ seed: int = 1234, ++ memmap_workers: int = 1, ++ max_train_samples: Optional[int] = None, ++ packed_sequence_specs: Optional[PackedSequenceSpecs] = None, ++ download_mode: Optional[str] = None, ++ val_proportion: Optional[float] = 0.05, ++ split_val_from_train: bool = True, ++ rewrite: bool = True, ++ delete_raw: bool = False, ++ hf_kwargs: Optional[dict[str, Any]] = None, ++ dataset_kwargs: Optional[dict[str, Any]] = None, ++ hf_filter_lambda: Optional[Callable] = None, ++ hf_filter_lambda_kwargs: Optional[dict[str, Any]] = None, ++ do_validation: bool = True, ++ do_test: bool = True, ++ ) -> None: ++ """Initializes the HFDatasetBuilder. ++ ++ Args: ++ dataset_name: Name of the dataset on Hugging Face Hub. ++ tokenizer: The tokenizer instance. ++ is_built_on_rank: Callable to determine if data should be built on the current rank. ++ process_example_fn: Function conforming to ProcessExampleFn protocol. ++ dataset_dict: Optional pre-loaded DatasetDict. ++ dataset_subset: Optional dataset subset name. ++ dataset_root: Optional root directory for data; defaults based on dataset_name. ++ split: Optional specific split to load. ++ seq_length: Sequence length for processing. ++ seed: Random seed. ++ memmap_workers: Number of workers for memmapping. ++ max_train_samples: Optional maximum number of training samples. ++ packed_sequence_specs: Optional PackedSequenceSpecs for packed sequence datasets. ++ download_mode: Download mode for `load_dataset`. ++ val_proportion: Proportion for validation split. ++ split_val_from_train: Whether to split validation from train set. ++ rewrite: Whether to rewrite existing processed files. ++ delete_raw: Whether to delete raw downloaded files. ++ hf_kwargs: Additional kwargs for `load_dataset`. ++ dataset_kwargs: Additional kwargs for the underlying dataset constructor. ++ hf_filter_lambda: Optional function to filter the dataset. ++ hf_filter_lambda_kwargs: Optional kwargs for the filter function. ++ do_validation: Whether to build the validation set. ++ do_test: Whether to build the test set. ++ """ ++ dataset_root = Path(dataset_root) if dataset_root else get_dataset_root(dataset_name) ++ ++ # Initialize the parent class with common parameters ++ super().__init__( ++ dataset_root=dataset_root, ++ tokenizer=tokenizer, ++ seq_length=seq_length, ++ seed=seed, ++ memmap_workers=memmap_workers, ++ dataset_kwargs=dataset_kwargs, ++ max_train_samples=max_train_samples, ++ packed_sequence_specs=packed_sequence_specs, ++ do_validation=do_validation, ++ do_test=do_test, ++ ) ++ ++ # HF-specific attributes ++ self.dataset_name = dataset_name ++ self.dataset_subset = dataset_subset ++ self.dataset_dict = dataset_dict ++ self.split = split ++ self.download_mode = download_mode ++ self.hf_kwargs = hf_kwargs or {} ++ self.val_proportion = val_proportion ++ self.split_val_from_train = split_val_from_train ++ self.delete_raw = delete_raw ++ self.process_example_fn = process_example_fn ++ self.hf_filter_lambda = hf_filter_lambda ++ self.hf_filter_lambda_kwargs = hf_filter_lambda_kwargs or {} ++ self.rewrite = rewrite ++ ++ if not val_proportion: ++ self.do_validation = False ++ self.do_test = False ++ ++ print_rank_0(f"Building HFDataset {self.dataset_name}") ++ ++ def prepare_data(self) -> None: ++ """Loads/downloads the dataset, filters it, preprocesses/splits it, and prepares memmaps.""" ++ if self.download_mode != "force_redownload" and self.hf_filter_lambda: ++ raise ValueError("`hf_filter_lambda` is not supported when `download_mode` is not `force_redownload`") ++ ++ if self.dataset_dict: ++ dataset = self.dataset_dict ++ else: ++ dataset = self._load_dataset() ++ ++ if self.hf_filter_lambda: ++ dataset = dataset.filter(self.hf_filter_lambda, **self.hf_filter_lambda_kwargs) ++ ++ preprocess_and_split_data( ++ dataset, ++ self.dataset_name, ++ self.dataset_root, ++ tokenizer=self.tokenizer, ++ process_example_fn=self.process_example_fn, ++ split_val_from_train=self.split_val_from_train, ++ val_proportion=self.val_proportion, ++ delete_raw=self.delete_raw, ++ seed=self.seed, ++ rewrite=self.rewrite, ++ do_test=self.do_test, ++ do_validation=self.do_validation, ++ ) ++ super().prepare_data() ++ ++ def _load_dataset(self) -> DatasetDict: ++ """Load the dataset from Hugging Face or use the provided dataset.""" ++ if isinstance(self.dataset_name, str): ++ logger.info(f"Loading HF dataset from {self.dataset_name} to {self.dataset_root}") ++ dataset = load_dataset( ++ self.dataset_name, ++ name=self.dataset_subset, ++ cache_dir=str(self.dataset_root), ++ split=self.split, ++ **self.hf_kwargs, ++ download_mode=self.download_mode, ++ ) ++ else: ++ raise ValueError("Expected `dataset_name` to be str, got " + str(type(self.dataset_name))) ++ ++ return dataset +diff --git a/megatron/bridge/data/datasets/__init__.py b/megatron/bridge/data/datasets/__init__.py +new file mode 100755 +index 0000000..e69de29 +diff --git a/megatron/bridge/data/datasets/fim_dataset.py b/megatron/bridge/data/datasets/fim_dataset.py +new file mode 100755 +index 0000000..1c197d1 +--- /dev/null ++++ b/megatron/bridge/data/datasets/fim_dataset.py +@@ -0,0 +1,286 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++ ++import logging ++from typing import Tuple ++ ++import numpy as np ++from megatron.core.datasets.gpt_dataset import GPTDataset ++from megatron.core.datasets.indexed_dataset import IndexedDataset ++from megatron.core.datasets.utils import Split ++ ++from megatron.bridge.training.config import GPTFIMDatasetConfig ++ ++ ++logger = logging.getLogger(__name__) ++ ++ ++class GPTFIMDataset(GPTDataset): ++ """FIM (Fill In The Middle) GPT Dataset ++ ++ Args: ++ indexed_dataset (IndexedDataset): The IndexedDataset around which to build the ++ MegatronDataset ++ ++ indexed_indices (np.ndarray): The set of the documents indices to expose ++ ++ num_samples (int): The number of samples to draw from the indexed dataset ++ ++ index_split (Split): The indexed_indices Split ++ ++ config (GPTFIMDatasetConfig): The GPT-specific container for all config sourced parameters ++ """ ++ ++ def __init__( ++ self, ++ indexed_dataset: IndexedDataset, ++ dataset_path: str, ++ indexed_indices: np.ndarray, ++ num_samples: int, ++ index_split: Split, ++ config: GPTFIMDatasetConfig, ++ ) -> None: ++ super().__init__(indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config) ++ ++ self.np_rng = np.random.RandomState(seed=self.config.random_seed) ++ logger.info(f"Initialized FIM RNG with seed = {self.config.random_seed}") ++ # get FIM params ++ self.fim_rate = self.config.fim_rate ++ self.fim_spm_rate = self.config.fim_spm_rate ++ self.fragment_fim_rate = self.config.fim_fragment_rate ++ fim_split_sample = self.config.fim_split_sample ++ self.no_fim_prefix = self.config.fim_no_prefix ++ if fim_split_sample: ++ fim_split_sample_ids = self.config.tokenizer._tokenizer.tokens_to_ids(fim_split_sample) ++ assert isinstance(fim_split_sample_ids, int) or len(fim_split_sample_ids) == 1 ++ self.fim_split_sample = ( ++ fim_split_sample_ids if isinstance(fim_split_sample_ids, int) else fim_split_sample_ids[0] ++ ) ++ else: ++ self.fim_split_sample = None ++ ++ # get extra tokens ids ++ fim_tokens = self.config.fim_extra_tokens ++ fim_tokens = [ ++ fim_tokens["prefix"], ++ fim_tokens["middle"], ++ fim_tokens["suffix"], ++ fim_tokens["pad"], ++ fim_tokens["eod"], ++ ] ++ fim_tokens_ids = self.config.tokenizer._tokenizer.tokens_to_ids(fim_tokens) ++ ( ++ self.prefix_tok_id, ++ self.middle_tok_id, ++ self.suffix_tok_id, ++ self.pad_tok_id, ++ self.eod_tok_id, ++ ) = fim_tokens_ids ++ ++ def _query_document_sample_shuffle_indices(self, idx: int) -> Tuple[np.ndarray, np.ndarray]: ++ """Get the text (token ids) and document ids for a given index ++ ++ Args: ++ idx (int): The index into the dataset ++ ++ Returns: ++ Tuple[np.ndarray, np.ndarray]: The text ids and document ids ++ """ ++ # Do the shuffle mapping ++ idx = self.shuffle_index[idx] ++ ++ # Get the beginning and end documents and offsets ++ doc_index_beg, doc_index_beg_offset = self.sample_index[idx] ++ doc_index_end, doc_index_end_offset = self.sample_index[idx + 1] ++ ++ document_ids = [] ++ sample_parts = [] ++ ++ # Sample spans a single document ++ if doc_index_beg == doc_index_end: ++ # Add the document id ++ document_ids.append(self.document_index[doc_index_beg]) ++ ++ # Add the entire sample ++ sample_parts.append( ++ self.dataset.get( ++ self.document_index[doc_index_beg], ++ offset=doc_index_beg_offset, ++ length=doc_index_end_offset - doc_index_beg_offset + 1, ++ ) ++ ) ++ ++ # Sample spans multiple documents ++ else: ++ for i in range(doc_index_beg, doc_index_end + 1): ++ # Add the document id ++ document_ids.append(self.document_index[i]) ++ ++ # Add the sample part ++ offset = 0 if i > doc_index_beg else doc_index_beg_offset ++ length = None if i < doc_index_end else doc_index_end_offset + 1 ++ sample_parts.append(self.dataset.get(self.document_index[i], offset=offset, length=length)) ++ ++ sample = np.concatenate(sample_parts) ++ ++ sample_len = sample.shape[0] ++ segment_breaks = np.argwhere(sample == self.eod_tok_id) ++ ++ if segment_breaks.shape != (0, 1): # then there is an EOD token in this example ++ curr_start_position = 0 ++ new_samples = [] ++ for loc in np.nditer(segment_breaks): ++ # Only permute non-empty segments. ++ if loc - curr_start_position > 0: ++ # permute {prefix, suffix, middle} or {suffix, prefix, middle} ++ permuted = self._fim_split_and_permute_sequence(sample[curr_start_position:loc]) ++ new_samples += [permuted, [self.eod_tok_id]] ++ ++ curr_start_position = loc + 1 # jump over the EOD token ++ # Permute the segment after the last EOD ++ permuted = self._fim_split_and_permute_sequence(sample[curr_start_position:]) ++ new_samples.append(permuted) ++ ++ sample = np.concatenate(new_samples) ++ else: ++ sample = self._fim_split_and_permute_sequence(sample) ++ ++ diff = sample.shape[0] - sample_len ++ if diff > 0: # too long ++ sample = sample[:sample_len] ++ elif diff < 0: # too short ++ sample = np.concatenate([sample, np.full((-1 * diff), self.pad_tok_id)]) ++ ++ assert sample.shape[0] == sample_len ++ ++ return (np.array(sample, dtype=np.int64), np.array(document_ids, dtype=np.int64)) ++ ++ def _fim_permute_sequence(self, sequence, rate): ++ return self._permute( ++ sequence, ++ rate, ++ self.fim_spm_rate, ++ self.config.tokenizer, ++ truncate_or_pad=False, ++ suffix_tok_id=self.suffix_tok_id, ++ prefix_tok_id=self.prefix_tok_id, ++ middle_tok_id=self.middle_tok_id, ++ pad_tok_id=self.pad_tok_id, ++ no_fim_prefix=self.no_fim_prefix, ++ ) ++ ++ def _fim_split_and_permute_sequence(self, sequence): ++ """ ++ If self.fim_split_sample is not None, split the sequence. ++ Then apply FIM on the fragments, or the whole sequence if self.fim_split_sample is None. ++ """ ++ if self.fim_split_sample is None: ++ return self._fim_permute_sequence(sequence, self.fim_rate) ++ # fim_split_sample is set: split the sample on this token and permute each fragment separately. ++ # Typically, if each sample is a repository, then we split again on the file level. ++ # Each fragment is a file, and we permute the files. ++ fragment_breaks = np.argwhere(sequence == self.fim_split_sample) ++ if fragment_breaks.shape == (0, 1): ++ # no split token in this sample ++ return self._fim_permute_sequence(sequence, self.fim_rate) ++ if not self.np_rng.binomial(1, self.fim_rate): ++ # don't do FIM preproc ++ return sequence ++ # Do FIM on each fragment ++ curr_start_position = 0 ++ new_samples = [] ++ for loc in np.nditer(fragment_breaks): ++ if loc - curr_start_position > 0: ++ permuted = self._fim_permute_sequence(sequence[curr_start_position:loc], self.fragment_fim_rate) ++ new_samples += [permuted, [self.fim_split_sample]] ++ curr_start_position = loc + 1 # Jump over the split token ++ # Permute the segment after the last split token ++ permuted = self._fim_permute_sequence(sequence[curr_start_position:], self.fragment_fim_rate) ++ new_samples.append(permuted) ++ ++ return np.concatenate(new_samples) ++ ++ def _permute( ++ self, ++ sample, ++ fim_rate, ++ fim_spm_rate, ++ tokenizer, ++ truncate_or_pad=True, ++ suffix_tok_id=None, ++ prefix_tok_id=None, ++ middle_tok_id=None, ++ pad_tok_id=None, ++ no_fim_prefix=None, ++ ): ++ """ ++ Take in a sample (np array w/ size (0,chunklength)) and perform a FIM transformation on it. ++ Maintain the same sample length (if transform creates a few extra tokens, drop them). ++ """ ++ if self.np_rng.binomial(1, fim_rate): # sample bernoulli dist ++ contents = tokenizer._tokenizer.ids_to_text(sample) ++ ++ # Do not apply FIM if the sample starts with no_fim_prefix ++ if no_fim_prefix is not None and contents.startswith(no_fim_prefix): ++ return sample ++ ++ try: ++ # A boundary can be =0 (prefix will be empty) ++ # a boundary can be =len(contents) (suffix will be empty) ++ # The two boundaries can be equal (middle will be empty) ++ boundaries = list(self.np_rng.randint(low=0, high=len(contents) + 1, size=2)) ++ boundaries.sort() ++ except ValueError as e: ++ print(len(contents), contents) ++ print(e) ++ raise e ++ ++ prefix = contents[: boundaries[0]] ++ middle = contents[boundaries[0] : boundaries[1]] ++ suffix = contents[boundaries[1] :] ++ ++ prefix = np.array([*tokenizer._tokenizer.text_to_ids(prefix)], dtype=np.int64) ++ middle = np.array([*tokenizer._tokenizer.text_to_ids(middle)], dtype=np.int64) ++ suffix = np.array([*tokenizer._tokenizer.text_to_ids(suffix)], dtype=np.int64) ++ ++ # here we truncate each given segment to fit the same length as it was before ++ # A consequence is that we never reach the end of a file? ++ # we should rather truncate at the context-level ++ if truncate_or_pad: ++ # need to make same length as the input. Take the 3 sentinel tokens into account ++ new_length = suffix.shape[0] + prefix.shape[0] + middle.shape[0] + 3 ++ diff = new_length - sample.shape[0] ++ if diff > 0: # too long ++ if ( ++ suffix.shape[0] <= diff ++ ): # if there's no space to truncate the suffix: stop and report it. atm i should have stopped this from happening ++ return sample ++ suffix = suffix[: suffix.shape[0] - diff] ++ elif diff < 0: # too short ++ suffix = np.concatenate([suffix, np.full((-1 * diff), pad_tok_id)]) ++ ++ if self.np_rng.binomial(1, fim_spm_rate): ++ # SPM (variant 2 from FIM paper) ++ new_sample = np.concatenate([[prefix_tok_id, suffix_tok_id], suffix, [middle_tok_id], prefix, middle]) ++ else: ++ # PSM ++ new_sample = np.concatenate( ++ [[prefix_tok_id], prefix, [suffix_tok_id], suffix, [middle_tok_id], middle] ++ ) ++ ++ else: ++ # don't do FIM preproc ++ new_sample = sample ++ ++ return new_sample +diff --git a/megatron/bridge/data/datasets/packed_sequence.py b/megatron/bridge/data/datasets/packed_sequence.py +new file mode 100755 +index 0000000..07c1a9c +--- /dev/null ++++ b/megatron/bridge/data/datasets/packed_sequence.py +@@ -0,0 +1,307 @@ ++# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++import json ++import logging ++import multiprocessing as mp ++from dataclasses import dataclass ++from multiprocessing import Pool ++from pathlib import Path ++ ++import numpy as np ++from megatron.core.msc_utils import MultiStorageClientFeature ++from tqdm import tqdm ++ ++from megatron.bridge.data.datasets.packing_utils import create_hist, create_packing_strategy, fill_packing_strategy ++from megatron.bridge.data.datasets.sft import create_sft_dataset ++from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer ++ ++ ++logger = logging.getLogger(__name__) ++ ++_shared_dataset = None ++ ++ ++def _tokenize_get_item(i): ++ return _shared_dataset[i] ++ ++ ++def _tokenize_init_worker(dataset): ++ global _shared_dataset ++ _shared_dataset = dataset ++ ++ ++def _retrieve_tokenized(dataset, num_workers): ++ if num_workers == 1: ++ return np.array([dataset[i] for i in tqdm(range(len(dataset)))]) ++ num_workers = num_workers if num_workers > 0 else mp.cpu_count() ++ with Pool(num_workers, initializer=_tokenize_init_worker, initargs=(dataset,)) as pool: ++ return np.array(list(tqdm(pool.imap(_tokenize_get_item, range(len(dataset))), total=len(dataset)))) ++ ++ ++def tokenize_dataset( ++ path: Path, ++ tokenizer: MegatronTokenizer, ++ max_seq_length: int, ++ seed: int, ++ dataset_kwargs: dict | None = None, ++ pad_seq_to_mult: int | None = 1, ++ num_tokenizer_workers: int = -1, ++): ++ """ ++ Tokenizes a dataset from the provided path using the specified tokenizer ++ and prepares it for further processing. ++ ++ Args: ++ path (Path): Path to the dataset file. ++ tokenizer (MegatronTokenizer): The tokenizer to use for tokenization. ++ max_seq_length (int): Maximum sequence length for the tokens. ++ seed (int): Random seed for shuffling the dataset. ++ dataset_kwargs (dict | None): Additional keyword arguments to pass to create_sft_dataset. ++ Can include 'chat', 'use_hf_tokenizer_chat_template', 'tool_schemas', etc. ++ pad_seq_to_mult (int | None): Optional multiple to pad each sequence to during packing ++ preparation (e.g., set to 2 * context_parallel_size for THD CP). ++ ++ Returns: ++ np.ndarray: A NumPy array containing the tokenized data. ++ """ ++ if not dataset_kwargs: ++ dataset_kwargs = {} ++ ++ # Handle tool_schemas - convert to JSON string if needed ++ ts = dataset_kwargs.get("tool_schemas") ++ if ts and not isinstance(ts, str): ++ dataset_kwargs["tool_schemas"] = json.dumps(ts) ++ ++ # Handle chat_template - set it on tokenizer if provided ++ chat_template = dataset_kwargs.pop("chat_template", None) ++ if chat_template: ++ # This is called during packing preparation (rank 0 only). ++ # The chat template is only needed to create the packed .npy files. ++ # Once created, all ranks load the pre-tokenized .npy files. ++ if hasattr(tokenizer, "_tokenizer"): ++ tokenizer._tokenizer.chat_template = chat_template ++ ++ if pad_seq_to_mult is not None and pad_seq_to_mult <= 0: ++ raise ValueError("pad_seq_to_mult must be a positive integer when provided.") ++ ++ # Keep the historical minimum of 16 unless a larger multiple is requested. ++ pad_seq_length_to_mult = 1 if pad_seq_to_mult is None else max(1, pad_seq_to_mult) ++ ++ dataset = create_sft_dataset( ++ path=path, ++ tokenizer=tokenizer, ++ seq_length=max_seq_length, ++ seed=seed, ++ is_test=True, ++ pad_seq_length_to_mult=pad_seq_length_to_mult, ++ **dataset_kwargs, ++ ) ++ ++ pad_id = dataset.tokenizer.eod ++ pad_seq_length_to_mult = dataset.pad_seq_length_to_mult ++ max_seq_length = dataset.max_seq_length ++ dataset = _retrieve_tokenized(dataset, num_tokenizer_workers) ++ ++ if pad_seq_to_mult > 1: ++ ++ def pre_pad_dataset(data, max_seq_length, max_length_to_pad, pad_id): ++ """ ++ Pad each individual data point to the length of max_length_to_pad. ++ This keeps packed samples divisible by the requested multiple (used for CP/THD). ++ """ ++ assert max_seq_length >= max_length_to_pad ++ for key, val in data.items(): ++ if key in {"input_ids", "context_ids"}: ++ if len(val) <= max_length_to_pad: ++ # input_ids are truncated by 1 for labels; add 1 extra pad token ++ val = val + [pad_id] * (max_length_to_pad - len(val) + 1) ++ elif len(val) > max_seq_length: ++ logging.info( ++ "Sequence length %d is larger than max_seq_length %d; truncating for packing.", ++ len(val), ++ max_seq_length, ++ ) ++ val = val[:max_seq_length] ++ data[key] = val ++ return ++ ++ ceil_to_nearest = lambda n, m: (n + m - 1) // m * m ++ for data in dataset: ++ max_length_to_pad = min(max_seq_length, ceil_to_nearest(len(data["input_ids"]), pad_seq_length_to_mult)) ++ pre_pad_dataset(data, max_seq_length, max_length_to_pad, pad_id) ++ ++ return dataset ++ ++ ++def prepare_packed_sequence_data( ++ input_path: Path, ++ output_path: Path, ++ output_metadata_path: Path, ++ packed_sequence_size: int, ++ tokenizer: MegatronTokenizer, ++ max_seq_length: int, ++ seed: int | None = 0, ++ packing_algorithm: str = "first_fit_shuffle", ++ dataset_kwargs: dict | None = None, ++ pad_seq_to_mult: int | None = 1, ++ num_tokenizer_workers: int = -1, ++): ++ """ ++ Prepares a packed sequence dataset from a given input file and saves it to an output file. ++ ++ Args: ++ input_path (Path): Path to the input dataset file. ++ output_path (Path): Path to save the packed sequence data. ++ output_metadata_path (Path): Path to save packing metadata. ++ packed_sequence_size (int): The maximum size for each packed sequence. ++ tokenizer (MegatronTokenizer): The tokenizer to use for tokenization. ++ max_seq_length (int): Maximum sequence length for the tokens. ++ seed (int | None): Random seed for shuffling (optional). ++ packing_algorithm (str): The algorithm used for packing sequences ++ currently supports "first_fit_shuffle" and "first_fit_decreasing". ++ dataset_kwargs (dict | None): Additional keyword arguments to pass to create_sft_dataset. ++ Enables packing with chat templates, tool schemas, etc. ++ pad_seq_to_mult (int | None): Optional multiple to pad each sequence to during packing ++ preparation (e.g., set to 2 * context_parallel_size for THD CP). ++ ++ Returns: ++ None: Saves the packed sequence data to the specified output path. ++ """ ++ logger.info(f"Preparing packed sequence from {input_path}") ++ dataset = tokenize_dataset( ++ input_path, ++ tokenizer, ++ max_seq_length, ++ seed, ++ dataset_kwargs, ++ pad_seq_to_mult=pad_seq_to_mult, ++ num_tokenizer_workers=num_tokenizer_workers, ++ ) ++ sequences, histogram = create_hist(dataset, max_seq_length) ++ ++ assignments, packing_metadata = create_packing_strategy(histogram, packed_sequence_size, packing_algorithm) ++ output_data = fill_packing_strategy(assignments, sequences, packed_sequence_size, tokenizer.eos_id) ++ ++ # save output data ++ if MultiStorageClientFeature.is_enabled(): ++ msc = MultiStorageClientFeature.import_package() ++ msc.numpy.save(output_path, output_data) ++ else: ++ np.save(output_path, output_data) ++ ++ # save packing metadata, packing_metadata is appended to the packing file if it exists ++ if output_metadata_path is not None: ++ try: ++ with output_metadata_path.open(mode="r") as f: ++ packing_metadata_file = json.load(f) ++ # 'packing_metadata_file' is expected to be a list of dicts: List[Dict[str, int]] ++ # Each dict corresponds to a packed dataset. Typically there will be two dicts, ++ # one each for the packed val and train datasets. ++ # Each dict records two values: 'max_samples_per_bin', the max ++ # number of samples per packed sequence, and 'dataset_max_seqlen', the max ++ # sequence length per sample in the packed dataset. ++ assert isinstance(packing_metadata_file, list), "invalid packing_metadata_file!" ++ except FileNotFoundError: ++ packing_metadata_file = [] ++ ++ packing_metadata_file.append(packing_metadata) ++ with output_metadata_path.open(mode="w") as f: ++ json.dump(packing_metadata_file, f) ++ ++ logger.info(f"Packed sequence is prepared and saved to {output_path}") ++ ++ ++@dataclass ++class PackedSequenceSpecs: ++ """ ++ Configuration class for packed sequence datasets. ++ ++ This class holds parameters related to sequence packing, including the size of the packed sequences, ++ tokenizer information, paths to packed data files, and other related settings. ++ """ ++ ++ packed_sequence_size: int = -1 ++ """ ++ If a positive integer, this arg enables training with sequence packing and specifies the pack size ++ If less than or equal to 0, sequence packing is disabled. Defaults to -1. ++ Note: This arg is distinct from `seq_length` because `seq_length` specifies the maximum length ++ of the original sequence (i.e. the length to truncate long sequences in the input data). ++ """ ++ ++ tokenizer_model_name: str = None ++ """ ++ Keep track of tokenizer model name, since each tokenizer produces a different packed sequence dataset file. ++ This field is set by llm.finetune api. ++ """ ++ ++ num_tokenizer_workers: int = -1 ++ """ ++ The number of worker processes to use for tokenization when preparing the packed sequence dataset. ++ If -1, the number of workers will be set to the number of CPU cores available ++ """ ++ ++ packed_train_data_path: str = None ++ """ ++ If specified, use this file for the packed training dataset instead of the default path. ++ """ ++ ++ packed_val_data_path: str = None ++ """ ++ If specified, use this file for the packed validation dataset instead of the default path. ++ """ ++ ++ packed_metadata_path: str = None ++ """ ++ If specified, use this file for the training and validation packing metadata file instead of the default path. ++ """ ++ ++ pad_cu_seqlens: bool = False ++ """ ++ If True, pad cu_seqlens to a constant size, which is required for use with cudagraphs. ++ """ ++ pad_seq_to_mult: int | None = 1 ++ """ ++ Optional multiple to pad each sample to when generating packed datasets. ++ For THD/context parallel, set to (context_parallel_size * 2) to keep samples divisible. ++ """ ++ ++ def __post_init__(self): ++ if self.packed_train_data_path is not None: ++ if MultiStorageClientFeature.is_enabled(): ++ msc = MultiStorageClientFeature.import_package() ++ self.packed_train_data_path = msc.Path(self.packed_train_data_path) ++ else: ++ self.packed_train_data_path = Path(self.packed_train_data_path) ++ assert self.packed_train_data_path.suffix == ".npy", ( ++ f"packed training data file must be a .npy file: {self.packed_train_data_path}" ++ ) ++ assert self.packed_train_data_path.exists(), ( ++ f"packed training data file does not exist: {self.packed_train_data_path}" ++ ) ++ ++ if self.packed_val_data_path is not None: ++ if MultiStorageClientFeature.is_enabled(): ++ msc = MultiStorageClientFeature.import_package() ++ self.packed_val_data_path = msc.Path(self.packed_val_data_path) ++ else: ++ self.packed_val_data_path = Path(self.packed_val_data_path) ++ assert self.packed_val_data_path.suffix == ".npy", ( ++ f"packed validation data file must be a .npy file: {self.packed_val_data_path}" ++ ) ++ assert self.packed_val_data_path.exists(), ( ++ f"packed validation data file does not exist: {self.packed_val_data_path}" ++ ) ++ ++ if self.pad_seq_to_mult is not None and self.pad_seq_to_mult <= 0: ++ raise ValueError("pad_seq_to_mult must be a positive integer when provided.") +diff --git a/megatron/bridge/data/datasets/packing_utils.py b/megatron/bridge/data/datasets/packing_utils.py +new file mode 100755 +index 0000000..8606eb5 +--- /dev/null ++++ b/megatron/bridge/data/datasets/packing_utils.py +@@ -0,0 +1,356 @@ ++# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++ ++import collections ++import logging ++from typing import Dict, List, Tuple ++ ++import numpy as np ++from tqdm import tqdm ++ ++ ++PACKING_ALGOS = ["first_fit_decreasing", "first_fit_shuffle"] ++ ++logger = logging.getLogger(__name__) ++ ++ ++def find_first_bin_that_fits(bin_sums: List[int], s: int, bin_size: int) -> int: ++ """ ++ Finds the first bin in a list of bins that has enough space to fit a sequence of size 's'. ++ ++ Args: ++ bins: A list of lists, where each inner list represents a bin and contains the current elements in that bin. ++ s: The size of the sequence to be placed in a bin. ++ bin_size: The maximum capacity of each bin. ++ ++ Returns: ++ The index of the first bin that can fit the sequence 's', or -1 if no such bin exists. ++ """ ++ for i, cur_sum in enumerate(bin_sums): ++ if cur_sum + s <= bin_size: ++ return i ++ return -1 ++ ++ ++def first_fit(seqlens: List[int], pack_size: int) -> List[List[int]]: ++ """ ++ Packs sequences of varying lengths into bins using the First-Fit algorithm. ++ ++ Args: ++ seqlens: A list of integers, representing the lengths of the sequences to be packed. ++ pack_size: The maximum capacity of each bin. ++ ++ Returns: ++ A list of lists, where each inner list represents a bin and contains the indices ++ of the sequences assigned to that bin. ++ """ ++ res = [] ++ res_sums = [] ++ for s in tqdm(seqlens): ++ first_bin = find_first_bin_that_fits(res_sums, s, pack_size) ++ if first_bin == -1: # open a new bin ++ res.append([s]) ++ res_sums.append(s) ++ else: ++ res[first_bin].append(s) ++ res_sums[first_bin] += s ++ return res ++ ++ ++def first_fit_decreasing(seqlens: List[int], pack_size: int) -> List[List[int]]: ++ """ ++ Packs sequences of varying lengths into bins using the First-Fit Decreasing algorithm. ++ ++ This is a variation of the First-Fit algorithm where the sequences are sorted by decreasing length before packing. ++ ++ Args: ++ seqlens: A list of integers, representing the lengths of the sequences to be packed. ++ pack_size: The maximum capacity of each bin. ++ ++ Returns: ++ A list of lists, similar to the output of the 'first_fit' function. ++ """ ++ sorted_seqlens = sorted(seqlens, reverse=True) ++ return first_fit(sorted_seqlens, pack_size) ++ ++ ++def first_fit_shuffle(seqlens: List[int], pack_size: int) -> List[List[int]]: ++ """ ++ Packs sequences of varying lengths into bins using the First-Fit with Shuffling algorithm. ++ ++ This variation shuffles the order of the sequences before applying the First-Fit algorithm. ++ ++ Args: ++ seqlens: A list of integers, representing the lengths of the sequences to be packed. ++ pack_size: The maximum capacity of each bin. ++ ++ Returns: ++ A list of lists, similar to the output of the 'first_fit' function. ++ """ ++ shuffled_seqlens = seqlens[:] ++ np.random.shuffle(shuffled_seqlens) ++ return first_fit(shuffled_seqlens, pack_size) ++ ++ ++def create_hist(dataset: np.array, truncate_seq_len: int) -> Tuple[Dict[int, List[Dict]], List[int]]: ++ """ ++ Creates a histogram of sequence lengths from a tokenized dataset. ++ ++ This function analyzes the tokenized dataset and creates a histogram showing the distribution of sequence lengths. ++ ++ Args: ++ dataset: A NumPy array containing the tokenized sequences. Each element is a dictionary that contains at minimum ++ the key `input_ids`. ++ truncate_seq_len: The maximum sequence length to consider in the histogram. ++ ++ Returns: ++ sequences: A dictionary where keys are sequence lengths and values are lists ++ of corresponding sequences from the dataset. ++ histogram: A list representing the histogram data (number of sequences for each length). ++ """ ++ logger.info("Creating histogram from tokenized dataset...") ++ ++ sequences = collections.defaultdict(list) ++ counts = [0] * (truncate_seq_len + 1) ++ ++ for item_dict in dataset: ++ # Minus 1 here to account for the fact that transformer input and label ++ # have one less token than the full sequence. ++ # Input is missing the last token and label is missing the first token ++ # (this way the tokens are aligned for next token prediction). ++ # We want pack size to be the length of the actual input and label, hence minus 1. ++ seq_len = len(item_dict["input_ids"]) - 1 ++ sequences[seq_len].append(item_dict) ++ counts[seq_len] += 1 ++ ++ logger.debug("Histogram of sequence lengths") ++ logger.debug(counts) ++ ++ histogram = [] ++ for seq_len in range(truncate_seq_len + 1): ++ histogram.append(len(sequences[seq_len])) ++ ++ return sequences, histogram ++ ++ ++def create_packing_strategy( ++ histogram: List[int], pack_size: int, packing_algorithm: str = "first_fit" ++) -> Tuple[List[List[int]], Dict[str, int]]: ++ """ ++ Packs sequences into bins using the specified packing algorithm. ++ ++ This function takes the histogram of sequence lengths, desired pack size, and a string representing the packing ++ algorithm to use. It then calls the corresponding function (e.g., 'first_fit_decreasing') and performs the ++ packing process using only sequence lengths as input (without the actual sequences). ++ ++ Args: ++ histogram: A list representing the histogram data (number of sequences for each length). ++ pack_size: The maximum capacity of each bin. ++ packing_algorithm: One of the supported packing algorithms from ['first_fit_decreasing', 'first_fit_shuffle'] ++ ++ Returns: ++ assignments: A list of lists, where each inner list represents a bin and contains the indices of the ++ sequence lengths assigned to that bin. ++ pack_metadata: A dict that records packing metadata, for instance the max number of samples per bin. ++ """ ++ ++ logger.info(f"Packing sequences to length {pack_size}...") ++ ++ all_seq_lens = [] ++ for i, count in enumerate(histogram): ++ all_seq_lens.extend([i] * count) ++ ++ packing_fn = globals()[packing_algorithm] ++ assignments: list[list[int]] = packing_fn(all_seq_lens, pack_size) ++ packed_seq_lens = [sum(x) for x in assignments] ++ packing_factor = len(all_seq_lens) / len(packed_seq_lens) ++ ++ max_seqlen = max(all_seq_lens) ++ max_samples_per_bin = max([len(b) for b in assignments]) ++ min_packed_seqlen = min(packed_seq_lens) ++ packing_efficiency = sum(packed_seq_lens) / len(packed_seq_lens) / pack_size * 100 ++ ++ packing_metadata = { ++ "dataset_max_seqlen": max_seqlen, ++ "max_samples_per_bin": max_samples_per_bin, ++ "packing_factor": round(packing_factor, 2), ++ "packing_efficiency": round(packing_efficiency, 2), ++ "pack_size": pack_size, ++ "min_packed_seqlen": min_packed_seqlen, ++ } ++ ++ logger.debug("Packed sequence lengths:") ++ logger.debug(packed_seq_lens) ++ logger.info(f"Packing is {packing_efficiency:.2f}% efficient") ++ logger.info( ++ f">>>>> For pack size {pack_size}, average number of sequences per pack is n = {packing_factor:.3f} <<<<<" ++ ) ++ return assignments, packing_metadata ++ ++ ++def fill_packing_strategy( ++ assignments: List[List[int]], ++ sequences: Dict[int, List[Dict]], ++ pack_size: int, ++ pad_id: int, ++) -> List[Dict]: ++ """ ++ Fills the packing strategy with actual sequence data based on assignments and sequence information. ++ This function takes the assignments generated by the packing algorithm (containing sequence length indices), ++ the original sequences data, and the pack size. It iterates through the assignments, retrieves the corresponding ++ sequences from the sequences dictionary, and constructs the final output data structure with input IDs, loss masks ++ (if available), and starting indices for each sequence in a packed sequence. ++ Args: ++ assignments: A list of lists, where each inner list represents a bin and contains the indices of the ++ sequence lengths assigned to that bin (output of 'create_packing_strategy'). ++ sequences: A dictionary where keys are sequence lengths and values are lists of corresponding sequences ++ from the dataset (output of 'create_hist'). ++ pack_size: The maximum capacity of each bin. ++ pad_id: The tokenizer's padding token. ++ Returns: ++ output_data: A list of dictionaries, where each dictionary represents a packed sequence with its input IDs, ++ loss mask (if available), and starting indices. ++ """ ++ ifile_handles = dict() ++ for seq_len in tqdm(range(pack_size + 1)): ++ per_seq_data = sequences[seq_len] ++ if len(per_seq_data) > 0: ++ perm = np.random.permutation(len(per_seq_data)) ++ input_ids = np.array([x["input_ids"] for x in per_seq_data])[perm].tolist() ++ try: ++ loss_mask = np.array([x["loss_mask"] for x in per_seq_data])[perm].tolist() ++ # roll loss mask by 1 to align with labels. We want to train on the output after the last context token ++ loss_mask = [x[1:] + [False] for x in loss_mask] ++ except KeyError: ++ try: ++ loss_mask = np.array( ++ [ ++ [ ++ # (x['answer_start_idx'] - 1) because we want to train on the output ++ # after the last context token ++ idx >= (x["answer_start_idx"] - 1) ++ for idx in range(len(x["input_ids"])) ++ ] ++ for x in per_seq_data ++ ] ++ )[perm].tolist() ++ except KeyError as err: ++ err_msg = "Key errors loss_mask and answer_start_idx missing in example - " ++ err_msg += f"{err} {per_seq_data[0]}" ++ logging.error(err_msg) ++ raise ValueError(err_msg) ++ ifile_handles[seq_len] = (input_ids, loss_mask) ++ input_ids, loss_mask, seq_start_id = {}, {}, {} ++ for oindex, assignment in tqdm(enumerate(assignments), total=len(assignments)): ++ _input_ids, _loss_mask, _seq_start_id = [], [], [0] ++ for seq_length in assignment: ++ _input_ids.extend(ifile_handles[seq_length][0].pop()) ++ _loss_mask.extend(ifile_handles[seq_length][1].pop()) ++ _seq_start_id.append(len(_input_ids)) ++ input_ids[oindex] = _input_ids ++ loss_mask[oindex] = _loss_mask ++ seq_start_id[oindex] = _seq_start_id[:-1] ++ output_data = [] ++ for i in range(len(input_ids)): ++ item_dict = { ++ "input_ids": input_ids[i], ++ "loss_mask": loss_mask[i], ++ "seq_start_id": seq_start_id[i], ++ } ++ output_data.append(item_dict) ++ assert all(not seq[0] for seq in ifile_handles.values()), "Error: There are items left over from the assignment" ++ assert all(not seq[1] for seq in ifile_handles.values()), "Error: There are items left over from the assignment" ++ return output_data ++ ++ ++def get_seqlen_list(elem: Dict) -> Tuple[List[int], int]: ++ """Extract per-sequence token counts from a packed dataset element. ++ ++ Args: ++ elem: A packed dataset element with 'input_ids' and 'seq_start_id' fields. ++ ++ Returns: ++ A tuple of (token_counts, tokens_minus_eos) where token_counts is a list of ++ per-sequence token counts (excluding EOS) and tokens_minus_eos is the total ++ token count excluding EOS tokens. ++ """ ++ num_seq = len(elem["seq_start_id"]) ++ tokens_total = len(elem["input_ids"]) ++ tokens_minus_eos = tokens_total - num_seq ++ ++ seq_boundaries = elem["seq_start_id"] + [tokens_total] ++ ++ # subtract 1 to account for removing eos token ++ token_counts = [seq_boundaries[i + 1] - seq_boundaries[i] - 1 for i in range(num_seq)] ++ ++ assert sum(token_counts) == tokens_minus_eos, (sum(token_counts), tokens_minus_eos) ++ ++ return token_counts, tokens_minus_eos ++ ++ ++def calculate_avg_seqlen( ++ dataset_file: str, gbs: int, max_seq_len: int, drop_remainder: bool ++) -> Tuple[float, float, float, float]: ++ """Calculate average sequence length statistics from a packed dataset. ++ ++ Args: ++ dataset_file: Path to the .npy packed dataset file. ++ gbs: Global batch size used to determine how many rows to process. ++ max_seq_len: Maximum sequence length (reserved for future use). ++ drop_remainder: If True, drop rows that don't fill a complete batch. ++ ++ Returns: ++ A tuple of (avg_seqlen_count, avg_seqlen_total, avg_seqlen_sq_individual, avg_seqlen_sq_per_row): ++ - avg_seqlen_count: Average number of sequences per row. ++ - avg_seqlen_total: Average total tokens (excluding EOS) per row. ++ - avg_seqlen_sq_individual: Average of squared per-sequence lengths. ++ - avg_seqlen_sq_per_row: Average of summed squared sequence lengths per row. ++ ++ Raises: ++ ValueError: If no rows remain after applying drop_remainder, or if no sequences are found. ++ """ ++ data = np.load(dataset_file, allow_pickle=True) ++ ++ total_len_accum = 0 ++ seqlen_sq_accum = 0 ++ seq_count_accum = 0 ++ ++ rows_total = len(data) ++ count = (rows_total // gbs) * gbs if drop_remainder else rows_total ++ ++ if count != rows_total: ++ logger.info(f"Dropping {rows_total - count}, total was {rows_total}") ++ ++ for i, elem in enumerate(data): ++ if i >= count: ++ break ++ seqlen_list, total_count = get_seqlen_list(elem) ++ seqlen_sq_list = [s * s for s in seqlen_list] ++ total_len_accum += total_count ++ seqlen_sq_accum += sum(seqlen_sq_list) ++ seq_count_accum += len(seqlen_list) ++ ++ if count == 0: ++ raise ValueError( ++ f"No rows to process: dataset has {rows_total} rows but gbs={gbs} with drop_remainder={drop_remainder}." ++ ) ++ if seq_count_accum == 0: ++ raise ValueError("No sequences found in dataset; cannot compute average sequence length.") ++ ++ avg_seqlen_count = seq_count_accum / count ++ avg_seqlen_total = total_len_accum / count ++ avg_seqlen_sq_individual = seqlen_sq_accum / seq_count_accum ++ avg_seqlen_sq_per_row = seqlen_sq_accum / count ++ ++ return avg_seqlen_count, avg_seqlen_total, avg_seqlen_sq_individual, avg_seqlen_sq_per_row +diff --git a/megatron/bridge/data/datasets/sft.py b/megatron/bridge/data/datasets/sft.py +new file mode 100755 +index 0000000..425c55a +--- /dev/null ++++ b/megatron/bridge/data/datasets/sft.py +@@ -0,0 +1,1233 @@ ++# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++ ++import json ++import logging ++import math ++import os ++import re ++from pathlib import Path ++from typing import Mapping ++ ++import datasets ++import numpy as np ++import torch ++from datasets import load_dataset ++from megatron.core.msc_utils import MultiStorageClientFeature ++from torch.utils.data import Dataset ++ ++from megatron.bridge.data.datasets.utils import ( ++ _chat_preprocess, ++ _get_samples_mapping, ++ _JSONLMemMapDataset, ++ _OnlineSampleMapping, ++ _preprocess, ++ _tokenize, ++) ++from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer ++ ++ ++DEFAULT_NEMO_CACHE_HOME = Path.home() / ".cache" / "nemo" ++NEMO_CACHE_HOME = Path(os.getenv("NEMO_HOME", DEFAULT_NEMO_CACHE_HOME)) ++DEFAULT_NEMO_DATASETS_CACHE = NEMO_CACHE_HOME / "datasets" ++NEMO_DATASETS_CACHE = Path(os.getenv("NEMO_DATASETS_CACHE", DEFAULT_NEMO_DATASETS_CACHE)) ++DEFAULT_NEMO_MODELS_CACHE = NEMO_CACHE_HOME / "models" ++NEMO_MODELS_CACHE = Path(os.getenv("NEMO_MODELS_CACHE", DEFAULT_NEMO_MODELS_CACHE)) ++ ++if os.getenv("TOKENIZERS_PARALLELISM") is None: ++ os.putenv("TOKENIZERS_PARALLELISM", "True") ++ ++logger = logging.getLogger(__name__) ++ ++# hack to avoid the "not enough disk space" error in some slurm cluster ++datasets.builder.has_sufficient_disk_space = lambda needed_bytes, directory=".": True ++ ++PREFIX_STR = ( ++ "\x00" # the prefix string used in the tokenizer to deal with the added empty token for some of the tokenizers ++) ++ ++__idx_version__ = "0.2" # index file version ++__idx_suffix__ = "idx" # index file suffix ++ ++ ++def get_dataset_root(name: str) -> Path: ++ """ ++ Returns the root directory for NeMo datasets, creating it if it doesn't exist. ++ ++ Args: ++ name (str): The name of the dataset, used to create a subdirectory within the NeMo datasets cache. ++ ++ Returns: ++ Path: The path to the dataset's root directory. ++ """ ++ output = Path(NEMO_DATASETS_CACHE) / name ++ output.mkdir(parents=True, exist_ok=True) ++ ++ return output ++ ++ ++def create_sft_dataset( ++ path: Path, ++ tokenizer: "MegatronTokenizer", ++ seq_length: int = 2048, ++ add_bos: bool = False, ++ add_eos: bool = True, ++ add_sep: bool = False, ++ seed: int = 1234, ++ label_key: str = "output", ++ answer_only_loss: bool = True, ++ truncation_field: str = "input", ++ pad_to_max_length: bool = False, ++ index_mapping_dir: str | None = None, ++ prompt_template: str = "{input} {output}", ++ truncation_method: str = "right", ++ memmap_workers: int = 2, ++ hf_dataset: bool = False, ++ global_sample_mapping: bool = False, ++ get_attention_mask_from_fusion: bool = True, ++ pack_metadata_file_path: Path = None, ++ pad_cu_seqlens: bool = False, ++ pad_seq_to_mult: int = 1, ++ chat: bool = False, ++ use_hf_tokenizer_chat_template: bool = False, ++ tool_schemas: str | dict | None = None, ++ **kwargs, ++) -> "GPTSFTDataset": ++ """ ++ Creates and returns a supervised fine-tuning (SFT) dataset instance. ++ ++ This function acts as a factory for different types of SFT datasets based on the ++ input parameters. It can create standard SFT datasets, chat-specific datasets, ++ or packed sequence datasets. ++ ++ Args: ++ path (Path): Path to the dataset file. For packed datasets, this should be a .npy file. ++ tokenizer (MegatronTokenizer): The tokenizer to use for tokenizing the data. ++ seq_length (int, optional): Maximum sequence length for each example. Defaults to 2048. ++ add_bos (bool, optional): Whether to add a beginning-of-sentence token. Defaults to False. ++ add_eos (bool, optional): Whether to add an end-of-sentence token. Defaults to True. ++ add_sep (bool, optional): Whether to add a separation token between prompt and answer. Defaults to False. ++ seed (int, optional): Random seed for data shuffling. Defaults to 1234. ++ label_key (str, optional): The key in the dataset corresponding to the label/output. Defaults to "output". ++ answer_only_loss (bool, optional): If True, compute loss only on the answer part. Defaults to True. ++ truncation_field (str, optional): Field(s) to truncate if the combined length exceeds `seq_length`. ++ Comma-separated if multiple. Defaults to "input". ++ pad_to_max_length (bool, optional): Whether to pad all samples to `max_seq_length`. Defaults to False. ++ index_mapping_dir (str | None, optional): Directory to store/load index mapping files. Defaults to None. ++ prompt_template (str, optional): F-string template for combining input fields. ++ Example: "{input} {output}". Defaults to "{input} {output}". ++ truncation_method (str, optional): Method for truncation ('left' or 'right'). Defaults to "right". ++ memmap_workers (int, optional): Number of workers for memory-mapped dataset loading. Defaults to 2. ++ hf_dataset (bool, optional): Whether to load the dataset using HuggingFace's `datasets` library. ++ Defaults to False. ++ global_sample_mapping (bool, optional): Whether to use a global sample mapping for shuffling across all data, ++ or shuffle within each epoch. Defaults to False. ++ get_attention_mask_from_fusion (bool): if true, lets attention kernel handle creation of causal mask instead ++ of adding it to the batch dict. ++ pack_metadata_file_path (Path, optional): Path to the metadata file for packed datasets. ++ Required if `pad_cu_seqlens` is True. Defaults to None. ++ pad_cu_seqlens (bool, optional): Whether to pad `cu_seqlens` for packed datasets, ++ required for cudagraphs. Defaults to False. ++ chat (bool, optional): If True, creates a `GPTSFTChatDataset`. Defaults to False. ++ use_hf_tokenizer_chat_template (bool, optional): If True, uses HuggingFace tokenizer's chat template ++ via `apply_chat_template` method. Only applies when `chat=True`. Defaults to False. ++ tool_schemas (str | dict | None, optional): Tool schemas for function calling support. ++ Can be a JSON string or a dict. Only applies when `chat=True` and ++ `use_hf_tokenizer_chat_template=True`. Defaults to None. ++ **kwargs: Additional keyword arguments passed to the specific dataset class constructor. ++ ++ Returns: ++ GPTSFTDataset | GPTSFTChatDataset | GPTSFTPackedDataset: An instance of the appropriate SFT dataset class. ++ """ ++ ++ gpt_sft_dataset_kwargs = { ++ "file_path": str(path), ++ "tokenizer": tokenizer, ++ "max_seq_length": seq_length, ++ "memmap_workers": memmap_workers, ++ "hf_dataset": hf_dataset, ++ "global_sample_mapping": global_sample_mapping, ++ "add_bos": add_bos, ++ "add_eos": add_eos, ++ "add_sep": add_sep, ++ "seed": seed, ++ "label_key": label_key, ++ "answer_only_loss": answer_only_loss, ++ "truncation_field": truncation_field, ++ "pad_to_max_length": pad_to_max_length, ++ "index_mapping_dir": index_mapping_dir, ++ "prompt_template": prompt_template, ++ "truncation_method": truncation_method, ++ "get_attention_mask_from_fusion": get_attention_mask_from_fusion, ++ } ++ ++ if path.suffix == ".npy": ++ return GPTSFTPackedDataset( ++ pack_metadata_file_path=pack_metadata_file_path, ++ pad_cu_seqlens=pad_cu_seqlens, ++ pad_seq_to_mult=pad_seq_to_mult, ++ **gpt_sft_dataset_kwargs, ++ **kwargs, ++ ) ++ elif chat: ++ return GPTSFTChatDataset( ++ **gpt_sft_dataset_kwargs, ++ use_hf_tokenizer_chat_template=use_hf_tokenizer_chat_template, ++ tool_schemas=tool_schemas, ++ **kwargs, ++ ) ++ else: ++ return GPTSFTDataset( ++ **gpt_sft_dataset_kwargs, ++ **kwargs, ++ ) ++ ++ ++class GPTSFTDataset(Dataset): ++ """ """ ++ ++ def __init__( ++ self, ++ file_path: str, ++ tokenizer: MegatronTokenizer, ++ max_seq_length: int = 1024, ++ min_seq_length: int = 1, ++ pad_seq_length_to_mult: int = 16, ++ add_bos: bool = False, ++ add_eos: bool = True, ++ add_sep: bool = False, ++ sep_id: int = None, ++ max_num_samples: int = None, ++ seed: int = 1234, ++ label_key: str = "answer", ++ answer_only_loss: bool = True, ++ truncation_field: str = "text", ++ pad_to_max_length: bool = False, # (@adithyare) allows for much faster training especially in PEFT settings. ++ index_mapping_dir: str = None, ++ prompt_template: str = None, ++ virtual_tokens: int = 0, ++ tokens_to_generate: int = 0, ++ memmap_workers: int | None = None, ++ hf_dataset: bool = False, ++ global_sample_mapping: bool = False, ++ truncation_method: str = "right", ++ special_tokens: Mapping[str, str] | None = None, # special tokens, a dictory of {token_type: token} ++ is_test: bool = False, ++ output_original_text: bool = False, ++ ceil_to_power_2: bool = False, ++ get_attention_mask_from_fusion: bool = True, ++ ): ++ """ ++ file_path: Path to a JSONL GPT supervised fine-tuning dataset. ++ Data is formatted as multiple JSON lines with each line formatted as follows: ++ { ++ 'input': 'John von Neumann\nVon Neumann made fundamental contributions ... ++ Q: What did the math of artificial viscosity do?', ++ 'output': 'smoothed the shock transition without sacrificing basic physics' ++ } ++ tokenizer: Tokenizer for the dataset. Instance of a class that inherits MegatronTokenizer (ex: SentencePiece). ++ max_seq_length (int): maximum sequence length for each dataset examples. ++ Examples will either be truncated to fit this length or dropped if they cannot be truncated. ++ min_seq_length (int): min length of each data example in the dataset. ++ Data examples will be dropped if they do not meet the min length requirements. ++ add_bos (bool): Whether to add a beginning of sentence token to each data example ++ add_eos (bool): Whether to add an end of sentence token to each data example ++ add_sep (bool): Whether to add a separation token to each data example (goes between prompt and answer) ++ tokens_to_generate (int): (inference only) Number of tokens to generate during inference ++ seed: Random seed for data shuffling. ++ max_num_samples: Maximum number of samples to load. ++ This can be > dataset length if you want to oversample data. If None, all samples will be loaded. ++ label_key: Key to use for the label in your JSONL file ++ answer_only_loss: If True, will compute the loss only on the answer part of the input. ++ If False, will compute the loss on the entire input. ++ truncation_field: Field to use for truncation. (Options: keys in prompt_template). ++ Field to be used for truncation if the combined length exceeds the max sequence length. ++ pad_to_max_length: Whether to pad the input to the max sequence length. ++ If False, will pad to the max length of the current batch. ++ index_mapping_dir: Directory to save the index mapping to. ++ If None, will write to the same folder as the dataset. ++ prompt_template: Prompt template to inject via an fstring. ++ Formatted like Q: {context_key}\n\nA: {label_key} ++ hf_dataset: Whether to load the json file with the HuggingFace dataset. ++ Otherwise, will load the jsonl file with the JSONLMemMapDataset. ++ global_sample_mapping: Whether to shuffle all data together, or shuffle the dataset within each epoch ++ truncation_method: Truncation from which position. Options: ['left', 'right'] ++ special_tokens: special tokens for the chat prompts, a dictionary of {token_type: token}. ++ Default: { ++ 'system_turn_start': '', ++ 'turn_start': '', ++ 'label_start': '', ++ 'end_of_turn': '\n', ++ 'end_of_name': '\n' ++ } ++ is_test: Whether this dataset is the test split. ++ output_original_text (bool): if true, will keep the original text in the output alongside the tokenized ids. ++ get_attention_mask_from_fusion (bool): if true, lets attention kernel handle creation of causal mask instead ++ of adding it to the batch dict. ++ """ ++ self.tokenizer = tokenizer ++ self.file_path = file_path ++ self.max_seq_length = max_seq_length ++ self.min_seq_length = min_seq_length ++ self.pad_seq_length_to_mult = pad_seq_length_to_mult ++ self.add_bos = add_bos ++ self.add_eos = add_eos ++ self.add_sep = add_sep ++ self.sep_id = sep_id ++ self.max_num_samples = max_num_samples ++ self.seed = seed ++ self.label_key = label_key ++ self.answer_only_loss = answer_only_loss ++ self.truncation_fields = truncation_field.split(",") if truncation_field is not None else [] ++ self.pad_to_max_length = pad_to_max_length ++ self.index_mapping_dir = index_mapping_dir ++ self.prompt_template = prompt_template ++ self.virtual_tokens = virtual_tokens ++ self.tokens_to_generate = tokens_to_generate ++ self.memmap_workers = memmap_workers ++ self.hf_dataset = hf_dataset ++ self.global_sample_mapping = global_sample_mapping ++ self.truncation_method = truncation_method ++ self.is_test = is_test ++ self.output_original_text = output_original_text ++ self.ceil_to_power_2 = ceil_to_power_2 ++ self.get_attention_mask_from_fusion = get_attention_mask_from_fusion ++ ++ if special_tokens is None: ++ self.special_tokens = { ++ "system_turn_start": "", ++ "turn_start": "", ++ "label_start": "", ++ "end_of_turn": "\n", ++ "end_of_name": "\n", ++ } ++ else: ++ self.special_tokens = special_tokens ++ ++ self._load_dataset() ++ ++ # Validate prompt template ++ self._maybe_validate_prompt_template() ++ ++ # Will be None after this call if `max_num_samples` is None ++ self._build_samples_mapping() ++ ++ def _load_dataset(self): ++ if self.hf_dataset: ++ self.indexed_dataset = load_dataset( ++ "json", ++ data_files=self.file_path, ++ cache_dir=self.index_mapping_dir, ++ num_proc=self.memmap_workers, ++ split="train", ++ ) ++ else: ++ self.indexed_dataset = _JSONLMemMapDataset( ++ dataset_paths=[self.file_path], ++ tokenizer=None, ++ header_lines=0, ++ index_mapping_dir=self.index_mapping_dir, ++ workers=self.memmap_workers, ++ ) ++ ++ def _maybe_validate_prompt_template(self): ++ assert self.prompt_template is not None, ( ++ f"we need prompt_template to combine contexts and label {self.label_key}" ++ ) ++ # When providing things like newlines in the prompt template via the CLI, they are escaped. ++ # This line unescapes them. ++ self.prompt_template = self.prompt_template.encode("utf-8").decode("unicode_escape") ++ self.prompt_template_keys = re.findall(r"{(.*?)}", self.prompt_template) ++ ++ label_placeholder = f"{{{self.label_key}}}" ++ assert self.prompt_template[-len(label_placeholder) :] == label_placeholder, ( ++ f"{label_placeholder} must be at the end of prompt_template." ++ ) ++ ++ # Legacy checkpoints has self.truncation_fields = ['context'] ++ # and self.prompt_template_keys = ['input', 'output'] ++ if len(self.truncation_fields) > 0: ++ if self.prompt_template_keys[0] == "input" and self.truncation_fields[0] == "context": ++ self.truncation_fields[0] = self.prompt_template_keys[0] ++ ++ assert set(self.truncation_fields).issubset(self.prompt_template_keys), ( ++ f"truncation_fields {self.truncation_fields} must in {self.prompt_template_keys}" ++ ) ++ ++ def _build_samples_mapping(self): ++ if self.max_num_samples is not None: ++ osm = ( ++ _OnlineSampleMapping(dataset_size=len(self.indexed_dataset), num_samples=self.max_num_samples) ++ if not self.global_sample_mapping ++ else None ++ ) ++ self.samples_mapping = _get_samples_mapping( ++ indexed_dataset=self.indexed_dataset, ++ data_prefix=self.file_path, ++ num_epochs=None, ++ max_num_samples=self.max_num_samples, ++ max_seq_length=self.max_seq_length - 2, ++ short_seq_prob=0, ++ seed=self.seed, ++ name=self.file_path.split("/")[-1], ++ binary_head=False, ++ index_mapping_dir=self.index_mapping_dir, ++ samples_mapping=osm, ++ ) ++ else: ++ self.samples_mapping = None ++ ++ def __len__(self): ++ """Return the total number of samples in this dataset.""" ++ if self.max_num_samples is None: ++ return len(self.indexed_dataset) ++ else: ++ return len(self.samples_mapping) ++ ++ def __getitem__(self, idx): ++ if isinstance(idx, np.int64): ++ idx = idx.item() ++ ++ if self.samples_mapping is not None: ++ assert idx < len(self.samples_mapping) ++ idx, _, _ = self.samples_mapping[idx] ++ if isinstance(idx, (np.uint32, np.int64)): ++ idx = idx.item() ++ ++ assert idx < len(self.indexed_dataset) ++ # idx may < 0 because we pad_samples_to_global_batch_size, e.g. id = -1 ++ if idx < 0: ++ idx = len(self) + idx ++ auto_gen_idx = True ++ else: ++ auto_gen_idx = False ++ try: ++ example = self.indexed_dataset[idx] ++ if auto_gen_idx: ++ example["__AUTOGENERATED__"] = True ++ except Exception as e: ++ logger.error(f"Error while loading example {idx} from dataset {self.file_path}") ++ raise e ++ return self._process_example(example) ++ ++ def _separate_template(self, prompt_template_values: list[str]): ++ """ ++ Combine contexts and label based on prompt_template into a list of strings and a list of keys. ++ ++ Args: ++ prompt_template_values (list[str]): the list of context and label strings ++ extrated from jsonl file with prompt_template_keys. ++ ++ Returns: ++ template_strings (list[str]): separated prompt_template with contexts/label ++ placeholder filled with corresponding strings ++ template_strings_keys (list[str]): strings point to placeholder keys or