Add docstrings to some modules and classes (#100)

cacheflow/model_executor/layers/layernorm.py

@@ -1,3 +1,4 @@
+"""Custom normalization layers."""
 import torch
 import torch.nn as nn
 
@@ -5,6 +6,11 @@ from cacheflow import layernorm_ops
 
 
 class RMSNorm(nn.Module):
+    """Root mean square normalization.
+
+    Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight.
+    Refer to https://arxiv.org/abs/1910.07467
+    """
 
     def __init__(
         self,