Gumbel-softmax adaptive quantization method and system for NVFP4 format

Abstract

The application relates to the technical field of artificial intelligence model compression, and discloses a Gumbel-Softmax adaptive quantization method and system for NVFP4 format and a storage medium. The method determines a tensor-level scaling factor, an E4M3 block-level scaling factor and an initial FP4 code, constructs a legal local candidate code set in an FP4 E2M1 amplitude grid around the initial code, optimizes logits under the driving of layer or block output reconstruction loss through Gumbel-Softmax or straight-through Gumbel-Softmax, and then hardens the legal FP4 code and projects the block-level scaling factor into E4M3 to output a legal NVFP4 weight representation. Thus, the optimizable space of the FP4 code is expanded without introducing a custom codebook, which helps to reduce reconstruction error and improve low-bit inference deployment compatibility.

Description

Technical Field

[0001] This application relates to the fields of artificial intelligence, deep learning model compression and low-bit quantization technology, and in particular to a Gumbel-Softmax adaptive quantization technique for the NVFP4 format. Background Technology

[0002] With the increasing application of Transformers, large language models, visual language models, and multimodal models in cloud inference, edge inference, and edge smart devices, the pressure on weight storage, GPU memory usage, and memory access bandwidth brought about by the scale of model parameters is becoming increasingly prominent. For example, in online services of large language models, the same model often needs to support concurrent inference for multiple users; in edge or edge devices, the model needs to process text, images, or multimodal inputs under limited GPU memory and power consumption conditions; in long context inference scenarios, weight access and activation propagation further amplify the engineering requirements of low-precision deployment. Therefore, how to reduce the representation cost of weights while maintaining the quality of model output as much as possible has become an important technical problem in the deployment of neural network models.

[0003] Low-bit quantization is a crucial means of reducing the aforementioned costs. Traditional INT4 quantization typically employs a uniform integer grid, which, while simple to implement, struggles to adequately match the long-tailed distribution and high-amplitude weights prevalent in large model weights. The NVFP4 format, composed of FP4 E2M1 amplitude codes, E4M3 block-level scaling factors, and tensor-level scaling factors, offers better adaptability to hardware inference deployments. However, existing NVFP4 quantization methods often employ fixed nearest-neighbor rounding or learn only the binary rounding direction between adjacent codes, resulting in a relatively limited optimization space. On the other hand, while ordinary differentiable quantization methods can handle discrete candidate selections, without constraints on the FP4 codes, E4M3 block-level scaling factors, and the final hardened format, they can easily generate quantization results that are difficult for NVFP4 inference engines to use directly. Furthermore, the determination of the block-level scaling factor and the rounding of weighted codes are often handled separately, making joint optimization difficult. Therefore, there is an urgent need for an adaptive quantization scheme that can expand the optimizable selection space of FP4 codes and coordinate block-level scaling factors while maintaining the legality of the NVFP4 hardware format. Summary of the Invention

[0004] The purpose of this application is to provide a Gumbel-Softmax adaptive quantization method, system, and storage medium for the NVFP4 format, which can alleviate the problems of limited optimization space of FP4 codes, disconnect between block-level scaling factor and weighted FP4 code selection, and inconsistency between optimization results during training and valid formats during deployment in existing round-last or binary rounding methods, while maintaining the validity of the NVFP4 format.

[0005] This application provides a Gumbel-Softmax adaptive quantization method for the NVFP4 format, including: S1. Obtain the weight tensor to be quantized and calibrate the activation; S2. According to the NVFP4 format, the tensor-level scaling factor and the E4M3 block-level scaling factor are determined by the weight tensor, and the normalized magnitude of each weight is taken as the nearest code in the FP4 E2M1 magnitude grid as the initial FP4 code of that weight; the normalized magnitude is obtained by dividing the magnitude of the weight by the product of the tensor-level scaling factor and the E4M3 block-level scaling factor; S3. Using the index of the initial FP4 code obtained in step S2 in the FP4 E2M1 amplitude grid as the center, construct a local candidate code set for the weight, wherein each candidate code in the local candidate code set belongs to the FP4 E2M1 amplitude grid. S4. Set trainable logits for the local candidate code set obtained in step S3, and perform differentiable selection on the local candidate code set through Gumbel-Softmax or pass-through Gumbel-Softmax to obtain the soft quantized code with the weight. S5. Soft quantization weights are constructed using the soft quantization code obtained in step S4, the tensor-level scaling factor obtained in step S2, and the E4M3 block-level scaling factor. Soft quantization output is calculated based on the calibration activation and the soft quantization weights. Full-precision output is calculated based on the calibration activation and the weight tensor. The logits in step S4 are optimized using the reconstruction loss between the soft quantization output and the full-precision output. S6. For the logits optimized in step S5, take the candidate code corresponding to the maximum value as the hard FP4 code of the weight, and determine or project the E4M3 block-level scaling factor as an E4M3 representable value to obtain the hard block-level scaling factor. S7. Determine the quantization weight value corresponding to the weight using the sign of the weight, the tensor scaling factor, the hard block scaling factor obtained in step S6, and the hard FP4 code, and output a valid NVFP4 weight representation including the sign of the weight, the hard FP4 code, the hard block scaling factor, and the tensor scaling factor.

[0006] In a preferred embodiment, the FP4 E2M1 amplitude grid is: Its code index is 0 to 7, and it is combined with the weight sign bit to form a signed FP4 value; The initial FP4 code mentioned in step S2 is determined as follows: the normalized magnitude of the weight is calculated. , and according to Determine the index of the initial FP4 code in the FP4 E2M1 amplitude grid. ; The quantization weight value corresponding to the legal NVFP4 weight representation in step S7 satisfies: In the formula, For the first Each weight, Its amplitude, The tensor scaling factor is... This is the E4M3 block-level scaling factor of the block to which the weight belongs. For normalized amplitude, The index in the FP4 E2M1 amplitude grid is The code, This indicates that the index that minimizes the absolute difference is returned. , This is the index of the initial FP4 code. This is the quantized weight value corresponding to this weight. The sign of this weight. This is a block-level scaling factor. It is a hard FP4 code and is taken from .

[0007] In a preferred embodiment, the local candidate code set described in step S3 is constructed as follows: In the formula, The initial FP4 code for this weight is the index in the FP4 E2M1 amplitude grid. The integer offset relative to that index. For the local search radius, This means truncating the input to a closed interval between 0 and 7. This is the set of candidate indices after deduplication. The index in the FP4 E2M1 amplitude grid is The code, The local candidate code set; Wherein, the local search radius When the value is 1, the local candidate code set contains at most 3 candidate codes. When the value is 2, the local candidate code set contains a maximum of 5 candidate codes.

[0008] In a preferred embodiment, the soft quantization code described in step S4 is obtained as follows: Gumbel noise is sampled for each candidate code in the local candidate code set. ,in obey Uniform distribution on the surface; calculate the selection probability of each candidate key. The soft quantization code is obtained by the following formula: In the formula, For the first The weight of the first Gumbel noise for candidate codes, For the corresponding uniformly distributed random numbers, For the first The logits corresponding to each candidate key. For temperature coefficient, It is a normalized exponential function, and its normalization is performed on candidate keys with the same weight. For the first The probability of choosing a candidate key. For the first One candidate code, This refers to the soft quantization code; And the temperature coefficient During the optimization process, the value decays from the initial value to the final value.

[0009] In a preferred embodiment, after determining the E4M3 block-level scaling factor in step S2, a learnable parameter is set for each block. Let the floating-point block scaling factor be 1. In step S5, the E4M3 pass-through estimator is used to determine the actual block-level scaling factor used in the forward computation. Therefore, in step S5, the tensor scaling factor is fixed, and the logits and the learnable parameters are jointly optimized. Furthermore, the loss used for optimization in step S5 is based on the reconstruction loss and also includes a scaling regularization term that constrains the E4M3 block-level scaling factor to not deviate from the E4M3 block-level scaling factor determined in step S2. In step S6, ( As the hard block-level scaling factor; In the formula, For the first Learnable parameters for each block, For the reason The resulting floating-point block-level scaling factor, It is a natural exponential function. This indicates that the input is projected as an E4M3 representable value. This indicates the gradient whose value within the brackets is truncated during backpropagation. This refers to the block-level scaling factor actually used in the forward computation. For the first Block-level scaling factor for each block, The E4M3 block-level scaling factor (i.e., determined in step S2) for this block. ), This is to take the average value over all blocks.

