ktransformers

This is a regression from PR #2181: prepare_decode_down_input() in kt-kernel/operators/amx/fp4-moe.hpp uses the optimized MXFP4 decode path even when __AVX512BF16__ is unavailable. On non-BF16 CPUs, down_ba_ is never populated, so the down projection reads uninitialized memory and produces garbage.

Fix by updating to the build containing PR fix(kt-kernel): restore MXFP4 decode fallback without AVX512 BF16, or gate the optimized path with __AVX512BF16__:

void prepare_decode_down_input(int expert_idx, int qlen) {
#if defined(__AVX512BF16__)
  if (qlen != 1 || config_.quant_config.group_size != 32 ||
      config_.intermediate_size % 32 != 0) {
    Base::prepare_decode_down_input(expert_idx, qlen);
    return;
  }
  assert(down_ba_[expert_idx]->natural_order);
#else
  Base::prepare_decode_down_input(expert_idx, qlen);
#endif
}

Then rebuild the kernel. This restores the base fallback for CPUs such as Xeon Platinum 8268/Cascade Lake that support AVX512F/VNNI but lack AVX512BF16.

This occurs in ktransformers 0.6.4 when PEFT/LoRA replaces the DeepSeek MoE gate with an nn.Linear, but _compute_routing still enters the deepseek_gate branch and treats raw logits as a (topk_ids, topk_weights, aux_loss) tuple.

Apply this one-line fix in kt_kernel/sft/layer.py:

# Before
if self.router_type == "deepseek_gate":

# After
if self.router_type == "deepseek_gate" and self._original_router is None:

This lets replaced Linear routers use the existing _original_router routing path, which performs softmax scoring and greedy top-k selection for DeepSeek-V2-Lite.

Also avoid lora_target: all in the LLaMA-Factory YAML, because it wraps the MoE router itself. Use the explicit target set:

lora_target: q_proj,kv_a_proj_with_mqa,kv_b_proj,o_proj,gate_proj,up_proj,down_proj

After the patch and config change, DeepSeek-V2-Lite LoRA SFT proceeds normally; the fix was verified with a 20-step training run and successful adapter export.

The model's config.json stores the correct value under rope_parameters (1000000), but transformers-kt 5.x exposes no top-level rope_theta, so the loader falls back to the default 10000. The incorrect RoPE base frequency causes perplexity to explode after roughly 4k tokens.

Affected setup reported with sglang-kt 0.7.0.post4, kt-kernel 0.7.0.post4, and transformers-kt 5.6.0.post5.

Fix the config/RoPE initialization to read the nested value:

if getattr(config, "rope_parameters", None):
    rope_theta = config.rope_parameters.get(
        "rope_theta", getattr(config, "rope_theta", 10000)
    )
else:
    rope_theta = getattr(config, "rope_theta", 10000)

Use this rope_theta when constructing the rotary embedding, then restart the server. The effective value must be 1000000, not 10000.

As a temporary local workaround, add the top-level value to the model's config.json:

import json, pathlib
p = pathlib.Path("/path/to/model/config.json")
c = json.loads(p.read_text())
c["rope_theta"] = c["rope_parameters"]["rope_theta"]
p.write_text(json.dumps(c, indent=2))

After the fix, long-context perplexity should remain approximately flat instead of rising sharply at positions 4096 and beyond.

This affects kt-kernel 0.6.4. Under FSDP2, LoRA-Expert parameters are DTensor shards, but save_lora_experts_to_adapter calls .data.cpu().clone() directly. That does not gather the shards, so safetensors receives a tensor with invalid storage and crashes after training, leaving only the standard PEFT LoRA weights in adapter_model.safetensors.

Patch kt_kernel/sft/lora.py around lines 797–799 to gather DTensors before saving:

def _materialise_for_save(t):
    try:
        from torch.distributed.tensor import DTensor
    except ImportError:
        try:
            from torch.distributed._tensor import DTensor
        except ImportError:
            DTensor = ()

    if DTensor and isinstance(t, DTensor):
        t = t.full_tensor()  # collective: must run on every FSDP rank

    return t.detach().cpu().clone()

state_dict[f"{base_key}.le_gate.weight"] = _materialise_for_save(expert.le_gate.weight)
state_dict[f"{base_key}.le_up.weight"]   = _materialise_for_save(expert.le_up.weight)
state_dict[f"{base_key}.le_down.weight"] = _materialise_for_save(expert.le_down.weight)

Do not call full_tensor() only on rank 0; it is a distributed collective and must be entered by all ranks. You may still guard the actual safetensors.torch.save_file call so only the main process writes the checkpoint.

After the fix, verify that adapter_model.safetensors contains le_gate, le_up, and le_down tensors and is larger than a run without LoRA Experts. As a temporary workaround, either disable kt_use_lora_experts or use a non-FSDP2/DDP configuration.

This occurs in ktransformers 0.6.4 with FSDP2 because FSDP2 wraps normal parameters as DTensor, while CPU-resident expert-LoRA parameters remain plain torch.Tensor. HF Trainer then passes a mixed list to _foreach_mul_. It happens in both real clipping (max_grad_norm > 0) and _get_grad_norm(), which calls clipping with inf just for logging, so setting max_grad_norm: 0.0 alone is insufficient.

Temporary workaround, with gradient clipping disabled:
1. In the training/YAML config, set:

max_grad_norm: 0.0

2. Patch the logging-only gradient-norm call before starting training:

from transformers import Trainer

_orig_get_grad_norm = Trainer._get_grad_norm

def _safe_get_grad_norm(self, model, grad_norm=None):
    try:
        return _orig_get_grad_norm(self, model, grad_norm)
    except RuntimeError as e:
        if "DTensor" in str(e):
            return None
        raise

Trainer._get_grad_norm = _safe_get_grad_norm

This allows LoRA SFT to run, but does not perform gradient clipping, so it is unsuitable for production-quality training.

Proper fix: register ktransformers expert-LoRA parameters as replicated DTensor parameters when FSDP2 is active, making model.parameters() type-homogeneous. Alternatively, exclude those expert-LoRA parameters from the Trainer/accelerate clipping list and clip them separately as plain tensors. Either approach handles both the clipping and gradient-norm logging call sites and restores max_grad_norm > 0.