audnai/Kimi-K2.6-transformers-5.x-patches

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

Patches 1–7: load + inference (transformers 5.x compat)

1. is_torch_fx_available shim (modeling_deepseek.py:47) — symbol removed from transformers.utils.import_utils in 5.0; provide a lambda: False fallback.
2. MoonViT3dEncoder.use_deterministic_attn = False (modeling_kimi_k25.py:573) — the encoder reads this attribute on self but never sets it; default to False.
3. 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.
4. 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.
5. 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.
6. Cache.to_legacy_cache removed (modeling_deepseek.py:1495) — pass the Cache object through directly when the method is unavailable.
7. 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.

8. MoEGate.forward assert 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.
9. 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)

1. 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).

2. 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.

3. 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.

4. 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),
   • call engine.restore_baseline() + apply_steering(...),
   • engine.export_merged() + merged.save_pretrained(out_dir),
   • 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).

5. 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:

1. 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).
2. Install nvidia driver-580-server-open + fabricmanager (must be matched version), CUDA 12.4+, Python 3.10, then:

   uv pip install torch==2.6.0+cu124 --index-url https://download.pytorch.org/whl/cu124
   uv pip install transformers~=5.6 accelerate~=1.13 peft==0.14.0 \
                  bitsandbytes>=0.49 abliterix==1.6.0 \
                  optuna huggingface_hub tiktoken blobfile
   

   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.
3. Download bullerwins/Kimi-K2.6-bf16 (2.05 TB, 64 shards) → e.g. /mnt/data/Kimi-K2.6-bf16.
4. 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).
5. 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.
6. Run abliterix:

   export OPENROUTER_API_KEY=...
   export ABLITERIX_AUTO_EXPORT_DIR=/mnt/data/Kimi-K2.6-abliterated
   export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
   abliterix --config configs/kimi_k2_6.toml --non-interactive --overwrite-checkpoint
   

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

modal run --detach modal_track2.py::train_sft --epochs 1
modal 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
3-B	GateBreaker τ=1 (724K total: 689K routed + 34K shared)	very aggressive	Refusal preserved
3-C	Residual rank-1 subtraction (heretic/Wollschläger style; 80 harmful + 80 benign calibration)	strength=1.0	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).
• abliterix: v1.4.0 (PyPI baseline) and v1.6.0 (forked with bf16 dequant cache fix). Upstream: https://github.com/wuwangzhang1216/abliterix.
• Original Kimi modeling code: copied from moonshotai/Kimi-K2.6 via bullerwins/Kimi-K2.6-bf16.
• All edits in this artifact are minimal, defensive, and version-tolerant — same files load cleanly under transformers 4.56+ and 5.6+.

License

Patches inherit the Modified MIT License from Kimi-K2.6 and the AGPL-3.0 of abliterix where applicable.