Need MPMD supporting for GPU to use pipeline parallelism training large scale model
Hi @MoFHeka ,

Tensorflow supports Data Parallelism only now. Model parallelism is yet to support fully for training. But with help of Dtensors we can achieve both data and model parallelism as per attached [tutorial](https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial).

As per my understanding of pipeline parallelism its hybrid approach of data and model parallelism. @SuryanarayanaY But apparently DTensor can't reach pipeline parallelism. Pipeline parallelism was based on sending and receiving collective operator. 
As tensorflow/compiler/xla/hlo/experimental/auto_sharding/auto_sharding.cc show, it seems that ZeRO stage 3 and pipeline parallel were already supported by XLA.  @MoFHeka I always felt like the XLA and Tensorflow/Jax actually support many hidden features but never mention or write the document for it :)). @dathudeptrai I really agree with that. Large projects often lead to difficulties in project management. 


> As tensorflow/compiler/xla/hlo/experimental/auto_sharding/auto_sharding.cc show, it seems that ZeRO stage 3 and pipeline parallel were already supported by XLA.

Hi @MoFHeka ,
 
I doubt it, correct me if I am wrong. I can see from Tf2.14v code [auto_sharding.cc](https://github.com/tensorflow/tensorflow/blob/v2.14.1/tensorflow/compiler/xla/hlo/experimental/auto_sharding/auto_sharding.cc) XLA supports SPMD which is for data parallelism which is supported by TF. Can you point exactly which part you are referring to that you feel that XLA supports Model parallelism or MPMD. This may help us to escalate the issue to SME and get confirmation. Thanks! > > As tensorflow/compiler/xla/hlo/experimental/auto_sharding/auto_sharding.cc show, it seems that ZeRO stage 3 and pipeline parallel were already supported by XLA.
> 
> Hi @MoFHeka ,
> 
> I doubt it, correct me if I am wrong. I can see from Tf2.14v code [auto_sharding.cc](https://github.com/tensorflow/tensorflow/blob/v2.14.1/tensorflow/compiler/xla/hlo/experimental/auto_sharding/auto_sharding.cc) XLA supports SPMD which is for data parallelism which is supported by TF. Can you point exactly which part you are referring to that you feel that XLA supports Model parallelism or MPMD. This may help us to escalate the issue to SME and get confirmation. Thanks!

@SuryanarayanaY Please check these comments. It said "This can result in a strategy similar to **ZeRO stage 3**. NOTE: The combination of this branch with **pipeline parallel** is not tested."
https://github.com/tensorflow/tensorflow/blob/99d80a9e254c9df7940b2902b14d15914dbbbcd9/tensorflow/compiler/xla/hlo/experimental/auto_sharding/auto_sharding.cc#L3249

And please check here, TPU already support MPMD for a long time. It said "If any of the inputs/outputs have maximal sharding, then fallback to MPMD. "
https://github.com/tensorflow/tensorflow/blob/d0321579d44b2a313df3389f95d77a480d1a0816/tensorflow/compiler/mlir/tensorflow/transforms/tpu_sharding_identification_pass.cc#L580
 @MoFHeka I suggest you use jax instead. SOmething like this or flash-attention, FP8 training, int8 training, ... all available in jax with support from XLA (natively). TF also used XLA but kinda hard to custom.  > @MoFHeka I suggest you use jax instead. SOmething like this or flash-attention, FP8 training, int8 training, ... all available in jax with support from XLA (natively). TF also used XLA but kinda hard to custom.

@dathudeptrai Unfortunately, JAX also doesn't support many features, such as sequence parallelism. And at the XLA level, even with JAX, many features are actually only supported by TPU. 
One more important thing, I can't train the CTR model with JAX, which lacks too many things. @MoFHeka https://github.com/NVIDIA/TransformerEngine/pull/602 > @MoFHeka [NVIDIA/TransformerEngine#602](https://github.com/NVIDIA/TransformerEngine/pull/602)

@dathudeptrai  This really surprised me. I always thought it was difficult to split the segmentation of sequence dimension in JAX sharding process.
But there is one thing I am not sure about, if I am going to use pipeline parallel training LLM with JAX, should I use a ray engine like alpa or a JAX native one? JAX doesn't seem to have a good software library that supports all accelerations right now. @MoFHeka Yeah. Generally speaking, coding in jax is harder than pytorch and a bit easier than TF. About low level customization, I think jax is better. Performance wise in my experiments showed that jax is better than pytorch :). Even deepspeed + Flash-attention-2 + Pytorch still not as good as jax :)). 

You can refer some opensource to see how you can custom the paralelism training in jax. https://github.com/alpa-projects/alpa, this I called `Deepspeed for Jax` :).  @dathudeptrai Unfortunately, due to the lack of sequence parallelism, the compute utilization of Alpa is lower than that of Megatron with the same tensor parallelism optimization. Because Alpa use too much device memory when using TP, which leads to a smaller batch size.

Besides, the CTR models really can't be trained with Jax. The Jax ecosystem of online services, data processing, and other components (such as Keras) is way too far behind TF. @MoFHeka Why not use both TF and Jax at the same time :), you can call jax code in TF code nowadays. Also please check out some advanced attention techniques recently introduced in jax (https://github.com/lhao499/large-sequence-modeling/tree/main). 

I personally think the biggest problem with both TF and Jax is documentations :)). > @MoFHeka Why not use both TF and Jax at the same time :), you can call jax code in TF code nowadays. Also please check out some advanced attention techniques recently introduced in jax (https://github.com/lhao499/large-sequence-modeling/tree/main).

@dathudeptrai Yes, that's right, the Jax kernel can be used in TF code, although there’s no big difference between Jax kernels and Keras layers with XLA. 
But the problem is that the **pipeline parallelism** capabilities of JAX cannot be used in TF. TF currently lacks pipeline parallelism components.
 
Even in recent updates, DTensor is used to support tensor parallelism. But in the training of the **CTR model**, one of the biggest usage scenarios of TF, what is more needed is the ability of pipeline parallelism. @SuryanarayanaY Hi~? Is there any way to implement pipeline training in tensorflow?
'tf.distribute.experimental.rpc.Server' with 'server.register' looks like a good choice, but I'm not sure.
### Issue type

Feature Request

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

tf 2.15

### Custom code

No

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Nowadays pipeline parallelism has been implemented in PyTorch for a long time. It's very useful for training a CTR model with Embedding pipeline or training a large language model between two machine.

### Standalone code to reproduce the issue

```shell
Maybe a easy send/recv construction like tpu with xla?
```


### Relevant log output

_No response_