Is it possible to finetune this on a custom dataset? about mixtral-offloading HOT 7 OPEN

  (model): MixtralModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0-31): 32 x MixtralDecoderLayer(
        (self_attn): MixtralAttention(
          (q_proj): HQQLinearTritonSavable()
          (k_proj): HQQLinearTritonSavable()
          (v_proj): HQQLinearTritonSavable()
          (o_proj): HQQLinearTritonSavable()
          (rotary_emb): MixtralRotaryEmbedding()
        )
        (block_sparse_moe): SparseMoeWrapper(
          (gate): Linear(in_features=4096, out_features=8, bias=False)
        )
        (input_layernorm): MixtralRMSNorm()
        (post_attention_layernorm): MixtralRMSNorm()
      )
    )
    (norm): MixtralRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, config)

import torch.nn as nn
from transformers.models.mixtral.modeling_mixtral import MixtralAttention

class LoRALayer(nn.Module):
    def __init__(self, module: nn.Linear, rank: int):
        super().__init__()
        self.module = module
        self.adapter_A = nn.Parameter(torch.empty(module.in_features, rank, device=module.weight.device))
        nn.init.kaiming_uniform_(self.adapter_A, a=5 ** 0.5)
        self.adapter_B = nn.Parameter(torch.zeros(rank, module.out_features, device=module.weight.device))

    def forward(self, input):
        bottleneck = F.linear(input, self.adapter_A.T)
        residual = F.linear(bottleneck, self.adapter_B.T)
        return self.module(input) + residual

def custom_get_peft_model(model, rank):
    for _, module in model.named_modules():
        if not isinstance(module, MixtralAttention):
            continue
        module.q_proj = LoRALayer(module.q_proj, rank)
        # TODO: {k, v, o}_proj
    return model

model = custom_get_peft_model(model, rank=8)

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in __getattr__(self, name)
   1693             if name in modules:
   1694                 return modules[name]
-> 1695         raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
   1696 
   1697     def __setattr__(self, name: str, value: Union[Tensor, 'Module']) -> None:

AttributeError: 'HQQLinearTritonSavable' object has no attribute 'in_features'

import torch.nn as nn
from transformers.models.mixtral.modeling_mixtral import MixtralAttention
from src.custom_layers import HQQLinearTritonSavable

class LoRALayer(nn.Module):
    def __init__(self, module: HQQLinearTritonSavable, rank: int):
        super().__init__()
        self.module = module
        in_features = module.meta['shape'][1]
        out_features = module.meta['shape'][0]
        self.adapter_A = nn.Parameter(torch.empty(in_features, rank, device=module.W_q.device))
        nn.init.kaiming_uniform_(self.adapter_A, a=5 ** 0.5)
        self.adapter_B = nn.Parameter(torch.zeros(rank, out_features, device=module.W_q.device))
       
    def forward(self, input):
        bottleneck = F.linear(input, self.adapter_A.T)
        residual = F.linear(bottleneck, self.adapter_B.T)
        return self.module(input) + residual

def custom_get_peft_model(model, rank):
    for _, module in model.named_modules():
        if not isinstance(module, MixtralAttention):
            continue
        module.q_proj = LoRALayer(module.q_proj, rank)
        # TODO: {k, v, o}_proj
    return model
model = custom_get_peft_model(model, rank=8)
´´´