K2-Horizon-0.9B / APPENDIX.md
chen11003's picture
Upload APPENDIX.md
9ac4303 verified
|
Raw
History Blame Contribute Delete
14.8 kB

K2-Horizon-0.9B Technical Appendix

This appendix contains the detailed architecture, checkpoint, training, deployment, hardware, evaluation, dataset, source-code, intended-use, and safety material moved from the full model card. The short README follows the common K2-Horizon release format.

Model Series Overview

K2-Horizon spans compact and larger dense models for transparent foundation-model research, staged checkpoint analysis, and practical deployment. The repositories share a common release structure while preserving model-specific training and serving guidance.

Model Architecture Parameters Context length Intended use
K2-Horizon-0.9B Dense decoder-only 1.08B including embeddings 131,072 Compact reasoning, local inference, and distillation research
K2-Horizon-3.7B Dense decoder-only 3.78B core 524,288 Efficient research, evaluation, and single-node serving
K2-Horizon-7B Dense decoder-only 7B core 524,288 General research, fine-tuning, and cost-conscious deployment
K2-Horizon-32B Dense decoder-only 32B core 524,288 Stronger long-context and reasoning experiments

This Repository

Property Value
Architecture K2HorizonForCausalLM (model_type: k2_horizon)
Parameters 1,078,285,824 (released as 0.9B; counted from the safetensors tensors)
Hidden size / layers 1,536 / 28
Attention heads / KV heads 32 / 8
Context length 131,072 tokens with YaRN RoPE scaling; original context length 8,192 tokens
Vocabulary size 64,256
Released weight dtype BF16
Format Hugging Face safetensors, one weight shard, with custom configuration and modeling code in the repository root
Distillation checkpoint Step 249 of a 500-step mOPD run

The main revision publishes K2HorizonForCausalLM, model_type: k2_horizon, and matching configuration_k2_horizon.py and modeling_k2_horizon.py modules. The Transformers and vLLM preflights below validate that public contract before loading weights.

Checkpoint Revisions

Revision Stage Max context
main Distilled release checkpoint, step 249 128K
mid2_47k Second context-extension stage 128K
mid1_75k First context-extension stage 40K
Base model Pre-distillation specialist merge 8K

Training Provenance

Training examples were routed to a math-and-code teacher, a STEM teacher, or an instruction-following teacher through the example's opd_domain metadata. The math-and-code teacher was used as the fallback when no recognized domain was present.

Domain Teacher checkpoint step
Math and code 2,739
STEM 499
Instruction following 1,499

The base context window was extended in stages from 8,192 to 40,960 and then to 131,072 tokens. The mid1_75k and mid2_47k repository revisions preserve the corresponding intermediate checkpoints. The distilled release checkpoint is on main; all stages use a vocabulary of 64,256 tokens.

Training Loss

K2-Horizon-0.9B begins with a task-arithmetic merge of three specialist checkpoints. mOPD then trains that merged student against math-and-code, STEM, and instruction-following teachers at the same time. The training objective combines an on-policy distillation loss with a reference-model KL term so the student can learn specialist behavior while remaining close to the merged base model.

The resulting checkpoint retains most of the specialist teachers' performance on the reported math and coding tasks. It also improves every reported IFEval submetric over the pre-distillation merge. This makes the model useful for research on compact reasoning models, local inference, distillation, and task-specific adaptation.

Deployment Guide

K2-Horizon-0.9B emits a reasoning segment before its final answer when the chat template is used. The template supports reasoning_effort values high, medium, and low, which select the model's full, fast, and faster reasoning modes respectively.

For a deterministic runtime check, use temperature=0 and generate 20 to 50 tokens. For general sampled generation, temperature=0.6 and top_p=0.95 reproduce the GPQA evaluation setting and are reasonable starting points. IFBench used temperature=0.8. Long math and coding tasks may require several thousand output tokens; choose limits from application measurements rather than treating an evaluation limit as a universal default.

vLLM

The validated serving image was reconstructed into the following manual runtime contract:

Component Validated value
Operating system Ubuntu 24.04, Linux x86-64
Python 3.12.13
CUDA toolkit 12.9
PyTorch 2.13.0+cu129
Transformers 5.16.1
Safetensors 0.8.0
FlashInfer 0.6.17
Attention backend vLLM FlashAttention 3; Triton 3.7.1
vLLM 0.26.1rc1.dev1212, PR #53806 source commit d9fd5f11

