AI News HubLIVE
站内改写6 分钟阅读

待翻译:Getting GLM-5.2 NVFP4 Post-Training off the ground

AI 服务暂时不可用,以下为来源摘要,待恢复后补全翻译:Patronus AI | Getting GLM-5.2 NVFP4 Post-Training off the ground Announcing our $50 Million Series B 🎉 Read Blog Post Here Docs Close Getting GLM-5.2 NVFP4 Post-Training off the ground The goal was deceptively simple t…

来源Hacker News AI作者: makaimc

AI 服务暂时不可用,以下为来源正文,待恢复后补全翻译。

Patronus AI | Getting GLM-5.2 NVFP4 Post-Training off the ground Announcing our $50 Million Series B 🎉 Read Blog Post Here Docs Close Getting GLM-5.2 NVFP4 Post-Training off the ground The goal was deceptively simple to state: take GLM-5.2, a 744B-parameter mixture-of-experts model quantized to 4-bit NVFP4, attach a bf16 LoRA adapter, and train it with reinforcement learning until it could play a level of Super Mario Bros., emitting button presses, reading the terrain ahead, and running for the flag. Realizing that specification required resolving a set of defects, which fall into three broad classes. The first is arithmetic: the 4-bit base did not fit within the memory available to it. The second is distributed-systems correctness: the warm-start adapter silently loaded the same expert block on every expert-parallel rank. The third, and the most resistant to diagnosis, concerns training stability: the reward collapse described below persisted under every standard remedy we applied, and was resolved only by removing a regularization term rather than introducing one. TL;DR: the entire setup GLM-5.2 is a 744B-parameter MoE (~40B active). NVFP4 base, frozen (the published nvidia/GLM-5.2-NVFP4 checkpoint), which lands as ~110 GB of weights per training GPU on 8×B200 once Transformer Engine loads it (Part I halves that) + bf16 LoRA (rank 64, MLP-only). Trainer: Megatron with tensor/expert parallelism (TP4·EP4·ETP2) on 1 node of 8×B200; for RL, rollouts are served by SGLang on a second node of 8×B200 (disaggregated, so serving never competes with the trainer for memory); orchestrated by miles/slime. RL algorithm: GRPO, 16 samples per prompt. Reward: how far Mario travels through level 1-1, minus a time penalty, plus a flag bonus.One caveat up front: "NVFP4" is selective, not applied everywhere. Only the routed MoE expert weights are actually in 4-bit. The attention layers, the dense early layers, the shared experts, embeddings, lm_head, the norms, and (of course) the LoRA adapter all stay in bf16. In other words, this is a mixed-precision model, and "the NVFP4 model" is shorthand. This is deliberate on NVIDIA's part: the same "keep the sensitive, low-volume layers in higher precision" logic that the related work below leans on. The typical post-training arc is two stages, SFT → policy-RL, and that is the sequence this article follows too. The toolchain Five pieces of open infrastructure do the work here, in two camps: a trainer that holds the weights and takes gradient steps, and a rollout engine that generates episodes, with an RL framework wiring them together. Megatron-LMTrainerNVIDIA's framework for training very large models.github.com/NVIDIA/Megatron-LM ↗Transformer EngineLow-precision kernelsNVIDIA's library of FP8/FP4 building blocks (quantized GEMMs, attention, LayerNorm) that actually execute the NVFP4 math on Blackwell. The "secretly-8-bit" memory bug lives here, in how it keeps a transposed copy of each weight.github.com/NVIDIA/TransformerEngine ↗SGLangRollout engineA fast inference/serving engine. During RL it hosts the current policy and generates the rollouts (the Mario episodes). Getting it to serve an NVFP4 MoE with a LoRA overlay on the experts took a specific combination of its runner and quantization backends.github.com/sgl-project/sglang ↗slimeRL frameworkAn open-source RL post-training framework that connects a training backend (Megatron) to a rollout engine (SGLang) over Ray, running the generate → score → learn → sync-weights loop. Notably, slime does not currently support LoRA: it assumes full-parameter training.github.com/THUDM/slime ↗ z.ai/blog/glm-5.2 ↗milesRL framework + LoRAThe slime-derived RL framework this project builds on. Its decisive difference from slime: miles supports LoRA: the reason this 744B model can be adapted with a small bf16 adapter instead of full-parameter RL. But miles does not officially support NVFP4; combining 4-bit quantization with LoRA (QLoRA) is what our fork adds. The stack is a chain of gaps filled, each layer supplying what the one below it lacks: slime → LoRA (miles) → NVFP4 QLoRA (our fork). That chain is the whole reason a 744B model is trainable here at all. Most fixes in this log (dropping the second quantized copy, EP/ETP-aware loading, DAPO wiring) live in that fork's patches to miles, Megatron, and SGLang.github.com/radixark/miles ↗ Getting 744 billion parameters to train at all Before any learning question could be asked, the model had to fit in memory, load successfully, and survive a forward and backward pass. A chain of infrastructure bugs stood in the way: in the environment, in the memory arithmetic, in the serving kernels, and in the distributed loader. Together they are the price of admission for QLoRA on a model this size. The standouts: The 4-bit base that was actually using 8-bit memory NVFP4 stores each weight in 4 bits, so after sharding across 8 GPUs the frozen base should occupy ~55 GB per GPU (744B × 4 bits ≈ 372 GB, so ~46 GB/GPU across 8 ranks, and ~55 once the per-16 block scales and the layers that stay in bf16 are counted). It consumed 110 GB per GPU, twice the expected memory. Even on a 180 GB B200, that leaves too little for activations, optimizer, and the colocated engine. The cause: the FP4 GEMM path keeps a persistent columnwise transpose copy of the weights alongside the rowwise one, so it can perform the matrix multiplications (matmuls) of both the forward and the backward in the kernel's preferred TN layout. These aren't one tensor read two ways, they're two independent NVFP4 quantizations of the same weight, each with its own 4-bit values and its own scale set (rowwise blocks 16 along in, columnwise blocks 16 along out), stored simultaneously because the scales genuinely differ between the two blockings. Two full copies of the weight matrix, double the memory footprint, and it overflowed a single node. The fix removes the columnwise copy, cutting the base from 110 GB to 56 GB, and reconstructs the transposed layout on the fly in a bf16 weight-gradient (dgrad) path during the backward. This is a deliberate memory-for-compute tradeoff, and the central one of the whole project. The columnwise copy exists purely to make the backward pass faster: it lets the weight-gradient GEMM use the kernel's preferred layout with no runtime transformation. Dropping it buys back ~54 GB, the difference between fitting on one node and not, and in exchange the backward must dequantize and transpose on the fly every step, spending extra FLOPs to reconstruct what used to be cached. It's the QLoRA bargain in a sentence: trade a slower backward for a model that fits. The same tradeoff logic appears elsewhere: the DSA chunked kernels (next) stream computation to cut activation memory, and the FP4 serving path (later) pays per-forward dequant compute rather than storing a bf16 copy. Going deeper: why a 4-bit weight needs two copies (on B200) Some notation first: let X∈ℝn×in denote the input and W∈ℝout×in the weight matrix. Blackwell's FP4 matmuls run fastest in a TN layout, BLAS shorthand (Transposed–Normal) for a matmul where both operands present the contraction (K) dimension as the innermost axis. In a linear layer the weight matrix is used in two different contractions: the forward (Fprop, Y = X·Wᵀ) reads W with the input dimension as the inner dimension, while the input-gradient (Dgrad, dX = dY·W) contracts over the output dimension and wants W in the transposed layout. Transposing a packed 4-bit tensor on the fly is awkward and slow, so Transformer Engine (v2.12.0) simply pre-quantizes and stores both a rowwise and a columnwise copy of every weight matrix. Two FP4 copies, and the 4-bit base is quietly back to ~8-bit, which now occupies ~110 GB of GPU RAM. Figure 1: why two copies. The same weight feeds two matmuls that contract different axes: the forward reduces over in (a row of W), the backward's Dgrad reduces over out (a column). Blackwell's FP4 MMA (matrix-multiply-accumulate, the tensor-core instruction that multiplies two tiles and accumulates the result) requires the contraction axis innermost (K-major) with the per-16 scales grouped along it, so the forward wants an in-blocked copy and the backward an out-blocked one, two independent NVFP4 quantizations of the same weight. That is the 2× footprint; our fix keeps only the rowwise copy and reconstructs the backward in bf16. Figure 2: the two GEMMs and their shapes. The weight W is reused, but each pass contracts a different dimension: the forward Y = X·Wᵀ reduces over in, the backward dX = dY·W reduces over out. That contracted dimension is K, the inner (matched) axis of the matmul; because the FP4 MMA requires K innermost with its 16-element scales aligned to K, and in ≠ out, the two passes need the weight in two different K-major layouts, which is why two copies exist. Why can't we just transpose the one copy? Because NVFP4 isn't plain 4-bit values, it's a block-scaled format, and the scales don't transpose. Each tensor carries its 4-bit E2M1 values plus a per-16-element-block FP8 (E4M3) scale and one FP32 global scale. An NVFP4 block is a contiguous run of 16 elements along one axis (a 1-D group, never a 2-D tile), and those blocks run along the contraction axis. The rowwise copy blocks the weight 1×16 along in (16 consecutive in values at a fixed out), while the columnwise copy blocks it 16×1 along out. A transpose changes the contraction axis, so the original scales can no longer be used: the transposed tensor requires new block groupings, and therefore new scale factors computed from the original higher-precision weights. Simply transposing the block layout along with the values does not work either, since the blocks would still be aligned to the wrong axis for the backward's contraction. The packed FP4 values are not a free transpose view either, because they are bit-packed two per byte and swizzled for the tensor-core tile. That is why the only options are to quantize along both axes up front or, as we do, dequantize the rowwise copy to bf16 for the backward and skip the FP4 transpose entirely. Transformer Engine documents this rowwise/columnwise scheme directly in its NVFP4 "handling transposes" notes. TN, NT, and what actually bindsTN is a memory-layout tag: the two GEMM inputs present the contraction dimension K as their innermost (contiguous) axis, i.e. one operand stored row-major and the other column-major. The exact letters are library-dependent. cuBLAS names layouts by the BLAS transa/transb flags (hence "TN"), whereas DeepGEMM names them relative to D = C + A @ B, where its default NT (non-transposed A = row-major, transposed B = column-major) is the very same physical arrangement, e.g. fp8_gemm_nt computes D = C + A @ Bᵀ. Layout support has widened across recent GPU generations, but that flexibility is about the data, not the scales, and it does not dissolve the two-copy problem: NVFP4's per-16 block scales are pinned to the quantization axis, so each pass still needs its scales grouped along its own contraction axis. The binding constraint is the scale axis rather than the data transpose, which is why even a layout-flexible SM100 kernel still wants a separate NVFP4 quantization per contraction. The bottom line. Take a tiny weight [[1,2,3],[4,5,6]] quantized rowwise, one scale per row: 1 is stored under scale 3, 4 under scale 6. Dgrad reads columns, so it needs a single scale for the block {1,4}, but those two elements were quantized under different row-scales, so the existing scales cannot produce the correct column scale max(1,4)=4. Recovering it means dequantizing each element back to bf16 and re-quantizing, which needs the full-precision values already discarded by FP4. The distinction is worth keeping straight: that Dgrad contracts over out is a mathematical requirement; the FP4 kernel then demanding per-16-block scales along that contraction axis is a hardwa [truncated for AI cost control]