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

weifengpy · weifengpy · commit cdbf57b25312 · 2024-05-20T16:31:57.000-07:00
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
diff --git a/test/test_fsdp2/test_fsdp2_eager.py b/test/test_fsdp2/test_fsdp2_eager.py
@@ -57,7 +57,7 @@ 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)

Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,7 @@ def init_multi_module(self) -> nn.Module:`
`57`	`57`	`def init_transformer(self, weight_tying: bool) -> nn.Module:`
`58`	`58`	`torch.manual_seed(42)`
`59`	`59`	`args = ModelArgs(`
`60`		`- n_layers=3, dim=768, n_heads=12, dropout_p=0.0, weight_tying=weight_tying`
	`60`	`+ n_layers=3, dim=768, n_heads=12, dropout_p=0.0, weight_tying=weight_tying, vocab_size=32,`
`61`	`61`	`)`
`62`	`62`	`module = Transformer(args).cuda()`
`63`	`63`	`self.broadcast_module(module)`