Open In Colab

In [1]:
!rm -rf outputs wandb sample_data
In [2]:
!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
!pip install -q datasets
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 92.2/92.2 MB 9.0 MB/s eta 0:00:00
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 224.5/224.5 kB 7.5 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 7.8/7.8 MB 94.3 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.3/1.3 MB 71.4 MB/s eta 0:00:00
  Building wheel for transformers (pyproject.toml) ... done
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 219.1/219.1 kB 8.8 MB/s eta 0:00:00
  Building wheel for peft (pyproject.toml) ... done
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
  Building wheel for accelerate (pyproject.toml) ... done
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 474.6/474.6 kB 14.0 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 110.5/110.5 kB 14.1 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 212.5/212.5 kB 15.4 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 134.3/134.3 kB 17.0 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.0/1.0 MB 48.2 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 114.5/114.5 kB 10.6 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 268.8/268.8 kB 32.2 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 149.6/149.6 kB 20.7 MB/s eta 0:00:00
In [3]:
!pip install einops scipy
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting einops
  Downloading einops-0.6.1-py3-none-any.whl (42 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 42.2/42.2 kB 2.6 MB/s eta 0:00:00
Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (1.10.1)
Requirement already satisfied: numpy<1.27.0,>=1.19.5 in /usr/local/lib/python3.10/dist-packages (from scipy) (1.22.4)
Installing collected packages: einops
Successfully installed einops-0.6.1
In [4]:
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

model_id = "tiiuae/falcon-rw-1b"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0}, trust_remote_code=True)
Downloading (…)okenizer_config.json:   0%|          | 0.00/255 [00:00<?, ?B/s]
Downloading (…)olve/main/vocab.json:   0%|          | 0.00/798k [00:00<?, ?B/s]
Downloading (…)olve/main/merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]
Downloading (…)/main/tokenizer.json:   0%|          | 0.00/2.11M [00:00<?, ?B/s]
Downloading (…)cial_tokens_map.json:   0%|          | 0.00/99.0 [00:00<?, ?B/s]
Downloading (…)lve/main/config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]
Downloading (…)/configuration_RW.py:   0%|          | 0.00/2.61k [00:00<?, ?B/s]
A new version of the following files was downloaded from https://huggingface.co/tiiuae/falcon-rw-1b:
- configuration_RW.py
. Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.
Downloading (…)main/modelling_RW.py:   0%|          | 0.00/47.5k [00:00<?, ?B/s]
A new version of the following files was downloaded from https://huggingface.co/tiiuae/falcon-rw-1b:
- modelling_RW.py
. Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.
Downloading pytorch_model.bin:   0%|          | 0.00/2.62G [00:00<?, ?B/s]
===================================BUG REPORT===================================
Welcome to bitsandbytes. For bug reports, please run

python -m bitsandbytes

 and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues
================================================================================
bin /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda118.so
CUDA_SETUP: WARNING! libcudart.so not found in any environmental path. Searching in backup paths...
CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so.11.0
CUDA SETUP: Highest compute capability among GPUs detected: 7.5
CUDA SETUP: Detected CUDA version 118
CUDA SETUP: Loading binary /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda118.so...
/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: /usr/lib64-nvidia did not contain ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] as expected! Searching further paths...
  warn(msg)
/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/sys/fs/cgroup/memory.events /var/colab/cgroup/jupyter-children/memory.events')}
  warn(msg)
/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('8013'), PosixPath('//172.28.0.1'), PosixPath('http')}
  warn(msg)
/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('--logtostderr --listen_host=172.28.0.12 --target_host=172.28.0.12 --tunnel_background_save_url=https'), PosixPath('//colab.research.google.com/tun/m/cc48301118ce562b961b3c22d803539adc1e0c19/gpu-t4-s-i2qu3r70uroy --tunnel_background_save_delay=10s --tunnel_periodic_background_save_frequency=30m0s --enable_output_coalescing=true --output_coalescing_required=true')}
  warn(msg)
/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/env/python')}
  warn(msg)
/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('//ipykernel.pylab.backend_inline'), PosixPath('module')}
  warn(msg)
/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/usr/local/cuda/lib64/libcudart.so.11.0'), PosixPath('/usr/local/cuda/lib64/libcudart.so')}.. We'll flip a coin and try one of these, in order to fail forward.
Either way, this might cause trouble in the future:
If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.
  warn(msg)
