[SPMD] Replace nn.Linear (huggingface#36)

Summary:
This pull request replaces the default nn.Linear with our patched version that doesn't flatten the high dimensional tensors.

Test Plan:
Tested on a V4-8.

examples/pytorch/language-modeling/run_clm.py

@@ -590,6 +590,10 @@ def get_mesh(ici_mesh_shape, dcn_mesh_shape=None):
       else:
         return xs.HybridMesh(ici_mesh_shape=ici_mesh_shape, dcn_mesh_shape=dcn_mesh_shape)
 
+    # Replace the linear layer
+    from torch_xla.distributed.fsdp.utils import apply_xla_patch_to_nn_linear
+    model = apply_xla_patch_to_nn_linear(model, xs.xla_patched_nn_linear_forward)
+
     # Convert the model from meta to XLA tensors one layer at a time to avoid
     # host-side OOM
     for name, param in model.named_parameters():