[train] Harden KTransformers MoE LoRA SFT integration (#10738)

src/llamafactory/hparams/model_args.py

@@ -470,10 +470,23 @@ class KTransformersArguments:
         default=False,
         metadata={"help": "Whether to use KTransformers AMX MoE backend for SFT training."},
     )
+    kt_cpu_activation: Literal["retain", "recompute"] | None = field(
+        default=None,
+        metadata={
+            "help": (
+                "Whether KTransformers retains CPU expert activations. Defaults to recompute while GPU "
+                "gradient checkpointing is enabled and retain otherwise."
+            )
+        },
+    )
     kt_weight_path: str | None = field(
         default=None,
         metadata={"help": "Path to pre-quantized INT8 expert weights (.kt files)."},
     )
+    kt_non_expert_weight_path: str | None = field(
+        default=None,
+        metadata={"help": "Path to the KT BF16 non-expert weight cache used with routed INT8 experts."},
+    )
     kt_expert_checkpoint_path: str | None = field(
         default=None,
         metadata={"help": "Path to expert checkpoint (safetensors) for online conversion."},
@@ -490,52 +503,202 @@ class KTransformersArguments:
         default=None,
         metadata={"help": "Intermediate size for GPU-side LoRA Experts."},
     )
+    _kt_inference_config: dict[str, Any] | None = field(default=None, init=False, repr=False)
+    _kt_config_handle: Any = field(default=None, init=False, repr=False)
+    _kt_adapter_artifact_path: str | None = field(default=None, init=False, repr=False)
 