Downloading (…)neration_config.json:   0%|          | 0.00/111 [00:00<?, ?B/s]
In [5]:
from peft import prepare_model_for_kbit_training

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
In [6]:
def print_trainable_parameters(model):
    """
    Prints the number of trainable parameters in the model.
    """
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
    )
In [7]:
from peft import LoraConfig, get_peft_model

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

model = get_peft_model(model, config)
print_trainable_parameters(model)
trainable params: 1572864 || all params: 709218304 || trainable%: 0.22177430998735193
In [8]:
from datasets import load_dataset

data = load_dataset("timdettmers/openassistant-guanaco")
data = data.map(lambda samples: tokenizer(samples["text"]), batched=True)
Downloading readme:   0%|          | 0.00/395 [00:00<?, ?B/s]
Downloading and preparing dataset json/timdettmers--openassistant-guanaco to /root/.cache/huggingface/datasets/timdettmers___json/timdettmers--openassistant-guanaco-c93588435bc90172/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4...
Downloading data files:   0%|          | 0/2 [00:00<?, ?it/s]
Downloading data:   0%|          | 0.00/20.9M [00:00<?, ?B/s]
Downloading data:   0%|          | 0.00/1.11M [00:00<?, ?B/s]
Extracting data files:   0%|          | 0/2 [00:00<?, ?it/s]
Generating train split: 0 examples [00:00, ? examples/s]
Generating test split: 0 examples [00:00, ? examples/s]
Dataset json downloaded and prepared to /root/.cache/huggingface/datasets/timdettmers___json/timdettmers--openassistant-guanaco-c93588435bc90172/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4. Subsequent calls will reuse this data.
  0%|          | 0/2 [00:00<?, ?it/s]
