10-605 / 10-405

Machine Learning from Large Datasets

Recap: the course so far

• Tool: Scalable computation on clusters (of CPU + disk)
• MapReduce, Spark and iteration, Spark workflows
• Tool: Learning as optimization
• MLE for binomials, linear regression, logistic regression, k-means clustering
• Tools: exact solutions, gradient descent (batch, parallel, stochastic), EM and EM-like methods (for k-means)
• “Craft” of ML
• Tool: feature extraction / augmentation/ kernels
• Tool: regularization and sparsity hyperparameter tuning
• Tool: the hash trick and efficient SGD
• Tool: distributed ML – interaction with sparsity, AllReduce
• Tool: Randomized algorithms
• Bloom filter (a randomized set)
• Count-Min sketch (a randomized mapping x 🡪R+ )
• Locality sensitive hashing (LSH)

2

This week

• Tool: Scalable computation on GPUs/clusters of GPUs
• MapReduce, Spark and iteration, Spark workflows
• Tool: Learning as optimization
• MLE for binomials, linear regression, logistic regression, k-means clustering optimizing non-convex / deep models
• Tools: exact solutions, gradient descent (mini-batch, parallel, stochastic), EM and EM-like methods (for k-means)
• “Craft” of ML
• Tool: distributed ML – interaction with sparsity, AllReduce
• Tool: foundation models and pre-training
• Tool: Randomized algorithms
• Bloom filter (a randomized set)
• Count-Min sketch (a randomized mapping x 🡪R+ )
• Locality sensitive hashing (LSH)

3

Today

1. Scaling up Mb-SGD: Communication schemes
2. AllReduce for GPU clusters
3. Challenges with mini-batched optimization
4. Fine-tuning
5. More on Modern optimizers

Follows Fall 2024 lectures from Prof. Smith closely

“Recap”: Autodiff for Logistic Regression

5

Generalized logistic

where:

Multiclass (K classes)

x.dot(W) : gives vector of K scalars

w is a D x K matrix

Compare predictions softmax(xW) to y

Extension: Deeper Logistic Regression

6

h1 = softmax(dot(x, w1))

h2 = softmax(dot(h1, w2))

…

h5 = softmax(dot(h4, w5))

r = l2NormSq(w1) + …. + l2NormSq(w5)

loss = crossEntropy(h5, y) + r

h1 = reLU(dot(x, w1) + b1)

h2 = reLU(dot(h1, w2) + b2)

…

h5 = reLU(dot(h4, w5) + b5)

r = l2NormSq(w1) + …. + l2NormSq(w5)

loss = crossEntropy(h5, y) + r

Not all data is linear …

REC ALL

Deeper models can learn non-linear hypotheses….

Non-convex loss

CH AL L E NGE

• Saddle point: a stationary point that is not a local minimum or maximum

Types of stationary points:

• Global minimum: actual minimizer, i.e., ^{f(x̂) ≤ f(x) ∀ x ∈ ℝk}

• Local minimum: _f(x̂_{) ≤ f(x) ∀ x ∈ ℝ}^k _:
• Local maximum: _f(x̂_{) ≥ f(x) ∀ x ∈ ℝ}^k _:

∥x − x^̂∥ ≤ ϵ

∥x − x^̂∥ ≤ ϵ

Non-convexity

Types of stationary points:

• Global minimum: finding these is hard (both in theory [NP-hard] and in practice)
• Local minimum: can often still provide reasonable solutions
• Saddle point: want to avoid these

CH AL L E NGE

Non-convexity

What makes non-convex optimization difficult?

• Potentially many local minima
• Saddle points
• Very flat regions
• Widely varying curvature

CH AL L E NGE

Most modern ML methods are parallel variants of mini-batch SGD

11

Streaming SGD:

​

• Iterative
• Sequential
• Fast
• Scale up by bounding memory

​

• You can handle very large datasets … but slowly

​

​

x1

x2

x3

w

x4

x5

x6

12

Streaming SGD:

​

• Iterative
• Sequential
• Fast
• Scale up by bounding memory

​