[0010] In a preferred embodiment, the search radius of the local candidate code set in step S3 is adaptively determined as an asymmetric radius according to the weights, including: Identify high-amplitude sparse regions ,in This is the index of the point where the adjacent code spacing begins to increase in the FP4 E2M1 amplitude grid. This is an indicator function; it takes the value 1 if the condition within the parentheses is true, and 0 otherwise. Determine the offset direction indication and , and when season ,in This is the normalized magnitude of the weight. This is the initial FP4 code for this weight. The preset threshold; The block to which this weight belongs satisfies the following conditions. weighting percentage ,when Exceeding the preset percentage threshold The base radius of the block will be determined at this time. Take as Otherwise, take as , Base radius; The asymmetric radii in the high-index direction and the low-index direction are synthesized according to the following formula: and with Construct a candidate index set for this weight, thereby redistributing the number of candidate codes from the low-amplitude dense region to the high-amplitude sparse region; In the formula, For the radius increment of the high amplitude region, For asymmetric expansion increment, The upper limit of the radius, To obtain the minimum value, and These are the search radii for the direction towards higher indexes and the direction towards lower indexes, respectively.

[0011] In a preferred embodiment, the loss used in step S5 to optimize the logits is further supplemented by a regularization term on top of the reconstruction loss. The regularization term includes at least one of an entropy regularization term and a moving regularization term; wherein: Entropy regularization term This is used to make the selection probability distribution of the differentiable selection approach the one-thermal distribution; Move regular expression This is used to limit the offset of the final FP4 code relative to the initial FP4 code; In the formula, For the first The probability distribution of the selection of each candidate key with each weight. Its entropy, For the first The probability of choosing a candidate key. For the first One candidate code, Its index in the FP4 E2M1 amplitude grid, This is the index of the initial FP4 code. This is to take the average of all weights.

[0012] In a preferred embodiment, the method is performed on the Transformer model layer by layer: after performing steps S1 to S7 on the current layer to obtain the valid NVFP4 weight representation of the layer, the current layer is replaced with the valid NVFP4 weight representation, and then the calibration activation of the next layer is collected and steps S1 to S7 are performed on the next layer, so that the calibration activation of the subsequent layer includes the quantized error of each upstream layer.

[0013] This application provides a Gumbel-Softmax adaptive quantization system for the NVFP4 format, comprising: The acquisition module is used to acquire the weight tensor to be quantized and the calibration activation. The determination module is used to determine the tensor-level scaling factor and the E4M3 block-level scaling factor from the weight tensor according to the NVFP4 format, and to take the nearest code in the FP4 E2M1 amplitude grid of each weight as the initial FP4 code of that weight. The normalized amplitude is obtained by dividing the amplitude of the weight by the product of the tensor-level scaling factor and the E4M3 block-level scaling factor. A construction module is used to construct a local candidate code set for the weight, centered on the index of the initial FP4 code in the FP4 E2M1 amplitude grid, wherein each candidate code in the local candidate code set belongs to the FP4 E2M1 amplitude grid. The selection module is used to set trainable logits for the local candidate code set and perform differentiable selection on the local candidate code set through Gumbel-Softmax or pass-through Gumbel-Softmax to obtain the soft quantized code with the weight. An optimization module is used to construct soft quantization weights from the soft quantization code, the tensor-level scaling factor, and the E4M3 block-level scaling factor; calculate a soft quantization output based on the calibration activation and the soft quantization weights; calculate a full-precision output based on the calibration activation and the weight tensor; and optimize the logits using the reconstruction loss between the soft quantization output and the full-precision output. The hardening module is used to take the candidate code corresponding to the maximum value of the optimized logits as the hard FP4 code of the weight, and to determine or project the E4M3 block-level scaling factor as the hard block-level scaling factor that can be represented by E4M3. The output module is used to determine the quantized weight value corresponding to the weight by the sign of the weight, the tensor scaling factor, the hard block scaling factor and the hard FP4 code, and output a valid NVFP4 weight representation including the sign of the weight, the hard FP4 code, the hard block scaling factor and the tensor scaling factor.

[0014] This application provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the above-described method.

[0015] This application uses a set of complementary technical means to expand the optimizable space of FP4 code and improve the quantization reconstruction effect while maintaining the legality of the NVFP4 hardware format. The technical effects can be explained from the following aspects.

[0016] Since the tensor scaling factor, E4M3 block scaling factor, and initial FP4 code are determined according to the NVFP4 format after obtaining the weight tensor and calibrating the activation, and a local candidate code set containing only valid FP4 codes is constructed around the initial code, and the reconstruction loss drives the differentiable selection of Gumbel-Softmax or pass-through Gumbel-Softmax, and finally hardens it into valid FP4 codes and determines or projects the block scaling factor to an E4M3 representable value, the entire optimization process is within the representation range allowed by the target hardware format from beginning to end. This avoids generating custom codebooks or high-precision scaling representations that are difficult to deploy during the training phase, and transforms the discrete code selection, which was originally non-differentiable due to format limitations, into a candidate allocation problem that can be optimized through backpropagation. At the same time, since the code selection is determined by the overall reconstruction state of the layer output or block output, rather than just by the point-to-point distance of a single weight, the candidate codes of multiple weights can cooperate with each other under the same output reconstruction target, which helps to suppress the accumulation of local quantization errors in the layer output.

[0017] Limit the FP4 amplitude code to a non-uniform grid. It is combined with the sign bit to form a signed FP4 value, which provides a value basis consistent with the hardware format for the construction and legal hardening of local candidate codes, so that the selected code and the code obtained by the final hardening must fall within the legal grid.

[0018] When the local candidate code set is constructed as a three-way or five-way candidate code according to the search radius, on the one hand, the weight is allowed to cross more than one adjacent code under the drive of the reconstruction target, so as to better adapt to the distribution characteristics of low amplitude dense and high amplitude sparse in the grid; on the other hand, the number of candidate codes is constrained to a small range, avoiding the memory overhead caused by setting logits for the entire grid.

[0019] By using Gumbel noise to calculate the selection probability and gradually decaying the temperature coefficient from the initial value to the final value, the optimization can maintain a certain exploratory capability in the early stage and tend to make discrete selection in the later stage. This helps the soft selection result to smoothly transition to hard selection, thereby reducing the difference between the soft training stage and the hard deployment stage.

[0020] When the block scaling factor is made through a learnable parameter When the block-level scaling factor participates in the same reconstruction objective as the E4M3 pass-through estimator and is fixed with a scaling regularization term, it can participate in optimization while maintaining the E4M3 deployment constraints, and can also be adapted to the weighted code selection, thereby alleviating the coupling error between high-amplitude codes and scaling factors and suppressing the degradation of mutual compensation between the two.

[0021] To address the non-uniformity of the grid, the search radius is determined to be an asymmetric radius with weighted increments and adjusted in conjunction with the proportion of high-amplitude weights within the block. This allows the candidate search budget to be appropriately shifted from low-amplitude dense areas to high-amplitude sparse areas, enabling wider code intervals to achieve more effective candidate coverage. Consequently, the effectiveness of local multi-candidate search is improved with limited memory and computational budget.

[0022] Entropy regularization helps make the candidate selection probability approach a one-hot distribution, while shift regularization is used to limit the excessive shift of the final FP4 code relative to the initial code. Together with temperature decay and hardening processes, these two factors allow the weights to explore locally when necessary and suppress degenerative shifts, thereby jointly reducing the difference between soft training results and hard deployment results.

[0023] Quantization is performed layer by layer, and the current layer is hardened and replaced with a valid NVFP4 representation before the calibration activation of the next layer is collected. This allows the optimization input of subsequent layers to reflect the activation distribution after upstream quantization, making the layer-by-layer optimization process closer to the actual state when the model is deployed.

[0024] The above method can be further implemented by carrying the corresponding computer program in the system's functional modules or in the storage medium, so that the aforementioned technical effects can be implemented in specific computer systems and engineering implementations, and the quantification results can be directly applied to the inference system that supports NVFP4, rather than just remaining at the numerical level.

[0025] In summary, the final output, which includes the weight sign bit, hard FP4 code, E4M3 hard block scaling factor, and tensor scaling factor, is a valid NVFP4 weight representation that can serve the inference deployment of neural networks such as large language models under conditions of low storage, low GPU memory usage, and low memory access bandwidth while maintaining format compatibility. Furthermore, through the combined use of techniques such as locally valid multiple candidate code construction, differentiable selection driven by reconstruction loss, joint adaptation of E4M3 block scaling factor, and valid hardened output, it brings corresponding improvements in terms of storage, GPU memory, and memory access bandwidth.

[0026] The specification of this application contains numerous technical features distributed across various technical solutions. Listing all possible combinations of these technical features (i.e., technical solutions) would make the specification excessively lengthy. To avoid this problem, the various technical features disclosed in the above-described invention, the various technical features disclosed in the following embodiments and examples, and the various technical features disclosed in the accompanying drawings can be freely combined to form various new technical solutions (all of which should be considered as described in this specification), unless such a combination of technical features is technically infeasible. For example, one example discloses feature A+B+C, and another example discloses feature A+B+D+E. Features C and D are equivalent technical means that serve the same function, and technically only one needs to be used; they cannot be used simultaneously. Feature E can technically be combined with feature C. Therefore, the solution A+B+C+D should not be considered as described because it is technically infeasible, while the solution A+B+C+E should be considered as described. Attached Figure Description

[0027] Figure 1 is a flowchart illustrating a Gumbel-Softmax adaptive quantization method for NVFP4 format according to an embodiment of this application.

[0028] Figure 2 is a schematic diagram of the structure of a Gumbel-Softmax adaptive quantization system for NVFP4 format according to an embodiment of this application. Detailed Implementation

[0029] In the following description, many technical details are presented to help the reader better understand this application. However, those skilled in the art will understand that the technical solutions claimed in this application can be implemented even without these technical details and various variations and modifications based on the following embodiments.

[0030] Explanation of some concepts: The NVFP4 format is a floating-point quantization representation for low-precision inference deployments. It is typically composed of a weight sign bit, an FP4 E2M1 amplitude code, an FP8 E4M3 block-level scaling factor shared by each microblock, and a tensor-level scaling factor. The amplitude grid it can represent is non-uniform.

[0031] The FP4 E2M1 amplitude grid refers to the set of non-negative amplitude codes that can be represented in the FP4 E2M1 format, exemplified by... The indices are 0 to 7, and each element in the grid is combined with the weight sign bit to form a signed FP4 value.

[0032] The FP8 E4M3 block-level scaling factor is a scaling factor that is represented in an octet floating-point format with a 4-bit exponent and a 3-bit mantissa, is shared by the same weighted block (microblock), and is limited to or projected to an E4M3 representable value.

[0033] Tensor-level scaling factor refers to the scaling factor used to determine the overall weight tensor to be quantized. It is used together with block-level scaling factor to determine the scale when the weight magnitude is mapped to the FP4 magnitude grid.

[0034] Microblocks refer to weight groups that share the same E4M3 block-level scaling factor, obtained by dividing the weight tensor according to a preset block size.

[0035] Calibration activation refers to the activation data obtained from calibration samples through model forward propagation during the quantization phase. It is used to calculate the reconstruction difference between the quantized output and the full-precision output, rather than to retrain the original model parameters.

[0036] Normalized magnitude is the value obtained by dividing the magnitude of a weight by the product of the tensor scaling factor and the E4M3 block scaling factor of its block. It is used to determine the candidate code for that weight in the FP4 E2M1 magnitude grid.

[0037] The initial FP4 code refers to the initial quantized code obtained by performing the nearest match in the FP4 E2M1 amplitude grid based on the normalized amplitude, which serves as the center for constructing the local candidate code set.

[0038] A local candidate code set refers to a set of several valid FP4 codes constructed by truncating and deduplicating the initial FP4 code within a preset search radius, centered on the index of the initial FP4 code in the amplitude grid.

[0039] logits refers to the trainable selection parameters set for each candidate code in a local candidate code set.

[0040] Gumbel-Softmax is a method that uses Gumbel noise and the Softmax function to make a differentiable approximation of discrete candidate options. It obtains the selection probability of each candidate option by setting trainable logits for the candidate options, superimposing Gumbel noise, and normalizing it with a temperature coefficient using Softmax, thereby enabling the discrete selection to be optimized through backpropagation.

[0041] Straight-through Gumbel-Softmax refers to a variation of Gumbel-Softmax that uses hard selection in one-hot form for forward propagation and soft probability approximation gradient for backpropagation, in order to make the forward computation closer to the hardened result of the deployment stage.

[0042] Soft quantization code refers to a continuous approximate quantization code obtained by weighted summation of local candidate codes based on the selection probability of candidate codes during the training phase.

[0043] Reconstruction loss refers to the loss calculated based on the difference between the output corresponding to the soft quantization weights and the output corresponding to the full precision weights, calculated separately based on the calibration activation. The object can be the layer output or the Transformer block output.

[0044] The E4M3 pass-through estimator refers to an estimation method that uses the E4M3 quantized value of the block-level scaling factor in the forward computation and preserves the approximate gradient of the corresponding learnable parameters in the backpropagation.

[0045] Hard FP4 code refers to the discrete FP4E2M1 code determined after optimization based on the candidate code corresponding to the maximum value of logits.

[0046] Hardening refers to the process of converting the soft selection results of the training phase into discrete FP4 codes and E4M3-representable block-level scaling factors for the deployment phase.

[0047] Layer-by-layer sequential quantization refers to the quantization process of quantizing the Transformer model layer by layer, and collecting the calibration activation of the next layer after hardening and replacing the current layer.

[0048] A valid NVFP4 weight representation is a weight representation that includes a weight sign bit, a hard FP4 code, an E4M3 hard block-level scaling factor, and a tensor-level scaling factor, and satisfies the NVFP4 format constraints.

[0049] The following is a brief summary of some of the innovative aspects of this application: In summary, the inventiveness of this application lies not in simply applying the differentiable discrete choice method of Gumbel-Softmax to low-bit quantization, nor in adding a trainable switch to the existing NVFP4 nearest-nearest rounding rule, but in: targeting the weight tensor of neural networks such as large language models. In the specific scenario of deployment on inference hardware supporting the NVFP4 format, the final output must consist of weight sign bits taken from a non-uniform FP4 E2M1 amplitude grid. Amplitude code, FP8 E4M3 block-level scaling factor shared by each microblock and tensor scaling factor Under the constraint that "the legal representation constituted by all these elements must not rely on any custom four-bit codebook," the construction of local candidate codes, differentiable class-based selection, reconstruction loss driving, E4M3 format constraints, and hardened output are organized into a consistent and mutually supportive quantization optimization closed loop. This application does not use fixed selection and normalization amplitudes. The closest code is not simply a choice between two adjacent lower and higher code points, but rather... Obtain the initial code index Then, by truncating and deduplicating around this index, a local candidate code set containing only valid FP4 E2M1 codes is constructed. This anchors the goal of "expanding the selection space" within the hardware format from the outset. Furthermore, through... Set logits The selection probability is obtained from Gumbel noise and temperature coefficient. And weighted into soft quantization code This allows the previously discrete and format-constrained FP4 code selection to accept layer output or block output reconstruction loss during the training phase. The reverse drive ensures that code selection serves the reconstruction of the model output rather than the minimization of pointwise distance. Furthermore, the block-level scaling factor is not isolated from weighted code selection but is instead integrated via learnable parameters. The E4M3 pass-through estimator is integrated with the same reconstruction objective for collaborative adaptation, and the four / six-block mode selection can further coordinate with the determination of the block-level scale in the high-amplitude range. Therefore, the aforementioned technical features are not independent stacked steps, but rather an organic whole formed around the technical problem of "improving quantization reconstruction performance without changing the NVFP4 format or introducing a custom codebook." This allows the soft selection during the training phase to be smoothly transformed into the hard FP4 code during the deployment phase. With block-level scaling factor and eventually formed The corresponding valid NVFP4 weight representation achieves a more targeted technical balance between model compression, memory usage, memory access bandwidth, and inference deployment compatibility. This overall effect, which is brought about by the interweaving of format constraints, locally valid multiple candidates, reconstruction-driven approach, and scale-code joint adaptation, helps to improve quantization reconstruction results while maintaining the validity of the hardware format.

[0050] Furthermore, through long-term in-depth research, the inventors of this application have discovered that the key difficulties encountered by existing low-bit quantization schemes in NVFP4 inference deployment are not merely the superficial phenomenon of "limited precision of four-bit representation," but rather the failure to form a unified optimization loop between the discrete code selection space, block-level scaling factor constraints, and the final hardware-readable format. For the nearest rounding method, the judgment is mainly based on the distance between the normalized magnitude of a single weight and a certain FP4 code point. After in-depth analysis, the inventors found that during model inference, the smallest point-to-point distance for a weight does not necessarily mean the smallest output reconstruction error for the layer or Transformer block containing that weight—especially after calibration activation participates in the calculation, the errors of multiple weights will collectively affect the output through matrix operations. This makes the locally nearest code not necessarily a more suitable choice under the output reconstruction target.

