diff --git a/cookbook/rl/grpo/dp_balance_grpo.py b/cookbook/rl/grpo/dp_balance_grpo.py
index d629fcd05..3b39613a3 100644
--- a/cookbook/rl/grpo/dp_balance_grpo.py
+++ b/cookbook/rl/grpo/dp_balance_grpo.py
@@ -1,14 +1,16 @@
-"""GRPO on GSM8K with DP workload balancing (`--balance-batch`).
+"""GRPO on GSM8K with DP workload balancing.
 
 `forward_backward` dispatches every mini-batch with `slice_dp`: each model DP rank trains on a contiguous
 slice, and the step waits for the slowest rank. Thinking completions vary a lot in length, so those
-slices are often uneven. `plan_dp_balance` reorders the mini-batch before dispatch so every rank keeps its
-sample count but gets a similar workload. It estimates each rank's cost by running the model's own
-micro-batch planner on that rank's slice, so it follows both the input layout (padded or packed) and the
-micro-batch options (fixed or dynamic). With `GRPOLoss` (a per-sequence mean) the update does not change;
-only the waiting shrinks. Training metrics gain `dp_balance/*` entries: the workload imbalance (busiest
-rank over the mean), per-rank tokens, micro-batches per rank and, for padded inputs, the padding ratio,
-before and after balancing.
+slices are often uneven. `TransformersModel` balances them by default: before dispatch it reorders the
+mini-batch, together with every per-sample argument such as `old_logps` and `advantages`, so every rank
+keeps its sample count but gets a similar workload, and it returns the outputs in the caller's order. It
+estimates each rank's cost by running the model's own micro-batch planner on that rank's slice, so it
+follows both the input layout (padded or packed) and the micro-batch options (fixed or dynamic). With
+`GRPOLoss` (a per-sequence mean) the update does not change; only the waiting shrinks. Training metrics
+gain `dp_balance/*` entries: the workload imbalance (busiest rank over the mean), per-rank tokens,
+micro-batches per rank and, for padded inputs, the padding ratio, before and after balancing.
+`--no-balance-batch` builds the model with `dp_balance=False`.
 
 Packed inputs (`--padding-free`) of linear-attention models such as Qwen3.5 need flash-linear-attention:
 the native GatedDeltaNet does not reset its state between packed sequences. On NPU that means
