The implications of check_vma in the shard_map API

[AakashKumarNain]

I think the documentation of shard_map API in general is good enough, but I think more examples on the implication of check_vma can be useful. That aside, let me present a simple example, where it is not clear to me whether it is the right thing to disable check_vma without having any implications.

Say I have an MLP with two linear layers. Here is the config:

# num_devices = 8
devices = np.array(jax.devices())
mesh_axis_names = ("feats",)
mesh = Mesh(devices, axis_names=mesh_axis_names)

# input shape: [batch_size, features], labels shape: [batch_size,]
# params = MLP(
#     fc1=Linear(
#         in_features=784,
#         out_features=128,
#         weight=float32[784,128],
#         bias=float32[128],
#         use_bias=True,
#     ),
#     fc2=Linear(
#         in_features=128,
#         out_features=10,
#         weight=float32[128,10],
#         bias=None,
#         use_bias=False,
#     )
# )

# forward function
@partial(jax.shard_map, mesh=mesh,
         in_specs=(P(None, 'feats'), param_specs),
         out_specs=P(),
         check_vma=False
        )
def forward_tp(inputs, params):
    # inputs shape: (256, 98)
    out = jnp.dot(inputs, params.fc1.weight)

    # our first layer is sharded for TP
    out = jax.lax.psum_scatter(out, 'feats', scatter_dimension=1, tiled=True)
    if params.fc1.bias is not None:
        out = out + params.fc1.bias
        
    # second layer is replicated, so we need to gather all features    
    out = jax.lax.all_gather(out, 'feats', axis=1, tiled=True)
    out = jnp.dot(out, params.fc2.weight)
    return out

Without check_vma=False, it will throw an error. I am not sure if turning it off is advisable or if there is any other workaround.

[yashk2810]

What is the error? If I guess what it is, you will need to use jax.lax.all_gather_invariant for out_specs to be P().

But usually no, you shouldn't need to disable check_vma.

[AakashKumarNain]

The error is this:

Check 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`.

[yashk2810]

Use jax.lax.all_gather_invariant if you want out_specs to be P()

If you want to use all_gather, then your out_specs should be P(None, feats)

[AakashKumarNain]

Got it. Thanks. BTW any doc to read about the varying->varying and varying ->invarying in detail?

[yashk2810]

https://docs.jax.dev/en/latest/notebooks/shard_map.html#tracking-how-values-vary-over-manual-mesh-axes-and-check-vma-true

https://docs.jax.dev/en/latest/jep/17111-shmap-transpose.html