Feat/forward only micro batch

src/twinkle/model/micro_batch.py

@@ -190,3 +190,130 @@ def select_batch(value: Any, indices: list[int], batch_size: int) -> Any:
     if hasattr(value, 'shape') and len(value.shape) > 0 and value.shape[0] == batch_size:
         return value[indices]
     return value
+
+
+def merge_micro_batch_outputs(
+    outputs: list[dict[str, Any]], plan: list[list[int]], *, task: str = 'causal_lm',
+    micro_batch_scales: list[float] | None = None,
+) -> dict[str, Any]:
+    """Restore sample outputs and reduce auxiliary loss with explicit training scales."""
+    merger = _MicroBatchOutputMerger(plan)
+    return merger.merge(outputs, task=task, micro_batch_scales=micro_batch_scales)
+
+
+class _MicroBatchOutputMerger:
+    """Keep sample ordering shared across field-specific and recursive merge operations."""
+
+    def __init__(self, plan):
+        if not plan or any(not indices for indices in plan):
+            raise ValueError('outputs must match a non-empty micro-batch plan')
+        merged_sample_indices = [index for indices in plan for index in indices]
+        if sorted(merged_sample_indices) != list(range(len(merged_sample_indices))):
+            raise ValueError('micro-batch plan must cover every sample exactly once')
+        self.plan = plan
+        self.restore_order_indices = sorted(range(len(merged_sample_indices)), key=merged_sample_indices.__getitem__)
+
+    @staticmethod
+    def _invalid(path, reason):
+        raise ValueError(f'{path}: {reason}')
+
+    def merge(self, outputs, *, task, micro_batch_scales):
+        if not outputs or len(outputs) != len(self.plan):
+            raise ValueError('outputs must match a non-empty micro-batch plan')
+        if any(set(output) != set(outputs[0]) for output in outputs):
+            self._invalid('outputs', 'inconsistent output keys across micro-batches')
+        pooled_tasks = ('embedding', 'generative_reranker', 'seq_cls', 'reranker')
+        known_sequence_dims = {
+            'logps': (1,), 'values': (1,), 'entropies': (1,),
+            'logits': () if task in pooled_tasks else (1,),
+            'embeddings': (), 'routed_experts': (1,),
+        }
+        return {
+            field: self._merge_field(field, [output[field] for output in outputs], known_sequence_dims, micro_batch_scales)
+            for field in outputs[0]
+        }
+
+    def _merge_field(self, field, values, known_sequence_dims, micro_batch_scales):
+        path = f'outputs.{field}'
+        if field == 'aux_loss':
+            return self._reduce_aux_loss(values, path, micro_batch_scales)
+        if field == 'past_key_values':
+            if any(value is not None for value in values):
+                self._invalid(path, 'forward_only expects None')
+            return None
+        if field in known_sequence_dims:
+            return self._merge_values(values, path, known_sequence_dims[field])
+        return self._merge_values(values, path)
+
+    @staticmethod
+    def _reduce_aux_loss(values, path, micro_batch_scales):
+        import torch
+
+        if all(value is None for value in values):
+            return None
+        if any(value is None for value in values):
+            raise ValueError(f'{path}: inconsistent optional outputs across micro-batches')
+        if micro_batch_scales is None or len(micro_batch_scales) != len(values):
+            raise ValueError(f'{path}: one training scale per micro-batch is required')
+        if any(not isinstance(value, (int, float, torch.Tensor))
+               or (isinstance(value, torch.Tensor) and value.ndim != 0) for value in values):
+            raise ValueError(f'{path}: expected scalar auxiliary losses')
+        return sum(value * scale for value, scale in zip(values, micro_batch_scales))
+
+    def _merge_values(self, values, path, sequence_dims=None):
+        import torch
+
+        first = values[0]
+        if all(value is None for value in values):
+            return None
+        if any(value is None for value in values):
+            self._invalid(path, 'inconsistent optional outputs across micro-batches')
+        if any(type(value) is not type(first) for value in values):
+            self._invalid(path, 'inconsistent output types across micro-batches')
+        if isinstance(first, torch.Tensor):
+            return self._merge_tensors(values, path, sequence_dims)
+        if isinstance(first, dict):
+            return self._merge_dicts(values, path, sequence_dims)
+        if isinstance(first, tuple):
+            return self._merge_tuples(values, path, sequence_dims)
+        if isinstance(first, list):
+            return self._merge_sample_lists(values, path)
+        self._invalid(path, f'unsupported micro-batch output type: {type(first).__name__}')
+
+    def _merge_tensors(self, values, path, sequence_dims):
+        from twinkle.utils import pad_and_stack_tensors
+
+        first = values[0]
+        if first.ndim == 0:
+            self._invalid(path, 'scalar forward outputs require an explicit reduction rule')
+        if any(value.ndim != first.ndim or value.dtype != first.dtype or value.device != first.device
+               for value in values):
+            self._invalid(path, 'tensor rank, dtype and device must agree across micro-batches')
+        if any(value.shape[0] != len(indices) for value, indices in zip(values, self.plan)):
+            self._invalid(path, f'expected a sample dimension matching the plan; got {[tuple(v.shape) for v in values]}')
+        # Known fields specify their sequence axes; unknown tensors use generic padding.
+        padding_dims = tuple(range(1, first.ndim)) if sequence_dims is None else sequence_dims
+        for dim in range(1, first.ndim):
+            if dim not in padding_dims and any(value.shape[dim] != first.shape[dim] for value in values):
+                self._invalid(path, f'non-sequence dimension {dim} must agree across micro-batches')
+        combined = pad_and_stack_tensors(values, pad_value=0)
+        return combined[self.restore_order_indices]
+
+    def _merge_dicts(self, values, path, sequence_dims):
+        first = values[0]
+        if any(set(value) != set(first) for value in values):
+            self._invalid(path, 'inconsistent output keys across micro-batches')
+        return {key: self._merge_values([value[key] for value in values], f'{path}.{key}', sequence_dims)
+                for key in first}
+
+    def _merge_tuples(self, values, path, sequence_dims):
+        if any(len(value) != len(values[0]) for value in values):
+            self._invalid(path, 'tuple lengths must agree across micro-batches')
+        return tuple(self._merge_values([value[index] for value in values], f'{path}[{index}]', sequence_dims)
+                     for index in range(len(values[0])))
+
+    def _merge_sample_lists(self, values, path):
+        if any(len(value) != len(indices) for value, indices in zip(values, self.plan)):
+            self._invalid(path, 'output list does not match micro-batch sample counts')
+        flattened = [item for value in values for item in value]
+        return [flattened[index] for index in self.restore_order_indices]

src/twinkle/model/transformers/transformers.py

@@ -39,7 +39,7 @@
 from twinkle.metric import Accuracy, LossMetric, Metric, TrainMetric
 from twinkle.model.base import ModelLoaderProtocol, TrainableModel, copy_checkpoint_args, rotate_checkpoints
 from twinkle.model.state_dict_utils import normalize_adapter_keys
-from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches, select_batch
+from twinkle.model.micro_batch import MicroBatchConfig, merge_micro_batch_outputs, plan_micro_batches, select_batch
 from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus
 from twinkle.model.transformers import feature_hooks
 from twinkle.model.transformers.moe import apply_expert_parallel
@@ -658,7 +658,6 @@ def __repr__(self) -> str:
         # into the exception string and that dump swallows the real traceback.
         return f'{type(self).__name__}(model_id={getattr(self, "model_id", None)!r})'
 
-    @staticmethod
     def _lazy_wrap_model(self):
         if not self._model_wrapped:
             optimizer_groups = [og for og in self.optimizer_group.values() if og.optimizer is not None]
@@ -940,10 +939,56 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             **kwargs:
                 adapter_name: Lora adapter name.
                 disable_lora: If True, disable LoRA and use base model for inference.
+                micro_batch_size: Maximum samples per local micro-batch.
+                dynamic_batching: Pack samples under max_tokens_per_micro_batch.
+                max_tokens_per_micro_batch: Token budget required for dynamic batching.
+                packing_algorithm: Dynamic packing strategy, 'ffd' (default) or 'kk'.
         Returns:
             The output of the model forward.
         """
+        config = MicroBatchConfig.from_kwargs(kwargs)
         adapter_name = kwargs.pop('adapter_name', self._get_default_group())
+        optimizer_config = self.optimizer_group[adapter_name]
+        inputs = self._prepare_forward_only_inputs(inputs, optimizer_config)
+        if config is None:
+            return self._forward_only_single(inputs=inputs, optimizer_config=optimizer_config, **kwargs)
+        return self._forward_only_micro_batches(
+            inputs=inputs, config=config, optimizer_config=optimizer_config, **kwargs)
+
+    def _prepare_forward_only_inputs(self, inputs, optimizer_config):
+        """Normalize and encode the local batch once, before optional micro-batching."""
+        if isinstance(inputs, dict):
+            inputs = [inputs]
+        if not inputs:
+            raise ValueError('inputs empty, check your DataLoader outputs')
+        if self._not_encoded(inputs[0]):
+            assert optimizer_config.template is not None, \
+                'Use set_template to add a template when trying to input `List[Trajectory]`'
+            inputs = optimizer_config.template.batch_encode(inputs)
+        if not inputs:
+            raise ValueError('encoded inputs empty, check your template outputs')
+        return inputs
+
+    def _forward_only_micro_batches(self, *, inputs, config, optimizer_config, **kwargs):
+        """Execute a local micro-batch plan and restore the complete return value."""
+        plan = self._build_micro_batch_plan(inputs, config, optimizer_config)
+        batches = []
+        for indices in plan:
+            micro_kwargs = {key: select_batch(value, indices, len(inputs)) for key, value in kwargs.items()}
+            outputs = self._forward_only_single(
+                inputs=[inputs[index] for index in indices], optimizer_config=optimizer_config,
+                _unpack_routing=True, **micro_kwargs)
+            batches.append(dict(outputs))
+        micro_batch_scales = None
+        if any(output.get('aux_loss') is not None for output in batches):
+            loss_instance = optimizer_config.get_loss(kwargs.get('loss_key'))
+            micro_batch_scales = [loss_instance.micro_batch_scale(inputs, indices) for indices in plan]
+        return merge_micro_batch_outputs(
+            batches, plan, task=kwargs.get('task', 'causal_lm'), micro_batch_scales=micro_batch_scales)
+
+    def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
+        """Run one evaluation micro-batch, retaining its state for metric accumulation."""
+        unpack_routing = kwargs.pop('_unpack_routing', False)
         disable_lora = kwargs.pop('disable_lora', False)
         temperature = float(kwargs.pop('temperature', 1.0))
         sampling_masks = kwargs.pop('sampling_masks', None)
@@ -954,19 +999,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
         # forward_only derives logps via selective_log_softmax and never calls the loss, so there is nothing
         # to skip here; popped so it does not linger in eval_status.forward_kwargs.
         kwargs.pop('compute_loss', None)
-        optimizer_config = self.optimizer_group[adapter_name]
         self._lazy_wrap_model()
-        if not inputs:
-            raise ValueError('inputs empty, check your DataLoader outputs')
         self.model.eval()
-        if (isinstance(inputs, dict) and self._not_encoded(inputs)) or (isinstance(inputs, list)
-                                                                        and self._not_encoded(inputs[0])):
-            # Trajectory or List[Trajectory]
-            assert optimizer_config.template is not None, \
-                'Use set_template to add a template when trying to input `List[Trajectory]`'
-            if isinstance(inputs, dict):
-                inputs = [inputs]
-            inputs = optimizer_config.template.batch_encode(inputs)  # noqa
         with torch.no_grad():
             processor: InputProcessor = optimizer_config.processor
             assert isinstance(processor, InputProcessor), 'Set InputProcessor correctly before forwarding'
@@ -1029,6 +1063,9 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
                 outputs = self.model(**inputs)
 
             recorded_routing = rr.finish()
+            if unpack_routing and recorded_routing is not None and processor._is_packed_position_ids(
+                    inputs.get('position_ids')):
+                recorded_routing, = processor._unpack_by_position_ids(inputs['position_ids'], recorded_routing)
 
             inputs['labels'] = labels
             if loss_scale is not None: