We originally did not include these layers in the tp plan because they have tied weights. Since that has been fixed (#666), we should be able to enable TP for these layers