api for model parallelism
`tf.keras.utils.multi_gpu_model` can be used to replicate a Keras model across multiple GPUs, using data parallelism.
`tf.split` and `tf.concat`are low-level TensorFlow operations that can be used to split tensors across multiple devices and combine their results.

As for training and inference of large language models on consumer-level GPUs with limited memory TensorFlow provides several solutions, such as:
**Gradient checkpointing**: This technique allows the model to be split into smaller pieces, which can be processed one at a time, reducing the amount of memory required.

**Mixed precision training**: This technique uses lower-precision floating point numbers (e.g., `float16`) for some parts of the computation, reducing memory usage and improving training speed.

**Model pruning:** This technique involves removing some of the model parameters that are less important, reducing the memory required for both training and inference.


while TensorFlow does not have built-in support for tensor parallelism like Hugging Face's PyTorch Transformers library, it is still possible to implement tensor parallelism in TensorFlow using the low-level APIs mentioned above, such as `tf.split` and `tf.concat`. 
 > tf.keras.utils.multi_gpu_model

Doesn't support. It is deprecated. 

> Mixed precision training

Good tools. But not a solutions for model parallel.

> Gradient checkpointing

Good tools. Doesn't support in tensorflow. (pytorch does, lol)

> Model pruning

Seriously! Any official proven example that shows effectiveness of this approach on LLM models? 

 In TensorFlow, you can use the `tf.distribute` API with `tf.keras` to implement model parallelism for distributed training. Here's a brief overview of how to use it:

1. First, you need to define a `tf.distribute.Strategy` object that represents the distribution strategy you want to use. The `tf.distribute.MirroredStrategy` is commonly used for model parallelism, as it replicates the model across multiple devices and synchronizes the gradients during training.

2. Next, you need to create a `tf.keras` model and compile it as usual. However, you should set the `tf.distribute.Strategy` object as the `tf.keras` optimizer, using the `tf.distribute.Strategy.scope()` context manager.

3. Finally, you can train the model using the `tf.keras` API as usual, but wrapped in a `tf.distribute.Strategy.scope()` context manager. This will ensure that the model is replicated and synchronized across multiple devices during training.

Here's an example code snippet that demonstrates how to implement model parallelism using `tf.distribute` with `tf.keras`:

```import tensorflow as tf

# Define the distribution strategy
strategy = tf.distribute.MirroredStrategy()

# Define the model inside the strategy scope
with strategy.scope():
    model = tf.keras.models.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(784,)),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    
    # Compile the model as usual, but use the strategy optimizer
    model.compile(optimizer=tf.keras.optimizers.Adam(),
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                  metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

# Train the model using the strategy scope
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(64)
with strategy.scope():
    model.fit(train_dataset, epochs=10)
```
This code defines a simple feedforward neural network with two hidden layers and trains it on the MNIST dataset using the `tf.keras` API with `tf.distribute.MirroredStrategy` for model parallelism. Note that the code is almost identical to the non-distributed version, except for the `tf.distribute.Strategy` object and the `tf.distribute.Strategy.scope()` context managers. @pogrushan 

What you did is called data-parallelism and not model parallelism.  I'm not sure if you can achieve Tensor parallel (TP) similar to hugging face implementation, closest one would be [tf.split ](https://www.tensorflow.org/api_docs/python/tf/split) which splits the Tensors into multiple tensors based on the provided axis, the challenge in this would be when using multiple devices to use the synchronous training. 
For this issue, you may find  [spatial parallel training](https://cloud.google.com/blog/products/ai-machine-learning/train-ml-models-on-large-images-and-3d-volumes-with-spatial-partitioning-on-cloud-tpus) useful.

For model parallelism and data parallelism, you can refer `DTensor`, a concept which enables synchronous distributed training. 
Below are the guides and tutorials which can be helpful for you.
1. [Distributed Training with DTensors ](https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial)
2. [DTensor Concepts](https://www.tensorflow.org/guide/dtensor_overview) This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60214">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60214">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Support

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

2.12

### Custom Code

Yes

### 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 Behaviour?

```shell
By using tf.distribute, only the data parallel is documented in the website. https://www.tensorflow.org/guide/distributed_training

Now, what are the APIs that can be used for model parallelism? 

In huggingface, it is documented that, there are many types of parallelism. Read here https://huggingface.co/docs/transformers/v4.15.0/parallelism

What the solutions tensorflow provides for large language model training/inference for consumer level GPU. GPU that has 16GB, 24GB v-ram.

Is it possible to do [Tensor-Parallel](https://huggingface.co/docs/transformers/v4.15.0/parallelism#tensor-parallelism) in Tensorflow.

How model parallism can be applied with APIs in Tensorflow and Keras?
```


### Standalone code to reproduce the issue

```shell
mirrored_strategy = tf.distribute.MirroredStrategy()
cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
```


### Relevant log output

_No response_</details>