LLM Pre-training & Fine-tuning

September 2023

Overview

Topics we’ll cover

​

• Evolution of LLMs
• LLM Pre-training
• Data collection
• Data cleaning
• Tokenization
• Causal Language Modeling
• LLM Training Considerations
• Memory requirements
• Mixed Precision training
• Distributed training
• Scaling laws
• Fine-tuning
• Types of fine-tuning
• Instruction datasets
• PEFT, LoRA, QLoRA

The Evolution of LLMs

A rapidly evolving ecosystem

https://www.coursera.org/learn/generative-ai-with-llms/

​

A rapidly evolving ecosystem

https://github.com/Mooler0410/LLMsPracticalGuide

Recipe for LLM Development

https://github.com/Mooler0410/LLMsPracticalGuide

LLM Pre-training

Recipe for LLM Development

Adapted from: “State of GPT” talk

LLM Pre-training Workflow

Data Collection

• Before we can train a model, we need to collect data
• Large text corpus
• Web scrapes (high quantity)
• Books (high quality)
• Academic literature
• Code
• Also can be domain specific data (e.g. BloombergGPT)
• Mix of public and private data

Adapted from: “State of GPT” talk

Llama Data Mixture

Data Cleaning

Goal is to increase quality, address bias, remove harmful content.

​

General steps

• Language detection - filter out undesired languages
• Quality Filtering
• Metric based - use perplexity to remove unnatural sentences
• Statistic based - punctuation distribution, symbol-to-word ratio, etc
• Keyword based - explicit content, HTML, hyperlinks, boilerplate,
• Deduplication - crucial to avoid imbalance distribution, improve performance, reduce training steps needed (especially large models)
• Exact deduplication (at document level)
• Near-deduplication with MinHash + LSH
• PII-removal

​

Cleaned corpora

• RefinedWeb (600B tokens)
• CCNet (360B tokens)
• The Pile (340B tokens)

​

Usually 1-10% of original data is actually used!

Appreciating the complexity of LLM data pipelines

https://arxiv.org/pdf/2303.18223.pdf

RefinedWeb

Tokenization

Convert all text into list of integers

​

Typical vocab size

• ~30-100k tokens
• 1 token ~= 0.75 words

​

Typical algorithm

• Byte Pair Encoding
• Tokenizer is learned on the corpus
• Learns rules to merge tokens based on most common character/sub-word pairs

Raw Text

Tokens

Inputs Ids

Pre-training: next token prediction / causal language modeling

Raw Text Dataset (batch_size, context_length)

Row 1: Here is an example document 1 <|endoftext|>

Row 2: Example document 2 <|endoftext|> Example document 3 <|endoftext|> Example�Row 3: Document 4 <|endoftext|> Example Document 5

​

Processed Dataset

Context Length (C)

Batch Size (B)

Pre-training: next token prediction / causal language modeling

Each cell only "sees" cells in its row (on the left of it)

Trying to predict the next cell (on the right of it)

• Yellow = its context
• Green = current position prediction
• Red = its target

​

Processed Dataset

​

Context Length (C)

Batch Size (B)

Pre-training: next token prediction / causal language modeling

Random tokens, just garbage

After 30 000 steps

Random words, wrong grammar

Sentence with correct grammar

Pre-training: Examples

49,152 vocabulary size

8192 context length

15.5B parameter

Trained on 1T tokens

Trained for �86,016 GPU hours

StarCoder

32,000 vocabulary size

4096 context length

7-70B parameter

Trained on 2T tokens Trained for �1,720,320 GPU hours (70B)

Llama2

LLM Training Considerations

Computational challenges

Approximate GPU RAM needed

1 parameter = 4 bytes (32-bit float)

1B parameters = 4x10^9 bytes = 4GB

• Llama2-7B → 28GB
• Llama2-13B → 52GB
• Llama2-70B → 280GB

For reference, common GPU Hardware:

Computational challenges

But that’s just to store model weights…

16 extra bytes per param

https://huggingface.co/docs/transformers/perf_train_gpu_one#anatomy-of-models-memory

Computational challenges

But that’s just to store model weights…

https://www.coursera.org/learn/generative-ai-with-llms

Computational challenges

• FP32 computations are 2x slower than FP16, but…

​

• With fp16 we have a smaller range of possible values, which causes issues with backward pass:
• Weight updates are imprecise (unstable loss)
• Gradients can underflow (very small numbers replaced by 0)
• Gradients can overflow (very large numbers replaced by nan/inf)

So why not just use half-precision?

https://docs.fast.ai/callback.fp16.html

Largest value represented by fp32 → 340,282,000,000,000,000,000,000,000,000,000,000,000

​

Largest value represented by fp16 → 65,504

​

Mixed-precision training

Switch between 32 bit and 16 bit operations during training.

​

Benefits

• Requires less memory which enables training larger models
• Requires less memory bandwidth with speeds up data transfer operations
• Math operations run faster in reduced precision - speeds up training
• No loss in task-specific accuracy

