Distillflow 0.2.0__py3-none-any.whl

distillflow/datasets/template/args.py

@@ -0,0 +1,56 @@
+from typing import Literal, Optional, Union
+
+from pydantic import BaseModel, Field, model_validator
+
+from . import ShareGptArgs, AlpacaArgs
+
+
+class TemplateArgs(BaseModel):
+    name: Literal["sharegpt", "alpaca"] = Field(
+        default="sharegpt",
+        description="Template used by the dataset. Based on the template, the data will be converted to a standard format before training"
+    )
+    args: Optional[Union[ShareGptArgs, AlpacaArgs]] = Field(
+        default=None,
+        description="Template args to map the corresponding dataset columns to the ones expected by sharegpt/alpaca"
+    )
+
+    # Validate and transform args based on template name
+    @model_validator(mode='after')
+    def validate_args(self) -> 'TemplateArgs':
+        if self.args is None:
+            # Set default args based on template
+            self.args = {
+                "sharegpt": ShareGptArgs(),
+                "alpaca": AlpacaArgs()
+            }[self.name]
+            return self
+
+        # If args is a dict, convert to appropriate type
+        if isinstance(self.args, dict):
+            template_classes = {
+                "sharegpt": ShareGptArgs,
+                "alpaca": AlpacaArgs
+            }
+
+            template_class = template_classes.get(self.name)
+            if template_class is None:
+                raise ValueError(f"Unknown template: {self.name}")
+
+            try:
+                self.args = template_class(**self.args)
+            except Exception as e:
+                raise ValueError(f"Invalid args for template '{self.name}': {str(e)}")
+
+        # Validate that args matches template
+        expected_type = ShareGptArgs if self.name == "sharegpt" else AlpacaArgs
+        if not isinstance(self.args, expected_type):
+            raise ValueError(
+                f"Template '{self.name}' expects {expected_type.__name__} "
+                f"but got {type(self.args).__name__}"
+            )
+        return self
+
+    model_config = {
+        "extra": "forbid"
+    }

distillflow/datasets/template/role.py

