optimum/amd/brevitas/accelerate_utils.py

# Copyright 2023 The HuggingFace Team. All rights reserved.
# Licensed under the MIT License.

import logging
from typing import Dict, Mapping, Optional, Union

import brevitas.config as config
import torch
from accelerate import dispatch_model, infer_auto_device_map
from accelerate.hooks import AlignDevicesHook, add_hook_to_module, remove_hook_from_module
from accelerate.utils import (
    check_tied_parameters_in_config,
    compute_module_sizes,
    find_tied_parameters,
    get_max_layer_size,
    get_max_memory,
    send_to_device,
)
from accelerate.utils.modeling import named_module_tensors
from brevitas.graph.utils import get_module
from brevitas.utils.python_utils import recurse_getattr
from psutil import virtual_memory


logger = logging.getLogger(__name__)


def align_input(model, device_map):
    if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
        main_device = "cpu"
    else:
        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
    hook = AlignDevicesHook(execution_device=main_device, io_same_device=True, skip_keys=None, tied_params_map=None)
    add_hook_to_module(model, hook)
    return model


# Adapted from accelerate.utils.modeling.infer_auto_device_map
def infer_fx_auto_device_map(
    model: torch.fx.GraphModule,
    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
    dtype: Optional[Union[str, torch.dtype]] = None,
    special_dtypes: Optional[Dict[str, Union[str, torch.dtype]]] = None,
    verbose: bool = False,
):
    """
    Extends accelerate's infer_auto_device_map function to be compatible with torch.fx.GraphModule.

    The main modifications are:
    - Work around the fact that module.__class__.__name__ is Module for everything
    - We do not need to keep entire blocks together anymore, since we add a functional equivalent of the AlignDeviceHook
    before every call function.
    """
    # TODO: Why no no_split_module_classes, clean_result parameters?

    # Get default / clean up max_memory
    max_memory = get_max_memory(max_memory)

    devices = list(max_memory.keys())
    if "disk" not in devices:
        devices.append("disk")
    gpus = [device for device in devices if device not in ["cpu", "disk"]]

    # Devices that need to keep space for a potential offloaded layer.
    if "mps" in gpus:
        main_devices = ["mps"]
    elif len(gpus) > 0:
        main_devices = [gpus[0], "cpu"]
    else:
        main_devices = ["cpu"]

    module_sizes = compute_module_sizes(model, dtype=dtype, special_dtypes=special_dtypes)
    tied_parameters = find_tied_parameters(model)

    if check_tied_parameters_in_config(model) and len(tied_parameters) == 0:
        logger.warn(
            "The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function."
        )

    device_map = {}
    current_device = 0
    current_memory_used = 0

    call_list = []
    buffers_attributes = [n for n, _ in list(named_module_tensors(model, recurse=True))]
    all_modules = [n.target for n in list(model.graph.nodes) if n.op == "call_module"]
    for node in model.graph.nodes:
        # If it's a module, we simply offload it or move it to the desired device
        if node.op == "call_module":
            name = node.target
            module = get_module(model, node.target)
            call_list.append((name, module))
        # If it's get_attr, we check what module it is attached to
        # In case the module is not part of call_module, we specifically allocate the buffer/parameter on some device
        # NB: This does NOT guarantee that it will be aligned with whatever input tensor it will be combined with
        # For that, there is a separate function
        if node.op == "get_attr":
            target = node.target
            if target in buffers_attributes:
                module_name = ".".join(target.split(".")[:-1])
                if module_name not in all_modules:
                    module = get_module(model, target)
                    call_list.append((target, module))

    # Direct submodules and parameters
    modules_to_treat = call_list

    # Initialize maximum largest layer, to know which space to keep in memory
    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, [])

    # Ready ? This is going to be a bit messy.
    while len(modules_to_treat) > 0:
        name, module = modules_to_treat.pop(0)
        if verbose:
            print(f"\nTreating module {name}.")
        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
        max_layer_names = [n for n in max_layer_names if n != name and not n.startswith(name + ".")]
        if len(max_layer_names) == 0:
            max_layer_size, max_layer_names = get_max_layer_size(
                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)], module_sizes, []
            )
        # Assess size needed
        module_size = module_sizes[name]

        # We keep relevant tied parameters only: one of the tied parameters in the group is inside the current module
        # and the other is not.
        tied_param_goups = [
            tied_group
            for tied_group in tied_parameters
            if any(name + "." in k + "." for k in tied_group) and not all(name + "." in k + "." for k in tied_group)
        ]

        if verbose and len(tied_param_goups) > 0:
            print(f"  Found the relevant tied param groups {tied_param_goups}")
        # Then we keep track of all the parameters that are tied to the current module, but not in the current module
        tied_params = sum(
            [[p for p in tied_group if name + "." not in p + "."] for tied_group in tied_param_goups], []
        )
        if verbose and len(tied_params) > 0:
            print(f"  So those parameters need to be taken into account {tied_params}")

        device = devices[current_device]
        current_max_size = max_memory[device] if device != "disk" else None
        # Reduce max size available by the largest layer.
        if devices[current_device] in main_devices:
            current_max_size = current_max_size - max_layer_size
        # Case 1 -> We're too big!
        if current_max_size is not None and current_memory_used + module_size > current_max_size:
            # For FX, we never split a leaf call_module
            if verbose:
                print(
                    f"Not enough space on {devices[current_device]} to put {name} (space available "
                    f"{current_max_size-current_memory_used}, module size {module_size})."
                )

            if verbose:
                print("This module cannot be split, going to the next device.")
            current_device += 1
            modules_to_treat = [(name, module)] + modules_to_treat
            current_memory_used = 0

        # Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
        elif len(tied_params) > 0:
            # First locate all tied modules
            tied_module_names = []
            tied_modules = []
            for tied_param in tied_params:
                tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
                tied_module_names.append(modules_to_treat[tied_module_index][0])
                tied_modules.append(modules_to_treat[tied_module_index][1])
            if verbose:
                print(
                    f"  It looks like {name} is going to fit on {devices[current_device]} but we have tied "
                    f"parameters to account for.\n  - Names {tied_params}\n  - Module names {tied_module_names}"
                )

            # Let's see if it all fits first
            module_size_with_ties = module_size
            for tied_param, tied_module_name in zip(tied_params, tied_module_names):
                module_size_with_ties += module_sizes[tied_module_name] - module_sizes[tied_param]

            if current_max_size is None or current_memory_used + module_size_with_ties <= current_max_size:
                # We really really fit!
                if verbose:
                    print(f"Putting {name} and {tied_module_names} on {devices[current_device]}.")
                current_memory_used += module_size_with_ties
                device_map[name] = devices[current_device]
                for tied_module_name in tied_module_names:
                    if tied_module_name in [m[0] for m in modules_to_treat]:
                        # The module may have been removed by a previous iteration of this loop.
                        tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n == tied_module_name][
                            0
                        ]
                        modules_to_treat.pop(tied_module_index)
                    device_map[tied_module_name] = devices[current_device]

            else:
                # We don't fit with the tied modules. Next question is: can we split one of the tied modules to make it
                # smaller or do we need to go on the next device?
                if verbose:
                    print(
                        f"Not enough space on {devices[current_device]} to put {name} and {tied_module_names} (space "
                        f"available {current_max_size-current_memory_used}, needed size {module_size_with_ties})."
                    )
                split_happened = False
                for tied_module_name, tied_module in zip(tied_module_names, tied_modules):
                    tied_module_children = list(tied_module.named_children())
                    if len(tied_module_children) == 0:
                        # can't break this one.
                        continue
                    if verbose:
                        print(f"Splitting {tied_module_name}.")
                    tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
                    tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
                    tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n == tied_module_name][0]

                    modules_to_treat = (
                        [(name, module)]
                        + modules_to_treat[:tied_module_index]
                        + tied_module_children
                        + modules_to_treat[tied_module_index + 1 :]
                    )
                    # Update the max layer size.
                    max_layer_size, max_layer_names = get_max_layer_size(
                        [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
                        module_sizes,
                        [],
                    )
                    split_happened = True
                    break

                if not split_happened:
                    # If the tied module is not split, we go to the next device
                    if verbose:
                        print("None of the tied module can be split, going to the next device.")
                    current_device += 1
                    modules_to_treat = [(name, module)] + modules_to_treat
                    current_memory_used = 0

        else:
            if verbose:
                if current_max_size is None:
                    print(f"Putting {name} (size={module_size}) on {devices[current_device]}.")
                else:
                    print(
                        f"Putting {name} (size={module_size}) on {devices[current_device]} "
                        f"(available={current_max_size-current_memory_used})."
                    )
            current_memory_used += module_size
            device_map[name] = devices[current_device]

    # If we have only one device, we simplify the device_map
    if len(set(device_map.values())) == 1:
        device_map = {"": list(device_map.values())[0]}
    return device_map


def offload_call_function(model: torch.fx.GraphModule, device_map: Dict):
    """
    Attaches AlignDevicesHook to fx.GraphModule call_function nodes. Although accelerate's `offload_model` attaches hooks
    to submodules, it is unable to detect call_function.
    """
    # If we only have one device, offloading is not needed
    if len(set(device_map.values())) == 1:
        return

    for node in model.graph.nodes:
        if node.op == "call_function":

            def new_func(*args, old_callable=node.target, **kwargs):
                args = list(args)
                device_mapping = {}

                # Identify the device for each tensor in args and kwargs
                for _, arg in enumerate(args):
                    all_devices = find_all_devices(arg)
                    if all_devices is not None:
                        device_mapping.update(dict(all_devices))

                for k, v in kwargs.items():
                    all_devices = find_all_devices(arg)
                    if all_devices is not None:
                        device_mapping.update(dict(all_devices))

                total_devices = [d for d in list(device_mapping.values()) if d is not None]

                # If there is only one device, no re-alignement is necessary
                if len(set(total_devices)) > 1:
                    # Pick the main device, i.e. the first device that is not 'cpu' or 'disk'
                    if set(device_mapping.values()) == {"cpu"} or set(device_mapping.values()) == {"cpu", "disk"}:
                        device = "cpu"
                    else:
                        device = [d for d in device_mapping.values() if d not in ["cpu", "disk"]][0]
                    # Align args and kwargs to the same device
                    args = send_to_device(args, device)
                    kwargs = send_to_device(kwargs, device)

                out = old_callable(*args, **kwargs)

                if len(set(total_devices)) > 1:
                    # Restore the original device to avoid memory leaks
                    for k, v in device_mapping.items():
                        k = k.to(v)

                return out

            node.meta["orig_target"] = node.target
            node.target = new_func

    model.recompile()
    model.graph.lint()


def remove_hooks(model: torch.nn.Module):
    for module in model.modules():
        if hasattr(module, "_hf_hook"):
            if hasattr(module, "allocate_params"):
                del module.allocate_params
            if hasattr(module, "offload_params"):
                del module.offload_params
    remove_hook_from_module(model, recurse=True)
    model.cpu()
    if hasattr(model, "graph"):
        for node in model.graph.nodes:
            if node.op == "call_function":
                if "orig_target" in node.meta:
                    node.target = node.meta["orig_target"]
                    del node.meta["orig_target"]
        model.recompile()
        model.graph.lint()


def update_internal_dict(module):
    prefix = module._hf_hook.weights_map.prefix
    for key in module.state_dict().keys():
        # It might happen that we call an quantization's inner modules, and this cause some parameters to be
        # already on meta device. This is not a problem for their value but we need to check here
        curr_device = (recurse_getattr(module, key + ".data")).device
        if str(curr_device) != "meta":
            module._hf_hook.weights_map.dataset.state_dict[prefix + key] = (
                recurse_getattr(module, key + ".data")
            ).cpu()


def find_all_devices(data):
    """
    Finds the device on which a nested dict/list/tuple of tensors lies (assuming they are all on the same device).
    Args:
        (nested list/tuple/dictionary of `torch.Tensor`): The data we want to know the device of.
    """
    if isinstance(data, Mapping):
        devices = []
        for obj in data.values():
            device = find_all_devices(obj)
            if device is not None:
                devices.extend(device)
        return devices
    elif isinstance(data, (tuple, list)):
        devices = []
        for obj in data:
            device = find_all_devices(obj)
            if device is not None:
                devices.extend(device)
        return devices
    elif isinstance(data, torch.Tensor):
        return [(data, str(data.device))]


def calc_gpu_device_map(absolute_mem_margin: float = 2.0 * 1e9, relative_mem_margin: float = 0.3) -> Dict[int, float]:
    torch.cuda.empty_cache()
    gpu_device_map = {
        i: (torch.cuda.mem_get_info(i)[0] - absolute_mem_margin) * (1.0 - relative_mem_margin)
        for i in range(torch.cuda.device_count())
    }
    return gpu_device_map


def calc_cpu_device_map(absolute_mem_margin: float = 2.0 * 1e9, relative_mem_margin: float = 0.3) -> Dict[str, float]:
    cpu_device_map = {"cpu": (virtual_memory().available - absolute_mem_margin) * (1.0 - relative_mem_margin)}
    return cpu_device_map


def offload_model(
    model: torch.nn.Module,
    gpu_device_map: Optional[Dict[int, float]] = None,
    cpu_device_map: Optional[Dict[str, float]] = None,
) -> torch.nn.Module:
    """
    Wraps accelerate's infer_auto_device_map and dispatch_model.

    This functions if compatible both with classic nn.Modules, and with torch.fx.GraphModule.
    """

    # FX vs non-FX model need different offloading
    config._FULL_STATE_DICT = True
    if gpu_device_map is None:
        gpu_device_map = calc_gpu_device_map()
    if cpu_device_map is None:
        cpu_device_map = calc_cpu_device_map()
    memory_map = {**cpu_device_map, **gpu_device_map}

    if isinstance(model, torch.fx.GraphModule):
        device_map = infer_fx_auto_device_map(model, memory_map)
        offload_call_function(model, device_map)
    else:
        device_map = infer_auto_device_map(model, memory_map, no_split_module_classes=model._no_split_modules)

    model = dispatch_model(model, device_map)

    # Fixes an asymetric behavior in Accelerate where hooks are not attached at all when a single device is used.
    # TODO: Fix directly in accelerate.
    if len(set(device_map.values())) == 1:
        model = align_input(model, device_map)

    config._FULL_STATE_DICT = False
    if "disk" in model.hf_device_map.values():
        raise ValueError("Disk offload is not supported with quantization.")

    # We attach these functions to the hooked modules for convenience when modifying parameters during PTQ (e.g. SmoothQuant).
    # Attaching these functions allows use to fix a bug in accelerate with offloading to RAM/disk where even though a submodule parameter is updated, it is actually not updated in the AlignDevicesHook `weights_map` and thus
    # the update is ignored elsewhere.
    # TODO: Fix this bug directly in accelerate. https://github.com/huggingface/accelerate/pull/2214 would fix the bug for RAM offliading.
    def allocate_params(module):
        """
        This function calls the pre_forward function of the _hf_hook, making sure parameters are on
        the selected device, rather than on the meta device.
        """
        if module._hf_hook.offload is False:
            return
        # When quantizing and retrieving parameters (e.g., during GPTQ), we want to recurse through
        # all the submodules
        for m in module.modules():
            if hasattr(m, "_hf_hook"):
                m._hf_hook.pre_forward(m)

    def offload_params(module):
        """
        This functions moves the parameters back to the meta device, after making sure to update the
        internal state dict with the most recent values.
        """
        if module._hf_hook.offload is False:
            return
        update_internal_dict(module)
        for m in module.modules():
            if hasattr(m, "_hf_hook"):
                m._hf_hook.post_forward(m, torch.tensor([]))

    for module in model.modules():
        if hasattr(module, "_hf_hook"):
            module.allocate_params = allocate_params
            module.offload_params = offload_params

    return model