pytorch - 💡(How to fix) Fix `torch.compile` produces different uint8 quantization output for E8M0 bit-manipulation pattern (`view(int32) → bitshift → clamp → uint8`) compared to eager mode [1 participants]

import torch
import torch.nn as nn
import torch.nn.functional as F

class E8M0QuantizationNetwork(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # Feature extraction: Conv + BN + ReLU
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool = nn.AdaptiveAvgPool2d(8)
        self.fc = nn.Linear(64 * 8 * 8, num_classes)

    def fp32_to_e8m0_rceil(self, x):
        """E8M0 quantization via bit manipulation with ceiling rounding."""
        x = x.contiguous()
        # Reinterpret float32 bits as int32
        x_bits = x.view(torch.int32)
        # Extract biased exponent: bits [30:23]
        biased_exp = (x_bits >> 23) & 0xFF
        # Extract mantissa: bits [22:0]
        mantissa = x_bits & 0x7FFFFF
        # Ceiling rounding: round up if any mantissa bit is set
        needs_round_up = (mantissa != 0).to(torch.int32)
        e8m0_biased = biased_exp + needs_round_up
        # Clamp to valid uint8 range
        e8m0_biased = torch.clamp(e8m0_biased, 0, 255)
        return e8m0_biased.to(torch.uint8)

    def forward(self, x):
        # Feature extraction
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = self.pool(x)

        # E8M0 quantization via bit manipulation
        quantized = self.fp32_to_e8m0_rceil(x)

        # Continue with quantized values as float for classification
        x = quantized.to(torch.float32)
        x = x.flatten(1)
        logits = self.fc(x)
        return logits, quantized

# NOTE: CPU, not CUDA
device = "cpu"
torch.manual_seed(42)
model = E8M0QuantizationNetwork().to(device).eval()
x = torch.randn(4, 3, 64, 64, device=device)

# Eager: deterministic
with torch.no_grad():
    ref_logits, ref_quant = model(x)
    ref_logits2, ref_quant2 = model(x)
print(f"Eager deterministic (logits): {(ref_logits - ref_logits2).abs().max().item():.6e}")
print(f"Eager deterministic (quant):  {(ref_quant.int() - ref_quant2.int()).abs().max().item()}")

# Compiled: different result
torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
with torch.no_grad():
    comp_logits, comp_quant = compiled(x)

logit_diff = (ref_logits.float() - comp_logits.float()).abs()
quant_diff = (ref_quant.int() - comp_quant.int()).abs()
print(f"Logits max_diff={logit_diff.max().item():.6e}")
print(f"Quant uint8 mismatches: {(quant_diff > 0).sum().item()} / {ref_quant.numel()}")
print(f"Quant max_diff={quant_diff.max().item()}")

---

Eager deterministic (logits): 0.000000e+00
Eager deterministic (quant):  0
Logits and quantized values differ systematically under torch.compile

---

PyTorch version: 2.12.0.dev20260327+cu126
Python: 3.10.12
OS: Ubuntu 22.04.5 LTS (WSL2)
GPU: NVIDIA GeForce RTX 3060 Laptop GPU
CUDA: 12.6