Make Dropout a no-op when p=0.0 (#259)

Co-authored-by: Dirk Groeneveld <dirkg@allenai.org>

olmo/model.py

@@ -45,6 +45,14 @@
 log = logging.getLogger(__name__)
 
 
+class Dropout(nn.Dropout):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if self.p == 0.0:
+            return input
+        else:
+            return F.dropout(input, self.p, self.training, self.inplace)
+
+
 class LayerNormBase(nn.Module):
     def __init__(self, config: ModelConfig):
         super().__init__()
@@ -283,7 +291,7 @@ def __init__(self, layer_id: int, config: ModelConfig):
         assert config.d_model % config.n_heads == 0
 
         # Dropout.
-        self.dropout = nn.Dropout(config.residual_dropout)
+        self.dropout = Dropout(config.residual_dropout)
 
         # Layer norms.
         self.k_norm: Optional[LayerNormBase] = None
@@ -639,7 +647,7 @@ def __init__(self, config: ModelConfig, init_params: bool = True):
                 wte=nn.Embedding(
                     config.embedding_size or config.vocab_size, config.d_model, device=config.init_device
                 ),
-                emb_drop=nn.Dropout(config.embedding_dropout),
+                emb_drop=Dropout(config.embedding_dropout),
                 blocks=nn.ModuleList([OlmoBlock.build(i, config) for i in range(config.n_layers)]),
                 ln_f=LayerNorm.build(config),
             )