From 434d2f5064e1bd0a793b604dc452da2d4a04ac2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=93=E9=9B=B7?= Date: Mon, 14 Jul 2025 10:39:15 +0800 Subject: [PATCH 01/16] PullRequest: 353 [Lite] Add gradient checkpointing to FSDPEngine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch mzy/add-gradient-ckpt of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/353 Reviewed-by: 博惟 * add gradient checkpointing --- arealite/engine/fsdp_engine.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index 8537e02b56..fc1acdc8c8 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -109,6 +109,11 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): attn_implementation=self.config.attn_impl, ) + if self.config.gradient_checkpointing: + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + # Simple auto wrap policy self.mixed_precision_policy = MixedPrecisionPolicy( param_dtype=torch.bfloat16, From d8038b26690f74b36820cca28aa8903766355a39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 14 Jul 2025 11:19:40 +0800 Subject: [PATCH 02/16] PullRequest: 354 [lite] GRPO pre-commit: minor changes in FSDP engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-fix1 of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/354 Reviewed-by: 晓雷 * . * . * . * . --- arealite/api/cli_args.py | 59 +++++++++++++++++++++++++++- arealite/engine/fsdp_engine.py | 70 +++++++++++++++++++++------------- arealite/tests/test_utils.py | 5 ++- arealite/utils/data.py | 53 +++++++++++-------------- arealite/utils/model.py | 8 ++++ 5 files changed, 135 insertions(+), 60 deletions(-) create mode 100644 arealite/utils/model.py diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index ca4b78bd8a..4d7302c8ad 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -12,7 +12,6 @@ from omegaconf import MISSING, OmegaConf from arealite.utils.fs import get_user_tmp -from realhf.api.cli_args import OptimizerConfig @dataclass @@ -84,6 +83,61 @@ def new(self, **kwargs): # Train Engine Configs +@dataclass +class OptimizerConfig: + """Configuration for model optimization during training. + Note: + Set type to "empty" for models that won't be trained. + """ + + type: str = field( + default="adam", + metadata={"help": "Optimizer type", "choices": ["adam", "empty"]}, + ) + lr: float = field(default=2e-5, metadata={"help": "Learning rate"}) + weight_decay: float = field(default=0.05, metadata={"help": "Weight decay"}) + beta1: float = field(default=0.9, metadata={"help": "Adam beta1 parameter"}) + beta2: float = field(default=0.95, metadata={"help": "Adam beta2 parameter"}) + eps: float = field(default=1e-5, metadata={"help": "Adam epsilon parameter"}) + min_lr_ratio: float = field( + default=0.0, + metadata={ + "help": "Minimum learning rate ratio after annealing", + }, + ) + lr_scheduler_type: str = field( + default="constant", + metadata={ + "help": "Learning rate scheduler type", + "choices": ["linear", "cosine", "constant"], + }, + ) + warmup_steps_proportion: float = field( + default=0.001, + metadata={ + "help": "Proportion of training steps for warmup", + }, + ) + offload: bool = field( + default=False, metadata={"help": "Enable optimizer state offloading"} + ) + initial_loss_scale: float = field( + default=2**32, metadata={"help": "Initial loss scaling factor"} + ) + min_loss_scale: float = field( + default=1.0, metadata={"help": "Minimum loss scaling factor"} + ) + loss_scale_window: float = field( + default=5, metadata={"help": "Window size for loss scaling adjustment"} + ) + hysteresis: int = field( + default=2, metadata={"help": "Hysteresis (scaling factor) for loss scaling"} + ) + gradient_clipping: float = field( + default=1.0, metadata={"help": "Gradient clipping threshold"} + ) + + @dataclass class FSDPWrapPolicy: transformer_layer_cls_to_wrap: Optional[List[str]] = field( @@ -127,10 +181,11 @@ class TrainEngineConfig: mb_spec: MicroBatchSpec = field(default_factory=MicroBatchSpec) # Training Backend Configuration + disable_dropout: bool = field(default=False) gradient_checkpointing: bool = field( default=True, metadata={"help": "Enable gradient checkpointing"} ) - bf16: bool = field(default=False, metadata={"help": "Use bf16 precision"}) + dtype: str = field(default="float16", metadata={"help": "Parameter dtype."}) optimizer: Optional[OptimizerConfig] = field( default=None, metadata={"help": "Optimizer configuration"} ) diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index fc1acdc8c8..35bd1f5b66 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -1,6 +1,7 @@ import gc import os import time +from datetime import datetime from typing import Any, Callable, Dict, List, Optional import torch @@ -32,7 +33,7 @@ pad_and_stack_tensors_along_first_dim, pad_mb_list, reorder_list, - split_packed_tensor_dict_into_mb_list, + split_padded_tensor_dict_into_mb_list, unpack_sequence, unsqueeze_mb_list, ) @@ -45,6 +46,7 @@ fsdp2_load_full_state_dict, get_cosine_schedule_with_warmup, ) +from arealite.utils.model import disable_dropout_in_model from arealite.utils.save_load import get_state_dict_from_repo_id_or_path from realhf.api.core.data_api import load_hf_tokenizer from realhf.base import logging, name_resolve, names, pkg_version @@ -95,19 +97,38 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) self.device = torch.device(int(os.environ["LOCAL_RANK"])) - dtype = torch.bfloat16 if self.config.bf16 else torch.float16 + dtype = getattr(torch, self.config.dtype) self.model_config = AutoConfig.from_pretrained( pretrained_model_name_or_path=self.config.path, trust_remote_code=True, ) self.tokenizer = load_hf_tokenizer(self.config.path) + tik = time.perf_counter() with torch.device("cuda"): - # initialize scratch model from config - model = AutoModelForCausalLM.from_config( - self.model_config, - torch_dtype=dtype, - attn_implementation=self.config.attn_impl, + if self.config.init_from_scratch: + # initialize scratch model from config + # NOTE: VLM cannot directly load state dict using this + # random initialized model, so otherwise we call + # from_pretrained rather than loading weights into this random model. + model = AutoModelForCausalLM.from_config( + self.model_config, + torch_dtype=dtype, + attn_implementation=self.config.attn_impl, + ) + else: + model = AutoModelForCausalLM.from_pretrained( + pretrained_model_name_or_path=self.config.path, + trust_remote_code=True, + torch_dtype=dtype, + attn_implementation=self.config.attn_impl, + ) + if self.config.disable_dropout: + disable_dropout_in_model(model) + if self.config.gradient_checkpointing: + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} ) + logger.info(f"Model creation and loading time: {time.perf_counter() - tik}") if self.config.gradient_checkpointing: model.gradient_checkpointing_enable( @@ -116,7 +137,7 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): # Simple auto wrap policy self.mixed_precision_policy = MixedPrecisionPolicy( - param_dtype=torch.bfloat16, + param_dtype=dtype, reduce_dtype=torch.float32, cast_forward_inputs=True, ) @@ -134,23 +155,14 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): } # Wrap with FSDP2 + tik = time.perf_counter() apply_fsdp2(model, fsdp_kwargs, self.config.fsdp.wrap_policy) + logger.info(f"Applying FSDP2 time: {time.perf_counter() - tik}") self.model = model - if not self.config.init_from_scratch: - # Load model from a initial checkpoint path, - # which should only be a huggingface checkpoint. - load_meta = SaveLoadMeta( - path=self.config.path, - weight_format="hf", - with_optim=False, - tokenizer=None, - base_model_path=self.config.path, - ) - self.load(load_meta) - # Set up optimizer if self.optimizer_config is not None: + tik = time.perf_counter() assert ( self.optimizer_config.type == "adam" ), "Only AdamW optimizer is supported in this engine." @@ -194,6 +206,7 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): raise ValueError( f"Unknown lr scheduler type {self.optimizer_config.lr_scheduler_type}" ) + logger.info(f"Create optimizer time: {time.perf_counter() - tik}") self.initialized = True @@ -328,15 +341,19 @@ def _prepare_mb_list(self, input_: TensorDict) -> MicroBatchList: if isinstance(input_, dict): input_ = TensorDict(input_, batch_size=[input_["input_ids"].shape[0]]) input_ = amend_position_ids(input_) - packed_input = pack_tensor_dict(input_) - mb_list = split_packed_tensor_dict_into_mb_list( - packed_input, - self.config.mb_spec, + mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec) + logger.info( + f"Microbatch #tokens (rank {dist.get_rank()}): {mb_list.group_lens}" ) + mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] mb_list = pad_mb_list(mb_list, pad_value=0.0) # NOTE: We unsqueeze here because huggingface transformer models requires # packed input to be of shape [1, total_seqlen]. mb_list = unsqueeze_mb_list(mb_list) + # FIXME: the resulting max_seqlen is a tensor rather than an integer + for mb in mb_list.mbs: + mb["max_seqlen"] = int(mb["max_seqlen"]) + mb["use_cache"] = False return mb_list def train_batch( @@ -361,9 +378,10 @@ def train_batch( dist.all_reduce(total_loss_weight) # Process microbatches with gradient accumulation - for pad_length, padded_mb_input, mb_input in zip( - mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs + for i, (pad_length, padded_mb_input, mb_input) in enumerate( + zip(mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs) ): + self.model.set_is_last_backward(i == len(mb_list.mbs) - 1) outputs = self.model(**padded_mb_input) logits = outputs.logits.squeeze(0) diff --git a/arealite/tests/test_utils.py b/arealite/tests/test_utils.py index 9b1c07378a..de341f2e0e 100644 --- a/arealite/tests/test_utils.py +++ b/arealite/tests/test_utils.py @@ -8,7 +8,7 @@ pad_and_stack_tensors_along_first_dim, pad_sequences_to_tensors, reorder_list, - split_packed_tensor_dict_into_mb_list, + split_padded_tensor_dict_into_mb_list, unpack_sequence, ) @@ -45,7 +45,8 @@ def test_micro_batch_split(mock_padded_data, n_mbs, max_tokens_per_mb): packed_data = pack_tensor_dict(mock_padded_data) original_lens = packed_data["cu_seqlens"][1:] - packed_data["cu_seqlens"][:-1] assert torch.allclose(original_lens, mock_padded_data["attention_mask"].sum(1)) - split_result = split_packed_tensor_dict_into_mb_list(packed_data, mb_spec) + split_result = split_padded_tensor_dict_into_mb_list(mock_padded_data, mb_spec) + split_result.mbs = [pack_tensor_dict(mb) for mb in split_result.mbs] reordered_lens = [original_lens[i] for i in split_result.forward_indices] # assert microbatch split result does not violate requirements diff --git a/arealite/utils/data.py b/arealite/utils/data.py index 0b322ad4bd..f6572d10d8 100644 --- a/arealite/utils/data.py +++ b/arealite/utils/data.py @@ -110,11 +110,11 @@ def pad_input(hidden_states, indices, batch, seqlen): def concat_padded_tensors( - tensor_dicts: List[Dict[str, torch.Tensor]], pad_value: float = 0.0 -) -> Dict[str, torch.Tensor]: + tensor_dicts: List[TensorDict], pad_value: float = 0.0 +) -> TensorDict: """Concatenate and pad tensors from multiple padded tensor dictionaries.""" if not tensor_dicts: - return {} + return TensorDict() # Find max sequence length across all dictionaries lens = [] @@ -156,7 +156,7 @@ def concat_padded_tensors( result[key] = torch.cat(tensors_to_concat, dim=0) if "attention_mask" not in result: result["attention_mask"] = attn_mask - return result + return TensorDict(result, batch_size=[len(lens)]) def to_device(data: Dict[str, torch.Tensor | Any], device) -> Dict[str, torch.Tensor]: @@ -290,13 +290,13 @@ class MicroBatchList: DEFAULT_MAX_TOKENS_PER_MB = int(1e12) -def split_packed_tensor_dict_into_mb_list( +def split_padded_tensor_dict_into_mb_list( data: TensorDict, mb_spec: MicroBatchSpec, group: Optional[dist.ProcessGroup] = None ) -> MicroBatchList: - """Split a packed tensordict into micro-batches based on the cumulative sequence lengths. + """Split a padded tensordict into micro-batches based on the attention mask. Args: - data (TensorDict): Dictionary containing packed tensors with "cu_seqlens" key. + data (TensorDict): Dictionary containing padded tensors. mb_spec (MicroBatchSpec): Specification for micro-batch splitting. group (Optional[dist.ProcessGroup]): Process group for distributed synchronization. @@ -304,24 +304,21 @@ def split_packed_tensor_dict_into_mb_list( MicroBatchList: A structure containing the split micro-batches and metadata. """ assert ( - "cu_seqlens" in data - ), "Input data must be packed and contain 'cu_seqlens' key." + "attention_mask" in data + ), "Input data must be padded and contain 'attention_mask' key." if mb_spec.max_tokens_per_mb is None: mb_spec = MicroBatchSpec.new( mb_spec, max_tokens_per_mb=DEFAULT_MAX_TOKENS_PER_MB ) - cu_seqlens = data["cu_seqlens"] - bs = cu_seqlens.shape[0] - 1 - total_lens = int(cu_seqlens[-1]) - input_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().numpy() + bs = data["attention_mask"].shape[0] + max_seqlen = data["attention_mask"].shape[1] + input_lens = data["attention_mask"].sum(1).long().cpu().numpy() # check tensor shape, split only 1d tensors with length "total_lens" to_split = {} not_to_split = {} for key, value in data.items(): - if key == "cu_seqlens" or key == "max_seqlen": - continue - if not torch.is_tensor(value) or value.numel() != total_lens: + if not torch.is_tensor(value) or value.numel() != bs * max_seqlen: not_to_split[key] = value else: to_split[key] = value @@ -331,6 +328,7 @@ def split_packed_tensor_dict_into_mb_list( splitted_lens = [ [input_lens[i] for i in group_index] for group_index in group_indices ] + group_n_seqs = [len(x) for x in splitted_lens] group_lens = [sum(x) for x in splitted_lens] forward_indices = datapack.flat2d(group_indices) @@ -340,12 +338,16 @@ def split_packed_tensor_dict_into_mb_list( def _split(tensor): """Split and pad a tensor based on forward indices and lens.""" # Unpack the sequence - unpacked = unpack_sequence(tensor, cu_seqlens=cu_seqlens) + unpacked = [tensor[i] for i in range(bs)] # Reorder according to forward indices reordered = reorder_list(unpacked, forward_indices) - reordered = torch.cat(reordered) + reordered = torch.stack(reordered) # Unpack again according to split lens - splitted = unpack_sequence(reordered, lens=group_lens) + splitted = [] + offset = 0 + for _n_seqs in group_n_seqs: + splitted.append(reordered[offset : offset + _n_seqs]) + offset += _n_seqs return splitted to_split = dict_map(to_split, lambda x: _split(x)) @@ -355,16 +357,7 @@ def _split(tensor): # organize splitted micro batches assert len(mbs) == len(splitted_lens), (len(mbs), len(splitted_lens)) for i, (mb, lens) in enumerate(zip(mbs, splitted_lens)): - max_seqlen = max(lens) - lens = torch.tensor(lens, device="cuda") - batch_cu_seqlens = torch.nn.functional.pad( - lens.cumsum(0, dtype=torch.int), (1, 0) - ) - results.append( - TensorDict( - **mb, **not_to_split, max_seqlen=max_seqlen, cu_seqlens=batch_cu_seqlens - ) - ) + results.append(TensorDict(**mb, **not_to_split)) return MicroBatchList( data=data, mbs=results, @@ -433,7 +426,7 @@ def pad_mb_list( # NOTE: GPU page size is 2MB # Take hidden size 4096 with bf16 dtype as an example, # the batch size of a page is 256 - pad_to_length = (l + 255) // 256 * 256 + pad_to_length = (int(l) + 255) // 256 * 256 padded_mb, pad_len = pad_packed_tensor_dict( mb, pad_to_length, pad_value=pad_value ) diff --git a/arealite/utils/model.py b/arealite/utils/model.py new file mode 100644 index 0000000000..5ba3254965 --- /dev/null +++ b/arealite/utils/model.py @@ -0,0 +1,8 @@ +import torch + + +# Copied from trl +def disable_dropout_in_model(model: torch.nn.Module) -> None: + for module in model.modules(): + if isinstance(module, torch.nn.Dropout): + module.p = 0 From 724628eaf0ce8ddce4f42e6120c698a1651ce04d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 14 Jul 2025 15:20:17 +0800 Subject: [PATCH 03/16] PullRequest: 355 [Lite] GRPO pre-commit 2: Refactor RemoteSGLangEngine thread and SGLang configuration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-fix1 of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/355?tab=commit Reviewed-by: 晓雷 * . * . * . * . * . * . * fix * . --- arealite/README.md | 4 +- arealite/api/cli_args.py | 90 +++++----- arealite/api/engine_api.py | 42 +++-- arealite/api/io_struct.py | 2 - arealite/engine/fsdp_engine.py | 9 +- arealite/engine/sglang_remote.py | 258 ++++++++++++++++----------- arealite/tests/test_sglang_engine.py | 58 +++--- arealite/utils/network.py | 100 +++++++++++ arealite/workflow/rlvr.py | 15 +- 9 files changed, 370 insertions(+), 208 deletions(-) create mode 100644 arealite/utils/network.py diff --git a/arealite/README.md b/arealite/README.md index f472c1bef5..80e659c0c0 100644 --- a/arealite/README.md +++ b/arealite/README.md @@ -92,7 +92,7 @@ def main_grpo(): future.result() # synchronous rollout - rollout_batch = rollout.rollout(batch, workflow=MyRolloutWorkflow(rollout_config.workflow)) + rollout_batch = rollout.rollout_batch(batch, workflow=MyRolloutWorkflow(rollout_config.workflow)) # or asynchronous rollout with filtering and off-policyness control # rollout_batch = rollout.prepare_batch(batch, # workflow=MyRolloutWorkflow(rollout_config.workflow), @@ -697,7 +697,7 @@ reward = TrainController(Critic()) rollout_controller = RolloutController(...) for _ in range(epochs): for _ in range(steps_per_epoch): - data = rollout_controller.rollout(prompt) + data = rollout_controller.rollout_batch(prompt) data['reward'] = reward.compute_values(data) ... ``` diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 4d7302c8ad..987056e2c5 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -199,6 +199,9 @@ class SGLangConfig: https://github.com/sgl-project/sglang for detailed documentation. """ + model_path: str = "" + random_seed: int = 1 + skip_tokenizer_init: bool = False disable_cuda_graph: bool = False disable_radix_cache: bool = False disable_cuda_graph_padding: bool = False @@ -234,10 +237,8 @@ class SGLangConfig: schedule_policy: str = "lpm" schedule_conservativeness: float = 1.0 cpu_offload_gb: int = 0 - dtype: str = "float16" kv_cache_dtype: str = "auto" - # logging log_level: str = "warning" log_level_http: Optional[str] = "warning" @@ -253,55 +254,60 @@ class SGLangConfig: @staticmethod def build_cmd( sglang_config: "SGLangConfig", - model_path, tp_size, base_gpu_id, + host, + port, dist_init_addr: Optional[str] = None, - served_model_name: Optional[str] = None, - skip_tokenizer_init: bool = True, + sglang_version: Optional[str] = None, ): - from realhf.base import network, pkg_version, seeding + from realhf.base import pkg_version from realhf.experiments.common.utils import asdict as conf_as_dict args: Dict = conf_as_dict(sglang_config) - args["random_seed"] = seeding.get_seed() - - if served_model_name is None: - served_model_name = model_path - host_ip = network.gethostip() - host = "localhost" if not sglang_config.enable_metrics else host_ip args = dict( host=host, - model_path=model_path, + port=port, # Model and tokenizer - tokenizer_path=model_path, + tokenizer_path=sglang_config.model_path, tokenizer_mode="auto", load_format="auto", trust_remote_code=True, device="cuda", - served_model_name=served_model_name, is_embedding=False, - skip_tokenizer_init=skip_tokenizer_init, # Other runtime options tp_size=tp_size, # Because we have set CUDA_VISIBLE_DEVICES to a single GPU in each process base_gpu_id=base_gpu_id, nnodes=1, node_rank=0, + # initialization addresses and ports dist_init_addr=dist_init_addr, **args, ) - - if pkg_version.is_version_less("sglang", "0.4.4"): + if sglang_version: + version_less_than_0_4_4 = ( + pkg_version.compare_versions(sglang_version, "0.4.4") < 0 + ) + version_less_than_0_4_3 = ( + pkg_version.compare_versions(sglang_version, "0.4.3") < 0 + ) + elif pkg_version.is_available("sglang"): + version_less_than_0_4_4 = pkg_version.is_version_less("sglang", "0.4.4") + version_less_than_0_4_3 = pkg_version.is_version_less("sglang", "0.4.3") + else: + raise ValueError( + "A installed SGLang package or a specific SGLang version should be provided to build SGLang server cmd." + ) + if version_less_than_0_4_4: args.pop("log_requests_level") - if pkg_version.is_version_less("sglang", "0.4.3"): + if version_less_than_0_4_3: args.pop("enable_nccl_nvls") args.pop("triton_attention_num_kv_splits") args.pop("cuda_graph_bs") args.pop("enable_memory_saver") args.pop("allow_auto_truncate") args.pop("file_storage_path") - flags = [] for k, v in args.items(): if v is None or v is False or v == "": @@ -320,8 +326,8 @@ def build_cmd( @dataclass class InferenceEngineConfig: - experiment_name: str - trial_name: str + experiment_name: str = MISSING + trial_name: str = MISSING max_concurrent_rollouts: None | int = field( default=None, metadata={ @@ -345,27 +351,20 @@ class InferenceEngineConfig: }, ) # Used by remote inference engines. - server_addrs: List[str] = field( - default_factory=list, - metadata={"help": "List of server addresses for inference."}, - ) + enable_rollout_tracing: bool = field(default=False) schedule_policy: str = field( default="round_robin", metadata={"help": "Request scheduling policy", "choices": ["round_robin"]}, ) + setup_timeout: float = field(default=90.0) request_timeout: float = field( - default=30.0, metadata={"help": "Timeout for HTTP requests."} + default=3600, metadata={"help": "Timeout for HTTP requests."} ) request_retries: int = field( default=3, metadata={"help": "Number of retries for failed requests."} ) -@dataclass -class SGLangEngineConfig: - pass - - @dataclass class _Timer: experiment_name: str = MISSING @@ -595,42 +594,53 @@ class BaseExperimentConfig: evaluator: EvaluatorConfig = field(default_factory=EvaluatorConfig) stats_logger: StatsLoggerConfig = field(default_factory=StatsLoggerConfig) + server_only: bool = False + sglang: SGLangConfig = field(default_factory=SGLangConfig) + @dataclass class SFTConfig(BaseExperimentConfig): model: TrainEngineConfig = field(default_factory=TrainEngineConfig) -def load_expr_config(argv: List[str], config_cls) -> Tuple[BaseExperimentConfig, str]: +def parse_cli_args(argv: List[str]): parser = argparse.ArgumentParser() parser.add_argument( "--config", help="The path of the main configuration file", required=True ) args, overrides = parser.parse_known_args(argv) - # Initialize hydra config config_file = Path(args.config).absolute() assert config_file.exists() # hydra only recognize relative paths - relpath = Path( - os.path.relpath(str(config_file), (Path(__file__).parent).absolute()) - ) + relpath = Path(os.path.relpath(str(config_file), Path(__file__).parent.absolute())) hydra_init(config_path=str(relpath.parent), job_name="app", version_base=None) cfg = hydra_compose( - config_name=str(relpath.name).rstrip(".yaml"), + config_name=str(relpath.name).split(".yaml")[0], overrides=overrides, ) + return cfg, config_file + +def to_structured_cfg(cfg, config_cls): # Merge with the default configuration. # The yaml and commandline can omit some default values defined in python dataclasses. default_cfg = OmegaConf.structured(config_cls) cfg = OmegaConf.merge(default_cfg, cfg) + return cfg + + +def load_expr_config(argv: List[str], config_cls): + cfg, config_file = parse_cli_args(argv) + cfg = to_structured_cfg(cfg, config_cls=config_cls) cfg = OmegaConf.to_object(cfg) assert isinstance(cfg, BaseExperimentConfig) - # Setup environment - from realhf.base import constants, name_resolve + from realhf.base import constants, name_resolve, names constants.set_experiment_trial_names(cfg.experiment_name, cfg.trial_name) name_resolve.reconfigure(cfg.cluster.name_resolve) + name_resolve.clear_subtree( + names.trial_root(experiment_name=cfg.experiment_name, trial_name=cfg.trial_name) + ) return cfg, str(config_file) diff --git a/arealite/api/engine_api.py b/arealite/api/engine_api.py index 26e124aea1..151117abe5 100644 --- a/arealite/api/engine_api.py +++ b/arealite/api/engine_api.py @@ -5,6 +5,7 @@ import torch from tensordict import TensorDict +from torchdata.stateful_dataloader import StatefulDataLoader from arealite.api.io_struct import ( FinetuneSpec, @@ -77,9 +78,9 @@ def step_lr_scheduler(self): def train_batch( self, - input_: Dict, - loss_fn: Callable[[torch.Tensor, Dict], torch.Tensor], - loss_weight_fn: Callable[[Dict], float], + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], ) -> Dict[str, float]: """Update the model with a batch of data and a loss function.""" raise NotImplementedError() @@ -87,9 +88,9 @@ def train_batch( @torch.no_grad() def eval_batch( self, - input_: Dict, - loss_fn: Callable[[torch.Tensor, Dict], torch.Tensor], - loss_weight_fn: Callable[[Dict], float], + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], ) -> torch.Tensor | None: """Evaluate the model using the forward pass and loss function.""" raise NotImplementedError() @@ -97,9 +98,9 @@ def eval_batch( @torch.no_grad() def forward( self, - input_: Dict, + input_: TensorDict, output_seqlens: List[List[int]] | None = None, - post_hook: Callable[[torch.Tensor, Dict], Any] | None = None, + post_hook: Callable[[torch.Tensor, TensorDict], Any] | None = None, aggregate_fn: Callable[[List[Any]], Any] = torch.cat, ) -> Any | None: """Run the forward pass or inference on the model. Note that it is gradient-free.""" @@ -127,12 +128,33 @@ def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: """Asynchronously submit a request to the inference engine. Exits immediately.""" raise NotImplementedError() - def wait(self, count: int, timeout: float) -> TensorDict: + def wait( + self, + count: int, + timeout: float | None = None, + should_accept: Callable | None = None, + ) -> TensorDict: """Wait for a specified number of requests to complete, with a timeout.""" raise NotImplementedError() - def rollout( + def rollout_batch( self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" ) -> TensorDict: """Submit a batch of requests to the inference engine and wait for the results.""" raise NotImplementedError() + + def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", + ): + """Asynchronously submit and wait until a full batch is ready.""" + raise NotImplementedError() + + def pause(self): + """Pause request submission for async rollout. Used during evaluation to prevent data over generation.""" + raise NotImplementedError() + + def resume(self): + """Resume request submission for async rollout.""" + raise NotImplementedError() diff --git a/arealite/api/io_struct.py b/arealite/api/io_struct.py index 3033af8c3a..b1f670a185 100644 --- a/arealite/api/io_struct.py +++ b/arealite/api/io_struct.py @@ -16,7 +16,6 @@ @dataclass class LLMRequest: rid: str = field(default_factory=lambda: str(uuid.uuid4())) - text: Optional[str] = None input_ids: List[int] = field(default_factory=list) gconfig: GenerationHyperparameters = field( default_factory=GenerationHyperparameters @@ -28,7 +27,6 @@ class LLMRequest: @dataclass class LLMResponse: # outputs - completions: str input_tokens: List[int] = field(default_factory=list) output_tokens: List[int] = field(default_factory=list) output_logprobs: List[float] = field(default_factory=list) diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index 35bd1f5b66..13a3491434 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -130,11 +130,6 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): ) logger.info(f"Model creation and loading time: {time.perf_counter() - tik}") - if self.config.gradient_checkpointing: - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - # Simple auto wrap policy self.mixed_precision_policy = MixedPrecisionPolicy( param_dtype=dtype, @@ -318,7 +313,9 @@ def upload_weights(self, meta: WeightUpdateMeta): self.config.trial_name, meta.model_version, ) - name_resolve.add(update_name, str(time.time_ns()), keepalive_ttl=120) + name_resolve.add( + update_name, str(datetime.now().timestamp()), keepalive_ttl=120 + ) else: raise ValueError(f"Unknown weight update type {meta.type}") diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index 7e967ac3e4..4069a8b1b2 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -1,23 +1,30 @@ import asyncio +import os +import random import threading import time import traceback from concurrent.futures import ThreadPoolExecutor +from datetime import datetime from queue import Empty, Full, Queue -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional import aiohttp +import requests import torch.distributed as dist from tensordict import TensorDict +from torchdata.stateful_dataloader import StatefulDataLoader from arealite.api.cli_args import InferenceEngineConfig from arealite.api.engine_api import InferenceEngine from arealite.api.io_struct import ( + FinetuneSpec, LLMRequest, LLMResponse, RolloutStat, WeightUpdateMeta, ) +from arealite.utils.padding import concat_padded_tensors from realhf.base import logging, name_resolve, names, pkg_version if TYPE_CHECKING: @@ -30,7 +37,7 @@ else: SGLANG_TOKEN_OUTPUT_IDENTIFIER = "token_ids" -ROLLOUT_POLL_WAIT_TIME = 0.4 +ROLLOUT_POLL_WAIT_TIME = 0.1 RID_CACHE_SIZE = 128 @@ -46,22 +53,51 @@ def __init__(self, config: InferenceEngineConfig): # Maintain the addresses for the recent 128 requests self.rid_queue = [] - self.addresses = config.server_addrs - self.server_idx = 0 + self.addresses = os.getenv("AREAL_LLM_SERVER_ADDRS").split(",") + if not self.addresses: + raise RuntimeError("No configured SGLang servers.") + logger.info("Waiting for server ready...") + for addr in self.addresses: + self._wait_for_server(addr) + logger.info("Servers are all ready!") - qsize = config.queue_size or config.max_concurrent_rollouts * 10 + self.server_idx = random.randint(0, len(self.addresses) - 1) + + qsize = config.queue_size or config.max_concurrent_rollouts * 16 self.input_queue = Queue(maxsize=qsize) self.output_queue = Queue(maxsize=qsize) self.result_cache = [] self.exiting = threading.Event() + self.paused = threading.Event() self.lock = threading.Lock() self.rollout_stat = RolloutStat() self._version = 0 - def initialize(self, addr: str | None, ft_spec: Optional[Dict[str, Any]] = None): + def _wait_for_server(self, address): + base_url = f"http://{address}" + tik = time.time() + while time.time() - tik < self.config.setup_timeout: + if self.check_health(base_url): + return + time.sleep(1) + raise RuntimeError("server launch failed") + + def check_health(self, base_url): + # Check server endpoint + try: + response = requests.get( + f"{base_url}/metrics", + timeout=30, + ) + return response.status_code == 200 + except requests.exceptions.RequestException as e: + return False + + def initialize(self, addr: str | None, ft_spec: FinetuneSpec = None): + self.rollout_tasks: Dict[str, asyncio.Task] = {} self.rollout_thread = threading.Thread(target=self._rollout_thread) self.rollout_thread.start() @@ -85,79 +121,45 @@ def _rollout_thread(self): traceback.print_exc() async def _rollout_thread_async(self): - data = None - - rollout_tasks: Dict[str, asyncio.Task] = {} + pending_data = [] + rollout_tasks = self.rollout_tasks rid = 0 - try: while not self.exiting.is_set(): # Load next data from controller - if data is None: + while True: try: data, workflow = self.input_queue.get_nowait() - logger.info(f"Get data from puller: {data}") + logger.debug(f"Get data from puller: {data}") + pending_data.append(data) except Empty: logger.debug(f"No data from puller stream.") + break # Check capacity - if dist.is_initialized(): - world_size = dist.get_world_size() - else: - world_size = 1 - - cannot_rollout_reason = [] - capacity = max(1, self.config.max_concurrent_rollouts // world_size) - can_rollout = len(rollout_tasks) < capacity - if not can_rollout: - cannot_rollout_reason.append( - f"Exceeding capacity: # running tasks {len(rollout_tasks)} >= capacity {capacity}" - ) - - # Staleness control - version = self.get_version() - ofp = self.config.max_head_offpolicyness - with self.lock: - sample_cnt = self.rollout_stat.accepted + self.rollout_stat.running - expected_version = sample_cnt // self.config.consumer_batch_size - not_staled = expected_version <= ofp + version - can_rollout &= not_staled - if not not_staled: - cannot_rollout_reason.append( - f"Staled: expected version ({expected_version}) = " - f"global sample cnt ({sample_cnt}) // batch size ({self.config.consumer_batch_size}), " - f"current latest version {version}, " - f"offpolicyness {self.config.max_head_offpolicyness}." - ) - - if not can_rollout: - logger.debug( - f"Cannot submit new rollouts. " - + "\n".join(cannot_rollout_reason) - ) - + capacity = self.get_capacity() # Create new rollout task - if can_rollout and data is not None: + while capacity > 0 and pending_data and not self.paused.is_set(): task = asyncio.create_task( - workflow.arun_episode(self, data), name=str(rid) + workflow.arun_episode(self, pending_data.pop(0)), name=str(rid) ) - rollout_tasks[str(rid)] = task - with self.lock: + rollout_tasks[str(rid)] = task self.rollout_stat.submitted += 1 self.rollout_stat.running += 1 - logger.info( - f"Submit rollout rid {rid}. " - f"Submit: {self.rollout_stat.submitted}, " - f"running: {self.rollout_stat.running}, " - f"accepted: {self.rollout_stat.accepted}." - ) - + if self.config.enable_rollout_tracing: + logger.info( + f"Submit rollout rid {rid}. " + f"Submit: {self.rollout_stat.submitted}, " + f"running: {self.rollout_stat.running}, " + f"accepted: {self.rollout_stat.accepted}." + ) + capacity -= 1 rid += 1 - data = None # Wait for rollout completion - tasks = list(rollout_tasks.values()) + with self.lock: + tasks = list(rollout_tasks.values()) done = [] if tasks: done, _ = await asyncio.wait( @@ -165,16 +167,19 @@ async def _rollout_thread_async(self): timeout=ROLLOUT_POLL_WAIT_TIME, return_when=asyncio.FIRST_COMPLETED, ) + if not done: + await asyncio.sleep(1) else: - await asyncio.sleep(ROLLOUT_POLL_WAIT_TIME) + await asyncio.sleep(1) # Collect done results for task in done: traj = await task traj: TensorDict task_rid = task.get_name() - rollout_tasks.pop(task_rid) - self.rollout_stat.accepted += 1 + with self.lock: + rollout_tasks.pop(task_rid) + self.rollout_stat.accepted += 1 try: self.output_queue.put_nowait(traj) @@ -185,21 +190,25 @@ async def _rollout_thread_async(self): with self.lock: self.rollout_stat.running -= 1 - logger.info( - f"Finish rollout {task_rid}. " - f"Submit: {self.rollout_stat.submitted}, " - f"running: {self.rollout_stat.running}, " - f"accepted: {self.rollout_stat.accepted}." - ) + if self.config.enable_rollout_tracing: + logger.info( + f"Finish rollout {task_rid}. " + f"Submit: {self.rollout_stat.submitted}, " + f"running: {self.rollout_stat.running}, " + f"accepted: {self.rollout_stat.accepted}." + ) + except Exception: + traceback.print_exc() finally: # Cancel remaining tasks - for task in rollout_tasks.values(): - if not task.done(): - task.cancel() - try: - await task - except asyncio.CancelledError: - pass + with self.lock: + for task in rollout_tasks.values(): + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass def choose_server(self) -> str: if self.config.schedule_policy == "round_robin": @@ -236,8 +245,7 @@ async def arequest_with_retry( async with aiohttp.ClientSession( timeout=aiohttp.ClientTimeout( total=timeout, - sock_connect=30, - sock_read=timeout, + sock_connect=timeout, ) ) as session: if method.upper() == "GET": @@ -252,7 +260,7 @@ async def arequest_with_retry( raise ValueError(f"Unsupported HTTP method: {method}") response.raise_for_status() - return response + return await response.json() except ( aiohttp.ClientError, @@ -288,15 +296,11 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: # NOTE: rid should NOT be passed in payload payload = { - "text": req.text, + "input_ids": req.input_ids.copy(), "sampling_params": sample_params, "return_logprob": True, "stream": False, } - if req.text: - payload["text"] = req.text - else: - payload["input_ids"] = req.input_ids # Make request start_time = time.perf_counter() @@ -324,7 +328,7 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: and len(accumulated_output_tokens) < gconfig.max_new_tokens ): # loop until the generation is complete - response = await self.arequest_with_retry( + result = await self.arequest_with_retry( endpoint="/generate", payload=payload, method="POST", @@ -332,10 +336,8 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: timeout=self.config.request_timeout, target_addr=server_addr, ) - result = await response.json() # Parse response - completions += result["text"] meta_info = result["meta_info"] output_tokens = [x[1] for x in meta_info["output_token_logprobs"]] output_logprobs = [x[0] for x in meta_info["output_token_logprobs"]] @@ -350,12 +352,11 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: finish_reason = meta_info["finish_reason"] stop_reason = finish_reason["type"] - payload["text"] += result["text"] + payload["input_ids"] += result[SGLANG_TOKEN_OUTPUT_IDENTIFIER] latency = time.perf_counter() - start_time return LLMResponse( - completions=completions, input_tokens=req.input_ids, output_tokens=accumulated_output_tokens, output_logprobs=accumulated_output_logprobs, @@ -376,10 +377,10 @@ def _update_weights(self, meta: WeightUpdateMeta): update_name = names.update_weights_from_disk( self.config.experiment_name, self.config.trial_name, meta.model_version ) - save_timestamp = int(name_resolve.wait(update_name, timeout=120)) - load_timestamp = time.time_ns() + save_timestamp = float(name_resolve.wait(update_name, timeout=120)) + load_timestamp = datetime.now().timestamp() logger.info( - f"Begin update weights from {meta.path}, responded in {(load_timestamp - save_timestamp)/1e6:.2f} ms" + f"Begin update weights from {meta.path}, responded in {(load_timestamp - save_timestamp):.2f}s" ) try: jobs = [ @@ -393,14 +394,14 @@ def _update_weights(self, meta: WeightUpdateMeta): finally: loop.close() logger.info( - f"Loading weights done in {(time.time_ns() - load_timestamp)/1e6:.2f} ms" + f"Loading weights done in {(datetime.now().timestamp() - load_timestamp):.2f}s" ) self.set_version(meta.model_version) else: raise NotImplementedError(f"Unsupported weight update type: {meta.type}") async def aupdate_weights_from_disk(self, addr, path: str): - response = await self.arequest_with_retry( + res = await self.arequest_with_retry( endpoint="/update_weights_from_disk", payload=dict(model_path=str(path), allow_interrupt=True), method="POST", @@ -408,7 +409,6 @@ async def aupdate_weights_from_disk(self, addr, path: str): timeout=self.config.request_timeout, target_addr=addr, ) - res = await response.json() assert res["success"] if "num_paused_requests" in res: logger.info( @@ -416,15 +416,40 @@ async def aupdate_weights_from_disk(self, addr, path: str): f"during updating weights for server {addr}" ) + def get_capacity(self): + if dist.is_initialized(): + world_size = dist.get_world_size() + else: + world_size = 1 + + max_concurrent_rollouts = max( + 1, self.config.max_concurrent_rollouts // world_size + ) + capacity = max_concurrent_rollouts - len(self.rollout_tasks) + # Staleness control + version = self.get_version() + ofp = self.config.max_head_offpolicyness + with self.lock: + sample_cnt = self.rollout_stat.accepted + self.rollout_stat.running + consumer_bs = max(1, self.config.consumer_batch_size // world_size) + capacity = min(capacity, (ofp + version + 1) * consumer_bs - sample_cnt) + return capacity + def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: try: self.input_queue.put_nowait((data, workflow)) except Full: raise RuntimeError("Input queue full. Please increase queue_size.") - def wait(self, count: int, timeout: float, should_accept: Callable) -> TensorDict: + def wait( + self, + count: int, + timeout: float | None = None, + should_accept: Callable | None = None, + ) -> TensorDict: tik = time.perf_counter() accepted = len(self.result_cache) + timeout = timeout or float(7 * 24 * 3600) while ( accepted < count and not self.exiting.is_set() @@ -432,14 +457,14 @@ def wait(self, count: int, timeout: float, should_accept: Callable) -> TensorDic ): try: result = self.output_queue.get(timeout=ROLLOUT_POLL_WAIT_TIME) - if should_accept(result): + if should_accept is None or should_accept(result): self.result_cache.append(result) accepted += 1 else: with self.lock: self.rollout_stat.accepted -= 1 except Empty: - time.sleep(ROLLOUT_POLL_WAIT_TIME) + pass if self.exiting.is_set(): raise RuntimeError("Rollout engine is exiting, cannot wait for results.") if accepted < count: @@ -450,16 +475,39 @@ def wait(self, count: int, timeout: float, should_accept: Callable) -> TensorDic self.result_cache[:count], self.result_cache[count:], ) - return TensorDict.cat(results, dim=0) + return concat_padded_tensors(results) - def rollout( + def rollout_batch( self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" ) -> TensorDict: """Submit a batch of requests to the inference engine and wait for the results.""" for item in data: self.submit(item, workflow) - return self.wait( - count=len(data), - timeout=self.config.request_timeout, - should_accept=lambda x: True, - ) + return self.wait(count=len(data)) + + def prepare_batch( + self, + data_generator: Iterator, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", + ): + assert dataloader.batch_size is not None + while True: + if self.get_capacity() + dataloader.batch_size > 0: + try: + data = next(data_generator) + except StopIteration: + data_generator = iter(dataloader) + data = next(data_generator) + for item in data: + self.submit(item, workflow=workflow) + try: + return self.wait(dataloader.batch_size, timeout=1) + except TimeoutError: + pass + + def pause(self): + self.paused.set() + + def resume(self): + self.paused.clear() diff --git a/arealite/tests/test_sglang_engine.py b/arealite/tests/test_sglang_engine.py index a4c2043317..62a1a1ff25 100644 --- a/arealite/tests/test_sglang_engine.py +++ b/arealite/tests/test_sglang_engine.py @@ -5,7 +5,6 @@ import uuid import pytest -import requests import torch from tensordict import TensorDict @@ -15,62 +14,43 @@ SGLangConfig, ) from arealite.api.io_struct import LLMRequest, LLMResponse, WeightUpdateMeta +from arealite.utils import network from realhf.api.core.data_api import load_hf_tokenizer -from realhf.base import network EXPR_NAME = "test_sglang_engine" TRIAL_NAME = "trial_0" MODEL_PATH = "/storage/testing/models/Qwen__Qwen3-1.7B/" if not os.path.exists(MODEL_PATH): MODEL_PATH = "Qwen/Qwen2-0.5B" -PORT = 13887 -DIST_PORT = 15887 +PORT, DIST_PORT = network.find_free_ports(2) HOST = network.gethostip() -def check_server_health(base_url): - # Check server endpoint - try: - response = requests.get( - f"{base_url}/metrics", - timeout=30, - ) - return response.status_code == 200 - except requests.exceptions.RequestException: - return False - - @pytest.fixture(scope="module") def sglang_server(): from realhf.base import seeding seeding.set_random_seed(1, EXPR_NAME) cmd = SGLangConfig.build_cmd( - sglang_config=SGLangConfig(mem_fraction_static=0.3), - model_path=MODEL_PATH, + sglang_config=SGLangConfig( + skip_tokenizer_init=True, + model_path=MODEL_PATH, + mem_fraction_static=0.3, + ), + host=HOST, + port=PORT, tp_size=1, base_gpu_id=0, dist_init_addr=f"{HOST}:{DIST_PORT}", - served_model_name=MODEL_PATH, - skip_tokenizer_init=False, ) # Launch process - full_command = f"{cmd} --port {PORT}" - full_command = full_command.replace("\\\n", " ").replace("\\", " ") + cmd = cmd.replace("\\\n", " ").replace("\\", " ") process = subprocess.Popen( - full_command.split(), + cmd.split(), text=True, stdout=sys.stdout, stderr=sys.stdout, ) - base_url = f"http://{HOST}:{PORT}" - tik = time.time() - while time.time() - tik < 90: - if check_server_health(base_url): - break - time.sleep(1) - if time.time() - tik > 90: - raise RuntimeError("server launch failed") yield process.terminate() @@ -80,11 +60,12 @@ async def test_remote_sglang_generate(sglang_server): from arealite.engine.sglang_remote import RemoteSGLangEngine config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) - config.server_addrs = [f"{HOST}:{PORT}"] + tokenizer = load_hf_tokenizer(MODEL_PATH) + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" engine = RemoteSGLangEngine(config) req = LLMRequest( rid=str(uuid.uuid4()), - text="hello! how are you today", + input_ids=tokenizer.encode("hello! how are you today"), gconfig=GenerationHyperparameters(max_new_tokens=16), ) resp = await engine.agenerate(req) @@ -95,7 +76,6 @@ async def test_remote_sglang_generate(sglang_server): == len(resp.output_tokens) == len(resp.output_versions) ) - assert isinstance(resp.completions, str) @pytest.mark.parametrize("n_samples", [1, 2, 4]) @@ -109,7 +89,7 @@ def test_remote_sglang_rollout(sglang_server, n_samples): max_concurrent_rollouts=2, consumer_batch_size=2, ) - config.server_addrs = [f"{HOST}:{PORT}"] + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" engine = RemoteSGLangEngine(config) engine.initialize(None, None) @@ -122,12 +102,13 @@ def test_remote_sglang_rollout(sglang_server, n_samples): reward_fn=lambda **kwargs: 1.0, # Dummy reward function gconfig=gconfig, tokenizer=tokenizer, + enable_thinking=False, ) data = { "messages": [{"role": "user", "content": "Hello, how are you?"}], } - result = engine.rollout([data] * 2, workflow=workflow) + result = engine.rollout_batch([data] * 2, workflow=workflow) assert isinstance(result, TensorDict) bs = result.batch_size assert bs == torch.Size([2 * n_samples]) @@ -147,7 +128,7 @@ def test_remote_sglang_staleness_control(sglang_server, bs, ofp, n_samples): consumer_batch_size=bs, max_head_offpolicyness=ofp, ) - config.server_addrs = [f"{HOST}:{PORT}"] + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" engine = RemoteSGLangEngine(config) engine.initialize(None, None) @@ -160,6 +141,7 @@ def test_remote_sglang_staleness_control(sglang_server, bs, ofp, n_samples): reward_fn=lambda **kwargs: 1.0, # Dummy reward function gconfig=gconfig, tokenizer=tokenizer, + enable_thinking=False, ) data = { "messages": [{"role": "user", "content": "Hello, how are you?"}], @@ -220,7 +202,7 @@ def test_disk_update_weights_from_fsdp_engine(tmp_path_factory, sglang_server): from arealite.engine.sglang_remote import RemoteSGLangEngine config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) - config.server_addrs = [f"{HOST}:{PORT}"] + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" inf_engine = RemoteSGLangEngine(config) # test update weights path = tmp_path_factory.mktemp("upload_weights_from_disk") diff --git a/arealite/utils/network.py b/arealite/utils/network.py new file mode 100644 index 0000000000..5e51826cbb --- /dev/null +++ b/arealite/utils/network.py @@ -0,0 +1,100 @@ +import random +import socket +from typing import List, Set + + +def gethostname(): + return socket.gethostname() + + +def gethostip(): + return socket.gethostbyname(socket.gethostname()) + + +def find_free_ports( + count: int, port_range: tuple = (1024, 65535), exclude_ports: Set[int] | None = None +) -> List[int]: + """ + Find multiple free ports within a specified range. + + Args: + count: Number of free ports to find + port_range: Tuple of (min_port, max_port) to search within + exclude_ports: Set of ports to exclude from search + + Returns: + List of free port numbers + + Raises: + ValueError: If unable to find requested number of free ports + """ + if exclude_ports is None: + exclude_ports = set() + + min_port, max_port = port_range + free_ports = [] + attempted_ports = set() + + # Calculate available port range + available_range = max_port - min_port + 1 - len(exclude_ports) + + if count > available_range: + raise ValueError( + f"Cannot find {count} ports in range {port_range}. " + f"Only {available_range} ports available." + ) + + max_attempts = count * 10 # Reasonable limit to avoid infinite loops + attempts = 0 + + while len(free_ports) < count and attempts < max_attempts: + # Generate random port within range + port = random.randint(min_port, max_port) + + # Skip if port already attempted or excluded + if port in attempted_ports or port in exclude_ports: + attempts += 1 + continue + + attempted_ports.add(port) + + if is_port_free(port): + free_ports.append(port) + + attempts += 1 + + if len(free_ports) < count: + raise ValueError( + f"Could only find {len(free_ports)} free ports " + f"out of {count} requested after {max_attempts} attempts" + ) + + return sorted(free_ports) + + +def is_port_free(port: int) -> bool: + """ + Check if a port is free by attempting to bind to it. + + Args: + port: Port number to check + + Returns: + True if port is free, False otherwise + """ + # Check TCP + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + sock.bind(("", port)) + sock.close() + except OSError: + return False + + # Check UDP + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + sock.bind(("", port)) + sock.close() + return True + except OSError: + return False diff --git a/arealite/workflow/rlvr.py b/arealite/workflow/rlvr.py index fbad7ccecd..a72d7b68cf 100644 --- a/arealite/workflow/rlvr.py +++ b/arealite/workflow/rlvr.py @@ -17,19 +17,24 @@ def __init__( reward_fn, gconfig: GenerationHyperparameters, tokenizer: PreTrainedTokenizerFast, + enable_thinking: bool, ): self.reward_fn = reward_fn self.gconfig = gconfig self.tokenizer = tokenizer + self.enable_thinking = enable_thinking async def arun_episode(self, engine, data): - text = self.tokenizer.apply_chat_template( - data["messages"], tokenize=False, add_generation_prompt=True + input_ids = self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, ) n_samples = self.gconfig.n_samples req = LLMRequest( rid=uuid.uuid4().hex, - text=text, + input_ids=input_ids, gconfig=self.gconfig.new(n_samples=1), ) resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) @@ -42,8 +47,8 @@ async def arun_episode(self, engine, data): versions = [-1] * resp.input_len + resp.output_versions reward = self.reward_fn( - prompt=req.text, - completions=resp.completions, + prompt=self.tokenizer.decode(input_ids), + completions=self.tokenizer.decode(resp.output_tokens), prompt_ids=resp.input_tokens, completion_ids=resp.output_tokens, **data, From 8a15551949ec8cd1989d1260a4bb95bc94403d7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 14 Jul 2025 18:41:22 +0800 Subject: [PATCH 04/16] PullRequest: 357 [lite] GRPO pre-commit 3: Fix typos and experiment utilities MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-fix2 of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/357?tab=comment Reviewed-by: 晓雷 * . * . * . * . * . * fix destroy process group --- arealite/api/cli_args.py | 2 +- arealite/api/engine_api.py | 6 ++++++ arealite/engine/fsdp_engine.py | 12 ++++++++++++ arealite/engine/sglang_remote.py | 4 ++++ arealite/utils/evaluator.py | 12 +++--------- arealite/utils/saver.py | 19 ++++++++++++++++--- arealite/utils/stats_logger.py | 23 +++++++++++++++-------- examples/arealite/configs/gsm8k_sft.yaml | 4 +--- examples/arealite/gsm8k_sft.py | 20 +++++++++++++++----- realhf/scheduler/client.py | 8 ++++---- realhf/system/generation_server.py | 4 ++-- 11 files changed, 79 insertions(+), 35 deletions(-) diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 987056e2c5..0ae41b2df3 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -2,7 +2,7 @@ import os from dataclasses import asdict, dataclass, field from pathlib import Path -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional import uvloop diff --git a/arealite/api/engine_api.py b/arealite/api/engine_api.py index 151117abe5..6334bf2248 100644 --- a/arealite/api/engine_api.py +++ b/arealite/api/engine_api.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional import torch +import torch.distributed as dist from tensordict import TensorDict from torchdata.stateful_dataloader import StatefulDataLoader @@ -41,6 +42,11 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): """Initialize environments for distributed training and load models.""" raise NotImplementedError() + @property + def parallelism_group(self) -> dist.ProcessGroup: + """The global communication group of this engine.""" + raise NotImplementedError() + def get_scheduling_config(self) -> Scheduling: """Get the scheduling configuration for the engine, e.g., image, cpu/gpu/memory size.""" raise NotImplementedError() diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index 13a3491434..b5e9877fb9 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -70,6 +70,8 @@ def __init__(self, config: TrainEngineConfig): self.cpu_offload = None # initialization self.initialized = False + self.own_global_group = False + self._parallelism_group = None self.weight_update_group_initialized = False # TODO: Handle the case when WORLD_SIZE is not set in launcher @@ -80,6 +82,11 @@ def train(self, mode: bool = True): self.model.train(mode=mode) return self + @property + def parallelism_group(self) -> dist.ProcessGroup: + assert self.initialized + return self._parallelism_group + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): # Initialize distributed enviroments and load model. assert addr is None, "FSDPEngine does not support remote initialization." @@ -92,6 +99,8 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): if not dist.is_initialized(): # TODO: Handle the condition when WORLD_SIZE and RANK is not set in launcher dist.init_process_group(backend="nccl") + self.own_global_group = True + self._parallelism_group = dist.new_group() # TODO: Handle the condition when LOCAL_RANK is not set in launcher torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) @@ -212,6 +221,9 @@ def destroy(self): gc.collect() torch.cuda.empty_cache() gc.collect() + dist.destroy_process_group(self.parallelism_group) + if self.own_global_group: + dist.destroy_process_group() self.initialized = False def save(self, meta: SaveLoadMeta): diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index 4069a8b1b2..c2401d386b 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -353,6 +353,10 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: stop_reason = finish_reason["type"] payload["input_ids"] += result[SGLANG_TOKEN_OUTPUT_IDENTIFIER] + sample_params["max_new_tokens"] = min( + sample_params["max_new_tokens"], + gconfig.max_new_tokens - len(output_tokens), + ) latency = time.perf_counter() - start_time diff --git a/arealite/utils/evaluator.py b/arealite/utils/evaluator.py index 9d559ea63b..8bbc48a33c 100644 --- a/arealite/utils/evaluator.py +++ b/arealite/utils/evaluator.py @@ -1,13 +1,9 @@ -from typing import TYPE_CHECKING, Any, Callable +from typing import Callable from arealite.api.cli_args import EvaluatorConfig from arealite.api.io_struct import FinetuneSpec from realhf.base import timeutil -if TYPE_CHECKING: - from tensordict import TensorDict - from torchdata.stateful_dataloader import StatefulDataLoader - class Evaluator: @@ -22,8 +18,7 @@ def __init__(self, config: EvaluatorConfig, ft_spec: FinetuneSpec): def evaluate( self, - valid_dataloader: "StatefulDataLoader", - evaluate_fn: Callable[["TensorDict"], Any], + evaluate_fn: Callable, epoch: int, step: int, global_step: int, @@ -32,5 +27,4 @@ def evaluate( epochs=int(step == self.ft_sepc.steps_per_epoch - 1), steps=1 ): return - for data in valid_dataloader: - evaluate_fn(data) + evaluate_fn() diff --git a/arealite/utils/saver.py b/arealite/utils/saver.py index 2692f0df82..f5a81dccc2 100644 --- a/arealite/utils/saver.py +++ b/arealite/utils/saver.py @@ -21,6 +21,18 @@ def __init__(self, config: SaverConfig, ft_spec: FinetuneSpec, for_recover: bool freq_sec=config.freq_secs, ) + @staticmethod + def get_save_checkpoint_root( + config: SaverConfig, + name: str = "default", + ): + path = os.path.join( + f"{config.fileroot}/checkpoints/{getpass.getuser()}/{config.experiment_name}/{config.trial_name}", + name, + ) + os.makedirs(path, exist_ok=True) + return path + @staticmethod def get_save_checkpoint_path( config: SaverConfig, @@ -30,8 +42,7 @@ def get_save_checkpoint_path( name: str = "default", ): path = os.path.join( - f"{config.fileroot}/checkpoints/{getpass.getuser()}/{config.experiment_name}/{config.trial_name}", - name, + Saver.get_save_checkpoint_root(config, name), f"epoch{epoch}epochstep{step}globalstep{globalstep}", ) os.makedirs(path, exist_ok=True) @@ -51,7 +62,9 @@ def save( epochs=int(step == self.ft_sepc.steps_per_epoch - 1), steps=1 ): return - path = self.get_save_checkpoint_path(epoch, step, global_step, name) + path = Saver.get_save_checkpoint_path( + self.config, epoch, step, global_step, name + ) weight_format = "hf" with_optim = False if self.for_recover: diff --git a/arealite/utils/stats_logger.py b/arealite/utils/stats_logger.py index 62cc0ba9ea..49c0d5594a 100644 --- a/arealite/utils/stats_logger.py +++ b/arealite/utils/stats_logger.py @@ -1,7 +1,7 @@ import getpass import os import time -from typing import Dict +from typing import Dict, List import torch.distributed as dist import wandb @@ -21,6 +21,8 @@ def __init__(self, config: StatsLoggerConfig, ft_spec: FinetuneSpec): self.ft_spec = ft_spec self.init() + self._last_commit_step = 0 + def init(self): if dist.is_initialized() and dist.get_rank() != 0: return @@ -61,7 +63,7 @@ def close(self): if self.summary_writer is not None: self.summary_writer.close() - def commit(self, epoch: int, step: int, global_step: int, data: Dict): + def commit(self, epoch: int, step: int, global_step: int, data: Dict | List[Dict]): if dist.is_initialized() and dist.get_rank() != 0: return self.info( @@ -69,12 +71,17 @@ def commit(self, epoch: int, step: int, global_step: int, data: Dict): f"Step {step+1}/{self.ft_spec.steps_per_epoch} " f"Train step {global_step + 1}/{self.ft_spec.total_train_steps} done." ) - self.info("Stats:") - self.print_stats(data) - wandb.log(data, step=global_step) - if self.summary_writer is not None: - for key, val in data.items(): - self.summary_writer.add_scalar(f"{key}", val, global_step) + if isinstance(data, Dict): + data = [data] + log_step = max(global_step, self._last_commit_step) + for i, item in enumerate(data): + self.info(f"Stats ({i+1}/{len(data)}):") + self.print_stats(item) + wandb.log(item, step=log_step + i) + if self.summary_writer is not None: + for key, val in item.items(): + self.summary_writer.add_scalar(f"{key}", val, log_step + i) + self._last_commit_step = log_step + len(data) - 1 def print_stats(self, stats: Dict[str, float]): self.info("\n" + tabulate_stats(stats)) diff --git a/examples/arealite/configs/gsm8k_sft.yaml b/examples/arealite/configs/gsm8k_sft.yaml index b9bad09568..c8aaa4bb09 100644 --- a/examples/arealite/configs/gsm8k_sft.yaml +++ b/examples/arealite/configs/gsm8k_sft.yaml @@ -16,7 +16,7 @@ model: path: /storage/openpsi/models/Qwen__Qwen3-1.7B/ init_from_scratch: false gradient_checkpointing: false - bf16: true + dtype: bfloat16 mb_spec: max_tokens_per_mb: 4096 optimizer: @@ -34,13 +34,11 @@ train_dataset: batch_size: 128 shuffle: true pin_memory: true - num_workers: 4 valid_dataset: batch_size: 128 shuffle: true pin_memory: true - num_workers: 4 # Utilities saver: diff --git a/examples/arealite/gsm8k_sft.py b/examples/arealite/gsm8k_sft.py index a9b5380a5a..70c67e3207 100644 --- a/examples/arealite/gsm8k_sft.py +++ b/examples/arealite/gsm8k_sft.py @@ -1,6 +1,7 @@ import os import sys +import torch.distributed as dist from datasets import Dataset, load_dataset from datasets.distributed import split_dataset_by_node from torchdata.stateful_dataloader import StatefulDataLoader @@ -95,22 +96,31 @@ def main_sft(): with stats_tracker.record_timing("save"): saver.save(engine, epoch, step, global_step) - with stats_tracker.record_timing("eval"), stats_tracker.scope("sft-eval"): + with stats_tracker.record_timing("eval"): # No need to log anything. Logging will be handled outside # via stats_tracker.export(). + def evaluate_fn(): + with stats_tracker.scope("sft-eval"): + for data in valid_dataloader: + engine.evaluate_lm(data) + evaluator.evaluate( - valid_dataloader, - engine.evaluate_lm, + evaluate_fn, epoch, step, global_step, ) - logger.commit(epoch, step, global_step, stats_tracker.export()) + logger.commit( + epoch, + step, + global_step, + stats_tracker.export(reduce_group=engine.parallelism_group), + ) global_step += 1 - engine.destroy() logger.close() + engine.destroy() if __name__ == "__main__": diff --git a/realhf/scheduler/client.py b/realhf/scheduler/client.py index 71cfdf8f21..3662b7c244 100644 --- a/realhf/scheduler/client.py +++ b/realhf/scheduler/client.py @@ -41,12 +41,12 @@ def __init__(self, run_name, worker_type, host, reason: JobState): class JobInfo: name: str state: JobState - host: str = ( + host: Optional[str] = ( None # The host on which the job is/was running. None if the job had not run. ) - submit_time: str = None - start_time: str = None - slurm_id: str = None # Slurm only. The Slurm id of the job. + submit_time: Optional[str] = None + start_time: Optional[str] = None + slurm_id: Optional[int] = None # Slurm only. The Slurm id of the job. class SchedulerClient: diff --git a/realhf/system/generation_server.py b/realhf/system/generation_server.py index 227554dad6..2a46141a1d 100644 --- a/realhf/system/generation_server.py +++ b/realhf/system/generation_server.py @@ -40,7 +40,7 @@ def execute_shell_command(command: str) -> subprocess.Popen: ) -def apply_sglang_path(): +def apply_sglang_patch(): p = Path(os.path.dirname(__file__)) patch_path = str( p.parent.parent @@ -75,7 +75,7 @@ def launch_server_cmd(command: str, port: int = 30000): If no port is specified, a free port is reserved. """ if not ray.is_initialized(): - apply_sglang_path() + apply_sglang_patch() assert port is not None full_command = f"{command} --port {port}" process = execute_shell_command(full_command) From 3f9596871126cbb354f8113cd007cfa248f0784c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Tue, 15 Jul 2025 12:24:24 +0800 Subject: [PATCH 05/16] PullRequest: 358 [lite] Support GRPO training locally with the GSM8k dataset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-fix3 of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/358 Reviewed-by: 晓雷 * . * . * . * . * fix loss mask * fix * . --- arealite/api/cli_args.py | 98 ++++++- arealite/engine/fsdp_engine.py | 8 +- arealite/engine/hf_engine.py | 2 +- arealite/engine/ppo/actor.py | 334 ++++++++++++++++++++++ arealite/engine/sft/lm_engine.py | 20 +- arealite/engine/sglang_remote.py | 186 +++++------- arealite/launcher/local.py | 307 ++++++++++++++++++++ arealite/tests/test_rlvr_workflow.py | 0 arealite/tests/test_sglang_engine.py | 6 +- arealite/tests/test_utils.py | 2 +- arealite/utils/device.py | 31 ++ arealite/utils/functional.py | 171 ++++++++++- arealite/utils/http.py | 56 ++++ arealite/workflow/rlvr.py | 8 +- examples/arealite/configs/gsm8k_grpo.yaml | 129 +++++++++ examples/arealite/gsm8k_grpo.py | 259 +++++++++++++++++ examples/arealite/gsm8k_sft.py | 7 +- 17 files changed, 1485 insertions(+), 139 deletions(-) create mode 100644 arealite/engine/ppo/actor.py create mode 100644 arealite/launcher/local.py delete mode 100644 arealite/tests/test_rlvr_workflow.py create mode 100644 arealite/utils/device.py create mode 100644 arealite/utils/http.py create mode 100644 examples/arealite/configs/gsm8k_grpo.yaml create mode 100644 examples/arealite/gsm8k_grpo.py diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 0ae41b2df3..1fb886d5f6 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -83,6 +83,8 @@ def new(self, **kwargs): # Train Engine Configs + + @dataclass class OptimizerConfig: """Configuration for model optimization during training. @@ -193,6 +195,85 @@ class TrainEngineConfig: fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) +@dataclass +class PPOActorConfig(TrainEngineConfig): + # Core PPO/GRPO Parameters + group_size: int = field( + default=1, metadata={"help": "Number of sequences in each group"} + ) + group_adv_norm: bool = field( + default=False, + metadata={ + "help": "Normalize advantages within each prompt group rather than globally" + }, + ) + ppo_n_minibatches: int = field( + default=4, metadata={"help": "Number of minibatches for each PPO update"} + ) + eps_clip: float = field( + default=0.2, metadata={"help": "Clipping factor for policy ratio"} + ) + c_clip: Optional[float] = field( + default=None, + metadata={ + "help": "Dual clipping factor for policy ratio, must > 1.0. None disables dual clipping." + }, + ) + temperature: float = field( + default=1.0, metadata={"help": "Temperature during generation."} + ) + # Reward + group_reward_norm: bool = field( + default=False, + metadata={ + "help": "Normalize final reward of each sequence (GRPO-style) to reduce length bias" + }, + ) + reward_scaling: float = field( + default=1.0, metadata={"help": "Reward scaling factor"} + ) + reward_bias: float = field(default=0.0, metadata={"help": "Reward bias"}) + reward_clip: float = field( + default=20.0, metadata={"help": "Maximum absolute value for reward clipping"} + ) + mask_no_eos_with_zero: bool = field( + default=False, + metadata={ + "help": "Mask truncated generations (no EOS token) and exclude from training" + }, + ) + + # Advantage Estimation + discount: float = field( + default=1.0, metadata={"help": "Discount factor for future rewards"} + ) + gae_lambda: float = field( + default=1.0, metadata={"help": "Lambda parameter for GAE"} + ) + adv_norm: bool = field( + default=True, metadata={"help": "Enable advantage normalization"} + ) + + # KL Control + kl_ctl: float = field(default=0.1, metadata={"help": "KL divergence coefficient"}) + + # Asynchronous RL + recompute_logprob: bool = field( + default=False, + metadata={"help": "Recompute logp and replace the logp returned by inference."}, + ) + use_decoupled_loss: bool = field( + default=False, + metadata={"help": "Use the decoupled loss. recompute_logprob must be True."}, + ) + behav_imp_weight_cap: Optional[float] = field( + default=None, + metadata={ + "help": "We filter out the tokens where behav_imp_weight exceeds behav_imp_weight_cap when computing the loss, must be > 1.0, use_decoupled_loss must be true" + }, + ) + + @dataclass class SGLangConfig: """Configuration for SGLang runtime. Refer to: @@ -350,7 +431,6 @@ class InferenceEngineConfig: "the request will not be accepted.", }, ) - # Used by remote inference engines. enable_rollout_tracing: bool = field(default=False) schedule_policy: str = field( default="round_robin", @@ -603,6 +683,17 @@ class SFTConfig(BaseExperimentConfig): model: TrainEngineConfig = field(default_factory=TrainEngineConfig) +@dataclass +class GRPOConfig(BaseExperimentConfig): + async_training: bool = field(default=True) + gconfig: GenerationHyperparameters = field( + default_factory=GenerationHyperparameters + ) + rollout: InferenceEngineConfig = field(default_factory=InferenceEngineConfig) + actor: PPOActorConfig = field(default_factory=PPOActorConfig) + ref: PPOActorConfig = field(default_factory=PPOActorConfig) + + def parse_cli_args(argv: List[str]): parser = argparse.ArgumentParser() parser.add_argument( @@ -636,11 +727,8 @@ def load_expr_config(argv: List[str], config_cls): cfg = OmegaConf.to_object(cfg) assert isinstance(cfg, BaseExperimentConfig) # Setup environment - from realhf.base import constants, name_resolve, names + from realhf.base import constants, name_resolve constants.set_experiment_trial_names(cfg.experiment_name, cfg.trial_name) name_resolve.reconfigure(cfg.cluster.name_resolve) - name_resolve.clear_subtree( - names.trial_root(experiment_name=cfg.experiment_name, trial_name=cfg.trial_name) - ) return cfg, str(config_file) diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index b5e9877fb9..07efc36fae 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -49,7 +49,7 @@ from arealite.utils.model import disable_dropout_in_model from arealite.utils.save_load import get_state_dict_from_repo_id_or_path from realhf.api.core.data_api import load_hf_tokenizer -from realhf.base import logging, name_resolve, names, pkg_version +from realhf.base import constants, logging, name_resolve, names, pkg_version logger = logging.getLogger("FSDPEngine") @@ -98,7 +98,11 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): """Initialize distributed communication and model.""" if not dist.is_initialized(): # TODO: Handle the condition when WORLD_SIZE and RANK is not set in launcher - dist.init_process_group(backend="nccl") + dist.init_process_group( + backend="nccl", + timeout=constants.NCCL_DEFAULT_TIMEOUT, + device_id=torch.device(int(os.environ["LOCAL_RANK"])), + ) self.own_global_group = True self._parallelism_group = dist.new_group() diff --git a/arealite/engine/hf_engine.py b/arealite/engine/hf_engine.py index cd685d0c43..6ddcd74f62 100644 --- a/arealite/engine/hf_engine.py +++ b/arealite/engine/hf_engine.py @@ -90,7 +90,7 @@ def init_distributed(self, config: ParallelismConfig, ft_spec: FinetuneSpec): ) torch.cuda.set_device("cuda:0") - dtype = torch.bfloat16 if self.engine_config.bf16 else torch.float16 + dtype = getattr(torch, self.config.dtype) self.model_config = AutoConfig.from_pretrained( pretrained_model_name_or_path=self.engine_config.path, trust_remote_code=True, diff --git a/arealite/engine/ppo/actor.py b/arealite/engine/ppo/actor.py new file mode 100644 index 0000000000..937f32dc6e --- /dev/null +++ b/arealite/engine/ppo/actor.py @@ -0,0 +1,334 @@ +import functools +from typing import Dict, List, Optional + +import torch +from tensordict import TensorDict + +from arealite.api.cli_args import MicroBatchSpec, PPOActorConfig +from arealite.api.engine_api import TrainEngine +from arealite.engine.fsdp_engine import FSDPEngine +from arealite.utils.data import split_padded_tensor_dict_into_mb_list +from arealite.utils.functional import ( + gather_logprobs, + gather_logprobs_entropy, + masked_normalization, + ppo_actor_loss_fn, +) +from realhf.base import stats_tracker + + +class PPOActor: + + def __init__(self, config: PPOActorConfig, engine: TrainEngine): + self.config = config + self.engine = engine + + self.reward_bias = config.reward_bias + self.reward_scaling = config.reward_scaling + self.reward_clip = config.reward_clip + + self.group_reward_norm = config.group_reward_norm + self.group_adv_norm = config.group_adv_norm + self.group_size = config.group_size + + self.kl_ctl = config.kl_ctl + + self.adv_norm = config.adv_norm + self.discount = config.discount + self.gae_lambda = config.gae_lambda + self.mask_no_eos_with_zero = config.mask_no_eos_with_zero + + self.temperature = config.temperature + + @torch.no_grad() + def compute_logp( + self, + data: TensorDict, + temperature: Optional[float] = None, + ) -> torch.Tensor | None: + + def calc_logprobs(logits, input_data): + labels = torch.roll(input_data["input_ids"], shifts=-1, dims=-1) + logprobs = gather_logprobs(logits, labels, temperature or 1.0) + return logprobs + + self.engine.eval() + return self.engine.forward( + input_=data, + post_hook=calc_logprobs, + aggregate_fn=lambda xs: torch.cat(xs, dim=-1), + ) + + def compute_advantages(self, data: TensorDict) -> None: + bs = data["input_ids"].shape[0] + max_seqlen = data["input_ids"].shape[1] + batch_indices = torch.arange( + bs, device=data["input_ids"].device, dtype=torch.long + ) + + # Compute rewards using the reward function in synchronous RLVR pipeline. + reward_score = data["rewards"] + reward_score = (reward_score + self.reward_bias) * self.reward_scaling + reward_score = torch.clip( + reward_score, max=self.reward_clip, min=-self.reward_clip + ) + if self.group_reward_norm: + for i in range(bs // self.group_size): + s = slice(i * self.group_size, (i + 1) * self.group_size) + r = reward_score[s] + reward_score[s] = (r - r.mean()) / (r.std() + 1e-9) + + loss_mask = data["loss_mask"].float() + loss_mask = torch.roll(loss_mask, shifts=-1, dims=-1) + # Apply the mask to log probabilities. + if not self.config.use_decoupled_loss and self.config.recompute_logprob: + # Overwrite logprobs produced by the inference engine + old_logp = data["logprobs"] = data["prox_logp"] + else: + old_logp = torch.roll(data["logprobs"], shifts=-1, dims=-1) + if not self.config.use_decoupled_loss: + # prox logp not available, use inferenced logp + data["prox_logp"] = old_logp + ref_logp = data.get("ref_logp", torch.zeros_like(old_logp)) + ref_logp *= loss_mask + old_logp *= loss_mask + + # Compute KL-regularized rewards. + attn_mask = data["attention_mask"] + seqlens = attn_mask.sum(-1).long() + seq_no_eos_mask = seqlens == attn_mask.shape[1] + rewards = -self.kl_ctl * (old_logp - ref_logp) + kl_rewards = rewards.clone() + # KL rewards at the next token after eos is zero. + rewards[batch_indices, seqlens - 1] = 0 + indices = torch.clip(seqlens - 2, min=0) + if self.mask_no_eos_with_zero: + rewards[batch_indices, indices] += torch.where( + seq_no_eos_mask, 0, reward_score + ) + else: + rewards[batch_indices, indices] += reward_score + + # Compute GAE. + if "values" not in data: + values = torch.zeros_like(rewards) + else: + values = data["values"] + advantages_reversed = [] + lastgaelam = 0 + for t in reversed(range(max_seqlen - 1)): + nextvalues = values[:, t + 1] + if t == max_seqlen - 2: + nextvalues *= seq_no_eos_mask + delta = rewards[:, t] + self.discount * nextvalues - values[:, t] + lastgaelam = delta + self.discount * self.gae_lambda * lastgaelam + advantages_reversed.append(lastgaelam) + advantages_reversed.append( + torch.zeros(bs, dtype=torch.float32, device=values.device) + ) + advantages = torch.stack(advantages_reversed[::-1], dim=1) + + # Optionally perform advantage normalization. + if self.adv_norm: + if self.group_adv_norm: + adv_list = [] + for i in range(0, bs, self.group_size): + s = slice(i * self.group_size, (i + 1) * self.group_size) + adv = advantages[s] + m = loss_mask[s] + adv_list.append(masked_normalization(adv, m, all_reduce=False)) + advantages = torch.cat(adv_list, 0) + else: + advantages = masked_normalization(advantages, loss_mask) + + # Store data in the dict. + data["advantages"] = advantages + data["kl_rewards"] = kl_rewards + data["tot_rewards"] = rewards + data["loss_mask"] = loss_mask + # because we have rolled old_logp by -1 + data["logprobs"] = old_logp + + def ppo_update(self, data: TensorDict) -> List[Dict[str, float]]: + attn_mask = data["attention_mask"] + loss_mask = data["loss_mask"] + reward_score = data["rewards"] + seqlens = attn_mask.sum(-1) + + all_stats = [] + ########## Logging code starts ########## + result_denominators = { + "correct_n_seqs": (reward_score > 0).bool(), + "incorrect_n_seqs": (reward_score <= 0).bool(), + } + global_denominators = dict( + n_seqs=torch.ones_like(reward_score, dtype=torch.bool), + n_tokens=torch.ones_like(loss_mask, dtype=torch.bool), + n_valid_tokens=loss_mask.bool(), + **result_denominators, + ) + stats_tracker.denominator(**global_denominators) + stats_tracker.stat( + correct_seq_len=seqlens.float(), denominator="correct_n_seqs" + ) + stats_tracker.stat( + incorrect_seq_len=seqlens.float(), denominator="incorrect_n_seqs" + ) + + stats = dict( + advantages=data["advantages"], + kl_rewards=data["kl_rewards"], + final_reward=data["tot_rewards"], + ) + stats_tracker.stat(**stats, denominator="n_valid_tokens") + + prompt_lens = [] + prompt_lens = data["attention_mask"].sum(-1) - data["loss_mask"].sum(-1) + seq_stats = dict( + no_eos_ratios=(seqlens == attn_mask.shape[-1]).float(), + task_reward=reward_score.float(), + prompt_len=prompt_lens.float(), + seq_len=seqlens.float(), + ) + stats_tracker.stat(**seq_stats, denominator="n_seqs") + scalars = dict( + mask_no_eos_with_zero=self.config.mask_no_eos_with_zero, + eps_clip=self.config.eps_clip, + ) + if self.config.c_clip is not None: + scalars["c_clip"] = self.config.c_clip + scalars["use_dual_clip"] = 1 + else: + scalars["use_dual_clip"] = 0 + if self.config.behav_imp_weight_cap is not None: + scalars["behav_imp_weight_cap"] = self.config.behav_imp_weight_cap + stats_tracker.scalar(**scalars) + + global_stats = stats_tracker.export(reduce_group=self.engine.parallelism_group) + for k in global_denominators: + keys = list(global_stats.keys()) + for k2 in keys: + if k2.endswith(k): + global_stats.pop(k2) + ########## Logging code ends ########## + + for key in ["rewards", "tot_rewards", "kl_rewards", "versions"]: + data.pop(key, None) + # NOTE: calling engine.train() is critical to enabling gradient checkpointing + self.engine.train() + mb_inputs = split_padded_tensor_dict_into_mb_list( + data, + mb_spec=MicroBatchSpec(n_mbs=self.config.ppo_n_minibatches), + ) + for mb in mb_inputs.mbs: + train_stat = self.engine.train_batch( + mb, + loss_fn=functools.partial( + grpo_loss_fn, + temperature=self.temperature, + eps_clip=self.config.eps_clip, + c_clip=self.config.c_clip, + behav_imp_weight_cap=self.config.behav_imp_weight_cap, + ), + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), + ) + stats_tracker.scalar(**train_stat) + all_stats.append( + stats_tracker.export(reduce_group=self.engine.parallelism_group) + ) + all_stats[0].update(global_stats) + return all_stats + + +class FSDPPPOActor(FSDPEngine): + + def __init__(self, config: PPOActorConfig): + super().__init__(config) + self.actor = PPOActor(config, self) + + @torch.no_grad() + def compute_logp(self, *args, **kwargs) -> torch.Tensor | None: + return self.actor.compute_logp(*args, **kwargs) + + @torch.no_grad() + def compute_advantages(self, *args, **kwargs) -> None: + self.actor.compute_advantages(*args, **kwargs) + + def ppo_update(self, *args, **kwargs) -> List[Dict[str, float]]: + return self.actor.ppo_update(*args, **kwargs) + + +def grpo_loss_fn( + logits: torch.Tensor, + input_data: Dict, + temperature: float, + eps_clip: float, + c_clip: float | None, + behav_imp_weight_cap: float | None, +): + """Loss function for actor step, all inputs should be splitted into + pipeline micro batches, returns loss and logging stats.""" + input_ids = input_data["input_ids"] + old_logp = input_data["logprobs"] + advantages = input_data["advantages"] + loss_mask = input_data["loss_mask"].bool() + prox_logp = input_data["prox_logp"] + + logprobs, entropy = gather_logprobs_entropy( + logits, torch.roll(input_ids, shifts=-1, dims=-1), temperature + ) + entropy = entropy.detach() + loss, stat = ppo_actor_loss_fn( + logprobs=logprobs, + old_logprobs=old_logp, + advantages=advantages, + eps_clip=eps_clip, + loss_mask=loss_mask, + c_clip=c_clip, + proximal_logprobs=prox_logp, + behav_imp_weight_cap=behav_imp_weight_cap, + ) + + # Log training statistics + stats_tracker.denominator( + n_tokens=torch.ones(logits.shape[0], dtype=torch.bool, device=logits.device), + n_valid_tokens=loss_mask.bool(), + clipped_tokens=stat["clip_mask"], + dual_clipped_tokens=stat["dual_clip_mask"], + ) + + stats_tracker.stat( + importance_weight=stat["importance_weight"], + approx_kl=stat["approx_kl"], + new_logp=logprobs.detach(), + old_logp=old_logp, + entropy=entropy.float(), + actor_loss=stat["loss"], + clip_ratio=stat["clip_mask"].float(), + dual_clip_ratio=stat["dual_clip_mask"].float(), + denominator="n_valid_tokens", + ) + if "behave_imp_weight" in stat: + stats_tracker.denominator(unclipped_behave_tokens=stat["behave_mask"]) + stats_tracker.stat( + behave_imp_weight=stat["behave_imp_weight"], + behave_approx_kl=stat["behave_approx_kl"], + denominator="unclipped_behave_tokens", + ) + vocab_min_logits = logits.detach().min(-1).values.float() + vocab_max_logits = logits.detach().max(-1).values.float() + stats_tracker.stat( + vocab_min_logits=vocab_min_logits, + vocab_max_logits=vocab_max_logits, + denominator="n_tokens", + ) + + clip_mask = stat["clip_mask"] + clipped_new_logp = torch.where(clip_mask, logprobs.detach(), 0.0) + clipped_old_logp = torch.where(clip_mask, old_logp, 0.0) + stats_tracker.stat( + clipped_new_logp=clipped_new_logp, + clipped_old_logp=clipped_old_logp, + denominator="clipped_tokens", + ) + return loss diff --git a/arealite/engine/sft/lm_engine.py b/arealite/engine/sft/lm_engine.py index 1b1715da28..0a6bec250e 100644 --- a/arealite/engine/sft/lm_engine.py +++ b/arealite/engine/sft/lm_engine.py @@ -20,7 +20,7 @@ def train_lm(self, data: TensorDict): return self.engine.train_batch( input_=data, loss_fn=compute_packed_sft_loss, - loss_weight_fn=lambda x: x["prompt_mask"].logical_not().count_nonzero(), + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) def evaluate_lm(self, data): @@ -28,7 +28,7 @@ def evaluate_lm(self, data): self.engine.eval_batch( input_=data, loss_fn=compute_packed_sft_loss, - loss_weight_fn=lambda x: x["prompt_mask"].logical_not().count_nonzero(), + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) @@ -49,26 +49,26 @@ def compute_packed_sft_loss( ) -> torch.Tensor: packed_input_ids: torch.Tensor = input_["input_ids"] cu_seqlens: torch.Tensor = input_["cu_seqlens"] - prompt_mask = input_["prompt_mask"].bool() + loss_mask = input_["loss_mask"].bool() logprobs = gather_logprobs(logits, torch.roll(packed_input_ids, shifts=-1, dims=-1)) - prompt_mask = torch.roll(prompt_mask, shifts=-1, dims=-1) - logprobs = torch.where(prompt_mask, 0, logprobs) + loss_mask = torch.roll(loss_mask, shifts=-1, dims=-1) + logprobs = torch.where(loss_mask, logprobs, 0) - loss = -logprobs.sum() / prompt_mask.logical_not().count_nonzero() + loss = -logprobs.sum() / loss_mask.count_nonzero() with torch.no_grad(): seqlogp = torch.zeros( cu_seqlens.shape[0] - 1, device=logits.device, dtype=torch.float64 ) for i in range(cu_seqlens.shape[0] - 1): - m = prompt_mask[cu_seqlens[i] - i : cu_seqlens[i + 1] - i - 1] + m = loss_mask[cu_seqlens[i] - i : cu_seqlens[i + 1] - i - 1] logp = logprobs[cu_seqlens[i] - i : cu_seqlens[i + 1] - i - 1] assert cu_seqlens[i + 1] - i - 1 <= logprobs.shape[0], ( cu_seqlens, logprobs.shape, ) - seqlogp[i] = torch.where(m, 0.0, logp.detach()).sum() / ( + seqlogp[i] = torch.where(m, logp.detach(), 0.0).sum() / ( m.numel() - m.count_nonzero() ) @@ -78,8 +78,8 @@ def compute_packed_sft_loss( cu_seqlens.shape[0] - 1, dtype=torch.bool, device=logprobs.device ), n_tokens=torch.ones(logits.shape[0], dtype=torch.bool, device=logits.device), - n_valid_tokens=prompt_mask.logical_not(), - prompt_tokens=prompt_mask, + n_valid_tokens=loss_mask, + prompt_tokens=loss_mask.logical_not(), ) stats_tracker.stat(ppl=(-seqlogp).exp().float(), denominator="n_seqs") stats_tracker.stat(loss=-logprobs.detach(), denominator="n_valid_tokens") diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index c2401d386b..c985fa6301 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -4,12 +4,11 @@ import threading import time import traceback -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ProcessPoolExecutor from datetime import datetime from queue import Empty, Full, Queue -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List -import aiohttp import requests import torch.distributed as dist from tensordict import TensorDict @@ -24,6 +23,7 @@ RolloutStat, WeightUpdateMeta, ) +from arealite.utils.http import arequest_with_retry from arealite.utils.padding import concat_padded_tensors from realhf.base import logging, name_resolve, names, pkg_version @@ -98,10 +98,12 @@ def check_health(self, base_url): def initialize(self, addr: str | None, ft_spec: FinetuneSpec = None): self.rollout_tasks: Dict[str, asyncio.Task] = {} + self.executor = ProcessPoolExecutor(max_workers=1) self.rollout_thread = threading.Thread(target=self._rollout_thread) self.rollout_thread.start() def destroy(self): + self.executor.shutdown() self.exiting.set() self.rollout_thread.join() @@ -217,64 +219,6 @@ def choose_server(self) -> str: return server raise NotImplementedError("Only round-robin scheduling is implemented.") - async def arequest_with_retry( - self, - endpoint: str, - payload: Optional[Dict[str, Any]] = None, - method: str = "POST", - max_retries: Optional[int] = None, - timeout: Optional[float] = None, - retry_delay: float = 1.0, - target_addr: Optional[str] = None, - ) -> aiohttp.ClientResponse: - timeout = timeout or self.config.request_timeout - last_exception = None - max_retries = max_retries or self.config.request_retries - - # Try with retries - for _ in range(max_retries): - if target_addr: - addr = target_addr - else: - addr = self.choose_server() - base_url = f"http://{addr}" - url = f"{base_url}{endpoint}" - - for attempt in range(max_retries): - try: - async with aiohttp.ClientSession( - timeout=aiohttp.ClientTimeout( - total=timeout, - sock_connect=timeout, - ) - ) as session: - if method.upper() == "GET": - response = await session.get(url) - elif method.upper() == "POST": - response = await session.post(url, json=payload) - elif method.upper() == "PUT": - response = await session.put(url, json=payload) - elif method.upper() == "DELETE": - response = await session.delete(url) - else: - raise ValueError(f"Unsupported HTTP method: {method}") - - response.raise_for_status() - return await response.json() - - except ( - aiohttp.ClientError, - aiohttp.ClientResponseError, - asyncio.TimeoutError, - ) as e: - last_exception = e - if attempt < max_retries - 1: - await asyncio.sleep(retry_delay) - continue - raise RuntimeError( - f"Failed after {max_retries} retries each. " f"Last error: {last_exception}" - ) - async def agenerate(self, req: LLMRequest) -> LLMResponse: """Async version of generate using aiohttp.""" # Prepare request payload @@ -328,13 +272,13 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: and len(accumulated_output_tokens) < gconfig.max_new_tokens ): # loop until the generation is complete - result = await self.arequest_with_retry( + result = await arequest_with_retry( + addr=self.choose_server(), endpoint="/generate", payload=payload, method="POST", max_retries=3, timeout=self.config.request_timeout, - target_addr=server_addr, ) # Parse response @@ -370,56 +314,29 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: ttft=latency, # Simplified for non-streaming ) - def update_weights(self, meta): - executor = ThreadPoolExecutor(max_workers=1) - return executor.submit(self._update_weights, meta) - - def _update_weights(self, meta: WeightUpdateMeta): + def update_weights(self, meta: WeightUpdateMeta): if meta.type == "disk": # Update weights from disk - # Wait for model checkpoints of meta.version - update_name = names.update_weights_from_disk( - self.config.experiment_name, self.config.trial_name, meta.model_version + # Use ProcessPool to bypass python GIL for running async coroutines + fut = self.executor.submit( + update_weights_from_disk, + self.config.experiment_name, + self.config.trial_name, + meta.model_version, + self.addresses, + meta.path, + self.config.request_retries, + self.config.request_timeout, ) - save_timestamp = float(name_resolve.wait(update_name, timeout=120)) - load_timestamp = datetime.now().timestamp() - logger.info( - f"Begin update weights from {meta.path}, responded in {(load_timestamp - save_timestamp):.2f}s" - ) - try: - jobs = [ - self.aupdate_weights_from_disk(addr, meta.path) - for addr in self.addresses - ] - loop = asyncio.new_event_loop() - # asyncio event loop should be manually set when running asyncio stuff in another thread - asyncio.set_event_loop(loop) - loop.run_until_complete(asyncio.gather(*jobs)) - finally: - loop.close() - logger.info( - f"Loading weights done in {(datetime.now().timestamp() - load_timestamp):.2f}s" - ) - self.set_version(meta.model_version) + + def callback(fut): + self.set_version(meta.model_version) + + fut.add_done_callback(callback) + return fut else: raise NotImplementedError(f"Unsupported weight update type: {meta.type}") - async def aupdate_weights_from_disk(self, addr, path: str): - res = await self.arequest_with_retry( - endpoint="/update_weights_from_disk", - payload=dict(model_path=str(path), allow_interrupt=True), - method="POST", - max_retries=3, - timeout=self.config.request_timeout, - target_addr=addr, - ) - assert res["success"] - if "num_paused_requests" in res: - logger.info( - f"{res['num_paused_requests']} requests are interrupted " - f"during updating weights for server {addr}" - ) - def get_capacity(self): if dist.is_initialized(): world_size = dist.get_world_size() @@ -515,3 +432,58 @@ def pause(self): def resume(self): self.paused.clear() + + +async def aupdate_weights_from_disk( + addr, path: str, request_retries: int, request_timeout: float +): + res = await arequest_with_retry( + addr=addr, + endpoint="/update_weights_from_disk", + payload=dict(model_path=str(path), allow_interrupt=True), + method="POST", + max_retries=request_retries, + timeout=request_timeout, + ) + assert res["success"] + if "num_paused_requests" in res: + logger.info( + f"{res['num_paused_requests']} requests are interrupted " + f"during updating weights for server {addr}" + ) + + +def update_weights_from_disk( + experiment_name, + trial_name, + model_version, + addresses, + path, + request_retries, + request_timeout, +): + async def _fn(): + # Wait for model checkpoints of meta.version + update_name = names.update_weights_from_disk( + experiment_name, trial_name, model_version + ) + save_timestamp = float(name_resolve.wait(update_name, timeout=120)) + load_timestamp = datetime.now().timestamp() + logger.info( + f"Begin update weights from {path}, responded in {(load_timestamp - save_timestamp):.2f}s" + ) + jobs = [ + aupdate_weights_from_disk( + addr, + path=path, + request_retries=request_retries, + request_timeout=request_timeout, + ) + for addr in addresses + ] + await asyncio.gather(*jobs) + logger.info( + f"Loading weights done in {(datetime.now().timestamp() - load_timestamp):.2f}s" + ) + + return asyncio.run(_fn()) diff --git a/arealite/launcher/local.py b/arealite/launcher/local.py new file mode 100644 index 0000000000..447e7db04b --- /dev/null +++ b/arealite/launcher/local.py @@ -0,0 +1,307 @@ +import getpass +import os +import re +import signal as signal_module +import subprocess +import sys +import time +from collections import defaultdict +from typing import Dict, List, Optional, Tuple, Union + +import psutil + +from arealite.api.cli_args import SGLangConfig, parse_cli_args, to_structured_cfg +from arealite.api.io_struct import AllocationMode, AllocationType +from arealite.utils.network import find_free_ports, gethostip +from realhf.base import gpu_utils, logging, name_resolve, names +from realhf.scheduler.client import JobException, JobInfo, JobState + +logger = logging.getLogger("Local Scheduler") + +JOB_STATE_TO_PROCESS_STATUS = { + JobState.NOT_FOUND: [], + JobState.PENDING: [psutil.STATUS_PARKED], + JobState.RUNNING: [ + psutil.STATUS_RUNNING, + psutil.STATUS_SLEEPING, + psutil.STATUS_DISK_SLEEP, + psutil.STATUS_TRACING_STOP, + psutil.STATUS_WAKING, + psutil.STATUS_WAITING, + psutil.STATUS_LOCKED, + psutil.STATUS_IDLE, + ], + JobState.COMPLETED: [ + psutil.STATUS_DEAD, + psutil.STATUS_STOPPED, + psutil.STATUS_ZOMBIE, + ], + JobState.FAILED: [], + JobState.CANCELLED: [], +} + +PROCESS_STATUS_TO_JOB_STATE = {} +for job_state, process_statuses in JOB_STATE_TO_PROCESS_STATUS.items(): + for process_status in process_statuses: + PROCESS_STATUS_TO_JOB_STATE[process_status] = job_state + + +def terminate_process_and_children(pid: int, signal: Optional[Union[str, int]] = None): + if signal is None: + signal = signal_module.SIGKILL + if isinstance(signal, str): + signal = getattr(signal_module, signal) + try: + parent = psutil.Process(pid) + children = parent.children(recursive=True) + for child in children: + terminate_process_and_children(child.pid) + parent.send_signal(signal) + except psutil.NoSuchProcess: + pass + + +class LocalLauncher: + def __init__(self, experiment_name: str, trial_name: str, fileroot: str): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.fileroot = fileroot + + self._jobs: Dict[str, subprocess.Popen] = {} + self._job_counter: Dict[str, int] = defaultdict(int) + self._job_states = {} + + self._gpu_counter = 0 + self._cuda_devices: List[str] = os.environ.get( + "CUDA_VISIBLE_DEVICES", ",".join(map(str, range(gpu_utils.gpu_count()))) + ).split(",") + if len(self._cuda_devices) < 1: + raise RuntimeError( + f"Local mode can only run when there is at least one GPU. " + f"CUDA_VISIBLE_DEVICES is currently set to {os.environ['CUDA_VISIBLE_DEVICES']}." + ) + + @property + def run_name(self): + return f"{self.experiment_name}_{self.trial_name}" + + def log_path_of(self, job_name: str) -> str: + log_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(log_path, exist_ok=True) + return os.path.join(log_path, f"{job_name}.log") + + def __del__(self): + self.wait() + + def submit_array( + self, + job_name: str, + cmd: str | List[str], + count: int = 1, + gpu: int = 0, + env_vars: Optional[Dict] = None, + ): + if env_vars is None: + env_vars = {} + if not isinstance(cmd, list): + cmd = [cmd] * count + offset = self._job_counter[job_name] + for i in range(count): + if gpu > 0: + # Allocate GPUs in a round-robin manner + visible_devices = [] + for _ in range(gpu): + available_device_id = self._gpu_counter % len(self._cuda_devices) + self._gpu_counter += 1 + visible_devices.append(available_device_id) + env_vars["CUDA_VISIBLE_DEVICES"] = ",".join( + str(self._cuda_devices[j]) for j in visible_devices + ) + c = ( + " ".join(str(k) + "=" + str(v) for k, v in env_vars.items()) + + " stdbuf -oL " + + cmd[i] + ) + c = f"{c} | tee -a {self.log_path_of(job_name)}" + logger.info("Starting local process with command: %s", c) + process = subprocess.Popen(c, shell=isinstance(c, str)) + self._jobs[f"{job_name}/{offset + i}"] = process + self._job_counter[job_name] += 1 + + def submit( + self, + job_name: str, + cmd: str | List[str], + gpu: int = 0, + env_vars: Optional[Dict] = None, + ): + self.submit_array(job_name=job_name, cmd=cmd, gpu=gpu, env_vars=env_vars) + + def stop(self, job_name, signal=None): + assert any(k.startswith(job_name) for k in self._jobs) + keys = [k for k, p in self._jobs.items() if k.startswith(job_name)] + procs = [p for k, p in self._jobs.items() if k.startswith(job_name)] + logger.info( + f"Stopping local process with signal {signal if signal else 'SIGKILL'}, " + f"pid: {[p.pid for p in procs]}" + ) + for p in procs: + terminate_process_and_children(p.pid, signal=signal) + for p in procs: + p.wait() + for k, p in zip(keys, procs): + self._jobs.pop(k) + del p + + def stop_all(self, signal=None): + # signal argument is ignored in local stop_all + for name in self._job_counter: + self.stop(name, signal=signal) + + def find(self, job_name): + if job_name in self._jobs: + return JobInfo(name=job_name, state=JobState.RUNNING, host="localhost") + else: + return JobInfo(name=job_name, state=JobState.NOT_FOUND) + + def find_all(self, job_name_regex=".*"): + rs = [] + for name in self._jobs: + if re.fullmatch(job_name_regex, name): + rs.append(self.find(name)) + return rs + + def wait( + self, + timeout=None, + check_status: Tuple[JobState, ...] = ( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + ), + remove_status: Tuple[JobState, ...] = (JobState.COMPLETED,), + update=False, + ): + deadline = None if timeout is None else time.time() + timeout + logger.info( + "Waiting for %d local running processes, pids: %s", + len(self._jobs), + " ".join(str(job.pid) for job in self._jobs.values()), + ) + left = set(self._jobs.keys()) + num_jobs_left = len(left) + + while len(left) > 0: + to_remove = [] + if len(left) < num_jobs_left: + num_jobs_left = len(left) + logger.info(f"Waiting for {num_jobs_left} jobs.") + if deadline is not None and time.time() > deadline: + raise TimeoutError( + f"Timeout waiting for {self.run_name}: {', '.join(sorted(left))}" + ) + # update job states + for job_name in list(left): + job = self._jobs[job_name] + pid = job.pid + process = psutil.Process(pid) + self._job_states[job_name] = PROCESS_STATUS_TO_JOB_STATE.get( + process.status(), JobState.NOT_FOUND + ) + + for job_name in list(left): + state = self._job_states[job_name] + if state in check_status: + raise JobException( + run_name=self.run_name, + worker_type=job_name.split("/")[0], + host="local", + reason=state, + ) + if state in remove_status: + logger.info(f"Job {job_name} is {state}.(Removed)") + left.remove(job_name) + to_remove.append(job_name) + + if update: + for k in to_remove: + self._jobs.pop(k) + worker_type = k.split("/")[0] + assert worker_type in self._job_counter + self._job_counter[worker_type] -= 1 + if self._job_counter[worker_type] <= 0: + self._job_counter.pop(worker_type) + + time.sleep(2) + + +def main_local(): + cfg, _ = parse_cli_args(sys.argv[2:]) + name_resolve.reconfigure(cfg.cluster.name_resolve) + name_resolve.clear_subtree( + names.trial_root(experiment_name=cfg.experiment_name, trial_name=cfg.trial_name) + ) + alloc_mode = AllocationMode.from_str(cfg.allocation_mode) + + launcher = LocalLauncher(cfg.experiment_name, cfg.trial_name, cfg.cluster.fileroot) + + server_cmd = [] + server_addrs = [] + if alloc_mode.type_ == AllocationType.DECOUPLED_SGLANG: + base_seed = cfg.sglang.random_seed + cfg.sglang = to_structured_cfg(cfg.sglang, SGLangConfig) + ports = find_free_ports(alloc_mode.gen_dp_size * 2, port_range=(10000, 50000)) + host_ip = gethostip() + host = "localhost" if not cfg.sglang.enable_metrics else host_ip + for i in range(alloc_mode.gen_dp_size): + cfg.sglang.random_seed = base_seed + i + cmd = SGLangConfig.build_cmd( + cfg.sglang, + host=host, + tp_size=alloc_mode.gen_tp_size, + base_gpu_id=0, + port=ports[i * 2], + dist_init_addr=f"localhost:{ports[i*2+1]}", + ) + server_cmd.append(cmd) + server_addrs.append(f"{host}:{ports[i * 2]}") + else: + raise NotImplementedError() + + # Launch inference servers. + launcher.submit_array( + job_name="llm_server", + cmd=server_cmd, + count=alloc_mode.gen_dp_size, + gpu=alloc_mode.gen_pp_size * alloc_mode.gen_tp_size, + ) + logger.info( + f"LLM inference server launched at: AREAL_LLM_SERVER_ADDRS={','.join(server_addrs)}" + ) + + # Launch trainer entrypoint + if not cfg.server_only: + launcher.submit( + job_name="trainer", + cmd=f"torchrun --nnodes 1 --nproc-per-node {alloc_mode.train_world_size} --standalone {' '.join(sys.argv[1:])}", + gpu=alloc_mode.train_world_size, + env_vars=dict(AREAL_LLM_SERVER_ADDRS=",".join(server_addrs)), + ) + + try: + launcher.wait( + check_status=( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + JobState.COMPLETED, + ), + remove_status=(), + ) + except (KeyboardInterrupt, JobException, TimeoutError) as e: + launcher.stop_all("SIGTERM") + raise e + + +if __name__ == "__main__": + main_local() diff --git a/arealite/tests/test_rlvr_workflow.py b/arealite/tests/test_rlvr_workflow.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/arealite/tests/test_sglang_engine.py b/arealite/tests/test_sglang_engine.py index 62a1a1ff25..a2d5694bf1 100644 --- a/arealite/tests/test_sglang_engine.py +++ b/arealite/tests/test_sglang_engine.py @@ -150,7 +150,7 @@ def test_remote_sglang_staleness_control(sglang_server, bs, ofp, n_samples): engine.submit(data, workflow=workflow) # wait for some time - time.sleep(15) + time.sleep(5) assert engine.output_queue.qsize() == min(bs * 2, bs * (ofp + 1)) # Update model version @@ -161,7 +161,7 @@ def test_remote_sglang_staleness_control(sglang_server, bs, ofp, n_samples): for _ in range(bs * 2): engine.submit(data, workflow=workflow) # wait for some time - time.sleep(15) + time.sleep(5) assert engine.output_queue.qsize() == min(bs * 4, bs * (ofp + 2)) # exit @@ -204,6 +204,7 @@ def test_disk_update_weights_from_fsdp_engine(tmp_path_factory, sglang_server): config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" inf_engine = RemoteSGLangEngine(config) + inf_engine.initialize(None, None) # test update weights path = tmp_path_factory.mktemp("upload_weights_from_disk") update_weight_meta = WeightUpdateMeta( @@ -213,3 +214,4 @@ def test_disk_update_weights_from_fsdp_engine(tmp_path_factory, sglang_server): engine.upload_weights(update_weight_meta) future.result() assert inf_engine.get_version() == 100 + inf_engine.destroy() diff --git a/arealite/tests/test_utils.py b/arealite/tests/test_utils.py index de341f2e0e..35281c6d05 100644 --- a/arealite/tests/test_utils.py +++ b/arealite/tests/test_utils.py @@ -28,7 +28,7 @@ def mock_padded_data(): ans_len = int(ans_len) seq = dict( input_ids=torch.randint(0, VOCAB_SIZE, size=(prompt_len + ans_len,)), - prompt_mask=torch.tensor([1] * prompt_len + [0] * ans_len), + loss_mask=torch.tensor([0] * prompt_len + [1] * ans_len), logprobs=torch.randn(prompt_len + ans_len), position_ids=torch.arange(prompt_len + ans_len), ) diff --git a/arealite/utils/device.py b/arealite/utils/device.py new file mode 100644 index 0000000000..840ca33d99 --- /dev/null +++ b/arealite/utils/device.py @@ -0,0 +1,31 @@ +from typing import Tuple + +import torch +import torch.distributed as dist + +from realhf.base import logging + +logger = logging.getLogger(__file__) + + +def _get_current_mem_info(unit: str = "GB", precision: int = 2) -> Tuple[str]: + """Get current memory usage.""" + assert unit in ["GB", "MB", "KB"] + divisor = 1024**3 if unit == "GB" else 1024**2 if unit == "MB" else 1024 + mem_allocated = torch.cuda.memory_allocated() + mem_reserved = torch.cuda.memory_reserved() + mem_free, mem_total = torch.cuda.mem_get_info() + mem_used = mem_total - mem_free + mem_allocated = f"{mem_allocated / divisor:.{precision}f}" + mem_reserved = f"{mem_reserved / divisor:.{precision}f}" + mem_used = f"{mem_used / divisor:.{precision}f}" + mem_total = f"{mem_total / divisor:.{precision}f}" + return mem_allocated, mem_reserved, mem_used, mem_total + + +# Adapted from verl +def log_gpu_stats(head: str, rank: int = 0): + if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank): + mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info() + message = f"{head}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, device memory used/total (GB): {mem_used}/{mem_total}" + logger.info(msg=message) diff --git a/arealite/utils/functional.py b/arealite/utils/functional.py index 0532958833..9ce736c97c 100644 --- a/arealite/utils/functional.py +++ b/arealite/utils/functional.py @@ -1,8 +1,175 @@ +from typing import Dict, Optional, Tuple + +import numpy as np import torch +import torch.distributed as dist @torch.compile -def gather_logprobs(logits: torch.Tensor, labels: torch.Tensor): - log_probs = torch.nn.functional.log_softmax(logits.float(), dim=-1) +def _gather_logprobs( + logits: torch.Tensor, labels: torch.Tensor, temperature: float = 1.0 +): + log_probs = torch.nn.functional.log_softmax(logits.float() / temperature, dim=-1) log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) return log_probs_labels + + +@torch.compile +def _gather_logprobs_entropy( + logits: torch.Tensor, labels: torch.Tensor, temperature: float = 1.0 +): + log_probs = torch.nn.functional.log_softmax(logits.float() / temperature, dim=-1) + entropy = -torch.sum(log_probs.exp() * log_probs, dim=-1) + log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) + return log_probs_labels, entropy + + +def gather_logprobs( + logits: torch.Tensor, + labels: torch.Tensor, + temperature: float = 1.0, + chunk_size: int = 1024, +): + batch_size = logits.shape[0] + + if batch_size <= chunk_size: + return _gather_logprobs(logits, labels, temperature) + + log_probs_labels_list = [] + + for i in range(0, batch_size, chunk_size): + end_idx = min(i + chunk_size, batch_size) + chunk_logits = logits[i:end_idx] + chunk_labels = labels[i:end_idx] + + chunk_log_probs = _gather_logprobs(chunk_logits, chunk_labels, temperature) + + log_probs_labels_list.append(chunk_log_probs) + + return torch.cat(log_probs_labels_list) + + +def gather_logprobs_entropy( + logits: torch.Tensor, + labels: torch.Tensor, + temperature: float = 1.0, + chunk_size: int = 1024, +): + batch_size = logits.shape[0] + + if batch_size <= chunk_size: + return _gather_logprobs_entropy(logits, labels, temperature) + + log_probs_labels_list = [] + entropy_list = [] + + for i in range(0, batch_size, chunk_size): + end_idx = min(i + chunk_size, batch_size) + chunk_logits = logits[i:end_idx] + chunk_labels = labels[i:end_idx] + + chunk_log_probs, chunk_entropy = _gather_logprobs_entropy( + chunk_logits, chunk_labels, temperature + ) + + log_probs_labels_list.append(chunk_log_probs) + entropy_list.append(chunk_entropy) + + return torch.cat(log_probs_labels_list), torch.cat(entropy_list) + + +@torch.no_grad() +def masked_normalization( + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + dim=None, + unbiased=False, + eps=1e-5, + high_precision=True, + all_reduce=True, + reduce_group=None, +): + dtype = torch.float64 if high_precision else torch.float32 + x = x.to(dtype) + if dim is None: + dim = tuple(range(len(x.shape))) + if mask is None: + factor = torch.tensor( + np.prod([x.shape[d] for d in dim]), dtype=dtype, device=x.device + ) + else: + mask = mask.to(dtype) + x = x * mask + factor = mask.sum(dim, keepdim=True) + x_sum = x.sum(dim=dim, keepdim=True) + x_sum_sq = x.square().sum(dim=dim, keepdim=True) + if dist.is_initialized() and all_reduce: + dist.all_reduce(factor, op=dist.ReduceOp.SUM, group=reduce_group) + dist.all_reduce(x_sum, op=dist.ReduceOp.SUM, group=reduce_group) + dist.all_reduce( + x_sum_sq, + op=dist.ReduceOp.SUM, + group=reduce_group, + ) + mean = x_sum / factor + meansq = x_sum_sq / factor + var = meansq - mean**2 + if unbiased: + var *= factor / (factor - 1) + return ((x - mean) / (var.sqrt() + eps)).float() + + +def ppo_actor_loss_fn( + logprobs: torch.Tensor, + old_logprobs: torch.Tensor, + advantages: torch.Tensor, + eps_clip: float, + loss_mask: torch.Tensor, + c_clip: Optional[float] = None, + proximal_logprobs: Optional[torch.Tensor] = None, + behav_imp_weight_cap: Optional[float] = None, +) -> Tuple[torch.Tensor, Dict]: + denorm_logprobs = ( + proximal_logprobs if proximal_logprobs is not None else old_logprobs + ) + loss_mask_count = loss_mask.count_nonzero() or 1 + ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) + clipped_ratio = torch.clamp(ratio, 1.0 - eps_clip, 1.0 + eps_clip) + pg_loss1 = -advantages * ratio + pg_loss2 = -advantages * clipped_ratio + clip_mask = pg_loss1.detach() < pg_loss2.detach() + pg_loss = torch.max(pg_loss1, pg_loss2) + if c_clip is not None: + assert c_clip > 1.0, c_clip + pg_loss3 = torch.sign(advantages) * c_clip * advantages + dual_clip_mask = pg_loss3.detach() < pg_loss.detach() + pg_loss = torch.min(pg_loss, pg_loss3) + else: + dual_clip_mask = torch.zeros_like(clip_mask) + if proximal_logprobs is not None: + behav_kl = proximal_logprobs - old_logprobs + behav_imp_weight = behav_kl.exp() + behav_mask = ( + (behav_imp_weight <= behav_imp_weight_cap).logical_and(loss_mask) + if behav_imp_weight_cap is not None + else loss_mask + ) + behav_kl = torch.where(behav_mask, behav_kl, 0.0) + behav_imp_weight = torch.where(behav_mask, behav_imp_weight, 0.0) + pg_loss = pg_loss * behav_imp_weight + logging_loss = pg_loss.detach() + pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count + clip_mask.logical_and_(loss_mask) + dual_clip_mask.logical_and_(loss_mask) + stat = dict( + loss=logging_loss, + importance_weight=ratio.detach(), + approx_kl=(logprobs - denorm_logprobs).detach(), + clip_mask=clip_mask, + dual_clip_mask=dual_clip_mask, + ) + if proximal_logprobs is not None: + stat["behave_imp_weight"] = behav_imp_weight + stat["behave_approx_kl"] = behav_kl + stat["behave_mask"] = behav_mask + return pg_loss, stat diff --git a/arealite/utils/http.py b/arealite/utils/http.py new file mode 100644 index 0000000000..29e4df4050 --- /dev/null +++ b/arealite/utils/http.py @@ -0,0 +1,56 @@ +import asyncio +from typing import Any, Dict, Optional + +import aiohttp + +DEFAULT_RETRIES = 1 +DEFAULT_REQUEST_TIMEOUT = 3600 + + +async def arequest_with_retry( + addr: str, + endpoint: str, + payload: Optional[Dict[str, Any]] = None, + method: str = "POST", + max_retries: Optional[int] = None, + timeout: Optional[float] = None, + retry_delay: float = 1.0, +) -> aiohttp.ClientResponse: + timeout = timeout or DEFAULT_REQUEST_TIMEOUT + last_exception = None + max_retries = max_retries or DEFAULT_RETRIES + base_url = f"http://{addr}" + url = f"{base_url}{endpoint}" + + for attempt in range(max_retries): + try: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout( + total=timeout, + sock_connect=timeout, + ) + ) as session: + if method.upper() == "GET": + response = await session.get(url) + elif method.upper() == "POST": + response = await session.post(url, json=payload) + elif method.upper() == "PUT": + response = await session.put(url, json=payload) + elif method.upper() == "DELETE": + response = await session.delete(url) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + response.raise_for_status() + return await response.json() + except ( + aiohttp.ClientError, + aiohttp.ClientResponseError, + asyncio.TimeoutError, + ) as e: + last_exception = e + if attempt < max_retries - 1: + await asyncio.sleep(retry_delay) + continue + raise RuntimeError( + f"Failed after {max_retries} retries each. " f"Last error: {last_exception}" + ) diff --git a/arealite/workflow/rlvr.py b/arealite/workflow/rlvr.py index a72d7b68cf..026f574925 100644 --- a/arealite/workflow/rlvr.py +++ b/arealite/workflow/rlvr.py @@ -42,8 +42,8 @@ async def arun_episode(self, engine, data): results = [] for resp in resps: seq = resp.input_tokens + resp.output_tokens - logprobs = [0] * resp.input_len + resp.output_logprobs - prompt_mask = [1] * resp.input_len + [0] * resp.output_len + logprobs = [0.0] * resp.input_len + resp.output_logprobs + loss_mask = [0] * resp.input_len + [1] * resp.output_len versions = [-1] * resp.input_len + resp.output_versions reward = self.reward_fn( @@ -56,10 +56,10 @@ async def arun_episode(self, engine, data): res = dict( # unsqueeze to add an additional batch dimension input_ids=torch.tensor(seq).unsqueeze(0), - prompt_mask=torch.tensor(prompt_mask).unsqueeze(0), + loss_mask=torch.tensor(loss_mask).unsqueeze(0), logprobs=torch.tensor(logprobs).unsqueeze(0), versions=torch.tensor(versions).unsqueeze(0), - attention_mask=torch.ones(len(seq)).unsqueeze(0), + attention_mask=torch.ones(len(seq), dtype=torch.bool).unsqueeze(0), # reward rewards=torch.tensor([reward]), ) diff --git a/examples/arealite/configs/gsm8k_grpo.yaml b/examples/arealite/configs/gsm8k_grpo.yaml new file mode 100644 index 0000000000..ac7488571a --- /dev/null +++ b/examples/arealite/configs/gsm8k_grpo.yaml @@ -0,0 +1,129 @@ +experiment_name: gsm8k-grpo +trial_name: trial0 +allocation_mode: sglang.d4p1t1+d4p1t1 +n_nodes: 1 +n_gpus_per_node: 8 +cluster: + fileroot: /tmp/arealite/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve +seed: 1 +total_train_epochs: 10 +tokenizer_path: ${actor.path} +async_training: true + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 4 + enable_rollout_tracing: false + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 512 + greedy: false + temperature: 1.0 + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: /storage/openpsi/models/Qwen__Qwen2-1.5B-Instruct/ + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: false + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + optimizer: + type: adam + lr: 2e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + backend: fsdp + + group_size: ${gconfig.n_samples} + group_adv_norm: false + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + behav_imp_weight_cap: 5.0 + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + optimizer: null + backend: fsdp + +# SGLang +server_only: false +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.9 + +# datasets +train_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + +valid_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +checkpointer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled \ No newline at end of file diff --git a/examples/arealite/gsm8k_grpo.py b/examples/arealite/gsm8k_grpo.py new file mode 100644 index 0000000000..07434a09ba --- /dev/null +++ b/examples/arealite/gsm8k_grpo.py @@ -0,0 +1,259 @@ +import os +import re +import sys + +import torch +import torch.distributed as dist +from datasets import Dataset, load_dataset +from datasets.distributed import split_dataset_by_node +from torchdata.stateful_dataloader import StatefulDataLoader + +from arealite.api.cli_args import GRPOConfig, load_expr_config +from arealite.api.io_struct import FinetuneSpec, WeightUpdateMeta +from arealite.engine.ppo.actor import FSDPPPOActor +from arealite.engine.sglang_remote import RemoteSGLangEngine +from arealite.utils.device import log_gpu_stats +from arealite.utils.evaluator import Evaluator +from arealite.utils.saver import Saver +from arealite.utils.stats_logger import StatsLogger +from arealite.workflow.rlvr import RLVRWorkflow +from realhf.api.core.data_api import load_hf_tokenizer +from realhf.base import stats_tracker + + +def process_gsm8k_rl_dataset(dataset: Dataset): + def process(sample): + messages = [{"role": "user", "content": sample["question"]}] + return {"messages": messages} + + dataset = dataset.map(process).remove_columns(["question"]) + return dataset + + +def get_gsm8k_dataset(split, rank, world_size): + dataset = load_dataset(path="openai/gsm8k", name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + return process_gsm8k_rl_dataset(dataset) + + +# Adapted from verl. +def extract_solution(solution_str, method="strict") -> str | None: + assert method in ["strict", "flexible"] + + final_answer = None + if method == "strict": + # this also tests the formatting of the model + solutions = re.findall("#### (\\-?[0-9\\.\\,]+)", solution_str) + if len(solutions) == 0: + final_answer = None + else: + # take the last solution + final_answer = solutions[-1].replace(",", "").replace("$", "") + elif method == "flexible": + answer = re.findall("(\\-?[0-9\\.\\,]+)", solution_str) + final_answer = None + if len(answer) == 0: + # no reward is there is no answer + pass + else: + invalid_str = ["", "."] + # find the last number that is not '.' + for final_answer in reversed(answer): + if final_answer not in invalid_str: + break + return final_answer + + +def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs): + from realhf.impl.dataset.math_parser import extract_answer + + sol = extract_answer(completions, data_name="math") + ans = extract_solution(solution_str=answer, method="strict") + if sol is None: + return 0 + if ans is None: + return 0 + return int(sol.strip() == ans.strip()) + + +def main_grpo(): + config, _ = load_expr_config(sys.argv[1:], GRPOConfig) + config: GRPOConfig + + rank = int(os.getenv("RANK")) + world_size = int(os.getenv("WORLD_SIZE")) + tokenizer = load_hf_tokenizer(config.tokenizer_path) + + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + get_gsm8k_dataset("train", rank, world_size), + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.train_dataset.drop_last, + ) + valid_dataloader = StatefulDataLoader( + get_gsm8k_dataset("test", rank, world_size), + batch_size=config.valid_dataset.batch_size // world_size, + shuffle=config.valid_dataset.shuffle, + num_workers=config.valid_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.valid_dataset.drop_last, + ) + ft_spec = FinetuneSpec( + total_train_epochs=config.total_train_epochs, + dataset_size=len(train_dataloader) * config.train_dataset.batch_size, + train_batch_size=config.train_dataset.batch_size, + ) + + # Initialize inference engine + rollout = RemoteSGLangEngine(config.rollout) + rollout.initialize(None, ft_spec) + eval_rollout = RemoteSGLangEngine(config.rollout) + eval_rollout.initialize(None, ft_spec) + # NOTE: set a large version such that eval does not have any offpolicyness control + eval_rollout.set_version(int(1e12)) + + # Initialize train engine + actor = FSDPPPOActor(config=config.actor) + actor.initialize(None, ft_spec) + ref = None + if config.actor.kl_ctl > 0 and config.ref is not None: + ref = FSDPPPOActor(config=config.ref) + ref.initialize(None, ft_spec) + + # Create rollout workflow + if tokenizer.pad_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.pad_token_id) + if tokenizer.eos_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.eos_token_id) + workflow = RLVRWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + enable_thinking=False, + ) + + # Run training. + saver = Saver(config.saver, ft_spec, for_recover=False) + logger = StatsLogger(config.stats_logger, ft_spec) + evaluator = Evaluator(config.evaluator, ft_spec) + + total_epochs = config.total_train_epochs + steps_per_epoch = len(train_dataloader) + max_steps = total_epochs * steps_per_epoch + + logger.info(f"total_epochs={total_epochs} step_per_epoch={steps_per_epoch}") + data_generator = iter(train_dataloader) + for global_step in range(max_steps): + epoch = global_step // steps_per_epoch + step = global_step % steps_per_epoch + + with stats_tracker.record_timing("rollout"): + if config.async_training: + batch = rollout.prepare_batch( + data_generator, + train_dataloader, + workflow=workflow, + ) + else: + try: + data = next(data_generator) + except StopIteration: + data_generator = iter(train_dataloader) + data = next(data_generator) + batch = rollout.rollout_batch(data, workflow=workflow) + + batch = batch.to(actor.device) + # Create barrier to synchronize all rollout processes. + dist.barrier() + torch.cuda.synchronize() + + if config.actor.recompute_logprob or config.actor.use_decoupled_loss: + with stats_tracker.record_timing("recompute_logp"): + logp = actor.compute_logp(batch) + batch["prox_logp"] = logp + log_gpu_stats("recompute logp") + + if ref is not None: + with stats_tracker.record_timing("ref_logp"): + batch["ref_logp"] = ref.compute_logp(batch) + log_gpu_stats("ref logp") + + with stats_tracker.record_timing("compute_advantage"): + actor.compute_advantages(batch) + log_gpu_stats("compute advantages") + + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("grpo_actor"), + ): + stats = actor.ppo_update(batch) + actor.step_lr_scheduler() + log_gpu_stats("ppo update") + + with stats_tracker.record_timing("update_weights"): + meta = WeightUpdateMeta( + type="disk", + path=os.path.join( + Saver.get_save_checkpoint_root(config.saver), + "update_weights", + str(global_step), + ), + alloc_mode=None, + comm_backend=None, + model_version=global_step + 1, + ) + if dist.get_rank() == 0: + future = rollout.update_weights(meta) + actor.upload_weights(meta) + if dist.get_rank() == 0: + future.result() + rollout.set_version(global_step + 1) + dist.barrier() + + with stats_tracker.record_timing("save"): + saver.save(actor, epoch, step, global_step) + + with stats_tracker.record_timing("eval"): + + def evaluate_fn(): + rollout.pause() + cnt = 0 + for data in valid_dataloader: + for item in data: + eval_rollout.submit(item, workflow) + cnt += 1 + batch = eval_rollout.wait(cnt, timeout=None) + rewards = batch["rewards"].float().to(actor.device) + with stats_tracker.scope("grpo-eval"): + stats_tracker.denominator( + n_seqs=torch.ones( + rewards.shape[0], + device=rewards.device, + dtype=torch.bool, + ) + ) + stats_tracker.stat(task_reward=rewards, denominator="n_seqs") + rollout.resume() + + evaluator.evaluate( + evaluate_fn, + epoch, + step, + global_step, + ) + + logger.commit(epoch, step, global_step, stats) + + logger.close() + eval_rollout.destroy() + rollout.destroy() + if ref is not None: + ref.destroy() + actor.destroy() + + +if __name__ == "__main__": + main_grpo() diff --git a/examples/arealite/gsm8k_sft.py b/examples/arealite/gsm8k_sft.py index 70c67e3207..c1d8735127 100644 --- a/examples/arealite/gsm8k_sft.py +++ b/examples/arealite/gsm8k_sft.py @@ -1,7 +1,6 @@ import os import sys -import torch.distributed as dist from datasets import Dataset, load_dataset from datasets.distributed import split_dataset_by_node from torchdata.stateful_dataloader import StatefulDataLoader @@ -23,10 +22,8 @@ def process(sample): sample["question"] + sample["answer"] + tokenizer.eos_token ) prompt_token = tokenizer.encode(sample["question"]) - prompt_mask = [1] * len(prompt_token) + [0] * ( - len(seq_token) - len(prompt_token) - ) - return {"input_ids": seq_token, "prompt_mask": prompt_mask} + loss_mask = [0] * len(prompt_token) + [1] * (len(seq_token) - len(prompt_token)) + return {"input_ids": seq_token, "loss_mask": loss_mask} dataset = dataset.map(process).remove_columns(["question", "answer"]) return dataset From b2bd639ac7beee6687f49eb804818b5b60d1d2cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Wed, 16 Jul 2025 16:26:52 +0800 Subject: [PATCH 06/16] PullRequest: 368 [lite] Refactor train engine after merging contributions from GitHub MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-train-engine of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/368 Reviewed-by: 晓雷 * . * . --- arealite/api/cli_args.py | 8 +- arealite/api/engine_api.py | 10 +- arealite/api/io_struct.py | 25 +- arealite/engine/autotp_engine.py | 128 +++++ arealite/engine/base_hf_engine.py | 360 ++++++++++++++ arealite/engine/fsdp_engine.py | 315 +----------- arealite/engine/hf_engine.py | 467 ------------------ arealite/engine/sft/lm_engine.py | 6 +- arealite/tests/test_sglang_local_engine.py | 4 - .../{test_engine.py => test_train_engine.py} | 13 +- 10 files changed, 540 insertions(+), 796 deletions(-) create mode 100644 arealite/engine/autotp_engine.py create mode 100644 arealite/engine/base_hf_engine.py delete mode 100644 arealite/engine/hf_engine.py rename arealite/tests/{test_engine.py => test_train_engine.py} (92%) diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 6dc197b307..906f8fdebb 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -18,7 +18,7 @@ class MicroBatchSpec: """Specification for splitting micro-batches during training.""" - n_mbs: int = field( + n_mbs: Optional[int] = field( default=1, metadata={ "help": "Number of micro-batches (or minimum number if max_tokens_per_mb is set). Used when max_tokens_per_mb is None or as minimum count", @@ -161,7 +161,7 @@ class FSDPEngineConfig: @dataclass -class HFEngineConfig: +class DeepSpeedAutoTPEngineConfig: autotp_size: Optional[int] = field( default=1, metadata={"help": "DeepSpeed AutoTP size"}, @@ -201,7 +201,9 @@ class TrainEngineConfig: ) backend: str = "" fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) - hf: HFEngineConfig = field(default_factory=HFEngineConfig) + ds_auto_tp: DeepSpeedAutoTPEngineConfig = field( + default_factory=DeepSpeedAutoTPEngineConfig + ) @dataclass diff --git a/arealite/api/engine_api.py b/arealite/api/engine_api.py index 6334bf2248..9fb24d222e 100644 --- a/arealite/api/engine_api.py +++ b/arealite/api/engine_api.py @@ -25,10 +25,10 @@ class Scheduling: cpu: int gpu: int mem: int - nodelist: str = None - exclude: str = None - partition: str = None - container_image: str = None + nodelist: Optional[str] = None + exclude: Optional[str] = None + partition: Optional[str] = None + container_image: Optional[str] = None env_vars: Dict[str, str] = field(default_factory=dict) # time utils from "https://slurm.schedmd.com/sbatch.html" time_limit: Optional[str] = None # see "--time" option for format @@ -105,7 +105,7 @@ def eval_batch( def forward( self, input_: TensorDict, - output_seqlens: List[List[int]] | None = None, + output_seqlens: List[int] | None = None, post_hook: Callable[[torch.Tensor, TensorDict], Any] | None = None, aggregate_fn: Callable[[List[Any]], Any] = torch.cat, ) -> Any | None: diff --git a/arealite/api/io_struct.py b/arealite/api/io_struct.py index 83d2cc63a7..28a369fd93 100644 --- a/arealite/api/io_struct.py +++ b/arealite/api/io_struct.py @@ -71,7 +71,7 @@ class AllocationType(enum.Enum): @dataclass class AllocationMode: type_: AllocationType - parallel_strat: None | Dict[str, Dict[str, int]] + parallel_strat: Dict[str, Dict[str, int]] @property def gen_tp_size(self) -> int: @@ -115,7 +115,7 @@ def from_str(cls, allocation_mode: str): raise NotImplementedError(f"Failed to parse allocation: {allocation_mode}") @staticmethod - def extract_3d_alloc(allocation_mode: str) -> Dict | None: + def extract_parallelism_strategy(allocation_mode: str) -> Dict: for x, y, z in itertools.permutations(["d", "t", "p"]): pattern = rf"{x}(\d+){y}(\d+){z}(\d+)" m = re.match(pattern, allocation_mode) @@ -130,29 +130,28 @@ def extract_3d_alloc(allocation_mode: str) -> Dict | None: z: c, } } + raise ValueError( + f"Unknown how to resolve parallelism strategy: {allocation_mode}" + ) @staticmethod - def extract_decoupled_alloc(allocation_mode: str) -> Dict | None: + def extract_decoupled_alloc(allocation_mode: str) -> Dict: pattern = re.compile( r"(?:(?:vllm|sglang)\.(.+?)\+(.+))|(?:(.+?)\+(?:vllm|sglang)\.(.+))" ) m = pattern.match(allocation_mode) if not m: - return + raise ValueError( + f"Unknown how to resolve decoupled allocation: {allocation_mode}" + ) if m.group(1): gen_alloc = m.group(1) other_alloc = m.group(2) else: gen_alloc = m.group(4) other_alloc = m.group(3) - gen_alloc = AllocationMode.extract_3d_alloc(gen_alloc) - if not gen_alloc: - return - other_alloc = AllocationMode.extract_3d_alloc( - other_alloc - ) or AllocationMode.extract_key_value_alloc(other_alloc) - if not other_alloc: - return + gen_alloc = AllocationMode.extract_parallelism_strategy(gen_alloc) + other_alloc = AllocationMode.extract_parallelism_strategy(other_alloc) other_alloc.update({"gen": gen_alloc["*"]}) return other_alloc @@ -171,7 +170,7 @@ class SaveLoadMeta: path: str weight_format: str with_optim: bool - tokenizer: PreTrainedTokenizerFast | None + tokenizer: Optional[PreTrainedTokenizerFast] base_model_path: str | None naive_distributed: bool = False diff --git a/arealite/engine/autotp_engine.py b/arealite/engine/autotp_engine.py new file mode 100644 index 0000000000..e068b9c1f7 --- /dev/null +++ b/arealite/engine/autotp_engine.py @@ -0,0 +1,128 @@ +import os + +import torch +import torch.distributed as dist +from safetensors.torch import save_file + +from arealite.api.cli_args import TrainEngineConfig +from arealite.api.engine_api import FinetuneSpec, SaveLoadMeta, WeightUpdateMeta +from arealite.engine.base_hf_engine import BaseHFEngine +from arealite.utils.save_load import ( + get_state_dict_from_repo_id_or_path, + is_existing_local_path, +) +from realhf.base import constants, logging + +logger = logging.getLogger("DeepSpeedAutoTPEngine") + + +class DeepSpeedAutoTPEngine(BaseHFEngine): + def __init__(self, config: TrainEngineConfig): + super().__init__(config) + + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): + """Initialize distributed communication and model.""" + assert ( + addr is None + ), "DeepSpeedAutoTPEngine does not support remote initialization." + import deepspeed + + self.create_process_group() + + world_size = int(os.environ.get("WORLD_SIZE")) + deepspeed.init_distributed( + dist_backend="nccl", + world_size=world_size, + timeout=constants.NCCL_DEFAULT_TIMEOUT, + ) + self.create_device_model() + # NOTE: the device context manager does not work here. + self.model = deepspeed.tp_model_init( + self.model, + tp_size=self.config.ds_auto_tp.autotp_size, + dtype=getattr(torch, self.config.dtype), + ).to(self.device) + self.create_optimizer(ft_spec) + self.initialized = True + + def _check_autotp(self): + tp_size = self.config.ds_auto_tp.autotp_size + config = self.model_config + num_attention_heads = config.num_attention_heads + num_key_value_heads = config.num_key_value_heads + hidden_size = config.hidden_size + intermediate_size = config.intermediate_size + + return ( + num_attention_heads % tp_size == 0 + and num_key_value_heads % tp_size == 0 + and hidden_size % tp_size == 0 + and intermediate_size % tp_size == 0 + ) + + def save(self, meta: SaveLoadMeta): + if meta.weight_format != "naive_distributed": + raise ValueError(f"Unknown weight format {meta.weight_format}. ") + if self.model is None: + raise RuntimeError("Model not initialized") + + rank = dist.get_rank() + world_size = dist.get_world_size() + if rank == 0: + os.makedirs(meta.path, exist_ok=True) + self.model_config.save_pretrained( + meta.path, + ) + if meta.tokenizer is not None: + meta.tokenizer.save_pretrained( + meta.path, + ) + + state_dict = self.model.state_dict() + if hasattr(self.model, "module"): + state_dict = { + k.replace("module.", "", 1) if k.startswith("module.") else k: v.cpu() + for k, v in state_dict.items() + } + else: + state_dict = {k: v.cpu() for k, v in state_dict.items()} + + # Only support store parameters from model partitions respectively + gathered_state_dicts = None + if rank == 0: + gathered_state_dicts = [None for _ in range(world_size)] + dist.gather_object( + obj=state_dict, object_gather_list=gathered_state_dicts, dst=0 + ) + if rank == 0: + for i, state_dict in enumerate(gathered_state_dicts): + save_file(state_dict, f"{meta.path}/rank_{i:02d}_model.safetensors") + if meta.with_optim: + self.save_optimizer_state(meta.path) + + def load(self, meta: SaveLoadMeta): + if meta.weight_format != "naive_distributed": + raise ValueError(f"Unknown weight format {meta.weight_format}. ") + rank = dist.get_rank() + # Only support load full model parameters from huggingface + # and load model partition locally + if rank == 0 or is_existing_local_path(meta.path): + path = f"{meta.path}/rank_{rank:02d}_model.safetensors" + full_state = get_state_dict_from_repo_id_or_path(meta.path) + + if hasattr(self.model, "module") and not hasattr(full_state): + full_state = { + f"module.{k}" if not k.startswith("module.") else k: v + for k, v in full_state.items() + } + self.model.load_state_dict( + full_state, strict=not self.model_config.tie_word_embeddings + ) + if self.model_config.tie_word_embeddings: + self.model.tie_weights() + + if meta.with_optim: + self.load_optimizer_state(meta.path) + + def upload_weights(self, meta: WeightUpdateMeta): + raise ValueError(f"update weight not implemented {meta.type}") diff --git a/arealite/engine/base_hf_engine.py b/arealite/engine/base_hf_engine.py new file mode 100644 index 0000000000..bcbda8ba27 --- /dev/null +++ b/arealite/engine/base_hf_engine.py @@ -0,0 +1,360 @@ +import gc +import os +import time +from typing import Any, Callable, Dict, List + +import torch +import torch.distributed as dist +from tensordict import TensorDict +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + PretrainedConfig, + PreTrainedTokenizerFast, + get_constant_schedule_with_warmup, + get_linear_schedule_with_warmup, +) + +from arealite.api.cli_args import TrainEngineConfig +from arealite.api.engine_api import FinetuneSpec, TrainEngine +from arealite.utils.data import ( + MicroBatchList, + amend_position_ids, + pack_tensor_dict, + pad_and_stack_tensors_along_first_dim, + pad_mb_list, + reorder_list, + split_padded_tensor_dict_into_mb_list, + unpack_sequence, + unsqueeze_mb_list, +) +from arealite.utils.fsdp import get_cosine_schedule_with_warmup +from arealite.utils.model import disable_dropout_in_model +from realhf.api.core.data_api import load_hf_tokenizer +from realhf.base import constants, logging + +logger = logging.getLogger("Base HF Engine") + + +class BaseHFEngine(TrainEngine): + def __init__(self, config: TrainEngineConfig): + self.config = config + self.optimizer_config = config.optimizer + + self.model: torch.nn.Module + self.optimizer: torch.optim.Optimizer + self.tokenizer: PreTrainedTokenizerFast + # huggingface model config + self.model_config: PretrainedConfig + + # initialization + self.initialized = False + self.own_global_group = False + self._parallelism_group: dist.ProcessGroup + self.weight_update_group_initialized = False + + self.world_size = int(os.environ["WORLD_SIZE"]) + + def train(self, mode: bool = True): + assert self.model is not None + self.model.train(mode=mode) + return self + + @property + def parallelism_group(self) -> dist.ProcessGroup: + assert self.initialized + return self._parallelism_group + + def create_process_group(self): + if not dist.is_initialized(): + # TODO: Handle the condition when WORLD_SIZE and RANK is not set in launcher + dist.init_process_group( + backend="nccl", + timeout=constants.NCCL_DEFAULT_TIMEOUT, + device_id=torch.device(int(os.environ["LOCAL_RANK"])), + ) + self.own_global_group = True + self._parallelism_group = dist.new_group() + + def create_device_model(self): + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + self.device = torch.device(int(os.environ["LOCAL_RANK"])) + + dtype = getattr(torch, self.config.dtype) + self.model_config = AutoConfig.from_pretrained( + pretrained_model_name_or_path=self.config.path, + trust_remote_code=True, + ) + self.tokenizer = load_hf_tokenizer(self.config.path) + tik = time.perf_counter() + with torch.device("cuda"): + if self.config.init_from_scratch: + # initialize scratch model from config + # NOTE: VLM cannot directly load state dict using this + # random initialized model, so otherwise we call + # from_pretrained rather than loading weights into this random model. + model = AutoModelForCausalLM.from_config( + self.model_config, + torch_dtype=dtype, + attn_implementation=self.config.attn_impl, + ) + else: + model = AutoModelForCausalLM.from_pretrained( + pretrained_model_name_or_path=self.config.path, + trust_remote_code=True, + torch_dtype=dtype, + attn_implementation=self.config.attn_impl, + ) + if self.config.disable_dropout: + disable_dropout_in_model(model) + if self.config.gradient_checkpointing: + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + logger.info(f"Model creation and loading time: {time.perf_counter() - tik}") + self.model = model + + def create_optimizer(self, ft_spec: FinetuneSpec): + if self.optimizer_config is None: + return + assert self.model is not None + # Set up optimizer + tik = time.perf_counter() + assert ( + self.optimizer_config.type == "adam" + ), "Only AdamW optimizer is supported in this engine." + lr = self.optimizer_config.lr + weight_decay = self.optimizer_config.weight_decay + beta1 = self.optimizer_config.beta1 + beta2 = self.optimizer_config.beta2 + eps = self.optimizer_config.eps + + self.optimizer = torch.optim.AdamW( + self.model.parameters(), + lr=lr, + weight_decay=weight_decay, + betas=(beta1, beta2), + eps=eps, + ) + total_train_steps = ft_spec.total_train_steps + num_warmup_steps = int( + self.optimizer_config.warmup_steps_proportion * total_train_steps + ) + + if self.optimizer_config.lr_scheduler_type == "cosine": + self.lr_scheduler = get_cosine_schedule_with_warmup( + self.optimizer, + num_warmup_steps, + total_train_steps, + min_lr_ratio=self.optimizer_config.min_lr_ratio, + ) + elif self.optimizer_config.lr_scheduler_type == "linear": + self.lr_scheduler = get_linear_schedule_with_warmup( + self.optimizer, + num_warmup_steps, + total_train_steps, + ) + elif self.optimizer_config.lr_scheduler_type == "constant": + self.lr_scheduler = get_constant_schedule_with_warmup( + self.optimizer, + num_warmup_steps, + ) + else: + raise ValueError( + f"Unknown lr scheduler type {self.optimizer_config.lr_scheduler_type}" + ) + logger.info(f"Create optimizer time: {time.perf_counter() - tik}") + + def destroy(self): + """Destroy the engine and release GPU memory.""" + del self.optimizer + del self.model + gc.collect() + torch.cuda.empty_cache() + gc.collect() + dist.destroy_process_group(self.parallelism_group) + if self.own_global_group: + dist.destroy_process_group() + self.initialized = False + + def save_optimizer_state(self, path: str): + # Save FSDP sharded state dict on each rank + assert self.optimizer is not None + assert dist.is_initialized() + rank = dist.get_rank() + shard_path = os.path.join( + path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" + ) + state_dict = self.optimizer.state_dict() + torch.save(state_dict, shard_path) + dist.barrier() + + def load_optimizer_state(self, path: str): + # Load FSDP sharded state dict + assert self.optimizer is not None + assert dist.is_initialized() + rank = dist.get_rank() + shard_path = os.path.join( + path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" + ) + optimizer_state_dict = torch.load(shard_path, weights_only=False) + self.optimizer.load_state_dict(optimizer_state_dict) + dist.barrier() + + def step_lr_scheduler(self): + assert self.lr_scheduler is not None + self.lr_scheduler.step() + + def prepare_mb_list(self, input_: TensorDict) -> MicroBatchList: + assert "attention_mask" in input_ and "input_ids" in input_ + if isinstance(input_, dict): + input_ = TensorDict(input_, batch_size=[input_["input_ids"].shape[0]]) + input_ = amend_position_ids(input_) + mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec) + logger.info( + f"Microbatch #tokens (rank {dist.get_rank()}): {mb_list.group_lens}" + ) + mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] + mb_list = pad_mb_list(mb_list, pad_value=0.0) + # NOTE: We unsqueeze here because huggingface transformer models requires + # packed input to be of shape [1, total_seqlen]. + mb_list = unsqueeze_mb_list(mb_list) + # FIXME: the resulting max_seqlen is a tensor rather than an integer + for mb in mb_list.mbs: + mb["max_seqlen"] = int(mb["max_seqlen"]) + mb["use_cache"] = False + for mb in mb_list.padded_mbs: + mb["max_seqlen"] = int(mb["max_seqlen"]) + mb["use_cache"] = False + return mb_list + + def train_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> Dict[str, float]: + """Train on a batch using gradient accumulation.""" + input_ = input_.to(self.device) + assert self.optimizer is not None + assert self.optimizer_config is not None + assert self.lr_scheduler is not None + + self.optimizer.zero_grad() + mb_list = self.prepare_mb_list(input_) + + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 + ) + assert total_loss_weight != 0 + dist.all_reduce(total_loss_weight) + + # Process microbatches with gradient accumulation + for i, (pad_length, padded_mb_input, mb_input) in enumerate( + zip(mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs) + ): + outputs = self.model(**padded_mb_input) + + logits = outputs.logits.squeeze(0) + logits = logits[:-pad_length] if pad_length > 0 else logits + loss = loss_fn(logits, mb_input) + loss_scale = loss_weight_fn(mb_input) / total_loss_weight + + # Scale loss for accumulation + # Revert gradient averaging across dp ranks + loss_scale *= self.world_size + + loss *= loss_scale + loss.backward() + + grad_norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + self.optimizer_config.gradient_clipping, + norm_type=2.0, + error_if_nonfinite=False, + foreach=None, + ) + if not torch.isfinite(grad_norm): + self.optimizer.zero_grad() + update_successful = False + else: + self.optimizer.step() + update_successful = True + + current_lr = self.lr_scheduler.get_last_lr()[0] + # Optimizer step + self.optimizer.step() + return dict( + update_successful=float(update_successful), + grad_norm=float(grad_norm) if grad_norm is not None else float("nan"), + lr=current_lr, + ) + + @torch.no_grad() + def eval_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> torch.Tensor | None: + """Evaluate on a batch.""" + input_ = input_.to(self.device) + mb_list = self.prepare_mb_list(input_) + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 + ) + assert total_loss_weight != 0 + + total_loss = 0.0 + total_weight = 0.0 + + for pad_length, padded_mb_input, mb_input in zip( + mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs + ): + outputs = self.model(**padded_mb_input) + logits = outputs.logits.squeeze(0) + logits = logits[:-pad_length] if pad_length > 0 else logits + loss = loss_fn(logits, mb_input) + + # Simple weight calculation (could be improved) + loss_scale = loss_weight_fn(mb_input) / total_loss_weight + total_loss += loss.item() * loss_scale + total_weight += loss_scale + + return torch.tensor(total_loss / total_weight) + + @torch.no_grad() + def forward( + self, + input_: TensorDict, + output_seqlens: List[int] | None = None, + post_hook: Callable[[torch.Tensor, TensorDict], Any] | None = None, + aggregate_fn: Callable[[List[Any]], Any] = torch.cat, + ) -> Any | None: + """Forward pass with optional post-processing.""" + input_ = input_.to(self.device) + cu_seqlens = pack_tensor_dict(input_)["cu_seqlens"] + mb_list = self.prepare_mb_list(input_) + + if output_seqlens is None: + output_seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().numpy().tolist() + + results = [] + for pad_length, padded_mb_input, mb_input in zip( + mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs + ): + outputs = self.model(**padded_mb_input) + logits = outputs.logits.squeeze(0) + logits = logits[:-pad_length] if pad_length > 0 else logits + + if post_hook: + result = post_hook(logits, mb_input) + results.append(result) + else: + results.append(logits) + + res = aggregate_fn(results) + output_seqlens = [output_seqlens[i] for i in mb_list.forward_indices] + unpacked = unpack_sequence(res, lens=output_seqlens, dim=0) + reordered = reorder_list(unpacked, mb_list.backward_indices) + return pad_and_stack_tensors_along_first_dim(reordered) diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index 07efc36fae..34131ec9b7 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -1,42 +1,20 @@ -import gc import os import time from datetime import datetime -from typing import Any, Callable, Dict, List, Optional +from typing import Callable, Dict, Optional import torch import torch.distributed as dist -import transformers from tensordict import TensorDict from torch.distributed.checkpoint.state_dict import ( StateDictOptions, get_model_state_dict, ) -from transformers import ( - AutoConfig, - AutoModelForCausalLM, - get_constant_schedule_with_warmup, - get_linear_schedule_with_warmup, -) +from transformers import PreTrainedTokenizerFast from arealite.api.cli_args import TrainEngineConfig -from arealite.api.engine_api import ( - FinetuneSpec, - SaveLoadMeta, - TrainEngine, - WeightUpdateMeta, -) -from arealite.utils.data import ( - MicroBatchList, - amend_position_ids, - pack_tensor_dict, - pad_and_stack_tensors_along_first_dim, - pad_mb_list, - reorder_list, - split_padded_tensor_dict_into_mb_list, - unpack_sequence, - unsqueeze_mb_list, -) +from arealite.api.engine_api import FinetuneSpec, SaveLoadMeta, WeightUpdateMeta +from arealite.engine.base_hf_engine import BaseHFEngine from arealite.utils.fsdp import ( CPUOffloadPolicy, MixedPrecisionPolicy, @@ -44,108 +22,35 @@ create_fsdp_device_mesh, fsdp2_clip_grad_norm_, fsdp2_load_full_state_dict, - get_cosine_schedule_with_warmup, ) -from arealite.utils.model import disable_dropout_in_model from arealite.utils.save_load import get_state_dict_from_repo_id_or_path -from realhf.api.core.data_api import load_hf_tokenizer -from realhf.base import constants, logging, name_resolve, names, pkg_version +from realhf.base import logging, name_resolve, names, pkg_version logger = logging.getLogger("FSDPEngine") -class FSDPEngine(TrainEngine): +class FSDPEngine(BaseHFEngine): def __init__(self, config: TrainEngineConfig): - self.config = config - self.optimizer_config = config.optimizer - - self.model = None - self.optimizer = None - self.tokenizer = None - # huggingface model config - self.model_config = None + super().__init__(config) # FSDP options self.mixed_precision_policy = None self.device_mesh = None self.cpu_offload = None - # initialization - self.initialized = False - self.own_global_group = False - self._parallelism_group = None - self.weight_update_group_initialized = False - - # TODO: Handle the case when WORLD_SIZE is not set in launcher - self.world_size = int(os.environ["WORLD_SIZE"]) - - def train(self, mode: bool = True): - assert self.model is not None - self.model.train(mode=mode) - return self - - @property - def parallelism_group(self) -> dist.ProcessGroup: - assert self.initialized - return self._parallelism_group def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): # Initialize distributed enviroments and load model. assert addr is None, "FSDPEngine does not support remote initialization." - assert pkg_version.is_version_greater_or_equal( "torch", "2.4.0" ), f"arealite only supports FSDP2, which requires torch>=2.4.0" - """Initialize distributed communication and model.""" - if not dist.is_initialized(): - # TODO: Handle the condition when WORLD_SIZE and RANK is not set in launcher - dist.init_process_group( - backend="nccl", - timeout=constants.NCCL_DEFAULT_TIMEOUT, - device_id=torch.device(int(os.environ["LOCAL_RANK"])), - ) - self.own_global_group = True - self._parallelism_group = dist.new_group() - - # TODO: Handle the condition when LOCAL_RANK is not set in launcher - torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) - self.device = torch.device(int(os.environ["LOCAL_RANK"])) - - dtype = getattr(torch, self.config.dtype) - self.model_config = AutoConfig.from_pretrained( - pretrained_model_name_or_path=self.config.path, - trust_remote_code=True, - ) - self.tokenizer = load_hf_tokenizer(self.config.path) - tik = time.perf_counter() - with torch.device("cuda"): - if self.config.init_from_scratch: - # initialize scratch model from config - # NOTE: VLM cannot directly load state dict using this - # random initialized model, so otherwise we call - # from_pretrained rather than loading weights into this random model. - model = AutoModelForCausalLM.from_config( - self.model_config, - torch_dtype=dtype, - attn_implementation=self.config.attn_impl, - ) - else: - model = AutoModelForCausalLM.from_pretrained( - pretrained_model_name_or_path=self.config.path, - trust_remote_code=True, - torch_dtype=dtype, - attn_implementation=self.config.attn_impl, - ) - if self.config.disable_dropout: - disable_dropout_in_model(model) - if self.config.gradient_checkpointing: - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - logger.info(f"Model creation and loading time: {time.perf_counter() - tik}") + self.create_process_group() + self.create_device_model() + # Wrap with FSDP2 # Simple auto wrap policy self.mixed_precision_policy = MixedPrecisionPolicy( - param_dtype=dtype, + param_dtype=getattr(torch, self.config.dtype), reduce_dtype=torch.float32, cast_forward_inputs=True, ) @@ -154,82 +59,19 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): self.cpu_offload = ( CPUOffloadPolicy() if self.config.fsdp.offload_params else None ) - fsdp_kwargs = { "mesh": self.device_mesh, "mp_policy": self.mixed_precision_policy, "offload_policy": self.cpu_offload, "reshard_after_forward": True, } - - # Wrap with FSDP2 tik = time.perf_counter() - apply_fsdp2(model, fsdp_kwargs, self.config.fsdp.wrap_policy) + apply_fsdp2(self.model, fsdp_kwargs, self.config.fsdp.wrap_policy) logger.info(f"Applying FSDP2 time: {time.perf_counter() - tik}") - self.model = model - - # Set up optimizer - if self.optimizer_config is not None: - tik = time.perf_counter() - assert ( - self.optimizer_config.type == "adam" - ), "Only AdamW optimizer is supported in this engine." - lr = self.optimizer_config.lr - weight_decay = self.optimizer_config.weight_decay - beta1 = self.optimizer_config.beta1 - beta2 = self.optimizer_config.beta2 - eps = self.optimizer_config.eps - - self.optimizer = torch.optim.AdamW( - self.model.parameters(), - lr=lr, - weight_decay=weight_decay, - betas=(beta1, beta2), - eps=eps, - ) - total_train_steps = ft_spec.total_train_steps - num_warmup_steps = int( - self.optimizer_config.warmup_steps_proportion * total_train_steps - ) - - if self.optimizer_config.lr_scheduler_type == "cosine": - self.lr_scheduler = get_cosine_schedule_with_warmup( - self.optimizer, - num_warmup_steps, - total_train_steps, - min_lr_ratio=self.optimizer_config.min_lr_ratio, - ) - elif self.optimizer_config.lr_scheduler_type == "linear": - self.lr_scheduler = get_linear_schedule_with_warmup( - self.optimizer, - num_warmup_steps, - total_train_steps, - ) - elif self.optimizer_config.lr_scheduler_type == "constant": - self.lr_scheduler = get_constant_schedule_with_warmup( - self.optimizer, - num_warmup_steps, - ) - else: - raise ValueError( - f"Unknown lr scheduler type {self.optimizer_config.lr_scheduler_type}" - ) - logger.info(f"Create optimizer time: {time.perf_counter() - tik}") + self.create_optimizer(ft_spec) self.initialized = True - def destroy(self): - """Destroy the engine and release GPU memory.""" - self.model = None - self.optimizer = None - gc.collect() - torch.cuda.empty_cache() - gc.collect() - dist.destroy_process_group(self.parallelism_group) - if self.own_global_group: - dist.destroy_process_group() - self.initialized = False - def save(self, meta: SaveLoadMeta): if meta.weight_format == "hf": self._save_model_to_hf(meta.path, meta.tokenizer) @@ -240,7 +82,7 @@ def save(self, meta: SaveLoadMeta): raise ValueError(f"Unknown weight format {meta.weight_format}. ") if meta.with_optim: - self._save_optimizer_state(meta.path) + self.save_optimizer_state(meta.path) def load(self, meta: SaveLoadMeta): if meta.weight_format == "hf": @@ -252,34 +94,10 @@ def load(self, meta: SaveLoadMeta): raise ValueError(f"Unknown weight format {meta.weight_format}. ") if meta.with_optim: - self._load_optimizer_state(meta.path) - - def _save_optimizer_state(self, path: str): - # Save FSDP sharded state dict on each rank - assert self.optimizer is not None - assert dist.is_initialized() - rank = dist.get_rank() - shard_path = os.path.join( - path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" - ) - state_dict = self.optimizer.state_dict() - torch.save(state_dict, shard_path) - dist.barrier() - - def _load_optimizer_state(self, path: str): - # Load FSDP sharded state dict - assert self.optimizer is not None - assert dist.is_initialized() - rank = dist.get_rank() - shard_path = os.path.join( - path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" - ) - optimizer_state_dict = torch.load(shard_path, weights_only=False) - self.optimizer.load_state_dict(optimizer_state_dict) - dist.barrier() + self.load_optimizer_state(meta.path) def _save_model_to_hf( - self, path: str, tokenizer: Optional[transformers.PreTrainedTokenizerFast] + self, path: str, tokenizer: Optional[PreTrainedTokenizerFast] ): """Save model in HuggingFace format.""" if self.model is None: @@ -345,35 +163,11 @@ def _update_weights_from_distributed(self): "Distributed weight update is not implemented for FSDPEngine yet. " ) - def step_lr_scheduler(self): - assert self.lr_scheduler is not None - self.lr_scheduler.step() - - def _prepare_mb_list(self, input_: TensorDict) -> MicroBatchList: - assert "attention_mask" in input_ and "input_ids" in input_ - if isinstance(input_, dict): - input_ = TensorDict(input_, batch_size=[input_["input_ids"].shape[0]]) - input_ = amend_position_ids(input_) - mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec) - logger.info( - f"Microbatch #tokens (rank {dist.get_rank()}): {mb_list.group_lens}" - ) - mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] - mb_list = pad_mb_list(mb_list, pad_value=0.0) - # NOTE: We unsqueeze here because huggingface transformer models requires - # packed input to be of shape [1, total_seqlen]. - mb_list = unsqueeze_mb_list(mb_list) - # FIXME: the resulting max_seqlen is a tensor rather than an integer - for mb in mb_list.mbs: - mb["max_seqlen"] = int(mb["max_seqlen"]) - mb["use_cache"] = False - return mb_list - def train_batch( self, input_: TensorDict, - loss_fn: Callable[[torch.Tensor, Dict], torch.Tensor], - loss_weight_fn: Callable[[Dict], float], + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], ) -> Dict[str, float]: """Train on a batch using gradient accumulation.""" input_ = input_.to(self.device) @@ -382,7 +176,7 @@ def train_batch( assert self.lr_scheduler is not None self.optimizer.zero_grad() - mb_list = self._prepare_mb_list(input_) + mb_list = self.prepare_mb_list(input_) total_loss_weight = torch.tensor( sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 @@ -394,7 +188,6 @@ def train_batch( for i, (pad_length, padded_mb_input, mb_input) in enumerate( zip(mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs) ): - self.model.set_is_last_backward(i == len(mb_list.mbs) - 1) outputs = self.model(**padded_mb_input) logits = outputs.logits.squeeze(0) @@ -409,6 +202,7 @@ def train_batch( loss *= loss_scale loss.backward() + # NOTE: grad norm clip function is different grad_norm = fsdp2_clip_grad_norm_( self.model.parameters(), max_norm=self.optimizer_config.gradient_clipping ) @@ -427,72 +221,3 @@ def train_batch( grad_norm=float(grad_norm) if grad_norm is not None else float("nan"), lr=current_lr, ) - - @torch.no_grad() - def eval_batch( - self, - input_: TensorDict, - loss_fn: Callable[[torch.Tensor, Dict], torch.Tensor], - loss_weight_fn: Callable[[Dict], float], - ) -> torch.Tensor | None: - """Evaluate on a batch.""" - input_ = input_.to(self.device) - mb_list = self._prepare_mb_list(input_) - total_loss_weight = torch.tensor( - sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 - ) - assert total_loss_weight != 0 - - total_loss = 0.0 - total_weight = 0.0 - - for pad_length, padded_mb_input, mb_input in zip( - mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs - ): - outputs = self.model(**padded_mb_input) - logits = outputs.logits.squeeze(0) - logits = logits[:-pad_length] if pad_length > 0 else logits - loss = loss_fn(logits, mb_input) - - # Simple weight calculation (could be improved) - loss_scale = loss_weight_fn(mb_input) / total_loss_weight - total_loss += loss.item() * loss_scale - total_weight += loss_scale - - return torch.tensor(total_loss / total_weight) - - @torch.no_grad() - def forward( - self, - input_: TensorDict, - output_seqlens: List[int] | None = None, - post_hook: Callable[[torch.Tensor, Dict], Any] | None = None, - aggregate_fn: Callable[[List[Any]], Any] = torch.cat, - ) -> Any | None: - """Forward pass with optional post-processing.""" - input_ = input_.to(self.device) - cu_seqlens = pack_tensor_dict(input_)["cu_seqlens"] - mb_list = self._prepare_mb_list(input_) - - if output_seqlens is None: - output_seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().numpy().tolist() - - results = [] - for pad_length, padded_mb_input, mb_input in zip( - mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs - ): - outputs = self.model(**padded_mb_input) - logits = outputs.logits.squeeze(0) - logits = logits[:-pad_length] if pad_length > 0 else logits - - if post_hook: - result = post_hook(logits, mb_input) - results.append(result) - else: - results.append(logits) - - res = aggregate_fn(results) - output_seqlens = [output_seqlens[i] for i in mb_list.forward_indices] - unpacked = unpack_sequence(res, lens=output_seqlens, dim=0) - reordered = reorder_list(unpacked, mb_list.backward_indices) - return pad_and_stack_tensors_along_first_dim(reordered) diff --git a/arealite/engine/hf_engine.py b/arealite/engine/hf_engine.py deleted file mode 100644 index ee93f9f59e..0000000000 --- a/arealite/engine/hf_engine.py +++ /dev/null @@ -1,467 +0,0 @@ -import gc -import os -import time -from typing import Any, Callable, Dict, List, Optional - -import torch -import torch.distributed as dist -import transformers -from safetensors.torch import save_file -from tensordict import TensorDict -from transformers import ( - AutoConfig, - AutoModelForCausalLM, - get_constant_schedule_with_warmup, - get_linear_schedule_with_warmup, -) - -from arealite.api.cli_args import TrainEngineConfig -from arealite.api.engine_api import ( - FinetuneSpec, - SaveLoadMeta, - TrainEngine, - WeightUpdateMeta, -) -from arealite.utils.data import ( - MicroBatchList, - amend_position_ids, - pack_tensor_dict, - pad_and_stack_tensors_along_first_dim, - pad_mb_list, - reorder_list, - split_packed_tensor_dict_into_mb_list, - unpack_sequence, - unsqueeze_mb_list, -) -from arealite.utils.fsdp import get_cosine_schedule_with_warmup -from arealite.utils.save_load import ( - get_state_dict_from_repo_id_or_path, - is_existing_local_path, -) -from realhf.api.core.data_api import load_hf_tokenizer -from realhf.base import logging, name_resolve, names - -logger = logging.getLogger("HFEngine") - - -class HFEngine(TrainEngine): - def __init__(self, config: TrainEngineConfig): - self.config = config - self.optimizer_config = config.optimizer - - self.model = None - self.optimizer = None - self.tokenizer = None - # huggingface model config - self.model_config = None - # initialization - self.initialized = False - self.weight_update_group_initialized = False - - def train(self, mode: bool = True): - assert self.model is not None - self.model.train(mode=mode) - return self - - def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): - """Initialize distributed communication and model.""" - assert addr is None, "HFEngine does not support remote initialization." - - world_size = int(os.environ.get("WORLD_SIZE", 0)) - if not dist.is_initialized() and world_size > 1: - try: - import deepspeed - except ImportError: - print( - "Warning: deepspeed is not installed. Some functionality may be disabled." - ) - deepspeed.init_distributed(dist_backend="nccl", world_size=world_size) - - local_rank = int(os.environ.get("LOCAL_RANK", 0)) - torch.cuda.set_device(local_rank) - self.device = torch.device(f"cuda:{local_rank}") - - dtype = getattr(torch, self.config.dtype) - self.model_config = AutoConfig.from_pretrained( - pretrained_model_name_or_path=self.config.path, - trust_remote_code=True, - ) - self.tokenizer = load_hf_tokenizer(self.config.path) - - self.model = AutoModelForCausalLM.from_config( - self.model_config, - torch_dtype=dtype, - attn_implementation=self.config.attn_impl, - ).to(f"cuda:{local_rank}") - - if not self.config.init_from_scratch: - # Load model from a initial checkpoint path, - # which should only be a huggingface checkpoint. - load_meta = SaveLoadMeta( - path=self.config.path, - weight_format="hf", - with_optim=False, - tokenizer=None, - base_model_path=self.config.path, - naive_distributed=False, - ) - - self.load(load_meta) - - if world_size > 1: - if self._check_autotp(): - self.model = deepspeed.tp_model_init( - self.model, tp_size=self.config.hf.autotp_size, dtype=dtype - ) - else: - raise RuntimeError("DeepSpeed AutoTP configuration error in HFEngine. ") - - # Set up optimizer - if self.optimizer_config is not None: - assert ( - self.optimizer_config.type == "adam" - ), "Only AdamW optimizer is supported in this engine." - lr = self.optimizer_config.lr - weight_decay = self.optimizer_config.weight_decay - beta1 = self.optimizer_config.beta1 - beta2 = self.optimizer_config.beta2 - eps = self.optimizer_config.eps - - self.optimizer = torch.optim.AdamW( - self.model.parameters(), - lr=lr, - weight_decay=weight_decay, - betas=(beta1, beta2), - eps=eps, - ) - total_train_steps = ft_spec.total_train_steps - num_warmup_steps = int( - self.optimizer_config.warmup_steps_proportion * total_train_steps - ) - - if self.optimizer_config.lr_scheduler_type == "cosine": - self.lr_scheduler = get_cosine_schedule_with_warmup( - self.optimizer, - num_warmup_steps, - total_train_steps, - min_lr_ratio=self.optimizer_config.min_lr_ratio, - ) - elif self.optimizer_config.lr_scheduler_type == "linear": - self.lr_scheduler = get_linear_schedule_with_warmup( - self.optimizer, - num_warmup_steps, - total_train_steps, - ) - elif self.optimizer_config.lr_scheduler_type == "constant": - self.lr_scheduler = get_constant_schedule_with_warmup( - self.optimizer, - num_warmup_steps, - ) - else: - raise ValueError( - f"Unknown lr scheduler type {self.optimizer_config.lr_scheduler_type}" - ) - - self.initialized = True - - def _check_autotp(self): - tp_size = self.config.hf.autotp_size - config = self.model_config - num_attention_heads = config.num_attention_heads - num_key_value_heads = config.num_key_value_heads - hidden_size = config.hidden_size - intermediate_size = config.intermediate_size - - return ( - num_attention_heads % tp_size == 0 - and num_key_value_heads % tp_size == 0 - and hidden_size % tp_size == 0 - and intermediate_size % tp_size == 0 - ) - - def destroy(self): - """Destroy the engine and release GPU memory.""" - self.model = None - self.optimizer = None - gc.collect() - torch.cuda.empty_cache() - gc.collect() - self.initialized = False - - def save(self, meta: SaveLoadMeta): - if meta.weight_format == "hf": - self._save_model_to_hf(meta.path, meta.tokenizer, meta.naive_distributed) - elif meta.weight_format == "dcp": - # TODO: implement DCP save/load for HF - raise NotImplementedError("DCP format saving is not implemented yet. ") - else: - raise ValueError(f"Unknown weight format {meta.weight_format}. ") - - if meta.with_optim: - self._save_optimizer_state(meta.path) - - def load(self, meta: SaveLoadMeta): - if meta.weight_format == "hf": - self._load_model_from_hf(meta.path, meta.naive_distributed) - elif meta.weight_format == "dcp": - # TODO: implement DCP save/load for HF - raise NotImplementedError("DCP format loading is not implemented yet. ") - else: - raise ValueError(f"Unknown weight format {meta.weight_format}. ") - - if meta.with_optim: - self._load_optimizer_state(meta.path) - - def _save_optimizer_state(self, path: str): - assert self.optimizer is not None - os.makedirs(path, exist_ok=True) - torch.save(self.optimizer.state_dict(), os.path.join(path, "optim.pt")) - - def _load_optimizer_state(self, path: str): - assert self.optimizer is not None - path = os.path.join(path, "optim.pt") - optimizer_state_dict = torch.load(path, weights_only=False) - self.optimizer.load_state_dict(optimizer_state_dict) - - def _save_model_to_hf( - self, - path: str, - tokenizer: Optional[transformers.PreTrainedTokenizerFast], - naive_distributed: bool, - ): - """Save model in HuggingFace format.""" - if self.model is None: - raise RuntimeError("Model not initialized") - - rank = dist.get_rank() - world_size = dist.get_world_size() - if rank == 0: - os.makedirs(path, exist_ok=True) - self.model_config.save_pretrained(path) - if tokenizer is not None: - tokenizer.save_pretrained(path) - - if world_size > 1: - dist.barrier() - - state_dict = self.model.state_dict() - - if hasattr(self.model, "module"): - state_dict = { - k.replace("module.", "", 1) if k.startswith("module.") else k: v.cpu() - for k, v in state_dict.items() - } - else: - state_dict = {k: v.cpu() for k, v in state_dict.items()} - - if world_size > 1 and naive_distributed: - # Only support store parameters from model partitions respectively - gathered_state_dicts = None - if rank == 0: - gathered_state_dicts = [None for _ in range(world_size)] - - dist.gather_object( - obj=state_dict, object_gather_list=gathered_state_dicts, dst=0 - ) - - if rank == 0: - for i, state_dict in enumerate(gathered_state_dicts): - save_file(state_dict, f"{path}/rank_{i:02d}_model.safetensors") - else: - self.model.save_pretrained(path, state_dict=state_dict) - - if world_size > 1: - dist.barrier() - - def _load_model_from_hf(self, path: str, naive_distributed: bool): - """Load model from HuggingFace format.""" - - rank = dist.get_rank() - # Only support load full model parameters from huggingface - # and load model partition locally - if rank == 0 or is_existing_local_path(path): - if naive_distributed: - path = f"{path}/rank_{rank:02d}_model.safetensors" - full_state = get_state_dict_from_repo_id_or_path(path) - - if hasattr(self.model, "module") and not hasattr(full_state): - full_state = { - f"module.{k}" if not k.startswith("module.") else k: v - for k, v in full_state.items() - } - self.model.load_state_dict( - full_state, strict=not self.model_config.tie_word_embeddings - ) - - if self.model_config.tie_word_embeddings: - self.model.tie_weights() - - def upload_weights(self, meta: WeightUpdateMeta): - if meta.type == "nccl": - if not self.weight_update_group_initialized: - self._init_distributed_weight_update(meta) - self._update_weights_from_distributed() - elif meta.type == "disk": - self._save_model_to_hf(meta.path, self.tokenizer, meta.naive_distributed) - update_name = names.update_weights_from_disk( - self.config.experiment_name, - self.config.trial_name, - meta.model_version, - ) - name_resolve.add(update_name, str(time.time_ns()), keepalive_ttl=120) - else: - raise ValueError(f"Unknown weight update type {meta.type}") - - def _init_distributed_weight_update(self, meta: WeightUpdateMeta): - raise NotImplementedError( - "Distributed weight update is not implemented for HFEngine yet. " - ) - - def _update_weights_from_distributed(self): - raise NotImplementedError( - "Distributed weight update is not implemented for HFEngine yet. " - ) - - def step_lr_scheduler(self): - assert self.lr_scheduler is not None - return self.lr_scheduler.step() - - def _prepare_mb_list(self, input_: TensorDict) -> MicroBatchList: - assert "attention_mask" in input_ and "input_ids" in input_ - if isinstance(input_, dict): - input_ = TensorDict(input_, batch_size=[input_["input_ids"].shape[0]]) - input_ = amend_position_ids(input_) - packed_input = pack_tensor_dict(input_) - mb_list = split_packed_tensor_dict_into_mb_list( - packed_input, - self.config.mb_spec, - ) - mb_list = pad_mb_list(mb_list, pad_value=0.0) - # NOTE: We unsqueeze here because huggingface transformer models requires - # packed input to be of shape [1, total_seqlen]. - mb_list = unsqueeze_mb_list(mb_list) - return mb_list - - def train_batch( - self, - input_: TensorDict, - loss_fn: Callable[[torch.Tensor, Dict], torch.Tensor], - loss_weight_fn: Callable[[Dict], float], - ) -> Dict[str, float]: - """Train on a batch using gradient accumulation.""" - input_ = input_.to(self.device) - assert self.optimizer is not None - assert self.optimizer_config is not None - assert self.lr_scheduler is not None - - self.optimizer.zero_grad() - mb_list = self._prepare_mb_list(input_) - - total_loss_weight = torch.tensor( - sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 - ) - assert total_loss_weight != 0 - - # Process microbatches with gradient accumulation - for pad_length, padded_mb_input, mb_input in zip( - mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs - ): - outputs = self.model(**padded_mb_input) - - logits = outputs.logits.squeeze(0) - logits = logits[:-pad_length] if pad_length > 0 else logits - loss = loss_fn(logits, mb_input) - loss_scale = loss_weight_fn(mb_input) / total_loss_weight - - loss *= loss_scale - loss.backward() - - grad_norm = torch.nn.utils.clip_grad_norm_( - self.model.parameters(), - self.optimizer_config.gradient_clipping, - norm_type=2.0, - error_if_nonfinite=False, - foreach=None, - ) - if not torch.isfinite(grad_norm): - self.optimizer.zero_grad() - update_successful = False - else: - self.optimizer.step() - update_successful = True - - current_lr = self.lr_scheduler.get_last_lr()[0] - # Optimizer step - self.optimizer.step() - return dict( - update_successful=float(update_successful), - grad_norm=float(grad_norm) if grad_norm is not None else float("nan"), - lr=current_lr, - ) - - @torch.no_grad() - def eval_batch( - self, - input_: TensorDict, - loss_fn: Callable[[torch.Tensor, Dict], torch.Tensor], - loss_weight_fn: Callable[[Dict], float], - ) -> torch.Tensor | None: - """Evaluate on a batch.""" - mb_list = self._prepare_mb_list(input_) - total_loss_weight = torch.tensor( - sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 - ) - assert total_loss_weight != 0 - - total_loss = 0.0 - total_weight = 0.0 - - for pad_length, padded_mb_input, mb_input in zip( - mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs - ): - outputs = self.model(**padded_mb_input) - logits = outputs.logits.squeeze(0) - logits = logits[:-pad_length] if pad_length > 0 else logits - loss = loss_fn(logits, mb_input) - - # Simple weight calculation (could be improved) - loss_scale = loss_weight_fn(mb_input) / total_loss_weight - total_loss += loss.item() * loss_scale - total_weight += loss_scale - - return torch.tensor(total_loss / total_weight) - - @torch.no_grad() - def forward( - self, - input_: TensorDict, - output_seqlens: List[int] | None = None, - post_hook: Callable[[torch.Tensor, Dict], Any] | None = None, - aggregate_fn: Callable[[List[Any]], Any] = torch.cat, - ) -> Any | None: - """Forward pass with optional post-processing.""" - cu_seqlens = pack_tensor_dict(input_)["cu_seqlens"] - mb_list = self._prepare_mb_list(input_) - - if output_seqlens is None: - output_seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().numpy().tolist() - - results = [] - for pad_length, padded_mb_input, mb_input in zip( - mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs - ): - outputs = self.model(**padded_mb_input) - logits = outputs.logits.squeeze(0) - logits = logits[:-pad_length] if pad_length > 0 else logits - - if post_hook: - result = post_hook(logits, mb_input) - results.append(result) - else: - results.append(logits) - - res = aggregate_fn(results) - output_seqlens = [output_seqlens[i] for i in mb_list.forward_indices] - unpacked = unpack_sequence(res, lens=output_seqlens, dim=0) - reordered = reorder_list(unpacked, mb_list.backward_indices) - return pad_and_stack_tensors_along_first_dim(reordered) diff --git a/arealite/engine/sft/lm_engine.py b/arealite/engine/sft/lm_engine.py index 0a6bec250e..18bee784b2 100644 --- a/arealite/engine/sft/lm_engine.py +++ b/arealite/engine/sft/lm_engine.py @@ -1,5 +1,3 @@ -from typing import Dict - import torch import torch.utils.data from tensordict import TensorDict @@ -44,9 +42,7 @@ def evaluate_lm(self, data): return self.lm_engine.evaluate_lm(data) -def compute_packed_sft_loss( - logits: torch.Tensor, input_: Dict[str, torch.Tensor] -) -> torch.Tensor: +def compute_packed_sft_loss(logits: torch.Tensor, input_: TensorDict) -> torch.Tensor: packed_input_ids: torch.Tensor = input_["input_ids"] cu_seqlens: torch.Tensor = input_["cu_seqlens"] loss_mask = input_["loss_mask"].bool() diff --git a/arealite/tests/test_sglang_local_engine.py b/arealite/tests/test_sglang_local_engine.py index 2cf4dce0bd..faebc7a4a3 100644 --- a/arealite/tests/test_sglang_local_engine.py +++ b/arealite/tests/test_sglang_local_engine.py @@ -67,7 +67,6 @@ async def test_local_sglang_generate(): gconfig=GenerationHyperparameters(max_new_tokens=16), ) resp = await engine.agenerate(req) - print(resp.completions) assert isinstance(resp, LLMResponse) assert resp.input_tokens == req.input_ids @@ -76,9 +75,6 @@ async def test_local_sglang_generate(): == len(resp.output_tokens) == len(resp.output_versions) ) - assert isinstance(resp.completions, str) - - time.sleep(5) engine.destroy() diff --git a/arealite/tests/test_engine.py b/arealite/tests/test_train_engine.py similarity index 92% rename from arealite/tests/test_engine.py rename to arealite/tests/test_train_engine.py index 6c5d07a3f1..9c208b07eb 100644 --- a/arealite/tests/test_engine.py +++ b/arealite/tests/test_train_engine.py @@ -13,7 +13,6 @@ from arealite.api.cli_args import MicroBatchSpec, OptimizerConfig, TrainEngineConfig from arealite.api.io_struct import FinetuneSpec, SaveLoadMeta -from arealite.engine.fsdp_engine import FSDPEngine VOCAB_SIZE = 100 MODEL_PATH = "/storage/testing/models/Qwen__Qwen3-1.7B/" @@ -53,10 +52,10 @@ def mock_input( def get_engine(engine_type: str, model_path: str): + from arealite.engine.autotp_engine import DeepSpeedAutoTPEngine from arealite.engine.fsdp_engine import FSDPEngine - from arealite.engine.hf_engine import HFEngine - engine_cls = {"hf": HFEngine, "fsdp": FSDPEngine}[engine_type] + engine_cls = {"auto_tp": DeepSpeedAutoTPEngine, "fsdp": FSDPEngine}[engine_type] engine_config = TrainEngineConfig( experiment_name=f"test-{engine_type}-engine", @@ -75,7 +74,7 @@ def mock_loss_fn(logits: torch.Tensor, input_data: Dict) -> torch.Tensor: return torch.mean(logits) -@pytest.fixture(scope="module", params=["fsdp", "hf"]) +@pytest.fixture(scope="module", params=["fsdp", "auto_tp"]) def engine(request): os.environ.update( { @@ -136,6 +135,12 @@ def test_train_batch(engine, mock_input): @torch.no_grad() def test_hf_save_load_weights(tmp_path_factory, engine, mock_input): + from arealite.engine.autotp_engine import DeepSpeedAutoTPEngine + + if isinstance(engine, DeepSpeedAutoTPEngine): + print("AutoTP engine does not support HF save/load for now.") + return + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) path = tmp_path_factory.mktemp("hf_engine_test") save_load_meta = SaveLoadMeta( From b56f5998ec5ed49df6073c726b2118155f858398 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Wed, 16 Jul 2025 17:22:54 +0800 Subject: [PATCH 07/16] PullRequest: 371 [lite] [fix] fix misc bugs in GRPO implementation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-fix0716 of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/371 Reviewed-by: 晓雷 * . --- arealite/engine/ppo/actor.py | 7 +- arealite/engine/sglang_remote.py | 108 +++++++++++++++++++------------ arealite/utils/data.py | 28 +++----- arealite/utils/functional.py | 38 ++++++----- arealite/utils/http.py | 71 ++++++++++++++------ arealite/workflow/rlvr.py | 4 +- examples/arealite/gsm8k_grpo.py | 9 +-- realhf/base/name_resolve.py | 4 +- 8 files changed, 159 insertions(+), 110 deletions(-) diff --git a/arealite/engine/ppo/actor.py b/arealite/engine/ppo/actor.py index 937f32dc6e..802c03e13c 100644 --- a/arealite/engine/ppo/actor.py +++ b/arealite/engine/ppo/actor.py @@ -114,7 +114,9 @@ def compute_advantages(self, data: TensorDict) -> None: values = torch.zeros_like(rewards) else: values = data["values"] - advantages_reversed = [] + advantages_reversed = [ + torch.zeros(bs, dtype=torch.float32, device=values.device) + ] lastgaelam = 0 for t in reversed(range(max_seqlen - 1)): nextvalues = values[:, t + 1] @@ -123,9 +125,6 @@ def compute_advantages(self, data: TensorDict) -> None: delta = rewards[:, t] + self.discount * nextvalues - values[:, t] lastgaelam = delta + self.discount * self.gae_lambda * lastgaelam advantages_reversed.append(lastgaelam) - advantages_reversed.append( - torch.zeros(bs, dtype=torch.float32, device=values.device) - ) advantages = torch.stack(advantages_reversed[::-1], dim=1) # Optionally perform advantage normalization. diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index c985fa6301..4b9eb861ab 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -7,10 +7,12 @@ from concurrent.futures import ProcessPoolExecutor from datetime import datetime from queue import Empty, Full, Queue -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List +from typing import TYPE_CHECKING, Any, Callable, Dict, List +import aiohttp import requests import torch.distributed as dist +import uvloop from tensordict import TensorDict from torchdata.stateful_dataloader import StatefulDataLoader @@ -23,8 +25,8 @@ RolloutStat, WeightUpdateMeta, ) -from arealite.utils.http import arequest_with_retry -from arealite.utils.padding import concat_padded_tensors +from arealite.utils.data import concat_padded_tensors +from arealite.utils.http import arequest_with_retry, get_default_connector from realhf.base import logging, name_resolve, names, pkg_version if TYPE_CHECKING: @@ -37,7 +39,7 @@ else: SGLANG_TOKEN_OUTPUT_IDENTIFIER = "token_ids" -ROLLOUT_POLL_WAIT_TIME = 0.1 +ROLLOUT_POLL_WAIT_TIME = 0.05 RID_CACHE_SIZE = 128 @@ -98,6 +100,7 @@ def check_health(self, base_url): def initialize(self, addr: str | None, ft_spec: FinetuneSpec = None): self.rollout_tasks: Dict[str, asyncio.Task] = {} + self.executor = ProcessPoolExecutor(max_workers=1) self.rollout_thread = threading.Thread(target=self._rollout_thread) self.rollout_thread.start() @@ -118,32 +121,39 @@ def get_version(self): def _rollout_thread(self): """Thread that runs the rollout loop.""" try: - asyncio.run(self._rollout_thread_async()) + uvloop.run(self._rollout_thread_async()) except Exception as e: traceback.print_exc() async def _rollout_thread_async(self): - pending_data = [] rollout_tasks = self.rollout_tasks rid = 0 + + # NOTE: session is not thread-safe, but we only submit requests in the sub-thread. + self.session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout( + total=self.config.request_timeout, + sock_connect=self.config.request_timeout, + connect=self.config.request_timeout, + ), + read_bufsize=1024 * 1024 * 10, + connector=get_default_connector(), + ) + try: while not self.exiting.is_set(): - # Load next data from controller - while True: - try: - data, workflow = self.input_queue.get_nowait() - logger.debug(f"Get data from puller: {data}") - pending_data.append(data) - except Empty: - logger.debug(f"No data from puller stream.") - break - # Check capacity capacity = self.get_capacity() # Create new rollout task - while capacity > 0 and pending_data and not self.paused.is_set(): + while ( + capacity > 0 + and not self.paused.is_set() + and self.input_queue.qsize() > 0 + ): + data, workflow = self.input_queue.get_nowait() + logger.debug(f"Get data from puller: {data}") task = asyncio.create_task( - workflow.arun_episode(self, pending_data.pop(0)), name=str(rid) + workflow.arun_episode(self, data), name=str(rid) ) with self.lock: rollout_tasks[str(rid)] = task @@ -158,7 +168,6 @@ async def _rollout_thread_async(self): ) capacity -= 1 rid += 1 - # Wait for rollout completion with self.lock: tasks = list(rollout_tasks.values()) @@ -169,11 +178,6 @@ async def _rollout_thread_async(self): timeout=ROLLOUT_POLL_WAIT_TIME, return_when=asyncio.FIRST_COMPLETED, ) - if not done: - await asyncio.sleep(1) - else: - await asyncio.sleep(1) - # Collect done results for task in done: traj = await task @@ -199,6 +203,7 @@ async def _rollout_thread_async(self): f"running: {self.rollout_stat.running}, " f"accepted: {self.rollout_stat.accepted}." ) + await asyncio.sleep(1) except Exception: traceback.print_exc() finally: @@ -213,10 +218,11 @@ async def _rollout_thread_async(self): pass def choose_server(self) -> str: - if self.config.schedule_policy == "round_robin": - server = self.addresses[self.server_idx] - self.server_idx = (self.server_idx + 1) % len(self.addresses) - return server + with self.lock: + if self.config.schedule_policy == "round_robin": + server = self.addresses[self.server_idx] + self.server_idx = (self.server_idx + 1) % len(self.addresses) + return server raise NotImplementedError("Only round-robin scheduling is implemented.") async def agenerate(self, req: LLMRequest) -> LLMResponse: @@ -253,7 +259,6 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: accumulated_versions = [] # Deal with rollout interruption - completions = "" stop_reason = "length" if req.rid in self.rid_to_address: @@ -273,11 +278,12 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: ): # loop until the generation is complete result = await arequest_with_retry( - addr=self.choose_server(), + session=self.session, + addr=server_addr, endpoint="/generate", payload=payload, method="POST", - max_retries=3, + max_retries=self.config.request_retries, timeout=self.config.request_timeout, ) @@ -297,10 +303,7 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: stop_reason = finish_reason["type"] payload["input_ids"] += result[SGLANG_TOKEN_OUTPUT_IDENTIFIER] - sample_params["max_new_tokens"] = min( - sample_params["max_new_tokens"], - gconfig.max_new_tokens - len(output_tokens), - ) + sample_params["max_new_tokens"] -= len(output_tokens) latency = time.perf_counter() - start_time @@ -408,18 +411,24 @@ def rollout_batch( def prepare_batch( self, - data_generator: Iterator, dataloader: StatefulDataLoader, workflow: "RolloutWorkflow", ): + if not hasattr(self, "data_generator"): + self.data_generator = iter(dataloader) assert dataloader.batch_size is not None while True: - if self.get_capacity() + dataloader.batch_size > 0: + # Submit at least two batches to allow maximum overlap + if ( + self.get_capacity() + dataloader.batch_size > 0 + and self.input_queue.qsize() + dataloader.batch_size + < self.input_queue.maxsize + ): try: - data = next(data_generator) + data = next(self.data_generator) except StopIteration: - data_generator = iter(dataloader) - data = next(data_generator) + self.data_generator = iter(dataloader) + data = next(self.data_generator) for item in data: self.submit(item, workflow=workflow) try: @@ -435,10 +444,12 @@ def resume(self): async def aupdate_weights_from_disk( - addr, path: str, request_retries: int, request_timeout: float + session, addr, path: str, request_retries: int, request_timeout: float ): + tik = time.time() res = await arequest_with_retry( addr=addr, + session=session, endpoint="/update_weights_from_disk", payload=dict(model_path=str(path), allow_interrupt=True), method="POST", @@ -472,9 +483,19 @@ async def _fn(): logger.info( f"Begin update weights from {path}, responded in {(load_timestamp - save_timestamp):.2f}s" ) + session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout( + total=request_timeout, + sock_connect=request_timeout, + connect=request_timeout, + ), + read_bufsize=1024 * 1024 * 10, + connector=get_default_connector(), + ) jobs = [ aupdate_weights_from_disk( - addr, + session=session, + addr=addr, path=path, request_retries=request_retries, request_timeout=request_timeout, @@ -482,8 +503,9 @@ async def _fn(): for addr in addresses ] await asyncio.gather(*jobs) + await session.close() logger.info( f"Loading weights done in {(datetime.now().timestamp() - load_timestamp):.2f}s" ) - return asyncio.run(_fn()) + return uvloop.run(_fn()) diff --git a/arealite/utils/data.py b/arealite/utils/data.py index f6572d10d8..b5f317416d 100644 --- a/arealite/utils/data.py +++ b/arealite/utils/data.py @@ -92,9 +92,7 @@ def unpad_input( seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad( - torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0) - ) + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) return ( rearrange(hidden_states, "b s ... -> (b s) ...")[indices], indices, @@ -116,16 +114,12 @@ def concat_padded_tensors( if not tensor_dicts: return TensorDict() - # Find max sequence length across all dictionaries - lens = [] - for tensor_dict in tensor_dicts: - for key, tensor in tensor_dict.items(): - if key != "attention_mask" and len(tensor.shape) == 2: - lens.append(tensor.shape[1]) - break - max_length = max(lens) - attn_mask = torch.arange(max_length).unsqueeze(0) < torch.tensor(lens).unsqueeze(1) + batch_sizes = [tuple(d.batch_size) for d in tensor_dicts] + new_batch_size = [sum(x[0] for x in batch_sizes), *batch_sizes[0][1:]] + # Find max sequence length across all dictionaries + assert all("attention_mask" in td for td in tensor_dicts) + max_length = max([x["attention_mask"].shape[1] for x in tensor_dicts]) result = {} # Process each key for key in tensor_dicts[0].keys(): @@ -154,9 +148,7 @@ def concat_padded_tensors( tensors_to_concat.append(tensor) result[key] = torch.cat(tensors_to_concat, dim=0) - if "attention_mask" not in result: - result["attention_mask"] = attn_mask - return TensorDict(result, batch_size=[len(lens)]) + return TensorDict(result, batch_size=new_batch_size) def to_device(data: Dict[str, torch.Tensor | Any], device) -> Dict[str, torch.Tensor]: @@ -231,7 +223,7 @@ def pack_tensor_dict(data: TensorDict): cu_seqlens = torch.cumsum(lens, dim=0) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) - total_length = cu_seqlens[-1].item() + total_length = int(cu_seqlens[-1].item()) # Pack tensors packed_data = {} for key, value in data.items(): @@ -246,7 +238,7 @@ def pack_tensor_dict(data: TensorDict): and value.shape[1] == seq_len ): packed_tensor = torch.empty( - total_length, *value.shape[2:], dtype=value.dtype, device=value.device + (total_length, *value.shape[2:]), dtype=value.dtype, device=value.device ) # Fill the packed tensor with values from the original tensor for i in range(bs): @@ -363,7 +355,7 @@ def _split(tensor): mbs=results, mb_spec=mb_spec, forward_indices=forward_indices, - backward_indices=backward_indices, + backward_indices=backward_indices.tolist(), group_lens=group_lens, ) diff --git a/arealite/utils/functional.py b/arealite/utils/functional.py index 9ce736c97c..3abafa8a6b 100644 --- a/arealite/utils/functional.py +++ b/arealite/utils/functional.py @@ -121,19 +121,24 @@ def masked_normalization( def ppo_actor_loss_fn( logprobs: torch.Tensor, + proximal_logprobs: torch.Tensor, old_logprobs: torch.Tensor, advantages: torch.Tensor, eps_clip: float, loss_mask: torch.Tensor, c_clip: Optional[float] = None, - proximal_logprobs: Optional[torch.Tensor] = None, behav_imp_weight_cap: Optional[float] = None, ) -> Tuple[torch.Tensor, Dict]: - denorm_logprobs = ( - proximal_logprobs if proximal_logprobs is not None else old_logprobs - ) + """ + When decoupled loss is disabled: + 1. if recompute logp, both old_logprobs and proximal_logprobs are recomputed logp; + 2. if no recomputation, both old_logp and proximal_logprobs are produced by the inference backend. + + When decoupled loss is enabled, proximal_logprobs is the recomputed logp, + old_logprobs is produced by the inference engine. + """ loss_mask_count = loss_mask.count_nonzero() or 1 - ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) + ratio = torch.where(loss_mask, torch.exp(logprobs - proximal_logprobs), 0) clipped_ratio = torch.clamp(ratio, 1.0 - eps_clip, 1.0 + eps_clip) pg_loss1 = -advantages * ratio pg_loss2 = -advantages * clipped_ratio @@ -146,17 +151,16 @@ def ppo_actor_loss_fn( pg_loss = torch.min(pg_loss, pg_loss3) else: dual_clip_mask = torch.zeros_like(clip_mask) - if proximal_logprobs is not None: - behav_kl = proximal_logprobs - old_logprobs - behav_imp_weight = behav_kl.exp() - behav_mask = ( - (behav_imp_weight <= behav_imp_weight_cap).logical_and(loss_mask) - if behav_imp_weight_cap is not None - else loss_mask - ) - behav_kl = torch.where(behav_mask, behav_kl, 0.0) - behav_imp_weight = torch.where(behav_mask, behav_imp_weight, 0.0) - pg_loss = pg_loss * behav_imp_weight + behav_kl = proximal_logprobs - old_logprobs + behav_imp_weight = behav_kl.exp() + behav_mask = ( + (behav_imp_weight <= behav_imp_weight_cap).logical_and(loss_mask) + if behav_imp_weight_cap is not None + else loss_mask + ) + behav_kl = torch.where(behav_mask, behav_kl, 0.0) + behav_imp_weight = torch.where(behav_mask, behav_imp_weight, 0.0) + pg_loss = pg_loss * behav_imp_weight logging_loss = pg_loss.detach() pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count clip_mask.logical_and_(loss_mask) @@ -164,7 +168,7 @@ def ppo_actor_loss_fn( stat = dict( loss=logging_loss, importance_weight=ratio.detach(), - approx_kl=(logprobs - denorm_logprobs).detach(), + approx_kl=(logprobs - proximal_logprobs).detach(), clip_mask=clip_mask, dual_clip_mask=dual_clip_mask, ) diff --git a/arealite/utils/http.py b/arealite/utils/http.py index 29e4df4050..a39a3614de 100644 --- a/arealite/utils/http.py +++ b/arealite/utils/http.py @@ -3,45 +3,74 @@ import aiohttp +from realhf.base import logging + DEFAULT_RETRIES = 1 DEFAULT_REQUEST_TIMEOUT = 3600 +logger = logging.getLogger(__file__) + + +def get_default_connector(): + return aiohttp.TCPConnector(limit=0, use_dns_cache=False, force_close=True) + + async def arequest_with_retry( addr: str, endpoint: str, payload: Optional[Dict[str, Any]] = None, + session: aiohttp.ClientSession | None = None, method: str = "POST", max_retries: Optional[int] = None, timeout: Optional[float] = None, retry_delay: float = 1.0, -) -> aiohttp.ClientResponse: + verbose=False, +) -> Dict: timeout = timeout or DEFAULT_REQUEST_TIMEOUT last_exception = None max_retries = max_retries or DEFAULT_RETRIES base_url = f"http://{addr}" url = f"{base_url}{endpoint}" + timeo = aiohttp.ClientTimeout( + total=timeout, + sock_connect=timeout, + connect=timeout, + ) + if session is None: + _session = aiohttp.ClientSession( + timeout=timeo, + read_bufsize=1024 * 1024 * 10, + connector=get_default_connector(), + ) + else: + _session = session + for attempt in range(max_retries): try: - async with aiohttp.ClientSession( - timeout=aiohttp.ClientTimeout( - total=timeout, - sock_connect=timeout, - ) - ) as session: - if method.upper() == "GET": - response = await session.get(url) - elif method.upper() == "POST": - response = await session.post(url, json=payload) - elif method.upper() == "PUT": - response = await session.put(url, json=payload) - elif method.upper() == "DELETE": - response = await session.delete(url) - else: - raise ValueError(f"Unsupported HTTP method: {method}") + if verbose: + logger.info("enter client session, start sending requests") + if method.upper() == "GET": + ctx = _session.get(url, timeout=timeo) + elif method.upper() == "POST": + ctx = _session.post(url, json=payload, timeout=timeo) + elif method.upper() == "PUT": + ctx = _session.put(url, json=payload, timeout=timeo) + elif method.upper() == "DELETE": + ctx = _session.delete(url, timeout=timeo) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + async with ctx as response: + if verbose: + logger.info("http requests return") response.raise_for_status() - return await response.json() + res = await response.json() + if verbose: + logger.info("get http result") + if session is None: + await _session.close() + return res except ( aiohttp.ClientError, aiohttp.ClientResponseError, @@ -51,6 +80,10 @@ async def arequest_with_retry( if attempt < max_retries - 1: await asyncio.sleep(retry_delay) continue + if session is None: + await _session.close() raise RuntimeError( - f"Failed after {max_retries} retries each. " f"Last error: {last_exception}" + f"Failed after {max_retries} retries each. " + f"Payload: {payload}. Addr: {addr}. Endpoint: {endpoint}. " + f"Last error: {last_exception}" ) diff --git a/arealite/workflow/rlvr.py b/arealite/workflow/rlvr.py index 026f574925..3ce55dfc00 100644 --- a/arealite/workflow/rlvr.py +++ b/arealite/workflow/rlvr.py @@ -8,7 +8,7 @@ from arealite.api.cli_args import GenerationHyperparameters from arealite.api.io_struct import LLMRequest from arealite.api.workflow_api import RolloutWorkflow -from arealite.utils.padding import concat_padded_tensors +from arealite.utils.data import concat_padded_tensors class RLVRWorkflow(RolloutWorkflow): @@ -61,7 +61,7 @@ async def arun_episode(self, engine, data): versions=torch.tensor(versions).unsqueeze(0), attention_mask=torch.ones(len(seq), dtype=torch.bool).unsqueeze(0), # reward - rewards=torch.tensor([reward]), + rewards=torch.tensor([float(reward)]), ) results.append(TensorDict(res, batch_size=[1])) diff --git a/examples/arealite/gsm8k_grpo.py b/examples/arealite/gsm8k_grpo.py index 07434a09ba..1841343da5 100644 --- a/examples/arealite/gsm8k_grpo.py +++ b/examples/arealite/gsm8k_grpo.py @@ -152,11 +152,7 @@ def main_grpo(): with stats_tracker.record_timing("rollout"): if config.async_training: - batch = rollout.prepare_batch( - data_generator, - train_dataloader, - workflow=workflow, - ) + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) else: try: data = next(data_generator) @@ -210,8 +206,9 @@ def main_grpo(): actor.upload_weights(meta) if dist.get_rank() == 0: future.result() - rollout.set_version(global_step + 1) dist.barrier() + torch.cuda.synchronize() + rollout.set_version(global_step + 1) with stats_tracker.record_timing("save"): saver.save(actor, epoch, step, global_step) diff --git a/realhf/base/name_resolve.py b/realhf/base/name_resolve.py index 9c6ec25ba5..ef9879e7aa 100644 --- a/realhf/base/name_resolve.py +++ b/realhf/base/name_resolve.py @@ -1360,7 +1360,9 @@ def _keepalive_thread_run(self): def make_repository(args: "NameResolveConfig"): if args.type == "nfs": - return NfsNameRecordRepository(args.nfs_record_root) + repo = NfsNameRecordRepository(args.nfs_record_root) + os.makedirs(repo.record_root, exist_ok=True) + return repo elif args.type == "etcd3": host, port = args.etcd3_addr.split(":") return Etcd3NameRecordRepository(host=host, port=int(port)) From ddabd9cc9dccc418c38928a9597db2fdbefb1d4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=93=E9=9B=B7?= Date: Mon, 21 Jul 2025 10:29:51 +0800 Subject: [PATCH 08/16] PullRequest: 370 [lite] Add Slurm Launcher and Ray Launcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch mzy/lite/launcher of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/370 Reviewed-by: 博惟 * . * . * . * fix --- arealite/api/cli_args.py | 81 ++++- arealite/engine/base_hf_engine.py | 3 +- arealite/launcher/ray.py | 410 ++++++++++++++++++++++ arealite/launcher/sglang_server.py | 197 +++++++++++ arealite/launcher/slurm.py | 528 +++++++++++++++++++++++++++++ arealite/utils/launcher.py | 111 ++++++ arealite/utils/ray.py | 23 ++ arealite/utils/slurm.py | 154 +++++++++ examples/arealite/gsm8k_grpo.py | 6 +- examples/arealite/gsm8k_sft.py | 6 +- 10 files changed, 1498 insertions(+), 21 deletions(-) create mode 100644 arealite/launcher/ray.py create mode 100644 arealite/launcher/sglang_server.py create mode 100644 arealite/launcher/slurm.py create mode 100644 arealite/utils/launcher.py create mode 100644 arealite/utils/ray.py create mode 100644 arealite/utils/slurm.py diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 906f8fdebb..4962cfab55 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -676,20 +676,6 @@ class ClusterSpecConfig: "help": "Root for logs and checkpoints. Should be available to all nodes." }, ) - gpu_type: str = field( - default="tesla", metadata={"help": "GPU type of the cluster. Used by slurm."} - ) - mount: str = field( - default="/storage:/storage", metadata={"help": "Mount path for slurm."} - ) - gpu_image: str = field(default="", metadata={"help": "slurm image for trainers."}) - cpu_image: str = field(default="", metadata={"help": "slurm image for CPU jobs."}) - gpu_infer_image: str = field( - default="", metadata={"help": "slurm image for LLM inference."} - ) - node_name_prefix: str = field( - default="slurmd-", metadata={"help": "Node prefix for a slurm cluster."} - ) n_nodes: int = field( default=32, metadata={ @@ -725,6 +711,72 @@ class DatasetConfig: drop_last: bool = field(default=True) +@dataclass +class SlurmLauncherConfig: + """Configuration for launching the SGLang server with Slurm.""" + + srun_additional_args: str = field( + default="--overlap --mpi=pmi2 -K --chdir $PWD", + metadata={"help": "Additional arguments to pass to the srun command."}, + ) + container_type: str = field( + default="apptainer", + metadata={ + "help": "Type of containers used in slurm", + "choices": ["apptainer", "none"], + }, + ) + mount: str = field( + default="/storage:/storage", metadata={"help": "Mount path for slurm."} + ) + trainer_image: str = field( + default="", metadata={"help": "slurm image for trainers."} + ) + inference_server_image: str = field( + default="", metadata={"help": "slurm image for LLM inference."} + ) + + +@dataclass +class LauncherConfig: + """Configuration for launching the SGLang server.""" + + inference_server_cpus_per_gpu: int = field( + default=4, + metadata={"help": "Number of CPUs allocated per GPU for inference server. "}, + ) + inference_server_mem_per_gpu: int = field( + default=32 * 1024, + metadata={"help": "Memory allocated per GPU for inference server in MB. "}, + ) + trainer_cpus_per_gpu: int = field( + default=4, + metadata={"help": "Number of CPUs allocated per GPU for training. "}, + ) + trainer_mem_per_gpu: int = field( + default=32 * 1024, + metadata={"help": "Memory allocated per GPU for training in MB. "}, + ) + inference_server_env_vars: str = field( + default="", + metadata={ + "help": "Environment variables for inference server, seperated by commas. " + "Example: 'ENV1=val1,ENV2=val2'. " + }, + ) + trainer_env_vars: str = field( + default="", + metadata={ + "help": "Environment variables for training, seperated by commas. " + "Example: 'ENV1=val1,ENV2=val2'. " + }, + ) + slurm: SlurmLauncherConfig = field( + default_factory=SlurmLauncherConfig, + metadata={"help": "Slurm launcher configuration."}, + ) + + @dataclass class BaseExperimentConfig: # NOTE: we need this unified config class because different experiments @@ -785,6 +837,7 @@ class BaseExperimentConfig: server_only: bool = False sglang: SGLangConfig = field(default_factory=SGLangConfig) + launcher: LauncherConfig = field(default_factory=LauncherConfig) @dataclass diff --git a/arealite/engine/base_hf_engine.py b/arealite/engine/base_hf_engine.py index bcbda8ba27..91cef08d03 100644 --- a/arealite/engine/base_hf_engine.py +++ b/arealite/engine/base_hf_engine.py @@ -67,7 +67,8 @@ def parallelism_group(self) -> dist.ProcessGroup: def create_process_group(self): if not dist.is_initialized(): - # TODO: Handle the condition when WORLD_SIZE and RANK is not set in launcher + # NOTE: environment variables such as WORLD_SIZE, LOCAL_RANK + # will be always set either by the launcher or torchrun dist.init_process_group( backend="nccl", timeout=constants.NCCL_DEFAULT_TIMEOUT, diff --git a/arealite/launcher/ray.py b/arealite/launcher/ray.py new file mode 100644 index 0000000000..52813aeec8 --- /dev/null +++ b/arealite/launcher/ray.py @@ -0,0 +1,410 @@ +import getpass +import importlib.util +import os +import pathlib +import sys +import time +from typing import Dict, List, Optional + +import ray +import ray.exceptions +from ray.runtime_env import RuntimeEnv +from ray.util.placement_group import PlacementGroup +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +import realhf.base.logging as logging +from arealite.api.cli_args import ( + ClusterSpecConfig, + LauncherConfig, + SGLangConfig, + parse_cli_args, + to_structured_cfg, +) +from arealite.api.io_struct import AllocationMode, AllocationType +from arealite.utils.launcher import ( + get_env_vars, + validate_config_for_distributed_launcher, + wait_sglang_server_addrs, +) +from arealite.utils.ray import get_placement_group_master_ip_and_port +from realhf.base import logging, name_resolve, names +from realhf.scheduler.client import JobException, JobState + +logger = logging.getLogger("RayLauncher") + +RAY_WAIT_CHECK_TIME_INTERVAL = 5 # seconds +DEFAULT_MAIN_FUNC_NAME = "main" + + +def run_func(file_path, function_name, *args, **kwargs): + # Convert the file path to a module name + module_name = file_path.replace("/", "_").replace(".", "_") + + # Load the module from file path + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + + # Get the function and execute it + try: + function = getattr(module, function_name) + except AttributeError as e: + raise ValueError( + f"Function '{function_name}' not found in module '{module_name}'. " + f"Please ensure the name of the main function in your entry point " + f"is '{function_name}'." + ) from e + return function(*args, **kwargs) + + +class RayLauncher: + def __init__(self, experiment_name: str, trial_name: str, fileroot: str): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.fileroot = fileroot + + # job_name to ray future + self.jobs = {} + + @property + def run_name(self): + return f"{self.experiment_name}_{self.trial_name}" + + def log_path_of(self, job_name: str) -> str: + log_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(log_path, exist_ok=True) + return os.path.join(log_path, f"{job_name}.log") + + def submit( + self, + job_name: str, + file_path: str, + func_name: str, + args: List[str], # arguments to pass to the function + gpus: int, + cpus: int, + mem: int, # MB + env_vars: Optional[Dict] = None, + placement_group: Optional[PlacementGroup] = None, + bundle_index: int = -1, + kwargs: Optional[ + Dict[str, str] + ] = None, # keyword arguments to pass to the function + ): + if kwargs is None: + kwargs = {} + runtime_env = RuntimeEnv( + env_vars=env_vars or dict(), + ) + scheduling_strategy = ( + PlacementGroupSchedulingStrategy( + placement_group=placement_group, + placement_group_bundle_index=bundle_index, + placement_group_capture_child_tasks=True, + ) + if placement_group is not None + else "DEFAULT" + ) + future = ray.remote( + num_cpus=cpus, + num_gpus=gpus, + memory=mem * 1024 * 1024, # Convert MB to bytes + runtime_env=runtime_env, + scheduling_strategy=scheduling_strategy, + )(run_func).remote(file_path, func_name, *args, **kwargs) + self.jobs[job_name] = future + return future + + def submit_array( + self, + job_name: str, + file_path: str, + func_name: str, + count: int, + nodes: int, + list_args: List[List], + gpus_per_task: int, + cpus_per_task: int, + mem_per_task: int, # MB + list_kwargs: List[Dict] | None = None, + env_vars: Optional[Dict] = None, + amend_torch_dist_env: bool = False, + ): + """Submit an array of jobs to Ray with ray placement groups. + + Note: Here we use `ray.remote` instead of `ray job submit` since `ray job submit` + does not support placement groups, and can not specify which node to run the job on. + Therefore we could not know the IP address of jobs for torch distributed initialization. + """ + + if count % nodes != 0: + raise ValueError( + f"Count {count} is not divisible by nodes {nodes}. " + "Please ensure that count is a multiple of nodes." + ) + assert ( + len(list_args) == count + ), f"Length of list_args {len(list_args)} does not match count {count}." + if list_kwargs is not None: + assert ( + len(list_kwargs) == count + ), f"Length of list_kwargs {len(list_kwargs)} does not match count {count}." + + tasks_per_node = count // nodes + gpus_per_node = gpus_per_task * tasks_per_node + cpus_per_node = cpus_per_task * tasks_per_node + mem_per_node = mem_per_task * tasks_per_node + + placement_group = ray.util.placement_group( + bundles=[ + { + "CPU": cpus_per_node, + "GPU": gpus_per_node, + "memory": mem_per_node * 1024 * 1024, # Convert MB to bytes + } + ] + * nodes, + strategy="STRICT_SPREAD", + ) + try: + ray.get(placement_group.ready(), timeout=30) + except ray.exceptions.GetTimeoutError as e: + logger.error( + "Ray placement group timeout, please check if the resource requirement " + "for your experiment exceeds the available resources in the cluster. \n" + f"ray.nodes(): {ray.nodes()} \n" + f"Placement Group bundles: " + f"cpus_per_node={cpus_per_node}, gpus_per_node={gpus_per_node}, " + f"mem_per_node={mem_per_node}MB, nodes={nodes}" + ) + raise e + + if amend_torch_dist_env: + host_ip, port = get_placement_group_master_ip_and_port(placement_group) + logger.info( + f"Amend torch distributed env vars: " + f"MASTER_ADDR={host_ip}, PORT={port}" + ) + + futures = [] + for i in range(count): + args = list_args[i] + kwargs = list_kwargs[i] if list_kwargs is not None else {} + + # manage environment variables + env_vars = env_vars or {} + if "CUDA_VISIBLE_DEVICES" in env_vars: + logger.warning( + "Setting CUDA_VISIBLE_DEVICES before running ray jobs may result in unexpected behavior." + ) + + node_id = i // tasks_per_node + _env_vars = { + **env_vars, + } + + if amend_torch_dist_env: + assert gpus_per_task == 1 + # NOTE: Here we only provide environment variables for torch distributed + # initialization, and LOCAL_RANK for torch.device. + # Other environment variables automatically set by torchrun are not set, and + # they should be never accessed in trainer code. + _env_vars.update( + { + "RANK": str(i), + "WORLD_SIZE": str(count), + # Ray will automatically isolate CUDA_VISIBLE_DEVICES for each GPU + "LOCAL_RANK": "0", + "MASTER_ADDR": str(host_ip), + "MASTER_PORT": str(port), + } + ) + future = self.submit( + job_name=f"{job_name}:{i}", + file_path=file_path, + func_name=func_name, + args=args, + gpus=gpus_per_task, + cpus=cpus_per_task, + mem=mem_per_task, + env_vars=_env_vars, + placement_group=placement_group, + bundle_index=node_id, + kwargs=kwargs, + ) + futures.append(future) + + return futures + + def stop(self, job_name: str, force: bool = False): + """Stop a job by name.""" + if job_name in self.jobs: + future = self.jobs[job_name] + try: + ray.cancel(future, force=force) + except Exception as e: + logger.error(f"Failed to cancel job {job_name}: {e}") + return + self.jobs.pop(job_name, None) + logger.info(f"Job {job_name} stopped.") + else: + logger.warning(f"Job {job_name} not found in running jobs.") + + def stop_all(self, force: bool = False): + """Stop all jobs.""" + for job_name in list(self.jobs.keys()): + self.stop(job_name, force=force) + logger.info("All jobs stopped.") + self.jobs.clear() + + def wait( + self, check_status=(JobState.FAILED,), remove_status=(JobState.COMPLETED,) + ): + """Check every RAY_WAIT_CHECK_TIME_INTERVAL seconds for the status of all jobs. + If a ray job returns, its status changes to JobState.COMPLETED. + If a ray job failed, its status changes to JobState.FAILED. + If any job is in check_status, stop all jobs at once. + If any job is in remove status, remove them from job list. + Return if all jobs are removed from job list, or some job is in check status. + """ + for status in list(check_status) + list(remove_status): + assert status in [ + JobState.COMPLETED, + JobState.FAILED, + ], "In RayLauncher.wait, we only check completed or failed jobs." + logger.info(f"Waiting for {len(self.jobs)} jobs.") + while self.jobs: + job_status = {} + for job_name, future in list(self.jobs.items()): + try: + r = ray.get(future, timeout=0.1) + logger.info(f"Job {job_name} completed with result: {r}") + job_status[job_name] = JobState.COMPLETED + except ray.exceptions.RayTaskError as e: + logger.error(f"Job {job_name} failed with error: {e}.") + job_status[job_name] = JobState.FAILED + except ray.exceptions.GetTimeoutError: + continue + + for job_name, status in job_status.items(): + if status in check_status: + logger.info(f"Job {job_name} is {status}, stopping all jobs.") + self.stop_all(force=True) + return + if status in remove_status: + logger.info(f"Job {job_name} is {status}, removed.") + self.jobs.pop(job_name) + + time.sleep(RAY_WAIT_CHECK_TIME_INTERVAL) + + +def ray_main(): + # usage: python -m arealite.launcher.ray --config [] + ray.init() + config, config_file = parse_cli_args(sys.argv[2:]) + config.launcher = to_structured_cfg(config.launcher, LauncherConfig) + config.cluster = to_structured_cfg(config.cluster, ClusterSpecConfig) + config.sglang = to_structured_cfg(config.sglang, SGLangConfig) + validate_config_for_distributed_launcher(config) + + name_resolve.reconfigure(config.cluster.name_resolve) + name_resolve.clear_subtree( + names.trial_root( + experiment_name=config.experiment_name, trial_name=config.trial_name + ) + ) + + n_nodes = config.n_nodes + n_gpus_per_node = config.n_gpus_per_node + launcher = RayLauncher( + experiment_name=config.experiment_name, + trial_name=config.trial_name, + fileroot=config.cluster.fileroot, + ) + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + sglang_cmds = [] + sglang_addrs = [] + n_sglang_nodes = 0 + if allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG: + # Launcher should launch SGLang servers according to allocation mode. + sglang_tp_size = allocation_mode.gen_tp_size + n_sglang_servers = allocation_mode.gen_dp_size + n_sglang_nodes = allocation_mode.gen_world_size // n_gpus_per_node + + base_seed = config.sglang.random_seed + sglang_args_list = [ + [sys.argv[2:] + [f"sglang.random_seed={base_seed + i}"]] + for i in range(n_sglang_servers) + ] + sglang_entry_point = str( + pathlib.Path(__file__).resolve().parent.joinpath("sglang_server.py") + ) + launcher.submit_array( + job_name="llm_server", + file_path=sglang_entry_point, + func_name=DEFAULT_MAIN_FUNC_NAME, + count=n_sglang_servers, + nodes=n_sglang_nodes, + list_args=sglang_args_list, + gpus_per_task=sglang_tp_size, + cpus_per_task=config.launcher.inference_server_cpus_per_gpu + * sglang_tp_size, + mem_per_task=config.launcher.inference_server_mem_per_gpu * sglang_tp_size, + env_vars=get_env_vars( + config.cluster.cluster_name, + config.launcher.inference_server_env_vars, + ), + ) + # Get SGLang server addresses via name_resolve + try: + sglang_addrs = wait_sglang_server_addrs( + config.experiment_name, + config.trial_name, + n_sglang_servers, + ) + except TimeoutError as e: + launcher.stop_all(force=True) + raise e + + trainer_n_nodes = n_nodes - n_sglang_nodes + trainer_entry_point = sys.argv[1] + n_trainer_processes = trainer_n_nodes * config.cluster.n_gpus_per_node + trainer_args_list = [[sys.argv[2:]] for _ in range(n_trainer_processes)] + if not config.server_only: + # In ray, we launch trainer in the granularity of processes (1 GPU per process) + # We amend environment variable similar to torchrun to ensure correct initialization of + # torch distributed. + launcher.submit_array( + job_name="trainer", + file_path=trainer_entry_point, + func_name=DEFAULT_MAIN_FUNC_NAME, + count=trainer_n_nodes * config.cluster.n_gpus_per_node, + nodes=trainer_n_nodes, + list_args=trainer_args_list, + gpus_per_task=1, + cpus_per_task=config.launcher.trainer_cpus_per_gpu, + mem_per_task=config.launcher.trainer_mem_per_gpu, + env_vars=dict( + **get_env_vars( + config.cluster.cluster_name, + config.launcher.trainer_env_vars, + ), + AREAL_LLM_SERVER_ADDRS=",".join(sglang_addrs), + ), + amend_torch_dist_env=True, + ) + + try: + launcher.wait(check_status=(JobState.COMPLETED, JobState.FAILED)) + except (KeyboardInterrupt, JobException, TimeoutError) as e: + launcher.stop_all(force=True) + raise e + + +if __name__ == "__main__": + # usage: python -m arealite.launcher.ray \ + # --config [] \ + # launcher.ray.main_func_name= + ray_main() diff --git a/arealite/launcher/sglang_server.py b/arealite/launcher/sglang_server.py new file mode 100644 index 0000000000..916d4dfea5 --- /dev/null +++ b/arealite/launcher/sglang_server.py @@ -0,0 +1,197 @@ +import os +import subprocess +import sys +import time +import uuid +from pathlib import Path +from typing import Optional + +import ray +import requests + +from arealite.api.cli_args import ( + ClusterSpecConfig, + NameResolveConfig, + SGLangConfig, + parse_cli_args, + to_structured_cfg, +) +from arealite.api.io_struct import AllocationMode, AllocationType +from arealite.utils.launcher import TRITON_CACHE_PATH +from arealite.utils.network import find_free_ports, gethostip +from realhf.base import logging, name_resolve, names, pkg_version + +logger = logging.getLogger("SGLangServer Wrapper") + + +def execute_shell_command(command: str) -> subprocess.Popen: + """ + Execute a shell command and return its process handle. + """ + # Replace newline continuations and split the command string. + command = command.replace("\\\n", " ").replace("\\", " ") + parts = command.split() + _env = os.environ.copy() + # To avoid DirectoryNotEmpty error caused by triton + triton_cache_path = _env.get("TRITON_CACHE_PATH", TRITON_CACHE_PATH) + unique_triton_cache_path = os.path.join(triton_cache_path, str(uuid.uuid4())) + _env["TRITON_CACHE_PATH"] = unique_triton_cache_path + return subprocess.Popen( + parts, + text=True, + env=_env, + stdout=sys.stdout, + stderr=subprocess.STDOUT, + ) + + +def apply_sglang_patch(): + p = Path(os.path.dirname(__file__)) + patch_path = str( + p.parent.parent + / "patch" + / "sglang" + / f"v{pkg_version.get_version('sglang')}.patch" + ) + + target_path = "" + sglang_meta = subprocess.check_output( + "python3 -m pip show sglang", shell=True + ).decode("ascii") + for line in sglang_meta.split("\n"): + line = line.strip() + if line.startswith("Editable project location: "): + target_path = str(Path(line.split(": ")[1]).parent) + + if target_path: + proc = subprocess.Popen( + ["git", "apply", patch_path], + cwd=target_path, + stderr=sys.stdout, + stdout=sys.stdout, + ) + proc.wait() + logger.info(f"Applied SGLang patch at {target_path}") + + +def launch_server_cmd(command: str): + """ + Launch the server using the given command. + If no port is specified, a free port is reserved. + """ + if not ray.is_initialized(): + apply_sglang_patch() + process = execute_shell_command(command) + return process + + +def wait_for_server(base_url: str, timeout: Optional[int] = None) -> None: + """Wait for the server to be ready by polling the /v1/models endpoint. + + Args: + base_url: The base URL of the server + timeout: Maximum time to wait in seconds. None means wait forever. + """ + start_time = time.time() + while True: + try: + response = requests.get( + f"{base_url}/v1/models", + headers={"Authorization": "Bearer None"}, + ) + if response.status_code == 200: + time.sleep(5) + break + + if timeout and time.time() - start_time > timeout: + raise TimeoutError("Server did not become ready within timeout period") + except requests.exceptions.RequestException: + time.sleep(1) + + +class SGLangServerWrapper: + def __init__( + self, + experiment_name: str, + trial_name: str, + sglang_config: SGLangConfig, + tp_size: int, + n_gpus_per_node: int, + ): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.config = sglang_config + self.tp_size = tp_size + self.server_process = None + self.n_gpus_per_node = n_gpus_per_node + + def run(self): + gpus_per_server = len(os.getenv("CUDA_VISIBLE_DEVICES").split(",")) + server_local_idx = ( + int(os.getenv("CUDA_VISIBLE_DEVICES").split(",")[0]) // gpus_per_server + ) + n_servers_per_node = max(1, self.n_gpus_per_node // gpus_per_server) + ports_per_server = 40000 // n_servers_per_node + port_range = ( + server_local_idx * ports_per_server + 10000, + (server_local_idx + 1) * ports_per_server + 10000, + ) + server_port, dist_init_port = find_free_ports(2, port_range) + + dist_init_addr = f"localhost:{dist_init_port}" + host_ip = gethostip() + + cmd = SGLangConfig.build_cmd( + self.config, + tp_size=self.tp_size, + base_gpu_id=0, + host=host_ip, + port=server_port, + dist_init_addr=dist_init_addr, + ) + self.server_process = launch_server_cmd(cmd) + wait_for_server(f"http://{host_ip}:{server_port}") + + name = names.gen_servers(self.experiment_name, self.trial_name) + name_resolve.add_subentry(name, f"{host_ip}:{server_port}") + + logger.info(f"SGLang server launched at: http://{host_ip}:{server_port}") + return_code = self.server_process.wait() + logger.info( + f"SGLang server at http://{host_ip}:{server_port} exits, returncode={return_code}" + ) + + def __del__(self): + if self.server_process and self.server_process.poll() is None: + logger.info("Terminating SGLang server process...") + self.server_process.terminate() + self.server_process.wait() + logger.info("SGLang server process terminated.") + + +def main(argv): + config, _ = parse_cli_args(argv) + config.sglang = to_structured_cfg(config.sglang, SGLangConfig) + config.cluster = to_structured_cfg(config.cluster, ClusterSpecConfig) + config.cluster.name_resolve = to_structured_cfg( + config.cluster.name_resolve, NameResolveConfig + ) + name_resolve.reconfigure(config.cluster.name_resolve) + + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + assert allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG + tp_size = allocation_mode.gen_tp_size + + sglang_server = SGLangServerWrapper( + config.experiment_name, + config.trial_name, + config.sglang, + tp_size, + n_gpus_per_node=config.n_gpus_per_node, + ) + sglang_server.run() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/arealite/launcher/slurm.py b/arealite/launcher/slurm.py new file mode 100644 index 0000000000..2bbf20f93d --- /dev/null +++ b/arealite/launcher/slurm.py @@ -0,0 +1,528 @@ +import getpass +import os +import re +import subprocess +import sys +import time +from typing import Dict, List, Optional, Tuple + +import realhf.base.logging as logging +from arealite.api.cli_args import ( + ClusterSpecConfig, + LauncherConfig, + SGLangConfig, + parse_cli_args, + to_structured_cfg, +) +from arealite.api.io_struct import AllocationMode, AllocationType +from arealite.utils.launcher import ( + get_env_vars, + validate_config_for_distributed_launcher, + wait_sglang_server_addrs, +) +from arealite.utils.slurm import ( + APPTAINER_CMD_TEMPLATE, + SBATCH_SCRIPT_TEMPLATE, + SRUN_CMD_TEMPLATE, + cancel_jobs, + query_jobs, +) +from realhf.base import logging, name_resolve, names +from realhf.scheduler.client import JobException, JobInfo, JobState + +logger = logging.getLogger("SlurmLauncher") + +SLURM_WAIT_CHECK_TIME_INTERVAL = 5 + + +class SlurmLauncher: + def __init__( + self, experiment_name: str, trial_name: str, fileroot: str, container_type: str + ): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.fileroot = fileroot + self.container_type = container_type + + # slurm_job_id -> JobInfo + self.jobs: Dict[int, JobInfo] = {} + self.job_names = [] + + @property + def run_name(self) -> str: + """Returns the run name of this launcher.""" + return f"{self.experiment_name}_{self.trial_name}" + + def slurm_name(self, job_name: str) -> str: + """Returns the slurm name of a job.""" + return f"{self.experiment_name}_{self.trial_name}:{job_name}" + + def log_path_of(self, job_name: str) -> str: + log_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(log_path, exist_ok=True) + return os.path.join(log_path, f"{job_name}.log") + + def sbatch_path_of(self, job_name: str) -> str: + sbatch_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(sbatch_path, exist_ok=True) + return os.path.join(sbatch_path, f"{job_name}.sh") + + def submit(self, job_name, cmd, **kwargs): + """Submits and launch a job with SBATCH. + + Args: + cmd (str or List[str]): The core command to be executed. + """ + return self.submit_array(job_name, cmd, count=1, **kwargs) + + def find_job_id(self, job_name: str): + job_name = self.slurm_name(job_name) + for job_id, job_info in self.jobs.items(): + if job_info.name == job_name: + return job_id + return None + + def submit_array( + self, + job_name: str, + cmd: List[str] | str, + count: int, + nodes: int, + n_gpus_per_node: int, + cpus_per_task: int, + mem_per_task: int, # MB + container_image: str, + srun_additional_args: str = "", + container_mounts: Optional[str] = None, + env_vars: Optional[Dict] = None, + nodelist: Optional[str] = None, + exclude: Optional[str] = None, + ): + """Submits and launch a job array with SBATCH. + Note that a job array has one (unique) slurm name, and one (unique) slurm id. + + Args: + job_name (str): The job name of the job array. The actual slurm name will be + `_:`. + cmd (str or List[str]): The core command to be executed. + count (int): The number of jobs in the array. + """ + assert job_name not in self.job_names, ( + f"Job {job_name} is already submitted. " + "Please use a different job name or stop the existing job." + ) + if isinstance(cmd, str): + cmd = [cmd] + assert len(cmd) == count, ( + f"Command length {len(cmd)} does not match the job count {count}. " + "Please provide a command for each job in the array." + ) + assert count % nodes == 0, ( + f"Job count {count} must be divisible by the number of nodes {nodes}. " + "Please adjust the job count or the number of nodes." + ) + ntasks_per_node = count // nodes + assert n_gpus_per_node % ntasks_per_node == 0, ( + "GPUs must be evenly distributed across tasks. " + f"Current #GPUs per node {n_gpus_per_node}, #tasks per node {ntasks_per_node}." + ) + + mem_per_cpu = mem_per_task // cpus_per_task # MB per CPU + mem_per_node = ( + mem_per_task * count // nodes + 1024 * 10 + ) # make sure slurm does not run out of resources + + sbatch_options = [ + f"--job-name={self.slurm_name(job_name)}", + f"--output={self.log_path_of(job_name)}", + "--open-mode=append", + "--no-requeue", + f"--nodes={nodes}-{nodes}", + f"--ntasks-per-node={ntasks_per_node}", + f"--gres=gpu:{n_gpus_per_node}", + f"--cpus-per-task={cpus_per_task}", + f"--mem={mem_per_node}M", + ] + + if nodelist: + sbatch_options.append(f"--nodelist={nodelist}") + if exclude: + sbatch_options.append(f"--exclude={exclude}") + + sbatch_options_str = "\n".join([f"#SBATCH {opt}" for opt in sbatch_options]) + + if env_vars is None: + env_vars = dict() + n_gpus_per_task = n_gpus_per_node // ntasks_per_node + assert ( + "CUDA_VISIBLE_DEVICES" not in env_vars + ), "CUDA_VISIBLE_DEVICES should be automatically resolved by Launcher instead of manually assigned." + + srun_cmds = [] + for i in range(count): + # resolve CUDA_VISIBLE_DEVICES for each task + gpu_id_start = (i % ntasks_per_node) * n_gpus_per_task + gpu_id_end = ((i % ntasks_per_node) + 1) * n_gpus_per_task + node_id = i // ntasks_per_node + _env_vars = { + **env_vars, + "CUDA_VISIBLE_DEVICES": ",".join( + str(x) for x in range(gpu_id_start, gpu_id_end) + ), + } + # Prepare the command for each job in the array + job_cmd = cmd[i] + + if self.container_type == "apptainer": + env_string = " ".join( + "--env {}={}".format(k, v) for k, v in _env_vars.items() + ) + apptainer_cmd = APPTAINER_CMD_TEMPLATE.format( + container_mounts=container_mounts or "", + container_env_strings=env_string, + container_image=container_image, + cmd=job_cmd, + ) + srun_cmd = SRUN_CMD_TEMPLATE.format( + additional_args=srun_additional_args, + nodes=1, + ntasks=1, + node_id=node_id, + n_gpus_per_node=n_gpus_per_task, + cpus_per_task=cpus_per_task, + mem_per_cpu=mem_per_cpu, + cmd=apptainer_cmd, + ) + elif self.container_type == "none": + env_string = "--export=" + ",".join( + "{}={}".format(k, v) for k, v in _env_vars.items() + ) + srun_additional_args = srun_additional_args + " " + env_string + srun_cmd = SRUN_CMD_TEMPLATE.format( + additional_args=srun_additional_args, + nodes=1, + ntasks=1, + node_id=node_id, + n_gpus_per_node=n_gpus_per_task, + cpus_per_task=cpus_per_task, + mem_per_cpu=mem_per_cpu, + cmd=job_cmd, + ) + else: + raise ValueError( + f"Unsupported container type: {self.container_type}. " + "Supported types are 'apptainer' and 'none'." + ) + srun_cmds.append(srun_cmd) + + srun_cmds = "\n".join(srun_cmds) + sbatch_script = SBATCH_SCRIPT_TEMPLATE.format( + sbatch_options=sbatch_options_str, + srun_additional_args=srun_additional_args, + srun_cmds=srun_cmds, + ) + sbatch_file_path = self.sbatch_path_of(f"{job_name}") + with open(sbatch_file_path, "w") as f: + f.write(sbatch_script) + + # Submit the job + try: + output = ( + subprocess.check_output(["sbatch", sbatch_file_path]) + .decode("utf-8") + .strip() + ) + logger.info( + f"Submitted Slurm job {self.slurm_name(job_name)} to scheduler. To check the output, run \n\t`tail -f {self.log_path_of(job_name)}`." + ) + except subprocess.CalledProcessError as e: + logger.warning( + f"Failed to submit job {self.slurm_name(job_name)}. " + f"For debugging, please make sure your sbatch command works " + f"and check generated sbatch file on {sbatch_file_path}." + ) + logger.error(f"Error message: {e}") + return + + match = re.search(r"Submitted batch job (\d+)", output) + slurm_job_id = int(match.group(1)) if match else None + if slurm_job_id is None: + logger.warning( + f"Failed to obtain job id for job {self.slurm_name(job_name)}. " + f"sbatch output: {output}" + ) + return + + assert isinstance(slurm_job_id, int) + self.jobs[slurm_job_id] = JobInfo( + name=self.slurm_name(job_name), + state=JobState.PENDING, + slurm_id=slurm_job_id, + ) + self._update_all() + + def stop(self, job_name, force=False): + """Stops a running job. + + Raises exception if there is no such job, but passes if the job + has stopped either successfully or not. + + Args: + job_name: The job name of the job array to stop. + The actual slurm job name will be `_:`. + """ + signal = "SIGKILL" if force else "SIGTERM" + job_id = self.find_job_id(job_name) + if not job_id: + return + return cancel_jobs(slurm_ids=[job_id], signal=signal) + + def stop_all(self, force=False): + """Stops all running jobs.""" + signal = "SIGKILL" if force else "SIGTERM" + return cancel_jobs(slurm_ids=list(self.jobs.keys()), signal=signal) + + def find(self, job_name) -> JobInfo | None: + """Gets the status of a job of this job. + + Args: + job_name: The job name of the job array to find. + The actual slurm job name will be `_:`. + + Returns: + A JobInfo if the job is found, or None otherwise. + """ + self._update_all() + job_id = self.find_job_id(job_name) + return self.jobs[job_id] if job_id else None + + def find_all(self, job_name_regex=".*") -> List[JobInfo]: + """Finds jobs. + + Args: + job_name_regex: job name regex. + + Returns: + A list of found JobInfo. + """ + self._update_all() + infos = [] + for r in self.jobs.values(): + job_name = r.name.split(":")[-1] # Extract the job name from slurm name + if re.fullmatch(job_name_regex, job_name): + infos.append(r) + return infos + + def _find_job_with_status( + self, + status: List[JobState], + ) -> List[JobInfo]: + """Finds jobs with the given status. + + Args: + status: A list of JobState to filter jobs. + + Returns: + A list of JobInfo with the given status. + """ + self._update_all() + return [r for r in self.jobs.values() if r.state in status] + + def wait( + self, + timeout=None, + check_status: Tuple[JobState, ...] = ( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + ), + remove_status: Tuple[JobState, ...] = (JobState.COMPLETED,), + update=False, + ): + """Waits until all jobs submitted via this client instance finish.""" + # begin wait + deadline = None if timeout is None else time.time() + timeout + + num_jobs_left = len(self.jobs) + left = list(self.jobs.keys()) + logger.info( + f"Waiting for {num_jobs_left} jobs. Jobs IDs: " + f"{','.join(sorted([str(x.slurm_id) for x in self.jobs.values()]))}." + ) + while len(left) > 0: + if len(left) < num_jobs_left: + num_jobs_left = len(left) + logger.info(f"Waiting for {num_jobs_left} jobs.") + if deadline is not None and time.time() > deadline: + raise TimeoutError( + f"Timeout waiting for {num_jobs_left} jobs. Job ID: " + f"{','.join(sorted([str(x.slurm_id) for x in self.jobs.values()]))}." + ) + self._update_all() + left = list(self.jobs.keys()) + for slurm_id in list(left): + slurm_info = self.jobs[slurm_id] + if slurm_info.slurm_id is None: + continue + if slurm_info.state in check_status: + raise JobException( + run_name=self.run_name, + worker_type=slurm_info.name, + host=slurm_info.host, + reason=slurm_info.state, + ) + if slurm_info.state in remove_status: + logger.info( + f"Job {slurm_info.name} is {slurm_info.state}. (Removed)" + ) + left.remove(slurm_id) + if update: + self.jobs.pop(slurm_info.slurm_id) + time.sleep(SLURM_WAIT_CHECK_TIME_INTERVAL) + + def _update_all(self): + """Updates the status of all jobs.""" + try: + slurm_infos = query_jobs(slurm_ids=list(self.jobs.keys())) + for slurm_info in slurm_infos: + assert slurm_info.slurm_id is not None + self.jobs[slurm_info.slurm_id] = slurm_info + except subprocess.CalledProcessError: + logger.warning( + "Calling squeue failed. Check slurm manually if you continue to see this warning." + ) + + +def slurm_main(): + config, _ = parse_cli_args(sys.argv[2:]) + config.launcher = to_structured_cfg(config.launcher, LauncherConfig) + config.cluster = to_structured_cfg(config.cluster, ClusterSpecConfig) + config.sglang = to_structured_cfg(config.sglang, SGLangConfig) + validate_config_for_distributed_launcher(config) + + name_resolve.reconfigure(config.cluster.name_resolve) + name_resolve.clear_subtree( + names.trial_root( + experiment_name=config.experiment_name, trial_name=config.trial_name + ) + ) + + n_nodes = config.n_nodes + n_gpus_per_node = config.n_gpus_per_node + + launcher = SlurmLauncher( + experiment_name=config.experiment_name, + trial_name=config.trial_name, + fileroot=config.cluster.fileroot, + container_type=config.launcher.slurm.container_type, + ) + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + sglang_cmds = [] + sglang_addrs = [] + n_sglang_nodes = 0 + if allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG: + # Launcher should launch SGLang servers according to allocation mode. + sglang_tp_size = allocation_mode.gen_tp_size + n_sglang_servers = allocation_mode.gen_dp_size + n_sglang_nodes = allocation_mode.gen_world_size // n_gpus_per_node + + base_seed = config.sglang.random_seed + sglang_server_cmd_template = f"python3 -m arealite.launcher.sglang_server {' '.join(sys.argv[2:])} sglang.random_seed={{seed}}" + for i in range(n_sglang_servers): + sglang_cmd = sglang_server_cmd_template.format( + seed=base_seed + i, + ) + sglang_cmds.append(sglang_cmd) + + launcher.submit_array( + job_name="llm_server", + cmd=sglang_cmds, + count=n_sglang_servers, + nodes=n_sglang_nodes, + n_gpus_per_node=config.cluster.n_gpus_per_node, + cpus_per_task=config.launcher.inference_server_cpus_per_gpu + * sglang_tp_size, + mem_per_task=config.launcher.inference_server_mem_per_gpu * sglang_tp_size, + srun_additional_args=config.launcher.slurm.srun_additional_args, + container_image=config.launcher.slurm.inference_server_image, + container_mounts=config.launcher.slurm.mount, + env_vars=get_env_vars( + config.cluster.cluster_name, + config.launcher.inference_server_env_vars, + ), + ) + # Get SGLang server addresses by name resolve + try: + sglang_addrs = wait_sglang_server_addrs( + config.experiment_name, + config.trial_name, + n_sglang_servers, + ) + except TimeoutError as e: + launcher.stop_all(force=True) + raise e + + trainer_n_nodes = n_nodes - n_sglang_nodes + # Here $head_node_ip is the IP address of the first node in the job array. + # $trainer_port is a free port on the head node. + # Both of them are obtained in by the SBATCH script. + trainer_cmd_template = ( + f"torchrun --nnodes={{nnodes}} --nproc-per-node={{nproc_per_node}} --node-rank {{node_rank}} " + f"--master-addr $head_node_ip --master-port $trainer_port {' '.join(sys.argv[1:])}" + ) + + trainer_cmds = [] + for i in range(trainer_n_nodes): + # In slurm, we launch trainer in the granularity of nodes with torchrun command. + trainer_cmds.append( + trainer_cmd_template.format( + nnodes=trainer_n_nodes, + nproc_per_node=config.cluster.n_gpus_per_node, + node_rank=i, + ) + ) + + if not config.server_only: + # launch trainers + launcher.submit_array( + job_name="trainer", + cmd=trainer_cmds, + count=trainer_n_nodes, + nodes=trainer_n_nodes, + n_gpus_per_node=config.cluster.n_gpus_per_node, + cpus_per_task=config.launcher.trainer_cpus_per_gpu + * config.cluster.n_gpus_per_node, + mem_per_task=config.launcher.trainer_mem_per_gpu + * config.cluster.n_gpus_per_node, + container_image=config.launcher.slurm.trainer_image, + srun_additional_args=config.launcher.slurm.srun_additional_args, + container_mounts=config.launcher.slurm.mount, + env_vars=dict( + **get_env_vars( + config.cluster.cluster_name, + config.launcher.trainer_env_vars, + ), + AREAL_LLM_SERVER_ADDRS=",".join(sglang_addrs), + ), + ) + + try: + launcher.wait( + check_status=( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + JobState.COMPLETED, + ), + remove_status=(), + ) + except (KeyboardInterrupt, JobException, TimeoutError) as e: + launcher.stop_all(force=True) + raise e + + +if __name__ == "__main__": + # usage: python -m arealite.launcher.slurm \ + # --config [] + slurm_main() diff --git a/arealite/utils/launcher.py b/arealite/utils/launcher.py new file mode 100644 index 0000000000..dfe7327d1f --- /dev/null +++ b/arealite/utils/launcher.py @@ -0,0 +1,111 @@ +import getpass +import os +import pathlib +import time +from typing import Dict, Optional + +from arealite.api.io_struct import AllocationMode, AllocationType +from realhf.base import logging, name_resolve, names + +logger = logging.getLogger("Launcher Utils") + +LOCAL_CACHE_DIR = "/tmp/arealite" +PYTORCH_KERNEL_CACHE_PATH = ( + f"{LOCAL_CACHE_DIR}/.cache/{getpass.getuser()}/torch/kernels/" +) +TRITON_CACHE_PATH = f"{LOCAL_CACHE_DIR}/.cache/{getpass.getuser()}/triton/" +os.makedirs(PYTORCH_KERNEL_CACHE_PATH, exist_ok=True) +os.makedirs(TRITON_CACHE_PATH, exist_ok=True) +BASE_ENVIRONS = { + "TOKENIZERS_PARALLELISM": "true", + "PYTORCH_KERNEL_CACHE_PATH": PYTORCH_KERNEL_CACHE_PATH, + "TRITON_CACHE_DIR": TRITON_CACHE_PATH, + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "PYTHONPATH": str(pathlib.Path(__file__).resolve().parent.parent.parent), +} +NA132_ENVIRONS = { + "NCCL_SOCKET_IFNAME": "bond0", + "NCCL_NET_PLUGIN": "", + "NCCL_IB_GID_INDEX": "3", + "NCCL_IB_TIMEOUT": "2", + "NCCL_IB_RETRY_CNT": "7", + "NCCL_IB_SL": "5", + "NCCL_IB_TC": "136", + "NCCL_IB_HCA": "mlx5_bond", + "NCCL_IB_QPS_PER_CONNECTION": "8", + "NCCL_SET_THREAD_NAME": "1", + "NCCL_DEBUG": "WARN", + "NCCL_DEBUG_SUBSYS": "INIT,TUNING,GRAPH", +} +SGLANG_SERVER_WAIT_TIMEOUT_SECONDS = 180 + + +def get_env_vars( + cluster_name: str, additional_env_vars: Optional[str] = None +) -> Dict[str, str]: + """Returns the environment variables for the cluster.""" + _additional_env_vars = ( + dict(item.split("=") for item in additional_env_vars.split(",")) + if additional_env_vars + else dict() + ) + if cluster_name == "na132": + return {**BASE_ENVIRONS, **NA132_ENVIRONS, **_additional_env_vars} + else: + return {**BASE_ENVIRONS, **_additional_env_vars} + + +def wait_sglang_server_addrs( + experiment_name: str, + trial_name: str, + n_sglang_servers: int, +): + # Get SGLang slurm nodes, find the hosts + name = names.gen_servers(experiment_name, trial_name) + start = time.perf_counter() + while True: + sglang_addrs = name_resolve.get_subtree(name) + if len(sglang_addrs) >= n_sglang_servers: + logger.info( + f"Found {len(sglang_addrs)} SGLang servers: {', '.join(sglang_addrs)}" + ) + break + + time.sleep(1) + if time.perf_counter() - start > SGLANG_SERVER_WAIT_TIMEOUT_SECONDS: + raise TimeoutError( + f"Timeout waiting for SGLang servers to be ready. " + f"Expected {n_sglang_servers} servers, found {len(sglang_addrs)}." + ) + return sglang_addrs + + +def validate_config_for_distributed_launcher(config): + n_nodes = config.n_nodes + n_gpus_per_node = config.n_gpus_per_node + if n_gpus_per_node < config.cluster.n_gpus_per_node: + raise ValueError( + f"Ray/Slurm Launcher requires at least {config.cluster.n_gpus_per_node} (#GPUs per node) GPU. For usecases of less GPUs, use LocalLauncher instead." + ) + elif n_gpus_per_node > config.cluster.n_gpus_per_node: + raise ValueError( + f"#GPU per node required by experiment ({n_gpus_per_node}) is larger than #GPU per node in the cluster ({config.cluster.n_gpus_per_node})." + ) + + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + assert ( + allocation_mode.gen_world_size + allocation_mode.train_world_size + == n_nodes * n_gpus_per_node + ), ( + f"#GPUs required for allocation mode {allocation_mode.gen_world_size + allocation_mode.train_world_size} " + f"is not equal to #GPUs in the config {n_nodes * n_gpus_per_node}." + ) + if allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG: + # Launcher should launch SGLang servers according to allocation mode. + assert ( + allocation_mode.gen_pp_size == 1 + ), "Pipeline generation in SGLang is not supported for now." + assert ( + allocation_mode.gen_tp_size <= config.cluster.n_gpus_per_node + ), "Currently only support SGLang TP size less <= #GPUs per node." diff --git a/arealite/utils/ray.py b/arealite/utils/ray.py new file mode 100644 index 0000000000..f049a72bad --- /dev/null +++ b/arealite/utils/ray.py @@ -0,0 +1,23 @@ +import ray +from ray.util.placement_group import PlacementGroup +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from arealite.utils.network import find_free_ports, gethostip + + +def get_placement_group_master_ip_and_port(placement_group: PlacementGroup): + def _master_ip_and_port(): + host_ip = gethostip() + port = find_free_ports(1, (10000, 60000))[0] + return host_ip, port + + future = ray.remote( + num_cpus=1, + num_gpus=0, + memory=10 * 1024 * 1024, # Convert MB to bytes + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=placement_group, + placement_group_bundle_index=0, + ), + )(_master_ip_and_port).remote() + return ray.get(future) diff --git a/arealite/utils/slurm.py b/arealite/utils/slurm.py new file mode 100644 index 0000000000..f876c5e84a --- /dev/null +++ b/arealite/utils/slurm.py @@ -0,0 +1,154 @@ +import subprocess +from typing import List, Literal, Optional + +from realhf.base import logging +from realhf.scheduler.client import JobInfo, JobState + +logger = logging.getLogger("Slurm Utils") + + +SQUEUE_FIELDS = [ + "JobID", + "State", + "SubmitTime", + "StartTime", + "Name", + "NodeList", + "UserName", + "MaxCPUs", + "cpus-per-task", + "NumTasks", + "tres-alloc", +] +STATUS_MAPPING = { + "RUNNING": JobState.RUNNING, + "COMPLETING": JobState.RUNNING, + "PENDING": JobState.PENDING, + "CANCELLED": JobState.CANCELLED, + "FAILED": JobState.FAILED, + "COMPLETED": JobState.COMPLETED, + "OUT_OF_MEMORY": JobState.FAILED, + "DEADLINE": JobState.COMPLETED, + "TIMEOUT": JobState.COMPLETED, +} + +SBATCH_SCRIPT_TEMPLATE = """#!/bin/bash +{sbatch_options} + +# Getting the node names +nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +echo nodes=$nodes + +nodes_array=($nodes) +echo node_array=$nodes_array + +head_node=${{nodes_array[0]}} +echo head_node=$head_node + +# Getting the head node IP address +head_node_ip=$(srun {srun_additional_args} --nodes=1 --ntasks=1 -n1 -c1 --mem=10M --nodelist="$head_node" hostname --ip-address) +echo head_node_ip=$head_node_ip + +# Find a free port on the head node +# Wonderful linux command to find a random free port (between 10000 and 60000) by deepseek +trainer_port=$(srun {srun_additional_args} --nodes=1 --ntasks=1 -n1 -c1 --mem=10M --nodelist="$head_node" bash -c "comm -23 <(seq 10000 60000 | sort) <(ss -tan | awk '{{print $4}}' | cut -d':' -f2 | grep '[0-9]\\{{1,5\\}}' | sort -u) | shuf | head -n 1") +echo trainer_port=$trainer_port + +# srun commands +{srun_cmds} + +wait +""" + +SRUN_CMD_TEMPLATE: str = """srun {additional_args} \\ + --nodelist=${{nodes_array[{node_id}]}} --nodes={nodes} --ntasks={ntasks} \\ + --gres=gpu:{n_gpus_per_node} --cpus-per-task={cpus_per_task} --mem-per-cpu={mem_per_cpu}M \\ + {cmd} & +""" + +APPTAINER_CMD_TEMPLATE: str = """singularity exec --no-home --writable-tmpfs --nv --pid \\ + --bind {container_mounts} \\ + {container_env_strings} \\ + {container_image} \\ + {cmd}""" + + +def cancel_jobs( + slurm_names: Optional[List[str]] = None, + slurm_ids: Optional[List[int]] = None, + signal: Literal["SIGINT", "SIGKILL"] = "SIGKILL", +): + assert ( + slurm_names is not None or slurm_ids is not None + ), "Must specify slurm_names or slurm_ids." + assert not ( + slurm_names and slurm_ids + ), "Cannot specify both slurm_names and slurm_ids." + cmd = ["scancel", "-s", signal] + if slurm_names is not None: + cmd += ["-n", ",".join(slurm_names)] + elif slurm_ids is not None: + cmd += [",".join(str(s) for s in slurm_ids)] + subprocess.check_call(cmd) + logger.info( + f"Cancelled Slurm job with signal {signal}: " + f"slurm identifiers {slurm_names if slurm_ids is None else slurm_ids}. CMD: {cmd}" + ) + + +def query_jobs( + slurm_names: Optional[List[str]] = None, + slurm_ids: Optional[List[int]] = None, + status: str = "all", + delimiter: str = "__PSI__", +) -> List[JobInfo]: + squeue_format = f":.{delimiter},".join(SQUEUE_FIELDS) + cmd = ["squeue", "-O", squeue_format, f"-t{status}"] + if slurm_names is not None: + cmd += ["-n", ",".join(slurm_names)] + if slurm_ids is not None: + cmd += ["-j", ",".join([str(s) for s in slurm_ids])] + + output = ( + subprocess.check_output(cmd, stderr=subprocess.DEVNULL).decode("ascii").strip() + ) + rs = [] + for line in output.split("\n")[1:]: + job_id, state, submit_time, start_time, slurm_name, nodelist, *_ = line.split( + delimiter + ) + rs.append( + JobInfo( + name=slurm_name, + state=STATUS_MAPPING[state], + host=nodelist, + submit_time=submit_time, + start_time=start_time, + slurm_id=int(job_id.strip()), + ) + ) + return rs + + +def parse_slurm_nodelist(nodelist: str) -> List[str]: + return ( + subprocess.check_output( + [ + "scontrol", + "show", + "hostnames", + nodelist, + ] + ) + .decode("utf-8") + .strip() + .split("\n") + ) + + +def get_slurm_host_ip(node: str, srun_addtional_args: str): + try: + cmd = f"srun {srun_addtional_args} --immediate=1 --nodes=1 --ntasks=1 -n1 -c1 --mem=10M --nodelist={node} hostname --ip-address" + return subprocess.check_output(cmd.split(" ")).decode("utf-8").strip() + except subprocess.CalledProcessError: + logger.warning(f"Get slurm host IP for node {node} failed.") diff --git a/examples/arealite/gsm8k_grpo.py b/examples/arealite/gsm8k_grpo.py index 1841343da5..15c24fc15a 100644 --- a/examples/arealite/gsm8k_grpo.py +++ b/examples/arealite/gsm8k_grpo.py @@ -76,8 +76,8 @@ def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **k return int(sol.strip() == ans.strip()) -def main_grpo(): - config, _ = load_expr_config(sys.argv[1:], GRPOConfig) +def main(args): + config, _ = load_expr_config(args, GRPOConfig) config: GRPOConfig rank = int(os.getenv("RANK")) @@ -253,4 +253,4 @@ def evaluate_fn(): if __name__ == "__main__": - main_grpo() + main(sys.argv[1:]) diff --git a/examples/arealite/gsm8k_sft.py b/examples/arealite/gsm8k_sft.py index c1d8735127..eaacd7564d 100644 --- a/examples/arealite/gsm8k_sft.py +++ b/examples/arealite/gsm8k_sft.py @@ -35,8 +35,8 @@ def get_gsm8k_dataset(split, tokenizer, rank, world_size): return process_gsm8k_sft_dataset(dataset, tokenizer) -def main_sft(): - config, _ = load_expr_config(sys.argv[1:], SFTConfig) +def main(args): + config, _ = load_expr_config(args, SFTConfig) config: SFTConfig rank = int(os.getenv("RANK")) @@ -121,4 +121,4 @@ def evaluate_fn(): if __name__ == "__main__": - main_sft() + main(sys.argv[1:]) From 2f1b679ad660acfed490ab821089e04cbb613246 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 21 Jul 2025 11:26:36 +0800 Subject: [PATCH 09/16] PullRequest: 392 [lite] Fix several bugs regarding RL learning and add an example to reproduce boba-math results. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-boba of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/392 Reviewed-by: 晓雷 * support fsdp engine and sglang remote engine * minor fix * . * refactor trainer * add close * rm mb_spec * . * fix * . * qwen2 grpo works * fix * fix * async works * fix * slurm launcher not tested * fix arg parse * . * sglang server wrapper * . * . * slurm run * ready for boba * debug * 32k run * . * . * fix * . * . * . * . * . * fix * . * fix * . * . * . * . * fix * . * . * . * . * . * . * . * refactor train engine * refactor train engine * . * fix update weight error * . * . * match train * format * . * fix * seems to work * . * . * . * . --- arealite/api/cli_args.py | 6 - arealite/engine/base_hf_engine.py | 3 +- arealite/engine/fsdp_engine.py | 3 +- arealite/engine/sft/lm_engine.py | 12 +- arealite/launcher/ray.py | 4 +- arealite/launcher/sglang_server.py | 2 +- arealite/launcher/slurm.py | 4 +- arealite/utils/data.py | 12 +- arealite/utils/launcher.py | 13 +- arealite/utils/padding.py | 48 ---- arealite/utils/stats_logger.py | 2 +- arealite/workflow/rlvr.py | 52 +++- examples/arealite/boba.py | 310 ++++++++++++++++++++++ examples/arealite/configs/boba.yaml | 141 ++++++++++ examples/arealite/configs/gsm8k_grpo.yaml | 2 +- examples/arealite/gsm8k_grpo.py | 61 ++--- 16 files changed, 541 insertions(+), 134 deletions(-) delete mode 100644 arealite/utils/padding.py create mode 100644 examples/arealite/boba.py create mode 100644 examples/arealite/configs/boba.yaml diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 4962cfab55..3b2afa55d0 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -794,12 +794,6 @@ class BaseExperimentConfig: default_factory=ClusterSpecConfig, metadata={"help": "Cluster specification. Mainly used by slurm."}, ) - n_nodes: int = field( - default=1, metadata={"help": "Number of nodes for experiment."} - ) - n_gpus_per_node: int = field( - default=8, metadata={"help": "Number of GPUs per node for this experiment."} - ) allocation_mode: str = field( default="", metadata={ diff --git a/arealite/engine/base_hf_engine.py b/arealite/engine/base_hf_engine.py index 91cef08d03..06a42d6c4d 100644 --- a/arealite/engine/base_hf_engine.py +++ b/arealite/engine/base_hf_engine.py @@ -263,6 +263,7 @@ def train_batch( # Scale loss for accumulation # Revert gradient averaging across dp ranks + # FIXME: should be DP size loss_scale *= self.world_size loss *= loss_scale @@ -283,8 +284,6 @@ def train_batch( update_successful = True current_lr = self.lr_scheduler.get_last_lr()[0] - # Optimizer step - self.optimizer.step() return dict( update_successful=float(update_successful), grad_norm=float(grad_norm) if grad_norm is not None else float("nan"), diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index 34131ec9b7..4f644967a0 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -188,6 +188,7 @@ def train_batch( for i, (pad_length, padded_mb_input, mb_input) in enumerate( zip(mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs) ): + self.model.set_requires_gradient_sync(i == len(mb_list.mbs) - 1) outputs = self.model(**padded_mb_input) logits = outputs.logits.squeeze(0) @@ -214,8 +215,6 @@ def train_batch( update_successful = True current_lr = self.lr_scheduler.get_last_lr()[0] - # Optimizer step - self.optimizer.step() return dict( update_successful=float(update_successful), grad_norm=float(grad_norm) if grad_norm is not None else float("nan"), diff --git a/arealite/engine/sft/lm_engine.py b/arealite/engine/sft/lm_engine.py index 18bee784b2..eaa8fc3499 100644 --- a/arealite/engine/sft/lm_engine.py +++ b/arealite/engine/sft/lm_engine.py @@ -58,15 +58,9 @@ def compute_packed_sft_loss(logits: torch.Tensor, input_: TensorDict) -> torch.T cu_seqlens.shape[0] - 1, device=logits.device, dtype=torch.float64 ) for i in range(cu_seqlens.shape[0] - 1): - m = loss_mask[cu_seqlens[i] - i : cu_seqlens[i + 1] - i - 1] - logp = logprobs[cu_seqlens[i] - i : cu_seqlens[i + 1] - i - 1] - assert cu_seqlens[i + 1] - i - 1 <= logprobs.shape[0], ( - cu_seqlens, - logprobs.shape, - ) - seqlogp[i] = torch.where(m, logp.detach(), 0.0).sum() / ( - m.numel() - m.count_nonzero() - ) + m = loss_mask[cu_seqlens[i] : cu_seqlens[i + 1]] + logp = logprobs[cu_seqlens[i] : cu_seqlens[i + 1]] + seqlogp[i] = torch.where(m, logp.detach(), 0.0).sum() / (m.count_nonzero()) ## Loggin stats stats_tracker.denominator( diff --git a/arealite/launcher/ray.py b/arealite/launcher/ray.py index 52813aeec8..0344e983e2 100644 --- a/arealite/launcher/ray.py +++ b/arealite/launcher/ray.py @@ -315,8 +315,8 @@ def ray_main(): ) ) - n_nodes = config.n_nodes - n_gpus_per_node = config.n_gpus_per_node + n_nodes = config.cluster.n_nodes + n_gpus_per_node = config.cluster.n_gpus_per_node launcher = RayLauncher( experiment_name=config.experiment_name, trial_name=config.trial_name, diff --git a/arealite/launcher/sglang_server.py b/arealite/launcher/sglang_server.py index 916d4dfea5..2e88c043ab 100644 --- a/arealite/launcher/sglang_server.py +++ b/arealite/launcher/sglang_server.py @@ -188,7 +188,7 @@ def main(argv): config.trial_name, config.sglang, tp_size, - n_gpus_per_node=config.n_gpus_per_node, + n_gpus_per_node=config.cluster.n_gpus_per_node, ) sglang_server.run() diff --git a/arealite/launcher/slurm.py b/arealite/launcher/slurm.py index 2bbf20f93d..ca91112185 100644 --- a/arealite/launcher/slurm.py +++ b/arealite/launcher/slurm.py @@ -407,8 +407,8 @@ def slurm_main(): ) ) - n_nodes = config.n_nodes - n_gpus_per_node = config.n_gpus_per_node + n_nodes = config.cluster.n_nodes + n_gpus_per_node = config.cluster.n_gpus_per_node launcher = SlurmLauncher( experiment_name=config.experiment_name, diff --git a/arealite/utils/data.py b/arealite/utils/data.py index b5f317416d..d6a6677805 100644 --- a/arealite/utils/data.py +++ b/arealite/utils/data.py @@ -399,9 +399,15 @@ def pad_packed_tensor_dict( padded_data[key] = new_max_seqlen elif torch.is_tensor(value) and value.numel() == total_length: # Pad the tensor to the new total length - padded_tensor = torch.nn.functional.pad( - value, (0, pad_length), value=pad_value - ) + if key == "position_ids": + # transformers will compute flash-attn arguments (e.g., cu_seqlens_q) + # according to this position ids. + pad = torch.arange(pad_length, dtype=torch.long, device=value.device) + padded_tensor = torch.cat([value, pad]) + else: + padded_tensor = torch.nn.functional.pad( + value, (0, pad_length), value=pad_value + ) padded_data[key] = padded_tensor else: padded_data[key] = value diff --git a/arealite/utils/launcher.py b/arealite/utils/launcher.py index dfe7327d1f..3931cc0e4e 100644 --- a/arealite/utils/launcher.py +++ b/arealite/utils/launcher.py @@ -81,17 +81,8 @@ def wait_sglang_server_addrs( def validate_config_for_distributed_launcher(config): - n_nodes = config.n_nodes - n_gpus_per_node = config.n_gpus_per_node - if n_gpus_per_node < config.cluster.n_gpus_per_node: - raise ValueError( - f"Ray/Slurm Launcher requires at least {config.cluster.n_gpus_per_node} (#GPUs per node) GPU. For usecases of less GPUs, use LocalLauncher instead." - ) - elif n_gpus_per_node > config.cluster.n_gpus_per_node: - raise ValueError( - f"#GPU per node required by experiment ({n_gpus_per_node}) is larger than #GPU per node in the cluster ({config.cluster.n_gpus_per_node})." - ) - + n_nodes = config.cluster.n_nodes + n_gpus_per_node = config.cluster.n_gpus_per_node allocation_mode = config.allocation_mode allocation_mode = AllocationMode.from_str(allocation_mode) assert ( diff --git a/arealite/utils/padding.py b/arealite/utils/padding.py deleted file mode 100644 index 615132248c..0000000000 --- a/arealite/utils/padding.py +++ /dev/null @@ -1,48 +0,0 @@ -from typing import List - -import torch -from tensordict import TensorDict - - -def concat_padded_tensors( - tensor_dicts: List[TensorDict], pad_value: float = 0.0 -) -> TensorDict: - """Concatenate and pad tensors from multiple padded tensor dictionaries.""" - if not tensor_dicts: - return TensorDict() - - batch_sizes = [tuple(d.batch_size) for d in tensor_dicts] - new_batch_size = [sum(x[0] for x in batch_sizes), *batch_sizes[0][1:]] - - # Find max sequence length across all dictionaries - assert all("attention_mask" in td for td in tensor_dicts) - max_length = max([x["attention_mask"].shape[1] for x in tensor_dicts]) - result = {} - # Process each key - for key in tensor_dicts[0].keys(): - tensors_to_concat = [] - for tensor_dict in tensor_dicts: - tensor = tensor_dict[key] - # Skip 1D tensors like rewards - if len(tensor.shape) == 1: - tensors_to_concat.append(tensor) - continue - current_length = tensor.shape[1] - if current_length < max_length: - # Pad tensor to max_length - pad_width = max_length - current_length - if key == "attention_mask": - # Pad attention mask with 0s - padding = torch.zeros( - (tensor.shape[0], pad_width), dtype=tensor.dtype - ) - else: - # Pad feature tensors with pad_value - padding = torch.full( - (tensor.shape[0], pad_width), pad_value, dtype=tensor.dtype - ) - tensor = torch.cat([tensor, padding], dim=1) - tensors_to_concat.append(tensor) - - result[key] = torch.cat(tensors_to_concat, dim=0) - return TensorDict(result, batch_size=new_batch_size) diff --git a/arealite/utils/stats_logger.py b/arealite/utils/stats_logger.py index 49c0d5594a..4c2612b7a7 100644 --- a/arealite/utils/stats_logger.py +++ b/arealite/utils/stats_logger.py @@ -73,7 +73,7 @@ def commit(self, epoch: int, step: int, global_step: int, data: Dict | List[Dict ) if isinstance(data, Dict): data = [data] - log_step = max(global_step, self._last_commit_step) + log_step = max(global_step, self._last_commit_step + 1) for i, item in enumerate(data): self.info(f"Stats ({i+1}/{len(data)}):") self.print_stats(item) diff --git a/arealite/workflow/rlvr.py b/arealite/workflow/rlvr.py index 3ce55dfc00..5404972b5f 100644 --- a/arealite/workflow/rlvr.py +++ b/arealite/workflow/rlvr.py @@ -1,11 +1,14 @@ import asyncio +import os import uuid +import colorama import torch from tensordict import TensorDict from transformers import PreTrainedTokenizerFast from arealite.api.cli_args import GenerationHyperparameters +from arealite.api.engine_api import InferenceEngine from arealite.api.io_struct import LLMRequest from arealite.api.workflow_api import RolloutWorkflow from arealite.utils.data import concat_padded_tensors @@ -18,13 +21,17 @@ def __init__( gconfig: GenerationHyperparameters, tokenizer: PreTrainedTokenizerFast, enable_thinking: bool, + dump_dir: str | None = None, ): self.reward_fn = reward_fn self.gconfig = gconfig self.tokenizer = tokenizer self.enable_thinking = enable_thinking + self.dump_dir = dump_dir + if self.dump_dir is not None and not os.path.exists(self.dump_dir): + os.makedirs(self.dump_dir, exist_ok=True) - async def arun_episode(self, engine, data): + async def arun_episode(self, engine: InferenceEngine, data): input_ids = self.tokenizer.apply_chat_template( data["messages"], tokenize=True, @@ -39,6 +46,12 @@ async def arun_episode(self, engine, data): ) resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) + version = engine.get_version() + prompt_strs = [] + completions_strs = [] + rewards = [] + seqlens = [] + results = [] for resp in resps: seq = resp.input_tokens + resp.output_tokens @@ -46,13 +59,19 @@ async def arun_episode(self, engine, data): loss_mask = [0] * resp.input_len + [1] * resp.output_len versions = [-1] * resp.input_len + resp.output_versions + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + prompt_strs.append(prompt_str) + completions_strs.append(completions_str) + seqlens.append(len(seq)) reward = self.reward_fn( - prompt=self.tokenizer.decode(input_ids), - completions=self.tokenizer.decode(resp.output_tokens), + prompt=prompt_str, + completions=completions_str, prompt_ids=resp.input_tokens, completion_ids=resp.output_tokens, **data, ) + rewards.append(reward) res = dict( # unsqueeze to add an additional batch dimension input_ids=torch.tensor(seq).unsqueeze(0), @@ -65,4 +84,31 @@ async def arun_episode(self, engine, data): ) results.append(TensorDict(res, batch_size=[1])) + if self.dump_dir is not None: + os.makedirs(os.path.join(self.dump_dir, str(version)), exist_ok=True) + # Get the unique identifier for this prompt + qid = None + for key in ["query_id", "id", "qid"]: + qid = data.get(key, None) + if qid is not None: + break + qid = qid or uuid.uuid4().hex + + # Dump rollout to file + with open( + os.path.join(self.dump_dir, str(version), f"{qid}.txt"), "a" + ) as f: + n_samples = self.gconfig.n_samples + for i, (p, c, r, sl) in enumerate( + zip(prompt_strs, completions_strs, rewards, seqlens) + ): + info = "\n".join( + [ + f"idx: {i + 1} / {n_samples}, seqlen: {sl}, reward is {r}.", + f"prompt is \n{colorama.Fore.YELLOW + colorama.Style.DIM}{p}{colorama.Style.RESET_ALL}", + f"sequence is: \n{colorama.Fore.YELLOW + colorama.Style.DIM}{c}{colorama.Style.RESET_ALL}", + ] + ) + f.write(info + "\n") + return concat_padded_tensors(results) diff --git a/examples/arealite/boba.py b/examples/arealite/boba.py new file mode 100644 index 0000000000..4211d11972 --- /dev/null +++ b/examples/arealite/boba.py @@ -0,0 +1,310 @@ +import asyncio +import os +import shutil +import sys +import uuid + +import colorama +import torch +import torch.distributed as dist +from datasets import load_dataset +from datasets.distributed import split_dataset_by_node +from tensordict import TensorDict +from torchdata.stateful_dataloader import StatefulDataLoader +from transformers import PreTrainedTokenizerFast + +from arealite.api.cli_args import ( + GenerationHyperparameters, + GRPOConfig, + load_expr_config, +) +from arealite.api.io_struct import FinetuneSpec, LLMRequest, WeightUpdateMeta +from arealite.api.workflow_api import RolloutWorkflow +from arealite.engine.ppo.actor import FSDPPPOActor +from arealite.engine.sglang_remote import RemoteSGLangEngine +from arealite.utils.data import concat_padded_tensors +from arealite.utils.device import log_gpu_stats +from arealite.utils.saver import Saver +from arealite.utils.stats_logger import StatsLogger +from realhf.api.core.data_api import load_hf_tokenizer +from realhf.base import logging, seeding, stats_tracker + +logger = logging.getLogger("boba math") + + +class RLVRWorkflow(RolloutWorkflow): + def __init__( + self, + reward_fn, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast, + dump_dir: str | None = None, + ): + self.reward_fn = reward_fn + self.gconfig = gconfig + self.tokenizer = tokenizer + self.dump_dir = dump_dir + if self.dump_dir is not None and not os.path.exists(self.dump_dir): + os.makedirs(self.dump_dir, exist_ok=True) + + async def arun_episode(self, engine, data): + input_ids = self.tokenizer.encode(data["prompt"]) + n_samples = self.gconfig.n_samples + req = LLMRequest( + rid=uuid.uuid4().hex, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) + + version = engine.get_version() + prompt_strs = [] + completions_strs = [] + rewards = [] + seqlens = [] + + results = [] + for resp in resps: + seq = resp.input_tokens + resp.output_tokens + logprobs = [0.0] * resp.input_len + resp.output_logprobs + loss_mask = [0] * resp.input_len + [1] * resp.output_len + versions = [-1] * resp.input_len + resp.output_versions + + prompt_str = data["prompt"] + completions_str = self.tokenizer.decode(resp.output_tokens) + prompt_strs.append(prompt_str) + completions_strs.append(completions_str) + seqlens.append(len(seq)) + reward = self.reward_fn( + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + rewards.append(reward) + res = dict( + # unsqueeze to add an additional batch dimension + input_ids=torch.tensor(seq).unsqueeze(0), + loss_mask=torch.tensor(loss_mask).unsqueeze(0), + logprobs=torch.tensor(logprobs).unsqueeze(0), + versions=torch.tensor(versions).unsqueeze(0), + attention_mask=torch.ones(len(seq), dtype=torch.bool).unsqueeze(0), + # reward + rewards=torch.tensor([float(reward)]), + ) + results.append(TensorDict(res, batch_size=[1])) + + if self.dump_dir is not None: + os.makedirs(os.path.join(self.dump_dir, str(version)), exist_ok=True) + # Get the unique identifier for this prompt + qid = None + for key in ["query_id", "id", "qid"]: + qid = data.get(key, None) + if qid is not None: + break + qid = qid or uuid.uuid4().hex + + # Dump rollout to file + with open( + os.path.join(self.dump_dir, str(version), f"{qid}.txt"), "a" + ) as f: + n_samples = self.gconfig.n_samples + for i, (p, c, r, sl) in enumerate( + zip(prompt_strs, completions_strs, rewards, seqlens) + ): + info = "\n".join( + [ + f"idx: {i + 1} / {n_samples}, seqlen: {sl}, reward is {r}.", + f"prompt is \n{colorama.Fore.YELLOW + colorama.Style.DIM}{p}{colorama.Style.RESET_ALL}", + f"sequence is: \n{colorama.Fore.YELLOW + colorama.Style.DIM}{c}{colorama.Style.RESET_ALL}", + ] + ) + f.write(info + "\n") + + return concat_padded_tensors(results) + + +def get_boba_math_dataset(tokenizer, rank, world_size): + dataset = load_dataset( + path="json", + split="train", + data_files="/storage/openpsi/users/xushusheng.xss/training_data/boba_106k_0319.jsonl", + ) + dataset = dataset.filter(lambda x: len(tokenizer.encode(x["prompt"])) <= 1024) + return split_dataset_by_node(dataset, rank=rank, world_size=world_size) + + +def boba_reward_fn( + prompt, completions, prompt_ids, completion_ids, query_id, solutions, **kwargs +): + from pebble import ProcessExpired, ProcessPool + + from realhf.impl.dataset.math_parser import process_results + + jobs = [] + with ProcessPool(max_workers=1) as executor: + for sol in solutions: + job = executor.schedule( + process_results, args=[completions, sol], timeout=15 + ) + jobs.append(job) + + label = 0 + for job in jobs: + try: + x = job.result() + except TimeoutError: + # print("[debug: timeout]") + logger.warning(f"Timeout occurred while justifying the math answer.") + x = (0, "timeout", "timeout") + except ProcessExpired as e: + logger.warning(f"Process terminated abnormally: {e}") + x = (0, "error", "error") + except Exception as e: + logger.warning(f"Other error occurred: {e.__class__.__name__}, {e}") + x = (0, "error", "error") + label = label or x[0] + return label + + +def main(args): + config, _ = load_expr_config(args, GRPOConfig) + config: GRPOConfig + + rank = int(os.getenv("RANK")) + world_size = int(os.getenv("WORLD_SIZE")) + tokenizer = load_hf_tokenizer(config.tokenizer_path) + + seeding.set_random_seed(config.seed, key=f"trainer{rank}") + + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + get_boba_math_dataset(tokenizer, rank, world_size), + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.train_dataset.drop_last, + ) + ft_spec = FinetuneSpec( + total_train_epochs=config.total_train_epochs, + dataset_size=len(train_dataloader) * config.train_dataset.batch_size, + train_batch_size=config.train_dataset.batch_size, + ) + + # Initialize inference engine + rollout = RemoteSGLangEngine(config.rollout) + rollout.initialize(None, ft_spec) + + # Initialize train engine + actor = FSDPPPOActor(config=config.actor) + actor.initialize(None, ft_spec) + ref = None + if config.actor.kl_ctl > 0 and config.ref is not None: + ref = FSDPPPOActor(config=config.ref) + ref.initialize(None, ft_spec) + + # Create rollout workflow + if tokenizer.pad_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.pad_token_id) + if tokenizer.eos_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.eos_token_id) + workflow = RLVRWorkflow( + reward_fn=boba_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + dump_dir=os.path.join( + StatsLogger.get_log_path(config.stats_logger), "generated" + ), + ) + + # Run training. + saver = Saver(config.saver, ft_spec, for_recover=False) + logger = StatsLogger(config.stats_logger, ft_spec) + + total_epochs = config.total_train_epochs + steps_per_epoch = len(train_dataloader) + max_steps = total_epochs * steps_per_epoch + + logger.info(f"total_epochs={total_epochs} step_per_epoch={steps_per_epoch}") + data_generator = iter(train_dataloader) + for global_step in range(max_steps): + epoch = global_step // steps_per_epoch + step = global_step % steps_per_epoch + + with stats_tracker.record_timing("rollout"): + if config.async_training: + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) + else: + try: + data = next(data_generator) + except StopIteration: + data_generator = iter(train_dataloader) + data = next(data_generator) + batch = rollout.rollout_batch(data, workflow=workflow) + + batch = batch.to(actor.device) + # Create barrier to synchronize all rollout processes. + dist.barrier() + torch.cuda.synchronize() + + if config.actor.recompute_logprob or config.actor.use_decoupled_loss: + with stats_tracker.record_timing("recompute_logp"): + logp = actor.compute_logp(batch) + batch["prox_logp"] = logp + log_gpu_stats("recompute logp") + + if ref is not None: + with stats_tracker.record_timing("ref_logp"): + batch["ref_logp"] = ref.compute_logp(batch) + log_gpu_stats("ref logp") + + with stats_tracker.record_timing("compute_advantage"): + actor.compute_advantages(batch) + log_gpu_stats("compute advantages") + + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("grpo_actor"), + ): + stats = actor.ppo_update(batch) + actor.step_lr_scheduler() + log_gpu_stats("ppo update") + + with stats_tracker.record_timing("update_weights"): + path = os.path.join( + Saver.get_save_checkpoint_root(config.saver), + "update_weights", + str(global_step + 1), + ) + meta = WeightUpdateMeta( + type="disk", + path=path, + alloc_mode=None, + comm_backend=None, + model_version=global_step + 1, + ) + if dist.get_rank() == 0: + future = rollout.update_weights(meta) + actor.upload_weights(meta) + if dist.get_rank() == 0: + future.result() + shutil.rmtree(path, ignore_errors=True) + dist.barrier() + torch.cuda.synchronize() + rollout.set_version(global_step + 1) + + with stats_tracker.record_timing("save"): + saver.save(actor, epoch, step, global_step) + + logger.commit(epoch, step, global_step, stats) + + logger.close() + rollout.destroy() + if ref is not None: + ref.destroy() + actor.destroy() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/arealite/configs/boba.yaml b/examples/arealite/configs/boba.yaml new file mode 100644 index 0000000000..3924a14cb9 --- /dev/null +++ b/examples/arealite/configs/boba.yaml @@ -0,0 +1,141 @@ +experiment_name: lite-boba-math +trial_name: run1 + +cluster: + n_nodes: 16 + n_gpus_per_node: 8 + cluster_name: na132 + fileroot: /storage/openpsi/experiments + name_resolve: + type: nfs + nfs_record_root: /storage/openpsi/experiments/name_resolve/lite-boba-math + etcd3_addr: etcd-client.openpsi-etcd.svc.sigma-na130-lingbo.na130.wl-robby.local:2379 + +seed: 1 +total_train_epochs: 10 +total_train_steps: null +tokenizer_path: ${actor.path} +allocation_mode: sglang.d96p1t1+d32p1t1 +async_training: true + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 400 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 4 + enable_rollout_tracing: true + +gconfig: + n_samples: 16 + min_new_tokens: 0 + max_new_tokens: 30720 + greedy: false + temperature: 1.0 + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: /storage/openpsi/models/deepseek-ai__DeepSeek-R1-Distill-Qwen-1.5B/ + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 32768 + optimizer: + type: adam + lr: 1e-5 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + backend: fsdp + + group_size: ${gconfig.n_samples} + group_adv_norm: false + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 4 + recompute_logprob: true + use_decoupled_loss: true + behav_imp_weight_cap: 5.0 + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 32768 + optimizer: null + backend: fsdp + +# SGLang +server_only: false +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.9 + +# datasets +train_dataset: + batch_size: 512 + shuffle: true + pin_memory: true + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +checkpointer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: online + +# Launcher +launcher: + inference_server_cpus_per_gpu: 15 + inference_server_mem_per_gpu: 153600 + trainer_cpus_per_gpu: 15 + trainer_mem_per_gpu: 153600 + slurm: + mount: /storage:/storage + trainer_image: /storage/openpsi/images/arealite-20250712-update-hf-xet.sif + inference_server_image: /storage/openpsi/images/arealite-20250712-update-hf-xet.sif \ No newline at end of file diff --git a/examples/arealite/configs/gsm8k_grpo.yaml b/examples/arealite/configs/gsm8k_grpo.yaml index ac7488571a..9f0107e6f8 100644 --- a/examples/arealite/configs/gsm8k_grpo.yaml +++ b/examples/arealite/configs/gsm8k_grpo.yaml @@ -41,7 +41,7 @@ actor: max_tokens_per_mb: 10240 optimizer: type: adam - lr: 2e-6 + lr: 1e-5 weight_decay: 0.01 beta1: 0.9 beta2: 0.999 diff --git a/examples/arealite/gsm8k_grpo.py b/examples/arealite/gsm8k_grpo.py index 15c24fc15a..6fbccb762e 100644 --- a/examples/arealite/gsm8k_grpo.py +++ b/examples/arealite/gsm8k_grpo.py @@ -1,5 +1,5 @@ import os -import re +import shutil import sys import torch @@ -18,7 +18,9 @@ from arealite.utils.stats_logger import StatsLogger from arealite.workflow.rlvr import RLVRWorkflow from realhf.api.core.data_api import load_hf_tokenizer -from realhf.base import stats_tracker +from realhf.base import logging, seeding, stats_tracker + +logger = logging.getLogger("GSM8K grpo") def process_gsm8k_rl_dataset(dataset: Dataset): @@ -36,44 +38,10 @@ def get_gsm8k_dataset(split, rank, world_size): return process_gsm8k_rl_dataset(dataset) -# Adapted from verl. -def extract_solution(solution_str, method="strict") -> str | None: - assert method in ["strict", "flexible"] - - final_answer = None - if method == "strict": - # this also tests the formatting of the model - solutions = re.findall("#### (\\-?[0-9\\.\\,]+)", solution_str) - if len(solutions) == 0: - final_answer = None - else: - # take the last solution - final_answer = solutions[-1].replace(",", "").replace("$", "") - elif method == "flexible": - answer = re.findall("(\\-?[0-9\\.\\,]+)", solution_str) - final_answer = None - if len(answer) == 0: - # no reward is there is no answer - pass - else: - invalid_str = ["", "."] - # find the last number that is not '.' - for final_answer in reversed(answer): - if final_answer not in invalid_str: - break - return final_answer - - def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs): - from realhf.impl.dataset.math_parser import extract_answer + from realhf.impl.dataset.math_parser import process_results - sol = extract_answer(completions, data_name="math") - ans = extract_solution(solution_str=answer, method="strict") - if sol is None: - return 0 - if ans is None: - return 0 - return int(sol.strip() == ans.strip()) + return int(process_results(completions, answer)[0]) def main(args): @@ -84,6 +52,8 @@ def main(args): world_size = int(os.getenv("WORLD_SIZE")) tokenizer = load_hf_tokenizer(config.tokenizer_path) + seeding.set_random_seed(config.seed, key=f"trainer{rank}") + # Create dataset and dataloaders train_dataloader = StatefulDataLoader( get_gsm8k_dataset("train", rank, world_size), @@ -133,6 +103,9 @@ def main(args): gconfig=config.gconfig, tokenizer=tokenizer, enable_thinking=False, + dump_dir=os.path.join( + StatsLogger.get_log_path(config.stats_logger), "generated" + ), ) # Run training. @@ -190,13 +163,14 @@ def main(args): log_gpu_stats("ppo update") with stats_tracker.record_timing("update_weights"): + path = os.path.join( + Saver.get_save_checkpoint_root(config.saver), + "update_weights", + str(global_step + 1), + ) meta = WeightUpdateMeta( type="disk", - path=os.path.join( - Saver.get_save_checkpoint_root(config.saver), - "update_weights", - str(global_step), - ), + path=path, alloc_mode=None, comm_backend=None, model_version=global_step + 1, @@ -206,6 +180,7 @@ def main(args): actor.upload_weights(meta) if dist.get_rank() == 0: future.result() + shutil.rmtree(path, ignore_errors=True) dist.barrier() torch.cuda.synchronize() rollout.set_version(global_step + 1) From ab5db3fa55837bde909d7cac8ffbfe387e96f8ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 21 Jul 2025 15:09:50 +0800 Subject: [PATCH 10/16] . --- arealite/workflow/multi_turn.py | 87 +++++++ docs/customization/agent.md | 398 ++++++++++++++++++++++++++++---- 2 files changed, 437 insertions(+), 48 deletions(-) create mode 100644 arealite/workflow/multi_turn.py diff --git a/arealite/workflow/multi_turn.py b/arealite/workflow/multi_turn.py new file mode 100644 index 0000000000..eec8aa1e77 --- /dev/null +++ b/arealite/workflow/multi_turn.py @@ -0,0 +1,87 @@ +import uuid + +import torch +from transformers import PreTrainedTokenizerFast + +from arealite.api.cli_args import GenerationHyperparameters +from arealite.api.engine_api import InferenceEngine +from arealite.api.io_struct import LLMRequest +from arealite.api.workflow_api import RolloutWorkflow +from arealite.utils.data import concat_padded_tensors + + +class MultiTurnWorkflow(RolloutWorkflow): + def __init__( + self, + reward_fn, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast, + max_turns: int, + turn_discount: float, + ): + self.reward_fn = reward_fn + self.gconfig = gconfig + self.tokenizer = tokenizer + self.max_turns = max_turns + self.turn_discount = turn_discount + + async def arun_episode(self, engine: InferenceEngine, data): + # Placeholders for the results + seq, logprobs, loss_mask, versions = [], [], [], [] + messages = data["messages"] + # Run multi-turn rollout until correct + t = reward = 0 + discount = 0 + rid = uuid.uuid4().hex + while reward == 0 and t < self.max_turns: + # Amend a prompt if the previous answer is incorrect + if t > 0: + messages += [ + {"role": "asistant", "content": completions_str}, + { + "role": "user", + "content": "Your answer is not correct. Please try to answer it again.", + }, + ] + # Convert the prompt into input_ids + input_ids = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + # Send generate request to get the response. + req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + resp = await engine.agenerate(req) + # compute reward: 1 for correct and 0 otherwise + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + # Amend results + input_len = len(resp.input_tokens) - len(seq) + seq += resp.input_tokens[-input_len:] + resp.output_tokens + logprobs += [0.0] * input_len + resp.output_logprobs + loss_mask += [0] * input_len + [1] * resp.output_len + versions += [-1] * input_len + resp.output_versions + # Increase counter + t += 1 + discount *= self.turn_discount + res = dict( + seq=torch.tensor(seq), + logprobs=torch.tensor(logprobs), + loss_mask=torch.tensor(loss_mask), + versions=torch.tensor(versions), + rewards=torch.tensor([float(reward * discount)]), + attetion_mask=torch.ones(len(seq), dtype=torch.bool), + ) + res = {k: v.unsqueeze(0) for k, v in res.items()} + return concat_padded_tensors([res]) diff --git a/docs/customization/agent.md b/docs/customization/agent.md index 8b13b440e5..2fb974d621 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -1,16 +1,283 @@ # Rollout and Agentic RL -This guide provides an example of modifying the rollout behavior for PPO training. +This guide demonstrates how to customize rollout behavior for PPO training by +implementing a multi-turn math agent that uses end-to-end reinforcement learning. Our +example agent will continuously try to solve a math problem until it reaches the correct +answer. -In particular, we implement a multi-turn math agent using end-to-end RL. The math agent will continuously attempt to think through and solve math problems until it reaches the correct answer. +## Approach: Using AReaLite (Recommended) -## Define Your Agent +The complete implementation is placed at `arealite/workflow/multi_turn.py`. -Create a new file under `realhf/impl/agent/`, for example, `math_multi_turn_agent.py`. Your `Agent` must implement the interface defined in `realhf/api/core/agent.py`, which requires implementing a single method: `collect_trajectory`. +### Step 1: Define Your Workflow + +AReaLite takes a flexible approach to agent definition. Rather than using rigid `Agent` +classes that might limit your agentic capabilities, AReaLite captures all rollout +behavior in a `RolloutWorkflow` class. This design gives you complete freedom to +customize your agent's behavior. + +```python +# arealite/api/workflow_api.py +class RolloutWorkflow: + async def arun_episode( + self, engine: InferenceEngine, data: Dict[str, Any] + ) -> TensorDict: + """Run a single episode of the workflow. + + See concrete example implementations under the `arealite/workflow` directory. + """ + raise NotImplementedError() +``` + +The workflow exposes a single `arun_episode` method that runs and collects data from a +single episode. This method takes two key arguments: + +1. **InferenceEngine**: Provides the `agenerate` method for generating responses to user + inputs +1. **data**: The prompt data loaded from your RL dataset + +Within this method, you have complete control over how your agent and environment +interact. + +#### Setting Up the Multi-Turn Math Workflow + +Let's build a multi-turn rollout workflow for solving math problems. First, we'll define +the `__init__` method to capture the utilities we need during rollout: + +> **Note**: You have complete flexibility in defining the `__init__` method. Pass any +> arguments needed to construct your workflow. If you want to use tools, pass the +> corresponding environment here so your agent can call it in the `arun_episode` method. + +```python +class MultiTurnWorkflow(RolloutWorkflow): + def __init__( + self, + reward_fn, + gconfig: GenerationHyperparameters, # aka sampling_params + tokenizer: PreTrainedTokenizerFast, + max_turns: int, + turn_discount: float, + ): + self.reward_fn = reward_fn + self.gconfig = gconfig + self.tokenizer = tokenizer + self.max_turns = max_turns + # Discount rewards if the agent takes longer to find the correct answer + self.turn_discount = turn_discount +``` + +#### Implementing the Episode Logic + +Now let's implement the `arun_episode` method. We'll start by tokenizing the prompt data +and converting it into an `LLMRequest` object for the inference engine: + +```python +class MultiTurnWorkflow(RolloutWorkflow): + # ... __init__ method above ... + + async def arun_episode(self, engine: InferenceEngine, data): + # Initialize result containers + seq, logprobs, loss_mask, versions = [], [], [], [] + messages = data["messages"] + # Run multi-turn rollout until we get the correct answer + t = reward = 0 + discount = 1.0 + rid = uuid.uuid4().hex + while reward == 0 and t < self.max_turns: + # Convert the conversation into input tokens + input_ids = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + # Generate response from the model + req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + resp = await engine.agenerate(req) + # ... continue processing ... +``` + +> **Note**: This example accesses the "messages" key from the prompt data to get +> OpenAI-compatible messages. This isn't mandatory—the key and prompt format depend +> entirely on your implementation. For instance, if your dataset stores prompt strings +> in a "prompt" column, you could get input token IDs via +> `self.tokenizer.encode(data["prompt"])`. + +> **Note**: The `rid` field in `LLMRequest` is the request ID. Requests with the same ID +> will reuse the LLM inference server's KV caches for efficiency. + +#### Handling Multi-Turn Conversations + +Next, we'll evaluate whether the current answer is correct using our `reward_fn`. This +function should return 1 for correct answers and 0 otherwise. When the answer is wrong, +we'll apply a discount, add feedback to the conversation, and let the model try again: + +```python +class MultiTurnWorkflow(RolloutWorkflow): + # ... previous methods ... + + async def arun_episode(self, engine: InferenceEngine, data): + # ... initialization code ... + while reward == 0 and t < self.max_turns: + # Add feedback if the previous answer was incorrect + if t > 0: + messages += [ + {"role": "assistant", "content": completions_str}, + { + "role": "user", + "content": "Your answer is not correct. Please try to answer it again." + }, + ] + # Generate response (code from above) + # ... + # Evaluate the response + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + # Update counters + t += 1 + discount *= self.turn_discount +``` + +#### Reward Function Signature + +For convenience when switching between different reward functions, we recommend +following this pre-defined signature: + +```python +def reward_fn( + prompt: str, + completions: str, + prompt_ids: List[int], + completion_ids: List[int], + **kwargs, +): + """Reward function for evaluating agent performance. + + This signature is recommended for compatibility with predefined workflows, + but you can modify it freely in custom implementations. + + Args: + prompt: The task description string + completions: The agent's response string + prompt_ids: Tokenized prompt + completion_ids: Tokenized response + **kwargs: Additional dataset attributes (solutions, input_outputs, etc.) + + Returns: + float: Reward value (typically 1.0 for correct, 0.0 for incorrect) + """ +``` + +While this signature is convenient, there are no strict restrictions on reward functions +in custom workflows—modify them as needed for your specific use case. + +#### Collecting Training Data + +Finally, let's complete the implementation by collecting trajectories in the +`TensorDict` format: + +```python +class MultiTurnWorkflow(RolloutWorkflow): + # ... previous methods ... + + async def arun_episode(self, engine: InferenceEngine, data): + # ... episode logic above ... + + while reward == 0 and t < self.max_turns: + # ... generation and evaluation ... + + # Collect trajectory data + input_len = len(resp.input_tokens) - len(seq) + seq += resp.input_tokens[-input_len:] + resp.output_tokens + logprobs += [0.0] * input_len + resp.output_logprobs + loss_mask += [0] * input_len + [1] * resp.output_len + versions += [-1] * input_len + resp.output_versions + + # Package results + res = dict( + input_ids=torch.tensor(seq), + logprobs=torch.tensor(logprobs), + loss_mask=torch.tensor(loss_mask), + versions=torch.tensor(versions), + rewards=torch.tensor(float(reward * discount)), + attention_mask=torch.ones(len(seq), dtype=torch.bool), + ) + res = {k: v.unsqueeze(0) for k, v in res.items()} + return concat_padded_tensors([res]) +``` + +> **Important**: The returned `TensorDict` must follow HuggingFace's padded data format, +> where each tensor has shape `[batch_size, sequence_length, *]`. This allows AReaLite +> to automatically batch multiple trajectories for the training engine. Since this +> example returns a single trajectory, we use `unsqueeze(0)` to create a size-1 batch. + +> **Note**: There are no restrictions on the keys in your `TensorDict`—different +> algorithms require different keys. This example targets the GRPO algorithm, so we +> include `input_ids`, `loss_mask`, `attention_mask`, and `logprobs` (needed for +> computing importance ratios). + +### Step 2: Training with Your Custom Workflow + +Using your custom workflow is straightforward—just construct it in your training script +and pass it to the `rollout_batch` or `prepare_batch` method: + +```python +def main(args): + # ... setup code ... + + # Create your custom workflow + workflow = MultiTurnWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + turn_discount=0.9, + max_turns=5, + ) + + # Run training—no other changes needed! + data_generator = iter(train_dataloader) + for global_step in range(max_steps): + with stats_tracker.record_timing("rollout"): + if config.async_training: + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) + else: + try: + data = next(data_generator) + except StopIteration: + data_generator = iter(train_dataloader) + data = next(data_generator) + batch = rollout.rollout_batch(data, workflow=workflow) + # ... continue with training loop ... +``` + +That's it! Your custom multi-turn math agent is now ready to train with reinforcement +learning. The workflow will automatically handle the multi-turn conversations, reward +computation, and data collection needed for effective RL training. + +## Alternative Approach: Using the Legacy Version (Not Recommended) + +While we strongly recommend using AReaLite for new projects, you might encounter legacy +code that uses the older Agent-based approach. Here's how it works for reference, though +we suggest migrating to the workflow-based system when possible. + +### Step 1: Define Your Agent Class + +Create a new file under `realhf/impl/agent/`, such as `math_multi_turn_agent.py`. Your +`Agent` must implement the interface defined in `realhf/api/core/agent.py`, which +requires a single method: `collect_trajectory`. ```python class MathMultiTurnAgent(Agent): - async def collect_trajectory( self, prompt: SequenceSample, @@ -18,65 +285,90 @@ class MathMultiTurnAgent(Agent): obs_queue: asyncio.Queue, act_queue: asyncio.Queue, ): + # Implementation goes here ... ``` -## Implement the `collect_trajectory` Logic +### Step 2: Implement the Trajectory Collection Logic -The `collect_trajectory` function takes a task prompt, an environment, and two queues as input, then produces several trajectories for the RL trainer. Within this function, you can create arbitrary data processing logic to produce the input for the inference engine (i.e., via `obs_queue`) and extract the action (i.e., via `act_queue`) from the generated tokens. +The `collect_trajectory` method takes a task prompt, an environment, and two +communication queues. Within this method, you control the data flow between your agent +and the inference engine using these queues: -In this example, the initial observation is the math problem itself. We put the token IDs and generation config into `obs_queue` and wait for the action produced by the inference engine from `act_queue`. After the inference engine returns, we extract the generated answers and send them to the environment. +- **obs_queue**: Send observations (token IDs and generation config) to the inference + engine +- **act_queue**: Receive actions (generated responses) from the inference engine + +Here's how the multi-turn conversation works: ```python for turn in range(self.num_turns): + # Send the current state to the inference engine await obs_queue.put((qid, token_ids, self.gconfig)) + + # Get the generated response act: BundledGenerationOutputs = await act_queue.get() - _, success, *_ = await env.step((qid, answers)) - ... + + # Evaluate the response through the environment + success, rewards = await env.step((qid, answers)) + # ... process results ... ``` -The environment is similar to a [gym environment](https://github.com/Farama-Foundation/Gymnasium), which defines two methods: `reset` and `step`. However, to maintain efficiency, we use an asynchronous implementation to avoid mutual blocking across different environment instances. +#### Environment Integration + +The environment follows a +[Gym-like interface](https://github.com/Farama-Foundation/Gymnasium) with `reset` and +`step` methods, but uses asynchronous implementations to prevent blocking across +different environment instances. -The math environment is stateless and essentially serves as a wrapper around the reward function: +For math problems, the environment is typically stateless and acts as a wrapper around +your reward function: ```python class MathCodeSingleStepEnv(EnvironmentService): - async def step(self, action: Tuple[str, List[str]]): qid, answers = action - ... - # Make `math_verify_call` async + # ... setup code ... + + # Run reward computation asynchronously format_rewards = await asyncio.to_thread( math_verify_call, answers, - ... + # ... other parameters ... ) return None, format_rewards, True, False, {} ``` -After `env.step` returns the reward for the current step, we can check whether the answer is correct. If not, we can append a user prompt and send it to `obs_queue` again to enter the next round. +#### Handling Multi-Turn Feedback + +After receiving the reward from `env.step`, check if the answer is correct. If not, +provide feedback and continue to the next turn: ```python for turn in range(self.num_turns): - ... - feedback = None + # ... generation and evaluation code ... + + # Provide feedback based on the result if success[0]: feedback = "Congratulations! You are correct!" else: feedback = "Unfortunately your answer is wrong. Let's try again." - + + # Format feedback as a user message feedback = "\n" + self.tokenizer.apply_chat_template( - [dict(content=feedback, role="user")], + [{"content": feedback, "role": "user"}], add_generation_prompt=True, tokenize=False, ) - feedback = self.tokenizer(feedback)["input_ids"] - token_ids.extend(feedback) + + # Add feedback tokens to the conversation + feedback_tokens = self.tokenizer(feedback)["input_ids"] + token_ids.extend(feedback_tokens) ``` -## Modify the Configuration +### Step 3: Register and Configure Your Agent -You're now close to running the end-to-end RL loop. The final step is to register and import your implementation, then modify the experiment configuration. +First, register your agent implementation: ```python # in realhf/impl/agent/math_multi_turn_agent.py @@ -88,48 +380,58 @@ register_agent("math-multi-turn", MathMultiTurnAgent) import realhf.impl.agent.math_multi_turn_agent ``` -In `realhf/experiments/async_exp/async_math_ppo.py`: +Then update your experiment configuration in +`realhf/experiments/async_exp/async_math_ppo.py`: -```diff +```python @dataclasses.dataclass class AsyncPPOMATHConfig(AsyncRLExperimentConfig, PPOMATHConfig): -+ # New CLI arguments are defined here -+ my_param: float = 1.0 + # Add any new CLI arguments your agent needs + my_param: float = 1.0 - # in realhf/experiments/async_exp/async_ppo_math_exp.py @property def agent(self) -> AgentAbstraction: return AgentAbstraction( -- "math-single-step", -+ "math-multi-turn", # Your registered name + "math-multi-turn", # Your registered agent name args=dict( -- ... -+ # Any configurations for your __init__ method -+ my_param=my_param, + # Pass any arguments needed for your __init__ method + my_param=self.my_param, + # ... other configuration ... ), ) @property def env(self) -> EnvServiceAbstraction: -- return EnvServiceAbstraction( -- "math-code-single-step", args=dict(dataset_path=self.dataset.path) -- ) -+ # Change to your customized environment if necessary -+ return EnvServiceAbstraction( -+ "my-env", args=dict(...) -+ ) + # Update to use your custom environment if needed + return EnvServiceAbstraction( + "math-code-single-step", + args=dict(dataset_path=self.dataset.path) + ) ``` -## Run Training +### Step 4: Run Training -Please follow the guide in [quickstart](../tutorial/quickstart.md). Generally, start your experiments by running: +Follow the standard training procedure outlined in the +[quickstart guide](../tutorial/quickstart.md). Launch your experiment with: ```bash -python3 training/main_async_ppo.py my_param=5.0 # and any additional CLI arguments +python3 training/main_async_ppo.py my_param=5.0 # plus any additional CLI arguments ``` -The training reward of our trial is shown below: +### Training Results + +Here's an example of the training reward curve from our multi-turn math agent: + +![Multi-turn Training Rewards](multiturn_reward.png) + +The agent successfully learns to solve math problems with improved accuracy over time, +demonstrating the effectiveness of the multi-turn approach. + +______________________________________________________________________ -![](multiturn_reward.png) +**Note**: While this legacy approach works, we strongly recommend using the AReaLite +workflow system for new projects. It provides better flexibility, cleaner abstractions, +and easier maintenance. Consider migrating existing legacy agents to the workflow-based +approach when possible. -Happy coding! \ No newline at end of file +Happy coding! From 415cc6b434f915db5e6ba37d4ca46184e1bf3d40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 21 Jul 2025 16:41:41 +0800 Subject: [PATCH 11/16] . --- docs/customization/agent.md | 87 +++--- docs/customization/algorithm.md | 481 +++++++++++++++++++++----------- 2 files changed, 364 insertions(+), 204 deletions(-) diff --git a/docs/customization/agent.md b/docs/customization/agent.md index 2fb974d621..5815e4c925 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -1,20 +1,20 @@ # Rollout and Agentic RL -This guide demonstrates how to customize rollout behavior for PPO training by -implementing a multi-turn math agent that uses end-to-end reinforcement learning. Our -example agent will continuously try to solve a math problem until it reaches the correct +This guide shows you how to create custom rollout behaviors for PPO training by building +a multi-turn math agent using end-to-end reinforcement learning. We'll walk through +creating an agent that keeps trying to solve math problems until it finds the correct answer. -## Approach: Using AReaLite (Recommended) +## Recommended Approach: AReaLite -The complete implementation is placed at `arealite/workflow/multi_turn.py`. +You can find the complete implementation in `arealite/workflow/multi_turn.py`. ### Step 1: Define Your Workflow -AReaLite takes a flexible approach to agent definition. Rather than using rigid `Agent` -classes that might limit your agentic capabilities, AReaLite captures all rollout -behavior in a `RolloutWorkflow` class. This design gives you complete freedom to -customize your agent's behavior. +AReaLite gives you flexibility in how you design your agents. Instead of rigid `Agent` +classes that might constrain your agent's capabilities, AReaLite captures all rollout +behavior in a `RolloutWorkflow` class. This approach lets you customize your agent's +behavior however you need. ```python # arealite/api/workflow_api.py @@ -29,8 +29,8 @@ class RolloutWorkflow: raise NotImplementedError() ``` -The workflow exposes a single `arun_episode` method that runs and collects data from a -single episode. This method takes two key arguments: +The workflow exposes an `arun_episode` method that runs and collects data from a single +episode. This method takes two key arguments: 1. **InferenceEngine**: Provides the `agenerate` method for generating responses to user inputs @@ -39,14 +39,19 @@ single episode. This method takes two key arguments: Within this method, you have complete control over how your agent and environment interact. +> **Note**: Each `arun_episode` call takes a single prompt and outputs the trajectories +> generated from that prompt—it's not batched. However, you can generate multiple +> trajectories from a single prompt (for example, with GRPO or tree search). + #### Setting Up the Multi-Turn Math Workflow Let's build a multi-turn rollout workflow for solving math problems. First, we'll define -the `__init__` method to capture the utilities we need during rollout: +the `__init__` method to set up what we need during rollout: -> **Note**: You have complete flexibility in defining the `__init__` method. Pass any -> arguments needed to construct your workflow. If you want to use tools, pass the -> corresponding environment here so your agent can call it in the `arun_episode` method. +> **Note**: You have complete flexibility in defining the `__init__` method. Pass +> whatever arguments you need to construct your workflow. If you want to use tools, pass +> the corresponding environment here so your agent can call it in the `arun_episode` +> method. ```python class MultiTurnWorkflow(RolloutWorkflow): @@ -100,20 +105,20 @@ class MultiTurnWorkflow(RolloutWorkflow): # ... continue processing ... ``` -> **Note**: This example accesses the "messages" key from the prompt data to get -> OpenAI-compatible messages. This isn't mandatory—the key and prompt format depend +> **Note**: This example uses the "messages" key from the prompt data to get +> OpenAI-compatible messages. This isn't required—the key and prompt format depend > entirely on your implementation. For instance, if your dataset stores prompt strings -> in a "prompt" column, you could get input token IDs via +> in a "prompt" column, you could get input token IDs with > `self.tokenizer.encode(data["prompt"])`. > **Note**: The `rid` field in `LLMRequest` is the request ID. Requests with the same ID -> will reuse the LLM inference server's KV caches for efficiency. +> will reuse the LLM inference server's KV caches for better efficiency. #### Handling Multi-Turn Conversations -Next, we'll evaluate whether the current answer is correct using our `reward_fn`. This -function should return 1 for correct answers and 0 otherwise. When the answer is wrong, -we'll apply a discount, add feedback to the conversation, and let the model try again: +Next, we'll check if the current answer is correct using our `reward_fn`. This function +should return 1 for correct answers and 0 otherwise. When the answer is wrong, we'll +apply a discount, add feedback to the conversation, and let the model try again: ```python class MultiTurnWorkflow(RolloutWorkflow): @@ -150,8 +155,8 @@ class MultiTurnWorkflow(RolloutWorkflow): #### Reward Function Signature -For convenience when switching between different reward functions, we recommend -following this pre-defined signature: +To make it easier to switch between different reward functions, we recommend following +this signature: ```python def reward_fn( @@ -178,8 +183,8 @@ def reward_fn( """ ``` -While this signature is convenient, there are no strict restrictions on reward functions -in custom workflows—modify them as needed for your specific use case. +While this signature is convenient, you're not restricted to it in custom +workflows—modify as needed for your specific use case. #### Collecting Training Data @@ -218,18 +223,18 @@ class MultiTurnWorkflow(RolloutWorkflow): > **Important**: The returned `TensorDict` must follow HuggingFace's padded data format, > where each tensor has shape `[batch_size, sequence_length, *]`. This allows AReaLite -> to automatically batch multiple trajectories for the training engine. Since this -> example returns a single trajectory, we use `unsqueeze(0)` to create a size-1 batch. +> to automatically batch multiple trajectories for training. Since this example returns +> a single trajectory, we use `unsqueeze(0)` to create a batch of size 1. -> **Note**: There are no restrictions on the keys in your `TensorDict`—different -> algorithms require different keys. This example targets the GRPO algorithm, so we -> include `input_ids`, `loss_mask`, `attention_mask`, and `logprobs` (needed for -> computing importance ratios). +> **Note**: You're not restricted to specific keys in your `TensorDict`—different +> algorithms need different keys. This example targets the GRPO algorithm, so we include +> `input_ids`, `loss_mask`, `attention_mask`, and `logprobs` (needed for computing +> importance ratios). ### Step 2: Training with Your Custom Workflow -Using your custom workflow is straightforward—just construct it in your training script -and pass it to the `rollout_batch` or `prepare_batch` method: +Using your custom workflow is straightforward—just create it in your training script and +pass it to the `rollout_batch` or `prepare_batch` method: ```python def main(args): @@ -260,15 +265,15 @@ def main(args): # ... continue with training loop ... ``` -That's it! Your custom multi-turn math agent is now ready to train with reinforcement -learning. The workflow will automatically handle the multi-turn conversations, reward +That's it! Your custom multi-turn math agent is now ready for reinforcement learning +training. The workflow will automatically handle the multi-turn conversations, reward computation, and data collection needed for effective RL training. -## Alternative Approach: Using the Legacy Version (Not Recommended) +## Legacy Approach: Agent-Based System -While we strongly recommend using AReaLite for new projects, you might encounter legacy -code that uses the older Agent-based approach. Here's how it works for reference, though -we suggest migrating to the workflow-based system when possible. +While we strongly recommend AReaLite for new projects, you might need to work with +legacy code that uses the older Agent-based approach. Here's how it works for reference, +though we suggest migrating to the workflow-based system when possible. ### Step 1: Define Your Agent Class @@ -433,5 +438,3 @@ ______________________________________________________________________ workflow system for new projects. It provides better flexibility, cleaner abstractions, and easier maintenance. Consider migrating existing legacy agents to the workflow-based approach when possible. - -Happy coding! diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index bf5030b767..8318e48d25 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -1,23 +1,215 @@ -# Training Algorithm +# Training Algorithm Implementation -An algorithm is encapsulated in a `ModelInterface`, which primarily defines three methods: +## Approach 1: Using AReaLite (Recommended) + +AReaLite structures RL algorithms around two core components: + +- **RolloutWorkflow**: Defines what data to generate during rollouts +- **TrainEngine**: Defines how to process the generated data for training + +We'll demonstrate this by implementing an RL algorithm similar to ReMax. + +### Step 1: Implementing the RolloutWorkflow + +The rollout workflow generates both greedy and sampled completions, then uses the reward +difference as the final training signal: + +```python +class ReMaxRLVRWorkflow(RolloutWorkflow): + async def arun_episode(self, engine: InferenceEngine, data): + # Prepare input tokens from chat messages + input_ids = self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + + n_samples = self.gconfig.n_samples + rid = uuid.uuid4().hex + + # Create requests for both sampled and greedy generation + sample_req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig, + ) + greedy_req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig.new(greedy=True), + ) + + # Generate both responses concurrently + resp, greedy_resp = await asyncio.gather( + engine.agenerate(sample_req), + engine.agenerate(greedy_req), + ) + + # Calculate rewards for both completions + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + + sample_reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + + greedy_completions = self.tokenizer.decode(greedy_resp.output_tokens) + greedy_reward = self.reward_fn( + prompt=prompt_str, + completions=greedy_completions, + prompt_ids=greedy_resp.input_tokens, + completion_ids=greedy_resp.output_tokens, + **data, + ) + + # Package results for training + res = dict( + # Add batch dimension + input_ids=torch.tensor(resp.input_tokens + resp.output_tokens).unsqueeze(0), + loss_mask=torch.tensor([0] * resp.input_len + [1] * resp.output_len).unsqueeze(0), + versions=torch.tensor([-1] * resp.input_len + resp.output_versions).unsqueeze(0), + attention_mask=torch.ones(resp.input_len + resp.output_len, dtype=torch.bool).unsqueeze(0), + # Use reward difference across all tokens + rewards=torch.tensor([float(sample_reward - greedy_reward)] * (resp.input_len + resp.output_len)), + ) + + return TensorDict(res, batch_size=[1]) +``` + +> **Note**: For detailed guidance on customizing rollout workflows, see the +> [agent customization guide](agent.md). + +### Step 2: Implementing the REINFORCE Training Algorithm + +Training algorithms are implemented by subclassing `TrainEngine` and using its atomic +operations like `forward`, `train_batch`, and `eval_batch`. + +First, let's define the REINFORCE loss function: ```python -# in realhf/api/core/model_api.py +def reinforce_loss_fn(logits, data): + input_ids = data["input_ids"] + loss_mask = data["loss_mask"].bool() + rewards = data["rewards"] + + logprobs = gather_logprobs( + logits, torch.roll(input_ids, shifts=-1, dims=-1) + ) + loss = -logprobs * rewards + loss = torch.where(loss_mask, loss, 0.0) + + return loss.sum() / loss_mask.count_nonzero() +``` + +Next, we implement the training engine. We use a two-class design to maintain backend +compatibility: + +```python +class ReinforceActor: + def __init__(self, engine: TrainEngine): + self.engine = engine + + def train_reinforce(self, data: TensorDict): + # Enable gradient checkpointing + self.engine.train() + return self.engine.train_batch( + data, + loss_fn=reinforce_loss_fn, + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), + ) + +class FSDPReinforceActor(FSDPEngine): + def __init__(self): + self.actor = ReinforceActor(self) + + def train_reinforce(self, *args, **kwargs): + return self.actor.train_reinforce(*args, **kwargs) +``` + +**Why two classes?** This design separates concerns: + +1. **Backend Agnostic Logic**: `ReinforceActor` contains the core REINFORCE algorithm + that works with any backend (FSDP, DeepSpeed, Megatron) since they share the same + `train_batch` API. + +1. **Backend-Specific Features**: `FSDPReinforceActor` inherits from `FSDPEngine` to + provide backend-specific utilities like `save`, `load`, and `upload_weights`. For + other backends, you'd create `MegatronReinforceActor`, etc. + +> **Note**: This pattern is similar to interfaces in Go or traits in Rust, adapted for +> Python's object model. + +### Step 3: Composing the Complete Training Loop + +The main training loop brings everything together: + +```python +def main(args): + # Initialize inference engine for rollouts + rollout = RemoteSGLangEngine(config.rollout) + rollout.initialize(None, ft_spec) + + # Initialize training engine + actor = FSDPReinforceActor(config=config.actor) + actor.initialize(None, ft_spec) + + # Create rollout workflow + workflow = ReMaxRLVRWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + ) + + # Main training loop + for global_step in range(max_steps): + # Generate training data + with stats_tracker.record_timing("rollout"): + try: + data = next(data_generator) + except StopIteration: + data_generator = iter(train_dataloader) + data = next(data_generator) + + batch = rollout.rollout_batch(data, workflow=workflow) + + batch = batch.to(actor.device) + + # Synchronize all processes + dist.barrier() + torch.cuda.synchronize() + + # Training step + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("actor"), + ): + stats = actor.train_reinforce(batch) + actor.step_lr_scheduler() + + # Update model weights + with stats_tracker.record_timing("update_weights"): + # Weight update logic here + ... +``` + +## Approach 2: Using Legacy ModelInterface (Not Recommended) + +The legacy approach encapsulates algorithms in a `ModelInterface` with three core +methods: + +```python +# From realhf/api/core/model_api.py class ModelInterface(abc.ABC): - """An interface for model training, inference, and generation. - - This interface is designed to follow the dependency injection pattern. - We pass the model to the interface and call its methods, ensuring that model APIs - and algorithms are fully decoupled. For example, REINFORCE and PPO can exhibit - different behaviors during training. Separate interfaces can be written for these - algorithms while using the same model that provides basic forward-backward-update - functionality (i.e., :class:`PipelinableEngine`). - - During runtime, the master worker requests model workers to execute a specific - interface type (e.g., generate) on a specific model. The model worker locates - the corresponding model, passes it into the requested interface, performs the - computation, and returns the result. + """Interface for model training, inference, and generation. + + This interface follows the dependency injection pattern, allowing + algorithms like REINFORCE and PPO to use the same underlying model + while exhibiting different training behaviors. """ def inference( @@ -45,144 +237,115 @@ class ModelInterface(abc.ABC): raise NotImplementedError() ``` -When the dataflow is fixed, it's usually sufficient to modify or add the file that defines the algorithm interface. +When the dataflow is fixed, you typically only need to modify the algorithm interface +file. -We provide two examples: (1) changing PPO's global advantage normalization to grouped normalization in GRPO, and (2) changing the original PPO loss to the decoupled PPO loss in AReaL's paper. +> **Note**: We recommend using asynchronous RL so you can customize generation behavior +> by [modifying your RL agent](agent.md) instead of the `generate` method. -```{note} -We recommend using asynchronous RL, so that you can customize the generation behavior by [modifying your RL agent](agent.md) and don't need to modify the `generate` method of model interfaces. -``` +### Example 1: Grouped Advantage Normalization + +Let's modify PPO's global advantage normalization to use grouped normalization (GRPO +approach). -## Grouped Advantage Normalization +#### Understanding Data Organization -The PPO algorithm is written in a single file `ppo_interface.py`. The method we are going to modify is the `train_step` method in `PPOActorInterface`. PPO's global advantage normalization looks like: +Each batch contains multiple prompts (batch size) and each prompt may have multiple +responses (group size). So total sequences = batch_size × group_size. + +Sequences have different lengths but are packed into a 1D tensor. We use `cu_seqlens` +(cumulative sequence lengths) to mark boundaries, similar to flash-attention. + +#### Implementation + +The standard PPO normalization looks like: ```python @dataclass class PPOActorInterface(ModelInterface): - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - ... + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + # ... if self.adv_norm: advantages = masked_normalization(advantages, loss_mask) - ... + # ... ``` -### An Additional Note on Data Management - -We need to explain how data in each batch is organized. - -Usually, each data batch (i.e., the `data` variable) includes multiple prompts. The number of prompts is called "batch size". Additionally, each prompt may have multiple corresponding answers. The number of answers is called "group_size". Therefore, there are batch_size × group_size sequences in each batch. - -These sequences have different lengths, but they are concatenated (or packed) together as a 1D tensor. The inner dimension is the "group" with the same prompt, and the outer dimension consists of answers from different prompts. Similar to flash-attention, we use `cu_seqlens` to mark the boundary of each sequence. `cu_seqlens` is the cumulative sum of sequence lengths across the batch. - -Each token in the sequence has a corresponding reward and advantage, so `advantages` is also a packed 1D tensor just like the tokens (i.e., `packed_input_ids`). However, the "sequences" of advantages are all one step shorter than tokens due to the auto-regressive nature of LLMs. We can only compute the loss on tokens except for the first one in each sequence. - -### Implementation - -For grouped advantage normalization, we need to partition the advantages into groups and run normalization within the tensor chunk of each group: +For grouped normalization, we partition advantages by group: -```diff +```python @dataclass class PPOActorInterface(ModelInterface): -+ group_adv_norm: bool = False + group_adv_norm: bool = False - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - ... + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + # ... if self.adv_norm: -- advantages = masked_normalization(advantages, loss_mask) -+ if not self.group_adv_norm: -+ advantages = masked_normalization(advantages, loss_mask) -+ else: -+ n_samples = data.bs -+ adv_list = [] -+ for i in range(0, n_samples, self.group_size): -+ # Start and end of the chunk -+ s = short1cu_seqlens[i] -+ e = short1cu_seqlens[i + self.group_size] -+ # Get advantages within each group of the same prompt -+ adv = advantages[s: e] -+ mask = loss_mask[s: e] -+ # Run normalization -+ advn = masked_normalization(adv, mask, all_reduce=False) -+ adv_list.append(advn) -+ advantages = torch.cat(adv_list, 0) - ... + if not self.group_adv_norm: + advantages = masked_normalization(advantages, loss_mask) + else: + n_samples = data.bs + adv_list = [] + for i in range(0, n_samples, self.group_size): + # Define chunk boundaries + s = short1cu_seqlens[i] + e = short1cu_seqlens[i + self.group_size] + + # Extract advantages for this group + adv = advantages[s:e] + mask = loss_mask[s:e] + + # Normalize within group + advn = masked_normalization(adv, mask, all_reduce=False) + adv_list.append(advn) + + advantages = torch.cat(adv_list, 0) + # ... ``` -### Modify Your Experiment Configuration +#### Configuration Changes -To make our new argument `group_adv_norm` effective in CLI args, we should make the following changes to the `PPOMathConfig` under `realhf/experiments/common/ppo_math_exp.py`: +Update the experiment configuration to expose the new parameter: -```diff +```python @dataclasses.dataclass class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): -+ group_adv_norm: bool = False + group_adv_norm: bool = False @property def rpcs(self): - ... - # interfaces + # ... actor_interface = ModelInterfaceAbstraction( "ppo_actor", args={ **copy.deepcopy(self.ppo_kwargs), -+ "group_adv_norm": self.group_adv_norm, - ... + "group_adv_norm": self.group_adv_norm, + # ... }, ) ``` -## The Decoupled PPO Loss +### Example 2: Decoupled PPO Loss -![decoupled loss](decoupled_loss.png) +The decoupled PPO loss (from AReaL's paper) recomputes probabilities before mini-batch +updates and uses this as π_prox: -As mentioned in AReaL's paper, we implement this loss by recomputing the probabilities before mini-batched updates, and use this value as π_prox to compute the above loss. +![decoupled loss](decoupled_loss.png) -### Probability Recomputation +#### Probability Recomputation -Recomputation involves a single forward pass, which has already been implemented by `PPOActorInterface.inference`. We need to call this method in the `train_step` method: +We recompute probabilities using the existing `inference` method: -```diff +```python @dataclass class PPOActorInterface(ModelInterface): -+ use_decoupled_loss: bool = False - - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: -+ if self.use_decoupled_loss: -+ s: SequenceSample = self.inference(model, data, mb_spec) -+ prox_logp = s.data["logprobs"] - ... -``` - -Next, we need to pass `prox_logp` to loss computation: + use_decoupled_loss: bool = False -```diff -@dataclass -class PPOActorInterface(ModelInterface): - ... + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + if self.use_decoupled_loss: + s: SequenceSample = self.inference(model, data, mb_spec) + prox_logp = s.data["logprobs"] - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - # Prepare data to be split into mini-batches. + # Prepare mini-batch data flat_data = dict( advantages=advantages, old_logp=old_logp, @@ -190,31 +353,32 @@ class PPOActorInterface(ModelInterface): packed_input_ids=input_.data["packed_input_ids"], kl_rewards=kl_rewards, ) -+ if self.use_decoupled_loss: -+ flat_data["prox_logp"] = prox_logp.float() + + if self.use_decoupled_loss: + flat_data["prox_logp"] = prox_logp.float() flat_input = SequenceSample.from_default( ids=list(range(input_.bs * self.group_size)), data=flat_data, seqlens=[int(x) for x in input_lens.cpu().numpy().tolist()], ) - ... + + # Split into mini-batches and train datas = flat_input.split_with_spec(spec) - ... for mb_i, data in enumerate(datas): train_stat = module.train_batch( input_=data, mb_spec=mb_spec, version_steps=model.version.global_step, loss_fn=_loss_fn, - loss_weight_fn=lambda x: x.data[ - "ppo_loss_mask" - ].count_nonzero(), + loss_weight_fn=lambda x: x.data["ppo_loss_mask"].count_nonzero(), token_normalize_scope=self.token_normalize_scope, ) ``` -The `flat_input` variable will be divided into mini-batches. Each mini-batch of data will be passed into the `train_batch` method to run distributed training. The data included in this `SequenceSample` object will all be passed into the `_loss_fn`. In this case, `_loss_fn` is a wrapper over `_ppo_actor_loss_from_model_outputs`: +#### Modifying the Loss Function + +Update the loss computation to use the recomputed probabilities: ```python def _ppo_actor_loss_from_model_outputs( @@ -222,19 +386,9 @@ def _ppo_actor_loss_from_model_outputs( input_: SequenceSample, ... ) -> torch.Tensor: - ... -``` + # ... + prox_logp = input_.data.get("prox_logp") -`logits` is the output of model forward, and `input_` is exactly the `input_` we passed into `train_batch`. So now we can retrieve the `prox_logp` via: - -```diff -def _ppo_actor_loss_from_model_outputs( - logits: torch.FloatTensor, # [tot_seqlen, vocab_size] - input_: SequenceSample, - ... -) -> torch.Tensor: - ... -+ prox_logp = input_.data["prox_logp"] loss, ppo_stat = ppo_functional.actor_loss_fn( logprobs=logprobs, old_logprobs=old_logp, @@ -242,16 +396,14 @@ def _ppo_actor_loss_from_model_outputs( eps_clip=eps_clip, loss_mask=ppo_loss_mask, c_clip=c_clip, -+ proximal_logprobs=prox_logp, + proximal_logprobs=prox_logp, behav_imp_weight_cap=behav_imp_weight_cap, ) ``` -We have successfully recomputed the probability and passed it into the loss function. Next we should revise the loss computation code. - -### Modifying the PPO Loss +And in the core loss function: -```diff +```python def actor_loss_fn( logprobs: torch.FloatTensor, old_logprobs: torch.FloatTensor, @@ -259,47 +411,52 @@ def actor_loss_fn( eps_clip: float, loss_mask: Optional[torch.BoolTensor] = None, c_clip: Optional[float] = None, -+ proximal_logprobs: Optional[torch.FloatTensor] = None, + proximal_logprobs: Optional[torch.FloatTensor] = None, behav_imp_weight_cap: Optional[torch.FloatTensor] = None, ) -> Tuple[torch.Tensor, Dict]: - ... -+ if proximal_logprobs is not None: -+ denorm_logprobs = proximal_logprobs -+ else: -+ denorm_logprobs = old_logprobs - ... + # Use proximal probabilities if available, otherwise use old probabilities + denorm_logprobs = proximal_logprobs if proximal_logprobs is not None else old_logprobs + loss_mask_count = loss_mask.count_nonzero() or 1 - # For numerical stability. -- ratio = torch.where(loss_mask, torch.exp(logprobs - old_logprobs), 0) -+ ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) - ... -+ if proximal_logprobs is not None: -+ behav_kl = proximal_logprobs - old_logprobs -+ behav_imp_weight = behav_kl.exp() -+ behav_kl = torch.where(loss_mask, behav_kl, 0.0) -+ behav_imp_weight = torch.where(loss_mask, behav_imp_weight, 0.0) -+ pg_loss = pg_loss * behav_imp_weight - ... + + # Compute importance weights + ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) + + # Apply behavioral importance weighting for decoupled loss + if proximal_logprobs is not None: + behav_kl = proximal_logprobs - old_logprobs + behav_imp_weight = behav_kl.exp() + behav_kl = torch.where(loss_mask, behav_kl, 0.0) + behav_imp_weight = torch.where(loss_mask, behav_imp_weight, 0.0) + pg_loss = pg_loss * behav_imp_weight + + # ... return pg_loss, stat ``` -### Modify the Experiment Configuration +#### Configuration Update -```diff +```python @dataclasses.dataclass class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): -+ use_decoupled_loss: bool = False + use_decoupled_loss: bool = False @property def rpcs(self): - ... - # interfaces + # ... actor_interface = ModelInterfaceAbstraction( "ppo_actor", args={ **copy.deepcopy(self.ppo_kwargs), -+ "use_decoupled_loss": self.use_decoupled_loss, - ... + "use_decoupled_loss": self.use_decoupled_loss, + # ... }, ) -``` \ No newline at end of file +``` + +______________________________________________________________________ + +This guide should help you implement custom RL training algorithms using either the +modern AReaLite approach or the legacy ModelInterface system. The AReaLite approach is +recommended for new implementations due to its cleaner separation of concerns and better +maintainability. From 9db11d1c233d970ae747f27125db6d48dac948fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 21 Jul 2025 17:18:48 +0800 Subject: [PATCH 12/16] . --- docs/customization/agent.md | 4 +- docs/customization/algorithm.md | 8 +- docs/customization/dataset.md | 220 ++++++++++++++++++++++++++++---- 3 files changed, 205 insertions(+), 27 deletions(-) diff --git a/docs/customization/agent.md b/docs/customization/agent.md index 5815e4c925..26e037b6ea 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -5,7 +5,7 @@ a multi-turn math agent using end-to-end reinforcement learning. We'll walk thro creating an agent that keeps trying to solve math problems until it finds the correct answer. -## Recommended Approach: AReaLite +## Option 1: Using AReaLite (Recommended) You can find the complete implementation in `arealite/workflow/multi_turn.py`. @@ -269,7 +269,7 @@ That's it! Your custom multi-turn math agent is now ready for reinforcement lear training. The workflow will automatically handle the multi-turn conversations, reward computation, and data collection needed for effective RL training. -## Legacy Approach: Agent-Based System +### Option 2: Using the Legacy Agent & Environment (NOT Recommended) While we strongly recommend AReaLite for new projects, you might need to work with legacy code that uses the older Agent-based approach. Here's how it works for reference, diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index 8318e48d25..5af4d0c9be 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -1,6 +1,6 @@ # Training Algorithm Implementation -## Approach 1: Using AReaLite (Recommended) +## Option 1: Using AReaLite (Recommended) AReaLite structures RL algorithms around two core components: @@ -106,6 +106,10 @@ def reinforce_loss_fn(logits, data): return loss.sum() / loss_mask.count_nonzero() ``` +```{note} +To decrease memory usage, AReaLite automatically packs multiple sequences in an 1D tensor before forward passes. Hence, the loss function should assume handling 1D *packed* tensors instead of *padded* tensors. +``` + Next, we implement the training engine. We use a two-class design to maintain backend compatibility: @@ -197,7 +201,7 @@ def main(args): ... ``` -## Approach 2: Using Legacy ModelInterface (Not Recommended) +## Option 2: Using the Legacy ModelInterface (NOT Recommended) The legacy approach encapsulates algorithms in a `ModelInterface` with three core methods: diff --git a/docs/customization/dataset.md b/docs/customization/dataset.md index 9d4b4e0221..6ea8ba164b 100644 --- a/docs/customization/dataset.md +++ b/docs/customization/dataset.md @@ -1,14 +1,180 @@ # Dataset -This guide provides detailed examples of how to create custom datasets in AReaL for model training. +## Option 1: Using AReaLite (Recommended) -## Define Your Dataset +AReaLite directly integrates with the `Dataset` class from the HuggingFace `datasets` +package. This gives you full flexibility to load, process, and filter your data before +training. -Create a new file under `realhf/impl/dataset/`, for example, `my_custom_dataset.py`. Your `Dataset` must implement the `torch.utils.data.Dataset` interface and follow the framework's conventions. +The required columns in your dataset depend on the specific implementation of the +`RolloutWorkflow` (for online reinforcement learning) or the training engines (for +offline training, such as `LMEngine` for Supervised Fine-Tuning (SFT)). + +Here are two concrete examples from the existing implementation: + +### SFT (Offline Training) + +In the SFT example, we see that the loaded data is directly passed to the `train_lm` +method: + +```python +# examples/arealite/gsm8k_sft.py +def main(args): + ... + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + get_gsm8k_dataset("train", tokenizer, rank, world_size), + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=pad_sequences_to_tensors, + drop_last=config.train_dataset.drop_last, + ) + ... + # Run training loop + for epoch in range(total_epochs): + for step, data in enumerate(train_dataloader): + stats = engine.train_lm(data) + engine.step_lr_scheduler() + stats_tracker.scalar(**stats) +``` + +In this case, the `train_lm` method requires the keys "input_ids", "attention_mask", and +"loss_mask" to function. We first tokenize the dataset to extract the "input_ids" and +"loss_mask". Then, the `pad_sequences_to_tensors` method is used to batch multiple +sequences and append the "attention_mask": + +```python +def process_gsm8k_sft_dataset(dataset: Dataset, tokenizer): + def process(sample): + seq_token = tokenizer.encode( + sample["question"] + sample["answer"] + tokenizer.eos_token + ) + prompt_token = tokenizer.encode(sample["question"]) + loss_mask = [0] * len(prompt_token) + [1] * (len(seq_token) - len(prompt_token)) + return {"input_ids": seq_token, "loss_mask": loss_mask} + + # Remove unnecessary columns to avoid errors during collation + dataset = dataset.map(process).remove_columns(["question", "answer"]) + return dataset + +def get_gsm8k_dataset(split, tokenizer, rank, world_size): + dataset = load_dataset(path="openai/gsm8k", name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + return process_gsm8k_sft_dataset(dataset, tokenizer) +``` + +### GRPO (Online Training) + +In the GRPO example, the loaded data is passed to the `InferenceEngine`, rather than the +`TrainEngine`: + +```python +# examples/arealite/gsm8k_ppo.py +def main(args): + ... + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + get_gsm8k_dataset("train", rank, world_size), + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.train_dataset.drop_last, + ) + # Initialize inference engine + rollout = RemoteSGLangEngine(config.rollout) + rollout.initialize(None, ft_spec) + workflow = RLVRWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + enable_thinking=False, + dump_dir=os.path.join( + StatsLogger.get_log_path(config.stats_logger), "generated" + ), + ) + # Run training loop + ... + for global_step in range(max_steps): + batch = rollout.rollout_batch(data, workflow=workflow) + ... +``` + +Note that the `collate_fn` here is an identity function, meaning it simply returns the +list of individual data items as a batch. In the `InferenceEngine`, the data is then +dispatched to multiple concurrent executions of `workflow.arun_episode`, where each +dispatched data corresponds to a single episode. + +The `RLVRWorkflow` implementation extracts the "messages" field from the data dictionary +as the prompt for generating a response. Additionally, this data is passed to the +`reward_fn` as keyword arguments, which allows the reward function to make use of other +dataset fields, like "answers". Here’s an example: + +```python +class RLVRWorkflow(RolloutWorkflow): + + async def arun_episode(self, engine: InferenceEngine, data): + input_ids = self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + n_samples = self.gconfig.n_samples + req = LLMRequest( + rid=uuid.uuid4().hex, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) + ... + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) +``` + +Thus, the "messages" field must be constructed when loading the dataset, and the reward +function should be defined to handle the dataset's specific fields. Here’s how you can +process the dataset for this example: + +```python +def process_gsm8k_rl_dataset(dataset: Dataset): + def process(sample): + messages = [{"role": "user", "content": sample["question"]}] + return {"messages": messages} + + # The dataset has two fields "messages" and "answer" + dataset = dataset.map(process).remove_columns(["question"]) + return dataset + +def get_gsm8k_dataset(split, rank, world_size): + dataset = load_dataset(path="openai/gsm8k", name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + return process_gsm8k_rl_dataset(dataset) + +def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs): + # "answer" is passed in through "**data" + from realhf.impl.dataset.math_parser import process_results + + return int(process_results(completions, answer)[0]) +``` + +## Option 2: Using the Legacy Dataset (NOT Recommended) + +### Define Your Dataset + +Create a new file under `realhf/impl/dataset/`, for example, `my_custom_dataset.py`. +Your `Dataset` must implement the `torch.utils.data.Dataset` interface and follow the +framework's conventions. ```python class MyCustomDataset(torch.utils.data.Dataset): - + def __init__( self, util: data_api.DatasetUtility, @@ -19,7 +185,7 @@ class MyCustomDataset(torch.utils.data.Dataset): custom_param: float = 1.0, ): """Custom dataset initialization - + Args: util: Dataset utility class containing tokenizer, seed, distributed info, etc. max_length: Maximum sequence length @@ -29,19 +195,19 @@ class MyCustomDataset(torch.utils.data.Dataset): """ self._util = util self.max_length = max_length - + # Load and split dataset data = data_api.load_shuffle_split_dataset(util, dataset_path, dataset_builder) - + # Your custom data processing logic ... ``` -## Implement Core Methods +### Implement Core Methods Every dataset class must implement the following two core methods: -### 1. `__len__` Method +#### 1. `__len__` Method Returns the size of the dataset: @@ -50,7 +216,7 @@ def __len__(self): return len(self.data_samples) ``` -### 2. `__getitem__` Method +#### 2. `__getitem__` Method Returns the sample at the specified index, must return a `SequenceSample` object: @@ -58,10 +224,10 @@ Returns the sample at the specified index, must return a `SequenceSample` object def __getitem__(self, idx): # Get raw data sample = self.data_samples[idx] - + # Process data ... - + # Return SequenceSample object return data_api.SequenceSample.from_default( ids=[sample["id"]], @@ -73,37 +239,44 @@ def __getitem__(self, idx): ) ``` -### Dataset Examples +#### Dataset Examples We provide some examples of dataset under `realhf/impl/dataset/`: + - For SFT, please refer `prompt_answer_dataset.py`. - For Reward model training, please refer `rw_paired_dataset.py` - For RL training, please refer `math_code_dataset.py` -## Data Format Requirements +### Data Format Requirements + +#### JSONL File Format -### JSONL File Format +Your data file should be in JSONL format, with one JSON object per line. If you are +using our PromptDataset implementation, your data should be like: -Your data file should be in JSONL format, with one JSON object per line. -If you are using our PromptDataset implementation, your data should be like: - Math Data + ```json {"qid": "sample_1", "prompt": "Solve this math problem: 2+2=", "solutions": ["\\boxed{4}"]} ``` + - Code Data + ```json {"qid": "sample_2", "prompt": "Code problem", "input_output": "{\"inputs\": [\"5\\n2 3 5 10 12\\n\"], \"outputs\": [\"17\\n\"]}"} ``` - `qid`: Unique identifier for the sample - `prompt`: Input prompt text -- `task`: Task type, used to distinguish how to calculate the reward. ("math" and "code" are supported now.) +- `task`: Task type, used to distinguish how to calculate the reward. ("math" and "code" + are supported now.) -Note: There is no format restriction for a customized dataset as long as it can be loaded by your custom code. +Note: There is no format restriction for a customized dataset as long as it can be +loaded by your custom code. -## Registration and Configuration +### Registration and Configuration -### Register Dataset +#### Register Dataset Register your dataset at the end of your dataset file: @@ -112,9 +285,10 @@ Register your dataset at the end of your dataset file: data_api.register_dataset("my-custom", MyCustomDataset) ``` -### Modify Experiment Configuration +#### Modify Experiment Configuration -Use your new dataset in the experiment configuration (refer to `realhf/experiments/common/*_exp.py`): +Use your new dataset in the experiment configuration (refer to +`realhf/experiments/common/*_exp.py`): ```python # in your experiment config file From 8a610f37f0347eb804f3caa5bdbac0004539c9d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 21 Jul 2025 17:30:03 +0800 Subject: [PATCH 13/16] . --- docs/_toc.yml | 2 +- docs/customization/agent.md | 7 +++---- docs/customization/algorithm.md | 5 ++++- docs/customization/dataset.md | 25 +++---------------------- 4 files changed, 11 insertions(+), 28 deletions(-) diff --git a/docs/_toc.yml b/docs/_toc.yml index fecabda369..6a197dfee8 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -16,9 +16,9 @@ parts: - file: references/reproduce - caption: Customization chapters: - - file: customization/dataset - file: customization/agent - file: customization/algorithm + - file: customization/dataset - caption: Code Walkthrough chapters: - file: developer/overview diff --git a/docs/customization/agent.md b/docs/customization/agent.md index 26e037b6ea..631af36d6d 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -1,9 +1,8 @@ # Rollout and Agentic RL -This guide shows you how to create custom rollout behaviors for PPO training by building -a multi-turn math agent using end-to-end reinforcement learning. We'll walk through -creating an agent that keeps trying to solve math problems until it finds the correct -answer. +This guide shows you how to create custom rollout behaviors for RL training by building +a multi-turn math agent, which keeps trying to solve math problems until it finds the +correct answer. ## Option 1: Using AReaLite (Recommended) diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index 5af4d0c9be..6b285f39e7 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -1,4 +1,7 @@ -# Training Algorithm Implementation +# Training Algorithm + +> **Note**: We recommend the user to first read the +> [agent customization guide](agent.md). ## Option 1: Using AReaLite (Recommended) diff --git a/docs/customization/dataset.md b/docs/customization/dataset.md index 6ea8ba164b..20b1a46d0b 100644 --- a/docs/customization/dataset.md +++ b/docs/customization/dataset.md @@ -24,19 +24,13 @@ def main(args): # Create dataset and dataloaders train_dataloader = StatefulDataLoader( get_gsm8k_dataset("train", tokenizer, rank, world_size), - batch_size=config.train_dataset.batch_size // world_size, - shuffle=config.train_dataset.shuffle, - num_workers=config.train_dataset.num_workers, collate_fn=pad_sequences_to_tensors, - drop_last=config.train_dataset.drop_last, ) ... # Run training loop for epoch in range(total_epochs): for step, data in enumerate(train_dataloader): stats = engine.train_lm(data) - engine.step_lr_scheduler() - stats_tracker.scalar(**stats) ``` In this case, the `train_lm` method requires the keys "input_ids", "attention_mask", and @@ -76,23 +70,13 @@ def main(args): # Create dataset and dataloaders train_dataloader = StatefulDataLoader( get_gsm8k_dataset("train", rank, world_size), - batch_size=config.train_dataset.batch_size // world_size, - shuffle=config.train_dataset.shuffle, - num_workers=config.train_dataset.num_workers, collate_fn=lambda x: x, - drop_last=config.train_dataset.drop_last, ) # Initialize inference engine rollout = RemoteSGLangEngine(config.rollout) - rollout.initialize(None, ft_spec) workflow = RLVRWorkflow( reward_fn=gsm8k_reward_fn, - gconfig=config.gconfig, - tokenizer=tokenizer, - enable_thinking=False, - dump_dir=os.path.join( - StatsLogger.get_log_path(config.stats_logger), "generated" - ), + ... ) # Run training loop ... @@ -102,7 +86,7 @@ def main(args): ``` Note that the `collate_fn` here is an identity function, meaning it simply returns the -list of individual data items as a batch. In the `InferenceEngine`, the data is then +list of individual data items as a batch. In `rollout_batch`, the data is then dispatched to multiple concurrent executions of `workflow.arun_episode`, where each dispatched data corresponds to a single episode. @@ -121,13 +105,10 @@ class RLVRWorkflow(RolloutWorkflow): add_generation_prompt=True, enable_thinking=self.enable_thinking, ) - n_samples = self.gconfig.n_samples req = LLMRequest( - rid=uuid.uuid4().hex, input_ids=input_ids, - gconfig=self.gconfig.new(n_samples=1), + ... ) - resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) ... reward = self.reward_fn( prompt=prompt_str, From 2d3c6ac3ffc87ca6593a11ab20e11dc76d0c5acf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Tue, 22 Jul 2025 15:15:36 +0800 Subject: [PATCH 14/16] . --- docs/_toc.yml | 5 + docs/customization/agent.md | 176 +----------- docs/customization/algorithm.md | 264 ------------------ docs/customization/dataset.md | 145 ---------- docs/legacy/customization/agent.md | 162 +++++++++++ docs/legacy/customization/algorithm.md | 259 +++++++++++++++++ docs/legacy/customization/dataset.md | 146 ++++++++++ .../customization/decoupled_loss.png | Bin .../customization/multiturn_reward.png | Bin 9 files changed, 574 insertions(+), 583 deletions(-) create mode 100644 docs/legacy/customization/agent.md create mode 100644 docs/legacy/customization/algorithm.md create mode 100644 docs/legacy/customization/dataset.md rename docs/{ => legacy}/customization/decoupled_loss.png (100%) rename docs/{ => legacy}/customization/multiturn_reward.png (100%) diff --git a/docs/_toc.yml b/docs/_toc.yml index 6a197dfee8..3062303e41 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -35,3 +35,8 @@ parts: - caption: Contributing chapters: - file: contrib + - caption: Customization (Legacy) + chapters: + - file: legacy/customization/agent + - file: legacy/customization/algorithm + - file: legacy/customization/dataset diff --git a/docs/customization/agent.md b/docs/customization/agent.md index 631af36d6d..c668f87f2d 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -1,10 +1,8 @@ # Rollout and Agentic RL This guide shows you how to create custom rollout behaviors for RL training by building -a multi-turn math agent, which keeps trying to solve math problems until it finds the -correct answer. - -## Option 1: Using AReaLite (Recommended) +a multi-turn math agent with **AReaLite**. This agent keeps trying to solve math +problems until it finds the correct answer. You can find the complete implementation in `arealite/workflow/multi_turn.py`. @@ -267,173 +265,3 @@ def main(args): That's it! Your custom multi-turn math agent is now ready for reinforcement learning training. The workflow will automatically handle the multi-turn conversations, reward computation, and data collection needed for effective RL training. - -### Option 2: Using the Legacy Agent & Environment (NOT Recommended) - -While we strongly recommend AReaLite for new projects, you might need to work with -legacy code that uses the older Agent-based approach. Here's how it works for reference, -though we suggest migrating to the workflow-based system when possible. - -### Step 1: Define Your Agent Class - -Create a new file under `realhf/impl/agent/`, such as `math_multi_turn_agent.py`. Your -`Agent` must implement the interface defined in `realhf/api/core/agent.py`, which -requires a single method: `collect_trajectory`. - -```python -class MathMultiTurnAgent(Agent): - async def collect_trajectory( - self, - prompt: SequenceSample, - env: EnvironmentService, - obs_queue: asyncio.Queue, - act_queue: asyncio.Queue, - ): - # Implementation goes here - ... -``` - -### Step 2: Implement the Trajectory Collection Logic - -The `collect_trajectory` method takes a task prompt, an environment, and two -communication queues. Within this method, you control the data flow between your agent -and the inference engine using these queues: - -- **obs_queue**: Send observations (token IDs and generation config) to the inference - engine -- **act_queue**: Receive actions (generated responses) from the inference engine - -Here's how the multi-turn conversation works: - -```python -for turn in range(self.num_turns): - # Send the current state to the inference engine - await obs_queue.put((qid, token_ids, self.gconfig)) - - # Get the generated response - act: BundledGenerationOutputs = await act_queue.get() - - # Evaluate the response through the environment - success, rewards = await env.step((qid, answers)) - # ... process results ... -``` - -#### Environment Integration - -The environment follows a -[Gym-like interface](https://github.com/Farama-Foundation/Gymnasium) with `reset` and -`step` methods, but uses asynchronous implementations to prevent blocking across -different environment instances. - -For math problems, the environment is typically stateless and acts as a wrapper around -your reward function: - -```python -class MathCodeSingleStepEnv(EnvironmentService): - async def step(self, action: Tuple[str, List[str]]): - qid, answers = action - # ... setup code ... - - # Run reward computation asynchronously - format_rewards = await asyncio.to_thread( - math_verify_call, - answers, - # ... other parameters ... - ) - return None, format_rewards, True, False, {} -``` - -#### Handling Multi-Turn Feedback - -After receiving the reward from `env.step`, check if the answer is correct. If not, -provide feedback and continue to the next turn: - -```python -for turn in range(self.num_turns): - # ... generation and evaluation code ... - - # Provide feedback based on the result - if success[0]: - feedback = "Congratulations! You are correct!" - else: - feedback = "Unfortunately your answer is wrong. Let's try again." - - # Format feedback as a user message - feedback = "\n" + self.tokenizer.apply_chat_template( - [{"content": feedback, "role": "user"}], - add_generation_prompt=True, - tokenize=False, - ) - - # Add feedback tokens to the conversation - feedback_tokens = self.tokenizer(feedback)["input_ids"] - token_ids.extend(feedback_tokens) -``` - -### Step 3: Register and Configure Your Agent - -First, register your agent implementation: - -```python -# in realhf/impl/agent/math_multi_turn_agent.py -register_agent("math-multi-turn", MathMultiTurnAgent) -``` - -```python -# in realhf/impl/agent/__init__.py -import realhf.impl.agent.math_multi_turn_agent -``` - -Then update your experiment configuration in -`realhf/experiments/async_exp/async_math_ppo.py`: - -```python -@dataclasses.dataclass -class AsyncPPOMATHConfig(AsyncRLExperimentConfig, PPOMATHConfig): - # Add any new CLI arguments your agent needs - my_param: float = 1.0 - - @property - def agent(self) -> AgentAbstraction: - return AgentAbstraction( - "math-multi-turn", # Your registered agent name - args=dict( - # Pass any arguments needed for your __init__ method - my_param=self.my_param, - # ... other configuration ... - ), - ) - - @property - def env(self) -> EnvServiceAbstraction: - # Update to use your custom environment if needed - return EnvServiceAbstraction( - "math-code-single-step", - args=dict(dataset_path=self.dataset.path) - ) -``` - -### Step 4: Run Training - -Follow the standard training procedure outlined in the -[quickstart guide](../tutorial/quickstart.md). Launch your experiment with: - -```bash -python3 training/main_async_ppo.py my_param=5.0 # plus any additional CLI arguments -``` - -### Training Results - -Here's an example of the training reward curve from our multi-turn math agent: - -![Multi-turn Training Rewards](multiturn_reward.png) - -The agent successfully learns to solve math problems with improved accuracy over time, -demonstrating the effectiveness of the multi-turn approach. - -______________________________________________________________________ - -**Note**: While this legacy approach works, we strongly recommend using the AReaLite -workflow system for new projects. It provides better flexibility, cleaner abstractions, -and easier maintenance. Consider migrating existing legacy agents to the workflow-based -approach when possible. diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index 6b285f39e7..a10ddfd681 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -203,267 +203,3 @@ def main(args): # Weight update logic here ... ``` - -## Option 2: Using the Legacy ModelInterface (NOT Recommended) - -The legacy approach encapsulates algorithms in a `ModelInterface` with three core -methods: - -```python -# From realhf/api/core/model_api.py -class ModelInterface(abc.ABC): - """Interface for model training, inference, and generation. - - This interface follows the dependency injection pattern, allowing - algorithms like REINFORCE and PPO to use the same underlying model - while exhibiting different training behaviors. - """ - - def inference( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> SequenceSample | None: - raise NotImplementedError() - - def generate( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> SequenceSample | None: - raise NotImplementedError() - - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - raise NotImplementedError() -``` - -When the dataflow is fixed, you typically only need to modify the algorithm interface -file. - -> **Note**: We recommend using asynchronous RL so you can customize generation behavior -> by [modifying your RL agent](agent.md) instead of the `generate` method. - -### Example 1: Grouped Advantage Normalization - -Let's modify PPO's global advantage normalization to use grouped normalization (GRPO -approach). - -#### Understanding Data Organization - -Each batch contains multiple prompts (batch size) and each prompt may have multiple -responses (group size). So total sequences = batch_size × group_size. - -Sequences have different lengths but are packed into a 1D tensor. We use `cu_seqlens` -(cumulative sequence lengths) to mark boundaries, similar to flash-attention. - -#### Implementation - -The standard PPO normalization looks like: - -```python -@dataclass -class PPOActorInterface(ModelInterface): - def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: - # ... - if self.adv_norm: - advantages = masked_normalization(advantages, loss_mask) - # ... -``` - -For grouped normalization, we partition advantages by group: - -```python -@dataclass -class PPOActorInterface(ModelInterface): - group_adv_norm: bool = False - - def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: - # ... - if self.adv_norm: - if not self.group_adv_norm: - advantages = masked_normalization(advantages, loss_mask) - else: - n_samples = data.bs - adv_list = [] - for i in range(0, n_samples, self.group_size): - # Define chunk boundaries - s = short1cu_seqlens[i] - e = short1cu_seqlens[i + self.group_size] - - # Extract advantages for this group - adv = advantages[s:e] - mask = loss_mask[s:e] - - # Normalize within group - advn = masked_normalization(adv, mask, all_reduce=False) - adv_list.append(advn) - - advantages = torch.cat(adv_list, 0) - # ... -``` - -#### Configuration Changes - -Update the experiment configuration to expose the new parameter: - -```python -@dataclasses.dataclass -class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): - group_adv_norm: bool = False - - @property - def rpcs(self): - # ... - actor_interface = ModelInterfaceAbstraction( - "ppo_actor", - args={ - **copy.deepcopy(self.ppo_kwargs), - "group_adv_norm": self.group_adv_norm, - # ... - }, - ) -``` - -### Example 2: Decoupled PPO Loss - -The decoupled PPO loss (from AReaL's paper) recomputes probabilities before mini-batch -updates and uses this as π_prox: - -![decoupled loss](decoupled_loss.png) - -#### Probability Recomputation - -We recompute probabilities using the existing `inference` method: - -```python -@dataclass -class PPOActorInterface(ModelInterface): - use_decoupled_loss: bool = False - - def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: - if self.use_decoupled_loss: - s: SequenceSample = self.inference(model, data, mb_spec) - prox_logp = s.data["logprobs"] - - # Prepare mini-batch data - flat_data = dict( - advantages=advantages, - old_logp=old_logp, - ppo_loss_mask=loss_mask, - packed_input_ids=input_.data["packed_input_ids"], - kl_rewards=kl_rewards, - ) - - if self.use_decoupled_loss: - flat_data["prox_logp"] = prox_logp.float() - - flat_input = SequenceSample.from_default( - ids=list(range(input_.bs * self.group_size)), - data=flat_data, - seqlens=[int(x) for x in input_lens.cpu().numpy().tolist()], - ) - - # Split into mini-batches and train - datas = flat_input.split_with_spec(spec) - for mb_i, data in enumerate(datas): - train_stat = module.train_batch( - input_=data, - mb_spec=mb_spec, - version_steps=model.version.global_step, - loss_fn=_loss_fn, - loss_weight_fn=lambda x: x.data["ppo_loss_mask"].count_nonzero(), - token_normalize_scope=self.token_normalize_scope, - ) -``` - -#### Modifying the Loss Function - -Update the loss computation to use the recomputed probabilities: - -```python -def _ppo_actor_loss_from_model_outputs( - logits: torch.FloatTensor, # [tot_seqlen, vocab_size] - input_: SequenceSample, - ... -) -> torch.Tensor: - # ... - prox_logp = input_.data.get("prox_logp") - - loss, ppo_stat = ppo_functional.actor_loss_fn( - logprobs=logprobs, - old_logprobs=old_logp, - advantages=advantages, - eps_clip=eps_clip, - loss_mask=ppo_loss_mask, - c_clip=c_clip, - proximal_logprobs=prox_logp, - behav_imp_weight_cap=behav_imp_weight_cap, - ) -``` - -And in the core loss function: - -```python -def actor_loss_fn( - logprobs: torch.FloatTensor, - old_logprobs: torch.FloatTensor, - advantages: torch.FloatTensor, - eps_clip: float, - loss_mask: Optional[torch.BoolTensor] = None, - c_clip: Optional[float] = None, - proximal_logprobs: Optional[torch.FloatTensor] = None, - behav_imp_weight_cap: Optional[torch.FloatTensor] = None, -) -> Tuple[torch.Tensor, Dict]: - # Use proximal probabilities if available, otherwise use old probabilities - denorm_logprobs = proximal_logprobs if proximal_logprobs is not None else old_logprobs - - loss_mask_count = loss_mask.count_nonzero() or 1 - - # Compute importance weights - ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) - - # Apply behavioral importance weighting for decoupled loss - if proximal_logprobs is not None: - behav_kl = proximal_logprobs - old_logprobs - behav_imp_weight = behav_kl.exp() - behav_kl = torch.where(loss_mask, behav_kl, 0.0) - behav_imp_weight = torch.where(loss_mask, behav_imp_weight, 0.0) - pg_loss = pg_loss * behav_imp_weight - - # ... - return pg_loss, stat -``` - -#### Configuration Update - -```python -@dataclasses.dataclass -class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): - use_decoupled_loss: bool = False - - @property - def rpcs(self): - # ... - actor_interface = ModelInterfaceAbstraction( - "ppo_actor", - args={ - **copy.deepcopy(self.ppo_kwargs), - "use_decoupled_loss": self.use_decoupled_loss, - # ... - }, - ) -``` - -______________________________________________________________________ - -This guide should help you implement custom RL training algorithms using either the -modern AReaLite approach or the legacy ModelInterface system. The AReaLite approach is -recommended for new implementations due to its cleaner separation of concerns and better -maintainability. diff --git a/docs/customization/dataset.md b/docs/customization/dataset.md index 20b1a46d0b..1382726315 100644 --- a/docs/customization/dataset.md +++ b/docs/customization/dataset.md @@ -1,7 +1,5 @@ # Dataset -## Option 1: Using AReaLite (Recommended) - AReaLite directly integrates with the `Dataset` class from the HuggingFace `datasets` package. This gives you full flexibility to load, process, and filter your data before training. @@ -144,146 +142,3 @@ def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **k return int(process_results(completions, answer)[0]) ``` - -## Option 2: Using the Legacy Dataset (NOT Recommended) - -### Define Your Dataset - -Create a new file under `realhf/impl/dataset/`, for example, `my_custom_dataset.py`. -Your `Dataset` must implement the `torch.utils.data.Dataset` interface and follow the -framework's conventions. - -```python -class MyCustomDataset(torch.utils.data.Dataset): - - def __init__( - self, - util: data_api.DatasetUtility, - max_length: Optional[int] = None, - dataset_path: Optional[str] = None, - dataset_builder: Optional[Callable[[], List[Dict]]] = None, - # Your custom parameters - custom_param: float = 1.0, - ): - """Custom dataset initialization - - Args: - util: Dataset utility class containing tokenizer, seed, distributed info, etc. - max_length: Maximum sequence length - dataset_path: Path to dataset file (optional) - dataset_builder: Data construction function (optional, alternative to dataset_path) - custom_param: Your custom parameter - """ - self._util = util - self.max_length = max_length - - # Load and split dataset - data = data_api.load_shuffle_split_dataset(util, dataset_path, dataset_builder) - - # Your custom data processing logic - ... -``` - -### Implement Core Methods - -Every dataset class must implement the following two core methods: - -#### 1. `__len__` Method - -Returns the size of the dataset: - -```python -def __len__(self): - return len(self.data_samples) -``` - -#### 2. `__getitem__` Method - -Returns the sample at the specified index, must return a `SequenceSample` object: - -```python -def __getitem__(self, idx): - # Get raw data - sample = self.data_samples[idx] - - # Process data - ... - - # Return SequenceSample object - return data_api.SequenceSample.from_default( - ids=[sample["id"]], - seqlens=[len(processed_data["input_ids"])], - data=dict( - packed_prompts=torch.tensor(processed_data["input_ids"], dtype=torch.long), - # Other necessary data fields - ), - ) -``` - -#### Dataset Examples - -We provide some examples of dataset under `realhf/impl/dataset/`: - -- For SFT, please refer `prompt_answer_dataset.py`. -- For Reward model training, please refer `rw_paired_dataset.py` -- For RL training, please refer `math_code_dataset.py` - -### Data Format Requirements - -#### JSONL File Format - -Your data file should be in JSONL format, with one JSON object per line. If you are -using our PromptDataset implementation, your data should be like: - -- Math Data - -```json -{"qid": "sample_1", "prompt": "Solve this math problem: 2+2=", "solutions": ["\\boxed{4}"]} -``` - -- Code Data - -```json -{"qid": "sample_2", "prompt": "Code problem", "input_output": "{\"inputs\": [\"5\\n2 3 5 10 12\\n\"], \"outputs\": [\"17\\n\"]}"} -``` - -- `qid`: Unique identifier for the sample -- `prompt`: Input prompt text -- `task`: Task type, used to distinguish how to calculate the reward. ("math" and "code" - are supported now.) - -Note: There is no format restriction for a customized dataset as long as it can be -loaded by your custom code. - -### Registration and Configuration - -#### Register Dataset - -Register your dataset at the end of your dataset file: - -```python -# in realhf/impl/dataset/my_custom_dataset.py -data_api.register_dataset("my-custom", MyCustomDataset) -``` - -#### Modify Experiment Configuration - -Use your new dataset in the experiment configuration (refer to -`realhf/experiments/common/*_exp.py`): - -```python -# in your experiment config file -@property -def datasets(self) -> List[DatasetAbstraction]: - return [ - DatasetAbstraction( - "my-custom", # Your registered name - args=dict( - dataset_path=self.dataset_path, - max_length=self.max_length, - custom_param=self.custom_param, - # Other initialization parameters - ), - ) - ] -``` diff --git a/docs/legacy/customization/agent.md b/docs/legacy/customization/agent.md new file mode 100644 index 0000000000..bbfe91d362 --- /dev/null +++ b/docs/legacy/customization/agent.md @@ -0,0 +1,162 @@ +# Rollout and Agentic RL (Legacy) + +> **Note**: While this legacy approach works, we strongly recommend using the AReaLite +> for new projects. It provides better flexibility, cleaner abstractions, and easier +> maintenance. + +### Step 1: Define Your Agent Class + +Create a new file under `realhf/impl/agent/`, such as `math_multi_turn_agent.py`. Your +`Agent` must implement the interface defined in `realhf/api/core/agent.py`, which +requires a single method: `collect_trajectory`. + +```python +class MathMultiTurnAgent(Agent): + async def collect_trajectory( + self, + prompt: SequenceSample, + env: EnvironmentService, + obs_queue: asyncio.Queue, + act_queue: asyncio.Queue, + ): + # Implementation goes here + ... +``` + +### Step 2: Implement the Trajectory Collection Logic + +The `collect_trajectory` method takes a task prompt, an environment, and two +communication queues. Within this method, you control the data flow between your agent +and the inference engine using these queues: + +- **obs_queue**: Send observations (token IDs and generation config) to the inference + engine +- **act_queue**: Receive actions (generated responses) from the inference engine + +Here's how the multi-turn conversation works: + +```python +for turn in range(self.num_turns): + # Send the current state to the inference engine + await obs_queue.put((qid, token_ids, self.gconfig)) + + # Get the generated response + act: BundledGenerationOutputs = await act_queue.get() + + # Evaluate the response through the environment + success, rewards = await env.step((qid, answers)) + # ... process results ... +``` + +#### Environment Integration + +The environment follows a +[Gym-like interface](https://github.com/Farama-Foundation/Gymnasium) with `reset` and +`step` methods, but uses asynchronous implementations to prevent blocking across +different environment instances. + +For math problems, the environment is typically stateless and acts as a wrapper around +your reward function: + +```python +class MathCodeSingleStepEnv(EnvironmentService): + async def step(self, action: Tuple[str, List[str]]): + qid, answers = action + # ... setup code ... + + # Run reward computation asynchronously + format_rewards = await asyncio.to_thread( + math_verify_call, + answers, + # ... other parameters ... + ) + return None, format_rewards, True, False, {} +``` + +#### Handling Multi-Turn Feedback + +After receiving the reward from `env.step`, check if the answer is correct. If not, +provide feedback and continue to the next turn: + +```python +for turn in range(self.num_turns): + # ... generation and evaluation code ... + + # Provide feedback based on the result + if success[0]: + feedback = "Congratulations! You are correct!" + else: + feedback = "Unfortunately your answer is wrong. Let's try again." + + # Format feedback as a user message + feedback = "\n" + self.tokenizer.apply_chat_template( + [{"content": feedback, "role": "user"}], + add_generation_prompt=True, + tokenize=False, + ) + + # Add feedback tokens to the conversation + feedback_tokens = self.tokenizer(feedback)["input_ids"] + token_ids.extend(feedback_tokens) +``` + +### Step 3: Register and Configure Your Agent + +First, register your agent implementation: + +```python +# in realhf/impl/agent/math_multi_turn_agent.py +register_agent("math-multi-turn", MathMultiTurnAgent) +``` + +```python +# in realhf/impl/agent/__init__.py +import realhf.impl.agent.math_multi_turn_agent +``` + +Then update your experiment configuration in +`realhf/experiments/async_exp/async_math_ppo.py`: + +```python +@dataclasses.dataclass +class AsyncPPOMATHConfig(AsyncRLExperimentConfig, PPOMATHConfig): + # Add any new CLI arguments your agent needs + my_param: float = 1.0 + + @property + def agent(self) -> AgentAbstraction: + return AgentAbstraction( + "math-multi-turn", # Your registered agent name + args=dict( + # Pass any arguments needed for your __init__ method + my_param=self.my_param, + # ... other configuration ... + ), + ) + + @property + def env(self) -> EnvServiceAbstraction: + # Update to use your custom environment if needed + return EnvServiceAbstraction( + "math-code-single-step", + args=dict(dataset_path=self.dataset.path) + ) +``` + +### Step 4: Run Training + +Follow the standard training procedure outlined in the +[quickstart guide](../tutorial/quickstart.md). Launch your experiment with: + +```bash +python3 training/main_async_ppo.py my_param=5.0 # plus any additional CLI arguments +``` + +### Training Results + +Here's an example of the training reward curve from our multi-turn math agent: + +![Multi-turn Training Rewards](multiturn_reward.png) + +The agent successfully learns to solve math problems with improved accuracy over time, +demonstrating the effectiveness of the multi-turn approach. diff --git a/docs/legacy/customization/algorithm.md b/docs/legacy/customization/algorithm.md new file mode 100644 index 0000000000..8c4b3991bc --- /dev/null +++ b/docs/legacy/customization/algorithm.md @@ -0,0 +1,259 @@ +# Training Algorithm + +> **Note**: The AReaLite approach is more recommended for new implementations due to its +> cleaner separation of concerns and better maintainability. + +The legacy approach encapsulates algorithms in a `ModelInterface` with three core +methods: + +```python +# From realhf/api/core/model_api.py +class ModelInterface(abc.ABC): + """Interface for model training, inference, and generation. + + This interface follows the dependency injection pattern, allowing + algorithms like REINFORCE and PPO to use the same underlying model + while exhibiting different training behaviors. + """ + + def inference( + self, + model: Model, + data: SequenceSample, + mb_spec: MicroBatchSpec, + ) -> SequenceSample | None: + raise NotImplementedError() + + def generate( + self, + model: Model, + data: SequenceSample, + mb_spec: MicroBatchSpec, + ) -> SequenceSample | None: + raise NotImplementedError() + + def train_step( + self, + model: Model, + data: SequenceSample, + mb_spec: MicroBatchSpec, + ) -> Dict | List[Dict]: + raise NotImplementedError() +``` + +When the dataflow is fixed, you typically only need to modify the algorithm interface +file. + +> **Note**: We recommend using asynchronous RL so you can customize generation behavior +> by [modifying your RL agent](agent.md) instead of the `generate` method. + +### Example 1: Grouped Advantage Normalization + +Let's modify PPO's global advantage normalization to use grouped normalization (GRPO +approach). + +#### Understanding Data Organization + +Each batch contains multiple prompts (batch size) and each prompt may have multiple +responses (group size). So total sequences = batch_size × group_size. + +Sequences have different lengths but are packed into a 1D tensor. We use `cu_seqlens` +(cumulative sequence lengths) to mark boundaries, similar to flash-attention. + +#### Implementation + +The standard PPO normalization looks like: + +```python +@dataclass +class PPOActorInterface(ModelInterface): + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + # ... + if self.adv_norm: + advantages = masked_normalization(advantages, loss_mask) + # ... +``` + +For grouped normalization, we partition advantages by group: + +```python +@dataclass +class PPOActorInterface(ModelInterface): + group_adv_norm: bool = False + + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + # ... + if self.adv_norm: + if not self.group_adv_norm: + advantages = masked_normalization(advantages, loss_mask) + else: + n_samples = data.bs + adv_list = [] + for i in range(0, n_samples, self.group_size): + # Define chunk boundaries + s = short1cu_seqlens[i] + e = short1cu_seqlens[i + self.group_size] + + # Extract advantages for this group + adv = advantages[s:e] + mask = loss_mask[s:e] + + # Normalize within group + advn = masked_normalization(adv, mask, all_reduce=False) + adv_list.append(advn) + + advantages = torch.cat(adv_list, 0) + # ... +``` + +#### Configuration Changes + +Update the experiment configuration to expose the new parameter: + +```python +@dataclasses.dataclass +class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): + group_adv_norm: bool = False + + @property + def rpcs(self): + # ... + actor_interface = ModelInterfaceAbstraction( + "ppo_actor", + args={ + **copy.deepcopy(self.ppo_kwargs), + "group_adv_norm": self.group_adv_norm, + # ... + }, + ) +``` + +### Example 2: Decoupled PPO Loss + +The decoupled PPO loss (from AReaL's paper) recomputes probabilities before mini-batch +updates and uses this as π_prox: + +![decoupled loss](decoupled_loss.png) + +#### Probability Recomputation + +We recompute probabilities using the existing `inference` method: + +```python +@dataclass +class PPOActorInterface(ModelInterface): + use_decoupled_loss: bool = False + + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + if self.use_decoupled_loss: + s: SequenceSample = self.inference(model, data, mb_spec) + prox_logp = s.data["logprobs"] + + # Prepare mini-batch data + flat_data = dict( + advantages=advantages, + old_logp=old_logp, + ppo_loss_mask=loss_mask, + packed_input_ids=input_.data["packed_input_ids"], + kl_rewards=kl_rewards, + ) + + if self.use_decoupled_loss: + flat_data["prox_logp"] = prox_logp.float() + + flat_input = SequenceSample.from_default( + ids=list(range(input_.bs * self.group_size)), + data=flat_data, + seqlens=[int(x) for x in input_lens.cpu().numpy().tolist()], + ) + + # Split into mini-batches and train + datas = flat_input.split_with_spec(spec) + for mb_i, data in enumerate(datas): + train_stat = module.train_batch( + input_=data, + mb_spec=mb_spec, + version_steps=model.version.global_step, + loss_fn=_loss_fn, + loss_weight_fn=lambda x: x.data["ppo_loss_mask"].count_nonzero(), + token_normalize_scope=self.token_normalize_scope, + ) +``` + +#### Modifying the Loss Function + +Update the loss computation to use the recomputed probabilities: + +```python +def _ppo_actor_loss_from_model_outputs( + logits: torch.FloatTensor, # [tot_seqlen, vocab_size] + input_: SequenceSample, + ... +) -> torch.Tensor: + # ... + prox_logp = input_.data.get("prox_logp") + + loss, ppo_stat = ppo_functional.actor_loss_fn( + logprobs=logprobs, + old_logprobs=old_logp, + advantages=advantages, + eps_clip=eps_clip, + loss_mask=ppo_loss_mask, + c_clip=c_clip, + proximal_logprobs=prox_logp, + behav_imp_weight_cap=behav_imp_weight_cap, + ) +``` + +And in the core loss function: + +```python +def actor_loss_fn( + logprobs: torch.FloatTensor, + old_logprobs: torch.FloatTensor, + advantages: torch.FloatTensor, + eps_clip: float, + loss_mask: Optional[torch.BoolTensor] = None, + c_clip: Optional[float] = None, + proximal_logprobs: Optional[torch.FloatTensor] = None, + behav_imp_weight_cap: Optional[torch.FloatTensor] = None, +) -> Tuple[torch.Tensor, Dict]: + # Use proximal probabilities if available, otherwise use old probabilities + denorm_logprobs = proximal_logprobs if proximal_logprobs is not None else old_logprobs + + loss_mask_count = loss_mask.count_nonzero() or 1 + + # Compute importance weights + ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) + + # Apply behavioral importance weighting for decoupled loss + if proximal_logprobs is not None: + behav_kl = proximal_logprobs - old_logprobs + behav_imp_weight = behav_kl.exp() + behav_kl = torch.where(loss_mask, behav_kl, 0.0) + behav_imp_weight = torch.where(loss_mask, behav_imp_weight, 0.0) + pg_loss = pg_loss * behav_imp_weight + + # ... + return pg_loss, stat +``` + +#### Configuration Update + +```python +@dataclasses.dataclass +class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): + use_decoupled_loss: bool = False + + @property + def rpcs(self): + # ... + actor_interface = ModelInterfaceAbstraction( + "ppo_actor", + args={ + **copy.deepcopy(self.ppo_kwargs), + "use_decoupled_loss": self.use_decoupled_loss, + # ... + }, + ) +``` diff --git a/docs/legacy/customization/dataset.md b/docs/legacy/customization/dataset.md new file mode 100644 index 0000000000..a4e03518a6 --- /dev/null +++ b/docs/legacy/customization/dataset.md @@ -0,0 +1,146 @@ +# Dataset (Legacy) + +> **Note**: While this legacy approach works, we strongly recommend using the AReaLite +> for new projects. It provides better flexibility, cleaner abstractions, and easier +> maintenance. + +### Define Your Dataset + +Create a new file under `realhf/impl/dataset/`, for example, `my_custom_dataset.py`. +Your `Dataset` must implement the `torch.utils.data.Dataset` interface and follow the +framework's conventions. + +```python +class MyCustomDataset(torch.utils.data.Dataset): + + def __init__( + self, + util: data_api.DatasetUtility, + max_length: Optional[int] = None, + dataset_path: Optional[str] = None, + dataset_builder: Optional[Callable[[], List[Dict]]] = None, + # Your custom parameters + custom_param: float = 1.0, + ): + """Custom dataset initialization + + Args: + util: Dataset utility class containing tokenizer, seed, distributed info, etc. + max_length: Maximum sequence length + dataset_path: Path to dataset file (optional) + dataset_builder: Data construction function (optional, alternative to dataset_path) + custom_param: Your custom parameter + """ + self._util = util + self.max_length = max_length + + # Load and split dataset + data = data_api.load_shuffle_split_dataset(util, dataset_path, dataset_builder) + + # Your custom data processing logic + ... +``` + +### Implement Core Methods + +Every dataset class must implement the following two core methods: + +#### 1. `__len__` Method + +Returns the size of the dataset: + +```python +def __len__(self): + return len(self.data_samples) +``` + +#### 2. `__getitem__` Method + +Returns the sample at the specified index, must return a `SequenceSample` object: + +```python +def __getitem__(self, idx): + # Get raw data + sample = self.data_samples[idx] + + # Process data + ... + + # Return SequenceSample object + return data_api.SequenceSample.from_default( + ids=[sample["id"]], + seqlens=[len(processed_data["input_ids"])], + data=dict( + packed_prompts=torch.tensor(processed_data["input_ids"], dtype=torch.long), + # Other necessary data fields + ), + ) +``` + +#### Dataset Examples + +We provide some examples of dataset under `realhf/impl/dataset/`: + +- For SFT, please refer `prompt_answer_dataset.py`. +- For Reward model training, please refer `rw_paired_dataset.py` +- For RL training, please refer `math_code_dataset.py` + +### Data Format Requirements + +#### JSONL File Format + +Your data file should be in JSONL format, with one JSON object per line. If you are +using our PromptDataset implementation, your data should be like: + +- Math Data + +```json +{"qid": "sample_1", "prompt": "Solve this math problem: 2+2=", "solutions": ["\\boxed{4}"]} +``` + +- Code Data + +```json +{"qid": "sample_2", "prompt": "Code problem", "input_output": "{\"inputs\": [\"5\\n2 3 5 10 12\\n\"], \"outputs\": [\"17\\n\"]}"} +``` + +- `qid`: Unique identifier for the sample +- `prompt`: Input prompt text +- `task`: Task type, used to distinguish how to calculate the reward. ("math" and "code" + are supported now.) + +Note: There is no format restriction for a customized dataset as long as it can be +loaded by your custom code. + +### Registration and Configuration + +#### Register Dataset + +Register your dataset at the end of your dataset file: + +```python +# in realhf/impl/dataset/my_custom_dataset.py +data_api.register_dataset("my-custom", MyCustomDataset) +``` + +#### Modify Experiment Configuration + +Use your new dataset in the experiment configuration (refer to +`realhf/experiments/common/*_exp.py`): + +```python +# in your experiment config file +@property +def datasets(self) -> List[DatasetAbstraction]: + return [ + DatasetAbstraction( + "my-custom", # Your registered name + args=dict( + dataset_path=self.dataset_path, + max_length=self.max_length, + custom_param=self.custom_param, + # Other initialization parameters + ), + ) + ] +``` diff --git a/docs/customization/decoupled_loss.png b/docs/legacy/customization/decoupled_loss.png similarity index 100% rename from docs/customization/decoupled_loss.png rename to docs/legacy/customization/decoupled_loss.png diff --git a/docs/customization/multiturn_reward.png b/docs/legacy/customization/multiturn_reward.png similarity index 100% rename from docs/customization/multiturn_reward.png rename to docs/legacy/customization/multiturn_reward.png From 8e966466f86aeb5b46b6933d5383e276e1ba8aac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Tue, 22 Jul 2025 15:18:01 +0800 Subject: [PATCH 15/16] . --- docs/customization/agent.md | 14 +++++++------- docs/customization/algorithm.md | 8 +++----- docs/customization/dataset.md | 4 ++-- docs/legacy/customization/agent.md | 16 ++++++++-------- docs/legacy/customization/algorithm.md | 16 ++++++++-------- docs/legacy/customization/dataset.md | 20 ++++++++++---------- 6 files changed, 38 insertions(+), 40 deletions(-) diff --git a/docs/customization/agent.md b/docs/customization/agent.md index c668f87f2d..c71aeab050 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -6,7 +6,7 @@ problems until it finds the correct answer. You can find the complete implementation in `arealite/workflow/multi_turn.py`. -### Step 1: Define Your Workflow +## Step 1: Define Your Workflow AReaLite gives you flexibility in how you design your agents. Instead of rigid `Agent` classes that might constrain your agent's capabilities, AReaLite captures all rollout @@ -40,7 +40,7 @@ interact. > generated from that prompt—it's not batched. However, you can generate multiple > trajectories from a single prompt (for example, with GRPO or tree search). -#### Setting Up the Multi-Turn Math Workflow +### Setting Up the Multi-Turn Math Workflow Let's build a multi-turn rollout workflow for solving math problems. First, we'll define the `__init__` method to set up what we need during rollout: @@ -68,7 +68,7 @@ class MultiTurnWorkflow(RolloutWorkflow): self.turn_discount = turn_discount ``` -#### Implementing the Episode Logic +### Implementing the Episode Logic Now let's implement the `arun_episode` method. We'll start by tokenizing the prompt data and converting it into an `LLMRequest` object for the inference engine: @@ -111,7 +111,7 @@ class MultiTurnWorkflow(RolloutWorkflow): > **Note**: The `rid` field in `LLMRequest` is the request ID. Requests with the same ID > will reuse the LLM inference server's KV caches for better efficiency. -#### Handling Multi-Turn Conversations +### Handling Multi-Turn Conversations Next, we'll check if the current answer is correct using our `reward_fn`. This function should return 1 for correct answers and 0 otherwise. When the answer is wrong, we'll @@ -150,7 +150,7 @@ class MultiTurnWorkflow(RolloutWorkflow): discount *= self.turn_discount ``` -#### Reward Function Signature +### Reward Function Signature To make it easier to switch between different reward functions, we recommend following this signature: @@ -183,7 +183,7 @@ def reward_fn( While this signature is convenient, you're not restricted to it in custom workflows—modify as needed for your specific use case. -#### Collecting Training Data +### Collecting Training Data Finally, let's complete the implementation by collecting trajectories in the `TensorDict` format: @@ -228,7 +228,7 @@ class MultiTurnWorkflow(RolloutWorkflow): > `input_ids`, `loss_mask`, `attention_mask`, and `logprobs` (needed for computing > importance ratios). -### Step 2: Training with Your Custom Workflow +## Step 2: Training with Your Custom Workflow Using your custom workflow is straightforward—just create it in your training script and pass it to the `rollout_batch` or `prepare_batch` method: diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index a10ddfd681..fe60a0171a 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -3,8 +3,6 @@ > **Note**: We recommend the user to first read the > [agent customization guide](agent.md). -## Option 1: Using AReaLite (Recommended) - AReaLite structures RL algorithms around two core components: - **RolloutWorkflow**: Defines what data to generate during rollouts @@ -12,7 +10,7 @@ AReaLite structures RL algorithms around two core components: We'll demonstrate this by implementing an RL algorithm similar to ReMax. -### Step 1: Implementing the RolloutWorkflow +## Step 1: Implementing the RolloutWorkflow The rollout workflow generates both greedy and sampled completions, then uses the reward difference as the final training signal: @@ -87,7 +85,7 @@ class ReMaxRLVRWorkflow(RolloutWorkflow): > **Note**: For detailed guidance on customizing rollout workflows, see the > [agent customization guide](agent.md). -### Step 2: Implementing the REINFORCE Training Algorithm +## Step 2: Implementing the REINFORCE Training Algorithm Training algorithms are implemented by subclassing `TrainEngine` and using its atomic operations like `forward`, `train_batch`, and `eval_batch`. @@ -151,7 +149,7 @@ class FSDPReinforceActor(FSDPEngine): > **Note**: This pattern is similar to interfaces in Go or traits in Rust, adapted for > Python's object model. -### Step 3: Composing the Complete Training Loop +## Step 3: Composing the Complete Training Loop The main training loop brings everything together: diff --git a/docs/customization/dataset.md b/docs/customization/dataset.md index 1382726315..8dd060956d 100644 --- a/docs/customization/dataset.md +++ b/docs/customization/dataset.md @@ -10,7 +10,7 @@ offline training, such as `LMEngine` for Supervised Fine-Tuning (SFT)). Here are two concrete examples from the existing implementation: -### SFT (Offline Training) +## SFT (Offline Training) In the SFT example, we see that the loaded data is directly passed to the `train_lm` method: @@ -56,7 +56,7 @@ def get_gsm8k_dataset(split, tokenizer, rank, world_size): return process_gsm8k_sft_dataset(dataset, tokenizer) ``` -### GRPO (Online Training) +## GRPO (Online Training) In the GRPO example, the loaded data is passed to the `InferenceEngine`, rather than the `TrainEngine`: diff --git a/docs/legacy/customization/agent.md b/docs/legacy/customization/agent.md index bbfe91d362..c926a6d8ee 100644 --- a/docs/legacy/customization/agent.md +++ b/docs/legacy/customization/agent.md @@ -4,7 +4,7 @@ > for new projects. It provides better flexibility, cleaner abstractions, and easier > maintenance. -### Step 1: Define Your Agent Class +## Step 1: Define Your Agent Class Create a new file under `realhf/impl/agent/`, such as `math_multi_turn_agent.py`. Your `Agent` must implement the interface defined in `realhf/api/core/agent.py`, which @@ -23,7 +23,7 @@ class MathMultiTurnAgent(Agent): ... ``` -### Step 2: Implement the Trajectory Collection Logic +## Step 2: Implement the Trajectory Collection Logic The `collect_trajectory` method takes a task prompt, an environment, and two communication queues. Within this method, you control the data flow between your agent @@ -48,7 +48,7 @@ for turn in range(self.num_turns): # ... process results ... ``` -#### Environment Integration +### Environment Integration The environment follows a [Gym-like interface](https://github.com/Farama-Foundation/Gymnasium) with `reset` and @@ -73,7 +73,7 @@ class MathCodeSingleStepEnv(EnvironmentService): return None, format_rewards, True, False, {} ``` -#### Handling Multi-Turn Feedback +### Handling Multi-Turn Feedback After receiving the reward from `env.step`, check if the answer is correct. If not, provide feedback and continue to the next turn: @@ -100,7 +100,7 @@ for turn in range(self.num_turns): token_ids.extend(feedback_tokens) ``` -### Step 3: Register and Configure Your Agent +## Step 3: Register and Configure Your Agent First, register your agent implementation: @@ -143,16 +143,16 @@ class AsyncPPOMATHConfig(AsyncRLExperimentConfig, PPOMATHConfig): ) ``` -### Step 4: Run Training +## Step 4: Run Training Follow the standard training procedure outlined in the -[quickstart guide](../tutorial/quickstart.md). Launch your experiment with: +[quickstart guide](../../tutorial/quickstart.md). Launch your experiment with: ```bash python3 training/main_async_ppo.py my_param=5.0 # plus any additional CLI arguments ``` -### Training Results +## Training Results Here's an example of the training reward curve from our multi-turn math agent: diff --git a/docs/legacy/customization/algorithm.md b/docs/legacy/customization/algorithm.md index 8c4b3991bc..e8289cf9fd 100644 --- a/docs/legacy/customization/algorithm.md +++ b/docs/legacy/customization/algorithm.md @@ -47,12 +47,12 @@ file. > **Note**: We recommend using asynchronous RL so you can customize generation behavior > by [modifying your RL agent](agent.md) instead of the `generate` method. -### Example 1: Grouped Advantage Normalization +## Example 1: Grouped Advantage Normalization Let's modify PPO's global advantage normalization to use grouped normalization (GRPO approach). -#### Understanding Data Organization +### Understanding Data Organization Each batch contains multiple prompts (batch size) and each prompt may have multiple responses (group size). So total sequences = batch_size × group_size. @@ -60,7 +60,7 @@ responses (group size). So total sequences = batch_size × group_size. Sequences have different lengths but are packed into a 1D tensor. We use `cu_seqlens` (cumulative sequence lengths) to mark boundaries, similar to flash-attention. -#### Implementation +### Implementation The standard PPO normalization looks like: @@ -106,7 +106,7 @@ class PPOActorInterface(ModelInterface): # ... ``` -#### Configuration Changes +### Configuration Changes Update the experiment configuration to expose the new parameter: @@ -128,14 +128,14 @@ class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): ) ``` -### Example 2: Decoupled PPO Loss +## Example 2: Decoupled PPO Loss The decoupled PPO loss (from AReaL's paper) recomputes probabilities before mini-batch updates and uses this as π_prox: ![decoupled loss](decoupled_loss.png) -#### Probability Recomputation +### Probability Recomputation We recompute probabilities using the existing `inference` method: @@ -180,7 +180,7 @@ class PPOActorInterface(ModelInterface): ) ``` -#### Modifying the Loss Function +### Modifying the Loss Function Update the loss computation to use the recomputed probabilities: @@ -238,7 +238,7 @@ def actor_loss_fn( return pg_loss, stat ``` -#### Configuration Update +### Configuration Update ```python @dataclasses.dataclass diff --git a/docs/legacy/customization/dataset.md b/docs/legacy/customization/dataset.md index a4e03518a6..5a2bcac652 100644 --- a/docs/legacy/customization/dataset.md +++ b/docs/legacy/customization/dataset.md @@ -4,7 +4,7 @@ > for new projects. It provides better flexibility, cleaner abstractions, and easier > maintenance. -### Define Your Dataset +## Define Your Dataset Create a new file under `realhf/impl/dataset/`, for example, `my_custom_dataset.py`. Your `Dataset` must implement the `torch.utils.data.Dataset` interface and follow the @@ -41,11 +41,11 @@ class MyCustomDataset(torch.utils.data.Dataset): ... ``` -### Implement Core Methods +## Implement Core Methods Every dataset class must implement the following two core methods: -#### 1. `__len__` Method +### 1. `__len__` Method Returns the size of the dataset: @@ -54,7 +54,7 @@ def __len__(self): return len(self.data_samples) ``` -#### 2. `__getitem__` Method +### 2. `__getitem__` Method Returns the sample at the specified index, must return a `SequenceSample` object: @@ -77,7 +77,7 @@ def __getitem__(self, idx): ) ``` -#### Dataset Examples +### Dataset Examples We provide some examples of dataset under `realhf/impl/dataset/`: @@ -85,9 +85,9 @@ We provide some examples of dataset under `realhf/impl/dataset/`: - For Reward model training, please refer `rw_paired_dataset.py` - For RL training, please refer `math_code_dataset.py` -### Data Format Requirements +## Data Format Requirements -#### JSONL File Format +### JSONL File Format Your data file should be in JSONL format, with one JSON object per line. If you are using our PromptDataset implementation, your data should be like: @@ -112,9 +112,9 @@ using our PromptDataset implementation, your data should be like: Note: There is no format restriction for a customized dataset as long as it can be loaded by your custom code. -### Registration and Configuration +## Registration and Configuration -#### Register Dataset +### Register Dataset Register your dataset at the end of your dataset file: @@ -123,7 +123,7 @@ Register your dataset at the end of your dataset file: data_api.register_dataset("my-custom", MyCustomDataset) ``` -#### Modify Experiment Configuration +### Modify Experiment Configuration Use your new dataset in the experiment configuration (refer to `realhf/experiments/common/*_exp.py`): From 5d58cb076f4af0125a573d71d55a267f4a327c71 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Tue, 22 Jul 2025 15:20:48 +0800 Subject: [PATCH 16/16] . --- docs/customization/algorithm.md | 2 +- docs/customization/dataset.md | 6 +++--- docs/legacy/customization/algorithm.md | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index fe60a0171a..f087f83dbe 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -3,7 +3,7 @@ > **Note**: We recommend the user to first read the > [agent customization guide](agent.md). -AReaLite structures RL algorithms around two core components: +**AReaLite** structures RL algorithms around two core components: - **RolloutWorkflow**: Defines what data to generate during rollouts - **TrainEngine**: Defines how to process the generated data for training diff --git a/docs/customization/dataset.md b/docs/customization/dataset.md index 8dd060956d..e808301c74 100644 --- a/docs/customization/dataset.md +++ b/docs/customization/dataset.md @@ -1,8 +1,8 @@ # Dataset -AReaLite directly integrates with the `Dataset` class from the HuggingFace `datasets` -package. This gives you full flexibility to load, process, and filter your data before -training. +**AReaLite** directly integrates with the `Dataset` class from the HuggingFace +`datasets` package. This gives you full flexibility to load, process, and filter your +data before training. The required columns in your dataset depend on the specific implementation of the `RolloutWorkflow` (for online reinforcement learning) or the training engines (for diff --git a/docs/legacy/customization/algorithm.md b/docs/legacy/customization/algorithm.md index e8289cf9fd..6e1e44b45e 100644 --- a/docs/legacy/customization/algorithm.md +++ b/docs/legacy/customization/algorithm.md @@ -1,4 +1,4 @@ -# Training Algorithm +# Training Algorithm (Legacy) > **Note**: The AReaLite approach is more recommended for new implementations due to its > cleaner separation of concerns and better maintainability.