Rank-reducing vs rank-preserving maps#

We can think ofvmap andpmap as unstacking each array input along an axis(e.g. unpacking a 2D matrix into its 1D rows), applying its body function toeach piece, and stacking the results back together, at least when collectivesaren’t involved:

defcheck_vmap(f,xs):ans=jax.vmap(f,in_axes=(0,),out_axes=0)(xs)expected=jnp.stack([f(x)forxinxs])# vmap reference semanticsprint(allclose(ans,expected))check_vmap(lambdax:x@x,jnp.arange(12).reshape(4,3))

True

For example, ifxs had shapef32[8,5] then eachx would have shapef32[5], and if eachf(x) had shapef32[3,7] then the final stacked resultvmap(f)(xs) would have shapef32[8,3,7]. That is, each application of thebody functionf takes as argument inputs with one fewer axis than thecorresponding argument tovmap(f). We can say these arerank-reducing mapswith unstacking/stacking of inputs/outputs.

The number of logical applications off, orinstances off, is determinedby the size of the input axis being mapped over: for example, if we map over aninput axis of size 8, semantically we get 8 logical applications of thefunction.

In contrast,shard_map does not have this rank-reducing behavior. Instead, wecan think of it as slicing (or “unconcatenating”) along input axes into blocks,applying the body function, and concatenating the results back together (againwhen collectives aren’t involved):

importnumpyasnpdevices=np.array(jax.devices()[:4])mesh=Mesh(devices,('i',))# mesh.shape['i'] = 4jax.set_mesh(mesh)defcheck_shmap(f,y):ans=jax.shard_map(f,in_specs=P('i'),out_specs=P('i'))(y)expected=jnp.concatenate([f(y_blk)fory_blkinjnp.split(y,mesh.shape['i'])])print(allclose(ans,expected))check_shmap(lambdax:x.T@x,jnp.arange(32).reshape(8,4))

True

Recall thatjnp.split slices its input into equally-sized blocks with the samerank, so that if in the above exampley had shapef32[8,5] then eachy_blk would have shapef32[2,5], and if eachf(y_blk) had shapef32[3,7] then the final concatenated resultshard_map(f,...)(y) would haveshapef32[12,7]. Soshard_map maps overshards, or blocks, of its inputs.We can say it’s arank-preserving map with unconcatenating/concatenating ofits inputs/outputs.

The number of logical applications off is determined by the mesh size, notby any input axis size: for example, if we have a mesh of total size 4 (i.e.over 4 devices) then semantically we get 4 logical applications of thefunction, corresponding to the 4 devices physically computing them.

Controlling how each input is split (unconcatenated) and tiled within_specs#

Each of thein_specs identifies some of the corresponding input array’s axeswith mesh axes by name usingPartitionSpecs, representing how to split (orunconcatenate) that input into the blocks to which the body function isapplied. That identification determines the shard sizes; when an input axis isidentified with a mesh axis, the input is split (unconcatenated) along thatlogical axis into a number of pieces equal to the corresponding mesh axis size.(It’s an error if the corresponding mesh axis size does not evenly divide theinput array axis size.) If an input’s pspec does not mention a mesh axis name,then there’s no splitting over that mesh axis. For example:

mesh=jax.make_mesh((4,2),('i','j'))jax.set_mesh(mesh)@jax.shard_map(in_specs=P('i',None),out_specs=P('i','j'))deff1(x_block):print(x_block.shape)# prints (3, 12)returnx_blockx1=jax.device_put(jnp.arange(12*12).reshape(12,12),P('i',None))y=f1(x1)

(3, 12)

Here, because the input pspec did not mention the mesh axis name'j', noinput array axis is split over that mesh axis; similarly, because the secondaxis of the input array is not identified with (and hence split over) any meshaxis, application off1 gets a full view of the input along that axis.

When a mesh axis is not mentioned in an input pspec, we can always rewrite to aless efficient program where all mesh axes are mentioned but the callerperforms ajnp.tile, for example:

@jax.shard_map(in_specs=P('i','j'),out_specs=P('i','j'))deff2(x_block):print(x_block.shape)returnx_blockx=jnp.arange(12*12).reshape(12,12)x_=jnp.tile(x,(1,mesh.shape['j']))# x_ has shape (12, 24)x_=jax.device_put(x,P('i','j'))y=f2(x_)# prints (3,12), and f1(x) == f2(x_)

(3, 6)

