refactor(train): share per-micro anchor preparation between train stacks

unirl/train/stack/base.py

@@ -6,7 +6,7 @@
 import math
 from contextlib import nullcontext
 from dataclasses import dataclass, replace
-from typing import Dict, List, Mapping, Optional, Tuple, Union
+from typing import Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 import torch
 
@@ -73,6 +73,34 @@ def _validate_anchor_contract(algorithm: StageAlgorithm) -> None:
         raise ValueError(f"{type(algorithm).__name__} recomputes its anchor but declares no anchor_fields.")
 
 
+def _prepare_segment_anchors(
+    algorithm: StageAlgorithm,
+    part: Part,
+    micro_slices: Sequence[Tuple[int, int]],
+    *,
+    caller: str,
+) -> None:
+    """Freeze declared anchors over planned micros and write them back onto ``part``."""
+    if part.segment is None:
+        return
+    if not algorithm.recomputes_anchor or len(micro_slices) == 1:
+        algorithm.prepare_segment(conditions=part.conditions, segment=part.segment)
+        return
+    collected: Dict[str, List[torch.Tensor]] = {field: [] for field in algorithm.anchor_fields}
+    for start, end in micro_slices:
+        micro = part.slice(start, end)
+        algorithm.prepare_segment(conditions=micro.conditions, segment=micro.segment)
+        for field in collected:
+            value = getattr(micro.segment, field, None)
+            if value is None:
+                raise RuntimeError(
+                    f"{caller}: {type(algorithm).__name__} declares anchor field {field!r} but a micro produced None."
+                )
+            collected[field].append(value)
+    for field, tensors in collected.items():
+        setattr(part.segment, field, torch.cat(tensors, dim=0))
+
+
 class TrainStack(Remote):
     """Single-stage stage-driven train stack — family-agnostic."""
 
@@ -111,30 +139,12 @@ def __init__(
 
     def prepare_segment(self, part: Part, *, plans: Plan) -> None:
         """Freeze the π_old anchor once, before the ``num_updates_per_batch`` loop."""
-        if part.segment is None:
-            return
-        algorithm = self.algorithm
-        if not algorithm.recomputes_anchor:
-            algorithm.prepare_segment(conditions=part.conditions, segment=part.segment)
-            return
-        micro_slices = [r for update in plans for r in update]
-        if len(micro_slices) == 1:
-            algorithm.prepare_segment(conditions=part.conditions, segment=part.segment)
-            return
-        collected: Dict[str, List[torch.Tensor]] = {field: [] for field in algorithm.anchor_fields}
-        for start, end in micro_slices:
-            micro = part.slice(start, end)
-            algorithm.prepare_segment(conditions=micro.conditions, segment=micro.segment)
-            for field in collected:
-                value = getattr(micro.segment, field, None)
-                if value is None:
-                    raise RuntimeError(
-                        f"{type(self).__name__}.prepare_segment: {type(algorithm).__name__} declares "
-                        f"anchor field {field!r} but a micro produced None."
-                    )
-                collected[field].append(value)
-        for field, parts in collected.items():
-            setattr(part.segment, field, torch.cat(parts, dim=0))
+        _prepare_segment_anchors(
+            self.algorithm,
+            part,
+            [r for update in plans for r in update],
+            caller=f"{type(self).__name__}.prepare_segment",
+        )
 
     def _run_update(
         self,

unirl/train/unified_model_stack.py

@@ -15,7 +15,7 @@
 from unirl.distributed.group.remote import Remote
 from unirl.train.backend.fsdp import FSDPBackend
 from unirl.train.stack import TrainStepResult, _build_micro_batch_slices
-from unirl.train.stack.base import _aggregate_update_results, _validate_anchor_contract
+from unirl.train.stack.base import _aggregate_update_results, _prepare_segment_anchors, _validate_anchor_contract
 from unirl.train.stack.planner.types import _positive_int, _update_ranges
 from unirl.types.sample import Part, Sample
 from unirl.types.sampling import ARSamplingParams, DiffusionSamplingParams
@@ -78,29 +78,15 @@ def _optimizer_step_slices(self, total: int) -> List[List[Tuple[int, int]]]:
 
     def prepare_segment(self, algorithm: StageAlgorithm, part: Part) -> None:
         """Freeze one algorithm's π_old anchor once, before the multi-update loop."""
-        if part.segment is None:
-            return
-        if not algorithm.recomputes_anchor:
-            algorithm.prepare_segment(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]
-        if len(micro_slices) == 1:
-            algorithm.prepare_segment(conditions=part.conditions, segment=part.segment)
-            return
-        collected: Dict[str, List[torch.Tensor]] = {field: [] for field in algorithm.anchor_fields}
-        for start, end in micro_slices:
-            micro = part.slice(start, end)
-            algorithm.prepare_segment(conditions=micro.conditions, segment=micro.segment)
-            for field in collected:
-                value = getattr(micro.segment, field, None)
-                if value is None:
-                    raise RuntimeError(
-                        f"UnifiedModelTrainStack.prepare_segment: {type(algorithm).__name__} declares "
-                        f"anchor field {field!r} but a micro-slice produced None."
-                    )
-                collected[field].append(value)
-        for field, parts in collected.items():
-            setattr(part.segment, field, torch.cat(parts, dim=0))
+        micro_slices: List[Tuple[int, int]] = []
+        if part.segment is not None and algorithm.recomputes_anchor:
+            micro_slices = [sl for step in self._optimizer_step_slices(int(part.batch_size)) for sl in step]
+        _prepare_segment_anchors(
+            algorithm,
+            part,
+            micro_slices,
+            caller="UnifiedModelTrainStack.prepare_segment",
+        )
 
     def _backward_part(
         self,