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

docs/en/advanced/ktransformers.md

@@ -0,0 +1,83 @@
+# KTransformers LoRA SFT
+
+KTransformers (KT) executes routed MoE experts on CPU while LLaMA-Factory remains responsible for data,
+LoRA arguments, and the training entry point. The production scope is routed-BF16 and routed-INT8 LoRA.
+
+KT has one user configuration source: the training YAML. Accelerate YAML contains FSDP2 settings only.
+LLaMA-Factory derives LoRA rank, alpha, dropout, activation policy, and local runtime capacity.
+
+```yaml
+finetuning_type: lora
+lora_rank: 8
+lora_alpha: 16
+lora_target: all
+
+use_kt: true
+disable_gradient_checkpointing: false
+kt_cpu_activation: retain
+kt_config:
+  kt_expert_weight_format: bf16
+  kt_backend: AMXBF16
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+```
+
+Routed INT8 additionally requires matching expert and BF16 non-expert artifacts:
+
+```yaml
+kt_weight_path: /abs/path/to/routed-int8-experts
+kt_non_expert_weight_path: /abs/path/to/bf16-non-expert-cache
+kt_config:
+  kt_expert_weight_format: int8
+  kt_backend: auto
+  kt_weight_lifecycle: persistent
+```
+
+Launch the standard training entry point through Accelerate:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 accelerate launch \
+  --config_file examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml \
+  src/train.py examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml
+```
+
+## Load a saved adapter
+
+Use a local, complete KT adapter directory for chat or evaluation. Repeat the training LoRA shape (`finetuning_type`,
+`lora_rank`, `lora_alpha`, and `lora_dropout`) and the KT base-weight settings. In particular, routed INT8 loading
+must use the same `kt_weight_path` and `kt_non_expert_weight_path` as training.
+
+```yaml
+model_name_or_path: /abs/path/to/base-model
+adapter_name_or_path: /abs/path/to/output/checkpoint-300
+finetuning_type: lora
+lora_rank: 8
+lora_alpha: 16
+lora_dropout: 0.0
+
+use_kt: true
+kt_cpu_activation: retain
+kt_config:
+  kt_expert_weight_format: bf16
+  kt_backend: AMXBF16
+  kt_num_threads: 96
+```
+
+```bash
+llamafactory-cli chat path/to/kt_adapter_infer.yaml
+llamafactory-cli eval path/to/kt_adapter_eval.yaml
+```
+
+The directory must contain the standard PEFT adapter files and, when fused routed-expert LoRA is used,
+`fused_expert_lora.safetensors` plus `kt_adapter_manifest.json`. LLaMA-Factory first loads the standard PEFT
+adapter, then KT validates and restores the fused artifact. `adapter_folder` may select a local subdirectory;
+paths outside the adapter root and Hub adapter IDs fail before model loading. Download a Hub bundle locally first.
+
+For training resume, keep the original training YAML and use `resume_from_checkpoint`. The optimizer checkpoint
+currently requires the same distributed world size. Missing, tampered, or mismatched artifacts fail closed instead
+of falling back to the source checkpoint.
+
+Do not combine KT with a second Transformers/FSDP checkpoint wrapper or Unsloth GC, and do not put `kt_config`
+in the Accelerate YAML. See the BF16 and INT8 examples under `examples/ktransformers/train_lora/`.

docs/en/index.rst

@@ -34,6 +34,7 @@ LlamaFactory Docs
 
    advanced/lora-and-quantization/lora
    advanced/lora-and-quantization/quantization
+   advanced/ktransformers
    advanced/distributed/fsdp
    advanced/distributed/deepspeed
    advanced/distributed/parallel-dp-tp-ep-sp-cp

docs/zh/advanced/ktransformers.md

@@ -0,0 +1,119 @@
+# KTransformers LoRA SFT
+
+KTransformers（KT）将 MoE routed experts 放在 CPU 执行，LLaMA-Factory 继续负责数据、LoRA 参数和训练入口。
+当前生产范围是 routed-BF16 LoRA 与 routed-INT8 LoRA；Accelerate 配置只负责 FSDP2，不再保存 KT 参数。
+
+## 安装检查
+
+必须同时安装带 KT 公共接口的 `ktransformers`、`transformers-kt` 和 `accelerate-kt`。启动前可检查：
+
+```bash
+python - <<'PY'
+from accelerate import Accelerator
+from kt_kernel.sft import resolve_kt_pretrained_artifacts
+from transformers import TrainingArguments
+
+assert hasattr(TrainingArguments, "update_kt_config")
+assert "adapter_only" in __import__("inspect").signature(Accelerator.get_state_dict).parameters
+print(resolve_kt_pretrained_artifacts)
+PY
+```
+
+## 配置
+
+KT 只有一个用户配置源：训练 YAML。LoRA rank、alpha、dropout 和 runtime capacity 由 LLaMA-Factory
+标准字段派生；不要在 `kt_config` 中重复填写。
+
+BF16 示例：
+
+```yaml
+finetuning_type: lora
+lora_rank: 8
+lora_alpha: 16
+lora_target: all
+
+use_kt: true
+disable_gradient_checkpointing: false
+kt_cpu_activation: retain
+kt_config:
+  kt_expert_weight_format: bf16
+  kt_backend: AMXBF16
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+```
+
+INT8 还需要相互匹配的 routed expert 与 BF16 non-expert cache：
+
+```yaml
+kt_weight_path: /abs/path/to/routed-int8-experts
+kt_non_expert_weight_path: /abs/path/to/bf16-non-expert-cache
+kt_config:
+  kt_expert_weight_format: int8
+  kt_backend: auto
+  kt_weight_lifecycle: persistent
+```
+
+完整配置见：
+
+- `examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml`
+- `examples/ktransformers/train_lora/deepseek_v3_int8_lora_sft_kt.yaml`
+
+Activation 策略：
+
+| `disable_gradient_checkpointing` | `kt_cpu_activation` | CPU / GPU |
+| --- | --- | --- |
+| `false` | `recompute` 或省略 | recompute / recompute |
+| `false` | `retain` | retain / recompute |
+| `true` | `retain` 或省略 | retain / retain |
+| `true` | `recompute` | 不支持，启动前报错 |
+
+## 启动与复用
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 accelerate launch \
+  --config_file examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml \
+  src/train.py examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml
+```
+
+输出 adapter 同时包含 standard PEFT 与 fused expert LoRA。
+
+## 新进程加载
+
+对话或评测必须使用本地的完整 KT adapter 目录，并重复训练时的 LoRA 形状配置：`finetuning_type`、
+`lora_rank`、`lora_alpha`、`lora_dropout`，以及相同的 KT base weight 配置。routed INT8 尤其要沿用训练时
+的 `kt_weight_path` 和 `kt_non_expert_weight_path`。
+
+```yaml
+model_name_or_path: /abs/path/to/base-model
+adapter_name_or_path: /abs/path/to/output/checkpoint-300
+finetuning_type: lora
+lora_rank: 8
+lora_alpha: 16
+lora_dropout: 0.0
+
+use_kt: true
+kt_cpu_activation: retain
+kt_config:
+  kt_expert_weight_format: bf16
+  kt_backend: AMXBF16
+  kt_num_threads: 96
+```
+
+```bash
+llamafactory-cli chat path/to/kt_adapter_infer.yaml
+llamafactory-cli eval path/to/kt_adapter_eval.yaml
+```
+
+目录必须包含 standard PEFT adapter 文件；使用 fused routed-expert LoRA 时，还必须包含
+`fused_expert_lora.safetensors` 和 `kt_adapter_manifest.json`。LLaMA-Factory 先加载 standard PEFT，随后由
+KT 校验并恢复 fused artifact。`adapter_folder` 可以选择本地子目录；越出 adapter 根目录的路径和 Hub
+adapter ID 会在加载模型前报错，Hub bundle 需要先完整下载到本地。
+
+续训应保留原训练 YAML，并使用 `resume_from_checkpoint`。分布式 optimizer checkpoint 暂要求相同 world
+size。artifact 缺失、hash 不匹配或来源模型不一致时会直接失败，不会退回源 checkpoint。
+
+不要同时启用 Transformers/FSDP activation checkpointing、Unsloth GC，也不要把 `kt_config` 放入
+Accelerate YAML。每次训练都应确认 loss/grad finite、base model 未修改，并验证 standard/router/fused LoRA
+均包含非零更新。

