diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 117642a246..803ccb6542 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -77,6 +77,11 @@ class NormConfig: default=1, metadata={"help": "Group size for group-level normalization"} ) + @property + def uses_group_statistics(self) -> bool: + """Whether normalization derives statistics from prompt groups.""" + return self.mean_level == "group" or self.std_level == "group" + def __post_init__(self): """Validate normalization configuration.""" valid_levels = {"batch", "group", None} @@ -1675,6 +1680,18 @@ class PPOActorConfig(TrainEngineConfig): default=None, metadata={"help": "Normalization configuration for advantages."} ) + # Partial rollout groups + min_usable_group_size: int | None = field( + default=None, + metadata={ + "help": "Minimum usable rollout slots a prompt group must keep to stay " + "trainable when some slots fail or are filtered. None derives the " + "minimum from reward_norm/adv_norm: 2 when either uses group " + "statistics (1 for a singleton target group), else 1. Only the v1 " + "rollout path consumes this option." + }, + ) + # KL Control kl_ctl: float = field(default=0.1, metadata={"help": "KL divergence coefficient"}) kl_estimator: str = field( @@ -1766,6 +1783,34 @@ class PPOActorConfig(TrainEngineConfig): metadata={"help": "Maximum number of new tokens to generate"}, ) + def _uses_group_statistics(self) -> bool: + for normalization in (self.reward_norm, self.adv_norm): + if normalization is None: + continue + if isinstance(normalization, (dict, DictConfig)): + if ( + normalization.get("mean_level") == "group" + or normalization.get("std_level") == "group" + ): + return True + elif normalization.uses_group_statistics: + return True + return False + + def resolve_min_usable_group_size(self, target_group_size: int) -> int: + """Minimum usable rollout slots a group must keep to stay trainable. + + An explicit ``min_usable_group_size`` wins. Otherwise group-relative + normalization needs at least two group members before partial groups + become a hazard; a singleton target group is complete by definition, + so it keeps the minimum of one. + """ + if self.min_usable_group_size is not None: + return self.min_usable_group_size + if self._uses_group_statistics(): + return min(2, target_group_size) + return 1 + def should_compute_prox_logp(self) -> bool: """Determine if forward pass is needed for proximal log-probabilities. @@ -1781,6 +1826,20 @@ def should_compute_prox_logp(self) -> bool: def __post_init__(self): """Validate PPO actor configuration.""" + if self.min_usable_group_size is not None: + if self.min_usable_group_size < 1: + raise ValueError( + "min_usable_group_size must be a positive integer, " + f"got {self.min_usable_group_size}" + ) + if self.min_usable_group_size < 2 and self._uses_group_statistics(): + raise ValueError( + "min_usable_group_size must be at least 2 when reward_norm or " + "adv_norm uses group statistics: a lone surviving rollout has " + "no group peers to normalize against. Leave it unset to derive " + "the minimum instead." + ) + reward_norm = self.reward_norm if isinstance(reward_norm, (dict, DictConfig)): reward_mean_level = reward_norm.get("mean_level") diff --git a/areal/api/engine_api.py b/areal/api/engine_api.py index deb73ef5a4..9bac97bc8b 100644 --- a/areal/api/engine_api.py +++ b/areal/api/engine_api.py @@ -202,6 +202,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Submit a batch of requests and wait for results. @@ -221,12 +222,16 @@ def rollout_batch( group_size : int, optional Number of times to run the workflow per input and concatenate results. Default is 1 (no grouping). + min_usable_group_size : int, optional + Estimator-owned minimum number of usable logical rollout slots. Must be + between 1 and ``group_size``. Default is 1. Returns ------- list[dict[str, Any]] A list of trajectory dictionaries, one per accepted rollout result. - Each trajectory contains tensors with shape [group_size, seqlen, ...]. + Each trajectory contains tensors whose leading dimension is the number + of usable slots, between ``min_usable_group_size`` and ``group_size``. """ raise NotImplementedError() @@ -241,6 +246,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Prepare a batch of data for training from a dataloader. @@ -261,6 +267,9 @@ def prepare_batch( If True, enables dynamic batch sizing. The method will stop collecting when (accepted + rejected) >= batch_size, returning only accepted results. This results in variable-sized batches of valid data. Default is False. + min_usable_group_size : int, optional + Estimator-owned minimum number of usable logical rollout slots. Must be + between 1 and ``group_size``. Default is 1. Returns ------- @@ -741,6 +750,7 @@ def submit( is_eval: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> int: """Submit a request to the inference engine and return immediately. @@ -773,6 +783,9 @@ def submit( is_eval : bool, optional Whether this is an evaluation workflow. Affects variables like trajectory dump path and statistics keys. By default False. + min_usable_group_size : int, optional + Estimator-owned minimum number of usable logical rollout slots. Must be + between 1 and ``group_size``. Default is 1. Returns ------- @@ -846,6 +859,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Submit a batch of requests to the inference engine and wait for the results. @@ -873,6 +887,9 @@ def rollout_batch( group_size : int, optional Number of times to run the workflow per input and concatenate results. Default is 1 (no grouping). + min_usable_group_size : int, optional + Estimator-owned minimum number of usable logical rollout slots. Must be + between 1 and ``group_size``. Default is 1. Returns ------- @@ -893,6 +910,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Asynchronously submit and wait until a full batch is ready with controlled staleness. @@ -902,9 +920,10 @@ def prepare_batch( This method caches an internal data generator on the first call. The ``dataloader``, ``workflow``, ``workflow_kwargs``, ``group_size``, - and ``should_accept_fn`` parameters are captured at the first invocation - and reused in all subsequent calls. Passing different arguments in - later calls will **not** take effect. + ``reward_normalization``, ``drop_incomplete_group``, + ``min_usable_group_size``, and ``should_accept_fn`` parameters are captured + at the first invocation and reused in all subsequent calls. Passing + different arguments in later calls will **not** take effect. If you need to switch configurations mid-training, consider: @@ -936,6 +955,9 @@ def prepare_batch( If True, enables dynamic batch sizing. The method will stop collecting when (accepted + rejected) >= batch_size, returning only accepted results. This results in variable-sized batches of valid data. Default is False. + min_usable_group_size : int, optional + Estimator-owned minimum number of usable logical rollout slots. Must be + between 1 and ``group_size``. Default is 1. Returns ------- diff --git a/areal/api/workflow_api.py b/areal/api/workflow_api.py index c16a9765af..9e996d0d85 100644 --- a/areal/api/workflow_api.py +++ b/areal/api/workflow_api.py @@ -22,6 +22,22 @@ async def arun_episode( ---- Returning `None` implies that this trajectory is rejected and will not be used for training. + When the resolved ``min_usable_group_size`` is at least 2 — derived + from group-relative normalization (``mean_level='group'`` or + ``std_level='group'`` in ``actor.reward_norm`` / ``actor.adv_norm``), + or set explicitly via ``actor.min_usable_group_size`` — and rollouts + are grouped (``n_samples >= 2``), each ``arun_episode`` call must contribute + exactly one training sample: a tensor dict with batch size 1, or a + dict containing a single interaction. Returning multiple samples per + episode (e.g. one row per turn, or tree-search branches) raises a + non-retryable ``WorkflowContractError`` that terminates training, + because group statistics would otherwise treat same-episode rows as + independent group members. Ungrouped rollouts (``n_samples=1``) + install no group wrapper and are not checked; a multi-sample episode + is then normalized as its own group of same-episode rows. To train + multi-sample workflows, use batch-level normalization or merge each + episode into a single sequence. + See concrete example implementations under the `areal/workflow` directory. Parameters diff --git a/areal/engine/core/__init__.py b/areal/engine/core/__init__.py index 499d02fd22..1334457ceb 100644 --- a/areal/engine/core/__init__.py +++ b/areal/engine/core/__init__.py @@ -4,12 +4,14 @@ from areal.engine.core.train_engine import ( aggregate_eval_losses, + compute_microbatch_loss_weight, compute_total_loss_weight, reorder_and_pad_outputs, ) __all__ = [ "aggregate_eval_losses", + "compute_microbatch_loss_weight", "compute_total_loss_weight", "reorder_and_pad_outputs", ] diff --git a/areal/engine/core/train_engine.py b/areal/engine/core/train_engine.py index 9bc83c2b77..f416818ad2 100644 --- a/areal/engine/core/train_engine.py +++ b/areal/engine/core/train_engine.py @@ -14,6 +14,7 @@ from areal.infra.platforms import current_platform from areal.utils.data import ( + TRANSPORT_DUMMY_KEY, MicroBatchList, pad_and_stack_tensors_along_first_dim, reorder_list, @@ -21,12 +22,29 @@ ) __all__ = [ + "compute_microbatch_loss_weight", "compute_total_loss_weight", "aggregate_eval_losses", "reorder_and_pad_outputs", ] +def compute_microbatch_loss_weight( + microbatch: dict[str, Any], + loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], +) -> torch.Tensor: + """Return zero without invoking an objective on transport-only data.""" + if microbatch.get(TRANSPORT_DUMMY_KEY) is not True: + return loss_weight_fn(microbatch) + reference = next( + (value for value in microbatch.values() if isinstance(value, torch.Tensor)), + None, + ) + if reference is None: + raise ValueError("Transport micro-batch does not contain a tensor") + return torch.zeros((), dtype=torch.float32, device=reference.device) + + def compute_total_loss_weight( mb_list: MicroBatchList, loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], @@ -52,7 +70,9 @@ def compute_total_loss_weight( The total loss weight (scalar tensor) after all_reduce. """ total_weight = ( - torch.stack([loss_weight_fn(mb) for mb in mb_list.mbs]) + torch.stack( + [compute_microbatch_loss_weight(mb, loss_weight_fn) for mb in mb_list.mbs] + ) .sum() .detach() .clone() @@ -138,7 +158,11 @@ def reorder_and_pad_outputs( The processed outputs, padded and stacked along batch dimension. """ res = aggregate_fn(outputs) + semantic_batch_size = len(output_seqlens) + output_seqlens = [*output_seqlens, *([1] * mb_list.transport_dummy_count)] seqlens = [output_seqlens[i] for i in mb_list.forward_indices] unpacked = unpack_sequence(res, lens=seqlens, dim=0) reordered = reorder_list(unpacked, mb_list.backward_indices) + if mb_list.transport_dummy_count: + reordered = reordered[:semantic_batch_size] return pad_and_stack_tensors_along_first_dim(reordered) diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index d4f1d343e1..1a3f9cee9f 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -60,6 +60,7 @@ from areal.api.io_struct import DeviceRuntimeInfo from areal.engine.core import ( aggregate_eval_losses, + compute_microbatch_loss_weight, compute_total_loss_weight, reorder_and_pad_outputs, ) @@ -582,6 +583,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: self._check_rollout_engine_connected() return self.rollout_coordinator.rollout_batch( @@ -589,6 +591,7 @@ def rollout_batch( workflow=workflow, workflow_kwargs=workflow_kwargs, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -603,6 +606,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: self._check_rollout_engine_connected() return self.rollout_coordinator.prepare_batch( @@ -611,6 +615,7 @@ def prepare_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, dynamic_bs=dynamic_bs, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, @@ -782,7 +787,9 @@ def train_batch( input_batched, _ = self._normalize_batch_input(input_) # Step 1: Prepare micro-batches - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) # Step 2: Compute total loss weight total_loss_weight = compute_total_loss_weight( @@ -822,7 +829,9 @@ def eval_batch( input_batched, _ = self._normalize_batch_input(input_) # Step 1: Prepare micro-batches - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) # Step 2: Compute total loss weight total_loss_weight = compute_total_loss_weight( @@ -880,7 +889,9 @@ def forward_batch( batch_size = len(output_seqlens) # Step 2: Prepare micro-batches - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) # Step 3: Forward using process_output_fn callback, collecting results outputs: list[torch.Tensor] = [] @@ -1878,7 +1889,12 @@ def _load_optimizer_state(self, path: str): self.optimizer.load_state_dict(optimizer_state_dict) dist.barrier(group=self.cpu_group) - def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: + def _prepare_mb_list( + self, + input_: dict[str, Any], + *, + allow_transport_padding: bool = False, + ) -> MicroBatchList: assert "attention_mask" in input_ and "input_ids" in input_ input_ = input_.copy() @@ -1941,7 +1957,12 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: else: input_ = amend_position_ids(input_) - mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec) + mb_list = split_padded_tensor_dict_into_mb_list( + input_, + self.config.mb_spec, + group=self.data_parallel_group if allow_transport_padding else None, + allow_transport_padding=allow_transport_padding, + ) mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] mb_list = pad_mb_list( mb_list, @@ -2150,7 +2171,7 @@ def _compute_logprobs_and_loss( loss_multiplier: float = 1.0, ) -> torch.Tensor: """Compute logprobs/entropy and return scaled loss.""" - local_weight = loss_weight_fn(ctx.mb_input) + local_weight = compute_microbatch_loss_weight(ctx.mb_input, loss_weight_fn) if local_weight == 0: return logits.mean() * 0.0 diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 7beb16ef69..40e25b4173 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -51,6 +51,7 @@ from areal.api.io_struct import DeviceRuntimeInfo from areal.engine.core import ( aggregate_eval_losses, + compute_microbatch_loss_weight, compute_total_loss_weight, reorder_and_pad_outputs, ) @@ -929,6 +930,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: self._check_rollout_engine_connected() return self.rollout_coordinator.rollout_batch( @@ -936,6 +938,7 @@ def rollout_batch( workflow=workflow, workflow_kwargs=workflow_kwargs, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -950,6 +953,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: self._check_rollout_engine_connected() return self.rollout_coordinator.prepare_batch( @@ -958,6 +962,7 @@ def prepare_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, dynamic_bs=dynamic_bs, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, @@ -1334,7 +1339,9 @@ def train_batch( input_batched, _ = self._normalize_batch_input(input_) # Step 1: Prepare micro-batches - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) # Step 2: Compute total loss weight. # Use DP+CP group: after CP all-gather each rank computes the full-sequence @@ -1395,7 +1402,9 @@ def eval_batch( input_batched, _ = self._normalize_batch_input(input_) # Step 1: Prepare micro-batches - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) # Step 2: Compute total loss weight (DP+CP, see train_batch comment). total_loss_weight = compute_total_loss_weight( @@ -1454,7 +1463,9 @@ def forward_batch( batch_size = len(output_seqlens) # Step 2: Prepare micro-batches - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) # Step 3: Forward using Megatron's pipeline function, collecting results outputs: list[torch.Tensor] = [] @@ -2640,7 +2651,12 @@ def _load_model_from_hf(self, path: str) -> None: fp8_direct_convert=self.fp8_direct_convert, ) - def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: + def _prepare_mb_list( + self, + input_: dict[str, Any], + *, + allow_transport_padding: bool = False, + ) -> MicroBatchList: assert "attention_mask" in input_ and "input_ids" in input_ # Parallel sizes pp_size = self.parallel_strategy.pipeline_parallel_size @@ -2691,6 +2707,7 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: input_, mb_spec, group=mpu.get_data_parallel_group(), + allow_transport_padding=allow_transport_padding, ) mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] # NOTE: Pad micro-batches to: @@ -2752,7 +2769,7 @@ def _compute_logprobs_and_loss( total_loss_weight: torch.Tensor, loss_multiplier: float = 1.0, ) -> torch.Tensor: - local_weight = loss_weight_fn(inputs) + local_weight = compute_microbatch_loss_weight(inputs, loss_weight_fn) if local_weight == 0: connected_output = ( output.logprobs if isinstance(output, ChunkedLMHeadOutput) else output diff --git a/areal/engine/sglang_remote.py b/areal/engine/sglang_remote.py index cd54bc8a27..de11203fa6 100644 --- a/areal/engine/sglang_remote.py +++ b/areal/engine/sglang_remote.py @@ -34,6 +34,7 @@ from areal.infra import RemoteInfEngine, RolloutController, WorkflowExecutor from areal.infra.platforms import current_platform from areal.infra.utils.launcher import TRITON_CACHE_PATH +from areal.infra.workflow_executor import WorkflowTaskResult from areal.utils import perf_tracer, stats_tracker from areal.utils.logging import getLogger from areal.utils.network import format_host_for_url @@ -569,6 +570,7 @@ def submit( proxy_addr: str | None = None, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> int: """Submit a request to the inference engine.""" return self._engine.submit( @@ -577,6 +579,7 @@ def submit( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, task_id=task_id, callback_addr=callback_addr, is_eval=is_eval, @@ -597,6 +600,11 @@ def wait_for_task( """Wait for a specific task to complete by task_id.""" return self._engine.wait_for_task(task_id, timeout, raise_timeout) + def _wait_for_task_result( + self, task_id: int, timeout: float | None = None, raise_timeout: bool = True + ) -> WorkflowTaskResult | None: + return self._engine._wait_for_task_result(task_id, timeout, raise_timeout) + def rollout_batch( self, data: list[dict[str, Any]], @@ -605,6 +613,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> dict[str, Any]: """Submit a batch of requests and wait for results. @@ -616,6 +625,7 @@ def rollout_batch( workflow=workflow, workflow_kwargs=workflow_kwargs, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -630,6 +640,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ): """Asynchronously submit and wait until a full batch is ready.""" return self._engine.prepare_batch( @@ -638,6 +649,7 @@ def prepare_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, dynamic_bs=dynamic_bs, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, diff --git a/areal/engine/vllm_remote.py b/areal/engine/vllm_remote.py index 2d76930a9d..a33d2eb72b 100644 --- a/areal/engine/vllm_remote.py +++ b/areal/engine/vllm_remote.py @@ -33,6 +33,7 @@ from areal.infra import RemoteInfEngine, RolloutController, WorkflowExecutor from areal.infra.platforms import current_platform from areal.infra.utils.launcher import TRITON_CACHE_PATH +from areal.infra.workflow_executor import WorkflowTaskResult from areal.utils import logging, perf_tracer, stats_tracker from areal.utils.network import format_host_for_url from areal.utils.vllm_response import parse_vllm_generation_response @@ -436,6 +437,7 @@ def submit( proxy_addr: str | None = None, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> int: """Submit a request to the inference engine.""" return self._engine.submit( @@ -444,6 +446,7 @@ def submit( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, task_id=task_id, callback_addr=callback_addr, is_eval=is_eval, @@ -464,6 +467,11 @@ def wait_for_task( """Wait for a specific task to complete by task_id.""" return self._engine.wait_for_task(task_id, timeout, raise_timeout) + def _wait_for_task_result( + self, task_id: int, timeout: float | None = None, raise_timeout: bool = True + ) -> WorkflowTaskResult | None: + return self._engine._wait_for_task_result(task_id, timeout, raise_timeout) + def rollout_batch( self, data: list[dict[str, Any]], @@ -472,6 +480,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> dict[str, Any]: """Submit a batch of requests and wait for results. @@ -483,6 +492,7 @@ def rollout_batch( workflow=workflow, workflow_kwargs=workflow_kwargs, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -497,6 +507,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ): """Asynchronously submit and wait until a full batch is ready.""" return self._engine.prepare_batch( @@ -505,6 +516,7 @@ def prepare_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, dynamic_bs=dynamic_bs, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, diff --git a/areal/experimental/engine/archon_engine.py b/areal/experimental/engine/archon_engine.py index b1fd64d1f6..94bd004c81 100644 --- a/areal/experimental/engine/archon_engine.py +++ b/areal/experimental/engine/archon_engine.py @@ -41,6 +41,7 @@ ) from areal.engine.core.train_engine import ( aggregate_eval_losses, + compute_microbatch_loss_weight, compute_total_loss_weight, reorder_and_pad_outputs, ) @@ -534,7 +535,9 @@ def train_batch( input_batched, _ = self._normalize_batch_input(input_) - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) total_loss_weight = compute_total_loss_weight( mb_list, loss_weight_fn, self.data_parallel_group @@ -571,7 +574,9 @@ def eval_batch( input_batched, _ = self._normalize_batch_input(input_) - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) total_loss_weight = compute_total_loss_weight( mb_list, loss_weight_fn, self.data_parallel_group @@ -629,7 +634,9 @@ def forward_batch( assert output_seqlens is not None batch_size = len(output_seqlens) - mb_list = self._prepare_mb_list(input_batched).to(self.device) + mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to( + self.device + ) def process_output( logits: torch.Tensor, ctx_dict: dict[str, Any] @@ -692,6 +699,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Perform rollout using connected inference engine.""" self._check_rollout_engine_connected() @@ -700,6 +708,7 @@ def rollout_batch( workflow=workflow, workflow_kwargs=workflow_kwargs, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -714,6 +723,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Prepare batch from dataloader with rollout.""" self._check_rollout_engine_connected() @@ -723,6 +733,7 @@ def prepare_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, dynamic_bs=dynamic_bs, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, @@ -1183,7 +1194,12 @@ def _normalize_batch_input( return concat_batch(input_) return input_, None - def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: + def _prepare_mb_list( + self, + input_: dict[str, Any], + *, + allow_transport_padding: bool = False, + ) -> MicroBatchList: assert "attention_mask" in input_ and "input_ids" in input_ input_ = input_.copy() @@ -1211,7 +1227,7 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: stages_per_rank = len(self.pp_stages) num_total_stages = pp_size * stages_per_rank n_seqs = input_["attention_mask"].shape[0] - if n_seqs < num_total_stages: + if n_seqs < num_total_stages and not allow_transport_padding: raise RuntimeError( f"Pipeline parallelism requires at least {num_total_stages} " f"sequences (pp_size={pp_size} * stages_per_rank=" @@ -1227,7 +1243,12 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: else: mb_spec = self.config.mb_spec - mb_list = split_padded_tensor_dict_into_mb_list(input_, mb_spec) + mb_list = split_padded_tensor_dict_into_mb_list( + input_, + mb_spec, + group=self.data_parallel_group if allow_transport_padding else None, + allow_transport_padding=allow_transport_padding, + ) mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] # LCM ensures page-aligned memory and exact CP slicing without extra padding. @@ -1275,7 +1296,7 @@ def _compute_logprobs_and_loss( loss_multiplier: float = 1.0, ) -> torch.Tensor: """Compute logprobs/entropy and return scaled loss.""" - local_weight = loss_weight_fn(ctx.mb_input) + local_weight = compute_microbatch_loss_weight(ctx.mb_input, loss_weight_fn) if local_weight == 0: return logits.mean() * 0.0 diff --git a/areal/infra/controller/rollout_controller.py b/areal/infra/controller/rollout_controller.py index 8878358dfe..14a84e2899 100644 --- a/areal/infra/controller/rollout_controller.py +++ b/areal/infra/controller/rollout_controller.py @@ -45,7 +45,15 @@ from areal.utils.perf_tracer import trace_perf from ..staleness_manager import StalenessManager -from ..workflow_executor import BatchTaskDispatcher, TaskIdGenerator +from ..workflow_executor import ( + BatchTaskDispatcher, + TaskIdGenerator, + WorkflowContractFailure, + WorkflowTaskResult, + get_workflow_result_error, + unwrap_workflow_result, + validate_rollout_group_sizes, +) logger = logging.getLogger("RolloutController") @@ -61,17 +69,12 @@ class _RemoteRolloutTaskInput: should_accept_fn: str | None is_eval: bool = False group_size: int = 1 + min_usable_group_size: int = 1 proxy_addr: str | None = None reward_normalization: bool = False drop_incomplete_group: bool = False -@dataclass -class _RemoteRolloutResult: - task_id: int - trajectory: dict[str, Any] - - class RolloutController: def __init__( self, @@ -107,7 +110,7 @@ def __init__( # Dispatcher will be initialized in initialize() after staleness_manager is ready self._dispatcher: ( - BatchTaskDispatcher[_RemoteRolloutTaskInput, _RemoteRolloutResult] | None + BatchTaskDispatcher[_RemoteRolloutTaskInput, WorkflowTaskResult] | None ) = None # HTTP callback server @@ -216,12 +219,13 @@ def initialize( # Create and initialize the dispatcher qsize = self.config.queue_size or max_concurrent_rollouts * 16 self._dispatcher = BatchTaskDispatcher[ - _RemoteRolloutTaskInput, _RemoteRolloutResult + _RemoteRolloutTaskInput, WorkflowTaskResult ]( max_queue_size=qsize, task_factory=self._create_submit_callback, staleness_manager=self._staleness_manager, enable_tracing=self.config.enable_rollout_tracing, + terminal_error_fn=get_workflow_result_error, ) # Initialize the dispatcher's async task runner self._dispatcher.initialize(logger=logger) @@ -848,7 +852,7 @@ def _rollout_stats(self) -> str: ) def _create_submit_callback(self, pending_task: _RemoteRolloutTaskInput): - async def _submit_then_wait() -> _RemoteRolloutResult | None: + async def _submit_then_wait() -> WorkflowTaskResult | None: # Choose worker via round-robin worker, rank = self._choose_worker() engine_name = self._engine_name(rank) @@ -880,6 +884,7 @@ async def _submit_then_wait() -> _RemoteRolloutResult | None: http_timeout=self.config.request_timeout, is_eval=pending_task.is_eval, group_size=pending_task.group_size, + min_usable_group_size=pending_task.min_usable_group_size, task_id=task_id, callback_addr=f"http://{self.callback_addr}/callback/rollout_complete", proxy_addr=proxy_addr, @@ -895,7 +900,7 @@ async def _submit_then_wait() -> _RemoteRolloutResult | None: # Fetch the result result = await self.scheduler.async_call_engine( worker.id, - "wait_for_task", + "_wait_for_task_result", engine_name=engine_name, task_id=engine_task_id, timeout=0.1, # A short time to prevent blocking other requests @@ -903,14 +908,19 @@ async def _submit_then_wait() -> _RemoteRolloutResult | None: http_timeout=self.config.request_timeout, ) - traj = result + if isinstance(result, WorkflowContractFailure): + manager.on_rollout_rejected() + return result + + traj = result.trajectory if result is not None else None if traj is not None: manager.on_rollout_accepted() if self.config.enable_rollout_tracing: logger.info( f"Finish and accept rollout. {self._rollout_stats()}" ) - return _RemoteRolloutResult(task_id=task_id, trajectory=traj) + assert result is not None + return result manager.on_rollout_rejected() if self.config.enable_rollout_tracing: @@ -949,7 +959,9 @@ def submit( proxy_addr: str | None = None, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> int: + validate_rollout_group_sizes(group_size, min_usable_group_size) workflow_str = self._resolve_workflow_str(workflow) should_accept_fn = self._resolve_should_accept_fn(should_accept_fn) if workflow_kwargs is None: @@ -968,6 +980,7 @@ def submit( task_id=task_id, is_eval=is_eval, group_size=group_size, + min_usable_group_size=min_usable_group_size, proxy_addr=proxy_addr, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, @@ -981,7 +994,10 @@ def wait( self, count: int, timeout: float | None = None, raise_timeout: bool = True ) -> list[dict[str, Any] | None]: # Delegate to dispatcher and extract trajectories - results = self.dispatcher.wait_results(count, timeout, raise_timeout) + results = [ + unwrap_workflow_result(result) + for result in self.dispatcher.wait_results(count, timeout, raise_timeout) + ] # Log and trace if self.config.enable_rollout_tracing: logger.info("Rollout results are ready!") @@ -998,6 +1014,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: perf_tracer.instant( "rollout_controller.rollout_batch", @@ -1011,6 +1028,7 @@ def rollout_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -1029,6 +1047,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Prepare a batch with controlled staleness. @@ -1038,6 +1057,7 @@ def prepare_batch( See :meth:`~areal.api.engine_api.InferenceEngine.prepare_batch` for parameters. """ + validate_rollout_group_sizes(group_size, min_usable_group_size) workflow_str = self._resolve_workflow_str(workflow) if workflow_kwargs is None: workflow_kwargs = {} @@ -1052,6 +1072,7 @@ def task_input_generator(): should_accept_fn=should_accept_fn, task_id=self._task_id_generator.next(), group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -1061,9 +1082,14 @@ def task_input_generator(): # Delegate to dispatcher assert dataloader.batch_size is not None - results = self.dispatcher.active_submit_and_wait( - self.data_generator, batch_size=dataloader.batch_size, dynamic_bs=dynamic_bs - ) + results = [ + unwrap_workflow_result(result) + for result in self.dispatcher.active_submit_and_wait( + self.data_generator, + batch_size=dataloader.batch_size, + dynamic_bs=dynamic_bs, + ) + ] # Return list of trajectories trajectories = [r.trajectory if r is not None else None for r in results] @@ -1247,7 +1273,7 @@ def staleness_manager(self): @property def dispatcher( self, - ) -> BatchTaskDispatcher[_RemoteRolloutTaskInput, _RemoteRolloutResult]: + ) -> BatchTaskDispatcher[_RemoteRolloutTaskInput, WorkflowTaskResult]: """Get the task dispatcher, ensuring initialization has been called.""" if self._dispatcher is None: raise RuntimeError( diff --git a/areal/infra/controller/train_controller.py b/areal/infra/controller/train_controller.py index 45f1c4b340..85e033b42d 100644 --- a/areal/infra/controller/train_controller.py +++ b/areal/infra/controller/train_controller.py @@ -94,6 +94,10 @@ def _dispatch_tensors( token_weights = [_item_weight(d) for d in item_list] n_groups = n // group_size + if n_groups < dp_size: + raise ValueError( + f"item group count ({n_groups}) must be at least dp_size ({dp_size})" + ) group_weights = [ sum(token_weights[g * group_size + k] for k in range(group_size)) @@ -127,10 +131,14 @@ def _pad_eval_batch( ) -> tuple[Any, ...]: """Pad the first tensor-like arg to a multiple of ``dp_size * group_size``. - Called before dispatch for explicit evaluation controller paths so that - ``balanced_greedy_partition`` always receives a divisible input. - Dummy items have zero attention/loss masks and contribute nothing - to metrics or loss. + Called before dispatch for explicit evaluation controller paths: eval + batches may be smaller than ``dp_size`` or not group-aligned, and padding + guarantees every DP rank receives at least one full group + (``_dispatch_tensors`` rejects ``n_groups < dp_size`` and + ``n % group_size != 0``). Training batches are intentionally not padded — + incomplete rollout groups make them ragged, and engines absorb ragged + shards via transport padding. Dummy items have zero attention/loss masks + and contribute nothing to metrics or loss. """ result = list(args) pad_target = dp_size * group_size @@ -773,6 +781,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: return self.rollout.prepare_batch( dataloader=dataloader, @@ -780,6 +789,7 @@ def prepare_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, dynamic_bs=dynamic_bs, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, @@ -794,6 +804,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: return self.rollout.rollout_batch( data=data, @@ -801,6 +812,7 @@ def rollout_batch( workflow_kwargs=workflow_kwargs, should_accept_fn=should_accept_fn, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) diff --git a/areal/infra/dist_rollout.py b/areal/infra/dist_rollout.py index 6fa86894f8..666ad3fd8d 100644 --- a/areal/infra/dist_rollout.py +++ b/areal/infra/dist_rollout.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import time from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -9,7 +10,12 @@ from areal.api import InferenceEngine, TrainEngine, WorkflowLike from areal.infra.platforms import current_platform +from areal.infra.workflow_executor import ( + MAX_CONSECUTIVE_EMPTY_ROLLOUT_ROUNDS, + ROLLOUT_COLLECTION_STALL_TIMEOUT_SECONDS, +) from areal.utils.data import ( + all_gather_ragged_tensor_container, all_gather_tensor_container, broadcast_tensor_container, split_and_unpad_tensor, @@ -26,6 +32,70 @@ class RedistributedData: group_indices: list[list[int]] +def _all_gather_ragged_trajectory_lists( + trajectories: list[dict[str, Any]], group=None +) -> list[list[dict[str, Any]]]: + """Gather variable-length trajectory lists without semantic placeholders.""" + world_size = dist.get_world_size(group) + lengths: list[int | None] = [None] * world_size + dist.all_gather_object(lengths, len(trajectories), group=group) + if all(length == len(trajectories) for length in lengths): + return all_gather_tensor_container(trajectories, group=group) + return all_gather_ragged_tensor_container(trajectories, group=group) + + +def _pack_gathered_trajectories( + all_gathered: list[list[dict[str, Any]]], + *, + world_size: int, + rank: int, + packing_algorithm: str = "ffd", +) -> RedistributedData: + """Pack gathered trajectories into per-rank shards without communicating. + + Every participating rank holds the same ``all_gathered`` input, so this + function is deterministic across ranks and issues no collectives; failures + raised here are symmetric across ranks. + """ + # Flatten the list of lists into a single list of trajectories + all_data = [] + for traj_list in all_gathered: + all_data.extend(traj_list) + + if len(all_data) < world_size: + raise RuntimeError( + f"Cannot redistribute {len(all_data)} trainable trajectory groups " + f"across {world_size} data-parallel ranks" + ) + + # Compute sequence lengths for load balancing + seqlens = [d["attention_mask"].sum().item() for d in all_data] + + # Remove pad positions from each trajectory (split_and_unpad_tensor + # auto-derives trim lengths from attention_mask when traj_seqlens=None) + all_data = [ + split_and_unpad_tensor( + d, n_trajs=1, traj_group_sizes=[d["attention_mask"].shape[0]] + )[0] + for d in all_data + ] + + allocate_fn = get_allocate_fn(packing_algorithm) + # Allocate trajectories to ranks using the configured packing algorithm + # No capacity limit leads to balanced partition across this group + group_indices = allocate_fn(seqlens, capacity=int(1e12), min_groups=world_size) + local_indices = group_indices[rank] + + # Select assigned trajectories for this rank (no concatenation — deferred to train side) + data = [all_data[i] for i in local_indices] + return RedistributedData( + all_data=all_data, + data=data, + rank=rank, + group_indices=group_indices, + ) + + def redistribute_trajectories( trajectories: list[dict[str, Any]], group=None, @@ -57,40 +127,12 @@ def redistribute_trajectories( - group_indices: Assignment of trajectory indices to each rank """ # All-gather trajectories from all ranks - all_gathered = all_gather_tensor_container(trajectories, group=group) - - # Flatten the list of lists into a single list of trajectories - all_data = [] - for traj_list in all_gathered: - all_data.extend(traj_list) - - # Compute sequence lengths for load balancing - seqlens = [d["attention_mask"].sum().item() for d in all_data] - - # Remove pad positions from each trajectory (split_and_unpad_tensor - # auto-derives trim lengths from attention_mask when traj_seqlens=None) - all_data = [ - split_and_unpad_tensor( - d, n_trajs=1, traj_group_sizes=[d["attention_mask"].shape[0]] - )[0] - for d in all_data - ] - - allocate_fn = get_allocate_fn(packing_algorithm) - # Allocate trajectories to ranks using the configured packing algorithm - # No capacity limit leads to balanced partition across this group - group_indices = allocate_fn( - seqlens, capacity=int(1e12), min_groups=dist.get_world_size(group) - ) - local_indices = group_indices[dist.get_rank(group=group)] - - # Select assigned trajectories for this rank (no concatenation — deferred to train side) - data = [all_data[i] for i in local_indices] - return RedistributedData( - all_data=all_data, - data=data, + all_gathered = _all_gather_ragged_trajectory_lists(trajectories, group=group) + return _pack_gathered_trajectories( + all_gathered, + world_size=dist.get_world_size(group), rank=dist.get_rank(group=group), - group_indices=group_indices, + packing_algorithm=packing_algorithm, ) @@ -99,9 +141,21 @@ def __init__(self, rollout_engine: InferenceEngine, train_engine: TrainEngine): self.rollout_engine = rollout_engine self.train_engine = train_engine + def _synchronize_head_error(self, error: str | None) -> str | None: + if not self.train_engine.is_data_parallel_head() or not dist.is_initialized(): + return error + head_errors: list[str | None] = [None] * dist.get_world_size( + self.train_engine.data_parallel_group + ) + dist.all_gather_object( + head_errors, error, group=self.train_engine.data_parallel_group + ) + return next((message for message in head_errors if message), None) + def _broadcast_and_redistribute_trajectories( self, trajectories: list[dict[str, Any]] | None, + preparation_error: str | None = None, ) -> list[dict[str, Any]]: """Broadcast and redistribute trajectories across distributed workers. @@ -122,18 +176,43 @@ def _broadcast_and_redistribute_trajectories( list[dict[str, Any]] Redistributed and broadcast batch available on all ranks (list of trajs) """ - if trajectories is not None: + error = self._synchronize_head_error(preparation_error) + batch = None + if trajectories is not None and error is None: + group = self.train_engine.data_parallel_group config = getattr(self.train_engine, "config", None) mb_spec = getattr(config, "mb_spec", None) packing_algorithm = getattr(mb_spec, "packing_algorithm", "ffd") - redist = redistribute_trajectories( - trajectories, - group=self.train_engine.data_parallel_group, - packing_algorithm=packing_algorithm, + # A rank-local failure inside the gather collectives is not + # recoverable: peers are parked inside those collectives, so + # posting the error all-gather instead would pair with their + # pending operations. Let it propagate without this rank issuing + # further communication. + all_gathered = _all_gather_ragged_trajectory_lists( + trajectories, group=group ) - batch = redist.data - else: - batch = None + try: + # Only collective-free work may run inside this recoverable + # block: every head holds the same gathered data, so failures + # are symmetric and every rank reaches the error sync below. + redist = _pack_gathered_trajectories( + all_gathered, + world_size=dist.get_world_size(group), + rank=dist.get_rank(group=group), + packing_algorithm=packing_algorithm, + ) + batch = redist.data + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + + error = self._synchronize_head_error(error) + error = broadcast_tensor_container( + error, + src_rank=self.train_engine.current_data_parallel_head(), + group=self.train_engine.context_and_model_parallel_group, + ) + if error is not None: + raise RuntimeError(f"Rollout batch preparation failed: {error}") current_platform.synchronize() dist.barrier(group=self.train_engine.cpu_group) @@ -149,6 +228,38 @@ def _broadcast_and_redistribute_trajectories( return batch + def _gather_collection_progress( + self, local_count: int, stalled_seconds: float + ) -> tuple[int, float]: + """Return the global trajectory count and the longest local stall. + + Both values derive from a single all-gather, so every data-parallel + head observes identical numbers and loop decisions based on them stay + collective-aligned. + """ + if not dist.is_initialized(): + return local_count, stalled_seconds + group = self.train_engine.data_parallel_group + progress: list[tuple[int, float] | None] = [None] * dist.get_world_size(group) + dist.all_gather_object(progress, (local_count, stalled_seconds), group=group) + counts, stalls = zip(*(entry for entry in progress if entry is not None)) + return sum(counts), max(stalls) + + def _prepare_on_data_parallel_head( + self, + prepare: Callable[[], list[dict[str, Any]]], + ) -> tuple[list[dict[str, Any]] | None, str | None]: + if not self.train_engine.is_data_parallel_head(): + return None, None + try: + trajectories = prepare() + return ( + tensor_container_to(trajectories, current_platform.current_device()), + None, + ) + except Exception as exc: + return None, f"{type(exc).__name__}: {exc}" + def rollout_batch( self, data: list[dict[str, Any]], @@ -157,6 +268,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Generate rollout batch with distributed coordination (synchronous). @@ -191,21 +303,21 @@ def rollout_batch( If rollout engine not connected via connect_engine() """ - trajectories = None - if self.train_engine.is_data_parallel_head(): - trajectories = self.rollout_engine.rollout_batch( + trajectories, preparation_error = self._prepare_on_data_parallel_head( + lambda: self.rollout_engine.rollout_batch( data, workflow=workflow, workflow_kwargs=workflow_kwargs, group_size=group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) - trajectories = tensor_container_to( - trajectories, current_platform.current_device() - ) + ) - return self._broadcast_and_redistribute_trajectories(trajectories) + return self._broadcast_and_redistribute_trajectories( + trajectories, preparation_error=preparation_error + ) def prepare_batch( self, @@ -217,6 +329,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Prepare async rollout batch with distributed coordination. @@ -252,20 +365,85 @@ def prepare_batch( If rollout engine not connected via connect_engine() """ - trajectories = None - if self.train_engine.is_data_parallel_head(): - trajectories = self.rollout_engine.prepare_batch( - dataloader, - workflow=workflow, - workflow_kwargs=workflow_kwargs, - should_accept_fn=should_accept_fn, - group_size=group_size, - dynamic_bs=dynamic_bs, - reward_normalization=reward_normalization, - drop_incomplete_group=drop_incomplete_group, - ) - trajectories = tensor_container_to( - trajectories, current_platform.current_device() + def _prepare_until_dispatchable() -> list[dict[str, Any]]: + trajectories: list[dict[str, Any]] = [] + min_global_batch_size = ( + dist.get_world_size(self.train_engine.data_parallel_group) + if dist.is_initialized() + else 1 ) + global_batch_size = 0 + empty_rounds = 0 + stall_started_at: float | None = None + while True: + local_error = None + prepared: list[dict[str, Any]] = [] + try: + prepared = self.rollout_engine.prepare_batch( + dataloader, + workflow=workflow, + workflow_kwargs=workflow_kwargs, + should_accept_fn=should_accept_fn, + group_size=group_size, + min_usable_group_size=min_usable_group_size, + dynamic_bs=dynamic_bs, + reward_normalization=reward_normalization, + drop_incomplete_group=drop_incomplete_group, + ) + except Exception as exc: + local_error = f"{type(exc).__name__}: {exc}" + coordinated_error = self._synchronize_head_error(local_error) + if coordinated_error is not None: + raise RuntimeError(coordinated_error) + trajectories.extend(prepared) + previous_global_batch_size = global_batch_size + stalled_seconds = ( + 0.0 + if stall_started_at is None + else time.monotonic() - stall_started_at + ) + global_batch_size, global_stalled_seconds = ( + self._gather_collection_progress(len(trajectories), stalled_seconds) + ) + if global_batch_size >= min_global_batch_size: + return trajectories + if not dynamic_bs: + raise RuntimeError( + "Fixed rollout preparation produced only " + f"{global_batch_size} trainable groups for " + f"{min_global_batch_size} data-parallel ranks" + ) + # The streak and stall duration are derived from the same + # all-gather, so every data-parallel head raises on the same + # iteration and the coordinated error path surfaces the + # stall on all ranks instead of a silent spin. + if global_batch_size > previous_global_batch_size: + empty_rounds = 0 + stall_started_at = None + else: + empty_rounds += 1 + if stall_started_at is None: + stall_started_at = time.monotonic() + if ( + empty_rounds >= MAX_CONSECUTIVE_EMPTY_ROLLOUT_ROUNDS + and global_stalled_seconds + >= ROLLOUT_COLLECTION_STALL_TIMEOUT_SECONDS + ): + raise RuntimeError( + "Dynamic rollout preparation added no trainable group " + f"for {empty_rounds} consecutive rounds over " + f"{global_stalled_seconds:.0f}s " + f"({global_batch_size}/{min_global_batch_size} across " + "data-parallel heads). Every group was rejected or " + "masked; check the reward function and " + "should_accept_fn, and the usable_slot_count / " + "fully_masked_group rollout statistics." + ) + + trajectories, preparation_error = self._prepare_on_data_parallel_head( + _prepare_until_dispatchable + ) - return self._broadcast_and_redistribute_trajectories(trajectories) + return self._broadcast_and_redistribute_trajectories( + trajectories, preparation_error=preparation_error + ) diff --git a/areal/infra/remote_inf_engine.py b/areal/infra/remote_inf_engine.py index e6d5ae8143..1113806593 100644 --- a/areal/infra/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -46,8 +46,8 @@ from areal.infra.utils.http import arequest_with_retry, get_default_connector from areal.infra.utils.launcher import wait_llm_server_addrs from areal.infra.utils.proc import kill_process_tree -from areal.utils import logging, name_resolve, names -from areal.utils.data import concat_padded_tensors +from areal.utils import logging, name_resolve, names, stats_tracker +from areal.utils.data import concat_padded_tensors, get_batch_size from areal.utils.dynamic_import import import_from_string from areal.utils.network import ( find_free_ports, @@ -57,7 +57,12 @@ ) from areal.utils.perf_tracer import trace_perf -from .workflow_executor import WorkflowExecutor +from .workflow_executor import ( + WorkflowContractError, + WorkflowExecutor, + WorkflowTaskResult, + validate_rollout_group_sizes, +) if TYPE_CHECKING: from areal.experimental.openai import InteractionWithTokenLogpReward @@ -75,15 +80,42 @@ def __init__( logger: Logger, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ): - if group_size < 1: - raise ValueError(f"group_size must be >= 1, got {group_size}") + validate_rollout_group_sizes(group_size, min_usable_group_size) self.workflow = workflow self.group_size = group_size + self.min_usable_group_size = min_usable_group_size self.logger = logger self.reward_normalization = reward_normalization self.drop_incomplete_group = drop_incomplete_group + def _record_group_stats(self, usable_slot_count: int, *, trainable: bool) -> None: + trainable_slot_count = usable_slot_count if trainable else 0 + stats_tracker.get(workflow_context.stat_scope()).scalar( + target_slot_count=self.group_size, + usable_slot_count=usable_slot_count, + trainable_slot_count=trainable_slot_count, + fully_masked_group=usable_slot_count == 0, + singleton_slot_group=usable_slot_count == 1, + pre_filter_usable_slot_yield=usable_slot_count / self.group_size, + pre_filter_trainable_slot_yield=trainable_slot_count / self.group_size, + ) + + def _validate_slot_cardinality(self, slot_sizes: list[int]) -> None: + if self.min_usable_group_size > 1 and any(size != 1 for size in slot_sizes): + raise WorkflowContractError( + "min_usable_group_size >= 2 (derived from group-relative " + "normalization, or set via actor.min_usable_group_size) requires " + "each rollout slot to contribute exactly one training sample; got " + f"slot sizes {slot_sizes}. Either return one sample per " + "arun_episode call (for agent workflows, set " + "agent.export_style='concat'), set mean_level/std_level to " + "'batch' in actor.reward_norm/actor.adv_norm, or unset " + "actor.min_usable_group_size. This contract violation is " + "non-retryable and stops training." + ) + async def arun_episode( self, engine: InferenceEngine, data: dict[str, Any] ) -> dict[str, Any] | None: @@ -94,9 +126,10 @@ async def arun_episode( ) valid_results = [r for r in results if r is not None] + usable_slot_count = len(valid_results) - # All results None -> return None if not valid_results: + self._record_group_stats(usable_slot_count, trainable=False) return None # Some results None -> drop entire group if requested. Reward @@ -110,13 +143,23 @@ async def arun_episode( "(drop_incomplete_group=True). prepare_batch will retry " "with a new prompt from the dataloader." ) + self._record_group_stats(usable_slot_count, trainable=False) return None if not self.reward_normalization: + action = ( + "dropping group below min_usable_group_size" + if usable_slot_count < self.min_usable_group_size + else "using remaining results" + ) self.logger.warning( f"GroupedRolloutWorkflow: {n_failed}/{len(results)} " - "trajectories returned None, using remaining results" + f"trajectories returned None, {action}" ) + if usable_slot_count < self.min_usable_group_size: + self._record_group_stats(usable_slot_count, trainable=False) + return None + # Check if results are InteractionWithTokenLogpReward dicts first = valid_results[0] if ( @@ -126,17 +169,24 @@ async def arun_episode( isinstance(v, InteractionWithTokenLogpReward) for v in first.values() ) ): + self._validate_slot_cardinality([len(result) for result in valid_results]) if self.reward_normalization and self.group_size > 1: if not self._normalize_group_rewards(results): + self._record_group_stats(usable_slot_count, trainable=False) return None # Merge dicts - each result is {completion_id: InteractionWithTokenLogpReward} merged: dict[str, InteractionWithTokenLogpReward] = {} for result in valid_results: merged.update(result) + self._record_group_stats(usable_slot_count, trainable=bool(merged)) return merged if merged else None # Otherwise, tensor dicts - concatenate + self._validate_slot_cardinality( + [get_batch_size(result) for result in valid_results] + ) concatenated = concat_padded_tensors(valid_results) + self._record_group_stats(usable_slot_count, trainable=bool(concatenated)) return concatenated if concatenated else None def _normalize_group_rewards( @@ -702,7 +752,9 @@ def _resolve_workflow( proxy_addr: str | None = None, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> RolloutWorkflow: + validate_rollout_group_sizes(group_size, min_usable_group_size) resolved: RolloutWorkflow # 0. None workflow = online mode (config-driven) @@ -723,6 +775,7 @@ def _resolve_workflow( self.logger, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, + min_usable_group_size=min_usable_group_size, ) return resolved @@ -820,6 +873,7 @@ def _resolve_workflow( self.logger, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, + min_usable_group_size=min_usable_group_size, ) return resolved @@ -1209,6 +1263,7 @@ def submit( proxy_addr: str | None = None, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> int: """Submit a request to the inference engine and return immediately. @@ -1225,6 +1280,9 @@ def submit( group_size : int Number of times to run the workflow per input and concatenate results. Default is 1 (no grouping). + min_usable_group_size : int + Estimator-owned minimum number of usable logical rollout slots. Must be + between 1 and ``group_size``. Default is 1. task_id : int, optional The task ID to use. If None, a new task ID will be generated internally. is_eval : bool, optional @@ -1248,6 +1306,7 @@ def submit( workflow, workflow_kwargs, group_size, + min_usable_group_size=min_usable_group_size, proxy_addr=proxy_addr, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, @@ -1292,6 +1351,13 @@ def wait_for_task( """Wait for a specific submitted task to complete.""" return self.workflow_executor.wait_for_task(task_id, timeout, raise_timeout) + def _wait_for_task_result( + self, task_id: int, timeout: float | None = None, raise_timeout: bool = True + ) -> WorkflowTaskResult | None: + return self.workflow_executor._wait_for_task_result( + task_id, timeout, raise_timeout + ) + def rollout_batch( self, data: list[dict[str, Any]], @@ -1300,6 +1366,7 @@ def rollout_batch( group_size: int = 1, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Submit a batch of requests and wait for results. @@ -1332,6 +1399,7 @@ def rollout_batch( workflow, workflow_kwargs, group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) @@ -1351,6 +1419,7 @@ def prepare_batch( dynamic_bs: bool = False, reward_normalization: bool = False, drop_incomplete_group: bool = False, + min_usable_group_size: int = 1, ) -> list[dict[str, Any]]: """Asynchronously submit and wait until a full batch is ready. @@ -1369,6 +1438,9 @@ def prepare_batch( Default is 1 (no grouping). dynamic_bs : bool, optional If True, enables dynamic batch sizing. Default is False. + min_usable_group_size : int + Estimator-owned minimum number of usable logical rollout slots. Must be + between 1 and ``group_size``. Default is 1. Returns ------- @@ -1384,6 +1456,7 @@ def prepare_batch( workflow, workflow_kwargs, group_size, + min_usable_group_size=min_usable_group_size, reward_normalization=reward_normalization, drop_incomplete_group=drop_incomplete_group, ) diff --git a/areal/infra/workflow_executor.py b/areal/infra/workflow_executor.py index b2870bf8df..433e9edd76 100644 --- a/areal/infra/workflow_executor.py +++ b/areal/infra/workflow_executor.py @@ -238,12 +238,53 @@ class _RolloutTaskInput: is_eval: bool = False -@dataclass +@dataclass(frozen=True) class _RolloutResult: task_id: int trajectory: dict[str, Any] +class WorkflowContractError(ValueError): + """Raised when a workflow violates a non-retryable output contract.""" + + +def validate_rollout_group_sizes(group_size: int, min_usable_group_size: int) -> None: + if group_size < 1: + raise ValueError(f"group_size must be >= 1, got {group_size}") + if not 1 <= min_usable_group_size <= group_size: + raise ValueError( + "min_usable_group_size must be between 1 and group_size " + f"({group_size}), got {min_usable_group_size}" + ) + + +@dataclass(frozen=True) +class WorkflowContractFailure: + message: str + + +WorkflowTaskResult = _RolloutResult | WorkflowContractFailure + + +def get_workflow_result_error( + result: WorkflowTaskResult, +) -> WorkflowContractError | None: + if isinstance(result, WorkflowContractFailure): + return WorkflowContractError(result.message) + return None + + +def unwrap_workflow_result( + result: WorkflowTaskResult | None, +) -> _RolloutResult | None: + if result is None: + return None + error = get_workflow_result_error(result) + if error is not None: + raise error + return result + + # Batch size for fetching from the async task runner _MAX_FETCH_BATCH_SIZE = 100 # Timeout for shutting down threads @@ -251,6 +292,13 @@ class _RolloutResult: # Timeout for "wait" and "wait_for_task" if timeout parameter is None _DEFAULT_WAIT_TIMEOUT_SECONDS = float(7 * 24 * 3600) +# Dynamic collection is declared stalled only when BOTH hold: this many +# consecutive rounds added zero trainable groups, and the streak lasted at +# least this long. Requiring both keeps fast legitimate all-reject bursts +# (e.g. a staleness flush after a weight update) from aborting the run. +MAX_CONSECUTIVE_EMPTY_ROLLOUT_ROUNDS = 8 +ROLLOUT_COLLECTION_STALL_TIMEOUT_SECONDS = 1800.0 + class WithTaskID(Protocol): task_id: int @@ -279,6 +327,7 @@ def __init__( task_factory: Callable[[TInput], Callable[[], Awaitable[TResult | None]]], staleness_manager: StalenessManager, enable_tracing: bool = False, + terminal_error_fn: Callable[[TResult], Exception | None] | None = None, ): self.runner = AsyncTaskRunner( max_queue_size=max_queue_size, @@ -287,6 +336,7 @@ def __init__( self.task_factory = task_factory self.staleness_manager = staleness_manager self.enable_tracing = enable_tracing + self.terminal_error_fn = terminal_error_fn self.logger: Logger # Unbounded deques for producer/consumer pattern @@ -571,7 +621,12 @@ def wait_results( timeout = _DEFAULT_WAIT_TIMEOUT_SECONDS with self._result_cv: - while len(self._pending_results) < count: + while True: + terminal_error = self._pop_terminal_error_locked() + if terminal_error is not None: + raise terminal_error + if len(self._pending_results) >= count: + break self._check_thread_exception() elapsed = time.perf_counter() - start_time @@ -604,7 +659,11 @@ def wait_results( return [r.data for r in selected] def wait_for_task( - self, task_id: int, timeout: float | None = None, raise_timeout: bool = True + self, + task_id: int, + timeout: float | None = None, + raise_timeout: bool = True, + raise_terminal_error: bool = True, ) -> TResult | None: """Wait for a specific task result by task_id.""" start_time = time.perf_counter() @@ -615,7 +674,13 @@ def wait_for_task( if task_id not in self._active_task_ids: raise ValueError(f"Task {task_id} is never submitted.") - while task_id not in self._pending_results: + while True: + if raise_terminal_error: + terminal_error = self._pop_terminal_error_locked(task_id) + if terminal_error is not None: + raise terminal_error + if task_id in self._pending_results: + break self._check_thread_exception() elapsed = time.perf_counter() - start_time @@ -632,6 +697,27 @@ def wait_for_task( self._result_cv.notify_all() return found_result.data + def _pop_terminal_error_locked( + self, task_id: int | None = None + ) -> Exception | None: + if self.terminal_error_fn is None: + return None + task_ids = ( + [task_id] if task_id is not None else list(self._pending_results.keys()) + ) + for pending_task_id in task_ids: + result = self._pending_results.get(pending_task_id) + if result is None or result.data is None: + continue + error = self.terminal_error_fn(result.data) + if error is None: + continue + self._pending_results.pop(pending_task_id) + self._active_task_ids.discard(pending_task_id) + self._result_cv.notify_all() + return error + return None + def active_submit_and_wait( self, input_generator: Generator[TInput, None, None], @@ -776,7 +862,7 @@ def __init__( # Dispatcher will be initialized in initialize() after staleness_manager is ready self._dispatcher: ( - BatchTaskDispatcher[_RolloutTaskInput, _RolloutResult] | None + BatchTaskDispatcher[_RolloutTaskInput, WorkflowTaskResult] | None ) = None self._task_id_generator = TaskIdGenerator() @@ -1061,11 +1147,12 @@ def initialize(self, logger=None, train_data_parallel_size: int | None = None): # Create and initialize the dispatcher qsize = self.config.queue_size or self.max_concurrent_rollouts * 16 - self._dispatcher = BatchTaskDispatcher[_RolloutTaskInput, _RolloutResult]( + self._dispatcher = BatchTaskDispatcher[_RolloutTaskInput, WorkflowTaskResult]( max_queue_size=qsize, task_factory=self._create_workflow_task, staleness_manager=self._staleness_manager, enable_tracing=self.config.enable_rollout_tracing, + terminal_error_fn=get_workflow_result_error, ) # Initialize the dispatcher's async task runner @@ -1107,7 +1194,7 @@ def _rollout_stats(self) -> str: def _create_workflow_task( self, pending_task: _RolloutTaskInput - ) -> Callable[[], Awaitable[_RolloutResult | None]]: + ) -> Callable[[], Awaitable[WorkflowTaskResult | None]]: """Wrapper to create an async function that will be executed by AsyncTaskRunner. This is a synchronous function that returns an async function, which allows @@ -1125,7 +1212,7 @@ def _create_workflow_task( filtering/validation. """ - async def _execute_workflow() -> _RolloutResult | None: + async def _execute_workflow() -> WorkflowTaskResult | None: """Execute workflow.arun_episode and apply AReaL-specific logic.""" task_id = pending_task.task_id @@ -1231,6 +1318,20 @@ async def _execute_workflow() -> _RolloutResult | None: ) return None + except WorkflowContractError as exc: + manager.on_rollout_rejected() + stats_tracker.get("rollout").scalar(rejected=1) + trace_session_event( + "mark_finalized", + task_id=task_id, + status="failed", + reason="workflow_contract_error", + ) + if self.logger is not None: + self.logger.error( + "Workflow contract violation: %s", exc, exc_info=True + ) + return WorkflowContractFailure(message=str(exc)) except Exception as exc: # pragma: no cover - workflow execution errors manager.on_rollout_rejected() stats_tracker.get("rollout").scalar(rejected=1) @@ -1289,7 +1390,10 @@ def wait( See :meth:`~areal.api.engine_api.InferenceEngine.wait` for parameters. """ # Delegate to dispatcher and extract trajectories - results = self.dispatcher.wait_results(count, timeout, raise_timeout) + results = [ + unwrap_workflow_result(result) + for result in self.dispatcher.wait_results(count, timeout, raise_timeout) + ] # Log and trace if self.config.enable_rollout_tracing: self.logger.info("Rollout results are ready!") @@ -1322,12 +1426,25 @@ def wait_for_task( -------- :meth:`~areal.api.engine_api.InferenceEngine.wait_for_task` """ - result = self.dispatcher.wait_for_task(task_id, timeout, raise_timeout) + result = unwrap_workflow_result( + self.dispatcher.wait_for_task(task_id, timeout, raise_timeout) + ) if result is not None and self.config.enable_rollout_tracing: self.logger.info(f"Task {task_id} completed successfully") return result.trajectory if result is not None else None + def _wait_for_task_result( + self, task_id: int, timeout: float | None = None, raise_timeout: bool = True + ) -> WorkflowTaskResult | None: + """Return the raw task result for controller-to-worker transport.""" + return self.dispatcher.wait_for_task( + task_id, + timeout, + raise_timeout, + raise_terminal_error=False, + ) + @trace_perf("workflow_executor.rollout_batch", category="scheduler") def rollout_batch( self, @@ -1416,9 +1533,14 @@ def task_input_generator(): # Delegate to dispatcher assert dataloader.batch_size is not None - results = self.dispatcher.active_submit_and_wait( - self.data_generator, batch_size=dataloader.batch_size, dynamic_bs=dynamic_bs - ) + results = [ + unwrap_workflow_result(result) + for result in self.dispatcher.active_submit_and_wait( + self.data_generator, + batch_size=dataloader.batch_size, + dynamic_bs=dynamic_bs, + ) + ] # Return list of trajectory dicts (filter out None) return [r.trajectory for r in results if r is not None] @@ -1452,7 +1574,9 @@ def staleness_manager(self) -> StalenessManager: return manager @property - def dispatcher(self) -> BatchTaskDispatcher[_RolloutTaskInput, _RolloutResult]: + def dispatcher( + self, + ) -> BatchTaskDispatcher[_RolloutTaskInput, WorkflowTaskResult]: """Get the task dispatcher, ensuring initialization has been called.""" if self._dispatcher is None: raise RuntimeError( diff --git a/areal/models/tree_attn/tree.py b/areal/models/tree_attn/tree.py index 7b655dc7a3..cf766f6d93 100644 --- a/areal/models/tree_attn/tree.py +++ b/areal/models/tree_attn/tree.py @@ -26,7 +26,7 @@ precompute_tree_attention_data, ) from areal.utils import logging, stats_tracker -from areal.utils.data import MicroBatchList +from areal.utils.data import TRANSPORT_DUMMY_KEY, MicroBatchList from areal.utils.perf_tracer import trace_perf, trace_scope logger = logging.getLogger("TreeAttentionCore") @@ -404,6 +404,7 @@ def build_packed_tree_batch( # Build packed outputs for each tree mbs: list[dict[str, Any]] = [] + padded_mbs: list[dict[str, Any]] = [] padding_lengths: list[int] = [] padded_to_lengths: list[int] = [] @@ -448,13 +449,17 @@ def build_packed_tree_batch( non_packable_keys, ) - mb = { + padded_mb = { "input_ids": input_ids, "position_ids": position_ids, "trie_node": trie, **extra_data, } + mb = dict(padded_mb) + if not trie.all_sequence_ids: + mb[TRANSPORT_DUMMY_KEY] = True mbs.append(mb) + padded_mbs.append(padded_mb) padding_lengths.append(padded_size - num_tokens) padded_to_lengths.append(padded_size) @@ -465,7 +470,7 @@ def build_packed_tree_batch( mb_spec=mb_spec, mbs=mbs, group_lens=[num for num in num_tokens_list], - padded_mbs=mbs, + padded_mbs=padded_mbs, padding_lengths=padding_lengths, padded_to_lengths=padded_to_lengths, _max_seqlen=max(padded_to_lengths), diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index 75dbbd7311..b2568c0f17 100644 --- a/areal/trainer/ppo/actor.py +++ b/areal/trainer/ppo/actor.py @@ -6,7 +6,7 @@ import torch from areal.api import TrainEngine -from areal.api.cli_args import MicroBatchSpec, PPOActorConfig, RejectionSamplingConfig +from areal.api.cli_args import PPOActorConfig, RejectionSamplingConfig from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.trainer.ppo.stats import infer_token_denominator @@ -26,7 +26,7 @@ Normalization, TrajBatchMeta, batched_call, - split_padded_tensor_dict_into_mb_list, + split_training_batch_into_microbatches, ) from areal.utils.functional import ( cispo_loss_fn, @@ -42,6 +42,33 @@ logger = logging.getLogger("PPOActor") +def _group_training_metrics( + loss_mask: torch.Tensor, group_sizes: list[int] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = loss_mask.shape[0] + if any(size < 1 for size in group_sizes) or sum(group_sizes) != batch_size: + raise ValueError( + f"group_sizes must be positive and sum to batch size {batch_size}, " + f"got {group_sizes}" + ) + + group_starts = torch.zeros(batch_size, dtype=torch.bool, device=loss_mask.device) + usable_group_sizes = torch.zeros( + batch_size, dtype=torch.float32, device=loss_mask.device + ) + group_loss_weights = torch.zeros_like(usable_group_sizes) + sizes = torch.tensor(group_sizes, dtype=torch.long, device=loss_mask.device) + ends = sizes.cumsum(0) + starts = ends - sizes + token_counts = loss_mask.reshape(batch_size, -1).sum(1, dtype=torch.float32) + cumulative_tokens = torch.nn.functional.pad(token_counts.cumsum(0), (1, 0)) + + group_starts[starts] = True + usable_group_sizes[starts] = sizes.to(usable_group_sizes.dtype) + group_loss_weights[starts] = cumulative_tokens[ends] - cumulative_tokens[starts] + return group_starts, usable_group_sizes, group_loss_weights + + def _infer_prompt_lens( attention_mask: torch.Tensor, loss_mask: torch.Tensor ) -> torch.Tensor: @@ -277,9 +304,11 @@ def _compute_advantages( @trace_perf("ppo_actor.ppo_update", category="compute") @stats_tracker.scope_func_wrapper("ppo_actor") def ppo_update(self, data: list[dict[str, Any]]) -> None: - batched_call(self._ppo_update, data, unpack=False) + batched_call(self._ppo_update, data, unpack=False, pass_meta=True) - def _ppo_update(self, data: dict[str, Any]) -> None: + def _ppo_update( + self, data: dict[str, Any], meta: TrajBatchMeta | None = None + ) -> None: attn_mask = data["attention_mask"] loss_mask = data["loss_mask"] reward_score = data["rewards"] @@ -312,6 +341,10 @@ def _ppo_update(self, data: dict[str, Any]) -> None: n_valid_tokens=loss_mask.bool(), **result_denominators, ) + group_metrics = None + if meta is not None: + group_metrics = _group_training_metrics(loss_mask, meta.traj_group_sizes) + global_denominators["n_groups"] = group_metrics[0] stats_tracker.denominator(**global_denominators) stats_tracker.stat( correct_seq_len=seqlens.float(), denominator="correct_n_seqs" @@ -335,6 +368,22 @@ def _ppo_update(self, data: dict[str, Any]) -> None: seq_len=seqlens.float(), ) stats_tracker.stat(**seq_stats, denominator="n_seqs") + if group_metrics is not None: + group_starts, usable_group_sizes, group_loss_weights = group_metrics + stats_tracker.stat( + usable_group_size=usable_group_sizes, + group_loss_weight=group_loss_weights, + denominator="n_groups", + ) + for group_size in sorted(set(meta.traj_group_sizes)): + denominator = f"n_groups_size_{group_size}" + stats_tracker.denominator( + **{denominator: group_starts & (usable_group_sizes == group_size)} + ) + stats_tracker.stat( + denominator=denominator, + **{f"group_loss_weight_size_{group_size}": group_loss_weights}, + ) scalars = dict( mask_no_eos_with_zero=self.config.mask_no_eos_with_zero, eps_clip=self.config.eps_clip, @@ -364,16 +413,17 @@ def _ppo_update(self, data: dict[str, Any]) -> None: data.pop(key, None) # NOTE: calling engine.train() is critical to enabling gradient checkpointing self.engine.train() - mb_inputs = split_padded_tensor_dict_into_mb_list( + mb_inputs = split_training_batch_into_microbatches( data, - mb_spec=MicroBatchSpec(n_mbs=self.config.ppo_n_minibatches), + n_mbs=self.config.ppo_n_minibatches, + group=self.engine.data_parallel_group, ) with stats_tracker.scope("update"): # Get current version for proximal approximation metrics current_version = self.engine.get_version() - for mb in mb_inputs.mbs: + for mb in mb_inputs: train_stat = self.engine.train_batch( mb, loss_fn=functools.partial( diff --git a/areal/trainer/ppo/critic.py b/areal/trainer/ppo/critic.py index 05f83c7441..965afbf66c 100644 --- a/areal/trainer/ppo/critic.py +++ b/areal/trainer/ppo/critic.py @@ -6,14 +6,14 @@ import torch from areal.api import TrainEngine -from areal.api.cli_args import MicroBatchSpec, PPOCriticConfig +from areal.api.cli_args import PPOCriticConfig from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.trainer.ppo.stats import infer_token_denominator from areal.utils import stats_tracker from areal.utils.data import ( batched_call, - split_padded_tensor_dict_into_mb_list, + split_training_batch_into_microbatches, ) from areal.utils.functional import ppo_critic_loss_fn from areal.utils.perf_tracer import trace_perf @@ -58,11 +58,12 @@ def _ppo_update(self, data: dict[str, Any]) -> None: # NOTE: calling engine.train() is critical to enabling gradient checkpointing self.engine.train() - mb_inputs = split_padded_tensor_dict_into_mb_list( + mb_inputs = split_training_batch_into_microbatches( data, - mb_spec=MicroBatchSpec(n_mbs=self.config.ppo_n_minibatches), + n_mbs=self.config.ppo_n_minibatches, + group=self.engine.data_parallel_group, ) - for mb in mb_inputs.mbs: + for mb in mb_inputs: train_stat = self.engine.train_batch( mb, loss_fn=functools.partial( diff --git a/areal/trainer/rl_trainer.py b/areal/trainer/rl_trainer.py index 117ff61c8c..2506c57913 100644 --- a/areal/trainer/rl_trainer.py +++ b/areal/trainer/rl_trainer.py @@ -4,6 +4,7 @@ import functools import os +import time from collections.abc import Callable from copy import deepcopy from typing import TYPE_CHECKING, Any, cast @@ -46,6 +47,10 @@ from areal.infra.data_service.controller.config import DataServiceConfig from areal.infra.data_service.rdataset import RDataset from areal.infra.utils.concurrent import call_maybe_async +from areal.infra.workflow_executor import ( + MAX_CONSECUTIVE_EMPTY_ROLLOUT_ROUNDS, + ROLLOUT_COLLECTION_STALL_TIMEOUT_SECONDS, +) from areal.utils import logging, perf_tracer, seeding, stats_tracker from areal.utils.dataloader import create_dataloader from areal.utils.environ import is_single_controller @@ -75,6 +80,85 @@ logger = logging.getLogger("RLTrainer") +def _collect_trainable_rollout_batch( + prepare_batch: Callable[[], list[dict[str, Any]]], + *, + dynamic_bs: bool, + min_batch_size: int = 1, + max_empty_rounds: int = MAX_CONSECUTIVE_EMPTY_ROLLOUT_ROUNDS, + stall_timeout: float = ROLLOUT_COLLECTION_STALL_TIMEOUT_SECONDS, +) -> list[dict[str, Any]]: + """Collect until every local DP consumer can receive a trajectory group. + + ``None`` results are objective-level rejections, so dynamic collection keeps + polling the ready queue. A streak of at least ``max_empty_rounds`` rounds + that add no trainable group AND lasts at least ``stall_timeout`` seconds + aborts the run instead of letting it spin silently; fast legitimate + all-reject bursts (e.g. a staleness flush after a weight update) stay + alive. Non-retryable workflow contract errors propagate through + ``prepare_batch`` without reaching this loop. + """ + if min_batch_size < 1: + raise ValueError(f"min_batch_size must be positive, got {min_batch_size}") + + batch: list[dict[str, Any]] = [] + empty_rounds = 0 + stall_started_at: float | None = None + last_warned_at = float("-inf") + while len(batch) < min_batch_size: + prepared = prepare_batch() + batch.extend(prepared) + if len(batch) >= min_batch_size: + break + if not dynamic_bs: + raise RuntimeError( + "Fixed rollout preparation produced only " + f"{len(batch)} trainable groups; at least {min_batch_size} are required" + ) + now = time.monotonic() + if prepared: + empty_rounds = 0 + stall_started_at = None + else: + empty_rounds += 1 + if stall_started_at is None: + stall_started_at = now + if ( + empty_rounds >= max_empty_rounds + and now - stall_started_at >= stall_timeout + ): + raise RuntimeError( + f"Dynamic rollout collection added no trainable group for " + f"{empty_rounds} consecutive rounds over " + f"{now - stall_started_at:.0f}s " + f"({len(batch)}/{min_batch_size} collected). Every group " + "was rejected or masked; check the reward function and " + "should_accept_fn, and the usable_slot_count / " + "fully_masked_group rollout statistics." + ) + if now - last_warned_at >= 60.0: + last_warned_at = now + logger.warning( + "Dynamic rollout batch has %d/%d trainable groups; collecting " + "from the ready queue again", + len(batch), + min_batch_size, + ) + return batch + + +def _minimum_consumer_batch_size(*consumers: Any | None) -> int: + """Return the largest DP degree among train engines consuming a rollout.""" + return max( + ( + consumer.parallel_strategy.dp_size + for consumer in consumers + if consumer is not None + ), + default=1, + ) + + class _EmptyDataLoader: """Minimal dataloader for online mode that yields empty dicts. @@ -645,6 +729,29 @@ def train( total_epochs: int | None = None, ): config = self.config + is_v1_rollout = config.rollout._version == "v1" + if not is_v1_rollout and config.actor.min_usable_group_size is not None: + raise ValueError( + "The v2 rollout path does not support actor.min_usable_group_size " + "yet; unset it or use a v1 rollout backend." + ) + min_usable_group_size = ( + config.actor.resolve_min_usable_group_size(config.gconfig.n_samples) + if is_v1_rollout + else 1 + ) + train_teacher = ( + self.teacher + if config.teacher is not None and config.teacher.engine_type == "train" + else None + ) + min_consumer_batch_size = ( + _minimum_consumer_batch_size( + self.actor, self.critic, self.ref, train_teacher + ) + if is_single_controller() + else 1 + ) start_step = ( self.recover_info.last_step_info.next().global_step if self.recover_info is not None @@ -694,16 +801,26 @@ def train( }, ), ): - rollout_batch = self.actor.prepare_batch( - self.train_dataloader, + prepare_kwargs = dict( workflow=workflow, workflow_kwargs=workflow_kwargs, should_accept_fn=dynamic_filter_fn, group_size=config.gconfig.n_samples, - dynamic_bs=self.config.dynamic_bs, + dynamic_bs=config.dynamic_bs, reward_normalization=config.gconfig.reward_normalization, drop_incomplete_group=config.gconfig.drop_incomplete_group, ) + if is_v1_rollout: + prepare_kwargs["min_usable_group_size"] = min_usable_group_size + rollout_batch = _collect_trainable_rollout_batch( + functools.partial( + self.actor.prepare_batch, + self.train_dataloader, + **prepare_kwargs, + ), + dynamic_bs=config.dynamic_bs, + min_batch_size=min_consumer_batch_size, + ) if self._should_offload_rollout: self._offload_rollout() diff --git a/areal/utils/data.py b/areal/utils/data.py index 1560742a1a..db0a7fa8ec 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -23,6 +23,8 @@ logger = logging.getLogger("DataUtils") +TRANSPORT_DUMMY_KEY = "_transport_dummy" + def get_batch_size(data: dict[str, Any]) -> int: if not data: @@ -621,6 +623,7 @@ class MicroBatchList: # sequence-level padding information align_to_lengths: list[int] | None = None old_cu_seqlens_list: list[torch.Tensor] | None = None + transport_dummy_count: int = 0 @property def max_seqlen(self) -> int: @@ -691,16 +694,77 @@ def to(self, *args, **kwargs): padded_to_lengths=self.padded_to_lengths, old_cu_seqlens_list=old_cu_seqlens_list, align_to_lengths=self.align_to_lengths, + transport_dummy_count=self.transport_dummy_count, ) DEFAULT_MAX_TOKENS_PER_MB = int(1e12) +def make_transport_dummy(template: dict[str, Any]) -> dict[str, Any]: + """Create one model-valid row for collective participation.""" + batch_size = get_batch_size(template) + if batch_size < 1: + raise ValueError("Cannot create transport padding from an empty batch") + + dummy: dict[str, Any] = {} + for key, value in template.items(): + if is_multi_modal_key(key) and isinstance(value, list): + dummy[key] = [{}] + elif ( + isinstance(value, torch.Tensor) + and value.ndim > 0 + and value.shape[0] == batch_size + ): + dummy[key] = torch.zeros_like(value[:1]) + elif isinstance(value, list) and len(value) == batch_size: + dummy[key] = [copy.deepcopy(value[0])] + else: + dummy[key] = copy.deepcopy(value) + + attention_mask = dummy.get("attention_mask") + if not isinstance(attention_mask, torch.Tensor) or attention_mask.ndim != 2: + raise ValueError("Transport padding requires a 2D attention_mask") + if attention_mask.shape[1] < 1: + raise ValueError("Transport padding requires sequence length >= 1") + attention_mask[:, 0] = 1 + if isinstance(dummy.get("loss_mask"), torch.Tensor): + dummy["loss_mask"].zero_() + return dummy + + +def make_transport_microbatch(template: dict[str, Any]) -> dict[str, Any]: + """Create one transport-only batch that arbitrary objectives must bypass.""" + dummy = make_transport_dummy(template) + dummy[TRANSPORT_DUMMY_KEY] = True + return dummy + + +def _pad_batch_to_min_groups( + data: dict[str, Any], + *, + min_groups: int, + granularity: int, +) -> tuple[dict[str, Any], int]: + batch_size = get_batch_size(data) + if batch_size % granularity != 0: + raise RuntimeError( + f"Batch size {batch_size} cannot divide granularity {granularity}." + ) + current_groups = batch_size // granularity + pad_count = max(min_groups - current_groups, 0) * granularity + if pad_count == 0: + return data, 0 + dummies = [make_transport_dummy(data) for _ in range(pad_count)] + return concat_padded_tensors([data, *dummies]), pad_count + + def split_padded_tensor_dict_into_mb_list( data: dict[str, Any], mb_spec: MicroBatchSpec, group: dist.ProcessGroup | None = None, + allow_transport_padding: bool = False, + synchronize: bool = True, ) -> MicroBatchList: """Split a padded dict of tensors into micro-batches based on the attention mask. @@ -708,6 +772,9 @@ def split_padded_tensor_dict_into_mb_list( data (Dict): Dictionary containing padded tensors. mb_spec (MicroBatchSpec): Specification for micro-batch splitting. group (Optional[dist.ProcessGroup]): Process group for distributed synchronization. + allow_transport_padding: Add model-valid rows when synchronized execution + requires more micro-batches than local semantic data can provide. + synchronize: Synchronize the micro-batch count across ``group``. Returns: MicroBatchList: A structure containing the split micro-batches and metadata. @@ -720,19 +787,56 @@ def split_padded_tensor_dict_into_mb_list( mb_spec, max_tokens_per_mb=DEFAULT_MAX_TOKENS_PER_MB ) granularity = mb_spec.granularity - bs = data["attention_mask"].shape[0] - if bs % granularity != 0: - raise RuntimeError(f"Batch size {bs} cannot divide granularity {granularity}.") - max_seqlen = data["attention_mask"].shape[1] - seq_lens = data["attention_mask"].sum(1).long().cpu().numpy().tolist() - input_lens = ( - data["attention_mask"] - .view(bs // granularity, granularity, -1) - .sum(dim=(1, 2)) - .long() - .cpu() - .numpy() - ) + semantic_batch_size = data["attention_mask"].shape[0] + allocation_spec = mb_spec + transport_dummy_count = 0 + target_n_mbs = max(mb_spec.n_mbs or 1, mb_spec.n_mbs_divisor) + + while True: + if allow_transport_padding: + data, added = _pad_batch_to_min_groups( + data, + min_groups=target_n_mbs, + granularity=granularity, + ) + transport_dummy_count += added + allocation_spec = MicroBatchSpec.new(mb_spec, n_mbs=target_n_mbs) + + bs = data["attention_mask"].shape[0] + if bs % granularity != 0: + raise RuntimeError( + f"Batch size {bs} cannot divide granularity {granularity}." + ) + max_seqlen = data["attention_mask"].shape[1] + seq_lens = data["attention_mask"].sum(1).long().cpu().numpy().tolist() + input_lens = ( + data["attention_mask"] + .view(bs // granularity, granularity, -1) + .sum(dim=(1, 2)) + .long() + .cpu() + .numpy() + ) + if transport_dummy_count: + input_lens[-transport_dummy_count // granularity :] = 0 + + if not allow_transport_padding: + group_indices = ( + allocate_balanced_mbs_synced(allocation_spec, input_lens, group=group) + if synchronize + else allocate_balanced_mbs(allocation_spec, input_lens) + ) + break + + group_indices = allocate_balanced_mbs(allocation_spec, input_lens) + if not synchronize or not dist.is_initialized(): + break + all_n_mbs: list[int | None] = [None] * dist.get_world_size(group) + dist.all_gather_object(all_n_mbs, len(group_indices), group=group) + synchronized_n_mbs = max(n for n in all_n_mbs if n is not None) + if all(n == synchronized_n_mbs for n in all_n_mbs): + break + target_n_mbs = synchronized_n_mbs # check for multimodal input data multimodal_keys = {key for key in data if is_multi_modal_key(key)} @@ -752,7 +856,6 @@ def split_padded_tensor_dict_into_mb_list( not_to_split[key] = value # split - group_indices = allocate_balanced_mbs_synced(mb_spec, input_lens, group=group) group_indices = [ seqpack.flat2d( [list(range(i * granularity, (i + 1) * granularity)) for i in group_index] @@ -803,19 +906,101 @@ def _split(tensor): results = [] # organize splitted micro batches assert len(mbs) == len(splitted_lens), (len(mbs), len(splitted_lens)) - for i, (mb, lens) in enumerate(zip(mbs, splitted_lens)): - results.append({**mb, **not_to_split}) + for mb, indices in zip(mbs, group_indices, strict=True): + has_transport_dummy = any(index >= semantic_batch_size for index in indices) + is_transport_dummy = has_transport_dummy and all( + index >= semantic_batch_size for index in indices + ) + if has_transport_dummy and not is_transport_dummy: + raise RuntimeError( + "Transport padding must not share a micro-batch with semantic rows" + ) + result = {**mb, **not_to_split} + if is_transport_dummy: + result[TRANSPORT_DUMMY_KEY] = True + results.append(result) return MicroBatchList( data=data, - mb_spec=mb_spec, + mb_spec=allocation_spec, mbs=results, forward_indices=forward_indices, backward_indices=backward_indices.tolist(), group_lens=group_lens, + transport_dummy_count=transport_dummy_count, ) +def split_training_batch_into_microbatches( + data: dict[str, Any], + n_mbs: int, + group: dist.ProcessGroup | None = None, +) -> list[dict[str, Any]]: + """Build a synchronized PPO schedule without all-dummy global steps.""" + if n_mbs < 1: + raise ValueError(f"n_mbs must be positive, got {n_mbs}") + batch_size = get_batch_size(data) + if batch_size < 1: + raise ValueError("Cannot split an empty training batch") + + local_n_mbs = min(batch_size, n_mbs) + local_mbs = split_padded_tensor_dict_into_mb_list( + data, + MicroBatchSpec(n_mbs=local_n_mbs), + synchronize=False, + ).mbs + if not dist.is_initialized(): + if local_n_mbs < n_mbs: + logger.warning( + "Reducing PPO minibatches from %d to %d for a batch of %d rows", + n_mbs, + local_n_mbs, + batch_size, + ) + return local_mbs + + counts: list[int | None] = [None] * dist.get_world_size(group) + dist.all_gather_object(counts, len(local_mbs), group=group) + concrete_counts = [count for count in counts if count is not None] + effective_n_mbs = max( + min(n_mbs, sum(concrete_counts)), + max(concrete_counts), + ) + if effective_n_mbs < n_mbs: + logger.warning( + "Reducing synchronized PPO minibatches from %d to %d for %d global " + "training microbatches", + n_mbs, + effective_n_mbs, + sum(concrete_counts), + ) + elif effective_n_mbs > n_mbs: + logger.warning( + "Increasing synchronized PPO minibatches from %d to %d because one " + "data-parallel rank produced that many local microbatches", + n_mbs, + effective_n_mbs, + ) + + group_rank = dist.get_rank(group=group) + offset = sum(concrete_counts[:group_rank]) + scheduled: list[dict[str, Any] | None] = [None] * effective_n_mbs + for index, microbatch in enumerate(local_mbs): + slot = (offset + index) % effective_n_mbs + if scheduled[slot] is not None: + raise RuntimeError( + "Microbatch scheduling collision at slot " + f"{slot} with {effective_n_mbs} synchronized slots" + ) + scheduled[slot] = microbatch + + dummy = make_transport_microbatch(data) + return [ + microbatch if microbatch is not None else copy.deepcopy(dummy) + for microbatch in scheduled + ] + + N_TOKENS_PER_PAGE = 256 @@ -1033,6 +1218,9 @@ def pad_mb_list( pad_value=pad_value, seq_align_to=seq_align_to, ) + padded_mb = { + key: value for key, value in padded_mb.items() if key != TRANSPORT_DUMMY_KEY + } padded_mb_inputs.append(padded_mb) pad_lengths.append(pad_len) pad_to_lengths.append(pad_to_length) @@ -1252,6 +1440,164 @@ def all_gather_tensor_container(data, group=None) -> list: return results +@dataclass(frozen=True) +class _TensorLeaf: + """Picklable stand-in for a tensor leaf inside a gathered container skeleton.""" + + shape: tuple[int, ...] + dtype: torch.dtype + device_type: str + + @property + def numel(self) -> int: + return int(np.prod(self.shape)) + + +def _deconstruct_tensor_container(value, out_tensors: list[torch.Tensor]): + """Split a container into a picklable skeleton and its tensor leaves (DFS order).""" + if torch.is_tensor(value): + out_tensors.append(value) + return _TensorLeaf(tuple(value.shape), value.dtype, value.device.type) + if isinstance(value, list): + return [_deconstruct_tensor_container(item, out_tensors) for item in value] + if isinstance(value, dict): + return { + key: _deconstruct_tensor_container(item, out_tensors) + for key, item in value.items() + } + return value + + +def _reconstruct_tensor_container(skeleton, tensors: Iterator[torch.Tensor]): + """Rebuild a container from its skeleton, consuming tensor leaves in DFS order.""" + if isinstance(skeleton, _TensorLeaf): + return next(tensors) + if isinstance(skeleton, list): + return [_reconstruct_tensor_container(item, tensors) for item in skeleton] + if isinstance(skeleton, dict): + return { + key: _reconstruct_tensor_container(item, tensors) + for key, item in skeleton.items() + } + return skeleton + + +def _skeleton_tensor_leaves(skeletons) -> list[_TensorLeaf]: + leaves: list[_TensorLeaf] = [] + + def _walk(value): + if isinstance(value, _TensorLeaf): + leaves.append(value) + elif isinstance(value, list): + for item in value: + _walk(item) + elif isinstance(value, dict): + for item in value.values(): + _walk(item) + + _walk(skeletons) + return leaves + + +def all_gather_ragged_tensor_container(items: list, group=None) -> list[list]: + """All-gather per-rank container lists whose lengths differ across ranks. + + Complements :func:`all_gather_tensor_container`, which requires every rank + to contribute the same number of items. One object all-gather exchanges + per-item skeletons (structure, non-tensor leaves, and tensor metadata); + tensor payloads then travel in one padded all-gather per (dtype, device + type) bucket. Buckets are derived from the gathered metadata, so every + rank — including ranks with no items — joins the same collectives. + """ + world_size = dist.get_world_size(group) + + local_tensors: list[torch.Tensor] = [] + local_skeletons = [ + _deconstruct_tensor_container(item, local_tensors) for item in items + ] + + all_skeletons: list[list | None] = [None] * world_size + dist.all_gather_object(all_skeletons, local_skeletons, group=group) + + leaves_by_rank = [_skeleton_tensor_leaves(skeletons) for skeletons in all_skeletons] + buckets = sorted( + { + (leaf.dtype, leaf.device_type) + for rank_leaves in leaves_by_rank + for leaf in rank_leaves + }, + key=str, + ) + + local_rank = dist.get_rank(group=group) + payloads: dict[tuple[torch.dtype, str], list[list[torch.Tensor]]] = {} + for bucket in buckets: + dtype, device_type = bucket + device = ( + torch.device("cpu") + if device_type == "cpu" + else current_platform.current_device() + ) + max_numel = max( + sum( + leaf.numel + for leaf in rank_leaves + if (leaf.dtype, leaf.device_type) == bucket + ) + for rank_leaves in leaves_by_rank + ) + local_bucket_tensors = [ + tensor + for tensor, leaf in zip( + local_tensors, leaves_by_rank[local_rank], strict=True + ) + if (leaf.dtype, leaf.device_type) == bucket + ] + flat = ( + torch.cat([tensor.reshape(-1) for tensor in local_bucket_tensors]) + if local_bucket_tensors + else torch.empty(0, dtype=dtype, device=device) + ) + padded = F.pad(flat, (0, max_numel - flat.numel())) + if max_numel > 0: + gathered = [torch.empty_like(padded) for _ in range(world_size)] + dist.all_gather(gathered, padded, group=group) + else: + # Every rank's payload is empty; slicing below yields 0-numel views. + gathered = [padded] * world_size + + bucket_payload: list[list[torch.Tensor]] = [] + for rank_leaves, buffer in zip(leaves_by_rank, gathered, strict=True): + offset = 0 + rank_tensors = [] + for leaf in rank_leaves: + if (leaf.dtype, leaf.device_type) != bucket: + continue + # Clone so results do not alias the padded gather buffers, + # which would otherwise pin world_size * max_rank_payload + # memory for the lifetime of the batch. + rank_tensors.append( + buffer.narrow(0, offset, leaf.numel).view(leaf.shape).clone() + ) + offset += leaf.numel + bucket_payload.append(rank_tensors) + payloads[bucket] = bucket_payload + + results: list[list] = [] + for rank_index, (skeletons, rank_leaves) in enumerate( + zip(all_skeletons, leaves_by_rank, strict=True) + ): + cursors = { + bucket: iter(bucket_payload[rank_index]) + for bucket, bucket_payload in payloads.items() + } + ordered = [ + next(cursors[(leaf.dtype, leaf.device_type)]) for leaf in rank_leaves + ] + results.append(_reconstruct_tensor_container(skeletons, iter(ordered))) + return results + + def broadcast_tensor_container(data, src_rank=0, group=None): if dist.get_rank() != src_rank: metadata = [None] @@ -1338,9 +1684,10 @@ def bcast_mb_list( mb_list.padding_lengths, mb_list.padded_to_lengths, mb_list.align_to_lengths, + mb_list.transport_dummy_count, ] if mb_list - else [None for _ in range(7)] + else [None for _ in range(8)] ) dist.broadcast_object_list(to_broadcast, src=src_rank, group=group) ( @@ -1351,6 +1698,7 @@ def bcast_mb_list( padding_lengths, padded_to_lengths, align_to_lengths, + transport_dummy_count, ) = to_broadcast return MicroBatchList( data=data, @@ -1364,6 +1712,7 @@ def bcast_mb_list( padded_to_lengths=padded_to_lengths, old_cu_seqlens_list=old_cu_seqlens_list, align_to_lengths=align_to_lengths, + transport_dummy_count=transport_dummy_count, ) @@ -1423,6 +1772,11 @@ def _build_group_slices( slices.append(slice(offset, offset + sz)) offset += sz return slices + if bs % self.group_size != 0: + raise ValueError( + f"batch size ({bs}) must be divisible by group_size " + f"({self.group_size}) when group_sizes is not provided" + ) return [ slice(i * self.group_size, (i + 1) * self.group_size) for i in range(bs // self.group_size) diff --git a/areal/utils/seqpack.py b/areal/utils/seqpack.py index 175461771f..0511185408 100644 --- a/areal/utils/seqpack.py +++ b/areal/utils/seqpack.py @@ -543,7 +543,9 @@ def balanced_greedy_partition(nums: list[int], K: int) -> list[list[int]]: Returns indices (not values) for each group. - Greedy with capacity-aware assignment. + Greedy with capacity-aware assignment. When ``len(nums)`` is not divisible + by ``K``, group cardinalities differ by at most one (the first + ``len(nums) % K`` groups take the extra item). Args: nums: List of values to partition @@ -553,14 +555,13 @@ def balanced_greedy_partition(nums: list[int], K: int) -> list[list[int]]: List of K lists, where each inner list contains the indices assigned to that group Raises: - ValueError: If len(nums) is not divisible by K or if len(nums) < K + ValueError: If len(nums) < K """ n = len(nums) if n < K: raise ValueError(f"Number of items ({n}) must be >= K ({K}).") - if n % K != 0: - raise ValueError("The length of nums must be divisible by K.") - m = n // K + min_items, extra_items = divmod(n, K) + capacities = [min_items + int(i < extra_items) for i in range(K)] # Sort indices by value in descending order sorted_indices = sorted(range(n), key=lambda i: -nums[i]) @@ -575,7 +576,7 @@ def balanced_greedy_partition(nums: list[int], K: int) -> list[list[int]]: chosen_group = -1 min_sum = float("inf") for i in range(K): - if counts[i] < m and sums[i] < min_sum: + if counts[i] < capacities[i] and sums[i] < min_sum: min_sum = sums[i] chosen_group = i diff --git a/areal/v2/training_service/data_proxy/dispatcher.py b/areal/v2/training_service/data_proxy/dispatcher.py index 7f829a74b6..04fe51d602 100644 --- a/areal/v2/training_service/data_proxy/dispatcher.py +++ b/areal/v2/training_service/data_proxy/dispatcher.py @@ -13,8 +13,9 @@ collectives via intra-group broadcast). - Collect results from DP heads and merge them back into the original trajectory order. -- Pad the batch to a multiple of ``dp_size * group_size`` when not - evenly divisible (eval-padding behaviour from PR 1109). +- Pad eval-route batches (``pad_eval_batch=True``) to a multiple of + ``dp_size * group_size`` (eval-padding behaviour from PR 1109); training + routes dispatch ragged shards whose cardinalities differ by at most one. Usage:: diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 2ab6520d6f..88cf64746c 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -402,6 +402,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `discount` | float | `1.0` | Discount factor for future rewards | | `gae_lambda` | float | `1.0` | Lambda parameter for GAE | | `adv_norm` | [`NormConfig`](section-norm) \| None | `None` | Normalization configuration for advantages. | +| `min_usable_group_size` | integer \| None | `None` | Minimum usable rollout slots a prompt group must keep to stay trainable when some slots fail or are filtered. None derives the minimum from reward_norm/adv_norm: 2 when either uses group statistics (1 for a singleton target group), else 1. Only the v1 rollout path consumes this option. | | `kl_ctl` | float | `0.1` | KL divergence coefficient | | `kl_estimator` | string | `"k1"` | KL divergence estimator **Choices:** `k1`, `k2`, `k3` | | `use_sapo_loss` | boolean | `False` | Use SAPO loss (mutually exclusive with PPO clipping) | diff --git a/docs/en/reference/agent_workflow.md b/docs/en/reference/agent_workflow.md index 5b24e8d61c..92d92b7014 100644 --- a/docs/en/reference/agent_workflow.md +++ b/docs/en/reference/agent_workflow.md @@ -563,6 +563,11 @@ The `ArealOpenAI` class extends `AsyncOpenAI` for direct engine integration: | `individual` | Returns all interactions as separate entries. Trajectories may share prefixes. | | `concat` | Builds conversation tree, returns only leaf nodes. Only valid for linear conversations with matched token sequences. | +`individual` exports several samples per episode, which on grouped rollouts +(`n_samples >= 2`) is incompatible with group-relative normalization and raises a +non-retryable `WorkflowContractError`; see the +[grouped-rollout contract](rollout_workflow.md#grouped-rollout). + ## Public API ```python diff --git a/docs/en/reference/rollout_workflow.md b/docs/en/reference/rollout_workflow.md index 9b0720d595..dae3d13340 100644 --- a/docs/en/reference/rollout_workflow.md +++ b/docs/en/reference/rollout_workflow.md @@ -206,13 +206,54 @@ When `group_size > 1`, the workflow is wrapped in `GroupedRolloutWorkflow`: 1. Results are merged based on their type: - **Tensor dictionaries**: Concatenated along the batch dimension - **InteractionWithTokenLogpReward dicts**: Merged into a single dictionary -1. If some runs return `None` (rejected), only valid results are kept -1. If all runs return `None`, the entire grouped result is `None` +1. Each slot returns its normal result type when usable and `None` when unusable. `None` + is intentionally opaque: classification and retry policy stay in the producer. +1. The wrapper waits only for the original slots. It neither retries an unusable slot + nor duplicates a usable result. +1. Usable slots are retained exactly once and concatenated. Their actual count remains a + prompt-group boundary during reward and advantage normalization. +1. `min_usable_group_size` defaults to `1`. The v1 RL trainer sets it to `2` when reward + or advantage normalization uses group statistics, because that statistic needs at + least two observations; a singleton target group (`n_samples: 1`) is complete by + definition and keeps the minimum of `1`. Setting `actor.min_usable_group_size` + replaces this derived value; explicit values below `2` are rejected while group + statistics are in use. Groups below the minimum return `None`; the asynchronous + collector then takes another ready prompt group. Batch-relative PPO and REINFORCE + retain a usable singleton. +1. Group statistics constrain the workflow contract. When the resolved + `min_usable_group_size` is at least `2` — derived from group statistics + (`mean_level: group` or `std_level: group`), or set explicitly via + `actor.min_usable_group_size` — and rollouts are grouped (`n_samples >= 2`), each + `arun_episode` call must contribute exactly one training sample — a tensor dict with + batch size 1, or a single exported interaction. All built-in workflows comply. A + workflow returning several samples per episode (one row per turn, tree-search + branches, or a multi-turn agent with `agent.export_style: individual`) raises a + non-retryable `WorkflowContractError` that terminates training, because group + mean/std would otherwise treat same-episode rows as independent group members. + Ungrouped rollouts (`n_samples: 1`) install no group wrapper and are not checked; a + multi-sample episode is then normalized as its own group of same-episode rows. To + train such workflows, switch `mean_level`/`std_level` to `batch` (or disable + normalization), or merge each episode into one sequence + (`agent.export_style: concat`). + +PPO-family actor loss remains globally token-weighted by default. Consequently, a +partial group with more valid response tokens contributes more loss weight than a +smaller or shorter group. This is the existing backward-compatible estimator, not an +implicit claim of equal prompt weighting. + +Grouped rollouts export `target_slot_count`, `usable_slot_count`, +`trainable_slot_count`, `fully_masked_group`, `singleton_slot_group`, +`pre_filter_usable_slot_yield`, and `pre_filter_trainable_slot_yield`. These count the +original rollout calls; final accepted and rejected counts remain collector metrics +after `should_accept_fn` runs. PPO training separately reports the physical usable +group-size and valid-token loss-weight distributions, including per-size +`group_loss_weight_size_` metrics. ### Output Shape -With `group_size=4` and a workflow returning `[1, seq_len]` tensors, the grouped output -has shape `[4, seq_len]` (4 samples concatenated). +With `group_size=4`, a workflow returning `[1, seq_len]` tensors, and all four slots +usable, the grouped output has shape `[4, seq_len]`. An incomplete accepted group uses +its actual usable count as the leading dimension. ### Implementation @@ -227,9 +268,9 @@ class GroupedRolloutWorkflow(RolloutWorkflow): for _ in range(self.group_size)] ) - # Filter None results + # A normal result is usable; None is unusable. valid_results = [r for r in results if r is not None] - if not valid_results: + if len(valid_results) < self.min_usable_group_size: return None # Merge based on result type diff --git a/docs/en/tutorial/agentic_rl.md b/docs/en/tutorial/agentic_rl.md index 1d099090fd..fe1ebf4b60 100644 --- a/docs/en/tutorial/agentic_rl.md +++ b/docs/en/tutorial/agentic_rl.md @@ -295,6 +295,12 @@ class CamelRLVRWorkflow(RolloutWorkflow): return client.export_interactions(style="individual") ``` +> **Note**: `style="individual"` exports several samples per episode. On grouped +> rollouts (`n_samples >= 2`) this is incompatible with group-relative normalization +> (`mean_level: group` / `std_level: group`) and raises a non-retryable +> `WorkflowContractError`; use `style="concat"` or batch-level normalization there. See +> the [grouped-rollout contract](../reference/rollout_workflow.md#grouped-rollout). + **Key points:** - **Parallel episode execution**: AReaL's training loop calls `arun_episode` in parallel diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 8214bce1b1..f714d52985 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -400,6 +400,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `discount` | float | `1.0` | Discount factor for future rewards | | `gae_lambda` | float | `1.0` | Lambda parameter for GAE | | `adv_norm` | [`NormConfig`](section-norm) \| None | `None` | Normalization configuration for advantages. | +| `min_usable_group_size` | integer \| None | `None` | Minimum usable rollout slots a prompt group must keep to stay trainable when some slots fail or are filtered. None derives the minimum from reward_norm/adv_norm: 2 when either uses group statistics (1 for a singleton target group), else 1. Only the v1 rollout path consumes this option. | | `kl_ctl` | float | `0.1` | KL divergence coefficient | | `kl_estimator` | string | `"k1"` | KL divergence estimator **Choices:** `k1`, `k2`, `k3` | | `use_sapo_loss` | boolean | `False` | Use SAPO loss (mutually exclusive with PPO clipping) | diff --git a/docs/zh/reference/agent_workflow.md b/docs/zh/reference/agent_workflow.md index 62c250b9e2..3621d00f2d 100644 --- a/docs/zh/reference/agent_workflow.md +++ b/docs/zh/reference/agent_workflow.md @@ -541,6 +541,10 @@ class OpenAIProxyWorkflow(RolloutWorkflow): | `individual` | 将所有交互作为单独条目返回。轨迹可能共享前缀。 | | `concat` | 构建对话树,仅返回叶子节点。仅对具有匹配 token 序列的线性对话有效。 | +`individual` 每个 episode 会导出多个样本;在分组 rollout(`n_samples >= 2`)下,这与 group-relative +normalization 不兼容,会抛出不可重试的 `WorkflowContractError`。参见 +[分组 rollout 契约](rollout_workflow.md#%E5%88%86%E7%BB%84-rollout)。 + ## 公共 API ```python diff --git a/docs/zh/reference/rollout_workflow.md b/docs/zh/reference/rollout_workflow.md index da3a791189..e40322790a 100644 --- a/docs/zh/reference/rollout_workflow.md +++ b/docs/zh/reference/rollout_workflow.md @@ -196,12 +196,39 @@ rollout: 1. 根据类型合并结果: - **张量字典**:沿批次维度连接 - **InteractionWithTokenLogpReward 字典**:合并为单个字典 -1. 如果某些运行返回 `None`(拒绝),仅保留有效结果 -1. 如果所有运行都返回 `None`,则整个分组结果为 `None` +1. 每个 slot 可用时返回正常结果类型,不可用时返回 `None`。`None` 有意保持不透明;原因分类与重试策略仍由 producer 负责。 +1. 包装器只等待最初提交的 slots,既不会重试不可用 slot,也不会复制可用结果。 +1. 可用 slots 会各保留一次并连接;实际数量会在 reward 和 advantage normalization 中继续作为 prompt-group 边界。 +1. `min_usable_group_size` 默认为 `1`。仅当 v1 RL trainer 的 reward 或 advantage normalization + 使用 group statistics(该统计量至少需要两个观测值)时,才会将其设为 `2`;singleton 目标 group + (`n_samples: 1`)本身即是完整的,因此下限保持为 `1`。设置 `actor.min_usable_group_size` 可覆盖该推导值;使用 group + statistics 时,低于 `2` 的显式值会被拒绝。低于下限的 group 返回 `None`,异步 collector 随后会接收另一个已就绪的 prompt + group。采用 batch-relative PPO 或 REINFORCE 时则会保留可用的 singleton。 +1. Group statistics 会约束 workflow 契约:当解析后的 `min_usable_group_size` 至少为 `2` ——由 group + statistics(`mean_level: group` 或 `std_level: group`)推导,或通过 + `actor.min_usable_group_size` 显式设置——且 rollout 分组(`n_samples >= 2`)时,每次 `arun_episode` + 调用必须恰好贡献一个训练样本——batch size 为 1 的张量字典,或单个导出的 interaction。所有内置 workflow 均满足该契约。如果 + workflow 每个 episode 返回多个样本(每 turn 一行、tree-search 分支,或使用 + `agent.export_style: individual` 的多轮 agent),会抛出不可重试的 `WorkflowContractError` 并终止训练,因为 + group mean/std 会把同一 episode 的多行当作独立的 group 成员。未分组的 rollout(`n_samples: 1`)不会安装 group + wrapper,因此不做该检查;此时多样本 episode 会作为由同一 episode 各行组成的独立 group 参与 normalization。若要训练此类 + workflow,请将 `mean_level`/`std_level` 改为 `batch`(或关闭 normalization),或把每个 episode + 合并为单条序列(`agent.export_style: concat`)。 + +PPO 系列 actor loss 默认仍按全局 token 加权。因此,有更多有效 response tokens 的 partial group 会比更小或更短的 +group 获得更高 loss weight。这是保持向后兼容的现有 estimator,并不表示各 prompt 隐式等权。 + +Grouped rollout 会导出 `target_slot_count`、`usable_slot_count`、 +`trainable_slot_count`、`fully_masked_group`、`singleton_slot_group`、 +`pre_filter_usable_slot_yield` 和 `pre_filter_trainable_slot_yield`。这些指标统计最初的 rollout +调用;`should_accept_fn` 运行后的最终 accepted 和 rejected 数量仍由 collector metrics 提供。PPO 训练会另外报告物理 +usable group size 与有效 token 的 loss-weight 分布,其中包括按 size 区分的 `group_loss_weight_size_` +指标。 ### 输出形状 -当 `group_size=4` 且工作流返回 `[1, seq_len]` 张量时,分组输出的形状为 `[4, seq_len]`(4 个样本连接)。 +当 `group_size=4`、工作流返回 `[1, seq_len]` 张量且 4 个 slot 均可用时,分组输出的形状为 +`[4, seq_len]`。如果接受的是不完整 group,则第一维是实际可用的 slot 数量。 ### 实现 @@ -216,9 +243,9 @@ class GroupedRolloutWorkflow(RolloutWorkflow): for _ in range(self.group_size)] ) - # 过滤 None 结果 + # 正常结果表示可用,None 表示不可用 valid_results = [r for r in results if r is not None] - if not valid_results: + if len(valid_results) < self.min_usable_group_size: return None # 根据结果类型合并 diff --git a/docs/zh/tutorial/agentic_rl.md b/docs/zh/tutorial/agentic_rl.md index 7d9c854f5e..45631c065a 100644 --- a/docs/zh/tutorial/agentic_rl.md +++ b/docs/zh/tutorial/agentic_rl.md @@ -271,6 +271,11 @@ class CamelRLVRWorkflow(RolloutWorkflow): return client.export_interactions(style="individual") ``` +> **注意**:`style="individual"` 每个 episode 会导出多个样本。在分组 rollout (`n_samples >= 2`)下,这与 +> group-relative normalization(`mean_level: group` / `std_level: group`)不兼容,会抛出不可重试的 +> `WorkflowContractError`;此时请改用 `style="concat"` 或 batch 级别的 normalization。参见 +> [分组 rollout 契约](../reference/rollout_workflow.md#%E5%88%86%E7%BB%84-rollout)。 + **关键点:** - **并行 episode 执行**:AReaL 的训练循环在多个样本上并行调用 `arun_episode` diff --git a/examples/openclaw/config.yaml b/examples/openclaw/config.yaml index 5075c17bd4..426d540c84 100644 --- a/examples/openclaw/config.yaml +++ b/examples/openclaw/config.yaml @@ -80,9 +80,8 @@ actor: metric: ratio upper: 5.0 reward_norm: - mean_level: group - std_level: group - group_size: ${gconfig.n_samples} + mean_level: batch + std_level: batch adv_norm: mean_level: batch std_level: batch diff --git a/tests/test_data_redistribution.py b/tests/test_data_redistribution.py index 9aa5584535..d7c445eded 100644 --- a/tests/test_data_redistribution.py +++ b/tests/test_data_redistribution.py @@ -1,3 +1,4 @@ +import os import pickle import subprocess import sys @@ -28,6 +29,83 @@ def assert_tensor_container_close(x1, x2): assert x1 == x2 +def _assert_redistribution_outputs(tmp_path, world_size): + redistributed_data = [] + for i in range(world_size): + with open(tmp_path / f"redistributed{i}.pkl", "rb") as f: + redistributed_data.append(pickle.load(f)) + + for x in redistributed_data[1:]: + assert_tensor_container_close(x.all_data, redistributed_data[0].all_data) + for x in redistributed_data[1:]: + assert_tensor_container_close( + x.group_indices, redistributed_data[0].group_indices + ) + + all_data = redistributed_data[0].all_data + group_indices = redistributed_data[0].group_indices + for x, indices in zip(redistributed_data, group_indices): + expected_data = [all_data[i] for i in indices] + assert_tensor_container_close(x.data, expected_data) + + +@pytest.mark.parametrize("ragged_flag", ["--ragged", "--ragged-empty"]) +def test_redistribute_ragged_cpu(tmp_path, ragged_flag): + world_size = 2 + port = find_free_ports(1)[0] + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={world_size}", + "--nnodes=1", + "--master-addr=localhost", + f"--master_port={port}", + "tests/torchrun/redistribute.py", + f"--dump-path={str(tmp_path)}", + "--backend=gloo", + ragged_flag, + ], + check=True, + text=True, + stderr=sys.stdout, + stdout=sys.stdout, + env={**os.environ, "PYTHONPATH": os.getcwd()}, + timeout=60, + ) + + _assert_redistribution_outputs(tmp_path, world_size) + with open(tmp_path / "redistributed0.pkl", "rb") as f: + assert len(pickle.load(f).all_data) == 3 + + +def test_hybrid_parallel_redistribution_error_is_collective_safe(tmp_path): + world_size = 4 + port = find_free_ports(1)[0] + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={world_size}", + "--nnodes=1", + "--master-addr=localhost", + f"--master_port={port}", + "tests/torchrun/redistribute.py", + f"--dump-path={str(tmp_path)}", + "--backend=gloo", + "--hybrid-error", + ], + check=True, + text=True, + stderr=sys.stdout, + stdout=sys.stdout, + env={**os.environ, "PYTHONPATH": os.getcwd()}, + timeout=60, + ) + + @pytest.mark.skipif(is_in_ci(), reason="CI machine will crash with all_gather_object") @pytest.mark.multi_gpu @pytest.mark.parametrize("world_size", [2, 4, 8]) @@ -54,20 +132,4 @@ def test_redistribute(world_size, tmp_path): except subprocess.CalledProcessError as e: pytest.fail(f"Test failed with error: {e.stderr}") - redistributed_data = [] - for i in range(world_size): - with open(tmp_path / f"redistributed{i}.pkl", "rb") as f: - redistributed_data.append(pickle.load(f)) - - for x in redistributed_data[1:]: - assert_tensor_container_close(x.all_data, redistributed_data[0].all_data) - for x in redistributed_data[1:]: - assert_tensor_container_close( - x.group_indices, redistributed_data[0].group_indices - ) - - all_data = redistributed_data[0].all_data - group_indices = redistributed_data[0].group_indices - for x, indices in zip(redistributed_data, group_indices): - expected_data = [all_data[i] for i in indices] - assert_tensor_container_close(x.data, expected_data) + _assert_redistribution_outputs(tmp_path, world_size) diff --git a/tests/test_eval_dispatch.py b/tests/test_eval_dispatch.py index 30c40706fa..ac2b83d96f 100644 --- a/tests/test_eval_dispatch.py +++ b/tests/test_eval_dispatch.py @@ -99,10 +99,15 @@ def test_pad_eval_batch_pads_when_n_less_than_dp(self): assert len(padded) == 4 assert _count_dummies(padded) == 2 - def test_dispatch_tensors_raises_when_not_divisible(self): + def test_dispatch_tensors_accepts_non_divisible_count(self): items = [_make_item(i) for i in range(7)] - with pytest.raises(ValueError, match="divisible"): - _dispatch_tensors(items, dp_size=4) + splits, _ = _dispatch_tensors(items, dp_size=4) + + assert sorted(len(group) for group in splits) == [1, 2, 2, 2] + + def test_dispatch_tensors_rejects_fewer_groups_than_ranks(self): + with pytest.raises(ValueError, match="item group count \\(2\\)"): + _dispatch_tensors([_make_item(0), _make_item(1)], dp_size=4) def test_pad_then_dispatch_end_to_end(self): items = [_make_item(i) for i in range(7)] diff --git a/tests/test_grouped_rollout_workflow.py b/tests/test_grouped_rollout_workflow.py index cee294d142..5adce40061 100644 --- a/tests/test_grouped_rollout_workflow.py +++ b/tests/test_grouped_rollout_workflow.py @@ -128,7 +128,7 @@ def test_dist_rollout_coordinator_forwards_reward_group_flags(monkeypatch): monkeypatch.setattr( coordinator, "_broadcast_and_redistribute_trajectories", - lambda trajectories: trajectories, + lambda trajectories, preparation_error=None: trajectories, ) monkeypatch.setattr(dist_rollout.current_platform, "current_device", lambda: "cpu") monkeypatch.setattr(dist_rollout, "tensor_container_to", lambda data, device: data) @@ -142,6 +142,7 @@ def test_dist_rollout_coordinator_forwards_reward_group_flags(monkeypatch): coordinator.rollout_batch( data=[{}], workflow=object(), + min_usable_group_size=2, reward_normalization=True, drop_incomplete_group=True, ) @@ -150,3 +151,80 @@ def test_dist_rollout_coordinator_forwards_reward_group_flags(monkeypatch): assert rollout_engine.prepare_kwargs["drop_incomplete_group"] is True assert rollout_engine.rollout_kwargs["reward_normalization"] is True assert rollout_engine.rollout_kwargs["drop_incomplete_group"] is True + assert rollout_engine.rollout_kwargs["min_usable_group_size"] == 2 + + +def test_ragged_trajectory_gather_routes_to_batched_ragged_gather(monkeypatch): + monkeypatch.setattr(dist_rollout.dist, "get_world_size", lambda _: 2) + monkeypatch.setattr( + dist_rollout.dist, + "all_gather_object", + lambda output, value, group=None: output.__setitem__(slice(None), [1, 2]), + ) + expected = [ + [{"rank": 0, "item": 0}], + [{"rank": 1, "item": 0}, {"rank": 1, "item": 1}], + ] + monkeypatch.setattr( + dist_rollout, + "all_gather_ragged_tensor_container", + lambda items, group=None: expected, + ) + + gathered = dist_rollout._all_gather_ragged_trajectory_lists( + [{"rank": 0, "item": 0}], group=object() + ) + + assert gathered == expected + + +def test_equal_trajectory_gather_keeps_fast_path(monkeypatch): + monkeypatch.setattr(dist_rollout.dist, "get_world_size", lambda _: 2) + monkeypatch.setattr( + dist_rollout.dist, + "all_gather_object", + lambda output, value, group=None: output.__setitem__(slice(None), [1, 1]), + ) + expected = [[{"rank": 0}], [{"rank": 1}]] + monkeypatch.setattr( + dist_rollout, + "all_gather_tensor_container", + lambda items, group=None: expected, + ) + + gathered = dist_rollout._all_gather_ragged_trajectory_lists( + [{"rank": 0}], group=object() + ) + + assert gathered == expected + + +def test_stalled_dynamic_preparation_fails_instead_of_spinning(monkeypatch): + class _EmptyRolloutEngine: + def __init__(self): + self.calls = 0 + + def prepare_batch(self, *args, **kwargs): + self.calls += 1 + return [] + + rollout_engine = _EmptyRolloutEngine() + coordinator = DistRolloutCoordinator(rollout_engine, _TrainEngine()) + monkeypatch.setattr(dist_rollout, "ROLLOUT_COLLECTION_STALL_TIMEOUT_SECONDS", 0.0) + captured = {} + monkeypatch.setattr( + coordinator, + "_broadcast_and_redistribute_trajectories", + lambda trajectories, preparation_error=None: captured.setdefault( + "error", preparation_error + ), + ) + + coordinator.prepare_batch( + dataloader=object(), + workflow=object(), + dynamic_bs=True, + ) + + assert rollout_engine.calls == 8 + assert "added no trainable group" in captured["error"] diff --git a/tests/test_incomplete_rollout_groups.py b/tests/test_incomplete_rollout_groups.py new file mode 100644 index 0000000000..f529aa27ef --- /dev/null +++ b/tests/test_incomplete_rollout_groups.py @@ -0,0 +1,369 @@ +from unittest.mock import MagicMock + +import pytest +import torch +from omegaconf import DictConfig + +from areal.api.cli_args import InferenceEngineConfig, NormConfig, PPOActorConfig +from areal.experimental.openai import InteractionWithTokenLogpReward +from areal.infra import workflow_context +from areal.infra.remote_inf_engine import GroupedRolloutWorkflow +from areal.infra.workflow_context import WorkflowContext +from areal.infra.workflow_executor import ( + WorkflowContractError, + WorkflowExecutor, + validate_rollout_group_sizes, +) +from areal.trainer.ppo.actor import _group_training_metrics +from areal.trainer.rl_trainer import _collect_trainable_rollout_batch +from areal.utils.functional import ppo_actor_loss_fn + + +class _SequenceWorkflow: + def __init__(self, results: list[dict | None]): + self.results = iter(results) + self.calls = 0 + + async def arun_episode(self, engine, data): + self.calls += 1 + return next(self.results) + + +class _CyclingDataLoader: + batch_size = 4 + + def __iter__(self): + return iter([[{}]]) + + +def _trajectory(token: int, batch_size: int = 1) -> dict[str, torch.Tensor]: + return { + "input_ids": torch.full((batch_size, 1), token), + "attention_mask": torch.ones(batch_size, 1, dtype=torch.bool), + } + + +@pytest.mark.parametrize( + ("group_size", "min_usable_group_size"), + [(0, 1), (2, 0), (2, 3)], +) +def test_rollout_group_size_contract_rejects_invalid_threshold( + group_size, min_usable_group_size +): + with pytest.raises(ValueError): + validate_rollout_group_sizes(group_size, min_usable_group_size) + + +@pytest.fixture +def rollout_stats(monkeypatch): + tracker = MagicMock() + monkeypatch.setattr( + "areal.infra.remote_inf_engine.stats_tracker.get", lambda _: tracker + ) + workflow_context.set(WorkflowContext(is_eval=False)) + return tracker + + +@pytest.mark.asyncio +async def test_grouped_rollout_keeps_each_usable_member_once(rollout_stats): + workflow = _SequenceWorkflow([_trajectory(11), None, _trajectory(22)]) + grouped = GroupedRolloutWorkflow(workflow, group_size=3, logger=MagicMock()) + + result = await grouped.arun_episode(MagicMock(), {}) + + assert result is not None + assert workflow.calls == 3 + torch.testing.assert_close(result["input_ids"], torch.tensor([[11], [22]])) + rollout_stats.scalar.assert_called_once_with( + target_slot_count=3, + usable_slot_count=2, + trainable_slot_count=2, + fully_masked_group=False, + singleton_slot_group=False, + pre_filter_usable_slot_yield=2 / 3, + pre_filter_trainable_slot_yield=2 / 3, + ) + + +@pytest.mark.asyncio +async def test_grouped_rollout_marks_empty_group_unusable(rollout_stats): + grouped = GroupedRolloutWorkflow( + _SequenceWorkflow([None, None, None]), + group_size=3, + logger=MagicMock(), + ) + + result = await grouped.arun_episode(MagicMock(), {}) + + assert result is None + rollout_stats.scalar.assert_called_once_with( + target_slot_count=3, + usable_slot_count=0, + trainable_slot_count=0, + fully_masked_group=True, + singleton_slot_group=False, + pre_filter_usable_slot_yield=0.0, + pre_filter_trainable_slot_yield=0.0, + ) + + +@pytest.mark.asyncio +async def test_grouped_rollout_applies_estimator_minimum(rollout_stats): + grouped = GroupedRolloutWorkflow( + _SequenceWorkflow([None, _trajectory(11), None]), + group_size=3, + logger=MagicMock(), + min_usable_group_size=2, + ) + + result = await grouped.arun_episode(MagicMock(), {}) + + assert result is None + rollout_stats.scalar.assert_called_once_with( + target_slot_count=3, + usable_slot_count=1, + trainable_slot_count=0, + fully_masked_group=False, + singleton_slot_group=True, + pre_filter_usable_slot_yield=1 / 3, + pre_filter_trainable_slot_yield=0.0, + ) + + +@pytest.mark.asyncio +async def test_grouped_rollout_keeps_singleton_by_default(rollout_stats): + grouped = GroupedRolloutWorkflow( + _SequenceWorkflow([None, _trajectory(11)]), + group_size=2, + logger=MagicMock(), + ) + + result = await grouped.arun_episode(MagicMock(), {}) + + assert result is not None + torch.testing.assert_close(result["input_ids"], torch.tensor([[11]])) + + +@pytest.mark.asyncio +async def test_group_relative_tensor_slot_must_produce_one_member(rollout_stats): + grouped = GroupedRolloutWorkflow( + _SequenceWorkflow([_trajectory(11, batch_size=2), _trajectory(22)]), + group_size=2, + logger=MagicMock(), + min_usable_group_size=2, + ) + + with pytest.raises(WorkflowContractError, match="slot sizes \\[2, 1\\]"): + await grouped.arun_episode(MagicMock(), {}) + + +@pytest.mark.asyncio +async def test_group_relative_interaction_slot_must_produce_one_member(rollout_stats): + grouped = GroupedRolloutWorkflow( + _SequenceWorkflow( + [ + { + "first": InteractionWithTokenLogpReward(), + "second": InteractionWithTokenLogpReward(), + }, + {"third": InteractionWithTokenLogpReward()}, + ] + ), + group_size=2, + logger=MagicMock(), + min_usable_group_size=2, + ) + + with pytest.raises(WorkflowContractError, match="slot sizes \\[2, 1\\]"): + await grouped.arun_episode(MagicMock(), {}) + + +@pytest.mark.parametrize("dynamic_bs", [False, True]) +def test_executor_propagates_group_contract_error_without_retry(dynamic_bs): + workflow = _SequenceWorkflow( + [_trajectory(11, batch_size=2), _trajectory(22, batch_size=2)] + ) + grouped = GroupedRolloutWorkflow( + workflow, + group_size=2, + logger=MagicMock(), + min_usable_group_size=2, + ) + inference_engine = MagicMock() + inference_engine.get_version.return_value = 0 + executor = WorkflowExecutor( + InferenceEngineConfig( + backend="sglang:d1", + consumer_batch_size=1, + max_concurrent_rollouts=1, + ), + inference_engine, + ) + executor.initialize() + + try: + with pytest.raises(WorkflowContractError, match="slot sizes \\[2, 2\\]"): + executor.prepare_batch( + _CyclingDataLoader(), + grouped, + dynamic_bs=dynamic_bs, + ) + finally: + executor.destroy() + + assert workflow.calls == 2 + + +def test_group_training_metrics_report_actual_size_and_token_weight(): + loss_mask = torch.tensor( + [ + [1, 1, 0], + [1, 0, 0], + [1, 1, 1], + [1, 1, 0], + [1, 0, 0], + ], + dtype=torch.bool, + ) + + group_starts, group_sizes, loss_weights = _group_training_metrics(loss_mask, [2, 3]) + + torch.testing.assert_close( + group_starts, torch.tensor([True, False, True, False, False]) + ) + torch.testing.assert_close(group_sizes, torch.tensor([2.0, 0.0, 3.0, 0.0, 0.0])) + torch.testing.assert_close(loss_weights, torch.tensor([3.0, 0.0, 6.0, 0.0, 0.0])) + + +def test_actor_loss_keeps_existing_token_weighted_group_reduction(): + advantages = torch.tensor([[1.0], [1.0], [3.0], [3.0], [3.0]]) + zeros = torch.zeros_like(advantages) + + loss, _ = ppo_actor_loss_fn( + logprobs=zeros, + proximal_logprobs=zeros, + old_logprobs=zeros, + advantages=advantages, + eps_clip=0.2, + loss_mask=torch.ones_like(advantages, dtype=torch.bool), + ) + + torch.testing.assert_close(loss, torch.tensor(-2.2), rtol=0.0, atol=1e-6) + + +@pytest.mark.parametrize( + "normalization", + [ + NormConfig(mean_level="group", std_level=None, group_size=4), + NormConfig(mean_level=None, std_level="group", group_size=4), + {"mean_level": "group", "std_level": None, "group_size": 4}, + DictConfig({"mean_level": None, "std_level": "group", "group_size": 4}), + ], +) +def test_min_usable_group_size_defaults_to_estimator_minimum(normalization): + group_relative = PPOActorConfig(adv_norm=normalization) + + assert group_relative.resolve_min_usable_group_size(target_group_size=4) == 2 + assert group_relative.resolve_min_usable_group_size(target_group_size=1) == 1 + + +def test_batch_relative_estimator_keeps_usable_singleton(): + actor = PPOActorConfig(adv_norm=NormConfig(mean_level="batch", std_level="batch")) + + assert actor.resolve_min_usable_group_size(target_group_size=4) == 1 + + +def test_explicit_min_usable_group_size_overrides_derivation(): + actor = PPOActorConfig( + adv_norm=NormConfig(mean_level="group", std_level="group", group_size=4), + min_usable_group_size=3, + ) + + assert actor.resolve_min_usable_group_size(target_group_size=4) == 3 + + +def test_group_statistics_reject_explicit_singleton_minimum(): + with pytest.raises(ValueError, match="at least 2"): + PPOActorConfig( + adv_norm=NormConfig(mean_level="group", std_level="group", group_size=4), + min_usable_group_size=1, + ) + + +def test_batch_relative_estimator_accepts_explicit_singleton_minimum(): + actor = PPOActorConfig( + adv_norm=NormConfig(mean_level="batch", std_level="batch"), + min_usable_group_size=1, + ) + + assert actor.resolve_min_usable_group_size(target_group_size=4) == 1 + + +def test_dynamic_collection_backfills_from_ready_groups(): + prepare_batch = MagicMock( + side_effect=[[{"group": "first"}], [], [{"group": "second"}]] + ) + + result = _collect_trainable_rollout_batch( + prepare_batch, + dynamic_bs=True, + min_batch_size=2, + ) + + assert result == [{"group": "first"}, {"group": "second"}] + assert prepare_batch.call_count == 3 + + +def test_dynamic_collection_aborts_after_consecutive_empty_rounds(): + prepare_batch = MagicMock(return_value=[]) + + with pytest.raises(RuntimeError, match="added no trainable group"): + _collect_trainable_rollout_batch( + prepare_batch, + dynamic_bs=True, + min_batch_size=1, + stall_timeout=0.0, + ) + + assert prepare_batch.call_count == 8 + + +def test_dynamic_collection_empty_streak_resets_on_progress(): + prepare_batch = MagicMock( + side_effect=[[], [{"group": "first"}], [], [{"group": "second"}]] + ) + + result = _collect_trainable_rollout_batch( + prepare_batch, + dynamic_bs=True, + min_batch_size=2, + max_empty_rounds=2, + stall_timeout=0.0, + ) + + assert result == [{"group": "first"}, {"group": "second"}] + assert prepare_batch.call_count == 4 + + +def test_dynamic_collection_tolerates_fast_all_reject_bursts(): + # A staleness flush drains many rejected rounds in near-zero time; the + # round bound alone must not abort while the stall timeout has not run. + prepare_batch = MagicMock(side_effect=[[]] * 20 + [[{"group": "fresh"}]]) + + result = _collect_trainable_rollout_batch( + prepare_batch, + dynamic_bs=True, + min_batch_size=1, + ) + + assert result == [{"group": "fresh"}] + assert prepare_batch.call_count == 21 + + +def test_fixed_collection_rejects_undersized_batch(): + with pytest.raises(RuntimeError, match="only 1 trainable groups"): + _collect_trainable_rollout_batch( + MagicMock(return_value=[{"group": "only"}]), + dynamic_bs=False, + min_batch_size=2, + ) diff --git a/tests/test_reward_norm_variable_group.py b/tests/test_reward_norm_variable_group.py index c6ce65232f..4757c7f134 100644 --- a/tests/test_reward_norm_variable_group.py +++ b/tests/test_reward_norm_variable_group.py @@ -127,6 +127,27 @@ def test_leave_one_out_singleton_group_outputs_zero(): assert out[0].item() == 0.0 +def test_leave_one_out_uses_only_usable_group_members(): + norm = Normalization( + NormConfig( + mean_level="group", + mean_leave1out=True, + std_level=None, + group_size=3, + ) + ) + rewards = torch.tensor([1.0, 3.0, 10.0, 20.0, 30.0]) + + result = norm(rewards, group_sizes=[2, 3]) + + torch.testing.assert_close( + result, + torch.tensor([-2.0, 2.0, -15.0, 0.0, 15.0]), + rtol=0.0, + atol=0.0, + ) + + def test_group_sizes_sum_mismatch_raises(): """Boundaries whose sizes do not sum to the batch size are rejected.""" norm = Normalization(_group_norm_config()) @@ -143,6 +164,14 @@ def test_group_sizes_non_positive_raises(): norm(x, group_sizes=[4, 0]) +def test_missing_group_sizes_rejects_ragged_batch(): + norm = Normalization(_group_norm_config()) + x = torch.tensor([0.0, 1.0, 2.0], dtype=torch.float32) + + with pytest.raises(ValueError, match="must be divisible by group_size"): + norm(x) + + def test_adv_norm_style_2d_variable_groups_normalize_per_group(): """Token-level (2D) advantages normalize per variable-size group on dim 0. diff --git a/tests/test_rollout_controller.py b/tests/test_rollout_controller.py index 749a4d193d..98899f7a44 100644 --- a/tests/test_rollout_controller.py +++ b/tests/test_rollout_controller.py @@ -21,6 +21,10 @@ ) from areal.infra import RolloutController from areal.infra.scheduler.local import LocalScheduler +from areal.infra.workflow_executor import ( + WorkflowContractError, + WorkflowContractFailure, +) from areal.utils.hf_utils import load_hf_tokenizer @@ -47,6 +51,7 @@ def __init__(self): self.engine_calls = [] self._pending_results = {} # worker_id -> dict[task_id -> result] self._task_counter = 0 + self.workflow_contract_error = None def create_workers(self, job, *args, **kwargs): """Create workers based on Job specification.""" @@ -100,15 +105,19 @@ async def async_call_engine(self, worker_id, method, *args, **kwargs): resp = requests.post(callback_addr, json=dict(task_id=task_id)) resp.raise_for_status() return task_id - # Handle wait_for_task method - elif method == "wait_for_task": + # Handle controller-safe task result retrieval + elif method == "_wait_for_task_result": task_id = kwargs.get("task_id") + if self.workflow_contract_error is not None: + return WorkflowContractFailure(message=self.workflow_contract_error) if ( worker_id in self._pending_results and task_id in self._pending_results[worker_id] ): result = self._pending_results[worker_id].pop(task_id) - return result + from areal.infra.workflow_executor import _RolloutResult + + return _RolloutResult(task_id=task_id, trajectory=result) return None elif method == "wait": # Return a result from pending results if available @@ -177,6 +186,13 @@ def __name__(cls): return "MockInferenceEngine" +class _CyclingDataLoader: + batch_size = 4 + + def __iter__(self): + return iter([[{"id": 0}]]) + + class TestRolloutControllerInitialization: def test_constructor(self): config = create_test_config(consumer_batch_size=16) @@ -486,6 +502,38 @@ def test_submit_passes_is_eval_and_group_size(self): class TestRolloutControllerBatchOperations: + @pytest.mark.parametrize("dynamic_bs", [False, True]) + def test_prepare_batch_propagates_workflow_contract_error(self, dynamic_bs): + config = create_test_config( + backend="sglang:d1", + consumer_batch_size=1, + max_concurrent_rollouts=1, + ) + scheduler = MockScheduler() + scheduler.workflow_contract_error = "logical slot produced two members" + controller = RolloutController( + inf_engine=MockInferenceEngine, + config=config, + scheduler=scheduler, + ) + controller.initialize(role="rollout", server_args={}) + + try: + with pytest.raises( + WorkflowContractError, match="logical slot produced two members" + ): + controller.prepare_batch( + _CyclingDataLoader(), + workflow="tests.utils.TestWorkflow", + workflow_kwargs={}, + dynamic_bs=dynamic_bs, + ) + finally: + controller.destroy() + + submit_calls = [call for call in scheduler.engine_calls if call[1] == "submit"] + assert len(submit_calls) == 1 + def test_rollout_batch_returns_list_of_dicts(self): """Verify RolloutController returns list of regular dicts, NOT RTensors. diff --git a/tests/test_seqpack.py b/tests/test_seqpack.py index 02cf3f9657..6ee6e3d7d5 100644 --- a/tests/test_seqpack.py +++ b/tests/test_seqpack.py @@ -221,13 +221,15 @@ def test_equal_group_sizes(self): for g in groups: assert len(g) == expected_size - def test_raises_on_non_divisible(self): - """Test error when n is not divisible by K.""" + def test_non_divisible_uses_nearly_equal_group_sizes(self): + """Test ragged partition cardinalities differ by at most one.""" nums = [1, 2, 3, 4, 5] K = 2 - with pytest.raises(ValueError, match="must be divisible by K"): - balanced_greedy_partition(nums, K) + groups = balanced_greedy_partition(nums, K) + + assert sorted(len(group) for group in groups) == [2, 3] + assert sorted(index for group in groups for index in group) == list(range(5)) def test_raises_on_too_few_items(self): """Test error when n < K.""" diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 94193fcc8d..673986cfa3 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -12,6 +12,7 @@ from tests.utils import get_model_path from areal.infra.rpc.serialization import deserialize_value, serialize_value +from areal.infra.workflow_executor import WorkflowContractFailure @dataclass @@ -117,6 +118,14 @@ def test_dataclass(self): assert deserialized.value == original.value assert torch.equal(deserialized.tensor, original.tensor) + def test_workflow_contract_failure(self): + """Workflow contract failures retain their terminal error identity.""" + original = WorkflowContractFailure(message="invalid group shape") + + deserialized = deserialize_value(serialize_value(original)) + + assert deserialized == original + def test_tokenizer(self): """Test Hugging Face tokenizer serialization.""" original = AutoTokenizer.from_pretrained( diff --git a/tests/test_train_controller.py b/tests/test_train_controller.py index 45ea2068f0..7ec381bc7c 100644 --- a/tests/test_train_controller.py +++ b/tests/test_train_controller.py @@ -593,6 +593,8 @@ def test_prepare_batch_delegates_to_rollout(self, train_controller, ft_spec): dataloader=mock_dataloader, workflow="test.workflow", workflow_kwargs={"key": "value"}, + group_size=2, + min_usable_group_size=2, ) mock_rollout.prepare_batch.assert_called_once_with( @@ -601,9 +603,10 @@ def test_prepare_batch_delegates_to_rollout(self, train_controller, ft_spec): workflow_kwargs={"key": "value"}, should_accept_fn=None, dynamic_bs=False, - group_size=1, + group_size=2, reward_normalization=False, drop_incomplete_group=False, + min_usable_group_size=2, ) def test_rollout_batch_delegates_to_rollout(self, train_controller, ft_spec): @@ -623,6 +626,8 @@ def test_rollout_batch_delegates_to_rollout(self, train_controller, ft_spec): data=data, workflow="test.workflow", workflow_kwargs={"key": "value"}, + group_size=2, + min_usable_group_size=2, ) mock_rollout.rollout_batch.assert_called_once_with( @@ -630,7 +635,8 @@ def test_rollout_batch_delegates_to_rollout(self, train_controller, ft_spec): workflow="test.workflow", workflow_kwargs={"key": "value"}, should_accept_fn=None, - group_size=1, + group_size=2, + min_usable_group_size=2, reward_normalization=False, drop_incomplete_group=False, ) diff --git a/tests/test_tree_transport.py b/tests/test_tree_transport.py new file mode 100644 index 0000000000..cd5000fcb2 --- /dev/null +++ b/tests/test_tree_transport.py @@ -0,0 +1,50 @@ +import torch +import torch.distributed as dist + +from areal.api.cli_args import MicroBatchSpec +from areal.engine.core.train_engine import compute_microbatch_loss_weight +from areal.models.tree_attn.tree import build_packed_tree_batch +from areal.utils.data import TRANSPORT_DUMMY_KEY + + +def test_tree_transport_dummy_bypasses_objective_weight(monkeypatch): + data = { + "input_ids": torch.arange(4).view(1, 4), + "attention_mask": torch.ones(1, 4, dtype=torch.bool), + "loss_mask": torch.ones(1, 4, dtype=torch.bool), + } + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda _group=None: 2) + + def _all_gather(outputs, local_count, group=None): + del group + outputs[0].copy_(local_count) + outputs[1].fill_(2) + + monkeypatch.setattr(dist, "all_gather", _all_gather) + + mb_list = build_packed_tree_batch( + data, + MicroBatchSpec(max_tokens_per_mb=128), + ) + semantic_mb, transport_mb = mb_list.mbs + + assert TRANSPORT_DUMMY_KEY not in semantic_mb + assert transport_mb[TRANSPORT_DUMMY_KEY] is True + assert mb_list.padded_mbs is not None + assert TRANSPORT_DUMMY_KEY not in mb_list.padded_mbs[1] + + callback_called = False + + def _loss_weight(_microbatch): + nonlocal callback_called + callback_called = True + return torch.tensor(1.0) + + torch.testing.assert_close( + compute_microbatch_loss_weight(transport_mb, _loss_weight), + torch.tensor(0.0), + rtol=0.0, + atol=0.0, + ) + assert callback_called is False diff --git a/tests/test_utils.py b/tests/test_utils.py index 0c82b47517..18fbcf3f24 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,13 +1,24 @@ import pytest import torch +import areal.utils.data as data_module from areal.api.cli_args import MicroBatchSpec +from areal.engine.core.train_engine import ( + compute_microbatch_loss_weight, + reorder_and_pad_outputs, +) +from areal.trainer.dpo.dpo_engine import _dpo_loss_weight +from areal.trainer.rw.rw_engine import _rw_loss_weight from areal.utils.data import ( + TRANSPORT_DUMMY_KEY, + MicroBatchList, pack_tensor_dict, pad_and_stack_tensors_along_first_dim, + pad_mb_list, pad_sequences_to_tensors, reorder_list, split_padded_tensor_dict_into_mb_list, + split_training_batch_into_microbatches, unpack_sequence, ) @@ -73,3 +84,253 @@ def test_micro_batch_split(mock_padded_data, n_mbs, max_tokens_per_mb, n_mbs_div assert torch.allclose(x, packed_data[key]) y = pad_and_stack_tensors_along_first_dim(xs) assert torch.allclose(mock_padded_data[key], y) + + +def _preference_batch() -> dict[str, torch.Tensor]: + return { + "input_ids": torch.arange(6).view(2, 3), + "attention_mask": torch.ones(2, 3, dtype=torch.bool), + } + + +@pytest.mark.parametrize("loss_weight_fn", [_dpo_loss_weight, _rw_loss_weight]) +def test_transport_padding_bypasses_objective_weight(loss_weight_fn): + mb_list = split_padded_tensor_dict_into_mb_list( + _preference_batch(), + MicroBatchSpec(n_mbs=2, granularity=2), + allow_transport_padding=True, + ) + mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] + semantic_mb, transport_mb = sorted( + mb_list.mbs, key=lambda mb: TRANSPORT_DUMMY_KEY in mb + ) + + # A model-valid preference pair has non-zero objective weight. The transport + # marker, rather than objective-specific fields, is what makes it weightless. + torch.testing.assert_close( + loss_weight_fn(transport_mb), torch.tensor(1.0), rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + compute_microbatch_loss_weight(semantic_mb, loss_weight_fn), + torch.tensor(1.0), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + compute_microbatch_loss_weight(transport_mb, loss_weight_fn), + torch.tensor(0.0), + rtol=0.0, + atol=0.0, + ) + + pad_mb_list(mb_list) + assert all( + TRANSPORT_DUMMY_KEY not in padded_mb for padded_mb in mb_list.padded_mbs or [] + ) + + +def test_noop_packed_padding_preserves_semantic_transport_marker(): + transport_mb = pack_tensor_dict( + { + "input_ids": torch.zeros(1, 1, dtype=torch.long), + "attention_mask": torch.ones(1, 1, dtype=torch.bool), + TRANSPORT_DUMMY_KEY: True, + } + ) + mb_list = MicroBatchList( + data=transport_mb, + mb_spec=MicroBatchSpec(max_tokens_per_mb=1), + mbs=[transport_mb], + group_lens=[1], + transport_dummy_count=1, + ) + + pad_mb_list(mb_list, pad_to_maximum=True) + + assert mb_list.padding_lengths == [0] + assert mb_list.mbs[0][TRANSPORT_DUMMY_KEY] is True + assert mb_list.padded_mbs is not None + assert TRANSPORT_DUMMY_KEY not in mb_list.padded_mbs[0] + + callback_called = False + + def _loss_weight(_microbatch): + nonlocal callback_called + callback_called = True + return torch.tensor(1.0) + + torch.testing.assert_close( + compute_microbatch_loss_weight(mb_list.mbs[0], _loss_weight), + torch.tensor(0.0), + rtol=0.0, + atol=0.0, + ) + assert callback_called is False + + +def test_forward_transport_padding_is_removed_from_outputs(): + data = { + "input_ids": torch.arange(3).view(1, 3), + "attention_mask": torch.ones(1, 3, dtype=torch.bool), + } + mb_list = split_padded_tensor_dict_into_mb_list( + data, + MicroBatchSpec(n_mbs=3), + allow_transport_padding=True, + ) + outputs = [mb["input_ids"][mb["attention_mask"]].float() for mb in mb_list.mbs] + + result = reorder_and_pad_outputs(outputs, [3], mb_list) + + torch.testing.assert_close( + result, torch.tensor([[0.0, 1.0, 2.0]]), rtol=0.0, atol=0.0 + ) + + +def test_transport_padding_converges_to_distributed_microbatch_count(monkeypatch): + monkeypatch.setattr(data_module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(data_module.dist, "get_world_size", lambda _group=None: 2) + + def _all_gather_counts(output, local_count, group=None): + del group + output[:] = [3, local_count] + + monkeypatch.setattr(data_module.dist, "all_gather_object", _all_gather_counts) + + mb_list = split_padded_tensor_dict_into_mb_list( + { + "input_ids": torch.arange(3).view(1, 3), + "attention_mask": torch.ones(1, 3, dtype=torch.bool), + }, + MicroBatchSpec(), + allow_transport_padding=True, + ) + + assert len(mb_list.mbs) == 3 + assert mb_list.transport_dummy_count == 2 + assert sum(TRANSPORT_DUMMY_KEY in mb for mb in mb_list.mbs) == 2 + + +def _training_batch(batch_size: int) -> dict[str, torch.Tensor]: + return { + "input_ids": torch.arange(batch_size * 3).view(batch_size, 3), + "attention_mask": torch.ones(batch_size, 3, dtype=torch.bool), + "loss_mask": torch.tensor([[0, 1, 1]], dtype=torch.bool).repeat(batch_size, 1), + "advantages": torch.ones(batch_size, 3), + } + + +@pytest.mark.parametrize( + ("rank", "batch_size", "expected_semantic_steps"), + [(0, 1, [True, False, False]), (1, 2, [False, True, True])], +) +def test_synchronized_training_schedule_has_semantic_global_member_per_step( + monkeypatch, rank, batch_size, expected_semantic_steps +): + monkeypatch.setattr(data_module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(data_module.dist, "get_world_size", lambda group=None: 2) + monkeypatch.setattr(data_module.dist, "get_rank", lambda group=None: rank) + + def _all_gather_counts(output, value, group=None): + del value, group + output[:] = [1, 2] + + monkeypatch.setattr(data_module.dist, "all_gather_object", _all_gather_counts) + + schedule = split_training_batch_into_microbatches( + _training_batch(batch_size), + n_mbs=4, + ) + + assert len(schedule) == 3 + assert [ + TRANSPORT_DUMMY_KEY not in microbatch for microbatch in schedule + ] == expected_semantic_steps + assert all( + bool(microbatch["loss_mask"].any()) == is_semantic + for microbatch, is_semantic in zip( + schedule, expected_semantic_steps, strict=True + ) + ) + + +@pytest.mark.parametrize( + ("rank", "batch_size", "expected_semantic_steps"), + [(0, 3, [True, True, True]), (1, 1, [True, False, False])], +) +def test_synchronized_training_schedule_preserves_extra_local_microbatches( + monkeypatch, rank, batch_size, expected_semantic_steps +): + monkeypatch.setattr(data_module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(data_module.dist, "get_world_size", lambda group=None: 2) + monkeypatch.setattr(data_module.dist, "get_rank", lambda group=None: rank) + + def _all_gather_counts(output, value, group=None): + del value, group + output[:] = [3, 1] + + def _split_into_local_microbatches(data, mb_spec, synchronize): + del synchronize + microbatches = [ + {key: value[index : index + 1] for key, value in data.items()} + for index in range(batch_size) + ] + return MicroBatchList( + data=data, + mb_spec=mb_spec, + mbs=microbatches, + group_lens=[1] * batch_size, + ) + + monkeypatch.setattr(data_module.dist, "all_gather_object", _all_gather_counts) + monkeypatch.setattr( + data_module, + "split_padded_tensor_dict_into_mb_list", + _split_into_local_microbatches, + ) + + schedule = split_training_batch_into_microbatches( + _training_batch(batch_size), + n_mbs=2, + ) + + assert len(schedule) == 3 + assert [ + TRANSPORT_DUMMY_KEY not in microbatch for microbatch in schedule + ] == expected_semantic_steps + + +def test_tensor_container_skeleton_round_trip(): + item = { + "input_ids": torch.arange(6, dtype=torch.long).view(2, 3), + "nested": [torch.ones(2, dtype=torch.bool), {"logprobs": torch.randn(4)}], + "reward": 1.5, + "task": "math", + } + + tensors: list[torch.Tensor] = [] + skeleton = data_module._deconstruct_tensor_container(item, tensors) + rebuilt = data_module._reconstruct_tensor_container(skeleton, iter(tensors)) + + assert torch.equal(rebuilt["input_ids"], item["input_ids"]) + assert torch.equal(rebuilt["nested"][0], item["nested"][0]) + assert torch.equal(rebuilt["nested"][1]["logprobs"], item["nested"][1]["logprobs"]) + assert rebuilt["reward"] == 1.5 + assert rebuilt["task"] == "math" + + +def test_skeleton_tensor_leaves_follow_depth_first_order(): + tensors: list[torch.Tensor] = [] + items = [ + {"a": torch.zeros(1, dtype=torch.long), "b": [torch.zeros(2)]}, + {"a": torch.zeros(3, dtype=torch.long)}, + ] + skeletons = [ + data_module._deconstruct_tensor_container(item, tensors) for item in items + ] + + leaves = data_module._skeleton_tensor_leaves(skeletons) + + assert [leaf.shape for leaf in leaves] == [(1,), (2,), (3,)] + assert [leaf.dtype for leaf in leaves] == [torch.long, torch.float32, torch.long] + assert [leaf.numel for leaf in leaves] == [1, 2, 3] diff --git a/tests/torchrun/redistribute.py b/tests/torchrun/redistribute.py index 6fca229f49..207261ef2a 100644 --- a/tests/torchrun/redistribute.py +++ b/tests/torchrun/redistribute.py @@ -6,18 +6,83 @@ import torch import torch.distributed as dist -from areal.infra.dist_rollout import redistribute_trajectories +from areal.infra.dist_rollout import DistRolloutCoordinator, redistribute_trajectories from areal.infra.platforms import current_platform from areal.utils.data import tensor_container_to +class _HybridTrainEngine: + def __init__(self, rank, dp_group, model_group): + self.rank = rank + self.data_parallel_group = dp_group + self.context_and_model_parallel_group = model_group + + def is_data_parallel_head(self): + return self.rank % 2 == 0 + + def current_data_parallel_head(self): + return self.rank - self.rank % 2 + + +class _HybridRolloutEngine: + def __init__(self, rank): + self.rank = rank + + def prepare_batch(self, *args, **kwargs): + if self.rank == 0: + raise ValueError("rank-local preparation failure") + return [] + + +def _test_hybrid_error(rank): + head_dp_group = dist.new_group([0, 2]) + non_head_dp_group = dist.new_group([1, 3]) + first_model_group = dist.new_group([0, 1]) + second_model_group = dist.new_group([2, 3]) + train_engine = _HybridTrainEngine( + rank, + head_dp_group if rank % 2 == 0 else non_head_dp_group, + first_model_group if rank < 2 else second_model_group, + ) + coordinator = DistRolloutCoordinator(_HybridRolloutEngine(rank), train_engine) + + try: + coordinator.prepare_batch(object(), object()) + except RuntimeError as exc: + assert "rank-local preparation failure" in str(exc) + else: + raise AssertionError("Expected coordinated rank-local preparation failure") + dist.barrier() + + trajectories = [] if train_engine.is_data_parallel_head() else None + try: + coordinator._broadcast_and_redistribute_trajectories(trajectories) + except RuntimeError as exc: + assert "Cannot redistribute 0 trainable trajectory groups" in str(exc) + else: + raise AssertionError("Expected coordinated rollout preparation failure") + dist.barrier() + + def main(args): - dist.init_process_group("nccl") + dist.init_process_group(args.backend) rank = int(os.environ["LOCAL_RANK"]) - current_platform.set_device(rank) - device = f"{current_platform.device_type}:{rank}" + if args.hybrid_error: + _test_hybrid_error(rank) + return + if args.backend == "nccl": + current_platform.set_device(rank) + device = f"{current_platform.device_type}:{rank}" + else: + device = "cpu" - bs = 16 + if args.ragged_empty: + # One rank contributes zero trajectories; the gather must still work. + bs = 3 * rank + elif args.ragged: + bs = rank + 1 + else: + bs = 16 prompt_lens = [random.randint(1, 10) for _ in range(bs)] ans_lens = [random.randint(1, 10) for _ in range(bs)] seqlens = [x + y for x, y in zip(prompt_lens, ans_lens)] @@ -59,5 +124,9 @@ def main(args): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--dump-path", type=str) + parser.add_argument("--backend", choices=["gloo", "nccl"], default="nccl") + parser.add_argument("--ragged", action="store_true") + parser.add_argument("--ragged-empty", action="store_true") + parser.add_argument("--hybrid-error", action="store_true") args = parser.parse_args() main(args) diff --git a/tests/v2/training_service/test_data_proxy_unit.py b/tests/v2/training_service/test_data_proxy_unit.py index 2be179225d..1c8643aacc 100644 --- a/tests/v2/training_service/test_data_proxy_unit.py +++ b/tests/v2/training_service/test_data_proxy_unit.py @@ -499,8 +499,14 @@ def _shard_handler(url, data, headers): } ) - with pytest.raises(ValueError, match="divisible by K"): - await dispatcher.dispatch("/train_batch").post(body) + result_bytes = await dispatcher.dispatch("/train_batch").post(body) + result_payload = orjson.loads(result_bytes) + merged = deserialize_value(result_payload["result"]) + + assert merged == [5, 11, 7] + assert sorted( + len(payload["args"][0]) for payload in session.captured_payloads + ) == [1, 2] @pytest.mark.asyncio async def test_dispatch_post_pads_eval_routes_only(self):