{"id":5647630,"date":"2023-03-28T13:47:39","date_gmt":"2023-03-28T17:47:39","guid":{"rendered":"https:\/\/lightning.ai\/pages\/?p=5647630"},"modified":"2023-04-13T07:01:16","modified_gmt":"2023-04-13T11:01:16","slug":"gradient-accumulation","status":"publish","type":"post","link":"https:\/\/lightning.ai\/pages\/blog\/gradient-accumulation\/","title":{"rendered":"Finetuning LLMs on a Single GPU Using Gradient Accumulation"},"content":{"rendered":"

Key takeaway<\/h3> Learn how to use gradient accumulation to train models with large batch sizes in order to work around hardware limitations when GPU memory is a concern. <\/div>\n

Previously, I shared an article using multi-GPU training strategies to speed up the finetuning of large language models<\/a>. Several of these strategies include mechanisms such as model or tensor sharding that distributes the model weights and computations across different devices to work around GPU memory limitations.<\/p>\n

However, many of us don’t have access to multi-GPU resources. This article therefore demonstrates a great workaround to train models with larger batch sizes when GPU memory is a concern: gradient accumulation.<\/p>\n

 <\/p>\n

Let’s Finetune BLOOM for Classification<\/strong><\/h2>\n

Let’s suppose we are interested in adopting a recent pretrained large language model for a downstream task such as text classification. We are going to work with BLOOM<\/a>, which is an open-source alternative to GPT-3. In particular, we are going to use a version of BLOOM that “only” has 560 million parameters — it should fit into the RAM of conventional GPUs without problems (for reference, the free tier of Google Colab has a GPU with 15 Gb of RAM.)<\/p>\n

Once we start, however, we bump into problems: our memory explodes during training or finetuning; we find that the only way to train this model is using a batch size of 1.<\/p>\n

\"\"<\/p>\n

 <\/p>\n

Note<\/h3> The code for finetuning BLOOM for a target classification task using a batch size of 1 is shown below. (You can also download the complete code from GitHub here<\/a>.)<\/p>\n

You can copy & paste this code directly into Google Colab. However, you also have to drag and drop the accompanying local_dataset_utilities.py<\/a> file into the same folder as we import some dataset utilities from this file.<\/p>\n

<\/div>\n

\n\n# pip install torch lightning matplotlib pandas torchmetrics watermark transformers datasets -U\n\nimport os
\nimport os.path as op
\nimport time\n\nfrom datasets import load_dataset
\nfrom lightning import Fabric
\nimport torch
\nfrom torch.utils.data import DataLoader
\nimport torchmetrics
\nfrom transformers import AutoTokenizer
\nfrom transformers import AutoModelForSequenceClassification
\nfrom watermark import watermark\n\nfrom local_dataset_utilities import download_dataset, load_dataset_into_to_dataframe, partition_dataset
\nfrom local_dataset_utilities import IMDBDataset\n\ndef tokenize_text(batch):
\n    return tokenizer(batch[\"text\"], truncation=True, padding=True, max_length=1024)\n\ndef train(num_epochs, model, optimizer, train_loader, val_loader, fabric):\n\n    for epoch in range(num_epochs):
\n        train_acc = torchmetrics.Accuracy(
\n            task=\"multiclass\", num_classes=2).to(fabric.device)\n\n        for batch_idx, batch in enumerate(train_loader):
\n            model.train()\n\n            ### FORWARD AND BACK PROP
\n            outputs = model(
\n                batch[\"input_ids\"],
\n                attention_mask=batch[\"attention_mask\"],
\n                labels=batch[\"label\"]
\n            ) \n\n            fabric.backward(outputs[\"loss\"])\n\n            ### UPDATE MODEL PARAMETERS
\n            optimizer.step()
\n            optimizer.zero_grad()\n\n            ### LOGGING
\n            if not batch_idx % 300:
\n                print(f\"Epoch: {epoch+1:04d}\/{num_epochs:04d} \"
\n                      f\"| Batch {batch_idx:04d}\/{len(train_loader):04d} \"
\n                      f\"| Loss: {outputs['loss']:.4f}\")\n\n            model.eval()
\n            with torch.no_grad():
\n                predicted_labels = torch.argmax(outputs[\"logits\"], 1)
\n                train_acc.update(predicted_labels, batch[\"label\"])\n\n        ### MORE LOGGING
\n        model.eval()
\n        with torch.no_grad():
\n            val_acc = torchmetrics.Accuracy(task=\"multiclass\", num_classes=2).to(fabric.device)
\n            for batch in val_loader:
\n                outputs = model(
\n                    batch[\"input_ids\"],
\n                    attention_mask=batch[\"attention_mask\"],
\n                    labels=batch[\"label\"]
\n                )
\n                predicted_labels = torch.argmax(outputs[\"logits\"], 1)
\n                val_acc.update(predicted_labels, batch[\"label\"])\n\n            print(f\"Epoch: {epoch+1:04d}\/{num_epochs:04d} \"
\n                  f\"| Train acc.: {train_acc.compute()*100:.2f}% \"
\n                  f\"| Val acc.: {val_acc.compute()*100:.2f}%\"
\n                  )
\n            train_acc.reset(), val_acc.reset()\n\nif __name__ == \"__main__\":\n\n    print(watermark(packages=\"torch,lightning,transformers\", python=True))
\n    print(\"Torch CUDA available?\", torch.cuda.is_available())
\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    torch.manual_seed(123)
\n    # torch.use_deterministic_algorithms(True)\n\n    ##########################
\n    ### 1 Loading the Dataset
\n    ##########################
\n    download_dataset()
\n    df = load_dataset_into_to_dataframe()
\n    if not (op.exists(\"train.csv\") and op.exists(\"val.csv\") and op.exists(\"test.csv\")):
\n        partition_dataset(df)\n\n    imdb_dataset = load_dataset(
\n        \"csv\",
\n        data_files={
\n            \"train\": \"train.csv\",
\n            \"validation\": \"val.csv\",
\n            \"test\": \"test.csv\",
\n        },
\n    )\n\n    #########################################
\n    ### 2 Tokenization and Numericalization
\n    #########################################\n\n    tokenizer = AutoTokenizer.from_pretrained(\"bigscience\/bloom-560m\", max_length=1024)
\n    print(\"Tokenizer input max length:\", tokenizer.model_max_length, flush=True)
\n    print(\"Tokenizer vocabulary size:\", tokenizer.vocab_size, flush=True)\n\n    print(\"Tokenizing ...\", flush=True)
\n    imdb_tokenized = imdb_dataset.map(tokenize_text, batched=True, batch_size=None)
\n    del imdb_dataset
\n    imdb_tokenized.set_format(\"torch\", columns=[\"input_ids\", \"attention_mask\", \"label\"])
\n    os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\n    #########################################
\n    ### 3 Set Up DataLoaders
\n    #########################################\n\n    train_dataset = IMDBDataset(imdb_tokenized, partition_key=\"train\")
\n    val_dataset = IMDBDataset(imdb_tokenized, partition_key=\"validation\")
\n    test_dataset = IMDBDataset(imdb_tokenized, partition_key=\"test\")\n\n    train_loader = DataLoader(
\n        dataset=train_dataset,
\n        batch_size=1,
\n        shuffle=True,
\n        num_workers=4,
\n        drop_last=True,
\n    )\n\n    val_loader = DataLoader(
\n        dataset=val_dataset,
\n        batch_size=1,
\n        num_workers=4,
\n        drop_last=True,
\n    )\n\n    test_loader = DataLoader(
\n        dataset=test_dataset,
\n        batch_size=1,
\n        num_workers=2,
\n        drop_last=True,
\n    )\n\n    #########################################
\n    ### 4 Initializing the Model
\n    #########################################\n\n    fabric = Fabric(accelerator=\"cuda\", devices=1, precision=\"16-mixed\")
\n    fabric.launch()\n\n    model = AutoModelForSequenceClassification.from_pretrained(
\n        \"bigscience\/bloom-560m\", num_labels=2)\n\n    optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)\n\n    model, optimizer = fabric.setup(model, optimizer)
\n    train_loader, val_loader, test_loader = fabric.setup_dataloaders(
\n        train_loader, val_loader, test_loader)\n\n    #########################################
\n    ### 5 Finetuning
\n    #########################################\n\n    start = time.time()
\n    train(
\n        num_epochs=1,
\n        model=model,
\n        optimizer=optimizer,
\n        train_loader=train_loader,
\n        val_loader=val_loader,
\n        fabric=fabric,
\n    )\n\n    end = time.time()
\n    elapsed = end-start
\n    print(f\"Time elapsed {elapsed\/60:.2f} min\")\n\n    with torch.no_grad():
\n        model.eval()
\n        test_acc = torchmetrics.Accuracy(task=\"multiclass\", num_classes=2).to(fabric.device)
\n        for batch in test_loader:
\n            outputs = model(
\n                batch[\"input_ids\"],
\n                attention_mask=batch[\"attention_mask\"],
\n                labels=batch[\"label\"]
\n            )
\n            predicted_labels = torch.argmax(outputs[\"logits\"], 1)
\n            test_acc.update(predicted_labels, batch[\"label\"])\n\n    print(f\"Test accuracy {test_acc.compute()*100:.2f}%\")
\n<\/code>