wehub-resource-sync
290 lines
10 KiB
Python
290 lines
10 KiB
Python
"""Driver-side helpers shared by the Ray Megatron pipeline.
|
|
|
|
This module owns the "plain Python" helpers that do not belong to any
|
|
particular trainer class:
|
|
|
|
* YAML → dict config parsing and merging
|
|
* structured config parsing from YAML group dicts
|
|
* driver-side dataset building (via dict, no argv round-trip)
|
|
* train/eval iteration bookkeeping
|
|
* extracting the canonical iteration from worker results
|
|
"""
|
|
import json
|
|
import os
|
|
import torch
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List
|
|
|
|
from swift.arguments import BaseArguments
|
|
from swift.utils import get_logger, parse_args, seed_everything, to_abspath
|
|
|
|
logger = get_logger()
|
|
|
|
_RAY_ONLY_KEYS = frozenset({'gpus', 'colocate_groups', 'nnodes'})
|
|
|
|
_PARALLEL_DEFAULTS: Dict[str, Any] = {
|
|
'tensor_model_parallel_size': 1,
|
|
'pipeline_model_parallel_size': 1,
|
|
'context_parallel_size': 1,
|
|
}
|
|
|
|
|
|
def parse_args_from_dict(class_type, cfg: Dict[str, Any]):
|
|
"""Construct a dataclass from a config dict via HfArgumentParser."""
|
|
argv = _dict_to_argv(cfg)
|
|
args, remaining_args = parse_args(class_type, argv)
|
|
if remaining_args:
|
|
logger.warning('parse_args_from_dict: unrecognised args: %s', remaining_args)
|
|
return args
|
|
|
|
|
|
def _dict_to_argv(cfg: Dict[str, Any]) -> List[str]:
|
|
argv: List[str] = []
|
|
for k, v in cfg.items():
|
|
if k in _RAY_ONLY_KEYS or v is None:
|
|
continue
|
|
flag = f'--{k}'
|
|
if isinstance(v, bool):
|
|
argv += [flag, str(v).lower()]
|
|
elif isinstance(v, (list, tuple)):
|
|
argv.append(flag)
|
|
argv += [str(item) for item in v]
|
|
elif isinstance(v, dict):
|
|
argv += [flag, json.dumps(v)]
|
|
else:
|
|
argv += [flag, str(v)]
|
|
return argv
|
|
|
|
|
|
@dataclass
|
|
class RayConfig:
|
|
rlhf_type: str = 'grpo'
|
|
colocate_groups: List[List[str]] = field(default_factory=list)
|
|
train_gpus: int = 0
|
|
rollout_gpus: int = 0
|
|
teacher_gpus: int = 0
|
|
sleep_level: int = 1
|
|
nnodes: int = 1
|
|
|
|
@property
|
|
def group_gpus(self) -> Dict[str, int]:
|
|
return {
|
|
'train': self.train_gpus,
|
|
'rollout': self.rollout_gpus,
|
|
'teacher': self.teacher_gpus,
|
|
}
|
|
|
|
def gpus_as_process_on_nodes(self, total_gpus: int) -> List[int]:
|
|
"""Split ``total_gpus`` evenly across ``nnodes`` for ResourcePool."""
|
|
if self.nnodes <= 1:
|
|
return [total_gpus]
|
|
per_node, remainder = divmod(total_gpus, self.nnodes)
|
|
if remainder != 0:
|
|
raise ValueError(f'total_gpus={total_gpus} is not evenly divisible by nnodes={self.nnodes}')
|
|
return [per_node] * self.nnodes
|
|
|
|
|
|
def parse_ray_yaml(config_path: str) -> 'tuple[RayConfig, Dict[str, Dict[str, Any]], Dict[str, Any]]':
|
|
"""Parse a Ray YAML config into (ray_config, group_dicts, shared_dict)."""
|
|
import yaml
|
|
with open(config_path) as f:
|
|
raw = yaml.safe_load(f)
|
|
|
|
rlhf_type = raw.get('rlhf_type')
|
|
colocate_groups = raw.pop('colocate_groups', [])
|
|
sleep_level = int(raw.pop('sleep_level', 1))
|
|
nnodes = int(raw.pop('nnodes', 1))
|
|
|
|
group_configs: Dict[str, dict] = {}
|
|
for g in KNOWN_GROUPS:
|
|
group_configs[g] = raw.pop(g, {}) or {}
|
|
|
|
gpu_counts = {g: int(cfg.pop('gpus', 0)) for g, cfg in group_configs.items()}
|
|
shared_config = dict(raw)
|
|
for key, default in _PARALLEL_DEFAULTS.items():
|
|
shared_config.setdefault(key, default)
|
|
|
|
ray_config = RayConfig(
|
|
rlhf_type=rlhf_type,
|
|
colocate_groups=colocate_groups,
|
|
train_gpus=gpu_counts.get('train', 0),
|
|
rollout_gpus=gpu_counts.get('rollout', 0),
|
|
teacher_gpus=gpu_counts.get('teacher', 0),
|
|
sleep_level=sleep_level,
|
|
nnodes=nnodes,
|
|
)
|
|
|
|
_validate_colocate_groups(colocate_groups, gpu_counts)
|
|
|
|
return ray_config, group_configs, shared_config
|
|
|
|
|
|
KNOWN_GROUPS = frozenset(('train', 'rollout', 'teacher'))
|
|
|
|
|
|
def _validate_colocate_groups(
|
|
colocate_groups: List[List[str]],
|
|
gpu_counts: Dict[str, int],
|
|
) -> None:
|
|
"""Validate colocate_groups: ≥2 roles, known, non-overlapping, each with gpus > 0."""
|
|
if not colocate_groups:
|
|
return
|
|
seen: set = set()
|
|
for idx, group in enumerate(colocate_groups):
|
|
if not isinstance(group, list) or len(group) < 2:
|
|
raise ValueError(f'colocate_groups[{idx}] must be a list of ≥2 roles, '
|
|
f'got {group!r}')
|
|
group_gpu_counts = set()
|
|
for role in group:
|
|
if role not in KNOWN_GROUPS:
|
|
raise ValueError(f'colocate_groups[{idx}] contains unknown role {role!r}; '
|
|
f'valid roles: {sorted(KNOWN_GROUPS)}')
|
|
if role in seen:
|
|
raise ValueError(f'Role {role!r} appears in multiple colocate groups')
|
|
seen.add(role)
|
|
n = gpu_counts.get(role, 0)
|
|
if n <= 0:
|
|
raise ValueError(f'Role {role!r} in colocate_groups[{idx}] has 0 GPUs; '
|
|
f'colocated roles must each have gpus > 0')
|
|
group_gpu_counts.add(n)
|
|
if len(group_gpu_counts) > 1:
|
|
raise ValueError(f'colocate_groups[{idx}] roles have different GPU counts '
|
|
f'{dict(zip(group, [gpu_counts[r] for r in group]))}; '
|
|
f'colocated roles must share the same GPU set')
|
|
|
|
|
|
def merge_group_dict(shared: Dict[str, Any], group: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Merge shared + group config, stripping Ray-only keys and None values."""
|
|
merged = {**shared, **group}
|
|
for k in _RAY_ONLY_KEYS:
|
|
merged.pop(k, None)
|
|
return {k: v for k, v in merged.items() if v is not None}
|
|
|
|
|
|
def estimate_dp_size(cfg: Dict[str, Any], gpus: int) -> int:
|
|
"""Estimate DP size from a merged group config dict."""
|
|
tp = cfg.get('tensor_model_parallel_size', 1)
|
|
pp = cfg.get('pipeline_model_parallel_size', 1)
|
|
cp = cfg.get('context_parallel_size', 1)
|
|
assert gpus % (tp * pp * cp) == 0
|
|
return gpus // (tp * pp * cp)
|
|
|
|
|
|
def build_dataset_from_dict(cfg: Dict[str, Any]):
|
|
"""Build dataset on the driver without instantiating a Megatron pipeline.
|
|
"""
|
|
from swift.megatron.arguments import MegatronRLHFArguments
|
|
from swift.rlhf_trainers.utils import identity_data_collator
|
|
cfg = dict(cfg)
|
|
cfg['skip_megatron_init'] = True
|
|
args = parse_args_from_dict(MegatronRLHFArguments, cfg)
|
|
|
|
if hasattr(args, 'seed'):
|
|
seed_everything(args.seed)
|
|
|
|
rlhf_type = args.rlhf_type
|
|
if rlhf_type in ('grpo', 'gkd'):
|
|
args.remove_unused_columns = False
|
|
if args.output_dir is None:
|
|
args.output_dir = f'megatron_output/{args.model_suffix}'
|
|
args.output_dir = to_abspath(args.output_dir)
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
with torch.device('meta'):
|
|
_, processor = args.get_model_processor(load_model=False, download_model=args.mcore_model is None)
|
|
|
|
template = _prepare_template(args, processor)
|
|
train_dataset, val_dataset = _prepare_dataset(args)
|
|
|
|
data_collator = identity_data_collator if rlhf_type in ('grpo', 'gkd') else template.data_collator
|
|
|
|
# TODO: integrate val_dataset / eval_iters into the training loop
|
|
return {
|
|
'train_dataset': train_dataset,
|
|
'val_dataset': val_dataset,
|
|
'data_collator': data_collator,
|
|
'micro_batch_size': args.micro_batch_size,
|
|
'global_batch_size': args.global_batch_size,
|
|
'padding_free': args.padding_free,
|
|
'num_train_epochs': args.num_train_epochs,
|
|
'train_iters': args.train_iters,
|
|
'save_strategy': args.save_strategy,
|
|
'eval_iters': args.eval_iters,
|
|
'num_generations': args.num_generations,
|
|
'template': template,
|
|
'_driver_args': args,
|
|
}
|
|
|
|
|
|
def _prepare_template(args, processor):
|
|
"""Create template from args and processor — no pipeline object needed."""
|
|
template = args.get_template(processor)
|
|
mode_mapping = {'grpo': 'train', 'gkd': 'train', 'kto': 'kto'}
|
|
template.set_mode(mode_mapping.get(args.rlhf_type, 'rlhf'))
|
|
template.use_megatron = True
|
|
return template
|
|
|
|
|
|
def _prepare_dataset(args: BaseArguments):
|
|
"""Load and optionally encode dataset — no pipeline object needed."""
|
|
# Ray pipeline has no validation/eval loop yet
|
|
if args.split_dataset_ratio and args.split_dataset_ratio > 0:
|
|
logger.warning(
|
|
'Ray pipeline has no validation loop yet; overriding split_dataset_ratio '
|
|
'%s -> 0.0 (no validation split).', args.split_dataset_ratio)
|
|
args.split_dataset_ratio = 0.0
|
|
if args.val_dataset:
|
|
logger.warning('Ray pipeline has no validation loop yet; ignoring val_dataset=%s.', args.val_dataset)
|
|
args.val_dataset = []
|
|
|
|
assert args.rlhf_type in ('grpo', 'gkd')
|
|
return args.load_dataset()
|
|
|
|
|
|
def compute_iter_params(data_info: Dict[str, Any], dp_size: int) -> Dict[str, Any]:
|
|
"""Compute train_iters / eval_iters / save_steps on the driver."""
|
|
mbs = data_info['micro_batch_size']
|
|
gbs = data_info['global_batch_size']
|
|
step_batch_size = mbs * dp_size
|
|
num_gen = data_info.get('num_generations', 1)
|
|
train_ds = data_info.get('train_dataset')
|
|
val_ds = data_info.get('val_dataset')
|
|
train_len = len(train_ds) if train_ds is not None and hasattr(train_ds, '__len__') else 0
|
|
val_len = len(val_ds) if val_ds is not None and hasattr(val_ds, '__len__') else 0
|
|
|
|
result: Dict[str, Any] = {}
|
|
|
|
if data_info.get('save_strategy') == 'epoch' and train_len > 0:
|
|
ds_sample = train_len // step_batch_size * step_batch_size * num_gen
|
|
result['save_steps'] = ds_sample // gbs
|
|
result['eval_steps'] = result['save_steps']
|
|
|
|
train_iters = data_info.get('train_iters')
|
|
if data_info.get('num_train_epochs') is not None and train_len > 0:
|
|
ds_sample = train_len // step_batch_size * step_batch_size * num_gen
|
|
train_iters = ds_sample * data_info['num_train_epochs'] // gbs
|
|
result['train_iters'] = train_iters
|
|
|
|
eval_iters = data_info.get('eval_iters', -1)
|
|
if eval_iters is not None and eval_iters < 0:
|
|
if val_len == 0:
|
|
eval_iters = 0
|
|
else:
|
|
ds_sample = val_len // step_batch_size * step_batch_size * num_gen
|
|
eval_iters = max(ds_sample // gbs, 1)
|
|
if val_len > 0 and val_len < step_batch_size:
|
|
eval_iters = 0
|
|
result['eval_iters'] = eval_iters or 0
|
|
|
|
return result
|
|
|
|
|
|
def extract_iteration(step_results) -> int:
|
|
"""Read the canonical iteration off ``WorkerGroup.execute`` results."""
|
|
if not step_results:
|
|
return 0
|
|
for r in step_results:
|
|
if isinstance(r, dict) and 'iteration' in r:
|
|
return int(r['iteration'])
|
|
return 0
|