[distributed][Tensor Parallelism] Implement early transforms for column-wise and row-wise linear and embedding

[crcrpar]

this implements a trace transform that converts one or more linear and/or embedding layers into column-wise or row-wise tensor parallel ones by (1) sharding their weight and bias and (2) inserting needed communication and/or scattering before and/or after the modified layers.

Out of four supported ops, row-wise parallel linear would lead to a BoundSymbol modification. The change caused is to omit the bias term and that bias is added to the result of communication (after post-processing).

---

example

class Model(nn.Module):
    def __init__(self, n_in: int, n_hidden: int, n_out: int) -> None:
        self.l1 = nn.Linear(n_in, n_hidden)
        self.l2 = nn.Linear(n_hidden, n_out)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.l2(F.gelu(self.l1(x))


device = torch.device(f"cuda:{rank}")

model = Model().to(device)
jitted_model = thunder.jit(model)
tp_jitted_model = thunder.distributed.column_parallel(jitted_model, ("l1",))
tp_jitted_model = thunder.distributed.row_parallel(tp_jitted_model, ("l2",))

x = torch.randn(..., device=device)
y = tp_jitted_model(x)
assert y.size(1) == n_out

cc @Borda @apaz-cli @carmocca @awaelchli @crcrpar

[crcrpar]

The failures as of f724a88 look related to #432.

[lantiga]

Amazing work @crcrpar

Mostly nitpicks, it was great fun reviewing this

[t-vi]

Supergood, LGTM. Very excited to see the test_tensor_parallel_both_column_and_row be the first actual example of composing early_transforms!
I added a few minor nits.

[lantiga]

Looks great! Ship it! 🚀