• You can handle very large datasets … but slowly

​

• You can speed it up by making the tasks in the stream bigger and doing them in parallel
• Minibatch SGD with a GPU is a good way of doing that

X1

X2

X3

w

X4

X5

X6

Minibatch examples X, y

Limitations are

• communication to GPU
• GPU memory

Methods for distributed optimization

less local computation, more communication

more local computation, less communication

stochastic gradient descent

gradient descent

divide-and-conquer / one-shot averaging

mini-batch SGD

local-updating methods

REC ALL

Mini-batch SGD

REC ALL

Gradient Descent Update:

w_i+1 = w_i − α_i ∇f(w_i)

Objective we want to solve:

min f(w)

w

n

where f(w) := _∑ f_j(w)

j=1

Stochastic Gradient Descent (SGD) Update:

w_i+1 = w_i − α_i ∇f_j(w_i)

with j sampled at random

Mini-batch SGD Update:

i

i+1 i i ℬ i

w = w − α ∇f (w )

with mini-batch B_i ⊆{1,… n} sampled at random

How to parallelize mini-batch SGD?

Parallelization scheme can vary based on:

• parallelization technique (what to parallelize?)
• synchronization model (when do we send updates?)
• communication strategy (where do we send them?)
• compression (what updates do we send?)

QU ES TION

*Note: While we motivate these approaches with mini-batch SGD, they are applicable to the other first-order optimization methods (AdaGrad, RMSProp, Adam, etc.)

Data parallel vs. model parallel

Data parallel

• Replicate the model on each machine
• Each machine accesses a portion of the dataset (shard)
• Most common approach: typically preferable in terms of implementation, fault tolerance, and parallelizability

Model parallel

• Replicate the data on each machine
• Each machine accesses a portion of the model
• An attractive approach for massive models

​

Also - hybrid schemes!

• Subcategories:
• Intra-layer parallel
• Tensor parallel
• Inter-layer parallel
• Pipeline parallel

​

Model parallel

intra-layer parallelism: units within one layer are partitioned

inter-layer parallelism: entire layers are partitioned

Tensor parallelism is a kind of intra-layer parallelism

• Tensor parallelism (intra-layer model parallelism): splits tensor operations (e.g., matrix matmul) across nodes

Pipeline parallelism schedules inter-layer parallelism

• Pipeline parallelism (inter-layer parallelism)

https://docs.graphcore.ai/projects/tf-model-parallelism/en/latest/pipelining.html

Pipeline parallelism schedules inter-layer parallelism

• Pipeline parallelism (inter-layer parallelism)

https://docs.graphcore.ai/projects/tf-model-parallelism/en/latest/pipelining.html

Example of combined/hybrid parallelism: FSDP

• Fully sharded data parallelism: data parallelism, but also shard the model parameters / optimization states / gradients to handle memory issues

Which strategy to use? (In Pytorch)

https://pytorch.org/tutorials/beginner/dist_overview.html

How to parallelize mini-batch SGD?

Parallelization scheme can vary based on:

• parallelization technique (what to parallelize?)
• synchronization model (when do we send updates?)
• communication strategy (where do we send them?)
• compression (what updates do we send?)

QU ES TION

Recap: Synchronous vs. asynchronous

asynchronous

• fast—no waiting!
• but: there may be errors

​

Image: Dimitris Papailiopoulos

synchronous

• guaranteed to be correct
• but: can be slow due to stragglers, overhead

Recap: Asynchronous ML

Image: Dimitris Papailiopoulos

if no conflicts (e.g., when the shared parameter, x, is very sparse), asynchrony is great!

​

recall: HogWild!

Recommended when:

• sparse updates that don’t overlap
• workers have low-bandwidth connections between them
• Not usually case with high-end GPU clusters

For GPUs: Synchronous vs. asynchronous

asynchronous

• fast—no waiting!
• but: there may be errors

​

Image: Dimitris Papailiopoulos

synchronous

• guaranteed to be correct
• but: can be slow due to stragglers, overhead

How to parallelize mini-batch SGD?

Parallelization scheme can vary based on:

