Optimizing Large Language Model Training Using FP4 Quantization

We sincerely thank the reviewer for the thorough and constructive feedback on our work. We appreciate your recognition of the framework’s potential and your insightful questions, which have helped clarify key aspects of our methodology and claims. Below, we address each point in detail.

W1: Speedup Estimation and Kernel-Level Latency Breakdown

A1: We completely agree with the importance of evaluating theoretical speedup. Since we employ a mixed-precision training framework, we analyze the impact of FP4 acceleration on each computational component of a standard Transformer layer. Given a hidden size h, batch size b, and sequence length s, the FLOP breakdown is as follows:

Since the backward pass requires approximately twice the FLOPs of the forward pass, the theoretical FP4 speedup (excluding DGE and OCC overhead) for a 7B model (h=4096, s=2048) is:

$\frac{3*(24bsh^2+5bs^2h+36bsh)}{3*(6bsh^2+5bs^2h+36bsh)}=\frac{24h+5s+36}{6h+5s+36}=3.12$

The reviewer also point out the importance of kernel-level latency breakdown. Currently, we use FP8 GEMM kernels for all matrix multiplications in the above table, leading to a theoretical speedup of:

$\frac{3*(24bsh^2+5bs^2h+36bsh)}{3*(12bsh^2+5bs^2h+36bsh)}=\frac{24h+5s+36}{12h+5s+36}=1.83$

However, in practice, the framework of FP8-LM reports only a 1.28× actual speedup for h=4096 (7B model). Our implementation incurs additional precision casting overhead due to extra precision casting from BF16 to FP4 (for simulation) and FP4 to FP8 (for FP8 GEMM computation). These conversions are unnecessary with specialized FP4 hardware, where quantization could be fused into the GEMM kernel. So although we write the FP4 quantization kernel, it will still cause heavy overhead in the training process, reduceing speedup from 1.28× to approximately 1.08× in real training.

W2: Perplexity (PPL) Evaluation

A2: We agree that PPL is a sensitive metric for language model training. Below, we present the evaluation results:

The results demonstrate that FP4 models achieve comparable or even slightly lower PPL than BF16 models. As expected, larger models achieve lower perplexity under the same training token budget.

Q3: For Equation (8), should there be a $\delta$ in the numerator of the constant term?

A3: Yes, it was a typo error and we're really sorry! We sincerely appreciate your attention to detail and will correct this in the final version.

Q4: Memory Consumption vs. Impact Statement Claim

A4: We appreciate the reviewer’s feedback on this potentially misleading claim. While our experiments confirm reduced GPU memory usage compared to BF16, this reduction primarily results from:

1. FP8 optimizer states in Mixed-precision optimizers.
2. FP8 activation storage (from FP8-LM framework).

The current FP4 online quantization strategy does not reduce GPU memory. Further reductions would require FP4 optimizer states and activation storage, which we leave for future work. To prevent confusion, we will remove the memory claim from the Impact Statement in the final version.

Q5: Quantization Granularity

A5: Yes, each token (activation) or each channel (weight) has a single scaling factor. Currently the provided quantization kernel in Appendix A only covers the quantization of the scaled tensor, and per-channel/per-token max value computation is not yet fused into the kernel.

We sincerely appreciate the reviewer’s insightful feedback, which has strengthened both our analysis and the clarity of our claims. We will incorporate all necessary revisions into the final manuscript.

We sincerely thank the reviewer for the thorough and constructive critique. We appreciate your recognition of the technical effort and alignment with hardware trends, as well as your insightful suggestions. Below, we address your concerns and questions:

W1: Hardware Dependency

A1: We acknowledge the lack of specific hardware as a limitation, as noted in Sec. 6. We provide a theoretical analysis of FP4 acceleration and demonstrate a theoretical 3× speedup (depending on model parameters). The detailed analysis is included in our response to reviewer gQsM (A1) due to space constraints. Additional analysis on DGE and OCC overhead, addressed in reviewer iCFi’s reply (A2), indicates the overhead remains acceptable while preserving accuracy. However, this is not speculative. As reviewer 1Hhi noted, while the absence of hardware testing affects empirical validation, it does not detract from our goal—to explore FP4 quantization for large-scale training in alignment with future hardware trends.

W2: Insufficient Scale (13B parameters, 100B tokens)

A2: We agree that scaling to trillion-parameter models is crucial. However, within academic compute constraints, our focus was to establish FP4’s feasibility. A 13B model trained on 100B tokens represents a reasonable scale in research, even if it is smaller than industrial models (e.g., GPT-4, LLaMA3). Scaling stability is essential but is best explored in future industry-driven work.

W3: Clarification on DGE

