transformers

[loading][TP] Fix device placement at loading-time, and simplify sharding primitives

#43003

Merged