That source revision contains the native K2HorizonForCausalLM implementation and the built-in k2_horizon reasoning and tool parsers. Other vLLM revisions have not been validated for this checkpoint. Pin the exact commit until the integration is available in an upstream release.

Use Linux x86-64 with a CUDA 12.9-compatible NVIDIA driver and Git. The setup uses vLLM's precompiled extension path while keeping the Python package on the exact reviewed source commit.

Show the pinned vLLM environment setup
git clone --filter=blob:none --no-checkout \
  https://github.com/vllm-project/vllm.git
cd vllm
git fetch origin pull/53806/head:refs/remotes/origin/pr-53806
git checkout --detach d9fd5f11423a1a5628fe29e7296ceb9de91aac3c
test "$(git rev-parse HEAD)" = \
  "d9fd5f11423a1a5628fe29e7296ceb9de91aac3c"

python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip uv
export UV_LINK_MODE=copy
VLLM_USE_PRECOMPILED=1 uv pip install --upgrade --editable . \
  --torch-backend=auto
uv pip install "transformers==5.16.1" "safetensors==0.8.0"
python -m pip check

python - <<'PY'
from vllm import ModelRegistry
from vllm.reasoning import ReasoningParserManager
from vllm.tool_parsers import ToolParserManager

assert "K2HorizonForCausalLM" in ModelRegistry.get_supported_archs()
assert ReasoningParserManager.get_reasoning_parser("k2_horizon") is not None
assert ToolParserManager.get_tool_parser("k2_horizon") is not None
PY

Download the repository chat template explicitly and start the server with one GPU. The 8,192-token profile below is a conservative starting point. Increase MAX_MODEL_LEN only after measuring KV-cache capacity; the checkpoint supports up to 131,072 tokens.

Show the vLLM serving command
source .venv/bin/activate

export MODEL_ID="IFM/K2-Horizon-0.9B"
export MODEL_REVISION="main"
export MAX_MODEL_LEN=8192
export CHAT_TEMPLATE="$(hf download "$MODEL_ID" chat_template.jinja \
  --revision "$MODEL_REVISION")"

vllm serve "$MODEL_ID" \
  --revision "$MODEL_REVISION" \
  --model-impl vllm \
  --trust-remote-code \
  --dtype bfloat16 \
  --tensor-parallel-size 1 \
  --max-model-len "$MAX_MODEL_LEN" \
  --max-num-seqs 1 \
  --gpu-memory-utilization 0.85 \
  --served-model-name "$MODEL_ID" \
  --chat-template "$CHAT_TEMPLATE" \
  --reasoning-parser k2_horizon \
  --tool-call-parser k2_horizon \
  --enable-auto-tool-choice

The reasoning parser moves <ifm|think>, <ifm|think_fast>, or <ifm|think_faster> text into the OpenAI-compatible response's reasoning_content field. The tool parser converts generated <ifm|tool_call> blocks into structured tool calls when tools are supplied in the request.

Show an OpenAI-compatible request
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused")
response = client.chat.completions.create(
    model="IFM/K2-Horizon-0.9B",
    messages=[{"role": "user", "content": "What is the square root of 2?"}],
    max_tokens=50,
    temperature=0,
    extra_body={"chat_template_kwargs": {"reasoning_effort": "high"}},
)
message = response.choices[0].message
print(getattr(message, "reasoning_content", None))
print(message.content)

SGLang

Native K2 Horizon support is provided by sgl-project/sglang#37654. Use a lmsysorg/sglang:dev image built after that PR is merged. Once support is included in a tagged SGLang release, use the corresponding versioned image.

Show the SGLang serving command
docker run --gpus all \
  --shm-size 32g \
  -p 30000:30000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --ipc=host \
  lmsysorg/sglang:dev \
  python3 -m sglang.launch_server \
    --model-path "IFM/K2-Horizon-0.9B" \
    --revision main \
    --tp 1 \
    --dtype bfloat16 \
    --context-length 8192 \
    --attention-backend fa3 \
    --reasoning-parser k2_horizon \
    --tool-call-parser k2_horizon \
    --mem-fraction-static 0.85 \
    --host 0.0.0.0 \
    --port 30000

