Distillflow 0.2.0__py3-none-any.whl

distillflow/common/__init__.py

@@ -0,0 +1,10 @@
+from .common import get_current_device, count_parameters, find_all_linear_modules, infer_optim_dtype
+from .logger import get_logger
+
+__all__ = [
+    get_current_device,
+    count_parameters,
+    find_all_linear_modules,
+    get_logger,
+    infer_optim_dtype
+]

distillflow/common/common.py

@@ -0,0 +1,98 @@
+import os
+from typing import Tuple, List
+
+from transformers import is_torch_xpu_available, is_torch_npu_available, PreTrainedModel
+from transformers.utils import is_torch_mps_available, is_torch_cuda_available, is_torch_bf16_gpu_available
+import torch
+from .logger import get_logger
+
+logger = get_logger(__name__)
+
+def get_current_device() -> torch.device:
+    local_rank = os.environ.get("LOCAL_RANK", "0")
+    device_map = {
+        is_torch_xpu_available: "xpu",
+        is_torch_npu_available: "npu",
+        is_torch_mps_available: "mps",
+        is_torch_cuda_available: "cuda"
+    }
+
+    device = "cpu"
+    for b, d in device_map.items():
+        if b():
+            device = "{}:{}".format(d, local_rank)
+
+    return torch.device(device)
+
+def count_parameters(model: "torch.nn.Module") -> Tuple[int, int]:
+    r"""
+    Returns the number of trainable parameters and number of all parameters in the model.
+    """
+    trainable_params, all_param = 0, 0
+    for param in model.parameters():
+        num_params = param.numel()
+        # if using DS Zero 3 and the weights are initialized empty
+        if num_params == 0 and hasattr(param, "ds_numel"):
+            num_params = param.ds_numel
+
+        # Due to the design of 4bit linear layers from bitsandbytes, multiply the number of parameters by itemsize
+        if param.__class__.__name__ == "Params4bit":
+            if hasattr(param, "quant_storage") and hasattr(param.quant_storage, "itemsize"):
+                num_bytes = param.quant_storage.itemsize
+            elif hasattr(param, "element_size"):  # for older pytorch version
+                num_bytes = param.element_size()
+            else:
+                num_bytes = 1
+
+            num_params = num_params * 2 * num_bytes
+
+        all_param += num_params
+        if param.requires_grad:
+            trainable_params += num_params
+
+    return trainable_params, all_param
+
+
+def find_all_linear_modules(model: PreTrainedModel) -> List[str]:
+    r"""
+    Finds all available modules to apply lora or galore.
+    """
+    model_type = getattr(model.config, "model_type", None)
+    forbidden_modules = {"lm_head"}
+    if model_type == "chatglm":
+        forbidden_modules.add("output_layer")
+    elif model_type == "internlm2":
+        forbidden_modules.add("output")
+    elif model_type in ["llava", "llava_next", "llava_next_video", "paligemma", "video_llava"]:
+        forbidden_modules.add("multi_modal_projector")
+    elif model_type == "qwen2_vl":
+        forbidden_modules.add("merger")
+
+    module_names = set()
+    for name, module in model.named_modules():
+        if any(forbidden_module in name for forbidden_module in forbidden_modules):
+            continue
+
+        if "Linear" in module.__class__.__name__ and "Embedding" not in module.__class__.__name__:
+            module_names.add(name.split(".")[-1])
+
+    logger.info("Found linear modules: {}".format(",".join(module_names)))
+    return list(module_names)
+
+_is_fp16_available = is_torch_npu_available() or is_torch_cuda_available()
+try:
+    _is_bf16_available = is_torch_bf16_gpu_available() or (is_torch_npu_available() and torch.npu.is_bf16_supported())
+except Exception:
+    _is_bf16_available = False
+
+
+def infer_optim_dtype(model_dtype: torch.dtype) -> torch.dtype:
+    r"""
+    Infers the optimal dtype according to the model_dtype and device compatibility.
+    """
+    if _is_bf16_available and model_dtype == torch.bfloat16:
+        return torch.bfloat16
+    elif _is_fp16_available:
+        return torch.float16
+    else:
+        return torch.float32

