[Doc] Update fsdp_utils type annotation based on PEP guide (#509)

Author: lancerts

slime/backends/fsdp_utils/actor.py

@@ -1,7 +1,7 @@
 from argparse import Namespace
+from collections.abc import Iterable
 from contextlib import nullcontext
 from itertools import accumulate
-from typing import Iterable, Optional
 
 import ray
 import torch
@@ -118,7 +118,7 @@ def init(self, args: Namespace, role: str, wandb_run_id: str, with_ref: bool = F
         self.micro_step = 0
         return 0
 
-    def sleep(self, tags: Optional[str | Iterable[str]]) -> None:
+    def sleep(self, tags: str | Iterable[str] | None) -> None:
         """Pause CUDA memory for tagged tensors via torch_memory_saver.
 
         When offloading is enabled, this forwards tags to
@@ -138,7 +138,7 @@ def sleep(self, tags: Optional[str | Iterable[str]]) -> None:
                 for tag in tags:
                     torch_memory_saver.pause(tag)
 
-    def wake_up(self, tags: Optional[str | Iterable[str]]) -> None:
+    def wake_up(self, tags: str | Iterable[str] | None) -> None:
         """Resume CUDA memory for tagged tensors via torch_memory_saver.
 
         When offloading is enabled, this forwards tags to
@@ -591,7 +591,7 @@ def update_gpu_params_dict(self, params_dict: dict[str, torch.Tensor]) -> None:
         self.model.load_state_dict(gpu_state_dict, strict=True)
         torch.cuda.synchronize()
 
-    def load_ref_model(self, ref_load_path: Optional[str]) -> None:
+    def load_ref_model(self, ref_load_path: str | None) -> None:
         """Load reference model weights once and cache them on CPU.
 
         Parameters:
@@ -654,7 +654,7 @@ def gather_log_probs(logits: torch.Tensor, input_ids: torch.Tensor, rollout_temp
 
 
 def gather_log_probs_packed(
-    logits: torch.Tensor, input_ids: torch.Tensor, cu_seqlens: Optional[torch.Tensor | float] = None
+    logits: torch.Tensor, input_ids: torch.Tensor, cu_seqlens: torch.Tensor | float | None = None
 ) -> torch.Tensor:
     """Gather next-token log probabilities for packed sequences.
 

slime/backends/fsdp_utils/arguments.py

@@ -1,7 +1,6 @@
 import argparse
 import dataclasses
 from dataclasses import dataclass
-from typing import Optional
 
 import yaml
 
@@ -25,13 +24,13 @@ class FSDPArgs:
 
     # Logging
     wandb_project: str = "slime-fsdp"
-    wandb_run_name: Optional[str] = None
+    wandb_run_name: str | None = None
 
     # Precision
     gradient_checkpointing: bool = False
 
     # YAML bookkeeping
-    config: Optional[str] = None
+    config: str | None = None
 
 
 def parse_fsdp_cli(extra_args_provider=None):
@@ -40,7 +39,9 @@ def parse_fsdp_cli(extra_args_provider=None):
     for f in dataclasses.fields(FSDPArgs):
         if f.name == "config":
             continue
-        arg_type = f.type if f.type != Optional[str] else str
+
+        arg_type = str if f.type == (str | None) else f.type
+
         if arg_type is bool:
             parser.add_argument(f"--{f.name.replace('_', '-')}", action="store_true")
         else:

slime/backends/fsdp_utils/data_packing.py

@@ -1,25 +1,24 @@
 """Data packing utilities for FSDP backend to reduce padding overhead."""
 
 import math
-from typing import Dict, List, Optional
 
 import torch
 
 from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
 
 
 def pack_sequences(
-    tokens: List[List[int]],
-    loss_masks: List[List[int]],
-    rewards: List[float],
-    raw_rewards: List,
-    response_lengths: List[int],
-    advantages: List[float],
-    returns: List[float],
-    rollout_log_probs: Optional[List[List[float]]] = None,
-    max_tokens_per_gpu: Optional[int] = None,
-    num_packs: Optional[int] = None,
-) -> List[Dict]:
+    tokens: list[list[int]],
+    loss_masks: list[list[int]],
+    rewards: list[float],
+    raw_rewards: list,
+    response_lengths: list[int],
+    advantages: list[float],
+    returns: list[float],
+    rollout_log_probs: list[list[float]] | None = None,
+    max_tokens_per_gpu: int | None = None,
+    num_packs: int | None = None,
+) -> list[dict]:
     """
     Pack sequences into dense batches with cumulative sequence lengths.
 
@@ -99,7 +98,7 @@ def pack_sequences(
     return result
 
 
-def unpack_sequences(packed_batch: Dict) -> List[Dict]:
+def unpack_sequences(packed_batch: dict) -> list[dict]:
     """
     Unpack sequences from a packed batch.
 

slime/backends/fsdp_utils/update_weight_utils.py

@@ -1,6 +1,6 @@
 import socket
 from argparse import Namespace
-from typing import Mapping, Optional, Sequence
+from collections.abc import Mapping, Sequence
 
 import ray
 import torch
@@ -90,7 +90,7 @@ def __init__(
         self,
         args: Namespace,
         model: torch.nn.Module,
-        weights: Optional[Mapping[str, Mapping[str, torch.Tensor]]],
+        weights: Mapping[str, Mapping[str, torch.Tensor]] | None,
         full_params: bool = False,
     ) -> None:
         self.args = args
@@ -116,7 +116,7 @@ def __init__(
     def connect_rollout_engines(
         self,
         rollout_engines: Sequence[ActorHandle],
-        rollout_engine_lock: Optional[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
     ) -> None:
         """Attach rollout engines and create per-engine IPC (Gloo) groups.
 
@@ -297,7 +297,7 @@ def __init__(self, args: Namespace, model: torch.nn.Module) -> None:
     def connect_rollout_engines(
         self,
         rollout_engines: Sequence[ActorHandle],
-        rollout_engine_lock: Optional[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
     ) -> None:
         """On rank 0, initialize a temporary NCCL group for parameter broadcast."""
         self.rollout_engines = rollout_engines