From e5a2172b00b3c0a408d402700c488fd936137a2e Mon Sep 17 00:00:00 2001
From: meichangsu1 <1484603386@qq.com>
Date: Wed, 21 Oct 2026 20:51:05 +0800
Subject: [PATCH 1/2] wip

---
 scripts/test_forward_only_micro_batches.py    | 149 ++++++++++
 src/twinkle/model/micro_batch.py              | 127 ++++++++
 .../model/transformers/transformers.py        |  63 +++-
 tests/model/test_micro_batch.py               | 277 ++++++++++++++++++
 4 files changed, 603 insertions(+), 13 deletions(-)
 create mode 100644 scripts/test_forward_only_micro_batches.py

diff --git a/scripts/test_forward_only_micro_batches.py b/scripts/test_forward_only_micro_batches.py
new file mode 100644
index 000000000..17e8a2080
--- /dev/null
+++ b/scripts/test_forward_only_micro_batches.py
@@ -0,0 +1,149 @@
+"""Compare real forward_only results through the public model/client API.
+
+Server (must run this branch with the Transformers backend):
+    python scripts/test_forward_only_micro_batches.py --model-id Qwen/Qwen3-0.6B
+Local model directory, without a server:
+    python scripts/test_forward_only_micro_batches.py --local --model-id /path/to/model --check-logits
+
+The server test creates a temporary adapter and removes it when finished. No
+backward or optimizer step is performed. Use --padding-free only with a model
+and attention backend that support packed sequences. aux_loss is not compared
+against the unsplit result: its contract is weighted micro-batch accumulation.
+"""
+from __future__ import annotations
+
+import argparse
+from copy import deepcopy
+import os
+import time
+import uuid
+
+import torch
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument('--model-id', required=True)
+    parser.add_argument('--base-url', default='http://localhost:9000')
+    parser.add_argument('--api-key', default=os.environ.get('TWINKLE_API_KEY', 'EMPTY_API_KEY'))
+    parser.add_argument('--local', action='store_true')
+    parser.add_argument('--samples', type=int, default=8,
+                        help='Use enough samples per DP rank to execute multiple micro-batches.')
+    parser.add_argument('--micro-batch-size', type=int, default=2)
+    parser.add_argument('--padding-free', action='store_true')
+    parser.add_argument('--check-logits', action='store_true', help='Also return full logits; increases memory/HTTP payload.')
+    parser.add_argument('--rtol', type=float, default=1e-4, help='For BF16/FP16, consider 0.01.')
+    parser.add_argument('--atol', type=float, default=1e-5, help='For BF16/FP16, consider 0.01.')
+    args = parser.parse_args()
+    if args.samples < 3 or not 0 < args.micro_batch_size < args.samples:
+        parser.error('require samples >= 3 and 0 < micro-batch-size < samples')
+    return args
+
+
+def build_inputs(count):
+    # Deliberately unsorted lengths and distinct token patterns expose reordering.
+    lengths = [7, 13, 5, 11, 9, 6, 12, 8]
+    inputs = []
+    for index in range(count):
+        length = lengths[index % len(lengths)]
+        tokens = [1 + (index + offset) % 6 for offset in range(length)]
+        inputs.append({'input_ids': tokens, 'labels': [-100, *tokens[1:]],
+                       'position_ids': list(range(length))})
+    return inputs
+
+
+def assert_same_outputs(expected, actual, lengths, *, check_logits, rtol, atol):
+    for field in ('logps', 'logits') if check_logits else ('logps',):
+        if expected.get(field) is None or actual.get(field) is None:
+            raise AssertionError(f'missing required output: {field}')
+        reference = torch.as_tensor(expected[field]).detach().cpu().float()
+        result = torch.as_tensor(actual[field]).detach().cpu().float()
+        if reference.ndim != result.ndim or result.shape[0] != len(lengths) or reference.shape[0] != len(lengths):
+            raise AssertionError(f'{field}: invalid batch shape {reference.shape} vs {result.shape}')
+        if reference.shape[2:] != result.shape[2:]:
+            raise AssertionError(f'{field}: feature dimensions differ')
+        max_error = 0.0
+        for index, length in enumerate(lengths):
+            if reference.shape[1] < length or result.shape[1] < length:
+                raise AssertionError(f'{field}: missing tokens for sample {index}')
+            expected_sample = reference[index, :length]
+            actual_sample = result[index, :length]
+            if not torch.isfinite(actual_sample).all():
+                raise AssertionError(f'{field}: non-finite output for sample {index}')
+            torch.testing.assert_close(actual_sample, expected_sample, rtol=rtol, atol=atol,
+                                       msg=lambda message: f'{field}, sample {index}: {message}')
+            max_error = max(max_error, (actual_sample - expected_sample).abs().max().item())
+        print(f'  {field}: samples={len(lengths)}, max_abs_error={max_error:.6g}', flush=True)
+    if not check_logits and actual.get('logits') is not None:
+        raise AssertionError('return_logits=False unexpectedly returned logits')
+
+
+def main():
+    args = parse_args()
+    adapter_name = None
+    hook = None
+    forward_calls = []
+    if args.local:
+        from twinkle.model import TransformersModel
+        model = TransformersModel(model_id=args.model_id, mixed_precision='no')
+        model.set_loss('CrossEntropyLoss')
+        model.set_processor('InputProcessor', padding_free=args.padding_free, padding_side='right')
+        model._lazy_wrap_model()
+        hook = model.model.register_forward_pre_hook(lambda module, inputs: forward_calls.append(1))
+    else:
+        from peft import LoraConfig
+        from twinkle import init_twinkle_client
+        from twinkle_client.model import MultiLoraTransformersModel
+        init_twinkle_client(base_url=args.base_url, api_key=args.api_key)
+        model = MultiLoraTransformersModel(model_id=args.model_id)
+        adapter_name = f'e2e-forward-only-{uuid.uuid4().hex[:10]}'
+        model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear'))
+
+    try:
+        if not args.local:
+            model.set_processor('InputProcessor', padding_free=args.padding_free, padding_side='right')
+            model.set_loss('CrossEntropyLoss')
+        inputs = build_inputs(args.samples)
+        lengths = [len(sample['input_ids']) for sample in inputs]
+        # A token budget that is valid even for a single longest sample.
+        budget = max(lengths) * args.micro_batch_size
+        options = {'return_logits': args.check_logits}
+
+        def forward(label, config):
+            start = time.monotonic()
+            previous_calls = len(forward_calls)
+            response = model.forward_only(inputs=deepcopy(inputs), **options, **config)
+            output = response if args.local else response.result
+            if not isinstance(output, dict):
+                raise AssertionError(f'{label}: expected an output dict, got {type(output).__name__}')
+            calls = len(forward_calls) - previous_calls
+            if args.local:
+                if (config and calls <= 1) or (not config and calls != 1):
+                    raise AssertionError(f'{label}: unexpected model forward count {calls}')
+            suffix = f', model_forwards={calls}' if args.local else ''
+            print(f'{label}: {time.monotonic() - start:.2f}s{suffix}', flush=True)
+            return output
+
+        baseline = forward('unsplit', {})
+        configs = [('static', {'micro_batch_size': args.micro_batch_size})]
+        configs.extend((algorithm, {
+            'micro_batch_size': args.micro_batch_size, 'dynamic_batching': True,
+            'max_tokens_per_micro_batch': budget, 'packing_algorithm': algorithm,
+        }) for algorithm in ('ffd', 'kk'))
+        for label, config in configs:
+            output = forward(label, config)
+            assert_same_outputs(baseline, output, lengths, check_logits=args.check_logits,
+                                rtol=args.rtol, atol=args.atol)
+        output = forward('unsplit-after-micro-batching', {})
+        assert_same_outputs(baseline, output, lengths, check_logits=args.check_logits,
+                            rtol=args.rtol, atol=args.atol)
+        print('ALL FORWARD_ONLY MICRO-BATCH CHECKS PASSED', flush=True)
+    finally:
+        if hook is not None:
+            hook.remove()
+        if adapter_name is not None:
+            model.remove_adapter(adapter_name)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/src/twinkle/model/micro_batch.py b/src/twinkle/model/micro_batch.py
index f0b47e050..bf64ceeb7 100644
--- a/src/twinkle/model/micro_batch.py
+++ b/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]
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 1ea0622fa..5f9b07e00 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/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:
diff --git a/tests/model/test_micro_batch.py b/tests/model/test_micro_batch.py
index ec4083e96..53af61161 100644
--- a/tests/model/test_micro_batch.py
+++ b/tests/model/test_micro_batch.py
@@ -247,3 +247,280 @@ def all_gather(states, local_state, *, group):
             MicroBatchConfig(micro_batch_size=1),
             OptimizerConfig(),
         )
+
+
+@pytest.mark.parametrize('algorithm', [None, 'ffd', 'kk'])
+def test_forward_only_micro_batches_restore_order_and_slice_masks(algorithm):
+    import torch
+
+    class ModelHarness:
+        _build_micro_batch_plan = TransformersModel._build_micro_batch_plan
+
+        def __init__(self):
+            self.optimizer_group = {'adapter': SimpleNamespace(
+                processor=InputProcessor(padding_free=False), _dp_group=None,
+                _ensure_dp_group=lambda: None,
+            )}
+            self.calls = []
+
+        def _get_default_group(self):
+            return 'adapter'
+
+        @staticmethod
+        def _not_encoded(_inputs):
+            return False
+
+        _prepare_forward_only_inputs = TransformersModel._prepare_forward_only_inputs
+        _forward_only_micro_batches = TransformersModel._forward_only_micro_batches
+
+        def _forward_only_single(self, *, inputs, sampling_masks, **kwargs):
+            ids = [item['sample_id'] for item in inputs]
+            self.calls.append(ids)
+            assert sampling_masks.tolist() == ids
+            assert kwargs['disable_lora'] is True
+            assert kwargs['task'] == 'embedding'
+            width = max(len(item['input_ids']) for item in inputs)
+            return {'logps': torch.tensor(ids)[:, None].expand(-1, width),
+                    'logits': None, 'routed_experts': torch.tensor(ids)[:, None, None, None]}
+
+    inputs = [{'input_ids': [1] * length, 'sample_id': index}
+              for index, length in enumerate([2, 5, 1, 4, 3])]
+    model = ModelHarness()
+    options = {'micro_batch_size': 2}
+    if algorithm is not None:
+        options.update(dynamic_batching=True, max_tokens_per_micro_batch=10, packing_algorithm=algorithm)
+    outputs = TransformersModel.forward_only.__wrapped__(
+        model, inputs=inputs, sampling_masks=torch.arange(5), disable_lora=True, task='embedding', **options)
+    assert len(model.calls) >= 3
+    assert outputs['logps'].shape == (5, 5)
+    assert outputs['logps'][:, 0].tolist() == list(range(5))
+    assert outputs['routed_experts'][:, 0, 0, 0].tolist() == list(range(5))
+    assert outputs['logits'] is None
+
+
+def test_forward_only_without_micro_batches_prepares_inputs_once():
+    class ModelHarness:
+        optimizer_group = {'adapter': object()}
+
+        def _get_default_group(self):
+            return 'adapter'
+
+        def _prepare_forward_only_inputs(self, inputs, optimizer_config):
+            assert optimizer_config is self.optimizer_group['adapter']
+            self.prepare_calls += 1
+            return [inputs]
+
+        def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
+            assert optimizer_config is self.optimizer_group['adapter']
+            return inputs, kwargs
+
+    model = ModelHarness()
+    model.prepare_calls = 0
+    inputs = {'input_ids': [1]}
+    result = TransformersModel.forward_only.__wrapped__(model, inputs=inputs, disable_lora=True)
+    assert result == ([inputs], {'disable_lora': True})
+    assert model.prepare_calls == 1
+
+
+def test_merge_micro_batch_outputs_handles_nested_results():
+    import torch
+    from twinkle.model.micro_batch import merge_micro_batch_outputs
+
+    result = merge_micro_batch_outputs([
+        {'hidden_states': (torch.tensor([[2., 2.], [0., 0.]]),)},
+        {'hidden_states': (torch.tensor([[1., 1.]]),)},
+    ], [[2, 0], [1]])
+    assert result['hidden_states'][0][:, 0].tolist() == [0., 1., 2.]
+
+
+@pytest.mark.parametrize('padding_free', [False, True])
+@pytest.mark.parametrize('algorithm', [None, 'ffd', 'kk'])
+def test_forward_only_tiny_model_matches_unsplit_outputs(padding_free, algorithm):
+    import torch
+    from unittest.mock import patch
+    from tests.model.test_value_model import _tiny_model_dir
+
+    model_dir = _tiny_model_dir()
+    if padding_free:
+        # GPT2 has no packed attention backend. Remove attention layers so this
+        # test isolates packing/unpacking without cross-document attention leakage.
+        from transformers import GPT2Config, GPT2LMHeadModel
+        config = GPT2Config.from_pretrained(model_dir)
+        config.n_layer = 0
+        GPT2LMHeadModel(config).save_pretrained(model_dir)
+    model = TransformersModel(model_id=model_dir, mixed_precision='no')
+    model.set_loss(CrossEntropyLoss())
+    model.set_processor(InputProcessor(padding_free=padding_free))
+    model._lazy_wrap_model()
+    device = str(next(model.model.parameters()).device)
+    inputs = [{'input_ids': [1, 2, 3, 4], 'labels': [-100, 2, 3, 4]},
+              {'input_ids': [1, 2], 'labels': [-100, 2]},
+              {'input_ids': [1, 2, 3], 'labels': [-100, 2, 3]}]
+    if padding_free:
+        for sample in inputs:
+            sample['position_ids'] = list(range(len(sample['input_ids'])))
+    options = {'micro_batch_size': 2}
+    if algorithm:
+        options.update(dynamic_batching=True, max_tokens_per_micro_batch=8, packing_algorithm=algorithm)
+    with patch('twinkle.processor.base.Platform.get_local_device', return_value=device):
+        expected = model.forward_only(inputs=inputs, return_logits=True)
+        actual = model.forward_only(inputs=inputs, return_logits=True, **options)
+    for index, sample in enumerate(inputs):
+        length = len(sample['input_ids'])
+        torch.testing.assert_close(actual['logps'][index, :length], expected['logps'][index, :length])
+        torch.testing.assert_close(actual['logits'][index, :length], expected['logits'][index, :length])
+    assert not actual['logps'].requires_grad
+
+
+def test_forward_only_micro_batches_match_value_model():
+    import torch
+    from unittest.mock import patch
+    from tests.model.test_value_model import _tiny_model_dir
+    from twinkle.loss import PPOValueLoss
+    from twinkle.model import TransformersValueModel
+
+    model = TransformersValueModel(model_id=_tiny_model_dir(), mixed_precision='no')
+    model.set_loss(PPOValueLoss())
+    model._lazy_wrap_model()
+    device = str(next(model.model.parameters()).device)
+    inputs = [{'input_ids': [1, 2, 3], 'labels': [-100, 2, 3]},
+              {'input_ids': [1, 2], 'labels': [-100, 2]},
+              {'input_ids': [1, 2, 3, 4], 'labels': [-100, 2, 3, 4]}]
+    with patch('twinkle.processor.base.Platform.get_local_device', return_value=device):
+        expected = model.forward_only(inputs=inputs)
+        actual = model.forward_only(inputs=inputs, micro_batch_size=1)
+    for index, sample in enumerate(inputs):
+        length = len(sample['input_ids'])
+        torch.testing.assert_close(actual['values'][index, :length], expected['values'][index, :length])
+    assert actual.get('logps') is None
+    assert not actual['values'].requires_grad
+
+
+@pytest.mark.parametrize('reduction, expected', [('mean', 3.5), ('sum', 6.0)])
+def test_forward_only_accumulates_aux_loss_using_training_micro_batch_scale(reduction, expected):
+    import torch
+
+    loss = CrossEntropyLoss(reduction=reduction)
+
+    class ModelHarness:
+        optimizer_group = {'adapter': SimpleNamespace(get_loss=lambda key: loss)}
+
+        def _get_default_group(self):
+            return 'adapter'
+
+        def _not_encoded(self, inputs):
+            return False
+
+        def _build_micro_batch_plan(self, inputs, config, optimizer_config):
+            return [[0], [1]]
+
+        _prepare_forward_only_inputs = TransformersModel._prepare_forward_only_inputs
+        _forward_only_micro_batches = TransformersModel._forward_only_micro_batches
+
+        def _forward_only_single(self, *, inputs, **kwargs):
+            value = float(inputs[0]['input_ids'][0])
+            return {'aux_loss': torch.tensor(value), 'logps': torch.tensor([[value]])}
+
+    inputs = [{'input_ids': [2], 'labels': [1]},
+              {'input_ids': [4, 4, 4], 'labels': [1, 1, 1]}]
+    outputs = TransformersModel.forward_only.__wrapped__(ModelHarness(), inputs=inputs, micro_batch_size=1)
+    assert outputs['aux_loss'].item() == expected
+    assert outputs['logps'][:, 0].tolist() == [2., 4.]
+
+
+@pytest.mark.parametrize('plan', [[[0], [0]], [[0], [2]], [[0], []]])
+def test_merge_rejects_invalid_sample_plan(plan):
+    import torch
+    from twinkle.model.micro_batch import merge_micro_batch_outputs
+
+    with pytest.raises(ValueError, match='plan'):
+        merge_micro_batch_outputs([{'logps': torch.ones(1, 2)}] * 2, plan)
+
+
+@pytest.mark.parametrize('field, task, first_shape, second_shape', [
+    ('logits', 'causal_lm', (1, 2, 4), (1, 3, 5)),
+    ('logits', 'seq_cls', (1, 2), (1, 3)),
+    ('embeddings', 'embedding', (1, 2), (1, 3)),
+    ('routed_experts', 'causal_lm', (1, 2, 3, 2), (1, 4, 3, 1)),
+])
+def test_merge_rejects_non_sequence_dimension_changes(field, task, first_shape, second_shape):
+    import torch
+    from twinkle.model.micro_batch import merge_micro_batch_outputs
+
+    with pytest.raises(ValueError, match=f'outputs.{field}: non-sequence dimension'):
+        merge_micro_batch_outputs([{field: torch.zeros(first_shape)}, {field: torch.zeros(second_shape)}],
+                                  [[0], [1]], task=task)
+
+
+def test_merge_pooled_outputs_and_unknown_nested_outputs():
+    import torch
+    from twinkle.model.micro_batch import merge_micro_batch_outputs
+
+    result = merge_micro_batch_outputs([
+        {'embeddings': torch.tensor([[2., 3.]]), 'extra': {'scores': torch.tensor([[20.]])}},
+        {'embeddings': torch.tensor([[0., 1.]]), 'extra': {'scores': torch.tensor([[0., 1.]])}},
+    ], [[1], [0]], task='embedding')
+    assert result['embeddings'].tolist() == [[0., 1.], [2., 3.]]
+    assert result['extra']['scores'].tolist() == [[0., 1.], [20., 0.]]
+
+
+def test_merge_fallback_validates_tuple_structure_and_pads_tensors():
+    import torch
+    from twinkle.model.micro_batch import merge_micro_batch_outputs
+
+    outputs = [{'extra_layers': (torch.ones(1, 2, 2, 2),)},
+               {'extra_layers': (torch.ones(1, 2, 3, 3) * 2,)}]
+    result = merge_micro_batch_outputs(outputs, [[1], [0]])
+    assert result['extra_layers'][0].shape == (2, 2, 3, 3)
+    assert result['extra_layers'][0][0, 0, 0, 0] == 2
+    outputs[1]['extra_layers'] = ()
+    with pytest.raises(ValueError, match='outputs.extra_layers: tuple lengths'):
+        merge_micro_batch_outputs(outputs, [[1], [0]])
+
+
+def test_merge_aux_loss_requires_explicit_training_scales():
+    import torch
+    from twinkle.model.micro_batch import merge_micro_batch_outputs
+
+    outputs = [{'aux_loss': torch.tensor(2.), 'logps': torch.tensor([[20.]])},
+               {'aux_loss': torch.tensor(4.), 'logps': torch.tensor([[10.]])}]
+    with pytest.raises(ValueError, match='outputs.aux_loss: one training scale'):
+        merge_micro_batch_outputs(outputs, [[1], [0]])
+    result = merge_micro_batch_outputs(outputs, [[1], [0]], micro_batch_scales=[.25, .75])
+    assert result['aux_loss'].item() == 3.5
+    assert result['logps'][:, 0].tolist() == [10., 20.]
+
+
+@pytest.mark.parametrize('micro_batch_size', [None, 1])
+def test_forward_only_encodes_trajectories_once(micro_batch_size):
+    import torch
+    from unittest.mock import Mock
+
+    encoded = [{'input_ids': [1, 2]}, {'input_ids': [3]}]
+    template = SimpleNamespace(batch_encode=Mock(return_value=encoded))
+    optimizer_config = SimpleNamespace(template=template)
+
+    class ModelHarness:
+        optimizer_group = {'adapter': optimizer_config}
+        _prepare_forward_only_inputs = TransformersModel._prepare_forward_only_inputs
+        _forward_only_micro_batches = TransformersModel._forward_only_micro_batches
+
+        def _get_default_group(self):
+            return 'adapter'
+
+        def _not_encoded(self, inputs):
+            return 'input_ids' not in inputs
+
+        def _build_micro_batch_plan(self, inputs, config, optimizer_config):
+            return [[0], [1]]
+
+        def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
+            assert optimizer_config.template is template
+            assert all('input_ids' in item for item in inputs)
+            return {'embeddings': torch.tensor([[item['input_ids'][0]] for item in inputs])}
+
+    inputs = [{'messages': []}, {'messages': []}]
+    options = {} if micro_batch_size is None else {'micro_batch_size': micro_batch_size}
+    result = TransformersModel.forward_only.__wrapped__(ModelHarness(), inputs=inputs, **options)
+    template.batch_encode.assert_called_once_with(inputs)
+    assert result['embeddings'].tolist() == [[1], [3]]

