[FEAT] Enhance DAPO for full dynamic sampling

areal/api/alloc_mode.py

@@ -839,6 +839,7 @@ def parse(self, expression: str):
             AllocationValidationError: When validation rules are violated
             ValueError: When parsing fails
         """
+
         try:
             tree = self.parser.parse(expression)
             transformer = _ParallelStrategyTransformer()

areal/api/cli_args.py

@@ -3,21 +3,19 @@
 import os
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
-from typing import Dict, List
 
 import uvloop
 import yaml
-
-from areal.utils.pkg_version import is_version_less
-
-uvloop.install()
 from hydra import compose as hydra_compose
 from hydra import initialize as hydra_init
 from hydra.core.global_hydra import GlobalHydra
 from omegaconf import MISSING, DictConfig, OmegaConf
 
 from areal.platforms import current_platform
 from areal.utils import name_resolve, pkg_version
+from areal.utils.pkg_version import is_version_less
+
+uvloop.install()
 
 
 @dataclass
@@ -129,11 +127,11 @@ class GenerationHyperparameters:
         default=1.0,
         metadata={"help": "Sampling temperature. Higher values increase diversity."},
     )
-    stop_token_ids: List[int] = field(
+    stop_token_ids: list[int] = field(
         default_factory=list,
         metadata={"help": "Stop generation when encountering these token IDs."},
     )
-    stop: List[str] | None = field(
+    stop: list[str] | None = field(
         default=None,
         metadata={
             "help": "One or multiple stop words. Generation will stop if one of these words is sampled."
@@ -232,7 +230,7 @@ class OptimizerConfig:
 class FSDPWrapPolicy:
     """Policy configuration for FSDP model layer wrapping. None defaults to wrapping transformer decoder layers defined by transformers."""
 
-    transformer_layer_cls_to_wrap: List[str] | None = field(
+    transformer_layer_cls_to_wrap: list[str] | None = field(
         default=None,
         metadata={"help": "A list of transformer layer names for FSDP to wrap."},
     )
@@ -310,7 +308,7 @@ class MegatronEngineConfig:
     recompute_method: str | None = "uniform"
     recompute_num_layers: int | None = 1
     distribute_saved_activations: bool | None = None
-    recompute_modules: List[str] | None = None
+    recompute_modules: list[str] | None = None
 
 
 @dataclass
@@ -378,7 +376,7 @@ class TrainEngineConfig:
     )
     lora_rank: int = field(default=32, metadata={"help": "lora rank"})
     lora_alpha: int = field(default=16, metadata={"help": "lora alpha"})
-    target_modules: List[str] = field(
+    target_modules: list[str] = field(
         default_factory=list,
         metadata={"help": "lora target_modules."},
     )
@@ -486,12 +484,10 @@ class PPOActorConfig(TrainEngineConfig):
         },
     )
     # Advanced Options
-    dynamic_sampling: bool = field(
-        default=False,
+    dynamic_sampling_strategy: str = field(
+        default="none",
         metadata={
-            "help": "Enable dynamic sampling (within DAPO). If enabled, groups with the same reward will be masked out. "
-            "Note that enabling this option will lead to variable batch sizes. If you want to use a constant batch size with dynamic filtering, "
-            "you should use the `should_accept` parameter in `rollout_batch` and `prepare_batch`."
+            "help": "Dynamic sampling strategy. Select from `none`, `dynamic` and `static`. See the doc for more details"
         },
     )
 
@@ -500,7 +496,7 @@ class PPOActorConfig(TrainEngineConfig):
         default=False,
         metadata={"help": "Log statistics for agent trajectories"},
     )
-    log_agent_stats_keys: List[str] = field(
+    log_agent_stats_keys: list[str] = field(
         default_factory=lambda: [],
         metadata={"help": "Keys for logging agent trajectory statistics"},
     )
@@ -574,7 +570,7 @@ def build_args(
         port,
         dist_init_addr: str | None = None,
     ):
-        args: Dict = conf_as_dict(vllm_config)
+        args: dict = conf_as_dict(vllm_config)
         args = dict(
             host=host,
             port=port,
@@ -608,11 +604,11 @@ def build_cmd(
             if v is None or v is False or v == "":
                 continue
             if v is True:
-                flags.append(f"--{k.replace('_','-')}")
+                flags.append(f"--{k.replace('_', '-')}")
             elif isinstance(v, list):
-                flags.append(f"--{k.replace('_','-')} {' '.join(map(str, v))}")
+                flags.append(f"--{k.replace('_', '-')} {' '.join(map(str, v))}")
             else:
-                flags.append(f"--{k.replace('_','-')} {v}")
+                flags.append(f"--{k.replace('_', '-')} {v}")
         return f"python3 -m areal.thirdparty.vllm.areal_vllm_server {' '.join(flags)}"
 
 
@@ -638,7 +634,7 @@ class SGLangConfig:
     enable_torch_compile: bool = False
     torch_compile_max_bs: int = 32
     cuda_graph_max_bs: int | None = None
-    cuda_graph_bs: List[int] | None = None
+    cuda_graph_bs: list[int] | None = None
     torchao_config: str = ""
     enable_nan_detection: bool = False
     enable_p2p_check: bool = False
@@ -667,8 +663,8 @@ class SGLangConfig:
     # lora
     enable_lora: bool | None = None
     max_lora_rank: int | None = None
-    lora_target_modules: List[str] | None = None
-    lora_paths: List[str] | None = None
+    lora_target_modules: list[str] | None = None
+    lora_paths: list[str] | None = None
     max_loaded_loras: int = 1
     max_loras_per_batch: int = 1
     lora_backend: str = "triton"
@@ -719,11 +715,11 @@ def build_cmd(
             if v is None or v is False or v == "":
                 continue
             if v is True:
-                flags.append(f"--{k.replace('_','-')}")
+                flags.append(f"--{k.replace('_', '-')}")
             elif isinstance(v, list):
-                flags.append(f"--{k.replace('_','-')} {' '.join(map(str, v))}")
+                flags.append(f"--{k.replace('_', '-')} {' '.join(map(str, v))}")
             else:
-                flags.append(f"--{k.replace('_','-')} {v}")
+                flags.append(f"--{k.replace('_', '-')} {v}")
         return f"python3 -m sglang.launch_server {' '.join(flags)}"
 
     @staticmethod
@@ -738,11 +734,11 @@ def build_args(
         node_rank: int = 0,
     ):
         # Map "all-linear" to "all"
-        args: Dict = conf_as_dict(sglang_config)
+        args: dict = conf_as_dict(sglang_config)
         if sglang_config.enable_multithread_load or sglang_config.enable_fast_load:
-            assert pkg_version.is_version_equal(
-                "sglang", "0.5.2"
-            ), f"Customized model loading requires exact SGLang version 0.5.2"
+            assert pkg_version.is_version_equal("sglang", "0.5.2"), (
+                "Customized model loading requires exact SGLang version 0.5.2"
+            )
             model_loader_extra_config = dict(
                 enable_multithread_load=sglang_config.enable_multithread_load,
                 enable_fast_load=sglang_config.enable_fast_load,
@@ -915,8 +911,8 @@ class WandBConfig:
     job_type: str | None = None
     group: str | None = None
     notes: str | None = None
-    tags: List[str] | None = None
-    config: Dict | None = None
+    tags: list[str] | None = None
+    config: dict | None = None
     id_suffix: str | None = "train"
 
 
@@ -926,7 +922,7 @@ class SwanlabConfig:
 
     project: str | None = None
     name: str | None = None
-    config: Dict | None = None
+    config: dict | None = None
     logdir: str | None = None
     mode: str | None = "disabled"
     api_key: str | None = os.getenv("SWANLAB_API_KEY", None)
@@ -1023,7 +1019,7 @@ class SchedulerConfig:
     endpoint: str = field(default="http://localhost:8081")
     deploy_mode: str = field(default="separation")
     functioncall_service_domain: str = field(default="http://localhost:8080")
-    reward_functioncall_config: Dict = field(default_factory=dict)
+    reward_functioncall_config: dict = field(default_factory=dict)
     reward_model_path: str = field(default="")
     reward_model_service_url: str = field(default="http://localhost:30000/classify")
 
@@ -1076,7 +1072,7 @@ class SlurmLauncherConfig:
         default="--mpi=pmi2 -K --chdir $PWD",
         metadata={"help": "Additional arguments to pass to the srun command."},
     )
-    additional_bash_cmds: List[str] | None = field(
+    additional_bash_cmds: list[str] | None = field(
         default=None,
         metadata={
             "help": "Additional bash commands to setup the container before running "
@@ -1244,7 +1240,7 @@ class PPOConfig(GRPOConfig):
     critic: PPOCriticConfig = field(default_factory=PPOCriticConfig)
 
 
-def parse_cli_args(argv: List[str]):
+def parse_cli_args(argv: list[str]):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--config", help="Path to the main configuration file", required=True
@@ -1277,7 +1273,7 @@ def to_structured_cfg(cfg, config_cls):
     return cfg
 
 
-def load_expr_config(argv: List[str], config_cls):
+def load_expr_config(argv: list[str], config_cls):
     cfg, config_file = parse_cli_args(argv)
     cfg = to_structured_cfg(cfg, config_cls=config_cls)
     cfg = OmegaConf.to_object(cfg)
@@ -1305,7 +1301,7 @@ def save_config(cfg, log_dir):
     os.makedirs(log_dir, exist_ok=True)
     config_save_path = os.path.join(log_dir, "config.yaml")
     with open(config_save_path, "w") as f:
-        config_dict: Dict = asdict(cfg)
+        config_dict: dict = asdict(cfg)
         yaml.dump(
             config_dict,
             f,

areal/api/workflow_api.py

@@ -1,6 +1,6 @@
 from __future__ import annotations  # noqa
 
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any
 
 from areal.experimental.openai.types import CompletionWithTokenLogpReward
 
@@ -9,10 +9,9 @@
 
 
 class RolloutWorkflow:
-
     async def arun_episode(
-        self, engine: "InferenceEngine", data: Dict[str, Any]
-    ) -> Dict[str, Any] | None | Dict[str, CompletionWithTokenLogpReward]:
+        self, engine: InferenceEngine, data: dict[str, Any]
+    ) -> dict[str, Any] | None | dict[str, CompletionWithTokenLogpReward]:
         """Run a single episode of the workflow.
 
         Note

