With Shared Microexponents,
A Little Shifting Goes a Long Way

To evaluate the numerical fidelity of BDR configurations, we adopt two approaches. The first is a statistical analysis to investigate the sensitivity of BDR to each design variable and how they contribute to the overall numerical robustness under different data distributions. This allows us to conduct an extensive search over the design space without running costly experiments for every configuration.

In addition to the statistical analysis, we conduct end-to-end training of various scale language models using different BDR configurations. We find a strong Pearson correlation [50, Ch. 5.1] between the results of our statistical analysis and the language model loss achieved in our end-to-end training runs in the narrow bit-width regime. This suggests that our statistical results are a reliable predictor of the relative numerical fidelity performance of different configurations in practice at single-digit bit-widths.

In our statistical analysis, we use Quantization Signal to Noise Ratio (QSNR) as a measure of numerical fidelity under different quantization schemes. QSNR is defined as the ratio of the power of the non-quantized signal (i.e., the original vector $\mathbf{X}=[x_{1},x_{2},\cdots,x_{k}]\in\mathbb{R}^{k}$) to the power of the quantization noise expressed in decibels. It is calculated as:

where $\lVert\mathbf{X}\rVert:=(\sum_{i=1}^{k}{x_{i}^{2}})^{1/2}$, and $Q(\cdot)$ is the quantization function which maps $k$-dimensional arrays with FP32 precision to the target quantization format. A higher QSNR indicates that the quantized vector better preserves the direction and magnitude of the original non-quantized vector in the space $\mathbb{R}^{k}$. The required QSNR level for a successful run varies depending on the task complexity, dataset, and model topology. We explore the trade-off between hardware efficiency and QSNR of different design points in Section IV-C.

In this section, we outline a hardware architecture that supports operations on BDR formats and describes our methodology for estimating the costs of these operations in circuit area, power consumption, and memory.

In Figure 6, when the pipeline is configured with $k_{1}=k_{2}=1$, it implements a standard scalar floating-point dot product. In this setting, the internal accumulation across the native reduction dimension is performed by normalizing the elements of a Hadamard product to the largest element and reducing in fixed-point. This is an optimistic approximation for scalar floating-point area overhead given that the data is serially accumulating in a floating-point type $f-2$ mantissa bits instead of building a full floating-point adder tree.

The BDR format is designed to be compatible with existing standard floating-point dot product pipelines in order to enable sub-circuit sharing and lower marginal costs. However, the cost of integrating BDR into a given design will depend on various factors, such as the underlying architecture, the set of data formats, and the relative throughputs of the existing design. To make the results of our synthesis and area estimation methodology more comparable, we normalize the standard-cell area of each BDR configuration to a baseline configurable FP8 dot product unit supporting both E4M3 and E5M2. This allows us to directly compare the relative costs of different BDR configurations and understand the trade-offs between numerical fidelity and hardware cost.

Due to software license usage restrictions, building hundreds of configurations and finding the optimal pipelining for each of them was not feasible. Our research on optimal register placements in smaller sweeps, however, shows that registers typically only accounted for about $10\%$ of the total area. Therefore, a design requiring half as many pipeline stages would only save around $5\%$ of the core area. To isolate the core area for a given configuration from differences in optimal pipeline stage placement or sub-circuit synthesis mapping, we synthesize each configuration with an easily achievable timing constraint of $10ns$ and with only inputs and outputs being registered. This ensures that synthesis implementation selection targets the minimum area in all designs. We use Synopsys Design Compiler and a leading-edge process node.

Figure 7 illustrates the trade-off between numerical fidelity and the hardware cost for different data format configurations. The x-axis indicates the implementation cost of each configuration as a function of the dot product area and memory efficiency. The reported QSNR (in decibels) on the y-axis is the averaged measured QSNR from over $10K$ independent vectors drawn from a normal Gaussian distribution with a variable variance (higher QSNR is better). We chose this distribution to cover a range of variances observed in gradient, error, weight, and activation tensors in a typical training cycle.

In our efficiency analysis, we normalize the dot product area to that of a baseline FP8 (supporting both E4M3 and E5M2) implementation and calculate the memory efficiency as the inverse of the ratio of the number of elements that can be packed into a $64B$ memory interface. Since the goal is to select a congruent data format for both training and inference, we give equal weight to the dot product area and memory efficiency as they are both important from a performance and power perspective for the two scenarios.

