code: def set_module_tensor_to_device_with_cache(
      module: nn.Module,
      tensor_name: str,
      device: Union[int, str, torch.device],
      value: Optional[torch.Tensor],
      dtype: Optional[Union[str, torch.dtype]] = None,
      cache_dict: Optional[Dict[str, FP4Params]] = None,
      max_cache_num: Optional[int] = 100,
  ):
      """
      A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
      `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).
  
      Args:
          module (`torch.nn.Module`):
              The module in which the tensor we want to move lives.
          param_name (`str`):
              The full name of the parameter/buffer.
          device (`int`, `str` or `torch.device`):
              The device on which to set the tensor.
          value (`torch.Tensor`):
              The value of the tensor (useful when going from the meta device to any other device).
          dtype (`torch.dtype`, *optional*):
              If passed along the value of the parameter will be cast to this `dtype`. Otherwise, `value` will be cast to
              the dtype of the existing parameter in the model.
          cache_dict (`Dict`, *optional*):
              The cache dict to save layer weights. This can improve the loading speed.
          max_cache_num (`int`, *optional*):
              The maximum number of weights saved in the cache_dict. You can adjust this number based on your GPU memory.
              Default is 100.
      """
      original_tensor_name = tensor_name
      assert value is not None
  
      # Recurse if needed
      if "." in tensor_name:
          splits = tensor_name.split(".")
          for split in splits[:-1]:
              new_module = getattr(module, split)
              if new_module is None:
                  raise ValueError(
                      f"{module} has no attribute {split}."
                  )
              module = new_module
          tensor_name = splits[-1]
  
      # Use cache to load weights
      if original_tensor_name in cache_dict:
          module._parameters[tensor_name] = cache_dict[
              original_tensor_name
          ]
          return
  
      if (
          tensor_name not in module._parameters
          and tensor_name not in module._buffers
      ):
          raise ValueError(
              f"{module} does not have a parameter or a buffer named {tensor_name}."
          )
      is_buffer = tensor_name in module._buffers
      old_value = getattr(module, tensor_name)
  
      if (
          old_value.device == torch.device("meta")
          and device not in ["meta", torch.device("meta")]
          and value is None
      ):
          raise ValueError(
              f"{tensor_name} is on the meta device, we need a `value` to put in on {device}."
          )
  
      if dtype is None:
          # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model
          value = value.to(old_value.dtype)
      elif not str(value.dtype).startswith(
          ("torch.uint", "torch.int", "torch.bool")
      ):
          value = value.to(dtype)
  
      param = (
          module._parameters[tensor_name]
          if tensor_name in module._parameters
          else None
      )
      param_cls = type(param)
  
      with torch.no_grad():
          if isinstance(module, LowBitLinear):
              # load cpu int4 weights
              new_value = FP4Params(
                  data=value,
                  requires_grad=False,
                  quantized=True,
                  _shape=(module.out_features, module.in_features),
                  convert_shape_only=False,
                  qtype=2,
              ).to(device)
              if len(cache_dict) < max_cache_num:
                  cache_dict.update({original_tensor_name: new_value})
          elif isinstance(value, torch.Tensor):
              new_value = value.to(device)
          else:
              new_value = torch.tensor(value, device=device)
          if is_buffer:
              module._buffers[tensor_name] = new_value
          elif isinstance(module, LowBitLinear):
              module._parameters[tensor_name] = new_value
          elif (
              value is not None
              or torch.device(device)
              != module._parameters[tensor_name].device
          ):
              param_cls = type(module._parameters[tensor_name])
              new_value = param_cls(
                  new_value, requires_grad=old_value.requires_grad
              ).to(device)
              module._parameters[tensor_name] = new_value
          else:
              raise NotImplementedError
  
imports: import argparse
  import gc
  import os
  import time
  import warnings
  from typing import Dict, List, Optional, Union
  import torch
  import torch.nn
  from accelerate import init_empty_weights
  from safetensors.torch import load_file, save_file
  from tqdm.auto import tqdm
  from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
  from transformers.generation import GenerationMixin
  from transformers.modeling_outputs import CausalLMOutputWithPast
  from bigdl.llm import optimize_model
  from bigdl.llm.transformers.low_bit_linear import FP4Params, LowBitLinear
label_desc: Use image node_id 1 for calling a global function with example usage: set_module_tensor_to_device_with_cache(module, tensor_name, device, value, dtype, cache_dict, max_cache_num) without return types