code: def load_model(
      cfg: DictDefault,
      tokenizer: PreTrainedTokenizerBase,
      inference: bool = False,
      reference_model: bool = False,
  ) -> Tuple[PreTrainedModel, Optional[PeftConfig]]:
      """
      Load a model for a given configuration and tokenizer.
      """
      base_model = cfg.base_model
      model_type = cfg.model_type
      model_config = load_model_config(cfg)
  
      # TODO refactor as a kwarg
      load_in_8bit = cfg.load_in_8bit
  
      if (
          hasattr(model_config, "model_type")
          and model_config.model_type == "btlm"
      ):
          if cfg.flash_attention:
              from axolotl.monkeypatch.btlm_attn_hijack_flash import (
                  replace_btlm_attn_with_flash_attn,
              )
  
              replace_btlm_attn_with_flash_attn(cfg.base_model)
  
      if (
          hasattr(model_config, "model_type")
          and model_config.model_type == "stablelm_epoch"
      ):
          if cfg.flash_attention and cfg.sample_packing:
              from axolotl.monkeypatch.stablelm_attn_hijack_flash import (
                  replace_stablelm_attn_with_flash_attn,
              )
  
              replace_stablelm_attn_with_flash_attn(cfg.base_model)
  
      if (
          cfg.is_llama_derived_model
          and cfg.flash_attention
          and cfg.sample_packing
      ):
          if cfg.device not in ["mps", "cpu"] and not inference:
              from axolotl.monkeypatch.llama_attn_hijack_flash import (
                  replace_llama_attn_with_flash_attn,
              )
  
              LOG.info(
                  "patching with flash attention for sample packing"
              )
              replace_llama_attn_with_flash_attn(
                  packed=cfg.sample_packing,
                  cross_entropy=cfg.flash_attn_cross_entropy,
                  rms_norm=cfg.flash_attn_rms_norm,
              )
      elif cfg.is_llama_derived_model and cfg.xformers_attention:
          from axolotl.monkeypatch.llama_attn_hijack_xformers import (
              hijack_llama_attention,
          )
  
          LOG.info("patching with xformers attention")
          hijack_llama_attention()
      elif cfg.is_llama_derived_model and cfg.sdp_attention:
          from axolotl.monkeypatch.llama_attn_hijack_sdp import (
              hijack_llama_sdp_attention,
          )
  
          LOG.info("patching with sdp attention")
          hijack_llama_sdp_attention()
  
      if (
          cfg.is_mistral_derived_model
          and cfg.flash_attention
          and cfg.sample_packing
      ):
          from axolotl.monkeypatch.mistral_attn_hijack_flash import (
              replace_mistral_attn_with_flash_attn,
          )
  
          LOG.info("patching with flash attention")
          replace_mistral_attn_with_flash_attn(
              packed=cfg.sample_packing
          )
  
      if (
          cfg.model_config_type == "mixtral"
          and cfg.flash_attention
          and cfg.sample_packing
      ):
          from axolotl.monkeypatch.mixtral import (
              replace_mixtral_attn_with_multipack_flash_attn,
          )
  
          LOG.info("patching with flash attention")
          replace_mixtral_attn_with_multipack_flash_attn()
  
      if (
          cfg.is_llama_derived_model
          and (cfg.max_packed_sequence_len or cfg.sample_packing)
          and not inference
      ):
          from axolotl.monkeypatch.llama_expand_mask import (
              hijack_expand_mask,
          )
  
          LOG.info("patching _expand_mask")
          hijack_expand_mask()
  
      model_kwargs = {}
  
      max_memory = cfg.max_memory
      device_map = cfg.device_map
  
      if cfg.gpu_memory_limit:
          gpu_memory_limit = (
              str(cfg.gpu_memory_limit) + "GiB"
              if isinstance(cfg.gpu_memory_limit, int)
              else cfg.gpu_memory_limit
          )
  
          max_memory = {}
          for i in range(torch.cuda.device_count()):
              max_memory[i] = gpu_memory_limit
          max_memory[
              "cpu"
          ] = "256GiB"  # something sufficiently large to fit anything
  
      if max_memory is not None:
          # Based on https://github.com/togethercomputer/OpenChatKit/blob/main/inference/bot.py
          from accelerate import (
              infer_auto_device_map,
              init_empty_weights,
          )
  
          with init_empty_weights():
              model_canvas = AutoModelForCausalLM.from_config(
                  model_config
              )
          model_canvas.tie_weights()
          device_map = infer_auto_device_map(
              model_canvas,
              max_memory=max_memory,
              dtype=cfg.torch_dtype,
          )
          # We can discard max_memory now as we have a device map set up for us
          max_memory = None
  
      model_kwargs["device_map"] = device_map
      model_kwargs["torch_dtype"] = cfg.torch_dtype
      # TODO can we put the reference model on it's own gpu? I think we have to move logits around to calculate loss
      # if cfg.rl:
      #     if torch.cuda.device_count() > 1:
      #         if reference_model:
      #             model_kwargs["device_map"] = "cuda:" + str(
      #                 torch.cuda.current_device() + 1
      #             )
      #         else:
      #             model_kwargs["device_map"] = "cuda:" + str(torch.cuda.current_device())
  
      if is_deepspeed_zero3_enabled():
          del model_kwargs["device_map"]
  
      if cfg.model_revision:
          model_kwargs["revision"] = cfg.model_revision
      if cfg.gptq:
          if not hasattr(model_config, "quantization_config"):
              LOG.warning(
                  "model config does not contain quantization_config information"
              )
          else:
              if cfg.gptq_disable_exllama is not None:
                  model_config.quantization_config[
                      "disable_exllama"
                  ] = cfg.gptq_disable_exllama
              model_kwargs["quantization_config"] = GPTQConfig(
                  **model_config.quantization_config
              )
      if cfg.adapter == "qlora" and cfg.load_in_4bit:
          bnb_config = {
              "load_in_4bit": True,
              "llm_int8_threshold": 6.0,
              "llm_int8_has_fp16_weight": False,
              "bnb_4bit_compute_dtype": cfg.torch_dtype,
              "bnb_4bit_use_double_quant": True,
              "bnb_4bit_quant_type": "nf4",
          }
  
          if cfg.bnb_config_kwargs:
              bnb_config.update(cfg.bnb_config_kwargs)
  
          model_kwargs["quantization_config"] = BitsAndBytesConfig(
              **bnb_config,
          )
      # sample packing uses custom FA2 patch
      if cfg.flash_attention:
          if not cfg.sample_packing:
              if (
                  cfg.is_llama_derived_model
                  or cfg.is_falcon_derived_model
                  or cfg.is_mistral_derived_model
                  or model_config.model_type == "mixtral"
              ):
                  model_kwargs[
                      "attn_implementation"
                  ] = "flash_attention_2"
                  model_config._attn_implementation = (  # pylint: disable=protected-access
                      "flash_attention_2"
                  )
          else:
              if model_config.model_type == "mixtral":
                  model_kwargs[
                      "attn_implementation"
                  ] = "flash_attention_2"
                  model_config._attn_implementation = (  # pylint: disable=protected-access
                      "flash_attention_2"
                  )
              else:
                  model_kwargs["attn_implementation"] = "eager"
                  model_config._attn_implementation = (  # pylint: disable=protected-access
                      "eager"
                  )
          if model_config.model_type == "phi-msft":
              model_config.flash_attn = True
              model_config.flash_rotary = True
              model_config.fused_dense = True
  
      try:
          if (
              cfg.is_llama_derived_model
              and not cfg.trust_remote_code
              and not cfg.gptq
          ):
              from transformers import LlamaForCausalLM
  
              model = LlamaForCausalLM.from_pretrained(
                  base_model,
                  config=model_config,
                  load_in_8bit=cfg.load_in_8bit
                  and cfg.adapter is not None,
                  load_in_4bit=cfg.load_in_4bit
                  and cfg.adapter is not None,
                  **model_kwargs,
              )
  
              if cfg.flash_attention and not inference:
                  from axolotl.monkeypatch.llama_attn_hijack_flash import (
                      replace_llama_mlp_with_swiglu,
                      replace_llama_qkv_with_fused,
                  )
  
                  if cfg.flash_attn_fuse_mlp:
                      LOG.info("patching with SwiGLU")
                      replace_llama_mlp_with_swiglu(model)
  
                  if cfg.flash_attn_fuse_qkv:
                      LOG.info("patching with fused QKV")
                      replace_llama_qkv_with_fused(model)
          # elif model_type == "GPTNeoXForCausalLM" and cfg.flash_attention:
          #     This is a WIP, still an issue with the backward pass
          #     RuntimeError: grad can be implicitly created only for scalar outputs
          #     TODO: try config.sequence_parallel = False
          #     # https://github.com/HazyResearch/flash-attention/blob/40a25c8ee7465cf547b929cfa2937034e37bfce9/tests/models/test_gpt_neox.py#L12
          #     # https://github.com/HazyResearch/flash-attention/tree/main/training#model-components
          #     # add `**kwargs` to https://github.com/HazyResearch/flash-attention/blob/40a25c8ee7465cf547b929cfa2937034e37bfce9/flash_attn/models/gpt.py#L442
          #     from flash_attn.utils.pretrained import state_dict_from_pretrained
          #     from flash_attn.models.gpt import GPTLMHeadModel
          #     from flash_attn.models.gpt_neox import remap_state_dict_hf_gpt_neox, gpt_neox_config_to_gpt2_config
          #     from transformers import GPTNeoXConfig
          #     config = gpt_neox_config_to_gpt2_config(GPTNeoXConfig.from_pretrained(base_model))
          #     config.use_flash_attn = True
          #     config.fused_bias_fc = True
          #     config.fused_mlp = True  # GPT-NeoX-20B uses "gelu_fast"
          #     config.activation_function = "gelu_fast"
          #     config.fused_dropout_add_ln = True
          #     # config.residual_in_fp32 = True
          #
          #     model: GPTLMHeadModel = GPTLMHeadModel.from_pretrained(
          #         base_model,
          #         config,
          #         dtype=torch_dtype,
          #         device=cfg.device,
          #     )
          #     model.train() # sets to train instead of eval mode
          elif (
              model_type == "PhiForCausalLM"
              or model_config.model_type == "phi-msft"
          ):
              from axolotl.models.phi import PhiForCausalLM
  
              model = PhiForCausalLM.from_pretrained(
                  base_model,
                  config=model_config,
                  load_in_8bit=cfg.load_in_8bit
                  and cfg.adapter is not None,
                  load_in_4bit=cfg.load_in_4bit
                  and cfg.adapter is not None,
                  **model_kwargs,
              )
          elif model_type == "MambaLMHeadModel":
              # FIXME this is janky at best and hacked together to make it work
              MambaLMHeadModel = (
                  fix_mamba_attn_for_loss()
              )  # pylint: disable=invalid-name
  
              model_kwargs["dtype"] = model_kwargs["torch_dtype"]
              model_kwargs["device"] = torch.cuda.current_device()
              del model_kwargs["torch_dtype"]
              del model_kwargs["device_map"]
  
              model = MambaLMHeadModel.from_pretrained(
                  base_model,
                  **model_kwargs,
              )
          elif model_type and not cfg.trust_remote_code:
              if cfg.gptq:
                  model = AutoModelForCausalLM.from_pretrained(
                      base_model,
                      config=model_config,
                      trust_remote_code=cfg.trust_remote_code or False,
                      **model_kwargs,
                  )
              else:
                  model = getattr(
                      transformers, model_type
                  ).from_pretrained(
                      base_model,
                      config=model_config,
                      load_in_8bit=cfg.load_in_8bit
                      and cfg.adapter is not None,
                      load_in_4bit=cfg.load_in_4bit
                      and cfg.adapter is not None,
                      trust_remote_code=cfg.trust_remote_code or False,
                      **model_kwargs,
                  )
          else:
              # Shouldn't be a problem most of the time. will obviously error if the model doesn't support this
              # when training starts
              if (
                  hasattr(model_config, "max_seq_len")
                  and model_config.max_seq_len
                  and cfg.sequence_len > model_config.max_seq_len
              ):
                  model_config.max_seq_len = cfg.sequence_len
                  LOG.warning(
                      f"increasing context length to {cfg.sequence_len}"
                  )
              elif (
                  hasattr(model_config, "max_sequence_length")
                  and model_config.max_sequence_length
                  and cfg.sequence_len
                  > model_config.max_sequence_length
              ):
                  model_config.max_sequence_length = cfg.sequence_len
                  LOG.warning(
                      f"increasing context length to {cfg.sequence_len}"
                  )
              if cfg.gptq:
                  model = AutoModelForCausalLM.from_pretrained(
                      base_model,
                      config=model_config,
                      trust_remote_code=cfg.trust_remote_code or False,
                      **model_kwargs,
                  )
              else:
                  model = AutoModelForCausalLM.from_pretrained(
                      base_model,
                      config=model_config,
                      load_in_8bit=cfg.load_in_8bit
                      and cfg.adapter is not None,
                      load_in_4bit=cfg.load_in_4bit
                      and cfg.adapter is not None,
                      trust_remote_code=cfg.trust_remote_code or False,
                      **model_kwargs,
                  )
      except Exception as err:  # pylint: disable=broad-exception-caught
          LOG.exception(err)
          raise err
  
      embeddings_len = (
          math.ceil(len(tokenizer) / 32) * 32
          if cfg.resize_token_embeddings_to_32x
          else len(tokenizer)
      )
      if (
          hasattr(model, "get_input_embeddings")
          and model.get_input_embeddings().num_embeddings
          < embeddings_len
      ):
          model.resize_token_embeddings(embeddings_len)
      else:
          model.tie_weights()
  
      if (
          hasattr(model, "config")
          and hasattr(model.config, "max_position_embeddings")
          and model.config.max_position_embeddings
          and cfg.sequence_len > model.config.max_position_embeddings
      ):
          LOG.warning(
              f"increasing model.config.max_position_embeddings from {model.config.max_position_embeddings} to {cfg.sequence_len}"
          )
          model.config.max_position_embeddings = cfg.sequence_len
  
      if (
          hasattr(model, "config")
          and hasattr(model.config, "bos_token_id")
          and model.config.bos_token_id
          and model.config.bos_token_id != tokenizer.bos_token_id
      ):
          model.config.bos_token_id = tokenizer.bos_token_id
  
      if (
          hasattr(model, "config")
          and hasattr(model.config, "eos_token_id")
          and model.config.eos_token_id
          and model.config.eos_token_id != tokenizer.eos_token_id
      ):
          model.config.eos_token_id = tokenizer.eos_token_id
  
      if hasattr(model, "device") and model.device.type == "cuda":
          log_gpu_memory_usage(LOG, "after model load", model.device)
  
      # make sure these are fp32 per Ramesh et al. (2021)
      embedding_modules = get_linear_embedding_layers(
          cfg.model_config_type
      )
      for name, module in model.named_modules():
          if any(m in name for m in ["norm", "gate"]):
              module.to(torch.float32)
          if model_config.model_type == "btlm":
              # don't upcast lm_head for btlm
              continue
          if any(m in name for m in embedding_modules):
              if hasattr(module, "weight"):
                  module.to(torch.float32)
  
      needs_fa2_dtype = cfg.adapter or cfg.fsdp
      skip_prepare_model_for_kbit_training = False
  
      if cfg.model_config_type == "qwen" and cfg.adapter == "lora":
          # Qwen doesn't play nicely with LoRA if this is enabled
          skip_prepare_model_for_kbit_training = True
  
      if (cfg.adapter == "lora" and load_in_8bit) or (
          cfg.adapter == "qlora" and cfg.load_in_4bit
      ):
          LOG.info(
              "converting PEFT model w/ prepare_model_for_kbit_training"
          )
          if cfg.gradient_checkpointing:
              model.gradient_checkpointing_enable()
          if not skip_prepare_model_for_kbit_training:
              model = prepare_model_for_kbit_training(
                  model,
                  use_gradient_checkpointing=cfg.gradient_checkpointing,
              )
          needs_fa2_dtype = True
  
      # LlamaRMSNorm layers are in fp32 after kbit_training or full finetune, so we need to
      # convert them back to fp16/bf16 for flash-attn compatibility.
      if needs_fa2_dtype or cfg.flash_attention:
          LOG.info(
              "converting modules to %s for flash attention",
              cfg.torch_dtype,
          )
          for name, module in model.named_modules():
              if "norm" in name:
                  module.to(cfg.torch_dtype)
              if any(m in name for m in embedding_modules):
                  if hasattr(module, "weight"):
                      module.to(cfg.torch_dtype)
  
      lora_config = None
      if not reference_model or cfg.lora_model_dir:
          model, lora_config = load_adapter(model, cfg, cfg.adapter)
  
      if (
          cfg.ddp
          and not load_in_8bit
          and not (cfg.rl and cfg.load_in_4bit)
      ):
          model.to(f"cuda:{cfg.local_rank}")
  
      if (
          torch.cuda.device_count() > 1
          and int(os.getenv("WORLD_SIZE", "1")) == 1
      ):
          setattr(model, "is_parallelizable", True)
          setattr(model, "model_parallel", True)
  
      requires_grad = []
      for name, param in model.named_parameters(recurse=True):
          if param.requires_grad:
              requires_grad.append(f"{name}: {param.requires_grad}")
      if len(requires_grad) == 0:
          LOG.warning(
              "there are no parameters that require gradient updates"
          )
      if hasattr(model, "config"):
          model.config.use_cache = False
  
      if cfg.flash_optimum:
          model = BetterTransformer.transform(model)
  
      if cfg.adapter is not None:
          log_gpu_memory_usage(LOG, "after adapters", model.device)
  
      # TODO resume_from_checkpoint handling
      return model, lora_config
  
imports: import logging
  import math
  import os
  from typing import Any, Optional, Tuple, Union
  import addict
  import bitsandbytes
  import torch
  import transformers
  from optimum.bettertransformer import BetterTransformer
  from peft import PeftConfig, prepare_model_for_kbit_training
  from peft.tuners.lora import QuantLinear
  from transformers import AddedToken, AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GPTQConfig, PreTrainedModel, PreTrainedTokenizerBase
  from transformers.deepspeed import is_deepspeed_zero3_enabled
  from axolotl.models.mamba import fix_mamba_attn_for_loss
  from axolotl.prompt_tokenizers import LLAMA_DEFAULT_EOS_TOKEN
  from axolotl.utils.bench import log_gpu_memory_usage
  from axolotl.utils.chat_templates import chat_templates
  from axolotl.utils.dict import DictDefault
  from axolotl.utils.lora_embeddings import get_linear_embedding_layers
label_desc: Use image node_id 4 for calling a global function with example usage: load_model(cfg, tokenizer, inference, reference_model) and returns: model,  lora_config