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

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

ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
Jouppi,
et
al.
2.5
OCS
Scheduling
Benefits
The
OCS
also
simplifies
scheduling,
which
increases
utilization.
For
TPU
v3,
a
256
chip
slice
meant
the
scheduler
had
to
find
256
contiguous
chips
that
were
idle.
For
TPU
v4,
it
can
pick
four
4
3
blocks
from
anywhere
in
the
supercomputer.
Slices
don’t
even
need
to
be
a
power
of
2;
they
can
be
4i×4j×4k,
where
0
<
i
≤
j
≤
k.
For
example,
a
user
could
request
a
192
TPU
v4
slice
with
a
geometry
of
4×4×12.
2.6
OCS
Modularity
and
Security
Benefits
Since
the
OCS
can
switch
circuits
in
milliseconds,
TPU
v4
can
easily
change
topology
to
match
the
application,
the
number
of
nodes,
and
the
system
that
runs
those
jobs.
TPU
v4
provides
wraparound
links
of
a
3D
Torus
for
most
slice
sizes,
which
doubles
both
the
bisection
bandwidth
and
the
bandwidth
of
important
collective
communication
operations
(e.g.,
all-reduce)
versus
the
mesh-like
alternative
[12],
yet
still
allowing
the
TPU
v4
to
scale
interconnect
bandwidth
up
to
16
3
(4096)
chips.
OCS
also
enables
an
air
gapped
network
isolation
between
different
slices,
which
enhances
the
security
of
multiple
customers
sharing
a
TPU
v4
supercomputer.
2.7
Tailoring
OCS
Topology
to
Improve
Performance
The
final
benefit
was
a
bonus
beyond
solving
problems
of
large
scale.
To
set
the
stage,
here
are
the
three
fundamental
types
of
parallelism
that
improve
the
training
time
of
DNNs:
1.
Data
Parallelism
:
Each
chip
computes
the
forward
and
backward
pass
on
a
subset
of
examples,
and
sends
the
gradients
that
it
calculates
for
its
subset
to
the
other
chips.
2.
Model
(or
Tensor)
Parallelism
:
Large
tensor
operations
and
their
weights
are
divided
across
multiple
chips,
so
that
each
chip
simultaneously
computes
a
subset
of
a
tensor
operation.
3.
Pipeline
Parallelism
:
For
a
DNN
with
many
layers,
each
chip
computes
a
subset
of
layers,
and
communicates
the
layer
results
to
chips
holding
the
adjacent
layers.
Users
can
change
the
TPU
v4
topology
to
match
the
type
of
parallelism
being
used,
e.g.,
for
a
512
slice,
pipeline
parallelism
might
want
a
cigar
shape
(4×4×32)
instead
of
the
conventional
8
3
cube
(8×8×8).
For
the
highest
bisection
bandwidth,
often
needed
by
embedding
heavy
applications,
the
conventional
8
3
cube
is
preferred.
ML
practitioners
often
combine
parallelism
types
to
get
the
best
performance,
such
as
data
plus
model
parallelism.
Model
parallelism
typically
has
two
parameters:
width
and
length.
To
take
full
advantage
of
the
bandwidth
available,
users
map
data
parallelism
along
one
dimension
of
the
3D
torus
and
the
two
model
parallel
parameters
on
the
other
dimensions.
Table
3
in
Section
4
gives
examples
of
performance
gains
of
1.2x
to
2.3x
by
varying
the
topology
and
the
hyperparameters.
TPU
v4
uses
a
single,
static
topology
for
each
training
job,
which
can
be
co-optimized
with
the
communication
requirements
of
the
training
job.
This
per-job
configuration
is
not
a
fundamental
limitation
of
the
OCS.
See
[55]
for
more
details
of
the
OCS
and
its
physical
construction.
2.8
Twisting
the
Torus
For
a
slice
with
a
number
of
TPUs
that
is
a
perfect
cube,
a
symmetric
torus
with
an
equal
number
of
TPUs
along
each
dimension
minimizes
latency
and
maximizes
bisection
bandwidth
(e.g.,
8
3
=
512
TPUs
organized
as
an
8x8x8
torus).
For
non-perfect-cube
slices,
a
rectangular
torus
can
be
built
with
different
numbers
of
TPUs
along
each
dimension.
Alternatively,
[7]
proposes
a
topology
that
outperforms
rectangular
tori
with
lower
latency
and
higher
bisection
bandwidth
without
increasing
switch
hardware.
The
twisted
torus
rewires
some
links
between
4
3
cubes
to
reduce
the
worst
case
latency.
Figure
5
shows
a
regular
topology
and
a
twisted
topology.
Since
TPU
v4
uses
an
OCS
to
connect
4
3
blocks,
the
“rewiring”
is
mostly
reprogramming
of
routing
in
the
OCS.
Using
all-to-all
communication
with
large
messages
as
a
microbenchmark,
Figure
6
below
shows
the
performance
gain
from
the
twisted
topology.
The
twisted
torus
improves
all-to-all
throughput
by
1.63x
and
1.31x
over
the
regular
torus
on
4×4×8
and
4×8×8
slices,
respectively.
While
the
popular
interconnect
topologies
today
are
Clos
networks
and
Dragonfly
networks,
the
twisting
option
reduces
the
worst
case
bisection
bandwidth
for
the
3D
tori
and
makes
them
more
attractive
for
today’s
supercomputers.
Figure
5:
Example
of
regular
(top)
and
twisted
torus
(bottom)
topologies
for
a
4×2
slice
of
TPU
v4
nodes.
The
TPU
v4
network
is
three-dimensional,
but
the
figure
uses
two
dimensions
for
ease
of
illustration.
Each
TPU
is
labeled
with
its
coordinates
in
the
slice.
The
electrical
connections
(red
dashed
lines)
remain
fixed.
By
utilizing
the
flexibility
of
the
OCSs,
the
optical
connections
(blue
solid
lines)
can
be
reconfigured
from
a
rectangular
torus
to
a
twisted
torus
without
any
physical
recabling
of
the
machine;
the
only
change
is
in
the
routing
tables.
TPU
v4
uses
a
k×k×2k
configuration
from
Camarero,
Martinez,
and
Beivide
[8].
Figure
6:
Measured
all-to-all
throughput
for
4×4×8
and
4×8×8
slices
using
regular
and
twisted
tori.
Measurements
are
steady
state
(large
aggregate
transfer
size)
with
individual
DMAs
being
4
KiB.
Each
column
also
shows
the
theoretical
delta
from
the
ideal
peak
as
a
stacked
bar
and
label
above
the
measured
performance.
4

TPU
v4:
An
Optically
Reconfigurable
Supercomputer
for
ML
with
Hardware
Support
for
Embeddings
Industrial
Product
ISCA
‘23,
June
17–21,
2023,
Orlando,
FL,
USA
2.9
Distribution
of
Topologies
Table
2
above
gives
a
sampling
of
slice
topologies
used
in
production.
Some
tasks
are
smaller
than
a
4
3
block,
so
they
can
only
use
a
2D
mesh;
29%
are
smaller
than
a
4
3
cube,
so
they
obviously
cannot
select
a
twisted
3D
torus.
Of
the
remaining
71%,
only
those
of
the
form
n×n×2n
or
n×2n×2n
(n≥4)
can
twist.
They
are
33%
(48%
of
71%).
The
actual
twisted
tori
are
28%
(86%
of
33%).
Stated
alternatively,
40%
of
the
topologies
that
are
4
3
blocks
or
larger
use
twisted
tori.
2.10
Cost
of
OCS
Flexibility
Remarkably,
given
all
the
benefits
of
OCSes,
their
cost
is
<5%
of
the
total
TPU
v4
supercomputer
capital
costs
and
<3%
of
total
power.
The
power
and
cost
accounting
includes
the
entire
optical
fabric,
including
the
optics
modules,
fiber,
and
OCS
infrastructure.
3.
SPARSECORE:
EMBEDDINGS
SUPPORT
Before
introducing
the
next
TPU
innovation,
let’s
review
recommendation
models,
embeddings,
and
distributed
training
to
set
the
stage.
3.1
Recommendation
Models
Deep
learning
recommendation
models
(DLRMs)
are
a
quarter
of
our
ML
workload
(Table
1
above).
DLRMs
are
used
in
advertising,
search
ranking,
YouTube,
and
Google
Play
applications
[1,
4,
11,
13].
Google’s
production
advertising
models
score
ads
for
billions
of
queries
daily,
and
consist
of
billions
of
weights,
and
train
on
more
than
one
trillion
examples,
and
are
required
to
perform
inference
at
well
over
one
hundred
thousand
requests
per
second
[1].
DLRM
model
sizes
are
determined
using
five
factors:
prediction
quality,
training
time,
total
training
cost,
serving
latency,
and
total
serving
cost.
Embeddings
are
a
key
component
of
DLRMs.
3.2
Embeddings
Inputs
for
DLRMs
consist
mainly
of
categorical
features.
Each
categorical
feature
contains
N
discrete
values
(N
is
commonly
referred
to
as
the
vocabulary
size).
For
example,
in
the
search
ranking
application,
the
search
query
is
a
categorical
feature,
and
N
is
the
number
of
words
in
the
English
language.
A
given
training
example
(query)
is
sparse,
and
contains
a
tiny
subset
of
words.
Neural
networks
typically
train
well
on
dense
vectors.
Embeddings
are
the
standard
and
effective
way
to
transform
categorical
feature
values
into
dense
vectors.
An
embedding
function
translates
from
the
large,
categorical
space
(e.g.,
words
in
the
English
language)
to
a
smaller,
dense
space
(e.g.,
a
100-vector
representing
each
word).
Embedding
functions
are
implemented
using
lookup
tables.
An
example
is
a
table
with
80,000
rows
(one
per
word)
of
width
100.
Each
training
example
can
look
up
either
one
row
(univalent
embedding)
or
a
small,
dynamic
number
of
rows
(multivalent
embedding,
typically
combined
by
summing).
A
neural
network
model
might
have
many
tables
of
many
sizes
for
different
categorical
features.
Embeddings
are
a
key
component
of
Google
DLRMs,
and
typically
form
the
first
layer
in
a
neural
network
model.
3.3
Distributed
Training
Embedding
tables
are
large,
and
can
range
in
size
from
O(10
MiB)
to
O(100
GiB).
In
aggregate,
all
the
embedding
tables
in
a
model
can
be
as
large
as
several
TiBs.
Hence,
such
tables
are
partitioned
across
the
memory
of
several
TPU
chips.
There
are
three
methods
for
partitioning:
(1)
column
sharding
splits
tables
along
their
width
across
multiple
chips,
(2)
row
sharding
splits
tables
along
their
vocabulary
size,
and
(3)
table
sharding
places
different
tables
on
different
chips.
These
distribution
strategies
are
collectively
termed
model
parallelism
in
the
context
of
neural
network
models.
For
small
embedding
tables,
replication
across
all
chips
(using
data
parallelism)
is
better
for
performance.
3.4
Key
Performance
Attributes
Embedding
lookup
operations
consist
mainly
of
small
gather
or
scatter
memory
accesses,
which
have
low
arithmetic
intensity.
As
opposed
to
dense
operations
(e.g.,
transformers,
fully
connected
networks)
where
the
chip
FLOPS/second
is
the
main
driver
of
end-to-end
performance,
embedding
lookup
operations
are
bottlenecked
by
the
memory
bandwidth,
memory
capacity,
and
VPU
(vector
processing
unit)
performance.
The
ICI
interconnec-
tion
network
(across
chips)
is
also
a
significant
performance
driver.
The
interconnect
bandwidth
and
performance
depends
on
the
type
of
parallelism
being
exploited.
For
model
parallelism
(common
case),
the
communication
pattern
consists
of
variable-length
all-to-all
exchange.
The
network
bisection
bandwidth
limits
performance.
For
data
parallelism,
the
communication
pattern
consists
of
all-reduce
operations,
which
injection
bandwidth
limits.
The
unstructured
sparsity
of
embeddings
is
also
prone
to
compute,
memory,
and
communication
load
imbalances
across
a
supercomputer.
To
reduce
load
imbalance,
deduplication
of
frequent
feature
values
is
commonly
used,
and
must
be
efficiently
supported
by
the
compute
substrate.
Deduplication
also
reduces
the
number
of
memory
accesses,
and
the
quantity
of
data
sent
over
the
interconnection
network,
further
improving
performance.
3.5
SparseCore
It’s
time
to
introduce
the
SparseCore
(
SC
).
For
the
training
phase,
embeddings
could
be
placed
on
the
TensorCore
or
the
host
CPUs
of
the
supercomputer.
The
TensorCore
has
wide
VPU
and
matrix
units,
and
is
optimized
for
dense
operations.
Placing
embeddings
on
the
TensorCore
would
be
suboptimal
due
to
small
gather/scatter
memory
accesses,
and
variable
length
data
exchange.
Placing
embeddings
on
the
host
CPUs
of
a
supercomputer
would
induce
an
Amdahl’s
Law
bottleneck
over
the
CPU
DRAM
interface,
amplified
by
the
4:1
TPU
v4
to
CPU
host
ratio.
Tail
latency
and
bandwidth
restrictions
of
data-center
networking
would
further
constrain
the
training
system.
Performance
could
be
optimized
using
the
total
HBM
capacity
of
a
TPU
supercomputer,
joined
by
a
dedicated
ICI
network,
and
with
fast
gather/scatter
memory
access
support.
This
insight
led
to
the
codesign
of
the
SparseCore
(SC).
The
SC
is
a
domain-specific
architecture
for
embedding
training
starting
with
TPU
v2,
with
later
improvements
in
TPU
v3
and
TPU
v4.
SCs
are
relatively
inexpensive,
at
a
total
of
only
~5%
of
the
die
area
and
~5%
of
the
power.
SCs
operate
in
a
sea-of-cores
configuration,
combining
supercomputer-scale
HBM
and
ICI
to
create
a
flat,
globally
addressable
memory
space
(128
TiB
in
TPU
5