docs/zh/index.rst

@@ -34,6 +34,7 @@ LlamaFactory 文档
 
    advanced/lora-and-quantization/lora
    advanced/lora-and-quantization/quantization
+   advanced/ktransformers
    advanced/distributed/fsdp
    advanced/distributed/deepspeed
    advanced/distributed/parallel-dp-tp-ep-sp-cp

examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml

@@ -13,13 +13,3 @@ num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
 rdzv_backend: static
 same_network: true
 use_cpu: false
-
-kt_config:
-  enabled: true
-  kt_backend: AMXBF16  # Use with original BF16 expert weights.
-  kt_num_threads: 96
-  kt_tp_enabled: true
-  kt_threadpool_count: 2
-  kt_max_cache_depth: 2
-  kt_share_backward_bb: true
-  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int4.yaml

@@ -13,13 +13,3 @@ num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
 rdzv_backend: static
 same_network: true
 use_cpu: false
-
-kt_config:
-  enabled: true
-  kt_backend: AMXINT4  # Use with online-converted INT4 expert weights
-  kt_num_threads: 96
-  kt_tp_enabled: true
-  kt_threadpool_count: 2
-  kt_max_cache_depth: 2
-  kt_share_backward_bb: true
-  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int8.yaml

@@ -13,13 +13,3 @@ num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
 rdzv_backend: static
 same_network: true
 use_cpu: false
