Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix Device Mismatch #12

Merged
merged 1 commit into from
Apr 16, 2024
Merged

Fix Device Mismatch #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions src/sparsetensors/quantization/lifecycle/forward.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ def fake_quantize(
zero_point: torch.Tensor,
args: QuantizationArgs,
) -> torch.Tensor:
max_q = torch.tensor(2**args.num_bits - 1)
max_q = torch.tensor(2**args.num_bits - 1, device=x.device)
Q = torch.zeros_like(x)
Q = quantize(x, scale, zero_point, max_q)
return dequantize(Q, scale, zero_point)
Expand Down Expand Up @@ -112,6 +112,7 @@ def _maybe_calibrate_or_quantize(
}:
return value

device = next(module.parameters()).device
scale = getattr(module, f"{base_name}_scale")
# zero_point = getattr(module, f"{base_name}_zero_point").data
zero_point = getattr(module, f"{base_name}_zero_point")
Expand All @@ -122,7 +123,7 @@ def _maybe_calibrate_or_quantize(
updated_scale, updated_zero_point = observer(value)

# update scale and zero point
scale.data = updated_scale
zero_point.data = updated_zero_point
scale.data = updated_scale.to(device)
zero_point.data = updated_zero_point.to(device)

return fake_quantize(value, scale, zero_point, args)
8 changes: 6 additions & 2 deletions src/sparsetensors/quantization/lifecycle/initialize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,15 @@ def initialize_module_for_quantization(module: Module, scheme: QuantizationSchem
def _initialize_scale_zero_point_observer(
module: Module, base_name: str, quantization_args: QuantizationArgs
):
device = next(module.parameters()).device

# initializes empty scale and zero point parameters for the module
init_scale = Parameter(torch.empty(0), requires_grad=False)
init_scale = Parameter(torch.empty(0, device=device), requires_grad=False)
module.register_parameter(f"{base_name}_scale", init_scale)

init_zero_point = Parameter(torch.empty(0, dtype=int), requires_grad=False)
init_zero_point = Parameter(
torch.empty(0, device=device, dtype=int), requires_grad=False
)
module.register_parameter(f"{base_name}_zero_point", init_zero_point)

# initialize observer module and attach as submodule
Expand Down
Loading