[FSDP2] set vocab_size=32 to avoid must be divisible by 16 error

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 6e80e7c
Pull Request resolved: #264

Author: weifengpy

test/test_fsdp2/test_fsdp2_eager.py

@@ -57,7 +57,12 @@ def init_multi_module(self) -> nn.Module:
     def init_transformer(self, weight_tying: bool) -> nn.Module:
         torch.manual_seed(42)
         args = ModelArgs(
-            n_layers=3, dim=768, n_heads=12, dropout_p=0.0, weight_tying=weight_tying
+            n_layers=3,
+            dim=768,
+            n_heads=12,
+            dropout_p=0.0,
+            weight_tying=weight_tying,
+            vocab_size=32,
         )
         module = Transformer(args).cuda()
         self.broadcast_module(module)