@@ -0,0 +1,10 @@
+from enum import Enum, unique
+
+
+@unique
+class Role(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+    FUNCTION = "function"
+    OBSERVATION = "observation"

distillflow/datasets/template/sharegpt.py

@@ -0,0 +1,84 @@
+import json
+from dataclasses import dataclass
+from typing import Dict, Any, Optional
+
+from .template import Template
+from .role import Role
+from ...common import get_logger
+
+logger = get_logger(__name__)
+
+@dataclass
+class ShareGptArgs:
+    messages: Optional[str] = "conversations"
+    # sharegpt tags
+    role_tag: Optional[str] = "from"
+    content_tag: Optional[str] = "value"
+    user_tag: Optional[str] = "human"
+    assistant_tag: Optional[str] = "gpt"
+    observation_tag: Optional[str] = "observation"
+    function_tag: Optional[str] = "function_call"
+    system_tag: Optional[str] = "system"
+
+class ShareGpt(Template):
+
+    def __init__(self, args: Optional[ShareGptArgs] = ShareGptArgs()):
+        self.args = args
+
+    def convert(self, example: Dict[str, Any]) -> Dict[str, Any]:
+        r"""
+        Converts sharegpt format dataset to the standard format.
+        """
+        tag_mapping = {
+            self.args.user_tag: Role.USER.value,
+            self.args.assistant_tag: Role.ASSISTANT.value,
+            self.args.observation_tag: Role.OBSERVATION.value,
+            self.args.function_tag: Role.FUNCTION.value,
+            self.args.system_tag: Role.SYSTEM.value,
+        }
+        odd_tags = (self.args.user_tag, self.args.observation_tag)
+        even_tags = (self.args.assistant_tag, self.args.function_tag)
+        accept_tags = (odd_tags, even_tags)
+        messages = example[self.args.messages]
+        system = []
+        if (
+                self.args.system_tag
+                and len(messages) != 0
+                and messages[0][self.args.role_tag] == self.args.system_tag
+        ):
+            system.append({"role": Role.SYSTEM.value, "content": messages[0][self.args.content_tag]})
+            messages = messages[1:]
+        else:
+            system.append({"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."})
+
+        aligned_messages = []
+        broken_data = False
+        for turn_idx, message in enumerate(messages):
+            if message[self.args.role_tag] not in accept_tags[turn_idx % 2]:
+                logger.warning("Invalid role tag in {}.".format(messages))
+                broken_data = True
+
+            aligned_messages.append(
+                {"role": tag_mapping[message.get(self.args.role_tag, '')], "content": message.get(self.args.content_tag, '')}
+            )
+
+        prompt = aligned_messages[:-1]
+        response = aligned_messages[-1:]
+
+        if isinstance(response, str):
+            try:
+                response = json.loads(response)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON in response: {response}") from e
+        if broken_data:
+            logger.warning("Skipping this abnormal example.")
+            prompt, response = [], []
+
+        output = {
+            "_prompt": prompt,
+            "_response": response,
+            "_system": system,
+        }
+
+        return output

distillflow/datasets/template/template.py

@@ -0,0 +1,6 @@
+from typing import Dict, Any
+
+class Template:
+    def convert(self, example: Dict[str, Any]) -> Dict[str, Any]:
+        raise NotImplementedError("Template must implement the convert method.")
+

distillflow/evaluation/rouge.py

@@ -0,0 +1,30 @@
+import evaluate
+
+def read_file(file_path):
+    """Reads text from a file and returns it as a list of lines."""
+    with open(file_path, 'r', encoding='utf-8') as file:
+        return [line.strip() for line in file.readlines()]
+
+
+def compute_rouge_scores(reference_file, generated_file):
+    # Load the ROUGE metric from the evaluate library
+    rouge = evaluate.load('rouge')
+
+    reference_texts = read_file(reference_file)
+    generated_texts = read_file(generated_file)
+
+    # Compute ROUGE scores for the provided texts
+    scores = rouge.compute(predictions=generated_texts, references=reference_texts)
+
+    # Return the scores
+    return scores
+
+
+# Example usage
+# reference_texts = ["The cat sat on the mat."]
+# generated_texts = ["The cat is sitting on the mat."]
+#
+# # Compute and print ROUGE scores
+# rouge_scores = compute_rouge_scores(reference_texts, generated_texts)
+# for metric, score in rouge_scores.items():
+#     print(f"{metric}: {score:.4f}")

distillflow/model/adapter.py

@@ -0,0 +1,257 @@
+import re
+
+import torch
+from peft import LoraConfig, LoraModel, PeftModel, TaskType, get_peft_model
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.modeling_utils import is_fsdp_enabled
+
+from .args import ModelArgs
+from .finetuning_args import FinetuningArgs
+from ..common.logger import get_logger
+from .unsloth import get_unsloth_peft_model, load_unsloth_peft_model
+
+from transformers import PretrainedConfig, PreTrainedModel
+
+from ..common import find_all_linear_modules
+
+logger = get_logger(__name__)
+
+
+def _setup_full_tuning(model: PreTrainedModel, is_trainable: bool, cast_trainable_params_to_fp32: bool) -> None:
+    if not is_trainable:
+        return
+
+    logger.info("Fine-tuning method: Full")
+    for name, param in model.named_parameters():
+        if cast_trainable_params_to_fp32:
+            param.data = param.data.to(torch.float32)
+
+
+def _setup_freeze_tuning(
+        model: PreTrainedModel,
+        finetuning_args: FinetuningArgs,
+        is_trainable: bool,
+        cast_trainable_params_to_fp32: bool,
+) -> None:
+    if not is_trainable:
+        return
+
+    logger.info("Fine-tuning method: Freeze")
+    if hasattr(model.config, "text_config"):  # composite models
+        config = getattr(model.config, "text_config")
+    else:
+        config = model.config
+
+    num_layers = (
+            getattr(config, "num_hidden_layers", None)
+            or getattr(config, "num_layers", None)
+            or getattr(config, "n_layer", None)
+    )
+    if not num_layers:
+        raise ValueError("Current model does not support freeze tuning.")
+
+    if finetuning_args.use_llama_pro:
+        if num_layers % finetuning_args.freeze_trainable_layers != 0:
+            raise ValueError(
+                "`num_layers` {} should be divisible by `num_layer_trainable` {}.".format(
+                    num_layers, finetuning_args.freeze_trainable_layers
+                )
+            )
+
+        stride = num_layers // finetuning_args.freeze_trainable_layers
+        trainable_layer_ids = range(stride - 1, num_layers + stride - 1, stride)
+    elif finetuning_args.freeze_trainable_layers > 0:  # fine-tuning the last n layers if num_layer_trainable > 0
+        trainable_layer_ids = range(max(0, num_layers - finetuning_args.freeze_trainable_layers), num_layers)
+    else:  # fine-tuning the first n layers if num_layer_trainable < 0
+        trainable_layer_ids = range(min(-finetuning_args.freeze_trainable_layers, num_layers))
+
+    hidden_modules = set()
+    non_hidden_modules = set()
+    for name, _ in model.named_parameters():
+        if ".0." in name:
+            hidden_modules.add(name.split(".0.")[-1].split(".")[0])
+        elif ".1." in name:  # MoD starts from layer 1
+            hidden_modules.add(name.split(".1.")[-1].split(".")[0])
+
+        if re.search(r"\.\d+\.", name) is None:
+            non_hidden_modules.add(name.split(".")[-2])
+
+    trainable_layers = []
+    for module_name in finetuning_args.freeze_trainable_modules:
+        if module_name != "all" and module_name not in hidden_modules:
+            raise ValueError(
+                "Module {} is not found, please choose from {}".format(module_name, ", ".join(hidden_modules))
+            )
+
+        for idx in trainable_layer_ids:
+            trainable_layers.append(".{:d}.{}".format(idx, module_name if module_name != "all" else ""))
+
+    if finetuning_args.freeze_extra_modules:
+        for module_name in finetuning_args.freeze_extra_modules:
+            if module_name not in non_hidden_modules:
+                raise ValueError(
+                    "Module {} is not found, please choose from {}".format(module_name, ", ".join(non_hidden_modules))
+                )
+
+            trainable_layers.append(module_name)
+
+    for name, param in model.named_parameters():
+        if any(trainable_layer in name for trainable_layer in trainable_layers):
+            if cast_trainable_params_to_fp32:
+                param.data = param.data.to(torch.float32)
+        else:
+            param.requires_grad_(False)
+
+    logger.info("Set trainable layers: {}".format(",".join(trainable_layers)))
+
+
+def _setup_lora_tuning(
+        config: PretrainedConfig,
+        model: PreTrainedModel,
+        model_args: ModelArgs,
+        is_trainable: bool,
+        cast_trainable_params_to_fp32: bool,
+) -> PeftModel:
+    finetuning_args = model_args.finetuning_args
+    if is_trainable:
+        logger.info("Fine-tuning method: {}".format("DoRA" if finetuning_args.use_dora else "LoRA"))
+
+    adapter_to_resume = None
+
+    if model_args.adapter_name_or_path is not None:
+        is_mergeable = True
+        if getattr(model, "quantization_method", None):  # merge lora in quantized model is unstable
+            assert len(model_args.adapter_name_or_path) == 1, "Quantized model only accepts a single adapter."
+            is_mergeable = False
+
+        if is_deepspeed_zero3_enabled():
+            assert len(model_args.adapter_name_or_path) == 1, "Cannot use multiple adapters in DeepSpeed ZeRO-3."
+            is_mergeable = False
+
+        if model_args.use_unsloth:
+            assert len(model_args.adapter_name_or_path) == 1, "Unsloth model only accepts a single adapter."
+            is_mergeable = False
+
+        if (is_trainable and not finetuning_args.create_new_adapter) or (not is_mergeable):
+            adapter_to_merge = model_args.adapter_name_or_path[:-1]
+            adapter_to_resume = model_args.adapter_name_or_path[-1]
+        else:
+            adapter_to_merge = model_args.adapter_name_or_path
+
+        init_kwargs = {
+            "subfolder": model_args.adapter_folder,
+            "offload_folder": model_args.offload_folder,
+            "cache_dir": model_args.cache_dir,
+            "revision": model_args.model_revision,
+            "token": model_args.hf_hub_token,
+        }
+
+        for adapter in adapter_to_merge:
+            model: LoraModel = PeftModel.from_pretrained(model, adapter, **init_kwargs)
+            model = model.merge_and_unload()
+
+        if len(adapter_to_merge) > 0:
+            logger.info("Merged {} adapter(s).".format(len(adapter_to_merge)))
+
+        if adapter_to_resume is not None:  # resume lora training
+            if model_args.use_unsloth:
+                model = load_unsloth_peft_model(config, model_args, is_trainable=is_trainable)
+            else:
+                model = PeftModel.from_pretrained(model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs)
+
+        logger.info("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))
+
+    if is_trainable and adapter_to_resume is None:  # create new lora weights while training
+        if len(finetuning_args.lora_target) == 1 and finetuning_args.lora_target[0] == "all":
+            target_modules = find_all_linear_modules(model)
+        else:
+            target_modules = finetuning_args.lora_target
+
+        if model_args.resize_vocab and finetuning_args.additional_target is None:
+            input_embeddings = model.get_input_embeddings()
+            output_embeddings = model.get_output_embeddings()
+            module_names = set()
+            for name, module in model.named_modules():
+                if module in [input_embeddings, output_embeddings]:
+                    module_names.add(name.split(".")[-1])
+
+            finetuning_args.additional_target = module_names
+            logger.warning("Vocab has been resized, add {} to trainable params.".format(",".join(module_names)))
+
+        peft_kwargs = {
+            "r": finetuning_args.lora_rank,
+            "target_modules": target_modules,
+            "lora_alpha": finetuning_args.lora_alpha,
+            "lora_dropout": finetuning_args.lora_dropout,
+            "use_rslora": finetuning_args.use_rslora,
+            "use_dora": finetuning_args.use_dora,
+            "modules_to_save": finetuning_args.additional_target,
+        }
+
+        if model_args.use_unsloth:
+            model = get_unsloth_peft_model(model, model_args, peft_kwargs)
+        else:
+            if finetuning_args.pissa_init:
+                if finetuning_args.pissa_iter == -1:
+                    logger.info("Using PiSSA initialization.")
+                    peft_kwargs["init_lora_weights"] = "pissa"
+                else:
+                    logger.info("Using PiSSA initialization with FSVD steps {}.".format(finetuning_args.pissa_iter))
+                    peft_kwargs["init_lora_weights"] = "pissa_niter_{}".format(finetuning_args.pissa_iter)
+
+            lora_config = LoraConfig(
+                task_type=TaskType.CAUSAL_LM,
+                inference_mode=False,
+                **peft_kwargs,
+            )
+            model = get_peft_model(model, lora_config)
+
+    if is_trainable and cast_trainable_params_to_fp32:
+        for param in filter(lambda p: p.requires_grad, model.parameters()):
+            param.data = param.data.to(torch.float32)
+
+    return model
+
+
+def init_adapter(
+        config: PretrainedConfig, model: PreTrainedModel, model_args: ModelArgs, is_trainable: bool) -> PreTrainedModel:
+    r"""
+    Initializes the adapters.
+
+    Support full-parameter, freeze and LoRA training.
+
+    Note that the trainable parameters must be cast to float32.
+    """
+    finetuning_args = model_args.finetuning_args
+    if is_trainable and getattr(model, "quantization_method", None) is not None:
+        if finetuning_args.finetuning_type != "lora":
+            raise ValueError("Quantized models can only be used for the LoRA tuning.")
+
+        if finetuning_args.pissa_init:
+            raise ValueError("Cannot initialize PiSSA adapter on quantized models.")
+
+    # cast trainable parameters to float32 if:
+    # 1. is_trainable and not pure_bf16 and not badam and quantization_bit is not None (qlora)
+    # 2. is_trainable and not pure_bf16 and not badam and not zero3 and not fsdp (zero3 or fsdp already in fp32)
+    cast_trainable_params_to_fp32 = False
+    quantization_args = model_args.quantization_args
+    if not is_trainable:
+        pass
+    elif quantization_args.quantization_bit is None and (is_deepspeed_zero3_enabled() or is_fsdp_enabled()):
+        logger.info("ZeRO3 / FSDP detected, remaining trainable params in float32.")
+    else:
+        logger.info("Upcasting trainable params to float32.")
+        cast_trainable_params_to_fp32 = True
+
+    if finetuning_args.finetuning_type == "full":
+        _setup_full_tuning(model, is_trainable, cast_trainable_params_to_fp32)
+    elif finetuning_args.finetuning_type == "freeze":
+        _setup_freeze_tuning(model, finetuning_args, is_trainable, cast_trainable_params_to_fp32)
+    elif finetuning_args.finetuning_type == "lora":
+        model = _setup_lora_tuning(
+            config, model, model_args, is_trainable, cast_trainable_params_to_fp32
+        )
+    else:
+        raise NotImplementedError("Unknown finetuning type: {}.".format(finetuning_args.finetuning_type))
+
+    return model

distillflow/model/args.py

@@ -0,0 +1,173 @@
+from typing import Optional, Literal
+
+from pydantic import Field, BaseModel
+
+from distillflow.model.finetuning_args import FinetuningArgs
+
+class ExportArguments(BaseModel):
+    r"""
+    Arguments pertaining to the model export.
+    """
+    export_quantization_bit: Optional[int] = Field(
+        default=None,
+        description="The number of bits to quantize the exported model."
+    )
+    export_quantization_dataset: Optional[str] = Field(
+        default=None,
+        description="Path to the dataset or dataset name to use in quantizing the exported model."
+    )
+    export_quantization_nsamples: int = Field(
+        default=128,
+        description="The number of samples used for quantization."
+    )
+    export_quantization_maxlen: int = Field(
+        default=1024,
+        description="The maximum length of the model inputs used for quantization."
+    )
+
+    
+class QuantizationArgs(ExportArguments, BaseModel):
+    r"""
+    Arguments pertaining to the quantization method.
+    """
+
+    quantization_method: Literal["bitsandbytes", "hqq", "eetq", "gptq"] = Field(
+        default="bitsandbytes",
+        description="Quantization method to use for on-the-fly quantization."
+    )
+    quantization_bit: Optional[int] = Field(
+        default=None,
+        description="The number of bits to quantize the model using on-the-fly quantization."
+    )
+    quantization_type: Literal["fp4", "nf4"] = Field(
+        default="nf4",
+        description="Quantization data type to use in bitsandbytes int4 training."
+    )
+    double_quantization: bool = Field(
+        default=True,
+        description="Whether or not to use double quantization in bitsandbytes int4 training."
+    )
+    quantization_device_map: Optional[Literal["auto"]] = Field(
+        default=None,
+        description="Device map used to infer the 4-bit quantized model, needs bitsandbytes>=0.43.0."
+    )
+
+    model_config = {
+        "extra": "forbid"
+    }
+
+
+class ModelArgs(BaseModel):
+    r"""
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune or infer.
+    """
+
+    model_name_or_path: str = Field(
+        description="Path to the model weight or identifier from huggingface.co/models or modelscope.cn/models.",
+        examples=["Qwen/Qwen2-0.5B"]
+    )
+    adapter_name_or_path: Optional[str] = Field(
+        default=None,
+        description="Path to the adapter weight or identifier from huggingface.co/models. "
+                "Use commas to separate multiple adapters."
+    )
+    adapter_folder: Optional[str] = Field(
+        default=None,
+        description="The folder containing the adapter weights to load."
+    )
+    cache_dir: Optional[str] = Field(
+        default=None,
+        description="Where to store the pre-trained models downloaded from huggingface.co or modelscope.cn."
+    )
+    resize_vocab: bool = Field(
+        default=False,
+        description="Whether or not to resize the tokenizer vocab and the embedding layers."
+    )
+    split_special_tokens: bool = Field(
+        default=False,
+        description="Whether or not the special tokens should be split during the tokenization process."
+    )
+    new_special_tokens: Optional[str] = Field(
+        default=None,
+        description="Special tokens to be added into the tokenizer. Use commas to separate multiple tokens."
+    )
+    model_revision: str = Field(
+        default="main",
+        description="The specific model version to use (can be a branch name, tag name or commit id)."
+    )
+    low_cpu_mem_usage: bool = Field(
+        default=True,
+        description="Whether or not to use memory-efficient model loading."
+    )
+    flash_attn: Literal["auto", "disabled", "sdpa", "fa2"] = Field(
+        default="auto",
+        description="Enable FlashAttention for faster training and inference."
+    )
+    use_unsloth: bool = Field(
+        default=False,
+        description="Whether or not to use unsloth's optimization for the LoRA training."
+    )
+    use_unsloth_gc: bool = Field(
+        default=False,
+        description="Whether or not to use unsloth's gradient checkpointing."
+    )
+    enable_liger_kernel: bool = Field(
+        default=False,
+        description="Whether or not to enable liger kernel for faster training."
+    )
+    disable_gradient_checkpointing: bool = Field(
+        default=False,
+        description="Whether or not to disable gradient checkpointing."
+    )
+    upcast_layernorm: bool = Field(
+        default=False,
+        description="Whether or not to upcast the layernorm weights in fp32."
+    )
+    upcast_lmhead_output: bool = Field(
+        default=False,
+        description="Whether or not to upcast the output of lm_head in fp32."
+    )
+    offload_folder: str = Field(
+        default="offload",
+        description="Path to offload model weights."
+    )
+    use_cache: bool = Field(
+        default=True,
+        description="Whether or not to use KV cache in generation."
+    )
+    infer_dtype: Literal["auto", "float16", "bfloat16", "float32"] = Field(
+        default="auto",
+        description="Data type for model weights and activations at inference."
+    )
+    hf_hub_token: Optional[str] = Field(
+        default=None,
+        description="Auth token to log in with Hugging Face Hub."
+    )
+    print_param_status: bool = Field(
+        default=False,
+        description="For debugging purposes, print the status of the parameters in the model."
+    )
+    output_attentions: bool = Field(
+        default=False,
+        description="Whether to output the attention mask for the model during forward pass"
+    )
+    quantization_args: Optional[QuantizationArgs] = Field(
+        default=QuantizationArgs(),
+        description="Arguments related to quantization"
+    )
+    finetuning_args: Optional[FinetuningArgs] = Field(
+        default=FinetuningArgs(),
+        description="Arguments related to finetuning (LoRA, Freeze etc.)"
+    )
+    deepspeed_config: str = Field(
+        default='./deepspeed/zero0.json',
+        description="Path to deepspeed config file. Defaults to stage-0 no optimization (Training/Inference)."
+    )
+    chat_template: Optional[str] = Field(
+        default=None,
+        description="Chat template to use when loading the model's tokenizer"
+    )
+
+    model_config = {
+        "extra": "forbid"
+    }