Found a small bug in load balancing. rank imbalance computed by sum instead of max - load balancing never took place.