Incorrect pjit definition of train_state object

[mattiasmar]

Is it obvious to anyone in this forum what explains the inconsistency of this collab's pjit out_axis_resources tree with that of the associated Flax train_state object's value tree PyTreeDef?

This is a duplicate of the discussion jax-ml/jax#9203.

ValueError: pjit out_axis_resources specification must be a tree prefix of the corresponding value, got specification 
((TraceState(trace=FrozenDict({
    Conv_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
    Dense_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
})), None),   #<-------Q1: Is the error due to this value?
 FrozenDict({
    Conv_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
    Dense_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
})) 

for value tree PyTreeDef
(
  (
    (
      CustomNode(namedtuple[<class 'optax._src.transform.TraceState'>],
      [CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()],
      [{'Conv_0': {'bias': *, 'kernel': *}, 'Dense_0': {'bias': *, 'kernel': *}}])]
      ), 
      CustomNode(namedtuple[<class 'optax._src.base.EmptyState'>], [])
    ),
    CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], 
    [{'Conv_0': {'bias': *, 'kernel': *}, 'Dense_0': {'bias': *, 'kernel': *}}])
  )
)

[jheek]

The eval_shape trick you used is clever but you fed the shapes as a tuple into eval_shape causing incorrect results:

params_shapes = jax.tree_map(lambda x: x.shape, model.params)
state_shapes = jax.eval_shape(get_initial_state, params_shapes)

Each shape here is a tuple which is itself a pytree you should simply pass model.params or jax.ShapeDtypeStruct instances to eval shape.

Some minor issues:

1. In the get_opt_spec you return a FrozenDict that comes from set_partitions while the leaf is just a dict
2. bias is rank 1 but you return a rank 2 partition spec P(None, None) instead of P(None)

[mattiasmar]

Thanks that solved the axis ValueError problem.

Now the Jax tells me:
One of pjit outputs was given the resource assignment of PartitionSpec(None, 'MP'), which implies that the size of its dimension 1 should be divisible by 8, but it is equal to 3.
Is there an automatic way to pad or limit the parallelization for the cases when the dimensions of a tensors don't stack up evenly with the number of accelerators, or to tell JAX to parallelize an particular op using fewer accelerators?

[mattiasmar]

I found a programmatic to achieve this that wasn't very cumbersome. Basically just passing the mesh' shape to the set_partition class and verifying that the intended parallelism is applicable for the associated dimension for each tensor.