diff --git a/cookbook/rl/grpo/dp_balance_grpo.py b/cookbook/rl/grpo/dp_balance_grpo.py
new file mode 100644
index 000000000..d629fcd05
--- /dev/null
+++ b/cookbook/rl/grpo/dp_balance_grpo.py
@@ -0,0 +1,249 @@
+"""GRPO on GSM8K with DP workload balancing (`--balance-batch`).
+
+`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.
+
+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
+flash-linear-attention >= 0.5.2, which ships the Ascend Triton backend; the script installs its kernels on
+the model workers.
+
+Usage:
+    sh dp_balance_grpo.sh                         # padded inputs, dynamic batching, FSDP2
+    sh dp_balance_grpo.sh --padding-free          # packed inputs
+    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
+from transformers import AutoConfig
+
+import twinkle
+from twinkle import DeviceGroup, DeviceMesh, Platform, get_device_placement, get_logger
+from twinkle.advantage import GRPOAdvantage
+from twinkle.checkpoint_engine import CheckpointEngineManager
+from twinkle.data_format import SamplingParams
+from twinkle.dataloader import DataLoader
+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.patch import Patch
+from twinkle.preprocessor.llm import GSM8KProcessor
+from twinkle.processor import InputProcessor
+from twinkle.reward import GSM8KAccuracyReward
+from twinkle.sampler import vLLMSampler
+
+logger = get_logger()
+
+SYSTEM_PROMPT = 'You are a helpful math assistant. Reason step by step and put your final answer within \\boxed{}.'
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument('--model-id', default='ms://Qwen/Qwen3.5-4B')
+    parser.add_argument('--dataset-id', default='ms://modelscope/gsm8k')
+    parser.add_argument('--subset-name', default='main')
+    parser.add_argument('--strategy', choices=('native_fsdp', 'accelerate'), default='native_fsdp')
+    parser.add_argument('--model-gpus', type=int, default=2)
+    parser.add_argument('--sampler-gpus', type=int, default=2)
+    parser.add_argument('--num-generations', type=int, default=8)
+    parser.add_argument('--max-tokens', type=int, default=2048, help='max new tokens per completion')
+    parser.add_argument('--batch-size', type=int, default=8, help='prompts per rollout batch')
+    parser.add_argument('--mini-batch-size', type=int, default=32, help='samples per optimizer step')
+    parser.add_argument('--micro-batch-size', type=int, default=16, help='row limit per micro-batch and rank')
+    parser.add_argument('--dynamic-batching', action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument('--max-tokens-per-micro-batch', type=int, default=8192)
+    parser.add_argument('--padding-free', action='store_true', help='packed inputs instead of padded ones')
+    parser.add_argument('--balance-batch', action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument('--max-steps', type=int, default=200)
+    parser.add_argument('--lr', type=float, default=1e-5)
+    parser.add_argument('--lora-r', type=int, default=16)
+    parser.add_argument('--attn-implementation', default='flash_attention_2')
+    return parser.parse_args()
+
+
+class NPULinearAttentionKernel(Patch):
+    """Install the flash-linear-attention GatedDeltaNet kernels on NPU.
+
+    Transformers reports them as available only on CUDA-like devices; the twinkle 'fla' kernel op installs
+    them on NPU. Under Ray the model lives on the workers, so this runs there through `apply_patch`.
+    """
+
+    def __call__(self, module, *args, **kwargs):
+        from twinkle.kernel import kernelize
+        from twinkle.kernel.config import DEFAULT_KERNEL_CONFIG
+        kernelize(module, {'fla': DEFAULT_KERNEL_CONFIG['fla']})
+
+
+def main():
+    args = parse_args()
+    device_type = Platform.get_platform().__name__
+    config_dir = HubOperation.download_model(args.model_id, ignore_model=True)
+    hf_config = AutoConfig.from_pretrained(config_dir, trust_remote_code=True)
+    is_qwen35 = hf_config.model_type.startswith('qwen3_5')
+    template = 'Qwen3_5Template' if is_qwen35 else 'Template'
+    num_gpus = args.model_gpus + args.sampler_gpus
+
+    device_groups = [
+        DeviceGroup(name='model', ranks=list(range(args.model_gpus)), device_type=device_type),
+        DeviceGroup(name='sampler', ranks=list(range(args.model_gpus, num_gpus)), device_type=device_type),
+    ]
+    if args.strategy == 'native_fsdp':
+        model_mesh = DeviceMesh.from_sizes(world_size=args.model_gpus, fsdp_size=args.model_gpus)
+    else:
+        model_mesh = DeviceMesh.from_sizes(world_size=args.model_gpus, dp_size=args.model_gpus)
+    sampler_mesh = DeviceMesh.from_sizes(world_size=args.sampler_gpus, dp_size=args.sampler_gpus)
+    twinkle.initialize(mode='ray', nproc_per_node=num_gpus, groups=device_groups, lazy_collect=False)
+    # 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)
+    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.
+    lora_config = LoraConfig(
+        target_modules=[
+            'q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj', 'in_proj_qkv', 'in_proj_z',
+            'in_proj_a', 'in_proj_b', 'out_proj'
+        ],
+        r=args.lora_r,
+        lora_alpha=args.lora_r * 2,
+        lora_dropout=0.05)
+    model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=1)
+    model.set_optimizer('AdamW', lr=args.lr)
+    model.set_lr_scheduler('CosineAnnealingLR', T_max=args.max_steps, eta_min=0)
+    model.set_loss('GRPOLoss', epsilon=0.2)
+    model.set_processor(InputProcessor, padding_free=args.padding_free)
+    model.set_template(template, model_id=args.model_id, enable_thinking=True)
+
+    sampler = vLLMSampler(
+        model_id=args.model_id,
+        engine_args={
+            'gpu_memory_utilization': 0.8,
+            'max_model_len': 8192,
+            'max_lora_rank': max(16, args.lora_r),
+            'enable_lora': True,
+        },
+        device_mesh=sampler_mesh,
+        remote_group='sampler',
+    )
+    sampler.set_template(template, model_id=args.model_id, enable_thinking=True)
+    ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler, platform=device_type)
+
+    def create_gsm8k_dataset():
+        dataset = Dataset()
+        dataset.add_dataset(DatasetMeta(args.dataset_id, subset_name=args.subset_name, split='train'))
+        # Thinking on and no length penalty: completion lengths differ a lot between prompts.
+        dataset.set_template(
+            template, model_id=args.model_id, max_length=4096, truncation_strategy='delete', enable_thinking=True)
+        dataset.map(GSM8KProcessor(system=SYSTEM_PROMPT))
+        dataset.encode(add_generation_prompt=True)
+        return dataset
+
+    dataloader = DataLoader(
+        dataset=create_gsm8k_dataset,
+        batch_size=args.batch_size,
+        min_batch_size=args.batch_size,
+        device_mesh=model_mesh,
+        remote_group='model',
+    )
+    accuracy_reward = GSM8KAccuracyReward()
+    advantage_fn = GRPOAdvantage()
+    metrics = CompletionRewardMetric()
+    sampling_params = SamplingParams(
+        max_tokens=args.max_tokens, num_samples=1, logprobs=1, temperature=1.0, top_p=0.95)
+
+    logger.info(get_device_placement())
+    logger.info(f'strategy={args.strategy}, padding_free={args.padding_free}, balance_batch={args.balance_batch}, '
+                f'model DP={dp_size}, micro-batch {micro_batch_kwargs}')
+    optim_step = 0
+    for batch in dataloader:
+        if optim_step >= args.max_steps:
+            break
+        metrics.reset()
+        expand_prompts = [prompt for prompt in batch for _ in range(args.num_generations)]
+        ckpt_manager.sync_weights(merge_and_sync=False)
+        sampler.reset_prefix_cache()
+        sample_responses = sampler.sample(expand_prompts, sampling_params)
+
+        all_inputs: List[Dict[str, Any]] = []
+        trajectories: List[Dict[str, Any]] = []
+        all_old_logps: List[List[float]] = []
+        completion_lengths: List[int] = []
+        for sample_response in sample_responses:
+            for sequence in sample_response.sequences:
+                all_inputs.append(sequence.new_input_feature)
+                # new_input_feature carries tokens only; the conversation, which rewards score, is in new_state.
+                trajectories.append({**sequence.new_input_feature, 'messages': list(sequence.new_state.messages)})
+                all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs])
+                completion_lengths.append(len(sequence.tokens))
+
+        rewards = accuracy_reward(trajectories)
+        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()
+            optim_step += 1
+            if optim_step >= args.max_steps:
+                break
+
+        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}')
+    model.save('dp-balance-grpo-final')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/cookbook/rl/grpo/dp_balance_grpo.sh b/cookbook/rl/grpo/dp_balance_grpo.sh
new file mode 100644
index 000000000..b71fbc50f
--- /dev/null
+++ b/cookbook/rl/grpo/dp_balance_grpo.sh
@@ -0,0 +1,28 @@
+#!/bin/sh
+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
+# 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
+#   sh dp_balance_grpo.sh --no-balance-batch        # baseline without balancing
+#   sh dp_balance_grpo.sh --model-id /path/to/Qwen3.5-4B --dataset-id /path/to/gsm8k/train.parquet
+
+python dp_balance_grpo.py \
+    --model-id ms://Qwen/Qwen3.5-4B \
+    --strategy native_fsdp \
+    --model-gpus 2 \
+    --sampler-gpus 2 \
+    --num-generations 8 \
+    --max-tokens 2048 \
+    --batch-size 8 \
+    --mini-batch-size 32 \
+    --micro-batch-size 16 \
+    --max-tokens-per-micro-batch 8192 \
+    --max-steps 200 \
+    --lr 1e-5 \
+    --lora-r 16 \
+    --balance-batch \
+    "$@"
diff --git a/src/twinkle/model/__init__.py b/src/twinkle/model/__init__.py
index a95d5a493..027afe7ef 100644
--- a/src/twinkle/model/__init__.py
+++ b/src/twinkle/model/__init__.py
@@ -5,6 +5,7 @@
 
 if TYPE_CHECKING:
     from .base import ModelLoaderProtocol, TrainableModel
+    from .dp_balance import DPBalancePlan, plan_dp_balance
     from .megatron import MegatronModel, MultiLoraMegatronModel
     from .micro_batch import MicroBatchConfig
     from .state_dict_utils import (MODULES_TO_SAVE_SEGMENT, collect_local_state, is_lora_state_key,
@@ -16,6 +17,7 @@
 else:
     _import_structure = {
         'base': ['ModelLoaderProtocol', 'TrainableModel'],
+        'dp_balance': ['DPBalancePlan', 'plan_dp_balance'],
         'micro_batch': ['MicroBatchConfig'],
         'state_dict_utils': [
             'MODULES_TO_SAVE_SEGMENT', 'collect_local_state', 'is_lora_state_key', 'is_saveable_adapter_key',
diff --git a/src/twinkle/model/dp_balance.py b/src/twinkle/model/dp_balance.py
new file mode 100644
index 000000000..43ba263ba
--- /dev/null
+++ b/src/twinkle/model/dp_balance.py
@@ -0,0 +1,249 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Data-parallel workload balancing for ``slice_dp`` dispatch.
+
+``slice_dp`` hands every DP rank one contiguous ``DeviceMesh.get_slice`` chunk, and the step waits for the
+slowest rank. ``plan_dp_balance`` permutes the batch so the chunks cost about the same, without changing how
+many samples each rank receives.
+
+Only batches that split evenly across the ranks are permuted. Losses such as ``GRPOLoss`` average over the
+samples of one rank and DP then averages the ranks, so a sample weighs ``1 / (dp_size * rank_size)``; when
+``get_slice`` gives some ranks one sample more, moving samples between ranks would change their weights.
+
+A rank's cost is estimated by running the same micro-batch planner the model runs on its chunk, so the
+estimate follows both the micro-batch layout and the input layout:
+
+- packed (``padding_free=True``): every sequence costs ``workload_fn(length)``;
+- padded: each micro-batch pads its rows to its longest sequence, so it costs ``rows * workload_fn(longest)``.
+
+With dynamic batching the ranks are re-planned with a common micro-batch count, as the model does before
+it starts the micro-batch loop.
+"""
+
+from __future__ import annotations
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+from .micro_batch import MicroBatchConfig, _batch_cost, _kk_partition, plan_micro_batches_for_lengths, sequence_length
+
+WorkloadFn = Callable[[int], int]
+
+
+def estimate_workload(length: int, hidden_size: int = 4096) -> int:
+    """Relative forward cost of one dense-transformer sequence.
+
+    Projections and MLP cost about ``12·h²·s`` FLOPs and attention ``2·h·s²``; dividing by ``2h`` keeps the
+    ratio and gives ``6·h·s + s²``.
+    """
+    return 6 * hidden_size * length + length * length
+
+
+def micro_batch_workload(lengths: Sequence[int], *, padding_free: bool, workload_fn: WorkloadFn = estimate_workload):
+    """Estimated cost of one micro-batch with the given sequence lengths."""
+    if not lengths:
+        return 0
+    if padding_free:
+        return sum(workload_fn(length) for length in lengths)
+    return len(lengths) * workload_fn(max(lengths))
+
+
+@dataclass(frozen=True)
+class _RankEstimate:
+    workloads: list[int]
+    tokens: list[int]
+    computed_tokens: list[int]
+    micro_batches: int
+
+
+def _estimate_ranks(
+    rank_lengths: list[list[int]],
+    config: MicroBatchConfig | None,
+    padding_free: bool,
+    workload_fn: WorkloadFn,
+) -> _RankEstimate:
+    if config is None:
+        plans = [[list(range(len(lengths)))] for lengths in rank_lengths]
+    else:
+        min_micro_batches = 1
+        while True:
+            plans = [
+                plan_micro_batches_for_lengths(
+                    lengths, config, padding_free=padding_free, min_micro_batches=min_micro_batches
+                )
+                for lengths in rank_lengths
+            ]
+            counts = [len(plan) for plan in plans]
+            if not config.dynamic_batching or len(set(counts)) == 1:
+                break
+            min_micro_batches = max(counts)
+
+    workloads, computed_tokens = [], []
+    for lengths, plan in zip(rank_lengths, plans, strict=True):
+        batches = [[lengths[index] for index in indices] for indices in plan]
+        workloads.append(
+            sum(micro_batch_workload(batch, padding_free=padding_free, workload_fn=workload_fn) for batch in batches)
+        )
+        computed_tokens.append(sum(_batch_cost(indices, lengths, padding_free) for indices in plan))
+    return _RankEstimate(
+        workloads=workloads,
+        tokens=[sum(lengths) for lengths in rank_lengths],
+        computed_tokens=computed_tokens,
+        micro_batches=max(len(plan) for plan in plans),
+    )
+
+
+def _contiguous(group_count: int, dp_size: int) -> list[list[int]]:
+    size = group_count // dp_size
+    return [list(range(rank * size, (rank + 1) * size)) for rank in range(dp_size)]
+
+
+def _karmarkar_karp(group_workloads: list[int], dp_size: int) -> list[list[int]]:
+    return [sorted(part) for part in _kk_partition(group_workloads, dp_size, equal_size=True)]
+
+
+def _snake(group_keys: list[tuple[int, int]], dp_size: int) -> list[list[int]]:
+    """Deal groups from the heaviest down, reversing direction every round.
+
+    Every rank then sees the same spread of lengths, so the micro-batch planner pads and splits all ranks
+    alike, which is what padded micro-batches need.
+    """
+    ordered = sorted(range(len(group_keys)), key=lambda group: (group_keys[group], -group), reverse=True)
+    ranks: list[list[int]] = [[] for _ in range(dp_size)]
+    for start in range(0, len(ordered), dp_size):
+        chunk = ordered[start : start + dp_size]
+        targets = range(dp_size) if (start // dp_size) % 2 == 0 else range(dp_size - 1, -1, -1)
+        for rank, group in zip(targets, chunk, strict=True):
+            ranks[rank].append(group)
+    return [sorted(groups) for groups in ranks]
+
+
+def _imbalance(values: Sequence[int]) -> float:
+    # max/mean estimates how much slower the step is than a perfectly even split.
+    mean = sum(values) / len(values)
+    return max(values) / mean if mean else 1.0
+
+
+@dataclass(frozen=True)
+class DPBalancePlan:
+    """A batch permutation for ``slice_dp`` dispatch plus before/after balance metrics."""
+
+    order: list[int]
+    strategy: str = "identity"
+    metrics: dict[str, float] = field(default_factory=dict)
+
+    def apply(self, values: Any) -> Any:
+        """Reorder one per-sample field (list, tuple, tensor or array) with the plan."""
+        if len(values) != len(self.order):
+            raise ValueError(f"cannot reorder {len(values)} values with a plan for {len(self.order)} samples")
+        if isinstance(values, list):
+            return [values[index] for index in self.order]
+        if isinstance(values, tuple):
+            return tuple(values[index] for index in self.order)
+        if hasattr(values, "shape"):
+            return values[self.order]
+        raise TypeError(f"cannot reorder a field of type {type(values).__name__}")
+
+
+def plan_dp_balance_for_lengths(
+    lengths: Sequence[int],
+    dp_size: int,
+    *,
+    padding_free: bool,
+    micro_batch_config: MicroBatchConfig | None = None,
+    group_size: int = 1,
+    workload_fn: WorkloadFn = estimate_workload,
+) -> DPBalancePlan:
+    """Same as :func:`plan_dp_balance`, for samples given by their sequence lengths."""
+    lengths = [int(length) for length in lengths]
+    total = len(lengths)
+    if dp_size <= 0:
+        raise ValueError(f"dp_size must be positive, got {dp_size}")
+    if group_size <= 0 or total % group_size:
+        raise ValueError(f"batch size {total} must be a positive multiple of group_size={group_size}")
+    group_count = total // group_size
+    identity = list(range(total))
+    if dp_size == 1 or group_count < dp_size:
+        return DPBalancePlan(order=identity)
+    if total % dp_size:
+        # Uneven rank sizes give samples rank-dependent loss weights; see the module docstring.
+        return DPBalancePlan(order=identity)
+    if group_size > 1 and group_count % dp_size:
+        raise ValueError(
+            f"keeping groups of {group_size} together needs the group count ({group_count}) "
+            f"to be divisible by dp_size={dp_size}"
+        )
+    sample_workloads = [workload_fn(length) for length in lengths]
+    if min(sample_workloads) < 0:
+        raise ValueError(f"workload_fn must return non-negative values, got {min(sample_workloads)}")
+
+    def members(group: int) -> range:
+        return range(group * group_size, (group + 1) * group_size)
+
+    group_workloads = [sum(sample_workloads[index] for index in members(group)) for group in range(group_count)]
+    group_longest = [max(lengths[index] for index in members(group)) for group in range(group_count)]
+    candidates = {
+        "identity": _contiguous(group_count, dp_size),
+        "karmarkar_karp": _karmarkar_karp(group_workloads, dp_size),
+        "snake": _snake(list(zip(group_longest, group_workloads, strict=True)), dp_size),
+    }
+
+    estimates: dict[str, tuple[list[int], _RankEstimate]] = {}
+    for name, rank_groups in candidates.items():
+        order = [index for groups in rank_groups for group in groups for index in members(group)]
+        rank_size = total // dp_size
+        rank_lengths = [
+            [lengths[index] for index in order[rank * rank_size : (rank + 1) * rank_size]] for rank in range(dp_size)
+        ]
+        try:
+            estimates[name] = order, _estimate_ranks(rank_lengths, micro_batch_config, padding_free, workload_fn)
+        except ValueError:
+            # The model rejects this layout too (for example a sequence over max_tokens_per_micro_batch).
+            continue
+    if not estimates:
+        return DPBalancePlan(order=identity)
+
+    # Ties keep the earlier candidate, so the batch is only permuted when the slowest rank gets lighter.
+    strategy = min(estimates, key=lambda name: (max(estimates[name][1].workloads), estimates[name][1].micro_batches))
+    order, after = estimates[strategy]
+    metrics = {"dp_balance/rank_tokens_mean": sum(lengths) / dp_size}
+    phases = [("after", after)]
+    if "identity" in estimates:
+        phases.insert(0, ("before", estimates["identity"][1]))
+    for phase, estimate in phases:
+        metrics[f"dp_balance/rank_tokens_max_{phase}"] = max(estimate.tokens)
+        metrics[f"dp_balance/rank_tokens_min_{phase}"] = min(estimate.tokens)
+        metrics[f"dp_balance/workload_imbalance_{phase}"] = _imbalance(estimate.workloads)
+        metrics[f"dp_balance/micro_batches_{phase}"] = estimate.micro_batches
+        if not padding_free:
+            metrics[f"dp_balance/padding_ratio_{phase}"] = 1 - sum(estimate.tokens) / sum(estimate.computed_tokens)
+    return DPBalancePlan(order=order, strategy=strategy, metrics=metrics)
+
+
+def plan_dp_balance(
+    inputs: Sequence[Any],
+    dp_size: int,
+    *,
+    padding_free: bool,
+    micro_batch_config: MicroBatchConfig | None = None,
+    group_size: int = 1,
+    workload_fn: WorkloadFn = estimate_workload,
+) -> DPBalancePlan:
+    """Plan a DP-balanced order for encoded ``inputs`` that will be dispatched with ``slice_dp``.
+
+    Pass the processor's ``padding_free`` and the micro-batch options given to ``forward_backward``
+    (``None`` when the model runs each rank's chunk as one batch). Apply the plan to ``inputs`` and every
+    aligned per-sample field (old logps, advantages, ...) before calling the model; results returned by the
+    model come back in the permuted order.
+
+    Every rank keeps the sample count ``get_slice`` gives it, and ``group_size`` consecutive samples stay on
+    one rank (for example chosen/rejected pairs). Within a rank, samples keep their original relative order.
+    A batch whose size is not a multiple of ``dp_size`` keeps its order (see the module docstring).
+    """
+    return plan_dp_balance_for_lengths(
+        [sequence_length(model_input) for model_input in inputs],
+        dp_size,
+        padding_free=padding_free,
+        micro_batch_config=micro_batch_config,
+        group_size=group_size,
+        workload_fn=workload_fn,
+    )
diff --git a/src/twinkle/model/micro_batch.py b/src/twinkle/model/micro_batch.py
index 93cc621f9..fa3942a7e 100644
--- a/src/twinkle/model/micro_batch.py
+++ b/src/twinkle/model/micro_batch.py
@@ -126,10 +126,18 @@ def __lt__(self, other: _KKState) -> bool:
         return self.spread > other.spread
 
 
-def _kk_partition(lengths: list[int], group_count: int) -> list[list[int]]:
+def _kk_partition(lengths: list[int], group_count: int, *, equal_size: bool = False) -> list[list[int]]:
+    ordered = sorted((value, index) for index, value in enumerate(lengths))
+    if equal_size:
+        if len(ordered) % group_count:
+            raise ValueError(f"equal-size KK needs len(lengths) % group_count == 0, got {len(ordered)} % {group_count}")
+        # Seeding each state with one item per group keeps every final group the same size.
+        seeds = [ordered[start : start + group_count] for start in range(0, len(ordered), group_count)]
+    else:
+        seeds = [[item] for item in ordered]
     queue = []
-    for value, index in sorted((value, index) for index, value in enumerate(lengths)):
-        heapq.heappush(queue, _KKState([(index, value)], group_count))
+    for seed in seeds:
+        heapq.heappush(queue, _KKState([(index, value) for value, index in seed], group_count))
     while len(queue) > 1:
         first = heapq.heappop(queue)
         second = heapq.heappop(queue)
@@ -169,13 +177,31 @@ def plan_micro_batches(
     if not inputs:
         raise ValueError("cannot plan micro-batches for empty inputs")
     if not config.dynamic_batching:
-        return [
-            list(range(start, min(start + config.micro_batch_size, len(inputs))))
-            for start in range(0, len(inputs), config.micro_batch_size)
-        ]
-    if min_micro_batches <= 0 or min_micro_batches > len(inputs):
-        raise ValueError(f"invalid min_micro_batches={min_micro_batches} for {len(inputs)} inputs")
+        return _fixed_plan(len(inputs), config.micro_batch_size)
     lengths = [sequence_length(model_input) for model_input in inputs]
+    return plan_micro_batches_for_lengths(
+        lengths, config, padding_free=padding_free, min_micro_batches=min_micro_batches
+    )
+
+
+def _fixed_plan(count: int, micro_batch_size: int) -> list[list[int]]:
+    return [list(range(start, min(start + micro_batch_size, count))) for start in range(0, count, micro_batch_size)]
+
+
+def plan_micro_batches_for_lengths(
+    lengths: list[int],
+    config: MicroBatchConfig,
+    *,
+    padding_free: bool,
+    min_micro_batches: int = 1,
+) -> list[list[int]]:
+    """Same as :func:`plan_micro_batches`, for samples given by their sequence lengths."""
+    if not lengths:
+        raise ValueError("cannot plan micro-batches for empty inputs")
+    if not config.dynamic_batching:
+        return _fixed_plan(len(lengths), config.micro_batch_size)
+    if min_micro_batches <= 0 or min_micro_batches > len(lengths):
+        raise ValueError(f"invalid min_micro_batches={min_micro_batches} for {len(lengths)} inputs")
     capacity = config.max_tokens_per_micro_batch
     oversized = [length for length in lengths if length > capacity]
     if oversized:
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 3b527678d..b7d6adfc0 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -975,7 +975,9 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         with _resolve_task_context(unwrapped_model, task, template=optimizer_config.template), \
                 _resolve_mm_graph_context(unwrapped_model, optimizer_config.template):
             with self.sp_strategy.context() if self.sp_strategy is not None else contextlib.nullcontext():
-                outputs = self.model(**inputs)
+                # Without gradient checkpointing use_cache stays True, and Transformers skips
+                # detecting padding-free packed sequences whenever a KV cache exists.
+                outputs = self.model(**inputs, use_cache=False)
 
         recorded_routing = rr.finish()
 
@@ -1147,7 +1149,9 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
             # unwrapped on purpose: task patches probe head/module attributes DDP does not forward.
             with _resolve_task_context(unwrapped_model, task, template=optimizer_config.template), lora_ctx:
                 with self.sp_strategy.context() if self.sp_strategy is not None else contextlib.nullcontext():
-                    outputs = self.model(**inputs)
+                    # In eval mode use_cache defaults to True, and Transformers skips detecting
+                    # padding-free packed sequences whenever a KV cache exists.
+                    outputs = self.model(**inputs, use_cache=False)
 
             recorded_routing = rr.finish()
 
diff --git a/src/twinkle/template/contract.py b/src/twinkle/template/contract.py
index cb22d6eb1..cceaafd94 100644
--- a/src/twinkle/template/contract.py
+++ b/src/twinkle/template/contract.py
@@ -1,6 +1,8 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 """Twinkle-side view of the shared template contract owned by ``swift.dev.template_contract``."""
 
+import numpy as np
+
 try:
     from swift.dev.template_contract import (
         SHIFTED_FIELDS,
@@ -38,7 +40,6 @@
         TemplateProtocol,
         TemplateSpec,
         Tool,
-        append_and_finalize,
         append_input_order,
         coerce_tool_call_arguments,
         feature_has_media,
@@ -54,8 +55,30 @@
         unshift,
         unshift_values,
     )
+    from swift.dev.template_contract import append_and_finalize as _swift_append_and_finalize
 except ImportError as error:
     raise ImportError(
         "twinkle requires a matching swift.dev.template_contract; install the pinned Swift revision before importing "
         "twinkle.template"
     ) from error
+
+
+def _batched_token_fields_as_arrays(row):
+    # swift extends numpy/torch token fields on their last axis but extends a Python list flat. A prompt that went
+    # through a dataset stores the processor's batched [1, L] mm_token_type_ids as a nested list, which a flat extend
+    # turns into the ragged [[0] * L, 0, 0, ...]; as an array it is extended to [1, L + n].
+    converted = None
+    for key in TOKEN_FILL:
+        value = row.get(key)
+        if isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], (list, tuple)):
+            if converted is None:
+                converted = dict(row)
+            converted[key] = np.asarray(value)
+    return row if converted is None else converted
+
+
+def append_and_finalize(row, tokens, *, scored, sampled, finalize):
+    """``swift.dev.template_contract.append_and_finalize``, extending batched list token fields on their last axis."""
+    return _swift_append_and_finalize(
+        _batched_token_fields_as_arrays(row), tokens, scored=scored, sampled=sampled, finalize=finalize
+    )
diff --git a/tests/model/test_dp_balance.py b/tests/model/test_dp_balance.py
new file mode 100644
index 000000000..1540a8237
--- /dev/null
+++ b/tests/model/test_dp_balance.py
@@ -0,0 +1,153 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import pytest
+import torch
+
+from twinkle.model.dp_balance import (
+    DPBalancePlan,
+    _estimate_ranks,
+    estimate_workload,
+    micro_batch_workload,
+    plan_dp_balance,
+    plan_dp_balance_for_lengths,
+)
+from twinkle.model.micro_batch import MicroBatchConfig, _kk_partition
+
+
+def _rank_chunks(values, dp_size):
+    base, remainder = divmod(len(values), dp_size)
+    chunks, start = [], 0
+    for rank in range(dp_size):
+        size = base + int(rank < remainder)
+        chunks.append(values[start : start + size])
+        start += size
+    return chunks
+
+
+def _peak(lengths, order, dp_size, config, padding_free):
+    rank_lengths = _rank_chunks([lengths[index] for index in order], dp_size)
+    return max(_estimate_ranks(rank_lengths, config, padding_free, estimate_workload).workloads)
+
+
+def test_equal_size_kk_partition_keeps_group_sizes():
+    groups = _kk_partition([9, 1, 8, 2, 7, 3], 2, equal_size=True)
+
+    assert sorted(len(group) for group in groups) == [3, 3]
+    assert sorted(index for group in groups for index in group) == list(range(6))
+    with pytest.raises(ValueError, match="equal-size"):
+        _kk_partition([1, 2, 3], 2, equal_size=True)
+
+
+def test_padded_micro_batch_costs_its_longest_row():
+    assert micro_batch_workload([4, 10], padding_free=False, workload_fn=lambda n: n) == 20
+    assert micro_batch_workload([4, 10], padding_free=True, workload_fn=lambda n: n) == 14
+
+
+@pytest.mark.parametrize(
+    "padding_free, config",
+    [
+        (True, None),
+        (False, MicroBatchConfig(micro_batch_size=2)),
+        (True, MicroBatchConfig(micro_batch_size=2)),
+        (False, MicroBatchConfig(micro_batch_size=4, dynamic_batching=True, max_tokens_per_micro_batch=4096)),
+        (True, MicroBatchConfig(micro_batch_size=4, dynamic_batching=True, max_tokens_per_micro_batch=4096)),
+    ],
+)
+def test_plan_keeps_rank_sizes_and_lowers_the_slowest_rank(padding_free, config):
+    lengths = [2000, 1900, 1800, 1700, 300, 250, 200, 150]
+
+    plan = plan_dp_balance_for_lengths(lengths, 2, padding_free=padding_free, micro_batch_config=config)
+
+    assert sorted(plan.order) == list(range(len(lengths)))
+    assert [len(chunk) for chunk in _rank_chunks(plan.order, 2)] == [4, 4]
+    assert plan.strategy != "identity"
+    assert _peak(lengths, plan.order, 2, config, padding_free) < _peak(lengths, range(8), 2, config, padding_free)
+    assert plan.metrics["dp_balance/workload_imbalance_after"] < plan.metrics["dp_balance/workload_imbalance_before"]
+    assert ("dp_balance/padding_ratio_after" in plan.metrics) == (not padding_free)
+
+
+def test_padded_single_batch_per_rank_cannot_beat_its_longest_sequence():
+    # Every rank pads its whole chunk to its longest row, so the rank holding 2000 costs 4 * w(2000) anyway.
+    plan = plan_dp_balance_for_lengths([2000, 1900, 1800, 1700, 300, 250, 200, 150], 2, padding_free=False)
+
+    assert plan.strategy == "identity"
+    assert plan.order == list(range(8))
+
+
+def test_padded_dynamic_batching_prefers_ranks_with_similar_lengths():
+    # Real-token balance can pair a long sequence with short ones, which padded micro-batches pay for.
+    lengths = [1000, 1000, 520, 480, 500, 500, 10, 10]
+    config = MicroBatchConfig(micro_batch_size=4, dynamic_batching=True, max_tokens_per_micro_batch=2000)
+
+    plan = plan_dp_balance_for_lengths(lengths, 2, padding_free=False, micro_batch_config=config)
+
+    peaks = {
+        name: _peak(lengths, order, 2, config, False)
+        for name, order in {"chosen": plan.order, "identity": list(range(8))}.items()
+    }
+    assert peaks["chosen"] < peaks["identity"]
+    assert plan.metrics["dp_balance/padding_ratio_after"] <= plan.metrics["dp_balance/padding_ratio_before"]
+
+
+def test_dynamic_batching_estimate_uses_the_model_common_micro_batch_count():
+    # On its own rank 1 fits in one padded micro-batch costing 3 * 30; split into rank 0's three, it costs 50.
+    config = MicroBatchConfig(micro_batch_size=4, dynamic_batching=True, max_tokens_per_micro_batch=100)
+    rank_lengths = [[90, 80, 70], [30, 10, 10]]
+
+    estimate = _estimate_ranks(rank_lengths, config, False, lambda n: n)
+
+    assert estimate.micro_batches == 3
+    assert estimate.workloads == [240, 50]
+
+
+def test_balanced_batch_keeps_its_order():
+    plan = plan_dp_balance_for_lengths([5, 1, 1, 5], 2, padding_free=True)
+
+    assert plan.order == [0, 1, 2, 3]
+    assert plan.strategy == "identity"
+
+
+def test_groups_stay_on_one_rank():
+    lengths = [900, 900, 100, 100, 800, 800, 200, 200]
+
+    plan = plan_dp_balance_for_lengths(lengths, 2, padding_free=True, group_size=2)
+
+    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)
+
+
+def test_uneven_batch_keeps_its_order():
+    # get_slice gives rank 0 three samples and rank 1 two, so a sample's loss weight depends on its rank.
+    plan = plan_dp_balance_for_lengths([900, 900, 900, 100, 100], 2, padding_free=True)
+
+    assert plan.order == [0, 1, 2, 3, 4]
+    assert plan.strategy == "identity"
+    assert plan.metrics == {}
+
+
+def test_uneven_grouped_batch_keeps_its_order():
+    plan = plan_dp_balance_for_lengths([900, 900, 100, 100] * 2 + [500, 500], 4, padding_free=True, group_size=2)
+
+    assert plan.order == list(range(10))
+    assert plan.metrics == {}
+
+
+def test_unplannable_batch_is_left_for_the_model_to_reject():
+    config = MicroBatchConfig(micro_batch_size=2, dynamic_batching=True, max_tokens_per_micro_batch=50)
+
+    plan = plan_dp_balance_for_lengths([100, 10, 10, 10], 2, padding_free=True, micro_batch_config=config)
+
+    assert plan.order == [0, 1, 2, 3]
+    assert plan.metrics == {}
+
+
+def test_plan_reads_encoded_inputs_and_reorders_aligned_fields():
+    inputs = [{"input_ids": torch.ones(length, dtype=torch.long)} for length in (800, 700, 50, 40)]
+
+    plan = plan_dp_balance(inputs, 2, padding_free=True)
+
+    assert isinstance(plan, DPBalancePlan)
+    assert plan.apply([0, 1, 2, 3]) == plan.order
+    assert plan.apply(tuple("abcd")) == tuple("abcd"[index] for index in plan.order)
+    assert plan.apply(torch.arange(4)).tolist() == plan.order
+    with pytest.raises(ValueError, match="cannot reorder"):
+        plan.apply([1, 2])
diff --git a/tests/template/test_contract_append.py b/tests/template/test_contract_append.py
new file mode 100644
index 000000000..5a213bad2
--- /dev/null
+++ b/tests/template/test_contract_append.py
@@ -0,0 +1,33 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+
+from twinkle.template.contract import append_and_finalize
+
+
+def _prompt(mm_token_type_ids):
+    return {"input_ids": [11, 12, 13], "labels": [-100, -100, -100], "mm_token_type_ids": mm_token_type_ids}
+
+
+def _append(prompt, tokens):
+    return append_and_finalize(prompt, tokens, scored=True, sampled=True, finalize=lambda rows: rows)
+
+
+def test_batched_list_token_type_ids_extend_on_last_axis():
+    # A dataset-encoded prompt stores the processor's [1, L] mm_token_type_ids as a nested list.
+    result = _append(_prompt([[0, 0, 0]]), [21, 22])
+
+    assert np.asarray(result["mm_token_type_ids"]).tolist() == [[0, 0, 0, 0, 0]]
+    assert list(result["input_ids"]) == [11, 12, 13, 21, 22]
+
+
+def test_flat_list_token_type_ids_stay_flat():
+    result = _append(_prompt([0, 0, 0]), [21, 22])
+
+    assert list(result["mm_token_type_ids"]) == [0, 0, 0, 0, 0]
+
+
+def test_prompt_feature_is_left_untouched():
+    prompt = _prompt([[0, 0, 0]])
+    _append(prompt, [21])
+
+    assert prompt["mm_token_type_ids"] == [[0, 0, 0]]
diff --git a/tests/transformers/test_padding_free_forward_only.py b/tests/transformers/test_padding_free_forward_only.py
new file mode 100644
index 000000000..4401aee0b
--- /dev/null
+++ b/tests/transformers/test_padding_free_forward_only.py
@@ -0,0 +1,64 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import tempfile
+from unittest.mock import patch
+
+import torch
+from transformers import AutoModelForCausalLM, Qwen3Config
+
+from twinkle.model import TransformersModel
+from twinkle.processor import InputProcessor
+
+
+def _tiny_qwen3_dir():
+    path = tempfile.mkdtemp()
+    config = Qwen3Config(
+        vocab_size=64,
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        head_dim=8,
+        max_position_embeddings=64,
+    )
+    torch.manual_seed(0)
+    AutoModelForCausalLM.from_config(config).save_pretrained(path)
+    return path
+
+
+def _sample(length, generator):
+    input_ids = torch.randint(0, 64, (length,), generator=generator)
+    return {
+        "input_ids": input_ids,
+        "labels": input_ids.clone(),
+        "attention_mask": torch.ones_like(input_ids),
+        "position_ids": torch.arange(length),
+    }
+
+
+def test_forward_only_keeps_padding_free_samples_apart():
+    """A sample packed after another must get the same log-probs as on its own."""
+    model = TransformersModel(model_id=_tiny_qwen3_dir(), mixed_precision="no")
+    model.set_processor(InputProcessor, padding_free=True)
+    model._lazy_wrap_model()
+    device = str(next(model.model.parameters()).device)
+    generator = torch.Generator().manual_seed(0)
+    first, second = _sample(12, generator), _sample(20, generator)
+    with patch("twinkle.processor.base.Platform.get_local_device", return_value=device):
+        packed = model.forward_only(inputs=[first, second])["logps"]
+        alone = model.forward_only(inputs=[second])["logps"]
+    torch.testing.assert_close(packed[1], alone[0])
+
+def test_forward_keeps_padding_free_samples_apart_without_checkpointing():
+    """Without gradient checkpointing the training forward must not build a KV cache either."""
+    model = TransformersModel(model_id=_tiny_qwen3_dir(), mixed_precision="no", gradient_checkpointing=False)
+    model.set_processor(InputProcessor, padding_free=True)
+    model._lazy_wrap_model()
+    model.model.train()
+    device = str(next(model.model.parameters()).device)
+    generator = torch.Generator().manual_seed(0)
+    first, second = _sample(12, generator), _sample(20, generator)
+    with patch("twinkle.processor.base.Platform.get_local_device", return_value=device):
+        packed = model.forward(inputs=[first, second])["logps"]
+        alone = model.forward(inputs=[second])["logps"]
+    torch.testing.assert_close(packed[1].detach(), alone[0].detach())