TP dtensor

[3outeille]

Summary

• Unified TP runtime in distributed/tensor_parallel.py: model.train() uses DTensor redistribute; model.eval() uses plain all-reduce / all-gather
• For testing purpose, only Cohere2Moe is able to easily switch between training (dtensor) & inference (plain all-reduce). Other models (Llama, Mistral, Mixtral, Qwen3) stay on the legacy TP backend
• apply_tensor_parallel() dispatches to the new DTensor backend or the legacy integrations/tensor_parallel path if enable_expert_parallel=False and style string in base_model_tp_plan use the new ALL_PARALLEL_STYLES registry
• Loading/saving hooks: DTensor shard-on-read loading; gather_tp_state_dict() for TP checkpoint save
• tp_size comes from the device_mesh, not from distributed_config field

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: cohere2_moe

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 29078143863:1
Result: failure | Jobs: 16 | Tests: 173,676 | Failures: 24 | Duration: 21h 47m