BitNet_Llama_model_test_huggingface_GPU.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8de45929",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bitnet_llama.model import LlamaForCausalLM, BitLinear\n",
    "from transformers import AutoTokenizer\n",
    "from transformers import LlamaConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3b4c5185",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"beomi/llama-2-ko-7b\")\n",
    "config = LlamaConfig(\n",
    "    vocab_size=len(tokenizer),\n",
    "    hidden_size=512,\n",
    "    intermediate_size=2048,\n",
    "    max_position_embeddings=512,\n",
    "    num_attention_heads=32,\n",
    "    num_hidden_layers=4,\n",
    "    num_key_value_heads=32,\n",
    "    pretraining_tp=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d0b17eaf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaConfig {\n",
       "  \"attention_bias\": false,\n",
       "  \"bos_token_id\": 1,\n",
       "  \"eos_token_id\": 2,\n",
       "  \"hidden_act\": \"silu\",\n",
       "  \"hidden_size\": 512,\n",
       "  \"initializer_range\": 0.02,\n",
       "  \"intermediate_size\": 2048,\n",
       "  \"max_position_embeddings\": 512,\n",
       "  \"model_type\": \"llama\",\n",
       "  \"num_attention_heads\": 32,\n",
       "  \"num_hidden_layers\": 4,\n",
       "  \"num_key_value_heads\": 32,\n",
       "  \"pretraining_tp\": 1,\n",
       "  \"rms_norm_eps\": 1e-06,\n",
       "  \"rope_scaling\": null,\n",
       "  \"rope_theta\": 10000.0,\n",
       "  \"tie_word_embeddings\": false,\n",
       "  \"transformers_version\": \"4.35.0.dev0\",\n",
       "  \"use_cache\": true,\n",
       "  \"vocab_size\": 46336\n",
       "}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "96eb9ae5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4bd84447",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'12.1'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.version.cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "02abd4c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8c079440",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Thu Oct 19 16:12:18 2023       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 535.104.12             Driver Version: 535.104.12   CUDA Version: 12.2     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  NVIDIA H100 PCIe               Off | 00000000:49:00.0 Off |                    0 |\n",
      "| N/A   39C    P0              85W / 350W |      7MiB / 81559MiB |      0%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "|   1  NVIDIA H100 PCIe               Off | 00000000:5A:00.0 Off |                    0 |\n",
      "| N/A   35C    P0              49W / 350W |      7MiB / 81559MiB |      0%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "|  No running processes found                                                           |\n",
      "+---------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f92dfc30",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LlamaForCausalLM(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d5381a3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.to(torch.bfloat16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "906e80fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(46336, 512)\n",
       "    (layers): ModuleList(\n",
       "      (0-3): 4 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaAttention(\n",
       "          (q_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (k_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (v_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (o_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): BitLinear(in_features=512, out_features=2048, bias=False)\n",
       "          (up_proj): BitLinear(in_features=512, out_features=2048, bias=False)\n",
       "          (down_proj): BitLinear(in_features=2048, out_features=512, bias=False)\n",
       "          (act_fn): SiLUActivation()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (lm_head): BitLinear(in_features=512, out_features=46336, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.to('cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f201d2de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Thu Oct 19 16:12:20 2023       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 535.104.12             Driver Version: 535.104.12   CUDA Version: 12.2     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  NVIDIA H100 PCIe               Off | 00000000:49:00.0 Off |                    0 |\n",
      "| N/A   39C    P0              84W / 350W |    603MiB / 81559MiB |      3%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "|   1  NVIDIA H100 PCIe               Off | 00000000:5A:00.0 Off |                    0 |\n",
      "| N/A   35C    P0              49W / 350W |      7MiB / 81559MiB |      0%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "|    0   N/A  N/A    573892      C   ...b.lee/bitnet-llama/.venv/bin/python      590MiB |\n",
      "+---------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "42a966ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "128590912"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.get_memory_footprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d3a212a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = model(**tokenizer('Hello world!', return_tensors='pt', return_token_type_ids=False).to('cuda:0'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c1cbf996",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1., -1.,  1.,  ..., -1.,  1.,  1.],\n",
       "        [ 1., -1., -1.,  ...,  1.,  1.,  1.],\n",
       "        [ 1., -1., -1.,  ..., -1.,  1., -1.],\n",
       "        ...,\n",
       "        [-1., -1., -1.,  ..., -1.,  1., -1.],\n",
       "        [ 1.,  1.,  1.,  ..., -1., -1.,  1.],\n",
       "        [-1.,  1.,  1.,  ..., -1., -1., -1.]], device='cuda:0',\n",
       "       dtype=torch.bfloat16, grad_fn=<CopyBackwards>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.model.layers[0].self_attn.q_proj.binarize_weights_groupwise()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "77a4bd28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method LlamaPreTrainedModel.save_pretrained of LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(46336, 512)\n",
       "    (layers): ModuleList(\n",
       "      (0-3): 4 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaAttention(\n",
       "          (q_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (k_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (v_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (o_proj): BitLinear(in_features=512, out_features=512, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): BitLinear(in_features=512, out_features=2048, bias=False)\n",
       "          (up_proj): BitLinear(in_features=512, out_features=2048, bias=False)\n",
       "          (down_proj): BitLinear(in_features=2048, out_features=512, bias=False)\n",
       "          (act_fn): SiLUActivation()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (lm_head): BitLinear(in_features=512, out_features=46336, bias=False)\n",
       ")>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.save_pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9e76bcb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained('test_1bit', save_binarized_weights=True)\n",
    "model.save_pretrained('test_full_bit', save_binarized_weights=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b65bbd73",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0032, -0.0339,  0.0150,  ...,  0.0041, -0.0048,  0.0061],\n",
       "        [-0.0105, -0.0049, -0.0586,  ..., -0.0092,  0.0188, -0.0084],\n",
       "        [-0.0383, -0.0109,  0.0031,  ..., -0.0410,  0.0211,  0.0223],\n",
       "        ...,\n",
       "        [ 0.0131, -0.0259,  0.0034,  ...,  0.0233, -0.0281, -0.0131],\n",
       "        [ 0.0062,  0.0198,  0.0085,  ...,  0.0129, -0.0205,  0.0050],\n",
       "        [ 0.0292,  0.0152, -0.0175,  ...,  0.0256,  0.0276,  0.0082]],\n",
       "       device='cuda:0', dtype=torch.bfloat16, requires_grad=True)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.model.layers[0].mlp.gate_proj.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3b9c83b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
       "        1., 1., 1., 1., 1., 1., 1., 1.], device='cuda:0', dtype=torch.bfloat16,\n",
       "       requires_grad=True)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.model.layers[2].post_attention_layernorm.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7d09b689",
   "metadata": {},
   "outputs": [],
   "source": [
    "def binarize_bitlinear_weights(self):\n",
    "    \"\"\"\n",
    "    Binarizes weights of BitLinear layers and returns the updated state_dict.\n",
    "    \"\"\"\n",
    "    state_dict = self.state_dict()\n",
    "    for name, module in self.named_modules():\n",
    "        if name == 'lm_head':\n",
    "            continue\n",
    "        if isinstance(\n",
    "            module, BitLinear\n",
    "        ):\n",
    "            \n",
    "            print(name)\n",
    "            state_dict[name+'.weight'] = module.binarize_weights_groupwise()\n",
    "    return state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "21c81a6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['model.embed_tokens.weight', 'model.layers.0.self_attn.q_proj.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.v_proj.weight', 'model.layers.0.self_attn.o_proj.weight', 'model.layers.0.mlp.gate_proj.weight', 'model.layers.0.mlp.up_proj.weight', 'model.layers.0.mlp.down_proj.weight', 'model.layers.0.input_layernorm.weight', 'model.layers.0.post_attention_layernorm.weight', 'model.layers.1.self_attn.q_proj.weight', 'model.layers.1.self_attn.k_proj.weight', 'model.layers.1.self_attn.v_proj.weight', 'model.layers.1.self_attn.o_proj.weight', 'model.layers.1.mlp.gate_proj.weight', 'model.layers.1.mlp.up_proj.weight', 'model.layers.1.mlp.down_proj.weight', 'model.layers.1.input_layernorm.weight', 'model.layers.1.post_attention_layernorm.weight', 'model.layers.2.self_attn.q_proj.weight', 'model.layers.2.self_attn.k_proj.weight', 'model.layers.2.self_attn.v_proj.weight', 'model.layers.2.self_attn.o_proj.weight', 'model.layers.2.mlp.gate_proj.weight', 'model.layers.2.mlp.up_proj.weight', 'model.layers.2.mlp.down_proj.weight', 'model.layers.2.input_layernorm.weight', 'model.layers.2.post_attention_layernorm.weight', 'model.layers.3.self_attn.q_proj.weight', 'model.layers.3.self_attn.k_proj.weight', 'model.layers.3.self_attn.v_proj.weight', 'model.layers.3.self_attn.o_proj.weight', 'model.layers.3.mlp.gate_proj.weight', 'model.layers.3.mlp.up_proj.weight', 'model.layers.3.mlp.down_proj.weight', 'model.layers.3.input_layernorm.weight', 'model.layers.3.post_attention_layernorm.weight', 'model.norm.weight', 'lm_head.weight'])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.state_dict().keys()#['model.layers.0.self_attn.q_proj']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "06756d83",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaRMSNorm()"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.model.layers[0].post_attention_layernorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9dd2034e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model.layers.0.self_attn.q_proj\n",
      "model.layers.0.self_attn.k_proj\n",
      "model.layers.0.self_attn.v_proj\n",
      "model.layers.0.self_attn.o_proj\n",
      "model.layers.0.mlp.gate_proj\n",
      "model.layers.0.mlp.up_proj\n",
      "model.layers.0.mlp.down_proj\n",
      "model.layers.1.self_attn.q_proj\n",
      "model.layers.1.self_attn.k_proj\n",
      "model.layers.1.self_attn.v_proj\n",
      "model.layers.1.self_attn.o_proj\n",
      "model.layers.1.mlp.gate_proj\n",
      "model.layers.1.mlp.up_proj\n",
      "model.layers.1.mlp.down_proj\n",
      "model.layers.2.self_attn.q_proj\n",
      "model.layers.2.self_attn.k_proj\n",
      "model.layers.2.self_attn.v_proj\n",
      "model.layers.2.self_attn.o_proj\n",
      "model.layers.2.mlp.gate_proj\n",
      "model.layers.2.mlp.up_proj\n",
      "model.layers.2.mlp.down_proj\n",
      "model.layers.3.self_attn.q_proj\n",
      "model.layers.3.self_attn.k_proj\n",
      "model.layers.3.self_attn.v_proj\n",
      "model.layers.3.self_attn.o_proj\n",
      "model.layers.3.mlp.gate_proj\n",
      "model.layers.3.mlp.up_proj\n",
      "model.layers.3.mlp.down_proj\n"
     ]
    }
   ],
   "source": [
    "bin_model = binarize_bitlinear_weights(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "35a25f59",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0032, -0.0339,  0.0150,  ...,  0.0041, -0.0048,  0.0061],\n",
       "        [-0.0105, -0.0049, -0.0586,  ..., -0.0092,  0.0188, -0.0084],\n",
       "        [-0.0383, -0.0109,  0.0031,  ..., -0.0410,  0.0211,  0.0223],\n",
       "        ...,\n",
       "        [ 0.0131, -0.0259,  0.0034,  ...,  0.0233, -0.0281, -0.0131],\n",
       "        [ 0.0062,  0.0198,  0.0085,  ...,  0.0129, -0.0205,  0.0050],\n",
       "        [ 0.0292,  0.0152, -0.0175,  ...,  0.0256,  0.0276,  0.0082]],\n",
       "       device='cuda:0', dtype=torch.bfloat16, requires_grad=True)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.model.layers[0].mlp.gate_proj.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "25a13794",
   "metadata": {},
   "outputs": [],
   "source": [
    "# bin_model.model.layers[0].mlp.gate_proj.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "40641bff",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.config.save_pretrained('./test-config')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb3972b3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bitnet-llama-venv",
   "language": "python",
   "name": "bitnet-llama-venv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}