distillflow/common/logger.py

@@ -0,0 +1,101 @@
+import logging
+import os
+import sys
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import Optional
+
+_thread_lock = threading.RLock()
+_default_handler: Optional["logging.Handler"] = None
+_default_log_level: "logging._Level" = logging.INFO
+
+
+class LoggerHandler(logging.Handler):
+    r"""
+    Redirects the logging output to the logging file for LLaMA Board.
+    """
+
+    def __init__(self, output_dir: str) -> None:
+        super().__init__()
+        formatter = logging.Formatter(
+            fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S"
+        )
+        self.setLevel(logging.INFO)
+        self.setFormatter(formatter)
+
+        os.makedirs(output_dir, exist_ok=True)
+        self.running_log = os.path.join(output_dir, "log.txt")
+        if os.path.exists(self.running_log):
+            os.remove(self.running_log)
+
+        self.thread_pool = ThreadPoolExecutor(max_workers=1)
+
+    def _write_log(self, log_entry: str) -> None:
+        with open(self.running_log, "a", encoding="utf-8") as f:
+            f.write(log_entry + "\n\n")
+
+    def emit(self, record) -> None:
+        if record.name == "httpx":
+            return
+
+        log_entry = self.format(record)
+        self.thread_pool.submit(self._write_log, log_entry)
+
+    def close(self) -> None:
+        self.thread_pool.shutdown(wait=True)
+        return super().close()
+
+
+def _get_default_logging_level() -> "logging._Level":
+    r"""
+    Returns the default logging level.
+    """
+    env_level_str = os.environ.get("LLAMAFACTORY_VERBOSITY", None)
+    if env_level_str:
+        if env_level_str.upper() in logging._nameToLevel:
+            return logging._nameToLevel[env_level_str.upper()]
+        else:
+            raise ValueError("Unknown logging level: {}.".format(env_level_str))
+
+    return _default_log_level
+
+
+def _get_library_name() -> str:
+    return __name__.split(".")[0]
+
+
+def _get_library_root_logger() -> "logging.Logger":
+    return logging.getLogger(_get_library_name())
+
+
+def _configure_library_root_logger() -> None:
+    r"""
+    Configures root logger using a stdout stream handler with an explicit format.
+    """
+    global _default_handler
+
+    with _thread_lock:
+        if _default_handler:
+            return
+
+        formatter = logging.Formatter(
+            fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+            datefmt="%m/%d/%Y %H:%M:%S",
+        )
+        _default_handler = logging.StreamHandler(sys.stdout)
+        _default_handler.setFormatter(formatter)
+        library_root_logger = _get_library_root_logger()
+        library_root_logger.addHandler(_default_handler)
+        library_root_logger.setLevel(_get_default_logging_level())
+        library_root_logger.propagate = False
+
+
+def get_logger(name: Optional[str] = None) -> "logging.Logger":
+    r"""
+    Returns a logger with the specified name. It it not supposed to be accessed externally.
+    """
+    if name is None:
+        name = _get_library_name()
+
+    _configure_library_root_logger()
+    return logging.getLogger(name)

distillflow/config/__init__.py

@@ -0,0 +1,6 @@
+from .config import Config
+
+__all__ = [
+    Config,
+]
+

distillflow/config/config.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from distillflow.datasets.args import DataArgs
+from distillflow.model.args import ModelArgs
+from distillflow.trainer.args import DistillArgs
+
+class Config(BaseModel):
+    student_model: ModelArgs = Field(
+        description="Details about the student model that we want to train"
+    )
+    teacher_model: ModelArgs = Field(
+        description="Details about the teacher model",
+        default=None
+    )
+    data: DataArgs = Field(
+        description="The datasets that we want to choose to run the training"
+    )
+    distill: DistillArgs = Field(
+        description="Distillation training parameters"
+    )
+
+    model_config = {
+        "extra": "forbid"
+    }

distillflow/config/validator.py

