[Doc] Include type annotation for cp_utils and model_provider (#468)

Author: lancerts

slime/backends/megatron_utils/cp_utils.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Callable, Union
 
 import torch
 import torch.distributed as dist
@@ -45,26 +45,26 @@ def get_logits_and_tokens_offset_with_cp(
 
 
 def get_sum_of_sample_mean(
-    total_lengths,
-    response_lengths,
-    loss_masks,
+    total_lengths: list[int],
+    response_lengths: list[int],
+    loss_masks: list[torch.Tensor],
     calculate_per_token_loss: bool = False,
-):
+) -> Callable[[torch.Tensor], torch.Tensor]:
     """
     Calculate correct sample mean for CP
     """
     cp_size = mpu.get_context_parallel_world_size()
     if cp_size == 1:
 
-        def sum_of_sample_mean(x: torch.Tensor):
+        def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor:
             return sum(
                 [
                     (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1)
                     for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks)
                 ]
             )
 
-        def sum_of_token(x: torch.Tensor):
+        def sum_of_token(x: torch.Tensor) -> torch.Tensor:
             return sum(
                 [(x_i * loss_mask_i).sum() for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks)]
             )
@@ -82,7 +82,7 @@ def sum_of_token(x: torch.Tensor):
             chunked_loss_masks.append(torch.cat([loss_mask_0, loss_mask_1], dim=0))
             cp_chunk_lengths.append(chunked_loss_masks[i].size(0))
 
-        def sum_of_sample_mean(x):
+        def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor:
             return sum(
                 [
                     (x_i * chunked_loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1)
@@ -92,7 +92,7 @@ def sum_of_sample_mean(x):
                 ]
             )
 
-        def sum_of_token(x: torch.Tensor):
+        def sum_of_token(x: torch.Tensor) -> torch.Tensor:
             return sum(
                 [
                     (x_i * chunked_loss_mask).sum()
@@ -103,7 +103,7 @@ def sum_of_token(x: torch.Tensor):
     return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token
 
 
-def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int):
+def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor:
     """
     Gather tensors across all ranks in the context parallel group.
     The first dimension of the output tensor will be the `response_length`.
@@ -122,7 +122,7 @@ def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length:
     chunk_1 = tensor[logits_offset[0][1] - logits_offset[0][0] :]
     assert chunk_1.shape[0] == logits_offset[1][1] - logits_offset[1][0]
 
-    def zero(len):
+    def zero(len: int) -> torch.Tensor:
         return torch.zeros(
             [len] + list(tensor.shape[1:]),
             dtype=tensor.dtype,
@@ -155,7 +155,7 @@ def zero(len):
     return full_tensor
 
 
-def slice_with_cp(tokens: torch.Tensor, pad_value):
+def slice_with_cp(tokens: torch.Tensor, pad_value: int) -> torch.Tensor:
     cp_rank = mpu.get_context_parallel_rank()
     cp_size = mpu.get_context_parallel_world_size()
 
@@ -172,7 +172,11 @@ def slice_with_cp(tokens: torch.Tensor, pad_value):
     return torch.cat([tokens[start_1:end_1], tokens[start_2:end_2]])
 
 
-def slice_log_prob_with_cp(log_prob: Union[list[float], torch.Tensor], total_length: int, response_length: int):
+def slice_log_prob_with_cp(
+    log_prob: Union[list[float], torch.Tensor],
+    total_length: int,
+    response_length: int,
+) -> Union[list[float], torch.Tensor]:
     assert len(log_prob) == response_length
 
     cp_size = mpu.get_context_parallel_world_size()

slime/backends/megatron_utils/model_provider.py

@@ -1,7 +1,8 @@
 # Adapt from https://github.com/NVIDIA/Megatron-LM/blob/b1efb3c7126ef7615e8c333432d76e08038e17ff/pretrain_gpt.py
+import argparse
 import inspect
 from contextlib import nullcontext
-from typing import Optional
+from typing import Literal, Optional
 
 import torch
 from megatron.core import tensor_parallel
@@ -12,19 +13,20 @@
     get_gpt_layer_with_transformer_engine_spec,
 )
 from megatron.core.transformer.spec_utils import import_module
+from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.training.arguments import core_transformer_config_from_args
 
 
 # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82
 class LinearForLastLayer(torch.nn.Linear):
     def __init__(
         self,
-        input_size,
-        output_size,
+        input_size: int,
+        output_size: int,
         *,
-        config,
-        bias=True,
-    ):
+        config: TransformerConfig,
+        bias: bool = True,
+    ) -> None:
         super().__init__(in_features=input_size, out_features=output_size, bias=bias)
         self.sequence_parallel = config.sequence_parallel
         if self.sequence_parallel:
@@ -36,19 +38,24 @@ def __init__(
 
     def forward(
         self,
-        input_,
-        weight=None,
-        runtime_gather_output=None,
-    ):
+        input_: torch.Tensor,
+        weight: Optional[torch.Tensor] = None,
+        runtime_gather_output: Optional[bool] = None,
+    ) -> tuple[torch.Tensor, None]:
         logits = super().forward(input_)
         logits = logits.float()
         if self.sequence_parallel:
             logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)
         return logits, None
 
 
-def get_model_provider_func(args, role: str = "actor"):
-    def model_provider(pre_process=True, post_process=True, vp_stage: Optional[int] = None) -> GPTModel:
+def get_model_provider_func(
+    args: argparse.Namespace,
+    role: Literal["actor", "critic"] = "actor",
+):
+    def model_provider(
+        pre_process: bool = True, post_process: bool = True, vp_stage: Optional[int] = None
+    ) -> GPTModel:
         """Builds the model.
 
         If you set the use_legacy_models to True, it will return the legacy GPT model and if not the mcore GPT model.
@@ -87,7 +94,7 @@ def oom_observer(device, alloc, device_alloc, device_free):
             torch._C._cuda_attach_out_of_memory_observer(oom_observer)
 
         # Experimental loading arguments from yaml
-        config = core_transformer_config_from_args(args)
+        config: TransformerConfig = core_transformer_config_from_args(args)
 
         if args.spec is not None:
             transformer_layer_spec = import_module(args.spec)
@@ -134,7 +141,7 @@ def oom_observer(device, alloc, device_alloc, device_free):
                 # Check if fp8_model_init supports preserve_high_precision_init_val
                 if "preserve_high_precision_init_val" in inspect.signature(fp8_model_init).parameters:
                     build_model_context_args["preserve_high_precision_init_val"] = True
-            except:
+            except Exception:
                 raise RuntimeError(
                     "--fp8-param-gather requires `fp8_model_init` from TransformerEngine, but not found."
                 )