A detailed proof of Theorem 1 can be found in Section IX. As also observed empirically in Figure 7, the QSNR has a linear relation with the number of mantissa bits $m$ and a logarithmic relation with the block granularity $k_{1}$ and $k_{2}$. Note that Eq. (4) provides a parameterized lower-bound for the BDR family of data formats. This in turn provides a guarantee of the worst-case numerical fidelity of BDR formats in the face of skewed distributions with correlated noise.

The specific parameters of these formats (see Table II) are chosen at the “knee” (shown as hashed points on the Pareto frontier in Figure 7) of a particular parameter’s sweep by trading off the increase in QSNR with the increase in the normalized area-memory efficiency product. For instance, across all the selected formats, increasing the value of $d_{2}$ from $1$-bit to $2$-bit increases the QSNR by only $0.5dB$, but with a $30-50\%$ increase in the normalized cost. Similarly, with a $d_{2}$ value of $1$, reducing $k_{2}$ from $8$ to $2$ increases the QSNR by approximately $2dB$, with only a marginal $3\%$ increase in the normalized cost. Whereas, further reducing the value of $k_{2}$ from $2$ to $1$ increases the QSNR only by $0.7dB$, but with a significant $30-40\%$ increase in the normalized cost. We pragmatically select a first-level block granularity of $16$ to avoid forcing a large reduction dimension on hardware architectures. In general, we also observe that the trade-off between the increase in the QSNR versus the increase in the normalized area-memory efficiency product reduces as we increase the number of bits per element.

Even though MX9 is not necessarily on the Pareto frontier in Figure 7, it is still a good choice for the basic data format for further evaluation as it allows all the selected formats to have the same settings (see Table II) except for their mantissa bit-widths, which allows for maximum hardware reuse. Moreover, the difference in QSNR ($\approx 0.8dB$) between the selected MX9 format and the configuration on the Pareto frontier is insignificant for the efficacy of the end-to-end training and inference workloads due to its already high QSNR value. Note that MX9 has approximately $3.6dB$ higher QSNR compared to that of MSFP16 [24]. Empirical measurements of proprietary generative inference models (results not reported) show a significant accuracy degradation in MSFP16 as a drop-in replacement option that is not occurring in MX9.

MX9 has a hardware efficiency close to that of FP8 with significantly higher numerical fidelity. For instance, in the case of a Gaussian distribution with variable variance, the QSNR of MX9 is about $16dB$ higher than FP8 (E4M3). A $16dB$ higher fidelity is roughly equivalent to having $2$ more mantissa bits in the scalar floating-point format. For the same distribution, MX6’s QSNR lies between the two FP8 variants E4M3 and E5M2 while providing an approximately $2\times$ advantage on the hardware cost as measured by the normalized area memory efficiency product. We qualitatively observed a similar trend under different data distributions as well.

To evaluate the quality of MX data formats, we created a comprehensive benchmark suite featuring over $20$ representative open-source and proprietary AI models. In the following, we provide an overview of the models in our benchmark suite.

• •

  Neural Machine Translation. We evaluate two model topologies in this category: LSTM-based and transformer-based. We use the MLPperf GNMT-V2 model [55] as the LSTM benchmark. The transformer models are based on the work of [52]. These models are trained using the WMT17-EN-DE data [19] and evaluated on the WMT14-EN-DE [20].

• •

  Text Encoding and Language Understanding. For this application, we evaluate the popular BERT [28] models: BERT-base and BERT-large. The models are trained on the Wikipedia data and evaluated on SQuAD-v1.1 [45].

• •

  Generative Language Modeling. Auto-regressive generative learning is probably the most significant and fastest-growing application in the NLP domain. In our benchmark suite, we evaluate MX performance on a variety of dense Generative Pre-Trained Transformers (GPT) models of size $6M$ up to $175B$. All models are trained with a sequence length of $1024$ to efficiency with the number of tokens processed estimated by GPT scaling power laws [35]. In addition to dense GPT models, we also evaluate an example Mixture-of-Experts (MoE) generative model from [44]. The MoE model in our study has $1.9B$ parameters trained for $40B$ tokens.