In other words, because each input pspec can mention each mesh axis name zeroor one times, rather than having to mention each name exactly once, we can saythat in addition to thejnp.split built into its input,shard_map also hasajnp.tile built into its input, at least logically (though the tiling maynot need to be carried out physically, depending on the arguments’ physicalsharding layout). The tiling to use is not unique; we could also have tiledalong the first axis, and used the pspecP(('j','i'),None).

Physical data movement is possible on inputs, as each device needs to have acopy of the appropriate data.

Controlling how each output assembled by concatenation, block transposition, and untiling usingout_specs#

Analogously to the input side, each of theout_specs identifies some of thecorresponding output array’s axes with mesh axes by name, representing how theoutput blocks (one for each application of the body function, or equivalentlyone for each physical device) should be assembled back together to form thefinal output value. For example, in both thef1 andf2 examples above theout_specs indicate we should form the final output by concatenating togetherthe block results along both axes, resulting in both cases an arrayy ofshape(12,24). (It’s an error if an output shape of the body function, i.e.an output block shape, has a rank too small for the concatenation described bythe corresponding output pspec.)

When a mesh axis name is not mentioned in an output pspec, it represents anun-tiling: when the user writes an output pspec which does not mention one ofthe mesh axis names, they promise that the output blocks are equal along thatmesh axis, and so only one block along that axis is used in the output (ratherthan concatenating all the blocks together along that mesh axis). For example,using the same mesh as above:

auto_mesh=jax.make_mesh((4,2),('i','j'),(Auto,Auto))withjax.set_mesh(auto_mesh):x=jnp.array([[3.]])z=jax.shard_map(lambda:x,in_specs=(),out_specs=P('i','j'))()print(z)# prints the same as jnp.tile(x, (4, 2))z=jax.shard_map(lambda:x,in_specs=(),out_specs=P('i',None))()print(z)# prints the same as jnp.tile(x, (4, 1)), or just jnp.tile(x, (4,))z=jax.shard_map(lambda:x,in_specs=(),out_specs=P(None,None))()print(z)# prints the same as jnp.tile(x, (1, 1)), or just x

[[3. 3.] [3. 3.] [3. 3.] [3. 3.]][[3.] [3.] [3.] [3.]][[3.]]

The body function closing over an array value is equivalent to passing it as anaugment with a corresponding input pspec of P(None, None). As another example,following more closely to the other examples above:

@jax.shard_map(in_specs=P('i','j'),out_specs=P('i',None))deff3(x_block):returnjax.lax.psum(x_block,'j')x=jax.device_put(jnp.arange(12*12).reshape(12,12),P('i','j'))y3=f3(x)print(y3.shape)

(12, 6)

The result has a second axis size of 6, half the size of the input’s secondaxis. In this case, the un-tile expressed by not mentioning the mesh axis name'j' in the output pspec was safe because of the collectivepsum, whichensures each output block is equal along the corresponding mesh axis. Here aretwo more examples where we vary which mesh axes are mentioned in the outputpspec:

@jax.shard_map(in_specs=P('i','j'),out_specs=P(None,'j'))deff4(x_block):returnjax.lax.psum(x_block,'i')x=jax.device_put(jnp.arange(12*12).reshape(12,12),P('i','j'))y4=f4(x)print(y4.shape)# (3,12)@jax.shard_map(in_specs=P('i','j'),out_specs=P(None,None))deff5(x_block):returnjax.lax.psum(x_block,('i','j'))y5=f5(x)print(y5.shape)# (3,6)

(3, 12)(3, 6)

On the physical side, not mentioning a mesh axis name in an output pspecassembles anArray from the output device buffers with replicated layoutalong that mesh axis.

There is no runtime check that the output blocks are actually equal along amesh axis to be un-tiled along, or equivalently that the corresponding physicalbuffers have equal values and thus can be interpreted as a replicated layoutfor a single logical array. But we can provide a static check mechanism whichraises an error on all potentially-incorrect programs.

Because theout_specs can mention mesh axis names zero or one times, andbecause they can be mentioned in any order, we can say that in addition to thejnp.concatenate built into its output,shard_map also has both anuntileand ablock transpose built into its output.

Physical data movement is not possible on outputs, no matter the output pspec.Instead,out_specs just encodes how to assemble the block outputs intoArrays, or physically how to interpret the buffers across devices as thephysical layout of a single logicalArray.

Tracking how values vary over manual mesh axes, andcheck_vma=True#