areal/engine/ppo/actor.py

@@ -14,7 +14,6 @@
     split_padded_tensor_dict_into_mb_list,
 )
 from areal.utils.functional import (
-    dynamic_sampling,
     gather_logprobs,
     gather_logprobs_entropy,
     ppo_actor_loss_fn,
@@ -46,7 +45,6 @@ def __init__(self, config: PPOActorConfig, engine: TrainEngine):
         self.mask_no_eos_with_zero = config.mask_no_eos_with_zero
 
         self.temperature = config.temperature
-        self.dynamic_sampling = config.dynamic_sampling
 
     @torch.no_grad()
     def compute_logp(
@@ -164,8 +162,6 @@ def compute_advantages(self, data: Dict[str, Any]) -> None:
         data["logprobs"] = old_logp
 
     def ppo_update(self, data: Dict[str, Any]) -> List[Dict[str, float]]:
-        if self.dynamic_sampling and len(data["rewards"]) % self.group_size == 0:
-            data, sampling_stat = dynamic_sampling(data, self.group_size)
 
         attn_mask = data["attention_mask"]
         loss_mask = data["loss_mask"]

areal/utils/data.py

@@ -221,6 +221,42 @@ def concat_padded_tensors(
     return result
 
 
+def truncate_dict_to_batch_size(
+    data: Dict[str, Any], batch_size: int
+) -> Dict[str, Any]:
+    """Truncate a dictionary containing tensors and numeric values to specified batch size.
+
+    This function handles different value types:
+    - Tensors: take first batch_size elements along the first dimension
+    - Numeric values: keep as is (no truncation)
+    - Other types: keep as is (no truncation)
+
+    Args:
+        data: Dictionary to truncate
+        batch_size: Target batch size for truncation
+
+    Returns:
+        Truncated dictionary
+    """
+    if not data:
+        return {}
+
+    result = {}
+
+    for key, value in data.items():
+        if torch.is_tensor(value) and len(value.shape) > 0:
+            # For tensors, take first batch_size elements along first dimension
+            if value.shape[0] > batch_size:
+                result[key] = value[:batch_size]
+            else:
+                result[key] = value
+        else:
+            # For numeric values and other types, keep as is
+            result[key] = value
+
+    return result
+
+
 def unpack_sequence(
     x: torch.Tensor,
     cu_seqlens: Optional[torch.Tensor] = None,