• parallelization technique (what to parallelize?)
• synchronization model (when do we send updates?)
• communication strategy (where do we send them?)
• compression (what updates do we send?)

QU ES TION

Recap: Tree AllReduce

Common pattern in ML:

• Reduce (e.g., aggregate gradients) information from many workers
• Distribute aggregation to all workers

​

​

1.

…

Tree vs Ring AllReduce

Common pattern in ML:

• Reduce (e.g., aggregate gradients) information from many workers
• Distribute aggregation to all workers

​

​

High-even clusters (e.g. TPU Pods) have high-bandwidth connections between some (not all pairs) of accelerators

​

Common topology: a ring / mesh / torus in 2D or 3D.

​

Ring all-reduce has different goals

Baidu, Horovod (Uber), 2017

• Use the bandwidth as effectively as possible:
• Don’t minimize number of messages
• Send smaller messages in parallel
• Use the high-bandwidth connections

​

Ring AllReduce

https://towardsdatascience.com/visual-intuition-on-ring-allreduce-for-distributed-deep-learning-d1f34b4911da

Ring AllReduce

Ring AllReduce

Everyone sends small messages to the next worker in the ring

Everyone picks a different subset of the parameters to send

Ring AllReduce

Workers reduce their local content with the incoming messages

a₀ + b₀

c₁ + b₁

c₂ + d₂

a₃ + d₃

Ring AllReduce

Everyone sends another slice to the next worker in the ring

a₀ + b₀

c₁ + b₁

c₂ + d₂

a₃ + d₃

…sending the most reduced slice they can

Ring AllReduce

Everyone sends another slice to the next worker in the ring

a₀ + b₀

c₁ + b₁

c₂ + d₂

a₃ + d₃

…sending the most reduced slice they can

Iterate until all slices are full reduced (all data has gone the way around the ring)

Ring AllReduce

Then continue passing the fully-reduced data until everybody has seen it

Ring AllReduce

Q: how many rounds of communication?

Q: how much is communicated per round per node?

2 (m-1)

k / m

m = ring cycle

k = size of param vector

Outline

1. Scaling up Mb-SGD: Communication schemes
2. AllReduce for GPU clusters
3. Challenges with mini-batched optimization
4. Fine-tuning
5. More on Modern optimizers

One major issue with mini-batch SGD

Systems perspective: large batches are good

• more computation, less communication Statistical perspective: large batches are bad
• why??

Two problems with large batches:

• diminishing returns (after a certain point, additional data on computing a single gradient doesn’t help much)
• poor generalization (for non-convex optimization)

Diminishing returns

large batches perform worse when they lack diversity, e.g., when data is replicated by some factor (R)

Yin et al., AISTATS 2018

a potential explanation: gradient diversity, i.e.,

the degree to which gradients differ from one another

Poor generalization

EVEN WORSE

large batch methods converge to lower test accuracy than small batches

Keskar, et al, ICLR 2017

Recap: Stochastic gradient descent

sample a “minibatch” of examples B<<N

approximate the gradient using them

“on average” this is the right direction

take a step in that direction

repeat….

43

For non-convex problems SGD/GD find a local minimum. SGD tends to find “flat”, not “sharp” local minima.

​

Variants of SGD are almost always used in deep learning.

Outline

1. Scaling up Mb-SGD: Communication schemes
2. AllReduce for GPU clusters
3. Challenges with mini-batched optimization
4. Fine-tuning
5. More on Modern optimizers

Non-convexity

RECALL

What makes non-convex optimization difficult?

• Potentially many local minima
• Saddle points
• Very flat regions
• Widely varying curvature

Techniques for large-scale optimization

Reduce computation:

•

•

•

Parallelize/distribute computation Use stochastic methods

Use first-order methods (i.e., no 2nd-order or higher gradients)

Reduce communication:

•

•

Keep large objects local Reduce iterations

REC ALL

Reducing iterations

We are interested in optimization techniques that reduce the total number of iterations until convergence

​

However, these techniques often come at a price: more computation

​

Will discuss two popular techniques for DL that aim to reduce iterations with little additional computation:

• Adaptive learning rate methods
• Momentum-based methods

Gradient descent

We can move anywhere in _R^d

Negative gradient is direction of

steepest descent!

2D Example:

•

•

Function values are in black/white and black represents higher values

Arrows are gradients

"Gradient2" by Sarang. Licensed under CC BY-SA 2.5 via Wikimedia Commons http://commons.wikimedia.org/wiki/File:Gradient2.svg#/media/File:Gradient2.svg

Update Rule: _{wt+1 = wt − α ∇f(wt)}

​

Negative Slope

Step Size

REC ALL

Motivation: Newton’s method

issue: gradient may be the same despite varying curvature

Newton’s method uses second order information to estimate the curvature of the function being minimized (i.e., the loss)

Motivation: Newton’s method

issue: gradient may be the same despite varying curvature

Newton’s method uses second order information to estimate the curvature of the function being minimized (i.e., the loss)

Motivation: Newton’s method

issue: gradient may be the same despite varying curvature

Newton step makes more progress

​

​

Newton’s method uses second order information to estimate the curvature of the function being minimized (i.e., the ERM objective)

Motivation: Newton’s method

Newton’s method

−1

w_t+1 = w_t − α ₍ ∇²f(w_t)₎ ∇f(w_t)

can we do better? idea: use second order information (i.e., the Hessian)

can significantly reduce number of iterations needed for convergence however, comes at a cost: need to invert k×k matrix → O(k³)

often too expensive for neural networks (k can be in the millions)

Gradient descent

_REC ALL

w_t+1 = w_t − α ∇f(w_t)

Motivation: Newton’s method

But – for functions with n inputs the “second derivative” is too large.

Hessian matrix:

Motivation: Newton’s method

Hessian is worse than the Jaccobian (from autodiff)

AdaGrad

A CHEAPER APPRO A CH

AdaGrad:

else

^wt+1

t

= w − αG

−1

t

t

∇f(w )

₍G_t)ii = ∑

t

j=1

∇f(w )

2

j i

, 0

where

• A cheap approximation of the Hessian: only estimate the diagonal elements
• Enforces a separate step size / learning rate for every coordinate of the weight vector
• Decreases the step size for coordinates with high gradient magnitudes
• Why does this make sense when the model is large?
• There may be high variability across the coordinates
• Potential issue: step size can become very small in some directions

Duchi et al, 2011

RMSProp

RMSProp:

w = w − αG

−1

t+1 t t

t

∇f(w )

₍G_t)ii = (g_t)_i , 0

where

else

i

₍g_t) = β(g_t−1)_i + (1 − β)( ∇f(w_t)_i)²

AN A LT ERN ATI VE T O AD A GRAD

• Take a moving average of the coordinates of the gradient, instead of a sum
• Puts more of an emphasis on the current gradients than the gradient history
• Can be more effective for DNNs
• Also see: AdaDelta

Hinton [Coursera], 2012

Momentum

Polyak momentum

w_t+1 = w_t − α ∇f(w_t) + β(w_t − w_t−1)

AN O T HER APPRO A CH

• Another cheap way to potentially improve gradient descent: use momentum
• Intuition: if current gradient step is in same direction as previous step, then move a little further in that direction; if not, move less far
• Also known as the heavy ball method

Polyak ‘64

Nesterov accelerated gradient

Nesterov momentum

w_t+1 = w_t − α ∇f(w_t + β(w_t − w_t−1)) + β(w_t − w_t−1)

• A “look-ahead” variant of momentum
• Provably faster convergence than GD for any convex function

Nesterov ’83

evaluate gradient at ‘look-ahead’ point

who is Adam:

ON E L A S T M E THOD: C OM B IN E EV ER Y THIN G!

Adam:

w = w − αG

−1

v

t+1 t t t

₍G_t)ii =

(g_t)_i , 0

where

else

i

₍g_t) = β (g

1 t−1 i 1 t i

) + (1 − β )( ∇f(w ) )²

v_t = β₂v_t−1 + (1 − β₂) ∇f(w_t)