@@ -0,0 +1,38 @@
+from pydantic import ValidationError
+
+from . import Config
+
+
+def print_validation_error(error: ValidationError):
+    """
+    Print detailed validation errors including field descriptions and examples
+    """
+    print("Validation Error:")
+    print("-" * 50)
+
+    for error in error.errors():
+        # Get the field path
+        field_path = " -> ".join(str(item) for item in error['loc'])
+
+        # Find the relevant model and field
+        current_model = Config
+        field_info = None
+
+        for item in error['loc']:
+            if isinstance(item, int):
+                continue
+            if hasattr(current_model, 'model_fields') and item in current_model.model_fields:
+                field_info = current_model.model_fields[item]
+                # If this field is another model, update current_model for nested fields
+                if hasattr(field_info.annotation, 'model_fields'):
+                    current_model = field_info.annotation
+
+        print(f"Field: {field_path}")
+        print(f"Error: {error['msg']}")
+
+        if field_info:
+            if field_info.description:
+                print(f"Description: {field_info.description}")
+            if field_info.examples:
+                print(f"Example: {field_info.examples}")
+        print("-" * 50)

distillflow/datasets/__init__.py

@@ -0,0 +1,6 @@
+from distillflow.datasets.loader import get_dataset
+
+__all__ = [
+    get_dataset
+]
+

distillflow/datasets/args.py

@@ -0,0 +1,99 @@
+from pydantic import BaseModel, Field, model_validator
+
+from typing import Optional, Literal, List
+
+from .template.args import TemplateArgs
+
+class DatasetArgs(BaseModel):
+    path: str = Field(
+        description="The path to huggingface dataset"
+    )
+    template: TemplateArgs = None
+    num_samples: Optional[int] = Field (
+        default=None,
+        description="Number of samples to pick from the given dataset",
+        examples=[1000, 100_000]
+    )
+    load_from_cache_file: bool = Field(
+        default=True,
+        description="Should load the dataset from cache if available, if disabled, the dataset will be re-synced from HF"
+    )
+    split: str = Field(
+        default="train",
+        description="Split to be used for this dataset (defaults to `train`)",
+        examples=["train", "test"]
+    )
+
+    model_config = {
+        "extra": "forbid"
+    }
+
+class DataArgs(BaseModel):
+    r"""
+    Arguments pertaining to what data we are going to input our model for training and evaluation.
+    """
+    seed: Optional[int] = Field(
+        default=0,
+        description="Seed to use when shuffling the dataset."
+    )
+    train_datasets: List[DatasetArgs] = Field(
+        description="The dataset(s) to use for training. Provide as a list of DatasetArgs."
+    )
+    eval_datasets: Optional[List[DatasetArgs]] = Field(
+        default=None,
+        description="The names of dataset(s) to use for evaluation. Provide as a list of DatasetArgs."
+    )
+    cache_dir: Optional[str] = Field(
+        default=None,
+        description="Where to store the pre-trained datasets downloaded."
+    )
+    hf_hub_token: Optional[str] = Field(
+        default=None,
+        description="Auth token to log in with Hugging Face Hub."
+    )
+    text_field: Optional[str] = Field(
+        default=None,
+        description="Name of the field key to convert the dataset."
+    )
+    streaming: bool = Field(
+        default=False,
+        description="Enable dataset streaming."
+    )
+    buffer_size: Optional[int] = Field(
+        default=16384,
+        description="Size of the buffer to randomly sample examples from in dataset streaming."
+    )
+    mix_strategy: Optional[Literal["concat", "interleave_under", "interleave_over"]] = Field(
+        default="concat",
+        description="Strategy to use in dataset mixing (concat/interleave) (undersampling/oversampling)."
+    )
+    interleave_probs: Optional[str] = Field(
+        default=None,
+        description="Probabilities to sample data from datasets. Use commas to separate multiple datasets."
+    )
+    max_samples: Optional[int] = Field(
+        default=None,
+        description="For debugging purposes, truncate the number of examples for each dataset."
+    )
+    test_size: float = Field(
+        default=0.0,
+        description="Size of the development set, should be an integer or a float in range `[0,1)`.",
+        examples=[1000, 100_000, 0.2, 0.5]
+    )
+
+    @model_validator(mode='after')
+    def validate_args(self) -> 'DataArgs':
+        if self.eval_datasets is not None and self.test_size > 0:
+            raise ValueError("Cannot specify `val_size` if `eval_dataset` is not None.")
+
+        if self.streaming and 1e-6 < self.test_size < 1:
+            raise ValueError("Streaming mode should have an integer test size.")
+
+        if self.streaming and self.max_samples is not None:
+            raise ValueError("`max_samples` is incompatible with `streaming`.")
+
+        return self
+
+    model_config = {
+        "extra": "forbid"
+    }

