feat(model): balance DP workloads inside TransformersModel by default

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}')

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

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.
 

docs/source_zh/组件/模型/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`中。
 

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__':

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

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,

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',

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: