models/lakonlab/utils/misc.py

# Copyright (c) 2025 Hansheng Chen

import gc
import contextlib
import functools
import torch
import torch.distributed as dist

from torch.distributed.tensor import DTensor
try:
    from torch.distributed.fsdp import FSDPModule
except:
    pass
from functools import partial
from six.moves import map, zip
from mmcv.parallel import is_module_wrapper
from peft.tuners.lora import LoraLayer


def multi_apply(func, *args, **kwargs):
    """Apply function to a list of arguments.

    Note:
        This function applies the ``func`` to multiple inputs and
        map the multiple outputs of the ``func`` into different
        list. Each list contains the same type of outputs corresponding
        to different inputs.

    Args:
        func (Function): A function that will be applied to a list of
            arguments

    Returns:
        tuple(list): A tuple containing multiple list, each list contains \
            a kind of returned results by the function
    """
    pfunc = partial(func, **kwargs) if kwargs else func
    map_results = map(pfunc, *args)
    return tuple(map(list, zip(*map_results)))


def reduce_mean(tensor):
    """Obtain the mean of tensor on different GPUs."""
    if not (dist.is_available() and dist.is_initialized()):
        return tensor
    tensor = tensor.clone()
    dist.all_reduce(tensor.div_(dist.get_world_size()), op=dist.ReduceOp.SUM)
    return tensor


def rgetattr(obj, attr, *args):
    def _getattr(obj, attr):
        if is_module_wrapper(obj):
            obj = obj.module
        return getattr(obj, attr, *args)
    return functools.reduce(_getattr, [obj] + attr.split('.'))


def rsetattr(obj, attr, val):
    pre, _, post = attr.rpartition('.')
    pre = rgetattr(obj, pre) if pre else obj
    if is_module_wrapper(pre):
        pre = pre.module
    return setattr(pre, post, val)


def rhasattr(obj, attr):
    return rgetattr(obj, attr, None) is not None


def rdelattr(obj, attr):
    pre, _, post = attr.rpartition('.')
    pre = rgetattr(obj, pre) if pre else obj
    if is_module_wrapper(pre):
        pre = pre.module
    return delattr(pre, post)


class module_requires_grad:
    def __init__(self, module, requires_grad=True):
        self.module = module
        self.requires_grad = requires_grad
        self.prev = []

    def __enter__(self):
        for p in self.module.parameters():
            self.prev.append(p.requires_grad)
            p.requires_grad = self.requires_grad

    def __exit__(self, exc_type, exc_value, traceback):
        for p, r in zip(self.module.parameters(), self.prev):
            p.requires_grad = r


class module_eval:
    def __init__(self, module):
        self.module = module
        self.prev = None

    def __enter__(self):
        self.prev = self.module.training
        self.module.train(False)

    def __exit__(self, exc_type, exc_value, traceback):
        self.module.train(self.prev)


def all_frozen(modules):
    for module in modules:
        for p in module.parameters():
            if p.requires_grad:
                return False
    return True


def tie_untrained_submodules(tgt_module, src_module, tie_tgt_lora_base_layer=False):
    for key, src_submodule in src_module._modules.items():
        if key in tgt_module._modules:
            if (tie_tgt_lora_base_layer
                    and isinstance(tgt_module._modules[key], LoraLayer)
                    and not isinstance(src_submodule, LoraLayer)):
                if all_frozen((tgt_module._modules[key]._modules['base_layer'], src_submodule)):
                    tgt_module._modules[key]._modules['base_layer'] = src_submodule
                else:
                    tie_untrained_submodules(
                        tgt_module._modules[key]._modules['base_layer'], src_submodule, tie_tgt_lora_base_layer)
            else:
                if all_frozen((tgt_module._modules[key], src_submodule)):
                    tgt_module._modules[key] = src_submodule
                else:
                    tie_untrained_submodules(
                        tgt_module._modules[key], src_submodule, tie_tgt_lora_base_layer)


def clone_params(tgt_module, src_module, recursive=True):
    """Clone parameters and buffers from src_module to tgt_module (sharing the same structure).
    Tied parameters/buffers are not cloned. Used for EMA model initialization.
    """
    for key, val in src_module._parameters.items():
        if (val is not None) \
                and (val is not tgt_module._parameters[key]):
            tgt_module._parameters[key] = val.clone()
    for key, val in src_module._buffers.items():
        if val is not tgt_module._buffers[key]:
            tgt_module._buffers[key] = val.clone()
    if recursive:
        for key, val in src_module._modules.items():
            clone_params(
                tgt_module._modules[key], val, recursive)


@torch.no_grad()
def kai_zhang_clip_grad(model, max_norm: float) -> float:
    """
    https://github.com/Kai-46/minFM/blob/385568691b021aa4888269dafa671900daf17cf6/utils/clip_grad.py#L9
    """
    shard_size, replicate_factor = 1, dist.get_world_size()
    if isinstance(model, FSDPModule):
        shard_size = model._get_fsdp_state()._fsdp_param_group.mesh_info.shard_mesh_size
        replicate_factor = dist.get_world_size() // shard_size

    # Separate DTensor and non-DTensor parameters
    all_param_grads = []
    dtensor_param_grads = []
    regular_param_grads = []

    for p in model.parameters():
        if (not p.requires_grad) or (p.grad is None):
            continue

        if isinstance(p.grad.data, DTensor):
            local_p_grad = p.grad.data.to_local()
            dtensor_param_grads.append(local_p_grad.ravel())
        else:
            local_p_grad = p.grad.data
            regular_param_grads.append(local_p_grad.ravel())

        all_param_grads.append(local_p_grad)

    # Compute local square sum for each group separately
    local_sq_sum = torch.tensor(0.0, device=all_param_grads[0].device)

    if dtensor_param_grads:
        dtensor_sq_sum = (torch.cat(dtensor_param_grads, dim=0) ** 2).float().sum()
        local_sq_sum = local_sq_sum + dtensor_sq_sum

    if regular_param_grads:
        regular_sq_sum = (torch.cat(regular_param_grads, dim=0) ** 2).float().sum()
        local_sq_sum = local_sq_sum + regular_sq_sum / shard_size

    # Single all-reduce operation
    global_sq_sum = local_sq_sum.clone()
    dist.all_reduce(global_sq_sum, op=dist.ReduceOp.SUM)
    global_sq_sum = global_sq_sum / replicate_factor

    total_norm = global_sq_sum.sqrt().item()

    # Only apply clipping when exceeding threshold
    if total_norm > max_norm:
        clip_factor = max_norm / total_norm
        torch._foreach_mul_(all_param_grads, clip_factor)

    return total_norm


def materialize_meta_states(module, device=None):
    if device is None:
        device = torch.get_default_device()
    assert device != torch.device('meta'), 'Please specify a non-meta device.'
    for mod in module.modules():
        for name, p in list(mod._parameters.items()):
            if p is not None and p.is_meta:
                new = torch.empty_like(p, device=device)
                new = torch.nn.Parameter(new, requires_grad=p.requires_grad)
                mod._parameters[name] = new
        for name, b in list(mod._buffers.items()):
            if b is not None and b.is_meta:
                new = torch.empty_like(b, device=device)
                mod._buffers[name] = new
    return module


@contextlib.contextmanager
def gc_context(enable=False):
    prev_enabled = gc.isenabled()
    if enable:
        gc.enable()
    else:
        gc.disable()
    try:
        yield
    finally:
        if prev_enabled:
            gc.enable()
        else:
            gc.disable()