提出在芯片上实现Transformer训练的软件架构,支持本地Adam优化,节省内存。
NeuronFabric: A Software Reference Architecture for On-Chip Transformer Training with Local Adam
- 设计可在FPGA/ASIC上运行的Transformer训练软件架构,支持本地Adam更新
- BF16W方案使334K参数模型内存需求从4.0MB降至3.34MB,腾出空间存激活值
- 适用于关注芯片级AI训练、硬件加速器设计的研究者与工程师
现有公开加速器架构通常将训练计算与优化器状态更新分离,或依赖外部存储和主机调度。本文提出NeuronFabric,一种面向未来FPGA和ASIC实现的Transformer训练软件参考架构,支持本地Adam更新。完整C#原型实现了前向传播、反向传播和Adam优化,无需外部机器学习框架。评估模型为334K参数自回归Transformer(d=88, H=4, f=264, L=4, vocab=256),在Shakespeare语料上训练。BF16W配置在80K样本后达到评估损失1.5426,接近FP32 GPU参考结果(1.5224),并生成连贯字符级文本。该方案将权重存储为BF16,而保留优化器动量为FP32,显著降低内存开销。334K参数模型在FP32下需约4.0MB内存,恰好匹配Xilinx ZCU102的BRAM容量;而BF16W仅需约3.34MB,为激活值预留空间。论文分析了词汇表预算约束,量化了内存节省效果,并规划下一步开展FPGA训练验证。本文未包含实际FPGA测量数据,旨在公开架构设计与软件实现,供后续FPGA/ASIC探索使用。
原文摘要 · Abstract (English)
Publicly documented accelerator architectures generally separate training computation from optimizer-state updates or rely on external memory and host orchestration. This paper presents NeuronFabric, a software reference architecture intended for future FPGA and ASIC implementations of transformer training with local Adam updates. A complete C# prototype implements forward pass, backpropagation, and Adam optimization without external machine-learning frameworks. The goal is to validate numerical correctness and memory requirements before hardware implementation. The evaluated model is a 334K-parameter autoregressive transformer (d=88, H=4, f=264, L=4, vocab=256) trained on the Shakespeare corpus. The BF16W configuration achieves evaluation loss 1.5426 after 80K samples, compared with 1.5224 for an FP32 GPU reference, while producing coherent character-level text. The paper introduces BF16W, which stores weights in BF16 while retaining Adam optimizer moments in FP32. This reduces memory requirements for on-chip training. A 334K-parameter FP32 model with Adam moments requires approximately 4.0 MB, matching the BRAM capacity of a Xilinx ZCU102 device. The BF16W variant requires approximately 3.34 MB, leaving memory available for activation storage. We describe the vocabulary-budget constraint observed during earlier experiments, quantify BF16W memory savings, and outline FPGA training as the next stage of development. No FPGA measurements are included in this paper. This publication serves as a public architectural disclosure and software reference implementation for future FPGA and ASIC exploration of the NeuronFabric architecture.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。