Under ashard_map, values can vary across function instances, or they can bethe same. For example, when we usein_specs to split an argument over a meshaxis, each function instance along that mesh axis gets a different value:

mesh=jax.make_mesh((2,),('i',))jax.set_mesh(mesh)@jax.shard_map(in_specs=P('i'),out_specs=P('i'))deff(x):print(x)return2*xx=jax.device_put(jnp.arange(6.),P('i'))f(x)

On TFRT_CPU_0 at mesh coordinates (i,) = (0,):[0. 1. 2.]On TFRT_CPU_1 at mesh coordinates (i,) = (1,):[3. 4. 5.]

Array([ 0.,  2.,  4.,  6.,  8., 10.], dtype=float32)

If insteadin_specs does not split the argument over a mesh axis, the valueis the same for each function instance along that axis:

@jax.shard_map(in_specs=P(),out_specs=P())deff(x):print(x)return2*xx=jnp.arange(6.)f(x)

On TFRT_CPU_0 at mesh coordinates (i,) = (0,):[0. 1. 2. 3. 4. 5.]On TFRT_CPU_1 at mesh coordinates (i,) = (1,):[0. 1. 2. 3. 4. 5.]

Array([ 0.,  2.,  4.,  6.,  8., 10.], dtype=float32)

A collective’s output may have a different variance than its input. Forexample, applying apsum produces the same output on each function instancealong an axis:

@jax.shard_map(in_specs=P('i'),out_specs=P())deff(x):y=jax.lax.psum(x,'i')print(y)returnyx=jax.device_put(jnp.arange(6.),P('i'))f(x)

On TFRT_CPU_0 at mesh coordinates (i,) = (0,):[3. 5. 7.]On TFRT_CPU_1 at mesh coordinates (i,) = (1,):[3. 5. 7.]

Array([3., 5., 7.], dtype=float32)

In general, each intermediate value in ashard_map can be either unvarying orpossibly-varying over each manual mesh axis. That information can be tracked inthe JAX type system, enabled by thecheck_vma=True argument toshard_map:

@jax.shard_map(in_specs=P('i'),out_specs=P())deff(x):print(jax.typeof(x))# f32[3]{i}y=jax.lax.psum(x,'i')print(jax.typeof(y))# f32[3]returnyx=jax.device_put(jnp.arange(6.),P('i'))f(x)

float32[3]{V:i}float32[3]

Array([3., 5., 7.], dtype=float32)

Here, the typef32[3]{i} means that the value ofx is varying over meshaxis'i'. The type ofy printing asf32[3] indicates it is unvarying overall mesh axes; that is, empty sets are not printed. We call this part of thetype thevarying manual axes (VMA), and it can be accessed viajax.typeof(x).vma.

In general, the VMA type of a value can include any subset of the manual meshaxes over which theshard_map is acting:

mesh=jax.make_mesh((4,2),('i','j'))jax.set_mesh(mesh)@jax.shard_map(in_specs=P('i','j'),out_specs=P('i'))deff(x):print(jax.typeof(x))# f32[2,2]{i,j}y=jax.lax.psum(x,'j')assertjax.typeof(y).vma=={'i'}print(jax.typeof(y))# f32[2,2]{i}returnyx=jax.device_put(jnp.arange(8*4.).reshape(8,4),P('i','j'))f(x)

float32[2,2]{V:(i,j)}float32[2,2]{V:i}

Array([[ 2.,  4.],       [10., 12.],       [18., 20.],       [26., 28.],       [34., 36.],       [42., 44.],       [50., 52.],       [58., 60.]], dtype=float32)

Tracking varying manual axes can be useful:

1. Your code can include prints, assertions, or conditionals about whethervalues are varying over expected mesh axes;

2. It enables efficient reverse-mode autodiff that doesn’t require defensivepsums (seeJEP);

3. The correctness ofout_specs can be checked, ruling out the potential bugexample below.

For example, thisout_specs bug is caught withcheck_vma=True, but uncaughtwithout it:

mesh=jax.make_mesh((2,),('i',))jax.set_mesh(mesh)x=jax.device_put(jnp.arange(6.),P('i'))try:y=jax.shard_map(lambdax:x,in_specs=P('i'),out_specs=P())(x)exceptExceptionase:print(e)