distillflow/datasets/loader.py

@@ -0,0 +1,191 @@
+from functools import partial
+from typing import TypedDict, Optional, Dict, Any, List
+
+import numpy as np
+from datasets import DatasetDict, load_dataset, Dataset, concatenate_datasets, interleave_datasets, IterableDataset
+from transformers import PreTrainedTokenizer
+
+from .args import DatasetArgs, DataArgs
+from .template import ShareGpt, Alpaca
+from ..common import get_logger
+
+logger = get_logger(__name__)
+
+class DatasetModule(TypedDict):
+    train_dataset: Optional[Dataset]
+    eval_dataset: Optional[Dataset]
+
+def _load_single_dataset(
+        dataset_args: DatasetArgs,
+        data_args: DataArgs,
+        tokenizer: PreTrainedTokenizer) -> Dataset:
+
+    if dataset_args is not None:
+        logger.info("Loading dataset {}...".format(dataset_args))
+        dataset = load_dataset(
+            path=dataset_args.path,
+            split=dataset_args.split,
+            cache_dir=data_args.cache_dir,
+            token=data_args.hf_hub_token,
+            streaming=data_args.streaming, # and (dataset_attr.load_from != "file")),
+            trust_remote_code=True
+        )
+        # Shuffle dataset with a pre-seed
+        dataset = dataset.shuffle(seed=data_args.seed)
+
+        # if data_args.streaming and (dataset_attr.load_from == "file"):  # faster than specifying streaming=True
+        #     dataset = dataset.to_iterable_dataset()  # TODO: add num shards parameter
+
+        if dataset_args.num_samples is not None:# and not data_args.streaming:
+            target_num = dataset_args.num_samples
+            indexes = np.random.permutation(len(dataset))[:target_num]  # all samples should be included
+            target_num -= len(indexes)
+            if target_num > 0:
+                expand_indexes = np.random.choice(len(dataset), target_num)
+                indexes = np.concatenate((indexes, expand_indexes), axis=0)
+
+            assert len(indexes) == dataset_args.num_samples, "Sample num mismatched."
+            dataset = dataset.select(indexes)
+            logger.info("Sampled {} examples from dataset {}.".format(dataset_args.num_samples, dataset_args))
+
+        if data_args.max_samples is not None:  # truncate dataset
+            max_samples = min(data_args.max_samples, len(dataset))
+            dataset = dataset.select(range(max_samples))
+
+        column_names = list(next(iter(dataset)).keys())
+
+        template = dataset_args.template
+        template_mapping = {
+            "sharegpt": ShareGpt(template.args),
+            "alpaca": Alpaca(template.args)
+        }
+        if data_args.streaming:
+            dataset = dataset.map(
+                partial(template_mapping[template.name].convert),
+                batched=False,
+                remove_columns=column_names,
+            )
+        else:
+            dataset = dataset.map(
+                partial(template_mapping[template.name].convert),
+                batched=False,
+                remove_columns=column_names,
+                load_from_cache_file=dataset_args.load_from_cache_file
+            )
+
+        if data_args.text_field is not None:
+            if isinstance(dataset, IterableDataset):
+                dataset = dataset.map(partial(to_text, data_args.text_field, tokenizer), batched=False,
+                                      remove_columns=dataset.column_names)
+            else:
+                dataset = dataset.map(partial(to_text, data_args.text_field, tokenizer), batched=False, remove_columns=dataset.column_names,
+                                  load_from_cache_file=dataset_args.load_from_cache_file)
+        return dataset
+
+def to_text(field_name, tokenizer: PreTrainedTokenizer, example: Dict[str, Any]) -> Dict[str, Any]:
+    system = example["_system"]
+    prompt = example["_prompt"]
+    response = example["_response"]
+    message = system + prompt + response
+    return {field_name: tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)}
+
+def split_dataset(dataset: Dataset, data_args: DataArgs, seed: int) -> DatasetDict:
+    r"""
+    Splits the dataset and returns a dataset dict containing train set and validation set.
+
+    Supports both map dataset and iterable dataset.
+    """
+    if data_args.streaming:
+        dataset = dataset.shuffle(buffer_size=data_args.buffer_size, seed=seed)
+        val_set = dataset.take(int(data_args.test_size))
+        train_set = dataset.skip(int(data_args.test_size))
+        return DatasetDict({"train": train_set, "validation": val_set})
+    else:
+        dataset = dataset.shuffle(seed=seed)
+        test_size = int(data_args.test_size) if data_args.test_size > 1 else data_args.test_size
+        dataset = dataset.train_test_split(test_size=test_size, seed=seed)
+        return DatasetDict({"train": dataset["train"], "validation": dataset["test"]})
+
+def get_dataset(data_args: DataArgs,
+                tokenizer: PreTrainedTokenizer,
+                tokenizer_function=None) -> DatasetModule:
+    # Load and preprocess dataset
+    # with training_args.main_process_first(desc="load dataset"):
+    dataset = _get_merged_dataset(data_args.train_datasets, data_args, tokenizer)
+    dataset = dataset.shuffle(seed=data_args.seed)
+    eval_dataset = None
+    if data_args.eval_datasets:
+        eval_dataset = _get_merged_dataset(data_args.eval_datasets, data_args, tokenizer)
+
+    dataset_dict = {}
+
+    if eval_dataset is None:
+        dataset_dict = split_dataset(dataset, data_args, data_args.seed)
+    else:
+        if data_args.streaming:
+            eval_dataset = eval_dataset.shuffle(buffer_size=data_args.buffer_size, seed=data_args.seed)
+
+        dataset_dict["validation"] = eval_dataset
+
+        if dataset is not None:
+            if data_args.streaming:
+                dataset = dataset.shuffle(buffer_size=data_args.buffer_size, seed=data_args.train_datasets.seed)
+
+            dataset_dict["train"] = dataset
+
+        dataset_dict = DatasetDict(dataset_dict)
+
+    dataset_module = {}
+    if "train" in dataset_dict:
+        dataset_module["train_dataset"] = dataset_dict["train"]
+
+    if "validation" in dataset_dict:
+        dataset_module["eval_dataset"] = dataset_dict["validation"]
+
+    if tokenizer_function is not None:
+        dataset_module["train_dataset"] = dataset_module["train_dataset"].map(tokenizer_function,
+                                              batched=True, num_proc=32, remove_columns=[data_args.text_field])
+
+        dataset_module["eval_dataset"] = dataset_module["eval_dataset"].map(tokenizer_function,
+                                              batched=True, num_proc=32, remove_columns=[data_args.text_field])
+    return dataset_module
+
+def _get_merged_dataset(
+    dataset_list: List[DatasetArgs],
+    data_args: DataArgs,
+    tokenizer: PreTrainedTokenizer
+) -> Optional[Dataset]:
+    r"""
+    Gets the merged datasets in the standard format.
+    """
+    datasets = []
+    for dataset_attr in dataset_list:
+        datasets.append(_load_single_dataset(dataset_attr, data_args, tokenizer))
+
+    return merge_dataset(datasets, data_args, data_args.seed)
+
+def merge_dataset(all_datasets: List[Dataset], data_args: DataArgs, seed) -> Dataset:
+    r"""
+    Merges multiple datasets to a unified dataset.
+    """
+    if len(all_datasets) == 1:
+        return all_datasets[0]
+    elif data_args.mix_strategy == "concat":
+        if data_args.streaming:
+            logger.warning("The samples between different datasets will not be mixed in streaming mode.")
+        return concatenate_datasets(all_datasets)
+    elif data_args.mix_strategy.startswith("interleave"):
+        if not data_args.streaming:
+            logger.warning("We recommend using `mix_strategy=concat` in non-streaming mode.")
+
+        return interleave_datasets(
+            datasets=all_datasets,
+            probabilities=data_args.interleave_probs,
+            seed=seed,
+            stopping_strategy="first_exhausted" if data_args.mix_strategy.endswith("under") else "all_exhausted",
+        )
+    else:
+        raise ValueError("Unknown mixing strategy: {}.".format(data_args.mix_strategy))
+
+
+