@@ -21,7 +23,6 @@
     sh dp_balance_grpo.sh --no-balance-batch      # the same run without balancing, for comparison
 """
 import argparse
-from functools import partial
 from typing import Any, Dict, List
 
 from peft import LoraConfig
@@ -36,8 +37,7 @@
 from twinkle.dataset import Dataset, DatasetMeta
 from twinkle.hub import HubOperation
 from twinkle.metric import CompletionRewardMetric
-from twinkle.model import MicroBatchConfig, TransformersModel, plan_dp_balance
-from twinkle.model.dp_balance import estimate_workload
+from twinkle.model import TransformersModel
 from twinkle.patch import Patch
 from twinkle.preprocessor.llm import GSM8KProcessor
 from twinkle.processor import InputProcessor
@@ -108,24 +108,24 @@ def main():
     # slice_dp splits by data_world_size (dp x fsdp), not dp_size: FSDP2 alone still has fsdp_size data ranks.
     dp_size = model_mesh.data_world_size
 
-    # Must match what forward_backward receives: the plan simulates the model's own micro-batch planner.
     micro_batch_kwargs = {
         'micro_batch_size': args.micro_batch_size,
         'dynamic_batching': args.dynamic_batching,
         'max_tokens_per_micro_batch': args.max_tokens_per_micro_batch if args.dynamic_batching else None,
         'packing_algorithm': 'ffd',
     }
-    micro_batch_config = MicroBatchConfig(**micro_batch_kwargs)
-    # The linear term of the workload estimate grows with the hidden size.
-    hidden_size = (getattr(hf_config, 'text_config', None) or hf_config).hidden_size
-    workload_fn = partial(estimate_workload, hidden_size=hidden_size)
 
     model_kwargs = {'attn_implementation': args.attn_implementation} if args.attn_implementation else {}
     if is_qwen35:
         from transformers import Qwen3_5ForConditionalGeneration
         model_kwargs['model_cls'] = Qwen3_5ForConditionalGeneration
     model = TransformersModel(
-        model_id=args.model_id, device_mesh=model_mesh, remote_group='model', strategy=args.strategy, **model_kwargs)
+        model_id=args.model_id,
+        device_mesh=model_mesh,
+        remote_group='model',
+        strategy=args.strategy,
+        dp_balance=args.balance_batch,
+        **model_kwargs)
     if device_type == 'NPU' and args.padding_free:
         model.apply_patch(NPULinearAttentionKernel)
     # Language-model projections only (full and linear attention, MLP); 'all-linear' would add the vision tower.
@@ -210,25 +210,10 @@ def create_gsm8k_dataset():
         metrics.accumulate(completion_lengths=completion_lengths, rewards={'accuracy': rewards})
         advantages = advantage_fn(rewards, num_generations=args.num_generations, scale='group').tolist()
 
-        balance_metrics: Dict[str, List[float]] = {}
         for mb_start in range(0, len(all_inputs), args.mini_batch_size):
             mb_inputs = all_inputs[mb_start:mb_start + args.mini_batch_size]
             mb_old_logps = all_old_logps[mb_start:mb_start + args.mini_batch_size]
             mb_advantages = advantages[mb_start:mb_start + args.mini_batch_size]
-            if args.balance_batch:
-                plan = plan_dp_balance(
-                    mb_inputs,
-                    dp_size,
-                    padding_free=args.padding_free,
-                    micro_batch_config=micro_batch_config,
-                    workload_fn=workload_fn,
-                )
-                # Every per-sample list is reordered with the same plan so they stay aligned.
-                mb_inputs, mb_old_logps, mb_advantages = (
-                    plan.apply(values) for values in (mb_inputs, mb_old_logps, mb_advantages))
-                for key, value in plan.metrics.items():
-                    balance_metrics.setdefault(key, []).append(value)
-
             model.forward_backward(
                 inputs=mb_inputs, old_logps=mb_old_logps, advantages=mb_advantages, **micro_batch_kwargs)
             model.clip_grad_and_step()
@@ -238,7 +223,6 @@ def create_gsm8k_dataset():
 
         log_dict = metrics.calculate()
         log_dict.update(model.calculate_metric(is_training=True))
-        log_dict.update({key: sum(values) / len(values) for key, values in balance_metrics.items()})
         logger.info(f'[Step {optim_step}/{args.max_steps}] {log_dict}')
 
     logger.info(f'Training completed. optim_steps={optim_step}')
diff --git a/cookbook/rl/grpo/dp_balance_grpo.sh b/cookbook/rl/grpo/dp_balance_grpo.sh
index b71fbc50f..d93e23046 100644
--- a/cookbook/rl/grpo/dp_balance_grpo.sh
+++ b/cookbook/rl/grpo/dp_balance_grpo.sh
@@ -2,8 +2,8 @@
 set -eu
 
 # GRPO on GSM8K with DP workload balancing across the model DP ranks.
-# Thinking completions vary a lot in length; --balance-batch evens out the work per rank before each
-# optimizer step. Defaults: Qwen3.5-4B, FSDP2 on 2 devices (2 data ranks) plus 2 vLLM devices, padded
+# Thinking completions vary a lot in length; TransformersModel evens out the work per rank before each
+# optimizer step (on by default, --no-balance-batch turns it off). Defaults: Qwen3.5-4B, FSDP2 on 2 devices (2 data ranks) plus 2 vLLM devices, padded
 # inputs, dynamic batching with at most 16 rows and 8192 tokens per micro-batch.
 # Override at invocation, e.g.:
 #   sh dp_balance_grpo.sh --padding-free            # packed inputs
diff --git a/docs/source_en/Components/Model/TransformersModel.md b/docs/source_en/Components/Model/TransformersModel.md
index 8c10e4b98..9a6698eec 100644
--- a/docs/source_en/Components/Model/TransformersModel.md
+++ b/docs/source_en/Components/Model/TransformersModel.md
@@ -16,6 +16,7 @@ class TransformersModel:
                  fsdp_config: Dict[str, Any] = None,
                  grad_scaler_config: Dict[str, Any] = None,
                  memory_efficient_init: bool = False,
+                 dp_balance: bool = True,
                  **kwargs):
         ...
 
@@ -49,6 +50,7 @@ class TransformersModel:
   Constraint: the effective `output_dtype` must equal `param_dtype` (or be omitted); otherwise construction raises `ValueError` — a mismatch introduces a dtype conversion at FSDP unit boundaries, and gradient checkpointing recompute then fails on mixed-dtype tensors. `reduce_dtype` only affects gradient reduction and is unconstrained.
 - grad_scaler_config: PyTorch's grad_scaler initialization configuration, see: [PyTorch's GradScaler constructor](https://github.com/pytorch/pytorch/blob/main/torch/cuda/amp/grad_scaler.py#L25)
 - memory_efficient_init: Whether to enable memory-efficient model initialization for FSDP. When enabled, only rank 0 loads full weights and broadcasts sharded parameters to other ranks, reducing peak memory usage during initialization. Default `False`. Note: The optimization currently only applies to transformers <= 4.57.6; for transformers >= 5.0.0, it may lead to negative performance impact.
+- dp_balance: Whether to balance the workload of the DP ranks, default `True`. In Ray mode, before `forward_backward` / `forward_only` split a batch across the DP ranks, the driver reorders it so every rank keeps its sample count but gets a similar workload, estimated with the same micro-batch planner the workers use. Every per-sample argument (`old_logps`, `advantages`, ...) is reordered with it, and the returned outputs are put back in the caller's order. In local (torchrun) mode, where each rank holds only its own slice, the ranks gather their sequence lengths, plan the same way, swap samples, and send each per-sample output back to the rank that holds its input. With a per-sample loss the update does not change; only the wait for the slowest rank shrinks. Training metrics gain `dp_balance/*` entries (workload imbalance, per-rank tokens, micro-batches and padding ratio, before and after). Losses that pair samples keep each pair together (`dp_balance_group_size`: 2 for DPO-style and reward losses, `num_generations` for REAL); embedding and listwise reranker losses, and plain callables, are never reordered. After a reordered `forward_only`, `calculate_loss` rejects per-sample lists, which would no longer line up; pass them to `forward_only` instead. Set `False` to keep the original split.
 - kwargs:
   - If you don't want to pass the model config field, you can put scattered configurations here. These parameters will be passed to `from_pretrained` or `from_config` later.
 
diff --git "a/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md" "b/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md"
index 85fc81588..0ae9aa679 100644
--- "a/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md"
+++ "b/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md"
@@ -16,6 +16,7 @@ class TransformersModel:
                  fsdp_config: Dict[str, Any] = None,
                  grad_scaler_config: Dict[str, Any] = None,
                  memory_efficient_init: bool = False,
+                 dp_balance: bool = True,
                  **kwargs):
         ...
 
@@ -49,6 +50,7 @@ class TransformersModel:
   约束：生效的`output_dtype`必须与`param_dtype`一致（或缺省），否则构造时抛出`ValueError`——两者不一致会在FSDP单元边界引入dtype转换，梯度检查点重算时张量dtype不一致会导致训练报错；`reduce_dtype`只作用于梯度归约，不受该约束。
 - grad_scaler_config: PyTorch的grad_scaler初始化配置，参见：[PyTorch的GradScaler构造](https://github.com/pytorch/pytorch/blob/main/torch/cuda/amp/grad_scaler.py#L25)
 - memory_efficient_init: 是否启用FSDP内存高效初始化。启用后仅rank 0加载完整权重，其余rank通过广播获取分片参数，降低初始化阶段的内存和显存峰值。默认`False`。注意：该优化目前仅适用于 transformers <= 4.57.6；对于 transformers >= 5.0.0，可能会导致负面性能影响。
+- dp_balance: 是否均衡各DP rank的计算量，默认`True`。Ray模式下，`forward_backward` / `forward_only`把batch切给各DP rank之前，driver会重排batch，使每个rank样本数不变、计算量接近（用与worker相同的micro-batch规划来估算）。每个样本对应的参数（`old_logps`、`advantages`等）随之一起重排，返回的输出会还原为调用方的顺序。local（torchrun）模式下每个rank只持有自己那份数据，各rank会先汇总序列长度、用同样的方法规划，再互换样本，算完把每个样本的输出送回持有它的rank。对按样本求均值的loss，参数更新不变，只是缩短了等最慢rank的时间。训练指标会多出`dp_balance/*`（均衡前后的负载不均衡度、各rank token数、micro-batch数和padding比例）。成对的loss会让每对样本留在一起（`dp_balance_group_size`：DPO类和reward loss为2，REAL为`num_generations`）；embedding、listwise reranker loss以及普通函数形式的loss不会重排。`forward_only`重排后，`calculate_loss`会拒绝按样本的list参数（顺序已对不上），请改为传给`forward_only`。设为`False`则保持原始切分。
 - kwargs:
   - 如果你不希望传递模型config字段，可以把零星的配置从这里放置进去。后续这些参数会传递到`from_pretrained`或者`from_config`中。
 
diff --git a/src/twinkle/infra/dispatch.py b/src/twinkle/infra/dispatch.py
index 185ae32ce..9775d674a 100644
--- a/src/twinkle/infra/dispatch.py
+++ b/src/twinkle/infra/dispatch.py
@@ -742,6 +742,15 @@ def _stash_channels(result: Any, channels: Sequence[Union[int, str]]) -> Any:
     raise TypeError(f'stash requires a tuple or dict return value, got {type(result).__name__}')
 
 
+def _collect_in_input_order(collect, plan):
+    """Collect as ``collect`` does, then undo the DP-balance reordering of the call's batch."""
+
+    def collect_func(result, device_mesh=None):
+        return plan.restore_outputs(_collect_func(collect, result, device_mesh))
+
+    return collect_func
+
+
 def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp', 'last_pp_first'], Callable] = 'slice',
                     execute: Literal['first', 'peer', 'all'] = 'all',
                     collect: Union[Literal['none', 'flatten', 'mean', 'sum', 'first', 'last_pp'], Callable] = 'none',