​

It works by identifying which steps require full precision for numerical stability, and using 16-bit everywhere else��This is standard method for pre-training models today

Forward pass and gradient computation are done in half precision == fast compute

Copy gradients back to full precision and update the model == numerical stability

Reducing floating point precision helps reduce model size and speed up computations

https://lightning.ai/pages/community/tutorial/accelerating-large-language-models-with-mixed-precision-techniques/

Distributed Training

When to use

• Use this if your model fits on one GPU
• But you want to speed up training

​

How it works

• Full model is replicated over multiple GPU’s
• Each is fed a slice of the data
• Processing is done in parallel, then synchronized after each train step

Distributed Data Parallel (DDP)

https://www.coursera.org/learn/generative-ai-with-llms

Distributed Training

When to use

• When model doesn’t fit in VRAM on single GPU

​

How it works

• Optimize memory by sharding/distributing model states across GPU’s with zero data overlap
• Different “stages” of sharding depending on what how much of the memory is distributed
• Trade-off in configuration complexity and I/O communication time to sync sharded outputs

Zero Redundancy Optimizer (ZeRO) - Microsoft

https://arxiv.org/pdf/1910.02054.pdf

ZeRO Stage 1�Shard Optimizer States (4x memory savings)

ZeRO Stage 2

Shard Optimizer States and gradients

(8x memory savings)

ZeRO Stage 3

Shard Optimizer States and gradients and model weights

