diff --git a/README.md b/README.md
index 5e53236..65a8449 100644
--- a/README.md
+++ b/README.md
@@ -1,3 +1,48 @@
+# 💫 2600-StarCoder QLoRA 4bit
+I followed this guide [Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA
+](https://huggingface.co/blog/4bit-transformers-bitsandbytes) by huggingface and implement the code in this repo 
+to load the model in 4bit int and train using the methods outlined in the paper.
+
+To make this work you're going to need the latese accelerate, transformers, and bitsandbytes libs. 
+```
+pip install -q -U bitsandbytes
+pip install -q -U git+https://github.com/huggingface/transformers.git
+pip install -q -U git+https://github.com/huggingface/peft.git
+pip install -q -U git+https://github.com/huggingface/accelerate.git
+```
+
+The model can beloaded in 11GB of VRAM at 4bit.
+```
+from transformers import BitsAndBytesConfig
+nf4_config = BitsAndBytesConfig(
+   load_in_4bit=True,
+   bnb_4bit_quant_type="nf4",
+   bnb_4bit_use_double_quant=True,
+   bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+model_nf4 = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=nf4_config)
+```
+
+We add a the new optimizer to TrainingArguments()  
+```     
+        ...
+        optim="paged_adamw_8bit"
+        ...
+```
+
+To utilize a second GPU for faster training we change the device map:
+```
+        device_map={"": Accelerator().process_index},
+```
+
+We launch the trainer with accelerate to make use of the second GPU.  We increase the max context to 2600 for 24GB, more GPT VRAM lets us train higher context lengths.
+```
+accelerate launch /home/gpu/code/starcoder/finetune/finetune.py --model_path=bigcode/starcoder --dataset_name=ArmelR/stack-exchange-instruction --subset=data/finetune --split=train --size_valid_set 10000 --streaming --seq_length 2600 --max_steps 1000 --batch_size 1 --input_column_name=question --output_column_name=response --save_freq=100 --learning_rate 0.0001  --lora_r 16
+```
+![image](https://github.com/binaryninja/starcoder/assets/5916066/cc159829-125a-48d2-8239-f65fdbd4ad91)
+
+
 # 💫 StarCoder
 
 [Paper](https://drive.google.com/file/d/1cN-b9GnWtHzQRoE7M7gAEyivY0kl4BYs/view) | [Model](https://huggingface.co/bigcode/starcoder) | [Playground](https://huggingface.co/spaces/bigcode/bigcode-playground) | [VSCode](https://marketplace.visualstudio.com/items?itemName=HuggingFace.huggingface-vscode) | [Chat](https://huggingface.co/spaces/HuggingFaceH4/starchat-playground)
diff --git a/finetune/finetune.py b/finetune/finetune.py
index 96ab961..704430a 100644
--- a/finetune/finetune.py
+++ b/finetune/finetune.py
@@ -4,11 +4,11 @@
 import torch
 from accelerate import Accelerator
 from datasets import load_dataset
-from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, set_peft_model_state_dict
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, set_peft_model_state_dict
 from torch.utils.data import IterableDataset
 from tqdm import tqdm
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, logging, set_seed
-from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
+from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl, BitsAndBytesConfig
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
 
 """
@@ -235,16 +235,37 @@ def create_datasets(tokenizer, args):
 
 
 def run_training(args, train_data, val_data):
+
     print("Loading the model")
     # disable caching mechanism when using gradient checkpointing
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     args.model_path,
+    #     use_auth_token=True,
+    #     use_cache=not args.no_gradient_checkpointing,
+    #     load_in_8bit=True,
+    #     device_map={"": Accelerator().process_index},
+    # )
+
+
     model = AutoModelForCausalLM.from_pretrained(
         args.model_path,
-        use_auth_token=True,
-        use_cache=not args.no_gradient_checkpointing,
-        load_in_8bit=True,
+        load_in_4bit=True,
         device_map={"": Accelerator().process_index},
+        #max_memory=max_memory,
+        #device_map="auto",
+        torch_dtype=torch.bfloat16,
+        quantization_config=BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        ),
     )
-    model = prepare_model_for_int8_training(model)
+    print("loaded")
+
+
+    model.gradient_checkpointing_enable()
+    model = prepare_model_for_kbit_training(model)
 
     lora_config = LoraConfig(
         r=args.lora_r,
@@ -281,9 +302,10 @@ def run_training(args, train_data, val_data):
         fp16=not args.no_fp16,
         bf16=args.bf16,
         weight_decay=args.weight_decay,
-        run_name="StarCoder-finetuned",
+        run_name="2600-StarCoder-finetuned",
         report_to="wandb",
         ddp_find_unused_parameters=False,
+        optim="paged_adamw_8bit"
     )
 
     trainer = Trainer(model=model, args=training_args, train_dataset=train_data, eval_dataset=val_data, callbacks=[SavePeftModelCallback, LoadBestPeftModelCallback])
@@ -296,7 +318,7 @@ def run_training(args, train_data, val_data):
 
 
 def main(args):
-    tokenizer = AutoTokenizer.from_pretrained(args.model_path, use_auth_token=True)
+    tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=args.model_path, use_auth_token=True)
     train_dataset, eval_dataset = create_datasets(tokenizer, args)
     run_training(args, train_dataset, eval_dataset)