seqax
seqax是基于JAX的序列建模框架，面向小型到中型LLM预训练研究。它以500行简洁代码实现模型、优化器、多主机FSDP和张量并行分区，具备良好扩展性和性能。seqax注重代码可读性和透明度，使数学计算、内存使用和芯片间通信等关键细节清晰可见，便于理解和修改。
