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

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

TPU
v4:
An
Optically
Reconfigurable
Supercomputer
for
Machine
Learning
with
Hardware
Support
for
Embeddings
Industrial
Product*
Norman
P.
Jouppi,
George
Kurian,
Sheng
Li,
Peter
Ma,
Rahul
Nagarajan,
Lifeng
Nai,
Nishant
Patil,
Suvinay
Subramanian,
Andy
Swing,
Brian
Towles,
Cliff
Young,
Xiang
Zhou,
Zongwei
Zhou,
and
David
Patterson
Google,
Mountain
View,
CA
{jouppi,gkurian,lsheng,pcma,rahulnagarajan,lnai,nishantpatil,suvinay,aswing,btowles,cliffy,zhoux,zongweiz}@google.com
and
pattrsn@cs.berkeley.edu
ABSTRACT
In
response
to
innovations
in
machine
learning
(ML)
models,
production
workloads
changed
radically
and
rapidly.
TPU
v4
is
the
fifth
Google
domain
specific
architecture
(DSA)
and
its
third
supercomputer
for
such
ML
models.
Optical
circuit
switches
(OCSes)
dynamically
reconfigure
its
interconnect
topology
to
improve
scale,
availability,
utilization,
modularity,
deployment,
security,
power,
and
performance;
users
can
pick
a
twisted
3D
torus
topology
if
desired.
Much
cheaper,
lower
power,
and
faster
than
Infiniband,
OCSes
and
underlying
optical
components
are
<5%
of
system
cost
and
<3%
of
system
power.
Each
TPU
v4
includes
SparseCores,
dataflow
processors
that
accelerate
models
that
rely
on
embeddings
by
5x–7x
yet
use
only
5%
of
die
area
and
power.
Deployed
since
2020,
TPU
v4
outperforms
TPU
v3
by
2.1x
and
improves
performance/Watt
by
2.7x.
The
TPU
v4
supercomputer
is
4x
larger
at
4096
chips
and
thus
nearly
10x
faster
overall,
which
along
with
OCS
flexibility
and
availability
allows
a
large
language
model
to
train
at
an
average
of
~60%
of
peak
FLOPS/second.
For
similar
sized
systems,
it
is
~4.3x–4.5x
faster
than
the
Graphcore
IPU
Bow
and
is
1.2x–1.7x
faster
and
uses
1.3x–1.9x
less
power
than
the
Nvidia
A100.
TPU
v4s
inside
the
energy-optimized
warehouse
scale
computers
of
Google
Cloud
use
~2-6x
less
energy
and
produce
~20x
less
CO2e
than
contemporary
DSAs
in
typical
on-premise
data
centers.
CCS
CONCEPTS
•
Computer
systems
organization
→
Architectures
→
Other
architectures
→
Neural
networks
KEYWORDS
Machine
learning,
domain
specific
architecture,
TPU,
GPU,
IPU,
supercomputer,
optical
interconnect,
reconfigurable,
embeddings,
large
language
model,
power
usage
effectiveness,
warehouse
scale
computer,
carbon
emissions,
energy,
CO2
equivalent
emissions
______________
*This
paper
is
part
of
the
Industry
Track
of
ISCA
2023's
program.
Permission
to
make
digital
or
hard
copies
of
part
or
all
of
this
work
for
personal
or
classroom
use
is
granted
without
fee
provided
that
copies
are
not
made
or
distributed
for
profit
or
commercial
advantage
and
that
copies
bear
this
notice
and
the
full
citation
on
the
first
page.
Copyrights
for
third-party
components
of
this
work
must
be
honored.
For
all
other
uses,
contact
the
Owner/Author.
ISCA
'23,
June
17–21,
2023,
Orlando,
FL,
USA
©
2023
Copyright
is
held
by
the
owner/author(s).
ACM
ISBN
979-8-4007-0095-8/23/06.
https://doi.org/10.1145/3579371.3589350
ACM
Reference
Format:
Norman
P.
Jouppi,
George
Kurian,
Sheng
Li,
Peter
Ma,
Rahul
Nagarajan,
Lifeng
Nai,
Nishant
Patil,
Suvinay
Subramanian,
Andy
Swing,
Brian
Towles,
Cliff
Young,
Xiang
Zhou,
Zongwei
Zhou,
David
Patterson.
2023.
TPU
v4:
An
Optically
Reconfigurable
Supercomputer
for
Machine
Learning
with
Hardware
Support
for
Embeddings:
Industrial
Product.
In
The
50th
Annual
International
Symposium
on
Computer
Architecture
(ISCA
’23),
June
17–21,
2023,
Orlando,
FL,
USA.
ACM,
New
York,
NY,
USA,
14
pages.
https://doi.org/10.1145/3579371.3589350
.
1
INTRODUCTION
Happily
for
architects,
machine
learning
(ML)
models
continue
to
evolve
in
challenging
ways,
both
in
scale
and
algorithmically
(see
Table
1
and
Section
7.7).
Examples
of
the
former
are
large
language
models
(
LLMs
)
and
examples
of
the
latter
are
the
embeddings
necessary
for
recommender
systems
(
deep
learning
recommendation
models
or
DLRMs
)
and
the
huge
calculations
of
Transformers
and
BERT.
The
incredible
scale
of
recent
LLMs
[6,
38,
54]
has
stretched
our
ML
supercomputer
scale
from
256
TPU
v2
nodes
to
4096
TPU
v4
nodes.
Reaching
such
a
scale
raises
reliability
problems
that
are
particularly
compounded
by
the
HPC-style,
checkpoint/restore,
everything-must-work
way
that
deep
neural
network
(DNN)
training
is
performed.
It
is
far
from
the
software
reliability
typical
of
mainline
Google
distributed
systems.
This
paper
describes
three
major
features
of
TPU
v4
that
respond
to
these
challenges:
1.
We
addressed
the
scale
and
reliability
obstacles
by
introducing
Optical
Circuit
Switches
(
OCSes
)
with
optical
data
links,
allowing
a
4K-node
supercomputer
through
reconfiguration
to
tolerate
1K
CPU
hosts
that
are
unavailable
0.1%–1.0%
of
the
time.
2.
We
disclose
the
hardware
support
for
embeddings
in
DLRMs
(
SparseCore
or
SC
),
part
of
TPUs
since
TPU
v2.
3.
Combining
the
first
two
capabilities,
embeddings
add
all-to-all
communication
patterns
to
the
demands
on
supercomputer-scale
interconnect.
Unlike
all-reduce
used
in
backpropagation,
which
maps
well
to
2D
and
3D
tori,
all-to-all
patterns
strain
bisection
bandwidth.
OCSes
enable
flexible
topology
configuration,
including
twisted
torus
[7],
which
has
better
bisection
properties.
LLMs
are
a
hot
topic
today
in
ML
circles.
While
scale
and
reliability
originally
motivated
OCSes
in
TPU
v4,
their
topological
flexibility
and
deployment
advantages
turned
out
to
improve
LLM
training
time
significantly.
Since
prior
papers
have
described
the
fundamentals
of
previous
TPUs
for
training
[26,
39]
and
for
inference
[25,
27],
this
paper
focuses
on
the
three
novel
features
of
TPU
v4
listed
above

ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
Jouppi,
et
al.
that
have
not
yet
been
described.
The
major
contributions
of
the
paper
are:
●
It
describes
and
evaluates
the
first
production
deployment
of
OCSes
in
a
supercomputer
and
the
first
to
allow
topology
reconfiguration
to
improve
performance.
●
It
describes
and
evaluates
the
first
accelerator
support
for
embeddings
in
a
commercial
ML
system.
●
It
documents
the
rapid
change
in
production
model
types
since
2016
for
the
fast
changing
ML
field
(Table
1).
●
It
shows
how
Google
uses
ML
to
co-optimize
DNN
models,
OCS
topology,
and
the
SparseCore.
The
next
section
introduces
OCSes
and
explains
their
many
benefits.
Section
3
motivates
the
SparseCore
and
shows
its
performance
gains.
Section
4
uses
ML
to
search
how
to
co-optimize
the
hardware
and
DNN
models.
The
next
two
sections
compare
performance
on
production
workloads
versus
TPU
v3
and
then
versus
the
Nvidia
A100
and
the
Graphcore
MK2
IPU
using
MLPerf.
The
paper
ends
with
a
discussion,
a
related
work
section,
and
a
summary.
Table
1
:
Workloads
by
DNN
model
type
(%
TPUs
used).
Over
90%
of
training
at
Google
is
on
TPUs.
The
parenthesized
entries
split
Transformer
models
into
the
subtypes
of
BERT
and
LLM.
Columns
2
to
4
show
workloads
for
inference
[25],
training
and
inference
[26],
and
inference
[27].
The
last
workload
is
for
training
on
TPU
v4s
over
30
days
in
October
2022.
DNN
Model
TPU
v1
7/2016
(Inference)
TPU
v3
4/2019
(Training
&
Inference)
TPU
v4
Lite
2/2020
(Inference)
TPU
v4
10/2022
(Training)
MLP/DLRM
61%
27%
25%
24%
RNN
29%
21%
29%
2%
CNN
5%
24%
18%
12%
Transformer
--
21%
28%
57%
(BERT)
--
--
(28%)
(26%)
(LLM)
--
--
--
(31%)
2
RECONFIGURABLE
OPTICAL
SWITCH
We
wanted
to
scale
up
the
number
of
chips
by
4x
versus
TPU
v3
just
as
TPU
v3
was
4x
TPU
v2.
Given
the
distance
between
TPU
v3
racks,
some
wrap-around
links
of
its
2D
torus
topology
were
so
long
that
they
had
to
be
optical
due
to
the
reach
limitation
of
electrical
interconnects.
Optical
links
are
>10x
more
expensive
than
electrical
links.
At
4x
the
scale,
there
would
be
even
more
optical
links.
Moreover,
there
were
concerns
about
the
bisection
bandwidth
of
a
2D
torus
of
that
size
and
the
availability
of
a
single
system
of
that
scale.
Using
a
3D
torus
increases
bisection
bandwidth
and
the
OCS
acts
like
a
plugboard
to
skip
failed
units.
2.1
Optical
Circuit
Switching
To
improve
data
center
networking,
Google
advanced
the
state-of-the-art
in
reliability
and
cost
of
optical
transceivers
and
OCSes
[43,54].
The
resulting
Google
Palomar
OCS
is
based
on
3D
Micro-Electro-Mechanical
Systems
(MEMS)
mirrors
that
switch
in
milliseconds.
They
employ
circulators
to
send
light
both
ways
in
a
fiber,
halving
the
number
of
required
ports
and
cables.
What
size
of
the
electrically-cabled
building
block
to
use?
Given
the
3D
torus,
3D
cubes
have
the
best
bisection
bandwidth,
suggesting
4×4×4
(64
chips)
or
8×8×8
(512).
With
4
TPU
v4s
per
CPU
host,
64
TPU
v4
chips
and
their
16
CPU
hosts
comfortably
fit
into
one
rack.
As
512
chips
need
multiple
racks,
a
4
3
building
block
was
chosen.
2.2
Construction
of
the
TPU
v4
Supercomputer
Figure
1
shows
the
links
from
the
6
“faces”
of
a
4
3
block.
There
are
16
links
per
face,
totaling
96
optical
links
per
block
that
connect
to
OCSes.
To
provide
the
wraparound
links
of
a
3D
torus,
the
links
on
the
opposing
sides
must
connect
to
the
same
OCS.
Thus,
each
4
3
block
connects
to
6
×
16
÷
2
=
48
OCSes.
The
Palomar
OCS
is
136×136
(128
ports
plus
8
spares
for
link
testing
and
repairs),
so
48
OCSes
connect
the
48
pairs
of
cables
from
64
4
3
blocks
(each
64
chips),
yielding
the
desired
total
of
4096
TPU
v4
chips.
Figure
1:
Connectivity
of
a
4×4×4
cube
(top)
to
3
OCSes
(bottom).
The
“+”
and
“–”
connections
with
the
same
dimension
and
index
are
connected
to
the
same
OCS;
48
of
these
in-out
pairs
each
connect
to
a
distinct
OCS.
Figure
2
below
shows
a
TPU
v4
package
and
four
of
them
mounted
on
the
printed
circuit
board.
Like
TPU
v3,
each
TPU
v4
contains
two
TensorCores
(
TC
).
Each
TC
contains
four
128x128
Matrix
Multiply
Units
(
MXUs
)
and
a
Vector
Processing
Unit
(
VPU
)
with
128
lanes
(16
ALUs
per
lane)
and
a
16
MiB
Vector
Memory
(
VMEM
).
The
two
TCs
share
a
128
MiB
Common
Memory
(
CMEM
).
The
PCB
embeds
4
Inter-Core
Interconnect
(
ICI
)
links,
connected
as
a
2×2
mesh;
16
external
ICI
links
go
to
other
trays
for
constructing
the
3D
torus.
Figure
3
below
shows
one
row
of
eight
racks,
where
each
rack
contains
16
tray-host
server
pairs.
Passive
electrical
cables
create
a
4×4×4
3D
mesh
in
a
rack.
Electrical-to
-optical
conversions
happen
at
the
fiber
connector
to
the
TPU
trays.
There
are
no
other
conversions
until
the
light
reaches
an
2

