构建了一个LLM训练框架，能在旧款GPU上运行而不崩溃

嘿，各位，我最近在摆弄Nanotron，发现它在模块级别导入了大量繁重的硬件特定依赖项（flash-attn、triton、functorch等），这让我非常沮丧。如果你试图在较旧或预算型GPU（如T4或V100）上运行它，它会在导入时直接崩溃。所以我编写了Picotron（https://github.com/Syntropy-AI-Labs/picotron）来解决这个问题。这是一个净室重写版，消除了所有强制性的GPU特定依赖。它几乎可以在任何支持PyTorch的GPU上运行（在计算能力低于8.0的旧卡上默认使用FP16，在较新的卡上使用BF16）。它默认回退到标准PyTorch SDPA，但如果检测到你已安装FlashAttention-2，它会在运行时连接到它。我使用AI助手编写了大量样板/代码模块，但已在本地运行成功，并在FineWeb-Edu上训练了一个很小的2M模型。此外还添加了以下配置： • GQA / MLA（多头潜在注意力） • QK-Norm 和 logit soft-capping（类似Gemma 2） • 并行FFN/Attn运行 • 在DDP上使用ZeRO-1包装 目前路线图很短：MoE准备（路由容量因子和负载均衡损失），使数据集准备比手动流式处理更容易。如果你一直在与CUDA依赖地狱作斗争，不妨看看：https://github.com/Syntropy-AI-Labs/picotron