[WIP] [CUDA] fsdp

[nastya236]

FSDP

Usage [it is questionable]:

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

The difference in peak memory is exactly as expected 14 GB [4B * 4 = 16GB -- master copy weights, therefore each rank holds only 16GB / 8 =2 GB 2 instead of 16GB].

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.