423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2026-08-17 05:25:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
yyj
2026-08-13 20:43:15 +08:00
committed by GitHub
parent 199b8873d7
commit bc4b42cefc
17 changed files with 573 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
docs/en/advanced/ktransformers.md Normal file
View File

@@ -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/`.

1
docs/en/index.rst
View File

@@ -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

119
docs/zh/advanced/ktransformers.md Normal file
View File

@@ -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
均包含非零更新。

1
docs/zh/index.rst
View File

@@ -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

10
examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml
View File

@@ -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

10
examples/ktransformers/accelerate/fsdp2_kt_int4.yaml
View File

@@ -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

10
examples/ktransformers/accelerate/fsdp2_kt_int8.yaml
View File

@@ -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

10
examples/ktransformers/accelerate/fsdp2_kt_int8_1gpu.yaml
View File

@@ -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

10
examples/ktransformers/accelerate/fsdp2_kt_int8_8gpu.yaml
View File

@@ -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

54
examples/ktransformers/train_lora/deepseek_v3_int8_lora_sft_kt.yaml Normal file
View File

@@ -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

13
examples/ktransformers/train_lora/deepseek_v3_lora_sft_kt.yaml
View File

@@ -40,6 +40,13 @@ resume_from_checkpoint: null
### ktransformers
use_kt: true
# Pair with fsdp2_kt_bf16.yaml for original BF16 checkpoints.
# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
# kt_weight_path: /path/to/DeepSeek-V3-AMXINT8
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
kt_share_backward_bb: true
# The Accelerate YAML contains FSDP settings only. KT has a single configuration owner here.

15
examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml
View File

@@ -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.
# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
# Pair the 397B path with fsdp2_kt_int8.yaml, tune cutoff_len to prepared weights and GPU memory.
# kt_weight_path: /path/to/Qwen3.5-MoE-AMXINT8
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
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.

236
src/llamafactory/hparams/model_args.py
View File

@@ -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":

56
src/llamafactory/hparams/parser.py
View File

@@ -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)

16
src/llamafactory/model/adapter.py
View File

@@ -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.")

19
src/llamafactory/model/model_utils/checkpointing.py
View File

@@ -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.")

11
src/llamafactory/model/model_utils/quantization.py
View File

@@ -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)