CoolFace
Modelpublic

ZengXiangyu/Qwen3-8b-HiCI-48k-500steps

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face
1likes11downloads
Model Card

Qwen3-8b-HiCI-48k-500steps

Model Description

This is a LoRA adapter for Qwen3-8B with HiCI (Hierarchical Construction-Integration) memory architecture, trained for long-context understanding up to 48K tokens.

Paper: HiCI (arXiv 2603.20843) Base: LongLoRA (ICLR 2024 Oral)

HiCI Architecture

Three-stage hierarchy per transformer layer:

  1. 1.Local Construction — M learnable query slots attend to each segment via bottleneck cross-attention → local summary L_i
  2. 2.Global Integration — multi-view statistics (mean/max/min/std/ℓ2-norm) → shared compression → attention-based selection → gated expansion → G
  3. 3.Top-down Broadcast — per-segment attention with augmented KV=[G, L_i, segment tokens]; queries from segment tokens only
Input (48K tokens) → 4 segments × 12K
  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.safetensors  (27 MB)
└── LoRA Adapters (r=8, alpha=16): q_proj, k_proj, v_proj, o_proj

trainable_params.bin  (~4 GB)
├── global_memory.*            — Local Construction modules (36 layers)
├── hierarchical_aggregator.*  — Global Integration modules (36 layers)
├── self_attn.q_norm / k_norm  — QK-Norm weights (Qwen3-specific, 36 layers)
├── input_layernorm / post_attention_layernorm — LayerNorm weights (36 layers)
├── model.embed_tokens.weight  — Token embeddings
└── model.norm.weight          — Final LayerNorm

Training Details

  • —Base Model: Qwen/Qwen3-8B
  • —Context Length: 49,152 tokens (48K)
  • —Segments: 8 × 6,144 tokens
  • —Local Memory Slots (M): 8 per segment
  • —Global Memory Slots (K): 4
  • —Memory Heads: 8, Bottleneck dim: 512
  • —LoRA: r=8, alpha=16, target: q/k/v/o_proj
  • —Checkpoint: step 500 / 1000
  • —Batch: perdevice=1, gradaccum=8 (effective batch=8)
  • —LR: 2e-5 (LoRA), 2e-4 (memory modules), grad clip=0.3
  • —Precision: bf16
  • —Hardware: 8× H200 141GB, DeepSpeed Stage 2

Usage

Requires `qwen3_attn_hici.py` from this repo.

python
import torch
import transformers
from peft import PeftModel
# Download qwen3_attn_hici.py from this repo first
import qwen3_attn_hici as hici_attn

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

# 2. Load base model
base_model = transformers.AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-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/Qwen3-8b-HiCI-48k-500steps")

# Load HiCI params (embed, norm, global_memory, hierarchical_aggregator)
import os
trainable_params_path = os.path.join(
    "ZengXiangyu/Qwen3-8b-HiCI-48k-500steps", "trainable_params.bin"
)
# (auto-loaded by PeftModel if using the HiCI-aware load script)

# 5. Tokenizer
tokenizer = transformers.AutoTokenizer.from_pretrained("ZengXiangyu/Qwen3-8b-HiCI-48k-500steps")

Citation

bibtex
@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

Apache 2.0 (follows Qwen3 license)