[WIP] [CUDA] fsdp

[nastya236]

Fully sharded data parallel.

dtype = mx.bfloat16
for i, layer in enumerate(model.layers):
    model.layers[i] = fully_shard(layer, group=groups["fsdp"], compute_dtype=dtype)
model.embed_tokens = fully_shard(model.embed_tokens, group=groups["fsdp"], compute_dtype =dtype)

It is incomplete, because bf16 weights that are gathered during forward are not resharded and used in backward. So we can cut another ~6GB by resharding and gathering back during backward.

Peak memory difference between sharded optimizer and this PR is 14 GB [4B * 4 = 16GB -- master copy weights, therefore each rank holds only 16GB / 8 =2 GB 2 instead of 16GB] as expected.