TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings — part 5/6

Cleaned-text code points [49744, 62098). Bibliography omitted here; source.md includes it.

TPU
v4:
An
Optically
Reconfigurable
Supercomputer
for
ML
with
Hardware
Support
for
Embeddings
Industrial
Product
ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
(
CFE
)
in
2021
for
the
US
was
40%,
but
it
was
88%
for
Google’s
Oklahoma
data
centers
[19].
They
exploit
the
plentiful
local
renewable
energy
[56].
Fortunately,
Oklahoma
hosts
all
TPU
v4s
in
Google
Cloud.
The
global
average
conversion
factor
from
electricity
to
CO2
equivalent
emissions
—
CO2e
,
including
greenhouse
gasses
like
methane—is
0.475
kg/kWh
[24].
After
acquiring
renewable
energy
in
Oklahoma,
matched
on
an
hourly
basis
with
our
energy
consumption,
it
dropped
to
0.074.
The
estimated
operational
CO2e
for
a
contemporary
DSA
in
an
average
on-premise
data
center
is
then
2.85
×
0.475
÷
0.074
or
~18.3x
higher
than
training
on
TPU
v4
in
Google
Cloud.
The
CFE
for
data
centers
depends
on
the
local
grid
plus
the
availability
and
acquisition
of
renewable
energy.
Given
such
a
large
impact
—crucial
for
anyone
building
or
using
ML
infrastructure—it’s
fortunate
that
Google
has
WSCs
with
very
high
CFE
to
house
TPU
v4
supercomputers.
7.7
How
fast
do
ML
workloads
change?
Table
1
above
shows
the
rapid
change
for
production
workloads
at
Google.
Note
the
drop
in
RNNs.
Like
RNNs,
Transformers
are
popular
for
natural
language
translation
and
text
summarization,
but
unlike
RNNs
they
process
the
input
all
at
once
rather
than
sequentially.
This
revised
model
architecture
means
the
operations
can
occur
in
parallel,
which
in
turn
means
Transformer
models
can
process
much
larger
data
sets.
Two
years
after
the
Transformer
paper
was
published
[58],
it
was
>20%
of
the
TPU
v3
workload.
The
Transformer
models
BERT
and
GPT-3
have
also
changed
the
workload
landscape.
Two
years
after
the
BERT
paper
was
published
[14],
it
was
>25%
of
Google’s
workload
on
TPU
v4,
and
it
remained
significant
in
2022.
Two
years
after
publication
of
GPT-3
[6],
LLMs
were
>30%
of
the
TPU
v4
production
workloads.
ML
workloads
can
change
dramatically
in
the
two
or
more
years
it
takes
to
design,
build,
and
deploy
a
new
ML
supercomputer.
Figure
17:
Change
in
size
of
DLRM0
over
time
measured
in
weights
and
embeddings.
Each
point
is
a
new
version
of
DLRM0
(43
total).
Each
weight
is
1
byte
and
each
embedding
is
4
bytes.
7.8
Do
individual
DNN
models
also
change?
Figure
17
shows
the
change
in
weights
and
embeddings
for
DLRM0
from
2017
to
2022.
Weights
grew
4.2x
and
embeddings
grew
3.8x.
Over
those
five
years
a
new
version
was
released
every
~6
weeks
(43
total).
DLRM0
ran
on
all
five
TPU
products
over
those
five
years.
The
velocity
of
change
of
the
models
and
of
the
architecture
highlights
the
importance
of
having
a
compiler
that
efficiently
leverages
the
features
of
the
underlying
DSA.
7.9
Is
MLPerf’s
DLRM
benchmark
(Figure
14
above)
realistic?
Production
DLRM
workloads
scale
much
better
[64]
than
MLPerf
DLRM
for
a
few
reasons.
MLPerf
DLRM
has
<2M
FP32
weights
while
DLRM0
has
137M
Int8
weights
(see
Figure
17).
Second,
the
global
batch
size
of
MLPerf
DLRM
is
capped
at
64k
for
optimal
model
quality
for
this
data
set
and
optimizer,
limiting
batch
size
to
128
per
SC
on
a
128-chip
system
(128
chips
×
4
SCs/chip
×
128
=
64k).
In
contrast,
internal
recommendation
models
often
reach
batch
sizes
of
2048–4096
and
usefully
scale
up
to
1024
chips
(see
Figure
11
above).
Third,
MLPerf
DLRM
has
only
26
univalent
features
(compared
to
100s
in
internal
workloads)
and
no
multivalent
features.
For
these
reasons,
fixed
overheads
per
batch
such
as
HBM
latency
and
CISC
instruction
generation
time
on
the
SC
core
sequencer
are
much
higher
on
MLPerf
DLRM
than
production
workloads.
Overheads
like
these
limit
its
useful
scalability
to
≤128
chips
for
TPU
v4
and
A100.
7.10
TPU
v4
has
less
HBM
capacity
than
A100;
could
that
limit
LLM
performance?
Our
autoML
LLM
configuration
search
(Section
4)
considers
HBM
capacity,
ICI-connected
aggregated
FLOPS/HBM
bandwidth,
and
other
model
parameters
(such
as
batch
size).
The
HBM
capacity
could
be
a
limiting
factor
in
some
cases,
but
typically
TPUv4
enables
larger
models
to
be
partitioned
across
more
chips
with
effective
compute-communication
overlap
[59]
with
little
overhead
(for
example,
two
case
studies
in
Table
3
above).
However,
given
the
higher
HBM
capacity
but
smaller
NVlink-connected
domain
of
Nvidia
GPUs,
Nvidia’s
best
LLM
configuration
might
be
very
different
from
the
best
one
for
TPU
v4.
7.11
How
can
DSAs
avoid
overspecialization?
Building
a
robust
DSA
with
a
multi-generational
lifetime
is
a
balancing
act
between
domain
specialization
and
generality
with
flexibility.
This
harmony
is
especially
important
given
ML
is
such
a
fast
evolving
field
(see
Sections
7.7
and
7.8).
For
example,
had
we
devoted
significant
resources
for
specialized
acceleration
of
RNNs,
that
effort
would
have
been
of
little
use
when
RNN
popularity
plummeted.
As
a
positive
example,
by
providing
a
flexible,
general,
and
balanced
design,
TPU
v4
proved
to
be
an
excellent
match
to
LLMs
that
were
popularized
[Bro22]
after
it
was
deployed.
8
RELATED
WORK
TPU
v4
uses
a
dedicated
3D
torus
interconnect.
Traditional
supercomputers
also
employ
tightly
connected
multiprocessors
over
3D
tori
with
a
high-bandwidth
interconnect
[15,
50].
Nvidia
GPU-based
systems
today
use
a
two-tier
network
hierarchy
with
NVLink
and
NVSwitch
among
groups
of
4
to
256
GPUs
and
Infiniband
beyond
that.
Twisted
tori
are
not
a
recent
invention.
The
ILLIAC-IV
twisted
one
dimension
of
its
wrap-around
links
for
the
2D
torus
that
it
used
[5].
Sequin
introduced
doubly-twisted
tori
for
mapping
binary
trees
onto
processor
arrays
[47].
Camara,
Moreto,
et
al.
11

ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
Jouppi,
et
al.
made
the
case
for
twisting
3D
tori
[7],
and
TPU
v4
follows
the
k×2k×2k
configuration
from
Camarero,
Martinez,
and
Beivide
[8].
Shalf
et
al.
and
Kamil
et
al.
proposed
a
mix
of
circuit
switching
and
packet
switching
for
use
in
traditional
supercomputers
[49,28]
and
Kamil
et
al.
later
suggested
that
a
MEMS-based
OCS
could
provide
circuit
switching
[29].
A
recent
paper
has
a
similar
investigation
plus
topology
and
parallelization
co-optimizations
to
accelerate
ML
[60].
For
data
centers,
there
have
been
many
proposals
for
OCS-based
networks
[17,
31,
35,
51,
53,
57].
Some
even
include
the
concept
of
topology
engineering.
However,
all
of
this
related
work
are
paper
designs
or
proof-of-concept,
testbed
scale
demonstrations,
in
contrast
to
the
widespread
deployment
of
OCSes
at
Google
for
data
center
networks
and
for
supercomputer
interconnect.
We
believe
that
TPU
v4
is
the
first
commercial
supercomputer
built
using
OCS
and
the
first
supercomputer
built
with
a
reconfigurable
interconnect
that
enhances
performance.
The
TPU
v4
supercomputer
implements
a
logically
shared
address
space
across
physical
chips.
Software
explicitly
controls
access
and
data
movement;
remote
memories
are
available
through
asynchronous
DMA
writes
only.
The
Cray
T3E
enabled
a
similar
logically
shared
address
space,
with
bulk
asynchronous
reads
and
writes,
load-store
access
to
remote
memories,
and
a
rich
suite
of
atomic
operations
[46].
The
TPU
v4
memory
system
is
tailored
for
high
performance,
with
each
chip
maintaining
tens
of
thousands
of
outstanding
memory
requests
(to
both
local
and
remote
memories).
Acceleration
of
embeddings
is
key
for
DLRMs
used
for
business-critical
applications,
and
we
believe
Google
was
the
first
to
include
on-ASIC
hardware
support
in
TPU
v2,
deployed
in
2017.
Neo
from
Facebook
(Meta)
[37]
trains
embedding
tables
with
up
to
12T
parameters
using
a
128-GPU
system.
Neo
also
exploits
table,
row,
column
and
data
parallelism;
overlaps
communication
and
compute;
and
improves
kernel
fusion.
Nvidia’s
MLPerf
entries
use
similar
techniques,
e.g.,
custom
fused
kernels
for
reduction
over
Infiniband.
Two
recent
papers
present
other
optimizations
for
embeddings
[18,
48].
9
SUMMARY
Two
major
architectural
features
of
TPU
v4
have
small
cost
but
outsized
advantages.
The
SparseCore
accelerates
embeddings
of
DLRM
models
by
5x-7x
by
providing
a
dataflow
sea-of-cores
architecture
that
allows
embeddings
to
be
placed
anywhere
in
the
128
TiB
physical
memory
of
the
TPU
v4
supercomputer.
This
gain
comes
at
the
cost
of
only
~5%
in
die
area
and
power.
The
OCSes
and
underlying
optical
components
at
the
heart
of
the
TPU
v4
supercomputer
are
relatively
inexpensive
at
<5%
of
overall
costs
and
<3%
of
overall
power
consumption,
yet
it
provides
a
remarkable
set
of
eight
benefits:
1.
Scalability.
2.
Improved
availability,
which
enables
the
TPU
v4
supercomputer
to
be
4x
larger
than
TPU
v3.
3.
Modularity,
allowing
the
faster
3D
torus
topology
from
64
to
3072
chips
and
novel
shapes
like
twisted
tori.
4.
Higher
performance,
as
users
can
pick
the
topology
that
is
best
for
their
application.
5.
Diminished
power,
as
MEMS
optical
circuit
switching
is
more
energy
efficient
than
electronic
packet
switching.
6.
Simplified
scheduling
to
improve
utilization.
7.
Faster
deployment,
for
better
return
on
investment.
8.
Enhanced
security,
which
encourages
different
organizations
to
share
use
of
TPU
v4
supercomputers.
Moreover,
replacing
OCS
and
ICI
with
Infiniband
increases
costs,
raises
power
consumption,
and
degrades
performance.
TPU
v4
is
faster
and
lower
power
than
contemporary
DSA
chips
made
using
similar
technologies
deployed
close
to
the
same
time
and
for
similar
sized
systems.
The
power
edge
might
be
even
larger
if
the
interconnects
are
included.
Training
time
of
LLMs
is
greatly
reduced
over
TPU
v3
by
using
3K
TPU
v4
slices
with
their
3D
torus
topology.
The
performance,
scalability,
and
availability
make
TPU
v4
supercomputers
the
workhorses
of
large
language
models
(LLMs)
like
LaMDA,
MUM,
and
PaLM
[54,
38,
9].
These
features
allowed
the
540B
parameter
PaLM
model
to
sustain
a
remarkable
57.8%
of
the
peak
hardware
floating
point
performance
over
50
days
while
training
on
TPU
v4
supercomputers
[9].
Google
has
deployed
dozens
of
TPU
v4
supercomputers,
including
eight
for
external
use
via
Google
Cloud.
Moreover,
the
large
size
of
the
TPU
v4
supercomputer
and
its
reliance
on
OCSes
looks
prescient
given
that
the
design
began
two
years
before
the
paper
was
published
that
has
stoked
the
enthusiasm
for
LLMs
[6].
Advances
by
computer
architects
in
the
state-of-the-art
of
warehouse
scale
computing
(WSCs)
save
energy
and
thus
help
reduce
the
carbon
footprint
of
ML.
When
energy-efficient
TPU
v4
supercomputers
are
housed
inside
energy-efficient
WSCs
that
rely
on
~90%
carbon
free
electricity,
they
can
consume
only
~
⅙
–½
of
the
energy
and
produce
only
~5%
of
the
operational
CO2e
from
training
on
a
typical
contemporary
ML
DSA
in
the
average
on-premise
data
center.
A
~20x
reduction
in
carbon
footprint
greatly
increases
the
chances
of
delivering
on
the
amazing
potential
of
ML
in
a
sustainable
manner
[42].
ACKNOWLEDGEMENTS
Special
thanks
go
to
the
Google
Platforms
Optics
team
that
kept
advancing
the
state-of-the-art
of
OCSes
and
optical
transceivers
after
other
organizations
gave
up.
We
also
thank
John
Hennessy,
Robert
Hundt,
Christos
Kozyrakis,
Sridhar
Lakshmanamurthy,
Jae
W.
Lee,
Hong
Liu,
Aamer
Mahmood,
Partha
Ranganathan,
Adrian
Sampson,
Amin
Vahdat,
Amir
Yazdanbakhsh,
George
Yuan,
and
the
anonymous
ISCA
reviewers
for
their
feedback
and
suggestions
on
this
paper.
We
finally
wish
to
thank
our
collaborators
who
helped
with
the
measurements
for
this
paper—Ben
Albrecht,
Jinliang
Wei,
Shibo
Wang,
Yuechao
Pan,
and
Ziqiang
Feng
for
Table
3
and
  Lluis-Miquel
