/ray_run_name`."},
+ )
+ ray_storage_path: str = field(
+ default="./saves",
+ metadata={"help": "The storage path to save training results to"},
+ )
+ ray_num_workers: int = field(
+ default=1,
+ metadata={"help": "The number of workers for Ray training. Default is 1 worker."},
+ )
+ resources_per_worker: Union[dict, str] = field(
+ default_factory=lambda: {"GPU": 1},
+ metadata={"help": "The resources per worker for Ray training. Default is to use 1 GPU per worker."},
+ )
+ placement_strategy: Literal["SPREAD", "PACK", "STRICT_SPREAD", "STRICT_PACK"] = field(
+ default="PACK",
+ metadata={"help": "The placement strategy for Ray training. Default is PACK."},
+ )
+ ray_init_kwargs: Optional[dict] = field(
+ default=None,
+ metadata={"help": "The arguments to pass to ray.init for Ray training. Default is None."},
+ )
+
+ def __post_init__(self):
+ self.use_ray = use_ray()
+ if isinstance(self.resources_per_worker, str) and self.resources_per_worker.startswith("{"):
+ self.resources_per_worker = _convert_str_dict(json.loads(self.resources_per_worker))
+
+
+@dataclass
+class TrainingArguments(RayArguments, Seq2SeqTrainingArguments):
+ r"""Arguments pertaining to the trainer."""
+
+ def __post_init__(self):
+ Seq2SeqTrainingArguments.__post_init__(self)
+ RayArguments.__post_init__(self)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/launcher.py b/post-training/LLaMA-Factory/src/llamafactory/launcher.py
new file mode 100644
index 0000000..169b042
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/launcher.py
@@ -0,0 +1,23 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from llamafactory.train.tuner import run_exp # use absolute import
+
+
+def launch():
+ run_exp()
+
+
+if __name__ == "__main__":
+ launch()
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/model/__init__.py
new file mode 100644
index 0000000..71d4f47
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from .loader import load_config, load_model, load_tokenizer
+from .model_utils.misc import find_all_linear_modules
+from .model_utils.quantization import QuantizationMethod
+from .model_utils.valuehead import load_valuehead_params
+
+
+__all__ = [
+ "QuantizationMethod",
+ "find_all_linear_modules",
+ "load_config",
+ "load_model",
+ "load_tokenizer",
+ "load_valuehead_params",
+]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/adapter.py b/post-training/LLaMA-Factory/src/llamafactory/model/adapter.py
new file mode 100644
index 0000000..bbc0056
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/adapter.py
@@ -0,0 +1,302 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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 re
+from typing import TYPE_CHECKING
+
+import torch
+from peft import LoraConfig, LoraModel, PeftModel, TaskType, get_peft_model
+from transformers.integrations import is_deepspeed_zero3_enabled
+
+from ..extras import logging
+from .model_utils.misc import find_all_linear_modules, find_expanded_modules
+from .model_utils.quantization import QuantizationMethod
+from .model_utils.unsloth import get_unsloth_peft_model, load_unsloth_peft_model
+from .model_utils.visual import get_forbidden_modules, patch_target_modules
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedModel
+
+ from ..hparams import FinetuningArguments, ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def _setup_full_tuning(
+ model: "PreTrainedModel",
+ finetuning_args: "FinetuningArguments",
+ is_trainable: bool,
+ cast_trainable_params_to_fp32: bool,
+) -> None:
+ if not is_trainable:
+ return
+
+ logger.info_rank0("Fine-tuning method: Full")
+ forbidden_modules = get_forbidden_modules(model.config, finetuning_args)
+ for name, param in model.named_parameters():
+ if not any(forbidden_module in name for forbidden_module in forbidden_modules):
+ if cast_trainable_params_to_fp32:
+ param.data = param.data.to(torch.float32)
+ else:
+ param.requires_grad_(False)
+
+
+def _setup_freeze_tuning(
+ model: "PreTrainedModel",
+ finetuning_args: "FinetuningArguments",
+ is_trainable: bool,
+ cast_trainable_params_to_fp32: bool,
+) -> None:
+ if not is_trainable:
+ return
+
+ logger.info_rank0("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(
+ f"`num_layers` {num_layers} should be "
+ f"divisible by `num_layer_trainable` {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)
+
+ forbidden_modules = get_forbidden_modules(model.config, finetuning_args)
+ for name, param in model.named_parameters():
+ if any(trainable_layer in name for trainable_layer in trainable_layers) and not any(
+ forbidden_module in name for forbidden_module in forbidden_modules
+ ):
+ if cast_trainable_params_to_fp32:
+ param.data = param.data.to(torch.float32)
+ else:
+ param.requires_grad_(False)
+
+ logger.info_rank0("Set trainable layers: {}".format(",".join(trainable_layers)))
+
+
+def _setup_lora_tuning(
+ config: "PretrainedConfig",
+ model: "PreTrainedModel",
+ model_args: "ModelArguments",
+ finetuning_args: "FinetuningArguments",
+ is_trainable: bool,
+ cast_trainable_params_to_fp32: bool,
+) -> "PeftModel":
+ if is_trainable:
+ logger.info_rank0("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_rank0(f"Merged {len(adapter_to_merge)} adapter(s).")
+
+ 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_rank0("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, finetuning_args.freeze_vision_tower)
+ else:
+ target_modules = finetuning_args.lora_target
+
+ if finetuning_args.use_llama_pro:
+ target_modules = find_expanded_modules(model, target_modules, finetuning_args.freeze_trainable_layers)
+
+ target_modules = patch_target_modules(model, finetuning_args, target_modules)
+
+ if (
+ finetuning_args.use_dora
+ and getattr(model, "quantization_method", None) is not None
+ and getattr(model, "quantization_method", None) != QuantizationMethod.BITS_AND_BYTES
+ ):
+ raise ValueError("DoRA is not compatible with PTQ-quantized models.")
+
+ 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_rank0("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_rank0("Using PiSSA initialization.")
+ peft_kwargs["init_lora_weights"] = "pissa"
+ else:
+ logger.info_rank0(f"Using PiSSA initialization with FSVD steps {finetuning_args.pissa_iter}.")
+ peft_kwargs["init_lora_weights"] = f"pissa_niter_{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: "ModelArguments",
+ finetuning_args: "FinetuningArguments",
+ is_trainable: bool,
+) -> "PreTrainedModel":
+ r"""Initialize the adapters.
+
+ Support full-parameter, freeze and LoRA training.
+
+ Note that the trainable parameters must be cast to float32.
+ """
+ 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 (zero3 already in fp32)
+ cast_trainable_params_to_fp32 = False
+ if not is_trainable:
+ pass
+ elif finetuning_args.pure_bf16 or finetuning_args.use_badam:
+ logger.info_rank0("Pure bf16 / BAdam detected, remaining trainable params in half precision.")
+ elif model_args.quantization_bit is None and is_deepspeed_zero3_enabled():
+ logger.info_rank0("DeepSpeed ZeRO3 detected, remaining trainable params in float32.")
+ else:
+ logger.info_rank0("Upcasting trainable params to float32.")
+ cast_trainable_params_to_fp32 = True
+
+ if finetuning_args.finetuning_type == "full":
+ _setup_full_tuning(model, finetuning_args, 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, finetuning_args, is_trainable, cast_trainable_params_to_fp32
+ )
+ else:
+ raise NotImplementedError(f"Unknown finetuning type: {finetuning_args.finetuning_type}.")
+
+ return model
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/loader.py b/post-training/LLaMA-Factory/src/llamafactory/model/loader.py
new file mode 100644
index 0000000..c87e796
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/loader.py
@@ -0,0 +1,219 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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 os
+from typing import TYPE_CHECKING, Any, Optional, TypedDict
+
+import torch
+from transformers import (
+ AutoConfig,
+ AutoModelForCausalLM,
+ AutoModelForSeq2SeqLM,
+ AutoModelForTextToWaveform,
+ AutoModelForVision2Seq,
+ AutoProcessor,
+ AutoTokenizer,
+)
+from trl import AutoModelForCausalLMWithValueHead
+
+from ..extras import logging
+from ..extras.misc import count_parameters, skip_check_imports, try_download_model_from_other_hub
+from ..extras.packages import is_transformers_version_greater_than
+from .adapter import init_adapter
+from .model_utils.liger_kernel import apply_liger_kernel
+from .model_utils.misc import register_autoclass
+from .model_utils.mod import convert_pretrained_model_to_mod, load_mod_pretrained_model
+from .model_utils.unsloth import load_unsloth_pretrained_model
+from .model_utils.valuehead import load_valuehead_params
+from .patcher import patch_config, patch_model, patch_processor, patch_tokenizer, patch_valuehead_model
+
+
+if is_transformers_version_greater_than("4.46.0"):
+ from transformers import AutoModelForImageTextToText
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
+
+ from ..hparams import FinetuningArguments, ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+class TokenizerModule(TypedDict):
+ tokenizer: "PreTrainedTokenizer"
+ processor: Optional["ProcessorMixin"]
+
+
+def _get_init_kwargs(model_args: "ModelArguments") -> dict[str, Any]:
+ r"""Get arguments to load config/tokenizer/model.
+
+ Note: including inplace operation of model_args.
+ """
+ skip_check_imports()
+ model_args.model_name_or_path = try_download_model_from_other_hub(model_args)
+ return {
+ "trust_remote_code": model_args.trust_remote_code,
+ "cache_dir": model_args.cache_dir,
+ "revision": model_args.model_revision,
+ "token": model_args.hf_hub_token,
+ }
+
+
+def load_tokenizer(model_args: "ModelArguments") -> "TokenizerModule":
+ r"""Load pretrained tokenizer and optionally loads processor.
+
+ Note: including inplace operation of model_args.
+ """
+ init_kwargs = _get_init_kwargs(model_args)
+ try:
+ tokenizer = AutoTokenizer.from_pretrained(
+ model_args.model_name_or_path,
+ use_fast=model_args.use_fast_tokenizer,
+ split_special_tokens=model_args.split_special_tokens,
+ padding_side="right",
+ **init_kwargs,
+ )
+ except ValueError: # try the fast one
+ tokenizer = AutoTokenizer.from_pretrained(
+ model_args.model_name_or_path,
+ use_fast=True,
+ padding_side="right",
+ **init_kwargs,
+ )
+ except Exception as e:
+ raise OSError("Failed to load tokenizer.") from e
+
+ patch_tokenizer(tokenizer, model_args)
+ try:
+ processor = AutoProcessor.from_pretrained(model_args.model_name_or_path, **init_kwargs)
+ patch_processor(processor, tokenizer, model_args)
+ except Exception as e:
+ logger.debug(f"Failed to load processor: {e}.")
+ processor = None
+
+ # Avoid load tokenizer, see:
+ # https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/auto/processing_auto.py#L324
+ if processor is not None and "Processor" not in processor.__class__.__name__:
+ logger.debug("The loaded processor is not an instance of Processor. Dropping it.")
+ processor = None
+
+ return {"tokenizer": tokenizer, "processor": processor}
+
+
+def load_config(model_args: "ModelArguments") -> "PretrainedConfig":
+ r"""Load model config."""
+ init_kwargs = _get_init_kwargs(model_args)
+ return AutoConfig.from_pretrained(model_args.model_name_or_path, **init_kwargs)
+
+
+def load_model(
+ tokenizer: "PreTrainedTokenizer",
+ model_args: "ModelArguments",
+ finetuning_args: "FinetuningArguments",
+ is_trainable: bool = False,
+ add_valuehead: bool = False,
+) -> "PreTrainedModel":
+ r"""Load pretrained model."""
+ init_kwargs = _get_init_kwargs(model_args)
+ config = load_config(model_args)
+ patch_config(config, tokenizer, model_args, init_kwargs, is_trainable)
+ apply_liger_kernel(config, model_args, is_trainable, require_logits=(finetuning_args.stage not in ["pt", "sft"]))
+
+ model = None
+ lazy_load = False
+ if model_args.use_unsloth:
+ if model_args.adapter_name_or_path is not None:
+ lazy_load = True
+ elif is_trainable:
+ model = load_unsloth_pretrained_model(config, model_args)
+
+ if model is None and not lazy_load:
+ init_kwargs["config"] = config
+ init_kwargs["pretrained_model_name_or_path"] = model_args.model_name_or_path
+
+ if model_args.mixture_of_depths == "load":
+ model = load_mod_pretrained_model(**init_kwargs)
+ else:
+ if type(config) in AutoModelForVision2Seq._model_mapping.keys(): # image-text
+ load_class = AutoModelForVision2Seq
+ elif (
+ is_transformers_version_greater_than("4.46.0")
+ and type(config) in AutoModelForImageTextToText._model_mapping.keys()
+ ): # image-text
+ load_class = AutoModelForImageTextToText
+ elif type(config) in AutoModelForSeq2SeqLM._model_mapping.keys(): # audio-text
+ load_class = AutoModelForSeq2SeqLM
+ elif type(config) in AutoModelForTextToWaveform._model_mapping.keys(): # audio hack for qwen2_5_omni
+ load_class = AutoModelForTextToWaveform
+ else:
+ load_class = AutoModelForCausalLM
+
+ if model_args.train_from_scratch:
+ model = load_class.from_config(config, trust_remote_code=model_args.trust_remote_code)
+ else:
+ model = load_class.from_pretrained(**init_kwargs)
+ if getattr(model.config, "model_type", None) == "qwen2_5_omni":
+ model = model.thinker # use part of Omni model
+
+ if model_args.mixture_of_depths == "convert":
+ model = convert_pretrained_model_to_mod(model, config, model_args)
+
+ if not lazy_load:
+ patch_model(model, tokenizer, model_args, is_trainable, add_valuehead)
+ register_autoclass(config, model, tokenizer)
+
+ model = init_adapter(config, model, model_args, finetuning_args, is_trainable)
+
+ if add_valuehead:
+ model = AutoModelForCausalLMWithValueHead.from_pretrained(model)
+ patch_valuehead_model(model)
+
+ if model_args.adapter_name_or_path is not None:
+ vhead_path = model_args.adapter_name_or_path[-1]
+ else:
+ vhead_path = model_args.model_name_or_path
+
+ vhead_params = load_valuehead_params(vhead_path, model_args)
+ if vhead_params is not None:
+ model.load_state_dict(vhead_params, strict=False)
+ logger.info_rank0(f"Loaded valuehead from checkpoint: {vhead_path}")
+
+ if not is_trainable:
+ model.requires_grad_(False)
+ for param in model.parameters():
+ if param.data.dtype == torch.float32 and model_args.compute_dtype != torch.float32:
+ param.data = param.data.to(model_args.compute_dtype)
+
+ model.eval()
+ else:
+ model.train()
+
+ trainable_params, all_param = count_parameters(model)
+ if is_trainable:
+ param_stats = (
+ f"trainable params: {trainable_params:,} || "
+ f"all params: {all_param:,} || trainable%: {100 * trainable_params / all_param:.4f}"
+ )
+ else:
+ param_stats = f"all params: {all_param:,}"
+
+ logger.info_rank0(param_stats)
+
+ if model_args.print_param_status and int(os.getenv("LOCAL_RANK", "0")) == 0:
+ for name, param in model.named_parameters():
+ print(f"name: {name}, dtype: {param.dtype}, device: {param.device}, trainable: {param.requires_grad}")
+
+ return model
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/attention.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/attention.py
new file mode 100644
index 0000000..0f95d3f
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/attention.py
@@ -0,0 +1,87 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from transformers.utils import is_flash_attn_2_available, is_torch_sdpa_available
+
+from ...extras import logging
+from ...extras.constants import AttentionFunction
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig
+
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def configure_attn_implementation(
+ config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool
+) -> None:
+ if getattr(config, "model_type", None) == "gemma2" and is_trainable:
+ if model_args.flash_attn == AttentionFunction.AUTO or model_args.flash_attn == AttentionFunction.FA2:
+ if is_flash_attn_2_available():
+ if model_args.flash_attn != AttentionFunction.FA2:
+ logger.warning_rank0("Gemma 2 should use flash attention 2, change `flash_attn` to fa2.")
+ model_args.flash_attn = AttentionFunction.FA2
+ else:
+ logger.warning_rank0("FlashAttention-2 is not installed, use eager attention.")
+ model_args.flash_attn = AttentionFunction.DISABLED
+ elif model_args.flash_attn == AttentionFunction.SDPA:
+ logger.warning_rank0(
+ "Gemma-2 should use soft-capping attention, while the SDPA attention does not support it."
+ )
+
+ if model_args.flash_attn == AttentionFunction.AUTO:
+ return
+
+ elif model_args.flash_attn == AttentionFunction.DISABLED:
+ requested_attn_implementation = "eager"
+
+ elif model_args.flash_attn == AttentionFunction.SDPA:
+ if not is_torch_sdpa_available():
+ logger.warning_rank0("torch>=2.1.1 is required for SDPA attention.")
+ return
+
+ requested_attn_implementation = "sdpa"
+ elif model_args.flash_attn == AttentionFunction.FA2:
+ if not is_flash_attn_2_available():
+ logger.warning_rank0("FlashAttention-2 is not installed.")
+ return
+
+ requested_attn_implementation = "flash_attention_2"
+ else:
+ raise NotImplementedError(f"Unknown attention type: {model_args.flash_attn}")
+
+ if getattr(config, "model_type", None) == "internlm2": # special case for custom models
+ setattr(config, "attn_implementation", requested_attn_implementation)
+ else:
+ setattr(config, "_attn_implementation", requested_attn_implementation)
+
+
+def print_attn_implementation(config: "PretrainedConfig") -> None:
+ if getattr(config, "model_type", None) == "internlm2": # special case for custom models
+ attn_implementation = getattr(config, "attn_implementation", None)
+ else:
+ attn_implementation = getattr(config, "_attn_implementation", None)
+
+ if attn_implementation == "flash_attention_2":
+ logger.info_rank0("Using FlashAttention-2 for faster training and inference.")
+ elif attn_implementation == "sdpa":
+ logger.info_rank0("Using torch SDPA for faster training and inference.")
+ else:
+ logger.info_rank0("Using vanilla attention implementation.")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/checkpointing.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/checkpointing.py
new file mode 100644
index 0000000..28e2a79
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/checkpointing.py
@@ -0,0 +1,174 @@
+# Copyright 2025 HuggingFace Inc., Daniel Han-Chen & the Unsloth team and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's Transformers and PEFT library,
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/modeling_utils.py
+# https://github.com/huggingface/peft/blob/v0.10.0/src/peft/utils/other.py
+# and the Unsloth library.
+# https://github.com/unslothai/unsloth/blob/July-2024/unsloth/models/_utils.py
+#
+# 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 inspect
+from functools import WRAPPER_ASSIGNMENTS, partial, wraps
+from types import MethodType
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+
+import torch
+
+from ...extras import logging
+from ...extras.constants import LAYERNORM_NAMES
+
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedModel
+
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def get_unsloth_gradient_checkpointing_func() -> Callable:
+ class UnslothGradientCheckpointing(torch.autograd.Function):
+ r"""Saves VRAM by smartly offloading to RAM."""
+
+ @staticmethod
+ @torch.cuda.amp.custom_fwd
+ def forward(
+ ctx: "torch.autograd.Function",
+ forward_function: "torch.Module",
+ hidden_states: "torch.Tensor",
+ *args: Union["torch.Tensor", Any],
+ ) -> "torch.Tensor":
+ saved_hidden_states = hidden_states.to("cpu", non_blocking=True)
+ with torch.no_grad():
+ output = forward_function(hidden_states, *args)
+
+ ctx.save_for_backward(saved_hidden_states)
+ ctx.forward_function = forward_function
+ ctx.args = args
+ return output
+
+ @staticmethod
+ @torch.cuda.amp.custom_bwd
+ def backward(ctx: "torch.autograd.Function", grad_output: "torch.Tensor") -> "torch.Tensor":
+ (hidden_states,) = ctx.saved_tensors
+ hidden_states = hidden_states.to("cuda", non_blocking=True).detach()
+ hidden_states.requires_grad_(True)
+ with torch.enable_grad():
+ (output,) = ctx.forward_function(hidden_states, *ctx.args)
+
+ torch.autograd.backward(output, grad_output)
+ return (None, hidden_states.grad) + (None,) * len(ctx.args)
+
+ return UnslothGradientCheckpointing.apply
+
+
+def get_custom_gradient_checkpointing_func(gradient_checkpointing_func: Callable) -> Callable:
+ r"""Only applies gradient checkpointing to trainable layers."""
+
+ @wraps(gradient_checkpointing_func, assigned=WRAPPER_ASSIGNMENTS + ("__self__",))
+ def custom_gradient_checkpointing_func(func: Callable, *args: Union["torch.Tensor", Any], **kwargs):
+ if isinstance(func, partial):
+ module: torch.nn.Module = func.func.__self__
+ else:
+ module: torch.nn.Module = func.__self__
+
+ has_grad = False
+ if any(param.requires_grad for param in module.parameters()):
+ has_grad = True
+ for arg in args:
+ if torch.is_tensor(arg) and torch.is_floating_point(arg):
+ arg.requires_grad_(True)
+ break # assume the first tensor is always the hidden states
+
+ if has_grad:
+ return gradient_checkpointing_func(func, *args, **kwargs)
+ else:
+ return func(*args, **kwargs)
+
+ return custom_gradient_checkpointing_func
+
+
+def _gradient_checkpointing_enable(
+ self: "PreTrainedModel",
+ gradient_checkpointing_kwargs: Optional[dict[str, Any]] = None,
+ use_unsloth_gc: bool = False,
+) -> None:
+ r"""Activates gradient checkpointing for the current model.
+
+ Modification of the original method to enable gradient checkpointing for block-wise optimizer.
+ """
+ from torch.utils.checkpoint import checkpoint
+
+ if not self.supports_gradient_checkpointing:
+ raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
+
+ if gradient_checkpointing_kwargs is None:
+ gradient_checkpointing_kwargs = {"use_reentrant": True}
+
+ if use_unsloth_gc:
+ gradient_checkpointing_func = get_unsloth_gradient_checkpointing_func()
+ else:
+ gradient_checkpointing_func = partial(checkpoint, **gradient_checkpointing_kwargs)
+
+ gradient_checkpointing_func = get_custom_gradient_checkpointing_func(gradient_checkpointing_func)
+ if "value" in inspect.signature(self._set_gradient_checkpointing).parameters: # old GC format
+ self.apply(partial(self._set_gradient_checkpointing, value=True))
+ self.enable_input_require_grads()
+ logger.warning_rank0_once("You are using the old GC format, some features (e.g. BAdam) will be invalid.")
+ else: # have already enabled input require gradients
+ self._set_gradient_checkpointing(enable=True, gradient_checkpointing_func=gradient_checkpointing_func)
+
+
+def _fp32_forward_post_hook(
+ module: "torch.nn.Module", args: tuple["torch.Tensor"], output: "torch.Tensor"
+) -> "torch.Tensor":
+ return output.to(torch.float32)
+
+
+def prepare_model_for_training(model: "PreTrainedModel", model_args: "ModelArguments") -> None:
+ r"""Prepare the model before training.
+
+ Include:
+ (1) cast the layernorm in fp32
+ (2) make output embedding layer require grads
+ (3) add the upcasting of the lm_head in fp32.
+ """
+ if model_args.upcast_layernorm:
+ logger.info_rank0("Upcasting layernorm weights in float32.")
+ for name, param in model.named_parameters():
+ if param.ndim == 1 and any(ln_name in name for ln_name in LAYERNORM_NAMES):
+ param.data = param.data.to(torch.float32)
+
+ if not model_args.disable_gradient_checkpointing:
+ if not getattr(model, "supports_gradient_checkpointing", False):
+ logger.warning_rank0("Current model does not support gradient checkpointing.")
+ else:
+ # use_reentrant=False might increase VRAM usage (have not been empirically verified yet)
+ # According to: https://github.com/huggingface/transformers/issues/28339
+ gradient_checkpointing_enable = partial(
+ _gradient_checkpointing_enable, use_unsloth_gc=model_args.use_unsloth_gc
+ )
+ model.gradient_checkpointing_enable = MethodType(gradient_checkpointing_enable, model)
+ model.gradient_checkpointing_enable(
+ gradient_checkpointing_kwargs={"use_reentrant": model_args.use_reentrant_gc}
+ )
+ setattr(model.config, "use_cache", False) # turn off when gradient checkpointing is enabled
+ logger.info_rank0("Gradient checkpointing enabled.")
+
+ if model_args.upcast_lmhead_output:
+ output_layer = model.get_output_embeddings()
+ if isinstance(output_layer, torch.nn.Linear) and output_layer.weight.dtype != torch.float32:
+ logger.info_rank0("Upcasting lm_head outputs in float32.")
+ output_layer.register_forward_hook(_fp32_forward_post_hook)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/embedding.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/embedding.py
new file mode 100644
index 0000000..c10e34f
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/embedding.py
@@ -0,0 +1,70 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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 math
+from contextlib import nullcontext
+from typing import TYPE_CHECKING
+
+import torch
+from transformers.integrations import is_deepspeed_zero3_enabled
+
+from ...extras import logging
+
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedModel, PreTrainedTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+
+def _noisy_mean_initialization(embed_weight: "torch.Tensor", num_new_tokens: int) -> None:
+ embedding_dim = embed_weight.size(1)
+ avg_weight = embed_weight[:-num_new_tokens].mean(dim=0, keepdim=True)
+ noise_weight = torch.empty_like(embed_weight[-num_new_tokens:])
+ noise_weight.normal_(mean=0, std=(1.0 / math.sqrt(embedding_dim)))
+ embed_weight[-num_new_tokens:] = avg_weight + noise_weight
+
+
+def resize_embedding_layer(model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer") -> None:
+ r"""Resize token embeddings."""
+ if is_deepspeed_zero3_enabled():
+ import deepspeed # type: ignore
+
+ params = [model.get_input_embeddings().weight]
+ if model.get_output_embeddings() is not None and not model.config.tie_word_embeddings:
+ params.append(model.get_output_embeddings().weight)
+
+ context_maybe_zero3 = deepspeed.zero.GatheredParameters(params, modifier_rank=0)
+ else:
+ context_maybe_zero3 = nullcontext()
+
+ with context_maybe_zero3:
+ current_embedding_size = model.get_input_embeddings().weight.size(0)
+
+ if len(tokenizer) > current_embedding_size:
+ if getattr(model, "quantization_method", None):
+ raise ValueError("Cannot resize embedding layers of a quantized model.")
+
+ if not isinstance(model.get_output_embeddings(), torch.nn.Linear):
+ raise ValueError("Current model does not support resizing embedding layers.")
+
+ model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64)
+ with context_maybe_zero3:
+ new_embedding_size = model.get_input_embeddings().weight.size(0)
+ num_new_tokens = new_embedding_size - current_embedding_size
+ _noisy_mean_initialization(model.get_input_embeddings().weight.data, num_new_tokens)
+ _noisy_mean_initialization(model.get_output_embeddings().weight.data, num_new_tokens)
+
+ logger.info_rank0(f"Resized token embeddings from {current_embedding_size} to {new_embedding_size}.")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/kv_cache.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/kv_cache.py
new file mode 100644
index 0000000..cd2c119
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/kv_cache.py
@@ -0,0 +1,44 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from ...extras import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig
+
+ from ...hparams import ModelArguments
+
+
+def configure_kv_cache(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None:
+ if not is_trainable:
+ setattr(config, "use_cache", model_args.use_cache)
+ if hasattr(config, "text_config"):
+ setattr(config.text_config, "use_cache", model_args.use_cache)
+
+ if model_args.use_cache:
+ logger.info_rank0("KV cache is enabled for faster generation.")
+ else:
+ logger.info_rank0("KV cache is disabled.")
+ else:
+ setattr(config, "use_cache", False)
+ if hasattr(config, "text_config"):
+ setattr(config.text_config, "use_cache", False)
+
+ logger.info_rank0("KV cache is disabled during training.")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/liger_kernel.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/liger_kernel.py
new file mode 100644
index 0000000..486c54e
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/liger_kernel.py
@@ -0,0 +1,77 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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 inspect
+from typing import TYPE_CHECKING
+
+from ...extras import logging
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig
+
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def apply_liger_kernel(
+ config: "PretrainedConfig",
+ model_args: "ModelArguments",
+ is_trainable: bool,
+ require_logits: bool,
+) -> None:
+ if not is_trainable or not model_args.enable_liger_kernel:
+ return
+
+ model_type = getattr(config, "model_type", None)
+ if model_type == "gemma":
+ from liger_kernel.transformers import apply_liger_kernel_to_gemma as apply_liger_kernel
+ elif model_type == "gemma2":
+ from liger_kernel.transformers import apply_liger_kernel_to_gemma2 as apply_liger_kernel
+ elif model_type == "gemma3":
+ from liger_kernel.transformers import apply_liger_kernel_to_gemma3 as apply_liger_kernel
+ elif model_type == "gemma3_text":
+ from liger_kernel.transformers import apply_liger_kernel_to_gemma3_text as apply_liger_kernel
+ elif model_type == "paligemma":
+ from liger_kernel.transformers import apply_liger_kernel_to_paligemma as apply_liger_kernel
+ elif model_type == "llama":
+ from liger_kernel.transformers import apply_liger_kernel_to_llama as apply_liger_kernel
+ elif model_type == "mistral":
+ from liger_kernel.transformers import apply_liger_kernel_to_mistral as apply_liger_kernel
+ elif model_type == "mixtral":
+ from liger_kernel.transformers import apply_liger_kernel_to_mixtral as apply_liger_kernel
+ elif model_type == "mllama":
+ from liger_kernel.transformers import apply_liger_kernel_to_mllama as apply_liger_kernel
+ elif model_type == "phi3":
+ from liger_kernel.transformers import apply_liger_kernel_to_phi3 as apply_liger_kernel
+ elif model_type == "qwen2":
+ from liger_kernel.transformers import apply_liger_kernel_to_qwen2 as apply_liger_kernel
+ elif model_type == "qwen2_vl":
+ from liger_kernel.transformers import apply_liger_kernel_to_qwen2_vl as apply_liger_kernel
+ elif model_type == "qwen2_5_vl":
+ from liger_kernel.transformers import apply_liger_kernel_to_qwen2_5_vl as apply_liger_kernel
+ else:
+ logger.warning_rank0("Current model does not support liger kernel.")
+ return
+
+ if require_logits and "fused_linear_cross_entropy" in inspect.signature(apply_liger_kernel).parameters:
+ logger.info_rank0("Current training stage does not support chunked cross entropy.")
+ kwargs = {"fused_linear_cross_entropy": False, "cross_entropy": True}
+ else:
+ kwargs = {}
+
+ apply_liger_kernel(**kwargs)
+ logger.info_rank0("Liger kernel has been applied to the model.")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/longlora.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/longlora.py
new file mode 100644
index 0000000..f7c36ee
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/longlora.py
@@ -0,0 +1,370 @@
+# Copyright 2025 EleutherAI, HuggingFace Inc., Yukang Chen, and the LlamaFactory team.
+#
+# This code is based on the EleutherAI's GPT-NeoX and the HuggingFace's Transformers libraries.
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llama/modeling_llama.py
+# This code is also inspired by the original LongLoRA implementation.
+# https://github.com/dvlab-research/LongLoRA/blob/main/llama_attn_replace.py
+#
+# 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 math
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch.nn as nn
+import transformers
+
+from ...extras import logging
+from ...extras.constants import SUPPORTED_CLASS_FOR_S2ATTN
+from ...extras.misc import check_version
+from ...extras.packages import is_transformers_version_greater_than
+
+
+if not is_transformers_version_greater_than("4.48.0"):
+ from transformers.modeling_flash_attention_utils import _flash_attention_forward
+ from transformers.models.llama.modeling_llama import (
+ Cache,
+ LlamaAttention,
+ LlamaFlashAttention2,
+ LlamaSdpaAttention,
+ apply_rotary_pos_emb,
+ repeat_kv,
+ )
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig
+
+ from ...hparams import ModelArguments
+
+
+transformers_logger = transformers.utils.logging.get_logger(__name__)
+
+
+# Modified from:
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llama/modeling_llama.py
+def llama_attention_forward(
+ self: "LlamaAttention",
+ hidden_states: "torch.Tensor",
+ attention_mask: Optional["torch.Tensor"] = None,
+ position_ids: Optional["torch.LongTensor"] = None,
+ past_key_value: Optional["Cache"] = None,
+ output_attentions: bool = False,
+ cache_position: Optional["torch.LongTensor"] = None,
+ position_embeddings: Optional[tuple["torch.Tensor", "torch.Tensor"]] = None,
+ **kwargs,
+) -> tuple["torch.Tensor", Optional["torch.Tensor"], Optional[tuple["torch.Tensor"]]]:
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states: torch.Tensor = self.q_proj(hidden_states)
+ key_states: torch.Tensor = self.k_proj(hidden_states)
+ value_states: torch.Tensor = self.v_proj(hidden_states)
+
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+ if position_embeddings is None:
+ cos, sin = self.rotary_emb(value_states, position_ids)
+ else:
+ cos, sin = position_embeddings
+
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+ if getattr(self.config, "group_size_ratio", None) and self.training: # shift
+ groupsz = int(q_len * getattr(self.config, "group_size_ratio"))
+ assert q_len % groupsz == 0, f"q_len {q_len} should be divisible by group size {groupsz}."
+ num_groups = q_len // groupsz
+
+ def shift(state: "torch.Tensor") -> "torch.Tensor":
+ state = state.transpose(1, 2) # output: (bsz, seq_len, n_heads, head_dim)
+ state = torch.cat(
+ (state[:, :, : self.num_heads // 2], state[:, :, self.num_heads // 2 :].roll(-groupsz // 2, dims=1)),
+ dim=2,
+ )
+ return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim).transpose(1, 2)
+
+ query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states)
+ if attention_mask is not None:
+ attention_mask = attention_mask[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1)
+
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+ if attention_mask is not None: # no matter the length, we just slice it
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ # upcast attention to fp32
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+ attn_output = torch.matmul(attn_weights, value_states) # (bsz, :, seq_len, :) or (bsz * n_group, :, groupsz, :)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ if getattr(self.config, "group_size_ratio", None) and self.training: # shift back
+ attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim)
+ attn_output = torch.cat(
+ (
+ attn_output[:, :, : self.num_heads // 2],
+ attn_output[:, :, self.num_heads // 2 :].roll(groupsz // 2, dims=1),
+ ),
+ dim=2,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+# Modified from:
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llama/modeling_llama.py
+def llama_flash_attention_2_forward(
+ self: "LlamaFlashAttention2",
+ hidden_states: "torch.Tensor",
+ attention_mask: Optional["torch.Tensor"] = None,
+ position_ids: Optional["torch.LongTensor"] = None,
+ past_key_value: Optional["Cache"] = None,
+ output_attentions: bool = False,
+ cache_position: Optional["torch.LongTensor"] = None,
+ position_embeddings: Optional[tuple["torch.Tensor", "torch.Tensor"]] = None,
+ **kwargs,
+) -> tuple["torch.Tensor", Optional["torch.Tensor"], Optional[tuple["torch.Tensor"]]]:
+ # LlamaFlashAttention2 attention does not support output_attentions
+ output_attentions = False
+
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states: torch.Tensor = self.q_proj(hidden_states)
+ key_states: torch.Tensor = self.k_proj(hidden_states)
+ value_states: torch.Tensor = self.v_proj(hidden_states)
+
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+ if position_embeddings is None:
+ cos, sin = self.rotary_emb(value_states, position_ids)
+ else:
+ cos, sin = position_embeddings
+
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+ # FlashAttention requires the input to have the shape (bsz, seq_len, n_heads, head_dim)
+ query_states = query_states.transpose(1, 2)
+ key_states = key_states.transpose(1, 2)
+ value_states = value_states.transpose(1, 2)
+
+ dropout_rate = self.attention_dropout if self.training else 0.0
+
+ input_dtype = query_states.dtype
+ if input_dtype == torch.float32:
+ if torch.is_autocast_enabled():
+ target_dtype = torch.get_autocast_gpu_dtype()
+ elif hasattr(self.config, "_pre_quantization_dtype"):
+ target_dtype = self.config._pre_quantization_dtype
+ else:
+ target_dtype = self.q_proj.weight.dtype
+
+ transformers_logger.warning_once("The input hidden states seems to be silently casted in float32.")
+ query_states = query_states.to(target_dtype)
+ key_states = key_states.to(target_dtype)
+ value_states = value_states.to(target_dtype)
+
+ if getattr(self.config, "group_size_ratio", None) and self.training: # shift
+ groupsz = int(q_len * getattr(self.config, "group_size_ratio"))
+ assert q_len % groupsz == 0, f"q_len {q_len} should be divisible by group size {groupsz}."
+ num_groups = q_len // groupsz
+
+ def shift(state: "torch.Tensor") -> "torch.Tensor":
+ state = torch.cat(
+ (state[:, :, : self.num_heads // 2], state[:, :, self.num_heads // 2 :].roll(-groupsz // 2, dims=1)),
+ dim=2,
+ )
+ return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim)
+
+ query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states)
+ if attention_mask is not None:
+ attention_mask = attention_mask[:, :groupsz].repeat(num_groups, 1)
+
+ attn_output: torch.Tensor = _flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ query_states.size(1),
+ dropout=dropout_rate,
+ sliding_window=getattr(self, "sliding_window", None),
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ is_causal=self.is_causal,
+ )
+
+ if getattr(self.config, "group_size_ratio", None) and self.training: # shift back
+ attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim)
+ attn_output = torch.cat(
+ (
+ attn_output[:, :, : self.num_heads // 2],
+ attn_output[:, :, self.num_heads // 2 :].roll(groupsz // 2, dims=1),
+ ),
+ dim=2,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+# Modified from:
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llama/modeling_llama.py
+def llama_sdpa_attention_forward(
+ self: "LlamaSdpaAttention",
+ hidden_states: "torch.Tensor",
+ attention_mask: Optional["torch.Tensor"] = None,
+ position_ids: Optional["torch.LongTensor"] = None,
+ past_key_value: Optional["Cache"] = None,
+ output_attentions: bool = False,
+ cache_position: Optional["torch.LongTensor"] = None,
+ position_embeddings: Optional[tuple["torch.Tensor", "torch.Tensor"]] = None,
+ **kwargs,
+) -> tuple["torch.Tensor", Optional["torch.Tensor"], Optional[tuple["torch.Tensor"]]]:
+ if output_attentions:
+ transformers_logger.warning_once(
+ "SDPA does not support `output_attentions=True`. Falling back to the vanilla attention"
+ )
+ return llama_attention_forward(
+ self,
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_value,
+ output_attentions=output_attentions,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states: torch.Tensor = self.q_proj(hidden_states)
+ key_states: torch.Tensor = self.k_proj(hidden_states)
+ value_states: torch.Tensor = self.v_proj(hidden_states)
+
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+ if position_embeddings is None:
+ cos, sin = self.rotary_emb(value_states, position_ids)
+ else:
+ cos, sin = position_embeddings
+
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+ if getattr(self.config, "group_size_ratio", None) and self.training: # shift
+ groupsz = int(q_len * getattr(self.config, "group_size_ratio"))
+ assert q_len % groupsz == 0, f"q_len {q_len} should be divisible by group size {groupsz}."
+ num_groups = q_len // groupsz
+
+ def shift(state: "torch.Tensor") -> "torch.Tensor":
+ state = state.transpose(1, 2) # output: (bsz, seq_len, n_heads, head_dim)
+ state = torch.cat(
+ (state[:, :, : self.num_heads // 2], state[:, :, self.num_heads // 2 :].roll(-groupsz // 2, dims=1)),
+ dim=2,
+ )
+ return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim).transpose(1, 2)
+
+ query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states)
+ if attention_mask is not None:
+ attention_mask = attention_mask[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1)
+
+ causal_mask = attention_mask
+ if attention_mask is not None:
+ causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
+
+ if query_states.device.type == "cuda" and causal_mask is not None: # avoid pytorch bug
+ query_states = query_states.contiguous()
+ key_states = key_states.contiguous()
+ value_states = value_states.contiguous()
+
+ is_causal = True if causal_mask is None and q_len > 1 else False
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query_states,
+ key_states,
+ value_states,
+ attn_mask=causal_mask,
+ dropout_p=self.attention_dropout if self.training else 0.0,
+ is_causal=is_causal,
+ )
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ if getattr(self.config, "group_size_ratio", None) and self.training: # shift back
+ attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim)
+ attn_output = torch.cat(
+ (
+ attn_output[:, :, : self.num_heads // 2],
+ attn_output[:, :, self.num_heads // 2 :].roll(groupsz // 2, dims=1),
+ ),
+ dim=2,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, None, past_key_value
+
+
+def _apply_llama_patch() -> None:
+ check_version("transformers>=4.45.0,<4.48.0", mandatory=True)
+ LlamaAttention.forward = llama_attention_forward
+ LlamaFlashAttention2.forward = llama_flash_attention_2_forward
+ LlamaSdpaAttention.forward = llama_sdpa_attention_forward
+
+
+def configure_longlora(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None:
+ if not is_trainable or not model_args.shift_attn:
+ return
+
+ logger = logging.get_logger(__name__)
+
+ if getattr(config, "model_type", None) in SUPPORTED_CLASS_FOR_S2ATTN:
+ setattr(config, "group_size_ratio", 0.25)
+ _apply_llama_patch()
+ logger.info_rank0("Using shift short attention with group_size_ratio=1/4.")
+ else:
+ logger.warning_rank0("Current model does not support shift short attention.")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/misc.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/misc.py
new file mode 100644
index 0000000..b0249b4
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/misc.py
@@ -0,0 +1,86 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from ...extras import logging
+from .visual import COMPOSITE_MODELS
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+
+def find_all_linear_modules(model: "PreTrainedModel", freeze_vision_tower: bool) -> list[str]:
+ r"""Find all available modules to apply LoRA, GaLore or APOLLO."""
+ 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")
+
+ if model_type in COMPOSITE_MODELS:
+ forbidden_modules.add(COMPOSITE_MODELS[model_type].projector_key)
+
+ if freeze_vision_tower and model_type in COMPOSITE_MODELS:
+ forbidden_modules.update(COMPOSITE_MODELS[model_type].vision_model_keys)
+
+ 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_rank0("Found linear modules: {}".format(",".join(module_names)))
+ return list(module_names)
+
+
+def find_expanded_modules(model: "PreTrainedModel", target_modules: list[str], num_layer_trainable: int) -> list[str]:
+ r"""Find the modules in the expanded blocks to apply lora."""
+ num_layers = getattr(model.config, "num_hidden_layers", None)
+ if not num_layers:
+ raise ValueError("Model was not supported.")
+
+ if num_layers % num_layer_trainable != 0:
+ raise ValueError(
+ f"`num_layers` {num_layers} should be divisible by `num_layer_trainable` {num_layer_trainable}."
+ )
+
+ stride = num_layers // num_layer_trainable
+ trainable_layer_ids = range(stride - 1, num_layers + stride - 1, stride)
+ trainable_layers = [f".{idx:d}." for idx in trainable_layer_ids]
+ module_names = []
+ for name, _ in model.named_modules():
+ if any(target_module in name for target_module in target_modules) and any(
+ trainable_layer in name for trainable_layer in trainable_layers
+ ):
+ module_names.append(name)
+
+ logger.info_rank0("Apply lora to layers: {}.".format(",".join(map(str, trainable_layer_ids))))
+ return module_names
+
+
+def register_autoclass(config: "PretrainedConfig", model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer"):
+ if "AutoConfig" in getattr(config, "auto_map", {}):
+ config.__class__.register_for_auto_class()
+ if "AutoModelForCausalLM" in getattr(config, "auto_map", {}):
+ model.__class__.register_for_auto_class()
+ if "AutoTokenizer" in tokenizer.init_kwargs.get("auto_map", {}):
+ tokenizer.__class__.register_for_auto_class()
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/mod.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/mod.py
new file mode 100644
index 0000000..5f67cd5
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/mod.py
@@ -0,0 +1,42 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from ...extras.constants import MOD_SUPPORTED_MODELS
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedModel
+
+ from ...hparams import ModelArguments
+
+
+def load_mod_pretrained_model(**init_kwargs) -> "PreTrainedModel":
+ from MoD import AutoMoDModelForCausalLM
+
+ return AutoMoDModelForCausalLM.from_pretrained(**init_kwargs)
+
+
+def convert_pretrained_model_to_mod(
+ model: "PreTrainedModel", config: "PretrainedConfig", model_args: "ModelArguments"
+) -> "PreTrainedModel":
+ from MoD import apply_mod_to_hf
+
+ if getattr(config, "model_type", None) not in MOD_SUPPORTED_MODELS:
+ raise ValueError("Current model is not supported by mixture-of-depth.")
+
+ model = apply_mod_to_hf(model)
+ model = model.to(model_args.compute_dtype)
+ return model
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/moe.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/moe.py
new file mode 100644
index 0000000..b3fca4f
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/moe.py
@@ -0,0 +1,87 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+import torch
+from transformers.integrations import is_deepspeed_zero3_enabled
+
+from ...extras.misc import check_version
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedModel
+
+ from ...hparams import ModelArguments
+
+
+def _set_z3_leaf_modules(model: "PreTrainedModel", leaf_modules: list["torch.nn.Module"]) -> None:
+ check_version("deepspeed>=0.13.0")
+ from deepspeed.utils import set_z3_leaf_modules # type: ignore
+
+ set_z3_leaf_modules(model, leaf_modules)
+
+
+def add_z3_leaf_module(model: "PreTrainedModel") -> None:
+ r"""Set module as a leaf module to skip partitioning in deepspeed zero3."""
+ if not is_deepspeed_zero3_enabled():
+ return
+
+ model_type = getattr(model.config, "model_type", None)
+ if model_type == "dbrx":
+ from transformers.models.dbrx.modeling_dbrx import DbrxFFN
+
+ _set_z3_leaf_modules(model, [DbrxFFN])
+
+ if model_type == "jamba":
+ from transformers.models.jamba.modeling_jamba import JambaSparseMoeBlock
+
+ _set_z3_leaf_modules(model, [JambaSparseMoeBlock])
+
+ if model_type == "jetmoe":
+ from transformers.models.jetmoe.modeling_jetmoe import JetMoeMoA, JetMoeMoE
+
+ _set_z3_leaf_modules(model, [JetMoeMoA, JetMoeMoE])
+
+ if model_type in ["kimi_vl", "deepseek_v3"]:
+ check_version("transformers>=4.51.1")
+ from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3MoE
+
+ _set_z3_leaf_modules(model, [DeepseekV3MoE])
+
+ if model_type == "mixtral":
+ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
+
+ _set_z3_leaf_modules(model, [MixtralSparseMoeBlock])
+
+ if model_type == "qwen2_moe":
+ from transformers.models.qwen2_moe.modeling_qwen2_moe import Qwen2MoeSparseMoeBlock
+
+ _set_z3_leaf_modules(model, [Qwen2MoeSparseMoeBlock])
+
+
+def configure_moe(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None:
+ model_type = getattr(config, "model_type", None)
+ if model_args.moe_aux_loss_coef is not None:
+ if model_type in ["jamba", "mixtral", "qwen2_moe"]:
+ setattr(config, "router_aux_loss_coef", model_args.moe_aux_loss_coef)
+
+ elif model_type == "deepseek":
+ setattr(config, "aux_loss_alpha", model_args.moe_aux_loss_coef)
+
+ elif model_type == "jetmoe":
+ setattr(config, "aux_loss_coef", model_args.moe_aux_loss_coef)
+
+ if model_type in ["dbrx", "jamba", "jetmoe", "mixtral", "qwen2_moe"]:
+ setattr(config, "output_router_logits", is_trainable)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/packing.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/packing.py
new file mode 100644
index 0000000..ed0d6c4
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/packing.py
@@ -0,0 +1,117 @@
+# Copyright 2025 Musab Gultekin and the LlamaFactory team.
+#
+# This code is based on the Musab Gultekin's functionary library.
+# https://github.com/MeetKai/functionary/blob/main/functionary/train/packing/monkey_patch_packing.py
+#
+# 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.
+#
+# MIT License
+#
+# Copyright (c) 2023 Musab Gultekin
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F
+
+from ...extras import logging
+
+
+if TYPE_CHECKING:
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def get_seqlens_in_batch(attention_mask: "torch.Tensor") -> "torch.Tensor":
+ r"""Get the sequnce lengths in the current batch.
+
+ e.g.
+ ```python
+ # input
+ [
+ [1, 1, 2, 2, 2, 0],
+ [1, 2, 2, 3, 3, 3],
+ ]
+ # output
+ [2, 3, 1, 2, 3]
+ ```
+ """
+ bsz = attention_mask.size(0)
+ dtype, device = attention_mask.dtype, attention_mask.device
+ max_num = torch.max(attention_mask).item()
+ counts: torch.Tensor = torch.zeros((bsz, max_num), dtype=dtype, device=device)
+ for i in range(max_num):
+ counts[:, i] = torch.sum(attention_mask == (i + 1), dim=-1)
+
+ counts = counts.flatten()
+ seqlens = counts[counts.nonzero().squeeze(dim=-1)]
+ return seqlens
+
+
+def get_unpad_data(attention_mask: "torch.Tensor") -> tuple["torch.Tensor", "torch.Tensor", int]:
+ r"""Prepare the indices and seqlens for flash attn varlen function.
+
+ Returns:
+ indices: indices of non-masked tokens from the flattened sequence.
+ cu_seqlens: the cumulative sequence lengths in the current batch, always starts from 0.
+ max_seqlen_in_batch: the largest seqlen in the current batch.
+
+ e.g.
+ ```python
+ # input
+ [
+ [1, 1, 2, 2, 2, 0],
+ [1, 2, 2, 3, 3, 3],
+ ]
+ # output
+ [0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11]
+ [0, 2, 5, 6, 8, 11]
+ 3
+ ```
+
+ """
+ seqlens_in_batch = get_seqlens_in_batch(attention_mask)
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+ return indices, cu_seqlens, max_seqlen_in_batch
+
+
+def configure_packing(model_args: "ModelArguments", is_trainable: bool) -> None:
+ if not is_trainable or not model_args.block_diag_attn:
+ return
+
+ import transformers.modeling_flash_attention_utils
+
+ transformers.modeling_flash_attention_utils._get_unpad_data = get_unpad_data
+ logger.info_rank0("Using block diagonal attention for sequence packing without cross-attention.")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/quantization.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/quantization.py
new file mode 100644
index 0000000..ffbf582
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/quantization.py
@@ -0,0 +1,183 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's Transformers and Optimum library.
+# https://github.com/huggingface/transformers/blob/v4.41.0/src/transformers/utils/quantization_config.py
+# https://github.com/huggingface/optimum/blob/v1.20.0/optimum/gptq/data.py
+#
+# 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 os
+import random
+from typing import TYPE_CHECKING, Any
+
+import torch
+from datasets import load_dataset
+from transformers import BitsAndBytesConfig, EetqConfig, GPTQConfig, HqqConfig
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.modeling_utils import is_fsdp_enabled
+
+from ...extras import logging
+from ...extras.constants import FILEEXT2TYPE, QuantizationMethod
+from ...extras.misc import check_version, get_current_device
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedTokenizer
+
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_quantization_dataset(tokenizer: "PreTrainedTokenizer", model_args: "ModelArguments") -> list[dict[str, Any]]:
+ r"""Prepare 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",
+ model_args: "ModelArguments",
+ init_kwargs: dict[str, Any],
+) -> None:
+ r"""Priority: PTQ-quantized (train/infer) > AutoGPTQ (export) > On-the-fly quantization (train/infer)."""
+ if getattr(config, "quantization_config", None): # ptq
+ if model_args.quantization_bit is not None:
+ logger.warning_rank0("`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:
+ check_version("auto_gptq>=0.5.0", mandatory=True)
+ quantization_config.pop("disable_exllama", None) # remove deprecated args
+ quantization_config["use_exllama"] = False # disable exllama
+
+ if quant_method == QuantizationMethod.AWQ:
+ check_version("autoawq", mandatory=True)
+
+ if quant_method == QuantizationMethod.AQLM:
+ check_version("aqlm>=1.1.0", mandatory=True)
+ quantization_config["bits"] = 2
+
+ quant_bits = quantization_config.get("bits", "?")
+ logger.info_rank0(f"Loading {quant_bits}-bit {quant_method.upper()}-quantized model.")
+
+ elif model_args.export_quantization_bit is not None: # auto-gptq
+ if model_args.export_quantization_bit not in [8, 4, 3, 2]:
+ raise ValueError("AutoGPTQ only accepts 2/3/4/8-bit quantization.")
+
+ check_version("optimum>=1.17.0", mandatory=True)
+ check_version("auto_gptq>=0.5.0", mandatory=True)
+ 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=model_args.export_quantization_bit,
+ dataset=_get_quantization_dataset(tokenizer, model_args),
+ )
+ init_kwargs["device_map"] = "auto"
+ init_kwargs["max_memory"] = get_max_memory()
+ logger.info_rank0(f"Quantizing model to {model_args.export_quantization_bit} bit with AutoGPTQ.")
+
+ elif model_args.quantization_bit is not None: # on-the-fly
+ if model_args.quantization_method == QuantizationMethod.BNB:
+ if model_args.quantization_bit == 8:
+ check_version("bitsandbytes>=0.37.0", mandatory=True)
+ init_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
+ elif model_args.quantization_bit == 4:
+ check_version("bitsandbytes>=0.39.0", mandatory=True)
+ init_kwargs["quantization_config"] = BitsAndBytesConfig(
+ load_in_4bit=True,
+ bnb_4bit_compute_dtype=model_args.compute_dtype,
+ bnb_4bit_use_double_quant=model_args.double_quantization,
+ bnb_4bit_quant_type=model_args.quantization_type,
+ bnb_4bit_quant_storage=model_args.compute_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 model_args.quantization_device_map == "auto":
+ if model_args.quantization_bit != 4:
+ raise ValueError("Only 4-bit quantized model can use fsdp+qlora or auto device map.")
+
+ check_version("bitsandbytes>=0.43.0", mandatory=True)
+ else:
+ init_kwargs["device_map"] = {"": get_current_device()} # change auto device map for inference
+
+ logger.info_rank0(f"Quantizing model to {model_args.quantization_bit} bit with bitsandbytes.")
+ elif model_args.quantization_method == QuantizationMethod.HQQ:
+ if model_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.")
+
+ check_version("hqq", mandatory=True)
+ init_kwargs["quantization_config"] = HqqConfig(
+ nbits=model_args.quantization_bit, quant_zero=False, quant_scale=False, axis=0
+ ) # use ATEN kernel (axis=0) for performance
+ logger.info_rank0(f"Quantizing model to {model_args.quantization_bit} bit with HQQ.")
+ elif model_args.quantization_method == QuantizationMethod.EETQ:
+ if model_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.")
+
+ check_version("eetq", mandatory=True)
+ init_kwargs["quantization_config"] = EetqConfig()
+ logger.info_rank0(f"Quantizing model to {model_args.quantization_bit} bit with EETQ.")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/rope.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/rope.py
new file mode 100644
index 0000000..30d0fdd
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/rope.py
@@ -0,0 +1,71 @@
+# Copyright 2025 LMSYS and the LlamaFactory team.
+# Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+# This code is inspired by the LMSYS's FastChat library.
+# https://github.com/lm-sys/FastChat/blob/v0.2.30/fastchat/train/train.py
+#
+# 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 math
+from typing import TYPE_CHECKING
+
+from ...extras import logging
+from ...extras.constants import RopeScaling
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig
+
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def configure_rope(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None:
+ if model_args.rope_scaling is None:
+ return
+
+ if not hasattr(config, "rope_scaling"):
+ logger.warning_rank0("Current model does not support RoPE scaling.")
+ return
+
+ rope_kwargs = {"rope_type": getattr(model_args.rope_scaling, "value", model_args.rope_scaling)} # handle enum
+ if model_args.model_max_length is not None:
+ if is_trainable and model_args.rope_scaling == RopeScaling.DYNAMIC:
+ logger.warning_rank0(
+ "Dynamic NTK scaling may not work well with fine-tuning. "
+ "See: https://github.com/huggingface/transformers/pull/24653"
+ )
+
+ current_max_length = getattr(config, "max_position_embeddings", None)
+ if (not current_max_length) or model_args.model_max_length <= current_max_length:
+ logger.warning_rank0("Input length is smaller than max length. Disabling rope scaling.")
+ return
+
+ logger.info_rank0(f"Enlarge max model length from {current_max_length} to {model_args.model_max_length}.")
+ setattr(config, "max_position_embeddings", model_args.model_max_length)
+ rope_kwargs["factor"] = float(math.ceil(model_args.model_max_length / current_max_length))
+ if model_args.rope_scaling == RopeScaling.DYNAMIC:
+ rope_kwargs["original_max_position_embeddings"] = current_max_length
+ elif model_args.rope_scaling == RopeScaling.LLAMA3:
+ rope_kwargs["original_max_position_embeddings"] = current_max_length
+ rope_kwargs["low_freq_factor"] = 1.0
+ rope_kwargs["high_freq_factor"] = 4.0
+ else:
+ rope_kwargs["factor"] = 2.0
+
+ setattr(config, "rope_scaling", rope_kwargs)
+ logger.info_rank0(
+ f"Using {rope_kwargs['rope_type']} scaling strategy and setting scaling factor to {rope_kwargs['factor']}."
+ )
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/unsloth.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/unsloth.py
new file mode 100644
index 0000000..8bb6aa6
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/unsloth.py
@@ -0,0 +1,96 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from ...extras import logging
+from ...extras.misc import get_current_device
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedModel
+
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_unsloth_kwargs(
+ config: "PretrainedConfig", model_name_or_path: str, model_args: "ModelArguments"
+) -> dict[str, Any]:
+ return {
+ "model_name": model_name_or_path,
+ "max_seq_length": model_args.model_max_length or 4096,
+ "dtype": model_args.compute_dtype,
+ "load_in_4bit": model_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": model_args.trust_remote_code,
+ "use_gradient_checkpointing": "unsloth",
+ }
+
+
+def load_unsloth_pretrained_model(
+ config: "PretrainedConfig", model_args: "ModelArguments"
+) -> Optional["PreTrainedModel"]:
+ r"""Optionally load pretrained model with unsloth. Used in training."""
+ from unsloth import FastLanguageModel # type: ignore
+
+ 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_rank0("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: "ModelArguments", peft_kwargs: dict[str, Any]
+) -> "PreTrainedModel":
+ r"""Get the peft model for the pretrained model with unsloth. Used in training."""
+ from unsloth import FastLanguageModel # type: ignore
+
+ unsloth_peft_kwargs = {
+ "model": model,
+ "max_seq_length": model_args.model_max_length,
+ "use_gradient_checkpointing": "unsloth",
+ }
+ return FastLanguageModel.get_peft_model(**peft_kwargs, **unsloth_peft_kwargs)
+
+
+def load_unsloth_peft_model(
+ config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool
+) -> "PreTrainedModel":
+ r"""Load peft model with unsloth. Used in both training and inference."""
+ from unsloth import FastLanguageModel # type: ignore
+
+ 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
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/valuehead.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/valuehead.py
new file mode 100644
index 0000000..137c6b7
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/valuehead.py
@@ -0,0 +1,72 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+import torch
+from transformers.utils import cached_file
+
+from ...extras import logging
+from ...extras.constants import V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME
+
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedModel
+
+ from ...hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def load_valuehead_params(path_or_repo_id: str, model_args: "ModelArguments") -> dict[str, torch.Tensor]:
+ r"""Load value head parameters from Hugging Face Hub or local disk.
+
+ Returns: dict with keys `v_head.summary.weight` and `v_head.summary.bias`.
+ """
+ kwargs = {"path_or_repo_id": path_or_repo_id, "cache_dir": model_args.cache_dir, "token": model_args.hf_hub_token}
+ err_text = ""
+
+ try:
+ from safetensors import safe_open
+
+ vhead_file = cached_file(filename=V_HEAD_SAFE_WEIGHTS_NAME, **kwargs)
+ with safe_open(vhead_file, framework="pt", device="cpu") as f:
+ return {key: f.get_tensor(key) for key in f.keys()}
+ except Exception as err:
+ err_text = str(err)
+
+ try:
+ vhead_file = cached_file(filename=V_HEAD_WEIGHTS_NAME, **kwargs)
+ return torch.load(vhead_file, map_location="cpu")
+ except Exception as err:
+ err_text = str(err)
+
+ logger.info_rank0(f"Provided path ({path_or_repo_id}) does not contain value head weights: {err_text}.")
+ logger.info_rank0("Ignore the above message if you are not resuming the training of a value head model.")
+ return None
+
+
+def prepare_valuehead_model(model: "PreTrainedModel") -> None:
+ if getattr(model.config, "model_type", None) == "llava":
+ setattr(model, "lm_head", model.language_model.get_output_embeddings())
+ setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])
+
+ if getattr(model.config, "model_type", None) == "chatglm":
+ setattr(model, "lm_head", model.transformer.output_layer)
+ setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])
+
+ if getattr(model.config, "model_type", None) == "internlm2":
+ setattr(model, "lm_head", model.output)
+ setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/visual.py b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/visual.py
new file mode 100644
index 0000000..901010c
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/model_utils/visual.py
@@ -0,0 +1,295 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's Transformers library.
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava/modeling_llava.py
+#
+# 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.
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import transformers
+import transformers.models
+from transformers.activations import ACT2FN
+
+from ...extras import logging
+
+
+if TYPE_CHECKING:
+ from transformers import LlavaConfig, PretrainedConfig, PreTrainedModel
+
+ from ...hparams import FinetuningArguments, ModelArguments
+
+
+logger = logging.get_logger(__name__)
+transformers_logger = transformers.utils.logging.get_logger(__name__)
+
+
+@dataclass
+class CompositeModel:
+ model_type: str
+ projector_key: str
+ vision_model_keys: list[str]
+ language_model_keys: list[str]
+ lora_conflict_keys: list[str]
+
+ def get_projector(self, module: "torch.nn.Module") -> "torch.nn.Module":
+ for key in self.projector_key.split("."):
+ module = getattr(module, key)
+
+ return module
+
+
+COMPOSITE_MODELS: dict[str, "CompositeModel"] = {}
+
+
+def _register_composite_model(
+ model_type: str,
+ projector_key: Optional[str] = None,
+ vision_model_keys: Optional[list[str]] = None,
+ language_model_keys: Optional[list[str]] = None,
+ lora_conflict_keys: Optional[list[str]] = None,
+):
+ r"""Register a new composite model.
+
+ Args:
+ model_type: model type
+ projector_key: multi_modal_projector
+ vision_model_keys: vision_tower
+ language_model_keys: language_model
+ lora_conflict_keys: None
+
+ """
+ COMPOSITE_MODELS[model_type] = CompositeModel(
+ model_type=model_type,
+ projector_key=projector_key or "multi_modal_projector",
+ vision_model_keys=vision_model_keys or ["vision_tower"],
+ language_model_keys=language_model_keys or ["language_model"],
+ lora_conflict_keys=lora_conflict_keys or [],
+ )
+
+
+class LlavaMultiModalProjectorForYiVL(torch.nn.Module):
+ def __init__(self, config: "LlavaConfig") -> None:
+ super().__init__()
+
+ self.config = config
+ if config is None:
+ return
+
+ self.linear_1 = torch.nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=True)
+ self.linear_2 = torch.nn.LayerNorm(config.text_config.hidden_size, bias=True)
+ self.linear_3 = torch.nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)
+ self.linear_4 = torch.nn.LayerNorm(config.text_config.hidden_size, bias=True)
+ self.act = ACT2FN[config.projector_hidden_act]
+
+ def forward(self, image_features: "torch.Tensor") -> "torch.Tensor":
+ hidden_states = self.linear_1(image_features)
+ hidden_states = self.linear_2(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.linear_3(hidden_states)
+ hidden_states = self.linear_4(hidden_states)
+ if hidden_states.dtype == torch.float32:
+ if torch.is_autocast_enabled():
+ target_dtype = torch.get_autocast_gpu_dtype()
+ elif hasattr(self.config, "_pre_quantization_dtype"):
+ target_dtype = self.config._pre_quantization_dtype
+ else:
+ target_dtype = self.linear_1.weight.dtype
+
+ transformers_logger.warning_once("The hidden states seems to be silently casted in float32.")
+ hidden_states = hidden_states.to(target_dtype)
+
+ return hidden_states
+
+
+class LlavaMultiModalProjectorForYiVLForVLLM(LlavaMultiModalProjectorForYiVL):
+ def __init__(self, vision_hidden_size: int, text_hidden_size: int, projector_hidden_act: str) -> None:
+ super().__init__(config=None)
+
+ self.linear_1 = torch.nn.Linear(vision_hidden_size, text_hidden_size, bias=True)
+ self.linear_2 = torch.nn.LayerNorm(text_hidden_size, bias=True)
+ self.linear_3 = torch.nn.Linear(text_hidden_size, text_hidden_size, bias=True)
+ self.linear_4 = torch.nn.LayerNorm(text_hidden_size, bias=True)
+ self.act = ACT2FN[projector_hidden_act]
+
+
+def autocast_projector_dtype(model: "PreTrainedModel", model_args: "ModelArguments") -> None:
+ r"""Cast projector output to half precision for fine-tuning quantized VLMs."""
+
+ def _mm_projector_forward_post_hook(
+ module: "torch.nn.Module", args: tuple["torch.Tensor"], output: "torch.Tensor"
+ ) -> "torch.Tensor":
+ return output.to(model_args.compute_dtype)
+
+ if getattr(model, "quantization_method", None):
+ model_type = getattr(model.config, "model_type", None)
+ if model_type in COMPOSITE_MODELS:
+ mm_projector = COMPOSITE_MODELS[model_type].get_projector(model)
+ else:
+ return
+
+ logger.info_rank0(f"Casting multimodal projector outputs in {model_args.compute_dtype}.")
+ mm_projector.register_forward_hook(_mm_projector_forward_post_hook)
+
+
+def configure_visual_model(config: "PretrainedConfig") -> None:
+ r"""Patch VLMs before loading them."""
+ if getattr(config, "text_config", None) and not getattr(config, "hidden_size", None):
+ # required for ds zero3 and valuehead models
+ setattr(config, "hidden_size", getattr(config.text_config, "hidden_size", None))
+
+ if getattr(config, "is_yi_vl_derived_model", None):
+ logger.info_rank0("Detected Yi-VL model, applying projector patch.")
+ transformers.models.llava.modeling_llava.LlavaMultiModalProjector = LlavaMultiModalProjectorForYiVL
+
+
+def get_forbidden_modules(config: "PretrainedConfig", finetuning_args: "FinetuningArguments") -> set[str]:
+ r"""Freeze vision tower and language model for VLM full/freeze tuning."""
+ model_type = getattr(config, "model_type", None)
+ forbidden_modules = set()
+ if model_type in COMPOSITE_MODELS:
+ if finetuning_args.freeze_vision_tower:
+ vision_model_keys = COMPOSITE_MODELS[model_type].vision_model_keys
+ logger.info_rank0(f"Set vision model not trainable: {vision_model_keys}.")
+ forbidden_modules.update(vision_model_keys)
+
+ if finetuning_args.freeze_multi_modal_projector:
+ projector_key = COMPOSITE_MODELS[model_type].projector_key
+ logger.info_rank0(f"Set multi model projector not trainable: {projector_key}.")
+ forbidden_modules.add(projector_key)
+
+ if finetuning_args.freeze_language_model:
+ language_model_keys = COMPOSITE_MODELS[model_type].language_model_keys
+ logger.info_rank0(f"Set language model not trainable: {language_model_keys}.")
+ forbidden_modules.update(language_model_keys)
+
+ return forbidden_modules
+
+
+def patch_target_modules(
+ model: "PreTrainedModel", finetuning_args: "FinetuningArguments", target_modules: list[str]
+) -> list[str]:
+ r"""Freeze vision tower for VLM LoRA tuning."""
+ model_type = getattr(model.config, "model_type", None)
+ if model_type in COMPOSITE_MODELS:
+ forbidden_modules = get_forbidden_modules(model.config, finetuning_args)
+ forbidden_modules.update(COMPOSITE_MODELS[model_type].lora_conflict_keys)
+ module_names = []
+ for name, _ in model.named_modules():
+ if any(target_module in name for target_module in target_modules) and not any(
+ forbidden_module in name for forbidden_module in forbidden_modules
+ ):
+ module_names.append(name)
+
+ return module_names
+ else:
+ return target_modules
+
+
+_register_composite_model(
+ model_type="internvl",
+)
+
+
+_register_composite_model(
+ model_type="gemma3",
+)
+
+
+_register_composite_model(
+ model_type="llama4",
+ vision_model_keys=["vision_model"],
+)
+
+
+_register_composite_model(
+ model_type="llava",
+)
+
+
+_register_composite_model(
+ model_type="llava_next",
+)
+
+
+_register_composite_model(
+ model_type="llava_next_video",
+)
+
+
+_register_composite_model(
+ model_type="minicpmv",
+ projector_key="resampler",
+ vision_model_keys=["vpm"],
+ language_model_keys=["llm"],
+)
+
+
+_register_composite_model(
+ model_type="minicpmo",
+ projector_key="resampler",
+ vision_model_keys=["vpm", "apm", "audio_avg_pooler", "audio_projection_layer", "tts"],
+ language_model_keys=["llm"],
+ lora_conflict_keys=["audio_projection_layer"],
+)
+
+
+_register_composite_model(
+ model_type="paligemma",
+)
+
+
+_register_composite_model(
+ model_type="video_llava",
+)
+
+
+_register_composite_model(
+ model_type="mllama",
+ vision_model_keys=["vision_model"],
+)
+
+
+_register_composite_model(
+ model_type="qwen2_audio",
+ vision_model_keys=["audio_tower"],
+)
+
+
+_register_composite_model(
+ model_type="qwen2_5_omni_thinker",
+ projector_key="visual.merger",
+ vision_model_keys=["visual.patch_embed", "visual.blocks", "audio_tower"],
+ language_model_keys=["model", "lm_head"],
+ lora_conflict_keys=["patch_embed"],
+)
+
+
+_register_composite_model(
+ model_type="qwen2_vl",
+ projector_key="visual.merger",
+ vision_model_keys=["visual.patch_embed", "visual.blocks"],
+ language_model_keys=["model", "lm_head"],
+ lora_conflict_keys=["patch_embed"],
+)
+
+
+_register_composite_model(
+ model_type="qwen2_5_vl",
+ projector_key="visual.merger",
+ vision_model_keys=["visual.patch_embed", "visual.blocks"],
+ language_model_keys=["model", "lm_head"],
+ lora_conflict_keys=["patch_embed"],
+)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/model/patcher.py b/post-training/LLaMA-Factory/src/llamafactory/model/patcher.py
new file mode 100644
index 0000000..28014de
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/model/patcher.py
@@ -0,0 +1,203 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from types import MethodType
+from typing import TYPE_CHECKING, Any
+
+import torch
+from peft import PeftModel
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.modeling_utils import is_fsdp_enabled
+
+from ..extras import logging
+from ..extras.misc import infer_optim_dtype
+from ..extras.packages import is_transformers_version_greater_than
+from .model_utils.attention import configure_attn_implementation, print_attn_implementation
+from .model_utils.checkpointing import prepare_model_for_training
+from .model_utils.embedding import resize_embedding_layer
+from .model_utils.kv_cache import configure_kv_cache
+from .model_utils.longlora import configure_longlora
+from .model_utils.moe import add_z3_leaf_module, configure_moe
+from .model_utils.packing import configure_packing
+from .model_utils.quantization import configure_quantization
+from .model_utils.rope import configure_rope
+from .model_utils.valuehead import prepare_valuehead_model
+from .model_utils.visual import autocast_projector_dtype, configure_visual_model
+
+
+if TYPE_CHECKING:
+ from transformers import PretrainedConfig, PreTrainedTokenizer, ProcessorMixin
+ from trl import AutoModelForCausalLMWithValueHead
+
+ from ..hparams import ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def patch_tokenizer(tokenizer: "PreTrainedTokenizer", model_args: "ModelArguments") -> None:
+ if "PreTrainedTokenizerBase" not in str(tokenizer._pad.__func__):
+ tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer)
+
+ if model_args.model_max_length is not None and tokenizer.model_max_length < model_args.model_max_length:
+ tokenizer.model_max_length = model_args.model_max_length # enlarge the tokenizer max length
+
+ if model_args.new_special_tokens is not None:
+ num_added_tokens = tokenizer.add_special_tokens(
+ dict(additional_special_tokens=model_args.new_special_tokens),
+ replace_additional_special_tokens=False,
+ )
+ logger.info_rank0("Add {} to special tokens.".format(",".join(model_args.new_special_tokens)))
+ if num_added_tokens > 0 and not model_args.resize_vocab:
+ model_args.resize_vocab = True
+ logger.warning_rank0("New tokens have been added, changed `resize_vocab` to True.")
+
+
+def patch_processor(
+ processor: "ProcessorMixin",
+ tokenizer: "PreTrainedTokenizer",
+ model_args: "ModelArguments",
+) -> None:
+ setattr(processor, "tokenizer", tokenizer)
+ setattr(processor, "image_max_pixels", model_args.image_max_pixels)
+ setattr(processor, "image_min_pixels", model_args.image_min_pixels)
+ setattr(processor, "image_do_pan_and_scan", model_args.image_do_pan_and_scan)
+ setattr(processor, "video_max_pixels", model_args.video_max_pixels)
+ setattr(processor, "video_min_pixels", model_args.video_min_pixels)
+ setattr(processor, "video_fps", model_args.video_fps)
+ setattr(processor, "video_maxlen", model_args.video_maxlen)
+ setattr(processor, "audio_sampling_rate", model_args.audio_sampling_rate)
+ setattr(processor, "use_audio_in_video", model_args.use_audio_in_video)
+
+
+def patch_config(
+ config: "PretrainedConfig",
+ tokenizer: "PreTrainedTokenizer",
+ model_args: "ModelArguments",
+ init_kwargs: dict[str, Any],
+ is_trainable: bool,
+) -> None:
+ if model_args.compute_dtype is None: # priority: bf16 > fp16 > fp32
+ if model_args.infer_dtype != "auto" and not is_trainable:
+ model_args.compute_dtype = getattr(torch, model_args.infer_dtype)
+ else:
+ model_args.compute_dtype = infer_optim_dtype(model_dtype=getattr(config, "torch_dtype", None))
+
+ configure_attn_implementation(config, model_args, is_trainable)
+ configure_rope(config, model_args, is_trainable)
+ configure_longlora(config, model_args, is_trainable)
+ configure_quantization(config, tokenizer, model_args, init_kwargs)
+ configure_moe(config, model_args, is_trainable)
+ configure_visual_model(config)
+ configure_packing(model_args, is_trainable)
+ configure_kv_cache(config, model_args, is_trainable)
+
+ if getattr(config, "model_type", None) == "qwen":
+ setattr(config, "use_flash_attn", model_args.flash_attn == "fa2")
+ for dtype_name, dtype in [("fp16", torch.float16), ("bf16", torch.bfloat16), ("fp32", torch.float32)]:
+ setattr(config, dtype_name, model_args.compute_dtype == dtype)
+
+ if getattr(config, "model_type", None) == "minicpmo":
+ setattr(config, "init_audio", True)
+ setattr(config, "init_tts", False)
+
+ # replace the top-k gating method
+ if getattr(config, "model_type", None) == "kimi_vl" and is_trainable:
+ setattr(config.text_config, "topk_method", "greedy")
+
+ if "LlavaLlamaForCausalLM" in getattr(config, "architectures", []):
+ raise ValueError("Please download llava models with hf-compatible format: https://huggingface.co/llava-hf")
+
+ if getattr(config, "model_type", None) == "internlm3" and not is_transformers_version_greater_than("4.47.1"):
+ raise RuntimeError("InternLM3 model requires transformers>=4.47.1, please upgrade it.")
+
+ # deepspeed zero3 is not compatible with low_cpu_mem_usage
+ init_kwargs["low_cpu_mem_usage"] = model_args.low_cpu_mem_usage and (not is_deepspeed_zero3_enabled())
+
+ # do not cast data type of the model deepspeed zero3 without qlora
+ if not (is_deepspeed_zero3_enabled() and model_args.quantization_bit is None):
+ init_kwargs["torch_dtype"] = model_args.compute_dtype
+
+ if init_kwargs["low_cpu_mem_usage"] and not is_fsdp_enabled(): # fsdp does not need device map
+ if "device_map" not in init_kwargs and model_args.device_map:
+ init_kwargs["device_map"] = model_args.device_map # device map requires low_cpu_mem_usage=True
+
+ if init_kwargs.get("device_map", None) == "auto":
+ init_kwargs["offload_folder"] = model_args.offload_folder
+
+
+def patch_model(
+ model: "PreTrainedModel",
+ tokenizer: "PreTrainedTokenizer",
+ model_args: "ModelArguments",
+ is_trainable: bool,
+ add_valuehead: bool,
+) -> None:
+ gen_config = model.generation_config # check and fix generation config
+ if not gen_config.do_sample and (
+ (gen_config.temperature is not None and gen_config.temperature != 1.0)
+ or (gen_config.top_p is not None and gen_config.top_p != 1.0)
+ or (gen_config.typical_p is not None and gen_config.typical_p != 1.0)
+ ):
+ gen_config.do_sample = True
+
+ if getattr(model.config, "model_type", None) not in ["minicpmv", "minicpmo"] and "GenerationMixin" not in str(
+ model.generate.__func__
+ ):
+ model.generate = MethodType(PreTrainedModel.generate, model)
+
+ if add_valuehead:
+ prepare_valuehead_model(model)
+
+ if model_args.resize_vocab:
+ resize_embedding_layer(model, tokenizer)
+
+ if is_trainable:
+ prepare_model_for_training(model, model_args)
+ autocast_projector_dtype(model, model_args)
+ add_z3_leaf_module(model)
+
+ if not model_args.use_unsloth:
+ print_attn_implementation(model.config)
+
+ try:
+ model.add_model_tags(["llama-factory"])
+ except Exception:
+ logger.warning_rank0("Cannot properly tag the model.")
+
+
+def patch_valuehead_model(model: "AutoModelForCausalLMWithValueHead") -> None:
+ def tie_weights(self: "AutoModelForCausalLMWithValueHead") -> None:
+ if isinstance(self.pretrained_model, PreTrainedModel):
+ self.pretrained_model.tie_weights()
+
+ def get_input_embeddings(self: "AutoModelForCausalLMWithValueHead") -> torch.nn.Module:
+ if isinstance(self.pretrained_model, PreTrainedModel):
+ return self.pretrained_model.get_input_embeddings()
+
+ def get_output_embeddings(self: "AutoModelForCausalLMWithValueHead") -> torch.nn.Module:
+ if isinstance(self.pretrained_model, PreTrainedModel):
+ return self.pretrained_model.get_output_embeddings()
+
+ def create_or_update_model_card(self: "AutoModelForCausalLMWithValueHead", output_dir: str) -> None:
+ if isinstance(self.pretrained_model, PeftModel):
+ self.pretrained_model.create_or_update_model_card(output_dir)
+
+ ignore_modules = [name for name, _ in model.named_parameters() if "pretrained_model" in name]
+ setattr(model, "_keys_to_ignore_on_save", ignore_modules)
+ setattr(model, "tie_weights", MethodType(tie_weights, model))
+ setattr(model, "get_input_embeddings", MethodType(get_input_embeddings, model))
+ setattr(model, "get_output_embeddings", MethodType(get_output_embeddings, model))
+ setattr(model, "create_or_update_model_card", MethodType(create_or_update_model_card, model))
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/train/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/callbacks.py b/post-training/LLaMA-Factory/src/llamafactory/train/callbacks.py
new file mode 100644
index 0000000..c1bd159
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/callbacks.py
@@ -0,0 +1,385 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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 os
+import signal
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+import transformers
+from peft import PeftModel
+from transformers import PreTrainedModel, ProcessorMixin, TrainerCallback
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, has_length
+from transformers.utils import (
+ SAFE_WEIGHTS_NAME,
+ WEIGHTS_NAME,
+ is_safetensors_available,
+)
+from typing_extensions import override
+
+from ..extras import logging
+from ..extras.constants import TRAINER_LOG, V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME
+from ..extras.misc import get_peak_memory, is_env_enabled, use_ray
+
+
+if is_safetensors_available():
+ from safetensors import safe_open
+ from safetensors.torch import save_file
+
+
+if TYPE_CHECKING:
+ from transformers import TrainerControl, TrainerState, TrainingArguments
+ from trl import AutoModelForCausalLMWithValueHead
+
+ from ..hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def fix_valuehead_checkpoint(
+ model: "AutoModelForCausalLMWithValueHead", output_dir: str, safe_serialization: bool
+) -> None:
+ r"""Fix the valuehead checkpoint files.
+
+ The model is already unwrapped.
+
+ There are three cases:
+ 1. full tuning without ds_zero3: state_dict = {"model.layers.*": ..., "v_head.summary.*": ...}
+ 2. lora tuning without ds_zero3: state_dict = {"v_head.summary.*": ...}
+ 3. under deepspeed zero3: state_dict = {"pretrained_model.model.layers.*": ..., "v_head.summary.*": ...}
+
+ We assume `stage3_gather_16bit_weights_on_model_save=true`.
+ """
+ if not isinstance(model.pretrained_model, (PreTrainedModel, PeftModel)):
+ return
+
+ if safe_serialization:
+ path_to_checkpoint = os.path.join(output_dir, SAFE_WEIGHTS_NAME)
+ with safe_open(path_to_checkpoint, framework="pt", device="cpu") as f:
+ state_dict: dict[str, torch.Tensor] = {key: f.get_tensor(key) for key in f.keys()}
+ else:
+ path_to_checkpoint = os.path.join(output_dir, WEIGHTS_NAME)
+ state_dict: dict[str, torch.Tensor] = torch.load(path_to_checkpoint, map_location="cpu")
+
+ os.remove(path_to_checkpoint)
+ decoder_state_dict, v_head_state_dict = {}, {}
+ for name, param in state_dict.items():
+ if name.startswith("v_head."):
+ v_head_state_dict[name] = param
+ else:
+ decoder_state_dict[name.replace("pretrained_model.", "", 1)] = param
+
+ model.pretrained_model.save_pretrained(
+ output_dir, state_dict=decoder_state_dict or None, safe_serialization=safe_serialization
+ )
+
+ if safe_serialization:
+ save_file(v_head_state_dict, os.path.join(output_dir, V_HEAD_SAFE_WEIGHTS_NAME), metadata={"format": "pt"})
+ else:
+ torch.save(v_head_state_dict, os.path.join(output_dir, V_HEAD_WEIGHTS_NAME))
+
+ logger.info_rank0(f"Value head model saved at: {output_dir}")
+
+
+class FixValueHeadModelCallback(TrainerCallback):
+ r"""A callback for fixing the checkpoint for valuehead models."""
+
+ @override
+ def on_save(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if args.should_save:
+ output_dir = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
+ fix_valuehead_checkpoint(
+ model=kwargs.pop("model"), output_dir=output_dir, safe_serialization=args.save_safetensors
+ )
+
+
+class SaveProcessorCallback(TrainerCallback):
+ r"""A callback for saving the processor."""
+
+ def __init__(self, processor: "ProcessorMixin") -> None:
+ self.processor = processor
+
+ @override
+ def on_save(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if args.should_save:
+ output_dir = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
+ self.processor.save_pretrained(output_dir)
+
+ @override
+ def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if args.should_save:
+ self.processor.save_pretrained(args.output_dir)
+
+
+class PissaConvertCallback(TrainerCallback):
+ r"""A callback for converting the PiSSA adapter to a normal one."""
+
+ @override
+ def on_train_begin(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if args.should_save:
+ model = kwargs.pop("model")
+ pissa_init_dir = os.path.join(args.output_dir, "pissa_init")
+ logger.info_rank0(f"Initial PiSSA adapter will be saved at: {pissa_init_dir}.")
+ if isinstance(model, PeftModel):
+ init_lora_weights = getattr(model.peft_config["default"], "init_lora_weights")
+ setattr(model.peft_config["default"], "init_lora_weights", True)
+ model.save_pretrained(pissa_init_dir, safe_serialization=args.save_safetensors)
+ setattr(model.peft_config["default"], "init_lora_weights", init_lora_weights)
+
+ @override
+ def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if args.should_save:
+ model = kwargs.pop("model")
+ pissa_init_dir = os.path.join(args.output_dir, "pissa_init")
+ pissa_backup_dir = os.path.join(args.output_dir, "pissa_backup")
+ pissa_convert_dir = os.path.join(args.output_dir, "pissa_converted")
+ logger.info_rank0(f"Converted PiSSA adapter will be saved at: {pissa_convert_dir}.")
+ # 1. save a pissa backup with init_lora_weights: True
+ # 2. save a converted lora with init_lora_weights: pissa
+ # 3. load the pissa backup with init_lora_weights: True
+ # 4. delete the initial adapter and change init_lora_weights to pissa
+ if isinstance(model, PeftModel):
+ init_lora_weights = getattr(model.peft_config["default"], "init_lora_weights")
+ setattr(model.peft_config["default"], "init_lora_weights", True)
+ model.save_pretrained(pissa_backup_dir, safe_serialization=args.save_safetensors)
+ setattr(model.peft_config["default"], "init_lora_weights", init_lora_weights)
+ model.save_pretrained(
+ pissa_convert_dir,
+ safe_serialization=args.save_safetensors,
+ path_initial_model_for_weight_conversion=pissa_init_dir,
+ )
+ model.load_adapter(pissa_backup_dir, "default", is_trainable=True)
+ model.set_adapter("default")
+ setattr(model.peft_config["default"], "init_lora_weights", init_lora_weights)
+
+
+class LogCallback(TrainerCallback):
+ r"""A callback for logging training and evaluation status."""
+
+ def __init__(self) -> None:
+ # Progress
+ self.start_time = 0
+ self.cur_steps = 0
+ self.max_steps = 0
+ self.elapsed_time = ""
+ self.remaining_time = ""
+ self.thread_pool: Optional[ThreadPoolExecutor] = None
+ # Status
+ self.aborted = False
+ self.do_train = False
+ # Web UI
+ self.webui_mode = is_env_enabled("LLAMABOARD_ENABLED")
+ if self.webui_mode and not use_ray():
+ signal.signal(signal.SIGABRT, self._set_abort)
+ self.logger_handler = logging.LoggerHandler(os.getenv("LLAMABOARD_WORKDIR"))
+ logging.add_handler(self.logger_handler)
+ transformers.logging.add_handler(self.logger_handler)
+
+ def _set_abort(self, signum, frame) -> None:
+ self.aborted = True
+
+ def _reset(self, max_steps: int = 0) -> None:
+ self.start_time = time.time()
+ self.cur_steps = 0
+ self.max_steps = max_steps
+ self.elapsed_time = ""
+ self.remaining_time = ""
+
+ def _timing(self, cur_steps: int) -> None:
+ cur_time = time.time()
+ elapsed_time = cur_time - self.start_time
+ avg_time_per_step = elapsed_time / cur_steps if cur_steps != 0 else 0
+ remaining_time = (self.max_steps - cur_steps) * avg_time_per_step
+ self.cur_steps = cur_steps
+ self.elapsed_time = str(timedelta(seconds=int(elapsed_time)))
+ self.remaining_time = str(timedelta(seconds=int(remaining_time)))
+
+ def _write_log(self, output_dir: str, logs: dict[str, Any]) -> None:
+ with open(os.path.join(output_dir, TRAINER_LOG), "a", encoding="utf-8") as f:
+ f.write(json.dumps(logs) + "\n")
+
+ def _create_thread_pool(self, output_dir: str) -> None:
+ os.makedirs(output_dir, exist_ok=True)
+ self.thread_pool = ThreadPoolExecutor(max_workers=1)
+
+ def _close_thread_pool(self) -> None:
+ if self.thread_pool is not None:
+ self.thread_pool.shutdown(wait=True)
+ self.thread_pool = None
+
+ @override
+ def on_init_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if (
+ args.should_save
+ and os.path.exists(os.path.join(args.output_dir, TRAINER_LOG))
+ and args.overwrite_output_dir
+ ):
+ logger.warning_rank0_once("Previous trainer log in this folder will be deleted.")
+ os.remove(os.path.join(args.output_dir, TRAINER_LOG))
+
+ @override
+ def on_train_begin(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if args.should_save:
+ self.do_train = True
+ self._reset(max_steps=state.max_steps)
+ self._create_thread_pool(output_dir=args.output_dir)
+
+ @override
+ def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ self._close_thread_pool()
+
+ @override
+ def on_substep_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if self.aborted:
+ control.should_epoch_stop = True
+ control.should_training_stop = True
+
+ @override
+ def on_step_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if self.aborted:
+ control.should_epoch_stop = True
+ control.should_training_stop = True
+
+ @override
+ def on_evaluate(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if not self.do_train:
+ self._close_thread_pool()
+
+ @override
+ def on_predict(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if not self.do_train:
+ self._close_thread_pool()
+
+ @override
+ def on_log(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if not args.should_save:
+ return
+
+ self._timing(cur_steps=state.global_step)
+ logs = dict(
+ current_steps=self.cur_steps,
+ total_steps=self.max_steps,
+ loss=state.log_history[-1].get("loss"),
+ eval_loss=state.log_history[-1].get("eval_loss"),
+ predict_loss=state.log_history[-1].get("predict_loss"),
+ reward=state.log_history[-1].get("reward"),
+ accuracy=state.log_history[-1].get("rewards/accuracies"),
+ lr=state.log_history[-1].get("learning_rate"),
+ epoch=state.log_history[-1].get("epoch"),
+ percentage=round(self.cur_steps / self.max_steps * 100, 2) if self.max_steps != 0 else 100,
+ elapsed_time=self.elapsed_time,
+ remaining_time=self.remaining_time,
+ )
+ if state.num_input_tokens_seen:
+ logs["throughput"] = round(state.num_input_tokens_seen / (time.time() - self.start_time), 2)
+ logs["total_tokens"] = state.num_input_tokens_seen
+
+ if is_env_enabled("RECORD_VRAM"):
+ vram_allocated, vram_reserved = get_peak_memory()
+ logs["vram_allocated"] = round(vram_allocated / (1024**3), 2)
+ logs["vram_reserved"] = round(vram_reserved / (1024**3), 2)
+
+ logs = {k: v for k, v in logs.items() if v is not None}
+ if self.webui_mode and all(key in logs for key in ("loss", "lr", "epoch")):
+ log_str = f"'loss': {logs['loss']:.4f}, 'learning_rate': {logs['lr']:2.4e}, 'epoch': {logs['epoch']:.2f}"
+ for extra_key in ("reward", "accuracy", "throughput"):
+ if logs.get(extra_key):
+ log_str += f", '{extra_key}': {logs[extra_key]:.2f}"
+
+ logger.info_rank0("{" + log_str + "}")
+
+ if self.thread_pool is not None:
+ self.thread_pool.submit(self._write_log, args.output_dir, logs)
+
+ @override
+ def on_prediction_step(
+ self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs
+ ):
+ if self.do_train:
+ return
+
+ if self.aborted:
+ sys.exit(0)
+
+ if not args.should_save:
+ return
+
+ eval_dataloader = kwargs.pop("eval_dataloader", None)
+ if has_length(eval_dataloader):
+ if self.max_steps == 0:
+ self._reset(max_steps=len(eval_dataloader))
+ self._create_thread_pool(output_dir=args.output_dir)
+
+ self._timing(cur_steps=self.cur_steps + 1)
+ if self.cur_steps % 5 == 0 and self.thread_pool is not None:
+ logs = dict(
+ current_steps=self.cur_steps,
+ total_steps=self.max_steps,
+ percentage=round(self.cur_steps / self.max_steps * 100, 2) if self.max_steps != 0 else 100,
+ elapsed_time=self.elapsed_time,
+ remaining_time=self.remaining_time,
+ )
+ self.thread_pool.submit(self._write_log, args.output_dir, logs)
+
+
+class ReporterCallback(TrainerCallback):
+ r"""A callback for reporting training status to external logger."""
+
+ def __init__(
+ self,
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ finetuning_args: "FinetuningArguments",
+ generating_args: "GeneratingArguments",
+ ) -> None:
+ self.model_args = model_args
+ self.data_args = data_args
+ self.finetuning_args = finetuning_args
+ self.generating_args = generating_args
+ os.environ["WANDB_PROJECT"] = os.getenv("WANDB_PROJECT", "llamafactory")
+
+ @override
+ def on_train_begin(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+ if not state.is_world_process_zero:
+ return
+
+ if "wandb" in args.report_to:
+ import wandb
+
+ wandb.config.update(
+ {
+ "model_args": self.model_args.to_dict(),
+ "data_args": self.data_args.to_dict(),
+ "finetuning_args": self.finetuning_args.to_dict(),
+ "generating_args": self.generating_args.to_dict(),
+ }
+ )
+
+ if self.finetuning_args.use_swanlab:
+ import swanlab # type: ignore
+
+ swanlab.config.update(
+ {
+ "model_args": self.model_args.to_dict(),
+ "data_args": self.data_args.to_dict(),
+ "finetuning_args": self.finetuning_args.to_dict(),
+ "generating_args": self.generating_args.to_dict(),
+ }
+ )
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/dpo/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/train/dpo/__init__.py
new file mode 100644
index 0000000..73c1a4a
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/dpo/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from .workflow import run_dpo
+
+
+__all__ = ["run_dpo"]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/dpo/trainer.py b/post-training/LLaMA-Factory/src/llamafactory/train/dpo/trainer.py
new file mode 100644
index 0000000..80f67c6
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/dpo/trainer.py
@@ -0,0 +1,297 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's TRL library.
+# https://github.com/huggingface/trl/blob/v0.8.0/trl/trainer/dpo_trainer.py
+#
+# 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 warnings
+from collections import defaultdict
+from contextlib import nullcontext
+from types import MethodType
+from typing import TYPE_CHECKING, Literal, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from transformers import Trainer
+from trl import DPOTrainer
+from trl.trainer import disable_dropout_in_model
+from typing_extensions import override
+
+from ...extras.constants import IGNORE_INDEX
+from ...extras.packages import is_transformers_version_greater_than
+from ..callbacks import SaveProcessorCallback
+from ..trainer_utils import create_custom_optimizer, create_custom_scheduler, get_batch_logps, nested_detach
+
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedModel, ProcessorMixin
+
+ from ...hparams import FinetuningArguments
+
+
+class CustomDPOTrainer(DPOTrainer):
+ def __init__(
+ self,
+ model: Union["PreTrainedModel", torch.nn.Module],
+ ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]],
+ finetuning_args: "FinetuningArguments",
+ processor: Optional["ProcessorMixin"],
+ disable_dropout: bool = True,
+ **kwargs,
+ ):
+ if is_transformers_version_greater_than("4.46"):
+ kwargs["processing_class"] = kwargs.pop("tokenizer")
+
+ if disable_dropout:
+ disable_dropout_in_model(model)
+ if ref_model is not None:
+ disable_dropout_in_model(ref_model)
+
+ self.finetuning_args = finetuning_args
+ self.f_divergence_type = "reverse_kl"
+ self.reference_free = False
+ self.use_dpo_data_collator = True # hack to avoid warning
+ self.generate_during_eval = False # disable at evaluation
+ self.label_pad_token_id = IGNORE_INDEX
+ self.padding_value = 0
+ self.is_encoder_decoder = model.config.is_encoder_decoder
+ self.precompute_ref_log_probs = False
+ self._precomputed_train_ref_log_probs = False
+ self._precomputed_eval_ref_log_probs = False
+ self._peft_has_been_casted_to_bf16 = False
+
+ self.ref_model = ref_model
+ self._stored_metrics = defaultdict(lambda: defaultdict(list))
+
+ # dpo hyperparams
+ self.beta = finetuning_args.pref_beta
+ self.loss_type = finetuning_args.pref_loss
+ self.ftx_gamma = finetuning_args.pref_ftx
+ self.label_smoothing = finetuning_args.dpo_label_smoothing
+ self.simpo_gamma = finetuning_args.simpo_gamma
+
+ Trainer.__init__(self, model=model, **kwargs)
+ self.model_accepts_loss_kwargs = False # overwrite trainer's default behavior
+ if not hasattr(self, "accelerator"):
+ raise AttributeError("Please update `transformers`.")
+
+ warnings.simplefilter("ignore") # remove gc warnings on ref model
+
+ if ref_model is not None:
+ if self.is_deepspeed_enabled:
+ if not (
+ getattr(ref_model, "is_loaded_in_8bit", False) or getattr(ref_model, "is_loaded_in_4bit", False)
+ ): # quantized models are already set on the correct device
+ self.ref_model = self._prepare_deepspeed(self.ref_model)
+ else:
+ self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+ self.ref_model.eval()
+
+ if processor is not None:
+ self.add_callback(SaveProcessorCallback(processor))
+
+ if finetuning_args.use_badam:
+ from badam import BAdamCallback, clip_grad_norm_old_version # type: ignore
+
+ self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
+ self.add_callback(BAdamCallback)
+
+ @override
+ def create_optimizer(self) -> "torch.optim.Optimizer":
+ if self.optimizer is None:
+ self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
+ return super().create_optimizer()
+
+ @override
+ def create_scheduler(
+ self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
+ ) -> "torch.optim.lr_scheduler.LRScheduler":
+ create_custom_scheduler(self.args, num_training_steps, optimizer)
+ return super().create_scheduler(num_training_steps, optimizer)
+
+ @override
+ def _get_train_sampler(self) -> Optional["torch.utils.data.Sampler"]:
+ if self.finetuning_args.disable_shuffling:
+ return torch.utils.data.SequentialSampler(self.train_dataset)
+
+ return super()._get_train_sampler()
+
+ @override
+ def get_batch_samples(self, *args, **kwargs):
+ r"""Replace the method of DPO Trainer with the one of the standard Trainer."""
+ return Trainer.get_batch_samples(self, *args, **kwargs)
+
+ def odds_ratio_loss(self, chosen_logps: "torch.Tensor", rejected_logps: "torch.Tensor") -> "torch.Tensor":
+ r"""Compute ORPO's odds ratio (OR) loss for batched log probabilities of the policy model."""
+ log_odds = (chosen_logps - rejected_logps) - (
+ torch.log1p(-torch.exp(chosen_logps)) - torch.log1p(-torch.exp(rejected_logps))
+ )
+ sft_loss = -chosen_logps
+ odds_ratio_loss = -F.logsigmoid(log_odds)
+ orpo_loss = sft_loss + self.beta * odds_ratio_loss
+ return orpo_loss
+
+ def simpo_loss(self, chosen_logps: "torch.Tensor", rejected_logps: "torch.Tensor") -> "torch.Tensor":
+ r"""Compute SimPO loss for batched log probabilities of the policy model."""
+ pi_logratios = chosen_logps - rejected_logps
+ gamma_logratios = self.simpo_gamma / self.beta
+ logits = pi_logratios - gamma_logratios
+ simpo_loss = -F.logsigmoid(self.beta * logits)
+ return simpo_loss
+
+ def compute_preference_loss(
+ self,
+ policy_chosen_logps: "torch.Tensor",
+ policy_rejected_logps: "torch.Tensor",
+ reference_chosen_logps: Optional["torch.Tensor"],
+ reference_rejected_logps: Optional["torch.Tensor"],
+ ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
+ r"""Compute loss for preference learning."""
+ if not self.finetuning_args.use_ref_model:
+ if self.loss_type == "orpo":
+ losses = self.odds_ratio_loss(policy_chosen_logps, policy_rejected_logps)
+ elif self.loss_type == "simpo":
+ losses = self.simpo_loss(policy_chosen_logps, policy_rejected_logps)
+ else:
+ raise NotImplementedError(f"Unknown loss type: {self.loss_type}.")
+
+ chosen_rewards = self.beta * policy_chosen_logps.to(self.accelerator.device).detach()
+ rejected_rewards = self.beta * policy_rejected_logps.to(self.accelerator.device).detach()
+ else:
+ losses, chosen_rewards, rejected_rewards = self.dpo_loss(
+ policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps
+ )
+
+ return losses, chosen_rewards, rejected_rewards
+
+ @override
+ def concatenated_forward(
+ self, model: "PreTrainedModel", batch: dict[str, "torch.Tensor"]
+ ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]:
+ r"""Compute the sum log probabilities of the labels under given logits if loss_type is not IPO, ORPO or SimPO.
+
+ Otherwise the average log probabilities.
+ """
+ if self.finetuning_args.use_ref_model:
+ batch = nested_detach(batch, clone=True) # avoid error
+
+ all_logits: torch.Tensor = model(**batch, return_dict=True, use_cache=False).logits.to(torch.float32)
+ all_logps, valid_length = get_batch_logps(logits=all_logits, labels=batch["labels"])
+ if self.loss_type in ["ipo", "orpo", "simpo"]:
+ all_logps = all_logps / valid_length
+
+ batch_size = batch["input_ids"].size(0) // 2
+ chosen_logps, rejected_logps = all_logps.split(batch_size, dim=0)
+ chosen_logits, rejected_logits = all_logits.split(batch_size, dim=0)
+ chosen_length, _ = valid_length.split(batch_size, dim=0)
+
+ if self.loss_type in ["ipo", "orpo", "simpo"]:
+ return chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_logps
+ else:
+ return chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_logps / chosen_length
+
+ @override
+ def compute_reference_log_probs(
+ self, model: "PreTrainedModel", batch: dict[str, "torch.Tensor"]
+ ) -> tuple[Optional["torch.Tensor"], Optional["torch.Tensor"]]:
+ r"""Compute log probabilities of the reference model."""
+ if not self.finetuning_args.use_ref_model:
+ return None, None
+
+ if self.ref_model is None:
+ ref_model = model
+ ref_context = self.accelerator.unwrap_model(model).disable_adapter()
+ else:
+ ref_model = self.ref_model
+ ref_context = nullcontext()
+
+ with torch.no_grad(), ref_context:
+ reference_chosen_logps, reference_rejected_logps, *_ = self.concatenated_forward(ref_model, batch)
+
+ return reference_chosen_logps, reference_rejected_logps
+
+ @override
+ def get_batch_loss_metrics(
+ self,
+ model: "PreTrainedModel",
+ batch: dict[str, "torch.Tensor"],
+ train_eval: Literal["train", "eval"] = "train",
+ ) -> tuple["torch.Tensor", dict[str, "torch.Tensor"]]:
+ r"""Compute the DPO loss and other metrics for the given batch of inputs for train or test."""
+ metrics = {}
+ (
+ policy_chosen_logps,
+ policy_rejected_logps,
+ policy_chosen_logits,
+ policy_rejected_logits,
+ policy_chosen_logps_avg,
+ ) = self.concatenated_forward(model, batch)
+
+ reference_chosen_logps, reference_rejected_logps = self.compute_reference_log_probs(model, batch)
+ losses, chosen_rewards, rejected_rewards = self.compute_preference_loss(
+ policy_chosen_logps,
+ policy_rejected_logps,
+ reference_chosen_logps,
+ reference_rejected_logps,
+ )
+ sft_loss = -policy_chosen_logps_avg
+ if self.ftx_gamma > 1e-6:
+ losses += self.ftx_gamma * sft_loss
+
+ prefix = "eval_" if train_eval == "eval" else ""
+ metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean().item()
+ metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean().item()
+ metrics[f"{prefix}rewards/accuracies"] = (chosen_rewards > rejected_rewards).float().mean().item()
+ metrics[f"{prefix}rewards/margins"] = (chosen_rewards - rejected_rewards).mean().item()
+ metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.mean().item()
+ metrics[f"{prefix}logps/rejected"] = policy_rejected_logps.mean().item()
+ metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.mean().item()
+ metrics[f"{prefix}logits/rejected"] = policy_rejected_logits.mean().item()
+ if self.loss_type == "orpo":
+ metrics[f"{prefix}sft_loss"] = sft_loss.mean().item()
+ metrics[f"{prefix}odds_ratio_loss"] = ((losses - sft_loss) / self.beta).mean().item()
+
+ return losses.mean(), metrics
+
+ @override
+ def compute_loss(
+ self, model: "PreTrainedModel", inputs: dict[str, "torch.Tensor"], return_outputs: bool = False, **kwargs
+ ) -> Union["torch.Tensor", tuple["torch.Tensor", list["torch.Tensor"]]]:
+ r"""Subclass and override to accept extra kwargs."""
+ return super().compute_loss(model, inputs, return_outputs)
+
+ @override
+ def log(self, logs: dict[str, float], *args, **kwargs) -> None:
+ r"""Log `logs` on the various objects watching training, including stored metrics."""
+ # logs either has "loss" or "eval_loss"
+ train_eval = "train" if "loss" in logs else "eval"
+ # Add averaged stored metrics to logs
+ key_list, metric_list = [], []
+ for key, metrics in self._stored_metrics[train_eval].items():
+ key_list.append(key)
+ metric_list.append(torch.tensor(metrics, dtype=torch.float).to(self.accelerator.device).mean().item())
+
+ del self._stored_metrics[train_eval]
+ if len(metric_list) < 10: # pad to for all reduce
+ for i in range(10 - len(metric_list)):
+ key_list.append(f"dummy_{i}")
+ metric_list.append(0.0)
+
+ metric_list = torch.tensor(metric_list, dtype=torch.float).to(self.accelerator.device)
+ metric_list = self.accelerator.reduce(metric_list, "mean").tolist()
+ for key, metric in zip(key_list, metric_list): # add remaining items
+ if not key.startswith("dummy_"):
+ logs[key] = metric
+
+ return Trainer.log(self, logs, *args, **kwargs)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/dpo/workflow.py b/post-training/LLaMA-Factory/src/llamafactory/train/dpo/workflow.py
new file mode 100644
index 0000000..97262ad
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/dpo/workflow.py
@@ -0,0 +1,113 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's TRL library.
+# https://github.com/huggingface/trl/blob/v0.8.0/examples/scripts/dpo.py
+#
+# 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.
+
+from typing import TYPE_CHECKING, Optional
+
+from ...data import PairwiseDataCollatorWithPadding, get_dataset, get_template_and_fix_tokenizer
+from ...extras.constants import IGNORE_INDEX
+from ...extras.misc import calculate_tps
+from ...extras.ploting import plot_loss
+from ...hparams import ModelArguments
+from ...model import load_model, load_tokenizer
+from ..trainer_utils import create_modelcard_and_push, create_ref_model
+from .trainer import CustomDPOTrainer
+
+
+if TYPE_CHECKING:
+ from transformers import Seq2SeqTrainingArguments, TrainerCallback
+
+ from ...hparams import DataArguments, FinetuningArguments
+
+
+def run_dpo(
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ callbacks: Optional[list["TrainerCallback"]] = None,
+):
+ tokenizer_module = load_tokenizer(model_args)
+ tokenizer = tokenizer_module["tokenizer"]
+ template = get_template_and_fix_tokenizer(tokenizer, data_args)
+ dataset_module = get_dataset(template, model_args, data_args, training_args, stage="rm", **tokenizer_module)
+ model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
+
+ data_collator = PairwiseDataCollatorWithPadding(
+ template=template,
+ model=model,
+ pad_to_multiple_of=8,
+ label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id,
+ **tokenizer_module,
+ )
+
+ # Create reference model
+ if finetuning_args.use_ref_model:
+ if finetuning_args.ref_model is None and (not training_args.do_train): # use the model itself
+ ref_model = model
+ else:
+ ref_model = create_ref_model(model_args, finetuning_args)
+ else:
+ ref_model = None
+
+ # Update arguments
+ training_args.remove_unused_columns = False # important for multimodal and pairwise dataset
+
+ # Initialize our Trainer
+ trainer = CustomDPOTrainer(
+ model=model,
+ ref_model=ref_model,
+ args=training_args,
+ finetuning_args=finetuning_args,
+ data_collator=data_collator,
+ callbacks=callbacks,
+ **dataset_module,
+ **tokenizer_module,
+ )
+
+ # Training
+ if training_args.do_train:
+ train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+ trainer.save_model()
+ if finetuning_args.include_effective_tokens_per_second:
+ train_result.metrics["effective_tokens_per_sec"] = calculate_tps(
+ dataset_module["train_dataset"], train_result.metrics, stage="rm"
+ )
+
+ trainer.log_metrics("train", train_result.metrics)
+ trainer.save_metrics("train", train_result.metrics)
+ trainer.save_state()
+ if trainer.is_world_process_zero() and finetuning_args.plot_loss:
+ keys = ["loss", "rewards/accuracies"]
+ if isinstance(dataset_module.get("eval_dataset"), dict):
+ keys += [f"eval_{key}_loss" for key in dataset_module["eval_dataset"].keys()]
+ else:
+ keys += ["eval_loss"]
+
+ plot_loss(training_args.output_dir, keys=keys)
+
+ # Evaluation
+ if training_args.do_eval:
+ metrics = trainer.evaluate(metric_key_prefix="eval")
+ if id(model) == id(ref_model): # unable to compute rewards if reference model is the model itself
+ remove_keys = [key for key in metrics.keys() if "rewards" in key]
+ for key in remove_keys:
+ metrics.pop(key)
+ trainer.log_metrics("eval", metrics)
+ trainer.save_metrics("eval", metrics)
+
+ # Create model card
+ create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/kto/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/train/kto/__init__.py
new file mode 100644
index 0000000..491b067
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/kto/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from .workflow import run_kto
+
+
+__all__ = ["run_kto"]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/kto/trainer.py b/post-training/LLaMA-Factory/src/llamafactory/train/kto/trainer.py
new file mode 100644
index 0000000..0323041
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/kto/trainer.py
@@ -0,0 +1,298 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's TRL library.
+# https://github.com/huggingface/trl/blob/v0.8.0/trl/trainer/kto_trainer.py
+#
+# 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 warnings
+from collections import defaultdict
+from contextlib import nullcontext
+from types import MethodType
+from typing import TYPE_CHECKING, Literal, Optional, Union
+
+import torch
+from transformers import Trainer
+from trl import KTOTrainer
+from trl.trainer import disable_dropout_in_model
+from typing_extensions import override
+
+from ...extras.constants import IGNORE_INDEX
+from ...extras.packages import is_transformers_version_greater_than
+from ..callbacks import SaveProcessorCallback
+from ..trainer_utils import create_custom_optimizer, create_custom_scheduler, get_batch_logps, nested_detach
+
+
+if TYPE_CHECKING:
+ import torch.utils.data
+ from transformers import PreTrainedModel, ProcessorMixin
+
+ from ...hparams import FinetuningArguments
+
+
+class CustomKTOTrainer(KTOTrainer):
+ def __init__(
+ self,
+ model: Union["PreTrainedModel", torch.nn.Module],
+ ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]],
+ finetuning_args: "FinetuningArguments",
+ processor: Optional["ProcessorMixin"],
+ disable_dropout: bool = True,
+ **kwargs,
+ ):
+ if is_transformers_version_greater_than("4.46"):
+ kwargs["processing_class"] = kwargs.pop("tokenizer")
+
+ if disable_dropout:
+ disable_dropout_in_model(model)
+ if ref_model is not None:
+ disable_dropout_in_model(ref_model)
+
+ self.finetuning_args = finetuning_args
+ self.reference_free = False
+ self.use_dpo_data_collator = True # hack to avoid warning
+ self.generate_during_eval = False # disable at evaluation
+ self.label_pad_token_id = IGNORE_INDEX
+ self.padding_value = 0
+ self.is_encoder_decoder = model.config.is_encoder_decoder
+ self.precompute_ref_log_probs = False
+ self._precomputed_train_ref_log_probs = False
+ self._precomputed_eval_ref_log_probs = False
+ self._peft_has_been_casted_to_bf16 = False
+
+ self.ref_model = ref_model
+ self._stored_metrics = defaultdict(lambda: defaultdict(list))
+
+ # kto hyperparams
+ self.beta = finetuning_args.pref_beta
+ self.desirable_weight = finetuning_args.kto_chosen_weight
+ self.undesirable_weight = finetuning_args.kto_rejected_weight
+ self.ftx_gamma = finetuning_args.pref_ftx
+
+ Trainer.__init__(self, model=model, **kwargs)
+ self.model_accepts_loss_kwargs = False # overwrite trainer's default behavior
+ if not hasattr(self, "accelerator"):
+ raise AttributeError("Please update `transformers`.")
+
+ warnings.simplefilter("ignore") # remove gc warnings on ref model
+
+ if ref_model is not None:
+ if self.is_deepspeed_enabled:
+ if not (
+ getattr(ref_model, "is_loaded_in_8bit", False) or getattr(ref_model, "is_loaded_in_4bit", False)
+ ): # quantized models are already set on the correct device
+ self.ref_model = self._prepare_deepspeed(self.ref_model)
+ else:
+ self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+ self.ref_model.eval()
+
+ if processor is not None:
+ self.add_callback(SaveProcessorCallback(processor))
+
+ if finetuning_args.use_badam:
+ from badam import BAdamCallback, clip_grad_norm_old_version # type: ignore
+
+ self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
+ self.add_callback(BAdamCallback)
+
+ @override
+ def create_optimizer(self) -> "torch.optim.Optimizer":
+ if self.optimizer is None:
+ self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
+ return super().create_optimizer()
+
+ @override
+ def create_scheduler(
+ self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
+ ) -> "torch.optim.lr_scheduler.LRScheduler":
+ create_custom_scheduler(self.args, num_training_steps, optimizer)
+ return super().create_scheduler(num_training_steps, optimizer)
+
+ @override
+ def _get_train_sampler(self) -> Optional["torch.utils.data.Sampler"]:
+ r"""Replace the sequential sampler of KTO Trainer created by trl with the random sampler."""
+ if self.finetuning_args.disable_shuffling:
+ return torch.utils.data.SequentialSampler(self.train_dataset)
+
+ return Trainer._get_train_sampler(self)
+
+ @override
+ def get_batch_samples(self, *args, **kwargs):
+ r"""Replace the method of KTO Trainer with the one of the standard Trainer."""
+ return Trainer.get_batch_samples(self, *args, **kwargs)
+
+ @override
+ def forward(
+ self, model: "PreTrainedModel", batch: dict[str, "torch.Tensor"], prefix: Literal["", "kl_"] = ""
+ ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
+ r"""Run forward pass and computes the log probabilities."""
+ batch = nested_detach(batch, clone=True) # avoid error
+ model_inputs = {
+ "input_ids": batch[f"{prefix}input_ids"],
+ "attention_mask": batch[f"{prefix}attention_mask"],
+ }
+ if f"{prefix}token_type_ids" in batch:
+ model_inputs["token_type_ids"] = batch[f"{prefix}token_type_ids"]
+
+ if "pixel_values" in batch:
+ model_inputs["pixel_values"] = batch["pixel_values"]
+
+ if "image_sizes" in batch:
+ model_inputs["image_sizes"] = batch["image_sizes"]
+
+ if "image_grid_thw" in batch:
+ model_inputs["image_grid_thw"] = batch["image_grid_thw"]
+
+ if "aspect_ratio_ids" in batch:
+ model_inputs["aspect_ratio_ids"] = batch["aspect_ratio_ids"]
+
+ if "aspect_ratio_mask" in batch:
+ model_inputs["aspect_ratio_mask"] = batch["aspect_ratio_mask"]
+
+ if f"{prefix}cross_attention_mask" in batch:
+ model_inputs["cross_attention_mask"] = batch[f"{prefix}cross_attention_mask"]
+
+ logits = model(**model_inputs, return_dict=True, use_cache=False).logits.to(torch.float32)
+ logps, valid_length = get_batch_logps(logits=logits, labels=batch[f"{prefix}labels"])
+ return logits, logps, logps / valid_length
+
+ @override
+ def concatenated_forward(
+ self, model: "PreTrainedModel", batch: dict[str, "torch.Tensor"]
+ ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]:
+ target_logits, target_logps, target_logps_avg = self.forward(model, batch)
+ with torch.no_grad():
+ _, kl_logps, _ = self.forward(model, batch, prefix="kl_")
+
+ if len(target_logps) != len(batch["kto_tags"]):
+ raise ValueError("Mismatched shape of inputs and labels.")
+
+ chosen_logits = target_logits[batch["kto_tags"]]
+ chosen_logps = target_logps[batch["kto_tags"]]
+ rejected_logits = target_logits[~batch["kto_tags"]]
+ rejected_logps = target_logps[~batch["kto_tags"]]
+ chosen_logps_avg = target_logps_avg[batch["kto_tags"]]
+ return chosen_logps, rejected_logps, chosen_logits, rejected_logits, kl_logps, chosen_logps_avg
+
+ @override
+ def compute_reference_log_probs(
+ self, model: "PreTrainedModel", batch: dict[str, "torch.Tensor"]
+ ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
+ r"""Compute log probabilities of the reference model."""
+ if self.ref_model is None:
+ ref_model = model
+ ref_context = self.accelerator.unwrap_model(model).disable_adapter()
+ else:
+ ref_model = self.ref_model
+ ref_context = nullcontext()
+
+ with torch.no_grad(), ref_context:
+ reference_chosen_logps, reference_rejected_logps, _, _, reference_kl_logps, _ = self.concatenated_forward(
+ ref_model, batch
+ )
+
+ return reference_chosen_logps, reference_rejected_logps, reference_kl_logps
+
+ @override
+ def get_batch_loss_metrics(
+ self,
+ model: "PreTrainedModel",
+ batch: dict[str, "torch.Tensor"],
+ ) -> tuple["torch.Tensor", dict[str, "torch.Tensor"]]:
+ r"""Compute the DPO loss and other metrics for the given batch of inputs for train or test."""
+ metrics = {}
+ (
+ policy_chosen_logps,
+ policy_rejected_logps,
+ policy_chosen_logits,
+ policy_rejected_logits,
+ policy_kl_logps,
+ policy_chosen_logps_avg,
+ ) = self.concatenated_forward(model, batch)
+ reference_chosen_logps, reference_rejected_logps, reference_kl_logps = self.compute_reference_log_probs(
+ model, batch
+ )
+ losses, chosen_rewards, rejected_rewards, kl = self.kto_loss(
+ policy_chosen_logps,
+ policy_rejected_logps,
+ policy_kl_logps,
+ reference_chosen_logps,
+ reference_rejected_logps,
+ reference_kl_logps,
+ )
+ losses = losses.nanmean()
+
+ if self.ftx_gamma > 1e-6 and len(policy_chosen_logps) > 0: # remember to rescale
+ sft_loss = -policy_chosen_logps_avg
+ losses += self.ftx_gamma * sft_loss.nanmean() / len(policy_chosen_logps) * len(batch["labels"])
+
+ num_chosen = len(chosen_rewards)
+ num_rejected = len(rejected_rewards)
+ if num_chosen > 0:
+ metrics["rewards/chosen_sum"] = chosen_rewards.nansum().item()
+ metrics["logps/chosen_sum"] = policy_chosen_logps.nansum().item()
+ metrics["logits/chosen_sum"] = policy_chosen_logits.nansum().item()
+ metrics["count/chosen"] = float(num_chosen)
+
+ if num_rejected > 0:
+ metrics["rewards/rejected_sum"] = rejected_rewards.nansum().item()
+ metrics["logps/rejected_sum"] = policy_rejected_logps.nansum().item()
+ metrics["logits/rejected_sum"] = policy_rejected_logits.nansum().item()
+ metrics["count/rejected"] = float(num_rejected)
+
+ metrics["kl"] = kl.item()
+ return losses, metrics
+
+ @override
+ def compute_loss(
+ self, model: "PreTrainedModel", inputs: dict[str, "torch.Tensor"], return_outputs: bool = False, **kwargs
+ ) -> Union["torch.Tensor", tuple["torch.Tensor", list["torch.Tensor"]]]:
+ r"""Subclass and override to accept extra kwargs."""
+ return super().compute_loss(model, inputs, return_outputs)
+
+ @override
+ def log(self, logs: dict[str, float], *args, **kwargs) -> None:
+ r"""Log `logs` on the various objects watching training, including stored metrics."""
+ # logs either has "loss" or "eval_loss"
+ train_eval = "train" if "loss" in logs else "eval"
+ prefix = "eval_" if train_eval == "eval" else ""
+ # Add averaged stored metrics to logs
+ key_list, metric_list = [], []
+ for key, metrics in self._stored_metrics[train_eval].items():
+ key_list.append(key)
+ metric_list.append(torch.tensor(metrics, dtype=torch.float).to(self.accelerator.device).sum().item())
+
+ del self._stored_metrics[train_eval]
+ if len(metric_list) < 9: # pad to for all reduce
+ for i in range(9 - len(metric_list)):
+ key_list.append(f"dummy_{i}")
+ metric_list.append(0.0)
+
+ metric_list = torch.tensor(metric_list, dtype=torch.float).to(self.accelerator.device)
+ metric_list = self.accelerator.reduce(metric_list, "sum").tolist()
+ metric_dict: dict[str, float] = dict(zip(key_list, metric_list))
+ for split in ["chosen", "rejected"]: # accumulate average metrics from sums and lengths
+ if f"count/{split}" in metric_dict:
+ for key in ("rewards", "logps", "logits"):
+ logs[f"{prefix}{key}/{split}"] = metric_dict[f"{key}/{split}_sum"] / metric_dict[f"count/{split}"]
+ del metric_dict[f"{key}/{split}_sum"]
+ del metric_dict[f"count/{split}"]
+
+ if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs: # calculate reward margin
+ logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"]
+
+ for key, metric in metric_dict.items(): # add remaining items
+ if not key.startswith("dummy_"):
+ logs[key] = metric
+
+ return Trainer.log(self, logs, *args, **kwargs)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/kto/workflow.py b/post-training/LLaMA-Factory/src/llamafactory/train/kto/workflow.py
new file mode 100644
index 0000000..7b16d1d
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/kto/workflow.py
@@ -0,0 +1,104 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's TRL library.
+# https://github.com/huggingface/trl/blob/v0.8.0/examples/scripts/kto.py
+#
+# 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.
+
+from typing import TYPE_CHECKING, Optional
+
+from ...data import KTODataCollatorWithPadding, get_dataset, get_template_and_fix_tokenizer
+from ...extras.constants import IGNORE_INDEX
+from ...extras.ploting import plot_loss
+from ...hparams import ModelArguments
+from ...model import load_model, load_tokenizer
+from ..trainer_utils import create_modelcard_and_push, create_ref_model
+from .trainer import CustomKTOTrainer
+
+
+if TYPE_CHECKING:
+ from transformers import Seq2SeqTrainingArguments, TrainerCallback
+
+ from ...hparams import DataArguments, FinetuningArguments
+
+
+def run_kto(
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ callbacks: Optional[list["TrainerCallback"]] = None,
+):
+ tokenizer_module = load_tokenizer(model_args)
+ tokenizer = tokenizer_module["tokenizer"]
+ template = get_template_and_fix_tokenizer(tokenizer, data_args)
+ dataset_module = get_dataset(template, model_args, data_args, training_args, stage="kto", **tokenizer_module)
+ model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
+
+ data_collator = KTODataCollatorWithPadding(
+ template=template,
+ model=model,
+ pad_to_multiple_of=8,
+ label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id,
+ **tokenizer_module,
+ )
+
+ # Create reference model
+ if finetuning_args.ref_model is None and (not training_args.do_train): # use the model itself
+ ref_model = model
+ else:
+ ref_model = create_ref_model(model_args, finetuning_args)
+
+ # Update arguments
+ training_args.remove_unused_columns = False # important for multimodal and pairwise dataset
+
+ # Initialize our Trainer
+ trainer = CustomKTOTrainer(
+ model=model,
+ ref_model=ref_model,
+ args=training_args,
+ finetuning_args=finetuning_args,
+ data_collator=data_collator,
+ callbacks=callbacks,
+ **dataset_module,
+ **tokenizer_module,
+ )
+
+ # Training
+ if training_args.do_train:
+ train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+ trainer.save_model()
+ trainer.log_metrics("train", train_result.metrics)
+ trainer.save_metrics("train", train_result.metrics)
+ trainer.save_state()
+ if trainer.is_world_process_zero() and finetuning_args.plot_loss:
+ keys = ["loss", "rewards/chosen"]
+ if isinstance(dataset_module.get("eval_dataset"), dict):
+ keys += [f"eval_{key}_loss" for key in dataset_module["eval_dataset"].keys()]
+ else:
+ keys += ["eval_loss"]
+
+ plot_loss(training_args.output_dir, keys=keys)
+
+ # Evaluation
+ if training_args.do_eval:
+ metrics = trainer.evaluate(metric_key_prefix="eval")
+ if id(model) == id(ref_model): # unable to compute rewards without a reference model
+ remove_keys = [key for key in metrics.keys() if "rewards" in key]
+ for key in remove_keys:
+ metrics.pop(key)
+ trainer.log_metrics("eval", metrics)
+ trainer.save_metrics("eval", metrics)
+
+ # Create model card
+ create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/ppo/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/__init__.py
new file mode 100644
index 0000000..ed9bc4d
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from .workflow import run_ppo
+
+
+__all__ = ["run_ppo"]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/ppo/ppo_utils.py b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/ppo_utils.py
new file mode 100644
index 0000000..9d462e7
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/ppo_utils.py
@@ -0,0 +1,80 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Literal, Optional
+
+import torch
+from transformers.integrations import is_deepspeed_zero3_enabled
+
+from ...extras.packages import is_requests_available
+
+
+if is_requests_available():
+ import requests
+
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedModel
+ from trl import AutoModelForCausalLMWithValueHead
+
+
+def get_rewards_from_server(server_url: str, messages: list[str]) -> list["torch.Tensor"]:
+ r"""Get reward scores from the API server."""
+ headers = {"Content-Type": "application/json"}
+ payload = {"model": "model", "messages": messages}
+ response = requests.post(server_url, json=payload, headers=headers)
+ rewards = json.loads(response.text)["scores"]
+ return torch.Tensor(rewards)
+
+
+def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["default", "reward"]) -> None:
+ r"""Replace the default/reward modules in the model. The model is already unwrapped."""
+ v_head_layer = model.v_head.summary
+ if is_deepspeed_zero3_enabled():
+ import deepspeed # type: ignore
+
+ params = [v_head_layer.weight, v_head_layer.bias]
+ context_maybe_zero3 = deepspeed.zero.GatheredParameters(params, modifier_rank=0)
+ else:
+ context_maybe_zero3 = nullcontext()
+
+ model.pretrained_model.set_adapter(target) # set the LoRA adapter to be active
+ with context_maybe_zero3:
+ if target == "reward": # save default head temporarily
+ setattr(model, "default_head_weight", v_head_layer.weight.data.detach().clone())
+ setattr(model, "default_head_bias", v_head_layer.bias.data.detach().clone())
+
+ device = v_head_layer.weight.device
+ v_head_layer.weight.data = model.get_buffer(f"{target}_head_weight").detach().clone().to(device)
+ v_head_layer.bias.data = model.get_buffer(f"{target}_head_bias").detach().clone().to(device)
+
+
+def dump_layernorm(model: "PreTrainedModel") -> dict[str, "torch.Tensor"]:
+ r"""Dump the layernorm parameters in the model. The model is already unwrapped (and gathered)."""
+ layer_norm_params = {}
+ for name, param in model.named_parameters():
+ if param.data.dtype == torch.float32:
+ layer_norm_params[name] = param.data.detach().clone()
+ param.data = param.data.to(model.config.torch_dtype)
+
+ return layer_norm_params
+
+
+def restore_layernorm(model: "PreTrainedModel", layernorm_params: Optional[dict[str, "torch.Tensor"]] = None) -> None:
+ r"""Restore the layernorm parameters in the model. The model is already unwrapped (and gathered)."""
+ for name, param in model.named_parameters():
+ if name in layernorm_params:
+ param.data = layernorm_params[name]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/ppo/trainer.py b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/trainer.py
new file mode 100644
index 0000000..1449745
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/trainer.py
@@ -0,0 +1,503 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's TRL library.
+# https://github.com/huggingface/trl/blob/v0.8.0/trl/trainer/ppo_trainer.py
+#
+# 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 math
+import os
+import sys
+import warnings
+from types import MethodType
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+from accelerate.utils import DistributedDataParallelKwargs
+from tqdm import tqdm
+from transformers import GenerationConfig, Trainer, TrainerControl, TrainerState
+from transformers.optimization import get_scheduler
+from transformers.trainer import DEFAULT_CALLBACKS
+from transformers.trainer_callback import CallbackHandler
+from transformers.trainer_pt_utils import remove_dummy_checkpoint
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+from transformers.utils import SAFE_WEIGHTS_NAME, WEIGHTS_NAME
+from trl import PPOConfig, PPOTrainer
+from trl.core import PPODecorators, logprobs_from_logits
+from trl.models.utils import unwrap_model_for_generation
+from typing_extensions import override
+
+from ...extras import logging
+from ...extras.misc import AverageMeter, count_parameters, get_current_device, get_logits_processor
+from ..callbacks import FixValueHeadModelCallback, SaveProcessorCallback
+from ..trainer_utils import create_custom_optimizer, create_custom_scheduler
+from .ppo_utils import dump_layernorm, get_rewards_from_server, replace_model, restore_layernorm
+
+
+if TYPE_CHECKING:
+ from datasets import Dataset
+ from transformers import (
+ DataCollatorWithPadding,
+ PreTrainedTokenizer,
+ ProcessorMixin,
+ Seq2SeqTrainingArguments,
+ TrainerCallback,
+ )
+ from trl import AutoModelForCausalLMWithValueHead
+
+ from ...hparams import FinetuningArguments, GeneratingArguments, ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+class CustomPPOTrainer(PPOTrainer, Trainer):
+ r"""Inherit PPOTrainer."""
+
+ def __init__(
+ self,
+ model_args: "ModelArguments",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ generating_args: "GeneratingArguments",
+ callbacks: Optional[list["TrainerCallback"]],
+ model: "AutoModelForCausalLMWithValueHead",
+ reward_model: Optional["AutoModelForCausalLMWithValueHead"],
+ ref_model: Optional["AutoModelForCausalLMWithValueHead"],
+ tokenizer: "PreTrainedTokenizer",
+ processor: Optional["ProcessorMixin"],
+ data_collator: "DataCollatorWithPadding",
+ train_dataset: Optional["Dataset"] = None,
+ eval_dataset: Optional["Dataset"] = None,
+ ) -> None:
+ if eval_dataset is not None:
+ raise NotImplementedError("PPOTrainer does not support eval dataset yet.")
+
+ backward_batch_size = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps
+ ppo_config = PPOConfig(
+ model_name=model_args.model_name_or_path,
+ learning_rate=training_args.learning_rate,
+ mini_batch_size=training_args.per_device_train_batch_size,
+ batch_size=backward_batch_size * finetuning_args.ppo_buffer_size,
+ gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+ ppo_epochs=finetuning_args.ppo_epochs,
+ max_grad_norm=training_args.max_grad_norm,
+ seed=training_args.seed,
+ optimize_device_cache=True,
+ target=finetuning_args.ppo_target,
+ use_score_scaling=finetuning_args.ppo_score_norm,
+ use_score_norm=finetuning_args.ppo_score_norm,
+ whiten_rewards=finetuning_args.ppo_whiten_rewards,
+ accelerator_kwargs={"step_scheduler_with_optimizer": False},
+ log_with=training_args.report_to[0] if training_args.report_to else None,
+ project_kwargs={"logging_dir": training_args.logging_dir},
+ )
+
+ # Add deepspeed config
+ if training_args.deepspeed_plugin is not None:
+ ppo_config.accelerator_kwargs["kwargs_handlers"] = [
+ DistributedDataParallelKwargs(find_unused_parameters=training_args.ddp_find_unused_parameters)
+ ]
+ ppo_config.accelerator_kwargs["deepspeed_plugin"] = training_args.deepspeed_plugin
+ if ppo_config.log_with is not None:
+ logger.warning_rank0("PPOTrainer cannot use external logger when DeepSpeed is enabled.")
+ ppo_config.log_with = None
+
+ # Create optimizer and scheduler
+ if training_args.max_steps > 0:
+ num_training_steps = training_args.max_steps
+ else:
+ total_train_batch_size = backward_batch_size * finetuning_args.ppo_buffer_size * training_args.world_size
+ num_training_steps = training_args.num_train_epochs * math.ceil(
+ len(train_dataset) / total_train_batch_size
+ )
+
+ optimizer = self.create_optimizer(model, training_args, finetuning_args)
+ scheduler = self.create_scheduler(training_args, num_training_steps, optimizer)
+
+ PPOTrainer.__init__(
+ self,
+ config=ppo_config,
+ model=model,
+ ref_model=ref_model,
+ tokenizer=tokenizer,
+ dataset=train_dataset,
+ optimizer=optimizer,
+ data_collator=data_collator,
+ lr_scheduler=scheduler,
+ )
+
+ self.args = training_args
+ self.model_args = model_args
+ self.finetuning_args = finetuning_args
+ self.reward_model = reward_model
+ self.current_device = get_current_device() # patch for deepspeed training
+
+ self.generation_config = GenerationConfig(
+ pad_token_id=self.tokenizer.pad_token_id,
+ eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids,
+ **generating_args.to_dict(),
+ )
+
+ self.state = TrainerState()
+ self.control = TrainerControl()
+ self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None
+ self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None
+ callbacks = DEFAULT_CALLBACKS if callbacks is None else DEFAULT_CALLBACKS + callbacks
+ self.callback_handler = CallbackHandler(
+ callbacks, self.accelerator.unwrap_model(self.model), self.tokenizer, self.optimizer, self.lr_scheduler
+ )
+ if self.args.max_steps > 0:
+ logger.info_rank0("max_steps is given, it will override any value given in num_train_epochs")
+
+ self.amp_context = torch.autocast(self.current_device.type)
+ warnings.simplefilter("ignore") # remove gc warnings on ref model
+
+ if finetuning_args.reward_model_type == "full":
+ if self.is_deepspeed_enabled:
+ if not (
+ getattr(reward_model.pretrained_model, "is_loaded_in_8bit", False)
+ or getattr(reward_model.pretrained_model, "is_loaded_in_4bit", False)
+ ): # quantized models are already set on the correct device
+ self.reward_model = self._prepare_deepspeed(self.reward_model)
+ else:
+ self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True)
+
+ self.add_callback(FixValueHeadModelCallback)
+
+ if processor is not None:
+ self.add_callback(SaveProcessorCallback(processor))
+
+ if finetuning_args.use_badam:
+ from badam import BAdamCallback, clip_grad_norm_old_version # type: ignore
+
+ self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
+ self.add_callback(BAdamCallback)
+
+ def ppo_train(self, resume_from_checkpoint: Optional[str] = None) -> None:
+ r"""Implement training loop for the PPO stage, like _inner_training_loop() in Huggingface's Trainer."""
+ if resume_from_checkpoint is not None:
+ raise ValueError("`resume_from_checkpoint` will be supported in the future version.")
+
+ total_train_batch_size = (
+ self.args.per_device_train_batch_size
+ * self.args.gradient_accumulation_steps
+ * self.finetuning_args.ppo_buffer_size
+ * self.args.world_size
+ )
+ if self.args.max_steps > 0:
+ num_examples = total_train_batch_size * self.args.max_steps
+ num_train_epochs = sys.maxsize
+ max_steps = self.args.max_steps
+ steps_in_epoch = self.args.max_steps
+ else:
+ len_dataloader = len(self.dataloader)
+ num_examples = len(self.dataset)
+ num_train_epochs = self.args.num_train_epochs
+ max_steps = math.ceil(num_train_epochs * len_dataloader)
+ steps_in_epoch = len_dataloader
+
+ self.state.max_steps = max_steps
+ self.state.num_train_epochs = num_train_epochs
+ self.state.is_local_process_zero = self.is_local_process_zero()
+ self.state.is_world_process_zero = self.is_world_process_zero()
+
+ logger.info_rank0("***** Running training *****")
+ logger.info_rank0(f" Num examples = {num_examples:,}")
+ logger.info_rank0(f" Num Epochs = {num_train_epochs:,}")
+ logger.info_rank0(f" Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}")
+ logger.info_rank0(
+ f" Total train batch size (w. parallel, buffer, distributed & accumulation) = {total_train_batch_size:,}"
+ )
+ logger.info_rank0(f" Gradient Accumulation steps = {self.args.gradient_accumulation_steps:,}")
+ logger.info_rank0(f" Num optimization epochs per batch = {self.finetuning_args.ppo_epochs:,}")
+ logger.info_rank0(f" Total training steps = {max_steps:,}")
+ logger.info_rank0(f" Number of trainable parameters = {count_parameters(self.model)[0]:,}")
+
+ dataiter = iter(self.dataloader)
+ loss_meter = AverageMeter()
+ reward_meter = AverageMeter()
+ self.callback_handler.on_train_begin(self.args, self.state, self.control)
+
+ for step in tqdm(range(max_steps), disable=not self.is_local_process_zero()):
+ try:
+ batch = next(dataiter)
+ except StopIteration:
+ dataiter = iter(self.dataloader)
+ batch = next(dataiter)
+
+ # Get inputs
+ self.model.eval()
+ self.tokenizer.padding_side = "right" # change padding side
+ queries, responses, rewards = [], [], []
+ for idx in range(0, self.config.batch_size, self.config.mini_batch_size):
+ mini_batch = {
+ "input_ids": batch["input_ids"][idx : idx + self.config.mini_batch_size],
+ "attention_mask": batch["attention_mask"][idx : idx + self.config.mini_batch_size],
+ }
+ mini_batch_queries, mini_batch_responses = self.get_inputs(mini_batch)
+ mini_batch_rewards = self.get_rewards(mini_batch_queries, mini_batch_responses)
+ queries.extend(mini_batch_queries)
+ responses.extend(mini_batch_responses)
+ rewards.extend(mini_batch_rewards)
+
+ # Run PPO step
+ self.model.train()
+ stats = self.step(queries, responses, rewards)
+ self.tokenizer.padding_side = "left" # restore padding side
+ loss_meter.update(float(stats["ppo/loss/total"]), n=len(rewards))
+ reward_meter.update(torch.stack(rewards).mean().item(), n=len(rewards))
+
+ if self.config.log_with is not None:
+ try:
+ batch["query"] = self.tokenizer.batch_decode(queries, skip_special_tokens=True)
+ batch["response"] = self.tokenizer.batch_decode(responses, skip_special_tokens=True)
+ self.log_stats(stats, batch, rewards)
+ except Exception:
+ logger.warning_rank0("Failed to save stats due to unknown errors.")
+
+ self.state.global_step += 1
+ self.callback_handler.on_step_end(self.args, self.state, self.control)
+
+ if self.is_local_process_zero() and (step + 1) % self.args.logging_steps == 0:
+ logs = dict(
+ loss=round(loss_meter.avg, 4),
+ reward=round(reward_meter.avg, 4),
+ learning_rate=stats["ppo/learning_rate"],
+ epoch=round(step / steps_in_epoch, 2),
+ )
+ tqdm.write(str(logs))
+ logs["step"] = step
+ self.state.log_history.append(logs)
+ self.callback_handler.on_log(self.args, self.state, self.control, logs)
+ loss_meter.reset()
+ reward_meter.reset()
+
+ if (step + 1) % self.args.save_steps == 0: # save checkpoint
+ self.save_model(
+ os.path.join(self.args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}")
+ )
+ self.callback_handler.on_save(self.args, self.state, self.control)
+
+ if self.control.should_epoch_stop or self.control.should_training_stop:
+ break
+
+ self.callback_handler.on_train_end(self.args, self.state, self.control)
+
+ @override
+ def create_optimizer(
+ self,
+ model: "AutoModelForCausalLMWithValueHead",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ ) -> "torch.optim.Optimizer":
+ optimizer = create_custom_optimizer(model, training_args, finetuning_args)
+ if optimizer is None:
+ decay_params, nodecay_params = [], []
+ decay_param_names = self.get_decay_parameter_names(model)
+ for name, param in model.named_parameters():
+ if param.requires_grad:
+ if name in decay_param_names:
+ decay_params.append(param)
+ else:
+ nodecay_params.append(param)
+
+ optim_class, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
+ param_groups = [
+ dict(params=nodecay_params),
+ dict(params=decay_params, weight_decay=training_args.weight_decay),
+ ]
+ optimizer = optim_class(param_groups, **optim_kwargs)
+
+ return optimizer
+
+ @override
+ def create_scheduler(
+ self, training_args: "Seq2SeqTrainingArguments", num_training_steps: int, optimizer: "torch.optim.Optimizer"
+ ) -> "torch.optim.lr_scheduler.LRScheduler":
+ create_custom_scheduler(training_args, num_training_steps, optimizer)
+ lr_scheduler = get_scheduler(
+ training_args.lr_scheduler_type,
+ optimizer=optimizer,
+ num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
+ num_training_steps=num_training_steps,
+ )
+ return lr_scheduler
+
+ @torch.no_grad()
+ def get_inputs(self, batch: dict[str, "torch.Tensor"]) -> tuple[list["torch.Tensor"], list["torch.Tensor"]]:
+ r"""Generate model's responses given queries."""
+ if batch["input_ids"].size(0) == 1: # handle llama2 ppo with gradient accumulation > 1
+ start_index = (batch["input_ids"][0] != self.tokenizer.pad_token_id).nonzero()[0].item()
+ for k, v in batch.items():
+ batch[k] = v[:, start_index:]
+
+ with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:
+ unwrapped_model: AutoModelForCausalLMWithValueHead = self.accelerator.unwrap_model(self.model)
+ if self.model_args.upcast_layernorm:
+ layernorm_params = dump_layernorm(unwrapped_model)
+
+ generate_output: torch.Tensor = unwrapped_model.generate(
+ generation_config=self.generation_config, logits_processor=get_logits_processor(), **batch
+ )
+ if self.model_args.upcast_layernorm:
+ restore_layernorm(unwrapped_model, layernorm_params)
+
+ query = batch["input_ids"].detach().cpu()
+ response = generate_output[:, batch["input_ids"].size(-1) :].detach().cpu()
+ queries, responses = [], []
+ for i in range(len(query)):
+ query_start_index = (query[i] != self.tokenizer.pad_token_id).nonzero()[0].item()
+ response_indexes = (response[i] != self.tokenizer.pad_token_id).nonzero()
+
+ if len(response_indexes) == 0: # allow empty response
+ response_length = 1
+ elif self.tokenizer.eos_token_id == self.tokenizer.pad_token_id: # include eos token
+ response_length = response_indexes[-1].item() + 2
+ else:
+ response_length = response_indexes[-1].item() + 1
+
+ queries.append(query[i, query_start_index:]) # remove padding from left
+ responses.append(response[i, :response_length]) # remove padding from right
+
+ return queries, responses
+
+ @torch.no_grad()
+ def get_rewards(
+ self,
+ queries: list["torch.Tensor"],
+ responses: list["torch.Tensor"],
+ ) -> list["torch.Tensor"]:
+ r"""Compute scores using given reward model.
+
+ Both inputs and outputs are put on CPU.
+ """
+ if self.finetuning_args.reward_model_type == "api":
+ token_ids = [torch.cat((q, r), dim=-1).tolist() for q, r in zip(queries, responses)]
+ messages = self.tokenizer.batch_decode(token_ids, skip_special_tokens=False)
+ return get_rewards_from_server(self.reward_model, messages)
+
+ batch: dict[str, torch.Tensor] = self.prepare_model_inputs(queries, responses)
+ unwrapped_model: AutoModelForCausalLMWithValueHead = self.accelerator.unwrap_model(self.model)
+
+ if self.finetuning_args.reward_model_type == "lora":
+ replace_model(unwrapped_model, target="reward")
+ reward_model = self.model
+ else:
+ reward_model = self.reward_model
+
+ with unwrap_model_for_generation(reward_model, self.accelerator), self.amp_context: # support bf16
+ values: torch.Tensor = reward_model(**batch, return_dict=True, use_cache=False)[-1]
+
+ if self.finetuning_args.reward_model_type == "lora":
+ replace_model(unwrapped_model, target="default")
+
+ rewards = values.gather(dim=-1, index=(batch["attention_mask"].sum(dim=-1, keepdim=True) - 1))
+ return rewards.float().detach() # use fp32 type
+
+ @override
+ @PPODecorators.empty_device_cache()
+ def batched_forward_pass(
+ self,
+ model: "AutoModelForCausalLMWithValueHead",
+ queries: "torch.Tensor",
+ responses: "torch.Tensor",
+ model_inputs: dict[str, Any],
+ return_logits: bool = False,
+ response_masks: Optional["torch.Tensor"] = None,
+ ) -> tuple["torch.Tensor", Optional["torch.Tensor"], "torch.Tensor", "torch.Tensor"]:
+ r"""Calculate model outputs in multiple batches.
+
+ Subclass and override to inject custom behavior.
+ """
+ bs = len(queries)
+ fbs = self.config.mini_batch_size
+ all_logprobs = []
+ all_logits = []
+ all_masks = []
+ all_values = []
+
+ for i in range(math.ceil(bs / fbs)):
+ input_kwargs = {key: value[i * fbs : (i + 1) * fbs] for key, value in model_inputs.items()}
+ query_batch = queries[i * fbs : (i + 1) * fbs]
+ response_batch = responses[i * fbs : (i + 1) * fbs]
+ if response_masks is not None:
+ response_masks_batch = response_masks[i * fbs : (i + 1) * fbs]
+ input_ids = input_kwargs["input_ids"]
+ attention_mask = input_kwargs["attention_mask"]
+
+ with self.amp_context: # support bf16
+ logits, _, values = model(**input_kwargs, return_dict=True, use_cache=False)
+
+ logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
+ masks = torch.zeros_like(attention_mask)
+ masks[:, :-1] = attention_mask[:, 1:]
+
+ for j in range(len(query_batch)):
+ start = len(query_batch[j]) - 1
+ if attention_mask[j, 0] == 0: # offset left padding
+ start += attention_mask[j, :].nonzero()[0].item()
+ end = start + len(response_batch[j])
+
+ if response_masks is not None:
+ response_masks_batch = torch.cat((torch.zeros_like(query_batch[j]), response_masks_batch[j]))[1:]
+
+ masks[j, :start] = 0
+ masks[j, end:] = 0
+ if response_masks is not None:
+ masks[j, start:end] = masks[j, start:end] * response_masks_batch[j][start:end]
+
+ if return_logits:
+ all_logits.append(logits)
+ else:
+ del logits
+
+ all_values.append(values)
+ all_logprobs.append(logprobs)
+ all_masks.append(masks)
+
+ return (
+ torch.cat(all_logprobs),
+ torch.cat(all_logits)[:, :-1] if return_logits else None,
+ torch.cat(all_values)[:, :-1],
+ torch.cat(all_masks)[:, :-1],
+ )
+
+ @override
+ def save_model(self, output_dir: Optional[str] = None) -> None:
+ r"""Save model checkpoint.
+
+ Subclass and override to inject custom behavior.
+ """
+ if output_dir is None:
+ output_dir = self.args.output_dir
+
+ if self.is_fsdp_enabled or self.is_deepspeed_enabled:
+ try:
+ state_dict = self.accelerator.get_state_dict(self.model) # must be called at all ranks
+ if self.args.should_save:
+ self._save(output_dir, state_dict=state_dict)
+ except ValueError:
+ logger.warning_rank0(
+ " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead,"
+ " use zero_to_fp32.py to recover weights"
+ )
+ if self.args.should_save:
+ self._save(output_dir, state_dict={})
+ # remove the dummy state_dict
+ remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
+ self.model.save_checkpoint(output_dir)
+
+ elif self.args.should_save:
+ unwrapped_model: AutoModelForCausalLMWithValueHead = self.accelerator.unwrap_model(self.model)
+ self._save(output_dir, state_dict=unwrapped_model.state_dict())
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/ppo/workflow.py b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/workflow.py
new file mode 100644
index 0000000..fa6629a
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/ppo/workflow.py
@@ -0,0 +1,79 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's TRL library.
+# https://github.com/huggingface/trl/blob/v0.8.0/examples/scripts/ppo.py
+#
+# 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.
+
+from typing import TYPE_CHECKING, Optional
+
+from ...data import MultiModalDataCollatorForSeq2Seq, get_dataset, get_template_and_fix_tokenizer
+from ...extras.ploting import plot_loss
+from ...model import load_model, load_tokenizer
+from ..callbacks import fix_valuehead_checkpoint
+from ..trainer_utils import create_ref_model, create_reward_model
+from .trainer import CustomPPOTrainer
+
+
+if TYPE_CHECKING:
+ from transformers import Seq2SeqTrainingArguments, TrainerCallback
+
+ from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments
+
+
+def run_ppo(
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ generating_args: "GeneratingArguments",
+ callbacks: Optional[list["TrainerCallback"]] = None,
+):
+ tokenizer_module = load_tokenizer(model_args)
+ tokenizer = tokenizer_module["tokenizer"]
+ template = get_template_and_fix_tokenizer(tokenizer, data_args)
+ dataset_module = get_dataset(template, model_args, data_args, training_args, stage="ppo", **tokenizer_module)
+ model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train, add_valuehead=True)
+
+ tokenizer.padding_side = "left" # use left-padding in generation while using right-padding in training
+ data_collator = MultiModalDataCollatorForSeq2Seq(template=template, model=model, **tokenizer_module)
+
+ # Create reference model and reward model
+ ref_model = create_ref_model(model_args, finetuning_args, add_valuehead=True)
+ reward_model = create_reward_model(model, model_args, finetuning_args)
+
+ # Initialize our Trainer
+ ppo_trainer: CustomPPOTrainer = CustomPPOTrainer(
+ model_args=model_args,
+ training_args=training_args,
+ finetuning_args=finetuning_args,
+ generating_args=generating_args,
+ callbacks=callbacks,
+ model=model,
+ reward_model=reward_model,
+ ref_model=ref_model,
+ data_collator=data_collator,
+ **dataset_module,
+ **tokenizer_module,
+ )
+
+ # Training
+ if training_args.do_train:
+ ppo_trainer.ppo_train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+ ppo_trainer.save_model()
+ if training_args.should_save:
+ fix_valuehead_checkpoint(model, training_args.output_dir, training_args.save_safetensors)
+
+ ppo_trainer.save_state() # must be called after save_model to have a folder
+ if ppo_trainer.is_world_process_zero() and finetuning_args.plot_loss:
+ plot_loss(training_args.output_dir, keys=["loss", "reward"])
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/pt/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/train/pt/__init__.py
new file mode 100644
index 0000000..1f5c289
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/pt/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from .workflow import run_pt
+
+
+__all__ = ["run_pt"]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/pt/trainer.py b/post-training/LLaMA-Factory/src/llamafactory/train/pt/trainer.py
new file mode 100644
index 0000000..8495bbb
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/pt/trainer.py
@@ -0,0 +1,81 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from types import MethodType
+from typing import TYPE_CHECKING, Optional
+
+import torch
+from transformers import Trainer
+from typing_extensions import override
+
+from ...extras.packages import is_transformers_version_greater_than
+from ..callbacks import SaveProcessorCallback
+from ..trainer_utils import create_custom_optimizer, create_custom_scheduler
+
+
+if TYPE_CHECKING:
+ from transformers import ProcessorMixin
+
+ from ...hparams import FinetuningArguments
+
+
+class CustomTrainer(Trainer):
+ r"""Inherit Trainer for custom optimizer."""
+
+ def __init__(
+ self, finetuning_args: "FinetuningArguments", processor: Optional["ProcessorMixin"], **kwargs
+ ) -> None:
+ if is_transformers_version_greater_than("4.46"):
+ kwargs["processing_class"] = kwargs.pop("tokenizer")
+
+ super().__init__(**kwargs)
+ if processor is not None:
+ # avoid wrong loss under gradient accumulation
+ # https://github.com/huggingface/transformers/pull/36044#issuecomment-2746657112
+ self.model_accepts_loss_kwargs = False
+
+ self.finetuning_args = finetuning_args
+
+ if processor is not None:
+ self.add_callback(SaveProcessorCallback(processor))
+
+ if finetuning_args.use_badam:
+ from badam import BAdamCallback, clip_grad_norm_old_version # type: ignore
+
+ self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
+ self.add_callback(BAdamCallback)
+
+ @override
+ def create_optimizer(self) -> "torch.optim.Optimizer":
+ if self.optimizer is None:
+ self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
+ return super().create_optimizer()
+
+ @override
+ def create_scheduler(
+ self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
+ ) -> "torch.optim.lr_scheduler.LRScheduler":
+ create_custom_scheduler(self.args, num_training_steps, optimizer)
+ return super().create_scheduler(num_training_steps, optimizer)
+
+ @override
+ def _get_train_sampler(self) -> Optional["torch.utils.data.Sampler"]:
+ if self.finetuning_args.disable_shuffling:
+ return torch.utils.data.SequentialSampler(self.train_dataset)
+
+ return super()._get_train_sampler()
+
+ @override
+ def compute_loss(self, model, inputs, *args, **kwargs):
+ return super().compute_loss(model, inputs, *args, **kwargs)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/pt/workflow.py b/post-training/LLaMA-Factory/src/llamafactory/train/pt/workflow.py
new file mode 100644
index 0000000..c6f4829
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/pt/workflow.py
@@ -0,0 +1,90 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's transformers library.
+# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/language-modeling/run_clm.py
+#
+# 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 math
+from typing import TYPE_CHECKING, Optional
+
+from transformers import DataCollatorForLanguageModeling
+
+from ...data import get_dataset, get_template_and_fix_tokenizer
+from ...extras.ploting import plot_loss
+from ...model import load_model, load_tokenizer
+from ..trainer_utils import create_modelcard_and_push
+from .trainer import CustomTrainer
+
+
+if TYPE_CHECKING:
+ from transformers import Seq2SeqTrainingArguments, TrainerCallback
+
+ from ...hparams import DataArguments, FinetuningArguments, ModelArguments
+
+
+def run_pt(
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ callbacks: Optional[list["TrainerCallback"]] = None,
+):
+ tokenizer_module = load_tokenizer(model_args)
+ tokenizer = tokenizer_module["tokenizer"]
+ template = get_template_and_fix_tokenizer(tokenizer, data_args)
+ dataset_module = get_dataset(template, model_args, data_args, training_args, stage="pt", **tokenizer_module)
+ model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
+ data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+ # Initialize our Trainer
+ trainer = CustomTrainer(
+ model=model,
+ args=training_args,
+ finetuning_args=finetuning_args,
+ data_collator=data_collator,
+ callbacks=callbacks,
+ **dataset_module,
+ **tokenizer_module,
+ )
+
+ # Training
+ if training_args.do_train:
+ train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+ trainer.save_model()
+ trainer.log_metrics("train", train_result.metrics)
+ trainer.save_metrics("train", train_result.metrics)
+ trainer.save_state()
+ if trainer.is_world_process_zero() and finetuning_args.plot_loss:
+ keys = ["loss"]
+ if isinstance(dataset_module.get("eval_dataset"), dict):
+ keys += [f"eval_{key}_loss" for key in dataset_module["eval_dataset"].keys()]
+ else:
+ keys += ["eval_loss"]
+
+ plot_loss(training_args.output_dir, keys=keys)
+
+ # Evaluation
+ if training_args.do_eval:
+ metrics = trainer.evaluate(metric_key_prefix="eval")
+ try:
+ perplexity = math.exp(metrics["eval_loss"])
+ except OverflowError:
+ perplexity = float("inf")
+
+ metrics["perplexity"] = perplexity
+ trainer.log_metrics("eval", metrics)
+ trainer.save_metrics("eval", metrics)
+
+ # Create model card
+ create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/rm/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/train/rm/__init__.py
new file mode 100644
index 0000000..f0e8a45
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/rm/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from .workflow import run_rm
+
+
+__all__ = ["run_rm"]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/rm/metric.py b/post-training/LLaMA-Factory/src/llamafactory/train/rm/metric.py
new file mode 100644
index 0000000..a7c3c43
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/rm/metric.py
@@ -0,0 +1,51 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import numpy as np
+
+from ...extras.misc import numpify
+
+
+if TYPE_CHECKING:
+ from transformers import EvalPrediction
+
+
+@dataclass
+class ComputeAccuracy:
+ r"""Compute reward accuracy and support `batch_eval_metrics`."""
+
+ def _dump(self) -> Optional[dict[str, float]]:
+ result = None
+ if hasattr(self, "score_dict"):
+ result = {k: float(np.mean(v)) for k, v in self.score_dict.items()}
+
+ self.score_dict = {"accuracy": []}
+ return result
+
+ def __post_init__(self):
+ self._dump()
+
+ def __call__(self, eval_preds: "EvalPrediction", compute_result: bool = True) -> Optional[dict[str, float]]:
+ chosen_scores, rejected_scores = numpify(eval_preds.predictions[0]), numpify(eval_preds.predictions[1])
+ if not chosen_scores.shape:
+ self.score_dict["accuracy"].append(chosen_scores > rejected_scores)
+ else:
+ for i in range(len(chosen_scores)):
+ self.score_dict["accuracy"].append(chosen_scores[i] > rejected_scores[i])
+
+ if compute_result:
+ return self._dump()
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/rm/trainer.py b/post-training/LLaMA-Factory/src/llamafactory/train/rm/trainer.py
new file mode 100644
index 0000000..8c14b0a
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/rm/trainer.py
@@ -0,0 +1,129 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's transformers library.
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/trainer.py
+#
+# 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 os
+from types import MethodType
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+from transformers import Trainer
+from typing_extensions import override
+
+from ...extras import logging
+from ...extras.packages import is_transformers_version_greater_than
+from ..callbacks import FixValueHeadModelCallback, SaveProcessorCallback
+from ..trainer_utils import create_custom_optimizer, create_custom_scheduler
+
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedModel, ProcessorMixin
+ from transformers.trainer import PredictionOutput
+
+ from ...hparams import FinetuningArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+class PairwiseTrainer(Trainer):
+ r"""Inherits Trainer to compute pairwise loss."""
+
+ def __init__(
+ self, finetuning_args: "FinetuningArguments", processor: Optional["ProcessorMixin"], **kwargs
+ ) -> None:
+ if is_transformers_version_greater_than("4.46"):
+ kwargs["processing_class"] = kwargs.pop("tokenizer")
+
+ super().__init__(**kwargs)
+ self.model_accepts_loss_kwargs = False # overwrite trainer's default behavior
+ self.finetuning_args = finetuning_args
+ self.can_return_loss = True # override property to return eval_loss
+ self.add_callback(FixValueHeadModelCallback)
+
+ if processor is not None:
+ self.add_callback(SaveProcessorCallback(processor))
+
+ if finetuning_args.use_badam:
+ from badam import BAdamCallback, clip_grad_norm_old_version # type: ignore
+
+ self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
+ self.add_callback(BAdamCallback)
+
+ @override
+ def create_optimizer(self) -> "torch.optim.Optimizer":
+ if self.optimizer is None:
+ self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
+ return super().create_optimizer()
+
+ @override
+ def create_scheduler(
+ self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
+ ) -> "torch.optim.lr_scheduler.LRScheduler":
+ create_custom_scheduler(self.args, num_training_steps, optimizer)
+ return super().create_scheduler(num_training_steps, optimizer)
+
+ @override
+ def _get_train_sampler(self) -> Optional["torch.utils.data.Sampler"]:
+ if self.finetuning_args.disable_shuffling:
+ return torch.utils.data.SequentialSampler(self.train_dataset)
+
+ return super()._get_train_sampler()
+
+ @override
+ def compute_loss(
+ self, model: "PreTrainedModel", inputs: dict[str, "torch.Tensor"], return_outputs: bool = False, **kwargs
+ ) -> Union["torch.Tensor", tuple["torch.Tensor", list["torch.Tensor"]]]:
+ r"""Compute pairwise loss. The first n examples are chosen and the last n examples are rejected.
+
+ Subclass and override to inject custom behavior.
+
+ Note that the first element will be removed from the output tuple.
+ See: https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/trainer.py#L3842
+ """
+ _, _, values = model(**inputs, output_hidden_states=True, return_dict=True, use_cache=False)
+ batch_size = inputs["input_ids"].size(0) // 2
+ chosen_masks, rejected_masks = torch.split(inputs["attention_mask"], batch_size, dim=0)
+ chosen_rewards, rejected_rewards = torch.split(values, batch_size, dim=0)
+ chosen_scores = chosen_rewards.gather(dim=-1, index=(chosen_masks.sum(dim=-1, keepdim=True) - 1))
+ rejected_scores = rejected_rewards.gather(dim=-1, index=(rejected_masks.sum(dim=-1, keepdim=True) - 1))
+ chosen_scores, rejected_scores = chosen_scores.squeeze(), rejected_scores.squeeze()
+
+ loss = -torch.nn.functional.logsigmoid(chosen_scores.float() - rejected_scores.float()).mean()
+ if return_outputs:
+ return loss, (loss, chosen_scores, rejected_scores)
+ else:
+ return loss
+
+ def save_predictions(self, predict_results: "PredictionOutput") -> None:
+ r"""Save model predictions to `output_dir`.
+
+ A custom behavior that not contained in Seq2SeqTrainer.
+ """
+ if not self.is_world_process_zero():
+ return
+
+ output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
+ logger.info_rank0(f"Saving prediction results to {output_prediction_file}")
+ chosen_scores, rejected_scores = predict_results.predictions
+
+ with open(output_prediction_file, "w", encoding="utf-8") as writer:
+ res: list[str] = []
+ for c_score, r_score in zip(chosen_scores, rejected_scores):
+ res.append(json.dumps({"chosen": round(float(c_score), 2), "rejected": round(float(r_score), 2)}))
+
+ writer.write("\n".join(res))
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/rm/workflow.py b/post-training/LLaMA-Factory/src/llamafactory/train/rm/workflow.py
new file mode 100644
index 0000000..18d562e
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/rm/workflow.py
@@ -0,0 +1,101 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's transformers library.
+# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/summarization/run_summarization.py
+#
+# 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.
+
+from typing import TYPE_CHECKING, Optional
+
+from ...data import PairwiseDataCollatorWithPadding, get_dataset, get_template_and_fix_tokenizer
+from ...extras.ploting import plot_loss
+from ...model import load_model, load_tokenizer
+from ..callbacks import fix_valuehead_checkpoint
+from ..trainer_utils import create_modelcard_and_push
+from .metric import ComputeAccuracy
+from .trainer import PairwiseTrainer
+
+
+if TYPE_CHECKING:
+ from transformers import Seq2SeqTrainingArguments, TrainerCallback
+
+ from ...hparams import DataArguments, FinetuningArguments, ModelArguments
+
+
+def run_rm(
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ callbacks: Optional[list["TrainerCallback"]] = None,
+):
+ tokenizer_module = load_tokenizer(model_args)
+ tokenizer = tokenizer_module["tokenizer"]
+ template = get_template_and_fix_tokenizer(tokenizer, data_args)
+ dataset_module = get_dataset(template, model_args, data_args, training_args, stage="rm", **tokenizer_module)
+ model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train, add_valuehead=True)
+ data_collator = PairwiseDataCollatorWithPadding(
+ template=template, model=model, pad_to_multiple_of=8, **tokenizer_module
+ )
+
+ # Update arguments
+ training_args.remove_unused_columns = False # important for multimodal and pairwise dataset
+
+ # Initialize our Trainer
+ trainer = PairwiseTrainer(
+ model=model,
+ args=training_args,
+ finetuning_args=finetuning_args,
+ data_collator=data_collator,
+ callbacks=callbacks,
+ compute_metrics=ComputeAccuracy(),
+ **dataset_module,
+ **tokenizer_module,
+ )
+
+ # Training
+ if training_args.do_train:
+ train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+ trainer.save_model()
+ if training_args.should_save:
+ fix_valuehead_checkpoint(model, training_args.output_dir, training_args.save_safetensors)
+
+ trainer.log_metrics("train", train_result.metrics)
+ trainer.save_metrics("train", train_result.metrics)
+ trainer.save_state()
+ if trainer.is_world_process_zero() and finetuning_args.plot_loss:
+ keys = ["loss"]
+ if isinstance(dataset_module.get("eval_dataset"), dict):
+ keys += sum(
+ [[f"eval_{key}_loss", f"eval_{key}_accuracy"] for key in dataset_module["eval_dataset"].keys()], []
+ )
+ else:
+ keys += ["eval_loss", "eval_accuracy"]
+
+ plot_loss(training_args.output_dir, keys=keys)
+
+ # Evaluation
+ if training_args.do_eval:
+ metrics = trainer.evaluate(metric_key_prefix="eval")
+ trainer.log_metrics("eval", metrics)
+ trainer.save_metrics("eval", metrics)
+
+ # Predict
+ if training_args.do_predict:
+ predict_results = trainer.predict(dataset_module["eval_dataset"], metric_key_prefix="predict")
+ trainer.log_metrics("predict", predict_results.metrics)
+ trainer.save_metrics("predict", predict_results.metrics)
+ trainer.save_predictions(predict_results)
+
+ # Create model card
+ create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/sft/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/train/sft/__init__.py
new file mode 100644
index 0000000..6107a9a
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/sft/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from .workflow import run_sft
+
+
+__all__ = ["run_sft"]
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/sft/metric.py b/post-training/LLaMA-Factory/src/llamafactory/train/sft/metric.py
new file mode 100644
index 0000000..f4f73ee
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/sft/metric.py
@@ -0,0 +1,134 @@
+# Copyright 2025 HuggingFace Inc., THUDM, and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's transformers library and the THUDM's ChatGLM implementation.
+# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/summarization/run_summarization.py
+# https://github.com/THUDM/ChatGLM-6B/blob/main/ptuning/main.py
+#
+# 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.
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import numpy as np
+import torch
+from transformers.utils import is_jieba_available, is_nltk_available
+
+from ...extras.constants import IGNORE_INDEX
+from ...extras.misc import numpify
+from ...extras.packages import is_rouge_available
+
+
+if TYPE_CHECKING:
+ from transformers import EvalPrediction, PreTrainedTokenizer
+
+
+if is_jieba_available():
+ import jieba # type: ignore
+
+
+if is_nltk_available():
+ from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu # type: ignore
+
+
+if is_rouge_available():
+ from rouge_chinese import Rouge # type: ignore
+
+
+def eval_logit_processor(logits: "torch.Tensor", labels: "torch.Tensor") -> "torch.Tensor":
+ r"""Compute the token with the largest likelihood to reduce memory footprint."""
+ if isinstance(logits, (list, tuple)):
+ if logits[0].dim() == 3: # (batch_size, seq_len, vocab_size)
+ logits = logits[0]
+ else: # moe models have aux loss
+ logits = logits[1]
+
+ if logits.dim() != 3:
+ raise ValueError("Cannot process the logits.")
+
+ return torch.argmax(logits, dim=-1)
+
+
+@dataclass
+class ComputeAccuracy:
+ r"""Compute accuracy and support `batch_eval_metrics`."""
+
+ def _dump(self) -> Optional[dict[str, float]]:
+ result = None
+ if hasattr(self, "score_dict"):
+ result = {k: float(np.mean(v)) for k, v in self.score_dict.items()}
+
+ self.score_dict = {"accuracy": []}
+ return result
+
+ def __post_init__(self):
+ self._dump()
+
+ def __call__(self, eval_preds: "EvalPrediction", compute_result: bool = True) -> Optional[dict[str, float]]:
+ preds, labels = numpify(eval_preds.predictions), numpify(eval_preds.label_ids)
+ for i in range(len(preds)):
+ pred, label = preds[i, :-1], labels[i, 1:]
+ label_mask = label != IGNORE_INDEX
+ self.score_dict["accuracy"].append(np.mean(pred[label_mask] == label[label_mask]))
+
+ if compute_result:
+ return self._dump()
+
+
+@dataclass
+class ComputeSimilarity:
+ r"""Compute text similarity scores and support `batch_eval_metrics`.
+
+ Wraps the tokenizer into metric functions, used in CustomSeq2SeqTrainer.
+ """
+
+ tokenizer: "PreTrainedTokenizer"
+
+ def _dump(self) -> Optional[dict[str, float]]:
+ result = None
+ if hasattr(self, "score_dict"):
+ result = {k: float(np.mean(v)) for k, v in self.score_dict.items()}
+
+ self.score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
+ return result
+
+ def __post_init__(self):
+ self._dump()
+
+ def __call__(self, eval_preds: "EvalPrediction", compute_result: bool = True) -> Optional[dict[str, float]]:
+ preds, labels = numpify(eval_preds.predictions), numpify(eval_preds.label_ids)
+
+ preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
+ labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)
+
+ decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True)
+ decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+ for pred, label in zip(decoded_preds, decoded_labels):
+ hypothesis = list(jieba.cut(pred))
+ reference = list(jieba.cut(label))
+
+ if len(" ".join(hypothesis).split()) == 0 or len(" ".join(reference).split()) == 0:
+ result = {"rouge-1": {"f": 0.0}, "rouge-2": {"f": 0.0}, "rouge-l": {"f": 0.0}}
+ else:
+ rouge = Rouge()
+ scores = rouge.get_scores(" ".join(hypothesis), " ".join(reference))
+ result = scores[0]
+
+ for k, v in result.items():
+ self.score_dict[k].append(round(v["f"] * 100, 4))
+
+ bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
+ self.score_dict["bleu-4"].append(round(bleu_score * 100, 4))
+
+ if compute_result:
+ return self._dump()
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/sft/trainer.py b/post-training/LLaMA-Factory/src/llamafactory/train/sft/trainer.py
new file mode 100644
index 0000000..fece1bd
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/sft/trainer.py
@@ -0,0 +1,165 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's transformers library.
+# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/trainer_seq2seq.py
+#
+# 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 os
+from types import MethodType
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import numpy as np
+import torch
+from transformers import Seq2SeqTrainer
+from typing_extensions import override
+
+from ...extras import logging
+from ...extras.constants import IGNORE_INDEX
+from ...extras.packages import is_transformers_version_greater_than
+from ..callbacks import SaveProcessorCallback
+from ..trainer_utils import create_custom_optimizer, create_custom_scheduler
+
+
+if TYPE_CHECKING:
+ from torch.utils.data import Dataset
+ from transformers import PreTrainedTokenizer, ProcessorMixin
+ from transformers.trainer import PredictionOutput
+
+ from ...hparams import FinetuningArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+class CustomSeq2SeqTrainer(Seq2SeqTrainer):
+ r"""Inherits Seq2SeqTrainer to compute generative metrics such as BLEU and ROUGE."""
+
+ def __init__(
+ self,
+ finetuning_args: "FinetuningArguments",
+ processor: Optional["ProcessorMixin"],
+ gen_kwargs: Optional[dict[str, Any]] = None,
+ **kwargs,
+ ) -> None:
+ if is_transformers_version_greater_than("4.46"):
+ kwargs["processing_class"] = kwargs.pop("tokenizer")
+ else:
+ self.processing_class: PreTrainedTokenizer = kwargs.get("tokenizer")
+
+ super().__init__(**kwargs)
+ if processor is not None:
+ # avoid wrong loss under gradient accumulation
+ # https://github.com/huggingface/transformers/pull/36044#issuecomment-2746657112
+ self.model_accepts_loss_kwargs = False
+
+ self.finetuning_args = finetuning_args
+ if gen_kwargs is not None:
+ # https://github.com/huggingface/transformers/blob/v4.45.0/src/transformers/trainer_seq2seq.py#L287
+ self._gen_kwargs = gen_kwargs
+
+ if processor is not None:
+ self.add_callback(SaveProcessorCallback(processor))
+
+ if finetuning_args.use_badam:
+ from badam import BAdamCallback, clip_grad_norm_old_version # type: ignore
+
+ self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
+ self.add_callback(BAdamCallback)
+
+ @override
+ def create_optimizer(self) -> "torch.optim.Optimizer":
+ if self.optimizer is None:
+ self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
+ return super().create_optimizer()
+
+ @override
+ def create_scheduler(
+ self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
+ ) -> "torch.optim.lr_scheduler.LRScheduler":
+ create_custom_scheduler(self.args, num_training_steps, optimizer)
+ return super().create_scheduler(num_training_steps, optimizer)
+
+ @override
+ def _get_train_sampler(self) -> Optional["torch.utils.data.Sampler"]:
+ if self.finetuning_args.disable_shuffling:
+ return torch.utils.data.SequentialSampler(self.train_dataset)
+
+ return super()._get_train_sampler()
+
+ @override
+ def compute_loss(self, model, inputs, *args, **kwargs):
+ return super().compute_loss(model, inputs, *args, **kwargs)
+
+ @override
+ def prediction_step(
+ self,
+ model: "torch.nn.Module",
+ inputs: dict[str, Union["torch.Tensor", Any]],
+ prediction_loss_only: bool,
+ ignore_keys: Optional[list[str]] = None,
+ **gen_kwargs,
+ ) -> tuple[Optional[float], Optional["torch.Tensor"], Optional["torch.Tensor"]]:
+ r"""Remove the prompt part in the generated tokens.
+
+ Subclass and override to inject custom behavior.
+ """
+ if self.args.predict_with_generate: # do not pass labels to model when generate
+ labels = inputs.pop("labels", None)
+ else:
+ labels = inputs.get("labels")
+
+ loss, generated_tokens, _ = super().prediction_step(
+ model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys, **gen_kwargs
+ )
+ if generated_tokens is not None and self.args.predict_with_generate:
+ generated_tokens[:, : inputs["input_ids"].size(-1)] = self.processing_class.pad_token_id
+ generated_tokens = generated_tokens.contiguous()
+
+ return loss, generated_tokens, labels
+
+ def save_predictions(
+ self, dataset: "Dataset", predict_results: "PredictionOutput", skip_special_tokens: bool = True
+ ) -> None:
+ r"""Save model predictions to `output_dir`.
+
+ A custom behavior that not contained in Seq2SeqTrainer.
+ """
+ if not self.is_world_process_zero():
+ return
+
+ output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
+ logger.info_rank0(f"Saving prediction results to {output_prediction_file}")
+
+ labels = np.where(
+ predict_results.label_ids != IGNORE_INDEX, predict_results.label_ids, self.processing_class.pad_token_id
+ )
+ preds = np.where(
+ predict_results.predictions != IGNORE_INDEX,
+ predict_results.predictions,
+ self.processing_class.pad_token_id,
+ )
+
+ for i in range(len(preds)):
+ pad_len = np.nonzero(preds[i] != self.processing_class.pad_token_id)[0]
+ if len(pad_len): # move pad token to last
+ preds[i] = np.concatenate((preds[i][pad_len[0] :], preds[i][: pad_len[0]]), axis=-1)
+
+ decoded_inputs = self.processing_class.batch_decode(dataset["input_ids"], skip_special_tokens=False)
+ decoded_preds = self.processing_class.batch_decode(preds, skip_special_tokens=skip_special_tokens)
+ decoded_labels = self.processing_class.batch_decode(labels, skip_special_tokens=skip_special_tokens)
+
+ with open(output_prediction_file, "w", encoding="utf-8") as f:
+ for text, pred, label in zip(decoded_inputs, decoded_preds, decoded_labels):
+ f.write(json.dumps({"prompt": text, "predict": pred, "label": label}, ensure_ascii=False) + "\n")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/sft/workflow.py b/post-training/LLaMA-Factory/src/llamafactory/train/sft/workflow.py
new file mode 100644
index 0000000..9e6c549
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/sft/workflow.py
@@ -0,0 +1,140 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's transformers library.
+# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/summarization/run_summarization.py
+#
+# 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.
+
+from typing import TYPE_CHECKING, Optional
+
+from ...data import SFTDataCollatorWith4DAttentionMask, get_dataset, get_template_and_fix_tokenizer
+from ...extras.constants import IGNORE_INDEX
+from ...extras.logging import get_logger
+from ...extras.misc import calculate_tps
+from ...extras.ploting import plot_loss
+from ...model import load_model, load_tokenizer
+from ..trainer_utils import create_modelcard_and_push
+from .metric import ComputeAccuracy, ComputeSimilarity, eval_logit_processor
+from .trainer import CustomSeq2SeqTrainer
+
+
+if TYPE_CHECKING:
+ from transformers import Seq2SeqTrainingArguments, TrainerCallback
+
+ from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments
+
+
+logger = get_logger(__name__)
+
+
+def run_sft(
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ training_args: "Seq2SeqTrainingArguments",
+ finetuning_args: "FinetuningArguments",
+ generating_args: "GeneratingArguments",
+ callbacks: Optional[list["TrainerCallback"]] = None,
+):
+ tokenizer_module = load_tokenizer(model_args)
+ tokenizer = tokenizer_module["tokenizer"]
+ template = get_template_and_fix_tokenizer(tokenizer, data_args)
+ dataset_module = get_dataset(template, model_args, data_args, training_args, stage="sft", **tokenizer_module)
+ model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
+
+ if getattr(model, "is_quantized", False) and not training_args.do_train:
+ setattr(model, "_hf_peft_config_loaded", True) # hack here: make model compatible with prediction
+
+ data_collator = SFTDataCollatorWith4DAttentionMask(
+ template=template,
+ model=model if not training_args.predict_with_generate else None,
+ pad_to_multiple_of=8 if training_args.do_train else None, # for shift short attention
+ label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id,
+ block_diag_attn=model_args.block_diag_attn,
+ attn_implementation=getattr(model.config, "_attn_implementation", None),
+ compute_dtype=model_args.compute_dtype,
+ **tokenizer_module,
+ )
+
+ # Override the decoding parameters of Seq2SeqTrainer
+ training_args.generation_max_length = training_args.generation_max_length or data_args.cutoff_len
+ training_args.generation_num_beams = data_args.eval_num_beams or training_args.generation_num_beams
+ training_args.remove_unused_columns = False # important for multimodal dataset
+
+ # Metric utils
+ metric_module = {}
+ if training_args.predict_with_generate:
+ metric_module["compute_metrics"] = ComputeSimilarity(tokenizer=tokenizer)
+ elif finetuning_args.compute_accuracy:
+ metric_module["compute_metrics"] = ComputeAccuracy()
+ metric_module["preprocess_logits_for_metrics"] = eval_logit_processor
+
+ # Keyword arguments for `model.generate`
+ gen_kwargs = generating_args.to_dict(obey_generation_config=True)
+ gen_kwargs["eos_token_id"] = [tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids
+ gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
+
+ # Initialize our Trainer
+ trainer = CustomSeq2SeqTrainer(
+ model=model,
+ args=training_args,
+ finetuning_args=finetuning_args,
+ data_collator=data_collator,
+ callbacks=callbacks,
+ gen_kwargs=gen_kwargs,
+ **dataset_module,
+ **tokenizer_module,
+ **metric_module,
+ )
+
+ # Training
+ if training_args.do_train:
+ train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+ trainer.save_model()
+ if finetuning_args.include_effective_tokens_per_second:
+ train_result.metrics["effective_tokens_per_sec"] = calculate_tps(
+ dataset_module["train_dataset"], train_result.metrics, stage="sft"
+ )
+
+ trainer.log_metrics("train", train_result.metrics)
+ trainer.save_metrics("train", train_result.metrics)
+ trainer.save_state()
+ if trainer.is_world_process_zero() and finetuning_args.plot_loss:
+ keys = ["loss"]
+ if isinstance(dataset_module.get("eval_dataset"), dict):
+ keys += sum(
+ [[f"eval_{key}_loss", f"eval_{key}_accuracy"] for key in dataset_module["eval_dataset"].keys()], []
+ )
+ else:
+ keys += ["eval_loss", "eval_accuracy"]
+
+ plot_loss(training_args.output_dir, keys=keys)
+
+ if training_args.predict_with_generate:
+ tokenizer.padding_side = "left" # use left-padding in generation
+
+ # Evaluation
+ if training_args.do_eval:
+ metrics = trainer.evaluate(metric_key_prefix="eval", **gen_kwargs)
+ trainer.log_metrics("eval", metrics)
+ trainer.save_metrics("eval", metrics)
+
+ # Predict
+ if training_args.do_predict:
+ logger.warning_rank0_once("Batch generation can be very slow. Consider using `scripts/vllm_infer.py` instead.")
+ predict_results = trainer.predict(dataset_module["eval_dataset"], metric_key_prefix="predict", **gen_kwargs)
+ trainer.log_metrics("predict", predict_results.metrics)
+ trainer.save_metrics("predict", predict_results.metrics)
+ trainer.save_predictions(dataset_module["eval_dataset"], predict_results, generating_args.skip_special_tokens)
+
+ # Create model card
+ create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/test_utils.py b/post-training/LLaMA-Factory/src/llamafactory/train/test_utils.py
new file mode 100644
index 0000000..6e4c4ff
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/test_utils.py
@@ -0,0 +1,119 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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.
+
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM
+from trl import AutoModelForCausalLMWithValueHead
+
+from ..data import get_dataset, get_template_and_fix_tokenizer
+from ..extras.misc import get_current_device
+from ..hparams import get_infer_args, get_train_args
+from ..model import load_model, load_tokenizer
+
+
+if TYPE_CHECKING:
+ from peft import LoraModel
+ from transformers import PreTrainedModel
+
+ from ..data.data_utils import DatasetModule
+
+
+def compare_model(model_a: "torch.nn.Module", model_b: "torch.nn.Module", diff_keys: list[str] = []) -> None:
+ state_dict_a = model_a.state_dict()
+ state_dict_b = model_b.state_dict()
+ assert set(state_dict_a.keys()) == set(state_dict_b.keys())
+ for name in state_dict_a.keys():
+ if any(key in name for key in diff_keys):
+ assert torch.allclose(state_dict_a[name], state_dict_b[name], rtol=1e-4, atol=1e-5) is False
+ else:
+ assert torch.allclose(state_dict_a[name], state_dict_b[name], rtol=1e-4, atol=1e-5) is True
+
+
+def check_lora_model(model: "LoraModel") -> tuple[set[str], set[str]]:
+ linear_modules, extra_modules = set(), set()
+ for name, param in model.named_parameters():
+ if any(module in name for module in ["lora_A", "lora_B"]):
+ linear_modules.add(name.split(".lora_", maxsplit=1)[0].split(".")[-1])
+ assert param.requires_grad is True
+ assert param.dtype == torch.float32
+ elif "modules_to_save" in name:
+ extra_modules.add(name.split(".modules_to_save", maxsplit=1)[0].split(".")[-1])
+ assert param.requires_grad is True
+ assert param.dtype == torch.float32
+ else:
+ assert param.requires_grad is False
+ assert param.dtype == torch.float16
+
+ return linear_modules, extra_modules
+
+
+def load_train_model(add_valuehead: bool = False, **kwargs) -> "PreTrainedModel":
+ model_args, _, _, finetuning_args, _ = get_train_args(kwargs)
+ tokenizer = load_tokenizer(model_args)["tokenizer"]
+ return load_model(tokenizer, model_args, finetuning_args, is_trainable=True, add_valuehead=add_valuehead)
+
+
+def load_infer_model(add_valuehead: bool = False, **kwargs) -> "PreTrainedModel":
+ model_args, _, finetuning_args, _ = get_infer_args(kwargs)
+ tokenizer = load_tokenizer(model_args)["tokenizer"]
+ return load_model(tokenizer, model_args, finetuning_args, is_trainable=False, add_valuehead=add_valuehead)
+
+
+def load_reference_model(
+ model_path: str,
+ lora_path: Optional[str] = None,
+ use_lora: bool = False,
+ use_pissa: bool = False,
+ is_trainable: bool = False,
+ add_valuehead: bool = False,
+) -> Union["PreTrainedModel", "LoraModel"]:
+ current_device = get_current_device()
+ if add_valuehead:
+ model: AutoModelForCausalLMWithValueHead = AutoModelForCausalLMWithValueHead.from_pretrained(
+ model_path, torch_dtype=torch.float16, device_map=current_device
+ )
+ if not is_trainable:
+ model.v_head = model.v_head.to(torch.float16)
+
+ return model
+
+ model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, device_map=current_device)
+ if use_lora or use_pissa:
+ model = PeftModel.from_pretrained(
+ model, lora_path, subfolder="pissa_init" if use_pissa else None, is_trainable=is_trainable
+ )
+ for param in filter(lambda p: p.requires_grad, model.parameters()):
+ param.data = param.data.to(torch.float32)
+
+ return model
+
+
+def load_dataset_module(**kwargs) -> "DatasetModule":
+ model_args, data_args, training_args, _, _ = get_train_args(kwargs)
+ tokenizer_module = load_tokenizer(model_args)
+ template = get_template_and_fix_tokenizer(tokenizer_module["tokenizer"], data_args)
+ dataset_module = get_dataset(template, model_args, data_args, training_args, kwargs["stage"], **tokenizer_module)
+ return dataset_module
+
+
+def patch_valuehead_model() -> None:
+ def post_init(self: "AutoModelForCausalLMWithValueHead", state_dict: dict[str, "torch.Tensor"]) -> None:
+ state_dict = {k[7:]: state_dict[k] for k in state_dict.keys() if k.startswith("v_head.")}
+ self.v_head.load_state_dict(state_dict, strict=False)
+ del state_dict
+
+ AutoModelForCausalLMWithValueHead.post_init = post_init
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/trainer_utils.py b/post-training/LLaMA-Factory/src/llamafactory/train/trainer_utils.py
new file mode 100644
index 0000000..89459a8
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/trainer_utils.py
@@ -0,0 +1,665 @@
+# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the original GaLore's implementation: https://github.com/jiaweizzhao/GaLore
+# and the original LoRA+'s implementation: https://github.com/nikhil-ghosh-berkeley/loraplus
+# and the original BAdam's implementation: https://github.com/Ledzy/BAdam
+# and the HuggingFace's TRL library: https://github.com/huggingface/trl
+#
+# 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 os
+from collections.abc import Mapping
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+
+import torch
+from transformers import Trainer
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.modeling_utils import is_fsdp_enabled
+from transformers.optimization import get_scheduler
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.trainer_pt_utils import get_parameter_names
+from typing_extensions import override
+
+from ..extras import logging
+from ..extras.constants import IGNORE_INDEX, SWANLAB_CONFIG
+from ..extras.packages import is_apollo_available, is_galore_available, is_ray_available
+from ..hparams import FinetuningArguments, ModelArguments
+from ..model import find_all_linear_modules, load_model, load_tokenizer, load_valuehead_params
+
+
+if is_galore_available():
+ from galore_torch import GaLoreAdafactor, GaLoreAdamW, GaLoreAdamW8bit # type: ignore
+
+
+if is_apollo_available():
+ from apollo_torch import APOLLOAdamW # type: ignore
+
+
+if is_ray_available():
+ import ray
+ from ray.train import RunConfig, ScalingConfig
+ from ray.train.torch import TorchTrainer
+
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedModel, TrainerCallback, TrainerState
+ from trl import AutoModelForCausalLMWithValueHead
+
+ from ..hparams import DataArguments, RayArguments, TrainingArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+class DummyOptimizer(torch.optim.Optimizer):
+ r"""A dummy optimizer used for the GaLore or APOLLO algorithm."""
+
+ def __init__(
+ self, lr: float = 1e-3, optimizer_dict: Optional[dict["torch.nn.Parameter", "torch.optim.Optimizer"]] = None
+ ) -> None:
+ dummy_tensor = torch.randn(1, 1)
+ self.optimizer_dict = optimizer_dict
+ super().__init__([dummy_tensor], {"lr": lr})
+
+ @override
+ def zero_grad(self, set_to_none: bool = True) -> None:
+ pass
+
+ @override
+ def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
+ pass
+
+
+def create_modelcard_and_push(
+ trainer: "Trainer",
+ model_args: "ModelArguments",
+ data_args: "DataArguments",
+ training_args: "TrainingArguments",
+ finetuning_args: "FinetuningArguments",
+) -> None:
+ kwargs = {
+ "tasks": "text-generation",
+ "finetuned_from": model_args.model_name_or_path,
+ "tags": ["llama-factory", finetuning_args.finetuning_type],
+ }
+ if data_args.dataset is not None:
+ kwargs["dataset"] = data_args.dataset
+
+ if model_args.use_unsloth:
+ kwargs["tags"] = kwargs["tags"] + ["unsloth"]
+
+ if not training_args.do_train:
+ pass
+ elif training_args.push_to_hub:
+ trainer.push_to_hub(**kwargs)
+ else:
+ trainer.create_model_card(license="other", **kwargs) # prevent from connecting to hub
+
+
+def create_ref_model(
+ model_args: "ModelArguments", finetuning_args: "FinetuningArguments", add_valuehead: bool = False
+) -> Optional[Union["PreTrainedModel", "AutoModelForCausalLMWithValueHead"]]:
+ r"""Create reference model for PPO/DPO training. Evaluation mode is not supported.
+
+ The valuehead parameter is randomly initialized since it is useless for PPO training.
+ """
+ if finetuning_args.ref_model is not None:
+ ref_model_args = ModelArguments.copyfrom(
+ model_args,
+ model_name_or_path=finetuning_args.ref_model,
+ adapter_name_or_path=finetuning_args.ref_model_adapters,
+ quantization_bit=finetuning_args.ref_model_quantization_bit,
+ )
+ ref_finetuning_args = FinetuningArguments()
+ tokenizer = load_tokenizer(ref_model_args)["tokenizer"]
+ ref_model = load_model(
+ tokenizer, ref_model_args, ref_finetuning_args, is_trainable=False, add_valuehead=add_valuehead
+ )
+ logger.info_rank0(f"Created reference model from {finetuning_args.ref_model}")
+ else:
+ if finetuning_args.finetuning_type == "lora":
+ ref_model = None
+ else:
+ ref_model_args = ModelArguments.copyfrom(model_args)
+ ref_finetuning_args = FinetuningArguments()
+ tokenizer = load_tokenizer(ref_model_args)["tokenizer"]
+ ref_model = load_model(
+ tokenizer, ref_model_args, ref_finetuning_args, is_trainable=False, add_valuehead=add_valuehead
+ )
+ logger.info_rank0("Created reference model from the model itself.")
+
+ return ref_model
+
+
+def create_reward_model(
+ model: "AutoModelForCausalLMWithValueHead", model_args: "ModelArguments", finetuning_args: "FinetuningArguments"
+) -> Optional["AutoModelForCausalLMWithValueHead"]:
+ r"""Create reward model for PPO training."""
+ if finetuning_args.reward_model_type == "api":
+ assert finetuning_args.reward_model.startswith("http"), "Please provide full url."
+ logger.info_rank0(f"Use reward server {finetuning_args.reward_model}")
+ return finetuning_args.reward_model
+ elif finetuning_args.reward_model_type == "lora":
+ model.pretrained_model.load_adapter(finetuning_args.reward_model, "reward")
+ for name, param in model.named_parameters(): # https://github.com/huggingface/peft/issues/1090
+ if "default" in name:
+ param.data = param.data.to(torch.float32) # trainable params should in fp32
+ vhead_params = load_valuehead_params(finetuning_args.reward_model, model_args)
+ assert vhead_params is not None, "Reward model is not correctly loaded."
+ model.register_buffer("reward_head_weight", vhead_params["v_head.summary.weight"], persistent=False)
+ model.register_buffer("reward_head_bias", vhead_params["v_head.summary.bias"], persistent=False)
+ model.register_buffer(
+ "default_head_weight", torch.zeros_like(vhead_params["v_head.summary.weight"]), persistent=False
+ )
+ model.register_buffer(
+ "default_head_bias", torch.zeros_like(vhead_params["v_head.summary.bias"]), persistent=False
+ )
+ logger.info_rank0(f"Loaded adapter weights of reward model from {finetuning_args.reward_model}")
+ return None
+ else:
+ reward_model_args = ModelArguments.copyfrom(
+ model_args,
+ model_name_or_path=finetuning_args.reward_model,
+ adapter_name_or_path=finetuning_args.reward_model_adapters,
+ quantization_bit=finetuning_args.reward_model_quantization_bit,
+ )
+ reward_finetuning_args = FinetuningArguments()
+ tokenizer = load_tokenizer(reward_model_args)["tokenizer"]
+ reward_model = load_model(
+ tokenizer, reward_model_args, reward_finetuning_args, is_trainable=False, add_valuehead=True
+ )
+ logger.info_rank0(f"Loaded full weights of reward model from {finetuning_args.reward_model}")
+ logger.warning_rank0("Please ensure the ppo model and reward model share SAME tokenizer and vocabulary.")
+ return reward_model
+
+
+def _get_decay_parameter_names(model: "PreTrainedModel") -> list[str]:
+ r"""Return a list of names of parameters with weight decay. (weights in non-layernorm layers)."""
+ decay_parameters = get_parameter_names(model, ALL_LAYERNORM_LAYERS)
+ decay_parameters = [name for name in decay_parameters if "bias" not in name]
+ return decay_parameters
+
+
+def _create_galore_optimizer(
+ model: "PreTrainedModel",
+ training_args: "TrainingArguments",
+ finetuning_args: "FinetuningArguments",
+) -> "torch.optim.Optimizer":
+ if len(finetuning_args.galore_target) == 1 and finetuning_args.galore_target[0] == "all":
+ galore_targets = find_all_linear_modules(model, finetuning_args.freeze_vision_tower)
+ else:
+ galore_targets = finetuning_args.galore_target
+
+ galore_params: list[torch.nn.Parameter] = []
+ for name, module in model.named_modules():
+ if isinstance(module, torch.nn.Linear) and any(target in name for target in galore_targets):
+ for param in module.parameters():
+ if param.requires_grad and len(param.shape) > 1:
+ galore_params.append(param)
+
+ galore_kwargs = {
+ "rank": finetuning_args.galore_rank,
+ "update_proj_gap": finetuning_args.galore_update_interval,
+ "scale": finetuning_args.galore_scale,
+ "proj_type": finetuning_args.galore_proj_type,
+ }
+
+ id_galore_params = {id(param) for param in galore_params}
+ decay_params, nodecay_params = [], [] # they are non-galore parameters
+ trainable_params: list[torch.nn.Parameter] = [] # galore_params + decay_params + nodecay_params
+ decay_param_names = _get_decay_parameter_names(model)
+ for name, param in model.named_parameters():
+ if param.requires_grad:
+ trainable_params.append(param)
+ if id(param) not in id_galore_params:
+ if name in decay_param_names:
+ decay_params.append(param)
+ else:
+ nodecay_params.append(param)
+
+ _, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
+
+ if training_args.optim == "adamw_torch":
+ optim_class = GaLoreAdamW
+ elif training_args.optim in ["adamw_bnb_8bit", "adamw_8bit", "paged_adamw_8bit"]:
+ optim_class = GaLoreAdamW8bit
+ elif training_args.optim == "adafactor":
+ optim_class = GaLoreAdafactor
+ else:
+ raise NotImplementedError(f"Unknown optim: {training_args.optim}.")
+
+ if finetuning_args.galore_layerwise:
+ logger.warning_rank0("The displayed gradient norm will be all zeros in layerwise GaLore.")
+ if training_args.gradient_accumulation_steps != 1:
+ raise ValueError("Per-layer GaLore does not support gradient accumulation.")
+
+ optimizer_dict: dict[torch.Tensor, torch.optim.Optimizer] = {}
+ for param in nodecay_params:
+ param_groups = [dict(params=[param], weight_decay=0.0)]
+ optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
+ for param in decay_params:
+ param_groups = [dict(params=[param], weight_decay=training_args.weight_decay)]
+ optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
+ for param in galore_params: # galore params have weight decay
+ param_groups = [dict(params=[param], weight_decay=training_args.weight_decay, **galore_kwargs)]
+ optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
+
+ def optimizer_hook(param: "torch.nn.Parameter"):
+ if param.grad is not None:
+ optimizer_dict[param].step()
+ optimizer_dict[param].zero_grad()
+
+ for param in trainable_params:
+ param.register_post_accumulate_grad_hook(optimizer_hook)
+
+ optimizer = DummyOptimizer(lr=training_args.learning_rate, optimizer_dict=optimizer_dict)
+ else:
+ param_groups = [
+ dict(params=nodecay_params, weight_decay=0.0),
+ dict(params=decay_params, weight_decay=training_args.weight_decay),
+ dict(params=galore_params, weight_decay=training_args.weight_decay, **galore_kwargs),
+ ]
+ optimizer = optim_class(param_groups, **optim_kwargs)
+
+ logger.info_rank0(
+ f"Using GaLore optimizer with args: {galore_kwargs}. "
+ "It may cause hanging at the start of training, wait patiently."
+ )
+ return optimizer
+
+
+def _create_apollo_optimizer(
+ model: "PreTrainedModel",
+ training_args: "TrainingArguments",
+ finetuning_args: "FinetuningArguments",
+) -> "torch.optim.Optimizer":
+ if len(finetuning_args.apollo_target) == 1 and finetuning_args.apollo_target[0] == "all":
+ apollo_targets = find_all_linear_modules(model, finetuning_args.freeze_vision_tower)
+ else:
+ apollo_targets = finetuning_args.apollo_target
+
+ apollo_params: list[torch.nn.Parameter] = []
+ for name, module in model.named_modules():
+ if isinstance(module, torch.nn.Linear) and any(target in name for target in apollo_targets):
+ for param in module.parameters():
+ if param.requires_grad and len(param.shape) > 1:
+ apollo_params.append(param)
+
+ apollo_kwargs = {
+ "rank": finetuning_args.apollo_rank,
+ "proj": finetuning_args.apollo_proj,
+ "proj_type": finetuning_args.apollo_proj_type,
+ "update_proj_gap": finetuning_args.apollo_update_interval,
+ "scale": finetuning_args.apollo_scale,
+ "scale_type": finetuning_args.apollo_scale_type,
+ "scale_front": finetuning_args.apollo_scale_front,
+ }
+
+ id_apollo_params = {id(param) for param in apollo_params}
+ decay_params, nodecay_params = [], [] # they are non-apollo parameters
+ trainable_params: list[torch.nn.Parameter] = [] # apollo_params + decay_params + nodecay_params
+ decay_param_names = _get_decay_parameter_names(model)
+ for name, param in model.named_parameters():
+ if param.requires_grad:
+ trainable_params.append(param)
+ if id(param) not in id_apollo_params:
+ if name in decay_param_names:
+ decay_params.append(param)
+ else:
+ nodecay_params.append(param)
+
+ _, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
+
+ if training_args.optim == "adamw_torch":
+ optim_class = APOLLOAdamW
+ else:
+ raise NotImplementedError(f"Unknown optim: {training_args.optim}.")
+
+ if finetuning_args.apollo_layerwise:
+ logger.warning_rank0("The displayed gradient norm will be all zeros in layerwise APOLLO.")
+ if training_args.gradient_accumulation_steps != 1:
+ raise ValueError("Per-layer APOLLO does not support gradient accumulation.")
+
+ optimizer_dict: dict[torch.Tensor, torch.optim.Optimizer] = {}
+ for param in nodecay_params:
+ param_groups = [dict(params=[param], weight_decay=0.0)]
+ optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
+ for param in decay_params:
+ param_groups = [dict(params=[param], weight_decay=training_args.weight_decay)]
+ optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
+ for param in apollo_params: # apollo params have weight decay
+ param_groups = [dict(params=[param], weight_decay=training_args.weight_decay, **apollo_kwargs)]
+ optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
+
+ def optimizer_hook(param: "torch.nn.Parameter"):
+ if param.grad is not None:
+ optimizer_dict[param].step()
+ optimizer_dict[param].zero_grad()
+
+ for param in trainable_params:
+ param.register_post_accumulate_grad_hook(optimizer_hook)
+
+ optimizer = DummyOptimizer(lr=training_args.learning_rate, optimizer_dict=optimizer_dict)
+ else:
+ param_groups = [
+ dict(params=nodecay_params, weight_decay=0.0),
+ dict(params=decay_params, weight_decay=training_args.weight_decay),
+ dict(params=apollo_params, weight_decay=training_args.weight_decay, **apollo_kwargs),
+ ]
+ optimizer = optim_class(param_groups, **optim_kwargs)
+
+ logger.info_rank0(f"Using APOLLO optimizer with args: {apollo_kwargs}.")
+ return optimizer
+
+
+def _create_loraplus_optimizer(
+ model: "PreTrainedModel",
+ training_args: "TrainingArguments",
+ finetuning_args: "FinetuningArguments",
+) -> "torch.optim.Optimizer":
+ default_lr = training_args.learning_rate
+ loraplus_lr = training_args.learning_rate * finetuning_args.loraplus_lr_ratio
+ embedding_lr = finetuning_args.loraplus_lr_embedding
+
+ decay_param_names = _get_decay_parameter_names(model)
+ param_dict: dict[str, list[torch.nn.Parameter]] = {
+ "lora_a": [],
+ "lora_b": [],
+ "lora_b_nodecay": [],
+ "embedding": [],
+ }
+ for name, param in model.named_parameters():
+ if param.requires_grad:
+ if "lora_embedding_B" in name:
+ param_dict["embedding"].append(param)
+ elif "lora_B" in name or param.ndim == 1:
+ if name in decay_param_names:
+ param_dict["lora_b"].append(param)
+ else:
+ param_dict["lora_b_nodecay"].append(param)
+ else:
+ param_dict["lora_a"].append(param)
+
+ optim_class, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
+ param_groups = [
+ dict(params=param_dict["lora_a"], lr=default_lr, weight_decay=training_args.weight_decay),
+ dict(params=param_dict["lora_b"], lr=loraplus_lr, weight_decay=training_args.weight_decay),
+ dict(params=param_dict["lora_b_nodecay"], lr=loraplus_lr, weight_decay=0.0),
+ dict(params=param_dict["embedding"], lr=embedding_lr, weight_decay=training_args.weight_decay),
+ ]
+ optimizer = optim_class(param_groups, **optim_kwargs)
+ logger.info_rank0(f"Using LoRA+ optimizer with loraplus lr ratio {finetuning_args.loraplus_lr_ratio:.2f}.")
+ return optimizer
+
+
+def _create_badam_optimizer(
+ model: "PreTrainedModel",
+ training_args: "TrainingArguments",
+ finetuning_args: "FinetuningArguments",
+) -> "torch.optim.Optimizer":
+ decay_params, nodecay_params = [], []
+ decay_param_names = _get_decay_parameter_names(model)
+ for name, param in model.named_parameters():
+ if param.requires_grad:
+ if name in decay_param_names:
+ decay_params.append(param)
+ else:
+ nodecay_params.append(param)
+
+ optim_class, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
+ param_groups = [
+ dict(params=nodecay_params, weight_decay=0.0),
+ dict(params=decay_params, weight_decay=training_args.weight_decay),
+ ]
+
+ if finetuning_args.badam_mode == "layer":
+ from badam import BlockOptimizer # type: ignore
+
+ base_optimizer = optim_class(param_groups, **optim_kwargs)
+ optimizer = BlockOptimizer(
+ base_optimizer=base_optimizer,
+ named_parameters_list=list(model.named_parameters()),
+ block_prefix_list=None,
+ switch_block_every=finetuning_args.badam_switch_interval,
+ start_block=finetuning_args.badam_start_block,
+ switch_mode=finetuning_args.badam_switch_mode,
+ verbose=finetuning_args.badam_verbose,
+ ds_zero3_enabled=is_deepspeed_zero3_enabled(),
+ )
+ logger.info_rank0(
+ f"Using BAdam optimizer with layer-wise update, switch mode is {finetuning_args.badam_switch_mode}, "
+ f"switch block every {finetuning_args.badam_switch_interval} steps, "
+ f"default start block is {finetuning_args.badam_start_block}"
+ )
+
+ elif finetuning_args.badam_mode == "ratio":
+ from badam import BlockOptimizerRatio # type: ignore
+
+ assert finetuning_args.badam_update_ratio > 1e-6
+ optimizer = BlockOptimizerRatio(
+ param_groups=param_groups,
+ named_parameters_list=list(model.named_parameters()),
+ update_ratio=finetuning_args.badam_update_ratio,
+ mask_mode=finetuning_args.badam_mask_mode,
+ verbose=finetuning_args.badam_verbose,
+ include_embedding=False,
+ **optim_kwargs,
+ )
+ logger.info_rank0(
+ f"Using BAdam optimizer with ratio-based update, update ratio is {finetuning_args.badam_update_ratio}, "
+ f"mask mode is {finetuning_args.badam_mask_mode}"
+ )
+
+ return optimizer
+
+
+def _create_adam_mini_optimizer(
+ model: "PreTrainedModel",
+ training_args: "TrainingArguments",
+) -> "torch.optim.Optimizer":
+ from adam_mini import Adam_mini # type: ignore
+
+ hidden_size = getattr(model.config, "hidden_size", None)
+ num_q_head = getattr(model.config, "num_attention_heads", None)
+ num_kv_head = getattr(model.config, "num_key_value_heads", None)
+
+ optimizer = Adam_mini(
+ named_parameters=model.named_parameters(),
+ lr=training_args.learning_rate,
+ betas=(training_args.adam_beta1, training_args.adam_beta2),
+ eps=training_args.adam_epsilon,
+ weight_decay=training_args.weight_decay,
+ model_sharding=is_fsdp_enabled() or is_deepspeed_zero3_enabled(),
+ dim=hidden_size,
+ n_heads=num_q_head,
+ n_kv_heads=num_kv_head,
+ )
+ logger.info_rank0("Using Adam-mini optimizer.")
+ return optimizer
+
+
+def create_custom_optimizer(
+ model: "PreTrainedModel",
+ training_args: "TrainingArguments",
+ finetuning_args: "FinetuningArguments",
+) -> Optional["torch.optim.Optimizer"]:
+ if finetuning_args.use_galore:
+ return _create_galore_optimizer(model, training_args, finetuning_args)
+
+ if finetuning_args.use_apollo:
+ return _create_apollo_optimizer(model, training_args, finetuning_args)
+
+ if finetuning_args.loraplus_lr_ratio is not None:
+ return _create_loraplus_optimizer(model, training_args, finetuning_args)
+
+ if finetuning_args.use_badam:
+ return _create_badam_optimizer(model, training_args, finetuning_args)
+
+ if finetuning_args.use_adam_mini:
+ return _create_adam_mini_optimizer(model, training_args)
+
+
+def create_custom_scheduler(
+ training_args: "TrainingArguments",
+ num_training_steps: int,
+ optimizer: Optional["torch.optim.Optimizer"] = None,
+) -> None:
+ if training_args.lr_scheduler_type == "warmup_stable_decay":
+ num_warmup_steps = training_args.get_warmup_steps(num_training_steps)
+ remaining_steps = num_training_steps - num_warmup_steps
+ num_stable_steps = remaining_steps // 3 # use 1/3 for stable by default
+ num_decay_steps = remaining_steps - num_stable_steps
+ scheduler_kwargs = training_args.lr_scheduler_kwargs or {}
+ default_kwargs = {
+ "num_stable_steps": num_stable_steps,
+ "num_decay_steps": num_decay_steps,
+ }
+ for key, value in default_kwargs.items():
+ if key not in scheduler_kwargs:
+ scheduler_kwargs[key] = value
+
+ training_args.lr_scheduler_kwargs = scheduler_kwargs
+
+ if optimizer is not None and isinstance(optimizer, DummyOptimizer):
+ optimizer_dict = optimizer.optimizer_dict
+ scheduler_dict: dict[torch.nn.Parameter, torch.optim.lr_scheduler.LRScheduler] = {}
+
+ for param in optimizer_dict.keys():
+ scheduler_dict[param] = get_scheduler(
+ training_args.lr_scheduler_type,
+ optimizer=optimizer_dict[param],
+ num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
+ num_training_steps=num_training_steps,
+ scheduler_specific_kwargs=training_args.lr_scheduler_kwargs,
+ )
+
+ def scheduler_hook(param: "torch.nn.Parameter"):
+ scheduler_dict[param].step()
+
+ for param in optimizer_dict.keys():
+ param.register_post_accumulate_grad_hook(scheduler_hook)
+
+
+def get_batch_logps(
+ logits: "torch.Tensor", labels: "torch.Tensor", label_pad_token_id: int = IGNORE_INDEX
+) -> tuple["torch.Tensor", "torch.Tensor"]:
+ r"""Compute the log probabilities of the given labels under the given logits.
+
+ Returns:
+ logps: A tensor of shape (batch_size,) containing the sum of log probabilities.
+ valid_length: A tensor of shape (batch_size,) containing the number of non-masked tokens.
+
+ """
+ if logits.shape[:-1] != labels.shape:
+ raise ValueError("Logits (batchsize x seqlen) and labels must have the same shape.")
+
+ labels = labels[:, 1:].clone()
+ logits = logits[:, :-1, :]
+ loss_mask = labels != label_pad_token_id
+ labels[labels == label_pad_token_id] = 0 # dummy token
+ per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
+ return (per_token_logps * loss_mask).sum(-1), loss_mask.sum(-1)
+
+
+def nested_detach(
+ tensors: Union["torch.Tensor", list["torch.Tensor"], tuple["torch.Tensor"], dict[str, "torch.Tensor"]],
+ clone: bool = False,
+):
+ r"""Detach `tensors` (even if it's a nested list/tuple/dict of tensors)."""
+ if isinstance(tensors, (list, tuple)):
+ return type(tensors)(nested_detach(t, clone=clone) for t in tensors)
+ elif isinstance(tensors, Mapping):
+ return type(tensors)({k: nested_detach(t, clone=clone) for k, t in tensors.items()})
+
+ if isinstance(tensors, torch.Tensor):
+ if clone:
+ return tensors.detach().clone()
+ else:
+ return tensors.detach()
+ else:
+ return tensors
+
+
+def get_swanlab_callback(finetuning_args: "FinetuningArguments") -> "TrainerCallback":
+ r"""Get the callback for logging to SwanLab."""
+ import swanlab # type: ignore
+ from swanlab.integration.transformers import SwanLabCallback # type: ignore
+
+ if finetuning_args.swanlab_api_key is not None:
+ swanlab.login(api_key=finetuning_args.swanlab_api_key)
+
+ if finetuning_args.swanlab_lark_webhook_url is not None:
+ from swanlab.plugin.notification import LarkCallback # type: ignore
+
+ lark_callback = LarkCallback(
+ webhook_url=finetuning_args.swanlab_lark_webhook_url,
+ secret=finetuning_args.swanlab_lark_secret,
+ )
+ swanlab.register_callbacks([lark_callback])
+
+ class SwanLabCallbackExtension(SwanLabCallback):
+ def setup(self, args: "TrainingArguments", state: "TrainerState", model: "PreTrainedModel", **kwargs):
+ if not state.is_world_process_zero:
+ return
+
+ super().setup(args, state, model, **kwargs)
+ try:
+ if hasattr(self, "_swanlab"):
+ swanlab_public_config = self._swanlab.get_run().public.json()
+ else: # swanlab <= 0.4.9
+ swanlab_public_config = self._experiment.get_run().public.json()
+ except Exception:
+ swanlab_public_config = {}
+
+ with open(os.path.join(args.output_dir, SWANLAB_CONFIG), "w") as f:
+ f.write(json.dumps(swanlab_public_config, indent=2))
+
+ swanlab_callback = SwanLabCallbackExtension(
+ project=finetuning_args.swanlab_project,
+ workspace=finetuning_args.swanlab_workspace,
+ experiment_name=finetuning_args.swanlab_run_name,
+ mode=finetuning_args.swanlab_mode,
+ config={"Framework": "🦙LlamaFactory"},
+ logdir=finetuning_args.swanlab_logdir,
+ )
+ return swanlab_callback
+
+
+def get_ray_trainer(
+ training_function: Callable,
+ train_loop_config: dict[str, Any],
+ ray_args: "RayArguments",
+) -> "TorchTrainer":
+ if not ray_args.use_ray:
+ raise ValueError("Ray was not enabled. Please set `USE_RAY=1` to enable ray.")
+
+ if ray_args.ray_init_kwargs is not None:
+ ray.init(**ray_args.ray_init_kwargs)
+
+ trainer = TorchTrainer(
+ training_function,
+ train_loop_config=train_loop_config,
+ scaling_config=ScalingConfig(
+ num_workers=ray_args.ray_num_workers,
+ resources_per_worker=ray_args.resources_per_worker,
+ placement_strategy=ray_args.placement_strategy,
+ use_gpu=True,
+ ),
+ run_config=RunConfig(
+ name=ray_args.ray_run_name,
+ storage_path=Path(ray_args.ray_storage_path).absolute().as_posix(),
+ ),
+ )
+ return trainer
diff --git a/post-training/LLaMA-Factory/src/llamafactory/train/tuner.py b/post-training/LLaMA-Factory/src/llamafactory/train/tuner.py
new file mode 100644
index 0000000..3adb382
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/train/tuner.py
@@ -0,0 +1,195 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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 os
+import shutil
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+import torch.distributed as dist
+from transformers import PreTrainedModel
+
+from ..data import get_template_and_fix_tokenizer
+from ..extras import logging
+from ..extras.constants import V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME
+from ..extras.misc import infer_optim_dtype
+from ..extras.packages import is_ray_available
+from ..hparams import get_infer_args, get_ray_args, get_train_args, read_args
+from ..model import load_model, load_tokenizer
+from .callbacks import LogCallback, PissaConvertCallback, ReporterCallback
+from .dpo import run_dpo
+from .kto import run_kto
+from .ppo import run_ppo
+from .pt import run_pt
+from .rm import run_rm
+from .sft import run_sft
+from .trainer_utils import get_ray_trainer, get_swanlab_callback
+
+
+if is_ray_available():
+ import ray
+ from ray.train.huggingface.transformers import RayTrainReportCallback
+
+
+if TYPE_CHECKING:
+ from transformers import TrainerCallback
+
+
+logger = logging.get_logger(__name__)
+
+
+def _training_function(config: dict[str, Any]) -> None:
+ args = config.get("args")
+ callbacks: list[Any] = config.get("callbacks")
+ model_args, data_args, training_args, finetuning_args, generating_args = get_train_args(args)
+
+ callbacks.append(LogCallback())
+ if finetuning_args.pissa_convert:
+ callbacks.append(PissaConvertCallback())
+
+ if finetuning_args.use_swanlab:
+ callbacks.append(get_swanlab_callback(finetuning_args))
+
+ callbacks.append(ReporterCallback(model_args, data_args, finetuning_args, generating_args)) # add to last
+
+ if finetuning_args.stage == "pt":
+ run_pt(model_args, data_args, training_args, finetuning_args, callbacks)
+ elif finetuning_args.stage == "sft":
+ run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
+ elif finetuning_args.stage == "rm":
+ run_rm(model_args, data_args, training_args, finetuning_args, callbacks)
+ elif finetuning_args.stage == "ppo":
+ run_ppo(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
+ elif finetuning_args.stage == "dpo":
+ run_dpo(model_args, data_args, training_args, finetuning_args, callbacks)
+ elif finetuning_args.stage == "kto":
+ run_kto(model_args, data_args, training_args, finetuning_args, callbacks)
+ else:
+ raise ValueError(f"Unknown task: {finetuning_args.stage}.")
+
+ if is_ray_available() and ray.is_initialized():
+ return # if ray is intialized it will destroy the process group on return
+
+ try:
+ if dist.is_initialized():
+ dist.destroy_process_group()
+ except Exception as e:
+ logger.warning(f"Failed to destroy process group: {e}.")
+
+
+def run_exp(args: Optional[dict[str, Any]] = None, callbacks: Optional[list["TrainerCallback"]] = None) -> None:
+ args = read_args(args)
+ if "-h" in args or "--help" in args:
+ get_train_args(args)
+
+ ray_args = get_ray_args(args)
+ callbacks = callbacks or []
+ if ray_args.use_ray:
+ callbacks.append(RayTrainReportCallback())
+ trainer = get_ray_trainer(
+ training_function=_training_function,
+ train_loop_config={"args": args, "callbacks": callbacks},
+ ray_args=ray_args,
+ )
+ trainer.fit()
+ else:
+ _training_function(config={"args": args, "callbacks": callbacks})
+
+
+def export_model(args: Optional[dict[str, Any]] = None) -> None:
+ model_args, data_args, finetuning_args, _ = get_infer_args(args)
+
+ if model_args.export_dir is None:
+ raise ValueError("Please specify `export_dir` to save model.")
+
+ if model_args.adapter_name_or_path is not None and model_args.export_quantization_bit is not None:
+ raise ValueError("Please merge adapters before quantizing the model.")
+
+ tokenizer_module = load_tokenizer(model_args)
+ tokenizer = tokenizer_module["tokenizer"]
+ processor = tokenizer_module["processor"]
+ template = get_template_and_fix_tokenizer(tokenizer, data_args)
+ model = load_model(tokenizer, model_args, finetuning_args) # must after fixing tokenizer to resize vocab
+
+ if getattr(model, "quantization_method", None) is not None and model_args.adapter_name_or_path is not None:
+ raise ValueError("Cannot merge adapters to a quantized model.")
+
+ if not isinstance(model, PreTrainedModel):
+ raise ValueError("The model is not a `PreTrainedModel`, export aborted.")
+
+ if getattr(model, "quantization_method", None) is not None: # quantized model adopts float16 type
+ setattr(model.config, "torch_dtype", torch.float16)
+ else:
+ if model_args.infer_dtype == "auto":
+ output_dtype = getattr(model.config, "torch_dtype", torch.float32)
+ if output_dtype == torch.float32: # if infer_dtype is auto, try using half precision first
+ output_dtype = infer_optim_dtype(torch.bfloat16)
+ else:
+ output_dtype = getattr(torch, model_args.infer_dtype)
+
+ setattr(model.config, "torch_dtype", output_dtype)
+ model = model.to(output_dtype)
+ logger.info_rank0(f"Convert model dtype to: {output_dtype}.")
+
+ model.save_pretrained(
+ save_directory=model_args.export_dir,
+ max_shard_size=f"{model_args.export_size}GB",
+ safe_serialization=(not model_args.export_legacy_format),
+ )
+ if model_args.export_hub_model_id is not None:
+ model.push_to_hub(
+ model_args.export_hub_model_id,
+ token=model_args.hf_hub_token,
+ max_shard_size=f"{model_args.export_size}GB",
+ safe_serialization=(not model_args.export_legacy_format),
+ )
+
+ if finetuning_args.stage == "rm":
+ if model_args.adapter_name_or_path is not None:
+ vhead_path = model_args.adapter_name_or_path[-1]
+ else:
+ vhead_path = model_args.model_name_or_path
+
+ if os.path.exists(os.path.join(vhead_path, V_HEAD_SAFE_WEIGHTS_NAME)):
+ shutil.copy(
+ os.path.join(vhead_path, V_HEAD_SAFE_WEIGHTS_NAME),
+ os.path.join(model_args.export_dir, V_HEAD_SAFE_WEIGHTS_NAME),
+ )
+ logger.info_rank0(f"Copied valuehead to {model_args.export_dir}.")
+ elif os.path.exists(os.path.join(vhead_path, V_HEAD_WEIGHTS_NAME)):
+ shutil.copy(
+ os.path.join(vhead_path, V_HEAD_WEIGHTS_NAME),
+ os.path.join(model_args.export_dir, V_HEAD_WEIGHTS_NAME),
+ )
+ logger.info_rank0(f"Copied valuehead to {model_args.export_dir}.")
+
+ try:
+ tokenizer.padding_side = "left" # restore padding side
+ tokenizer.init_kwargs["padding_side"] = "left"
+ tokenizer.save_pretrained(model_args.export_dir)
+ if model_args.export_hub_model_id is not None:
+ tokenizer.push_to_hub(model_args.export_hub_model_id, token=model_args.hf_hub_token)
+
+ if processor is not None:
+ processor.save_pretrained(model_args.export_dir)
+ if model_args.export_hub_model_id is not None:
+ processor.push_to_hub(model_args.export_hub_model_id, token=model_args.hf_hub_token)
+
+ except Exception as e:
+ logger.warning_rank0(f"Cannot save tokenizer, please copy the files manually: {e}.")
+
+ ollama_modelfile = os.path.join(model_args.export_dir, "Modelfile")
+ with open(ollama_modelfile, "w", encoding="utf-8") as f:
+ f.write(template.get_ollama_modelfile(tokenizer))
+ logger.info_rank0(f"Ollama modelfile saved in {ollama_modelfile}")
diff --git a/post-training/LLaMA-Factory/src/llamafactory/webui/__init__.py b/post-training/LLaMA-Factory/src/llamafactory/webui/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/post-training/LLaMA-Factory/src/llamafactory/webui/chatter.py b/post-training/LLaMA-Factory/src/llamafactory/webui/chatter.py
new file mode 100644
index 0000000..a2242bb
--- /dev/null
+++ b/post-training/LLaMA-Factory/src/llamafactory/webui/chatter.py
@@ -0,0 +1,230 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# 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 os
+from collections.abc import Generator
+from typing import TYPE_CHECKING, Any, Optional
+
+from transformers.utils import is_torch_npu_available
+
+from ..chat import ChatModel
+from ..data import Role
+from ..extras.constants import PEFT_METHODS
+from ..extras.misc import torch_gc
+from ..extras.packages import is_gradio_available
+from .common import get_save_dir, load_config
+from .locales import ALERTS
+
+
+if TYPE_CHECKING:
+ from ..chat import BaseEngine
+ from .manager import Manager
+
+
+if is_gradio_available():
+ import gradio as gr
+
+
+def _escape_html(text: str) -> str:
+ r"""Escape HTML characters."""
+ return text.replace("<", "<").replace(">", ">")
+
+
+def _format_response(text: str, lang: str, escape_html: bool, thought_words: tuple[str, str]) -> str:
+ r"""Post-process the response text.
+
+ Based on: https://huggingface.co/spaces/Lyte/DeepSeek-R1-Distill-Qwen-1.5B-Demo-GGUF/blob/main/app.py
+ """
+ if thought_words[0] not in text:
+ return _escape_html(text) if escape_html else text
+
+ text = text.replace(thought_words[0], "")
+ result = text.split(thought_words[1], maxsplit=1)
+ if len(result) == 1:
+ summary = ALERTS["info_thinking"][lang]
+ thought, answer = text, ""
+ else:
+ summary = ALERTS["info_thought"][lang]
+ thought, answer = result
+
+ if escape_html:
+ thought, answer = _escape_html(thought), _escape_html(answer)
+
+ return (
+ f"{summary} \n{thought}\n
\nLLaMA Board: A One-stop Web UI for Getting Started with LLaMA Factory Real Estate Agency \n \n \n
\n
Our Services \n
Whether you're looking for a new home, a rental property, or a commercial property, we have the expertise and resources to help you.
\n
\n
\n \n
\n\nWelcome to our Real Estate Agency \n\n At our Real Estate Agency, we are dedicated to providing our clients with the best possible service. We understand that buying or selling a property can be a big decision, which is why we offer a wide range of services to help you find the perfect property. Whether you're looking for a new home, a rental property, or a commercial property, we have the expertise and resources to help you.\n
\n \n\n\n \n\n\n
\n
\n
Real Estate Agency \n\n\n
\n
Follow Us \n\n\n
\n
\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original web page you provided. I've included placeholders for images and links where necessary.\n\n```html\n\n\n\n Healthcare Provider \n \n \n
\n
About Us \n
We are committed to providing the best care possible. Our team of experienced healthcare professionals is dedicated to providing high-quality, compassionate care to our patients. We believe in treating everyone with respect and empathy, regardless of their background or condition.
\n
\n
\n
Services \n
We offer a wide range of services to meet the needs of our patients. These include primary care, specialty care, and more. Our team of healthcare professionals is here to help you every step of the way.
\n
\n
\n
Doctors \n
Our team of healthcare professionals is some of the best in the industry. They are experts in their fields and are dedicated to providing the best care possible.
\n
\n
\n\nHealthcare Provider \n
\n\n \n\nAbout Us \n\n At Healthcare Provider, we are committed to providing the best care possible. Our team of experienced healthcare professionals is dedicated to providing high-quality, compassionate care to our patients. We believe in treating everyone with respect and empathy, regardless of their background or condition.\n
\n \n\nServices \n\n We offer a wide range of services to meet the needs of our patients. These include primary care, specialty care, and more. Our team of healthcare professionals is here to help you every step of the way.\n
\n \n\nDoctors \n\n Our team of healthcare professionals is some of the best in the industry. They are experts in their fields and are dedicated to providing the best care possible.\n
\n \n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot. I've included placeholders for images and links where necessary.\n\n```html\n\n\n\n Tech Company \n \n Welcome to Tech Company \n At Tech Company, we are dedicated to providing the best technology solutions for our clients. Our team of experts is committed to delivering high-quality, reliable, and innovative products and services. We are passionate about technology and are always looking for new ways to improve and innovate. Whether you're a small business or a large corporation, we're here to help you achieve your technology goals.
\n \n \n Home \n About \n Contact \n \n \n At Tech Company, we are dedicated to providing the best technology solutions for our clients. Our team of experts is committed to delivering high-quality, reliable, and innovative products and services. We are passionate about technology and are always looking for new ways to improve and innovate. Whether you're a small business or a large corporation, we're here to help you achieve your technology goals.
\n \n \n\n\n```\n\nThis code uses Tailwind CSS for styling and includes placeholders for images and links where necessary. The layout and design should closely resemble the original screenshot.", "label": "\n\n\nTech Company \n
\n\n \n\n
Welcome to Tech Company \n
\n At Tech Company, we are dedicated to providing the best technology solutions for our clients. Our team of experts is committed to delivering high-quality, reliable, and innovative products and services. We are passionate about technology and are always looking for new ways to improve and innovate. Whether you're a small business or a large corporation, we're here to help you achieve your technology goals.\n
\n
Our Educational Institution \n \n \n
\n
At our educational institution, we are committed to providing a high-quality education that prepares students for success in their chosen fields. Our experienced faculty, state-of-the-art facilities, and comprehensive curriculum ensure that our students are well-equipped to excel in their studies and beyond. \n
Our mission is to provide a supportive and enriching learning environment where students can thrive academically and personally. We believe in the power of education to shape the future and are dedicated to helping our students achieve their goals.
\n
\n \n
\n
\n
\n
Welcome to Our Educational Institution \n
\n At our educational institution, we are committed to providing a high-quality education that prepares students for success in their chosen fields. Our experienced faculty, state-of-the-art facilities, and comprehensive curriculum ensure that our students are well-equipped to excel in their studies and beyond.\n
\n
\n Learn More\n \n
\n
Travel Agency \n \n \n
\n
Explore the World \n
We offer a wide range of travel packages to destinations all over the world. Whether you're looking for a relaxing beach vacation, an adventurous trek, or a cultural experience, we've got you covered.
\n
\n
\n
\n
\n
\n
Travel Agency \n
\n Welcome to our travel agency, where we specialize in creating unforgettable travel experiences. We offer a wide range of travel packages to destinations all over the world. Whether you're looking for a relaxing beach vacation, an adventurous trek, or a cultural experience, we've got you covered.\n
\n
\n
\n
Welcome to Our Real Estate Agency \n \n \n
\n
\n
At our real estate agency, we are committed to providing our clients with the highest level of service and expertise. We understand that buying or selling a home is a big decision, and we are here to help you every step of the way. Whether you're looking for a home, selling your current home, or investing in real estate, we have the resources and expertise to help you.
\n
\n
\n
Our Team of Experienced Real Estate Professionals \n
Our team of experienced real estate professionals is dedicated to providing you with comprehensive market analysis, property research, and expert advice. We understand that buying or selling a home can be a complex process, which is why we offer a wide range of services to meet your unique needs.
\n
We are committed to providing you with the best possible service and we are always here to help you. Whether you're looking for a home, selling your home, or investing in real estate, we are here to support you.
\n
\n
\n \n
\n
\n
\n
Welcome to Our Real Estate Agency \n
At our real estate agency, we are committed to providing our clients with the highest level of service and expertise. We understand that buying or selling a home is a big decision, and we are here to help you every step of the way. Whether you're looking for your dream home, selling your current home, or investing in real estate, we have the resources and expertise to help you.
\n
Our team of experienced real estate professionals is dedicated to providing you with comprehensive market analysis, property research, and expert advice. We understand that buying or selling a home can be a complex process, which is why we offer a wide range of services to meet your unique needs.
\n
We are committed to providing you with the best possible service and we are always here to help you. Whether you're looking for a home, selling your home, or investing in real estate, we are here to support you.
\n
\n
\n
Fashion Brand \n \n \n
\n
Our Collections \n
Our collections are designed to cater to a wide range of tastes and preferences, from classic to modern, from casual to formal, and everything in between. We believe that fashion is just as much about what you wear as it is about how you live.
\n
\n
\n
\n
Product 1 \n
Product description 1
\n
\n
\n
\n
Product 2 \n
Product description 2
\n
\n
\n
\n
Product 3 \n
Product description 3
\n
\n
\n
\n \n
\n
\n
Welcome to our fashion brand, where we believe in the power of style and self-expression. We are a team of passionate individuals who are dedicated to creating unique and timeless pieces that are not only fashionable but also comfortable and wearable. Our collections are designed to cater to a wide range of tastes and preferences, from classic to modern, from casual to formal, and everything in between. We believe that fashion is not just about what you wear, but also about how you live. That's why we offer a wide range of products that are not only stylish but also functional. Whether you're looking for a wardrobe upgrade or simply want to express your individuality, we have something for everyone.
\n
\n
\n
At our brand, we are committed to sustainability and ethical production. We use only the highest quality materials and we work closely with our suppliers to ensure that our products are not only stylish but also ethically sourced. We believe that fashion should not only be about what you wear, but also about how you live. That's why we offer a wide range of products that are not only stylish but also functional. Whether you're looking for a wardrobe upgrade or simply want to express your individuality, we have something for everyone.
\n
\n
Online Education Platform \n \n \n
\n
About Us \n
Our online education platform is designed to provide flexible, interactive, and engaging learning experiences. We offer a wide range of courses in various subjects, from science and art to technology and business, and everything in between. Our courses are taught by experienced professionals in their respective fields, and they are designed to be accessible to all, regardless of your previous knowledge or experience.
\n
\n
\n
Our Courses \n
We offer a variety of courses to suit your learning needs. Whether you're just starting out or looking to advance your skills, we have something for everyone. Our courses are designed to be interactive and engaging, with a mix of video lectures, quizzes, and practical exercises.
\n
\n
\n
Why Choose Us? \n
Our online education platform is committed to providing high-quality education to everyone, anywhere, at any time. We believe in the power of education to change lives, and we're committed to making that a reality.
\n
\n
\nAbout Us \nOur online education platform is designed to provide flexible, interactive, and engaging learning experiences. We offer a wide range of courses in various subjects, from science to art, from technology to business, and everything in between. Our courses are taught by experienced professionals in their respective fields, and they are designed to be accessible to all, regardless of your previous knowledge or experience.
\n \n\nOur Courses \nWe offer a variety of courses to suit your learning needs. Whether you're just starting out or looking to advance your skills, we have something for everyone. Our courses are designed to be interactive and engaging, with a mix of video lectures, quizzes, and practical exercises.
\n \n\nWhy Choose Us? \nOur online education platform is committed to providing high-quality education to everyone, anywhere, at any time. We believe in the power of education to change lives, and we're committed to making that a reality.
\n \n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot. I've included placeholders for images and links where necessary.\n\n```html\n\n\n\n Tech Company \n \n \n
Welcome to Tech Company \n
At Tech Company, we are dedicated to providing the best technology solutions for our clients. Our team of experts is committed to delivering high-quality, reliable, and innovative products and services. We are passionate about technology and are always looking for new ways to improve and innovate.
\n
\n\n\n```\n\nThis code uses Tailwind CSS for styling and includes placeholders for images and links. You can replace the `placeholder.jpg` with your actual image URL and adjust the links as needed.", "label": "\n\n\nTech Company \n
\n\n \n\nWelcome to Tech Company \n\n At Tech Company, we are dedicated to providing the best technology solutions for our clients. Our team of experts is committed to delivering high-quality, reliable, and innovative products and services. We are passionate about technology and are always looking for new ways to improve and innovate.\n
\n\n \n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "```html\n\n\n\n Art Gallery \n \n \n
\n
\n
Card 1 \n
Explore the world of art with our collection of stunning pieces. From paintings to sculptures, we have something for everyone.
\n
\n
\n
\n
Card 2 \n
Discover the beauty of nature through our collection of nature-inspired artwork. Each piece is a work of art in itself.
\n
\n
\n
\n
Card 3 \n
Explore the power of abstract art with our collection of abstract pieces. Each piece is a unique expression of the artist's soul.
\n
\n
\n
\n \n
\n
\n
Art Gallery \n
\n Welcome to our Art Gallery, where we showcase a wide variety of artwork pieces from various artists. Our mission is to provide a platform for artists to showcase their work and for art enthusiasts to discover new and exciting pieces. We believe in the power of art to inspire, educate, and bring people together.\n
\n
\n\n Explore Artwork\n \n
\n
\n
Our Tech Company \n \n
\n
\n
\n
\n
\n
Welcome to Our Tech Company \n
\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n
\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n
\n
\n
\n
\n
About \n
\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n
\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n
\n
\n
Services \n
\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n
\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n
\n
\n
\n
\n
Contact \n
\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n
\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n
\n
\n
Phone \n
\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n
\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n
\n
\n
\n
\n
Email \n
\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n
\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n
\n
\n
Address \n
\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n
\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n
\n
\n
\n \n\n\n \n\n\nWelcome to Our Tech Company \n\n At our tech company, we are dedicated to providing innovative and reliable technology solutions to meet your business needs. Our team of experts is committed to delivering high-quality products and services that exceed your expectations.\n
\n\n We specialize in a wide range of services, including software development, IT consulting, and cybersecurity. Whether you're a small business or a large corporation, we have the expertise to help you achieve your goals.\n
\n \n\nContact Us \n\n Phone: 123-456-7890\n
\n\n Email: info@techcompany.com\n
\n\n Address: 123 Tech Street, City, State, Zip\n
\n \n \n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot. Note that I've used placeholder images and links for the elements that are not specified in the original screenshot.\n\n```html\n\n\n\n Travel Agency \n \n \n
\n
Welcome to our travel agency \n
We are a professional travel agency that offers a wide range of services to make your travel experience unforgettable.
\n
Discover the world with us and create memories that last a lifetime.
\n
\n
\n
\n
\n
\n
\n
\n
Travel Agency \n
\n
\nBook Now \n
\n
\n
\n
\n
\n
Welcome to our travel agency \n
We are a professional travel agency that offers a wide range of services to make your travel experience unforgettable.
\n
Learn More \n
\n
\n
Non-Profit Organization \n \n \n
\n
\n
Success Stories \n
We have been successful in raising funds for numerous charities and causes. We are proud to have helped so many people in need.
\n
\n
\n
Volunteer Opportunities \n
We are always looking for volunteers to help us with our mission. If you are interested in volunteering, please contact us.
\n
\n
\n
Donation \n
If you would like to donate, please use the form below.
\n
\n
\n
Contact Information \n
Phone: 123-456-7890 Email: info@nonprofit.org
\n
\n
\n
Volunteer Sign-up \n
If you are interested in volunteering, please sign up below.
\n
\n
\n \n
\n
\n
\n
\n
Non-Profit Organization \n
Mission \n
\n Our mission is to provide a platform for individuals to donate to causes they care about. We believe in the power of community and the importance of giving back.\n
\n
Success Stories \n
\n We have been successful in raising funds for numerous charities and causes. We are proud to have helped so many people in need.\n
\n
Volunteer Opportunities \n
\n We are always looking for volunteers to help us with our mission. If you are interested in volunteering, please contact us.\n
\n
\n
\n
\n
Donation \n
\n If you would like to donate, please use the form below.\n
\n\n
\n
\n
Contact Information \n
\n Phone: 123-456-7890\n Email: info@nonprofit.org\n
\n
\n
\n
Volunteer Sign-up \n
\n If you are interested in volunteering, please sign up below.\n
\n\n
\n
Healthcare Provider \n \n \n
\n
Our Services \n
We offer a wide range of services including physical therapy, occupational therapy, speech therapy, and more.
\n
\n
\n
About Us \n
We are a team of experienced healthcare professionals dedicated to providing the best care possible.
\n
\n \n
\nOur Services \nWe offer a wide range of services including physical therapy, occupational therapy, speech therapy, and more.
\n \n\nAbout Us \nWe are a team of experienced healthcare professionals dedicated to providing the best care possible.
\n \n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "To create an HTML and Tailwind CSS version of the web page, we need to ensure that the layout, size, text, and color of all elements match the original screenshot. Here's a step-by-step guide to achieve this:\n\n1. **HTML Structure**: We'll start by creating the basic HTML structure.\n2. **Tailwind CSS Classes**: We'll use Tailwind CSS classes to style the elements according to the screenshot.\n\nHere's the complete code:\n\n```html\n\n\n\n Web Page \n \n
\n
\n
Company Name \n
Categories
\n
Category 1
\n
Category 2
\n
Category 3
\n
\n
Product Image
\n
\n
\n
\n
Product Image \n
Categories
\n
Category 1
\n
Category 2
\n
Category 3
\n
\n
Product Image
\n
\n
\n
\n
\n
Customer Reviews \n
This is a great product. Highly recommend it!
\n
John Doe
\n
\n
\n\n
\n
Company Name \n
\n
\n \n\n\n
Categories \n
\nCategory 1 \nCategory 2 \nCategory 3 \n \n
\n
Product Image \n
\n
\n
Customer Reviews \n
\n
\"This is a great product. I highly recommend it.\"
\n
- John Doe
\n
\n
\n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot.\n\n```html\n\n\n\n Food Delivery Service \n \n Welcome to Our Food Delivery Service \n At our food delivery service, we are committed to providing you with the best service and quality food. We use only the freshest ingredients and our experienced chefs to prepare your meals. Our delivery service is fast, reliable, and convenient. We are proud to offer a wide variety of dishes from classic comfort food to unique and exotic dishes. We also offer a wide range of promotions and discounts to make your meal experience even more enjoyable.
\n \n \n Featured Dishes \n Customer Reviews \n \n \n \n Featured Dishes \n Discover our delicious and diverse menu of dishes. From classic comfort food to unique and exotic dishes, we have something for everyone.
\n \n \n Customer Reviews \n Read what our satisfied customers have to say about our food delivery service.
\n \n \n \n\n\n```\n\nThis code uses Tailwind CSS for styling and ensures that the layout and elements match the original screenshot. The `placeholder.jpg` images are placeholders and should be replaced with actual images.", "label": "\n\n
Welcome to Our Food Delivery Service \n
\n At our food delivery service, we are committed to providing you with the best service and quality food. We use only the freshest ingredients and our experienced chefs to prepare your meals. Our delivery service is fast, reliable, and convenient. We are proud to offer a wide variety of dishes, from classic comfort food to unique and exotic dishes. We also offer a wide range of promotions and discounts to make your meal experience even more enjoyable.\n
\n
\n
\n
Featured Dishes \n\n\n
\n
Customer Reviews \n\n\n
\n
Company Website \n \n
\n
\n
Welcome to Our Company \n
We offer a wide range of high-quality products and services to our customers.
\n
\n
\n
\n
Our Services \n
We provide top-notch services to our customers. Learn more about our company and our mission.
\n
\n
\n \n\n\nHome \nProducts \nServices \nAbout Us \n \n \n\n\n
\n
Products \n
We offer a wide range of high-quality products.
\n
\n
\n
Services \n
We provide top-notch services to our customers.
\n
\n
\n
About Us \n
Learn more about our company and our mission.
\n
\n
\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "```html\n\n\n\n Pharmaceutical Company \n \n Pharmaceutical Company \n \n \n
\n About Us \n We are a leading pharmaceutical company dedicated to providing high-quality, safe, and effective medicines to our patients. Our mission is to improve the health and well-being of our communities through our innovative medicines.
\n \n
\n Our Services \n We offer a wide range of services including prescription drugs, over-the-counter medicines, and health and wellness products. We also provide expert advice and consultation on a range of health issues.
\n \n
\n Request a Consultation \n If you're not sure what you need, or if you're not sure how to take your medicines, we're here to help. Contact us today to schedule a consultation.
\n \n
\nPharmaceutical Company \n \n\n\nAbout Us \n\n We are a leading pharmaceutical company dedicated to providing high-quality, safe, and effective medicines to our patients. Our mission is to improve the health and well-being of our communities through our innovative medicines.\n
\n \n\nOur Services \n\n We offer a wide range of services including prescription drugs, over-the-counter medicines, and health and wellness products. We also provide expert advice and consultation on a range of health issues.\n
\n \n\nRequest a Consultation \n\n If you're not sure what you need, or if you're not sure how to take your medicines, we're here to help. Contact us today to schedule a consultation.\n
\nRequest a Consultation \n \n \n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "To create an HTML and Tailwind CSS version of the web page, we need to ensure that the layout, text, and colors match the original screenshot. Here's a step-by-step guide to achieve this:\n\n1. **HTML Structure**: We'll start with the basic HTML structure.\n2. **Tailwind CSS**: We'll use Tailwind CSS classes to style the elements.\n\nHere's the complete code:\n\n```html\n\n\n\n Home Decor Store \n \n \n
\n
Welcome to Home Decor Store \n
\n At Home Decor Store, we offer a wide range of high-quality home decor items to make your living space more beautiful and functional. From stylish lamps and cozy throws to elegant furniture and unique accessories, we have something for everyone.\n
\n
\n Shop Now\n \n
\n
\n
Featured Products \n
\n
\n
\n
\n
Lampshade \n
Elegant and stylish, our lampshades are perfect for adding a touch of sophistication to any room.
\n
Price: $50
\n
\n Buy Now\n \n
\n
\n
\n
\n
\n
Throw Blanket \n
Soft and cozy, our throw blankets are perfect for adding warmth and comfort to any space.
\n
Price: $30
\n
\n Buy Now\n \n
\n
\n
\n
\n
\n\n \n\n\n```\n\n### Explanation:\n1. **Header**: The header section includes navigation links for \"Home,\" \"Inspiration,\" \"Products,\" and \"About Us.\" The logo is placed on the left, and the navigation links are on the right.\n2. **Main Content**: The main content section includes a welcome message, a featured product section, and a \"Shop Now\" button.\n3. **Footer**: The footer section is similar to the header, with navigation links for \"Home,\" \"Inspiration,\" \"Products,\" and \"About Us.\"\n\nThis code uses Tailwind CSS classes to style the elements, ensuring that the layout and design match the original screenshot.", "label": "\n\n\n\n \n \n\n\n\n\n\n \n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "To create an HTML and Tailwind CSS version of the web page that matches the original screenshot, you can use the following code. This code includes the necessary HTML structure, Tailwind CSS classes, and placeholder images and links.\n\n```html\n\n\n\n Welcome to our Tech Company \n \n \n
\n
Our Services \n
Our team of experts is dedicated to delivering the best possible service, ensuring your business runs smoothly and efficiently.
\n
\n \n
\nWelcome to our Tech Company \n\n At our tech company, we are committed to providing innovative and efficient solutions for your business needs. Our team of experts is dedicated to delivering the best possible service, ensuring your business runs smoothly and efficiently.\n
\nLearn More \n \n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "```html\n\n\n\n My Fashion Brand \n \n \n
\n
\n
\n
Product 1 \n
Product 1 description goes here.
\n
\n
\n
\n
Product 2 \n
Product 2 description goes here.
\n
\n
\n
\n
Product 3 \n
Product 3 description goes here.
\n
\n
\n \n
\n\n \n\n\n Welcome to our fashion brand, where we believe in the power of style to elevate your everyday. Our mission is to create timeless pieces that are not just fashionable, but also comfortable, durable, and made with the finest materials. We are committed to providing our customers with high-quality, stylish clothing that is not only fashionable but also functional. We believe that fashion should not only be about what you wear, but also about how you live. That's why we offer a wide range of products that cater to all tastes and styles. From casual wear to formal attire, our collection is designed to make you feel confident and beautiful. We are constantly updating our collection with the latest trends and styles, ensuring that you always have something new to wear. So, come and explore our fashion brand and discover the perfect pieces that will elevate your style and your everyday.\n
\n \n\n\n
\n
Contact Us
\n
123 Street, City, Country
\n
Phone: 123-456-7890
\n
Email: info@fashionbrand.com
\n
\n
\n
\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "```html\n\n\n\n Travel Agency \n \n Welcome to Educational Institution \n At Educational Institution, we are committed to providing the best education possible. Our experienced and dedicated staff are here to help you succeed in your studies. We offer a wide range of programs, from online courses to in-person classes. Our flexible schedules and affordable tuition make us a great choice for students of all ages and backgrounds.
\n \n \n Home \n About \n Contact \n \n \n We believe in the power of education to change lives. That's why we are dedicated to providing a supportive and inclusive learning environment. Our programs are designed to help you reach your full potential, no matter where you are in your education journey.
\n \n \n\n\n```", "label": "\n\n\nEducational Institution \n
\n\n \n\n
Welcome to Educational Institution \n
\n At Educational Institution, we are committed to providing the best education possible. Our experienced and dedicated staff are here to help you succeed in your studies. We offer a wide range of programs, from online courses to in-person classes. Our flexible schedules and affordable tuition make us a great choice for students of all ages and backgrounds.\n
\n
\n We believe in the power of education to change lives. That's why we are dedicated to providing a supportive and inclusive learning environment. Our programs are designed to help you reach your full potential, no matter where you are in your education journey.\n
\n
Gaming Company \n \n \n
\n
Our Games \n
At Gaming Company, we are constantly working on new and exciting games. From fast-paced action to relaxing puzzle games, we have something for everyone. Check out our latest releases and stay tuned for future updates.
\n
\n
\n
Our Games \n
At Gaming Company, we are constantly working on new and exciting games. From fast-paced action to relaxing puzzle games, we have something for everyone. Check out our latest releases and stay tuned for future updates.
\n
\n \n \n
\n\nWelcome to Gaming Company \n\n We are a leading provider of innovative and exciting gaming experiences. Our mission is to bring the joy of gaming to everyone, no matter where they are. We believe in the power of games to inspire, engage, and connect people from all walks of life.\n
\n \n\nOur Games \n\n At Gaming Company, we are constantly working on new and exciting games. From fast-paced action to relaxing puzzle games, we have something for everyone. Check out our latest releases and stay tuned for future updates.\n
\n \n\nJoin Our Community \n\n We believe in the power of community. That's why we have a vibrant and active community of gamers who are always ready to share their experiences, ask questions, and help each other out. Join us today and become a part of our community.\n
\n \n \n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "```html\n\n\n\n Financial Advisory Firm \n \n \n
\n
\n
Personal Financial Planning \n
Personal Financial Planning is a comprehensive approach to managing your finances. We help you create a financial plan that aligns with your goals and objectives.
\n
\n
\n
Investment Planning \n
Investment Planning involves selecting and managing investments to achieve your financial goals. We help you make informed investment decisions.
\n
\n
\n
Retirement Planning \n
Retirement Planning is the process of planning for your retirement. We help you create a retirement plan that ensures you have a comfortable retirement.
\n
\n
\n
Tax Planning \n
Tax Planning involves managing your taxes to minimize your tax liability. We help you understand your tax obligations and make informed tax decisions.
\n
\n
\n \n
\n
\n
\n
\n
\n
\n
Welcome to Our Financial Advisory Firm \n
\n At our firm, we are committed to providing our clients with the best possible financial advice. We understand that financial planning is a complex process that requires a deep understanding of your individual needs and circumstances. That's why we offer a wide range of services to help you achieve your financial goals.\n
\n
\n
\n
\n
\n
\n
Our Services \n
\nPersonal Financial Planning \nInvestment Planning \nRetirement Planning \nTax Planning \n \n
\n
\n
\n
\n
\n
\n
\n
\n Travel Destinations\n
\n
\n Welcome to our Travel Agency\n
\n
\n We are a travel agency that offers a wide range of services to make your travel experience unforgettable.\n
\n
\n
\n
\n
\n
\n
\n
Popular Destinations \n
\n Discover the best destinations in the world. From bustling cities to serene beaches, our agency has you covered.\n
\n
\n
\n
\n
\n
\n
\n
\n
Beach Destinations \n
\n Escape to the sun and sand with our beach destinations. From tropical islands to coastal towns, we have you covered.\n
\n
\n
\n
\n
\n
\n
\n
\n
Welcome to our Travel Agency \n
We are a travel agency that offers a wide range of services to make your travel experience unforgettable.
\n
\n
\n
\n
Popular Destinations \n
\n
\n
\n
\n
Beach Destinations \n
Discover the best beaches in the world.
\n
\n
\n\n
\n
Travel Agency \n \n Welcome to our Travel Agency \n We are a dynamic travel agency that offers a wide range of travel destinations. Whether you're looking for a relaxing beach vacation, an adventurous mountain trek, or a cultural city break, we've got you covered.
\n \n \n Home \n Destinations \n About Us \n Contact Us \n \n \n Explore Destinations \n Discover the world with us. From exotic beaches to breathtaking mountains, we have something for everyone.
\n \n \n\n\n```\n\nThis code uses Tailwind CSS for styling and ensures that the layout and design closely match the original screenshot. The `header`, `nav`, `main`, and `footer` sections are styled to match the original layout and content.", "label": "\n\n\nWelcome to our Travel Agency \nWe are a dynamic travel agency that offers a wide range of travel destinations. Whether you're looking for a relaxing beach vacation, an adventurous mountain trek, or a cultural city break, we've got you covered.
\nExplore Destinations \n \n \n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot. I've included placeholders for images and links where necessary.\n\n```html\n\n\n\n Marketing Agency \n \n
\n
\n
\n
\n
Our Services \n
We are a marketing firm that specializes in providing top-notch services to businesses of all sizes.
\n
\n
\n
\n\nWe are a marketing firm that specializes in providing top-notch services to businesses of all sizes. \n\n Get Started\n \n \n \n\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot. I'll use placeholder images and links where necessary.\n\n```html\n\n\n\n Real Estate Company \n \n \n Featured Listings \n \n \n\n \n Search Properties \n \n \n\n \n Featured Listings \n \n \n\n \n\n \n \n\n \n\n\n```\n\nThis code uses Tailwind CSS for styling and includes placeholders for images and links. You can replace the placeholder text and images with your actual content.", "label": "\nEducation Platform \n \n \n
\n
\n
Featured Courses \n
We offer a wide range of courses to help you reach your educational goals.
\n
View Courses \n
\n
\n
Course Categories \n
Explore our diverse course categories to find the one that suits you best.
\n
View Categories \n
\n
\n
Pricing Information \n
Discover our flexible pricing options to fit your budget.
\n
View Pricing \n
\n
\n
\n\n
\n
Welcome to Our Education Platform \n
We offer a wide range of courses to help you reach your educational goals.
\n
Enroll Now \n
\n
\n\n
\n
\n
Featured Courses \n\n\n
\n
Course Categories \n\nPricing Information \n\n\n
\n
Non-profit Organization \n \n \n
Your tagline here
\n
\n
Story Title 1
\n
Your story content goes here
\n
\n
\n
Our Impact
\n
\n
Story Title 2
\n
Your story content goes here
\n
\n
\n
Story Title 3
\n
Your story content goes here
\n
\n
\n
\nNon-profit Organization \nYour tagline goes here
\nDonate or Get Involved \n \n\nOur Impact \n\n
Story Title 1 \n
Your story content goes here
\n
\n
Story Title 2 \n
Your story content goes here
\n
\n
Story Title 3 \n
Your story content goes here
\n
\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Sure! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot.\n\n```html\n\n\n\n Tech Startup \n \n \n Home \n About \n Contact \n \n Welcome to Tech Startup \n At Tech Startup, we are passionate about creating innovative and playful designs. Our team of experts is dedicated to delivering the best possible user experience. We believe in the power of design to bring joy and excitement to our users.
\n \n\n \n Tech Startup \n We are committed to using our skills and knowledge to create a better future for all.
\n Learn More \n \n\n \n\n\n```\n\nThis code uses Tailwind CSS for styling and includes placeholders for images and links. You can replace the placeholder images and links with your actual content.", "label": "\n\n\nTech Startup \n
\n\n\n \n\nWelcome to Tech Startup \nAt Tech Startup, we are passionate about creating innovative and playful designs. Our team of experts is dedicated to delivering the best possible user experience. We believe in the power of design to bring joy and excitement to our users.
\nOur mission is to make the world a better place, one design at a time. We are committed to using our skills and knowledge to create a better future for all.
\n \n\n\n
\n
Tech Startup
\n
123 Main St, Anytown, USA
\n
(123) 456-7890
\n
info@techstartup.com
\n
\n
\n
Follow us on:
\n
Facebook
\n
Twitter
\n
Instagram
\n
\n
\n\n\n"}
+{"prompt": "system\nYou are a helpful assistant.\nuser\nHere is a screenshot of a web page. Please write a HTML and Tailwind CSS to make it look exactly like the original web page. Pay attention to things like size, text, position, and color of all the elements, as well as the overall layout. If it involves any images, use 'placeholder.jpg' as the placeholder. If it involves any links, use '#' as the placeholder.\nassistant\n", "predict": "Certainly! Below is the HTML and Tailwind CSS code to create a web page that looks exactly like the original screenshot. Note that I've used placeholder images and links for the elements that are not directly provided in the screenshot.\n\n```html\n\n\n\n Our Educational Institution \n \n
\n
Welcome to Our Educational Institution \n
We provide top-notch education with a modern, clean design.
\n
\n