DataParallel/ModelParallel tests needed for TCA/TGA/TRA

Bug description Though we have a few MP tests already, we have had a few bugs crop up now, indicating to me that we're not exhaustively testing.

cc @emilydinan @jxmsML