• Combines momentum and RMSprop
• Default choice in most DL frameworks
• However, does not converge (even for convex problems) without modifications
• [Reddi, Kale, Kumar, ICLR ’19 best paper]

RMSProp update

Momentum

Early stopping

• Recall: what we actually care about is generalization performance (i.e., how will the trained model perform on new, unseen data?)
• Many techniques to regularize DNNs and prevent overfitting
• Early stopping is a simple (and cheap!) way to improve generalization and prevent overfitting

​

• Idea: stop training when validation accuracy starts to reduce

ON E L A S T TRIC K …

Outline

1. Scaling up Mb-SGD: Communication schemes
2. AllReduce for GPU clusters
3. Challenges with mini-batched optimization
4. Fine-tuning
5. More on Modern optimizers

Training in the Foundation Model Era: Fine-Tuning

How are foundation models changing deep learning?

https://dataforest.ai/blog/ai-foundation-models-for-big-business-innovation

Classical ML

model

data

training

ML in the foundation model era

pretrained

lots of data!

pretraining

aligned

alignment

finetuned

flnetuning

How do we ‘finetune’ models?

Start with a model trained on a lot of data

Using the pretrained model as a starting point, update the weights by fitting the model to a small amount of finetuning data (e.g., running a small number of steps of gradient descent)

Credit: AssemblyAI

How do we ‘finetune’ models?

Start with a model trained on a lot of data

Using the pretrained model as a starting point, update the weights by fitting the model to a small amount of finetuning data (e.g., running a small number of steps of gradient descent)

Credit: Prasad Mahamulkar

Gradient descent for model training

Start at a random point, w₀

Repeat

Determine a descent direction Choose a step size

Update model using gradient calculated from training data

Until stopping criterion is satisfied

Update Rule:

w_i+1 = w_i — a_i∇f(w_i)

REC ALL

f(w)

w₁ w₀

_w* … _w2

Gradient descent for model finetuning

with a pretrained model, w_pre

Start at a random point, w₀

Repeat

Determine a descent direction Choose a step size

Update model using gradient

calculated from training finetuning data Until stopping criterion is satisfied

w

f(w)

w₁ w₀

_w* … _w2

Update Rule:

w_i+1 = w_i — a_i∇f(w_i)

Fine-tuning example: BERT (2019)

Training in the Foundation Model Era: Efficient Fine-Tuning

Potential issue: pre-trained models are getting large!

340M parameters

175B parameters

???

https://www.dailydoseofds.com/implementing-lora-from-scratch-for-fine-tuning-llms/

How do we ‘finetune’ large models?

When finetuning, we expect the change in weights (orange) to be small

​

However, updating all of these weights can be expensive!

​

How can we reduce the number of weights we need to update?

… enter parameter-efficient finetuning

Types of PEFT methods

• specification: pick a small subset of the parameters and tune those
• addition: introduce additional trainable parameters, and just train those
• reparameterization: re-parameterize the network into something more efficient to train

https://www.andrew.cmu.edu/course/11-667/lectures/W4L2_PETM.pptx.pdf

Linear probing

Credit: Ke et al, 2024, https://arxiv.org/pdf/2402.18905v1

Just update a portion of the weights, e.g., just the last layer (linear probing)

aka Foundation models as features

Adapters

Adapters are new modules are added between layers of a pre-trained network. Original model weights are fixed — just the adapter modules are tuned

Recap (again):�Recovering latent factors in a matrix

80

K * m

n * K

~

LoRA (low rank adaptation) finetuning

Key idea: don’t learn a full d x d matrix of weights, instead learn a low-rank approximation to this matrix!

full finetuning

LoRA

finetuning

LoRA (low rank adaptation) finetuning

LoRa scaling: default value of alpha/r is 1. However, can be tuned to trade-off emphasis on pretrained weights vs. updated weights

QLoRA: Quantized LoRA

Key idea: quantize pretrained weights so they are very small / memory efficient Hope is that LoRA finetuning can help to offset issues with lower precision

https://github.com/bitsandbytes-foundation/bitsandbytes