[0051] The inventors further discovered that the FP4 E2M1 amplitude grid of NVFP4 itself has non-uniform characteristics, with denser code points in low-amplitude regions and sparser code points in high-amplitude regions. If the binary rounding approach between adjacent lower and higher codes is still used, the optimizable space will be limited to an excessively narrow local range, making it difficult to cover the wider code intervals in the sparse high-amplitude regions. Meanwhile, the block-level scaling factor and the FP4 code for each weight are not independent of each other: the E4M3 scaling factor within the same block will change the normalization position of multiple weights when mapped to the grid, and the FP4 code ultimately selected by each weight will in turn affect the output reconstruction state of that block. Existing methods struggle to effectively solve this problem because they typically determine the scaling factor before rounding, or use codebooks and scaling values ​​unconstrained by the NVFP4 format in ordinary differentiable quantization, leading to differences between the optimization results during training and the valid format during deployment. Based on the above in-depth research, the inventors of this application propose a technical solution for differentiable selection of local multi-candidate FP4 codes under NVFP4 format constraints: First, an initial FP4 code is determined based on the tensor scaling factor and the E4M3 block scaling factor. Then, a local candidate code set containing only valid FP4 E2M1 codes is constructed around the initial code. The candidate selection parameters are optimized under reconstruction loss driven by Gumbel-Softmax or pass-through Gumbel-Softmax. In an optional implementation, the E4M3 block scaling factor is also involved in joint adaptation under the same objective and hardened to a valid NVFP4 weight representation after optimization. The implementation process of this application is described in detail below through specific embodiments.

[0052] The technical solutions of this application will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only for explaining this application and are not intended to limit the scope of protection of this application; the technical features in the following embodiments and the technical features between different embodiments can be combined with each other without conflict. For training frameworks, optimizers, data loading methods, tensor arrangement methods and inference engine interfaces not specifically limited in this application, conventional implementation methods in the art can be adopted; for FP4 E2M1 codes, E4M3 block-level scaling factors, tensor-level scaling factors and weight sign bits related to the NVFP4 format, the final result should be able to form a valid NVFP4 weight representation. To facilitate understanding by those skilled in the art, some English technical terms are reserved in this embodiment, and their meanings shall be as commonly understood in the art.

[0053] It should be noted that this application deals with neural network weights in the low-precision floating-point format NVFP4 for modern accelerated hardware. The NVFP4 weight representation is exemplarily composed of weight symbols, FP4 E2M1 amplitude codes, a shared FP8 E4M3 block-level scaling factor for each micro-block, and a tensor-level scaling factor, and its representable amplitude grid is non-uniform. The core of this application is not to propose a new four-bit numerical format, but rather to extend the selection of the FP4 code for each weight from a fixed nearest-last rounding or binary selection between two adjacent codes to a differentiable discrete selection among several valid FP4 codes, under the constraints of existing valid NVFP4 formats. This selection can be further coordinated with the E4M3 block-level scaling factor, thereby improving low-bit quantization accuracy while maintaining hardware format compatibility. In an illustrative comparison, three strategies for determining FP4 codes can be compared: round-to-nearest rounding always selects the nearest code; binary format-aware rounding selects only one of the adjacent lower and upper codes; while this application performs differentiable categorical assignment among multiple valid FP4 codes around the initial code.

[0054] Example 1: Gumbel-Softmax Adaptive Quantization Method for NVFP4 Format This embodiment provides a Gumbel-Softmax adaptive quantization method for the NVFP4 format, which can be applied to post-training quantization of weight tensors in linear layers, attention projection layers, and feedforward network projection layers in neural network models such as Transformer, large language models, visual language models, and multimodal models. Referring to the overall training and hardening process shown in Figure 1, for ease of explanation, the following description uses the weight tensor of a linear layer as an example; when applied to a Transformer model, this linear layer can be a query, key, or value projection layer or an attention output projection layer in the attention module, or an up, gate, or down projection layer in the feedforward network module. The method exemplarily includes the following steps. Steps 100 to 700 correspond to steps S1 to S7, respectively.

[0055] Step 100: Obtain the weight tensor to be quantized and the calibration activation.

[0056] Specifically, the weight tensor to be quantized is denoted as , of which Each weight is denoted as Calibration activation is recorded as This can be derived from the activation output of calibration samples used for post-training quantization in the current model or a partially quantized model. It's important to note that calibration activation is not used to retrain the original model parameters, but rather to measure the difference between the layer or block output produced by the soft-quantized weights and the full-precision output during the quantization phase. Therefore, subsequent FP4 code selection is not solely based on the point-by-point distance from a single weight to the quantized point, but rather reflects the contribution of the layer or module containing that weight to the reconstruction of the model output.

[0057] Step 200: According to the NVFP4 format, determine the tensor-level scaling factor and the E4M3 block-level scaling factor from the weight tensor, and take the nearest code in the FP4 E2M1 amplitude grid of the normalized magnitude of each weight as the initial FP4 code of that weight.

[0058] More specifically, this step can be further subdivided into the following sub-steps.

[0059] Step 210: Determine the FP4 E2M1 amplitude grid. The FP4 E2M1 amplitude grid is a discrete set of non-negative amplitudes, exemplarily represented as... in, Each element in the code represents a non-negative amplitude value that can be represented by FP4 E2M1, with indices ranging from 0 to 7. Indicates that the index is The code. The sign bit of the weight and The amplitude codes in the grid are combined to form a signed FP4 value. It should be noted that this grid is non-uniform: in the low-amplitude region, the interval between adjacent codes is 0.5, while in the high-amplitude region, the interval gradually increases to 1 or even 2, for example, the intervals from 2 to 3, 3 to 4, and 4 to 6 increase sequentially. It is precisely this non-uniform characteristic that makes simple nearest rounding or binary selection between adjacent codes difficult to fully utilize the grid's expressive power in the high-amplitude sparse region.

[0060] Step 220: Determine the tensor scaling factor In one exemplary implementation, The scaling factor can be determined based on the maximum magnitude of the weight tensor; in another exemplary implementation, it can also be determined based on the quantiles of the weight magnitudes to reduce the impact of extreme outlier weights on the overall scaling ratio. The above methods are merely examples, as long as they provide a tensor-level scale for subsequent NVFP4 quantization.

[0061] Step 230: Determine the E4M3 block-level scaling factor. Divide the weight tensor into multiple blocks according to a preset block size, with each block corresponding to an E4M3 block-level scaling factor. For the first... For each block, first determine the weight magnitude and tensor scaling factor within that block. Determine the initial block-level scaling factor, then quantize or project it into an E4M3 representable value to obtain... That is, the first The E4M3 block-level scaling factor used during the initialization phase of each block. For cases where the tail is less than a complete block, the block-level scaling factor can be calculated only for the actual weighted elements and then packaged.

[0062] Step 240: Determine the normalized magnitude of each weight and the initial FP4 code. For the... Weights Let its block be the th If there are several blocks, then its normalized magnitude is... in, Indicates the first The normalized magnitude of each weight, Indicates its amplitude, This represents the tensor scaling factor. This represents the E4M3 block-level scaling factor of the block to which it belongs. The purpose of normalized magnitude is to map the original weight magnitude to the scale of the FP4 E2M1 magnitude grid in order to determine candidate codes within that grid. Subsequently, codes are selected from the grid that correspond to... The closest code is used as the initial FP4 code, and its index is... in, Indicates the first The index of the initial FP4 code in the grid for each weight. This indicates that the index that minimizes the absolute difference within the parentheses is returned. It should be noted that the initial FP4 code here can be derived from the nearest rounding, or in other implementations from GPTQ, AWQ, format-aware adaptive rounding, or other quantization initialization methods capable of generating the initial NVFP4 code. This initial FP4 code will serve as the center for subsequently constructing the local candidate code set.

[0063] Step 300: Construct a local candidate code set for the weight, centered on the index of the initial FP4 code in the amplitude grid, wherein each candidate code in the candidate code set belongs to the FP4 E2M1 amplitude grid.

[0064] Specifically, this step does not involve binary selection only between adjacent lower and higher codes in the initial FP4 code, but rather at the initial index. Construct multiple valid candidate codes in the vicinity. Let the local search radius be . Then the first The candidate index set and local candidate code set with weights can be represented as follows: in, Represents the set of candidate indices. Indicates relative initial index Integer offset, This means truncating the input to a closed interval between 0 and 7. Indicates a candidate index. This represents a local candidate code set; if boundary truncation results in duplicate indices, then duplicate indices are removed. For example, when When 1 is selected, the local candidate code set contains at most 3 candidate codes; when When the value is 2, the local candidate code set contains at most 5 candidate codes. For or In the boundary case, the number of candidate codes may be less than the maximum number mentioned above, but each candidate code belongs to... Therefore, it does not introduce custom codebooks or non-NVFP4 valid codes. This construction avoids the memory overhead of full-mesh logits while allowing weights to span more than one adjacent code, driven by the output reconstruction target, thus better adapting to non-uniform meshes.

