Set the default of use_custom_sort_vjp to False

PiperOrigin-RevId: 808631664

Author: jesselu-google

src/MaxText/configs/base.yml

@@ -167,7 +167,7 @@ sparse_matmul: True
 capacity_factor: -1.0 # a factor to decide expert capacity for token dropping, and no dropping by default
 load_balance_loss_weight: 0.01 # weight for the load balance loss
 use_random_routing: False # whether to use random routing for debug/test purpose
-use_custom_sort_vjp: True # whether to use a custom sort vjp for sparse matmul ops
+use_custom_sort_vjp: False # whether to use a custom sort vjp for sparse matmul ops
 use_ring_of_experts: False # whether to use ring of experts for sparse matmul expert parallelism
 # Tunable tiling dimensions used for Megablox
 tile_batch_seq: 512