[GH-ISSUE #2145] kt-sft + FSDP2: DTensor _foreach_mul_ crash when kt expert-LoRA params mix with DTensor lists in grad clipping #1288

New issue

Closed

opened 2026-10-29 14:50:53 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:50:53 +00:00

Owner

Originally created by @PUMA-SPz on GitHub (Aug 8, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2145

[BUG] kt-sft under FSDP2: expert-LoRA params stay plain `torch.Tensor` while FSDP2 params are `DTensor`, so any `_foreach_` op over `model.parameters()` fails

Component: kt-sft expert-LoRA parameter registration (interaction with FSDP2 + HF Trainer)
Version: kt-kernel / ktransformers 0.6.4, torch 2.9.1+cu128, transformers-kt 5.6.0,
accelerate-kt 1.14.0, LLaMA-Factory 0.9.6.dev0
Config: the shipped examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml
(also reproduced with num_processes: 1)

Symptom

RuntimeError: aten._foreach_mul_.Tensor: got mixed torch.Tensor and DTensor,
need to convert all torch.Tensor to DTensor before calling distributed operators!
  File "torch/nn/utils/clip_grad.py", line 174, in _clip_grads_with_norm_
  File "accelerate/accelerator.py", line 3030, in clip_grad_norm_
  File "transformers/trainer.py", line 2580, in _get_grad_norm

Why it is unavoidable today

kt keeps MoE experts on CPU (the entire point of the AMX backend).
Therefore DDP cannot be used — it rejects CPU-resident params outright:
module must have its parameters and buffers on device cuda:0 but found one of them on device: cpu
So FSDP2 is mandatory. FSDP2 converts model params to DTensor, but kt's expert-LoRA
params remain plain torch.Tensor.
Any _foreach_ op over the combined model.parameters() list then raises.

There are TWO call sites, and both must be handled

This is the part that makes it look unfixable — disabling one leaves the other:

Trainer._clip_grad_norm -> real clipping. Only runs when max_grad_norm > 0.
Trainer._get_grad_norm -> calls clip_grad_norm_(model.parameters(), float("inf"))
purely to log the gradient norm, and therefore fires even when clipping is disabled:

def _get_grad_norm(self, model, grad_norm=None):
    if grad_norm is None:
        # Compute norm without clipping (inf means no actual clipping happens)
        grad_norm = self.accelerator.clip_grad_norm_(model.parameters(), float("inf"))

Setting max_grad_norm: 0.0 alone still crashes (site 2). Patching site 2 alone still crashes
once clipping is on (site 1).

Current workaround (not a fix)

max_grad_norm: 0.0, and
wrap the _get_grad_norm call in try/except RuntimeError and return None on a DTensor error.

With both, LoRA SFT trains successfully: DeepSeek-V2-Lite, 20 steps in 474 s, ~12.5 s/step steady
state, loss 11.78 -> 10.81, 305 MB adapter written. But gradient clipping is disabled, which is
not acceptable for real training runs.

Proposed proper fix

Register kt's expert-LoRA parameters as DTensor (replicated placement is sufficient) when
running under FSDP2, so model.parameters() is type-homogeneous. That fixes both call sites at
once, restores gradient clipping, and should also unblock multi-GPU (currently only
num_processes: 1 has been validated here).

Alternative/complementary: have kt exclude its expert-LoRA params from the parameter list handed
to the Trainer/accelerator, and clip them separately.

Note for `transformers`

_get_grad_norm performing a distributed collective purely for logging is arguably fragile in
its own right — a try/except degrade-to-None there would make mixed-param setups survivable
regardless of what the backend does.

Originally created by @PUMA-SPz on GitHub (Aug 8, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2145 # [BUG] kt-sft under FSDP2: expert-LoRA params stay plain `torch.Tensor` while FSDP2 params are `DTensor`, so any `_foreach_` op over `model.parameters()` fails **Component:** kt-sft expert-LoRA parameter registration (interaction with FSDP2 + HF Trainer) **Version:** kt-kernel / ktransformers **0.6.4**, torch 2.9.1+cu128, transformers-kt 5.6.0, accelerate-kt 1.14.0, LLaMA-Factory 0.9.6.dev0 **Config:** the shipped `examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml` (also reproduced with `num_processes: 1`) ## Symptom ``` RuntimeError: aten._foreach_mul_.Tensor: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators! File "torch/nn/utils/clip_grad.py", line 174, in _clip_grads_with_norm_ File "accelerate/accelerator.py", line 3030, in clip_grad_norm_ File "transformers/trainer.py", line 2580, in _get_grad_norm ``` ## Why it is unavoidable today - kt keeps MoE experts on **CPU** (the entire point of the AMX backend). - Therefore **DDP cannot be used** — it rejects CPU-resident params outright: `module must have its parameters and buffers on device cuda:0 but found one of them on device: cpu` - So **FSDP2 is mandatory**. FSDP2 converts model params to `DTensor`, but kt's expert-LoRA params remain plain `torch.Tensor`. - Any `_foreach_` op over the combined `model.parameters()` list then raises. ## There are TWO call sites, and both must be handled This is the part that makes it look unfixable — disabling one leaves the other: 1. `Trainer._clip_grad_norm` -> real clipping. Only runs when `max_grad_norm > 0`. 2. `Trainer._get_grad_norm` -> calls `clip_grad_norm_(model.parameters(), float("inf"))` **purely to log the gradient norm**, and therefore fires **even when clipping is disabled**: ```python def _get_grad_norm(self, model, grad_norm=None): if grad_norm is None: # Compute norm without clipping (inf means no actual clipping happens) grad_norm = self.accelerator.clip_grad_norm_(model.parameters(), float("inf")) ``` Setting `max_grad_norm: 0.0` alone still crashes (site 2). Patching site 2 alone still crashes once clipping is on (site 1). ## Current workaround (not a fix) - `max_grad_norm: 0.0`, **and** - wrap the `_get_grad_norm` call in `try/except RuntimeError` and return `None` on a DTensor error. With both, LoRA SFT trains successfully: DeepSeek-V2-Lite, 20 steps in 474 s, ~12.5 s/step steady state, loss 11.78 -> 10.81, 305 MB adapter written. **But gradient clipping is disabled**, which is not acceptable for real training runs. ## Proposed proper fix Register kt's expert-LoRA parameters as **DTensor** (replicated placement is sufficient) when running under FSDP2, so `model.parameters()` is type-homogeneous. That fixes both call sites at once, restores gradient clipping, and should also unblock multi-GPU (currently only `num_processes: 1` has been validated here). Alternative/complementary: have kt exclude its expert-LoRA params from the parameter list handed to the Trainer/accelerator, and clip them separately. ## Note for `transformers` `_get_grad_norm` performing a distributed collective purely for **logging** is arguably fragile in its own right — a `try/except` degrade-to-`None` there would make mixed-param setups survivable regardless of what the backend does.

bsakaguchi closed this issue

2026-10-29 14:50:53 +00:00