Fix docstring args names

build/builder.py

@@ -400,7 +400,7 @@ def _maybe_parellelize_model(
     if the user specifies using distributed inference. If not, this is a no-op.
 
     Args:
-        module (:class:`nn.Module`):
+        model (:class:`nn.Module`):
             Module to be parallelized.
         builder_args (:class:`BuilderArgs`):
             Command args for model building.

distributed/checkpoint.py

@@ -108,7 +108,7 @@ def load_checkpoints_to_model(
     We parallelize the module and load the distributed checkpoint to the model.
 
     Args:
-        module (:class:`nn.Module`):
+        model (:class:`nn.Module`):
             Module to be parallelized.
         builder_args (:class:`BuilderArgs`):
             Command args for model building.

distributed/parallelize_llama.py

@@ -28,7 +28,7 @@ def apply_tp(
 
 
     Args:
-        module (:class:`nn.Module`):
+        model (:class:`nn.Module`):
             Module to be parallelized.
         world_mesh (:class:`DeviceMesh`):
             Object which describes the mesh topology
@@ -104,7 +104,7 @@ def parallelize_llama(
     the model must fit on GPU or CPU memory.
 
     Args:
-        module (:class:`nn.Module`):
+        model (:class:`nn.Module`):
             Module to be parallelized.
         world_mesh (:class:`DeviceMesh`):
             Object which describes the mesh topology

distributed/world_maker.py

@@ -24,7 +24,7 @@ def launch_distributed(
     using distributed inference. If not, this is a no-op.
 
     Args:
-        config: str:
+        toml_config: str:
             toml file for the inference config.
     Returns:
         Tuple[Optional[DeviceMesh], Optional[ParallelDims]]: 

eval.py

@@ -167,7 +167,7 @@ def eval(
     Args:
         model (Transformer): The pre-trained language model to evaluate.
         tokenizer: The tokenizer to use for encoding/decoding text.
-        task (str): The name of the evaluation task to perform.
+        tasks (Optional[list]): The names of the evaluation tasks to perform.
         limit (Optional[int]): The maximum number of samples to evaluate (None for all available).
         max_seq_length (Optional[int]): The maximum sequence length allowed for input text.
 
@@ -210,7 +210,7 @@ def main(args) -> None:
     Args:
         checkpoint_path (Path): The path to the model checkpoint file to load.
         compile (bool): Whether or not to compile the model for optimization.
-        task (Optional[str]): The name of the evaluation task or a list of tasks to perform.
+        tasks (Optional[list]): The names of the evaluation tasks to perform.
         limit (Optional[int]): The maximum number of samples to evaluate (None for all available).
         max_seq_length (Optional[int]): The maximum sequence length allowed for input text.
 

tokenizer/tiktoken.py

@@ -116,16 +116,16 @@ def encode(
             s (str): The input string to be encoded.
             bos (bool): Whether to prepend the beginning-of-sequence token.
             eos (bool): Whether to append the end-of-sequence token.
-            allowed_tokens ("all"|set[str]): allowed special tokens in string
-            disallowed_tokens ("all"|set[str]): special tokens that raise an error when in string
+            allowed_special ("all"|set[str]): allowed special tokens in string
+            disallowed_special ("all"|set[str]): special tokens that raise an error when in string
 
         Returns:
             list[int]: A list of token IDs.
 
         By default, setting disallowed_special=() encodes a string by ignoring
         special tokens. Specifically:
         - Setting `disallowed_special` to () will cause all text corresponding
-          to special tokens to be encoded as natural text (insteading of raising
+          to special tokens to be encoded as natural text (instead of raising
           an error).
         - Setting `allowed_special` to "all" will treat all text corresponding
           to special tokens to be encoded as special tokens.