class_docstr: Tools for reconstructing a model from a remote LoRA model.
  (Transferring only LoRA data costs much less!)
  Usage:
      Step 1 (Sender):
          filter_state_dict_lora()
  
      Step 2 (Sender, Optional):
          extract_lora_config()
  
      Step 3 (Sender):
          send state_dict_lora and lora_config_dict
  
      Step 4 (Receiver):
          reconstruct_increase()
  
      Step 5 (Receiver):
          load_state_dict_increase()
docstr: if keep_non_lora, also return non_lora state_dict
docstr_tok: ["if","keep_non_lora",",","also","return","non_lora","state_dict"]
code: def filter_state_dict_lora(
      state_dict: Dict[str, Any], keep_non_lora=False
  ):
      """
      if keep_non_lora, also return non_lora state_dict
      """
      state_dict_lora = OrderedDict()
      state_dict_non_lora = OrderedDict()
      for k, v in state_dict.items():
          if "lora_A" in k or "lora_B" in k:
              state_dict_lora[k] = v
          elif keep_non_lora:
              state_dict_non_lora[k] = v
      if keep_non_lora:
          return state_dict_lora, state_dict_non_lora
      else:
          return state_dict_lora, None
  
file_path: colossalai/applications/Chat/coati/ray/lora_constructor.py
imports: from collections import OrderedDict
  from dataclasses import dataclass
  from typing import Any, Dict
  import torch.nn
  from coati.models.lora import LoraLinear
label_desc: Use image node_id 6 for calling the LoRAConstructor obj's underlying member method code with example usage: obj.filter_state_dict_lora(state_dict, keep_non_lora) and returns: state_dict_lora,  state_dict_non_lora, state_dict_lora,  None