diff --git a/docs/en/advanced/ktransformers.md b/docs/en/advanced/ktransformers.md new file mode 100644 index 000000000..89963014b --- /dev/null +++ b/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/`. diff --git a/docs/en/index.rst b/docs/en/index.rst index 6750c4994..407502481 100644 --- a/docs/en/index.rst +++ b/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 diff --git a/docs/zh/advanced/ktransformers.md b/docs/zh/advanced/ktransformers.md new file mode 100644 index 000000000..58b666ff3 --- /dev/null +++ b/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 +均包含非零更新。 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index f2cc60fcf..6cf685285 100644 --- a/docs/zh/index.rst +++ b/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 diff --git a/examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml b/examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml index d5548289a..32f2f4569 100644 --- a/examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml +++ b/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 diff --git a/examples/ktransformers/accelerate/fsdp2_kt_int4.yaml b/examples/ktransformers/accelerate/fsdp2_kt_int4.yaml index ec20e76c6..32f2f4569 100644 --- a/examples/ktransformers/accelerate/fsdp2_kt_int4.yaml +++ b/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 diff --git a/examples/ktransformers/accelerate/fsdp2_kt_int8.yaml b/examples/ktransformers/accelerate/fsdp2_kt_int8.yaml index 733f2c39a..32f2f4569 100644 --- a/examples/ktransformers/accelerate/fsdp2_kt_int8.yaml +++ b/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 diff --git a/examples/ktransformers/accelerate/fsdp2_kt_int8_1gpu.yaml b/examples/ktransformers/accelerate/fsdp2_kt_int8_1gpu.yaml index a76ceae3e..8cec01593 100644 --- a/examples/ktransformers/accelerate/fsdp2_kt_int8_1gpu.yaml +++ b/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 diff --git a/examples/ktransformers/accelerate/fsdp2_kt_int8_8gpu.yaml b/examples/ktransformers/accelerate/fsdp2_kt_int8_8gpu.yaml index 7ba10cd48..ef5a65623 100644 --- a/examples/ktransformers/accelerate/fsdp2_kt_int8_8gpu.yaml +++ b/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 diff --git a/examples/ktransformers/train_lora/deepseek_v3_int8_lora_sft_kt.yaml b/examples/ktransformers/train_lora/deepseek_v3_int8_lora_sft_kt.yaml new file mode 100644 index 000000000..236601e5c --- /dev/null +++ b/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 diff --git a/examples/ktransformers/train_lora/deepseek_v3_lora_sft_kt.yaml b/examples/ktransformers/train_lora/deepseek_v3_lora_sft_kt.yaml index c3282391b..04c3ed34d 100644 --- a/examples/ktransformers/train_lora/deepseek_v3_lora_sft_kt.yaml +++ b/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. -# 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. diff --git a/examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml b/examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml index 40f75442c..fbba31c5b 100644 --- a/examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml +++ b/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. -# 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. diff --git a/src/llamafactory/hparams/model_args.py b/src/llamafactory/hparams/model_args.py index 4212a46df..e5ee103fd 100644 --- a/src/llamafactory/hparams/model_args.py +++ b/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": diff --git a/src/llamafactory/hparams/parser.py b/src/llamafactory/hparams/parser.py index c22c1e375..1f686d641 100644 --- a/src/llamafactory/hparams/parser.py +++ b/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) diff --git a/src/llamafactory/model/adapter.py b/src/llamafactory/model/adapter.py index 3872d3dd4..ce4a8d1b3 100644 --- a/src/llamafactory/model/adapter.py +++ b/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.") diff --git a/src/llamafactory/model/model_utils/checkpointing.py b/src/llamafactory/model/model_utils/checkpointing.py index 0ba7ec96a..6620bd077 100644 --- a/src/llamafactory/model/model_utils/checkpointing.py +++ b/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.") diff --git a/src/llamafactory/model/model_utils/quantization.py b/src/llamafactory/model/model_utils/quantization.py index 6ad607db3..745237553 100644 --- a/src/llamafactory/model/model_utils/quantization.py +++ b/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)