{
  "source_file": "0930",
  "source_platform": "个人博客",
  "author": "Philipp Schmid",
  "title_listed": "DeepSpeed + LoRA + Flash Attention",
  "keywords": "DeepSpeed,LoRA,Flash Attention",
  "domain": "大模型算法",
  "content_type": "教程",
  "is_paid_flag": "否",
  "url": "https://www.philschmid.de/deepspeed-lora-flash-attention",
  "publish_date": "2023-09",
  "corpus_stage": "预训练/SFT",
  "notes": "DeepSpeed+LoRA+Flash Attention组合训练",
  "id": "133476e52106",
  "collected_at": "2026-10-02T14:29:12",
  "title": "Fine-tune Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention",
  "body_markdown": "Fine-tune Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention\nSeptember 20, 2023\n9\nminute read\nView Code\nFalcon 180B is the newest version of Falcon LLM family. It is the biggest open source model with 180B parameter and trained on more data - 3.5T tokens with context length window upto 4K tokens. In this example we will show how to fine-tune Falcon 180B using DeepSpeed, Hugging Face Transformers, LoRA with Flash Attention on a multi-GPU machine.\nIn detail you will learn how to:\nSetup Development Environment\nLoad and prepare the dataset\nFine-Tune Falcon 180B using DeepSpeed, Hugging Face Transformers, LoRA with Flash Attention\nBefore we get into the code lets take a quick look on the technologies and methods we are going to use:\nWhat is DeepSpeed ZeRO?\nDeepSpeed ZeRO focuses on efficient large-scale training of Transformers. ZeRO, or Zero Redundancy Optimizer, reduces memory footprint by partitioning model states across devices instead of basic data parallelism. This saves significant memory - ZeRO-Infinity can reduce usage 100x vs data parallelism. ZeRO-Offload further reduces memory by offloading parts of model and optimizer to CPU, enabling 10B+ parameter models on 1 GPU. ZeRO\nintegrates with HuggingFace Transformers through a configuration file\n.\nWhat is LoRA?\nLoRA\nenables efficient fine-tuning of large language models. It decomposes weight matrices into smaller, trainable update matrices that adapt while keeping original weights frozen. This drastically reduces trainable parameters for faster, lower-memory tuning. LoRA integrates into\nTransformers via Hugging Face's PEFT\n. It combines well with methods like DeepSpeed. Key advantages are efficient tuning, portable models, and no inference latency when merging trained weights. LoRA allows adaptively training massive models with limited resources.\nWhat is Flash Attention?\nFlash Attention is an algorithm that speeds up the core attention mechanism in Transformer language models by restructuring computations. It uses techniques like tiling and recomputation to reduce the high memory costs of attention, enabling models to process longer text sequences. Flash Attention 2 optimizes parallelism and work partitioning for 2x speedup over the previous version, reaching 230 TFLOPS/s on A100 GPUs.\nAccess Falcon 180B\nBefore we can start training we have to make sure that we accepted the license\ntiiuae/falcon-180B\nto be able to use it. You can accept the license by clicking on the Agree and access repository button on the model page at:\ntiiuae/falcon-180B\nThe example was created and run a DGX A100 8-GPU machine with 80GB GPU memory per GPU.\n1. Setup Development Environment\nconda create --name hf python=3.10 -c conda-forge\nPython\n# install torch with the correct cuda version, check nvcc --version\n!\npip install torch\n--\nextra-index-url https://download.pytorch.org/whl/cu118\n--\nupgrade\n# install Hugging Face Libraries and additional dependencies\n!\npip install\n\"transformers==4.33.1\"\n\"datasets==2.14.5\"\n\"accelerate==0.22.0\"\n\"evaluate==0.4.0\"\n\"peft==0.5.0\"\ntensorboard packaging\n--\nupgrade\n# install deepspeed and ninja for jit compilations of kernels\n!\npip install\n\"deepspeed==0.10.3\"\nninja\n--\nupgrade\n# install additional Flash Attention\n!\npip install flash-attn\n--\nno-build-isolation\n--\nupgrade\nTo access any Falcon 180B asset we need to login into our hugging face account. We can do this by running the following command:\nPython\n!\nhuggingface-cli login\n--\ntoken\nYOUR_TOKEN\n2. Load and prepare the dataset\nwe will use the\ndolly\nan open source dataset of instruction-following records generated by thousands of Databricks employees in several of the behavioral categories outlined in the\nInstructGPT paper\n, including brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization.\nPython\n{\n\"instruction\"\n:\n\"What is world of warcraft\"\n,\n\"context\"\n:\n\"\"\n,\n\"response\"\n:\n\"World of warcraft is a massive online multi player role playing game. It was released in 2004 by bizarre entertainment\"\n}\nTo load the\nsamsum\ndataset, we use the\nload_dataset()\nmethod from the 🤗 Datasets library.\nPython\nfrom\ndatasets\nimport\nload_dataset\nfrom\nrandom\nimport\nrandrange\n# Load dataset from the hub\ndataset =\nload_dataset\n(\n\"databricks/databricks-dolly-15k\"\n,\nsplit\n=\n\"train\"\n)\nprint(\nf\n\"dataset size:\n{\nlen(dataset)\n}\n\"\n)\nprint(dataset[\nrandrange\n(len(dataset))])\n# dataset size: 15011\nTo instruct tune our model we need to convert our structured examples into a collection of tasks described via instructions. We define a\nformatting_function\nthat takes a sample and returns a string with our format instruction.\nPython\ndef\nformat_dolly\n(\nsample\n):\ninstruction =\nf\n\"### Instruction\n\\n\n{\nsample[\n'instruction'\n]\n}\n\"\ncontext =\nf\n\"### Context\n\\n\n{\nsample[\n'context'\n]\n}\n\"\nif\nlen(sample[\n\"context\"\n]) >\n0\nelse\nNone\nresponse =\nf\n\"### Answer\n\\n\n{\nsample[\n'response'\n]\n}\n\"\n# join all the parts together\nprompt =\n\"\n\\n\\n\n\"\n.\njoin\n([i\nfor\ni\nin\n[instruction, context, response]\nif\ni\nis\nnot\nNone\n])\nreturn\nprompt\nlets test our formatting function on a random example.\nPython\nfrom\nrandom\nimport\nrandrange\nprint(\nformat_dolly\n(dataset[\nrandrange\n(len(dataset))]))\nIn addition, to formatting our samples we also want to pack multiple samples to one sequence to have a more efficient training.\nPython\nfrom\ntransformers\nimport\nAutoTokenizer\nmodel_id =\n\"tiiuae/falcon-180B\"\ntokenizer = AutoTokenizer.\nfrom_pretrained\n(model_id)\ntokenizer.pad_token = tokenizer.eos_token\nWe define some helper functions to pack our samples into sequences of a given length and then tokenize them.\nPython\nfrom\nrandom\nimport\nrandint\nfrom\nitertools\nimport\nchain\nfrom\nfunctools\nimport\npartial\n# template dataset to add prompt to each sample\ndef\ntemplate_dataset\n(\nsample\n):\nsample[\n\"text\"\n] =\nf\n\"\n{\nformat_dolly\n(sample)\n}{\ntokenizer.eos_token\n}\n\"\nreturn\nsample\n# apply prompt template per sample\ndataset = dataset.\nmap\n(template_dataset,\nremove_columns\n=list(dataset.features))\n# print random sample\nprint(dataset[\nrandint\n(\n0\n, len(dataset))][\n\"text\"\n])\n# empty list to save remainder from batches to use in next batch\nremainder = {\n\"input_ids\"\n: [],\n\"attention_mask\"\n: [],\n\"token_type_ids\"\n: []}\ndef\nchunk\n(\nsample\n,\nchunk_length\n=\n2048\n):\n# define global remainder variable to save remainder from batches to use in next batch\nglobal\nremainder\n# Concatenate all texts and add remainder from previous batch\nconcatenated_examples = {k: list(\nchain\n(*sample[k]))\nfor\nk\nin\nsample.\nkeys\n()}\nconcatenated_examples = {k: remainder[k] + concatenated_examples[k]\nfor\nk\nin\nconcatenated_examples.\nkeys\n()}\n# get total number of tokens for batch\nbatch_total_length = len(concatenated_examples[list(sample.\nkeys\n())[\n0\n]])\n# get max number of chunks for batch\nif\nbatch_total_length >= chunk_length:\nbatch_chunk_length = (batch_total_length // chunk_length) * chunk_length\n# Split by chunks of max_len.\nresult = {\nk: [t[i : i + chunk_length]\nfor\ni\nin\nrange(\n0\n, batch_chunk_length, chunk_length)]\nfor\nk, t\nin\nconcatenated_examples.\nitems\n()\n}\n# add remainder to global variable for next batch\nremainder = {k: concatenated_examples[k][batch_chunk_length:]\nfor\nk\nin\nconcatenated_examples.\nkeys\n()}\n# prepare labels\nresult[\n\"labels\"\n] = result[\n\"input_ids\"\n].\ncopy\n()\nreturn\nresult\n# tokenize and chunk dataset\nlm_dataset = dataset.\nmap\n(\nlambda\nsample\n:\ntokenizer\n(sample[\n\"text\"\n]),\nbatched\n=\nTrue\n,\nremove_columns\n=list(dataset.features)\n).\nmap\n(\npartial\n(chunk,\nchunk_length\n=\n2048\n),\nbatched\n=\nTrue\n,\n)\n# Print total number of samples\nprint(\nf\n\"Total number of samples:\n{\nlen(lm_dataset)\n}\n\"\n)\nAfter we processed the datasets we want to save it to disk to be able to use the processed dataset later during training.\nPython\nlm_dataset.\nsave_to_disk\n(\n\"dolly-processed\"\n)\n3. Fine-Tune Falcon 180B using DeepSpeed, Hugging Face Transformers, LoRA with Flash Attention\nDeepSpeed ZeRO is natively integrated into the\nHugging Face Transformers Trainer\n. The integration enables leveraging ZeRO by simply providing a DeepSpeed config file, and the Trainer takes care of the rest. We created 2 deepspeed configurations for the experiments we ran, including\nCPU offloading\n:\nds_falcon_180b_z3.json\nds_falcon_180b_z3_offload.json\nAs mentioned in the beginning, we ran those example using a 8x NVIDIA A100 80GB. This means we can leverage\nbf16\n, which reduces the memory footprint of the model by almost ~2x, which allows us to train without offloading efficiently. We are going to use the\nds_falcon_180b_z3.json\n. If you are irritated by the\nauto\nvalues, check the\ndocumentation\n.\nIn addition to the deepspeed configuration we also need a training script, which implements LoRA and patches our model to use flash-attention. We created a\nrun_ds_lora.py\nscript, which patches the falcon model using the\nfalcon_patch.py\nutils and implements LoRA using\npeft_utils.py\n.\nWhen you run make sure that you have the same folder structure and utils/configs available. The easiest way is to clone the whole repository. Go into the\ntraining\ndirectory and start the training.\nOnce we made sure that we have the right configuration and training script we can start the training using\ntorchrun\n.\nPython\n!\ntorchrun\n--\nnproc_per_node\n8\nrun_ds_lora.py \\\n--\nmodel_id tiiuae/falcon-\n180B\n\\\n--\ndataset_path dolly-processed \\\n--\noutput_dir falcon-\n180b\n-lora-fa \\\n--\nnum_train_epochs\n3\n\\\n--\nper_device_train_batch_size\n1\n\\\n--\nlearning_rate\n4e-3\n\\\n--\ngradient_checkpointing\nTrue\n\\\n--\ngradient_accumulation_steps\n8\n\\\n--\nbf16\nTrue\n\\\n--\ntf32\nTrue\n\\\n--\nuse_flash_attn\nTrue\n\\\n--\nlr_scheduler_type\n\"constant_with_warmup\"\n\\\n--\nlogging_steps\n25\n\\\n--\nsave_steps\n100\n\\\n--\nsave_total_limit\n3\n\\\n--\ndeepspeed configs/ds_falcon_180b_z3.json\nNote: Since we are using LoRA we are only saving the \"trained\" adapter weights, to save some storage. If you want to merge the adapters back into the base model and save the merged model you can add\n--merge_adapters True\nor use the\nmerge_adapter_weights.py\nscript.\nIn our example for Falcon 180B, the training time was\n153 minutes\nor ~2 hours for 3 epochs. For comparison the pretraining cost of Falcon 180B was ~7,000,000 GPU hours, which is 3,500,000 time more than fine-tuning.\nConclusion\nIn the blog post you learn how to fine-tune Falcon 180B model using DeepSpeed, Hugging Face Transformers, and LoRA with Flash Attention on a multi-GPU machine. We used:\nDeepSpeed ZeRO for memory optimization, enabling training models with up to trillions of parameters on limited GPU memory. We used stage 3 (ZeRO-Infinity) to optimize memory usage.\nHugging Face Transformers and Datasets for easily loading and preparing the text dataset as well as providing an intuitive Trainer API.\nLoRA, a method to efficiently fine-tune large language models by only updating a small percentage of parameters each iteration. This drastically reduces memory usage and computational costs.\nFlash Attention - a highly optimized attention implementation that further reduces the memory footprint.\nCompining all of those methods allows us to fine-tune LLMs with over 100B+ parameter with limited resources. The example provides a template for efficiently tuning the largest publicly available models.\nThanks for reading! If you have any questions, feel free to contact me on\nTwitter\nor\nLinkedIn\n.",
  "metadata": {
    "title": "Fine-tune Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention",
    "description": "In this example we will show how to fine-tune Falcon 180B using DeepSpeed, Hugging Face Transformers, LoRA with Flash Attention on a multi-GPU machine.",
    "keywords": null,
    "author": "Philipp Schmid",
    "og:title": "Fine-tune Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention",
    "og:description": "In this example we will show how to fine-tune Falcon 180B using DeepSpeed, Hugging Face Transformers, LoRA with Flash Attention on a multi-GPU machine.",
    "og:url": "https://www.philschmid.de/deepspeed-lora-flash-attention",
    "og:image": "https://www.philschmid.de/static/blog/deepspeed-lora-flash-attention/thumbnail.jpg",
    "og:image:width": "1200",
    "og:image:height": "630",
    "og:image:alt": "Fine-tune Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention",
    "og:type": "article",
    "twitter:card": "summary_large_image",
    "twitter:title": "Fine-tune Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention",
    "twitter:description": "In this example we will show how to fine-tune Falcon 180B using DeepSpeed, Hugging Face Transformers, LoRA with Flash Attention on a multi-GPU machine.",
    "twitter:image": "https://www.philschmid.de/static/blog/deepspeed-lora-flash-attention/thumbnail.jpg",
    "article:published_time": "2023-09-20"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/133476e52106/code/body-000.txt",
      "chars": 557,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-001.txt",
      "chars": 41,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-002.txt",
      "chars": 201,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-003.txt",
      "chars": 268,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-004.txt",
      "chars": 371,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-005.txt",
      "chars": 84,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-006.txt",
      "chars": 166,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-007.txt",
      "chars": 2031,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-008.txt",
      "chars": 42,
      "language": "text"
    },
    {
      "file": "attachments/133476e52106/code/body-009.txt",
      "chars": 526,
      "language": "text"
    }
  ],
  "images": [],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://github.com/philschmid/deep-learning-pytorch-huggingface/blob/main/training/deepseed-falcon-180b-lora-fa.ipynb",
      "label": "View Code",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/docs/transformers/v4.26.1/en/main_classes/deepspeed",
      "label": "integrates with HuggingFace Transformers through a configuration file",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2106.09685",
      "label": "LoRA",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/docs/peft/conceptual_guides/lora",
      "label": "Transformers via Hugging Face's PEFT",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/tiiuae/falcon-180B",
      "label": "tiiuae/falcon-180B",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/datasets/databricks/databricks-dolly-15k",
      "label": "dolly",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2203.02155",
      "label": "InstructGPT paper",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/docs/transformers/v4.33.1/en/main_classes/deepspeed",
      "label": "Hugging Face Transformers Trainer",
      "type": "a"
    },
    {
      "href": "https://github.com/philschmid/deep-learning-pytorch-huggingface/tree/main/training/configs/ds_falcon_180b_z3.json",
      "label": "ds_falcon_180b_z3.json",
      "type": "a"
    },
    {
      "href": "https://github.com/philschmid/deep-learning-pytorch-huggingface/tree/main/training/configs/ds_falcon_180b_z3_offload.json",
      "label": "ds_falcon_180b_z3_offload.json",
      "type": "a"
    },
    {
      "href": "https://github.com/philschmid/deep-learning-pytorch-huggingface/tree/main/training/run_ds_lora.py",
      "label": "run_ds_lora.py",
      "type": "a"
    },
    {
      "href": "https://github.com/philschmid/deep-learning-pytorch-huggingface/tree/main/training/utils/falcon_patch.py",
      "label": "falcon_patch.py",
      "type": "a"
    },
    {
      "href": "https://github.com/philschmid/deep-learning-pytorch-huggingface/tree/main/training/utils/peft_utils.py",
      "label": "peft_utils.py",
      "type": "a"
    },
    {
      "href": "https://github.com/philschmid/deep-learning-pytorch-huggingface/tree/main/training/scripts/merge_adapter_weights.py",
      "label": "merge_adapter_weights.py",
      "type": "a"
    },
    {
      "href": "https://twitter.com/_philschmid",
      "label": "Twitter",
      "type": "a"
    },
    {
      "href": "https://www.linkedin.com/in/philipp-schmid-a6a2bb196/",
      "label": "LinkedIn",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:52:47Z",
  "source_url": "https://www.philschmid.de/deepspeed-lora-flash-attention",
  "papers": [
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2106.09685",
      "bytes": 1609513,
      "reused": true,
      "local_path": "papers/repair/e7d2fed957472d7677ec.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2203.02155",
      "bytes": 1797405,
      "reused": true,
      "local_path": "papers/repair/4b8cf107df7118e0ffe8.pdf"
    }
  ],
  "actual_video_files": [],
  "body_chars": 11232,
  "expected_body_images": 0,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/133476e52106/pages/source.html",
  "local_html": "attachments/133476e52106/pages/article.html",
  "body_base_path": "attachments/133476e52106/pages/",
  "body_markdown_path": "attachments/133476e52106/pages/article.md"
}