Distillflow 0.2.0__py3-none-any.whl

distillflow/model/quantization.py

@@ -0,0 +1,193 @@
+import os
+from enum import Enum, unique
+from random import random
+from typing import Any, Dict, Optional, List
+
+import torch
+from datasets import load_dataset
+from transformers import BitsAndBytesConfig, EetqConfig, HqqConfig, GPTQConfig
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.modeling_utils import is_fsdp_enabled
+from transformers.utils.versions import require_version
+
+from .args import QuantizationArgs
+from ..common.logger import get_logger
+from ..common import get_current_device
+
+from transformers import PretrainedConfig, PreTrainedTokenizer
+
+
+logger = get_logger(__name__)
+
+
+@unique
+class QuantizationMethod(str, Enum):
+    r"""
+    Borrowed from `transformers.utils.quantization_config.QuantizationMethod`.
+    """
+
+    BITS_AND_BYTES = "bitsandbytes"
+    GPTQ = "gptq"
+    AWQ = "awq"
+    AQLM = "aqlm"
+    QUANTO = "quanto"
+    EETQ = "eetq"
+    HQQ = "hqq"
+
+FILEEXT2TYPE = {
+    "arrow": "arrow",
+    "csv": "csv",
+    "json": "json",
+    "jsonl": "json",
+    "parquet": "parquet",
+    "txt": "text",
+}
+
+def _get_quantization_dataset(tokenizer: "PreTrainedTokenizer", model_args: "ModelArguments") -> List[Dict[str, Any]]:
+    r"""
+    Prepares the tokenized dataset to perform AutoGPTQ. Do not use tensor output for JSON serialization.
+    """
+    if os.path.isfile(model_args.export_quantization_dataset):
+        data_path = FILEEXT2TYPE.get(model_args.export_quantization_dataset.split(".")[-1], None)
+        data_files = model_args.export_quantization_dataset
+    else:
+        data_path = model_args.export_quantization_dataset
+        data_files = None
+
+    dataset = load_dataset(
+        path=data_path,
+        data_files=data_files,
+        split="train",
+        cache_dir=model_args.cache_dir,
+        token=model_args.hf_hub_token,
+    )
+
+    samples = []
+    maxlen = model_args.export_quantization_maxlen
+    for _ in range(model_args.export_quantization_nsamples):
+        n_try = 0
+        while True:
+            if n_try > 100:
+                raise ValueError("Cannot find satisfying example, considering decrease `export_quantization_maxlen`.")
+
+            sample_idx = random.randint(0, len(dataset) - 1)
+            sample: Dict[str, "torch.Tensor"] = tokenizer(dataset[sample_idx]["text"], return_tensors="pt")
+            n_try += 1
+            if sample["input_ids"].size(1) > maxlen:
+                break  # TODO: fix large maxlen
+
+        word_idx = random.randint(0, sample["input_ids"].size(1) - maxlen - 1)
+        input_ids = sample["input_ids"][:, word_idx : word_idx + maxlen]
+        attention_mask = sample["attention_mask"][:, word_idx : word_idx + maxlen]
+        samples.append({"input_ids": input_ids.tolist(), "attention_mask": attention_mask.tolist()})
+
+    return samples
+
+
+def _configure_quantization(
+        config: PretrainedConfig,
+        tokenizer: PreTrainedTokenizer,
+        quantization_args: QuantizationArgs,
+        init_kwargs: Dict[str, Any],
+        torch_dtype: Optional[torch.dtype]
+) -> None:
+    r"""
+    Priority: PTQ-quantized (train/infer) > AutoGPTQ (export) > On-the-fly quantization (train/infer)
+    """
+
+    if quantization_args is None:
+        return
+    if getattr(config, "quantization_config", None):  # ptq
+        if quantization_args.quantization_bit is not None:
+            logger.warning("`quantization_bit` will not affect on the PTQ-quantized models.")
+
+        if is_deepspeed_zero3_enabled() or is_fsdp_enabled():
+            raise ValueError("DeepSpeed ZeRO-3 or FSDP is incompatible with PTQ-quantized models.")
+
+        quantization_config: Dict[str, Any] = getattr(config, "quantization_config", None)
+        quant_method = quantization_config.get("quant_method", "")
+
+        if quant_method == QuantizationMethod.GPTQ:
+            require_version("auto_gptq>=0.5.0", "To fix: pip install auto_gptq>=0.5.0")
+            quantization_config.pop("disable_exllama", None)  # remove deprecated args
+            quantization_config["use_exllama"] = False  # disable exllama
+
+        if quant_method == QuantizationMethod.AWQ:
+            require_version("autoawq", "To fix: pip install autoawq")
+
+        if quant_method == QuantizationMethod.AQLM:
+            require_version("aqlm>=1.1.0", "To fix: pip install aqlm[gpu]>=1.1.0")
+            quantization_config["bits"] = 2
+
+        quant_bits = quantization_config.get("bits", "?")
+        logger.info("Loading {}-bit {}-quantized model.".format(quant_bits, quant_method.upper()))
+
+    elif quantization_args.export_quantization_bit is not None:  # auto-gptq
+        if quantization_args.export_quantization_bit not in [8, 4, 3, 2]:
+            raise ValueError("AutoGPTQ only accepts 2/3/4/8-bit quantization.")
+
+        require_version("optimum>=1.17.0", "To fix: pip install optimum>=1.17.0")
+        require_version("auto_gptq>=0.5.0", "To fix: pip install auto_gptq>=0.5.0")
+        from accelerate.utils import get_max_memory
+
+        if getattr(config, "model_type", None) == "chatglm":
+            raise ValueError("ChatGLM model is not supported yet.")
+
+        init_kwargs["quantization_config"] = GPTQConfig(
+            bits=quantization_args.export_quantization_bit,
+            dataset=_get_quantization_dataset(tokenizer, quantization_args),
+        )
+        init_kwargs["device_map"] = "auto"
+        init_kwargs["max_memory"] = get_max_memory()
+        logger.info("Quantizing model to {} bit with AutoGPTQ.".format(quantization_args.export_quantization_bit))
+    elif quantization_args.quantization_bit is not None:  # on-the-fly
+        if quantization_args.quantization_method == QuantizationMethod.BITS_AND_BYTES.value:
+            if quantization_args.quantization_bit == 8:
+                require_version("bitsandbytes>=0.37.0", "To fix: pip install bitsandbytes>=0.37.0")
+                init_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
+            elif quantization_args.quantization_bit == 4:
+                require_version("bitsandbytes>=0.39.0", "To fix: pip install bitsandbytes>=0.39.0")
+                init_kwargs["quantization_config"] = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch_dtype,
+                    bnb_4bit_use_double_quant=quantization_args.double_quantization,
+                    bnb_4bit_quant_type=quantization_args.quantization_type,
+                    bnb_4bit_quant_storage=torch_dtype,  # crucial for fsdp+qlora
+                )
+            else:
+                raise ValueError("Bitsandbytes only accepts 4-bit or 8-bit quantization.")
+
+            # Do not assign device map if:
+            # 1. deepspeed zero3 or fsdp (train)
+            # 2. auto quantization device map (inference)
+            if is_deepspeed_zero3_enabled() or is_fsdp_enabled() or quantization_args.quantization_device_map == "auto":
+                if quantization_args.quantization_bit != 4:
+                    raise ValueError("Only 4-bit quantized model can use fsdp+qlora or auto device map.")
+
+                require_version("bitsandbytes>=0.43.0", "To fix: pip install bitsandbytes>=0.43.0")
+            else:
+                init_kwargs["device_map"] = {"": get_current_device()}  # change auto device map for inference
+
+            logger.info("Quantizing model to {} bit with bitsandbytes.".format(quantization_args.quantization_bit))
+        elif quantization_args.quantization_method == QuantizationMethod.HQQ.value:
+            if quantization_args.quantization_bit not in [8, 6, 5, 4, 3, 2, 1]:
+                raise ValueError("HQQ only accepts 1/2/3/4/5/6/8-bit quantization.")
+
+            if is_deepspeed_zero3_enabled() or is_fsdp_enabled():
+                raise ValueError("HQQ quantization is incompatible with DeepSpeed ZeRO-3 or FSDP.")
+
+            require_version("hqq", "To fix: pip install hqq")
+            init_kwargs["quantization_config"] = HqqConfig(
+                nbits=quantization_args.quantization_bit, quant_zero=False, quant_scale=False, axis=0
+            )  # use ATEN kernel (axis=0) for performance
+            logger.info("Quantizing model to {} bit with HQQ.".format(quantization_args.quantization_bit))
+        elif quantization_args.quantization_method == QuantizationMethod.EETQ.value:
+            if quantization_args.quantization_bit != 8:
+                raise ValueError("EETQ only accepts 8-bit quantization.")
+
+            if is_deepspeed_zero3_enabled() or is_fsdp_enabled():
+                raise ValueError("EETQ quantization is incompatible with DeepSpeed ZeRO-3 or FSDP.")
+
+            require_version("eetq", "To fix: pip install eetq")
+            init_kwargs["quantization_config"] = EetqConfig()
+            logger.info("Quantizing model to {} bit with EETQ.".format(quantization_args.quantization_bit))

distillflow/model/tokenizer.py

@@ -0,0 +1,55 @@
+from types import MethodType
+from typing import Dict, Any
+
+from transformers import AutoTokenizer, PreTrainedTokenizerBase, PreTrainedTokenizer
+
+from distillflow.common import get_logger
+from distillflow.model.args import ModelArgs
+
+logger = get_logger(__name__)
+
+def tokenizer_init_kwargs(model_args: ModelArgs) -> Dict[str, Any]:
+    r"""
+    Gets arguments to load config/tokenizer/model.
+
+    Note: including inplace operation of model_args.
+    """
+    return {
+        "trust_remote_code": True,
+        "cache_dir": model_args.cache_dir,
+        "revision": model_args.model_revision,
+        "token": model_args.hf_hub_token,
+    }
+
+def load_tokenizer(model_args: ModelArgs, template: str = None, padding_side='right') -> PreTrainedTokenizer:
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path,
+            split_special_tokens=model_args.split_special_tokens,
+            padding_side=padding_side,
+            **tokenizer_init_kwargs(model_args),
+        )
+    except Exception as e:
+        raise OSError("Failed to load tokenizer.") from e
+
+    if model_args.new_special_tokens is not None:
+        print(model_args.new_special_tokens)
+        # exit()
+        num_added_tokens = tokenizer.add_special_tokens(
+            dict(additional_special_tokens=model_args.new_special_tokens.split(',')),
+            replace_additional_special_tokens=False,
+        )
+        logger.info("Add {} to special tokens.".format(model_args.new_special_tokens))
+        if num_added_tokens > 0 and not model_args.resize_vocab:
+            model_args.resize_vocab = True
+            logger.warning("New tokens have been added, changed `resize_vocab` to True.")
+    else:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    if "PreTrainedTokenizerBase" not in str(tokenizer._pad.__func__):
+        tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer)
+
+    if model_args.chat_template is not None:
+        tokenizer.chat_template = model_args.chat_template
+
+    return tokenizer

distillflow/model/unsloth.py

@@ -0,0 +1,92 @@
+from typing import Any, Dict, Optional
+
+import torch
+
+from ..common import get_current_device, infer_optim_dtype
+from ..common.logger import get_logger
+
+from transformers import PretrainedConfig, PreTrainedModel
+
+from .args import ModelArgs
+
+logger = get_logger(__name__)
+
+
+def _get_unsloth_kwargs(
+        config: PretrainedConfig, model_name_or_path: str, model_args: ModelArgs
+) -> Dict[str, Any]:
+    if model_args.infer_dtype != "auto":
+        torch_dtype = getattr(torch, model_args.infer_dtype)
+    else:
+        torch_dtype = infer_optim_dtype(model_dtype=getattr(config, "torch_dtype", None))
+
+    return {
+        "model_name": model_name_or_path,
+        "max_seq_length": 4096,
+        "dtype": torch_dtype,
+        "load_in_4bit": model_args.quantization_args.quantization_bit == 4,
+        "token": model_args.hf_hub_token,
+        "device_map": {"": get_current_device()},
+        "rope_scaling": getattr(config, "rope_scaling", None),
+        "fix_tokenizer": False,
+        "trust_remote_code": True,
+        "use_gradient_checkpointing": "unsloth",
+    }
+
+
+def load_unsloth_pretrained_model(
+        config: PretrainedConfig, model_args: ModelArgs
+) -> Optional[PreTrainedModel]:
+    r"""
+    Optionally loads pretrained model with unsloth. Used in training.
+    """
+    from unsloth import FastLanguageModel
+
+    unsloth_kwargs = _get_unsloth_kwargs(config, model_args.model_name_or_path, model_args)
+    try:
+        model, _ = FastLanguageModel.from_pretrained(**unsloth_kwargs)
+    except NotImplementedError:
+        logger.warning("Unsloth does not support model type {}.".format(getattr(config, "model_type", None)))
+        model = None
+        model_args.use_unsloth = False
+
+    return model
+
+
+def get_unsloth_peft_model(
+        model: "PreTrainedModel", model_args: "ModelArgs", peft_kwargs: Dict[str, Any]
+) -> "PreTrainedModel":
+    r"""
+    Gets the peft model for the pretrained model with unsloth. Used in training.
+    """
+    from unsloth import FastLanguageModel
+
+    unsloth_peft_kwargs = {
+        "model": model,
+        "max_seq_length": 4096,
+        "use_gradient_checkpointing": "unsloth",
+    }
+    return FastLanguageModel.get_peft_model(**peft_kwargs, **unsloth_peft_kwargs)
+
+
+def load_unsloth_peft_model(
+        config: "PretrainedConfig", model_args: "ModelArgs", is_trainable: bool
+) -> "PreTrainedModel":
+    r"""
+    Loads peft model with unsloth. Used in both training and inference.
+    """
+    from unsloth import FastLanguageModel
+
+    unsloth_kwargs = _get_unsloth_kwargs(config, model_args.adapter_name_or_path[0], model_args)
+    try:
+        if not is_trainable:
+            unsloth_kwargs["use_gradient_checkpointing"] = False
+
+        model, _ = FastLanguageModel.from_pretrained(**unsloth_kwargs)
+    except NotImplementedError:
+        raise ValueError("Unsloth does not support model type {}.".format(getattr(config, "model_type", None)))
+
+    if not is_trainable:
+        FastLanguageModel.for_inference(model)
+
+    return model

distillflow/trainer/AdaptationLayer.py

@@ -0,0 +1,82 @@
+from typing import List, Dict
+
+import torch
+
+class AdaptationLayer(torch.nn.Module):
+    def __init__(self, student_dim,
+                 teacher_dim,
+                 num_student_layers: int,
+                 num_teacher_layers: int,
+                 strategy="interpolate",
+                 dtype=torch.bfloat16,
+                 selection_indices: List[int]=None,
+                 weights:List[List[int]] =None):
+        super().__init__()
+        self.projections = torch.nn.ModuleList([
+            torch.nn.Linear(student_dim, teacher_dim, dtype=dtype)
+            for _ in range(num_student_layers)
+        ])
+        # self.layer_mapping = self.create_layer_mapping(num_student_layers, num_teacher_layers)
+        self.layer_mapping = self.map_teacher_to_student_layers(num_student_layers, num_teacher_layers, strategy, selection_indices, weights)
+        self.dtype = dtype
+
+    def map_teacher_to_student_layers(self, num_student_layers, num_teacher_layers, strategy="select",
+                                      selection_indices:List[int]=None, weights:List[List[int]]=None) -> {}:
+        """
+        Maps teacher model layers to student model layers based on the specified strategy.
+
+        Args:
+            num_student_layers (int): Number of layers in the student model.
+            num_teacher_layers (int): Number of layers in the teacher model.
+            strategy (str): Layer mapping strategy ("direct", "select", "interpolate", "weighted").
+            selection_indices (list): Specific teacher layers to select (used when strategy="select").
+            weights (list of lists): Weights for combining teacher layers for each student layer
+                                     (used when strategy="weighted").
+
+        Returns:
+            list: List of mapping indices or weights from teacher layers to align with student layers.
+        """
+        if strategy == "direct":
+            # Direct one-to-one mapping
+            return {
+                i: i
+                for i in range(num_student_layers)
+            }
+
+        elif strategy == "select":
+            # Use specific teacher layers for mapping
+            if selection_indices is None:
+                raise ValueError("selection_indices must be provided for 'select' strategy.")
+            if len(selection_indices) != num_student_layers:
+                raise ValueError("Number of selection_indices must match num_student_layers.")
+            return {
+                i: teacher_layer
+                for i, teacher_layer in enumerate(selection_indices)
+            }
+
+        elif strategy == "interpolate":
+            # Interpolate teacher layers to match student layers
+            return {
+                i: round(i * (num_teacher_layers - 1) / (num_student_layers - 1))
+                for i in range(num_student_layers)
+            }
+        elif strategy == "weighted":
+            # Weighted combination of teacher layers for each student layer
+            if weights is None:
+                raise ValueError("weights must be provided for 'weighted' strategy.")
+            if len(weights) != num_student_layers:
+                raise ValueError("Number of weight sets must match num_student_layers.")
+            return {
+                i: {j: weight for j, weight in enumerate(weight_set) if weight > 0}
+                for i, weight_set in enumerate(weights)
+            }
+        else:
+            raise ValueError(f"Unknown strategy: {strategy}")
+
+    def forward(self, student_hidden_states):
+        adapted_hidden_states = []
+        for i, hidden_state in enumerate(student_hidden_states):
+            if i >= len(self.projections):
+                break
+            adapted_hidden_states.append(self.projections[i](hidden_state.to(self.dtype)))
+        return adapted_hidden_states