A3: We apologize for not including the full mathematical derivation of Eq. 7 due to paper space constraints. We choose power function $f(x)=x^a$ (a<1) as the base differentiable estimation function because it saturates at 1 as x tends to 1, and quickly drops to 0 as x tends to 0, very much in line with the right half of the quantization function. To fully simulate the quantization function, we take absolute values of x and apply sign function for central symmetry: $f(x)=sign(x)\cdot|x|^a$. We then translate the function on the x-axis and the y-axis to move to the quantization range of [0, $\delta$]: $f(x)=\delta(1+sign(x-\frac{\delta}{2})\cdot|x-\frac{\delta}{2}|^a)$. Finally, we adjust the value of the power exponent to control the steepness and write a(a<1) as 1/k(k>1). We will clarify this in the final version to avoid possible confusion.

W4: Memory Usage and Training Speed

A4: We thank the reviewer for the suggestion. Below, we report real-time training memory usage and throughput:

Note that our method is based on FP8-LM framework, where its speed up for three models are 1.03x, 1.24x and 1.27x (paper reported). In our implementation, we need to do extra precision casting for BF16 to FP4 (for simulation) and FP4 to FP8 (for FP8 GEMM computation), leading to large overhead since native FP4 hardware is inaccessible. In other words, evaluating speed performance during a real training process without dedicated hardware support can only serve as a point of reference.

W5: OCC’s Novelty and Computational Burden

A5: We thank the reviewer for pointing out that we did not state the OCC methodology very well in the paper. OCC’s novelty lies in its effective clamping strategies for characterizing numerical ranges and sparse compensation which avoids dense high-precision residuals. This method is suitable for online use and better matches the pre-training task. For the computation burden of $\Delta Y$, please refer to reviwer iCFi's reply (A2) due to the charactor limit. We'll refine these statements in the final version.

Q6: Perplexity (PPL) Metrics

A6: We acknowledge PPL’s importance as evaluation metric. Due to the character limitation, we kindly refer to reviwer gQsM's reply (A2) for detailed PPL results.

Q7: Convergence with Increasing Tokens

A7: Our preliminary analyses of existing scaling trends suggest that extended token training would likely achieve stable convergence. While scaling analysis is crucial for understanding model convergence, resource limitations currently prevent comprehensive token scalability studies. We highlight this as a critical research direction and commit to open-sourcing our framework to support collaborative exploration of token scaling dynamics.

Q8: Open-Sourcing the Code

A8: We sincerely appreciate the reviewer’s interest in implementation details. We are fully committed to open-sourcing the code and will release it once it has been thoroughly organized to ensure clarity and usability.

Thank you again for your rigorous feedback. We hope these clarifications address your concerns.

We sincerely thank the reviewer for the thoughtful evaluation of our work and for recognizing the significance of our contributions. We appreciate your positive feedback on the effectiveness of DGE and OCC, as well as the thoroughness of our experiments. Additionally, we are grateful for your understanding of the hardware constraints that necessitated our simulation-based approach, which, while unavoidable, slightly limits empirical hardware-level assessments.

Q1: OCC’s Computational Overhead During Actual Training Scenarios

A1: The computational overhead of OCC primarily arises from additional sparse matrix multiplications. Specifically, the input activation tensor Y is decomposed as: $Y=Y_c+\Delta Y$, where $\Delta Y$ contains outlier values processed in higher-precision FP8 GeMM. Since $Y_c=\text{clamp}(Y, max=\alpha,min=1-\alpha)$ and $\alpha$ is very close to 1, the $\Delta Y$ matrix remains highly sparse, with a sparsity ratio of $2*(1-\alpha)$.

A detailed theoretical analysis of the computational cost of $\Delta Y$ GeMM is provided in our response to reviewer iCFi (A2). We kindly refer you to those details if you're interested. The results show that under the training configuration of hidden_size=4096, sequence length=2048 and $\alpha=0.99$ (as used inour 7B model training), the theoretical speedup of FP4 training decrease slightly from 3.12 to 2.95 due to OCC. Considering the large accuracy benefits of the OCC method, we think this performance drop is acceptable. Given OCC’s significant accuracy benefits, we consider this tradeoff acceptable. However, careful tuning of $\alpha$ is crucial to maintain high sparsity, as current hardware struggles with efficient sparse matrix multiplications at lower sparsity levels.

However, the situation in real training scenarios is a little different. In our current training setup, FP8 GeMM is used to simulate FP4 behavior, and $\Delta Y$ is naturally handled in FP8 format. Consequently, both $Y_c$ and $\Delta Y$ are computed within FP8 GeMM, allowing their results to be directly summed without requiring an additional GeMM operation, thereby eliminating OCC-related overhead in this simulation setup.

At present, the primary computational overhead stems from:

1. DGE computations (negligible).
2. Precision casting between BF16, FP4, and FP8 (significant).