(#GPUx memory savings)

Full model copy

Distributed Training

When to use

• When model doesn’t fit in VRAM on single GPU

​

How it works

• Optimize memory by sharding/distributing model states across GPU’s with zero data overlap
• Different “stages” of sharding depending on what how much of the memory is distributed
• Stage3 aka Fully Sharded Data Parallel (FSDP)

Fully Sharded Data Parallel (FSDP)

https://www.coursera.org/learn/generative-ai-with-llms

Chinchilla Scaling Laws for Compute-optimal Models (2022)

Experiments

• Previously thought (from Open AI paper) that scaling model size was they key to improve model loss
• Trained >400 models from 70M to 16B parameters on 5-500B tokens
• Then, trained Chinchilla-70B with same compute budget as Gopher-280B, but 4x more data

​

Findings

• Chinchilla outperforms Gopher (and GPT3) on a large range of tasks
• So…
• Many existing large models are over-parameterized and under-trained
• Smaller models trained on more data could perform as well as larger models
• Compute optimal training datasize is ~20x the number of model parameters

Deepmind investigated the optimal model size and # tokens for training LLMs under a given compute budget

https://www.coursera.org/learn/generative-ai-with-llms

https://arxiv.org/pdf/2203.15556.pdf

Fine-tuning

Recipe for LLM Development

Adapted from: “State of GPT” talk

Why supervised fine-tuning?

Base model is not trained to be an assistant (e.g. chatty)

Base models are optimized to predict the next word based on the corpus they were trained on… which alone isn’t all that useful

Why supervised fine-tuning?

Fine tuned models are.

Instruction fine-tuning makes it easier to access models knowledge in a familiar, conversational way

Types of fine-tuning

1. Continued pre-training: With domain-specific data, apply the same pre-training regime (next token prediction, masked language modeling) on the base model → Allows for domain adaptation.

​

1. Single-task fine-tuning: The pre-trained model is honed for a narrow and specific task such as toxicity detection or summarization, similar to BERT and T5 → Allows you to use smaller models with fewer labeled examples.
1. Possibility for catastrophic forgetting (aka alignment tax)
2. This may not be a bad thing!

​

1. Multi-task / Instruction fine-tuning: The pre-trained (base) model is fine-tuned on examples of instruction-output pairs to follow instructions, answer questions, and be conversationally be helpful → Allows you to interact with model naturally.

​

1. Reinforcement learning with human feedback (RLHF): This combines instruction fine-tuning with reinforcement learning. It requires collecting human preferences (e.g., pairwise comparisons) which are then used to train a reward model. The reward model is then used to further fine-tune the instructed LLM via RL techniques such as proximal policy optimization (PPO).

Multi-task fine-tuning

• T5 from Google (2019) was the first model that enabled multi-task capabilities from a single model
• Previously, needed a specific fine-tuned model per task
• First example of “prompt engineering” (in retrospect)

https://arxiv.org/pdf/1910.10683.pdf

Instruction fine-tuning

• FLAN-T5 (2022)
• Fine tuning on instruction-context-answer pairs from a variety of tasks
• Improves usability via natural language
• Helps to generalize to unseen tasks

https://arxiv.org/pdf/2210.11416.pdf

Instruction datasets

What makes a good instruction dataset?

• Clear specific instructions
• Diverse set of topics / tasks
• Consistent formatting
• Human feedback

​

Human-written

• Dolly - 15k crowdsourced examples

​

Imitation learning (simple instructions)

• Alpaca - used self-instruct, 52k examples
• Vicuna - 70k examples from ShareGPT
• GPT4All - 800k examples from ChatGPT
• LIMA - 1k examples, high quality

​

Better imitation learning (complex instructions)

• WizardLM - EvolInstruct
• Orca - Explanation traces

https://arxiv.org/pdf/2210.11416.pdf

Human-written

LLM Generated (imitation learning)

Parameter Efficient Fine-tuning (PEFT)

Full model training is inaccessible without significant compute

Parameter Efficient Fine-tuning (PEFT)

What?

Parameter-Efficient Fine-Tuning (PEFT) is a technique that allows us to fine-tune a large pretrained model on a specific downstream task while requiring significantly fewer parameters than full fine-tuning.

Why?

• Recipe: Pre-training on generic data + fine-tuning on specific downstream task
• Large LLMs: full fine-tuning becomes infeasible to train on consumer hardware
• Catastrophic forgetting: tuning all model parameters is prone to overfitting
• Storage: storing and deploying fine-tuned models independently for each downstream task becomes very expensive

How?

PEFT approaches only fine-tune a small number of (extra) model parameters while freezing most parameters of the pretrained LLMs, thereby greatly decreasing the computational and storage costs, being portable, avoiding catastrophic forgetting and better in low data regimes.

Methods

• Prefix tuning
• Prompt tuning
• LoRA
• etc.

​

​

LoRA Overview

​

Pros:

• Fine tuning of LLMs using a fraction of the memory requirements
• Fine-tuning is accessible on common GPU’s
• Tiny checkpoints (example here)
• Performance comparable to full fine-tuning
• No Inference latency addition

​

Cons:

• (only during training) The forward and backward pass is approximately twice as slow, due to the additional matrix multiplications in the adapter layers.

​

Figure explaining how LoRA layers works: extra parameters (in blue) are added on top of frozen layers (in orange)

Adapted from the original paper, figure 1

LoRA Overview

The theory behind it:

​

• The rank of the matrix is the number of linearly independent column vectors.
• Full rank means all columns are linearly independent.
• Low rank means some columns are linear combinations of the other columns.

​

​

• We take inspiration from Li et al. (2018a); Aghajanyan et al. (2020) which show that the learned over-parametrized models in fact reside on a low intrinsic dimension. We hypothesize that the change in weights during model adaptation also has a low “intrinsic rank”

​

LoRA Overview

During training:

LoRA Overview

1. Freeze original self-attention weights (for specified layers)

​

During training:

W: (768x768)

LoRA Overview

1. Freeze original self-attention weights (for specified layers)
2. Inject 2 rank decomposition matrices whose product is same size matrix as original

​

During training:

W: (768x768)

A: (16x768)

B: (768x16)

​

BA: (768x768)

​

r = 16

LoRA Overview

1. Freeze original self-attention weights (for specified layers)
2. Inject 2 rank decomposition matrices whose product is same size matrix as original
3. Train the weights of just these two low-rank matrices
1. Pass embedding matrix through each separately and add together

​

During training:

W: (768x768)

A: (16x768)

B: (768x16)

​

BA: (768x768)

​

r = 16

LoRA Overview

1. Freeze original self-attention weights (for specified layers)
2. Inject 2 rank decomposition matrices whose product is same size matrix as original
3. Train the weights of just these two low-rank matrices
1. Pass embedding matrix through each separately and add together

​

During training:

W: (768x768)

A: (16x768)

B: (768x16)��BA: (768x768)

​

r = 16

Number of Trainable Params:

• W: 590k
• A, B: 12k

​

95% less trainable params

​

LoRA Overview

At inference:

W: (768x768)

A: (16x768)

B: (768x16)��BA: (768x768)

​

r = 16

h = W(x) + BA(x)

​

h = (W + BA)(x)

​

LoRA Overview

At inference:

W: (768x768)

A: (16x768)

B: (768x16)��BA: (768x768)

​

r = 16

h = W(x) + BA(x)

​

h = (W + BA)(x)

​

Easy to swap out LoRA adapters for different tasks!

LoRA Overview

Performance

LoRA Overview

Choosing the rank parameter

LoRA Overview

Best practices

Highlights:

• Comparable evaluation performance than full fine-tuning on a variety of tasks
• Images (diffusion models for different “styles”)
• Audio (fine-tuning whisper large on a new language)
• Text (fine-tuning language models)
• No additional inference latency
• Hyper-parameters:
• Higher learning rate than full fine-tuning (order of magnitude of 10-100x)
• Rank between 4-16
• Lora alpha: usually 2-4x the LoRA rank

More details:

• Empirical findings and best practices with PEFT approaches

​

QLoRA Overview

Frozen model is 4-bit quantized

• 30B model on 24GB GPU
• 65B model on 48GB GPU

Uses Paged Optimizers that allow offload to CPU RAM when spikes happen

• Prevents error and loss of training in middle of run

Only for fine-tuning, not for pre-training

Empirically demonstrated to preserve full fine-tuning task performance

​

Thank you!