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

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

ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
Jouppi,
et
al.
v4).
In
contrast
to
all-reduces
of
large
parameter
tensors
in
dense
training,
the
all-to-all
transfers
of
smaller
embedding
vectors
use
HBM
and
ICI
with
finer-grained
access
patterns
for
scatter/gather.
As
separate
cores,
SCs
allow
parallelization
across
dense
compute,
SC,
and
ICI
communications.
Figure
7
shows
the
SC
block
diagram,
which
we
consider
a
“dataflow”
architecture
because
data
flows
from
memory
to
a
variety
of
directly
connected
specialized
compute
units.
The
most
general
SC
units
are
the
16
compute
tiles
(dark
blue
boxes
in
Figure
7).
Each
tile
has
an
associated
HBM
channel
and
supports
multiple
outstanding
memory
accesses.
Each
tile
has
a
Fetch
Unit
,
a
programmable
8-wide
SIMD
Vector
Processing
Unit
(
scVPU
,
not
to
be
confused
with
VPU
of
the
TC
in
TPU
v4),
and
a
Flush
Unit
.
The
Fetch
Unit
reads
activations
and
parameters
from
the
HBM
into
the
tile’s
slice
of
a
2.5
MiB
Sparse
Vector
Memory
(Spmem)
.
The
scVPU
uses
the
same
ALUs
as
TC's
VPU.
The
Flush
Unit
writes
updated
parameters
to
HBM
during
the
backward
pass.
In
addition,
the
five
Cross-Channel
Units
(gold
boxes
in
Figure
7)
perform
specific
embedding
operations,
which
their
names
explain.
Like
TPU
v1,
the
units
execute
CISC-like
instructions
and
operate
on
variable-length
inputs,
where
the
run-time
of
each
instruction
is
data-dependent.
The
cross-channel
units
operate
across
all
16
banks
of
Spmem
collectively.
Figure
7:
SparseCore
(SC)
Hardware
Architecture.
Figure
8:
Bisection
bandwidth
ratio
of
TPU
v4
to
TPU
v3
and
performance
sensitivity
to
bisection
bandwidth.
The
model
used
is
a
DLRM
with
~100M
dense
parameters
in
fully
connected
layers,
~20B
embedding
parameters
(~300
features
mapped
to
~150
tables),
and
1-100
average
valency
per
feature.
The
global
batch
size
is
scaled
proportionately
to
the
number
of
chips.
3.6
SparseCore
Performance
The
end-to-end
embedding
lookup
performance
is
essentially
proportional
to
the
bisection
bandwidth
due
to
the
all-to-all
transfers
of
small
embedding
vectors.
For
the
2D
torus
used
in
TPU
v2
and
TPU
v3,
this
bandwidth
scales
as
N
1/2
for
N
chips.
The
3D
torus
in
TPU
v4
scales
as
N
2/3
[12].
Figure
8
shows
that
the
TPU
v3/v4
bisection
bandwidth
ratio
is
2–4x
higher
at
a
given
chip
count
and
accelerates
embeddings
by
1.1x–2.0x.
At
1024
chips,
SC
overheads
start
to
dominate,
so
bisection
bandwidth
is
less
important.
Figure
9
below
shows
performance
of
an
internal
production
recommendation
model
(DLRM0,
see
Sections
7.8
and
7.9)
across
the
two
TPU
generations
for
128
chips.
The
standalone
CPU
configuration
has
576
Skylake
sockets
(400
for
learners
and
176
for
variable
servers).
The
bottom
two
bars
show
TPU
v4
without
SC,
where
the
embeddings
are
placed
in
CPU
memory.
The
“Emb
on
CPU”
bar
places
embeddings
in
CPU
host
memory
and
the
“Emb
on
Variable
Server”
bar
places
embeddings
on
64
external
variable
servers.
TPU
v3
is
faster
than
CPUs
by
9.8x.
TPU
v4
beats
TPU
v3
by
3.1x
and
CPUs
by
30.1x.
When
embeddings
are
placed
in
CPU
memory
for
TPU
v4,
performance
drops
by
5x–7x,
with
bottlenecks
due
to
CPU
memory
bandwidth.
Figure
9:
Performance
of
an
internal
recommendation
model
(DLRM0)
on
CPUs,
TPU
v3,
TPU
v4,
and
TPU
v4
with
embeddings
in
CPU
memory
(not
using
SparseCore).
The
numbers
in
parenthesis
indicate
the
number
of
sockets
(CPU
or
TPU)
used
for
training.
4
USING
ML
TO
TAILOR
THE
DNN
TO
THE
TPU
AND
THE
TPU
TOPOLOGY
TO
THE
DNN
To
enable
Pareto-optimizations
over
quality
and
performance
for
DNN
models,
we
developed
platform-aware
neural
architecture
search
(PA-NAS)
at
scale
to
tailor
DNN
models
for
TPU
v4
supercomputers
automatically
[32].
A
PA-NAS
designed
CNN1
achieves
~1.6X
better
performance
(QPS
and
latency)
than
the
baseline
designed
by
generic
NAS,
with
comparable
accuracy
[32].
Unlike
[33],
here
we
show
how
PA-NAS
improved
the
performance
of
DLRM0
on
TPU
v4.
DLRMs
are
twice
as
popular
as
CNNs
for
our
ML
workloads
(see
Table
1
above).
Unlike
CNNs,
which
mostly
use
TensorCores
(TCs),
DLRMs
use
both
SCs
and
TCs.
PA-NAS
can
shift
computation
load
between
sparse
layers
(running
on
SCs)
and
dense
layers
(running
on
TCs)
for
Pareto-optimal
performance
and
quality.
Figure
10
below
shows
PA-NAS
on
a
production-scale
DLRM
model.
Despite
having
been
optimized
manually
and
by
using
a
generic
NAS,
the
original
DLRM0
idled
the
SC
~25%
of
the
execution
time
(top
blue
bar
in
Figure
10)
because
of
the
load
imbalance
between
SCs
and
TCs.
PA-NAS
enables
the
end-to-end
Pareto-optimizations
on
quality
and
performance
over
both
6

