transformers

a44dcbe5 - Fixes needed for n-d parallelism and TP (#39562)

Fixes needed for n-d parallelism and TP (#39562) Handle non-DTensors cases in TP Layers Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>