pytorch - 💡(How to fix) Fix `torch.compile` produces different classification output for model using E8M0 bit-manipulation quantization (`view(int32) → bitshift → clamp → uint8`) [1 comments, 1 participants]

import os
os.environ["TRITON_BACKENDS_IN_TREE"] = "1"

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

class E8M0QuantizationBlock(nn.Module):
    def __init__(self, in_channels=16, out_channels=32):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def e8m0_quantize(self, x):
        x = x.contiguous()
        x_bits = x.view(torch.int32)
        biased_exp = (x_bits >> 23) & 0xFF
        mantissa = x_bits & 0x7FFFFF
        needs_round_up = mantissa != 0
        e8m0_biased = biased_exp + needs_round_up.to(torch.int32)
        e8m0_biased = torch.clamp(e8m0_biased, 0, 255)
        return e8m0_biased.to(torch.uint8)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        quantized = self.e8m0_quantize(x)
        x = quantized.to(torch.float32)
        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)
        return x, quantized


class MultiScaleE8M0Model(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.quant_block1 = E8M0QuantizationBlock(32, 64)
        self.pool1 = nn.MaxPool2d(2)
        self.quant_block2 = E8M0QuantizationBlock(64, 128)
        self.pool2 = nn.MaxPool2d(2)
        self.quant_block3 = E8M0QuantizationBlock(128, 256)
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(256, num_classes)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x, q1 = self.quant_block1(x)
        x = self.pool1(x)
        x, q2 = self.quant_block2(x)
        x = self.pool2(x)
        x, q3 = self.quant_block3(x)
        x = self.global_pool(x)
        x = x.view(x.size(0), -1)
        logits = self.fc(x)
        return logits, (q1, q2, q3)


device = "cuda"
torch.manual_seed(42)
model = MultiScaleE8M0Model().to(device).eval()
x = torch.randn(4, 3, 64, 64, device=device)

# Eager
with torch.no_grad():
    eager_logits, eager_quants = model(x)

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

# Compare
diff = (eager_logits - comp_logits).abs().max().item()
print(f"Logits max diff: {diff}")
print(f"Match: {torch.allclose(eager_logits, comp_logits, atol=1e-5)}")
# Expected: max_diff > 0.1 — significant mismatch

---

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