[0065] Step 400: Set trainable logits for the local candidate code set, and perform differentiable selection through Gumbel-Softmax or pass-through Gumbel-Softmax to obtain the soft quantized code with the weight.

[0066] This step can be further subdivided into the following sub-steps.

[0067] Step 410, for the local candidate code set The first in Candidate codes Configure the corresponding trainable logits For example, logits can be initialized with all candidates having the same value, or the candidate corresponding to the initial code can be assigned a relatively large initial value so that the optimization tends to retain the most recently rounded choice in the early stages.

[0068] Step 420: Sample Gumbel noise and calculate soft quantization code. For the... The weight of the first weight Each candidate code is sampled with Gumbel noise. in, Indicates the first The weight of the first The Gumbel noise corresponding to each candidate code, To show obedience Uniformly distributed random numbers on an open interval. Represent the natural logarithm function; use open intervals. Alternatively, truncating the random numbers at the boundaries can avoid... This causes numerical anomalies. Based on this, the selection probability of each candidate key is calculated. in, Indicates the first The probability of choosing a candidate key. Indicates the first Trainable logits corresponding to each candidate key Indicates the temperature coefficient. This represents a normalized exponential function, with normalization performed on each candidate key with the same weight. Subsequently, the candidate keys are weighted and summed according to their selection probabilities to obtain the soft-quantized code. in, Indicates the first A soft quantization code with weights, Indicates the first in its local candidate code set There are several candidate codes. It should be noted that in one implementation, a soft Gumbel-Softmax can be used, employing continuous selection probabilities for both forward and backward propagation; in another preferred implementation, a straight-through Gumbel-Softmax can be used, where forward propagation uses one-hot hard selection and backward propagation uses a gradient approximation with soft probabilities. This makes the forward computation closer to the hardened results of the deployment phase, thereby reducing the difference between soft training and hard deployment. Furthermore, the temperature coefficient... During the optimization process, the initial value is gradually reduced from a large value to a smaller final value, so as to maintain a certain exploratory ability in the early stage of training and make the selection probability gradually tend to discrete selection in the later stage of training.

[0069] Step 500: Construct soft quantization weights from the soft quantization code, calculate the soft quantization output and full-precision output based on calibration activation, and optimize the logits using the reconstruction loss between the two.

[0070] Specifically, for the first Each weight can be determined based on its sign and tensor scaling factor. The block-level scaling factor and soft quantization code of the block to which it belongs. We obtain the soft quantization weights, which are then used to construct the soft quantization weight tensor. Taking a linear layer as an example, the full-precision output and the soft-quantized output can be expressed as follows: in, Indicates activation based on calibration With full-precision weight tensor The obtained full-precision output, Indicates activation based on calibration With soft quantization weight tensor The resulting soft quantization output, and Let represent the transposes of the corresponding weight tensors, respectively. Based on the difference between them, a reconstruction loss is constructed. in, Indicates the reconstruction loss. This indicates taking the average. This represents the L2 norm. The reconstruction loss measures the reconstruction error of soft-quantized weights relative to full-precision weights at the layer output level. It's important to emphasize that, compared to simply minimizing the pointwise error between a single weight and the nearest quantized point, this approach allows the FP4 code selection to serve the overall reconstruction effect of either the layer output or the module output; from the perspective of layer output reconstruction, the locally nearest FP4 code may not be the optimal code for the model output. Optionally, the object of the reconstruction loss can also be extended to the output of the entire Transformer block, i.e., simultaneously quantizing the attention projection, attention output projection, or feedforward network projection under the same objective, and comparing the output of the entire block to further improve accuracy.

[0071] Furthermore, when optimizing logits, at least one of the entropy regularization term and the moving regularization term can be added to the reconstruction loss. The entropy regularization term can be written as... in, Represents the entropy regularization term. Indicates the first The probability distribution of the selection of each weight on each candidate key. This represents the entropy of the distribution. This represents averaging over all weights. The entropy regularization term helps to make the selection probability distribution approximate a one-hot distribution, thus reducing the difference that arises when transitioning from soft to hard selection after optimization. The shift regularization term can be written as... in, Indicates a move regular expression. Candidate code Index in the grid, This represents the index of the initial FP4 code. The shift regularization term is used to limit the offset of the final FP4 code relative to the initial FP4 code, avoiding uncontrolled excessive shifts after expanding the candidate space. It should be noted that constructing local multiple candidate codes "allows" weights to explore across adjacent codes when necessary, while shift regularization is used to "suppress" unnecessary excessive shifts. The two are not contradictory but rather form a cooperative relationship of "allowing local exploration while suppressing degenerate shifts." Therefore, the total loss can be expressed as... (Formula 14) in, Indicates the total loss. and These represent the weight coefficients of the entropy regularization term and the moving regularization term, respectively. These weight coefficients can be set according to the model size, the scale of the calibration data, and the number of candidate codes. In some implementations, the weights of one or more regularization terms can be set to 0. Backpropagation of the total loss updates the logits corresponding to each candidate code, allowing the candidate code selection to gradually adapt to the output reconstruction target under calibration activation. The optimizer can, exemplarily, employ an adaptive gradient optimizer commonly used in the art.

[0072] Step 600: Take the candidate code corresponding to the maximum value of the optimized logits as the hard FP4 code of the weight, and determine or project the E4M3 block-level scaling factor as an E4M3 representable value to obtain the hard block-level scaling factor.

[0073] Specifically, for the first Each weight is used, and after optimization, the candidate code corresponding to the maximum logits value is selected as the hard FP4 code. in, Indicates the first A weighted hard FP4 code, This represents the candidate key number corresponding to the maximum value of logits. This indicates that the candidate key number that maximizes logits should be returned. Because... All candidate keys in the set belong to Therefore, hardening yields It must be a valid FP4 E2M1 amplitude code. Regarding the block-level scaling factor, if the current implementation uses a fixed block-level scaling factor, then the result obtained in step 200 will be... As a hard block-level scaling factor; if the implementation method of learnable block-level scaling factor described later is adopted, the optimized floating-point block-level scaling factor is projected into an E4M3 representable value, resulting in After hardening, a validity check can be performed, including, for example, checking whether all hardened FP4 codes belong to... Whether all hard block-level scaling factors are E4M3 representable values, whether the block size and tail block processing conform to the packing rules, whether there are non-numbers (NaN) or infinity (Inf) in the quantization weights, and whether the output format can be read by the NVFP4 simulation layer or an inference engine that supports NVFP4.

[0074] Step 700: Determine the quantized weight value corresponding to the weight by the sign of the weight, the tensor scaling factor, the hard block scaling factor and the hard FP4 code, and output a valid NVFP4 weight representation.

[0075] Specifically, no. The quantized weight value corresponding to each weight can be expressed as: (Formula 16) in, Indicates the first The quantized weight value corresponding to each weight. Indicate its symbol, This indicates the hard block-level scaling factor of the block to which it belongs. This represents the hard FP4 code. The final output of a valid NVFP4 weight representation includes the weight symbol, hard FP4 code, hard block-level scaling factor, and tensor-level scaling factor. Therefore, this embodiment uses differentiable discrete selection to optimize the FP4 code during the training phase, while still outputting the standard NVFP4 representation during the deployment phase, without introducing a custom four-bit codebook. This allows it to be read by the NVFP4 simulation layer or an NVFP4-enabled inference engine and used for inference deployment, thereby reducing weight storage, GPU memory usage, and memory access bandwidth.

[0076] Example 2: Implementation of Joint Optimization of E4M3 Block-Level Scaling Factor This embodiment, based on Embodiment 1, further sets learnable parameters for each block, enabling weighted FP4 code selection and the E4M3 block-level scaling factor to adapt collaboratively under the reconstruction loss drive. Specifically, for the first... Each block sets learnable parameters. And set the floating-point block scaling factor to in, Indicates the first The floating-point block-level scaling factor for each block. Indicates the first Learnable parameters for each block, Representing the natural exponential function, logarithmic domain parameterization ensures the scaling factor remains consistently positive. To maintain consistency between the training forward computation and the deployment format, this embodiment employs the E4M3 pass-through estimator in step 500 to determine the actual block-level scaling factor used in the forward computation. (Formula 18) in, This indicates the block-level scaling factor actually used in the forward computation. This indicates that the input will be projected as an E4M3 representable value. This indicates that the gradient of the values ​​within the brackets is truncated during backpropagation. Therefore, the forward computation uses E4M3 quantization values ​​close to the deployment format, while the backpropagation... Preserve the approximate gradient so that the block-level scaling factor can participate in the reconstruction loss optimization.