This uses SGLang's native K2HorizonForCausalLM implementation; no --trust-remote-code, source patch, or external parser plugin is required. The 8,192-token limit is a conservative starting point; increase it only after measuring KV-cache capacity.

Transformers

The checkpoint can also be loaded directly from the Hugging Face repository. Its configuration_k2_horizon.py and modeling_k2_horizon.py files are loaded through trust_remote_code=True. Use a clean environment so the direct path does not inherit vLLM's build dependencies. Transformers 4.57.x is not compatible with this remote configuration class; use the validated 5.14.1 version below. Transformers may print nonfatal cache_position documentation diagnostics while loading the remote code, but BF16 loading and generation complete normally.

Show the Transformers environment setup
python3.12 -m venv .venv-transformers
source .venv-transformers/bin/activate
python -m pip install --upgrade pip
python -m pip install "torch==2.11.0" \
  --index-url https://download.pytorch.org/whl/cu128
python -m pip install \
  "transformers==5.14.1" \
  "safetensors==0.8.0"
python -m pip check

This deterministic sample loads the released weights as BF16 on one CUDA GPU and generates only 50 tokens, making it suitable as an end-to-end smoke test.

Show the Transformers inference example
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "IFM/K2-Horizon-0.9B"
REVISION = "main"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    dtype=torch.bfloat16,
    trust_remote_code=True,
).to("cuda").eval()

messages = [{"role": "user", "content": "What is the square root of 2?"}]
inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    reasoning_effort="high",
    return_dict=True,
    return_tensors="pt",
)
inputs = {name: value.to(model.device) for name, value in inputs.items()}
inputs.pop("token_type_ids", None)

with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        do_sample=False,
        pad_token_id=tokenizer.pad_token_id,
    )

new_tokens = outputs[0, inputs["input_ids"].shape[-1]:]
print(tokenizer.decode(new_tokens, skip_special_tokens=True))

Hardware Planning

The BF16 weights occupy approximately 2.0 GiB. One CUDA GPU is sufficient for short-context inference; 8 GiB is a practical minimum for a one-request smoke test, while 16 GiB or more provides useful room for longer prompts and runtime workspaces. The BF16 KV cache is approximately 0.44 GiB per request at 8,192 tokens and approximately 7 GiB at 131,072 tokens, before allocator, activation, CUDA-graph, and framework overhead. Start with a short context and one sequence, then increase context length and concurrency from measured memory headroom.

Source Code

The model repository stores the checkpoint directly at its root. It includes:

  • model.safetensors and model.safetensors.index.json
  • config.json, generation_config.json, and the K2 architecture code
  • tokenizer files and three chat-template variants
  • README.md and LICENSE

Use chat_template.jinja for ordinary chat and OpenAI-compatible serving. chat_template_generation.jinja and chat_template_asst_tool_gen.jinja are specialized generation and assistant-tool-generation variants.

Evaluation

  • Training: multi-teacher on-policy distillation with OPD loss weight 0.1, reference-KL weight 0.01, learning rate 1e-7, and a 500-step schedule. The selected checkpoint is step 249.
  • Evaluation: the AIME, coding, and instruction-following evaluations used the training evaluation path with SGLang as the rollout engine. GPQA-Diamond was evaluated with Eval360-V2 at revision f5081bf.
  • Export: the distributed training checkpoint was converted to Hugging Face safetensors and checked for tensor parity; all 255 expected weight tensors matched. The architecture label was later updated from XllmForCausalLM to K2HorizonForCausalLM without changing the weights.

Intended Use

K2-Horizon-0.9B is intended for compact reasoning research, local inference, evaluation dry runs, distillation studies, and task-specific adaptation. It is a research release and should be evaluated on representative prompts before deployment.

Limitations and Safety

  • AIME 2025 remains below the math-and-code teacher, and AIME results have high sampling uncertainty because each benchmark contains only 30 problems.
  • GPQA-Diamond is statistically close to the pre-distillation base result, so the reported run does not demonstrate a clear STEM improvement.
  • Long-horizon tool use remains substantially weaker than single-turn tool calling in the BFCL v4 breakdown.
  • Benchmark scores depend on prompt templates, reasoning effort, sampling parameters, framework versions, and evaluation harness details. Validate the model on representative prompts before deployment.
  • As with other language models, K2-Horizon-0.9B can produce inaccurate, biased, or unsafe text. Applications should use task-specific evaluation, input and output controls, monitoring, and human review where appropriate.