Argon-0.5B

Argon-0.5B 模型报告

龚钰轩

Beihang University
Email:[email protected]

模型网页：http://argonai.cn/Argon0.5B/ welcome to visit!

Argon-0.5B 一个自研的基线模型，复刻 DeepSeek-V4 模型训练典型优化器，并加入 Engram 模块。

本仓库计划上传模型权重、切分后的训练数据、训练代码、tokenizer 资产和完整配置，使 Argon-0.5B 成为一个可审计、可复现、可继续训练的研究型预训练样例。

1. 项目初衷

这个项目的初衷是复刻 DeepSeek 技术栈中的关键训练流程，并尝试在 500M 参数规模上实现一个完整的预训练闭环。

Argon-0.5B 的设计重点是：

• 复刻 DeepSeek 风格 tokenizer、数据 pack、bf16 多卡预训练流程；
• 使用 DeepSeek 128K tokenizer 路径，并区分 raw token 与 canonical token；
• 构建从原始语料、统一 JSONL、tokenizer 切分、ctx4096 pack 到最终 checkpoint 的全流程；
• 复现 Engram 模块，并验证其在小模型中的接入方式；
• 实验 MTP、MHC、Muon/AdamW 混合优化器等训练组件；
• 形成一个可以公开审计的全栈训练样例，而不仅是单独的模型权重。

2. 模型概览

模型结构配置文件：

configs/model/dense_500m_main.yaml

3. 架构特性

Argon-0.5B 不是标准 Transformers 模型，不能直接用 AutoModelForCausalLM 加载。它依赖本仓库中的自定义 PyTorch 代码。

主要结构包括：

• 24 层 dense Transformer；
• RMSNorm；
• RoPE；
• clamped SwiGLU；
• tied embedding；
• Engram memory 模块；
• MTP 辅助预测头；
• MHC 模块；
• DeepSeek 128K tokenizer 路径；
• raw token 主 LM loss + canonical token Engram lookup。

Engram 配置：

layers: [2, 12]
ngram_orders: [2, 3]
num_hash_heads: 4
memory_dim: 128
slots_per_layer: 40000

4. Tokenizer

本模型使用 DeepSeek 风格 128K tokenizer。

相关文件：

configs/tokenizer/deepseek_tokenizer.yaml
artifacts/tokenizer/deepseek_main/
artifacts/tokenizer/engram_compression_deepseek_v1/raw_to_canonical.uint32.npy

训练时：

• 主模型 CE loss 使用 raw token ids；
• Engram lookup 使用压缩后的 canonical token ids；
• context length 为 4096。

5. 训练数据

本次训练使用的是已经切分并 pack 好的数据，而不是直接读取原始 JSONL。

packed 数据目录：

artifacts/data/packed/deepseek_128k_ctx4096/current_train/

主要文件：

raw_input_ids.bin           约 32G
canonical_input_ids.bin     约 32G
loss_mask.bin              约 7.8G
metadata.jsonl             约 967M
packing_report.json        约 1.6K

packed 数据总大小约 72G。

训练语料设计偏向：

• 中文教育网页；
• 数学核心数据；
• 科学、论文、PDF、教材；
• 书籍、百科、公开长文；
• 英文教育网页；
• Markdown、技术文档和少量代码说明；
• 长上下文平衡数据；
• 中文通用高质量覆盖；
• Ultra-FineWeb 中英高分子集；
• 少量中文教材式合成知识；
• 少量弱问答/说明型数据。

本轮训练最终使用：

8,373,338,112 tokens

6. 训练配置

训练硬件：

4 x NVIDIA RTX PRO 6000 Blackwell Server Edition

训练命令：

NPROC_PER_NODE=4 bash scripts/05_train_500m_engram_pro6000.sh \
  --batch_size 4 \
  --gradient_accumulation_steps 3 \
  --compile false

有效 batch：

4 GPUs × batch_size 4 × grad_accum 3 × ctx4096
= 196,608 tokens / optimizer step

主要配置文件：

configs/train/s2_b_500m_engram_pro6000.yaml
configs/data/pack_ctx4096.yaml
configs/optim/hybrid_muon_500m.yaml

优化器设计：

• backbone matrix：AdamW warmup 后切换到 Muon；
• MTP matrix：AdamW warmup 后切换到 Muon；
• embedding/head/norm/bias：AdamW；
• Engram memory：Adam；
• Engram gate/proj/conv：AdamW；
• MHC params：AdamW；
• cosine scheduler；
• warmup tokens：3M；
• Muon switch after：5M tokens；
• grad clip：1.0。

7. 训练结果

最终 checkpoint：

checkpoint_step_42589.pt

最终训练状态：

本地下载后的最终 checkpoint SHA256：

e6377fb533dbfa36ef14d28774e8b8a58f026e3e5ef94f02826e31b2a7de30dc

8. 训练图表

以下图表由完整训练日志 train_log.jsonl 生成，曲线使用 200-step rolling mean 平滑。

8.1 Loss 曲线

8.2 训练吞吐

8.3 Engram Gate 曲线

8.4 Engram / MHC 健康指标

8.5 优化器更新范数

9. 使用方式

该模型需要本仓库代码加载。

CPU 续写测试：

python3 scripts/29_pure_continue_eval_500m.py \
  --checkpoint checkpoints/checkpoint_step_42589.pt \
  --device cpu \
  --max-new-tokens 64 \
  --temperature 0.35 \
  --top-k 40 \
  "强化学习是"

CUDA 续写测试：

python3 scripts/29_pure_continue_eval_500m.py \
  --checkpoint checkpoints/checkpoint_step_42589.pt \
  --device cuda \
  --max-new-tokens 128 \
  "The main idea of reinforcement learning is"

也可以用 chat 模板做 smoke test：

python3 -m engram_v4.eval.chat_generate \
  --checkpoint checkpoints/checkpoint_step_42589.pt \
  --device cuda \
  --prompt "你好，简单介绍一下你自己。" \
  --max-new-tokens 128

但当前模型没有经过 SFT，因此 chat 模板输出不代表最终对话能力。

10. 初步测试

模型可以正常生成中文和英文，tokenizer、checkpoint、Engram 路径均已验证可用。

当前 checkpoint 的初步续写现象：

• 中文和英文都能生成；
• 基本语法正常；
• 教育/论文风格已经出现；
• 在部分抽象概念上会出现重复；
• 数学、科学解释仍需要进一步评测；
• 该 checkpoint 主要用于观察 base pretraining 阶段的语言建模效果。

示例：

Prompt: 机器学习的核心思想是
Output: 通过数据训练,使模型能够学习到数据中的规律,并能够预测数据中的未知数据...

Prompt: 矩阵乘法可以理解为
Output: “乘数相加”,即“乘数相加”...

这些结果说明模型已完成基础语言建模训练，并可以作为后续技术实验的 checkpoint。

11. 引用

@misc{argon05b2026,
  title = {Argon-0.5B: A 520M Parameter DeepSeek-Style Pretraining and Engram Reproduction Experiment},
  author = {Gong, Yuxuan},
  year = {2026},
  note = {Custom PyTorch pretraining stack with DeepSeek-style tokenizer path and Engram memory components}
}

12. 声明

Argon-0.5B 是实验性研究模型。当前 checkpoint 代表 base pretraining 阶段结果，主要用于技术复现、训练流程研究和后续继续训练。使用者需要自行评估模型行为，并确认数据、代码和权重的使用方式符合自己的应用场景和许可证要求。