[0077] Furthermore, to avoid uncontrolled scale-code degeneracy between the block-level scaling factor and the FP4 code, this embodiment fixes the tensor-level scaling factor. And add scaling regular expression. (Formula 19) in, This indicates the scaling regular expression. Indicates the first step in the optimization process Block-level scaling factor for each block, This represents the block-level scaling factor (i.e., the value obtained during initialization in step 200) of the E4M3 block. ), This represents taking the average value over all blocks. The total loss can then be further written as... (Formula 20) in, This represents the weighting coefficient of the scaling regularization term. After optimization, [the following will be applied]. This serves as the block-level scaling factor. In this approach, instead of fixing the scaling factor first and then adjusting the FP4 code in isolation, or using a high-precision scaling factor that is difficult to deploy, this embodiment coordinates the block-level scaling factor with the selection of FP4 codes for each weight under the E4M3 deployment constraints, thereby mitigating the coupling error between the high-amplitude code and the scaling factor.

[0078] Example 3: Implementation of Weighted Asymmetric Candidate Radius This embodiment, based on Embodiment 1, further provides a candidate budget adaptive allocation method for NVFP4 non-uniform amplitude grids, as a preferred implementation of step 300. Because... In the low-to-medium amplitude region, the adjacent code intervals are small, while in the high-amplitude region, the adjacent code intervals are large. If the same symmetric radius is applied to all weights, too many candidate budgets may be allocated in the low-amplitude dense region, while insufficient to cover the wide code intervals in the high-amplitude sparse region. Therefore, in this embodiment, the asymmetric search radius is determined weight by weight based on the grid interval where the weight is located, the offset direction of the normalized amplitude relative to the initial code, and the proportion of high-amplitude weights within the block.

[0079] Specifically, first, identify the high-amplitude sparse region indicator. in, Indicates the first An indicator value for whether a weight is located in a high-amplitude sparse region. This indicates an indicator function that takes the value 1 if the condition within the parentheses is true and 0 otherwise. This indicates the index corresponding to the point where the interval between adjacent codes in the amplitude grid begins to increase. It should be noted that... It is not derived from experience, but determined by the mesh structure—in one exemplary implementation, based on the mesh shown in Equation 1, it can be... The above regions are designated as the starting area for high-amplitude sparse regions, but this setting can be adjusted based on the specific spacing structure of the target NVFP4 compatible mesh. Next, the offset direction indicator is determined. in, Indicates normalized amplitude relative to the initial FP4 code Is it biased towards higher indexes? This indicates whether there is a bias towards lower indexes. Specifically, when... When less than the preset small threshold This is to avoid frequent switching of candidate expansion directions due to minor perturbations. Again, statistical analysis... Within each block, satisfy weighting percentage When it exceeds the preset percentage threshold This indicates that high-amplitude weights are relatively concentrated within the block, and the basic radius of the block can be reduced accordingly. Take as Otherwise, take the base radius. This allows for increased search budgets only for blocks with high outlier weights. Based on this, the asymmetric radii towards higher indexes and lower indexes can be expressed as... in, and These represent the search radii in the direction of higher index and lower index, respectively. This indicates the radius increment of the high-amplitude region. Indicates the asymmetric expansion increment. Indicates the upper limit of the radius (used to control video memory). This indicates taking the minimum value. Therefore, the first... The candidate index set with weights can be represented as in, This represents the candidate index set obtained after using asymmetric radii, and then the local candidate code set is obtained according to Formula 5. The candidate code set is still only composed of It consists of valid codes, so it does not violate NVFP4 format constraints.

[0080] For example, take exemplarily , , , , When a certain weight is normalized, the magnitude... At that time, its initial code index (Corresponding code 3), because Therefore ,and Therefore ,but , The candidate index was truncated to The candidate set is Compared to fixed Only at that time In this embodiment, under the same triggering conditions, a high-amplitude code 6 is additionally incorporated, allowing the weight to span a wide interval from 4 to 6 under the drive of reconstruction loss. For example, when... hour, (Corresponding code 0.5) , ,but , The candidate set is In areas with dense low-amplitude values, no additional candidate budget is consumed. The above example is uniquely determined by the grid structure and decision rules, and can be reproduced without relying on any model experiments, which can be used to directly verify the effect of the structural technique of "redistributing candidate budgets from dense low-amplitude areas to sparse high-amplitude areas".

[0081] It should be pointed out that, , , and The value can be determined based on the NVFP4 mesh structure, intra-block statistics, and logits memory budget, and should not be considered a unique fixed value. When the average number of candidates exceeds the preset budget, the asymmetric expansion increment can be reduced first, followed by the radius increment of the high-amplitude region, or the fixed radius candidate construction method in Embodiment 1 can be reverted to. Furthermore, the fixed radius construction in step 300 of Embodiment 1 and the adaptive radius construction in this embodiment essentially share the technical feature of "constructing a local legal candidate set around the initial FP4 code," and this embodiment is only one preferred implementation of it.

[0082] Example 4: Model-level Quantization Implementation Method Performed in Layer Sequence This embodiment, based on any of the above embodiments, performs layer-by-layer quantization on the Transformer model. Specifically, calibration activations are first collected for the current layer, the full-precision output is calculated using the full-precision weights of the current layer, and NVFP4 quantization optimization is performed on the layer according to the methods in Embodiments 1 to 3. After the current layer is optimized, it is hardened and replaced with a valid NVFP4 simulation layer or a packed NVFP4 layer. Subsequently, calibration activations for the next layer are collected based on the replaced model, and the same quantization process is performed on the next layer. For example, the quantization order can be attention projection, attention output projection, and the up projection, gated projection, and down projection of the feedforward network. The significance of this sequential quantization method is that the calibration activations received by subsequent layers already include the errors after quantization of the upstream layers, so the reconstruction target of subsequent layers is closer to the actual input distribution when the model is deployed. It should be noted that this method does not guarantee the elimination of all inter-layer error accumulation, but it can make the layer-by-layer optimization process closer to the final deployment state, thereby helping to reduce the distribution bias caused by independent layer-by-layer quantization.

[0083] Example 5: Optional Four / Six Block Mode Selection Implementation This embodiment, based on any of the above embodiments, further sets a block-level mode selection for each block, allowing the block to discretely select between six-range scaling and four-range scaling, in order to improve the representation error of the high-amplitude range of NVFP4. Specifically, for the first... Let there be a block, and let the maximum magnitude statistic of this block after normalization by the tensor scaling factor be... Two candidate block-level scaling factors can then be constructed. in, This indicates the block-level scaling factor determined by 6 as the range corresponding to the high-end amplitude code. This represents the block-level scaling factor determined by a range corresponding to the high-end amplitude code, where 4 is the value. Indicates the first The normalized maximum amplitude statistics for each block are calculated. Block-level logits are set for two modes, and differentiable selection between the two modes is made through Gumbel-Softmax or pass-through Gumbel-Softmax. During the hardening stage, the mode corresponding to the maximum value of the block-level logits is taken, and the selected block-level scaling factor is projected to an E4M3 representable value. This implementation can share statistical information with the intra-block high amplitude statistics in Embodiment 3, so that the candidate budget allocation and the block-level scaling mode selection can complement each other. For blocks with concentrated high amplitude weights, the four / six-block mode selection helps to improve the representation error of the high amplitude range while maintaining valid NVFP4 hardening. It should be noted that this implementation is an optional enhancement method and is not required for every embodiment.

[0084] Example 6: System, Device, and Storage Medium Implementation This embodiment also provides a Gumbel-Softmax adaptive quantization system for the NVFP4 format. Referring to Figure 2, the system may include an acquisition module, a determination module, a construction module, a selection module, an optimization module, a hardening module, and an output module. The system comprises the following modules: an acquisition module for acquiring the weight tensor to be quantized and the calibration activation; a determination module for determining the tensor-level scaling factor, the E4M3 block-level scaling factor, and the initial FP4 code based on the NVFP4 format; a construction module for constructing a local candidate code set around the index of the initial FP4 code, ensuring that all candidate codes in the local candidate code set belong to the FP4 E2M1 magnitude grid; a selection module for setting trainable logits for the local candidate code set and obtaining soft-quantized codes through Gumbel-Softmax or pass-through Gumbel-Softmax; an optimization module for optimizing the logits based on the reconstruction loss between the calibration activation, soft-quantized weights, and full-precision weight output, and optionally jointly optimizing the block-level scaling factor; a hardening module for hardening the optimized logits into hard FP4 codes and determining or projecting the block-level scaling factor to an E4M3 representable value; and an output module for outputting a valid NVFP4 weight representation including the weight symbol, hard FP4 code, hard block-level scaling factor, and tensor-level scaling factor.

