Kimi-K2.6 — transformers 5.x compatibility patches + abliterix run scaffolding
This artifact documents the engineering work to make moonshotai/Kimi-K2.6 loadable + runnable under modern transformers~=5.6 and to integrate it with abliterix for MoE-routing-aware abliteration. The original work was done on a p5.48xlarge (8× H100 80GB SXM, NVSwitch). All abliteration prerequisites (steering vectors, safety-expert profiling) ran end-to-end on H100. The Optuna trial loop hit CUDA OOM at trial 0 on H100 with abliterix v1.4.0 — bigger hardware (or vLLM tensor parallelism, or abliterix v1.6.0's bf16 dequant cache patch) is required.
v2 update (2026-05-06): added two more modeling patches required for SFT, the Track-2 SFT pipeline (which is currently in-flight on Modal), the abliterix v1.6.0 dequant-cache fix, and a full negative-results table for the activation-engineering technique class on Kimi-K2.6 (Tracks 1/3/4). See results/.
What's in here
patches/ ← Kimi modeling + config patches (drop-in for /mnt/<MODEL>)
modeling_kimi_k25.py patched (7 transformers-5.x patches)
modeling_deepseek.py patched
config.json patched (vision_config._attn_implementation = eager)
configuration_*.py unchanged (here for completeness)
diffs/*.diff unified diffs against upstream
training_mode_fixes.md + 2 patches needed only for SFT (gate + MoE eval mode)
abliterix/ ← three patched abliterix files (replace in venv site-packages)
cli.py + non-interactive auto-export at end of run
core/engine.py + Kimi wrapper transformer_layers path
eval/detector.py + max_tokens reasoning headroom for thinking judges
abliterix-v1.6.0-notes.md + bf16 dequant cache patch (steering.py) — fixes the H100 OOM
configs/
kimi_k2_6.toml abliterix config: bs=4, eager attn, manual device_map, MLA-disabled
scripts/
test_kimi_bundled.py **fast (3-second) reproducer**: tiny synthetic Kimi using bundled
modeling code on CPU. Exercises forward + generate(use_cache=True)
+ batch hidden_states. Use this to validate any modeling patch
before the 7-min real-model load.
validate_kimi_load.py full-model validation on 8× H100 (~10 min)
export_best.py post-run export script (untested; abliterix's auto-export in
the patched cli.py is the primary path)
results/ ← Tracks 1/3/4 negative results (full prompt × response tables)
track1_gatebreaker_results.json GateBreaker τ=2 (115K safety neurons) — refusal preserved
track3_stack_results.json GateBreaker τ=1 + residual rank-1 stacked — refusal preserved
track4_dose_results.json dose-response sweep, residual strength 2.0/3.0 — broken at 3.0
modal/ ← end-to-end Modal pipeline (alternative path to bare-metal)
modal_track2.py SFT + abliterix Modal entrypoints with all 9 patches inline
TRACK5_ABLITERIX_MODAL_HANDOFF.md full design doc (patches with citations, decision log)
The nine Kimi modeling patches
Each patch is small and defensive (uses try/except or hasattr checks so the same file works on transformers 4.x or 5.x).
is_torch_fx_available shim (modeling_deepseek.py:47) — symbol removed from transformers.utils.import_utils in 5.0; provide a lambda: False fallback.
MoonViT3dEncoder.use_deterministic_attn = False (modeling_kimi_k25.py:573) — the encoder reads this attribute on self but never sets it; default to False.
tie_weights(*args, **kwargs) (modeling_kimi_k25.py:880) — transformers 5.x calls tie_weights(recompute_mapping=False), the original signature didn't accept kwargs. Forward kwargs and fall back on TypeError.
apply_rotary_pos_emb cos/sin slice (modeling_deepseek.py:377) — transformers 5.x cache contract: position_ids may span the full kv_seq_len during incremental generation; slice cos/sin to q.shape[-2] so q_pe/k_pe stay shape-correct in the assignment at line 806.
DynamicCache.from_legacy_cache removed (modeling_deepseek.py:1416) — fall back to a fresh DynamicCache() when the helper isn't available; also harden get_usable_length call.
Cache.to_legacy_cache removed (modeling_deepseek.py:1495) — pass the Cache object through directly when the method is unavailable.
Explicit GenerationMixin inheritance on both DeepseekV3ForCausalLM (modeling_deepseek.py:1494) and KimiK25ForConditionalGeneration (modeling_kimi_k25.py:834) — transformers 5.x no longer auto-inherits, so .generate() would otherwise raise.
Patches 8–9: training-mode fixes (required for SFT, not abliterix inference)
These bite ANY LoRA SFT pipeline that calls model.train(). They're not in the bundled modeling_*.py because they're not safe to apply unconditionally — they force the affected modules to permanent eval mode, which is correct for inference + LoRA-on-attention training but would silently degrade quality if you ever wanted to train the routed experts. Apply at runtime via the snippet in patches/training_mode_fixes.md.
MoEGate.forwardassert not self.training (modeling_deepseek.py:468, the noaux_tc topk path) — router was designed inference-only; assert fires on first training step. Fix: monkey-patch MoEGate.train() to a no-op + eval(). Gate weights aren't a LoRA target so they don't need training mode anyway.
DeepseekV3MoE.forward only-assigns-y-in-inference branch (modeling_deepseek.py:551) — if not self.training: y = self.moe_infer(...); line 555 then references y → NameError in train mode. Fix: same eval-mode override on DeepseekV3MoE. Forces the moe_infer path (which is @torch.no_grad) for routed experts; gradients still flow into shared_experts because that submodule is invoked separately on identity outside the no_grad block.
Plus one config edit: vision_config._attn_implementation set from flash_attention_2 to eager in config.json (Kimi doesn't actually need flash-attn for vision and we don't have it built for our torch/CUDA combo).
The abliterix patches (v1.4.0 baseline)
core/engine.py — transformer_layers — added a fallback branch m.language_model.model.layers for the KimiK25ForConditionalGeneration wrapper layout (Kimi inverts the standard Mistral3/Qwen-VL m.model.language_model.layers path).
core/engine.py — resolve_model_class — when auto_map advertises AutoModelForCausalLM but not AutoModelForImageTextToText, prefer AutoModelForCausalLM. Required because Kimi's custom wrapper class isn't in transformers' built-in VL registry.
eval/detector.py — judge max_tokens — bumped from len(uncached) * 5 + 50 to + 1500 so thinking models (Gemini-3.x-pro/flash, Claude-3.7-thinking, o-series) have reasoning headroom before the JSON labels get emitted. Without this the judge truncates and JSON parsing fails.
cli.py — non-interactive auto-export — the upstream non-interactive flow prints "finished" and returns without saving the abliterated weights. Patched to:
pick the Pareto-best trial (min refusals, tiebreak min KL),
copy modeling files so the saved repo is self-contained.
Output dir is configurable via ABLITERIX_AUTO_EXPORT_DIR env var (default /mnt/nvme3/Kimi-K2.6-abliterated-bf16).
cli.py — _try exception surfacing — abliterix's auto-batch-size probe silently returned None on RuntimeError/CUDA OOM, hiding the actual failure. Patched to print the exception before returning. Critical when a single iteration costs 7+ minutes of model load.
abliterix v1.6.0 — additional patch for H100 (see abliterix-v1.6.0-notes.md)
abliterix v1.6.0 added a bf16 dequant cache to core/steering.py. The v1.4.0 cache used torch.float32 for the projection math, and on a 1T MoE with 244 LoRA modules × ~200 MB f32 cache, this exceeds GPU 2's budget on 8× H100 (it's tighter than GPU 0/7 because it carries both layers AND profiling buffers). This is THE error that crashed trial 0 in the H100 run. v1.6.0's .to(torch.bfloat16) halves the cache to ~24 GB, plus a v.to(W.dtype) cast before (v @ W) for downstream dtype consistency. We never re-tested v1.6.0 on H100 (jumped straight to B200 on Modal) — likely fits, worth trying.
Two validated execution paths
Path A: bare-metal 8× H200 / B200 + abliterix CLI
Use this if you have direct hardware. The 7+5 patches above + configs/kimi_k2_6.toml + the standard abliterix invocation:
Spin up a host with ≥ 1× B200 192GB or ≥ 1× H200 141GB worth of usable VRAM beyond the 4-bit weights (~125 GB extra to absorb LoRA + KV cache for 384-expert MoE × 60 layers).
Practical recipe: 8× B200 192GB (1.5 TB total) or 8× H200 141GB (1.13 TB total).
Install nvidia driver-580-server-open + fabricmanager (must be matched version), CUDA 12.4+, Python 3.10, then:
Note: torch 2.6.0+ (not 2.5.1) is required if you want HuggingFace Trainer's resume-from-checkpoint — see "torch.load CVE-2025-32434" below.
Download bullerwins/Kimi-K2.6-bf16 (2.05 TB, 64 shards) → e.g. /mnt/data/Kimi-K2.6-bf16.
Apply patches:
Copy patches/{modeling_kimi_k25,modeling_deepseek}.py and patches/config.json over the model dir.
Copy abliterix/{cli,core/engine,eval/detector}.py over the venv's site-packages/abliterix/ files.
For SFT pipelines also: apply patches 8–9 at runtime (snippet in patches/training_mode_fixes.md).
Run the fast reproducer first (scripts/test_kimi_bundled.py) to confirm the modeling patches still apply cleanly. Should print === ALL CHECKS PASSED === in ~3 seconds.
Total wall time end-to-end on B200×8: ~2 hours load + 12 trials × ~50 min = ~12 h.
Path B: Modal serverless (8× H200 or 8× B200)
Use this if you don't have bare-metal access. modal/modal_track2.py is the working pipeline that has all 9 modeling patches + 5 abliterix patches inlined as runtime monkey-patches; the patched abliterix v1.6.0 is pip install -e'd from a local fork.
bash
1modal run --detach modal_track2.py::train_sft --epochs 12modal run --detach modal_track2.py::run_abliterix --num-trials 12
Both functions auto-push their adapter to HF on completion. See modal/TRACK5_ABLITERIX_MODAL_HANDOFF.md for the full design rationale, decision log, and recovery procedures.
Pinned versions that work end-to-end on Modal H200:8 (validated with full SFT loss curve, 1 epoch, 5,768 train + 289 val rows, max_seq_len=4096):
torch==2.6.0+cu124 # NOT 2.5.1 — see torch.load gating below
transformers==5.6.2
accelerate==1.13.0
peft==0.14.0
trl==0.13.0 # caveat: needs `text` column not `messages` auto-detect
datasets==3.0.0
bitsandbytes==0.49.2
LoRA targets: ["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "o_proj", "shared_experts.{gate,up,down}_proj"] → 87,569,408 trainable params / 1,026,966,945,776 total = 0.0085% trainable. Routed experts are NOT LoRA targets — they're bnb_4bit-packed and not directly editable, plus the moe_infer path is @torch.no_grad.
trl 0.13 gotcha: SFTTrainer expects a text column. Pre-render messages via tok.apply_chat_template(...) before passing to the trainer — auto-detection from messages is trl 0.16+.
Kimi tokenizer gotcha:chat_template.jinja ships as a separate file (not embedded in tokenizer_config.json). If your from_pretrained loader doesn't auto-pick it up, set tok.chat_template = open(...)/chat_template.jinja).read() manually.
torch.load gating (CVE-2025-32434)
transformers ≥ 5.6 enforces the CVE-2025-32434 mitigation: torch.load is blocked on torch < 2.6, even with weights_only=True. This affects HuggingFace Trainer's resume_from_checkpoint=... path because optimizer state (optimizer.pt) is saved as a torch pickle. Bake torch==2.6.0+ into the image if you want resume-from-checkpoint to work; we discovered this the hard way mid-flight on Modal SFT and bumped the image.
Inference / abliterix doesn't hit this — only Trainer resume.
Hardware sizing — empirical floor
Working set during steady-state training/abliteration:
Component
Size
Kimi-K2.6 NF4 model
506 GB
bnb dequant cache (bf16, abliterix v1.6.0)
~24 GB
244 LoRA adapters (rank 16)
~1 GB
Activations + KV cache (batch=1, seq=4096)
~50–80 GB
Total practical floor
~620 GB
Hardware
Total VRAM
Verdict
8× H100 80GB
640 GB
Just under threshold; OOMs on abliterix v1.4.0 dequant cache. v1.6.0 fix may make it fit, untested.
8× H200 141GB
1,128 GB
Comfortable. Validated end-to-end SFT (loaded in 31 min on warm Modal Volume; 71 s/step training).
8× B200 192GB
1,536 GB
Very comfortable. Validated for abliterix v1.6.0 on Modal.
What's known to NOT work
8× H100 80GB SXM at bs=4 with abliterix v1.4.0 — OOMs at trial 0 inside apply_steering when LoRA adapters are added on top of NF4 weights. Specifically GPU 1 (or 3, alternating) hits 79.12 / 79.18 GiB. The model alone takes ~64 GB/GPU at NF4, which leaves only ~15 GiB for activations + LoRA adapters across 60 layers × 23K down_proj instances. The abliterix v1.6.0 bf16 dequant cache patch should help here — untested.
Downgrading transformers to 4.56.x — abliterix v1.4.0 metadata pins transformers~=5.3, and 4.56's huggingface-hub<1.0 requirement collides with the kernels package's >=1.0 requirement. Dependency-wedged.
Loading moonshotai/Kimi-K2.6 (compressed-tensors int4) directly on 8× H100 80GB or 8× H200 141GB — compressed-tensors 0.15.0.1 decompresses to BF16 on first forward pass; the BF16-equivalent peaks at ~1.7 TB which exceeds aggregate VRAM. Workaround = go via bullerwins/Kimi-K2.6-bf16 + bnb_4bit re-quant.
gpu_kwargs={"timeout": 7*24*60*60} on Modal @app.function — Modal hard-caps function timeout at 24h regardless of plan. For multi-day work, design for resume-from-checkpoint and chain function invocations.
Negative results — activation-engineering attacks on Kimi-K2.6
Issue #221 asked whether MoE-aware abliteration techniques like GateBreaker (arXiv:2512.21008) and polyhedral-cone refusal subtraction (arXiv:2502.17420) generalize to Kimi-K2.6. Tracks 1/3/4 in results/ give the conclusive answer: no operating point exists where refusal breaks while coherence holds.
Track
Method
Strength
Outcome
1
GateBreaker τ=2 (115K safety neurons in expert gate_proj/up_proj zeroed)
aggressive
Refusal preserved on harmful prompts; benign quality preserved
Refusal preserved (initial "softening" was a max_new_tokens=100 truncation artifact)
3-D
Stacked: GateBreaker τ=1 + residual subtraction
strength=1.0
Same as residual alone (residual dominates at 1.0)
4
Residual subtraction dose sweep
strength=2.0
Same as baseline — clean refusals
4
Residual subtraction dose sweep
strength=3.0
Model collapses — outputs degenerate "777..." token loops on every prompt, including benign
4
Stacked GateBreaker + residual
strength=2.0
Verbose policy-deliberation but ends on refusal
4
Stacked GateBreaker + residual
strength=3.0
Model collapses
Full per-prompt response tables in results/track{1,3,4}_*_results.json. The conclusion holds across the entire activation-engineering technique class on Kimi-K2.6's 384-expert MoE — confirming hamsaOmar's K2.5 finding extends to K2.6. Path forward: gradient-driven approaches (LoRA SFT, DPO, abliterix v1.6.0's Optuna trial loop). The Modal pipeline in modal/ runs both.
Provenance
Base weights: bullerwins/Kimi-K2.6-bf16 (community decompression of moonshotai/Kimi-K2.6 from int4-compressed-tensors → BF16, 2.05 TB / 64 shards).