@@ -755,7 +764,8 @@ def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp', 'last_pp
                     control: bool = False,
                     replay: bool = False,
                     worker_selector: Optional[str] = None,
-                    stash: Optional[Sequence[Union[int, str]]] = None):
+                    stash: Optional[Sequence[Union[int, str]]] = None,
+                    dp_balance: bool = False):
     """Patch each method called from remote(which class should be decorated with `remote_class`) with this decorator.
 
     Args:
@@ -820,6 +830,13 @@ def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp', 'last_pp
             whole. The caller's collect then materialises only the rest of the container, and the
             next remote call dereferences the stashed references at its worker. Local mode has no
             remote path, so the natural value passes through untouched.
+        dp_balance: Only with ``dispatch='slice_dp'``. Before slicing, the driver asks the instance's
+            ``_dp_balance_call(name, args, kwargs)`` for a ``DPBalancePlan`` that evens out the
+            per-rank workload, reorders the batch-aligned arguments with it, and puts the
+            batch-aligned fields of the collected result back in the caller's order. Calls whose
+            arguments hold object refs are dispatched unchanged. In local mode, where each rank holds
+            only its own slice, the instance's ``_dp_balance_local(run, args, kwargs)`` runs the call
+            instead and swaps samples between the ranks.
 
     Wrapped continuous-work calls also accept ``_session_ids``, a list or tuple of
     non-empty session strings (or ``None`` for ordinary load balancing), aligned
@@ -887,6 +904,9 @@ def _finish_dispatch():
                         raise ValueError(f'{func.__name__}: no in-flight request for sticky key '
                                          f'{_sticky_key_value!r}')
                 if infra._mode == 'local':
+                    balance_local = getattr(self, '_dp_balance_local', None) if dp_balance else None
+                    if balance_local is not None and dispatch == 'slice_dp':
+                        return balance_local(functools.partial(func, self), args, kwargs)
                     return func(self, *args, **kwargs)
                 elif infra._mode == 'ray':
                     check_unsafe(*args, **kwargs)
@@ -995,6 +1015,15 @@ def _finish_dispatch():
                                                            _batch_len,
                                                            _rgt,
                                                            on_dispatched=_finish_dispatch, session_ids=session_ids)
+                        _call_collect = collect
+                        if (dp_balance and _dispatch == 'slice_dp' and not _subset
+                                and not RayHelper.has_ref(args, kwargs)):
+                            balance_call = getattr(self, '_dp_balance_call', None)
+                            balanced = balance_call(func.__name__, args, kwargs) if balance_call else None
+                            if balanced is not None:
+                                _plan, args, kwargs = balanced
+                                if _plan is not None:
+                                    _call_collect = _collect_in_input_order(collect, _plan)
                         if RayHelper.has_ref(args, kwargs):
                             # If has any object-ref, dispatch in worker, because we don't know the structure in the ref.
                             # for example, dataloader returns any data list.
@@ -1028,7 +1057,7 @@ def _track_submission(worker, result):
                             result = ray.get(result)
                         # This is a result future, call it to get the actual result
                         result_func = RayHelper.do_get_and_collect_func(
-                            _collect_func, collect, result, device_mesh, timeout=_rgt)
+                            _collect_func, _call_collect, result, device_mesh, timeout=_rgt)
                         _local_lazy_collect = infra._lazy_collect
                         recorder = getattr(self, '_record_worker_configuration', None)
                         if func.__name__ == '__iter__':
diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py
index bc02ebe92..730b7c5a8 100644
--- a/src/twinkle/loss/base.py
+++ b/src/twinkle/loss/base.py
@@ -1,4 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import Optional
+
 from twinkle.data_format import InputFeature, LossOutput, ModelOutput
 
 
@@ -8,6 +10,9 @@ class Loss:
     require_entropy = False
     require_logps = True
     require_values = False
+    # Consecutive samples that must stay together, in order, on one DP rank when the model reorders a
+    # batch to balance DP workloads (2 for interleaved chosen/rejected pairs). None forbids reordering.
+    dp_balance_group_size: Optional[int] = 1
 
     def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput:
         ...
@@ -35,3 +40,10 @@ def micro_batch_kwargs(self, kwargs, indices, batch_size):
     def batch_normalizer(self, inputs: list[InputFeature]):
         """Return this batch's additive reduction denominator, or ``None`` for mean-reduced losses."""
         return None
+
+
+def dp_balance_group_size(loss) -> Optional[int]:
+    """The ``dp_balance_group_size`` of ``loss``; a plain callable may depend on order, so it gets None."""
+    if loss is None:
+        return 1
+    return loss.dp_balance_group_size if isinstance(loss, Loss) else None
diff --git a/src/twinkle/loss/dpo.py b/src/twinkle/loss/dpo.py
index b50eb0f11..c24fcbbc0 100644
--- a/src/twinkle/loss/dpo.py
+++ b/src/twinkle/loss/dpo.py
@@ -40,6 +40,8 @@ def _cap_exp(value: 'torch.Tensor', cap: float = -1) -> 'torch.Tensor':
 class PreferenceLossBase(Loss):
     """Base class for preference optimization losses with shared utilities."""
 
+    dp_balance_group_size = 2
+
     def __init__(self, ignore_index: int = -100):
         self.ignore_index = ignore_index
 
@@ -478,6 +480,9 @@ class KTOLoss(PreferenceLossBase):
         ignore_index: Index to ignore in labels (default: -100).
     """
 
+    # Unpaired: every completion stands alone.
+    dp_balance_group_size = 1
+
     def __init__(
         self,
         beta: float = 0.1,
diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py
index 130edb03f..142aa1dc7 100644
--- a/src/twinkle/loss/grpo.py
+++ b/src/twinkle/loss/grpo.py
@@ -1047,6 +1047,11 @@ def __init__(self, real_tau: float = 0.5, num_generations: int = 1, **kwargs):
         self.real_tau = real_tau
         self.num_generations = num_generations
 
+    @property
+    def dp_balance_group_size(self) -> int:
+        # Each generation group is read off consecutive rows of the local batch.
+        return self.num_generations
+
     def _reduce_loss(
         self,
         logps: 'torch.Tensor',
diff --git a/src/twinkle/loss/infonce.py b/src/twinkle/loss/infonce.py
index 928a36766..c24a1ed2e 100644
--- a/src/twinkle/loss/infonce.py
+++ b/src/twinkle/loss/infonce.py
@@ -117,6 +117,8 @@ class EmbeddingLoss(Loss):
     """
 
     supports_mrl = True