shard_map applied to the function '_rem_singleton' was given out_specs which require replication which can't be statically inferred given the mesh:The mesh given has shape (2,) with corresponding axis names ('i',).out_specs is P() which implies that the corresponding output value is replicated across mesh axis 'i', but could not infer replication over any axesCheck if these output values are meant to be replicated over those mesh axes. If not, consider revising the corresponding out_specs entries. If so, consider disabling the check by passing the check_vma=False argument to `jax.shard_map`.

Here theout_specs incorrectly promise that each function instance along meshaxis'i' produces the same value and thus we can choose just one of them.Withcheck_vma=True (the default) it raises an exception, while withcheck_vma=False there is no exception and instead we get silent undefinedbehavior.

Sometimes we want to treat a value that is unvarying over a mesh axis asvarying over that mesh axis. That’s whatjax.lax.pcast does:

@jax.shard_map(in_specs=P(),out_specs=None)deff(x):print(jax.typeof(x))# f32[6]y=jax.lax.pcast(x,'i',to='varying')print(jax.typeof(y))# f32[6]{i}x=jnp.arange(6.)f(x)

float32[6]float32[6]{V:i}

Think ofjax.lax.pcast(...,to='varying') as applying atype cast: it’s a no-op at runtime,though under reverse-mode autodiff it transposes to ajax.lax.psum (seeJEP). Thatmakes sense because they do opposite things to the VMA: wherey:f32[3]{i}=jax.lax.pcast(x:f32[3],'i',to='varying'),we correspondingly havex_grad:f32[3]=jax.lax.psum(y_grad:f32[3]{i},'i').

JAX implicitly insertsjax.lax.pcast(...,to='varying') calls in many cases,especially for binary operations:

@jax.shard_map(in_specs=(P('i'),P()),out_specs=P('i'))deff(x,y):returnx*yx=jax.device_put(jnp.arange(6.),P('i'))y=jnp.arange(3.)print(jax.make_jaxpr(f)(x,y))

{lambda; a:f32[6@i] b:f32[3].letc:f32[6@i] = shard_map[      check_vma=True      in_specs=(P('i',), P())      jaxpr={lambda; d:f32[3]{V:i} e:f32[3].letf:f32[3]{V:i} = pvary[axes=('i',)] e          g:f32[3]{V:i} = mul d fin(g,) }      manual_axes=frozenset({'i'})      mesh=AbstractMesh('i': 2, axis_types=(Explicit,), device_kind=cpu, num_cores=None)      out_specs=(P('i',),)    ] a bin(c,) }

In a jaxpr, the multiplication operation requires the VMA types of itsarguments to match, but for convenience thejax.numpy andjax.lax APIsautomatically applyjax.lax.pcast(...,to='varying') to make argument VMAtypes agree. In a jaxpr, thesejax.lax.pcast calls show up aspvary sincejax.lax.pcast(...,to='varying') dispatches tolax.pvary.

In some cases, like withjax.lax.scan, you might need to applyjax.lax.pcast yourself to ensure VMA types match as required. For example,this code raises an error:

mesh=jax.make_mesh((2,),('i',))jax.set_mesh(mesh)@jax.shard_map(in_specs=(P('i'),P()),out_specs=P('i'))deff(x,y):defbody(carry,_):c1,c2=carryreturn(c2,c1),()# swap the carry(x_,y_),_=jax.lax.scan(body,(x,y),(),length=2)returnx_,y_x=jnp.arange(6.)y=jnp.arange(3.)try:f(x,y)exceptExceptionase:print(e)

in_specs passed to shard_map: P('i',) does not match the specs of the input: P(None,) for arg: float32[6]. `in_specs` is an optional argument so you can omit specifying it and shard_map will infer the in_specs from the arguments. If you want to reshard your inputs, you can use `jax.reshard` on the arguments and then pass those args to shard_map.

To make the types match, we need to applyjax.lax.pcast to some arguments tothescan:

mesh=jax.make_mesh((2,),('i',))jax.set_mesh(mesh)@jax.shard_map(in_specs=(P('i'),P()),out_specs=P('i'))deff(x,y):defbody(carry,_):c1,c2=carryreturn(c2,c1),()# swap the carryy=jax.lax.pcast(y,'i',to='varying')# apply pcast to fix the error(x_,y_),_=jax.lax.scan(body,(x,y),(),length=2)returnx_,y_x=jax.device_put(jnp.arange(6.),P('i'))y=jnp.arange(3.)f(x,y)

(Array([0., 1., 2., 3., 4., 5.], dtype=float32), Array([0., 1., 2., 0., 1., 2.], dtype=float32))

