Fix docstring Args entries that name a parameter the function does not take

deepspeed/autotuning/scheduler.py

@@ -212,7 +212,8 @@ def parse_results(self, metric):
         """ Parses the metric file of the finished experiments to select the optimal DeepSpeed configuration.
 
         Args:
-            finished_experiments (dcit): a dictionary of experiment id and experiment description.
+            metric (str): the key to read from each experiment's metrics file when
+                comparing configurations.
 
         Returns:
             The path to the result folder of the experiment with the optimal configuration.

deepspeed/autotuning/utils.py

@@ -231,7 +231,7 @@ def prune_config(config, ignored_keys=[]):
     """ Prunes the input configurations
 
     Args:
-        configs (dict): A configuration dictionary.
+        config (dict): A configuration dictionary.
         ignored_keys (list, optional): the keys of the sections to delete. Defaults to [].
 
     Returns:

deepspeed/inference/v2/engine_v2.py

@@ -165,7 +165,7 @@ def query(self, uid: int, max_request_tokens: int, max_request_blocks) -> Tuple[
             uid (int): The UID of the sequence (as tracked by the scheduling entity). If
                 this is a new sequence (with a UID unknown to the inference engine), then
                 an empty placeholder is created to pass to the occupancy logic.
-            n_tokens (int): The number of tokens to hypothetically send.
+            max_request_tokens (int): The number of tokens to hypothetically send.
 
         Returns:
             Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks
@@ -253,7 +253,7 @@ def serialize(self, save_path: str) -> None:
         Serialize the model to a file.
 
         Arguments:
-            path (str): Path to the file to serialize to.
+            save_path (str): Path to the file to serialize to.
         """
         param_file_name = make_param_filename(save_path, self._model.tp_rank, self._model.tp_size)
         metadata_file_name = make_metadata_filename(save_path, self._model.tp_rank, self._model.tp_size)

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_pre_norm.py

@@ -30,7 +30,6 @@ def __call__(self, z_res: torch.Tensor, z_hid: torch.Tensor, x_res: torch.Tensor
             x_res (torch.Tensor): Input residual.
             y_hid (torch.Tensor): Input hidden states.
             gamma (torch.Tensor): Gamma tensor.
-            beta (torch.Tensor): Beta tensor.
 
         Returns:
             output (torch.Tensor): Output tensor.

deepspeed/inference/v2/kernels/ragged_ops/embed/embed.py

@@ -26,7 +26,7 @@ class RaggedEmbeddingKernel(DSKernelBase):
     def __init__(self, embed_dtype: torch.dtype, token_dtype: torch.dtype, embed_dim: int) -> None:
         """
         Args:
-            fp_dtype (torch.dtype): Data type of the embedding table and output dtype.
+            embed_dtype (torch.dtype): Data type of the embedding table and output dtype.
                 Supported values are torch.float16, torch.bfloat16, and torch.float32.
             token_dtype (torch.dtype): Data type of the token ids. Supported values are
                 torch.int32 and torch.int64.

deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py

@@ -26,7 +26,6 @@ def __init__(self, head_size: int, n_q_heads: int, n_kv_heads: int, dtype: torch
         """
         Args:
             head_size: The size of the attention head.
-            q_ratio: Ratio of q heads to kv heads (for GQA)
             dtype: Data type for the input/output. Supported values are torch.float16 and torch.bfloat16.
         """
 

deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.py

@@ -43,7 +43,7 @@ def __call__(self, moe_input: torch.Tensor, expert_cumsum: torch.Tensor, mapped_
             moe_input (torch.Tensor): The direct input for the MoE GEMM of shape [n_tokens * n_top_k, hidden_size].
             expert_cumsum (torch.Tensor): The cumulative sum of the expert counts of shape [n_experts].
             mapped_slots (torch.Tensor): The index of the token in the expert's input of shape [n_tokens, n_top_k].
-            hidden_states (torch.Tensor): The hidden states of shape [n_tokens, hidden_size].
+            activations (torch.Tensor): The hidden states of shape [n_tokens, hidden_size].
             expert_counts (torch.Tensor): The number of tokens assigned to each expert of shape [n_experts].
             assignments (torch.Tensor): The expert assignments of shape [n_tokens, n_top_k].
             offsets (torch.Tensor): The offsets into the expert for a given token of shape [n_tokens, n_top_K].

deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.py

@@ -45,9 +45,9 @@ def __call__(self, expert_counts: torch.Tensor, scores: torch.Tensor, assignment
                 it is recommended to write to 0 during the MoE output remapping.
             scores (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place expert scaling
                 value.
-            expert_assignment (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place
+            assignments (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place
                 which expert a token has been assigned to.
-            expert_offset (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place which
+            offsets (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place which
                 offset within an experts group a token is.
             logits (torch.Tensor): Raw logits of gating function.
             batch (RaggedBatchWrapper): Batch information for ragged tensor.

deepspeed/inference/v2/model_implementations/sharding/attn_out.py

@@ -74,7 +74,7 @@ def attn_out_in_features(out_features: int,
     Helper to calculate the expected output projection dimension of a QKV projection matrix.
 
     Args:
-        in_features (int): The model dimension.
+        out_features (int): The model dimension.
         shard_rank (int): Which rank to return the corresponding size for.
         num_shards (int): The total number of shards the parameter is distributed across.
         head_size (int): The size of each attention head.

deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py

@@ -94,7 +94,7 @@ def forward(self,
             hidden_states (torch.Tensor): The hidden states from the model. This is the output of the
                 final layer of the model.
             vocab_embedding (torch.Tensor): The vocab embedding table.
-            raged_metadata (RaggedBatchWrapper): The ragged batch metadata.
+            ragged_metadata (RaggedBatchWrapper): The ragged batch metadata.
             gamma (Optional[torch.Tensor]): The gamma tensor for normalization.
             beta (Optional[torch.Tensor]): The beta tensor for normalization.
         """

deepspeed/inference/v2/ragged/kv_cache.py

@@ -67,11 +67,8 @@ def __init__(self,
         blocked KV-caches.
 
         Parameters:
-            config (KVCacheConfig): The configuration of the KV-cache.
-            slack (int): The amount of slack space to reserve in GPU memory for the cache.
-            enable_offload (bool): Whether to enable offloading of the cache to the host.
-            blocks (int): The number of blocks to pre-allocate for the cache. If this is set,
-                slack will be ignored.
+            configs (Tuple[KVCacheConfig, ...]): The configurations of the KV-cache.
+            offload (bool): Whether to enable offloading of the cache to the host.
         """
         self._configs = configs
         self._memory_config = memory_config

deepspeed/inference/v2/ragged/ragged_manager.py

@@ -59,8 +59,6 @@ def __init__(self,
         """
         The key
 
-        Parameters:
-            block_size (int): The number of tokens to allocate in each block.
         """
         self._config = config
         self._kv_configs = kv_configs

deepspeed/model_implementations/transformers/ds_transformer.py

@@ -26,8 +26,6 @@
 class DeepSpeedTransformerInference(nn.Module):
     """Initialize the DeepSpeed Transformer Layer.
         Arguments:
-            layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers,
-                layer_id will be 0,1,2...23 when each layer object is instantiated
             config: An object of DeepSpeedInferenceConfig
             mp_group: Model parallelism group initialized on the modeling side.
             quantize_scales: This argument groups all the layers' scales used for quantization

deepspeed/module_inject/module_quantize.py

@@ -16,7 +16,6 @@ def quantize_transformer_layer(orig_layer_impl, model, megatron=False, preln=Fal
         megatron (bool): megatron model-parallel implementation (this is supported for inference only)
         preln (bool): does the original layer implementation do pre or post layer norm?
 
-        Note: For Bert kind of models, we inject based on the DeepSpeed-Example models, if not setting huggingface flag.
 
     Returns:
         Updated nn.module with quantized transformer layers

deepspeed/moe/sharded_moe.py

@@ -629,7 +629,7 @@ class MOELayer(Base):
     Args:
         gate (torch.nn.Module):
             gate network
-        expert (torch.nn.Module):
+        experts (torch.nn.Module):
             expert network
     """
 

deepspeed/ops/lion/cpu_lion.py

@@ -32,7 +32,7 @@ def __init__(self, model_params, lr=1e-3, betas=(0.9, 0.999), weight_decay=0, fp
             betas (Tuple[float, float], optional): coefficients used for computing
                 running averages of gradient and its square. (default: (0.9, 0.999))
             weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
-            full_precision_optimizer_states: creates momentum and variance in full precision regardless of
+            fp32_optimizer_states: creates momentum and variance in full precision regardless of
                         the precision of the parameters (default: True)
         """
 

deepspeed/ops/sparse_attention/bert_sparse_self_attention.py

@@ -53,7 +53,7 @@ def forward(self, hidden_states, attention_mask):
 
         Arguments:
             hidden_states: required: hidden_states tensor of the bert model
-            attn_mask: required: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported
+            attention_mask: required: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported
 
         Return:
              context_layer: a dense tensor containing attention context

deepspeed/ops/sparse_attention/sparse_self_attention.py

@@ -105,8 +105,6 @@ def forward(self, query, key, value, rpe=None, key_padding_mask=None, attn_mask=
             rpe: optional: a tensor same dimension as x that is used as relative position embedding
             key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength)
             attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported
-            key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied
-            attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied
 
         Return:
              attn_output: a dense tensor containing attention context

deepspeed/ops/sparse_attention/sparsity_config.py

@@ -71,7 +71,6 @@ def __init__(self, num_heads, block=16, different_layout_per_head=False):
 
         Arguments:
              num_heads: required: an integer determining number of attention heads of the layer.
-             seq_len: required: an integer determining number of attention heads of the layer.
              different_layout_per_head: optional: this is just for the sake of consistency with other sparsity formats; can ignore it for DenseSparsityConfig
         """
 
@@ -269,7 +268,6 @@ def __init__(self,
              local_window_blocks: optional: a list of integers determining the number of blocks in each local attention window. It assumes first number determines # of blocks in the first local window, second the second window, ..., and the last number determines the number of blocks in the remaining local windows.
              global_block_indices: optional: a list of integers determining which blocks are considered as global attention. Given indices, determine the blocks that all other token blocks attend to and they attend to all other token blocks. Default value is only index 0. Notice that if global_block_end_indices parameter is set, this parameter is used as starting index of each global window.
              global_block_end_indices: optional: a list of integers determining end indices of global window blocks. By default this is not used. But if it is set, it must have the same size of global_block_indices parameter, and combining this two parameters, for each index i, blocks from global_block_indices[i] to global_block_end_indices[i] (exclusive) are considered as global attention.
-             num_global_blocks: optional: an integer determining how many consecutive blocks in a local window is used as the representative of the window for global attention.
              attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure.
              horizontal_global_attention: optional: a boolean determining if blocks that are global representative of a local window, also attend to all other blocks. This is valid only if attention type is `bidirectional`. Looking at the attention matrix, that means global attention not only includes the vertical blocks, but also horizontal blocks.
         """

deepspeed/ops/transformer/inference/diffusers_attention.py

@@ -99,8 +99,6 @@ def backward(ctx, grad_output, grad_output1, grad_output2, grad_output3):
 class DeepSpeedDiffusersAttention(nn.Module):
     """Initialize the DeepSpeed Transformer Layer.
         Arguments:
-            layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers,
-                layer_id will be 0,1,2...23 when each layer object is instantiated
             config: An object of DeepSpeedInferenceConfig
     """
     layer_id = 0

deepspeed/ops/transformer/inference/moe_inference.py

@@ -159,8 +159,6 @@ def forward(self, input, async_op=False):
 class DeepSpeedMoEInference(nn.Module):
     """Initialize the DeepSpeed MoE Transformer Layer.
         Arguments:
-            layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers,
-                layer_id will be 0,1,2...23 when each layer object is instantiated
             config: An object of DeepSpeedInferenceConfig
             mp_group: Model parallelism group initialized on the modeling side.
             quantize_scales: This argument groups all the layers' scales used for quantization

deepspeed/profiling/flops_profiler/profiler.py

@@ -61,7 +61,7 @@ class FlopsProfiler(object):
     To profile a trained model in inference, use the `get_model_profile` API.
 
     Args:
-        object (torch.nn.Module): The PyTorch model to profile.
+        model (torch.nn.Module): The PyTorch model to profile.
     """
 
     def __init__(self, model, ds_engine=None, recompute_fwd_factor=0.0):

deepspeed/runtime/fp16/loss_scaler.py

@@ -206,7 +206,6 @@ class DynamicLossScaler(LossScalerBase):
 
     Args:
         init_scale (float, optional, default=2**32):  Initial loss scale attempted by :class:`DynamicLossScaler.`
-        scale_factor (float, optional, default=2.0):  Factor used when adjusting the loss scale. If an overflow is encountered, the loss scale is readjusted to loss scale/``scale_factor``.  If ``scale_window`` consecutive iterations take place without an overflow, the loss scale is readjusted to loss_scale*``scale_factor``.
         scale_window (int, optional, default=1000):  Number of consecutive iterations without an overflow to wait before increasing the loss scale.
         consecutive_hysteresis (bool, optional, default=False): Whether to refill hysteresis if we reach an iteration that doesn't overflow
     """

deepspeed/runtime/fp16/onebit/adam.py

@@ -120,8 +120,6 @@ def step(self, closure=None, grads=None):
             output params (list of tensors, optional): A reduced precision copy
                 of the updated weights written out in addition to the regular
                 updated weights. Have to be of same type as gradients. (default: None)
-            scale (float, optional): factor to divide gradient tensor values
-                by before applying to weights. (default: 1)
         """
         loss = None
         if closure is not None:

deepspeed/runtime/fp16/onebit/zoadam.py

@@ -135,8 +135,6 @@ def step(self, closure=None, grads=None):
             output params (list of tensors, optional): A reduced precision copy
                 of the updated weights written out in addition to the regular
                 updated weights. Have to be of same type as gradients. (default: None)
-            scale (float, optional): factor to divide gradient tensor values
-                by before applying to weights. (default: 1)
         """
         loss = None
         if closure is not None:

deepspeed/runtime/pipe/engine.py

@@ -1361,7 +1361,7 @@ def load_module_state_dict(self,
         If ``state_dict`` is not ``None`` or a ``str``, we revert to ``super()`` expecting a ``dict``.
 
         Args:
-            state_dict (str, None): unused
+            checkpoint (dict): the checkpoint whose module state is loaded
             strict (bool, optional): Strict state loading. Defaults to True.
         """
         assert custom_load_fn is None, "custom_load_fn not supported w. pipeline parallelism"

deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py

@@ -255,9 +255,6 @@ def _process_selected_fp32_groups_grad(self, tensor, total_size, communication_d
         """
         Process gradients for selected columns in FP32 groups
 
-        Args:
-            param: The parameter to process
-            param_id: ID of the parameter
         """
 
         curr_size = 0

deepspeed/runtime/zero/stage3.py

@@ -2200,9 +2200,8 @@ def get_grad_norm_direct(self, gradients, params, norm_type=2):
         the gradients are modified in place.
 
         Arguments:
-            parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            params (Iterable[Tensor] or Tensor): an iterable of Tensors or a
                 single Tensor that will have gradients normalized
-            max_norm (float or int): max norm of the gradients
             norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
                 infinity norm.
 

deepspeed/runtime/zero/stage_1_and_2.py

@@ -1997,9 +1997,8 @@ def get_grad_norm_direct(self, gradients, params, norm_type=2):
         the gradients are modified in place.
 
         Arguments:
-            parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            params (Iterable[Tensor] or Tensor): an iterable of Tensors or a
                 single Tensor that will have gradients normalized
-            max_norm (float or int): max norm of the gradients
             norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
                 infinity norm.
 

deepspeed/runtime/zero/utils.py

@@ -150,7 +150,7 @@ def apply_to_tensors_only(function, value, warning_msg_fn=None):
     Apply `function` to every Tensor in `value`.
 
     Args:
-        functional: The function class to apply.
+        function: The function class to apply.
         value (Any): Target object to apply `function` to.
 
     Returns:

deepspeed/utils/groups.py

@@ -257,7 +257,7 @@ def _create_model_parallel(model_parallel_size_):
     Initialize model data parallel groups.
 
     Arguments:
-        model_parallel_size: number of GPUs used to parallelize model.
+        model_parallel_size_: number of GPUs used to parallelize model.
 
     Returns:
         Tuple of data parallel group and model parallel group