However, if native FP4 hardware becomes available and is used for $Y_c$ computations, then OCC’s computational cost must be carefully reassessed.

Once again, we appreciate your constructive feedback and your recognition of this work’s broader relevance in advancing resource-efficient LLM training.

We sincerely thank the reviewer for the thorough and insightful evaluation of our work, as well as for recognizing the theoretical and practical contributions of our FP4 training framework. Below, we address the questions and suggestions:

Q1 (Theroretical Claims): Clarification on OCC’s $\alpha$ value inconsistency (Sec. 3.2 vs. Sec. 4.3)

A1: Thank you very much for pointing out this issue. This is indeed a typo in Section 4.3—the correct statement should be that a lower $\alpha$ leads to better model accuracy. Theoretically, a lower $\alpha$ results in reduced quantization loss, as stated in Section 3.2 and shown in Table 1. Experimentally, Figure 6(c) in Section 4.3 confirms that lower $\alpha$ corresponds to a lower training loss curve. We will correct this in the final version. Again, we appreciate your keen observation!

Q2 (Experimental Designs Or Analyses): Details on speed evaluation / Theoretical analysis on the computation overhead or other cost

A2: Thank you for the suggestion. We recognize the importance of analyzing FP4’s theoretical speedup and computational overhead. The theoretical speedup of FP4 (excluding DGE and OCC overhead) is:

$\frac{3*(24bsh^2+5bs^2h+36bsh)}{3*(6bsh^2+5bs^2h+36bsh)}=\frac{24h+5s+36}{6h+5s+36}=3.12 \quad (h=4096, s=2048)$

For detailed analysis on the decomposition of computational components, we kindly refer the reviewer to the response to reviewer gQsM (A1) due to space constraints.

Regarding computational overhead:

• DGE overhead: DGE introduces an additional nonlinear function during GEMM backward for weight updates, adding ~8 FLOPs per input element (Eq. 8). Accumulating all GEMM operation, this will cause a total overhead of $8*(3bsh+bsh+4bsh+4bsh)=96bsh$ (Four additions in brackets are weight shapes of attn qkv proj, attn out proj, MLP up proj and MLP down proj). Note that this overhead occurs only once per forward-backward iteration.

• OCC overhead: OCC incurs extra sparse matrix multiplication. FP8 sparse GEMM is used with an activation sparsity of $2(1-\alpha)$. These FP8 sparse GEMM are added into every GEMM computation, adding an extra $2*(1-\alpha)(12bsh^2)$ FLOPs, where $12bsh^2$ comes from the cumulative sum of all GeMM calculations in the decomposition table that can be accelerated by FP4 (attn qkv proj, attn out proj, MLP up proj and MLP down proj), but in FP8 format with twice as many FLOPs as in FP4 $\bigg(2*(1.5bsh^2+0.5bsh^2+2bsh^2+2bsh^2)\bigg)$. Since we set $\alpha=0.99$, meaning these GEMMs are highly sparse, the overhead remains small, though hardware inefficiencies in sparse matrix multiplication necessitate choosing a larger $\alpha$ to ensure the high sparsity of the $\Delta Y$ matrix.

Accounting for these overheads, the adjusted theoretical speedup is:

$\frac{3*(24bsh^2+5bs^2h+36bsh)}{3*(6bsh^2+5bs^2h+36bsh+2*(1-\alpha)(12bsh^2))+96bsh}=\frac{24h+5s+36}{6h+24*(1-\alpha)h+5s+68}=2.95 \quad (h=4096, s=2048, \alpha=0.99)$

The overhead from DGE and OCC acounts for $32/(6h+5s+36)=0.1\\%$ and $24(1-\alpha)h/(6h+5s+36)=5.6\\%$ of the total computation, respectively, reducing the theoretical speedup ratio for FP4 compared to BF16 reduced from 3.12 to 2.95. We believe this is an acceptable trade-off between accuracy and efficiency since they can largely reduce quantization errors during training.

Q3 (Experimental Designs Or Analyses): Loss curve spikes in Fig. 5 (larger models):

A3: Thank you for the insightful observation. Larger models (e.g., 13B) exhibit more pronounced loss spikes compared to smaller models (e.g., 7B, 1.3B). While similar spikes occur in BF16, they are more frequent and severe under FP4 due to its limited representation range and the larger number of elements per FP4 quantization vector, increasing quantization error. Addressing this may require more aggressive accuracy compensation strategies, like decreasing the parameter $\alpha$. Additionally, the issue may stem from the optimizer, as we currently use a FP8-based mix-precision optimizer. Larger models may demand higher precision for the optimizer.

Thank you again for your valuable comments. We appreciate your constructive feedback, which has helped us clarify key aspects of our methodology and presentation.