From 0f805bcb879740b3863c7bfa228fd4f1a1dfc3b8 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Tue, 5 Nov 2024 20:33:04 -0500 Subject: [PATCH 1/5] enable compile on hpu Signed-off-by: yiliu30 --- auto_round/autoround.py | 13 +++--------- auto_round/utils.py | 47 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 49 insertions(+), 11 deletions(-) diff --git a/auto_round/autoround.py b/auto_round/autoround.py index 513bde7c5..129189db9 100644 --- a/auto_round/autoround.py +++ b/auto_round/autoround.py @@ -51,6 +51,7 @@ mv_module_from_gpu, unsupport_meta_device, detect_device_count, clear_memory, get_multimodal_block_names, get_library_version, + compile_func, ) from .low_cpu_mem.utils import get_layers_before_block @@ -386,11 +387,7 @@ def quant_layers(self, layer_names, layer_inputs): self.model = mv_module_from_gpu(self.model, self.low_cpu_mem_usage) clear_memory() - torch_version = get_library_version("torch") - if version.parse(torch_version) >= version.parse("2.5.99"): - quant_layer = torch.compile(self.quant_layer) - else: - quant_layer = self.quant_layer + quant_layer = compile_func(self.quant_layer) for layer_name in layer_names: layer_input = layer_inputs[layer_name] layer_input = to_device(layer_input, self.cache_device) @@ -1110,11 +1107,7 @@ def quant_blocks( elif isinstance(input_others[key], list): for i in range(len(input_others[key])): to_dtype(input_others[key][i], tmp_dtype) - torch_version = get_library_version("torch") - if version.parse(torch_version) >= version.parse("2.5.99"): - quant_block = torch.compile(self.quant_block) - else: - quant_block = self.quant_block + quant_block = compile_func(self.quant_block, device) pbar = tqdm(range(0, len(block_names), nblocks)) for i in pbar: diff --git a/auto_round/utils.py b/auto_round/utils.py index 9b21a15cf..30c4f5b6a 100644 --- a/auto_round/utils.py +++ b/auto_round/utils.py @@ -18,7 +18,7 @@ import sys import subprocess from collections import UserDict - +import re # for cpu usage import cpuinfo import numpy as np @@ -865,3 +865,48 @@ def clear_memory(tensor=None): gc.collect() torch.cuda.empty_cache() + +# Copied from TorchAO +def parse_version(version_string): + # Extract just the X.Y.Z part from the version string + match = re.match(r"(\d+\.\d+\.\d+)", version_string) + if match: + version = match.group(1) + return [int(x) for x in version.split(".")] + else: + raise ValueError(f"Invalid version string format: {version_string}") + + +def compare_versions(v1, v2): + v1_parts = parse_version(v1) + v2_parts = parse_version(v2) + return (v1_parts > v2_parts) - (v1_parts < v2_parts) + + +def torch_version_at_least(min_version): + return compare_versions(torch.__version__, min_version) >= 0 + + +TORCH_VERSION_AT_LEAST_2_6 = torch_version_at_least("2.6.0") +TORCH_VERSION_AT_LEAST_2_5 = torch_version_at_least("2.5.0") +TORCH_VERSION_AT_LEAST_2_4 = torch_version_at_least("2.4.0") + + +def compile_func_on_hpu(func): + if TORCH_VERSION_AT_LEAST_2_4: + return torch.compile(func, backend="hpu_backend") + return func + + +def compile_func_on_cuda_or_cpu(func): + if TORCH_VERSION_AT_LEAST_2_6: + return torch.compile(func) + else: + return func + + +def compile_func(fun, device): + if "hpu" in str(device): + return compile_func_on_hpu(fun, device) + else: + return compile_func_on_cuda_or_cpu(fun, device) From 3d58ec6263f1d9290d59203e56777ab7106e609e Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Tue, 5 Nov 2024 22:38:38 -0500 Subject: [PATCH 2/5] refine clean memory Signed-off-by: yiliu30 --- auto_round/utils.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/auto_round/utils.py b/auto_round/utils.py index 30c4f5b6a..aae8dade8 100644 --- a/auto_round/utils.py +++ b/auto_round/utils.py @@ -751,14 +751,10 @@ def is_autoround_exllamav2_available(): res = False return res - +@lru_cache(None) def is_hpu_supported(): # pragma: no cover try: - import subprocess import habana_frameworks.torch.core as htcore # pylint: disable=E0401 - hqt_version = subprocess.check_output(['pip', 'show', \ - 'habana_quantization_toolkit']).decode().split('\n')[1].split(': ')[1] - assert (hqt_version >= "1.17") except ImportError as e: return False return True @@ -859,13 +855,21 @@ def get_autogptq_packing_qlinear(backend, bits=4, group_size=128, sym=False): return QuantLinear -def clear_memory(tensor=None): +def _clear_memory_for_cpu_and_cuda(tensor=None): if tensor is not None: del tensor gc.collect() torch.cuda.empty_cache() +def clear_memory(tensor=None): + if is_hpu_supported(): + # hpu does not have empty_cache + return + else: + _clear_memory_for_cpu_and_cuda(tensor) + + # Copied from TorchAO def parse_version(version_string): # Extract just the X.Y.Z part from the version string From 0135666dd53804be59d06fcf7d05cbdfe15925dd Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Wed, 6 Nov 2024 21:41:15 -0500 Subject: [PATCH 3/5] refine version compare Signed-off-by: yiliu30 --- auto_round/utils.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/auto_round/utils.py b/auto_round/utils.py index aae8dade8..a39b28fe8 100644 --- a/auto_round/utils.py +++ b/auto_round/utils.py @@ -870,27 +870,15 @@ def clear_memory(tensor=None): _clear_memory_for_cpu_and_cuda(tensor) -# Copied from TorchAO -def parse_version(version_string): - # Extract just the X.Y.Z part from the version string - match = re.match(r"(\d+\.\d+\.\d+)", version_string) - if match: - version = match.group(1) - return [int(x) for x in version.split(".")] - else: - raise ValueError(f"Invalid version string format: {version_string}") - - def compare_versions(v1, v2): - v1_parts = parse_version(v1) - v2_parts = parse_version(v2) - return (v1_parts > v2_parts) - (v1_parts < v2_parts) + return version.parse(v1) >= version.parse(v2) -def torch_version_at_least(min_version): - return compare_versions(torch.__version__, min_version) >= 0 +def torch_version_at_least(version_string): + return compare_versions(torch.__version__, version_string) +TORCH_VERSION_AT_LEAST_2_6_PRE_RELEASE = torch_version_at_least("2.5.99") TORCH_VERSION_AT_LEAST_2_6 = torch_version_at_least("2.6.0") TORCH_VERSION_AT_LEAST_2_5 = torch_version_at_least("2.5.0") TORCH_VERSION_AT_LEAST_2_4 = torch_version_at_least("2.4.0") @@ -903,7 +891,7 @@ def compile_func_on_hpu(func): def compile_func_on_cuda_or_cpu(func): - if TORCH_VERSION_AT_LEAST_2_6: + if TORCH_VERSION_AT_LEAST_2_6_PRE_RELEASE: return torch.compile(func) else: return func From 0bb39ed1c3657dc60980b8c5529d8f97958fb0e1 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Wed, 6 Nov 2024 21:46:19 -0500 Subject: [PATCH 4/5] update the compile mode environment vars Signed-off-by: yiliu30 --- auto_round/utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/auto_round/utils.py b/auto_round/utils.py index a39b28fe8..87018c220 100644 --- a/auto_round/utils.py +++ b/auto_round/utils.py @@ -884,8 +884,14 @@ def torch_version_at_least(version_string): TORCH_VERSION_AT_LEAST_2_4 = torch_version_at_least("2.4.0") +def check_hpu_compile_mode(): + assert os.environ["PT_HPU_LAZY_MODE"] == "0", "Please set `PT_HPU_LAZY_MODE=0` to use HPU compile mode" + # Note: this is a temporary solution, will be removed in the future + assert os.environ["PT_ENABLE_INT64_SUPPORT"] == "1", "Please set `PT_ENABLE_INT64_SUPPORT=1` to use HPU compile mode" + def compile_func_on_hpu(func): if TORCH_VERSION_AT_LEAST_2_4: + check_hpu_compile_mode() return torch.compile(func, backend="hpu_backend") return func From b2ae935777acb8f9de663dc3be70e0f77d7ce3c1 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Wed, 6 Nov 2024 22:08:41 -0500 Subject: [PATCH 5/5] fix pylints Signed-off-by: yiliu30 --- auto_round/autoround.py | 3 ++- auto_round/utils.py | 13 +++++++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/auto_round/autoround.py b/auto_round/autoround.py index 129189db9..180fd5fff 100644 --- a/auto_round/autoround.py +++ b/auto_round/autoround.py @@ -387,7 +387,8 @@ def quant_layers(self, layer_names, layer_inputs): self.model = mv_module_from_gpu(self.model, self.low_cpu_mem_usage) clear_memory() - quant_layer = compile_func(self.quant_layer) + device = next(self.model.parameters()).device + quant_layer = compile_func(self.quant_layer, device) for layer_name in layer_names: layer_input = layer_inputs[layer_name] layer_input = to_device(layer_input, self.cache_device) diff --git a/auto_round/utils.py b/auto_round/utils.py index 87018c220..783b2d19b 100644 --- a/auto_round/utils.py +++ b/auto_round/utils.py @@ -885,9 +885,14 @@ def torch_version_at_least(version_string): def check_hpu_compile_mode(): - assert os.environ["PT_HPU_LAZY_MODE"] == "0", "Please set `PT_HPU_LAZY_MODE=0` to use HPU compile mode" + assert ( + os.environ["PT_HPU_LAZY_MODE"] == "0" + ), "Please set `PT_HPU_LAZY_MODE=0` to use HPU compile mode" # Note: this is a temporary solution, will be removed in the future - assert os.environ["PT_ENABLE_INT64_SUPPORT"] == "1", "Please set `PT_ENABLE_INT64_SUPPORT=1` to use HPU compile mode" + assert ( + os.environ["PT_ENABLE_INT64_SUPPORT"] == "1" + ), "Please set `PT_ENABLE_INT64_SUPPORT=1` to use HPU compile mode" + def compile_func_on_hpu(func): if TORCH_VERSION_AT_LEAST_2_4: @@ -905,6 +910,6 @@ def compile_func_on_cuda_or_cpu(func): def compile_func(fun, device): if "hpu" in str(device): - return compile_func_on_hpu(fun, device) + return compile_func_on_hpu(fun) else: - return compile_func_on_cuda_or_cpu(fun, device) + return compile_func_on_cuda_or_cpu(fun)