transformers

82603b6c - Allow `device_mesh` have multiple dim (#38949)

Allow `device_mesh` have multiple dim (#38949) * Feat: something * Feat: initial changes * tmp changes to unblock * Refactor * remove todo * Feat: docstring --------- Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>