From ab41f648cb8b19a36bad2a21dde18290716c53e1 Mon Sep 17 00:00:00 2001
From: meichangsu1 <1484603386@qq.com>
Date: Wed, 21 Oct 2026 21:35:33 +0800
Subject: [PATCH 2/2] wip

---
 scripts/test_forward_only_micro_batches.py | 149 ---------------------
 tests/model/test_micro_batch.py            |  13 +-
 2 files changed, 12 insertions(+), 150 deletions(-)
 delete mode 100644 scripts/test_forward_only_micro_batches.py

diff --git a/scripts/test_forward_only_micro_batches.py b/scripts/test_forward_only_micro_batches.py
deleted file mode 100644
index 17e8a2080..000000000
--- a/scripts/test_forward_only_micro_batches.py
+++ /dev/null
@@ -1,149 +0,0 @@
-"""Compare real forward_only results through the public model/client API.
-
-Server (must run this branch with the Transformers backend):
-    python scripts/test_forward_only_micro_batches.py --model-id Qwen/Qwen3-0.6B
-Local model directory, without a server:
-    python scripts/test_forward_only_micro_batches.py --local --model-id /path/to/model --check-logits
-
-The server test creates a temporary adapter and removes it when finished. No
-backward or optimizer step is performed. Use --padding-free only with a model
-and attention backend that support packed sequences. aux_loss is not compared
-against the unsplit result: its contract is weighted micro-batch accumulation.
-"""
-from __future__ import annotations
-
-import argparse
-from copy import deepcopy
-import os
-import time
-import uuid
-
-import torch
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument('--model-id', required=True)
-    parser.add_argument('--base-url', default='http://localhost:9000')
-    parser.add_argument('--api-key', default=os.environ.get('TWINKLE_API_KEY', 'EMPTY_API_KEY'))
-    parser.add_argument('--local', action='store_true')
-    parser.add_argument('--samples', type=int, default=8,
-                        help='Use enough samples per DP rank to execute multiple micro-batches.')
-    parser.add_argument('--micro-batch-size', type=int, default=2)
-    parser.add_argument('--padding-free', action='store_true')
-    parser.add_argument('--check-logits', action='store_true', help='Also return full logits; increases memory/HTTP payload.')
-    parser.add_argument('--rtol', type=float, default=1e-4, help='For BF16/FP16, consider 0.01.')
-    parser.add_argument('--atol', type=float, default=1e-5, help='For BF16/FP16, consider 0.01.')
-    args = parser.parse_args()
-    if args.samples < 3 or not 0 < args.micro_batch_size < args.samples:
-        parser.error('require samples >= 3 and 0 < micro-batch-size < samples')
-    return args
-
-
-def build_inputs(count):
-    # Deliberately unsorted lengths and distinct token patterns expose reordering.
-    lengths = [7, 13, 5, 11, 9, 6, 12, 8]
-    inputs = []
-    for index in range(count):
-        length = lengths[index % len(lengths)]
-        tokens = [1 + (index + offset) % 6 for offset in range(length)]
-        inputs.append({'input_ids': tokens, 'labels': [-100, *tokens[1:]],
-                       'position_ids': list(range(length))})
-    return inputs
-
-
-def assert_same_outputs(expected, actual, lengths, *, check_logits, rtol, atol):
-    for field in ('logps', 'logits') if check_logits else ('logps',):
-        if expected.get(field) is None or actual.get(field) is None:
-            raise AssertionError(f'missing required output: {field}')
-        reference = torch.as_tensor(expected[field]).detach().cpu().float()
-        result = torch.as_tensor(actual[field]).detach().cpu().float()
-        if reference.ndim != result.ndim or result.shape[0] != len(lengths) or reference.shape[0] != len(lengths):
-            raise AssertionError(f'{field}: invalid batch shape {reference.shape} vs {result.shape}')
-        if reference.shape[2:] != result.shape[2:]:
-            raise AssertionError(f'{field}: feature dimensions differ')
-        max_error = 0.0
-        for index, length in enumerate(lengths):
-            if reference.shape[1] < length or result.shape[1] < length:
-                raise AssertionError(f'{field}: missing tokens for sample {index}')
-            expected_sample = reference[index, :length]
-            actual_sample = result[index, :length]
-            if not torch.isfinite(actual_sample).all():
-                raise AssertionError(f'{field}: non-finite output for sample {index}')
-            torch.testing.assert_close(actual_sample, expected_sample, rtol=rtol, atol=atol,
-                                       msg=lambda message: f'{field}, sample {index}: {message}')
-            max_error = max(max_error, (actual_sample - expected_sample).abs().max().item())
-        print(f'  {field}: samples={len(lengths)}, max_abs_error={max_error:.6g}', flush=True)
-    if not check_logits and actual.get('logits') is not None:
-        raise AssertionError('return_logits=False unexpectedly returned logits')
-
-
-def main():
-    args = parse_args()
-    adapter_name = None
-    hook = None
-    forward_calls = []
-    if args.local:
-        from twinkle.model import TransformersModel
-        model = TransformersModel(model_id=args.model_id, mixed_precision='no')
-        model.set_loss('CrossEntropyLoss')
-        model.set_processor('InputProcessor', padding_free=args.padding_free, padding_side='right')
-        model._lazy_wrap_model()
-        hook = model.model.register_forward_pre_hook(lambda module, inputs: forward_calls.append(1))
-    else:
-        from peft import LoraConfig
-        from twinkle import init_twinkle_client
-        from twinkle_client.model import MultiLoraTransformersModel
-        init_twinkle_client(base_url=args.base_url, api_key=args.api_key)
-        model = MultiLoraTransformersModel(model_id=args.model_id)
-        adapter_name = f'e2e-forward-only-{uuid.uuid4().hex[:10]}'
-        model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear'))
-
-    try:
-        if not args.local:
-            model.set_processor('InputProcessor', padding_free=args.padding_free, padding_side='right')
-            model.set_loss('CrossEntropyLoss')
-        inputs = build_inputs(args.samples)
-        lengths = [len(sample['input_ids']) for sample in inputs]
-        # A token budget that is valid even for a single longest sample.
-        budget = max(lengths) * args.micro_batch_size
-        options = {'return_logits': args.check_logits}
-
-        def forward(label, config):
-            start = time.monotonic()
-            previous_calls = len(forward_calls)
-            response = model.forward_only(inputs=deepcopy(inputs), **options, **config)
-            output = response if args.local else response.result
-            if not isinstance(output, dict):
-                raise AssertionError(f'{label}: expected an output dict, got {type(output).__name__}')
-            calls = len(forward_calls) - previous_calls
-            if args.local:
-                if (config and calls <= 1) or (not config and calls != 1):
-                    raise AssertionError(f'{label}: unexpected model forward count {calls}')
-            suffix = f', model_forwards={calls}' if args.local else ''
-            print(f'{label}: {time.monotonic() - start:.2f}s{suffix}', flush=True)
-            return output
-
-        baseline = forward('unsplit', {})
-        configs = [('static', {'micro_batch_size': args.micro_batch_size})]
-        configs.extend((algorithm, {
-            'micro_batch_size': args.micro_batch_size, 'dynamic_batching': True,
-            'max_tokens_per_micro_batch': budget, 'packing_algorithm': algorithm,
-        }) for algorithm in ('ffd', 'kk'))
-        for label, config in configs:
-            output = forward(label, config)
-            assert_same_outputs(baseline, output, lengths, check_logits=args.check_logits,
-                                rtol=args.rtol, atol=args.atol)
-        output = forward('unsplit-after-micro-batching', {})
-        assert_same_outputs(baseline, output, lengths, check_logits=args.check_logits,
-                            rtol=args.rtol, atol=args.atol)
-        print('ALL FORWARD_ONLY MICRO-BATCH CHECKS PASSED', flush=True)
-    finally:
-        if hook is not None:
-            hook.remove()
-        if adapter_name is not None:
-            model.remove_adapter(adapter_name)
-
-
-if __name__ == '__main__':
-    main()
diff --git a/tests/model/test_micro_batch.py b/tests/model/test_micro_batch.py
index 53af61161..685e5c78a 100644
--- a/tests/model/test_micro_batch.py
+++ b/tests/model/test_micro_batch.py
@@ -6,6 +6,7 @@
 from twinkle.loss import CrossEntropyLoss, GRPOLoss
 from twinkle.loss.base import Loss
 from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches
+from twinkle.model.optimizer_group import BaseOptimizerGroup
 from twinkle.model.transformers.transformers import TransformersModel
 from twinkle.processor import InputProcessor
 
@@ -78,8 +79,15 @@ def test_loss_without_micro_batch_semantics_fails_when_split():
 
 def test_transformers_forward_backward_keeps_original_default_path():
     class ModelHarness:
+        _forward_backward_single = TransformersModel._forward_backward_single
+        _not_encoded = staticmethod(lambda inputs: False)
+
+        def _get_default_group(self):
+            return 'adapter'
+
         def __init__(self):
             self.calls = []
+            self.optimizer_group = {'adapter': BaseOptimizerGroup(loss_instance=CrossEntropyLoss())}
 
         def forward(self, *, inputs, **_kwargs):
             self.calls.append(('forward', inputs))
@@ -103,8 +111,9 @@ def backward(self, **_kwargs):
 
 
 def test_transformers_forward_backward_executes_real_micro_batches():
-    class OptimizerConfig:
+    class OptimizerConfig(BaseOptimizerGroup):
         def __init__(self):
+            super().__init__()
             self.processor = InputProcessor(padding_free=False)
             self.template = None
             self.train_status = SimpleNamespace(loss_value=None, num_tokens=0.0)
@@ -117,6 +126,7 @@ def _ensure_dp_group(self):
             pass
 
     class ModelHarness:
+        _forward_backward_single = TransformersModel._forward_backward_single
         _build_micro_batch_plan = TransformersModel._build_micro_batch_plan
         _forward_backward_micro_batch = TransformersModel._forward_backward_micro_batch
         _forward_backward_micro_batches = TransformersModel._forward_backward_micro_batches
@@ -164,6 +174,7 @@ def backward(self, *, sync_gradients, **_kwargs):
         sync_gradients=True,
     )
 
+    assert outputs['loss'] == 2.0
     assert model.forward_batches == [[0, 1], [2, 3]]
     assert model.backward_calls == [(False, 1.0), (True, 1.0)]
     assert model.optimizer_group['adapter'].train_status.num_tokens == 1.0