Here’s a summary of collective primitives and how they affect varying manual axis types:

A few notes on the table:

• The functionjax.lax.psum is a convenience wrapper aroundpsum_invariant.

• It’s surprising thatall_gather isVarying->Varying, but that’s becauseit’s really the transpose ofpsum_scatter which isVarying->Varying.

• Neitherpscatter norall_gather_invariant have user APIs at the time ofwriting, but they’re described here for completeness.

Example 1:all-gather on one side#

Consider performing a matrix multiplication where we shard the left-hand sideargument (can think: parameters) on its leading (non-contracting) dimension:

lhs_spec=P('i',None)lhs=device_put(jax.random.normal(jax.random.key(0),(8,8)),lhs_spec)

And we shard the right-hand side argument (can think: activations) on itscontracting dimension, with a similar sharding for the output:

rhs_spec=P('i',None)rhs=device_put(jax.random.normal(jax.random.key(1),(8,4)),rhs_spec)

To perform this matrix multiplication, we can first all-gather the right-handside and then perform local matrix multiplies against the sharded left-handside:

@jax.jit@jax.shard_map(in_specs=(lhs_spec,rhs_spec),out_specs=rhs_spec)defmatmul_allgather(lhs_block,rhs_block):rhs=jax.lax.all_gather(rhs_block,'i',tiled=True)returnlhs_block@rhs

out=matmul_allgather(lhs,rhs)print(jnp.allclose(out,lhs@rhs,atol=1e-3,rtol=1e-3))

True

That’s great, but we’re not getting any compute/communication overlaphere: before we can start the matmul, we need theall_gather to complete.Here’s a profile using the same code, but on larger example shapes ((8192,8192) forlhs and(8192,1024) forrhs):

We can get compute/communication overlap if instead of callingall_gather webasically inline our above implementation ofall_gather in terms ofppermute, then interleave steps of the gather permutation with local matrixmultiplies:

@jax.jit@jax.shard_map(in_specs=(lhs_spec,rhs_spec),out_specs=rhs_spec)defmatmul_allgather_overlapped(lhs_block,rhs_block):size=jax.lax.axis_size('i')idx=jax.lax.axis_index('i')shift=partial(jax.lax.ppermute,axis_name='i',perm=[(i,(i+1)%size)foriinrange(size)])B=lhs_block.shape[1]//sizelhs_blocks=lambdai:lax.dynamic_slice_in_dim(lhs_block,i*B,B,1)out_block=lhs_blocks(idx)@rhs_blockforiinrange(1,size):rhs_block=shift(rhs_block)out_block+=lhs_blocks((idx-i)%size)@rhs_blockreturnout_block

out=matmul_allgather_overlapped(lhs,rhs)print(jnp.allclose(out,lhs@rhs,atol=1e-3,rtol=1e-3))

True

This implementation allows overlap between communication and computation, andalso avoids gathering a large intermediate onto each device. But on TPU it usesonly half the interconnect bandwidth by permuting in only one direction alongthe ring. To permute bidirectionally, we just split the blocks in half and sendeach half in each direction:

@jax.jit@jax.shard_map(in_specs=(lhs_spec,rhs_spec),out_specs=rhs_spec)defmatmul_allgather_overlapped_bidi(lhs_block,rhs_block):size=jax.lax.axis_size('i')idx=jax.lax.axis_index('i')shift_up=partial(jax.lax.ppermute,axis_name='i',perm=[(i,(i+1)%size)foriinrange(size)])shift_dn=partial(jax.lax.ppermute,axis_name='i',perm=[(i,(i-1)%size)foriinrange(size)])B=lhs_block.shape[1]//size//2# half-size blockslhs_blocks=lambdai,hi:lax.dynamic_slice_in_dim(lhs_block,(2*i+hi)*B,B,1)rhs_block_lo,rhs_block_hi=jnp.split(rhs_block,2,axis=0)out_block=lhs_blocks(idx,0)@rhs_block_loout_block+=lhs_blocks(idx,1)@rhs_block_hiforiinrange(1,size):rhs_block_lo=shift_up(rhs_block_lo)rhs_block_hi=shift_dn(rhs_block_hi)out_block+=lhs_blocks((idx-i)%size,0)@rhs_block_loout_block+=lhs_blocks((idx+i)%size,1)@rhs_block_hireturnout_block