distillflow/datasets/template/__init__.py

@@ -0,0 +1,12 @@
+from .alpaca import Alpaca, AlpacaArgs
+from .sharegpt import ShareGptArgs, ShareGpt
+from .template import Template
+
+__all__ = [
+    Alpaca,
+    AlpacaArgs,
+    ShareGpt,
+    ShareGptArgs,
+    Template,
+]
+

distillflow/datasets/template/alpaca.py

@@ -0,0 +1,85 @@
+from dataclasses import dataclass
+from typing import Dict, Any, Optional
+
+from .template import Template
+from .role import Role
+
+@dataclass
+class AlpacaArgs:
+    prompt: Optional[str] = "instruction"
+    query: Optional[str] = "input"
+    response: Optional[str] = "output"
+    history: Optional[str] = None
+    system_tag: Optional[str] = None
+    tools: Optional[str] = None
+
+'''
+Expected data format
+[
+  {
+    "instruction": "human instruction (required)",
+    "input": "human input (optional)",
+    "output": "model response (required)",
+    "system": "system prompt (optional)",
+    "history": [
+      ["human instruction in the first round (optional)", "model response in the first round (optional)"],
+      ["human instruction in the second round (optional)", "model response in the second round (optional)"]
+    ]
+  }
+]
+'''
+class Alpaca(Template):
+    def __init__(self, args: Optional[AlpacaArgs] = AlpacaArgs()):
+        self.args = args
+
+    def convert(self, example: Dict[str, Any]) -> Dict[str, Any]:
+        r"""
+        Converts alpaca format dataset to the standard format.
+        """
+        prompt = []
+        if self.args.history and isinstance(example[self.args.history], list):
+            for old_prompt, old_response in example[self.args.history]:
+                prompt.append({"role": Role.USER.value, "content": old_prompt})
+                prompt.append({"role": Role.ASSISTANT.value, "content": old_response})
+
+        query = []
+        if self.args.prompt and example[self.args.prompt]:
+            query.append(example[self.args.prompt])
+
+        if self.args.query and example[self.args.query]:
+            query.append(example[self.args.query])
+
+        prompt.append({"role": Role.USER.value, "content": "\n".join(query)})
+
+        # if args.kto_tag and isinstance(example[args.kto_tag], bool):  # kto example
+        #     response = [{"role": Role.ASSISTANT.value, "content": example[args.response]}]
+        #     if example[self.args.kto_tag]:
+        #         response = response + [{"role": Role.ASSISTANT.value, "content": ""}]
+        #     else:
+        #         response = [{"role": Role.ASSISTANT.value, "content": ""}] + response
+        # elif (
+        #         args.ranking
+        #         and isinstance(example[args.chosen], str)
+        #         and isinstance(example[args.rejected], str)
+        # ):  # pairwise example
+        #     response = [
+        #         {"role": Role.ASSISTANT.value, "content": example[args.chosen]},
+        #         {"role": Role.ASSISTANT.value, "content": example[args.rejected]},
+        #     ]
+        if self.args.response and isinstance(example[self.args.response], str):  # normal example
+            response = [{"role": Role.ASSISTANT.value, "content": example[self.args.response]}]
+        else:  # unsupervised
+            response = []
+
+        system = [{"role": Role.SYSTEM.value,
+                   "content": example[self.args.system_tag]} if self.args.system_tag in example else {
+            "role": Role.SYSTEM.value,
+            "content": "You are a helpful assistant."}]
+
+        output = {
+            "_prompt": prompt,
+            "_response": response,
+            "_system": system,
+            # "_tools": example[self.args.tools] if self.args.tools else "",
+        }
+        return output