diff --git a/deepspeed/autotuning/scheduler.py b/deepspeed/autotuning/scheduler.py
index 14e9541d03a3..150c880afbca 100755
--- a/deepspeed/autotuning/scheduler.py
+++ b/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.
diff --git a/deepspeed/autotuning/utils.py b/deepspeed/autotuning/utils.py
index c874bf631c46..0ea9529ad72a 100644
--- a/deepspeed/autotuning/utils.py
+++ b/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:
diff --git a/deepspeed/inference/v2/engine_v2.py b/deepspeed/inference/v2/engine_v2.py
index 4a358310377f..080a81b34b21 100644
--- a/deepspeed/inference/v2/engine_v2.py
+++ b/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)
diff --git a/deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_pre_norm.py b/deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_pre_norm.py
index 3b040d88b50f..e7c2cb100d5e 100644
--- a/deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_pre_norm.py
+++ b/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.
diff --git a/deepspeed/inference/v2/kernels/ragged_ops/embed/embed.py b/deepspeed/inference/v2/kernels/ragged_ops/embed/embed.py
index 0443ce3fdd8e..5a33f0e8a21c 100644
--- a/deepspeed/inference/v2/kernels/ragged_ops/embed/embed.py
+++ b/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.
diff --git a/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py b/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py
index aacbec0bd3ae..6ea537d2c1b9 100644
--- a/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py
+++ b/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.
         """
 
diff --git a/deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.py b/deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.py
index 7efcedb4e880..ddb1a1b04cb4 100644
--- a/deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.py
+++ b/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].
diff --git a/deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.py b/deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.py
index 72ba2b6019bb..ce1310b02292 100644
--- a/deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.py
+++ b/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.
diff --git a/deepspeed/inference/v2/model_implementations/sharding/attn_out.py b/deepspeed/inference/v2/model_implementations/sharding/attn_out.py
index ce7c105531ea..be52e3a4ae97 100644
--- a/deepspeed/inference/v2/model_implementations/sharding/attn_out.py
+++ b/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.
diff --git a/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py b/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py
index 36130902c665..80f02ea4ece0 100644
--- a/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py
+++ b/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.
         """
diff --git a/deepspeed/inference/v2/ragged/kv_cache.py b/deepspeed/inference/v2/ragged/kv_cache.py
index ceba3190b93c..b210d7373a3c 100644
--- a/deepspeed/inference/v2/ragged/kv_cache.py
+++ b/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
diff --git a/deepspeed/inference/v2/ragged/ragged_manager.py b/deepspeed/inference/v2/ragged/ragged_manager.py
index ecc3c52a5834..04bdd434516d 100644
--- a/deepspeed/inference/v2/ragged/ragged_manager.py
+++ b/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
diff --git a/deepspeed/model_implementations/transformers/ds_transformer.py b/deepspeed/model_implementations/transformers/ds_transformer.py
index 7d1f2a96b591..8cae74b69cdf 100644
--- a/deepspeed/model_implementations/transformers/ds_transformer.py
+++ b/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
diff --git a/deepspeed/module_inject/module_quantize.py b/deepspeed/module_inject/module_quantize.py
index 1f5b2f8a1d28..a031a71365a1 100755
--- a/deepspeed/module_inject/module_quantize.py
+++ b/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
diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py
index 1d706b6f94c4..1744f5edb3c1 100644
--- a/deepspeed/moe/sharded_moe.py
+++ b/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
     """
 
diff --git a/deepspeed/ops/lion/cpu_lion.py b/deepspeed/ops/lion/cpu_lion.py
index 03342a3fcd34..251dbcceaca3 100755
--- a/deepspeed/ops/lion/cpu_lion.py
+++ b/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)
         """
 
diff --git a/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py b/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py
index 37f065e48631..d426ac1883c4 100755
--- a/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py
+++ b/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
diff --git a/deepspeed/ops/sparse_attention/sparse_self_attention.py b/deepspeed/ops/sparse_attention/sparse_self_attention.py
index b673c4561902..30eb69e9759c 100644
--- a/deepspeed/ops/sparse_attention/sparse_self_attention.py
+++ b/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
diff --git a/deepspeed/ops/sparse_attention/sparsity_config.py b/deepspeed/ops/sparse_attention/sparsity_config.py
index b5d9be073bae..9029cf60d931 100644
--- a/deepspeed/ops/sparse_attention/sparsity_config.py
+++ b/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.
         """
diff --git a/deepspeed/ops/transformer/inference/diffusers_attention.py b/deepspeed/ops/transformer/inference/diffusers_attention.py
index 3c2340ccfc6f..613132b5d3af 100644
--- a/deepspeed/ops/transformer/inference/diffusers_attention.py
+++ b/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
diff --git a/deepspeed/ops/transformer/inference/moe_inference.py b/deepspeed/ops/transformer/inference/moe_inference.py
index 3a9785985d19..62ebf73919e9 100644
--- a/deepspeed/ops/transformer/inference/moe_inference.py
+++ b/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
diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py
index a07a773aa6b2..883326a22302 100644
--- a/deepspeed/profiling/flops_profiler/profiler.py
+++ b/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):
diff --git a/deepspeed/runtime/fp16/loss_scaler.py b/deepspeed/runtime/fp16/loss_scaler.py
index b0fc8d942d99..01fd9bd84319 100755
--- a/deepspeed/runtime/fp16/loss_scaler.py
+++ b/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
     """
diff --git a/deepspeed/runtime/fp16/onebit/adam.py b/deepspeed/runtime/fp16/onebit/adam.py
index fa817573f734..d2229ac1b7f4 100644
--- a/deepspeed/runtime/fp16/onebit/adam.py
+++ b/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:
diff --git a/deepspeed/runtime/fp16/onebit/zoadam.py b/deepspeed/runtime/fp16/onebit/zoadam.py
index 70282ec41714..0521b2a82922 100644
--- a/deepspeed/runtime/fp16/onebit/zoadam.py
+++ b/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:
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index b9e110e68ba7..b70b262b14b0 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/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"
diff --git a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py b/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py
index a4f2c249610e..f6d06c543e9d 100644
--- a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py
+++ b/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
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index e300a1d369c4..d19cce8cbfe2 100644
--- a/deepspeed/runtime/zero/stage3.py
+++ b/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.
 
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
index 85dd6ffb46b5..1d159a576a0e 100755
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/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.
 
diff --git a/deepspeed/runtime/zero/utils.py b/deepspeed/runtime/zero/utils.py
index acbce2a8a41d..627c2070875e 100755
--- a/deepspeed/runtime/zero/utils.py
+++ b/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:
diff --git a/deepspeed/utils/groups.py b/deepspeed/utils/groups.py
index ac0caa03fd16..7b387095ab0d 100644
--- a/deepspeed/utils/groups.py
+++ b/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