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

Cleaned-text code points [39074, 49744). 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
interpolations
based
on
the
number
of
chips.
The
published
MLPerf
results
for
TPU
v4
and
A100
both
scale
to
much
larger
systems
than
the
IPU
(4096
vs
256
chips).
For
similar
sized
systems,
TPU
v4
is
1.15x
faster
for
BERT
than
the
A100
and
~4.3x
faster
than
the
IPU.
For
ResNet,
TPU
v4
is
1.67x
and
~4.5x
faster,
respectively.
Table
6
shows
the
power
we
measured
running
MLPerf
benchmarks;
A100s
use
on
average
1.3x–1.9x
more
power.
Figure
14:
Reported
MLPerf
Training
2.0
highest
performance
[36]
relative
to
A100.
Each
column
labels
the
number
of
chips
per
system.
Graphcore
submitted
results
for
BERT
and
ResNet.
TPU
v4
DLRM
is
in
the
research
category.
The
MLPerf
DLRM
is
not
representative
of
production
DLRMs
[64]
(see
Section
7.9).
Figure
15:
Reported
MLPerf
training
2.0
performance
for
BERT
(top)
and
ResNet
(bottom)
[36]
relative
to
an
8-Way
A100
GPU
on
a
log-log
scale.
To
include
IPUs,
we
only
show
BERT
and
ResNet
in
this
figure.
At
the
largest
scale
of
4096
chips,
TPU
v4
is
1.15x
as
fast
as
the
Nvidia
A100
for
BERT.
At
256
chips,
the
maximum
IPU
size
in
MLPerf,
TPU
v4
is
~4.3x
as
fast
as
the
MK2
IPU
Bow.
At
the
largest
scale,
4096
TPU
v4s
are
1.67x
as
fast
as
4216
Nvidia
A100s
for
ResNet.
At
256
chips,
TPU
v4
is
~4.5x
as
fast
as
the
MK2
IPU
Bow.
The
points
are
the
reported
results,
and
the
dashed
lines
are
interpolations
for
intermediate
sizes
systems.
For
TPU
v4,
the
results
for
≤2048
chips
are
from
MLPerf
Training
1.0;
all
the
other
points
for
all
systems
are
from
MLPerf
Training
2.0.
Table
6:
Mean
power
for
DSA
plus
HBM
for
64
chip
systems
running
MLPerf.
Adding
the
switches
would
show
even
higher
efficiency
for
TPU
v4
(Section
7.3).
We
use
nvidia-smi
reported
power
measurement
on
Azure
Standard_ND96amsr_A100_v4
VMs
during
a
rerun
of
Nvidia's
MLPerf
2.0
64-chip
submission.
TPU
v4
power
measurements
are
done
by
running
Google
MLPerf
2.0
benchmark
code
on
64-chip
scale
in
Google
data
center.
The
TPU
v4
mean
power
measurement
is
2%–8%
higher
than
in
Table
4,
but
the
workloads
differ.
MLPerf
3.0
may
add
power
measurements
to
performance
in
the
October
2023
round.
MLPerf
Benchmark
A100
TPU
v4
Ratio
BERT
380
W
197
W
1.93
ResNet
273
W
206
W
1.33
7
DISCUSSION
We
comment
next
on
11
questions
readers
might
have
from
our
analysis
of
TPU
v4
and
the
other
DSAs.
7.1
Do
peak
FLOPS/second
predict
real
performance?
Many
in
the
ML
community
think
peak
FLOPS/second
are
a
good
performance
proxy
[45],
but
they
are
not.
For
example,
TPU
v4
is
4.3x–4.5x
faster
on
two
MLPerf
benchmarks
than
IPU
Bow
on
equal
sized
systems
despite
only
having
a
1.10x
edge
in
peak
FLOPS/second.
Another
example
is
that
the
A100
peak
FLOPS/second
rate
is
1.13x
TPU
v4,
but
TPU
v4
is
1.15x–1.67x
faster
for
the
same
number
of
chips.
Figure
16
gives
the
relationship
between
peak
FLOPS/sec
and
memory
bandwidth
using
the
roofline
model
[61].
Figure
16:
Roofline
models
for
TPU
v4,
TPU
v3,
and
A100
plus
DNN
models
[61].
Operational
intensity
is
in
parentheses.
The
A100
higher
peak
performance
is
for
Boost
Mode
of
up
to
1410
MHz;
the
base
clock
is
1095
MHz.
If
the
average
rate
was
1243
MHz,
the
peak
performance
of
the
A100
and
TPU
v4
would
be
equal
(same
ceiling
in
the
roofline
model).
We
measured
the
A100
running
MLPerf
BERT,
and
the
average
clock
rate
was
1280
MHz
due
to
capping
of
power
use
by
Nvidia
GPU
software.
Amdahl’s
Law
reminds
us
that
system
balance—in
compute,
memory,
and
interconnect—with
sufficient
power
and
cooling
to
keep
everything
executing
is
still
important.
The
integrated
ICI
network
lets
TPU
supercomputers
scale
performance
gracefully
and
the
OCSes
let
users
tailor
the
topology
to
the
application
to
improve
performance.
9

ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
Jouppi,
et
al.
7.2
How
does
OCS
differ
from
NVLink
and
NVSwitch?
Every
TPU
since
2017
has
had
its
own
built-in
router
between
links
(ICI)
to
neighboring
chips
in
the
torus.
Like
NVlink,
it
enables
“glueless”
connection
of
TPUs,
but
at
a
larger
scale
than
4-8
GPUs:
256
TPU
v2s
and
1024
TPU
v3s.
We
think
of
optical
circuit
switching
as
the
next
generation
of
ICI
versus
a
response
to
the
latest
NVSwitch,
which
uses
electrical
packet
switching
for
8
GPUs
with
one
switch
and
up
to
256
GPUs
with
two
levels
of
switches.
OCSes
are
just
fibers
connected
by
mirrors,
so
any
bandwidth
running
through
a
fiber
can
be
switched
between
input
and
output
fibers
by
the
OCS
across
4096
chips
today
(or
even
more
in
the
future).
For
example,
an
OCS
could
handle
multiple
terabits/second
per
link
by
using
wavelength
multiplexing.
Moreover,
all
inputs
can
be
connected
to
all
outputs,
but
the
connections
must
be
1:1.
7.3
What
if
TPU
v4
used
IB
versus
OCS?
Let’s
start
with
Infiniband
(IB)
versus
OCS
switches.
Just
as
NVLink
connects
8
GPUs
in
a
DGX,
8
TPUs
would
use
ICI.
We
follow
Nvidia’s
guidance
by
using
a
full
3-level
fat
tree
for
the
hybrid
IB/ICI
network
[41].
At
an
average
of
one
NIC
per
GPU,
a
1120
A100
superpod
needs
164
Mellanox
QM8790
40-port
IB
switches
[41],
each
priced
at
~$15k–$18k
[10,
23].
The
1120
IB
NICs
are
extra.
To
replace
the
48
128-port
OCSes,
4096
TPU
v4s
need
568
IB
switches.
An
OCS
is
no
more
expensive
per
port
than
an
IB
switch,
but
it
can
support
higher
bandwidth
because
it
is
passively
reflecting
light
encoded
at
the
source.
The
hybrid
IB/ICI
option
is
substantially
more
expensive
and
harder
for
software.
Furthermore,
active
packet
processing
for
an
IB
switch
is
far
more
power
hungry
than
the
tiny
amount
of
power
required
to
hold
the
MEMS
mirrors
to
their
configured
orientation
in
an
OCS.
ICI
link
bandwidth
is
2x
IB—400
vs
200
Gbit/s—but
system
speed
is
harder
to
judge.
An
internal
event-driven
simulator
that
operates
at
the
TensorFlow
graph
operation
level
evaluated
a
hybrid
ICI/IB
network.
(It
ignores
protocol
processing
on
the
CPU,
which
can
be
significant.)
Depending
on
the
slice
size,
an
optimized
all-reduce
would
run
1.8x–2.4x
slower
and
an
all-to-all
would
be
1.2x–2.4x
slower.
This
network
heterogeneity
is
also
a
software
challenge.
As
communication
is
only
a
portion
of
training
time,
overall
IB
slowdown
for
a
DNN
might
be
as
little
as
~10%.
However,
the
biggest
impact
is
losing
the
benefits
that
originally
inspired
the
use
of
OCS
(Section
2):
availability,
scale,
utilization,
modularity,
power
efficiency,
deployability,
and
so
on.
7.4
Nvidia
announced
the
H100,
the
successor
to
A100,
in
2022.
Why
not
compare
TPU
v4
to
it?
After
systems
are
running
production
applications
in
the
field,
the
Google
tradition
is
to
write
retrospective,
peer-reviewed
papers
for
prominent
conferences.
The
rationale
is
that
the
intellectual
incentives
and
deadlines
for
a
prestigious
publication
encourage
architects
working
on
the
next
generation
to
take
the
time
to
reflect
and
to
make
careful,
detailed,
apples-to-apples
comparisons
to
the
previous
chip
and
contemporary
alternatives
that
can
pass
peer
review.
The
lessons
learned
improve
future
designs.
The
good
news
is
that
these
retrospective,
peer-reviewed,
apples-to-apples
papers
are
widely
read,
e.g.,
[25]
has
>4000
citations.
If
Google
sold
chips
versus
using
them
internally,
we
might
instead
need
to
publish
unreviewed
whitepapers
much
earlier
in
the
chip
lifecycle.
Speaking
of
apples-to-apples,
both
TPU
v4s
and
A100s
were
deployed
in
2020
and
both
use
7
nm
technology.
The
newer,
700W
H100s
were
not
available
at
AWS,
Azure,
or
Google
Cloud
when
we
did
the
research
in
2022
or
even
when
we
submitted
the
final
camera
ready
paper
in
2023.
The
appropriate
H100
match
would
be
a
successor
to
TPU
v4
widely
deployed
in
a
similar
time
frame
and
technology
(e.g.,
in
2023
and
4
nm).
7.5
Why
30%–90%
more
power
for
A100
(Table
6)?
It
is
hard
to
find
a
complete
quantitative
answer
for
these
two
complex
designs.
The
4x
larger
on-chip
SRAM
(160
MB
versus
40
MB)
allows
memory
transfers
to
DRAM
to
be
in
larger
blocks,
improving
energy
efficiency.
Figure
13
above
shows
turning
on
the
CMEM
local
memory,
which
increases
on-chip
SRAM
from
32
MB
to
160
MB,
improves
performance
by
1.18x
and
performance/Watt
by
1.24x.
The
following
three
qualitative
factors
could
explain
the
rest
of
the
gap,
but
we
can't
say
definitively
without
additional
work.
Support
for
multithreading
on
the
GPU
leads
to
a
100x
larger
register
file
(27
MiB
versus
0.25
MiB),
which
likely
requires
more
energy
for
register
accesses—even
though
the
GPU
uses
a
single
port
SRAM—as
power
generally
increases
with
the
square
root
of
memory
capacity
[22].
Second,
the
128x128
MXUs
of
TPU
v4
mean
each
128
entry
input
gets
reused
128
times,
whereas
the
4x4
FP16
array
multipliers
of
the
A100
only
get
reused
4
times,
leading
to
more
on-chip
SRAM
accesses.
Finally,
the
~40%
larger
A100
chip
may
have
longer
data
buses,
increasing
data
transmission
energy.
7.6
What
is
the
CO2e
from
TPU
v4
vs
other
DSAs?
There
is
considerable
concern
about
the
carbon
footprint
of
ML
[42,
45].
Let’s
compare
the
cloud-only
TPU
v4
to
a
hypothetical
recent
DSA
in
an
on-premise
data
center.
Practitioners
can
reduce
operational
energy
use
and
CO2
emissions
by
optimizing
the
“4Ms”
[42]:
1.
Let’s
assume
TPU
v4
and
another
DSA
are
training
the
same
models,
so
the
Model
parameter
is
1.0
in
this
case.
2.
The
Machine
parameter
is
measured
in
perform-
ance/Watt.
TPU
v4
is
2.7x
TPU
v3.
The
MLPerf
power
plan
is
in
progress,
so
we
estimate
for
others.
TPU
v4
is
1.2x–1.7x
faster
and
1.3x–1.9x
lower
power
than
the
A100
and
4.3x-4.5x
faster
than
the
IPU,
whose
TDP
is
300
W.
TPU
v4
chip
performance/Watt
is
thus
~2x–6x
versus
a
contemporary
DSA;
to
be
conservative,
we
assume
2x
for
this
calculation.
3.
The
Mechanization
parameter
is
data
center
power
efficiency,
measured
as
Power
Usage
Effectiveness
(
PUE
).
For
on
premise
data
centers,
PUE
is
often
high.
Computer
architects
improved
PUE
by
helping
advance
the
state-of-the-art
of
warehouse
scale
computers
(
WSCs
)
[2,
3,
16,
30,
44,
62].
For
example,
Google
halved
its
average
energy
overhead
from
21%
(PUE
=
1.21)
in
2008
to
10%
(PUE
=
1.10)
[20].
Worldwide
average
PUE
fell
from
2.50
in
2008
to
1.57
[52]
as
users
closed
their
older
data
centers
and
switched
to
WSCs
in
the
cloud
[34].
The
relative
energy
consumption
is
then
2
×
1.57
÷
1.10
or
2.85x
more
energy
(kWh)
on
a
contemporary
DSA
in
an
average
on-premise
data
center
versus
TPU
v4
in
Google
Cloud.
4.
Map
factors
in
the
cleanliness
of
the
energy
supply,
which
varies
considerably
by
location.
WSCs
can
be
placed
anywhere,
while
on-premise
data
centers
depend
on
the
local
grid.
For
example,
the
average
portion
of
carbon
free
electrical
energy
10