-
-kt_config:
-  enabled: true
-  kt_backend: AMXINT8  # Use with online-converted INT8 expert weights
-  kt_num_threads: 96
-  kt_tp_enabled: true
-  kt_threadpool_count: 2
-  kt_max_cache_depth: 2
-  kt_share_backward_bb: true
-  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int8_1gpu.yaml

@@ -13,13 +13,3 @@ num_processes: 1 # Adjust based on your GPU count; 1 is suitable for 1 GPU
 rdzv_backend: static
 same_network: true
 use_cpu: false
-
-kt_config:
-  enabled: true
-  kt_backend: AMXINT8  # Use with online-converted INT8 expert weights
-  kt_num_threads: 96
-  kt_tp_enabled: true
-  kt_threadpool_count: 2
-  kt_max_cache_depth: 2
-  kt_share_backward_bb: true
-  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int8_8gpu.yaml

@@ -13,13 +13,3 @@ num_processes: 8 # Adjust based on your GPU count; 8 is suitable for 8 GPUs
 rdzv_backend: static
 same_network: true
 use_cpu: false
-
-kt_config:
-  enabled: true
-  kt_backend: AMXINT8  # Use with online-converted INT8 expert weights
-  kt_num_threads: 96
-  kt_tp_enabled: true
-  kt_threadpool_count: 2
-  kt_max_cache_depth: 2
-  kt_share_backward_bb: true
-  lora_rank: 8

examples/ktransformers/train_lora/deepseek_v3_int8_lora_sft_kt.yaml

@@ -0,0 +1,54 @@
+### model
+model_name_or_path: /path/to/DeepSeek-V3.1-source
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_rank: 8
+lora_alpha: 16
+lora_target: all
+
+### dataset
+dataset: identity, alpaca_en_demo
+template: deepseek3
+cutoff_len: 2048
+max_samples: 100000
+overwrite_cache: true
+preprocessing_num_workers: 16
+dataloader_num_workers: 4
+
+### output
+output_dir: saves/KT_FT_deepseekV3_int8
+logging_steps: 10
+save_steps: 500
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1.0e-4
+num_train_epochs: 3.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 180000000
+
+### ktransformers
+use_kt: true
+kt_cpu_activation: retain
+kt_weight_path: /path/to/routed-int8-experts
+kt_non_expert_weight_path: /path/to/bf16-non-expert-cache
+kt_config:
+  kt_expert_weight_format: int8
+  kt_backend: auto
+  kt_weight_lifecycle: persistent
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_share_backward_bb: true

examples/ktransformers/train_lora/deepseek_v3_lora_sft_kt.yaml

@@ -40,6 +40,13 @@ resume_from_checkpoint: null
 
 ### ktransformers
 use_kt: true
-# Pair with fsdp2_kt_bf16.yaml for original BF16 checkpoints.
+kt_cpu_activation: retain
-# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
+kt_config:
-# kt_weight_path: /path/to/DeepSeek-V3-AMXINT8
+  kt_expert_weight_format: bf16
+  kt_backend: AMXBF16
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_share_backward_bb: true
+# The Accelerate YAML contains FSDP settings only. KT has a single configuration owner here.

examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml

@@ -40,7 +40,14 @@ resume_from_checkpoint: null
 
 ### ktransformers
 use_kt: true
-# For original BF16 checkpoints, start with examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml.
+kt_cpu_activation: retain
-# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
+kt_config:
-# Pair the 397B path with fsdp2_kt_int8.yaml, tune cutoff_len to prepared weights and GPU memory.
+  kt_expert_weight_format: bf16
-# kt_weight_path: /path/to/Qwen3.5-MoE-AMXINT8
+  kt_backend: AMXBF16
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_model_max_length: 2176  # Includes the text-only template's dummy-image tokens.
+  kt_share_backward_bb: true
+# The Accelerate YAML contains FSDP settings only. KT has a single configuration owner here.

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]:
+    _KT_DERIVED_KEYS = frozenset(
-        r"""Build KT config values from LLaMA-Factory model and LoRA arguments."""
+        {
-        kt_config = {
+            "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": model_max_length,
+                "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)
+        self.configure_kt_checkpointing(training_args)
-        env_mapping = {
+        kt_config = self.get_kt_config_dict(
-            "kt_weight_path": "ACCELERATE_KT_WEIGHT_PATH",
+            finetuning_args,
-            "kt_expert_checkpoint_path": "ACCELERATE_KT_EXPERT_CHECKPOINT_PATH",
+            model_max_length,
-            "kt_model_max_length": "ACCELERATE_KT_MODEL_MAX_LENGTH",
+            self._get_advanced_kt_config(training_args),
-            "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
         )
-        if gc_enabled:
+        update_kt_config = getattr(training_args, "update_kt_config", None)
-            hf_kt._kt_config.setdefault("kt_share_cache_pool", True)
+        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)