distillflow/trainer/args.py

@@ -0,0 +1,56 @@
+from typing import Literal, Optional, List, Dict
+
+from pydantic import BaseModel, Field
+from trl import SFTConfig
+
+class DistillArgs(BaseModel):
+    sft_config: SFTConfig = Field(
+        description="SFT Config, with hyperparameters required for training"
+    )
+    type: Literal["logits", "layers", "attention", "fine-tune"] = Field(
+        default="logits",
+        description="Type of distillation to perform on the given student model with the given dataset (default: logits)"
+    )
+    max_seq_length: Optional[int] = Field (
+        default=4096,
+        description="Maximum sequence length to use during training (default: 4096)",
+        examples=[1024, 2048, 4096]
+    )
+    dataset_text_field: Optional[str] = Field(
+        default="text",
+        description="The key to which the data is mapped",
+    )
+    temperature: Optional[float] = Field(
+        default=0.5,
+        description="Temperature"
+    )
+    alpha: Optional[float] = Field(
+        default=2.0,
+        description="alpha"
+    )
+    resume_from_checkpoint: Optional[str] = Field(
+        default=None,
+        description="Training checkpoint folder path to resume the training from which to resume the training"
+    )
+    strategy : Literal["direct", "select", "interpolate", "weighted"] = Field(
+        default = "interpolate",
+        description="Strategy to select when mapping the teacher and student layers/attention map"
+    )
+    selection_indices: Optional[List[int]] = Field(
+        default=None,
+        description="If selected strategy `select`, provides mapping between student and teacher model layer/attention map",
+        examples=[[1,2,0]] # mapping 0th student layer to 1st layer of teacher and so on
+    )
+    weights: Optional[List[List[int]]] = Field(
+        default=None,
+        description="If selected strategy `weighted`, provides the weights for each of the teacher layers to be used when computing student layer",
+        examples=[[[100,10,40,50], [10,5,50,25]]] # 0th of student is computed with teacher layer weights 100 for 0th layer, 10 for 1st layer and so on
+    )
+
+    remove_unused_columns: Optional[bool] = Field(
+        default=False,
+    )
+
+    model_config = {
+        "extra": "forbid"
+    }

distillflow/trainer/attention_distillation.py

