gilesthomas-com

作者以JAX框架从零构建并训练了一个GPT-2 Small模型（参数量与原始论文一致），全程仅凭笔记，不参考任何现有代码。训练耗时37小时15分钟，在测试集上的损失（3.418784）优于等效PyTorch模型（3.538161）和原版GPT-2 Small（3.499677）。文章详细记录了从最简单的"映射到自身"模型开始，逐项添加位置编码、LayerNorm、多头注意力、Transformer块等组件，并验证每一步确实降低了损失的过程。

作者以Sebastian Raschka的《从零构建大语言模型》为教材学习现代AI，在独立完成PyTorch模型训练后，为了验证理解而非机械复现代码，选择用JAX框架编写LLM训练循环。本文记录了从零构建A-to-A（输入等于输出）的最小模型与训练流程的过程，涉及数据加载、设备内存管理、优化器配置等关键技术细节，以及从PyTorch迁移到JAX/Flax NNX/Optax生态的经验教训。

本文探讨了"角色混淆"这一概念，即大型语言模型（LLM）在推理时会忽略特定的角色标签（如<system>、<user>、<think>），而更倾向于根据文本的语气推断角色身份。这种现象解释了许多越狱攻击（jailbreak）的工作原理——例如，用户可以通过模仿模型自身推理痕迹的语气来欺骗模型，使其违反安全策略。文章还讨论了一种编号猜谜游戏的变体仍能成功愚弄ChatGPT 5.5，并提出了通过直接修改嵌入向量来区分角色、从而修复此问题的几种可能方案。

本文介绍了在调试JAX/Flax NNX训练循环时的一个实用技巧：通过对模型参数进行哈希处理，来检测参数是否真的在更新。作者在训练一个拥有7700万参数的LLM时，发现损失值一直停留在10.82不变，通过打印梯度值看起来正常，但无法直接观察参数变化。解决方案是将NumPy数组转换为字节后计算哈希值，这样即使是微小的参数变化也会导致哈希值完全改变。最终发现问题是使用了`@jax.jit`而不是Flax NNX专用的`@nnx.jit`装饰器，导致参数无法进行原位更新。

作者此前将家庭局域网升级到 10Gb/s，但在炎热天气下，基于 Marvell 芯片的老款 10GBASE-T SFP+ 模块因过热而反复断开连接。为解决这一问题，作者更换为采用 Broadcom BCM84891 PHY 芯片的 10Gtek ASF-10G-T80-INT 模块。更换后链路稳定性显著提升，交换机 CPU 温度下降约 5°C，但新模块的 EEPROM 伪装成 Intel 光纤模块，导致无法通过 SNMP 直接监测温度。

当你在JAX中使用`default_device`上下文管理器在CPU上创建数组时，数组实际上并未"提交"（committed）到该设备，JAX会自由地将其移动到其他设备（如GPU）。这会导致看似简单的数组切片操作耗时数秒——因为JAX可能每次都在移动整个父数组。解决方案是使用`jax.device_put`显式将数组提交到目标设备，这样后续操作的速度可以从1.2秒降至0.0002秒以下。对于需要在GPU上训练但将数据集保存在CPU内存中的场景（如LLM训练），理解这一机制至关重要。

作者在将 PyTorch 的 LLM 代码移植到 JAX 时，遇到加载 19GiB 数据集时 CUDA 显存不足的问题。通过探究发现，JAX 默认将数据分配到最快的可用后端（GPU），而不会自动利用 CPU 内存。文章详细介绍了 JAX 的 backend 与 device 概念，并展示了如何使用 `jax.default_device` 上下文管理器临时切换到 CPU 设备加载大数组，以绕过 GPU 显存限制。

文章介绍了如何在 Flax（基于 JAX 的神经网络库）中使用 Safetensors 保存和加载模型检查点。作者发现 Safetensors 官方提供的 Flax/JAX API 仅支持平铺的字典结构（字符串→JAX 数组），而 Flax 的 nnx.State.to_pure_dict 会产生嵌套字典，直接传入会导致报错。解决方案是先通过 nnx.to_flat_state 将状态转换为扁平结构，再拼接成点分隔键名的简单字典，即可正常使用 safetensors.flax 的 save_file 和 load_file 函数。

这是一篇面向PyTorch用户的JAX框架入门介绍。作者通过对比两种框架的训练循环实现，阐述了JAX的核心特点：更贴近数学表达的函数式设计、基于JIT编译而非逐段优化的执行方式，以及通过GradTracer自动追踪求导的优雅机制。文章还讨论了JAX的PyTree数据结构如何灵活处理复杂的梯度计算，并预告将对JAX的潜在局限性进行分析。

作者尝试在 MikroTik 10GBASE-T SFP+ 模块上安装 Raspberry Pi 常用的迷你散热片来降低温度。经过 24 小时测试，温度下降了约 3.5°C，效果不算显著但有一定改善。文章还提到此类 SFP+ 模块有新旧两代之分，旧款采用 Marvell 芯片、额定传输距离 30 米，新款采用 Broadcom 芯片、可达 100 米，作者计划未来考虑更换为新款模块。

