You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
E0531 10:54:04.832741 2609008 spmd_partitioner.cc:569] [spmd] Involuntary full rematerialization. The compiler was not able to go from sharding {devices=[2,1,4]<=[4,2]T(1,0)} to {devices=[1,8,1]<=[8]} without doing a full rematerialization of the tensor. You probably want to enrich the sharding annotations to prevent this from happening.
E0531 10:54:04.832805 2609008 spmd_partitioner.cc:569] [spmd] Involuntary full rematerialization. The compiler was not able to go from sharding {devices=[2,1,4]<=[4,2]T(1,0)} to {devices=[1,8,1]<=[8]} without doing a full rematerialization of the tensor. You probably want to enrich the sharding annotations to prevent this from happening.
We've had to write our own resharding logic instead of solely relying on with_sharding_constraint to avoid this issue.
The text was updated successfully, but these errors were encountered:
This means there was some conflict in sharding - potentially sharding that was propagated from an intermediate.
The XLA compiler had to rematerialize the full tensor to reshard it.
This PR improves logging to show you which Hlo Instruction has the conflict #15402
If you want to further debug this, look at the Hlo dump after ShardingPropagation pass but before SPMD partitioning. Then using my logging PR look at the HloSharding metadata of that HloInstruction and the instructions around it. Most likely you will see there is a conflict like (4,8)->(8,4) triggered the reshard.
Ported this issue from google/jax#21562
This code
Gives this warning
We've had to write our own resharding logic instead of solely relying on
with_sharding_constraint
to avoid this issue.The text was updated successfully, but these errors were encountered: