wehub-resource-sync
410 lines
16 KiB
Python
410 lines
16 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import importlib
|
|
import os
|
|
import ray
|
|
from contextlib import contextmanager
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from swift.utils import get_logger
|
|
from .base_trainer import BaseRayTrainer
|
|
from .driver_utils import (build_dataset_from_dict, compute_iter_params, estimate_dp_size, merge_group_dict,
|
|
parse_ray_yaml)
|
|
|
|
logger = get_logger()
|
|
|
|
_TRAINER_REGISTRY: Dict[str, Dict[str, Any]] = {
|
|
'grpo': {
|
|
'trainer': 'swift.ray.megatron.grpo_trainer.GRPOTrainer',
|
|
'loss': 'swift.ray.megatron.loss.grpo.GRPOLoss',
|
|
},
|
|
'gkd': {
|
|
'trainer': 'swift.ray.megatron.gkd_trainer.GKDTrainer',
|
|
'loss': 'swift.ray.megatron.loss.gkd.GKDLoss',
|
|
},
|
|
}
|
|
|
|
|
|
def register_ray_trainer(
|
|
rlhf_type: str,
|
|
trainer: str,
|
|
loss: Optional[str] = None,
|
|
):
|
|
"""Register a custom algorithm for the Ray pipeline.
|
|
|
|
Args:
|
|
rlhf_type: Algorithm identifier (e.g. ``'grpo'``).
|
|
trainer: Dotted path to the driver-side trainer class.
|
|
The class must accept ``(worker_groups, rollout_replicas)``
|
|
and expose ``set_data_info()`` / ``train()`` methods.
|
|
Example: ``'swift.ray.megatron.grpo_trainer.GRPOTrainer'``
|
|
loss: Dotted path to a ``Loss`` subclass that defines
|
|
``forward_step`` + ``loss_func``.
|
|
Pass ``None`` to use the internal trainer's forward_step.
|
|
"""
|
|
_TRAINER_REGISTRY[rlhf_type] = {'trainer': trainer, 'loss': loss}
|
|
|
|
|
|
class MegatronRayPipeline:
|
|
|
|
def __init__(self, config_path: str):
|
|
self.ray_config, group_configs, shared_config = parse_ray_yaml(config_path)
|
|
shared_config['use_ray'] = True
|
|
|
|
self.rlhf_type = self.ray_config.rlhf_type
|
|
if self.rlhf_type not in _TRAINER_REGISTRY:
|
|
raise ValueError('Unknown rlhf_type %r. Available: %s' % (self.rlhf_type, list(_TRAINER_REGISTRY)))
|
|
|
|
self.group_cfgs: Dict[str, Dict[str, Any]] = {
|
|
g: merge_group_dict(shared_config, gd)
|
|
for g, gd in group_configs.items()
|
|
}
|
|
self.shared_cfg = {k: v for k, v in shared_config.items() if v is not None}
|
|
|
|
self._entry = _TRAINER_REGISTRY[self.rlhf_type]
|
|
self.resource_pool_manager = None
|
|
self.worker_groups: Dict[str, Any] = {}
|
|
self.rollout_replicas: List[Any] = []
|
|
self.teacher_replicas: List[Any] = []
|
|
|
|
def init(self) -> None:
|
|
# Initialize Ray, create resource pools, spawn workers and replicas.
|
|
self._data_info = self._build_dataset()
|
|
self._compute_train_iters()
|
|
ray.init(ignore_reinit_error=True)
|
|
self._create_pools()
|
|
self._init_worker_groups()
|
|
with self._colocate_offload_ctx():
|
|
self._init_rollout_replicas()
|
|
self._init_teacher_replicas()
|
|
self._driver_trainer: BaseRayTrainer = self._create_trainer()
|
|
self._driver_trainer.set_data_info(self._data_info)
|
|
|
|
def train(self) -> Any:
|
|
"""Run the training loop. Requires ``init()`` to have been called."""
|
|
if not hasattr(self, '_driver_trainer'):
|
|
raise RuntimeError('MegatronRayPipeline.train(): call init() first')
|
|
return self._driver_trainer.train()
|
|
|
|
def run(self) -> Any:
|
|
"""Convenience: ``init()`` + ``train()`` + ``shutdown()``."""
|
|
self.init()
|
|
try:
|
|
return self.train()
|
|
finally:
|
|
self._shutdown()
|
|
|
|
def _build_dataset(self) -> Dict[str, Any]:
|
|
# Merge train group config (tuner_type, lora_rank, etc.) into shared cfg so that
|
|
# _check_teacher can detect LoRA self-distillation (teacher_model == model + lora).
|
|
cfg = {**self.shared_cfg, **self.group_cfgs.get('train', {})}
|
|
data_info = build_dataset_from_dict(cfg)
|
|
return data_info
|
|
|
|
def _compute_train_iters(self):
|
|
train_cfg = self.group_cfgs.get('train')
|
|
gpus = self.group_gpus.get('train', 0)
|
|
assert train_cfg is not None and gpus > 0
|
|
dp_size = estimate_dp_size(train_cfg, gpus)
|
|
iter_params = compute_iter_params(self._data_info, dp_size)
|
|
train_iters = iter_params.get('train_iters')
|
|
assert train_iters is not None and train_iters > 0
|
|
|
|
self.group_cfgs['train']['train_iters'] = train_iters
|
|
|
|
def _create_pools(self):
|
|
from .resource_pool import ResourcePool, ResourcePoolManager
|
|
|
|
colocated_sets = {frozenset(g) for g in self.colocate_groups}
|
|
pool_mapping: Dict[str, ResourcePool] = {}
|
|
assigned: set = set()
|
|
|
|
for colocated in colocated_sets:
|
|
# Colocated roles share one GPU set, so they must request the same number
|
|
# of gpus. Validate explicitly (and avoid relying on frozenset order).
|
|
gpus_by_role = {g: self.group_gpus.get(g, 0) for g in colocated}
|
|
distinct = set(gpus_by_role.values())
|
|
if distinct == {0}:
|
|
continue
|
|
if len(distinct) > 1:
|
|
raise ValueError(f'Colocated roles must request the same number of gpus, but got '
|
|
f'{gpus_by_role}. Set an equal `gpus` for all roles in {sorted(colocated)}.')
|
|
gpus = distinct.pop()
|
|
pon = self.ray_config.gpus_as_process_on_nodes(gpus)
|
|
shared = ResourcePool(pon, max_colocate_count=len(colocated))
|
|
for g in colocated:
|
|
pool_mapping[g] = shared
|
|
assigned.add(g)
|
|
|
|
for name, gpus in self.group_gpus.items():
|
|
if name in assigned or gpus <= 0:
|
|
continue
|
|
pon = self.ray_config.gpus_as_process_on_nodes(gpus)
|
|
pool_mapping[name] = ResourcePool(pon)
|
|
|
|
self.resource_pool_manager = ResourcePoolManager(pool_mapping)
|
|
self.resource_pool_manager.create_all()
|
|
|
|
def _is_rollout_hybrid(self) -> bool:
|
|
"""True if rollout shares its pool with train (HYBRID mode)."""
|
|
return any('rollout' in cg and 'train' in cg for cg in self.colocate_groups)
|
|
|
|
def _init_worker_groups(self):
|
|
self._validate_grpo_train_batch_params()
|
|
|
|
self._spawn_train_group('train')
|
|
|
|
train_wg = self.worker_groups['train']
|
|
padding_vals = train_wg.broadcast('get_padding_to')
|
|
self._data_info['_padding_to'] = next((v for v in padding_vals if v is not None), None)
|
|
|
|
def _validate_grpo_train_batch_params(self) -> None:
|
|
"""Early validation of GRPO batch params before spawning workers.
|
|
|
|
Resolves generation_batch_size / steps_per_generation, then validates
|
|
DP alignment — all via static helpers in RLHFMegatronArgumentsMixin.
|
|
"""
|
|
if self.rlhf_type != 'grpo':
|
|
return
|
|
train_cfg = self.group_cfgs.get('train')
|
|
train_gpus = self.group_gpus.get('train', 0)
|
|
if not train_cfg or train_gpus <= 0:
|
|
return
|
|
|
|
cfg = dict(train_cfg)
|
|
global_batch_size = cfg.get('global_batch_size')
|
|
if global_batch_size <= 0:
|
|
return
|
|
|
|
micro_batch_size = cfg.get('micro_batch_size', 1)
|
|
num_generations = cfg.get('num_generations', 8)
|
|
|
|
from swift.megatron.arguments.megatron_args import RLHFMegatronArgumentsMixin
|
|
generation_batch_size, _ = RLHFMegatronArgumentsMixin.resolve_generation_batch_size(
|
|
cfg.get('generation_batch_size'), cfg.get('steps_per_generation'), global_batch_size, num_generations)
|
|
|
|
dp_size = estimate_dp_size(cfg, train_gpus)
|
|
RLHFMegatronArgumentsMixin.validate_batch_dp_alignment(generation_batch_size, num_generations, dp_size,
|
|
micro_batch_size, train_gpus)
|
|
|
|
def _spawn_train_group(self, role: str) -> None:
|
|
from .megatron_worker import MegatronWorker
|
|
from .worker_group import WorkerGroup
|
|
|
|
pool = self.resource_pool_manager.get_pool(role)
|
|
cfg = dict(self.group_cfgs.get(role, {}))
|
|
cfg.setdefault('rlhf_type', self.ray_config.rlhf_type)
|
|
worker_cls = ray.remote(num_gpus=0)(MegatronWorker)
|
|
wg = WorkerGroup.from_pool(role, pool, worker_cls=worker_cls)
|
|
|
|
loss_cls = self._entry.get('loss')
|
|
rollout_config = self._build_rollout_config_for_workers() if self._is_rollout_hybrid() else None
|
|
wg.broadcast('init_actor', cfg, loss_cls_path=loss_cls, rollout_config=rollout_config)
|
|
wg.build_dispatch_info(worker_cls=MegatronWorker)
|
|
|
|
self.worker_groups[role] = wg
|
|
logger.info('MegatronWorker group [%s] on %d GPUs', role, pool.world_size)
|
|
|
|
@contextmanager
|
|
def _colocate_offload_ctx(self):
|
|
"""Offload train workers during vLLM init (colocate only)."""
|
|
need = self._is_rollout_hybrid() and bool(self.shared_cfg.get('offload_model', True))
|
|
colocated_wgs = [
|
|
wg for role, wg in self.worker_groups.items() if need and any(role in g for g in self.colocate_groups)
|
|
]
|
|
for wg in colocated_wgs:
|
|
wg.broadcast('offload_to_cpu')
|
|
try:
|
|
yield
|
|
finally:
|
|
for wg in colocated_wgs:
|
|
wg.broadcast('reload_to_gpu')
|
|
|
|
def _init_rollout_replicas(self) -> None:
|
|
rollout_gpus = self.group_gpus.get('rollout', 0)
|
|
if rollout_gpus <= 0:
|
|
self.rollout_replicas = []
|
|
return
|
|
|
|
from .rollout.replica import RolloutReplica
|
|
|
|
rollout_cfg = self._with_router_replay_rollout_config(self.group_cfgs.get('rollout', {}))
|
|
is_hybrid = self._is_rollout_hybrid()
|
|
pool = self.resource_pool_manager.get_pool('train' if is_hybrid else 'rollout')
|
|
|
|
template_kwargs = self._get_template_kwargs_for_rollout()
|
|
self.rollout_replicas = RolloutReplica.create_replicas(
|
|
rollout_cfg=rollout_cfg,
|
|
rollout_gpus=rollout_gpus,
|
|
pool=pool,
|
|
is_hybrid=is_hybrid,
|
|
sleep_level=self.ray_config.sleep_level,
|
|
template_kwargs=template_kwargs,
|
|
)
|
|
|
|
def _init_teacher_replicas(self) -> None:
|
|
teacher_gpus = self.group_gpus.get('teacher', 0)
|
|
if teacher_gpus <= 0:
|
|
self.teacher_replicas = []
|
|
return
|
|
|
|
from .rollout.replica import RolloutReplica
|
|
|
|
teacher_cfg = self.group_cfgs.get('teacher', {})
|
|
pool = self.resource_pool_manager.get_pool('teacher')
|
|
template_kwargs = self._get_template_kwargs_for_rollout()
|
|
args = self._data_info.get('_driver_args')
|
|
if args is not None:
|
|
base_len = template_kwargs.get('max_length') or getattr(args, 'max_length')
|
|
template_kwargs = dict(template_kwargs)
|
|
template_kwargs['max_length'] = base_len + args.max_completion_length
|
|
self.teacher_replicas = RolloutReplica.create_replicas(
|
|
rollout_cfg=teacher_cfg,
|
|
rollout_gpus=teacher_gpus,
|
|
pool=pool,
|
|
is_hybrid=False,
|
|
sleep_level=0,
|
|
template_kwargs=template_kwargs,
|
|
actor_name_prefix='swift_teacher_server',
|
|
)
|
|
|
|
def _with_router_replay_rollout_config(self, rollout_cfg: Dict[str, Any]) -> Dict[str, Any]:
|
|
cfg = dict(rollout_cfg or {})
|
|
args = self._data_info.get('_driver_args')
|
|
router_mode = getattr(args, 'router_replay_mode', 'disabled') if args is not None else 'disabled'
|
|
if router_mode != 'R3':
|
|
return cfg
|
|
|
|
from swift.rlhf_trainers.utils import check_vllm_version_ge
|
|
if not check_vllm_version_ge('0.14.0'):
|
|
raise ValueError('router_replay_mode=R3 requires vLLM>=0.14.0 to return routed_experts.')
|
|
|
|
engine_kwargs = dict(cfg.get('vllm_engine_kwargs') or {})
|
|
engine_kwargs.setdefault('enable_return_routed_experts', True)
|
|
|
|
# https://github.com/vllm-project/vllm/pull/39917
|
|
import vllm
|
|
from packaging import version
|
|
vllm_version = vllm.__version__
|
|
if vllm_version is not None and version.parse('0.21.0rc1') <= version.parse(vllm_version) <= version.parse(
|
|
'0.21.0'):
|
|
engine_kwargs.setdefault('async_scheduling', False)
|
|
|
|
cfg['vllm_engine_kwargs'] = engine_kwargs
|
|
return cfg
|
|
|
|
def _get_template_kwargs_for_rollout(self) -> Dict[str, Any]:
|
|
"""Extract template config for vLLM alignment (padding_free=False, sp=1)."""
|
|
args = self._data_info.get('_driver_args')
|
|
if args is None:
|
|
return {}
|
|
kwargs = args.get_template_kwargs()
|
|
kwargs['padding_free'] = False
|
|
kwargs['sequence_parallel_size'] = 1
|
|
return kwargs
|
|
|
|
def _build_rollout_config_for_workers(self) -> Optional[Dict[str, Any]]:
|
|
"""Build rollout config dict for MegatronWorker._init_rollout_adapter.
|
|
|
|
Returns None if no rollout GPUs are configured.
|
|
"""
|
|
rollout_gpus = self.group_gpus.get('rollout', 0)
|
|
if rollout_gpus <= 0:
|
|
return None
|
|
rollout_cfg = self.group_cfgs.get('rollout', {})
|
|
bucket_mb = int(os.environ.get('SWIFT_RAY_WEIGHT_BUCKET_MB', '2048'))
|
|
return {
|
|
'rollout_tp_size': rollout_cfg.get('vllm_tensor_parallel_size', 1),
|
|
'rollout_dp_size': rollout_cfg.get('vllm_data_parallel_size', 1),
|
|
'bucket_size_mb': bucket_mb,
|
|
}
|
|
|
|
def _create_trainer(self):
|
|
cls_path = self._entry['trainer']
|
|
mod_path, cls_name = cls_path.rsplit('.', 1)
|
|
mod = importlib.import_module(mod_path)
|
|
trainer_cls = getattr(mod, cls_name)
|
|
weight_sync_mode = self._get_weight_sync_mode()
|
|
sleep_level = self._resolve_sleep_level()
|
|
return trainer_cls(
|
|
self.worker_groups,
|
|
self.rollout_replicas,
|
|
weight_sync_mode=weight_sync_mode,
|
|
sleep_level=sleep_level,
|
|
teacher_replicas=self.teacher_replicas)
|
|
|
|
def _resolve_sleep_level(self) -> int:
|
|
"""Colocate: honor user config; separate: force 0."""
|
|
user_level = self.ray_config.sleep_level
|
|
if self._is_rollout_hybrid():
|
|
return user_level
|
|
if user_level != 0:
|
|
logger.warning('sleep_level=%d ignored in separate mode (vLLM stays resident). '
|
|
'Overriding to 0.', user_level)
|
|
return 0
|
|
|
|
def _get_weight_sync_mode(self) -> str:
|
|
"""Colocate: naive (IPC); separate: nccl (broadcast)."""
|
|
if self._is_rollout_hybrid():
|
|
return 'naive'
|
|
return 'nccl'
|
|
|
|
@property
|
|
def group_gpus(self) -> Dict[str, int]:
|
|
return self.ray_config.group_gpus
|
|
|
|
@property
|
|
def colocate_groups(self) -> List[List[str]]:
|
|
return self.ray_config.colocate_groups
|
|
|
|
def _shutdown(self):
|
|
"""Best-effort teardown — each step swallows exceptions so a
|
|
failure in one stage does not skip the remaining cleanup."""
|
|
for replica in self.rollout_replicas:
|
|
try:
|
|
replica.shutdown()
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning('RolloutReplica shutdown failed: %s', e)
|
|
self.rollout_replicas = []
|
|
|
|
for replica in self.teacher_replicas:
|
|
try:
|
|
replica.shutdown()
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning('TeacherReplica shutdown failed: %s', e)
|
|
self.teacher_replicas = []
|
|
|
|
seen: set = set()
|
|
for wg in self.worker_groups.values():
|
|
if id(wg) not in seen:
|
|
seen.add(id(wg))
|
|
try:
|
|
wg.shutdown()
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning('WorkerGroup shutdown failed: %s', e)
|
|
self.worker_groups.clear()
|
|
|
|
if self.resource_pool_manager is not None:
|
|
try:
|
|
self.resource_pool_manager.destroy_all()
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning('destroy_all placement groups failed: %s', e)
|
|
ray.shutdown()
|
|
|
|
|
|
def main():
|
|
import sys
|
|
argv = sys.argv[1:]
|
|
config_path = None
|
|
for i, arg in enumerate(argv):
|
|
if arg == '--config' and i + 1 < len(argv):
|
|
config_path = argv[i + 1]
|
|
break
|
|
if config_path is None:
|
|
raise ValueError('Usage: python -m swift.ray.megatron.pipeline --config <yaml>')
|
|
|
|
return MegatronRayPipeline(config_path).run()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|