modify type check ignores

Author: ITcarrot

RoundPipe/RoundPipe.py

@@ -1,6 +1,6 @@
 """The RoundPipe model wrapper and execution runtime."""
 
-from beartype.typing import * # type: ignore[reportWildcardImportFromLibrary]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 from beartype import beartype
 import traceback
 import copy
@@ -66,7 +66,7 @@ def __init__(self,
 
         self.num_layers: int = len(self.layers)
         self.layer_workload: List[float] = []
-        self.layer_gradient_ready_events: List[torch.cuda.Event] = [torch.cuda.Event() for _ in range(self.num_layers)] # type: ignore[reportAttributeAccessIssue]
+        self.layer_gradient_ready_events: List[torch.cuda.Event] = [torch.cuda.Event() for _ in range(self.num_layers)] # pyright: ignore[reportAttributeAccessIssue]
         for layer in self.layers:
             self.layer_workload.append(get_model_size(layer))
         self.model_timer: ModelTimer = ModelTimer(self.num_layers)
@@ -76,13 +76,13 @@ def __init__(self,
             pinned_tensor = torch.empty_like(parm.data, dtype=torch.float16 if use_fp16 and parm.is_floating_point() else None, pin_memory=True)
             pinned_tensor.copy_(parm.data)
             parm.data = pinned_tensor
-            parm.data_cpu = pinned_tensor # type: ignore[attr-defined]
+            parm.data_cpu = pinned_tensor # pyright: ignore[reportAttributeAccessIssue]
         for buffer in tqdm.tqdm(self.model.buffers(), total=sum(1 for _ in self.model.buffers()),
                                 desc=f'Roundpipe: Process buffers in {self.name}', leave=False):
             pinned_tensor = torch.empty_like(buffer.data, dtype=torch.float16 if use_fp16 and buffer.is_floating_point() else None, pin_memory=True)
             pinned_tensor.copy_(buffer.data)
             buffer.data = pinned_tensor
-            buffer.data_cpu = pinned_tensor # type: ignore[attr-defined]
+            buffer.data_cpu = pinned_tensor # pyright: ignore[reportAttributeAccessIssue]
 
         self.RoundPipe_initialized: bool = True
 
@@ -157,7 +157,7 @@ def forward(self, *args: Any,
                 tag = backward_schedule_simulator.get_next_tag()
                 for context in reversed(run_context):
                     tag, output_require_grad_idx, *output_require_grad \
-                        = RoundPipeMicrobatchBackward.apply(context, batch, tag, *context.flatten_inputs[0]) # type: ignore
+                        = RoundPipeMicrobatchBackward.apply(context, batch, tag, *context.flatten_inputs[0]) # pyright: ignore[reportGeneralTypeIssues]
                     for idx, item in zip(output_require_grad_idx, output_require_grad):
                         batch.flatten_states[context.microbatch_id][idx] = item
                 backward_schedule_simulator.update_current_tag(tag)
@@ -166,7 +166,7 @@ def forward(self, *args: Any,
                 # ensuring gradients to be calculated even if inputs do not require grad.
                 all_inputs = [item for batch_context in run_context
                             for item in batch_context.flatten_inputs[0]]
-                output_require_grad_idx, *output_require_grad = RoundPipeBatchedBackward.apply(run_context, batch, gradient_anchor, *all_inputs) # type: ignore
+                output_require_grad_idx, *output_require_grad = RoundPipeBatchedBackward.apply(run_context, batch, gradient_anchor, *all_inputs) # pyright: ignore[reportGeneralTypeIssues]
                 for (batch_idx, idx), item in zip(output_require_grad_idx, output_require_grad):
                     batch.flatten_states[batch_idx][idx] = item
 
@@ -207,7 +207,7 @@ def train_iter(self, input_args: Tuple[Any, ...] = (),
             context.input_backward_events = batch.backward_events[batch_idx]
 
         all_inputs = [item for batch_input in batch.flatten_states for item in batch_input]
-        input_backward_handle: torch.Tensor = RoundPipeInputBackward.apply(run_context, *all_inputs) # type: ignore
+        input_backward_handle: torch.Tensor = RoundPipeInputBackward.apply(run_context, *all_inputs) # pyright: ignore[reportAssignmentType]
 
         for layer_group_id in range(len(execute_plan.fwd_plan)):
             device = get_next_device()
@@ -225,7 +225,9 @@ def train_iter(self, input_args: Tuple[Any, ...] = (),
         if isinstance(batch.loss_list[0], torch.Tensor):
             loss = torch.zeros_like(batch.loss_list[0], device=torch.device('cpu'))
             for batch_loss in batch.loss_list:
-                loss = loss + batch_loss.cpu() # type: ignore[reportOperatorIssue]
+                assert isinstance(batch_loss, torch.Tensor), \
+                    "Inconsistent loss types across microbatches."
+                loss = loss + batch_loss.cpu()
         else:
             loss = [torch.zeros_like(t, device=torch.device('cpu')) for t in batch.loss_list[0]]
             for batch_loss in batch.loss_list:

RoundPipe/RunConfig.py

@@ -1,6 +1,6 @@
 """Runtime configuration objects shared across RoundPipe components."""
 
-from beartype.typing import * # type: ignore[reportWildcardImportFromLibrary]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 from beartype import beartype
 
 import torch

RoundPipe/batch.py

@@ -9,7 +9,7 @@
     avg_reducer: Predefined reducer that averages scalar losses across microbatches.
 """
 
-from beartype.typing import * # type: ignore[reportWildcardImportFromLibrary]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 import warnings
 
 import torch
@@ -90,14 +90,14 @@ def get_avg_reducer_args() -> Tuple[torch.Tensor, Callable[[torch.Tensor, torch.
         a reducer callable that updates the running average.
     """
     init_val = torch.tensor(0)
-    init_val.roundpipe_avg_reducer_sum = torch.tensor(0) # type: ignore[reportAttributeAccessIssue]
-    init_val.roundpipe_avg_reducer_count = 0 # type: ignore[reportAttributeAccessIssue]
+    init_val.roundpipe_avg_reducer_sum = torch.tensor(0) # pyright: ignore[reportAttributeAccessIssue]
+    init_val.roundpipe_avg_reducer_count = 0 # pyright: ignore[reportAttributeAccessIssue]
     def reduce(reduced_val: torch.Tensor, new_val: torch.Tensor) -> torch.Tensor:
-        val_sum = reduced_val.roundpipe_avg_reducer_sum + new_val # type: ignore[reportAttributeAccessIssue]
-        val_count = reduced_val.roundpipe_avg_reducer_count + 1 # type: ignore[reportAttributeAccessIssue]
+        val_sum = reduced_val.roundpipe_avg_reducer_sum + new_val # pyright: ignore[reportAttributeAccessIssue]
+        val_count = reduced_val.roundpipe_avg_reducer_count + 1 # pyright: ignore[reportAttributeAccessIssue]
         new_reduced = val_sum / val_count
-        new_reduced.roundpipe_avg_reducer_sum = val_sum # type: ignore[reportAttributeAccessIssue]
-        new_reduced.roundpipe_avg_reducer_count = val_count # type: ignore[reportAttributeAccessIssue]
+        new_reduced.roundpipe_avg_reducer_sum = val_sum # pyright: ignore[reportAttributeAccessIssue]
+        new_reduced.roundpipe_avg_reducer_count = val_count # pyright: ignore[reportAttributeAccessIssue]
         return new_reduced
     return init_val, reduce
 avg_reducer: _CustomReducer = _CustomReducer(*get_avg_reducer_args())
@@ -153,7 +153,7 @@ def __init__(self, args: Tuple, kwargs: Dict[str, Any],
         for batch_idx, args_kwargs in enumerate(zip(args_list, kwargs_list)):
             forward_event: Set[torch.cuda.Event] = set()
             backward_event: Set[torch.cuda.Event] = set()
-            cpu_tensor_backward_event: torch.cuda.Event = torch.cuda.Event() # type: ignore[reportAssignmentType]
+            cpu_tensor_backward_event: torch.cuda.Event = torch.cuda.Event() # pyright: ignore[reportAssignmentType]
             flatten_input, flatten_spec = tree_flatten(args_kwargs)
             for idx, item in enumerate(flatten_input):
                 if isinstance(item, torch.Tensor):

RoundPipe/device.py

@@ -5,7 +5,7 @@
     cur_device: Index tracking the next device to schedule work on.
 """
 
-from beartype.typing import * # type: ignore[reportWildcardImport]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 from enum import Enum
 import threading
 import itertools
@@ -142,10 +142,10 @@ def __init__(self, id: int, device: torch.device):
         self.id: int = id
         self.device: torch.device = device
 
-        self.param_upstream: torch.cuda.Stream = torch.cuda.Stream(device) # type: ignore[reportAttributeAccessIssue]
-        self.upstream: torch.cuda.Stream = torch.cuda.Stream(device) # type: ignore[reportAttributeAccessIssue]
+        self.param_upstream: torch.cuda.Stream = torch.cuda.Stream(device) # pyright: ignore[reportAttributeAccessIssue]
+        self.upstream: torch.cuda.Stream = torch.cuda.Stream(device) # pyright: ignore[reportAttributeAccessIssue]
         self.compute_stream: torch.cuda.Stream = torch.cuda.default_stream(self.device)
-        self.downstream: torch.cuda.Stream = torch.cuda.Stream(device) # type: ignore[reportAttributeAccessIssue]
+        self.downstream: torch.cuda.Stream = torch.cuda.Stream(device) # pyright: ignore[reportAttributeAccessIssue]
         self.mem_manager: InterStreamMemManager = InterStreamMemManager(
             self.param_upstream, self.upstream, self.compute_stream, self.downstream
         )
@@ -174,7 +174,7 @@ def mark_upload(self) -> None:
         Returns:
             The recorded event is appended to ``upload_mark``.
         """
-        event: torch.cuda.Event = torch.cuda.Event() # type: ignore[reportAttributeAccessIssue]
+        event: torch.cuda.Event = torch.cuda.Event() # pyright: ignore[reportAssignmentType]
         with torch.cuda.stream(self.upstream):
             event.record()
         self.upload_mark.append(event)

RoundPipe/models/__init__.py

@@ -9,7 +9,7 @@
         to their corresponding wrapper module paths.
 """
 
-from beartype.typing import * # type: ignore[reportWildcardImportFromLibrary]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 import importlib
 
 from ..RoundPipe import RoundPipe

RoundPipe/models/function.py

@@ -1,4 +1,4 @@
-from beartype.typing import * # type: ignore[reportWildcardImportFromLibrary]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 import types
 import math
 
@@ -28,7 +28,7 @@ def forward(ctx: Any, logits: torch.Tensor, labels: torch.Tensor,
         return loss
 
     @staticmethod
-    def backward(ctx: Any, grad_loss: torch.Tensor) -> Tuple[Optional[torch.Tensor], ...]: # type: ignore[override]
+    def backward(ctx: Any, grad_loss: torch.Tensor) -> Tuple[Optional[torch.Tensor], ...]: # pyright: ignore[reportIncompatibleMethodOverride]
         logits, labels, num_items_in_batch = ctx.saved_tensors
         grad_logits = torch.empty_like(logits)
 

RoundPipe/models/qwen3.py

@@ -1,4 +1,4 @@
-from beartype.typing import * # type: ignore[reportWildcardImportFromLibrary]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 import warnings
 
 import torch
@@ -24,7 +24,7 @@ def forward(
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         past_key_values: Optional[Any] = None,
-        inputs_embeds: Optional[torch.Tensor] = None, # type: ignore[reportRedeclaration]
+        inputs_embeds: Optional[torch.Tensor] = None, # pyright: ignore[reportRedeclaration]
         labels: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,

RoundPipe/profile.py

@@ -4,7 +4,7 @@
     PROFILER_TYPE: Type of profiler detected from environment variables.
 """
 
-from beartype.typing import * # type: ignore[reportWildcardImportFromLibrary]
+from beartype.typing import * # pyright: ignore[reportWildcardImportFromLibrary]
 
 import os
 import contextlib