[Docs] Simple example of using MLX distributed

[stefpi]

related issue: #2930

• added distributed layers into python/nn/layers for auto summarization
• added tensor parallelism inference example in docs
  • described purpose and usage of ShardedToAllLinear and AllToShardedLinear layers
  • showed simple example of combining layers together and its benefit
  • TP applied simple Llama model inference script (#1403 on mlx-examples)
• extracted data parallel example and moved it to dedicated doc

[stefpi]

@awni where would be the ideal place for layers.distributed.shard_linear and layers.distributed.shard_inplace functions to be indexed in the docs? I placed the distributed layers under nn/layers, but I can’t find a natural place for the shard functions.

[awni]

I think we should make a new subsection under "Neural Networks", call it say "Distributed" and add all the helper routines and maybe also distributed layers there.

[awni]

That's pretty cool.

[stefpi]

For a data parallel training example, should I pull the existing example from usage/distributed doc into an example/data parallelism doc and expand on it? Or should I leave that and make something new?

[awni]

We might not need an additional example given we already have something in the usage section. Wdyt?

[stefpi]

I think the existing example is good. Maybe we could put the existing example in a dedicated doc under examples section and then reference both DP and TP examples in the Distributed Communication doc? That way the DP and TP examples are next to each other and it is easier to expand on it in the future without cluttering the usage doc.

[stefpi]

I have extracted the data parallel training example into examples/data parallelism and then referenced both DP and TP examples in usage/distributed communication. Let me know if thats fine, I can revert it back if not.

[awni]

This is super nice, thanks for the contribution!