-    def get_kt_config_dict(self, finetuning_args: Any, model_max_length: int | None) -> dict[str, Any]:
-        r"""Build KT config values from LLaMA-Factory model and LoRA arguments."""
-        kt_config = {
-            "kt_lora_rank": getattr(finetuning_args, "lora_rank", None),
-            "kt_lora_alpha": getattr(finetuning_args, "lora_alpha", None),
-            "kt_weight_path": self.kt_weight_path,
-            "kt_expert_checkpoint_path": self.kt_expert_checkpoint_path,
-            "kt_model_max_length": model_max_length,
-            "kt_use_lora_experts": self.kt_use_lora_experts,
-            "kt_lora_expert_num": self.kt_lora_expert_num,
-            "kt_lora_expert_intermediate_size": self.kt_lora_expert_intermediate_size,
+    _KT_DERIVED_KEYS = frozenset(
+        {
+            "enabled",
+            "kt_activation_policy",
+            "kt_expert_checkpoint_path",
+            "kt_full_weight_grad",
+            "kt_lora_alpha",
+            "kt_lora_dropout",
+            "kt_lora_expert_intermediate_size",
+            "kt_lora_expert_num",
+            "kt_lora_rank",
+            "kt_non_expert_weight_path",
+            "kt_skip_expert_loading",
+            "kt_train_mode",
+            "kt_use_lora_experts",
+            "kt_weight_path",
         }
+    )
+
+    def __post_init__(self) -> None:
+        if self.kt_cpu_activation not in {None, "retain", "recompute"}:
+            raise ValueError("`kt_cpu_activation` must be `retain` or `recompute`.")
+        if not self.use_kt and self.kt_cpu_activation is not None:
+            raise ValueError("`kt_cpu_activation` is only valid when `use_kt: true`.")
+
+    def get_kt_activation_policy(self) -> dict[str, str]:
+        r"""Resolve LF's GPU checkpoint switch and KT's CPU activation setting."""
+        gpu_activation = "retain" if self.disable_gradient_checkpointing else "recompute"
+        cpu_activation = self.kt_cpu_activation or gpu_activation
+        if cpu_activation == "recompute" and gpu_activation == "retain":
+            raise ValueError(
+                "`kt_cpu_activation: recompute` requires GPU gradient checkpointing. "
+                "Set `disable_gradient_checkpointing: false` or use `kt_cpu_activation: retain`."
+            )
+
+        return {"cpu": cpu_activation, "gpu": gpu_activation}
+
+    @staticmethod
+    def _get_accelerator_kt_config(training_args: Any) -> Any:
+        accelerator_config = getattr(training_args, "accelerator_config", None)
+        if isinstance(accelerator_config, dict):
+            return accelerator_config.get("kt_config")
+        return getattr(accelerator_config, "kt_config", None)
+
+    def _normalize_advanced_kt_config(self, raw_config: Any) -> dict[str, Any]:
+        if raw_config is None:
+            return {}
+        if not isinstance(raw_config, dict):
+            raise TypeError("LLaMA-Factory `kt_config` must be a flat mapping.")
+
+        config = dict(raw_config)
+        conflicts = sorted(set(config) & self._KT_DERIVED_KEYS)
+        if conflicts:
+            raise ValueError(f"These `kt_config` values are derived from LLaMA-Factory arguments: {conflicts}.")
+        return config
+
+    def _get_advanced_kt_config(self, training_args: Any) -> dict[str, Any]:
+        raw_config = getattr(training_args, "kt_config", None)
+        accelerator_config = self._get_accelerator_kt_config(training_args)
+        if raw_config is None:
+            if accelerator_config is not None:
+                raise ValueError(
+                    "Put KTransformers settings in the LLaMA-Factory training YAML `kt_config`; "
+                    "remove `kt_config` from the Accelerate config."
+                )
+            return {}
+        if accelerator_config is not None and accelerator_config != raw_config:
+            raise ValueError("LLaMA-Factory YAML and Accelerate config cannot define different KT settings.")
+        return self._normalize_advanced_kt_config(raw_config)
+
+    def configure_kt_checkpointing(self, training_args: Any) -> None:
+        r"""Keep LLaMA-Factory as the single gradient-checkpointing entry point."""
+        if self.use_unsloth or self.use_unsloth_gc:
+            raise ValueError("KTransformers cannot be combined with Unsloth checkpoint wrapping.")
+        if getattr(training_args, "gradient_checkpointing", False):
+            raise ValueError(
+                "KTransformers uses LLaMA-Factory's `disable_gradient_checkpointing`; "
+                "remove `gradient_checkpointing: true`."
+            )
+        if getattr(training_args, "gradient_checkpointing_kwargs", None) is not None:
+            raise ValueError("KTransformers supplies its checkpoint context; remove `gradient_checkpointing_kwargs`.")
+
+        fsdp_config = getattr(training_args, "fsdp_config", None)
+        if isinstance(fsdp_config, dict) and fsdp_config.get("activation_checkpointing"):
+            raise ValueError("Disable FSDP activation checkpointing when using KTransformers.")
+        if os.environ.get("FSDP_ACTIVATION_CHECKPOINTING", "false").lower() in {"1", "true", "yes"}:
+            raise ValueError("Disable FSDP activation checkpointing when using KTransformers.")
+
+        self.get_kt_activation_policy()
+        if not self.disable_gradient_checkpointing:
+            self.use_reentrant_gc = False
+        training_args.gradient_checkpointing = False
+        training_args.gradient_checkpointing_kwargs = None
+
+    def get_kt_config_dict(
+        self,
+        finetuning_args: Any,
+        model_max_length: int | None,
+        advanced_config: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        r"""Map LLaMA-Factory-owned training values to the public KT configuration."""
+        if getattr(finetuning_args, "finetuning_type", None) != "lora":
+            raise ValueError("KTransformers thin integration currently supports LoRA finetuning only.")
+
+        kt_config = dict(advanced_config or {})
+        configured_capacity = kt_config.pop("kt_model_max_length", None)
+        if configured_capacity is not None:
+            try:
+                configured_capacity = int(configured_capacity)
+            except (TypeError, ValueError) as exc:
+                raise ValueError("`kt_model_max_length` must be a positive integer.") from exc
+            if configured_capacity <= 0:
+                raise ValueError("`kt_model_max_length` must be a positive integer.")
+
+        kt_config.update(
+            {
+                "kt_lora_rank": getattr(finetuning_args, "lora_rank", None),
+                "kt_lora_alpha": getattr(finetuning_args, "lora_alpha", None),
+                "kt_lora_dropout": getattr(finetuning_args, "lora_dropout", None),
+                "kt_weight_path": self.kt_weight_path,
+                "kt_non_expert_weight_path": self.kt_non_expert_weight_path,
+                "kt_expert_checkpoint_path": self.kt_expert_checkpoint_path,
+                "kt_model_max_length": max(model_max_length or 0, configured_capacity or 0) or None,
+                "kt_use_lora_experts": self.kt_use_lora_experts,
+                "kt_lora_expert_num": self.kt_lora_expert_num,
+                "kt_lora_expert_intermediate_size": self.kt_lora_expert_intermediate_size,
+                "kt_activation_policy": self.get_kt_activation_policy(),
+                "kt_train_mode": "lora",
+                "kt_full_weight_grad": False,
+            }
+        )
         return {key: value for key, value in kt_config.items() if value is not None}
 
+    def _resolve_kt_adapter_artifact_dir(self, operation: str) -> str | None:
+        if not self.adapter_name_or_path:
+            return None
+        if len(self.adapter_name_or_path) != 1:
+            raise ValueError("KTransformers accepts a single `adapter_name_or_path`.")
+
+        adapter_root = os.path.realpath(os.path.expanduser(self.adapter_name_or_path[0]))
+        adapter_dir = adapter_root
+        if self.adapter_folder:
+            adapter_dir = os.path.realpath(os.path.join(adapter_root, self.adapter_folder))
+            if os.path.commonpath((adapter_root, adapter_dir)) != adapter_root:
+                raise ValueError("`adapter_folder` must stay inside the KT adapter directory.")
+        if not os.path.isdir(adapter_dir):
+            raise ValueError(f"KTransformers {operation} requires a local adapter directory.")
+        return adapter_dir
+
     def apply_kt_config(self, finetuning_args: Any, training_args: Any, model_max_length: int | None) -> None:
         r"""Apply LLaMA-Factory KT args to transformers/accelerate KT integration points."""
         if not self.use_kt:
             return
 
-        kt_config = self.get_kt_config_dict(finetuning_args, model_max_length)
-        env_mapping = {
-            "kt_weight_path": "ACCELERATE_KT_WEIGHT_PATH",
-            "kt_expert_checkpoint_path": "ACCELERATE_KT_EXPERT_CHECKPOINT_PATH",
-            "kt_model_max_length": "ACCELERATE_KT_MODEL_MAX_LENGTH",
-            "kt_lora_rank": "ACCELERATE_KT_LORA_RANK",
-            "kt_lora_alpha": "ACCELERATE_KT_LORA_ALPHA",
-            "kt_use_lora_experts": "ACCELERATE_KT_USE_LORA_EXPERTS",
-            "kt_lora_expert_num": "ACCELERATE_KT_LORA_EXPERT_NUM",
-            "kt_lora_expert_intermediate_size": "ACCELERATE_KT_LORA_EXPERT_INTERMEDIATE_SIZE",
-        }
-        for key, env_key in env_mapping.items():
-            value = kt_config.get(key)
-            if value is not None:
-                os.environ[env_key] = str(value)
-
-        hf_kt = getattr(training_args, "hf_kt_config", None)
-        if hf_kt is None or not hasattr(hf_kt, "_kt_config") or not isinstance(hf_kt._kt_config, dict):
-            return
-
-        hf_kt._kt_config.update(kt_config)
-        gc_enabled = getattr(training_args, "gradient_checkpointing", False) or not getattr(
-            self, "disable_gradient_checkpointing", True
+        self.configure_kt_checkpointing(training_args)
+        kt_config = self.get_kt_config_dict(
+            finetuning_args,
+            model_max_length,
+            self._get_advanced_kt_config(training_args),
         )
-        if gc_enabled:
-            hf_kt._kt_config.setdefault("kt_share_cache_pool", True)
+        update_kt_config = getattr(training_args, "update_kt_config", None)
+        if not callable(update_kt_config):
+            raise RuntimeError(
+                "The installed Transformers-KT does not provide `TrainingArguments.update_kt_config()`."
+            )
+
+        adapter_dir = self._resolve_kt_adapter_artifact_dir("training")
+        update_kt_config(kt_config, adapter_name_or_path=adapter_dir)
+
+    def configure_kt_loading(self, finetuning_args: Any, model_max_length: int | None) -> None:
+        r"""Configure KT model loading for inference and evaluation."""
+        if not self.use_kt:
+            if self._kt_inference_config is not None:
+                raise ValueError("`kt_config` requires `use_kt: true`.")
+            return
+        if self.infer_backend != EngineName.HF:
+            raise ValueError("KTransformers inference requires `infer_backend: huggingface`.")
+
+        adapter_dir = self._resolve_kt_adapter_artifact_dir("inference")
+
+        try:
+            from transformers.integrations.kt import configure_kt
+        except (ImportError, ModuleNotFoundError) as exc:
+            raise RuntimeError("The installed Transformers-KT does not provide `configure_kt()`.") from exc
+
+        kt_config = self.get_kt_config_dict(
+            finetuning_args,
+            model_max_length,
+            self._normalize_advanced_kt_config(self._kt_inference_config),
+        )
+        self._kt_adapter_artifact_path = adapter_dir
+        self._kt_config_handle = configure_kt(kt_config)
 
 
 @dataclass
@@ -580,6 +743,7 @@ class ModelArguments(
         ExportArguments.__post_init__(self)
         VllmArguments.__post_init__(self)
         SGLangArguments.__post_init__(self)
+        KTransformersArguments.__post_init__(self)
 
     @classmethod
     def copyfrom(cls, source: "Self", **kwargs) -> "Self":

src/llamafactory/hparams/parser.py

@@ -18,6 +18,7 @@
 import json
 import os
 import sys
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Optional
 
@@ -48,6 +49,14 @@ logger = logging.get_logger(__name__)
 check_dependencies()
 
 
+@dataclass
+class _KTransformersRuntimeArguments:
+    kt_config: dict[str, Any] | None = field(
+        default=None,
+        metadata={"help": "Advanced KTransformers settings used during inference or evaluation."},
+    )
+
+
 _TRAIN_ARGS = [
     ModelArguments,
     DataArguments,
@@ -56,9 +65,9 @@ _TRAIN_ARGS = [
     GeneratingArguments,
 ]
 _TRAIN_CLS = tuple[ModelArguments, DataArguments, TrainingArguments, FinetuningArguments, GeneratingArguments]
-_INFER_ARGS = [ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
+_INFER_ARGS = [ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments, _KTransformersRuntimeArguments]
 _INFER_CLS = tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
-_EVAL_ARGS = [ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
+_EVAL_ARGS = [ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments, _KTransformersRuntimeArguments]
 _EVAL_CLS = tuple[ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
 
 if is_mcore_adapter_available() and is_env_enabled("USE_MCA"):
@@ -117,6 +126,26 @@ def read_args(args: dict[str, Any] | list[str] | None = None) -> dict[str, Any]
         return sys.argv[1:]
 
 
+def _get_kt_runtime_capacity(
+    data_args: "DataArguments",
+    training_args: "TrainingArguments",
+    finetuning_args: "FinetuningArguments",
+) -> int:
+    r"""Return the largest local token batch submitted to a KT expert."""
+    tokens_per_sample = data_args.cutoff_len
+    if finetuning_args.stage == "sft" and data_args.packing:
+        tokens_per_sample += 1
+    if finetuning_args.stage == "sft" and training_args.do_train:
+        tokens_per_sample = ((tokens_per_sample + 7) // 8) * 8
+
+    local_batch_sizes = [1]
+    if training_args.do_train:
+        local_batch_sizes.append(training_args.per_device_train_batch_size)
+    if training_args.do_eval or training_args.do_predict:
+        local_batch_sizes.append(training_args.per_device_eval_batch_size)
+    return tokens_per_sample * max(local_batch_sizes)
+
+
 def _parse_args(
     parser: "HfArgumentParser", args: dict[str, Any] | list[str] | None = None, allow_extra_keys: bool = False
 ) -> tuple[Any]:
@@ -340,13 +369,21 @@ def _configure_mbridge_training_args(training_args, data_args, finetuning_args)
 def _parse_infer_args(args: dict[str, Any] | list[str] | None = None) -> _INFER_CLS:
     parser = HfArgumentParser(_INFER_ARGS)
     allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
-    return _parse_args(parser, args, allow_extra_keys=allow_extra_keys)
+    model_args, data_args, finetuning_args, generating_args, kt_args = _parse_args(
+        parser, args, allow_extra_keys=allow_extra_keys
+    )
+    model_args._kt_inference_config = kt_args.kt_config
+    return model_args, data_args, finetuning_args, generating_args
 
 
 def _parse_eval_args(args: dict[str, Any] | list[str] | None = None) -> _EVAL_CLS:
     parser = HfArgumentParser(_EVAL_ARGS)
     allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
-    return _parse_args(parser, args, allow_extra_keys=allow_extra_keys)
+    model_args, data_args, eval_args, finetuning_args, kt_args = _parse_args(
+        parser, args, allow_extra_keys=allow_extra_keys
+    )
+    model_args._kt_inference_config = kt_args.kt_config
+    return model_args, data_args, eval_args, finetuning_args
 
 
 def get_ray_args(args: dict[str, Any] | list[str] | None = None) -> RayArguments:
@@ -605,10 +642,10 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
     elif training_args.fp16:
         model_args.compute_dtype = torch.float16
 
+    data_args.packing = data_args.packing if data_args.packing is not None else finetuning_args.stage == "pt"
     model_args.device_map = {"": get_current_device()}
     model_args.model_max_length = data_args.cutoff_len
     model_args.block_diag_attn = data_args.neat_packing
-    data_args.packing = data_args.packing if data_args.packing is not None else finetuning_args.stage == "pt"
 
     # Log on each process the small summary
     logger.info(
@@ -620,7 +657,11 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
     transformers.set_seed(training_args.seed)
 
     if model_args.use_kt:
-        model_args.apply_kt_config(finetuning_args, training_args, model_args.model_max_length)
+        model_args.apply_kt_config(
+            finetuning_args,
+            training_args,
+            _get_kt_runtime_capacity(data_args, training_args, finetuning_args),
+        )
 
     return model_args, data_args, training_args, finetuning_args, generating_args
 
@@ -657,6 +698,8 @@ def get_infer_args(args: dict[str, Any] | list[str] | None = None) -> _INFER_CLS
     else:
         model_args.device_map = "auto"
 
+    model_args.configure_kt_loading(finetuning_args, data_args.cutoff_len)
+
     return model_args, data_args, finetuning_args, generating_args
 
 
@@ -675,6 +718,7 @@ def get_eval_args(args: dict[str, Any] | list[str] | None = None) -> _EVAL_CLS:
     _check_extra_dependencies(model_args, finetuning_args)
 
     model_args.device_map = "auto"
+    model_args.configure_kt_loading(finetuning_args, data_args.cutoff_len)
 
     transformers.set_seed(eval_args.seed)
 

src/llamafactory/model/adapter.py

@@ -138,6 +138,12 @@ def _setup_freeze_tuning(
     logger.info_rank0("Set trainable layers: {}".format(",".join(trainable_layers)))
 
 
+def _load_kt_inference_adapter_artifacts(model: "PreTrainedModel", adapter_path: str) -> None:
+    from kt_kernel.sft import load_kt_adapter_artifacts
+
+    load_kt_adapter_artifacts(model, adapter_path)
+
+
 def _setup_lora_tuning(
     config: "PretrainedConfig",
     model: "PreTrainedModel",
@@ -185,6 +191,8 @@ def _setup_lora_tuning(
             "revision": model_args.model_revision,
             "token": model_args.hf_hub_token,
         }
+        if model_args.use_kt:
+            init_kwargs["autocast_adapter_dtype"] = cast_trainable_params_to_fp32
 
         for adapter in adapter_to_merge:
             model: LoraModel = PeftModel.from_pretrained(model, adapter, **init_kwargs)
@@ -209,6 +217,12 @@ def _setup_lora_tuning(
                         model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs
                     )
 
+            if model_args.use_kt and not is_trainable:
+                adapter_path = model_args._kt_adapter_artifact_path
+                if adapter_path is None:
+                    raise RuntimeError("KT adapter artifacts were not resolved before model loading.")
+                _load_kt_inference_adapter_artifacts(model, adapter_path)
+
         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
@@ -264,7 +278,7 @@ def _setup_lora_tuning(
                 raise ValueError("KTransformers only supports LoRA finetuning.")
 
             peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, **peft_kwargs)
-            model = get_peft_model(model, peft_config)
+            model = get_peft_model(model, peft_config, autocast_adapter_dtype=cast_trainable_params_to_fp32)
         elif model_args.use_unsloth:
             if finetuning_args.finetuning_type == "oft":
                 raise ValueError("Unsloth is currently not supported for OFT.")

src/llamafactory/model/model_utils/checkpointing.py

@@ -40,6 +40,23 @@ if TYPE_CHECKING:
 logger = logging.get_logger(__name__)
 
 
+def _get_gradient_checkpointing_kwargs(model_args: "ModelArguments") -> dict[str, Any]:
+    r"""Build checkpoint kwargs through KT's public activation-context provider."""
+    if not model_args.use_kt:
+        return {"use_reentrant": model_args.use_reentrant_gc}
+
+    policy = model_args.get_kt_activation_policy()
+    if policy["gpu"] != "recompute":
+        return {"use_reentrant": False}
+
+    try:
+        from kt_kernel.sft import get_activation_checkpoint_context_fn
+    except (ImportError, ModuleNotFoundError) as exc:
+        raise RuntimeError("The installed kt-kernel does not provide the activation checkpoint context API.") from exc
+
+    return {"use_reentrant": False, "context_fn": get_activation_checkpoint_context_fn()}
+
+
 def get_unsloth_gradient_checkpointing_func() -> Callable:
     class UnslothGradientCheckpointing(torch.autograd.Function):
         r"""Saves VRAM by smartly offloading to RAM."""
@@ -172,7 +189,7 @@ def prepare_model_for_training(model: "PreTrainedModel", model_args: "ModelArgum
             )
             model.gradient_checkpointing_enable = MethodType(gradient_checkpointing_enable, model)
             model.gradient_checkpointing_enable(
-                gradient_checkpointing_kwargs={"use_reentrant": model_args.use_reentrant_gc}
+                gradient_checkpointing_kwargs=_get_gradient_checkpointing_kwargs(model_args)
             )
             setattr(model.config, "use_cache", False)  # turn off when gradient checkpointing is enabled
             logger.info_rank0("Gradient checkpointing enabled.")

src/llamafactory/model/model_utils/quantization.py

@@ -40,6 +40,10 @@ if TYPE_CHECKING:
 logger = logging.get_logger(__name__)
 
 
+def _uses_kt_non_expert_cache(model_args: "ModelArguments") -> bool:
+    return model_args.use_kt and bool(model_args.kt_non_expert_weight_path)
+
+
 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):
@@ -108,6 +112,13 @@ def configure_quantization(
             init_kwargs["ignore_mismatched_sizes"] = True
 
         if quant_method == QuantizationMethod.FP8:
+            if _uses_kt_non_expert_cache(model_args):
+                if model_args.quantization_bit is not None:
+                    raise ValueError("`quantization_bit` cannot be combined with KT weight caches.")
+
+                logger.info_rank0("Skipping source FP8 dequantization because KT weight caches are configured.")
+                return
+
             from transformers import FineGrainedFP8Config
 
             quant_config = FineGrainedFP8Config(dequantize=True)