Revert "[style] minor: remove subclass" (#1441)

Author: zhuzilin

slime/backends/fsdp_utils/actor.py

@@ -31,7 +31,7 @@
 from ..training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, loss_function
 from . import checkpoint
 from .lr_scheduler import get_lr_scheduler
-from .parallel import create_fsdp_parallel_state
+from .parallel import FSDPParallelState
 from .update_weight_utils import UpdateWeightFromDistributed, UpdateWeightFromTensor
 
 logger = logging.getLogger(__name__)
@@ -55,7 +55,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
         super().init(args, role, with_ref)
 
         # Setup ParallelState for both CP and non-CP cases
-        self.parallel_state = create_fsdp_parallel_state(args)
+        self.parallel_state = FSDPParallelState(args)
 
         torch.manual_seed(args.seed)
 

slime/backends/fsdp_utils/parallel.py

@@ -12,47 +12,45 @@
 logger = logging.getLogger(__name__)
 
 
-def create_fsdp_parallel_state(args: Namespace) -> ParallelState:
-    """Create a ParallelState instance for FSDP configuration."""
-    world_size = dist.get_world_size()
-    rank = dist.get_rank()
-
-    cp_size = args.context_parallel_size
-    dp_rank = rank // cp_size
-    cp_rank = rank % cp_size
-
-    mesh = init_device_mesh("cuda", mesh_shape=(world_size // cp_size, cp_size), mesh_dim_names=("dp", "cp"))
-
-    logger.info(
-        f"[Rank {rank}] Device mesh (2D): world_size={world_size}, "
-        f"cp_size={cp_size}, dp_size={world_size // cp_size}"
-    )
-    logger.info(f"[Rank {rank}] Mesh shape: {mesh.shape}, " f"dp_rank={dp_rank}, cp_rank={cp_rank}")
-
-    # Setup Ring Flash Attention with CP group from mesh (only when cp_size > 1)
-    if cp_size > 1:
-        substitute_hf_flash_attn(mesh.get_group("cp"), heads_k_stride=1)
-        logger.info(f"[Rank {rank}] CP initialized via device mesh")
-    else:
-        logger.info(f"[Rank {rank}] Pure DP mode (cp_size=1)")
-
-    parallel_state = ParallelState(
-        dp_rank=dp_rank,
-        dp_src_rank=dp_rank // world_size,
-        dp_size=world_size // cp_size,
-        cp_rank=cp_rank,
-        cp_size=cp_size,
-        dp_cp_rank=rank,
-        dp_cp_size=world_size,
-        dp_group=mesh.get_group("dp"),
-        dp_cp_group=dist.group.WORLD,
-        dp_cp_group_gloo=get_gloo_group(),
-        cp_group=mesh.get_group("cp"),
-        tp_size=1,
-        tp_rank=0,
-        tp_group=dist.new_group([rank]),
-    )
-
-    parallel_state.dp_mesh = mesh["dp"]
-
-    return parallel_state
+class FSDPParallelState(ParallelState):
+    def __init__(self, args: Namespace):
+        super().__init__()
+
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+
+        self.cp_size = args.context_parallel_size
+        self.dp_size = world_size // self.cp_size
+        self.dp_cp_size = world_size
+
+        self.dp_rank = rank // self.cp_size
+        self.cp_rank = rank % self.cp_size
+        self.dp_cp_rank = rank
+        self.dp_src_rank = self.dp_rank // world_size
+
+        self.tp_size = 1
+        self.tp_rank = 0
+        self.tp_group = dist.new_group([rank])
+
+        self.mesh = init_device_mesh(
+            "cuda", mesh_shape=(world_size // self.cp_size, self.cp_size), mesh_dim_names=("dp", "cp")
+        )
+        self.dp_mesh = self.mesh["dp"]
+
+        self.dp_group = self.mesh.get_group("dp")
+        self.cp_group = self.mesh.get_group("cp")
+        self.dp_cp_group = dist.group.WORLD
+        self.dp_cp_group_gloo = get_gloo_group()
+
+        logger.info(
+            f"[Rank {rank}] Device mesh (2D): world_size={world_size}, "
+            f"cp_size={self.cp_size}, dp_size={world_size // self.cp_size}"
+        )
+        logger.info(f"[Rank {rank}] Mesh shape: {self.mesh.shape}, " f"dp_rank={self.dp_rank}, cp_rank={self.cp_rank}")
+
+        # Setup Ring Flash Attention with CP group from mesh (only when cp_size > 1)
+        if self.cp_size > 1:
+            substitute_hf_flash_attn(self.cp_group, heads_k_stride=1)
+            logger.info(f"[Rank {rank}] CP initialized via device mesh")
+        else:
+            logger.info(f"[Rank {rank}] Pure DP mode (cp_size=1)")

slime/backends/megatron_utils/actor.py

@@ -33,7 +33,7 @@
 from .checkpoint import load_checkpoint
 from .initialize import init, is_megatron_main_rank
 from .model import forward_only, initialize_model_and_optimizer, save, train
-from .parallel import create_megatron_parallel_state
+from .parallel import MegatronParallelState
 from .update_weight.common import named_params_and_buffers
 from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed
 from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor
@@ -92,7 +92,7 @@ def init(
             args, role
         )
 
-        self.parallel_state = create_megatron_parallel_state(model=self.model)
+        self.parallel_state = MegatronParallelState(model=self.model)
 
         if role == "critic":
             if self.args.offload_train:

slime/backends/megatron_utils/model.py

@@ -28,10 +28,9 @@
 from ..training_utils.data import DataIterator, get_batch
 from ..training_utils.log_utils import aggregate_forward_results, aggregate_train_losses, log_train_step
 from ..training_utils.loss import loss_function
-from ..training_utils.parallel import ParallelState
 from .checkpoint import load_checkpoint, save_checkpoint
 from .model_provider import get_model_provider_func
-from .parallel import get_packed_seq_params
+from .parallel import MegatronParallelState, get_packed_seq_params
 
 logger = logging.getLogger(__name__)
 
@@ -157,7 +156,7 @@ def forward_only(
     model: Sequence[DDP],
     data_iterator: Sequence[DataIterator],
     num_microbatches: Sequence[int],
-    parallel_state: ParallelState,
+    parallel_state: MegatronParallelState,
     store_prefix: str = "",
 ) -> dict[str, list[torch.Tensor]]:
     """Run forward passes only and collect non-loss outputs (e.g., logprobs).
@@ -297,7 +296,7 @@ def train_one_step(
     optimizer: MegatronOptimizer,
     opt_param_scheduler: OptimizerParamScheduler,
     num_microbatches: int,
-    parallel_state: ParallelState,
+    parallel_state: MegatronParallelState,
 ) -> tuple[dict[str, float], float]:
     """Execute a single pipeline-parallel training step.
 
@@ -482,7 +481,7 @@ def train(
     opt_param_scheduler: OptimizerParamScheduler,
     data_iterator: Sequence[DataIterator],
     num_microbatches: Sequence[int],
-    parallel_state: ParallelState,
+    parallel_state: MegatronParallelState,
 ) -> None:
     """Run training over a rollout consisting of multiple steps.
 

slime/backends/megatron_utils/parallel.py

@@ -12,44 +12,45 @@
 logger = logging.getLogger(__name__)
 
 
-def create_megatron_parallel_state(
-    model: torch.nn.Module | Sequence[torch.nn.Module] | None = None,
-) -> ParallelState:
-    vpp_size_value = mpu.get_virtual_pipeline_model_parallel_world_size()
-    if vpp_size_value is None:
-        vpp_size = 1
-        microbatch_group_size_per_vp_stage = None
-    elif vpp_size_value > 1:
-        assert model is not None
-        model_to_check = model[0] if isinstance(model, Sequence) else model
-        config = get_model_config(model_to_check)
-        vpp_size = vpp_size_value
-        microbatch_group_size_per_vp_stage = config.microbatch_group_size_per_vp_stage
-    else:
-        vpp_size = 1
-        microbatch_group_size_per_vp_stage = None
+class MegatronParallelState(ParallelState):
+    """
+    ParallelState for Megatron backend, initialized from mpu module.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module | Sequence[torch.nn.Module] | None = None,
+    ):
+        super().__init__()
+
+        self.dp_rank = mpu.get_data_parallel_rank(with_context_parallel=False)
+        self.cp_rank = mpu.get_context_parallel_rank()
+        self.tp_rank = mpu.get_tensor_model_parallel_rank()
+        self.dp_cp_rank = mpu.get_data_parallel_rank(with_context_parallel=True)
+        self.dp_src_rank = mpu.get_data_parallel_src_rank(with_context_parallel=True)
+
+        self.dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False)
+        self.dp_cp_size = mpu.get_data_parallel_world_size(with_context_parallel=True)
+        self.cp_size = mpu.get_context_parallel_world_size()
+        self.tp_size = mpu.get_tensor_model_parallel_world_size()
 
-    parallel_state = ParallelState(
-        dp_rank=mpu.get_data_parallel_rank(with_context_parallel=False),
-        dp_src_rank=mpu.get_data_parallel_src_rank(with_context_parallel=True),
-        dp_size=mpu.get_data_parallel_world_size(with_context_parallel=False),
-        cp_rank=mpu.get_context_parallel_rank(),
-        cp_size=mpu.get_context_parallel_world_size(),
-        dp_cp_rank=mpu.get_data_parallel_rank(with_context_parallel=True),
-        dp_cp_size=mpu.get_data_parallel_world_size(with_context_parallel=True),
-        dp_group=mpu.get_data_parallel_group(with_context_parallel=False),
-        dp_cp_group=mpu.get_data_parallel_group(with_context_parallel=True),
-        dp_cp_group_gloo=mpu.get_data_parallel_group_gloo(with_context_parallel=True),
-        cp_group=mpu.get_context_parallel_group(),
-        tp_size=mpu.get_tensor_model_parallel_world_size(),
-        tp_rank=mpu.get_tensor_model_parallel_rank(),
-        tp_group=mpu.get_tensor_model_parallel_group(),
-        is_pp_last_stage=mpu.is_pipeline_last_stage(),
-        vpp_size=vpp_size,
-        microbatch_group_size_per_vp_stage=microbatch_group_size_per_vp_stage,
-    )
+        self.dp_group = mpu.get_data_parallel_group(with_context_parallel=False)
+        self.dp_cp_group = mpu.get_data_parallel_group(with_context_parallel=True)
+        self.dp_cp_group_gloo = mpu.get_data_parallel_group_gloo(with_context_parallel=True)
+        self.cp_group = mpu.get_context_parallel_group()
+        self.tp_group = mpu.get_tensor_model_parallel_group()
 
-    return parallel_state
+        self.is_pp_last_stage = mpu.is_pipeline_last_stage()
+        vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()
+        if vpp_size is None:
+            self.vpp_size = 1
+            self.microbatch_group_size_per_vp_stage = None
+        elif vpp_size > 1:
+            assert model is not None
+            model_to_check = model[0] if isinstance(model, Sequence) else model
+            config = get_model_config(model_to_check)
+            self.vpp_size = vpp_size
+            self.microbatch_group_size_per_vp_stage = config.microbatch_group_size_per_vp_stage
 
 
 def get_packed_seq_params(batch: dict[str, torch.Tensor], args: Namespace) -> PackedSeqParams:

slime/backends/training_utils/parallel.py

@@ -4,10 +4,6 @@
 
 @dataclass
 class ParallelState:
-    """Core parallel state shared across all backends.
-    Required by the general training utils.
-    """
-
     dp_rank: int
     dp_src_rank: int
     dp_size: int
@@ -22,6 +18,13 @@ class ParallelState:
     tp_size: int
     tp_rank: int
     tp_group: dist.ProcessGroup | None
-    is_pp_last_stage: bool = True
-    vpp_size: int | None = 1
-    microbatch_group_size_per_vp_stage: int | None = None
+    dp_mesh: dist.DeviceMesh | None
+    cp_mesh: dist.DeviceMesh | None
+    is_pp_last_stage: bool
+    vpp_size: int | None
+    microbatch_group_size_per_vp_stage: int | None
+
+    def __init__(self):
+        self.vpp_size = 1
+        self.microbatch_group_size_per_vp_stage = None
+        self.is_pp_last_stage = True