out=matmul_allgather_overlapped_bidi(lhs,rhs)print(jnp.allclose(out,lhs@rhs,atol=1e-3,rtol=1e-3))

True

In practice, to reduce compile times we would probably roll this into ajax.lax.fori_loop. We might also have additional axes of parallelisminvolved.

Example 2:psum_scatter the result#

Another sharding we might start with has bothlhs andrhs sharded alongtheir contracting dimensions, with the output sharded likerhs again:

lhs_spec=P(None,'i')lhs=device_put(lhs,lhs_spec)rhs_spec=P('i',None)rhs=device_put(rhs,rhs_spec)

Here we can use areduce_scatter to perform the contraction sum over shards:

@jax.shard_map(in_specs=(lhs_spec,rhs_spec),out_specs=rhs_spec)defmatmul_psumscatter(lhs_block,rhs_block):out_summand=lhs_block@rhs_blockreturnjax.lax.psum_scatter(out_summand,'i',tiled=True)out=matmul_psumscatter(lhs,rhs)print(jnp.allclose(out,lhs@rhs,atol=1e-3,rtol=1e-3))

True

But the scattering communication must wait for the entire local matrix multiplyto finish before it can start. To get communication/computation overlap, we caninline an implementation ofpsum_scatter in terms ofppermute, theninterleave the communication steps with local matrix multiplies:

@jax.shard_map(in_specs=(lhs_spec,rhs_spec),out_specs=rhs_spec)defmatmul_psumscatter_overlapped(lhs_block,rhs_block):size=jax.lax.axis_size('i')idx=jax.lax.axis_index('i')shift=partial(jax.lax.ppermute,axis_name='i',perm=[(i,(i-1)%size)foriinrange(size)])lhs_block=lhs_block.reshape(size,-1,lhs_block.shape[1])# split 1st axisout_summand=lhs_block[(idx+1)%size]@rhs_blockforiinrange(1,size):out_summand=shift(out_summand)out_summand+=lhs_block[(idx+i+1)%size]@rhs_blockreturnout_summand

out=matmul_psumscatter_overlapped(lhs,rhs)print(jnp.allclose(out,lhs@rhs,atol=1e-3,rtol=1e-3))

True

As in the previous example, to fully utilize interconnects on TPU, we’d run abidirectional version:

@jax.shard_map(in_specs=(lhs_spec,rhs_spec),out_specs=rhs_spec)defmatmul_psumscatter_overlapped_bidi(lhs_block,rhs_block):size=jax.lax.axis_size('i')idx=jax.lax.axis_index('i')shift_up=partial(jax.lax.ppermute,axis_name='i',perm=[(i,(i+1)%size)foriinrange(size)])shift_dn=partial(jax.lax.ppermute,axis_name='i',perm=[(i,(i-1)%size)foriinrange(size)])B=lhs_block.shape[0]//size//2# half-size blockslhs_blocks=lambdai,hi:lax.dynamic_slice_in_dim(lhs_block,(2*i+hi)*B,B,0)out_summand_lo=lhs_blocks((idx-1)%size,0)@rhs_blockout_summand_hi=lhs_blocks((idx+1)%size,1)@rhs_blockforiinrange(1,size):out_summand_lo=shift_up(out_summand_lo)out_summand_hi=shift_dn(out_summand_hi)out_summand_lo+=lhs_blocks((idx-i-1)%size,0)@rhs_blockout_summand_hi+=lhs_blocks((idx+i+1)%size,1)@rhs_blockreturnjnp.concatenate([out_summand_lo,out_summand_hi])

out=matmul_psumscatter_overlapped_bidi(lhs,rhs)print(jnp.allclose(out,lhs@rhs,atol=1e-3,rtol=1e-3))

True

8-way batch data parallelism#

The simplest multi-device parallelism strategy is to shard the batch of inputsand targets over multiple devices, replicate the parameters over those devices,and apply the model in parallel to those shards of data. To evaluate the totalloss, the devices need only communicate with a scalar-sized all-reduce-sum atthe end. (To evaluate the gradient of the loss, the devices must performall-reduce-sums of parameter gradients in the backward pass.)