+    # Anchors, positives and negatives are laid out by position in groups of varying size.
+    dp_balance_group_size = None
 
     def __init__(self, mrl_dims: Optional[Dict[int, float]] = None, **kwargs):
         if mrl_dims and not self.supports_mrl:
diff --git a/src/twinkle/loss/reranker.py b/src/twinkle/loss/reranker.py
index 533290981..8193bf2d6 100644
--- a/src/twinkle/loss/reranker.py
+++ b/src/twinkle/loss/reranker.py
@@ -73,6 +73,8 @@ class ListwiseRerankerLoss(Loss):
     require_logits = True
     require_entropy = False
     require_logps = False
+    # A query's candidates form a group of varying size, opened by its positive.
+    dp_balance_group_size = None
 
     def __init__(self, temperature: float = 1.0, min_group_size: int = 2, **kwargs):
         if temperature <= 0:
diff --git a/src/twinkle/loss/reward.py b/src/twinkle/loss/reward.py
index a3d60bf2f..e6b5dece9 100644
--- a/src/twinkle/loss/reward.py
+++ b/src/twinkle/loss/reward.py
@@ -34,6 +34,7 @@ class RewardLoss(Loss):
     # Reads outputs['logits'] (the [B, 1] score head), not per-token logps.
     require_logits = True
     require_logps = False
+    dp_balance_group_size = 2
 
     def __init__(self, center_rewards_coefficient: float = 0.0, **kwargs):
         super().__init__()
diff --git a/src/twinkle/model/base.py b/src/twinkle/model/base.py
index 4d971b32e..3f8a192a4 100644
--- a/src/twinkle/model/base.py
+++ b/src/twinkle/model/base.py
@@ -347,6 +347,82 @@ def _not_encoded(inputs):
             k in ('input_ids', 'input_embedding') or k.endswith('_input_ids') or k.endswith('_input_embedding')
             for k in inputs)
 
