wehub-resource-sync
88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
from .resources import Dependency, app, get_image
|
|
|
|
img = get_image(dependencies=[Dependency.transformer_engine, Dependency.fouroversix])
|
|
|
|
with img.imports():
|
|
import torch
|
|
from fouroversix import AdaptiveBlockScalingRule, QuantizeBackend, quantize_to_fp4
|
|
from fouroversix.quantize import from_blocked
|
|
|
|
|
|
@app.function(image=img, gpu="B200")
|
|
def create_test_case(
|
|
backend_a: str = "cuda",
|
|
backend_b: str = "transformer_engine",
|
|
scale_rule: str = "mse",
|
|
) -> None:
|
|
M, N = 1024, 1024 # noqa: N806
|
|
|
|
torch.set_printoptions(precision=10)
|
|
|
|
backend_a = QuantizeBackend(backend_a)
|
|
backend_b = QuantizeBackend(backend_b)
|
|
scale_rule = AdaptiveBlockScalingRule(scale_rule)
|
|
|
|
for random_seed in range(10):
|
|
torch.manual_seed(random_seed)
|
|
|
|
x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda")
|
|
out_a = quantize_to_fp4(
|
|
x,
|
|
backend=backend_a,
|
|
scale_rule=scale_rule,
|
|
)
|
|
out_b = quantize_to_fp4(
|
|
x,
|
|
backend=backend_b,
|
|
scale_rule=scale_rule,
|
|
)
|
|
x_sf_a = from_blocked(out_a.scale_factors.bfloat16(), (M, N // 16))
|
|
x_sf_b = from_blocked(out_b.scale_factors.bfloat16(), (M, N // 16))
|
|
|
|
print(f"x absmax: {x.abs().max()}")
|
|
|
|
if not torch.allclose(out_a.amax, out_b.amax):
|
|
print("Backends A and B have different amax values!")
|
|
print(f"{backend_a}: {out_a.amax}")
|
|
print(f"{backend_b}: {out_b.amax}")
|
|
return
|
|
|
|
if not torch.allclose(x_sf_a.bfloat16(), x_sf_b.bfloat16()):
|
|
mismatch_prop = (x_sf_a != x_sf_b).sum() / x_sf_a.numel()
|
|
print(
|
|
"Backends A and B have different scale factors! "
|
|
f"{mismatch_prop:.2%} mismatch",
|
|
)
|
|
|
|
[i, *_], [j, *_] = torch.where(x_sf_a != x_sf_b)
|
|
print(backend_a)
|
|
print("sf", x_sf_a[i, j])
|
|
print("e2m1", out_a.e2m1_values[i, 8 * j : 8 * (j + 1)])
|
|
print(backend_b)
|
|
print("sf", x_sf_b[i, j])
|
|
print("e2m1", out_b.e2m1_values[i, 8 * j : 8 * (j + 1)])
|
|
print("original")
|
|
print("x", x[i, 16 * j : 16 * (j + 1)])
|
|
return
|
|
|
|
if not torch.allclose(out_a.e2m1_values, out_b.e2m1_values):
|
|
mismatch_prop = (
|
|
out_a.e2m1_values != out_b.e2m1_values
|
|
).sum() / out_a.e2m1_values.numel()
|
|
print(
|
|
"Backends A and B have different e2m1 values! "
|
|
f"{mismatch_prop:.2%} mismatch",
|
|
)
|
|
|
|
[i, *_], [j, *_] = torch.where(out_a.e2m1_values != out_b.e2m1_values)
|
|
print(i, j)
|
|
print("normconst", out_a.amax)
|
|
print("sf", x_sf_a[i, j // 8])
|
|
print(backend_a)
|
|
print("e2m1", out_a.e2m1_values[i, 8 * (j // 8) : 8 * (j // 8 + 1)])
|
|
print(backend_b)
|
|
print("e2m1", out_b.e2m1_values[i, 8 * (j // 8) : 8 * (j // 8 + 1)])
|
|
print("original")
|
|
print("x", x[i, 16 * (j // 8) : 16 * (j // 8 + 1)])
|
|
return
|