[HARD] Support arbitrary tensor parallel splits

[rltakashige]

Context:
Tensor parallelism allows models to be split among multiple devices by distributing weights among devices.
auto_parallel strategies for tensor parallelism use MLX LM's shard_linear and shard_inplace functions to split tensors evenly. These functions assume that the sharding dimensions are exactly divisible by the group size.

Issue:
There are several limitations to this approach:

1. Certain models are unable to be supported by tensor parallelism, particularly quantized models (see GLM Air 4.5 config.json with gs 32 quantization).
2. Most models cannot be parallelised across 3 nodes, since their intermediate size is not divisible by 3.
3. This issue is a blocker for smarter tensor parallelism on heterogeneous devices.

A potential start is smart padding.

[DevKiDCosmo]

I would also add a small benchmark test to see how to split the model among the devices instead of splitting it equally like I understand. I would like to work on this problem but with someone else to get to know a little bit more about the code base.

The dimensions don't need to be exactly divisible by the group size if we implement a system where percentages or fraction of the model are distributed. To this day, all I know is that the models are always evenly divided, right?

[Evanev7]

yup! we have serious issues running TP on 3 devices because most model hidden dimensions aren't divisible by it.

We do have a fixed split we need to implement, but our TP sharding code is a lot more involved than our Pipeline code so it's not just an issue of rounding some percentages better.

[DevKiDCosmo]

Then I don't see a problem. Probably I misunderstood something.

By TP a layer or multiple layer is splitted across devices. But how are they splitted across those devices? I think I might get it? Instead of splitting the model vertically (by layer) it split horizontally (by neurons). But not every layer has the exact same amount of neurons in a layer. Is this the problem? And is every layer split horizontally?

[DevKiDCosmo]

Is there also something hybrid like powerful devices are taking layers and less powerful devices doing it by TP?

[rltakashige]

Correct me if I misunderstood your reply - a model is composed of multiple transformer blocks (as well as some other modules we don't care too much about).

In pipeline parallelism, some of the transformer blocks are assigned to one device, some others to a second device, etc.
In tensor parallelism, we shard the linear modules inside the transformer blocks. All of the linear modules are split across their intermediate dimension (although which dimension this is depends on the linear module).

The problem isn't that layers have different numbers of parameters, but rather that the tensor parallel code expects computation and communication on tensors of the same size.

Is there also something hybrid like powerful devices are taking layers and less powerful devices doing it by TP?

There is an argument to be made for this (combining pipeline + tensor parallelism). However, as it stands, in practice, the machines capable of doing effective TP have RDMA, and you would gain better performance from simply doing TP on all the machines. While it's outside the scope of this issue, it's a nice thing to explore!

[DevKiDCosmo]

For Pipeline I see by using a 4gb model a rise of 4gb in ram usage which is strange. Why is ist not 1.3gb per device bye 3 device?