Tensor parallel support for LLM training.

[czkkkkkk]

Feature request

Hi HF team,

I'm wondering about the current status of tensor parallelism (TP) support in Hugging Face. I've noticed that some standard models, such as llama4 and mixtral, include TP sharding plans, and .from_pretrained appears to support loading models with a TP plan. So it seems that TP is supported for inference.

However, I'm curious about training support. Does the transformers library support TP combined with data parallelism (DP) during training? Also, it looks like .save_pretrained doesn't currently support saving TP-sharded models—can you confirm if that's the case, or if there's a workaround?

Thanks in advance!

Motivation

To support large-scale LLM training with TP.

Your contribution

Happy to contribute if there is a specific way to support TP+DP training.

[S1ro1]

Hello, we currently work on supporting up-to 3D parallelism (TP+DP+FSDP) with Accelerate, you can follow the progress and discussion in this PR. Going to open matching PR to transformers soon.

As for TP-only, this should now work on main of Accelerate and Transformers for training AND inference with these 2 PRs (#37054) and (#3457) merged.

RE model saving: I'm not entirely sure on this, I think it's being worked on in #36434 .