DSV2-Lite-IndexCache-Warmup-400

A modified derivative of DeepSeek-V2-Lite, converted from a distributed training checkpoint to BF16 safetensors. Method: indexcache; stage: warmup; checkpoint step: 400. The layer pattern is FSS repeated nine times. Each full layer computes routes reused by its two shared layers. The causal key selection budget is 2048 tokens.

This warmup checkpoint retains the learned indexer and frozen backbone. The default attention_mode="sparse" evaluates the learned selector. Set config.attention_mode="dense" before loading to evaluate the dense teacher backbone.

Loading

Install the versions in requirements.txt. Tested with Transformers 4.57.3. Replace the path below with this repository ID or the downloaded directory.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "fxmeng/DSV2-Lite-IndexCache-Warmup-400"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_path, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda:0"
).eval()
inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda:0")
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=32, do_sample=False,
                            pad_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(output[0], skip_special_tokens=True))

All four exports share byte-identical modeling_dots_dsa.py; their configurations and weights determine the method and checkpoint. The runtime also supports all six method/stage configurations. checkpoint_stage records provenance and does not switch the inference attention mode. No Megatron environment or custom CUDA extension is required. Use left padding for batched generation.

Validation and limitations

validation.json records strict loading, exact tensor names and shapes, full parameter finiteness, English/Chinese/code generation, beam search, cached versus uncached decoding, finite LM loss, and 4097-token sparse prefill/decode/cache-crop checks. The shared implementation passed 20 unit tests and 4 subtests. BF16 and FP32 results are both retained, including numerical differences. These are conversion and inference checks, not benchmark quality measurements.

The initial BF16 logit-difference thresholds (maximum < 0.75, mean < 0.05) were exceeded. Across batch, cache, and cache-crop comparisons, the largest BF16 difference was 0.5; the largest FP32 difference was 9.15527e-05. All FP32 comparisons met tighter limits (maximum < 0.001, mean < 0.0001). The sampled next-token choices agreed in both precisions, and six-token greedy generation agreed with and without cache in each precision. This supports correct inference logic with BF16 numerical sensitivity. It does not guarantee identical generations across precisions or for arbitrary prompts.

For comparison, the original DeepSeek-V2-Lite loaded with Transformers' native implementation also exceeded the initial BF16 limits: maximum/mean differences were 0.4375/0.06381 for batching and 0.625/0.11797 for 4097-token cached decoding. Its sampled next-token choices also agreed. This reference diagnostic is included in validation.json.

The reference runtime scores selectors and applies masked SDPA in query chunks. It still performs dense work and does not provide optimized sparse CUDA serving. The training sequence length was 65536; this release test covers 4097 tokens and does not establish full 64K accuracy or performance. The inherited tokenizer length metadata is 16384 and the model's RoPE configuration allows larger positions; neither value establishes a tested context length. BF16 batching and decoding can change numerical results and discrete MoE/attention routes; see the measured differences in validation.json. The warmup KL training objective is not implemented in this inference runtime.

License and modifications

The base weights and tokenizer originate from DeepSeek-V2-Lite, copyright DeepSeek. The model weights are governed by the included DeepSeek License Agreement v1.0. This derivative adds learned sparse indexers and a shared Hugging Face model implementation. Conversion changes tensor layout and storage without retraining. Checkpoint details are in provenance.json.