1D sharding config is very different from 2d sharding config

[tengyifei]

In order to enable 1D sharding, we need to add these to the CLI:

--fsdp "full_shard" --fsdp_config ~/fsdp_config.json

and the fsdp_config.json file is

{
    "fsdp_transformer_layer_cls_to_wrap": [
        "LlamaDecoderLayer"
    ],
    "xla": true,
    "xla_fsdp_v2": true,
    "xla_fsdp_grad_ckpt": true
}

However, in order to enable 2D sharding, we need to instead put this on the CLI:

--spmd_2d_sharding 2

(remember to remove fsdp related flags)

This is very confusing.