+    def _dp_balance_call(self, method: str, args: tuple, kwargs: Dict[str, Any]):
+        """Driver side of ``remote_function(dp_balance=True)``.
+
+        Returns ``(plan, args, kwargs)`` with the batch-aligned arguments reordered and a ``_dp_balance``
+        keyword for the workers, or ``None`` to dispatch the call unchanged. ``plan`` is ``None`` when the
+        batch keeps its order, so the result needs no restoring.
+        """
+        from .dp_balance import call_lengths, plan_call
+        if getattr(self, '_dp_balance_off', False):
+            return None
+        mesh = getattr(self, 'device_mesh', None)
+        if mesh is None or mesh.data_world_size <= 1:
+            return None
+        lengths = call_lengths(args, kwargs)
+        if lengths is None:
+            return None
+        spec = self._dp_balance_spec(adapter_name=kwargs.get('adapter_name'), loss_key=kwargs.get('loss_key'))
+        if not spec['enabled']:
+            # Fixed at construction, so later calls skip the blocking spec round trip.
+            self._dp_balance_off = True
+            return None
+        if spec['group_size'] is None:
+            return None
+        plan = plan_call(lengths, mesh.data_world_size, spec, kwargs)
+        if plan is None:
+            return None
+        if plan.reorders:
+            args, kwargs = plan.reorder_call(args, kwargs)
+        kwargs = dict(kwargs, _dp_balance={'reordered': plan.reorders, 'metrics': plan.metrics})
+        return (plan if plan.reorders else None), args, kwargs
+
+    def _dp_balance_local(self, run: Callable, args: tuple, kwargs: Dict[str, Any]):
+        """Local-mode side of ``remote_function(dp_balance=True)``, where each rank holds only its own slice.
+
+        The ranks gather their sequence lengths and plan the joined batch exactly as the Ray driver would.
+        When the plan reorders, they swap samples so each rank runs its share of the reordered batch, then
+        send every per-sample output back to the rank that holds its input. Each decision below is made
+        alike on all ranks, from configuration or gathered data, so the collectives never diverge.
+        """
+        import torch.distributed as dist
+        from .dp_balance import call_leaves, call_lengths, exchange_call, leaf_layout, plan_call, restore_rank_outputs
+        from .dp_balance import to_cpu
+        mesh = getattr(self, 'device_mesh', None)
+        if (getattr(self, '_dp_balance_off', False) or mesh is None or mesh.data_world_size <= 1
+                or not dist.is_initialized()
+                or sorted(mesh.mesh.flatten().tolist()) != list(range(dist.get_world_size()))):
+            return run(*args, **kwargs)
+        spec = self._dp_balance_spec(adapter_name=kwargs.get('adapter_name'), loss_key=kwargs.get('loss_key'))
+        if not spec['enabled']:
+            self._dp_balance_off = True
+            return run(*args, **kwargs)
+        if spec['group_size'] is None:
+            return run(*args, **kwargs)
+
+        def all_gather(value):
+            gathered = [None] * dist.get_world_size()
+            dist.all_gather_object(gathered, value)
+            # One rank per data rank, in data-rank order; TP / SP peers hold the same slice.
+            return [gathered[rank] for rank in mesh.get_collect_ranks()]
+
+        lengths = call_lengths(args, kwargs)
+        summaries = all_gather(None if lengths is None else (lengths, leaf_layout(call_leaves(args, kwargs))))
+        if (any(summary is None for summary in summaries) or len({len(summary[0]) for summary in summaries}) != 1
+                or any(summary[1] != summaries[0][1] for summary in summaries)):
+            # Joining the slices back into the batch get_slice cut needs equal slices with the same fields.
+            return run(*args, **kwargs)
+        plan = plan_call([length for summary in summaries for length in summary[0]], len(summaries), spec, kwargs)
+        if plan is None:
+            return run(*args, **kwargs)
+        balance = {'reordered': plan.reorders, 'metrics': plan.metrics}
+        if not plan.reorders:
+            return run(*args, **dict(kwargs, _dp_balance=balance))
+        rank_leaves = all_gather(to_cpu(call_leaves(args, kwargs)))
+        args, kwargs = exchange_call(plan, rank_leaves, mesh.data_rank, args, kwargs)
+        output = run(*args, **dict(kwargs, _dp_balance=balance))
+        return restore_rank_outputs(output, all_gather(to_cpu(output)), plan, mesh.data_rank)
     def upload_to_hub(self,
                       checkpoint_dir: str,
                       hub_model_id: str,
diff --git a/src/twinkle/model/dp_balance.py b/src/twinkle/model/dp_balance.py
index 43ba263ba..2869f6efa 100644
--- a/src/twinkle/model/dp_balance.py
+++ b/src/twinkle/model/dp_balance.py
@@ -20,6 +20,8 @@
 """
 
 from __future__ import annotations
+import copy
+import functools
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 from typing import Any
@@ -123,6 +125,228 @@ def _imbalance(values: Sequence[int]) -> float:
     return max(values) / mean if mean else 1.0
 
 
+def _reorder(value: Any, order: list[int], batch_size: int) -> Any:
+    # Mirrors what slice_dp slices: lists and tensors along dim 0, and dicts of them, key by key.
+    import torch
+
+    if isinstance(value, dict):
+        return {key: _reorder(item, order, batch_size) for key, item in value.items()}
+    if isinstance(value, list) and len(value) == batch_size:
+        return [value[index] for index in order]
+    if isinstance(value, torch.Tensor) and value.dim() > 0 and value.shape[0] == batch_size:
+        return value[order]
+    return value
+
+
+def _sliced_with_other_size(value: Any, batch_size: int) -> bool:
+    import torch
+
+    if isinstance(value, dict):
+        return any(_sliced_with_other_size(item, batch_size) for item in value.values())
+    if isinstance(value, list):
+        return len(value) != batch_size
+    if isinstance(value, torch.Tensor) and value.dim() > 0:
+        return value.shape[0] != batch_size
+    return False
+
+
+def batch_aligned(args: tuple, kwargs: dict[str, Any], batch_size: int) -> bool:
+    """Whether every argument ``slice_dp`` slices has one entry per sample, so a reorder keeps them aligned."""
+    return not any(_sliced_with_other_size(value, batch_size) for value in (*args, *kwargs.values()))
+
+
+def call_lengths(args: tuple, kwargs: dict[str, Any]) -> list[int] | None:
+    """Sequence lengths of a ``slice_dp`` call's ``inputs``, or None when the call cannot be reordered."""
+    inputs = kwargs.get("inputs")
+    if kwargs.get("return_logits"):
+        # Raw logits of packed inputs are laid out per rank, so they cannot be put back in order.
+        return None
+    if not isinstance(inputs, list) or not inputs or not all(
+        isinstance(model_input, dict) and "input_ids" in model_input for model_input in inputs
+    ):
+        # Trajectories, packed rows and group rows: their lengths are only known after encoding.
+        return None
+    if not batch_aligned(args, kwargs, len(inputs)):
+        # slice_dp also slices this list on its own, so its entries cannot be matched to samples.
+        return None
+    return [sequence_length(model_input) for model_input in inputs]
+
+
+def plan_call(lengths: list[int], dp_size: int, spec: dict[str, Any], kwargs: dict[str, Any]) -> DPBalancePlan | None:
+    """Plan a call from the model's ``_dp_balance_spec``, or None when the call keeps its order untouched."""
+    workload_fn = estimate_workload
+    if spec["hidden_size"] is not None:
+        workload_fn = functools.partial(estimate_workload, hidden_size=spec["hidden_size"])
+    try:
+        plan = plan_dp_balance_for_lengths(
+            lengths,
+            dp_size,
+            padding_free=spec["padding_free"],
+            micro_batch_config=MicroBatchConfig.from_kwargs(dict(kwargs)),
+            group_size=spec["group_size"],
+            workload_fn=workload_fn,
+        )
+    except ValueError:
+        # Options or batch shapes the model itself rejects; let it report them.
+        return None
+    return plan if plan.metrics else None
+
+
+# Local mode: every rank holds only its own slice of the batch, so the ranks swap samples instead of a
+# driver reordering the whole batch. A "leaf" is one value slice_dp would slice (a list, or a tensor along
+# dim 0), addressed by its path through the call's args and kwargs.
+
+
+def _leaves(value: Any, path: tuple) -> dict[tuple, Any]:
+    import torch
+
+    if isinstance(value, dict):
+        leaves = {}
+        for key, item in value.items():
+            leaves.update(_leaves(item, (*path, key)))
+        return leaves
+    if isinstance(value, list) or (isinstance(value, torch.Tensor) and value.dim() > 0):
+        return {path: value}
+    return {}
+
+
+def call_leaves(args: tuple, kwargs: dict[str, Any]) -> dict[tuple, Any]:
+    """The values of a call that ``slice_dp`` would slice, keyed by their path."""
+    leaves = {}
+    for index, arg in enumerate(args):
+        leaves.update(_leaves(arg, ("args", index)))
+    for key, value in kwargs.items():
+        leaves.update(_leaves(value, ("kwargs", key)))
+    return leaves
+
+
+def leaf_layout(leaves: dict[tuple, Any]) -> dict[tuple, Any]:
+    """What the ranks must agree on before their leaves can be joined into one batch."""
+    return {path: "list" if isinstance(value, list) else (value.dim(), str(value.dtype)) for path, value in leaves.items()}
+
+
+def to_cpu(value: Any) -> Any:
+    """``value`` with every tensor detached onto the CPU, for ``all_gather_object``."""
+    import torch
+
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, dict):
+        moved = copy.copy(value)
+        for key, item in value.items():
+            moved[key] = to_cpu(item)
+        return moved
+    if isinstance(value, (list, tuple)):
+        return type(value)(to_cpu(item) for item in value)
+    return value
+
+
+def _device(value: Any):
+    import torch
+
+    if isinstance(value, torch.Tensor):
+        return value.device if value.device.type != "cpu" else None
+    items = value.values() if isinstance(value, dict) else value if isinstance(value, (list, tuple)) else ()
+    for item in items:
+        device = _device(item)
+        if device is not None:
+            return device
+    return None
+
+
+def _to(value: Any, device) -> Any:
+    import torch
+
+    if device is None:
+        return value
+    if isinstance(value, torch.Tensor):
+        return value.to(device)
+    if isinstance(value, dict):
+        return {key: _to(item, device) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return type(value)(_to(item, device) for item in value)
+    return value
+
+
+def _take(value: Any, indices: list[int]) -> Any:
+    if isinstance(value, list):
+        return [value[index] for index in indices]
+    return value[indices]
+
+
+def _with_leaves(value: Any, path: tuple, leaves: dict[tuple, Any]) -> Any:
+    if path in leaves:
+        return leaves[path]
+    if isinstance(value, dict):
+        return {key: _with_leaves(item, (*path, key), leaves) for key, item in value.items()}
+    return value
+
+
+def exchange_call(
+    plan: DPBalancePlan, rank_leaves: list[dict[tuple, Any]], data_rank: int, args: tuple, kwargs: dict[str, Any]
+) -> tuple[tuple, dict[str, Any]]:
+    """This rank's share of the reordered batch, built from every data rank's leaves (in data-rank order).
+
+    Tensors are joined the way ``collect_tensor_dict`` joins them for the Ray driver, padding the
+    trailing dims to the longest rank, and land on the device the local value was on.
+    """
+    from twinkle.utils import pad_and_stack_tensors
+
+    size = len(plan.order) // len(rank_leaves)
+    mine = plan.order[data_rank * size : (data_rank + 1) * size]
+    swapped = {}
+    for path, local in call_leaves(args, kwargs).items():
+        parts = [leaves[path] for leaves in rank_leaves]
+        joined = [item for part in parts for item in part] if isinstance(local, list) else pad_and_stack_tensors(parts)
+        swapped[path] = _to(_take(joined, mine), _device(local))
+    return (
+        tuple(_with_leaves(arg, ("args", index), swapped) for index, arg in enumerate(args)),
+        {key: _with_leaves(value, ("kwargs", key), swapped) for key, value in kwargs.items()},
+    )
+
+
+def _rows(value: Any) -> int | None:
+    import torch
+
+    if isinstance(value, list):
+        return len(value)
+    if isinstance(value, torch.Tensor) and value.dim() > 0:
+        return value.shape[0]
+    return None
+
+
+def _restore_rank(value: Any, rank_values: list[Any], sources: list[tuple[int, int]]) -> Any:
+    from twinkle.utils import pad_and_stack_tensors
+
+    if isinstance(value, dict):
+        restored = copy.copy(value)
+        for key, item in value.items():
+            restored[key] = _restore_rank(
+                item, [other.get(key) if isinstance(other, dict) else None for other in rank_values], sources
+            )
+        return restored
+    size = len(sources)
+    if _rows(value) != size or any(_rows(other) != size or type(other) is not type(value) for other in rank_values):
+        # Not one entry per sample (a loss, a count): it stays this rank's own value.
+        return value
+    if isinstance(value, list):
+        return _to([rank_values[rank][row] for rank, row in sources], _device(value))
+    return pad_and_stack_tensors([rank_values[rank][row : row + 1] for rank, row in sources]).to(value.device)
+
+
+def restore_rank_outputs(output: Any, rank_outputs: list[Any], plan: DPBalancePlan, data_rank: int) -> Any:
+    """Put this rank's own samples' per-sample outputs back, fetched from the ranks that computed them.
+
+    ``rank_outputs`` holds every data rank's output, in data-rank order. Values that are not per-sample
+    keep this rank's own value.
+    """
+    size = len(plan.order) // len(rank_outputs)
+    position = [0] * len(plan.order)
+    for index, sample in enumerate(plan.order):
+        position[sample] = index
+    sources = [divmod(position[sample], size) for sample in range(data_rank * size, (data_rank + 1) * size)]
+    return _restore_rank(output, rank_outputs, sources)
+
 @dataclass(frozen=True)
 class DPBalancePlan:
     """A batch permutation for ``slice_dp`` dispatch plus before/after balance metrics."""
@@ -143,6 +367,25 @@ def apply(self, values: Any) -> Any:
             return values[self.order]
         raise TypeError(f"cannot reorder a field of type {type(values).__name__}")
 
+    @property
+    def reorders(self) -> bool:
+        return self.order != list(range(len(self.order)))
+
+    def reorder_call(self, args: tuple, kwargs: dict[str, Any]) -> tuple[tuple, dict[str, Any]]:
+        """Reorder every batch-aligned argument of a ``slice_dp`` call, the way ``slice_dp`` would slice it."""
+        size = len(self.order)
+        return (
+            tuple(_reorder(arg, self.order, size) for arg in args),
+            {key: _reorder(value, self.order, size) for key, value in kwargs.items()},
+        )
+
+    def restore_outputs(self, outputs: Any) -> Any:
+        """Put the batch-aligned fields of a collected result back in the caller's order."""
+        inverse = [0] * len(self.order)
+        for position, index in enumerate(self.order):
+            inverse[index] = position
+        return _reorder(outputs, inverse, len(self.order))
+
 
 def plan_dp_balance_for_lengths(
     lengths: Sequence[int],
diff --git a/src/twinkle/model/optimizer_group.py b/src/twinkle/model/optimizer_group.py
index 566bc03ca..8d6a939f6 100644
--- a/src/twinkle/model/optimizer_group.py
+++ b/src/twinkle/model/optimizer_group.py
@@ -53,6 +53,9 @@ class BaseOptimizerGroup:
     pending_backward: bool = False
     pending_failed: bool = False
     pending_loss: float = 0.0
+    # Set by forward_only when the driver reordered its batch for DP balance.
+    dp_balance_reordered: bool = False
+    dp_balance_metrics: Dict[str, List[float]] = field(default_factory=dict)
 
     def set_loss(self, loss: Loss, loss_key: Optional[str] = None) -> None:
         if loss_key is None:
@@ -161,6 +164,10 @@ def accumulate_metrics(self, is_training, include_default: bool = True, metric_s
                 continue
             metric.accumulate(status.inputs, status.outputs, **metric_kwargs)
 
+    def record_dp_balance(self, metrics: Dict[str, float]) -> None:
+        for key, value in metrics.items():
+            self.dp_balance_metrics.setdefault(key, []).append(float(value))
+
     def calculate_metrics(self, is_training):
         """Calculate and return metrics."""
         self.accumulate_metrics(is_training)
@@ -183,6 +190,9 @@ def calculate_metrics(self, is_training):
         results = {}
         for metric in status.metrics:
             results.update(metric.calculate())
+        if is_training and self.dp_balance_metrics:
+            results.update({key: sum(values) / len(values) for key, values in self.dp_balance_metrics.items()})
+            self.dp_balance_metrics.clear()
         status.inputs = None
         status.outputs = None
         return results
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index b7d6adfc0..83f64a8ed 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -35,6 +35,7 @@
 from twinkle.hub import HubOperation
 from twinkle.infra import collect_dp_and_flatten, collect_tensor_dict
 from twinkle.loss import CrossEntropyLoss, Loss
+from twinkle.loss.base import dp_balance_group_size
 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
@@ -345,6 +346,7 @@ def __init__(
             init_strategy: Optional[str] = None,
             use_liger_kernel: bool = False,
             liger_kernel_config: Optional[Dict[str, Any]] = None,
+            dp_balance: bool = True,
             **kwargs):
         os.environ['TOKENIZERS_PARALLELISM'] = 'true'
         # Opt-out of the cuDNN SDPA backend (falls back to flash/mem-efficient, numerically
@@ -357,6 +359,8 @@ def __init__(
         # The Default tokenizer will be used to save with a model if no template was set.
         self._default_tokenizer = None
         self.device_mesh = device_mesh
+        # Read through _dp_balance_spec: reorder slice_dp batches to even out DP workloads.
+        self._dp_balance_enabled = dp_balance
         self.mixed_precision = mixed_precision
         self._fsdp_config = dict(fsdp_config or {})
         self._ddp_config = ddp_config or {}
@@ -904,6 +908,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         task = kwargs.pop('task', 'causal_lm')
         loss_key = kwargs.pop('loss_key', None)
         optimizer_config = self.optimizer_group[adapter_name]
+        optimizer_config.dp_balance_reordered = False
         self._lazy_wrap_model()
         if not inputs:
             raise ValueError('inputs empty, check your DataLoader outputs')
@@ -1010,7 +1015,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
     # lazy_collect follows the global default (True): a caller that reads the output itself calls
     # ``.result()`` first, while a ref-only pipeline hands the handle straight to a downstream remote
     # call, which resolves and collects it where the data is consumed (see ``LazyCollect``).
-    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict, dp_balance=True)
     @timed
     def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajectory]], **kwargs):
         """Call forward function without grad and record the inputs and outputs.
@@ -1030,6 +1035,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
         config = MicroBatchConfig.from_kwargs(kwargs)
         adapter_name = kwargs.pop('adapter_name', self._get_default_group())
         optimizer_config = self.optimizer_group[adapter_name]
+        balance = kwargs.pop('_dp_balance', None)
+        optimizer_config.dp_balance_reordered = bool(balance and balance['reordered'])
         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)
@@ -1333,6 +1340,11 @@ def calculate_loss(self,
         """Calculate one selected loss, optionally pre-normalized for a multi-RPC transaction."""
         adapter_name = kwargs.pop('adapter_name', self._get_default_group())
         optimizer_config = self.optimizer_group[adapter_name]
+        if optimizer_config.dp_balance_reordered and any(isinstance(value, list) for value in kwargs.values()):
+            raise ValueError(
+                'forward_only reordered this batch across DP ranks to balance their workload, so per-sample '
+                'lists passed to calculate_loss would no longer line up with it. Pass them to forward_only, or '
+                'build the model with dp_balance=False.')
         loss_instance: Loss = optimizer_config.get_loss(loss_key)
         assert isinstance(loss_instance, Loss), 'Set a loss_instance before calculating loss'
         status = optimizer_config.train_status if self.model.training else optimizer_config.eval_status
@@ -1599,7 +1611,7 @@ def _forward_backward_single(
             **kwargs,
         )
 
-    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict, dp_balance=True)
     @timed
     def forward_backward(
         self,
@@ -1613,6 +1625,10 @@ def forward_backward(
         """Run one homogeneous loss RPC and optionally finalize a two-RPC backward transaction."""
         adapter_name = kwargs.get('adapter_name', self._get_default_group())
         optimizer_config = self.optimizer_group[adapter_name]
+        balance = kwargs.pop('_dp_balance', None)
+        optimizer_config.dp_balance_reordered = False
+        if balance is not None:
+            optimizer_config.record_dp_balance(balance['metrics'])
         completing_pending = optimizer_config.prepare_backward_transaction(finalize_gradients, loss_weight)
         transactional = not finalize_gradients or completing_pending
         requested_sync = kwargs.pop('sync_gradients', None)
@@ -2482,6 +2498,23 @@ def calculate_metric(self, is_training, **kwargs):
         optimizer_config.require_finalized('calculate metrics')
         return optimizer_config.calculate_metrics(is_training)
 
+    @remote_function(execute='first', collect='first', lazy_collect=False)
+    def _dp_balance_spec(self, adapter_name: Optional[str] = None, loss_key: Optional[str] = None) -> Dict[str, Any]:
+        """What the driver needs to plan a DP-balanced order for ``forward_backward`` / ``forward_only``."""
+        optimizer_config = self.optimizer_group[adapter_name or self._get_default_group()]
+        try:
+            group_size = dp_balance_group_size(optimizer_config.get_loss(loss_key))
+        except ValueError:
+            # The call itself will fail on this loss_key; nothing tells which samples belong together.
+            group_size = None
+        text_config = self.hf_config.get_text_config() if hasattr(self.hf_config, 'get_text_config') else self.hf_config
+        return {
+            'enabled': getattr(self, '_dp_balance_enabled', False),
+            'padding_free': bool(getattr(optimizer_config.processor, 'padding_free', False)),
+            'hidden_size': getattr(text_config, 'hidden_size', None),
+            'group_size': group_size,
+        }
+
     def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs):
         assert adapter_name, 'Use a different adapter_name, current is empty.'
         freeze_llm = kwargs.pop('freeze_llm', False)
diff --git a/tests/model/test_dp_balance.py b/tests/model/test_dp_balance.py
index 1540a8237..ca666bbf9 100644
--- a/tests/model/test_dp_balance.py
+++ b/tests/model/test_dp_balance.py
@@ -1,7 +1,13 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+from types import SimpleNamespace
+
 import pytest
 import torch
 
+from twinkle import remote_function
+from twinkle.infra.collectors import collect_tensor_dict
+from twinkle.model.base import TrainableModel
+
 from twinkle.model.dp_balance import (
     DPBalancePlan,
     _estimate_ranks,
@@ -151,3 +157,179 @@ def test_plan_reads_encoded_inputs_and_reorders_aligned_fields():
     assert plan.apply(torch.arange(4)).tolist() == plan.order
     with pytest.raises(ValueError, match="cannot reorder"):
         plan.apply([1, 2])
+
+
+def test_reorder_call_moves_every_batch_aligned_argument_together():
+    plan = DPBalancePlan(order=[2, 0, 3, 1], strategy="kk", metrics={})
+    inputs = [{"input_ids": [index] * (index + 1)} for index in range(4)]
+    old_logps = [[float(index)] for index in range(4)]
+
+    args, kwargs = plan.reorder_call(
+        (inputs,),
+        {
+            "old_logps": old_logps,
+            "advantages": torch.arange(4.0),
+            "extra": {"ref_logps": torch.arange(8).view(4, 2)},
+            "micro_batch_size": 2,
+            "short": [1, 2],
+        },
+    )
+
+    assert args[0] == [inputs[index] for index in plan.order]
+    assert kwargs["old_logps"] == [old_logps[index] for index in plan.order]
+    assert kwargs["advantages"].tolist() == [2.0, 0.0, 3.0, 1.0]
+    assert kwargs["extra"]["ref_logps"][:, 0].tolist() == [4, 0, 6, 2]
+    assert kwargs["micro_batch_size"] == 2
+    assert kwargs["short"] == [1, 2]
+
+
+def test_restore_outputs_undoes_the_reorder():
+    plan = DPBalancePlan(order=[2, 0, 3, 1], strategy="kk", metrics={})
+    outputs = {"logps": torch.arange(12).view(4, 3), "loss": 1.5, "sequences": list("abcd")}
+
+    _, shuffled = plan.reorder_call((), {"outputs": outputs})
+    restored = plan.restore_outputs(shuffled["outputs"])
+
+    assert torch.equal(restored["logps"], outputs["logps"])
+    assert restored["sequences"] == outputs["sequences"]
+    assert restored["loss"] == 1.5
+
+
+class _FakeModel:
+    _dp_balance_call = TrainableModel._dp_balance_call
+
+    def __init__(self, enabled=True, group_size=1, data_world_size=2):
+        self.device_mesh = SimpleNamespace(data_world_size=data_world_size)
+        self.spec = {"enabled": enabled, "padding_free": True, "hidden_size": 1024, "group_size": group_size}
+        self.spec_calls = 0
+
+    def _dp_balance_spec(self, adapter_name=None, loss_key=None):
+        self.spec_calls += 1
+        return self.spec
+
+
+def _call_inputs(lengths):
+    return [{"input_ids": torch.ones(length, dtype=torch.long)} for length in lengths]
+
+
+def test_model_reorders_the_call_and_tells_the_workers():
+    inputs = _call_inputs([900, 800, 700, 600, 50, 40, 30, 20])
+    advantages = list(range(8))
+
+    plan, args, kwargs = _FakeModel()._dp_balance_call(
+        "forward_backward", (), {"inputs": inputs, "advantages": advantages}
+    )
+
+    assert plan is not None and plan.reorders
+    assert [len(item["input_ids"]) for item in kwargs["inputs"]] == [len(inputs[i]["input_ids"]) for i in plan.order]
+    assert kwargs["advantages"] == plan.order
+    assert kwargs["_dp_balance"]["reordered"] is True
+    assert "dp_balance/workload_imbalance_after" in kwargs["_dp_balance"]["metrics"]
+
+
+def test_model_keeps_preference_pairs_together():
+    inputs = _call_inputs([900, 900, 800, 800, 50, 50, 30, 30])
+
+    plan, _, _ = _FakeModel(group_size=2)._dp_balance_call("forward_backward", (), {"inputs": inputs})
+
+    pairs = [plan.order[index : index + 2] for index in range(0, 8, 2)]
+    assert all(first % 2 == 0 and second == first + 1 for first, second in pairs)
+
+
+@pytest.mark.parametrize(
+    "model, extra",
+    [
+        (_FakeModel(enabled=False), {}),
+        (_FakeModel(group_size=None), {}),
+        (_FakeModel(data_world_size=1), {}),
+        (_FakeModel(), {"return_logits": True}),
+        (_FakeModel(), {"group_scores": [0.5, 0.25]}),
+    ],
+)
+def test_model_leaves_the_call_alone(model, extra):
+    kwargs = {"inputs": _call_inputs([900, 800, 700, 600, 50, 40, 30, 20]), **extra}
+
+    assert model._dp_balance_call("forward_backward", (), kwargs) is None
+
+
+def test_disabled_model_is_asked_once():
+    model = _FakeModel(enabled=False)
+    kwargs = {"inputs": _call_inputs([900, 800, 700, 600, 50, 40, 30, 20])}
+
+    assert model._dp_balance_call("forward_backward", (), kwargs) is None
+    assert model._dp_balance_call("forward_backward", (), kwargs) is None
+    assert model.spec_calls == 1
+
+
+def test_group_size_follows_the_loss():
+    from twinkle.loss import CrossEntropyLoss, DPOLoss
+    from twinkle.loss.base import dp_balance_group_size
+
+    assert dp_balance_group_size(None) == 1
+    assert dp_balance_group_size(CrossEntropyLoss()) == 1
+    assert dp_balance_group_size(DPOLoss()) == 2
+    assert dp_balance_group_size(lambda inputs, outputs: 0) is None
+
+
+class _LocalModel:
+    _dp_balance_local = TrainableModel._dp_balance_local
+
+    def __init__(self, mesh):
+        self.device_mesh = mesh
+        self.seen = None
+
+    def _dp_balance_spec(self, adapter_name=None, loss_key=None):
+        return {"enabled": True, "padding_free": True, "hidden_size": 1024, "group_size": 1}
+
+    @remote_function(dispatch="slice_dp", collect=collect_tensor_dict, dp_balance=True)
+    def forward_backward(self, *, inputs, advantages, _dp_balance=None):
+        lengths = [len(model_input["input_ids"]) for model_input in inputs]
+        self.seen = (lengths, advantages.tolist(), _dp_balance)
+        return {
+            "loss": torch.tensor(float(sum(lengths))),
+            "lengths": torch.tensor(lengths),
+            "logps": [torch.full((length,), float(length)) for length in lengths],
+        }
+
+
+def _local_worker(rank, rendezvous, result_dir):
+    import os
+
+    import torch.distributed as dist
+
+    from twinkle import infra
+    from twinkle.utils import DeviceMesh
+
+    os.environ.update(RANK=str(rank), WORLD_SIZE="2")
+    infra._mode = "local"
+    dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2)
+    try:
+        lengths = [900, 800, 700, 600, 50, 40, 30, 20][rank * 4 : (rank + 1) * 4]
+        model = _LocalModel(DeviceMesh.from_sizes(world_size=2, dp_size=2, device_type="cpu"))
+        output = model.forward_backward(
+            inputs=_call_inputs(lengths), advantages=torch.tensor(lengths, dtype=torch.float)
+        )
+        seen_lengths, seen_advantages, balance = model.seen
+        torch.save(
+            {
+                "seen": seen_lengths,
+                "advantages_aligned": seen_advantages == [float(length) for length in seen_lengths],
+                "reordered": balance["reordered"],
+                "restored": output["lengths"].tolist() == lengths
+                and [len(logps) for logps in output["logps"]] == lengths,
+                "own_loss": float(output["loss"]) == float(sum(seen_lengths)),
+            },
+            f"{result_dir}/{rank}.pt",
+        )
+    finally:
+        dist.destroy_process_group()
+
+
+def test_local_mode_swaps_samples_and_returns_outputs_to_their_rank(tmp_path):
+    torch.multiprocessing.spawn(_local_worker, args=(str(tmp_path / "gloo"), str(tmp_path)), nprocs=2, join=True)
+    results = [torch.load(tmp_path / f"{rank}.pt") for rank in range(2)]
+
+    assert all(result["reordered"] and result["advantages_aligned"] for result in results)
+    assert all(result["restored"] and result["own_loss"] for result in results)
+    assert sorted(results[0]["seen"] + results[1]["seen"]) == [20, 30, 40, 50, 600, 700, 800, 900]
+    assert max(sum(result["seen"]) for result in results) < 900 + 800 + 700 + 600
diff --git a/tests/model/test_micro_batch.py b/tests/model/test_micro_batch.py
index 65d85ef84..e37fb3699 100644
--- a/tests/model/test_micro_batch.py
+++ b/tests/model/test_micro_batch.py
@@ -442,7 +442,7 @@ def _forward_only_single(self, *, inputs, sampling_masks, **kwargs):
 
 def test_forward_only_without_micro_batches_prepares_inputs_once():
     class ModelHarness:
-        optimizer_group = {'adapter': object()}
+        optimizer_group = {'adapter': SimpleNamespace()}
 
         def _get_default_group(self):
             return 'adapter'