• •

  Image Classification. We evaluate three main families of models for image classification: ResNet [32], MobileNet-v2 [46], and DeiT [51]. These benchmarks are trained and validated with ImageNet ILSVRC12 data [25].

• •

  Image Generation. We use the DDPM model [39] as a representative model in this category. We evaluate two variations of the DDPM model in our study. Both models are trained with $4000$ diffusion steps using ImageNet-64.

Tables III to V summarize the MX inferencing results in both direct casting and quantization-aware fine-tuning settings across a variety of model topologies, scales, and modalities. As shown, a direct cast to MX9 can be mostly used as a drop-in replacement to high-precision formats.

Direct casting to the more efficient MX6 format sometimes causes model accuracy to degrade. In these cases, the loss in accuracy can be recovered through a modest amount of quantization-aware fine-tuning. We report both the direct cast and fine-tuning results for MX6 in Tables III and V. We generally saw the best quantization-aware fine-tuning results when we reset the optimizer, adjusted the initial learning rate, and eliminated rate decay, dropout, and momentum. In all our fine-tuning experiments, we use FP32 for the backward pass. The amount of fine-tuning necessary varied across benchmarks but was always much shorter than the original training duration.

The MX data format also presents opportunities for optimizing zero-shot and few-shot generative inferencing with no task-specific fine-tuning. See Table IV for quantization results of OpenAI GPT3-175B model with MX formats. Consistent with the results for discriminative inferencing reported in Table III, MX9 provides a drop-in replacement alternative for efficient generative inferencing. Generative inferencing at large-scale allows the use of even narrower data formats such as MX4 and MX6 without the need for much quantization-aware fine-tuning. To the best of our knowledge, this is the first work reporting high-quality generative inferencing with both weights and activations quantized to ultra-narrow bit-width through a direct cast. In prior quantized generative inferencing work, the activation is typically kept at high precision BF16/FP16 and only weights are being quantized [27, 42]. This is because INT8/INT4 quantization of activations usually requires careful offline outlier suppression [56] to achieve reasonable results.

MX9 provides a robust drop-in replacement for high bit-width AI training (i.e., FP32/BF16/FP16) without requiring an outlier detection heuristic or any change in the training recipe, hyper-parameters, or model topology. Table III provides an overview of MX9 training results for a variety of discriminative benchmarks across different model topologies and data modalities. The results show that the models trained with MX9 match FP32 baseline accuracy within the error margin of run-to-run FP32-based training when initialized with different random seeds. The work in [37] reports FP8 results on a subset of these benchmarks.

In addition to benchmarks listed in Table III, we further evaluated the robustness of MX9 training on three proprietary recommendation models. Table VI summarizes the results of this evaluation. Consistent with the results of our open-sourced benchmark suite, MX9 presents a robust training alternative for production-scale recommendation models. In our experience, FP8-based training led to gradient explosion in PR-rec3 models while MX9 had a robust convergence trend.

Internally, we use a tight threshold of $0.02\%$ for production models, which is set based on the training variations from run to run with the FP32 format [26]. Training PR-rec1 with MX9 and MX6 can both meet this requirement but MX4 results in $\sim 0.07\%$ gap with the same number of training iterations and hyper-parameters as the baseline FP32. MX6 and MX4 have roughly $2\times$ and $4\times$ lower circuitry compared to FP8. We leave further evaluation of MX4 and MX6 for training recommendation models to future work. PR-rec2 and PR-rec3 are more challenging benchmarks to quantize. As shown in Table VI, certain layers (e.g., first and last layer) should be kept in high bit-width to meet the $0.02\%$ threshold.

Table VII shows the result of training generative transformers of various sizes with MX9 format. MX9 matches the quality of FP32 without the need for any heuristic-based outlier detection (e.g., transformer engine in [37]) or change in the training recipe (number of training iterations, hyper-parameters, or model topology).

Our assessment of training generative language models with the more efficient MX6 format shows that the models trained with this data format are also able to converge robustly. Figure 9 shows the Language Model (LM) loss as a function of normalized cost per training iteration for MX9 versus MX6. MX6 requires more training iterations compared to the baseline FP32 (and by extension MX9) to converge to the same LM loss. Given the relative throughput of MX6, however, the model can still converge to the same quality as the baseline with an overall lower training cost. Quantifying the exact impact of the MX6 data format on the underlying generative training power laws is out of the scope of this paper.