|
25 | 25 | from __future__ import annotations |
26 | 26 |
|
27 | 27 | import torch |
| 28 | +from compressed_tensors.offload import disable_onloading |
28 | 29 | from compressed_tensors.quantization import QuantizationScheme, QuantizationStrategy |
29 | 30 | from compressed_tensors.utils import match_named_modules |
30 | 31 |
|
31 | 32 | __all__ = [ |
32 | | - "_layer_indivisible", |
33 | 33 | "get_layers_indivisible_by_group_size", |
34 | 34 | "validate_group_size_divisibility", |
35 | 35 | ] |
@@ -77,16 +77,19 @@ def get_layers_indivisible_by_group_size( |
77 | 77 | :return: List of (fqn, columns, group_size) for each layer that would |
78 | 78 | fail at save/forward due to indivisibility. |
79 | 79 | """ |
80 | | - indivisible: list[tuple[str, int, int]] = [] |
81 | | - for name, module in match_named_modules(model, resolved_targets, ignore): |
82 | | - scheme: QuantizationScheme | None = getattr(module, "quantization_scheme", None) |
83 | | - if scheme is None or scheme.weights is None: |
84 | | - continue |
85 | | - result = _layer_indivisible(module, scheme.weights) |
86 | | - if result is not None: |
87 | | - columns, group_size = result |
88 | | - indivisible.append((name, columns, group_size)) |
89 | | - return indivisible |
| 80 | + with disable_onloading(): |
| 81 | + indivisible: list[tuple[str, int, int]] = [] |
| 82 | + for name, module in match_named_modules(model, resolved_targets, ignore): |
| 83 | + scheme: QuantizationScheme | None = getattr( |
| 84 | + module, "quantization_scheme", None |
| 85 | + ) |
| 86 | + if scheme is None or scheme.weights is None: |
| 87 | + continue |
| 88 | + result = _layer_indivisible(module, scheme.weights) |
| 89 | + if result is not None: |
| 90 | + columns, group_size = result |
| 91 | + indivisible.append((name, columns, group_size)) |
| 92 | + return indivisible |
90 | 93 |
|
91 | 94 |
|
92 | 95 | def validate_group_size_divisibility( |
|
0 commit comments