From 6cd3c09a90e34643d167ae428b6d6aa5751186b3 Mon Sep 17 00:00:00 2001 From: leviking98z-rgb Date: Mon, 3 Aug 2026 06:25:15 +0800 Subject: [PATCH] perf(unified-model): per-track micro-batch geometry for AR and image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The unified-model stack slices both lineage levels with one shared ``micro_batch_size``, but the two tracks are shaped differently: AR rows are variable-length responses whose activation peak is set by the longest row in a micro, while image rows are fixed-shape latents where a larger micro is clean parallelism. One number cannot suit both, so the image track is pinned to whatever the AR track can afford. Add optional ``ar_micro_batch_size`` / ``image_micro_batch_size``, both defaulting to the shared ``micro_batch_size`` so omitting them is a no-op. ``_optimizer_step_slices`` and ``prepare_segment`` take the size explicitly, which keeps a replay-anchored track's π_old anchor frozen at exactly the geometry training will use (an anchor recomputed at a different micro size would break the on-policy ratio). Also weight each micro's metrics by its sample share in ``_backward_part``, so a ragged final micro no longer skews the reported mean toward its fewer rows. This matches the sample-share ``loss_scale`` already used for backward; equal micros reduce to the previous plain average. ``aggregate_numeric_metrics`` grows a keyword-only ``weights`` argument, renormalized per key so a metric present in only some micros still averages correctly. --- .../bagel_trainside_unigrpo.yaml | 4 + examples/unified_model/hi3_vllmomni.yaml | 7 ++ .../unified_model/hi3_vllmomni_veomni_ep.yaml | 6 ++ unirl/train/unified_model_stack.py | 102 +++++++++++++++--- unirl/utils/misc.py | 33 +++++- 5 files changed, 133 insertions(+), 19 deletions(-) diff --git a/examples/unified_model/bagel_trainside_unigrpo.yaml b/examples/unified_model/bagel_trainside_unigrpo.yaml index a9dfff863..a6fcac537 100644 --- a/examples/unified_model/bagel_trainside_unigrpo.yaml +++ b/examples/unified_model/bagel_trainside_unigrpo.yaml @@ -191,6 +191,10 @@ algorithm: stack: _target_: unirl.train.unified_model_stack.UnifiedModelTrainStack micro_batch_size: 1 + # Optional per-track override of micro_batch_size (null = use the shared value). + # See hi3_vllmomni.yaml for the AR-vs-image trade-off and the loss_agg_mode caveat. + ar_micro_batch_size: null + image_micro_batch_size: null max_grad_norm: 1.0 num_updates_per_batch: 2 diff --git a/examples/unified_model/hi3_vllmomni.yaml b/examples/unified_model/hi3_vllmomni.yaml index 07499e965..8c8592c0f 100644 --- a/examples/unified_model/hi3_vllmomni.yaml +++ b/examples/unified_model/hi3_vllmomni.yaml @@ -184,6 +184,13 @@ algorithm: stack: _target_: unirl.train.unified_model_stack.UnifiedModelTrainStack micro_batch_size: 1 + # Optional per-track override of micro_batch_size (null = use the shared value). + # The image track is fixed-shape latents, so a larger micro is clean parallelism; + # the AR track is variable-length responses, where a larger micro pads to the + # longest row. Raising ar_micro_batch_size also needs a seq-mean-* loss_agg_mode + # (GRPO's default token-mean weights by token count, not by row count). + ar_micro_batch_size: null + image_micro_batch_size: null max_grad_norm: 1.0 data_source: diff --git a/examples/unified_model/hi3_vllmomni_veomni_ep.yaml b/examples/unified_model/hi3_vllmomni_veomni_ep.yaml index 011e1d24a..c0af1e2fc 100644 --- a/examples/unified_model/hi3_vllmomni_veomni_ep.yaml +++ b/examples/unified_model/hi3_vllmomni_veomni_ep.yaml @@ -191,6 +191,12 @@ algorithm: stack: _target_: unirl.train.unified_model_stack.UnifiedModelTrainStack micro_batch_size: 1 + # Optional per-track override of micro_batch_size (null = use the shared value). + # See hi3_vllmomni.yaml: the image track is fixed-shape latents (a larger micro is + # clean parallelism), the AR track is variable-length responses (a larger micro + # pads, and needs a seq-mean-* loss_agg_mode to stay loss-equivalent). + ar_micro_batch_size: null + image_micro_batch_size: null max_grad_norm: 1.0 data_source: diff --git a/unirl/train/unified_model_stack.py b/unirl/train/unified_model_stack.py index 8d6050cf5..2628d7717 100644 --- a/unirl/train/unified_model_stack.py +++ b/unirl/train/unified_model_stack.py @@ -27,6 +27,12 @@ clip / ratio trust region actually engages (the UniGRPO / FlowGRPO PPO schedule). Mirrors :class:`~unirl.train.stack.TrainStack` but for the two-algorithm backbone. +``ar_micro_batch_size`` / ``image_micro_batch_size`` (both defaulting to the shared +``micro_batch_size``) let the two tracks use different gradient-accumulation +geometry. The lineage levels differ in row count (AR is ``P*N``, image is +``P*N*M``) and in row shape (variable-length AR responses vs fixed-shape latents), +so the micro size that fits one is rarely the best for the other. + This is the multi-stage train stack — several stage algorithms share one optimizer step, in contrast to the single-stage ``TrainStack``. """ @@ -36,7 +42,7 @@ import logging from contextlib import nullcontext from dataclasses import replace -from typing import Dict, List, Mapping, Tuple +from typing import Dict, List, Mapping, Optional, Tuple import torch @@ -76,6 +82,8 @@ def __init__( micro_batch_size: int, max_grad_norm: float, num_updates_per_batch: int = 1, + ar_micro_batch_size: Optional[int] = None, + image_micro_batch_size: Optional[int] = None, ) -> None: super().__init__() if int(micro_batch_size) < 1: @@ -86,6 +94,30 @@ def __init__( self.ar_algorithm = ar_algorithm self.image_algorithm = image_algorithm self.micro_batch_size = int(micro_batch_size) + # Per-track micro-batch size. The two lineage levels have different shapes: + # AR rows are variable-length responses (a larger micro pads to the longest + # and raises the activation peak), while image rows are fixed-shape latents + # (a larger micro is clean parallelism). Both default to the shared + # ``micro_batch_size``, so omitting them keeps the previous behavior. + # + # Micro-batching is only loss-equivalent to one full-batch backward when the + # per-micro ``loss_scale`` (sample share) reconstructs the whole-batch mean. + # That holds for a per-row mean (FlowGRPO's flat latent mean; GRPO's + # ``seq-mean-*`` modes, which average per sequence then over sequences), but + # NOT for GRPO's default ``token-mean`` when responses differ in length — + # there the exact objective weights by token count, not by row count. Raise + # ``ar_micro_batch_size`` above 1 only with a ``seq-mean-*`` + # ``loss_agg_mode``, or accept that shift. + self.micro_batch_sizes = { + "ar": _positive_int( + name="UnifiedModelTrainStack.ar_micro_batch_size", + value=self.micro_batch_size if ar_micro_batch_size is None else ar_micro_batch_size, + ), + "image": _positive_int( + name="UnifiedModelTrainStack.image_micro_batch_size", + value=self.micro_batch_size if image_micro_batch_size is None else image_micro_batch_size, + ), + } self.max_grad_norm = float(max_grad_norm) # PPO-style multi-update: split each rollout shard into this many disjoint # mini-batches and run ONE optimizer step per mini-batch, with the π_old @@ -106,27 +138,45 @@ def __init__( f"num_updates_per_batch=1." ) - def _optimizer_step_slices(self, total: int) -> List[List[Tuple[int, int]]]: + def _optimizer_step_slices( + self, + total: int, + *, + micro_batch_size: Optional[int] = None, + ) -> List[List[Tuple[int, int]]]: """Per-optimizer-step lists of absolute ``(start, end)`` micro-batch slices. One inner list per ``num_updates_per_batch`` mini-batch (one optimizer step), - each split into ``micro_batch_size`` micro-batches. Shared by + each split into the requested micro-batch size (or the shared + ``micro_batch_size`` fallback). Shared by :meth:`prepare_segment` (to freeze the anchor at the exact geometry) and the train loop. Mirrors :meth:`unirl.train.stack.TrainStack._optimizer_step_slices`. """ + resolved_micro_batch_size = ( + self.micro_batch_size + if micro_batch_size is None + else _positive_int(name="micro_batch_size", value=micro_batch_size) + ) steps: List[List[Tuple[int, int]]] = [] for mini_start, mini_end in _update_ranges(total_size=total, num_updates=self.num_updates_per_batch): steps.append( [ (mini_start + ms, mini_start + me) for ms, me in _build_micro_batch_slices( - total_size=mini_end - mini_start, micro_batch_size=self.micro_batch_size + total_size=mini_end - mini_start, + micro_batch_size=resolved_micro_batch_size, ) ] ) return steps - def prepare_segment(self, algorithm: StageAlgorithm, part: Part) -> None: + def prepare_segment( + self, + algorithm: StageAlgorithm, + part: Part, + *, + micro_batch_size: Optional[int] = None, + ) -> None: """Freeze one algorithm's π_old anchor once, before the multi-update loop. No-op if ``segment`` is None or the algorithm has no ``prepare_segment``. If @@ -134,7 +184,9 @@ def prepare_segment(self, algorithm: StageAlgorithm, part: Part) -> None: — e.g. FlowGRPO under ``old_logp_source='replay'``), the declared ``anchor_fields`` are recomputed at the SAME (mini, micro) slices training will use, so the on-policy ratio is exactly 1 (mirrors - :meth:`TrainStack.prepare_segment`). Rollout-anchored algorithms (the BAGEL + :meth:`TrainStack.prepare_segment`). ``micro_batch_size`` must therefore be the + SAME per-track value the train loop will use for this Part; passing None falls + back to the shared ``micro_batch_size``. Rollout-anchored algorithms (the BAGEL UniGRPO recipe: AR GRPO + image ``old_logp_source='rollout'``) take the one-shot path — the anchor is the rollout emission, geometry-independent. """ @@ -147,7 +199,11 @@ def prepare_segment(self, algorithm: StageAlgorithm, part: Part) -> None: if recomputes is None or not recomputes(): prepare(conditions=part.conditions, segment=part.segment) return - micro_slices = [sl for step in self._optimizer_step_slices(int(part.batch_size)) for sl in step] + micro_slices = [ + sl + for step in self._optimizer_step_slices(int(part.batch_size), micro_batch_size=micro_batch_size) + for sl in step + ] if len(micro_slices) == 1: prepare(conditions=part.conditions, segment=part.segment) return @@ -195,6 +251,7 @@ def _backward_part( bs = int(part.batch_size) update_total = sum(end - start for start, end in micro_slices) micros: List[AlgorithmStepResult] = [] + micro_weights: List[float] = [] total_loss = 0.0 has_backward = False @@ -210,10 +267,19 @@ def _backward_part( loss_scale=loss_scale, ) micros.append(result) + micro_weights.append(loss_scale) total_loss += result.loss * loss_scale has_backward = has_backward or result.has_backward - aggregated: Mapping[str, object] = aggregate_numeric_metrics([r.metrics for r in micros if r.metrics]) + # Weight each micro's metrics by its sample share, matching the weighted + # objective used for backward. A ragged final micro (batch not divisible by + # the micro size) otherwise skews an unweighted mean toward its fewer rows. + # Equal-sized micros reduce to the previous plain average. + metric_items = [(result.metrics, weight) for result, weight in zip(micros, micro_weights) if result.metrics] + aggregated: Mapping[str, object] = aggregate_numeric_metrics( + [metrics for metrics, _ in metric_items], + weights=[weight for _, weight in metric_items], + ) # grad_norm / lr are filled by ``_train_one_step`` after the shared optimizer step. partial = TrainStepResult( loss=total_loss, @@ -337,13 +403,19 @@ def train_track( profiler = self._train_step_profiler() if scope == "train" else None with profiler.record("train_track") if profiler is not None else nullcontext(): # Freeze each Part's π_old anchor once, before the multi-update loop. - self.prepare_segment(self.ar_algorithm, ar_part) - self.prepare_segment(self.image_algorithm, image_part) - - # N optimizer steps over disjoint mini-batches (each Part sliced by the same - # shared _optimizer_step_slices; M=1 keeps ar/image 1:1 and equally sized). - ar_steps = self._optimizer_step_slices(int(ar_part.batch_size)) - image_steps = self._optimizer_step_slices(int(image_part.batch_size)) + # Each track's anchor geometry must match the micro size training uses. + self.prepare_segment(self.ar_algorithm, ar_part, micro_batch_size=self.micro_batch_sizes["ar"]) + self.prepare_segment(self.image_algorithm, image_part, micro_batch_size=self.micro_batch_sizes["image"]) + + # N optimizer steps over disjoint mini-batches. Each Part is sliced at its + # own micro size (the two lineage levels differ in both row count and row + # shape), so ar/image micro counts need not match. + ar_steps = self._optimizer_step_slices( + int(ar_part.batch_size), micro_batch_size=self.micro_batch_sizes["ar"] + ) + image_steps = self._optimizer_step_slices( + int(image_part.batch_size), micro_batch_size=self.micro_batch_sizes["image"] + ) per_update: List[Dict[str, TrainStepResult]] = [] for u in range(self.num_updates_per_batch): per_update.append( diff --git a/unirl/utils/misc.py b/unirl/utils/misc.py index 418771622..9b0b907e9 100644 --- a/unirl/utils/misc.py +++ b/unirl/utils/misc.py @@ -142,11 +142,27 @@ def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict: return dict(items) -def aggregate_numeric_metrics(metrics_list: List[Dict[str, Any]]) -> Dict[str, float]: - """Average numeric metric keys across repeated metric dictionaries.""" +def aggregate_numeric_metrics( + metrics_list: List[Dict[str, Any]], + *, + weights: Optional[List[float]] = None, +) -> Dict[str, float]: + """Average numeric metric keys across repeated metric dictionaries. + + ``weights`` (same length as ``metrics_list``) gives a weighted mean instead of a + plain one — used when the entries cover unequal sample counts, e.g. gradient- + accumulation micro-batches whose final micro is ragged. Weights are renormalized + per key over the entries that actually carry it, so a key present in only some + dictionaries still averages correctly. Non-positive total weight for a key falls + back to the unweighted mean. + """ aggregated: Dict[str, float] = {} if not metrics_list: return aggregated + if weights is not None and len(weights) != len(metrics_list): + raise ValueError( + f"aggregate_numeric_metrics: weights has length {len(weights)} but metrics_list has {len(metrics_list)}." + ) all_keys = set() for metrics in metrics_list: @@ -154,7 +170,8 @@ def aggregate_numeric_metrics(metrics_list: List[Dict[str, Any]]) -> Dict[str, f for key in all_keys: values: List[float] = [] - for metrics in metrics_list: + value_weights: List[float] = [] + for index, metrics in enumerate(metrics_list): if key not in metrics: continue value = metrics[key] @@ -164,7 +181,15 @@ def aggregate_numeric_metrics(metrics_list: List[Dict[str, Any]]) -> Dict[str, f values.append(float(value)) elif isinstance(value, (int, float)): values.append(float(value)) - if values: + else: + continue + value_weights.append(1.0 if weights is None else float(weights[index])) + if not values: + continue + total_weight = sum(value_weights) + if total_weight > 0.0: + aggregated[key] = sum(v * w for v, w in zip(values, value_weights)) / total_weight + else: aggregated[key] = sum(values) / len(values) return aggregated