@@ -0,0 +1,113 @@
+from typing import Union, Optional
+
+import torch
+from accelerate import Accelerator
+from datasets import IterableDataset
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from trl import SFTTrainer
+
+from distillflow.common import get_current_device
+from distillflow.datasets.loader import DatasetModule
+from distillflow.trainer.AdaptationLayer import AdaptationLayer
+from distillflow.trainer.args import DistillArgs
+
+class AttentionTrainer(SFTTrainer):
+    def __init__(self,
+                 accelerator: Accelerator,
+                 distill_args: DistillArgs,
+                 teacher_model: PreTrainedModel,
+                 model: PreTrainedModel,
+                 dataset_module: DatasetModule,
+                 tokenizer: PreTrainedTokenizerBase
+                 ):
+        self.teacher_model = teacher_model
+        self.distill_args = distill_args
+        train_dataset = dataset_module["train_dataset"]
+        eval_dataset = dataset_module["eval_dataset"]
+        self.device = get_current_device()
+        self.adaptation_layer = AdaptationLayer(
+            model.config.hidden_size,
+            teacher_model.config.hidden_size,
+            model.config.num_hidden_layers,
+            teacher_model.config.num_hidden_layers,
+            dtype=torch.float16,
+            strategy=distill_args.strategy,
+            selection_indices=distill_args.selection_indices,
+            weights=distill_args.weights
+        ).to(self.device)
+
+        if isinstance(train_dataset, IterableDataset) and distill_args.sft_config.max_steps == -1:
+            raise ValueError("max steps should be specified when using dataset with streaming mode enabled.")
+
+        distill_args.sft_config.max_length = distill_args.max_seq_length
+        distill_args.sft_config.dataset_text_field = distill_args.dataset_text_field
+
+        super().__init__(model=model, args=distill_args.sft_config, train_dataset=train_dataset,
+                         eval_dataset=eval_dataset, processing_class=tokenizer)
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # Forward pass for the student model
+        student_outputs = model(**inputs, output_attentions=True)
+
+        self.teacher_model = self.teacher_model.to(self.device) if self.device.type == "mps" else self.teacher_model
+
+        teacher_model = self.teacher_model.module if hasattr(self.teacher_model, 'module') else self.teacher_model
+
+        # Forward pass for the teacher model
+        with torch.no_grad():
+            teacher_outputs = teacher_model(**inputs, output_attentions=True)
+
+        # Primary task loss (e.g., cross-entropy loss)
+        loss = student_outputs.loss
+
+        # Attention-based distillation loss
+        attention_loss = self.compute_attention_loss(
+            teacher_outputs.attentions,
+            student_outputs.attentions
+        )
+
+        # Combine losses
+        total_loss = ((1 - self.distill_args.alpha) * loss + self.distill_args.alpha * attention_loss) / self.args.gradient_accumulation_steps
+
+        return (total_loss, student_outputs) if return_outputs else total_loss
+
+    def compute_attention_loss(self, teacher_attentions, student_attentions):
+        """
+        Compute attention-based distillation loss.
+
+        Args:
+            teacher_attentions: List of teacher model attention maps.
+            student_attentions: List of student model attention maps.
+
+        Returns:
+            Total attention distillation loss.
+        """
+        loss = 0.0
+        num_layers = len(student_attentions)
+
+        self.adaptation_layer = self.adaptation_layer.to(self.device)
+
+        for student_idx, teacher_idx in self.adaptation_layer.layer_mapping.items():
+            if self.distill_args.strategy == "weighted":
+                teacher_attention = torch.zeros_like(teacher_attentions[0])
+                for idx, weight in teacher_idx.items():
+                    teacher_attention = weight * teacher_attentions[idx]
+            else:
+                teacher_attention = teacher_attentions[teacher_idx]
+
+            student_attention = student_attentions[student_idx]
+            # Align dimensions if needed
+            if teacher_attention.size() != student_attention.size():
+                teacher_attention = teacher_attention.mean(dim=1, keepdim=True).expand(-1, student_attention.size(1), -1, -1)
+                teacher_attention = F.interpolate(
+                    teacher_attention,
+                    size=student_attention.size()[-2:],  # Resize spatial dimensions
+                    mode="bilinear",
+                    align_corners=False,
+                )
+
+            # MSE Loss for attention maps
+            loss += F.mse_loss(student_attention, teacher_attention)
+
+        return loss / num_layers

distillflow/trainer/fine_tuning.py

@@ -0,0 +1,51 @@
+from accelerate import Accelerator
+from datasets import IterableDataset
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from trl import SFTTrainer
+import torch
+import torch.nn.functional as F
+
+from .args import DistillArgs
+from ..common import get_current_device
+from ..datasets.loader import DatasetModule
+
+class FineTuning(SFTTrainer):
+    def __init__(self,
+                 accelerator: Accelerator,
+                 distill_args: DistillArgs,
+                 teacher_model: PreTrainedModel,
+                 model: PreTrainedModel,
+                 dataset_module: DatasetModule,
+                 tokenizer: PreTrainedTokenizerBase
+                 ):
+        self.accelerator = accelerator
+        self.distill_args = distill_args
+        train_dataset = dataset_module["train_dataset"]
+        eval_dataset = dataset_module["eval_dataset"]
+        self.device = get_current_device()
+        if self.device.type == 'mps':
+            # Explicitly place the models on device since accelerate prepare does not work on MPS.
+            model = model.to(self.device)
+
+        if isinstance(train_dataset, IterableDataset) and distill_args.sft_config.max_steps == -1:
+            raise ValueError("max steps should be specified when using dataset with streaming mode enabled.")
+
+        distill_args.sft_config.max_length = distill_args.max_seq_length
+        distill_args.sft_config.dataset_text_field = distill_args.dataset_text_field
+
+        super().__init__(model=model, args=distill_args.sft_config, train_dataset=train_dataset,
+                         eval_dataset=eval_dataset, processing_class=tokenizer)
+
+
+    def output(self, model, inputs, no_grad):
+        if no_grad:
+            with torch.no_grad():
+                return model(**inputs)
+        else:
+            return model(**inputs)
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        student_model = model.module if hasattr(model, 'module') else model
+        student_outputs = self.output(student_model, inputs, False)
+
+        return student_outputs.loss