fromjax.shardingimportMesh,PartitionSpecasPmesh=jax.make_mesh((8,),('batch',))jax.set_mesh(mesh)# replicate initial params on all devices, shard data batch over devicesbatch=jax.device_put(batch,NamedSharding(mesh,P('batch')))params=jax.device_put(params,NamedSharding(mesh,P()))# adapt the loss function to sum the losses across devices@jax.shard_map(out_specs=P())defloss_dp(params,local_batch):inputs,targets=local_batchpredictions=predict(params,inputs)# use reference 'predict`local_loss=jnp.mean(jnp.sum((predictions-targets)**2,axis=-1))returnjax.lax.pmean(local_loss,'batch')

We can check that the loss and its gradients match the reference (base) model:

print(jax.jit(loss)(params,batch))print(jax.jit(loss_dp)(params,batch))

11.920311.9203

defallclose(a,b):returntree_all(tree_map(partial(jnp.allclose,atol=1e-2,rtol=1e-2),a,b))print(allclose(jax.jit(jax.grad(loss))(params,batch),jax.jit(jax.grad(loss_dp))(params,batch)))

True

We can print the compiler IR to inspect the gradient computation and verifythat the collective all-reduce-sum operations happen where we’d expect: at theend of the forward pass to compute the loss value, and in the backward pass tocompute the total parameter gradients.

8-way fully sharded data parallelism (FSDP)#

Another strategy is to additionally shard the parameters over the devices,all-gathering each one when the full value is needed for thejnp.dot or biasaddition. Since we only have one full parameter in local device memory at atime, rather than keeping all parameters in all device memories as in thepreceding DP example, we free up significant memory that we can use for largermodels or larger batch sizes. And because XLA will overlap computation andinter-device communication, the wall-clock time doesn’t suffer.

So now we need collectives in two places: the model prediction functionpredict needs to all-gather the parameters before they’re used, and as in theDP case the loss function needs to sum the local losses to compute the totalloss.

There’s one other ingredient we need: we don’t want to store the fully gatheredparameters from the forward pass for use on the backward pass. Instead, we wantto gather them again on the backward pass. We can express that by usingjax.remat with acustompolicy(or acustom_vjp), though XLA typically does that rematerializationautomatically.

This generalFSDPapproach is similartoweight update sharding (WUS) andZeRO-3.

# shard data batch *and params* over devicesmesh=jax.make_mesh((4,),('batch',))jax.set_mesh(mesh)batch=jax.device_put(batch,P('batch'))params=jax.device_put(params,P('batch'))# adapt the prediction function to gather weights just before their use,# and to re-gather them on the backward pass (rather than saving them)@partial(jax.remat,policy=lambdaop,*_,**__:str(op)!='all_gather')defpredict_fsdp(params_frag,inputs):forW_frag,b_fraginparams_frag:W=jax.lax.all_gather(W_frag,'batch',tiled=True)b=jax.lax.all_gather(b_frag,'batch',tiled=True)outputs=jnp.dot(inputs,W)+binputs=jax.nn.relu(outputs)returnoutputs@jax.shard_map(out_specs=P())defloss_fsdp(local_params,local_batch):inputs,targets=local_batchpredictions=predict_fsdp(local_params,inputs)local_loss=jnp.mean(jnp.sum((predictions-targets)**2,axis=-1))returnjax.lax.pmean(local_loss,'batch')

Again we can check that the loss and its gradients match the reference model:

repl_params=jax.device_put(params,P())repl_batch=jax.device_put(batch,P())print(jax.jit(loss)(repl_params,repl_batch))print(jax.jit(loss_fsdp)(params,batch))print(allclose(jax.jit(jax.grad(loss))(repl_params,repl_batch),jax.jit(jax.grad(loss_fsdp))(params,batch)))

11.92029811.920298True

8-way tensor parallelism (TP)#

Usually we don’t use tensor model parallelism by itself, but seeing it inisolation is a good warmup on parallel matrix multiplication. It’s also a goodexample of usingshard_map in a library function, called in a largerjit-based computation.

The parallelization idea is that we’ll keep the data/activations sharded overits feature axis (rather than its batch axis), and we’ll similarly shard weightmatrices over their input-feature axis (and biases over their feature axis).Then to perform the parallel matrix multiplication, we’ll perform local matrixmultiplications followed by apsum_scatter to sum the local results andefficiently scatter the result’s shards.

