A blog about my thoughts on ML Sys and LLMs
Intro At my previous job I ran into a problem: our JAX trainer in an RL job would OOM after running for a while, which is pretty rare. Most of my OOM debugging experience came from Torch, where there are basically two moves: The careful way: use the Torch Memory Visualizer to break down where memory goes and optimize from there. Often it’s just a tensor you forgot to free, or an allocation that…
引言 在之前的工作中我遇到一个问题：在我们的RL Job里，Jax Trainer跑着跑着OOM了，这很罕见。 我之前大部分的OOM Debug经验都来自Torch——对付Torch的OOM，基本就是两个套路： 细致点的话，用 Torch Memory Visualizer 把显存占用的来源一点点分析出来，然后针对性优化，很多时候其实是忘记释放或者申请了没有必要的大内存； 懒一点的话，在疑似Memory Leak的角落加上 torch.cuda.empty_cache() 。 所以碰到Jax OOM这个问题的时候我想到了两个更基本的问题： 为什么我很少在Jax中遇到OOM？ 我能不能使用 jax.cuda.empty_cache() ？ 想回答这两个问题，就得先搞明白Jax到底是怎么分配内存的——于是就有了这篇blog。 这篇文章也同步分享在知乎： JAX 是怎么分配内存的 。 概览…
Authored by Biao He Zilin Zhu Ji Li 1. What is slime? slime is a LLM post-training framework aiming for RL Scaling, it was designed to be: Versatile – with a fully customizable rollout interface and flexible training setups (colocated or decoupled, synchronous or asynchronous, RL or SFT cold start). Performant - integrating SGLang for inference and Megatron-LM for training, natively. Maintainable…
本文也在我的知乎专栏中发布， 知乎链接 作者 何标 朱子霖 李冀 1. 什么是slime？ slime 是一个强化学习大规模训练框架，提供以下核心能力： 多功能 – 拥有完全可定制的rollout接口和灵活的训练设置（同卡或分离、同步或异步、RL或SFT）。 高性能 - 原生集成Megatron和SGLang进行训练和推理。 易维护 - 轻量级代码库，并可从Megatron预训练平滑过渡到SGLang部署。 1 大规模验证 - 最近发布的 zai-org/GLM-4.5(355B) 和 zai-org/GLM-4.5-Air(106B) 都是通过slime做的RL训练。 slime主要由三个核心模块组成 2 ： 训练模块（Megatron） – 处理主要的训练过程，从数据缓冲区读取数据，并在训练后与rollout模块同步参数 Rollout模块（SGLang + Router） –…
Authored by Biao He Ata Fatahi 1. Introduction Reinforcement learning (RL) for large language models (LLMs) presents unique challenges due to its integration of inference and training in each step, demanding significant scalability and resource efficiency. The verl library, designed for RL training of LLMs, combines advanced training strategies like Fully Sharded Data Parallel ( FSDP ) and…
Authored by Biao He Qingquan Song 0x0. Introduction In the past few weeks, we’ve implemented the Flash Attention Backend end-to-end in SGLang, which is now the default attention backend as of SGLang 0.4.6 release . Throughout this journey, we learned a lot about how Attention Backend functions in modern LLM serving engines and developed a deeper understanding of Flash Attention itself. In this…
Introduction Flash Attention 1 is an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. It has been widely used in LLM inference and training, and is the default attention backend in modern serving engines like SGLang, vLLM, etc. Naive Attention Calculation Before we figure out how Flash…
Authored by Binyao Jiang This guide explains how to calculate the parameter size of a Mixture of Experts (MoE) large language model (LLM) using its architecture and configuration file. We’ll use the Qwen3-30B-A3B model as an example to demonstrate the process. 1. Understand the Model Architecture To calculate a model’s parameter size, you first need to understand its architecture. Initially, I…
This guide explains how to calculate the parameter size of a dense large language model (LLM) using its architecture and configuration file. We’ll use the Qwen3-32B model as an example to demonstrate the process. 1. Understand the Model Architecture To calculate a model’s parameter size, you first need to understand its architecture. Initially, I considered technical reports as a primary source,…
Knowledge distillation is a model compression technique whereby a small network (student) is taught by a larger trained neural network (teacher). I. What is model distillation? Model distillation is a technique used to transfer knowledge from a larger, more complex model (the “teacher” model ) to a smaller, simpler model (the “student” model ) in order to improve the performance of the smaller…
NOTE: This markdown cheatsheet is a typography demo for this theme. Check out this post to learn more about this markdown usage when you want to get started with this theme. Enjoy! Typography Elements in One Let’s start with a informative paragraph. This text is bolded. But not this one! How about italic text? Cool right? Ok, let’s combine them together. Yeah, that’s right! I have code to…