TPU
v4:
An
Optically
Reconfigurable
Supercomputer
for
ML
with
Hardware
Support
for
Embeddings
Industrial
Product
ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
optical-to-electrical
converter
at
the
fiber
connector
of
the
destination
tray.
The
48
OCSes
join
eight
rows
together
to
form
the
complete
64-rack
system.
Figure
2:
The
TPU
v4
package
(ASIC
in
center
plus
4
HBM
stacks)
and
printed
circuit
board
with
4
liquid-cooled
packages.
The
board's
front
panel
has
4
top-side
PCIe
connectors
and
16
bottom-side
OSFP
connectors
for
inter-tray
ICI
links.
Figure
3:
Eight
of
64
racks
for
one
4096-chip
supercomputer
.
Figure
4:
Impact
of
OCS
connected
versus
a
statically
con-
nected
supercomputer
on
goodput
(i.e.,
effective
throughput)
as
CPU
availability
and
slice
size
varies
on
a
log
scale
.
Goodput
is
counterintuitive
at
large
slices.
At
¼
of
the
4K
chips,
goodput
for
both
99.0%
and
99.5%
is
75%,
as
3
slices
occupy
¾
of
the
chips.
Spares
are
needed
to
allow
scheduling
jobs
despite
some
failed
nodes,
so
you
can’t
realistically
schedule
two
2k
node
slices
from
4k
nodes.
With
one
2k
node
slice
(50%
of
4k),
you
have
50%
spares,
so
it
will
have
50%
goodput.
With
3k
nodes
(75%
of
4k),
you
have
25%
spares,
and
therefore
75%
goodput.
2.3
OCS
Availability
Benefits
An
OCS
raises
availability
by
routing
around
failures.
The
main
problem
is
the
CPU
host;
each
host
has
4
TPU
v4s,
so
1K
hosts
per
supercomputer.
Like
HPC
supercomputers,
the
workload
consists
of
a
range
of
scale
sizes,
called
slices
:
64
chips,
128
chips,
256
chips,
and
so
on.
Figure
4
shows
the
“goodput”
of
slice
sizes
as
host
availability
varies
from
99.0%
to
99.9%
with
and
without
OCSes.
Without
OCSes,
host
availability
must
be
99.9%
to
offer
reasonable
slice
goodput.
OCSes
also
have
fair
goodput
for
99.0%
and
99.5%
for
most
slice
sizes.
Figure
4
assumes
all
slice
size
requests
are
equal,
but
workloads
have
many
sizes
(Table
2).
Table
2:
Sampling
of
popularity
of
TPU
v4
slices
for
a
day
in
November
2022.
This
table
includes
all
slices
used
≥
0.1%.
Twistable
(Section
2.8)
but
not
twisted
means
the
slice
geometry
allows
twisting
(n×n×2n
or
n×2n×2n),
but
the
user
picks
the
regular
topology.
The
software
scheduler
requires
that
slices
have
dimensions
x
≤
y
≤
z.
Half
of
the
slices
have
x,
y,
and
z
as
either
4
or
8.
Chips
<64
64
Regular
Tori
1×1×1
(1)
2.1%
4×4×4
(64)
13.9%
1×1×2
(2)
0.4%
1×2×2
(4)
6.7%
2×2×2
(8)
4.7%
2×2×4
(16)
6.4%
2×4×4
(32)
8.9%
Total
%
29%
14%
Chips
128-192
256-384
Twisted
Tori
4×4×8_T
(128)
16.0%
4×8×8_T
(256)
9.2%
Twistable,
not
twisted
Tori
4×4×8_NT
(128)
1.5%
4×8×8_NT
(256)
1.5%
Regular
Tori
4×4×12n
(192)
0.7%
4×4×16
(256)
1.0%
4×8×12
(384)
0.1%
Total
%
18%
12%
Chips
512-768
1024-1536
Twisted
Tori
8×8×16_T
(1K)
1.8%
Twistable,
not
twisted
Tori
8×8×16_NT
(1K)
1.4%
Regular
Tori
4×16×16
(1K)
0.3%
8×8×8
(512)
9.6%
4×4×64
(1K)
0.1%
4×8×16
(512)
1.7%
4×8×32
(1K)
0.1%
4×4×32
(512)
0.6%
8×12×16
(1.5K)
0.1%
8×8×12
(768)
0.7%
4×4×96
(1.5K)
0.1%
8×8×24
(1.5K)
0.1%
Total
%
13%
4%
Chips
2048-3072
Twisted
Tori
8×16×16_T
(2K)
1.4%
Twistable,
not
twisted
Tori
8×16×16_NT
(2K)
0.3%
Regular
Tori
12×16×16
(3K)
5.7%
4×4×192
(3K)
0.4%
Total
%
8%
2.4
OCS
Deployment
Benefits
.
The
OCSes
also
shrank
deployment
time.
TPU
v3
systems
were
not
usable
until
all
1024
chips
and
all
cables
were
installed
and
tested.
Delivery
delays
for
any
component
held
up
the
entire
supercomputer.
For
TPU
v4,
OCSes
made
each
rack
independent,
so
each
4
3
block
was
put
into
production
as
soon
as
64
chips
and
the
necessary
cables
were
installed
and
tested.
Incremental
deployment
greatly
improved
the
time
to
production
use
and
thus
cost
effectiveness
of
the
TPU
v4
supercomputers.
3