pytorch - 💡(How to fix) Fix `torch.compile` produces different uint8 quantization values for E8M0 `log2 → ceil → clamp → uint8` pipeline compared to eager mode [1 comments, 1 participants]

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

class E8M0LogQuantizer(nn.Module):
    """E8M0 quantization via log-domain: abs → log2 → ceil → clamp → uint8."""
    def __init__(self, eps=1e-7, bias=127):
        super().__init__()
        self.eps = eps
        self.bias = bias

    def forward(self, x):
        x_abs = torch.abs(x) + self.eps
        log2_val = torch.log2(x_abs)
        ceil_val = torch.ceil(log2_val)
        clamped = torch.clamp(ceil_val, min=-127, max=127)
        encoded = (clamped + self.bias).to(torch.uint8)
        return encoded


class LogDomainCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # Feature extraction
        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.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool = nn.AdaptiveAvgPool2d(4)

        # E8M0 log-domain quantizer
        self.quantizer = E8M0LogQuantizer()

        # Classifier
        self.fc = nn.Linear(128 * 4 * 4, num_classes)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.pool(x)

        # E8M0 log-domain quantization
        encoded = self.quantizer(x)

        # Classification from float features
        cls = x.flatten(1)
        logits = self.fc(cls)
        return logits, encoded


# NOTE: CPU — the log2→ceil pipeline bug reproduces on CPU
device = "cpu"
torch.manual_seed(42)
model = LogDomainCNN().to(device).eval()
# Positive-biased input (abs + offset) to ensure valid log2 range
x = torch.abs(torch.randn(4, 3, 64, 64, device=device)) + 0.1

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

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

logit_diff = (ref_logits.float() - comp_logits.float()).abs()
enc_diff = (ref_enc.int() - comp_enc.int()).abs()
print(f"Logits max_diff: {logit_diff.max().item():.6e}")
print(f"Encoded uint8 mismatches: {(enc_diff > 0).sum().item()} / {ref_enc.numel()}")
print(f"Encoded max_diff: {enc_diff.max().item()}")

---

Eager deterministic (logits): 0.000000e+00
Eager deterministic (encoded): 0
Logits and encoded uint8 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