Add support to Fairscale Parallel Layers

[loretoparisi]

Description & Motivation

Add support to Fairscale parallel layers:

from fairscale.nn.model_parallel.layers import (
    ColumnParallelLinear,
    RowParallelLinear,
    VocabParallelEmbedding,
)

that requires _MODEL_PARALLEL_GROUP to be initialized

    self.tok_embeddings = VocabParallelEmbedding(
  File "/home/coder/.local/lib/python3.8/site-packages/fairscale/nn/model_parallel/layers.py", line 118, in __init__
    self.num_embeddings, get_model_parallel_rank(), get_model_parallel_world_size()
  File "/home/coder/.local/lib/python3.8/site-packages/fairscale/nn/model_parallel/initialize.py", line 157, in get_model_parallel_rank
    return torch.distributed.get_rank(group=get_model_parallel_group())
  File "/home/coder/.local/lib/python3.8/site-packages/fairscale/nn/model_parallel/initialize.py", line 128, in get_model_parallel_group
    assert _MODEL_PARALLEL_GROUP is not None, "model parallel group is not initialized"
AssertionError: model parallel group is not initialized

See #20234 for details.

Pitch

To support Llama3.1 initialization that requires Fariscale in the original release with parallel layers.

Alternatives

Get rid of parallel layers or wrap into Lightning (if it is not available already)

Additional context

No response

cc @Borda