TPU
v4:
An
Optically
Reconfigurable
Supercomputer
for
ML
with
Hardware
Support
for
Embeddings
Industrial
Product
ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
embedding
layers
(running
on
SC)
and
hidden
layers
(running
on
TC)
for
DLRM0,
which
approaches
perfect
SC-TC
load-balance
(lower
blue
and
red
bars
in
Figure
10)
and
improves
DLRM0
end-to-end
performance
by
>10%.
This
performance
uplift
is
equivalent
to
improvements
historically
achieved
by
a
team
of
>10
experts
over
about
half
a
year,
further
demonstrating
PA-NAS’s
capability
of
increasing
accelerator
flexibility
and
performance
gain.
We
can
also
use
search
to
tailor
the
TPU
v4
topology
to
the
DNN
Model.
Table
3
shows
gains
in
performance
from
searching
for
configurations
for
two
models.
The
first
example
increased
performance
2.3x
by
changing
the
geometry
for
a
512
TPU
v4
slice
over
the
novice
user’s
initial
design
for
an
LLM.
The
second
example
shows
a
harder
task,
demonstrating
an
improvement
of
1.2x
over
an
expert’s
design
for
the
pre-training
phase
of
GPT-3.
Figure
10:
Performance
improvements
by
PA-NAS
for
DLRM0.
Since
DLRMs
use
both
SCs
and
TCs,
the
maximum
of
SparseComputingTime
and
DenseComputingTime
is
the
end-to-
end
training
step
time
of
a
DLRM.
All
computing
times
in
the
figure
(SparseComputingTime
and
DenseComputingTime
for
both
original
and
optimized
DLRM0)
are
normalized
against
the
DenseComputingTime
of
the
original
DLRM0,
since
the
original
DLRM0
is
bottlenecked
on
Dense
Computing.
Embedding
sizes
and
FLOPs
are
normalized
against
those
of
the
original
DLRM0.
Table
3:
Improvements
in
performance
as
we
vary
the
topology
of
a
512
TPU
v4
slice
for
training
an
LLM
and
the
GPT-3
pre-training
stage
.
The
original
LLM
used
model
parallelism
of
dimensions
16×32
and
no
pipeline
or
data
parallelism
for
a
4×8×16
topology.
The
revision
changed
model
parallelism
to
64×8
and
the
topology
to
8×8×8.
The
“1D/2D
activation/weight
partitioning”
option
is
a
typical
way
of
partitioning
tensors
in
a
large
model
graph
(see
Figure
7
in
[63]).
For
GPT-3,
the
original
row
used
a
8×8×8
topology,
pipeline
parallelism
of
depth
8,
no
data
parallelism,
and
model
parallelism
of
dimensions
8×8.
The
revision
changed
the
topology
to
4×8×16,
pipeline
depth
to
16,
data
parallelism
to
4,
model
parallelism
parameters
to
1×8.
Case
Versions
Through-
put
(seqs/sec)
Hyper-Parameters
(topology,
partition
spec
[pipeline,
data,
model1,
model2],
1D/2D
activation/weight
partitioning)
LLM
Novice’s
pick
17.9
(1.0x)
4×8×16,
[1,
1,
16,
32],
2D/2D
Best
perf.
41.3
(2.3x)
8×8×8,
[1,
1,
64,
8],
1D/2D
GPT-3
Pre-
training
Expert’s
pick
21.0
(1.0x)
8×8×8,
[8,
1,
8,
8],
2D/2D
Best
perf.
25.0
(1.2x)
4×8×16,
[16,
4,
1,
8],
1D/1D
5
PRODUCTION
WORKLOAD
PERFORMANCE
Table
1
above
shows
the
workload
mix
for
2022
plus
the
history
of
how
it
has
changed
over
time;
Section
7.7
discusses
the
workload
changes.
We
use
8
applications
to
capture
the
production
workload
to
compare
TPU
v3
to
TPU
v4.
Figure
11
shows
how
efficiently
eight
production
workloads
scale
up
on
TPU
v4.
We
are
encouraged
that
half
of
the
workloads
(CNN0,
RNN0,
RNN1,
and
BERT1)
scale
well
to
3K
chips.
(Figure
4
above
shows
that
given
CPU
hosts
with
availability
of
99.0%
to
99.9%,
in
practice
it
is
much
easier
to
schedule
a
3K
slice
than
a
4K
slice.)
For
the
remainder,
our
production
teams
are
aware
of
where
the
scaling
limits
are
and
are
constructing
solutions,
but
those
solutions
are
not
yet
implemented
to
allow
measuring
performance
at
full
3K
scale
Figure
11:
Scalability
of
TPU
v4
production
workloads
on
a
log-log
scale.
Infrastructural
limitations
currently
hinder
getting
the
last
few
data
points:
BERT0
scales
to
2K,
DLRM0/1
to
1K.
Table
4:
TPU
v4
and
TPU
v3
[26]
features.
Measured
power
is
for
the
ASIC
and
HBM
running
production
applications.
Google
TPUv4
TPUv3
Production
deployment
2020
2018
Peak
TFLOPS
275
(bf16
or
int8)
123
(bf16)
Clock
Rate
1050
MHz
940
MHz
Tech.
node,
Die
size
7
nm,
<600
mm2
16
nm,
<
700
mm2
Transistor
count
22
billion
10
billion
Chips
per
CPU
host
4
8
TDP
N.A.
N.A.
Idle,
min/mean/max
power
90,
121/170/192
W
123,
175/220/262
W
Inter
Chip
Interconnect
6
links
@
50
GB/s
4
links
@
70
GB/s
Largest
scale
configuration
4096
chips
1024
chips
Processor
Style
Single
Instruction
2D
Data
Single
Instruction
2D
Data
Processors
/
Chip
2
2
Threads
/
Core
1
1
SparseCores
/
Chip
4
2
On
Chip
Memory
128
(CMEM)
+
32
MiB
(VMEM)
+
10
MiB
(spMEM)
32
MiB
(VMEM)
+
5
MiB
(spMEM)
Register
File
Size
0.25
MiB
0.25
MiB
HBM2
capacity,
BW
32
GiB,
1200
GB/s
32
GiB,
900
GB/s
7

ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
Jouppi,
et
al.
Table
4
compares
key
features
of
TPU
v3
and
TPU
v4.
Manufactured
in
7
nm
instead
of
16
nm,
TPU
v4
has
twice
the
matrix
multipliers
(enabled
by
the
increased
process
density)
and
an
11%
faster
clock—this
drives
the
2.2X
gain
in
peak
performance.
About
40%
of
the
performance/Watt
improvement
was
from
technology
and
the
rest
was
from
design
improvements
(e.g.,
balancing
the
pipeline,
implementing
clock
gating).
The
HBM
memory
bandwidth
is
1.3x
higher.
Depending
on
the
slice
size,
the
bisection
bandwidth
of
TPU
v4
is
2x–4x
(see
Figure
8
above).
It
also
has
the
128
MB
on-chip
CMEM
scratchpad
memory
not
found
in
TPU
v3.
Figure
12
shows
how
much
faster
TPU
v4
supercomputers
are
than
TPU
v3
supercomputers
at
the
same
slice
size.
Given
the
comparisons
in
Table
4,
it’s
not
surprising
that
at
the
same
slice
size
most
applications
run
1.5x-2.0x
faster
on
TPU
v4
than
on
TPU
v3.
DLRM0
is
3.0-3.5x
faster
and
DLRM1
is
2.8x
at
512
chips
as
TPU
v4
has
twice
as
many
SCs
and
their
clock
rate
is
faster.
The
surprise
is
RNN1;
it
runs
3.3x
faster
on
TPU
v4.
RNN1’s
small
weights
and
small
batch
size
benefit
significantly
from
CMEM
bandwidth
versus
HBM.
Figure
12
:
Speedup
of
TPU
v4
vs
v3
for
the
same
slice
sizes.
Figure
13:
Per
chip
performance
(top)
and
package-level
performance/Watt
(bottom)
for
production
applications
relative
to
TPU
v3
for
CMEM
turned
on
and
off
for
smaller
systems
(e.g,
32
chips).
DLRM1
is
much
faster
at
512
chips.
DLRMs
here
are
different
from
MLPerf
DLRM
(see
Section
7.9).
Figure
13
shows
results
with
CMEM
turned
off
on
TPU
v4;
it
contributes
to
1.2x
performance
gain
overall
but
2x
for
RNN1.
It
also
shows
that
TPU
v4
has
2.1x
the
performance
and
2.7x
the
performance/Watt
of
TPU
v3;
as
mentioned
above,
~40%
of
the
gain
was
from
the
technology
and
the
rest
from
design.
LLM
training
will
become
a
benchmark
in
a
future
MLPerf
release.
We
omit
performance
of
internal
LLMs—31%
of
the
workload
in
Table
1—on
TPU
v3
because
it
is
unoptimized.
TPUv3’s
2D
fixed
topology
hinders
high-performance
model
partitioning
needed
for
LLMs.
We
also
lack
TPUv3
chip
capacity
to
train
large
LLMs
within
reasonable
time-to-convergence
SLOs
given
their
lower
FLOPS/second
and
suboptimal
model
partitioning
performance.
Table
5:
Features
of
the
two
DSAs
[21,
40]
reporting
MLPerf
2.0
Training
results
besides
TPU
v4.
The
A100
has
32×108
=
3456
threads
and
the
IPU
has
6×1472
=
8832
threads.
Nvidia
A100
Graphcore
MK2
IPU
Production
deployment
2020
2021
Peak
TFLOPS
312
(bf16),
624
(i8)
250
(bf16)
Clock
Rate
Base/Boost
1095
/1410
MHz
1850
MHz
Tech.
node,
Die
size
7
nm,
826
mm
2
7
nm,
832
mm
2
Transistor
count
54
billion
59
billion
Chips
per
CPU
host
4
4
TDP
400
W
300
W
Inter
Chip
Interconnect
12
links
@
25
GB/s
3
links
@
64
GB/s
Largest
scale
MLPerf
2.0
configuration
4216
chips
256
chips
Processor
Style
Single
Instruction
Multiple
Threads
Multiple
Instruction
Multiple
Data
Processors
/
Chip
108
1472
Threads
/
Core
32
6
On
Chip
Memory
40
MiB
900
MiB
Register
File
Size
27
MiB
1.40
MiB
HBM2
capacity,
BW
80
GiB,
2039
GB/s
0
6
MLPERF
BENCHMARK
PERFORMANCE
This
section
compares
DSAs
based
on
the
already
published
results
for
MLPerf
Training
[36].
Table
5
compares
the
features
of
two
of
the
entries:
the
NVIDIA
A100
and
Graphcore
MK2
IPU.
While
the
top
six
rows
are
quite
similar
to
TPU
v4
in
Table
4,
the
remaining
rows
show
the
diverse
choices
of
the
architects
in
terms
of
processor
style,
number
of
processors
per
chip,
number
of
threads
per
processor,
register
file
size,
and
on-chip
vs.
off-chip
memory.
First,
TPU
v4
has
only
2
threads
while
A100
has
32×108
=
3456
threads
and
6×1472
=
8832
threads
for
the
IPU.
The
striking
features
of
the
IPU
Bow
are
the
~1500
cores,
900
MB
of
on-chip
SRAM,
and
zero
attached
HBM
or
other
DRAM.
The
striking
features
of
the
A100
are
a
27
MB
register
file
(to
support
multithreading)
and
only
40
MB
of
on-chip
SRAM.
Both
chips
use
full-reticle
dies,
making
them
~40%
larger
than
the
TPU
v4
die.
Figure
14
below
shows
the
fastest
performance
per
DSA
for
five
MLPerf
benchmarks.
(Graphcore
ran
two
of
the
five.)
Vendors
are
free
to
pick
the
size
of
the
system
for
which
they
want
to
report
results.
Ideally,
MLPerf
would
benchmark
systems
of
equal
size
or
cost
or
power,
but
that
is
not
required.
Figure
15
below
shows
the
reported
results
for
ResNet
and
BERT
as
large
points
while
the
dashed
lines
between
the
points
are
8