[0085] The modules in the above system embodiments can be implemented by software, hardware, or a combination of both. The module division described above is merely an illustrative functional division; other division methods may be used in actual implementation. Each module can be integrated into a processing unit or exist separately. Furthermore, this application may also provide a quantization device, which includes a processor and a memory. The memory stores a computer program, and when the processor executes the computer program, it implements the steps in any of the above method embodiments. This application may also provide a computer-readable storage medium storing a computer program thereon, which, when executed by a processor, implements the steps in any of the above method embodiments. The above system, device, and storage medium embodiments are based on the same inventive concept as the method embodiments. Related technical features, technical effects, and optional implementation methods can be referred to each other, and will not be repeated here.

[0086] Abnormal situations and legality handling To improve the stability of the engineering implementation, several abnormal situations can be handled in the above embodiments. For all-zero blocks or near-all-zero blocks, adaptive selection can be skipped and the weights within the block can be assigned to zero codes to avoid division by zero or instability in the normalized magnitude calculation. For boundary indexes... or The weights, after truncation and deduplication of the candidate index set, may only have one candidate code remaining. In this case, the weight can degenerate into a deterministic assignment, without needing to set logits or participate in Gumbel-Softmax sampling, thus avoiding the logical idling of randomly selecting a single candidate. For weights whose normalized magnitude falls within the neighborhood of the midpoint between two codes, symmetrical expansion can be used to avoid frequent jittering in the candidate expansion direction. For weights with less than a complete block at the end, candidate construction, hardening, and packing can be performed only on the actual weights. For the implementation of jointly optimizing the block-level scaling factor, if the hardened block-level scaling factor fails the E4M3 validity check, it can revert to the initialized E4M3 block-level scaling factor. The hard FP4 code selection is then re-executed based on the trained logits. It should be noted that the above processing does not change the main line of this application: "local legal candidate code set - differentiable selection - reconstruction loss optimization - legal NVFP4 hardening".

[0087] Optional implementation details and legality checks In one implementation, the projection process of the E4M3 block-level scaling factor includes numerical range limiting and rounding. Specifically, the block-level scaling factor to be projected can first be limited to a finite range of positive numbers that the target E4M3 format can represent; for non-zero scaling factors less than the minimum positive representable value, they can be set to the minimum positive representable value or zeroed out; for scaling factors greater than the maximum representable value, they can be saturated to the maximum representable value; for non-numerical or infinite values, they can fall back to the E4M3 block-level scaling factor initialized in step 200. The rounding method of the projection can be nearest-neighbor rounding or rounding to even numbers. The above processing is used to ensure that the forward calculation, hardened output, and final packaged result all satisfy the E4M3 representable value constraints.

[0088] In one implementation, the weight tensor is expanded according to the contiguous memory order used for matrix multiplication calculations or weight storage in the inference engine, and divided into multiple micro-blocks according to a preset block size; the preset block size can be 16, or it can be determined according to the packing rules of the target NVFP4 inference engine. For weights with less than a complete block at the end, only the block-level scaling factor is calculated for the actual weight elements, a candidate code set is constructed, and hardening is performed; during packing, the weights can be padded according to the tail block rules of the target inference engine, or the number of valid elements in the tail block can be recorded.

[0089] In one implementation, the trainable logits set for the local candidate code set can be initialized with the same value; alternatively, relatively large initial values ​​can be set for the candidates corresponding to the initial FP4 code, and relatively small initial values ​​can be set for other candidates, so that the initial optimization is close to the NVFP4 initial representation obtained by the most recent rounding. Temperature coefficient The initial temperature can be gradually decreased to the final temperature using cosine decay, exponential decay, or piecewise decay methods. The optimization process can stop when a preset number of optimization steps is reached, the change in reconstruction loss is less than a preset threshold, or the output error of the verification layer no longer decreases.

[0090] In one implementation, the random number used when sampling Gumbel noise... Take from open interval If the random number is generated from a closed interval If the random source is generated, it can be truncated to... ,in The minimum positive number is preset to avoid or This causes numerical anomalies.

[0091] In one implementation, calibration activation is obtained by forward propagation of representative calibration samples through the model. These calibration samples can be text samples, image-text samples, or input samples from the deployment scenario. For cases where quantization is performed layer-by-layer, the calibration activation of the current layer can come from the output of an upstream layer that has been replaced with an NVFP4 simulation layer or a packaged NVFP4 layer, so that the optimized input of subsequent layers can reflect the upstream quantized error.

[0092] In one implementation, a validity check is performed after the hardened output, the validity check including: checking whether all hardened FP4 codes belong to the FP4 E2M1 amplitude grid. The process involves checking whether the hard block-level scaling factor is an E4M3 representable value, whether the tail block meets the packing rules, whether there are non-numerical or infinite values ​​in the quantization weights, and whether the output weight symbols, hard FP4 codes, hard block-level scaling factors, and tensor-level scaling factors can be parsed according to the target NVFP4 representation. For blocks that fail the validity check, the process can fall back to the initialized E4M3 block-level scaling factor and re-execute the hard FP4 code selection based on the optimized logits.

[0093] Technical effects of the above embodiments The above embodiments have the following technical effects: First, by constructing multiple local legal candidate codes around the initial FP4 code in the non-uniform amplitude grid of FP4 E2M1, the traditional nearest rounding or binary adjacent rounding is extended to a differentiable selection of local multiple candidate codes. This helps to expand the optimizable space of FP4 codes without introducing a custom codebook and obtain cross-code search capability in high amplitude sparse regions. Second, by optimizing logits with the reconstruction loss of layer output or block output under calibration activation, the FP4 code selection for each weight is no longer determined solely by the pointwise weight error, but is closer to the reconstruction target of the model inference output, which helps to reduce the reconstruction error of layer / block output. Third, through pass-through Gumbel-Softmax, temperature annealing, entropy regularization, and hardening legality checks, the difference between the soft training stage and the hard deployment stage is reduced. Furthermore, through the E4M3 pass-through estimator and scaling regularization term, the block-level scaling factor and weightwise FP4 code selection can be synergistically adapted while maintaining the E4M3 deployment constraints, alleviating the coupling error caused by the separate optimization of the two. Furthermore, by redistributing the candidate budget from low-amplitude dense regions to high-amplitude sparse regions through weighted asymmetric candidate radii, the effectiveness of local multi-candidate search can be improved under limited GPU memory and computational budget. Finally, through hardened validity checks and NVFP4 representation output, a valid NVFP4 weight representation including weight symbols, hard FP4 codes, E4M3 hard block-level scaling factors, and tensor-level scaling factors can be formed. This helps reduce weight storage, GPU memory usage, and memory access bandwidth requirements in large model inference, and improves deployment compatibility with inference engines or simulation layers that support NVFP4.

[0094] It should be noted that in this patent application, relational terms such as "first" and "second" are used only to distinguish one entity or operation from another, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Furthermore, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitations, an element defined by the phrase "comprising one" does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes said element. In this patent application, if it refers to performing an action according to an element, it means performing the action at least according to that element, including two cases: performing the action only according to that element, and performing the action according to that element and other elements. Expressions such as "multiple," "repeatedly," and "various" include two, two times, two kinds, and more than two, more than two times, and more than two kinds.

[0095] All documents mentioned in this application are considered to be incorporated in their entirety into the disclosure of this application so that they can serve as a basis for modifications if necessary. Furthermore, it should be understood that after reading the foregoing disclosure of this application, those skilled in the art can make various alterations or modifications to this application, and these equivalent forms also fall within the scope of protection claimed in this application.

Claims