mesh=jax.make_mesh((8,),('feats',))jax.set_mesh(mesh)batch=jax.device_put(batch,NamedSharding(mesh,P(None,'feats')))params=jax.device_put(params,NamedSharding(mesh,P('feats')))defpredict_tp(params,inputs):forW,binparams:outputs=gemm_tp(inputs,W,b)inputs=jax.nn.relu(outputs)returnoutputs@jax.shard_map(in_specs=(P(None,'feats'),P('feats',None),P('feats')),out_specs=P(None,'feats'))defgemm_tp(inputs,W,b):block_result=jnp.dot(inputs,W)returnjax.lax.psum_scatter(block_result,'feats',scatter_dimension=1,tiled=True)+bdefloss_tp(params,batch):inputs,targets=batchpredictions=predict_tp(params,inputs)returnjnp.mean(jnp.sum((predictions-targets)**2,axis=-1))# NOTE psum!

FSDP + TP, withshard_map at the top level#

We can compose these strategies together, using multiple axes of parallelism.

mesh=jax.make_mesh((4,2),('batch','feats'))jax.set_mesh(mesh)batch=jax.device_put(batch,NamedSharding(mesh,P('batch','feats')))params=jax.device_put(params,NamedSharding(mesh,P(('feats','batch'))))# mostly same as previous predict_fsdp definition, except we call gemm_tp@partial(jax.remat,policy=lambdaop,*_,**__:str(op)!='all_gather')defpredict_fsdp_tp(params_frag,inputs):forW_frag,b_fraginparams_frag:W=jax.lax.all_gather(W_frag,'batch',tiled=True)b=jax.lax.all_gather(b_frag,'batch',tiled=True)block_result=jnp.dot(inputs,W)outputs=jax.lax.psum_scatter(block_result,'feats',scatter_dimension=1,tiled=True)+binputs=jax.nn.relu(outputs)returnoutputs@jax.shard_map(in_specs=(P(('feats','batch')),P('batch','feats')),out_specs=P())defloss_fsdp_tp(local_params,local_batch):inputs,targets=local_batchpredictions=predict_fsdp_tp(local_params,inputs)sq_err=jax.lax.psum(jnp.sum((predictions-targets)**2,axis=-1),'feats')returnjax.lax.pmean(jnp.mean(sq_err),'batch')

Notice how we have to dotwo collective reductions: one over'feats' andone over'batch'. In the pure TP example, we didn’t write the'feats'reduction explicitly because we only usedshard_map withingemm_tp; in thecallerloss_tp, the compiler automatically translated our use ofjnp.sum toperform apsum as needed given the sharded result returned bypredict_tp.

repl_params=jax.device_put(params,P())repl_batch=jax.device_put(batch,P())print(jax.jit(loss)(repl_params,repl_batch))print(jax.jit(loss_fsdp_tp)(params,batch))print(allclose(jax.jit(jax.grad(loss))(repl_params,repl_batch),jax.jit(jax.grad(loss_fsdp_tp))(params,batch)))

11.92029811.920298True

SPMD pipeline parallelism (PP)#

With pipeline parallelism we aim to parallelize the evaluation of layers atdifferent depths in our network. For example, one device might compute theapplication of the first layer while another device computes the application ofthe second; when they finish, the first device passes its results to the secondwhile the second passes its results to the device responsible for the thirdlayer, and the process repeats. In general the number of pipeline stages may bedifferent from the number of layers, as each stage may be responsible formultiple layers.

With SPMD pipelining, we exploit the fact that most layers in the network applythe computation, just with different parameter values. In particular, we canstack together all the parameters except for those for the first and lastlayers, then use ashard_map to map over blocks of those layer parameters,where each block of parameters corresponds to a pipeline stage. We then use thejax.lax.ppermute collective to shift data down the parallel pipeline.

This particular pipelining strategy is essentiallythe GPipestrategy. There are several variants, aswell as quite different strategies, and which is appropriate can depend on thespeed of the networking between stages and batch sizes. But for this tutorialwe’ll focus on just one strategy.

First, we choose some pipeline parameters:

L=len(params)-2# num layers, excluding first and lastN=batch_size# batch sizeF=params[0][0].shape[1]# num features# choose some pipeline parametersS=2# number of stagesB=8# size of each microbatchassertL%S==0,"S (number of stages) must divide L (number of inner layers)"# compute some useful quantitiesM,ragged=divmod(N,B)# M is number of microbatchesassertnotragged,"B (size of each microbatch) must divide total batch size"K,ragged=divmod(M,S)# K is microbatches per stageassertnotragged,"S (number of stages) must divide number of microbatches"print(f'{S} stages,{L//S} layer(s) per stage,{L} pipelined layers total')print(f'{B} examples per microbatch,{M} microbatches total')