From c680155b4e63ee66d3e1a7f2f9ac478afd527e27 Mon Sep 17 00:00:00 2001 From: weixin_46040161 Date: Tue, 3 Nov 2026 14:31:29 +0800 Subject: [PATCH 1/3] feat(model): balance DP workloads inside TransformersModel by default forward_backward and forward_only now reorder each slice_dp batch on the driver (Ray mode) so every DP rank gets a similar workload, reorder every batch-aligned argument with it, and return the outputs in the caller's order. TransformersModel(dp_balance=False) turns it off. - remote_function(dp_balance=True) asks the model's _dp_balance_call for a plan; calls with object refs, raw logits, unencoded inputs or lists of another length are dispatched unchanged. - Loss.dp_balance_group_size keeps paired samples together (preference and reward losses: 2, REAL: num_generations); embedding and listwise reranker losses, and plain callables, opt out. - calculate_loss rejects per-sample lists after a reordered forward_only. - dp_balance/* metrics are reported by calculate_metric. - The GRPO cookbook no longer calls plan_dp_balance itself. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- cookbook/rl/grpo/dp_balance_grpo.py | 50 +++------ cookbook/rl/grpo/dp_balance_grpo.sh | 4 +- src/twinkle/infra/dispatch.py | 29 ++++- src/twinkle/loss/base.py | 12 ++ src/twinkle/loss/dpo.py | 5 + src/twinkle/loss/grpo.py | 5 + src/twinkle/loss/infonce.py | 2 + src/twinkle/loss/reranker.py | 2 + src/twinkle/loss/reward.py | 1 + src/twinkle/model/base.py | 50 +++++++++ src/twinkle/model/dp_balance.py | 49 ++++++++ src/twinkle/model/optimizer_group.py | 10 ++ .../model/transformers/transformers.py | 36 +++++- tests/model/test_dp_balance.py | 106 ++++++++++++++++++ tests/model/test_micro_batch.py | 2 +- 15 files changed, 323 insertions(+), 40 deletions(-) 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/src/twinkle/infra/dispatch.py b/src/twinkle/infra/dispatch.py index 185ae32ce..a1d3a0be8 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,11 @@ 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. Wrapped continuous-work calls also accept ``_session_ids``, a list or tuple of non-empty session strings (or ``None`` for ordinary load balancing), aligned @@ -995,6 +1010,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 +1052,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__': @@ -1141,6 +1165,7 @@ def _notify_collect_error(_e, _caller=_caller, _ctx=_ctx, _infra=infra): wrapper._control = control wrapper._replay = replay wrapper._worker_selector = worker_selector + wrapper._dp_balance = dp_balance if enable_continous_work: # Attached to the class by remote_class, and called instead of this # method when the driver routes requests worker by worker. 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..fab6c199f 100644 --- a/src/twinkle/model/base.py +++ b/src/twinkle/model/base.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import functools import os import re import shutil @@ -22,6 +23,8 @@ from torch.optim.lr_scheduler import LRScheduler from transformers import PreTrainedModel, PretrainedConfig + from .micro_batch import MicroBatchConfig + def copy_checkpoint_args(output_dir: str, checkpoint_dir: str) -> None: """Copy the run-level CLI metadata into a completed checkpoint on the master rank.""" @@ -347,6 +350,53 @@ 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_micro_batch_config(self, kwargs: Dict[str, Any]) -> Optional['MicroBatchConfig']: + """The micro-batch layout the workers will plan for a call with these keyword arguments.""" + from .micro_batch import MicroBatchConfig + return MicroBatchConfig.from_kwargs(dict(kwargs)) + + 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 batch_aligned, estimate_workload, plan_dp_balance + mesh = getattr(self, 'device_mesh', None) + inputs = kwargs.get('inputs') + if mesh is None or mesh.data_world_size <= 1 or 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 + spec = self._dp_balance_spec(adapter_name=kwargs.get('adapter_name'), loss_key=kwargs.get('loss_key')) + if not spec['enabled'] or spec['group_size'] is None: + return None + try: + plan = plan_dp_balance( + inputs, + mesh.data_world_size, + padding_free=spec['padding_free'], + micro_batch_config=self._dp_balance_micro_batch_config(kwargs), + group_size=spec['group_size'], + workload_fn=functools.partial(estimate_workload, hidden_size=spec['hidden_size']), + ) + except ValueError: + # Options or batch shapes the model itself rejects; let it report them. + return None + if not plan.metrics: + 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 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..dc24e75af 100644 --- a/src/twinkle/model/dp_balance.py +++ b/src/twinkle/model/dp_balance.py @@ -123,6 +123,36 @@ 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())) + + @dataclass(frozen=True) class DPBalancePlan: """A batch permutation for ``slice_dp`` dispatch plus before/after balance metrics.""" @@ -143,6 +173,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..a081edec1 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 by the driver 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,22 @@ 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: + loss = optimizer_config.get_loss(loss_key) + except ValueError: + loss = 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) or 4096, + 'group_size': dp_balance_group_size(loss), + } + 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..b42c10972 100644 --- a/tests/model/test_dp_balance.py +++ b/tests/model/test_dp_balance.py @@ -1,7 +1,11 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +from types import SimpleNamespace + import pytest import torch +from twinkle.model.base import TrainableModel + from twinkle.model.dp_balance import ( DPBalancePlan, _estimate_ranks, @@ -151,3 +155,105 @@ 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 + _dp_balance_micro_batch_config = TrainableModel._dp_balance_micro_batch_config + + 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} + + def _dp_balance_spec(self, adapter_name=None, loss_key=None): + 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_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 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' From 2dcfb40098a1d3084b2ec8aa06faaf3926fb0cab Mon Sep 17 00:00:00 2001 From: weixin_46040161 Date: Tue, 3 Nov 2026 16:29:04 +0800 Subject: [PATCH 2/3] fix(model): address DP balance review findings - Skip the blocking _dp_balance_spec round trip once a model reports dp_balance off; the flag is fixed at construction. - Do not reorder when the loss for the call cannot be resolved. - Leave the hidden-size default to estimate_workload. - Drop the unread wrapper._dp_balance and the one-line micro-batch config helper. - Document dp_balance in the TransformersModel docs (en/zh). Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../Components/Model/TransformersModel.md | 2 ++ .../TransformersModel.md" | 2 ++ src/twinkle/infra/dispatch.py | 1 - src/twinkle/model/base.py | 23 +++++++++++-------- .../model/transformers/transformers.py | 9 ++++---- tests/model/test_dp_balance.py | 12 +++++++++- 6 files changed, 33 insertions(+), 16 deletions(-) diff --git a/docs/source_en/Components/Model/TransformersModel.md b/docs/source_en/Components/Model/TransformersModel.md index 8c10e4b98..1d76cf92c 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. 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..2bb42122d 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`等）随之一起重排，返回的输出会还原为调用方的顺序。对按样本求均值的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 a1d3a0be8..aa828af45 100644 --- a/src/twinkle/infra/dispatch.py +++ b/src/twinkle/infra/dispatch.py @@ -1165,7 +1165,6 @@ def _notify_collect_error(_e, _caller=_caller, _ctx=_ctx, _infra=infra): wrapper._control = control wrapper._replay = replay wrapper._worker_selector = worker_selector - wrapper._dp_balance = dp_balance if enable_continous_work: # Attached to the class by remote_class, and called instead of this # method when the driver routes requests worker by worker. diff --git a/src/twinkle/model/base.py b/src/twinkle/model/base.py index fab6c199f..11aac4b9f 100644 --- a/src/twinkle/model/base.py +++ b/src/twinkle/model/base.py @@ -23,8 +23,6 @@ from torch.optim.lr_scheduler import LRScheduler from transformers import PreTrainedModel, PretrainedConfig - from .micro_batch import MicroBatchConfig - def copy_checkpoint_args(output_dir: str, checkpoint_dir: str) -> None: """Copy the run-level CLI metadata into a completed checkpoint on the master rank.""" @@ -350,11 +348,6 @@ 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_micro_batch_config(self, kwargs: Dict[str, Any]) -> Optional['MicroBatchConfig']: - """The micro-batch layout the workers will plan for a call with these keyword arguments.""" - from .micro_batch import MicroBatchConfig - return MicroBatchConfig.from_kwargs(dict(kwargs)) - def _dp_balance_call(self, method: str, args: tuple, kwargs: Dict[str, Any]): """Driver side of ``remote_function(dp_balance=True)``. @@ -363,6 +356,9 @@ def _dp_balance_call(self, method: str, args: tuple, kwargs: Dict[str, Any]): batch keeps its order, so the result needs no restoring. """ from .dp_balance import batch_aligned, estimate_workload, plan_dp_balance + from .micro_batch import MicroBatchConfig + if getattr(self, '_dp_balance_off', False): + return None mesh = getattr(self, 'device_mesh', None) inputs = kwargs.get('inputs') if mesh is None or mesh.data_world_size <= 1 or kwargs.get('return_logits'): @@ -376,16 +372,23 @@ def _dp_balance_call(self, method: str, args: tuple, kwargs: Dict[str, Any]): # slice_dp also slices this list on its own, so its entries cannot be matched to samples. return None spec = self._dp_balance_spec(adapter_name=kwargs.get('adapter_name'), loss_key=kwargs.get('loss_key')) - if not spec['enabled'] or spec['group_size'] is None: + 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 + 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( inputs, mesh.data_world_size, padding_free=spec['padding_free'], - micro_batch_config=self._dp_balance_micro_batch_config(kwargs), + micro_batch_config=MicroBatchConfig.from_kwargs(dict(kwargs)), group_size=spec['group_size'], - workload_fn=functools.partial(estimate_workload, hidden_size=spec['hidden_size']), + workload_fn=workload_fn, ) except ValueError: # Options or batch shapes the model itself rejects; let it report them. diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index a081edec1..35dd67ac2 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -2503,15 +2503,16 @@ def _dp_balance_spec(self, adapter_name: Optional[str] = None, loss_key: Optiona """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: - loss = optimizer_config.get_loss(loss_key) + group_size = dp_balance_group_size(optimizer_config.get_loss(loss_key)) except ValueError: - loss = None + # 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) or 4096, - 'group_size': dp_balance_group_size(loss), + '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): diff --git a/tests/model/test_dp_balance.py b/tests/model/test_dp_balance.py index b42c10972..8937a2a2a 100644 --- a/tests/model/test_dp_balance.py +++ b/tests/model/test_dp_balance.py @@ -195,13 +195,14 @@ def test_restore_outputs_undoes_the_reorder(): class _FakeModel: _dp_balance_call = TrainableModel._dp_balance_call - _dp_balance_micro_batch_config = TrainableModel._dp_balance_micro_batch_config 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 @@ -249,6 +250,15 @@ def test_model_leaves_the_call_alone(model, 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 From 415c1beffad08101918e71ea99515ac8b897db8b Mon Sep 17 00:00:00 2001 From: weixin_46040161 Date: Tue, 3 Nov 2026 17:15:38 +0800 Subject: [PATCH 3/3] feat(model): balance DP workloads in local (torchrun) mode Each rank holds only its own slice, so the ranks all-gather their sequence lengths, plan the joined batch as the Ray driver would, swap samples when the plan reorders, and send per-sample outputs back to their rank. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../Components/Model/TransformersModel.md | 2 +- .../TransformersModel.md" | 2 +- src/twinkle/infra/dispatch.py | 7 +- src/twinkle/model/base.py | 79 ++++--- src/twinkle/model/dp_balance.py | 194 ++++++++++++++++++ .../model/transformers/transformers.py | 2 +- tests/model/test_dp_balance.py | 66 ++++++ 7 files changed, 320 insertions(+), 32 deletions(-) diff --git a/docs/source_en/Components/Model/TransformersModel.md b/docs/source_en/Components/Model/TransformersModel.md index 1d76cf92c..9a6698eec 100644 --- a/docs/source_en/Components/Model/TransformersModel.md +++ b/docs/source_en/Components/Model/TransformersModel.md @@ -50,7 +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. 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. +- 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 2bb42122d..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" @@ -50,7 +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`等）随之一起重排，返回的输出会还原为调用方的顺序。对按样本求均值的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`则保持原始切分。 +- 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 aa828af45..9775d674a 100644 --- a/src/twinkle/infra/dispatch.py +++ b/src/twinkle/infra/dispatch.py @@ -834,7 +834,9 @@ def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp', 'last_pp ``_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. + 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 @@ -902,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) diff --git a/src/twinkle/model/base.py b/src/twinkle/model/base.py index 11aac4b9f..3f8a192a4 100644 --- a/src/twinkle/model/base.py +++ b/src/twinkle/model/base.py @@ -1,5 +1,4 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -import functools import os import re import shutil @@ -355,21 +354,14 @@ def _dp_balance_call(self, method: str, args: tuple, kwargs: Dict[str, Any]): 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 batch_aligned, estimate_workload, plan_dp_balance - from .micro_batch import MicroBatchConfig + from .dp_balance import call_lengths, plan_call if getattr(self, '_dp_balance_off', False): return None mesh = getattr(self, 'device_mesh', None) - inputs = kwargs.get('inputs') - if mesh is None or mesh.data_world_size <= 1 or kwargs.get('return_logits'): - # Raw logits of packed inputs are laid out per rank, so they cannot be put back in order. + if mesh is None or mesh.data_world_size <= 1: 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. + 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']: @@ -378,28 +370,59 @@ def _dp_balance_call(self, method: str, args: tuple, kwargs: Dict[str, Any]): return None if spec['group_size'] is None: return None - 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( - inputs, - mesh.data_world_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 - if not plan.metrics: + 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 dc24e75af..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 @@ -153,6 +155,198 @@ def batch_aligned(args: tuple, kwargs: dict[str, Any], batch_size: int) -> bool: 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.""" diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 35dd67ac2..83f64a8ed 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -359,7 +359,7 @@ 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 by the driver through _dp_balance_spec: reorder slice_dp batches to even out DP workloads. + # 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 {}) diff --git a/tests/model/test_dp_balance.py b/tests/model/test_dp_balance.py index 8937a2a2a..ca666bbf9 100644 --- a/tests/model/test_dp_balance.py +++ b/tests/model/test_dp_balance.py @@ -4,6 +4,8 @@ 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 ( @@ -267,3 +269,67 @@ def test_group_size_follows_the_loss(): 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