codeclm/utils/offload_profiler.py

import torch
from torch.func import functional_call
import queue
import threading
from typing import Dict, List, Any
import omegaconf
from pydantic import BaseModel, validator
from typing import Optional
from functools import wraps

def _callable_once(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        method_called_flag = f"_called_once_{func.__name__}"
        if getattr(self, method_called_flag, False):
            raise RuntimeError(f"{func.__name__} can only be called once.")
        setattr(self, method_called_flag, True)
        return func(self, *args, **kwargs)
    return wrapper

class OffloadCleanCacheWrapperParam(BaseModel):
    module: Any 
    method_name: str
    diff_mem_gb_thre: float

class OffloadParam(BaseModel):
    offload_module: Any 
    cpu_mem_gb: float
    pre_copy_step: Optional[int] = None
    clean_cache_after_forward: Optional[bool] = None
    dtype: Optional[str] = None 
    offload_layer_dict: Dict[str, int] = {}
    ignore_layer_list: List[str] = []
    clean_cache_wrapper: Optional[OffloadCleanCacheWrapperParam] = None
    debug: Optional[bool] = None

    @validator('dtype')
    def parse_dtype(cls, value):
        if value is None:
            return None
        dtype_map = {
            'torch.float16': torch.float16,
            'torch.float32': torch.float32,
            'torch.float64': torch.float64,
            'torch.int64': torch.int64,
        }
        if value not in dtype_map:
            raise ValueError(f"Unsupported dtype: {value}")
        return dtype_map[value]
    
    def init_param_dict(self):
        param_dict = {}
        param_dict['cpu_mem_gb'] = self.cpu_mem_gb
        if self.pre_copy_step is not None:
            param_dict['pre_copy_step'] = self.pre_copy_step
        if self.clean_cache_after_forward is not None:
            param_dict['clean_cache_after_forward'] = self.clean_cache_after_forward
        if self.debug is not None:
            param_dict['debug'] = self.debug
        
        return param_dict
        
    def offload_layer_param_dict(self):
        param_dict = {}
        param_dict['module'] = self.offload_module
        param_dict['offload_layer_dict'] = self.offload_layer_dict
        param_dict['ignore_layer_list'] = self.ignore_layer_list
        param_dict['dtype'] = self.dtype

        return param_dict
    
    def clean_cache_param_dict(self):
        param_dict = {}
        if self.clean_cache_wrapper is not None:
            param_dict['module'] = self.clean_cache_wrapper.module
            param_dict['method_name'] = self.clean_cache_wrapper.method_name
            param_dict['diff_mem_gb_thre'] = self.clean_cache_wrapper.diff_mem_gb_thre

        return param_dict
    
    @staticmethod
    def recursive_print(model, indent=0):
        for field_name, field_info in model.__fields__.items():
            field_value = getattr(model, field_name)
            print(" " * indent + f"{field_name}:")

            if issubclass(type(field_value), BaseModel):
                print(" " * (indent + 2) + f"--- Nested model: {field_value.__class__.__name__}")
                OffloadParam.recursive_print(field_value, indent + 4) 
            else:
                print(" " * (indent + 2) + f"class: {field_value.__class__.__name__}")
                if isinstance(field_value, torch.nn.Module):
                    pass
                else:
                    print(" " * (indent + 2) + f"value: {field_value}")

    def show(self):
        print("-"*20 + "[OffloadParam]" + "-"*20)
        OffloadParam.recursive_print(self)
        print("-"*40)


class OffloadParamParse:
    def __init__(self):
        pass

    @staticmethod
    def _get_model(root_model: torch.nn.Module, model_dir: str):
        assert(model_dir.startswith("self")), f"model_dir {model_dir} must startswith `self`"
        model = root_model
        for layer in model_dir.split('.'):
            if layer == "self":
                continue
            assert(hasattr(model, layer)), f"model not has layer [{layer}]!"
            model = getattr(model, layer)
        return model

    @staticmethod
    def parse_config(root_model: torch.nn.Module, cfg: omegaconf.DictConfig)->OffloadParam:
        assert(hasattr(cfg, "offload_module") and hasattr(cfg, "cpu_mem_gb") and hasattr(cfg, "dtype"))
        
        offload_module = OffloadParamParse._get_model(root_model, cfg.offload_module)
        cpu_mem_gb = cfg.cpu_mem_gb
        dtype = cfg.dtype

        pre_copy_step = cfg.pre_copy_step \
            if hasattr(cfg, "pre_copy_step") else None

        clean_cache_after_forward = cfg.clean_cache_after_forward \
            if hasattr(cfg, "clean_cache_after_forward") else None
            
        offload_layer_dict = {k: v for k, v in cfg.offload_layer_dict.items()} \
            if hasattr(cfg, "offload_layer_dict") else {}

        ignore_layer_list = cfg.ignore_layer_list \
            if hasattr(cfg, "ignore_layer_list") else []
        
        debug = cfg.debug if hasattr(cfg, "debug") else None
        
        clean_cache_wrapper = None
        if hasattr(cfg, "clean_cache_wrapper"):
            clean_cache_cfg = cfg.clean_cache_wrapper
            cc_module = OffloadParamParse._get_model(root_model, clean_cache_cfg.module)
            cc_method_name = clean_cache_cfg.method_name
            diff_mem_gb_thre = clean_cache_cfg.diff_mem_gb_thre
            clean_cache_wrapper = OffloadCleanCacheWrapperParam(
                                        module=cc_module, 
                                        method_name=cc_method_name, 
                                        diff_mem_gb_thre=diff_mem_gb_thre)
        
        return OffloadParam(
            offload_module=offload_module,
            cpu_mem_gb=cpu_mem_gb,
            pre_copy_step=pre_copy_step,
            clean_cache_after_forward=clean_cache_after_forward,
            dtype=dtype,
            offload_layer_dict=offload_layer_dict,
            ignore_layer_list=ignore_layer_list,
            clean_cache_wrapper=clean_cache_wrapper,
            debug=debug
            )


class LayerParamStruct:
    def __init__(self):
        self.count = 0
        self.device_state = None


class OffloadProfiler:
    def __init__(self, device_index=0, cpu_mem_gb=-1, pre_copy_step=1, clean_cache_after_forward=False, debug=False):
        self.clean_cache_after_forward = clean_cache_after_forward
        self.cpu_mem_gb = cpu_mem_gb
        self.cpu_mem_b_count = 0
        self.device_index = device_index
        self.execution_order = []
        self.execution_order_idx = {} 
        self.pin_memory = False
        test_data = torch.rand(1,1, device='cpu')
        pin_data = test_data.pin_memory()
        self.pin_memory = pin_data.is_pinned()
        print(f"pin:{self.pin_memory}")
        self.copy_stream = torch.cuda.Stream() 
        self.copy_queue = queue.Queue() 
        self.layer_param:Dict[str, LayerParamStruct] = {} 
        self.model_map = {}
        self.stop_flag = False
        self.copy_condition = threading.Condition()
        self.queue_condition = threading.Condition()
        self.mem_line_b = 0

        self.copy_thread = threading.Thread(target=self._copy_thread_fun)
        self.copy_thread.daemon = True
        self.copy_thread.start()

        self.cur_copy_idx = 0 
        self.execute_over = False
        self.pre_copy_step = pre_copy_step

        self.tmp_state_list = []
        self.tmp_state_idx = 0
        for i in range(pre_copy_step + 2):
            self.tmp_state_list.append(None)

        self.debug = debug

    def stop(self):
        self.stop_flag = True
        with self.queue_condition:
            self.queue_condition.notify()
        self.copy_thread.join()

        del self.layer_param
        del self.model_map
        del self.copy_stream

    def _copy_thread_fun(self):
        while self.stop_flag == False:
            layer_name = "--"
            with self.queue_condition:
                while self.copy_queue.qsize() == 0 and self.stop_flag == False:
                    self.queue_condition.wait()
                if self.stop_flag == True:
                    break
                layer_name = self.copy_queue.get()
            with torch.cuda.stream(self.copy_stream):
                if layer_name in self.model_map:
                    model = self.model_map[layer_name]
                    self.tmp_state_list[self.tmp_state_idx] = {
                        k: v.to(torch.device(f"cuda:{self.device_index}"), non_blocking=False)
                        for k, v in model.state_dict().items()
                    }
                    self.copy_stream.synchronize()

                    device_state = self.tmp_state_list[self.tmp_state_idx]
                    self.tmp_state_idx = (self.tmp_state_idx + 1) % len(self.tmp_state_list)

                    with self.copy_condition:
                        if layer_name in self.layer_param:
                            self.layer_param[layer_name].count += 1
                        else:
                            self.layer_param[layer_name] = LayerParamStruct()
                            self.layer_param[layer_name].count = 1
                        self.layer_param[layer_name].device_state = device_state
                        self.copy_condition.notify()
                else:
                    print(f"get model error! {layer_name}")
        print("copy thread stop..")

    def _get_new_step_copy_begin_end(self, tag_name):
        
        pre_copy_step = self.pre_copy_step
        pre_copy_step = min(pre_copy_step, len(self.execution_order) // 2)
        
        cur_exe_idx = self.execution_order_idx[tag_name]
        copy_begin = self.cur_copy_idx
        copy_end = cur_exe_idx + pre_copy_step + 1
        if copy_end - copy_begin > len(self.execution_order):
            copy_end %= len(self.execution_order)
        if copy_end - copy_begin > pre_copy_step + 1 or copy_end - copy_begin < 0:
            # jump
            self.cur_copy_idx = cur_exe_idx
            copy_begin, copy_end = self._get_new_step_copy_begin_end(tag_name=tag_name)
        return copy_begin, copy_end
    
    def make_forward_wrapper(self, module, tag_name, ignore_layer_list=[]):
        original_forward = module.forward
        layer_param_size = 0
        for name, param in module.named_parameters():
            layer_param_size += param.data.numel() * param.data.element_size() / 1024 / 1024 #MB
        
        taget_cpu_mem_b = self.cpu_mem_gb * 1024 * 1024 * 1024
        offload = False
        for name, param in module.named_parameters():
            p_name = f"{tag_name}.{name}" if tag_name else name
            for i_layer in ignore_layer_list:
                if p_name.startswith(i_layer):
                    if self.debug:
                        print(f"ignore layer param: {p_name}")
                    continue

            if taget_cpu_mem_b >= 0 and self.cpu_mem_b_count >= taget_cpu_mem_b:
                break
            cpu_data = torch.empty_strided(size=param.data.size(),
                                        stride=param.data.stride(),
                                        dtype=param.data.dtype,
                                        layout=param.data.layout,
                                        device='cpu',
                                        pin_memory=self.pin_memory)
            cpu_data.copy_(param.data)
            param.data = cpu_data

            param_size = param.data.numel() * param.data.element_size()
            self.cpu_mem_b_count += param_size
            offload = True
        if self.debug:
            print(f"layer: {tag_name}, type: {module.__class__.__name__}, size(MB): {layer_param_size}, offload: {offload}, sum_offload_size(MB): {self.cpu_mem_b_count/1024/1024}")
        
        if offload:
            copy_condition = self.copy_condition
            queue_condition = self.queue_condition
            copy_queue = self.copy_queue
            layer_param = self.layer_param
            def forward_wrapper(*args, **kwargs):
                module.forward = original_forward

                execute_over = False if tag_name not in self.execution_order_idx else True
                if execute_over == False:
                    self.model_map[tag_name] = module
                    self.execution_order.append(tag_name)
                    self.execution_order_idx[tag_name] = len(self.execution_order) - 1
                    copy_queue.put(tag_name)
                    with queue_condition:
                        queue_condition.notify()
                else: 
                
                    copy_begin, copy_end = self._get_new_step_copy_begin_end(tag_name=tag_name)
                    if copy_end > copy_begin:
                        for idx in range(copy_begin, copy_end):
                            idx = idx % len(self.execution_order)
                            copy_tag_name = self.execution_order[idx]
                            copy_queue.put(copy_tag_name)
                            with queue_condition:
                                queue_condition.notify()

                        self.cur_copy_idx = copy_end % len(self.execution_order)
                
                run_state = None
                with self.copy_condition:
                    while tag_name not in self.layer_param:
                        copy_condition.wait()
                    run_state = self.layer_param[tag_name].device_state
                    self.layer_param[tag_name].count -= 1
                    
                module.eval()
                with torch.no_grad():
                    output = functional_call(module, run_state, args=args, kwargs=kwargs)
                with self.copy_condition:
                    if self.layer_param[tag_name].count == 0:
                        del self.layer_param[tag_name]
                diff_mem_b_thre = 1 * (1024 ** 3)
                if self.clean_cache_after_forward:
                    reserved = torch.cuda.memory_reserved()
                    if reserved > self.mem_line_b:
                        torch.cuda.empty_cache()
                        cur_reserved = torch.cuda.memory_reserved()
                        diff_mem = reserved - cur_reserved
                        if diff_mem > diff_mem_b_thre:
                            self.mem_line_b = cur_reserved + (reserved - cur_reserved) / 2 + 10
                        else:
                            self.mem_line_b = reserved + 10
                        if self.debug:
                            print(f"child mem line update, clean cache:{reserved/1024/1024}, cur mem: {cur_reserved/1024/1024}  new limit: {self.mem_line_b / 1024 / 1024}, child name: {tag_name}")
                    
                module.forward = forward_wrapper
                return output
            module.forward = forward_wrapper
        
        torch.cuda.empty_cache()
        return module
    
    def reset_empty_cache_mem_line(self):
        self.mem_line_b = 0
        torch.cuda.empty_cache()
    
    def clean_cache_wrapper(self, module, method_name='', diff_mem_gb_thre=1):
        if not hasattr(module, method_name) or not callable(getattr(module, method_name)):
            print(f"no this method {method_name}")
            return module
        
        original_fun = getattr(module, method_name)
        diff_mem_b_thre = diff_mem_gb_thre * (1024 ** 3)
        self.reset_empty_cache_mem_line()

        def clean_wrapper(*args, **kwargs):
            setattr(module, method_name, original_fun)
            output = original_fun(*args, **kwargs)
            reserved = torch.cuda.memory_reserved()
            if reserved > self.mem_line_b:
                torch.cuda.empty_cache()
                cur_reserved = torch.cuda.memory_reserved()
                diff_mem = reserved - cur_reserved
                if diff_mem > diff_mem_b_thre:
                    self.mem_line_b = cur_reserved + (reserved - cur_reserved) / 2 + 10
                else:
                    self.mem_line_b = reserved + 10

                if self.debug:
                    print(f"mem line update, clean cache:{reserved/1024/1024}, cur mem: {cur_reserved/1024/1024}  new limit: {self.mem_line_b / 1024 / 1024}")
            setattr(module, method_name, clean_wrapper)
            return output
        
        setattr(module, method_name, clean_wrapper)
        return module
    
    @_callable_once
    def offload_layer(self, module, offload_layer_dict={},  ignore_layer_list=[], dtype:torch.dtype = None):
        return self._offload_layer(
                                    module=module,
                                    tag="",
                                    offload_layer_dict=offload_layer_dict,
                                    ignore_layer_list=ignore_layer_list,
                                    dtype=dtype
                                    )
    
    def _offload_layer(self, module, tag="", offload_layer_dict={},  ignore_layer_list=[], dtype:torch.dtype = None):
        """
            Offload specific layers of a PyTorch model to a specified depth.
            A model can only be offloaded once.

            Args:
                module (torch.nn.Module): 
                    The PyTorch model containing the layers to offload. This is the model that will be modified in place.
                
                tag (str, optional): 
                    A string identifier for the model. 
                    Default is an empty string.
                
                offload_layer_dict (dict, optional): 
                    A dictionary where keys are layer names and values represent the depth at which the offloading should occur. 
                    For example, 
                    ```offload_layer_dict = {'cfm_wrapper': 5, 'hubert': 4}``` means that the `cfm_wrapper` layer should 
                    be offloaded at depth 5, and the `hubert` layer should be offloaded at depth 4.
                    Default is an empty dictionary.
                
                ignore_layer_list (list, optional): 
                    A list of layer names or parameter identifiers to be ignored during the offloading process. 
                    Layers in this list will not be offloaded, even if they are present in the `offload_layer_dict`. 
                     For example, 
                    ```ignore_layer_list = ['cfm_wrapper.estimator.h', 'cfm_wrapper.estimator.adaln_single']```
                    means that layers starting with `cfm_wrapper.estimator.h` or  'cfm_wrapper.estimator.adaln_single' will not be offload.
                    Default is an empty list.
                
                dtype (torch.dtype, optional): 
                    The data type (e.g., `torch.float16`, `torch.float32`) to which the offloaded layers should be converted. 
                    If `None`, the data type of the layers will remain unchanged. Default is `None`.

            Returns:
                None
        """
        for p in module._parameters.values():
            if p is not None:
                p.data = p.data.to(torch.device(f"cuda:{self.device_index}"))
                if dtype is not None:
                    p.data = p.data.to(dtype)
        for b in module._buffers.values():
            if b is not None:
                b.data = b.data.to(torch.device(f"cuda:{self.device_index}"))
                if dtype is not None:
                    b.data = b.data.to(dtype)
        for attr_name, attr in module.__dict__.items():
            if isinstance(attr, torch.Tensor) and not attr_name.startswith('_'):
                attr.data = attr.data.to(torch.device(f"cuda:{self.device_index}"))
                if dtype is not None:
                    attr.data = attr.data.to(dtype)

        for name, child in module.named_children():
            current_tag = f"{tag}.{name}" if tag else name
            child = child.to(torch.device(f"cuda:{self.device_index}"))
            if dtype is not None:
                child = child.to(dtype)

            torch.cuda.empty_cache()
            setattr(module, name, child)
            pre_name = current_tag.split('.')[0]
            if pre_name not in offload_layer_dict:
                param_size = 0
                for p in child.parameters():
                    param_size += p.data.numel() * p.data.element_size()
                param_size = param_size / 1024 / 1024
                if self.debug:
                    print(f"not offload layer {current_tag}, size: {param_size}MB")
                continue
            
            has_children = any(child.named_children())
            layer_count = current_tag.count('.') + 1
            
            layer_deep = offload_layer_dict[pre_name]
            if layer_count >= layer_deep:
                has_children = False 
            
            if has_children:
                self._offload_layer(module=child, 
                                   tag=current_tag, 
                                   offload_layer_dict=offload_layer_dict, 
                                   ignore_layer_list=ignore_layer_list,
                                   dtype=dtype)
                continue

            ignore = False
            for i_layer in ignore_layer_list:
                if current_tag.startswith(i_layer):
                    ignore = True
                    if self.debug:
                        print(f"ignore layer offload: {current_tag}")
                    break
    
            if hasattr(child, "forward") and not ignore:
                child = self.make_forward_wrapper(
                    child, current_tag, ignore_layer_list=ignore_layer_list
                )
        return module
    
    def get_execution_order(self):
        return self.execution_order