Map:   0%|          | 0/9846 [00:00<?, ? examples/s]
Token indices sequence length is longer than the specified maximum sequence length for this model (1897 > 1024). Running this sequence through the model will result in indexing errors
Map:   0%|          | 0/518 [00:00<?, ? examples/s]
In [9]:
!pip install wandb
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting wandb
  Downloading wandb-0.15.3-py3-none-any.whl (2.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.0/2.0 MB 34.7 MB/s eta 0:00:00
Requirement already satisfied: Click!=8.0.0,>=7.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (8.1.3)
Collecting GitPython!=3.1.29,>=1.0.0 (from wandb)
  Downloading GitPython-3.1.31-py3-none-any.whl (184 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 184.3/184.3 kB 22.2 MB/s eta 0:00:00
Requirement already satisfied: requests<3,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (2.27.1)
Requirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (5.9.5)
Collecting sentry-sdk>=1.0.0 (from wandb)
  Downloading sentry_sdk-1.24.0-py2.py3-none-any.whl (206 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 206.5/206.5 kB 10.6 MB/s eta 0:00:00
Collecting docker-pycreds>=0.4.0 (from wandb)
  Downloading docker_pycreds-0.4.0-py2.py3-none-any.whl (9.0 kB)
Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from wandb) (6.0)
Collecting pathtools (from wandb)
  Downloading pathtools-0.1.2.tar.gz (11 kB)
  Preparing metadata (setup.py) ... done
Collecting setproctitle (from wandb)
  Downloading setproctitle-1.3.2-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (30 kB)
Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from wandb) (67.7.2)
Requirement already satisfied: appdirs>=1.4.3 in /usr/local/lib/python3.10/dist-packages (from wandb) (1.4.4)
Requirement already satisfied: protobuf!=4.21.0,<5,>=3.19.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (3.20.3)
Requirement already satisfied: six>=1.4.0 in /usr/local/lib/python3.10/dist-packages (from docker-pycreds>=0.4.0->wandb) (1.16.0)
Collecting gitdb<5,>=4.0.1 (from GitPython!=3.1.29,>=1.0.0->wandb)
  Downloading gitdb-4.0.10-py3-none-any.whl (62 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 62.7/62.7 kB 8.2 MB/s eta 0:00:00
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (1.26.15)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (2022.12.7)
Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (2.0.12)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (3.4)
Collecting smmap<6,>=3.0.1 (from gitdb<5,>=4.0.1->GitPython!=3.1.29,>=1.0.0->wandb)
  Downloading smmap-5.0.0-py3-none-any.whl (24 kB)
Building wheels for collected packages: pathtools
  Building wheel for pathtools (setup.py) ... done
  Created wheel for pathtools: filename=pathtools-0.1.2-py3-none-any.whl size=8791 sha256=b27691c28003b69fe677abef24f386b27cd0ced353c004c65f8ee7af91cc0dd3
  Stored in directory: /root/.cache/pip/wheels/e7/f3/22/152153d6eb222ee7a56ff8617d80ee5207207a8c00a7aab794
Successfully built pathtools
Installing collected packages: pathtools, smmap, setproctitle, sentry-sdk, docker-pycreds, gitdb, GitPython, wandb
Successfully installed GitPython-3.1.31 docker-pycreds-0.4.0 gitdb-4.0.10 pathtools-0.1.2 sentry-sdk-1.24.0 setproctitle-1.3.2 smmap-5.0.0 wandb-0.15.3
In [10]:
import wandb
wandb.login()
wandb: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)
wandb: You can find your API key in your browser here: https://wandb.ai/authorize
wandb: Paste an API key from your profile and hit enter, or press ctrl+c to quit:
 ··········
wandb: Appending key for api.wandb.ai to your netrc file: /root/.netrc
Out[10]:
True
In [11]:
%env WANDB_PROJECT=qlora-falcon-guanaco
env: WANDB_PROJECT=qlora-falcon-guanaco
In [12]:
import transformers

# needed for falcon tokenizer
tokenizer.pad_token = '<|endoftext|>'
tokenizer.bos_token = '>>ABSTRACT<<'
tokenizer.eos_token = '<|endoftext|>'

output_dir = "outputs"

trainer = transformers.Trainer(
    model=model,
    train_dataset=data["train"],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=40,
        max_steps=100,
        learning_rate=2e-4,
        fp16=True,
        save_strategy="steps",
        save_steps=10,
        logging_steps=1,
        output_dir=output_dir,
        optim="paged_adamw_8bit"
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
trainer.train()
model.save_pretrained(output_dir)
Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to the W&B docs.
wandb: Currently logged in as: utensil. Use `wandb login --relogin` to force relogin
Tracking run with wandb version 0.15.3
Run data is saved locally in /content/wandb/run-20230528_174740-fdsvtgrs
Syncing run sweet-snow-1 to Weights & Biases (docs)
View project at https://wandb.ai/utensil/qlora-falcon-guanaco
View run at https://wandb.ai/utensil/qlora-falcon-guanaco/runs/fdsvtgrs
You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
[ 16/100 00:57 < 05:47, 0.24 it/s, Epoch 0.01/1]
Step Training Loss
1 1.997600
2 2.265900
3 1.809600
4 2.020400
5 3.371100
6 1.832700
7 2.383400
8 1.730600
9 2.520400
10 2.620500
11 2.613400
12 2.683500
13 2.354000
14 1.917100

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ in <cell line: 29>:29                                                                            │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:1696 in train                    │
│                                                                                                  │
│   1693 │   │   inner_training_loop = find_executable_batch_size(                                 │
│   1694 │   │   │   self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size  │
│   1695 │   │   )                                                                                 │
│ ❱ 1696 │   │   return inner_training_loop(                                                       │
│   1697 │   │   │   args=args,                                                                    │
│   1698 │   │   │   resume_from_checkpoint=resume_from_checkpoint,                                │
│   1699 │   │   │   trial=trial,                                                                  │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:1973 in _inner_training_loop     │
│                                                                                                  │
│   1970 │   │   │   │   │   with model.no_sync():                                                 │
│   1971 │   │   │   │   │   │   tr_loss_step = self.training_step(model, inputs)                  │
│   1972 │   │   │   │   else:                                                                     │
│ ❱ 1973 │   │   │   │   │   tr_loss_step = self.training_step(model, inputs)                      │
│   1974 │   │   │   │                                                                             │
│   1975 │   │   │   │   if (                                                                      │
│   1976 │   │   │   │   │   args.logging_nan_inf_filter                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:2787 in training_step            │
│                                                                                                  │
│   2784 │   │   │   return loss_mb.reduce_mean().detach().to(self.args.device)                    │
│   2785 │   │                                                                                     │
│   2786 │   │   with self.compute_loss_context_manager():                                         │
│ ❱ 2787 │   │   │   loss = self.compute_loss(model, inputs)                                       │
│   2788 │   │                                                                                     │
│   2789 │   │   if self.args.n_gpu > 1:                                                           │
│   2790 │   │   │   loss = loss.mean()  # mean() to average on multi-gpu parallel training        │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:2819 in compute_loss             │
│                                                                                                  │
│   2816 │   │   │   labels = inputs.pop("labels")                                                 │
│   2817 │   │   else:                                                                             │
│   2818 │   │   │   labels = None                                                                 │
│ ❱ 2819 │   │   outputs = model(**inputs)                                                         │
│   2820 │   │   # Save past state if it exists                                                    │
│   2821 │   │   # TODO: this needs to be fixed and made cleaner later.                            │
│   2822 │   │   if self.args.past_index >= 0:                                                     │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501 in _call_impl            │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/peft/peft_model.py:686 in forward                        │
│                                                                                                  │
│    683 │   ):                                                                                    │
│    684 │   │   peft_config = self.active_peft_config                                             │
│    685 │   │   if not isinstance(peft_config, PromptLearningConfig):                             │
│ ❱  686 │   │   │   return self.base_model(                                                       │
│    687 │   │   │   │   input_ids=input_ids,                                                      │
│    688 │   │   │   │   attention_mask=attention_mask,                                            │
│    689 │   │   │   │   inputs_embeds=inputs_embeds,                                              │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501 in _call_impl            │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/accelerate/hooks.py:165 in new_forward                   │
│                                                                                                  │
│   162 │   │   │   with torch.no_grad():                                                          │
│   163 │   │   │   │   output = old_forward(*args, **kwargs)                                      │
│   164 │   │   else:                                                                              │
│ ❱ 165 │   │   │   output = old_forward(*args, **kwargs)                                          │
│   166 │   │   return module._hf_hook.post_forward(module, output)                                │
│   167 │                                                                                          │
│   168 │   module.forward = new_forward                                                           │
│                                                                                                  │
│ /root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/104655c0c067936f1ae2b4 │
│ 73625fe5161375591a/modelling_RW.py:753 in forward                                                │
│                                                                                                  │
│    750 │   │                                                                                     │
│    751 │   │   return_dict = return_dict if return_dict is not None else self.config.use_return  │
│    752 │   │                                                                                     │
│ ❱  753 │   │   transformer_outputs = self.transformer(                                           │
│    754 │   │   │   input_ids,                                                                    │
│    755 │   │   │   past_key_values=past_key_values,                                              │
│    756 │   │   │   attention_mask=attention_mask,                                                │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501 in _call_impl            │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/accelerate/hooks.py:165 in new_forward                   │
│                                                                                                  │
│   162 │   │   │   with torch.no_grad():                                                          │
│   163 │   │   │   │   output = old_forward(*args, **kwargs)                                      │
│   164 │   │   else:                                                                              │
│ ❱ 165 │   │   │   output = old_forward(*args, **kwargs)                                          │
│   166 │   │   return module._hf_hook.post_forward(module, output)                                │
│   167 │                                                                                          │
│   168 │   module.forward = new_forward                                                           │
│                                                                                                  │
│ /root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/104655c0c067936f1ae2b4 │
│ 73625fe5161375591a/modelling_RW.py:640 in forward                                                │
│                                                                                                  │
│    637 │   │   │   │   │                                                                         │
│    638 │   │   │   │   │   return custom_forward                                                 │
│    639 │   │   │   │                                                                             │
│ ❱  640 │   │   │   │   outputs = torch.utils.checkpoint.checkpoint(                              │
│    641 │   │   │   │   │   create_custom_forward(block),                                         │
│    642 │   │   │   │   │   hidden_states,                                                        │
│    643 │   │   │   │   │   alibi,                                                                │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:249 in checkpoint              │
│                                                                                                  │
│   246 │   │   raise ValueError("Unexpected keyword arguments: " + ",".join(arg for arg in kwar   │
│   247 │                                                                                          │
│   248 │   if use_reentrant:                                                                      │
│ ❱ 249 │   │   return CheckpointFunction.apply(function, preserve, *args)                         │
│   250 │   else:                                                                                  │
│   251 │   │   return _checkpoint_without_reentrant(                                              │
│   252 │   │   │   function,                                                                      │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/autograd/function.py:506 in apply                  │
│                                                                                                  │
│   503 │   │   if not torch._C._are_functorch_transforms_active():                                │
│   504 │   │   │   # See NOTE: [functorch vjp and autograd interaction]                           │
│   505 │   │   │   args = _functorch.utils.unwrap_dead_wrappers(args)                             │
│ ❱ 506 │   │   │   return super().apply(*args, **kwargs)  # type: ignore[misc]                    │
│   507 │   │                                                                                      │
│   508 │   │   if cls.setup_context == _SingleLevelFunction.setup_context:                        │
│   509 │   │   │   raise RuntimeError(                                                            │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:107 in forward                 │
│                                                                                                  │
│   104 │   │   ctx.save_for_backward(*tensor_inputs)                                              │
│   105 │   │                                                                                      │
│   106 │   │   with torch.no_grad():                                                              │
│ ❱ 107 │   │   │   outputs = run_function(*args)                                                  │
│   108 │   │   return outputs                                                                     │
│   109 │                                                                                          │
│   110 │   @staticmethod                                                                          │
│                                                                                                  │
│ /root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/104655c0c067936f1ae2b4 │
│ 73625fe5161375591a/modelling_RW.py:636 in custom_forward                                         │
│                                                                                                  │
│    633 │   │   │   │   def create_custom_forward(module):                                        │
│    634 │   │   │   │   │   def custom_forward(*inputs):                                          │
│    635 │   │   │   │   │   │   # None for past_key_value                                         │
│ ❱  636 │   │   │   │   │   │   return module(*inputs, use_cache=use_cache, output_attentions=ou  │
│    637 │   │   │   │   │                                                                         │
│    638 │   │   │   │   │   return custom_forward                                                 │
│    639                                                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501 in _call_impl            │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/accelerate/hooks.py:165 in new_forward                   │
│                                                                                                  │
│   162 │   │   │   with torch.no_grad():                                                          │
│   163 │   │   │   │   output = old_forward(*args, **kwargs)                                      │
│   164 │   │   else:                                                                              │
│ ❱ 165 │   │   │   output = old_forward(*args, **kwargs)                                          │
│   166 │   │   return module._hf_hook.post_forward(module, output)                                │
│   167 │                                                                                          │
│   168 │   module.forward = new_forward                                                           │
│                                                                                                  │
│ /root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/104655c0c067936f1ae2b4 │
│ 73625fe5161375591a/modelling_RW.py:385 in forward                                                │
│                                                                                                  │
│    382 │   │   residual = hidden_states                                                          │
│    383 │   │                                                                                     │
│    384 │   │   # Self attention.                                                                 │
│ ❱  385 │   │   attn_outputs = self.self_attention(                                               │
│    386 │   │   │   layernorm_output,                                                             │
│    387 │   │   │   layer_past=layer_past,                                                        │
│    388 │   │   │   attention_mask=attention_mask,                                                │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501 in _call_impl            │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/accelerate/hooks.py:165 in new_forward                   │
│                                                                                                  │
│   162 │   │   │   with torch.no_grad():                                                          │
│   163 │   │   │   │   output = old_forward(*args, **kwargs)                                      │
│   164 │   │   else:                                                                              │
│ ❱ 165 │   │   │   output = old_forward(*args, **kwargs)                                          │
│   166 │   │   return module._hf_hook.post_forward(module, output)                                │
│   167 │                                                                                          │
│   168 │   module.forward = new_forward                                                           │
│                                                                                                  │
│ /root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/104655c0c067936f1ae2b4 │
│ 73625fe5161375591a/modelling_RW.py:306 in forward                                                │
│                                                                                                  │
│    303 │   │   │   │   attention_scores = attention_scores.to(torch.float32)                     │
│    304 │   │   │   # attn_weights = torch.masked_fill(attention_scores, attention_mask, torch.f  │
│    305 │   │   │   attention_probs = F.softmax(                                                  │
│ ❱  306 │   │   │   │   (attention_scores + alibi.view(batch_size, self.num_heads, 1, -1)) * sel  │
│    307 │   │   │   │   dim=-1,                                                                   │
│    308 │   │   │   │   dtype=hidden_states.dtype,                                                │
│    309 │   │   │   )                                                                             │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
OutOfMemoryError: CUDA out of memory. Tried to allocate 2.74 GiB (GPU 0; 14.75 GiB total capacity; 8.16 GiB already
allocated; 2.63 GiB free; 11.06 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try 
setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and 
PYTORCH_CUDA_ALLOC_CONF
In [ ]:
trainer = transformers.Trainer(
    model=model,
    train_dataset=data["train"],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=40,
        max_steps=1000,
        learning_rate=2e-4,
        fp16=True,
        save_strategy="steps",
        save_steps=100,
        logging_steps=1,
        output_dir=output_dir,
        optim="paged_adamw_8bit"
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
trainer.train(resume_from_checkpoint=True)
model.save_pretrained(output_dir)
In [ ]: