Programming TPUs in JAX | How To Scale Your Model

[jacobaustin123]

JAX programming, broadly construed!

[eitanporat]

Hi! Thanks for the great book.

I believe that Problem 1's attached solution doesn't correspond to the problem description (it's only on one axis).

I have been getting an issue with my solution and have written a Github Issue. I would really appreciate if you could help me out

I also have gotten another error I couldn't solve. I run this line (as described in this chapter):

  replicated = jax.lax.with_sharding_constraint(replicated, P(None, 'Y'))

and I keep getting an error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipython-input-215-3887179796.py in <cell line: 0>()
     10 
     11 with mesh:
---> 12   average_A_shmap = average_shard(A)
     13 
     14 average_A_shmap.addressable_shards

    [... skipping hidden 13 frame]

1 frames
    [... skipping hidden 15 frame]

/usr/local/lib/python3.11/dist-packages/jax/_src/util.py in __hash__(self)
    520 
    521   def __hash__(self):
--> 522     return hash(
    523       (
    524         self.f.__code__,

TypeError: unhashable type: 'set'

has the API perhaps changed?

[eitanporat]

For more context, here is my Colab notebook.

[jacobaustin123]

Hi Eitan. The question wasn't super well phrased and had a few mistakes in it, so thank you for pointing this out. I've gone ahead and clarified the question and also given full solutions to both parts. Let me know if this makes more sense.

I think the sharding constraint may need to be a NamedSharding and not PartitionSpec. I've also updated the book.

[eitanporat]

average_shmap = shard_map.shard_map(
    lambda x: x.sum(keepdims=True), 
    mesh=mesh, 
    in_specs=P('X','Y'), out_specs=P('X','Y')
)

def average(x):
  X, Y = mesh.axis_sizes
  return x.reshape(X, x.shape[0] // X, Y, x.shape[1] // Y).sum(axis=(1, 3))

maybe mean instead of sum?

[jacobaustin123]

Ah whoops, yes. Fixed now.

[timabramski]

First we create In and W sharded across our devices1. W is sharded 2 way along the contracting dimension, while In is sharded 4-ways (along both the contracting and output dimensions). This corresponds to a sharding W[DX, F] and In[BX, DY], aka a kind of model and data parallelism.

I'm pretty sure this W[D_X, F] should be W[D_Y, F]

[jacobaustin123]

Yes, fixed.