Llama-3-8b-HiCI-16k

Llama-3-8b-HiCI-16k

Model Description

This is a HiCI adapter checkpoint for Llama-3-8B, extending its context window to 16K tokens. It contains three components: LoRA adapters (q/k/v/o_proj), HiCI module weights (LocalConstructor + GlobalIntegrator), and fine-tuned embedding + LayerNorm weights.

Paper: HiCI (arXiv 2603.20843)

HiCI Architecture

Three-stage hierarchy per transformer layer:

1. Local Construction — M learnable query slots attend to each segment via bottleneck cross-attention → local summary L_i
2. Global Integration — multi-view statistics (mean/max/min/std/ℓ2-norm) → shared compression → attention-based selection → gated expansion → G
3. Top-down Broadcast — per-segment attention with augmented KV=[G, L_i, segment tokens]; queries from segment tokens only

Input (16K tokens) → 4 segments × 4K
  Stage 1: 8 local slots per segment → L_i
  Stage 2: multi-view stats → K=4 global slots G
  Stage 3: Q=[chunk], KV=[G, L_i, chunk] → Flash Attention

Trainable Components

adapter_model.bin  (25 MB)
└── LoRA Adapters (r=8, alpha=16): q_proj, k_proj, v_proj, o_proj

trainable_params.bin  (~3.5 GB)
├── local_constructor.*            — Local Construction modules (32 layers)
├── global_integrator.*  — Global Integration modules (32 layers)
├── input_layernorm / post_attention_layernorm — LayerNorm weights (32 layers)
├── model.embed_tokens.weight  — Token embeddings (vocab=128,258)
└── model.norm.weight          — Final LayerNorm

Note on Llama-3 GQA: Llama-3-8B uses Grouped Query Attention (8 KV heads vs 32 query heads). The base k_proj / v_proj output dim is 1024 (not 4096). HiCI modules are unaffected — they use their own bottleneck projections (dim=512) independent of the base attention head structure.

Training Details

• Base Model: meta-llama/Meta-Llama-3-8B
• Context Length: 16,384 tokens (16K)
• Segments: 4 × 4,096 tokens
• Local Representation Slots (M): 8 per segment
• Global Representation Slots (K): 4
• HiCI Attention Heads: 8, Bottleneck dim: 512, Shared compress dim: 128
• LoRA: r=8, alpha=16, target: q/k/v/o_proj
• Checkpoint: step 1000
• Batch: per_device=1, grad_accum=16 (effective batch=16)
• LR: 2e-5 (LoRA), 2e-4 (HiCI modules), grad clip=0.3
• Precision: bf16
• Hardware: 4× H200 141GB, DeepSpeed Stage 2, transformers 4.40.0

Usage

Requires llama3_attn_hici.py from this repo (transformers >= 4.40.0).

import torch
import transformers
from peft import PeftModel
import llama3_attn_hici as hici_attn

# 1. Replace attention with HiCI BEFORE loading model
hici_attn.MIXED_GROUP_TRAINING = False
hici_attn.replace_llama_attn(use_flash_attn=True, use_full=False, use_hierarchical_forward=True)

# 2. Load base model
base_model = transformers.AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B", torch_dtype=torch.bfloat16, device_map="auto",
)

# 3. Register HiCI modules (must match training config)
hici_attn.register_hici_to_model(base_model, num_memory_slots=8, global_slots=4, num_heads=8, bottleneck_dim=512)

# 4. Load LoRA adapter + trainable_params
model = PeftModel.from_pretrained(base_model, "ZengXiangyu/Llama-3-8b-HiCI-16k")

# 5. Tokenizer (tiktoken-based, no tokenizer.model needed)
tokenizer = transformers.AutoTokenizer.from_pretrained("ZengXiangyu/Llama-3-8b-HiCI-16k")

Citation

@article{zeng2026hici,
  title={HiCI: Hierarchical Construction-Integration for Long-Context Attention},
  author={Zeng, Xiangyu and Xu, Qi and Wang, Yunke and Xu, Chang},
  journal={arXiv preprint arXiv:2603.20843},
  year={2026}
}

License

This model follows the Meta Llama 3 Community License.