Munguia,
Hao
Wu,
and
Yuechao
Pan
for
Table
6—and
the
many
engineers
who
developed
the
TPU
v4
chip,
hardware,
software,
and
the
many
teams
contributing
to
its
deployment,
including
but
not
limited
to:
Zach
Cotton,
Pedram
Dashti,
Bill
Edwards,
Hema
Hariharan,
Chetan
Kale,
Georgios
Konstadinidis,
Alan
Kulawik,
Justin
Lee,
Hong
Liu,
Erji
Mao,
Omkar
Pathak,
Erick
Tuttle,
Daoyi
Wang,
Kevin
Yasumura,
and
Sara
Zebian.
REFERENCES
[1]
Anil,
R.,
Gadanho,
S.,
Huang,
D.,
Jacob,
N.,
Li,
Z.,
Lin,
D.,
Phillips,
T.,
Pop,
C.,
Regan,
K.,
Shamir,
G.I.
and
Shivanna,
R.,
2022.
On
the
Factory
Floor:
ML
Engineering
for
Industrial-Scale
Ads
Recommendation
Models
.
16th
ACM
Conference
on
Recommender
Systems.
[2]
Barroso,
L.A.,
and
Hölzle,
U.,
2009,
First
Edition.
The
datacenter
as
a
computer:
An
introduction
to
the
design
of
warehouse-scale
machines.
Synthesis
lectures
on
computer
architecture
,
6(3),
pp.1-120.
12