1. A Gumbel-Softmax adaptive quantization method for NVFP4 format, characterized in that, include: S1. Obtain the weight tensor to be quantized and calibrate the activation; S2. According to the NVFP4 format, the tensor-level scaling factor and the E4M3 block-level scaling factor are determined by the weight tensor, and the normalized magnitude of each weight is taken as the nearest code in the FP4 E2M1 magnitude grid as the initial FP4 code of that weight; the normalized magnitude is obtained by dividing the magnitude of the weight by the product of the tensor-level scaling factor and the E4M3 block-level scaling factor; S3. Using the index of the initial FP4 code obtained in step S2 in the FP4 E2M1 amplitude grid as the center, construct a local candidate code set for the weight, wherein each candidate code in the local candidate code set belongs to the FP4 E2M1 amplitude grid. S4. Set trainable logits for the local candidate code set obtained in step S3, and perform differentiable selection on the local candidate code set through Gumbel-Softmax or pass-through Gumbel-Softmax to obtain the soft quantized code with the weight. S5. Soft quantization weights are constructed using the soft quantization code obtained in step S4, the tensor-level scaling factor obtained in step S2, and the E4M3 block-level scaling factor. Soft quantization output is calculated based on the calibration activation and the soft quantization weights. Full-precision output is calculated based on the calibration activation and the weight tensor. The logits in step S4 are optimized using the reconstruction loss between the soft quantization output and the full-precision output. S6. For the logits optimized in step S5, take the candidate code corresponding to the maximum value as the hard FP4 code of the weight, and determine or project the E4M3 block-level scaling factor as an E4M3 representable value to obtain the hard block-level scaling factor. S7. Determine the quantization weight value corresponding to the weight using the sign of the weight, the tensor scaling factor, the hard block scaling factor obtained in step S6, and the hard FP4 code, and output a valid NVFP4 weight representation including the sign of the weight, the hard FP4 code, the hard block scaling factor, and the tensor scaling factor.

2. The method according to claim 1, characterized in that, The FP4 E2M1 amplitude grid is Its code index is 0 to 7, and it is combined with the weight sign bit to form a signed FP4 value; The initial FP4 code mentioned in step S2 is determined as follows: the normalized magnitude of the weight is calculated. , and according to Determine the index of the initial FP4 code in the FP4 E2M1 amplitude grid. ; The quantization weight value corresponding to the legal NVFP4 weight representation in step S7 satisfies: In the formula, For the first Each weight, Its amplitude, The tensor scaling factor is... This is the E4M3 block-level scaling factor of the block to which the weight belongs. For normalized amplitude, The index in the FP4 E2M1 amplitude grid is The code, This indicates that the index that minimizes the absolute difference is returned. , This is the index of the initial FP4 code. This is the quantized weight value corresponding to this weight. The sign of this weight. This is a block-level scaling factor. It is a hard FP4 code and is taken from .

3. The method according to claim 1, characterized in that, The local candidate code set mentioned in step S3 is constructed in the following manner: In the formula, The initial FP4 code for this weight is the index in the FP4 E2M1 amplitude grid. The integer offset relative to that index. For the local search radius, This means truncating the input to a closed interval between 0 and 7. This is the set of candidate indices after deduplication. The index in the FP4 E2M1 amplitude grid is The code, The local candidate code set; Wherein, the local search radius When the value is 1, the local candidate code set contains at most 3 candidate codes. When the value is 2, the local candidate code set contains a maximum of 5 candidate codes.

4. The method according to claim 1, characterized in that, The soft quantization code mentioned in step S4 is obtained in the following way: Gumbel noise is sampled for each candidate code in the local candidate code set. ,in obey Uniform distribution on the surface; calculate the selection probability of each candidate key. The soft quantization code is obtained by the following formula: In the formula, For the first The weight of the first Gumbel noise for candidate codes, For the corresponding uniformly distributed random numbers, For the first The logits corresponding to each candidate key. For temperature coefficient, It is a normalized exponential function, and its normalization is performed on candidate keys with the same weight. For the first The probability of choosing a candidate key. For the first One candidate code, This refers to the soft quantization code; And the temperature coefficient During the optimization process, the value decays from the initial value to the final value.

5. The method according to claim 1, characterized in that, After determining the E4M3 block-level scaling factor in step S2, learnable parameters are set for each block. Let the floating-point block scaling factor be 1. ; In step S5, the E4M3 pass-through estimator is used to determine the actual block-level scaling factor used in the forward computation: Therefore, in step S5, the tensor scaling factor is fixed, and the logits and the learnable parameters are jointly optimized. Furthermore, the loss used for optimization in step S5 is based on the reconstruction loss and also includes a scaling regularization term that constrains the E4M3 block-level scaling factor to not deviate from the E4M3 block-level scaling factor determined in step S2. ; In step S6, As the hard block-level scaling factor; In the formula, For the first Learnable parameters for each block, For the reason The resulting floating-point block-level scaling factor, It is a natural exponential function. This indicates that the input is projected as an E4M3 representable value. This indicates the gradient whose value within the brackets is truncated during backpropagation. This refers to the block-level scaling factor actually used in the forward computation. For the first Block-level scaling factor for each block, The E4M3 block-level scaling factor (i.e., determined in step S2) for this block. ), This is to take the average value over all blocks.

6. The method according to claim 2, characterized in that, The search radius of the local candidate code set mentioned in step S3 is adaptively determined as an asymmetric radius according to the weights, including: Identify high-amplitude sparse regions ,in This is the index of the point where the adjacent code spacing begins to increase in the FP4 E2M1 amplitude grid. This is an indicator function; it takes the value 1 if the condition within the parentheses is true, and 0 otherwise. Determine the offset direction indication and , and when season ,in This is the normalized magnitude of the weight. This is the initial FP4 code for this weight. The preset threshold; The block to which this weight belongs satisfies the following conditions. weighting percentage ,when Exceeding the preset percentage threshold The base radius of the block will be determined at this time. Take as Otherwise, take as , Base radius; The asymmetric radii in the high-index direction and the low-index direction are synthesized according to the following formula: and with Construct a candidate index set for this weight, thereby redistributing the number of candidate codes from the low-amplitude dense region to the high-amplitude sparse region; In the formula, For the radius increment of the high amplitude region, For asymmetric expansion increment, The upper limit of the radius, To obtain the minimum value, and These are the search radii for the direction towards higher indexes and the direction towards lower indexes, respectively.

7. The method according to claim 1, characterized in that, In step S5, the loss used to optimize the logits is based on the reconstruction loss plus a regularization term, which includes at least one of an entropy regularization term and a moving regularization term; wherein: Entropy regularization term This is used to make the selection probability distribution of the differentiable selection approach the one-thermal distribution; Move regular expression This is used to limit the offset of the final FP4 code relative to the initial FP4 code; In the formula, For the first The probability distribution of the selection of each candidate key with each weight. Its entropy, For the first The probability of choosing a candidate key. For the first One candidate code, Its index in the FP4 E2M1 amplitude grid, This is the index of the initial FP4 code. This is to take the average of all weights.

8. The method according to claim 1, characterized in that, The method is executed layer by layer on the Transformer model: after performing steps S1 to S7 on the current layer to obtain the legal NVFP4 weight representation of the layer, the current layer is replaced with the legal NVFP4 weight representation, and then the calibration activation of the next layer is collected and steps S1 to S7 are performed on the next layer, so that the calibration activation of the subsequent layers includes the quantized error of each upstream layer.

9. A Gumbel-Softmax adaptive quantization system for NVFP4 format, characterized in that, include: The acquisition module is used to acquire the weight tensor to be quantized and the calibration activation. The determination module is used to determine the tensor-level scaling factor and the E4M3 block-level scaling factor from the weight tensor according to the NVFP4 format, and to take the nearest code in the FP4 E2M1 amplitude grid of each weight as the initial FP4 code of that weight. The normalized amplitude is obtained by dividing the amplitude of the weight by the product of the tensor-level scaling factor and the E4M3 block-level scaling factor. A construction module is used to construct a local candidate code set for the weight, centered on the index of the initial FP4 code in the FP4 E2M1 amplitude grid, wherein each candidate code in the local candidate code set belongs to the FP4 E2M1 amplitude grid. The selection module is used to set trainable logits for the local candidate code set and perform differentiable selection on the local candidate code set through Gumbel-Softmax or pass-through Gumbel-Softmax to obtain the soft quantized code with the weight. An optimization module is used to construct soft quantization weights from the soft quantization code, the tensor-level scaling factor, and the E4M3 block-level scaling factor; calculate a soft quantization output based on the calibration activation and the soft quantization weights; calculate a full-precision output based on the calibration activation and the weight tensor; and optimize the logits using the reconstruction loss between the soft quantization output and the full-precision output. The hardening module is used to take the candidate code corresponding to the maximum value of the optimized logits as the hard FP4 code of the weight, and to determine or project the E4M3 block-level scaling factor as the hard block-level scaling factor that can be represented by E4M3. The output module is used to determine the quantized weight value corresponding to the weight by the sign of the weight, the tensor scaling factor, the hard block scaling factor and the hard FP4 code, and output a valid NVFP4 weight representation including the sign of the weight, the hard FP4 code, the hard block scaling factor and the tensor scaling factor.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by a processor, it implements the method of any one of claims 1 to 8.

Images