Nandi-Mini-600M-Early-Checkpoint

简介

Nandi-Mini-600M-Early-Checkpoint 是即将推出的 Nandi-Mini-600M 模型系列的早期 checkpoint（在训练了 2500 亿 tokens 之后）。这并非最终模型，而是一款紧凑型多语言语言模型，专注于高效能、部署灵活性以及对印度语言的支持。

该模型完全从头开始训练，旨在以较低的计算和内存成本提供强大性能。分享此 checkpoint 是为了让大家提前了解模型的缩放行为和训练进展。

本次发布的是早期 checkpoint，并非最终收敛模型。随着训练的持续进行和规模的扩大，其性能有望进一步提升。

📢 我们即将发布技术博客！敬请关注！

架构亮点

Nandi-Mini-600M 引入了多项以效率为中心的架构优化，专为紧凑且功能强大的语言模型而设计。

共享 KV（共享键值向量）

共享 KV 是 Nandi-Mini 中探索的核心架构理念之一。不再分别计算键（Key）和值（Value）的投影，而是两者复用一个共享的潜在表示，同时专门为注意力计算应用一个轻量级的键归一化步骤。

这种设计在推理过程中将 KV 缓存的内存使用量减少了约 50%，而计算开销仅略有增加，因为旋转位置编码（RoPE）和键归一化是在注意力计算过程中动态应用的。

Nandi 支持两种 KV 缓存模式：

"kv_cache_mode": "shared"

采用共享KV，将KV缓存内存减少约50%，计算开销略有增加。

"kv_cache_mode": "vanilla"

采用标准的独立键值缓存（Key-Value caching），以实现最大的推理兼容性并降低计算开销。

键值缓存（KV-Cache）内存对比

• Vanilla KV → 标准键值缓存内存使用量
• Shared KV → 键值缓存占用空间降低约50%

共享键值缓存（Shared KV）是我们对可部署基础模型更广泛关注的一部分，这些模型针对以下场景进行了优化：

• 本地人工智能系统（On-premise AI systems）
• 内存受限的部署环境
• 边缘设备
• 长上下文推理工作负载

这仍是Nandi模型系列中一个活跃的研究领域，我们计划在即将发布的技术博客中分享更深入的技术细节。

模型详情

• 类型：因果语言模型（Causal Language Model）
• 训练阶段：早期预训练检查点（2500亿 tokens）
• 参数：约6亿
• 架构：Transformer解码器
• 位置编码：RoPE
• 归一化：RMSNorm + QK Norm
• 激活函数：SwiGLU
• 注意力机制：GQA + Shared KV
• 嵌入层：带因式分解设计的绑定嵌入（Tied embeddings with factorized design）
• 上下文长度：2,048 tokens（计划扩展至32,000 tokens）
• 词汇表大小：131,072

📊 基准测试结果

这并非最终模型，而是一个早期检查点。因此，结果并非最终版本。目前仅完成了20%的训练。

通用基准测试

各语言的分词生育率得分

🌍 支持的语言

该模型在英语和多种印度语言上进行了训练，包括：

Hindi, Bengali, Tamil, Telugu, Marathi, Gujarati, Kannada, Malayalam, Punjabi, Odia

🚀 使用方法

!pip install transformers=='5.4.0'

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "FrontiersMind/Nandi-Mini-600M-Early-Checkpoint"

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    trust_remote_code=True
)

device = "cuda" if torch.cuda.is_available() else "cpu"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    dtype=torch.bfloat16
).to(device).eval()


#model.config.kv_cache_mode = "shared" # Use this one if wants to save 50% KV cache, but this will slight more compute
model.config.kv_cache_mode = "vanilla"

prompt = """The night was quiet and the streets were empty"""

model_inputs = tokenizer(
    [prompt],
    return_tensors="pt"
).to(model.device)

outputs = model.generate(
        **model_inputs,
        max_new_tokens=50,
        do_sample=True,
        temperature=0.3,
        top_k=20,
        top_p=0.95,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
        use_cache=True,  
    )

response = tokenizer.decode(
    outputs[0],
    skip_special_tokens=True
)

print(response)