作者分享了在家中将网络从 2.5Gb/s 升级到 10Gb/s 的完整实操经历，涵盖设备选型（MikroTik 交换机、Asus 网卡、Protectli 路由器）、结构化布线兼容性测试、DAC 与 RJ45 模块的散热挑战，以及最终实现接近 10Gb/s 全屋网络的实际效果。文章还讨论了 SFP+ 模块高温（高达 93°C）的监控数据与散热优化思路，并对未来升级至光纤网络的前景做了展望。

作者因 ISP 推出 10Gb 家庭宽带而升级家中有线网络，发现从 2.5Gb/s 跃升至 10Gb/s 需要重新面对许多早已被遗忘的挑战：散热管理（10GBASE-T 发热极严重）、线缆标准（CAT-5E 无法稳定支持 10Gb，需升级至 CAT-6 或 CAT-6A），以及 SFP+、DAC 直连铜缆和 AOC 有源光缆等替代方案。文章回顾了从 10BASE2 同轴电缆到现代交换机的网络技术演进，解释了为何家庭用户现在也需要关注此前只有企业网络管理员才需关心的物理层问题。

作者在完成《从零开始构建大语言模型》主体内容后，深入研究了书中的附录部分，发现其中包含了许多能节省时间的内容，但通过自己探索解决问题反而加深了对知识的理解。附录涵盖了PyTorch基础、分布式训练、梯度裁剪、学习率调度和LoRA等实用主题。

作者完成了从零训练GPT-2基础模型的旅程，通过一系列干预措施将训练时间缩短至44小时，最终模型性能接近GPT-2 small。文章总结了各项技术调整的效果，包括权重绑定、混合精度训练、梯度裁剪等技术对模型损失的影响。

本文作者在基于Sebastian Raschka的书籍构建GPT-2小型风格LLM后，通过一系列干预措施尝试提升模型性能，并采用改进的评估方法对多个模型进行指令微调测试。结果显示，测试集损失与指令遵循能力之间存在复杂关系，某些模型表现超出预期，而训练配置差异（如梯度累积与分布式数据并行）对结果产生了不一致的影响。

作者通过训练一个1.63亿参数的GPT-2小型模型，展示了LLM从初始的"词符沙拉"到生成连贯文本的演变过程。在训练了约10亿词符后，模型开始产生有意义的句子，最终能够生成类似商业和自我激励内容的文本。

作者在本地训练GPT-2小型风格LLM时，为了匹配云端训练的最佳批次大小效果，采用了梯度累积技术。通过将多个小批次的前向-后向传播梯度累积后再进行优化器更新，实现了等效大批次训练的稳定性优势，最终在本地RTX 3090上成功复现了云端训练的质量。

作者在云端对163M参数的GPT-2风格模型进行了多项干预措施的组合测试，包括梯度裁剪、移除dropout、调整学习率调度等。通过三次训练实验，最终获得了3.577761的测试损失，显著优于基线模型，但仍未达到原始GPT-2权重的性能水平。

作者通过一系列实验探索了不同训练干预措施对GPT-2风格模型性能的影响，发现学习率调度等改进措施能带来约2%的性能提升。进一步研究发现，随机种子变化导致的模型性能差异（标准差约0.015）与某些干预措施的效果相当，表明训练过程中的随机性影响不容忽视。

本文探讨了在从头训练GPT-2小基础模型时，关闭PyTorch的自动混合精度(AMP)和降低矩阵乘法精度对模型性能的影响。实验结果显示，虽然使用完整的float32精度带来了微小的测试损失改进，但训练时间增加了两倍多，成本增加了三倍，性价比极低。

作者开发了lambda-manager工具，通过轮询API自动检测Lambda Labs上8x A100实例的可用性，并在发现可用时自动启动实例并发送Telegram通知，以解决GPU资源紧张问题。

本文探讨了权重绑定技术，该技术通过共享输入嵌入和输出投影层的参数来减少模型参数量。作者基于Sebastian Raschka的代码，在一个163M参数的小型模型上测试了这一经典技术，发现虽然它能降低参数规模，但可能限制模型表达能力，导致性能下降。

本文探讨了在从头训练GPT-2小模型时权重衰减的作用，解释了权重衰减作为正则化技术如何通过惩罚模型权重的大小来防止过拟合，并分析了AdamW优化器中weight_decay参数的最佳设置。

本文探讨了在从头训练GPT-2小型基础模型时如何设置学习率，分析了固定学习率的局限性，并介绍了学习率调度策略，特别是余弦衰减和预热机制。作者还讨论了如何确定合适的初始学习率值。

本文探讨了在基于Sebastian Raschka书籍构建的GPT-2小型基础模型中，为注意力权重矩阵添加偏置项的实验。结果显示，尽管现代LLM通常不使用QKV偏置，但在这个特定规模的模型上，添加偏置使测试损失降低了0.023，训练稳定性也有所改善。