Add docstrings to some modules and classes (#100)

cacheflow/model_executor/models/gpt2.py

@@ -15,7 +15,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""1D GPT-2 model compatible with HuggingFace weights."""
+"""Inference-only GPT-2 model compatible with HuggingFace weights.
+
+The input of the model is flattened to a 1D tensor of tokens. The model uses
+InputMetadata to extract the original 2D shape of the input.
+"""
 from typing import Dict, List, Optional, Tuple
 
 import torch

cacheflow/model_executor/models/gpt_neox.py

@@ -14,7 +14,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""1D GPT-NeoX model compatible with HuggingFace weights."""
+"""Inference-only GPT-NeoX model compatible with HuggingFace weights.
+
+The input of the model is flattened to a 1D tensor of tokens. The model uses
+InputMetadata to extract the original 2D shape of the input.
+"""
 from typing import Dict, List, Optional, Tuple
 
 import torch
@@ -79,6 +83,7 @@ class GPTNeoXAttention(nn.Module):
 
 
 class GPTNeoXMLP(nn.Module):
+
     def __init__(self, config: GPTNeoXConfig):
         super().__init__()
         self.dense_h_to_4h = ColumnParallelLinear(config.hidden_size,

cacheflow/model_executor/models/llama.py

@@ -19,7 +19,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""1D LLaMA model compatible with HuggingFace weights."""
+"""Inference-only LLaMA model compatible with HuggingFace weights.
+
+The input of the model is flattened to a 1D tensor of tokens. The model uses
+InputMetadata to extract the original 2D shape of the input.
+"""
 from typing import Dict, List, Optional, Tuple
 
 import torch

cacheflow/model_executor/models/opt.py

@@ -14,7 +14,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""1D OPT model compatible with HuggingFace weights."""
+"""Inference-only OPT model compatible with HuggingFace weights.
+
+The input of the model is flattened to a 1D tensor of tokens. The model uses
+InputMetadata to extract the original 2D shape of the input.
+"""
 from typing import Dict, List, Optional, Tuple
 
 import torch