Merge pull request #19 from b-chu/patch-1

Enable FSDP sharding for bias

Author: tgale96

megablocks/layers/moe.py

@@ -119,7 +119,7 @@ def __init__(self, args : Arguments):
         # Note that the output bias is not parallelized with expert
         # model parallelism.
         self.bias = torch.nn.Parameter(torch.empty(
-            1, 1, args.hidden_size,
+            args.hidden_size,
             device=args.device,
             dtype=common.dtype(args)))
         torch.nn.init.zeros_(self.bias)