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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] . --- 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 21bd032eb9872ccbd6e1aac140734945390a9e1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Thu, 24 Jul 2025 13:46:01 +0800 Subject: [PATCH 11/20] PullRequest: 408 [Feature] Bump SGLang version to v0.4.9.post2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/sgl049 of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/408 Reviewed-by: 晓雷 * . * bump arealite to sglang 0.4.9.post2 * . * PullRequest: 412 [lite] Minor refactor on `UpdateWeightMeta` --- arealite/README.md | 1 - arealite/api/cli_args.py | 79 ------ arealite/api/engine_api.py | 29 ++- arealite/api/io_struct.py | 67 +++-- arealite/engine/base_hf_engine.py | 11 +- arealite/engine/fsdp_engine.py | 43 ++-- arealite/engine/sglang_remote.py | 248 ++++++++++--------- arealite/launcher/local.py | 2 +- arealite/launcher/sglang_server.py | 46 +--- arealite/tests/test_fsdp_engine_nccl.py | 41 ++- arealite/tests/test_sglang_engine.py | 13 +- examples/arealite/boba.py | 45 ++-- examples/arealite/configs/boba.yaml | 6 +- examples/arealite/configs/gsm8k_grpo.yaml | 4 +- examples/arealite/gsm8k_grpo.py | 39 +-- examples/env/scripts/setup-container-deps.sh | 11 - examples/env/scripts/setup-eval-pip-deps.sh | 6 +- examples/env/scripts/setup-pip-deps.sh | 16 +- patch/sglang/v0.4.6.post2.patch | 144 ----------- patch/sglang/v0.4.6.post4.patch | 144 ----------- pyproject.toml | 3 +- realhf/system/generation_server.py | 31 --- realhf/system/gserver_manager.py | 49 ++-- requirements.txt | 1 - 24 files changed, 336 insertions(+), 743 deletions(-) delete mode 100644 examples/env/scripts/setup-container-deps.sh delete mode 100644 patch/sglang/v0.4.6.post2.patch delete mode 100644 patch/sglang/v0.4.6.post4.patch diff --git a/arealite/README.md b/arealite/README.md index 65c49fccd0..4a23ce60a4 100644 --- a/arealite/README.md +++ b/arealite/README.md @@ -163,7 +163,6 @@ class WeightUpdateMeta: type: str path: str | None alloc_mode: AllocationMode | None - comm_backend: str | None @dataclass class SaveLoadMeta: diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 3b2afa55d0..7dc9be6dcf 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -285,85 +285,6 @@ class PPOActorConfig(TrainEngineConfig): ) -@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: diff --git a/arealite/api/engine_api.py b/arealite/api/engine_api.py index 9fb24d222e..bd00eb5757 100644 --- a/arealite/api/engine_api.py +++ b/arealite/api/engine_api.py @@ -12,6 +12,7 @@ FinetuneSpec, LLMRequest, LLMResponse, + ParamSpec, SaveLoadMeta, WeightUpdateMeta, ) @@ -63,7 +64,19 @@ def eval(self): return self.train(False) def upload_weights(self, meta: WeightUpdateMeta): - """Upload weights to the inference engine.""" + """Upload weights to the inference engine (in a blocking manner).""" + raise NotImplementedError() + + def get_param_specs(self) -> List[ParamSpec]: + """Get the parameter specifications for the model.""" + raise NotImplementedError() + + def set_version(self, version: int): + """Set the current weight version in the train engine.""" + raise NotImplementedError() + + def get_version(self) -> int: + """Get the current weight version in the train engine.""" raise NotImplementedError() def save(self, meta: SaveLoadMeta): @@ -122,12 +135,20 @@ def initialize(self, addr: str | None, ft_spec): def destroy(self): """Destroy the engine and release GPU memory.""" + async def agenerate(self, req: LLMRequest) -> LLMResponse: + """Asynchronously generate a response for the given request.""" + raise NotImplementedError() + def update_weights(self, meta: WeightUpdateMeta) -> Future: - """Update weights in the inference engine.""" + """Update weights in the inference engine in a non-blocking manner.""" raise NotImplementedError() - async def agenerate(self, req: LLMRequest) -> LLMResponse: - """Asynchronously generate a response for the given request.""" + def set_version(self, version: int) -> None: + """Set the current weight version in the inference engine.""" + raise NotImplementedError() + + def get_version(self) -> int: + """Get the current weight version in the inference engine.""" raise NotImplementedError() def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: diff --git a/arealite/api/io_struct.py b/arealite/api/io_struct.py index bd11d02753..ceb59535fc 100644 --- a/arealite/api/io_struct.py +++ b/arealite/api/io_struct.py @@ -3,13 +3,18 @@ import enum import itertools import re +import os import uuid from dataclasses import dataclass, field -from typing import Any, Dict, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple from transformers import PreTrainedTokenizerFast -from arealite.api.cli_args import GenerationHyperparameters +from arealite.api.cli_args import GenerationHyperparameters, SaverConfig +from arealite.utils.network import find_free_ports, gethostip + +if TYPE_CHECKING: + from arealite.api.engine_api import TrainEngine @dataclass @@ -155,20 +160,54 @@ def extract_decoupled_alloc(allocation_mode: str) -> Dict: return other_alloc +@dataclass +class ParamSpec: + name: str + shape: Tuple + dtype: str + + @dataclass class WeightUpdateMeta: - type: str - path: str | None - alloc_mode: AllocationMode | None - comm_backend: str | None - model_version: int = 0 - tp_size: int = 1 - master_address: str = "127.0.0.1" - master_port: int = 29500 - world_size: int = 1 - group_name: str = "aupdate_weights_from_distributed" - parameter_names: List[str] = field(default_factory=list) - state_dict_key_to_shape: Dict[str, Tuple[int]] = field(default_factory=dict) + type: Literal["disk", "nccl"] + path: str | None = None + alloc_mode: AllocationMode | None = None + + nccl_master_address: str = "127.0.0.1" + nccl_master_port: int = 29500 + nccl_param_specs: List[ParamSpec] = field(default_factory=list) + nccl_group_name: str = "update_weight_group" + + @classmethod + def from_disk( + cls, + saver_config: SaverConfig, + ) -> "WeightUpdateMeta": + from arealite.utils.saver import Saver + path = os.path.join( + Saver.get_save_checkpoint_root(saver_config), + "weight_update", + ) + return cls( + type="disk", + path=path, + ) + + @classmethod + def from_fsdp_nccl( + cls, + allocation_mode: AllocationMode, + fsdp_engine: "TrainEngine", + nccl_group_name: str = "update_weight_group", + ): + return cls( + type="nccl", + alloc_mode=allocation_mode, + nccl_master_address=gethostip(), + nccl_master_port=find_free_ports(1)[0], + nccl_param_specs=fsdp_engine.get_param_specs(), + nccl_group_name=nccl_group_name, + ) @dataclass diff --git a/arealite/engine/base_hf_engine.py b/arealite/engine/base_hf_engine.py index 8d0b1a80f8..f6b9e2a99f 100644 --- a/arealite/engine/base_hf_engine.py +++ b/arealite/engine/base_hf_engine.py @@ -46,6 +46,7 @@ def __init__(self, config: TrainEngineConfig): self.tokenizer: PreTrainedTokenizerFast # huggingface model config self.model_config: PretrainedConfig + self._version: int = 0 # initialization self.initialized = False @@ -55,6 +56,12 @@ def __init__(self, config: TrainEngineConfig): self.world_size = int(os.environ["WORLD_SIZE"]) + def set_version(self, version: int): + self._version = version + + def get_version(self) -> int: + return self._version + def train(self, mode: bool = True): assert self.model is not None self.model.train(mode=mode) @@ -191,7 +198,7 @@ def save_optimizer_state(self, path: str): ) state_dict = self.optimizer.state_dict() torch.save(state_dict, shard_path) - dist.barrier() + dist.barrier(device_ids=[self.device.index]) def load_optimizer_state(self, path: str): # Load FSDP sharded state dict @@ -203,7 +210,7 @@ def load_optimizer_state(self, path: str): ) optimizer_state_dict = torch.load(shard_path, weights_only=False) self.optimizer.load_state_dict(optimizer_state_dict) - dist.barrier() + dist.barrier(device_ids=[self.device.index]) def step_lr_scheduler(self): assert self.lr_scheduler is not None diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index a0926c3fc9..4338b6d55d 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -1,7 +1,7 @@ import os import time from datetime import datetime -from typing import Callable, Dict, Optional, Tuple +from typing import Callable, Dict, List, Optional import torch import torch.distributed as dist @@ -14,7 +14,8 @@ from transformers import PreTrainedTokenizerFast from arealite.api.cli_args import TrainEngineConfig -from arealite.api.engine_api import FinetuneSpec, SaveLoadMeta, WeightUpdateMeta +from arealite.api.engine_api import FinetuneSpec +from arealite.api.io_struct import ParamSpec, SaveLoadMeta, WeightUpdateMeta from arealite.engine.base_hf_engine import BaseHFEngine from arealite.utils.distributed import init_custom_process_group from arealite.utils.fsdp import ( @@ -119,7 +120,7 @@ def _save_model_to_hf( if tokenizer is not None: tokenizer.save_pretrained(path) - dist.barrier() + dist.barrier(device_ids=[self.device.index]) def _load_model_from_hf(self, path: str): """Load model from HuggingFace format.""" @@ -140,7 +141,7 @@ def upload_weights(self, meta: WeightUpdateMeta): if not self.weight_update_group_initialized: self._init_distributed_weight_update(meta) self._update_weights_from_distributed() - dist.barrier() + dist.barrier(device_ids=[self.device.index]) torch.cuda.synchronize() elif meta.type == "disk": self._save_model_to_hf(meta.path, self.tokenizer) @@ -149,7 +150,7 @@ def upload_weights(self, meta: WeightUpdateMeta): update_name = names.update_weights_from_disk( self.config.experiment_name, self.config.trial_name, - meta.model_version, + self.model_version, ) name_resolve.add( update_name, str(datetime.now().timestamp()), keepalive_ttl=120 @@ -158,16 +159,18 @@ def upload_weights(self, meta: WeightUpdateMeta): raise ValueError(f"Unknown weight update type {meta.type}") def _init_distributed_weight_update(self, meta: WeightUpdateMeta): + # NOTE: Processes launched with torchrun will set the following env var to True, + # which blocks creating another TCP store for weight update. + os.environ["TORCHELASTIC_USE_AGENT_STORE"] = str(False) if dist.get_rank() == 0: self.weight_update_group = init_custom_process_group( backend="nccl", - world_size=meta.world_size, - init_method=f"tcp://{meta.master_address}:{meta.master_port}", + world_size=meta.alloc_mode.gen_world_size + 1, + init_method=f"tcp://{meta.nccl_master_address}:{meta.nccl_master_port}", rank=0, - group_name=meta.group_name, + group_name=meta.nccl_group_name, ) - # NOTE: synchronizing with sglang's barrier - dist.barrier(group=self.weight_update_group, device_ids=[self.device.index]) + # NOTE: sglang v0.4.9.post2 or later does not have the barrier call self.weight_update_group_initialized = True def _update_weights_from_distributed(self): @@ -179,23 +182,29 @@ def _update_weights_from_distributed(self): else: tensor = param.data if dist.get_rank() == 0: - print(f"Broadcasting {name} with shape {tensor.shape}", flush=True) + logger.debug( + f"Broadcasting {name} with shape {tensor.shape}", flush=True + ) dist.broadcast(tensor, src=0, group=self.weight_update_group) - dist.barrier() del tensor # optional, for memory hygiene torch.cuda.empty_cache() - def get_param_meta_for_distributed_update(self) -> Dict[str, Tuple[int]]: - """Return a dict mapping param name to its shape (expanded if DTensor).""" - param_shapes = {} + def get_param_specs(self) -> List[ParamSpec]: + param_specs = [] for name, param in self.model.named_parameters(): if isinstance(param.data, DTensor): tensor = param.data.full_tensor() else: tensor = param.data - param_shapes[name] = tuple(tensor.shape) + param_specs.append( + ParamSpec( + name=name, + shape=tuple(tensor.shape), + dtype=str(tensor.dtype).split("torch.")[1], + ) + ) del tensor # free memory if full_tensor was created - return param_shapes + return param_specs def train_batch( self, diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index 4c838f4370..bac9de346f 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -1,10 +1,11 @@ import asyncio import os import random +import shutil import threading import time import traceback -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from concurrent.futures import Future, ProcessPoolExecutor from datetime import datetime from queue import Empty, Full, Queue from typing import TYPE_CHECKING, Any, Callable, Dict, List @@ -27,17 +28,12 @@ ) 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 +from realhf.base import logging, name_resolve, names if TYPE_CHECKING: from arealite.api.workflow_api import RolloutWorkflow logger = logging.getLogger(__name__) -if pkg_version.is_available("sglang"): - if pkg_version.is_version_greater_or_equal("sglang", "0.4.4"): - SGLANG_TOKEN_OUTPUT_IDENTIFIER = "output_ids" - else: - SGLANG_TOKEN_OUTPUT_IDENTIFIER = "token_ids" ROLLOUT_POLL_WAIT_TIME = 0.05 RID_CACHE_SIZE = 128 @@ -91,10 +87,7 @@ def _wait_for_server(self, address): def check_health(self, base_url): # Check server endpoint try: - response = requests.get( - f"{base_url}/metrics", - timeout=30, - ) + response = requests.get(f"{base_url}/health", timeout=30) return response.status_code == 200 except requests.exceptions.RequestException as e: return False @@ -123,7 +116,7 @@ def _rollout_thread(self): """Thread that runs the rollout loop.""" try: uvloop.run(self._rollout_thread_async()) - except Exception as e: + except Exception: traceback.print_exc() async def _rollout_thread_async(self): @@ -259,9 +252,7 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: accumulated_output_logprobs = [] accumulated_versions = [] - # Deal with rollout interruption - stop_reason = "length" - + # A single "rid" shares the same sever to allow KV cache reuse if req.rid in self.rid_to_address: server_addr = self.rid_to_address[req.rid] else: @@ -273,10 +264,19 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: self.rid_to_address[req.rid] = server_addr self.rid_queue.append(req.rid) + # Deal with rollout interruption + # "abort" is the stop reason for later v0.4.9.post2 after + # we call the pause_generation endpoint + stop_reason = None while ( stop_reason != "stop" and len(accumulated_output_tokens) < gconfig.max_new_tokens ): + # Request is interrupted, wait for some time to avoid interfering + # with update weights requests + if stop_reason is not None: + await asyncio.sleep(0.5) + # loop until the generation is complete result = await arequest_with_retry( session=self.session, @@ -288,8 +288,17 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: timeout=self.config.request_timeout, ) - # Parse response meta_info = result["meta_info"] + # Check if generation is complete + finish_reason = meta_info["finish_reason"] + stop_reason = finish_reason["type"] + if ( + stop_reason == "abort" + and finish_reason.get("message") == "Abort before prefill" + ): + continue + + # Parse response output_tokens = [x[1] for x in meta_info["output_token_logprobs"]] output_logprobs = [x[0] for x in meta_info["output_token_logprobs"]] @@ -299,11 +308,7 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: # FIXME: Update with actual server versions accumulated_versions.extend([-1] * len(output_tokens)) - # Check if generation is complete - finish_reason = meta_info["finish_reason"] - stop_reason = finish_reason["type"] - - payload["input_ids"] += result[SGLANG_TOKEN_OUTPUT_IDENTIFIER] + payload["input_ids"] += result["output_ids"] sample_params["max_new_tokens"] -= len(output_tokens) latency = time.perf_counter() - start_time @@ -318,30 +323,24 @@ 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): + for addr in self.addresses: + res = requests.post(f"http://{addr}/pause_generation") + res.raise_for_status() + fut = Future() if meta.type == "nccl": - if not self.distributed_weight_update_initialized: - self._init_distributed_weight_update(meta) - tik = time.perf_counter() - try: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - for param in meta.parameter_names: - jobs = [ - self.aupdate_weights_from_distributed(addr, meta, param) - for addr in self.addresses - ] - loop.run_until_complete(asyncio.gather(*jobs)) - finally: - loop.close() - logger.info( - f"Distributed update weights done in {time.perf_counter() - tik}s" + fut = self.executor.submit( + update_weights_from_distributed, + meta, + self.addresses, + self.config.request_timeout, + not self.distributed_weight_update_initialized, ) - self.set_version(meta.model_version) + + def callback(fut): + self.distributed_weight_update_initialized = True + + fut.add_done_callback(callback) elif meta.type == "disk": # Update weights from disk # Use ProcessPool to bypass python GIL for running async coroutines @@ -349,7 +348,7 @@ def _update_weights(self, meta: WeightUpdateMeta): update_weights_from_disk, self.config.experiment_name, self.config.trial_name, - meta.model_version, + self.get_version(), self.addresses, meta.path, self.config.request_retries, @@ -357,64 +356,19 @@ def _update_weights(self, meta: WeightUpdateMeta): ) def callback(fut): - self.set_version(meta.model_version) + shutil.rmtree(meta.path, ignore_errors=True) fut.add_done_callback(callback) - return fut else: raise NotImplementedError(f"Unsupported weight update type: {meta.type}") - def _init_distributed_weight_update(self, meta: WeightUpdateMeta): - try: - # Initialize weights update group - jobs = [ - self.ainit_weights_update_group(addr, meta) 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)) - self.distributed_weight_update_initialized = True - logger.info(f"Distributed update weights initialized") - finally: - loop.close() + def callback(fut): + for addr in self.addresses: + res = requests.post(f"http://{addr}/continue_generation") + res.raise_for_status() - async def ainit_weights_update_group(self, addr: str, meta: WeightUpdateMeta): - rank_offset = 1 + self.addresses.index(addr) * meta.tp_size - payload = { - "master_address": meta.master_address, - "master_port": str(meta.master_port), - "rank_offset": rank_offset, - "world_size": meta.world_size, - "group_name": meta.group_name, - "backend": "nccl", - } - res = await arequest_with_retry( - addr=addr, - endpoint="/init_weights_update_group", - payload=payload, - method="POST", - max_retries=1, - timeout=self.config.request_timeout, - ) - assert res["success"] - - async def aupdate_weights_from_distributed( - self, addr: str, meta: WeightUpdateMeta, parameter_name: str - ): - res = await arequest_with_retry( - addr=addr, - endpoint="/update_weights_from_distributed", - payload={ - "name": parameter_name, - "dtype": "bfloat16", - "shape": meta.state_dict_key_to_shape[parameter_name], - }, - method="POST", - max_retries=1, - timeout=self.config.request_timeout, - ) - assert res["success"] + fut.add_done_callback(callback) + return fut def get_capacity(self): if dist.is_initialized(): @@ -519,27 +473,6 @@ def resume(self): self.paused.clear() -async def aupdate_weights_from_disk( - 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", - 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, @@ -569,12 +502,14 @@ async def _fn(): connector=get_default_connector(), ) jobs = [ - aupdate_weights_from_disk( - session=session, + arequest_with_retry( addr=addr, - path=path, - request_retries=request_retries, - request_timeout=request_timeout, + session=session, + endpoint="/update_weights_from_disk", + payload=dict(model_path=str(path)), + method="POST", + max_retries=request_retries, + timeout=request_timeout, ) for addr in addresses ] @@ -585,3 +520,72 @@ async def _fn(): ) return uvloop.run(_fn()) + + +def update_weights_from_distributed( + meta: WeightUpdateMeta, + addresses: List[str], + request_timeout, + init_group: bool, +): + async def _fn(): + tik = time.perf_counter() + if init_group: + await asyncio.gather( + *[ + ainit_weights_update_group(addr, i, meta, request_timeout) + for i, addr in enumerate(addresses) + ] + ) + await asyncio.gather( + *[ + arequest_with_retry( + addr=addr, + endpoint="/update_weights_from_distributed", + payload={ + "names": [pspec.name for pspec in meta.nccl_param_specs], + "dtypes": [pspec.dtype for pspec in meta.nccl_param_specs], + "shapes": [pspec.shape for pspec in meta.nccl_param_specs], + "group_name": meta.nccl_group_name, + }, + method="POST", + max_retries=1, + timeout=request_timeout, + ) + for addr in addresses + ] + ) + logger.info(f"Distributed update weights done in {time.perf_counter() - tik}s") + + return uvloop.run(_fn()) + + +async def ainit_weights_update_group( + addr: str, + server_idx: int, + meta: WeightUpdateMeta, + request_timeout: float, +): + assert meta.alloc_mode is not None + if meta.alloc_mode.gen_pp_size != 1: + raise NotImplementedError( + "NCCL weight update with PP size > 1 is not implemented yet." + ) + rank_offset = 1 + server_idx * meta.alloc_mode.gen_tp_size + payload = { + "master_address": meta.nccl_master_address, + "master_port": str(meta.nccl_master_port), + "rank_offset": rank_offset, + "world_size": meta.alloc_mode.gen_world_size + 1, + "backend": "nccl", + "group_name": meta.nccl_group_name, + } + res = await arequest_with_retry( + addr=addr, + endpoint="/init_weights_update_group", + payload=payload, + method="POST", + max_retries=1, + timeout=request_timeout, + ) + assert res["success"] diff --git a/arealite/launcher/local.py b/arealite/launcher/local.py index 447e7db04b..27cc0ba2a3 100644 --- a/arealite/launcher/local.py +++ b/arealite/launcher/local.py @@ -283,7 +283,7 @@ def main_local(): 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:])}", + cmd=f"torchrun --nnodes 1 --nproc-per-node {alloc_mode.train_world_size} --master-addr localhost --master-port {find_free_ports(1, (10000, 50000))[0]} {' '.join(sys.argv[1:])}", gpu=alloc_mode.train_world_size, env_vars=dict(AREAL_LLM_SERVER_ADDRS=",".join(server_addrs)), ) diff --git a/arealite/launcher/sglang_server.py b/arealite/launcher/sglang_server.py index 2e88c043ab..938ce5c744 100644 --- a/arealite/launcher/sglang_server.py +++ b/arealite/launcher/sglang_server.py @@ -3,10 +3,8 @@ import sys import time import uuid -from pathlib import Path from typing import Optional -import ray import requests from arealite.api.cli_args import ( @@ -19,12 +17,12 @@ 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 +from realhf.base import logging, name_resolve, names logger = logging.getLogger("SGLangServer Wrapper") -def execute_shell_command(command: str) -> subprocess.Popen: +def launch_server_cmd(command: str) -> subprocess.Popen: """ Execute a shell command and return its process handle. """ @@ -45,46 +43,6 @@ def execute_shell_command(command: str) -> subprocess.Popen: ) -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. diff --git a/arealite/tests/test_fsdp_engine_nccl.py b/arealite/tests/test_fsdp_engine_nccl.py index 7d71fd6c41..fde448a40e 100644 --- a/arealite/tests/test_fsdp_engine_nccl.py +++ b/arealite/tests/test_fsdp_engine_nccl.py @@ -12,10 +12,9 @@ SGLangConfig, TrainEngineConfig, ) -from arealite.api.io_struct import FinetuneSpec, WeightUpdateMeta +from arealite.api.io_struct import AllocationMode, FinetuneSpec, WeightUpdateMeta from arealite.engine.fsdp_engine import FSDPEngine from arealite.engine.sglang_remote import RemoteSGLangEngine -from arealite.utils.network import find_free_ports from realhf.base import network EXPR_NAME = "test_fsdp_engine_nccl" @@ -33,7 +32,7 @@ def check_server_health(base_url): try: - response = requests.get(f"{base_url}/metrics", timeout=30) + response = requests.get(f"{base_url}/health", timeout=30) return response.status_code == 200 except requests.exceptions.RequestException: return False @@ -82,14 +81,14 @@ def sglang_server_nccl(): def test_fsdpengine_nccl_weight_update_to_remote(tmp_path_factory, sglang_server_nccl): - # 设置分布式环境变量 + # Set environment variables for torch distributed os.environ["WORLD_SIZE"] = "1" os.environ["RANK"] = "0" os.environ["LOCAL_RANK"] = "0" os.environ["MASTER_ADDR"] = HOST os.environ["MASTER_PORT"] = str(MASTER_PORT) - # 启动本地FSDPEngine + # Initialize FSDPEngine engine_config = TrainEngineConfig( experiment_name=EXPR_NAME, trial_name=TRIAL_NAME, @@ -99,39 +98,31 @@ def test_fsdpengine_nccl_weight_update_to_remote(tmp_path_factory, sglang_server engine = FSDPEngine(engine_config) ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=100, train_batch_size=2) engine.initialize(None, ft_spec) + engine.model_version = 123 - # 启动远端RemoteSGLangEngine + # Initialize RemoteSGLangEngine config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) config.server_addrs = [f"{HOST}:{PORT}"] remote_engine = RemoteSGLangEngine(config) remote_engine.initialize(None, None) - # 构造WeightUpdateMeta（type=nccl） - param_meta = engine.get_param_meta_for_distributed_update() - meta = WeightUpdateMeta( - type="nccl", - path=None, - alloc_mode=None, - comm_backend="nccl", - model_version=123, - tp_size=1, - master_address="localhost", - master_port=find_free_ports(1)[0], - world_size=2, - group_name=GROUP_NAME, - parameter_names=list(param_meta.keys()), - state_dict_key_to_shape=param_meta, + # Get WeightUpdateMeta + meta = WeightUpdateMeta.from_fsdp_nccl( + AllocationMode.from_str("sglang.d1p1t1+d1p1t1"), + engine, + nccl_group_name=GROUP_NAME, ) - # 本地engine广播参数 - future = remote_engine.update_weights(meta) + # Broadcast weights + future = remote_engine.update_weights(meta, engine.model_version) print("got future", flush=True) engine.upload_weights(meta) print("uploaded wexights to remote engine", flush=True) - # 远端engine拉取参数 + # Wait for remote engine to finish future.result(timeout=120) print("got result", flush=True) - # 检查远端参数版本 + # Check model_version in meta and remote engine + remote_engine.set_version(engine.get_version()) assert remote_engine.get_version() == 123 remote_engine.destroy() engine.destroy() diff --git a/arealite/tests/test_sglang_engine.py b/arealite/tests/test_sglang_engine.py index 71a6f2f5dd..73ce724c97 100644 --- a/arealite/tests/test_sglang_engine.py +++ b/arealite/tests/test_sglang_engine.py @@ -31,10 +31,7 @@ def check_server_health(base_url): try: - response = requests.get( - f"{base_url}/metrics", - timeout=30, - ) + response = requests.get(f"{base_url}/health", timeout=30) return response.status_code == 200 except requests.exceptions.RequestException as e: return False @@ -85,6 +82,7 @@ async def test_remote_sglang_generate(sglang_server): tokenizer = load_hf_tokenizer(MODEL_PATH) os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" engine = RemoteSGLangEngine(config) + engine.initialize(None, None) req = LLMRequest( rid=str(uuid.uuid4()), input_ids=tokenizer.encode("hello! how are you today"), @@ -211,6 +209,7 @@ def test_disk_update_weights_from_fsdp_engine(tmp_path_factory, sglang_server): engine = FSDPEngine(engine_config) ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=100, train_batch_size=2) engine.initialize(None, ft_spec) + engine.model_version = 100 # setup name resolve import realhf.base.name_resolve as name_resolve @@ -229,10 +228,8 @@ def test_disk_update_weights_from_fsdp_engine(tmp_path_factory, sglang_server): inf_engine.initialize(None, None) # test update weights path = tmp_path_factory.mktemp("upload_weights_from_disk") - update_weight_meta = WeightUpdateMeta( - type="disk", path=path, alloc_mode=None, comm_backend=None, model_version=100 - ) - future = inf_engine.update_weights(update_weight_meta) + update_weight_meta = WeightUpdateMeta.from_disk(path=path) + future = inf_engine.update_weights(update_weight_meta, engine.get_version()) engine.upload_weights(update_weight_meta) future.result() assert inf_engine.get_version() == 100 diff --git a/examples/arealite/boba.py b/examples/arealite/boba.py index 4211d11972..d975a7b09c 100644 --- a/examples/arealite/boba.py +++ b/examples/arealite/boba.py @@ -1,6 +1,5 @@ import asyncio import os -import shutil import sys import uuid @@ -18,7 +17,12 @@ GRPOConfig, load_expr_config, ) -from arealite.api.io_struct import FinetuneSpec, LLMRequest, WeightUpdateMeta +from arealite.api.io_struct import ( + AllocationMode, + FinetuneSpec, + LLMRequest, + WeightUpdateMeta, +) from arealite.api.workflow_api import RolloutWorkflow from arealite.engine.ppo.actor import FSDPPPOActor from arealite.engine.sglang_remote import RemoteSGLangEngine @@ -154,7 +158,6 @@ def boba_reward_fn( 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: @@ -204,6 +207,18 @@ def main(args): ref = FSDPPPOActor(config=config.ref) ref.initialize(None, ft_spec) + # NOTE: Weight update meta only requires address and free port of rank 0, + # but `WeightUpdateMeta.from_fsdp_nccl` has to be executed on all ranks + # due to `engine.get_param_specs()`. + # Therefore, we create weight update meta on all ranks, then broadcast the one on rank 0. + weight_update_meta = [ + WeightUpdateMeta.from_fsdp_nccl( + AllocationMode.from_str(config.allocation_mode), actor + ) + ] + dist.broadcast_object_list(weight_update_meta, src=0) + weight_update_meta = weight_update_meta[0] + # 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) @@ -245,7 +260,7 @@ def main(args): batch = batch.to(actor.device) # Create barrier to synchronize all rollout processes. - dist.barrier() + dist.barrier(device_ids=[actor.device.index]) torch.cuda.synchronize() if config.actor.recompute_logprob or config.actor.use_decoupled_loss: @@ -272,26 +287,16 @@ 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=path, - alloc_mode=None, - comm_backend=None, - model_version=global_step + 1, - ) + rollout.pause() if dist.get_rank() == 0: - future = rollout.update_weights(meta) - actor.upload_weights(meta) + future = rollout.update_weights(weight_update_meta) + actor.upload_weights(weight_update_meta) if dist.get_rank() == 0: future.result() - shutil.rmtree(path, ignore_errors=True) - dist.barrier() + dist.barrier(device_ids=[actor.device.index]) torch.cuda.synchronize() + rollout.resume() + actor.set_version(global_step + 1) rollout.set_version(global_step + 1) with stats_tracker.record_timing("save"): diff --git a/examples/arealite/configs/boba.yaml b/examples/arealite/configs/boba.yaml index 3924a14cb9..3b9c5b27e1 100644 --- a/examples/arealite/configs/boba.yaml +++ b/examples/arealite/configs/boba.yaml @@ -87,7 +87,7 @@ sglang: random_seed: ${seed} skip_tokenizer_init: true dtype: ${actor.dtype} - max_running_requests: null + max_running_requests: null context_length: 32768 mem_fraction_static: 0.9 @@ -128,7 +128,7 @@ stats_logger: fileroot: ${cluster.fileroot} wandb: mode: online - + # Launcher launcher: inference_server_cpus_per_gpu: 15 @@ -138,4 +138,4 @@ launcher: 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 + inference_server_image: /storage/openpsi/images/sglang-v0.4.9.post2-cu126-v2.sif diff --git a/examples/arealite/configs/gsm8k_grpo.yaml b/examples/arealite/configs/gsm8k_grpo.yaml index 9f0107e6f8..18b5fedaf9 100644 --- a/examples/arealite/configs/gsm8k_grpo.yaml +++ b/examples/arealite/configs/gsm8k_grpo.yaml @@ -1,9 +1,9 @@ experiment_name: gsm8k-grpo trial_name: trial0 allocation_mode: sglang.d4p1t1+d4p1t1 -n_nodes: 1 -n_gpus_per_node: 8 cluster: + n_nodes: 1 + n_gpus_per_node: 8 fileroot: /tmp/arealite/experiments name_resolve: type: nfs diff --git a/examples/arealite/gsm8k_grpo.py b/examples/arealite/gsm8k_grpo.py index 6fbccb762e..f44c0d35d0 100644 --- a/examples/arealite/gsm8k_grpo.py +++ b/examples/arealite/gsm8k_grpo.py @@ -1,5 +1,4 @@ import os -import shutil import sys import torch @@ -9,7 +8,7 @@ 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.api.io_struct import AllocationMode, FinetuneSpec, WeightUpdateMeta from arealite.engine.ppo.actor import FSDPPPOActor from arealite.engine.sglang_remote import RemoteSGLangEngine from arealite.utils.device import log_gpu_stats @@ -93,6 +92,18 @@ def main(args): ref = FSDPPPOActor(config=config.ref) ref.initialize(None, ft_spec) + # NOTE: Weight update meta only requires address and free port of rank 0, + # but `WeightUpdateMeta.from_fsdp_nccl` has to be executed on all ranks + # due to `engine.get_param_specs()`. + # Therefore, we create weight update meta on all ranks, then broadcast the one on rank 0. + weight_update_meta = [ + WeightUpdateMeta.from_fsdp_nccl( + AllocationMode.from_str(config.allocation_mode), actor + ) + ] + dist.broadcast_object_list(weight_update_meta, src=0) + weight_update_meta = weight_update_meta[0] + # 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) @@ -136,7 +147,7 @@ def main(args): batch = batch.to(actor.device) # Create barrier to synchronize all rollout processes. - dist.barrier() + dist.barrier(device_ids=[actor.device.index]) torch.cuda.synchronize() if config.actor.recompute_logprob or config.actor.use_decoupled_loss: @@ -163,26 +174,16 @@ 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=path, - alloc_mode=None, - comm_backend=None, - model_version=global_step + 1, - ) + rollout.pause() if dist.get_rank() == 0: - future = rollout.update_weights(meta) - actor.upload_weights(meta) + future = rollout.update_weights(weight_update_meta) + actor.upload_weights(weight_update_meta) if dist.get_rank() == 0: future.result() - shutil.rmtree(path, ignore_errors=True) - dist.barrier() + dist.barrier(device_ids=[actor.device.index]) torch.cuda.synchronize() + rollout.resume() + actor.set_version(global_step + 1) rollout.set_version(global_step + 1) with stats_tracker.record_timing("save"): diff --git a/examples/env/scripts/setup-container-deps.sh b/examples/env/scripts/setup-container-deps.sh deleted file mode 100644 index 2b07e31601..0000000000 --- a/examples/env/scripts/setup-container-deps.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/sh -AREAL_PATH=$PWD -cd /sglang -git apply $AREAL_PATH/patch/sglang/v0.4.6.post4.patch -cd $AREAL_PATH - -# Package used for calculating math reward -pip install -e evaluation/latex2sympy - -# Install AReaL -pip install -e . \ No newline at end of file diff --git a/examples/env/scripts/setup-eval-pip-deps.sh b/examples/env/scripts/setup-eval-pip-deps.sh index fbdf4ee4f5..659ed73460 100644 --- a/examples/env/scripts/setup-eval-pip-deps.sh +++ b/examples/env/scripts/setup-eval-pip-deps.sh @@ -1,9 +1,9 @@ #/bin/bash # basic dependencies pip install -U pip -pip uninstall deepspeed flash-attn pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y -pip install nvidia-ml-py +pip uninstall pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y +pip install pynvml nvidia-ml-py pip install -e evaluation/latex2sympy pip install vllm==0.8.5 --no-build-isolation -pip install flash_attn --no-build-isolation +pip install "flash-attn<=2.7.3" --no-build-isolation pip install -r evaluation/requirements.txt \ No newline at end of file diff --git a/examples/env/scripts/setup-pip-deps.sh b/examples/env/scripts/setup-pip-deps.sh index 8fa203bf97..1599170a38 100644 --- a/examples/env/scripts/setup-pip-deps.sh +++ b/examples/env/scripts/setup-pip-deps.sh @@ -1,24 +1,14 @@ #!/bin/bash # basic dependencies pip install -U pip -pip uninstall torch deepspeed flash-attn pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y -pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 -pip install "sglang[all]==0.4.6.post4" +pip uninstall pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y +pip install torch==2.7.1 torchaudio==2.7.1 torchvision==0.22.1 "deepspeed>=0.17.2" pynvml +pip install "sglang[all]==0.4.9.post2" pip install megatron-core==0.11.0 nvidia-ml-py pip install git+https://github.com/garrett4wade/cugae --no-build-isolation --verbose pip install "flash-attn<=2.7.3" --no-build-isolation # Package used for calculating math reward pip install -e evaluation/latex2sympy - -# Install an editable sglang -rm -rf ./sglang -git clone -b v0.4.6.post4 https://github.com/sgl-project/sglang -AREAL_PATH=$PWD -cd sglang -git apply ../patch/sglang/v0.4.6.post4.patch -pip install -e "python[all]" --no-deps -cd $AREAL_PATH - # Install AReaL pip install -e . diff --git a/patch/sglang/v0.4.6.post2.patch b/patch/sglang/v0.4.6.post2.patch deleted file mode 100644 index 6bf47bf7b4..0000000000 --- a/patch/sglang/v0.4.6.post2.patch +++ /dev/null @@ -1,144 +0,0 @@ -diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 174656b2..33fe0a5f 100644 ---- a/python/sglang/srt/managers/io_struct.py -+++ b/python/sglang/srt/managers/io_struct.py -@@ -687,10 +687,21 @@ class FlushCacheReqOutput: - success: bool - - -+@dataclass -+class InterruptAllReqInput: -+ pass -+ -+ -+@dataclass -+class InterruptAllReqOutput: -+ num_interrupted_requests: int -+ -+ - @dataclass - class UpdateWeightFromDiskReqInput: - # The model path with the new weights - model_path: str -+ allow_interrupt: bool = False - # The format to load the weights - load_format: Optional[str] = None - -diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 8891115c..843a8a82 100644 ---- a/python/sglang/srt/managers/scheduler.py -+++ b/python/sglang/srt/managers/scheduler.py -@@ -70,6 +70,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -419,6 +421,7 @@ class Scheduler( - # Init request dispatcher - self._request_dispatcher = TypeBasedDispatcher( - [ -+ (InterruptAllReqInput, self.interrupt_all_requests), - (TokenizedGenerateReqInput, self.handle_generate_request), - (TokenizedEmbeddingReqInput, self.handle_embedding_request), - (FlushCacheReqInput, self.flush_cache_wrapped), -@@ -1938,6 +1941,15 @@ class Scheduler( - def _pause_engine(self) -> Tuple[List[Req], int]: - raise NotImplementedError() - -+ def interrupt_all_requests(self, recv_req: InterruptAllReqInput): -+ num = len(self.waiting_queue) + len(self.running_batch.reqs) -+ for req in self.waiting_queue: -+ req.sampling_params.max_new_tokens = 0 -+ for req in self.running_batch.reqs: -+ req.sampling_params.max_new_tokens = len(req.output_ids) -+ logger.info(f"Interrupt {num} requests.") -+ return InterruptAllReqOutput(num) -+ - def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): - """In-place update of the weights from disk.""" - success, message = self.tp_worker.update_weights_from_disk(recv_req) -diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 82709b09..bfab3ce7 100644 ---- a/python/sglang/srt/managers/tokenizer_manager.py -+++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -76,6 +76,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -265,6 +267,9 @@ class TokenizerManager: - self.resume_memory_occupation_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -+ self.interrupt_requests_communicator = _Communicator( -+ self.send_to_scheduler, server_args.dp_size -+ ) - self.flush_cache_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -@@ -294,6 +299,10 @@ class TokenizerManager: - UpdateWeightFromDiskReqOutput, - self._handle_update_weights_from_disk_req_output, - ), -+ ( -+ InterruptAllReqOutput, -+ self.interrupt_requests_communicator.handle_recv, -+ ), - ( - InitWeightsUpdateGroupReqOutput, - self.init_weights_update_group_communicator.handle_recv, -@@ -767,6 +776,13 @@ class TokenizerManager: - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() - -+ if obj.allow_interrupt: -+ num_interrupted_requests = await self.interrupt_all_requests( -+ InterruptAllReqInput() -+ ) -+ # Set a break point to wait for the interrupt to finish -+ await asyncio.sleep(0.1) -+ - # default the load format to the server_args - if obj.load_format is None: - obj.load_format = self.server_args.load_format -@@ -776,7 +792,12 @@ class TokenizerManager: - # Hold the lock if it is not async. This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: -- return await self._wait_for_model_update_from_disk(obj) -+ success, message, n_paused = ( -+ await self._wait_for_model_update_from_disk(obj) -+ ) -+ if obj.allow_interrupt: -+ return success, message, num_interrupted_requests -+ return success, message, n_paused - - async def _wait_for_model_update_from_disk( - self, obj: UpdateWeightFromDiskReqInput -@@ -849,6 +870,18 @@ class TokenizerManager: - result = (await self.update_weights_from_tensor_communicator(obj))[0] - return result.success, result.message - -+ async def interrupt_all_requests( -+ self, -+ obj: InterruptAllReqInput, -+ request: Optional[fastapi.Request] = None, -+ ) -> Tuple[bool, str]: -+ self.auto_create_handle_loop() -+ result = await self.interrupt_requests_communicator(obj) -+ if self.server_args.dp_size == 1: -+ return result[0].num_interrupted_requests -+ else: -+ return [r.num_interrupted_requests for r in result] -+ - async def get_weights_by_name( - self, obj: GetWeightsByNameReqInput, request: Optional[fastapi.Request] = None - ): diff --git a/patch/sglang/v0.4.6.post4.patch b/patch/sglang/v0.4.6.post4.patch deleted file mode 100644 index b7dbd09cf7..0000000000 --- a/patch/sglang/v0.4.6.post4.patch +++ /dev/null @@ -1,144 +0,0 @@ -diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 5390668c..db370d19 100644 ---- a/python/sglang/srt/managers/io_struct.py -+++ b/python/sglang/srt/managers/io_struct.py -@@ -687,10 +687,21 @@ class FlushCacheReqOutput: - success: bool - - -+@dataclass -+class InterruptAllReqInput: -+ pass -+ -+ -+@dataclass -+class InterruptAllReqOutput: -+ num_interrupted_requests: int -+ -+ - @dataclass - class UpdateWeightFromDiskReqInput: - # The model path with the new weights - model_path: str -+ allow_interrupt: bool = False - # The format to load the weights - load_format: Optional[str] = None - -diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 1178eec5..318dee33 100644 ---- a/python/sglang/srt/managers/scheduler.py -+++ b/python/sglang/srt/managers/scheduler.py -@@ -73,6 +73,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -427,6 +429,7 @@ class Scheduler( - # Init request dispatcher - self._request_dispatcher = TypeBasedDispatcher( - [ -+ (InterruptAllReqInput, self.interrupt_all_requests), - (TokenizedGenerateReqInput, self.handle_generate_request), - (TokenizedEmbeddingReqInput, self.handle_embedding_request), - (FlushCacheReqInput, self.flush_cache_wrapped), -@@ -1971,6 +1974,15 @@ class Scheduler( - def _pause_engine(self) -> Tuple[List[Req], int]: - raise NotImplementedError() - -+ def interrupt_all_requests(self, recv_req: InterruptAllReqInput): -+ num = len(self.waiting_queue) + len(self.running_batch.reqs) -+ for req in self.waiting_queue: -+ req.sampling_params.max_new_tokens = 0 -+ for req in self.running_batch.reqs: -+ req.sampling_params.max_new_tokens = len(req.output_ids) -+ logger.info(f"Interrupt {num} requests.") -+ return InterruptAllReqOutput(num) -+ - def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): - """In-place update of the weights from disk.""" - success, message = self.tp_worker.update_weights_from_disk(recv_req) -diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index b646fae1..c668728b 100644 ---- a/python/sglang/srt/managers/tokenizer_manager.py -+++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -80,6 +80,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -279,6 +281,9 @@ class TokenizerManager: - self.slow_down_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -+ self.interrupt_requests_communicator = _Communicator( -+ self.send_to_scheduler, server_args.dp_size -+ ) - self.flush_cache_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -@@ -309,6 +314,10 @@ class TokenizerManager: - UpdateWeightFromDiskReqOutput, - self._handle_update_weights_from_disk_req_output, - ), -+ ( -+ InterruptAllReqOutput, -+ self.interrupt_requests_communicator.handle_recv, -+ ), - ( - InitWeightsUpdateGroupReqOutput, - self.init_weights_update_group_communicator.handle_recv, -@@ -799,6 +808,13 @@ class TokenizerManager: - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() - -+ if obj.allow_interrupt: -+ num_interrupted_requests = await self.interrupt_all_requests( -+ InterruptAllReqInput() -+ ) -+ # Set a break point to wait for the interrupt to finish -+ await asyncio.sleep(0.1) -+ - # default the load format to the server_args - if obj.load_format is None: - obj.load_format = self.server_args.load_format -@@ -808,7 +824,12 @@ class TokenizerManager: - # Hold the lock if it is not async. This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: -- return await self._wait_for_model_update_from_disk(obj) -+ success, message, n_paused = ( -+ await self._wait_for_model_update_from_disk(obj) -+ ) -+ if obj.allow_interrupt: -+ return success, message, num_interrupted_requests -+ return success, message, n_paused - - async def _wait_for_model_update_from_disk( - self, obj: UpdateWeightFromDiskReqInput -@@ -881,6 +902,18 @@ class TokenizerManager: - result = (await self.update_weights_from_tensor_communicator(obj))[0] - return result.success, result.message - -+ async def interrupt_all_requests( -+ self, -+ obj: InterruptAllReqInput, -+ request: Optional[fastapi.Request] = None, -+ ) -> Tuple[bool, str]: -+ self.auto_create_handle_loop() -+ result = await self.interrupt_requests_communicator(obj) -+ if self.server_args.dp_size == 1: -+ return result[0].num_interrupted_requests -+ else: -+ return [r.num_interrupted_requests for r in result] -+ - async def get_weights_by_name( - self, obj: GetWeightsByNameReqInput, request: Optional[fastapi.Request] = None - ): diff --git a/pyproject.toml b/pyproject.toml index cd4e26dae3..c779fa77c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,10 +31,9 @@ dependencies = [ "huggingface_hub", "datasets", "accelerate", - "transformers==4.51.1", + "transformers==4.53.0", # Scientific computing - "numpy<2.0.0", "scipy", "pandas", "matplotlib", diff --git a/realhf/system/generation_server.py b/realhf/system/generation_server.py index 2a46141a1d..5d09707b36 100644 --- a/realhf/system/generation_server.py +++ b/realhf/system/generation_server.py @@ -40,42 +40,11 @@ def execute_shell_command(command: str) -> subprocess.Popen: ) -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, port: int = 30000): """ 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() assert port is not None full_command = f"{command} --port {port}" process = execute_shell_command(full_command) diff --git a/realhf/system/gserver_manager.py b/realhf/system/gserver_manager.py index e746fe976b..a789ae7fe0 100644 --- a/realhf/system/gserver_manager.py +++ b/realhf/system/gserver_manager.py @@ -155,39 +155,22 @@ def check_new_params(self) -> str | None: return None - async def flush_requests_and_update_weights( - self, server_url, new_param_path, update_weights_retries=5 - ): - server_index = self.server_urls.index(server_url) - success = False - for _ in range(update_weights_retries): - async with aiohttp.ClientSession( - server_url, - timeout=aiohttp.ClientTimeout( - total=self.config.flush_request_timeout, - sock_connect=self.config.flush_request_timeout, - ), - ) as session: - async with session.post( - f"/update_weights_from_disk", - json=dict(model_path=new_param_path, allow_interrupt=True), - ) as resp: - if resp.status == 200: - res = await resp.json() - success = res["success"] - if success: - if "num_paused_requests" in res: - logger.info( - f"{res['num_paused_requests']} requests are interrupted " - f"during updating weights for server {server_index}: {server_url}" - ) - return - logger.warning( - f"Update weights failed: {res['message']}. Retrying." - ) - logger.warning(f"Update weights failed: {resp.reason}. Retrying.") - time.sleep(0.1) - raise RuntimeError("Update weights failed.") + async def flush_requests_and_update_weights(self, server_url, new_param_path): + async with aiohttp.ClientSession( + server_url, + timeout=aiohttp.ClientTimeout( + total=self.config.flush_request_timeout, + sock_connect=self.config.flush_request_timeout, + ), + ) as session: + (await session.post("/pause_generation")).raise_for_status() + async with session.post( + f"/update_weights_from_disk", + json=dict(model_path=new_param_path), + ) as resp: + resp.raise_for_status() + assert (await resp.json())["success"] + (await session.post("/continue_generation")).raise_for_status() def _round_robin_schedule(self, req_meta: GenReqMeta) -> int: if not hasattr(self, "round_robin_idx"): diff --git a/requirements.txt b/requirements.txt index 178e958355..0c21f9b0a6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,7 +25,6 @@ numba packaging pandas pybind11>=2.10.0 -numpy<2.0.0 psutil pynvml pytest From 06d93704f433c3e2ec9ffcb365597f088dbd1677 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Thu, 24 Jul 2025 14:26:01 +0800 Subject: [PATCH 12/20] PullRequest: 422 [lite] Fix tests and scripts after updating sgl to 0.4.9 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/sgl049 of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/422 Reviewed-by: 晓雷 * . * bump arealite to sglang 0.4.9.post2 * . * PullRequest: 412 [lite] Minor refactor on `UpdateWeightMeta` * . --- arealite/api/io_struct.py | 3 ++- arealite/tests/test_fsdp_engine_nccl.py | 6 +---- arealite/tests/test_sglang_engine.py | 33 +++---------------------- examples/env/validate_installation.py | 9 ++++--- 4 files changed, 12 insertions(+), 39 deletions(-) diff --git a/arealite/api/io_struct.py b/arealite/api/io_struct.py index ceb59535fc..136c15fc38 100644 --- a/arealite/api/io_struct.py +++ b/arealite/api/io_struct.py @@ -2,8 +2,8 @@ # Licensed under the Apache License, Version 2.0 import enum import itertools -import re import os +import re import uuid from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple @@ -184,6 +184,7 @@ def from_disk( saver_config: SaverConfig, ) -> "WeightUpdateMeta": from arealite.utils.saver import Saver + path = os.path.join( Saver.get_save_checkpoint_root(saver_config), "weight_update", diff --git a/arealite/tests/test_fsdp_engine_nccl.py b/arealite/tests/test_fsdp_engine_nccl.py index fde448a40e..e06b94bf8b 100644 --- a/arealite/tests/test_fsdp_engine_nccl.py +++ b/arealite/tests/test_fsdp_engine_nccl.py @@ -98,7 +98,6 @@ def test_fsdpengine_nccl_weight_update_to_remote(tmp_path_factory, sglang_server engine = FSDPEngine(engine_config) ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=100, train_batch_size=2) engine.initialize(None, ft_spec) - engine.model_version = 123 # Initialize RemoteSGLangEngine config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) @@ -114,15 +113,12 @@ def test_fsdpengine_nccl_weight_update_to_remote(tmp_path_factory, sglang_server ) # Broadcast weights - future = remote_engine.update_weights(meta, engine.model_version) + future = remote_engine.update_weights(meta) print("got future", flush=True) engine.upload_weights(meta) print("uploaded wexights to remote engine", flush=True) # Wait for remote engine to finish future.result(timeout=120) print("got result", flush=True) - # Check model_version in meta and remote engine - remote_engine.set_version(engine.get_version()) - assert remote_engine.get_version() == 123 remote_engine.destroy() engine.destroy() diff --git a/arealite/tests/test_sglang_engine.py b/arealite/tests/test_sglang_engine.py index 73ce724c97..d660315342 100644 --- a/arealite/tests/test_sglang_engine.py +++ b/arealite/tests/test_sglang_engine.py @@ -2,7 +2,6 @@ import subprocess import sys import time -import uuid import pytest import requests @@ -14,7 +13,7 @@ InferenceEngineConfig, SGLangConfig, ) -from arealite.api.io_struct import LLMRequest, LLMResponse, WeightUpdateMeta +from arealite.api.io_struct import WeightUpdateMeta from arealite.utils import network from realhf.api.core.data_api import load_hf_tokenizer @@ -74,30 +73,6 @@ def sglang_server(): process.terminate() -@pytest.mark.asyncio -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) - tokenizer = load_hf_tokenizer(MODEL_PATH) - os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" - engine = RemoteSGLangEngine(config) - engine.initialize(None, None) - req = LLMRequest( - rid=str(uuid.uuid4()), - input_ids=tokenizer.encode("hello! how are you today"), - gconfig=GenerationHyperparameters(max_new_tokens=16), - ) - resp = await engine.agenerate(req) - assert isinstance(resp, LLMResponse) - assert resp.input_tokens == req.input_ids - assert ( - len(resp.output_logprobs) - == len(resp.output_tokens) - == len(resp.output_versions) - ) - - @pytest.mark.parametrize("n_samples", [1, 2, 4]) def test_remote_sglang_rollout(sglang_server, n_samples): from arealite.engine.sglang_remote import RemoteSGLangEngine @@ -226,11 +201,11 @@ def test_disk_update_weights_from_fsdp_engine(tmp_path_factory, sglang_server): os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" inf_engine = RemoteSGLangEngine(config) inf_engine.initialize(None, None) + inf_engine.set_version(100) # test update weights path = tmp_path_factory.mktemp("upload_weights_from_disk") - update_weight_meta = WeightUpdateMeta.from_disk(path=path) - future = inf_engine.update_weights(update_weight_meta, engine.get_version()) + update_weight_meta = WeightUpdateMeta(type="disk", path=str(path)) + future = inf_engine.update_weights(update_weight_meta) engine.upload_weights(update_weight_meta) future.result() - assert inf_engine.get_version() == 100 inf_engine.destroy() diff --git a/examples/env/validate_installation.py b/examples/env/validate_installation.py index 61ef6f7c5d..67ae740cdf 100644 --- a/examples/env/validate_installation.py +++ b/examples/env/validate_installation.py @@ -67,6 +67,7 @@ def test_torch_cuda(self, torch_module): def test_flash_attn_functionality(self, flash_attn_module): """Test flash attention functionality.""" # Try to import key functions + import flash_attn_2_cuda from flash_attn import flash_attn_func, flash_attn_varlen_func print(" - Flash attention functions imported successfully") @@ -79,12 +80,12 @@ def test_sglang_functionality(self, sglang_module): """Test SGLang basic functionality.""" # Basic import test is sufficient for CI import sgl_kernel - from sglang import launch_server - assert Version(get_version("sglang")) == Version("0.4.6.post4") + from sglang import Engine, launch_server + assert Version(get_version("sglang")) == Version("0.4.9.post2"), "SGLang version should be v0.4.9.post2" print(" - SGLang imported successfully") def test_transformers(self, transformers_module): - assert Version(get_version("transformers")) == Version("4.51.1") + assert Version(get_version("transformers")) == Version("4.53.1"), "transformers version should be 4.53.1" print(" - transformers imported successfully") def validate_critical_dependencies(self): @@ -140,7 +141,7 @@ def validate_optional_dependencies(self): self.test_import("flashattn_hopper", required=False) # Optional utilities - self.test_import("tensorboardx", required=False) + self.test_import("tensorboardX", required=False) self.test_import("swanlab", required=False) self.test_import("matplotlib", required=False) self.test_import("seaborn", required=False) From d6dbd9013a4255f36b566f51eb791277c2afed1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Thu, 24 Jul 2025 14:27:18 +0800 Subject: [PATCH 13/20] PullRequest: 423 [lite] Remove the boba example for github release. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/remove-boba of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/423 Reviewed-by: 晓雷 * . * . --- examples/arealite/boba.py | 315 ---------------------- examples/arealite/configs/boba.yaml | 141 ---------- examples/arealite/configs/gsm8k_grpo.yaml | 2 +- examples/arealite/configs/gsm8k_sft.yaml | 2 +- 4 files changed, 2 insertions(+), 458 deletions(-) delete mode 100644 examples/arealite/boba.py delete mode 100644 examples/arealite/configs/boba.yaml diff --git a/examples/arealite/boba.py b/examples/arealite/boba.py deleted file mode 100644 index d975a7b09c..0000000000 --- a/examples/arealite/boba.py +++ /dev/null @@ -1,315 +0,0 @@ -import asyncio -import os -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 ( - AllocationMode, - 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: - 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) - - # NOTE: Weight update meta only requires address and free port of rank 0, - # but `WeightUpdateMeta.from_fsdp_nccl` has to be executed on all ranks - # due to `engine.get_param_specs()`. - # Therefore, we create weight update meta on all ranks, then broadcast the one on rank 0. - weight_update_meta = [ - WeightUpdateMeta.from_fsdp_nccl( - AllocationMode.from_str(config.allocation_mode), actor - ) - ] - dist.broadcast_object_list(weight_update_meta, src=0) - weight_update_meta = weight_update_meta[0] - - # 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(device_ids=[actor.device.index]) - 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"): - rollout.pause() - if dist.get_rank() == 0: - future = rollout.update_weights(weight_update_meta) - actor.upload_weights(weight_update_meta) - if dist.get_rank() == 0: - future.result() - dist.barrier(device_ids=[actor.device.index]) - torch.cuda.synchronize() - rollout.resume() - actor.set_version(global_step + 1) - 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 deleted file mode 100644 index 3b9c5b27e1..0000000000 --- a/examples/arealite/configs/boba.yaml +++ /dev/null @@ -1,141 +0,0 @@ -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/sglang-v0.4.9.post2-cu126-v2.sif diff --git a/examples/arealite/configs/gsm8k_grpo.yaml b/examples/arealite/configs/gsm8k_grpo.yaml index 18b5fedaf9..9b0362d20e 100644 --- a/examples/arealite/configs/gsm8k_grpo.yaml +++ b/examples/arealite/configs/gsm8k_grpo.yaml @@ -32,7 +32,7 @@ gconfig: actor: experiment_name: ${experiment_name} trial_name: ${trial_name} - path: /storage/openpsi/models/Qwen__Qwen2-1.5B-Instruct/ + path: Qwen/Qwen2-1.5B-Instruct init_from_scratch: false disable_dropout: true gradient_checkpointing: false diff --git a/examples/arealite/configs/gsm8k_sft.yaml b/examples/arealite/configs/gsm8k_sft.yaml index c8aaa4bb09..8d05abe89e 100644 --- a/examples/arealite/configs/gsm8k_sft.yaml +++ b/examples/arealite/configs/gsm8k_sft.yaml @@ -13,7 +13,7 @@ tokenizer_path: ${model.path} model: experiment_name: ${experiment_name} trial_name: ${trial_name} - path: /storage/openpsi/models/Qwen__Qwen3-1.7B/ + path: Qwen/Qwen3-1.7B init_from_scratch: false gradient_checkpointing: false dtype: bfloat16 From 209aec6aa56ffd27e943edb2af2bd38512c0d012 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Thu, 24 Jul 2025 18:57:55 +0800 Subject: [PATCH 14/20] update readme --- arealite/README.md | 1057 +++++++++++++++++--------------------------- 1 file changed, 404 insertions(+), 653 deletions(-) diff --git a/arealite/README.md b/arealite/README.md index 4a23ce60a4..1db957f36f 100644 --- a/arealite/README.md +++ b/arealite/README.md @@ -1,52 +1,173 @@ -# AReaL v1.0.0 Design Doc +# AReaLite Design Document + +## Motivation + +AReaL presents several challenges that make it difficult for AI researchers to adopt, +understand, and develop with effectively. The primary issue stems from its +*system-centric* rather than *AI-centric* architecture. The reinforcement learning +algorithm workflow is built around multiple *workers* executing consecutive *model +function calls* — concepts that are unfamiliar to most AI researchers. This forces users +to first master these system-level abstractions before they can implement workflows and +algorithms for their specific research needs. + +Beyond architectural concerns, AReaL suffers from accumulated technical debt. The +codebase contains substantial legacy code inherited from previous projects that no +longer serves a purpose but significantly increases complexity for both users and +developers. Even experienced core developers sometimes struggle with debugging due to +this accumulated complexity. + +The landscape of RL workflow development tools has matured considerably, making it +possible to achieve comparable efficiency with significantly fewer lines of code. This +presents an ideal opportunity to redesign the API and distill the massive codebase into +something clean and maintainable. Rather than pursuing maximum efficiency, our goal is +to deliver 90% of AReaL's functionality while dramatically reducing code complexity and +user burden. This philosophy drives AReaLite — the lightweight version of AReaL. + +AReaLite serves as the first phase in AReaL's broader refactoring initiative. It +functions both as a standalone training library with intuitive interfaces and as the +foundation for AReaL's future core API definitions. The plan is to transform AReaL's +current worker-based architecture into an AI-centric architecture similar to AReaLite, +where AReaL will **extend** AReaLite's APIs and implementations to support additional +backends for efficient large-scale training. + +## Design Principles + +Our design is guided by seven core principles: + +1. **Native asynchronous RL training support** — Built from the ground up for + disentangled generation and training +1. **AI-centric design** — Minimize exposure to system concepts like "PlacementGroup" +1. **PyTorch-centric approach** — Use raw PyTorch types without unnecessary abstractions +1. **Transparent algorithm orchestration** — Make the flow of operations clear and + understandable +1. **Developer-friendly navigation** — Enable easy access to implementation details via + Ctrl+click in IDEs +1. **Ecosystem compatibility** — Integrate smoothly with existing ML/RL tools +1. **Single-file customization** — Allow RL pipeline modifications within a single file + +## Architecture + +### Core Directory Structure -______________________________________________________________________ +``` +arealite/ +├── api/ # Abstract interfaces and dataclasses +├── engine/ # Training and inference engines +├── launcher/ # Launcher for different backends +├── tests/ # Standalone test scripts +└── workflow/ # Custom RL rollout workflows +``` + +### Component Overview + +#### 1. API Layer (`api/`) + +The API layer establishes clean contracts between components through abstract interfaces +and data classes: + +- **`engine_api.py`**: Defines `TrainEngine` for SPMD-based distributed training + backends and `InferenceEngine` for streaming LLM inference +- **`workflow.py`**: Defines `RolloutWorkflow` for RL data collection within a unified + method interface +- **`cli_args.py`**: Contains configuration dataclasses for all system components + +The workflow object invokes `InferenceEngine` to complete data collection following +customized patterns, providing flexibility while maintaining consistency. + +#### 2. Backend Layer (`engine/`) + +The backend layer provides adapters for third-party libraries, ensuring they conform to +the APIs defined in `engine_api.py`. These components deliver core inference and +training capabilities: + +- **`fsdp_engine.py`**: FSDP-based training engine utilizing PyTorch FSDP2 +- **`sglang_remote.py`**: Client interface for generation with remote SGLang servers -SFT example: +#### 3. Customization Layer (`engine/ppo/`, `workflow/`) + +This layer applies backend capabilities to implement specific reinforcement learning +pipelines: + +- **`engine/ppo/actor.py`**: PPO algorithm implementation that leverages a `TrainEngine` +- **`workflow/rlvr.py`**: RLVR workflow that utilizes an `InferenceEngine` to sample + multiple responses for each prompt + +#### 4. Entry Point Layer (`examples/arealite/`) + +The entry point layer composes customization layer implementations to create complete RL +training pipelines. While we provide several reference examples, users have complete +freedom to adapt these to their specific use cases. + +Entry points can be launched using launchers from `arealite/launcher/`, similar to other +distributed launch tools like `torchrun`: ```bash -torchrun --nnodes 1 --nproc-per-node 8 examples/arealite/gsm8k_sft.py --config examples/arealite/configs/gsm8k_sft.yaml +python3 -m arealite.launcher.ray entrypoint.py --config my-config.yaml ``` -GRPO example: +### Data Flow Architecture + +*TODO: Add architectural diagram* + +## Usage Examples + +### Basic RL Training + +Users must provide a YAML configuration file, though they can override configuration +parameters for hyperparameter searches or other experimental needs: ```bash -python3 -m arealite.launcher.local examples/arealite/gsm8k_grpo.py --config examples/arealite/configs/gsm8k_grpo.yaml +# Launch with Ray launcher: 4 nodes (4 GPUs each), 3 nodes for generation, 1 node for training +python3 -m arealite.launcher.ray examples/arealite/gsm8k_grpo.py \ + --config examples/arealite/configs/gsm8k_grpo.yaml \ + experiment_name= \ + trial_name= \ + allocation_mode=sglang.d12p1t1+d4p1t1 \ + cluster.n_nodes=4 \ + cluster.n_gpus_per_node=4 + +# Launch with Slurm launcher: 16 nodes (8 GPUs each), 12 nodes for generation, 4 nodes for training +python3 -m arealite.launcher.slurm examples/arealite/gsm8k_grpo.py \ + --config examples/arealite/configs/gsm8k_grpo.yaml \ + experiment_name= \ + trial_name= \ + allocation_mode=sglang.d96p1t1+d32p1t1 \ + cluster.n_nodes=16 \ + cluster.n_gpus_per_node=8 ``` -______________________________________________________________________ +## Customization Guide -We will provide both single-controller and SPMD user interfaces. The SPMD interface will -be delivered with AReaLite, which is the paradigm most users are familiar with, just -like using `torchrun` or `deepspeed`. However, this paradigm may lack some flexibility -over global scheduling and control. To unlock the full potential with customized -distributed execution, we will also provide a single-controller mode just like using Ray ---- but our scheduler backend will not be restricted to Ray. Our code will be able to -run with any scheduler in the cluster, such as native SLURM and K8S. +For detailed customization instructions, please refer to our documentation: -However, we want the user code to stay the same for both modes. The following is a -simple usage example: +- [Agent Customization](../docs/customization/agent.md) +- [Dataset Customization](../docs/customization/dataset.md) +- [Algorithm Customization](../docs/customization/algorithm.md) -```python -def get_current_time(): - ... +## Implementation Details + +### Entry Point Design Philosophy + +We considered two primary design patterns for entry points, each with distinct +trade-offs: + +#### Single-Controller Pattern + +The most modular approach uses a single-controller pattern where only one process in the +cluster executes the main coordination logic: +```python def my_reward_fn(prompt, completion, prompt_ids, completion_ids, **kwargs): return len(completion_ids) class MyRolloutWorkflow: - def __init__(self, config: Any): - self.config = config - self.env = LocalToolingEnv() - self.env.register_tool(get_current_time) - async def arun_episode(self, engine: InferenceEngine, data: Dict[str, Any]) -> Dict[str, Tensor]: - ... message = [ {"role": "system", "message": ...}, {"role": "user", "message": ...}, ] + for _ in range(self.config.num_turns): text = tokenizer.apply_chat_template(message, tools=self.env.list_tools()) req = LLMRequest(text=text, ...) @@ -54,463 +175,335 @@ class MyRolloutWorkflow: tool_name, tool_args = parse_tool(resp) cur_time = await self.env.aexecute(tool_name, tool_args) message += [{"role": "user", "message": f"The current time is {cur_time}"}] + reward = my_reward_fn(None, None, None, req.input_ids, **data) - ... return output def main_grpo(): - dataset = load_dataset("openai/gsm8k", split="train") + config, _ = load_expr_config(args, GRPOConfig) - rollout_config, training_config = load_expr_config(sys.argv[1:]) + # Create rollout workflow + workflow = MyRolloutWorkflow() # Single-controller mode initialization scheduler = SlurmScheduler() - rollout = RolloutController( - SGLangEngine(rollout_config.engine), - rollout_config.controller, - scheduler, - ) - actor = TrainController( - MegatronGRPOActor(training_config.actor), - config.training_controller_config, - scheduler, - ) - ref = TrainController( - MegatronGRPOActor(training_config.ref), - config.training_controller_config, - scheduler, - ) - # SPMD mode initialization - # rollout = RemoteSGLangEngine(rollout_config.engine) - # actor = MegatronGRPOActor(training_config.actor) - # ref = MegatronGRPOActor(training_config.ref) + rollout = RolloutController(SGLangEngine(config.rollout), scheduler) + actor = TrainController(MegatronGRPOActor(config.actor), scheduler) - rollout.initialize() - actor.initialize() - ref.initialize() - - # Synchronous RL + # Training loop dataloader = StatefulDataloader(dataset) - for epoch in range(config.epoches): - data_generator = iter(dataloader) - for prompt in range(steps_per_epoch): - prompt = next(data_generator) + for _ in range(max_steps): + # Collect trajectories using rollout workflow + batch = rollout.rollout_batch(next(dataloader), workflow=workflow) + batch: DistributedBatch + + # Prepare training inputs + adv_batch = actor.compute_advantages_and_returns(batch) + batch['advantages'] = adv_batch['advantages'] + + # Execute PPO update + stats = actor.ppo_update(batch) + + # Update inference engine weights (non-blocking to prevent NCCL blocking) + future = rollout.update_weights(wcfg) + actor.upload_weights(wcfg) + future.result() +``` - # Update inference engine weights - future = rollout.update_weights(wcfg) - actor.upload_weights(wcfg) - future.result() +**Advantages:** - # synchronous rollout - 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), - # should_accept=lambda x: x['rewards'].mean() > 0) - - # In the single-controller mode - rollout_batch: DistributedBatch - x: TensorDict = rollout_batch.load_data() - # In the SPMD mode - # rollout_batch: TensorDict - - batch['input_ids'] = rollout_batch['input_ids'] - batch['rewards'] = rollout_batch['rewards'] - - # prepare train inputs - batch['ref_logp'] = ref.compute_logp(batch) - adv_batch = actor.compute_advantages_and_returns(batch) - batch['advantages'] = adv_batch['advantages'] - - # PPO update - stats = actor.ppo_update(batch) - print(stats) - -if __name__ == "__main__": - main_grpo() -``` +- Provides maximum flexibility for device allocation, scheduling, and data arrangement -The launch commands will be: +**Disadvantages:** -```bash -# Single-controller mode -python3 main_grpo.py --config my_config.yaml rollout.workflow.x=1 -# SPMD mode -CUDA_VISIBLE_DEVICES=0,1 nohup python3 -m sglang.launch_server \ - --seed 1 --host x.x.x.x --port 7777 --dp_size 2 > server.out 2>&1 & -CUDA_VISIBLE_DEVICES=3,4 \ - torchrun --nnodes 1 --nproc-per-node 2 \ - main_grpo.py --config my_config.yaml \ - rollout.workflow.x=1 \ - rollout.engine.addresses="\[x.x.x.x, y.y.y.y\]" +- Introduces multiple abstractions (`TrainController`, `Scheduler`, `DistributedBatch`) + that complicate the script + +#### SPMD Pattern + +Given that AI researchers are more familiar with the SPMD (Single Program, Multiple +Data) pattern used in standard model training, we also provide an entry point following +this approach. With N GPUs dedicated to training, N processes execute the following +code: + +```python +def main_grpo(): + config, _ = load_expr_config(args, GRPOConfig) + + # Create rollout workflow + workflow = MyRolloutWorkflow() + + # SPMD mode initialization + rollout = RemoteSGLangEngine(config.rollout) + actor = MegatronGRPOActor(config.actor) + + # Training loop + dataloader = StatefulDataloader(dataset) + for _ in range(max_steps): + # Data collection (only on data parallel head) + if is_dp_head: + batch = rollout.rollout_batch(next(dataloader), workflow=workflow) + batch_list = [batch] + else: + batch_list = [None] + + # Broadcast data to all processes + batch = dist.broadcast(batch_list, src=0, group=model_parallel_group)[0] + batch: TensorDict + + # Prepare training inputs + adv_batch = actor.compute_advantages_and_returns(batch) + batch['advantages'] = adv_batch['advantages'] + + # Execute PPO update + stats = actor.ppo_update(batch) + + # Update weights (coordinated across processes) + if rank == 0: + future = rollout.update_weights(wcfg) + actor.upload_weights(wcfg) + if rank == 0: + future.result() ``` -## Core API +The SPMD pattern uses only concepts familiar to AI researchers, though it requires some +control flow branching based on parallelism strategy. -- A specific algorithm must use these core components. -- Concrete implementations must follow the API definition. +### Training Engine Architecture -### TrainEngine +The training engine operates at two abstraction levels to balance flexibility with ease +of use. -TrainEngine is a thin wrapper around existing training frameworks (FSDP, Megatron), -providing a unified interface for RL algorithms for computation, parameter saving and -loading, and providing a unified weight update interface for inference engines. +#### Basic Level: Backend Adapters -```python -############################################# -@dataclass -class WeightUpdateMeta: - type: str - path: str | None - alloc_mode: AllocationMode | None - -@dataclass -class SaveLoadMeta: - path: str - weight_format: str - with_optim: bool - tokenizer: PreTrainedTokenizerFast | None - base_model_path: str | None +The foundational level provides unified interfaces for RL algorithms, handling +computation, parameter management, and weight updates for inference engines. Each RL +training experiment must use one of the implemented backends: +```python class TrainEngine(abc.ABC): - # control api - def __init__(self, para_config) - self.para_config = para_config - - def initialize(self, addr: str|None, ft_spec|None): - # Initialize distributed environment, initialize and load model - # addr is the corresponding service address when deploying sglang or fsdp/megatron remotely - # The controller passes the master addr when calling - pass - def get_scheduling_config(self): - # Get the resource configuration information required by the scheduler to schedule the engine, - # such as the engine's image, cpu/gpu/memory size - pass + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): + """Initialize distributed training environment and load models.""" + raise NotImplementedError() def destroy(self): - """Destroy the engine and release GPU memory.""" + """Clean up engine resources and release GPU memory.""" pass - async def upload_weights(self, meta: WeightUpdateMeta): - pass + def upload_weights(self, meta: WeightUpdateMeta): + """Upload weights to inference engine (blocking operation).""" + raise NotImplementedError() def save(self, meta: SaveLoadMeta): - pass + """Save model weights and optimizer states for checkpointing.""" + raise NotImplementedError() def load(self, meta: SaveLoadMeta): - pass - - # data api - def step_lr_scheduler(self): - """Step learning rate scheduler.""" - # Due to PPO minibatch updates, multiple train batches may need to be called - # before calling step_lr_scheduler once, so this api needs to be separated + """Load model weights and optimizer states from checkpoint.""" raise NotImplementedError() 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]: - pass + """Update model parameters using provided batch and loss function.""" + raise NotImplementedError() @torch.no_grad() - def eval_batch( - self, - input_: Dict, - loss_fn: Callable[[torch.Tensor, Dict], torch.Tensor], - loss_weight_fn: Callable[[Dict], float], - ) -> torch.Tensor | None: - pass - - @torch.no_grad() def forward( self, - input_: Dict, - output_seqlens: List[List[int]] | None = None, - post_hook: Callable[[torch.Tensor, Dict], Any] | None = None, + 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: - pass -############################################# - -# Implementation example -class FSDPEngine(TrainEngine): - def __init__(self, config: EngineConfig): - self.config = config - self.weight_update_group_initialized = False - - def initialize(self, addr: str|None, ft_spec: FinetuneSpec): - self.model_config = AutoConfig.from_pretrained(self.config.path) - with torch.device("meta"): - model = AutoModelForCausalLM.from_config(self.model_config) - self.model = FSDP(model) - self.optimizer = Adam(self.model.parameters()) - - def destroy(self): - del self.optimizer - del self.model - gc.collect() - torch.cuda.empty_cache() - gc.collect() - - async def upload_weights(self, meta: WeightUpdateMeta): - if meta.type == 'nccl': - if not self.weight_update_group_initialized: - await self.init_distributed_weight_update(meta) - return await self.aupdate_weights_from_distributed() - if meta.type == 'disk': - self.save_to_hf(meta.path) - return - - def save(self, meta): - if meta.weight_format == 'hf': - self.save_to_hf(meta.path, meta.tokenizer, meta.base_model_path) - elif meta.weight_format == 'dcp': - self.save_model_to_dcp(meta.path) - - if meta.with_optim: - self.save_optimizer_state(meta.path) - - def load(self, meta): - ... - - - ############# Helper methods start ############# - - def load_from_hf(self, path): - sd = load_hf_state_dict(path) - load_fsdp_state_dict(self.model, full_sd=sd) - - def save_to_hf(self, path, - tokenizer: Optional[transformers.PreTrainedTokenizerFast] = None, - base_model_path: Optional[str] = None, - ): - if dist.rank() == 0: - sd = {} - for n, p in self.model.named_parameters(): - sd[n] = p.data.gather() - torch.save(sd, path) - if tokenizer is not None: - tokenizer.save_pretrained(path) - if base_model_path is not None: - copy_hf_configs(base_model_path, path) - dist.barrier() - - def save_for_recover(self, path: str): - self.save_model_to_dcp(path) - self.save_optimizer_state(path) - - def load_from_recover(self, path): - self.load_model_from_dcp(path) - self.load_optimizer_state(path) - - async def init_distributed_weight_update(self, meta: WeightUpdateMeta): - # Initialize NCCL communication group for weight updates - ... - - async def aupdate_weights_from_distributed(self): - # Update inference weights through NCCL - # Different engines (FSDP, Megatron) may have different weight aggregation, - # splitting and communication methods - # Keep this high-level interface instead of defining subdivided interfaces - # to facilitate different engines implementing the most efficient weight communication path - ... - - def load_model_from_dcp(self, path: str): - # Load pytorch distributed checkpoint model from disk during recovery - ... - - def save_model_to_dcp(self, path: str): - # Save model in dcp format for recovery - ... - - def save_optimizer_state(self, path: str): - # Save optimizer state for recovery - raise NotImplementedError() - - def load_optimizer_state(self, path: str): - # Load optimizer state during recovery + """Execute gradient-free forward pass for inference.""" raise NotImplementedError() ``` -### Algorithm-Specific TrainEngine API +#### Algorithm Level: Extended Engines -Extended engines (such as Actor in PPO) provide convenient organization and calling of -computational interfaces specific to algorithms. These computational interfaces maintain -single-process computational logic, but can be called by controllers in top-level -training scripts to complete distributed semantic computational orchestration. +Extended engines like PPO Actor provide algorithm-specific organization and +computational interfaces. They leverage backend core methods (such as `forward`) to +generate algorithm-required tensors and execute specialized model updates. The produced +objects (e.g., `FSDPPPOActor`) are also instances of `TrainEngine`s, but affiliated with +methods that specially designed for the algorithm (e.g., `ppo_update`). ```python -class Actor(Engine): +class PPOActor: + + def __init__(self, config: PPOActorConfig, engine: TrainEngine): + self.config = config + self.engine = engine @torch.no_grad() - def compute_logps(self, input_: Dict[str, Tensor]) -> torch.Tensor: - ... # unpad - logps = self.forward(xxx) - ... # pad back - return logps + def compute_logp( + self, + data: TensorDict, + temperature: Optional[float] = None, + ) -> torch.Tensor | None: - def compute_advantages_and_returns(self, input_: Dict) -> Dict: + 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: + """Compute advantages for PPO training.""" + # Implementation details... pass - def ppo_update(self, input_: Dict) -> List[Dict[str, float]]: - ... - all_stats = [] - for _ in range(self.ppo_n_minibatches): - stats = self.train_batch(xxx, loss_fn=actor_loss_fn) - all_stats.append(stats) - return all_stats - -class Critic(Engine): - - @torch.no_grad() - def compute_values(self, input_: Dict) -> torch.Tensor: + def ppo_update(self, data: TensorDict) -> List[Dict[str, float]]: + """Execute PPO policy update.""" + # Implementation details... pass - def ppo_update(self, input_: Dict) -> List[Dict[str, float]]: - ... - all_stats = [] - for _ in range(self.ppo_n_minibatches): - stats = self.engine.train_batch(xxx, loss_fn=critic_loss_fn) - all_stats.append(stats) - return all_stats +class FSDPPPOActor(FSDPEngine): + """FSDP-backed PPO Actor implementation.""" -class FSDPActor(FSDPEngine, Actor): - pass + def __init__(self, config: PPOActorConfig): + super().__init__(config) + self.actor = PPOActor(config, self) -class MegatronActor(FSDPEngine, Actor): - pass + @torch.no_grad() + def compute_logp(self, *args, **kwargs) -> torch.Tensor | None: + return self.actor.compute_logp(*args, **kwargs) -class FSDPCritic(MegatronEngine, Critic): - pass + @torch.no_grad() + def compute_advantages(self, *args, **kwargs) -> None: + self.actor.compute_advantages(*args, **kwargs) -class MegatronCritic(MegatronEngine, Critic): - pass + def ppo_update(self, *args, **kwargs) -> List[Dict[str, float]]: + return self.actor.ppo_update(*args, **kwargs) ``` -### Inference Engine API - -Define the InferenceEngine API in a local-like mode, rather than a client-server -separated form, mainly for user-centered convenience in using the InferenceEngine as a -tool. +### Inference Engine Design -InferenceEngine can internally start a SGLang subprocess (SGLangEngine), or call a -remotely deployed service (RemoteSGLangEngine). +The inference engine's core functionality revolves around `generate` and +`update_weights` methods. These methods can interface with HTTP server APIs or invoke +local LLM engines: ```python -class InferenceEngine(ABC): - def __init__(self, config) - self.config = config +class InferenceEngine(abc.ABC): - @abstractmethod - def initialize(self, addr: str|None, ft_spec): - """Start SGLang Engine, starting the engine will call model loading by default - """ - self.tasks = [] + def initialize(self, addr: str | None, ft_spec): + """Initialize distributed inference environment and load models.""" + raise NotImplementedError() - async def update_weights(self, meta) -> None: - # Update model weights based on meta information + def destroy(self): + """Clean up engine resources and release GPU memory.""" pass async def agenerate(self, req: LLMRequest) -> LLMResponse: - # Given a prompt, generate a response with LLM - pass + """Generate response asynchronously for the given request.""" + raise NotImplementedError() + + def update_weights(self, meta: WeightUpdateMeta) -> Future: + """Update inference engine weights asynchronously.""" + raise NotImplementedError() +``` + +#### Workflow Integration - def submit(self, data: Dict[str, Any], workflow): - """ - Asynchronously submit rollout request - """ - task = asyncio.create_task(workflow.arun_episode(self, data)) - self.tasks.append(task) +User-defined rollout workflows utilize the `InferenceEngine` to generate trajectories. +The workflow's `arun_episode` method produces one or more trajectories from a single +prompt. The generation process is streaming rather than batched, with each dataset item +processed independently. Here's a simplified RLVR 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), + ) + + # Generate multiple responses concurrently + resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) - async def _wait_async(self, count: int, timeout: int) -> DistributedBatch: - tik = time.time() - n = 0 results = [] - while time.time() - tik < timeout and n < count: - done, _ = await asyncio.wait(self.tasks, return_when=FIRST_COMPLETED) - for task in done: - results.append(await task) - if n < count: - raise TimeoutError() - return DistributedBatch(results) - - def wait(self, count: int, timeout: int) -> DistributedBatch: - """ - Synchronous wait interface, until the request returns count records - """ - return asyncio.run(self._wait_async(count, timeout)) - - @abstractmethod - def rollout(self, data: List[Dict[str, Any]], workflow) -> DistributedBatch: - """ - Synchronously submit rollout request, until all inference requests are completed and returned - """ - pass + for resp in resps: + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + + results.append(TensorDict(res, batch_size=[1])) + + return concat_padded_tensors(results) +``` -################################### -# Implementation example -class SGLangEngine(InferenceEngine): +#### Batch Processing and Asynchronous Operations - def __init__(self, config: InfEngineConfig): - self.config = config - self.weight_update_group_initialized = False +While individual trajectory collection is straightforward, batching and asynchronous +execution require additional infrastructure. The `InferenceEngine` provides high-level +methods: `submit`, `wait`, `rollout_batch`, and `prepare_batch`. - async def update_weights(self, meta) -> None: - if meta.type == 'nccl': - if not self.weight_update_group_initialized: - await self.init_distributed_weight_update(meta) - return await self.aupdate_weights_from_distributed() - if meta.type == 'disk': - self.update_weights_from_disk(meta.path) - return +The `rollout_batch` method submits multiple `workflow.arun_episode` jobs to an +asynchronous thread pool using `submit`, then waits for completion using `wait`. The +`prepare_batch` method separates submission and waiting to enable asynchronous rollout: - async def agenerate(self, req: LLMRequest) -> LLMResponse: - # Given a prompt, generate a response with LLM - return await self.llm.generate_async(xxx) - - # Weight update - @abstractmethod - def update_weights_from_disk(self, path) -> None: - """Update model weights from disk""" - ... - - @abstractmethod - async def ainit_distributed_weights_update(self, meta_info: WeightUpdateMeta): - # Initialize **all** needed weight synchronization communication groups and communication plans - # (which communication type and which parameters to communicate at each step) - # Depending on the engine's partitioning method, multiple communication groups may be initialized - # for weight synchronization - # Since both inference and training engines need to enter this function, - # it needs to be defined as an async function - ... - - @abstractmethod - async def aupdate_weights_from_distributed(self) -> None: - """Use the initialized weight communication plan and communication groups to update model weights with NCCL - - Since both inference and training engines need to enter this function, - it needs to be defined as an async function - """ - ... - - @abstractmethod - def check_health(self) -> bool: - """Check server health status - - Returns: - bool: Whether the server is healthy - """ - pass +```python +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. Consider increasing queue_size.") + +def wait( + self, + count: int, + timeout: float | None = None, + should_accept: Callable | None = None, +) -> TensorDict: + """Wait for specified number of results with optional filtering.""" + # Implementation details... + pass + +def rollout_batch( + self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" +) -> TensorDict: + """Submit batch requests and wait for all results.""" + for item in data: + self.submit(item, workflow) + return self.wait(count=len(data)) + +def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", +): + """Prepare batch for asynchronous processing.""" + # Implementation details... + pass ``` -### RolloutWorkflow +### RolloutWorkflow Interface -RolloutWorkflow is a class that provides the arun_episode method. This method has a -fixed signature, used to collect one agent trajectory. +The `RolloutWorkflow` class provides the `arun_episode` method with a standardized +signature for collecting agent trajectories: ```python class MyRolloutWorkflow: @@ -521,263 +514,21 @@ class MyRolloutWorkflow: async def arun_episode(self, engine: InferenceEngine, data: Dict[str, Any]) -> Dict[str, Tensor]: - ... req = LLMRequest(input_ids=data['input_ids'], ...) + for _ in range(self.config.num_turns): resp = await engine.agenerate(req) res = await self.tool_executor.aexecute_tool(resp.completion) req.input_ids += res + reward = my_reward_fn(None, None, None, req.input_ids, **data) - ... return output ``` -### RolloutController & TrainController +### Controller Architecture -They have the same API as `InferenceEngine` and `TrainEngine`, respectively. - -## Other Components - -1. Algorithm workflows don't necessarily use these components; other replaceable - components such as implementations in rllm or verl-agent can also be used -1. Mainly for internal implementation and division of labor, as a thin wrapper to - facilitate adaptation of external packages - -### Environment - -1. Support multi-tool management and unified execution interface -1. Support local calling and mcp service calling -1. "Tools" belong to an instance rather than a class, register_tool is defined as a - method rather than a static method, this is (1) to prevent tools from subclasses - being registered to the base class, causing potential naming or calling conflicts; - (2) to support multiple tasks for the same service (e.g., browser), with each task - having a different toolset - -```python -from abc import ABC, abstractmethod -from typing import Any, Dict, List - -class ToolingEnv: - def __init__(self): - self.is_initialized = False - - self.tool_registry: Dict[str, Callable] = {} - self.tool_schemas: List[Dict[str, Any]] = [] - - async def ainitialize(self): - """ - Performs the initialization logic for the environment. - For stateful environments, this is where resources are created and - prepared (e.g., launching a browser). - """ - pass - - def list_tools(self) -> List[Dict[str, Any]]: - pass - - async def aexecute(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: - pass - - async def aclose(self): - """ - Destroys the environment, releasing all held resources. - This method is critical for stateful environments (e.g., a browser session). - """ - pass - -class MCPToolingEnv(ToolingEnv): - def __init__(self, config: MCPToolingEnvConfig): - self.config = config - # init a mcp client - self.mcp_client = mcp.session_client(config.url) - self.tool_registry: Dict[str, Callable] = mcp.session_client.list_tools() - self.tool_schemas: List[Dict[str, Any]] = [] - - def list_tools(self) -> List[Dict[str, Any]]: - return self.tool_schemas - - def aexecute(self, tool_name: str, tool_args: Dict[str, Any]): - pass - - -class LocalToolingEnv(ToolingEnv): - - @staticmethod - def generate_schema(func: Callable) -> Dict[str, Any]: - """ - Generates a JSON schema for a function using introspection. - """ - # Use the function's docstring as the tool's description. - description = inspect.getdoc(func) or "No description provided." - sig = inspect.signature(func) - parameters = { - "type": "object", - "properties": {}, - "required": [], - } - - # Mapping from Python types to JSON Schema types - type_mapping = { - str: "string", - int: "integer", - float: "number", - bool: "boolean", - } - - for name, param in sig.parameters.items(): - # Default to string type if type hint is missing or complex - param_type = type_mapping.get(param.annotation, "string") - parameters["properties"][name] = {"type": param_type} - - # If a parameter has no default value, it is considered required. - if param.default is inspect.Parameter.empty: - parameters["required"].append(name) - - return { - "type": "function", - "function": { - "name": func.__name__, - "description": description, - "parameters": parameters, - } - } - - def register_tool(self, func: Callable) -> Callable: - """ - A decorator that registers a Python function as a tool in this environment. - """ - if not callable(func): - raise TypeError("The provided object must be a callable function.") - - tool_name = func.__name__ - if tool_name in self.tool_registry: - raise ValueError(f"Tool with name '{tool_name}' is already registered.") - - # Add the function to the registry and its schema to the schema list. - self.tool_registry[tool_name] = func - self.tool_schemas.append(self.generate_schema(func)) - - def list_tools(self) -> List[Dict[str, Any]]: - """ - Lists all available tools provided by this environment and their descriptions. - - Returns: - A list of dictionaries, where each dictionary describes a tool. - """ - return self.tool_schemas - - async def aexecute(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: - """ - Executes a specified tool. - - Args: - tool_name (str): The name of the tool to execute. - tool_args (Dict[str, Any]): The arguments required to call the tool. - - Returns: - Any: The result of the tool's execution, typically a string or - structured JSON. - """ - if tool_name not in self._tool_registry: - return f"Error: Tool '{tool_name}' is not registered." - - tool_func = self._tool_registry[tool_name] - - try: - result = tool_func(**tool_args) - return result - except TypeError as e: - # This exception is often triggered by missing or incorrect argument types. - return f"Error executing '{tool_name}': Invalid arguments. Details: {e}" - except Exception as e: - return f"Error executing '{tool_name}': An unexpected error occurred. Details: {e}" - -``` - -### Reward - -1. Workflows for computing rewards using models and computing rewards based on rules - should be separated -1. Rule-based reward computation is defined as a function with a predefined signature, - which can be local or remote, and is generally wrapped in the rollout workflow; - user-written reward functions can also not use this signature as long as they provide - a workflow -1. Computing rewards using models requires users to initialize a controller/engine - themselves and explicitly call it in the algorithm workflow - -```python - -############ Rule-based reward ############ - -# The signature is just a reference, can be defined arbitrarily -def reward_fn(prompt: str, completion: str, prompt_ids: List[int], - completion_ids: List[int], **kwargs): - # prompt: prompt string (the task this data needs to complete) - # completion: trajectory string generated by the model based on the task - # prompt_ids: token ids of the prompt - # completion_ids: token ids of the trajectory generated by the model - # kwargs: all other attributes of this data in the dataset, - # for example, solutions, input_outputs, etc. - pass - -############ Model-based reward ############ - -reward = TrainController(Critic()) -rollout_controller = RolloutController(...) -for _ in range(epochs): - for _ in range(steps_per_epoch): - data = rollout_controller.rollout_batch(prompt) - data['reward'] = reward.compute_values(data) - ... -``` - -### Dataset - -Use huggingface datasets and pytorch torchdata. In Single-Controller mode, only one -process per experiment is responsible for data loading. - -```python -from datasets import Dataset, load_dataset -from datasets.distributed import split_dataset_by_node - -dataset = load_dataset( - path, - name=name, - split=split, - data_files=data_files, -) -dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) - -# Access the first data item -data: Dict = dataset[0] - -# Access the "prompt" column -data: List = data['prompt'] - -# Data processing -def process_example(example, idx): - # Add query_id column - example["query_id"] = str(idx) - example["prompt"] = example["question"] - - # used by the reward function - example["method"] = reward_mode - return example - -dataset = dataset.map( - lambda example, idx: process_example(example, idx), - with_indices=True, -) - -# Data loading and shuffle -from torchdata.stateful_dataloader import StatefulDataLoader -dataloader = StatefulDataLoader( - dataset=dataset, - batch_size=batch_size, - shuffle=True, - drop_last=True, - collate_fn=lambda x: x, # Can change the data batch packing method by changing this parameter -) -for data in dataloader: - assert isinstance(data, list) -``` +`RolloutController` and `TrainController` mirror the APIs of `InferenceEngine` and +`TrainEngine` respectively. Controllers handle engine deployment across the cluster and +manage data distribution, invoking engine methods through remote procedure calls (RPCs). +This architecture enables distributed operation while maintaining familiar interfaces +for users. From 86ee6876362987a6b3881a9917fc0b7c884b981f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Fri, 25 Jul 2025 11:12:30 +0800 Subject: [PATCH 15/20] PullRequest: 431 [Fix] Fix environment of lite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-dev of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/431 Reviewed-by: 晓雷 * change requirements * . * . * . --- examples/arealite/configs/gsm8k_grpo.yaml | 2 +- examples/env/scripts/setup-pip-deps.sh | 2 +- pyproject.toml | 4 +++- requirements.txt | 3 +++ 4 files changed, 8 insertions(+), 3 deletions(-) diff --git a/examples/arealite/configs/gsm8k_grpo.yaml b/examples/arealite/configs/gsm8k_grpo.yaml index 9b0362d20e..e1630bd339 100644 --- a/examples/arealite/configs/gsm8k_grpo.yaml +++ b/examples/arealite/configs/gsm8k_grpo.yaml @@ -83,7 +83,7 @@ sglang: dtype: ${actor.dtype} max_running_requests: null context_length: 32768 - mem_fraction_static: 0.9 + mem_fraction_static: 0.8 # datasets train_dataset: diff --git a/examples/env/scripts/setup-pip-deps.sh b/examples/env/scripts/setup-pip-deps.sh index 1599170a38..e907109e7b 100644 --- a/examples/env/scripts/setup-pip-deps.sh +++ b/examples/env/scripts/setup-pip-deps.sh @@ -6,7 +6,7 @@ pip install torch==2.7.1 torchaudio==2.7.1 torchvision==0.22.1 "deepspeed>=0.17. pip install "sglang[all]==0.4.9.post2" pip install megatron-core==0.11.0 nvidia-ml-py pip install git+https://github.com/garrett4wade/cugae --no-build-isolation --verbose -pip install "flash-attn<=2.7.3" --no-build-isolation +pip install "flash-attn<=2.8.2" --no-build-isolation # Package used for calculating math reward pip install -e evaluation/latex2sympy diff --git a/pyproject.toml b/pyproject.toml index c779fa77c7..4e9e558ef7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "huggingface_hub", "datasets", "accelerate", - "transformers==4.53.0", + "transformers==4.53.1", # Scientific computing "scipy", @@ -56,6 +56,8 @@ dependencies = [ "torchdata", "autoflake", "tensordict", + "pybase64", + "msgspec", # Monitoring and logging "wandb", diff --git a/requirements.txt b/requirements.txt index 0c21f9b0a6..7cde02b42e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -74,3 +74,6 @@ torchdata autoflake tensordict deepspeed>=0.17.2 +pybase64 +msgspec +transformers==4.53.1 \ No newline at end of file From 04c8da4eb0bbc2b8b31de05920e6797c27f9116a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=B0=E4=B8=B4?= Date: Fri, 25 Jul 2025 22:48:24 +0800 Subject: [PATCH 16/20] PullRequest: 440 [FIX] fix update weight from disk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch sxj/lite-fix-disk-update of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/440 Reviewed-by: 博惟 * [FIX] fix update weight from disk --- arealite/engine/fsdp_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index 4338b6d55d..25b48031f5 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -150,7 +150,7 @@ def upload_weights(self, meta: WeightUpdateMeta): update_name = names.update_weights_from_disk( self.config.experiment_name, self.config.trial_name, - self.model_version, + self.get_version(), ) name_resolve.add( update_name, str(datetime.now().timestamp()), keepalive_ttl=120 From 7998055c4c7eca1ede7c178cbb849e443fe111da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=93=E9=9B=B7?= Date: Tue, 29 Jul 2025 10:21:42 +0800 Subject: [PATCH 17/20] PullRequest: 442 [lite] Refactor `RemoteSGLangEngine` into two parts: `RemoteSGLangEngine` and `WorkflowExecutor`. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch mzy/workflow-executor of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/442 Reviewed-by: 博惟 * refactor workflow executor * . * fix tests and eval * . * . * revert workflow executor into remote sglang engine * . --- arealite/api/engine_api.py | 1 + arealite/api/workflow_api.py | 249 +++++++++++++++++++- arealite/engine/sglang_remote.py | 270 ++++------------------ arealite/utils/http.py | 1 + examples/arealite/configs/gsm8k_grpo.yaml | 6 +- 5 files changed, 299 insertions(+), 228 deletions(-) diff --git a/arealite/api/engine_api.py b/arealite/api/engine_api.py index bd00eb5757..db0f4f5018 100644 --- a/arealite/api/engine_api.py +++ b/arealite/api/engine_api.py @@ -174,6 +174,7 @@ def prepare_batch( self, dataloader: StatefulDataLoader, workflow: "RolloutWorkflow", + should_accept: Callable | None = None, ): """Asynchronously submit and wait until a full batch is ready.""" raise NotImplementedError() diff --git a/arealite/api/workflow_api.py b/arealite/api/workflow_api.py index 9141399c1d..02248463c1 100644 --- a/arealite/api/workflow_api.py +++ b/arealite/api/workflow_api.py @@ -1,10 +1,29 @@ -from typing import TYPE_CHECKING, Any, Dict +import asyncio +import queue +import threading +import time +import traceback +from typing import TYPE_CHECKING, Any, Callable, Dict, List +import torch.distributed as dist +import uvloop 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 RolloutStat +from arealite.utils.data import concat_padded_tensors +from realhf.base import logging if TYPE_CHECKING: from arealite.api.engine_api import InferenceEngine +logger = logging.getLogger("arealite.workflow_api") + + +ROLLOUT_POLL_WAIT_TIME = 0.05 + class RolloutWorkflow: @@ -16,3 +35,231 @@ async def arun_episode( See concrete example implementations under the `arealite/workflow` directory. """ raise NotImplementedError() + + +class WorkflowExecutor: + + def __init__( + self, + config: InferenceEngineConfig, + inference_engine: "InferenceEngine", + ): + config.max_concurrent_rollouts = ( + config.max_concurrent_rollouts or config.consumer_batch_size + ) + self.config = config + self.exiting = threading.Event() + self.paused = threading.Event() + self.lock = threading.Lock() + + self.inference_engine = inference_engine + + qsize = config.queue_size or config.max_concurrent_rollouts * 16 + self.input_queue = queue.Queue(maxsize=qsize) + self.output_queue = queue.Queue(maxsize=qsize) + self.result_cache: List[TensorDict] = [] + + self.rollout_stat = RolloutStat() + + def initialize(self): + self.rollout_tasks: Dict[str, asyncio.Task] = {} + self.rollout_thread = threading.Thread(target=self._rollout_thread) + self.rollout_thread.start() + + def destroy(self): + self.exiting.set() + self.rollout_thread.join() + + 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 + with self.lock: + version = self.inference_engine.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 _rollout_thread(self): + """Thread that runs the rollout loop.""" + try: + uvloop.run(self._rollout_thread_async()) + except Exception: + traceback.print_exc() + + async def _rollout_thread_async(self): + rollout_tasks = self.rollout_tasks + rid = 0 + try: + while not self.exiting.is_set(): + # Check capacity + capacity = self.get_capacity() + # Create new rollout task + 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.inference_engine, data), + name=str(rid), + ) + with self.lock: + rollout_tasks[str(rid)] = task + self.rollout_stat.submitted += 1 + self.rollout_stat.running += 1 + 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 + # Wait for rollout completion + with self.lock: + tasks = list(rollout_tasks.values()) + done = [] + if tasks: + done, _ = await asyncio.wait( + tasks, + timeout=ROLLOUT_POLL_WAIT_TIME, + return_when=asyncio.FIRST_COMPLETED, + ) + # Collect done results + for task in done: + traj = await task + traj: TensorDict + task_rid = task.get_name() + with self.lock: + rollout_tasks.pop(task_rid) + self.rollout_stat.accepted += 1 + + try: + self.output_queue.put_nowait(traj) + except queue.Full: + raise RuntimeError( + "Output queue full. Please increase queue_size." + ) + + with self.lock: + self.rollout_stat.running -= 1 + 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}." + ) + await asyncio.sleep(1) + except Exception: + traceback.print_exc() + finally: + # Cancel remaining tasks + with self.lock: + for task in rollout_tasks.values(): + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: + try: + self.input_queue.put_nowait((data, workflow)) + except queue.Full: + raise RuntimeError("Input queue full. Please increase queue_size.") + + 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() + and time.perf_counter() - tik < timeout + ): + try: + result = self.output_queue.get(timeout=ROLLOUT_POLL_WAIT_TIME) + 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 queue.Empty: + pass + if self.exiting.is_set(): + raise RuntimeError("Rollout engine is exiting, cannot wait for results.") + if accepted < count: + raise TimeoutError( + f"Timed out waiting for {count} rollouts, " f"only received {accepted}." + ) + results, self.result_cache = ( + self.result_cache[:count], + self.result_cache[count:], + ) + return concat_padded_tensors(results) + + 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)) + + def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", + should_accept: Callable | None = None, + ): + if not hasattr(self, "data_generator"): + self.data_generator = iter(dataloader) + assert dataloader.batch_size is not None + while True: + # 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(self.data_generator) + except StopIteration: + self.data_generator = iter(dataloader) + data = next(self.data_generator) + for item in data: + self.submit(item, workflow=workflow) + try: + return self.wait( + dataloader.batch_size, timeout=1, should_accept=should_accept + ) + except TimeoutError: + pass + + def pause(self): + self.paused.set() + + def resume(self): + self.paused.clear() diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index bac9de346f..3db53d6c56 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -2,17 +2,13 @@ import os import random import shutil -import threading import time -import traceback from concurrent.futures import Future, ProcessPoolExecutor from datetime import datetime -from queue import Empty, Full, Queue -from typing import TYPE_CHECKING, Any, Callable, Dict, List +from typing import 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,28 +19,21 @@ FinetuneSpec, LLMRequest, LLMResponse, - RolloutStat, WeightUpdateMeta, ) -from arealite.utils.data import concat_padded_tensors +from arealite.api.workflow_api import RolloutWorkflow, WorkflowExecutor from arealite.utils.http import arequest_with_retry, get_default_connector from realhf.base import logging, name_resolve, names -if TYPE_CHECKING: - from arealite.api.workflow_api import RolloutWorkflow logger = logging.getLogger(__name__) -ROLLOUT_POLL_WAIT_TIME = 0.05 RID_CACHE_SIZE = 128 class RemoteSGLangEngine(InferenceEngine): def __init__(self, config: InferenceEngineConfig): - config.max_concurrent_rollouts = ( - config.max_concurrent_rollouts or config.consumer_batch_size - ) self.config = config self.rid_to_address = {} @@ -52,29 +41,20 @@ def __init__(self, config: InferenceEngineConfig): self.rid_queue = [] self.addresses = os.getenv("AREAL_LLM_SERVER_ADDRS").split(",") + self.session = None + 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!") 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.distributed_weight_update_initialized = False - self._version = 0 + self.workflow_executor = WorkflowExecutor( + config=config, + inference_engine=self, + ) + def _wait_for_server(self, address): base_url = f"http://{address}" tik = time.time() @@ -92,135 +72,44 @@ def check_health(self, base_url): 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] = {} - + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None = None): + logger.info("Waiting for server ready...") + for addr_ in self.addresses: + self._wait_for_server(addr_) + logger.info("Servers are all ready!") self.executor = ProcessPoolExecutor(max_workers=1) - self.rollout_thread = threading.Thread(target=self._rollout_thread) - self.rollout_thread.start() + self.workflow_executor.initialize() def destroy(self): self.executor.shutdown() - self.exiting.set() - self.rollout_thread.join() def set_version(self, version): - with self.lock: - self._version = version + self._version = version def get_version(self): - with self.lock: - return self._version - - def _rollout_thread(self): - """Thread that runs the rollout loop.""" - try: - uvloop.run(self._rollout_thread_async()) - except Exception: - traceback.print_exc() - - async def _rollout_thread_async(self): - 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(): - # Check capacity - capacity = self.get_capacity() - # Create new rollout task - 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, data), name=str(rid) - ) - with self.lock: - rollout_tasks[str(rid)] = task - self.rollout_stat.submitted += 1 - self.rollout_stat.running += 1 - 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 - # Wait for rollout completion - with self.lock: - tasks = list(rollout_tasks.values()) - done = [] - if tasks: - done, _ = await asyncio.wait( - tasks, - timeout=ROLLOUT_POLL_WAIT_TIME, - return_when=asyncio.FIRST_COMPLETED, - ) - # Collect done results - for task in done: - traj = await task - traj: TensorDict - task_rid = task.get_name() - with self.lock: - rollout_tasks.pop(task_rid) - self.rollout_stat.accepted += 1 - - try: - self.output_queue.put_nowait(traj) - except Full: - raise RuntimeError( - "Output queue full. Please increase queue_size." - ) - - with self.lock: - self.rollout_stat.running -= 1 - 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}." - ) - await asyncio.sleep(1) - except Exception: - traceback.print_exc() - finally: - # Cancel remaining tasks - with self.lock: - for task in rollout_tasks.values(): - if not task.done(): - task.cancel() - try: - await task - except asyncio.CancelledError: - pass + return self._version def choose_server(self) -> str: - 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 + 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: """Async version of generate using aiohttp.""" + if self.session is None: + # NOTE: Lazily initialize aiohttp.ClientSession since it needs to be initialized + # inside asyncio loop in WorkflowExecutor + 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(), + ) # Prepare request payload gconfig = req.gconfig stop_token_ids = gconfig.stop_token_ids @@ -370,30 +259,8 @@ def callback(fut): fut.add_done_callback(callback) return fut - 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 submit(self, data: Dict[str, Any], workflow: RolloutWorkflow) -> None: + return self.workflow_executor.submit(data, workflow) def wait( self, @@ -401,76 +268,31 @@ def wait( 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() - and time.perf_counter() - tik < timeout - ): - try: - result = self.output_queue.get(timeout=ROLLOUT_POLL_WAIT_TIME) - 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: - pass - if self.exiting.is_set(): - raise RuntimeError("Rollout engine is exiting, cannot wait for results.") - if accepted < count: - raise TimeoutError( - f"Timed out waiting for {count} rollouts, " f"only received {accepted}." - ) - results, self.result_cache = ( - self.result_cache[:count], - self.result_cache[count:], + return self.workflow_executor.wait( + count, + timeout=timeout, + should_accept=should_accept, ) - return concat_padded_tensors(results) 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)) + return self.workflow_executor.rollout_batch(data, workflow) def prepare_batch( self, dataloader: StatefulDataLoader, - workflow: "RolloutWorkflow", + workflow: RolloutWorkflow, ): - if not hasattr(self, "data_generator"): - self.data_generator = iter(dataloader) - assert dataloader.batch_size is not None - while True: - # 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(self.data_generator) - except StopIteration: - self.data_generator = iter(dataloader) - data = next(self.data_generator) - for item in data: - self.submit(item, workflow=workflow) - try: - return self.wait(dataloader.batch_size, timeout=1) - except TimeoutError: - pass + return self.workflow_executor.prepare_batch(dataloader, workflow) def pause(self): - self.paused.set() + """Pause request submission for async rollout. Used during evaluation to prevent data over generation.""" + return self.workflow_executor.pause() def resume(self): - self.paused.clear() + """Resume request submission for async rollout.""" + return self.workflow_executor.resume() def update_weights_from_disk( diff --git a/arealite/utils/http.py b/arealite/utils/http.py index a39a3614de..0ddca140da 100644 --- a/arealite/utils/http.py +++ b/arealite/utils/http.py @@ -61,6 +61,7 @@ async def arequest_with_retry( 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") diff --git a/examples/arealite/configs/gsm8k_grpo.yaml b/examples/arealite/configs/gsm8k_grpo.yaml index e1630bd339..6613784488 100644 --- a/examples/arealite/configs/gsm8k_grpo.yaml +++ b/examples/arealite/configs/gsm8k_grpo.yaml @@ -81,7 +81,7 @@ sglang: random_seed: ${seed} skip_tokenizer_init: true dtype: ${actor.dtype} - max_running_requests: null + max_running_requests: null context_length: 32768 mem_fraction_static: 0.8 @@ -91,7 +91,7 @@ train_dataset: shuffle: true pin_memory: true -valid_dataset: +valid_dataset: batch_size: 256 shuffle: true pin_memory: true @@ -126,4 +126,4 @@ stats_logger: trial_name: ${trial_name} fileroot: ${cluster.fileroot} wandb: - mode: disabled \ No newline at end of file + mode: disabled From e825b99a1cc53245a9467f8fbe79d6de29790ac7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=93=E9=9B=B7?= Date: Thu, 31 Jul 2025 12:59:21 +0800 Subject: [PATCH 18/20] PullRequest: 456 [lite] [Bug] Use `ProcessPoolExecutor` to calculate reward to avoid rollout slow down MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch mzy/lite/fix-reward of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/456?tab=comment Reviewed-by: 博惟 * fix reward * . * . * . --- arealite/workflow/rlvr.py | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/arealite/workflow/rlvr.py b/arealite/workflow/rlvr.py index 5404972b5f..07bb1b06fd 100644 --- a/arealite/workflow/rlvr.py +++ b/arealite/workflow/rlvr.py @@ -1,6 +1,8 @@ import asyncio +import functools import os import uuid +from concurrent.futures import ProcessPoolExecutor import colorama import torch @@ -13,6 +15,8 @@ from arealite.api.workflow_api import RolloutWorkflow from arealite.utils.data import concat_padded_tensors +REWARD_TIMEOUT_SECONDS = 15 + class RLVRWorkflow(RolloutWorkflow): def __init__( @@ -28,6 +32,7 @@ def __init__( self.tokenizer = tokenizer self.enable_thinking = enable_thinking self.dump_dir = dump_dir + self.rw_executor = ProcessPoolExecutor(max_workers=4) if self.dump_dir is not None and not os.path.exists(self.dump_dir): os.makedirs(self.dump_dir, exist_ok=True) @@ -53,6 +58,7 @@ async def arun_episode(self, engine: InferenceEngine, data): seqlens = [] results = [] + loop = asyncio.get_event_loop() for resp in resps: seq = resp.input_tokens + resp.output_tokens logprobs = [0.0] * resp.input_len + resp.output_logprobs @@ -64,13 +70,23 @@ async def arun_episode(self, engine: InferenceEngine, data): prompt_strs.append(prompt_str) completions_strs.append(completions_str) seqlens.append(len(seq)) - reward = self.reward_fn( - prompt=prompt_str, - completions=completions_str, - prompt_ids=resp.input_tokens, - completion_ids=resp.output_tokens, - **data, - ) + try: + reward = await asyncio.wait_for( + loop.run_in_executor( + self.rw_executor, + functools.partial( + self.reward_fn, + prompt_str, + completions_str, + resp.input_tokens, + resp.output_tokens, + **data, + ), + ), + timeout=REWARD_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + reward = 0 rewards.append(reward) res = dict( # unsqueeze to add an additional batch dimension From 7a77fab1a7b74066d48e30f13946439efa1a9a14 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Thu, 31 Jul 2025 15:47:06 +0800 Subject: [PATCH 19/20] PullRequest: 460 [lite][fix] add a warning when reward computation timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-fix of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/460 Reviewed-by: 晓雷 * add a warning when reward computation timeout --- arealite/experimental/sglang_engine.py | 1 + arealite/launcher/ray.py | 1 - arealite/workflow/rlvr.py | 6 ++++++ 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/arealite/experimental/sglang_engine.py b/arealite/experimental/sglang_engine.py index 777091e8d2..9ee7bdc706 100644 --- a/arealite/experimental/sglang_engine.py +++ b/arealite/experimental/sglang_engine.py @@ -94,6 +94,7 @@ def _rollout_thread(self): asyncio.run(self._rollout_thread_async()) except Exception as e: traceback.print_exc() + raise e async def _rollout_thread_async(self): data = None diff --git a/arealite/launcher/ray.py b/arealite/launcher/ray.py index 0344e983e2..0a0fdf525c 100644 --- a/arealite/launcher/ray.py +++ b/arealite/launcher/ray.py @@ -324,7 +324,6 @@ def ray_main(): ) 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: diff --git a/arealite/workflow/rlvr.py b/arealite/workflow/rlvr.py index 07bb1b06fd..fa97bcae73 100644 --- a/arealite/workflow/rlvr.py +++ b/arealite/workflow/rlvr.py @@ -14,6 +14,9 @@ from arealite.api.io_struct import LLMRequest from arealite.api.workflow_api import RolloutWorkflow from arealite.utils.data import concat_padded_tensors +from realhf.base import logging + +logger = logging.getLogger("RLVR workflow") REWARD_TIMEOUT_SECONDS = 15 @@ -86,6 +89,9 @@ async def arun_episode(self, engine: InferenceEngine, data): timeout=REWARD_TIMEOUT_SECONDS, ) except asyncio.TimeoutError: + logger.warning( + f"Computing reward timeout after {REWARD_TIMEOUT_SECONDS}s. Set reward to 0." + ) reward = 0 rewards.append(reward) res = dict( From 7d1c852abaffacf3e27ad217e0db2ed0503d8793 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Thu, 31 Jul 2025 18:45:27 +0800 Subject: [PATCH 20/20] PullRequest: 465 [lite][fix] Fix issues raised by tsao MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch fw/lite-fix of git@code.alipay.com:inclusionAI/AReaL.git into lite https://code.alipay.com/inclusionAI/AReaL/pull_requests/465 Reviewed-by: 晓雷 * fix --- arealite/api/cli_args.py | 42 ++++---------- arealite/api/workflow_api.py | 68 +++++++++++------------ arealite/engine/fsdp_engine.py | 2 +- arealite/engine/ppo/actor.py | 2 +- arealite/workflow/multi_turn.py | 2 +- docs/arealite/gsm8k_grpo.md | 26 ++------- docs/customization/agent.md | 15 ++--- docs/customization/algorithm.md | 9 +-- examples/arealite/configs/gsm8k_grpo.yaml | 1 + examples/arealite/gsm8k_grpo.py | 10 +--- 10 files changed, 64 insertions(+), 113 deletions(-) diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 7dc9be6dcf..8b90d3bcb4 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -195,7 +195,10 @@ class TrainEngineConfig: gradient_checkpointing: bool = field( default=True, metadata={"help": "Enable gradient checkpointing"} ) - dtype: str = field(default="float16", metadata={"help": "Parameter dtype."}) + dtype: str = field(default="bfloat16", metadata={"help": "Parameter dtype."}) + grad_reduce_dtype: str = field( + default="float32", metadata={"help": "Gradient reduce dtype."} + ) optimizer: Optional[OptimizerConfig] = field( default=None, metadata={"help": "Optimizer configuration"} ) @@ -262,7 +265,7 @@ class PPOActorConfig(TrainEngineConfig): default=1.0, metadata={"help": "Lambda parameter for GAE"} ) adv_norm: bool = field( - default=True, metadata={"help": "Enable advantage normalization"} + default=True, metadata={"help": "Enable advantage normalization globally"} ) # KL Control @@ -275,7 +278,9 @@ class PPOActorConfig(TrainEngineConfig): ) use_decoupled_loss: bool = field( default=False, - metadata={"help": "Use the decoupled loss. recompute_logprob must be True."}, + metadata={ + "help": "Use the decoupled loss. Implicitly enable recompute_logprob." + }, ) behav_imp_weight_cap: Optional[float] = field( default=None, @@ -329,7 +334,7 @@ class SGLangConfig: schedule_policy: str = "lpm" schedule_conservativeness: float = 1.0 cpu_offload_gb: int = 0 - dtype: str = "float16" + dtype: str = "bfloat16" kv_cache_dtype: str = "auto" # logging log_level: str = "warning" @@ -407,31 +412,8 @@ def build_args( dist_init_addr=dist_init_addr, **args, ) - sglang_version = pkg_version.get_version("sglang") - 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 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") - + if not pkg_version.is_version_greater_or_equal("sglang", "0.4.9.post2"): + raise RuntimeError("Needs sglang>=0.4.9.post2 to run the code.") return args @@ -466,7 +448,7 @@ class InferenceEngineConfig: default="round_robin", metadata={"help": "Request scheduling policy", "choices": ["round_robin"]}, ) - setup_timeout: float = field(default=90.0) + setup_timeout: float = field(default=120.0) request_timeout: float = field( default=3600, metadata={"help": "Timeout for HTTP requests."} ) diff --git a/arealite/api/workflow_api.py b/arealite/api/workflow_api.py index 02248463c1..fbb24ddc88 100644 --- a/arealite/api/workflow_api.py +++ b/arealite/api/workflow_api.py @@ -1,4 +1,5 @@ import asyncio +import itertools import queue import threading import time @@ -76,18 +77,17 @@ def get_capacity(self): 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 with self.lock: + max_concurrent_rollouts = max( + 1, self.config.max_concurrent_rollouts // world_size + ) + capacity = max_concurrent_rollouts - len(self.rollout_tasks) + # Staleness control version = self.inference_engine.get_version() - ofp = self.config.max_head_offpolicyness - with self.lock: + ofp = self.config.max_head_offpolicyness 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) + consumer_bs = max(1, self.config.consumer_batch_size // world_size) + capacity = min(capacity, (ofp + version + 1) * consumer_bs - sample_cnt) return capacity def _rollout_thread(self): @@ -105,6 +105,7 @@ async def _rollout_thread_async(self): # Check capacity capacity = self.get_capacity() # Create new rollout task + self.lock.acquire() while ( capacity > 0 and not self.paused.is_set() @@ -116,22 +117,22 @@ async def _rollout_thread_async(self): workflow.arun_episode(self.inference_engine, data), name=str(rid), ) - with self.lock: - rollout_tasks[str(rid)] = task - self.rollout_stat.submitted += 1 - self.rollout_stat.running += 1 - 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}." - ) + rollout_tasks[str(rid)] = task + self.rollout_stat.submitted += 1 + self.rollout_stat.running += 1 + 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 + tasks = list(rollout_tasks.values()) + self.lock.release() + # Wait for rollout completion - with self.lock: - tasks = list(rollout_tasks.values()) done = [] if tasks: done, _ = await asyncio.wait( @@ -148,14 +149,6 @@ async def _rollout_thread_async(self): rollout_tasks.pop(task_rid) self.rollout_stat.accepted += 1 - try: - self.output_queue.put_nowait(traj) - except queue.Full: - raise RuntimeError( - "Output queue full. Please increase queue_size." - ) - - with self.lock: self.rollout_stat.running -= 1 if self.config.enable_rollout_tracing: logger.info( @@ -164,6 +157,13 @@ async def _rollout_thread_async(self): f"running: {self.rollout_stat.running}, " f"accepted: {self.rollout_stat.accepted}." ) + try: + self.output_queue.put_nowait(traj) + except queue.Full: + raise RuntimeError( + "Output queue full. Please increase queue_size." + ) + await asyncio.sleep(1) except Exception: traceback.print_exc() @@ -235,7 +235,7 @@ def prepare_batch( should_accept: Callable | None = None, ): if not hasattr(self, "data_generator"): - self.data_generator = iter(dataloader) + self.data_generator = itertools.cycle(dataloader) assert dataloader.batch_size is not None while True: # Submit at least two batches to allow maximum overlap @@ -244,11 +244,7 @@ def prepare_batch( and self.input_queue.qsize() + dataloader.batch_size < self.input_queue.maxsize ): - try: - data = next(self.data_generator) - except StopIteration: - self.data_generator = iter(dataloader) - data = next(self.data_generator) + data = next(self.data_generator) for item in data: self.submit(item, workflow=workflow) try: diff --git a/arealite/engine/fsdp_engine.py b/arealite/engine/fsdp_engine.py index 25b48031f5..112f9cc371 100644 --- a/arealite/engine/fsdp_engine.py +++ b/arealite/engine/fsdp_engine.py @@ -54,7 +54,7 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): # Simple auto wrap policy self.mixed_precision_policy = MixedPrecisionPolicy( param_dtype=getattr(torch, self.config.dtype), - reduce_dtype=torch.float32, + reduce_dtype=getattr(torch, self.config.grad_reduce_dtype), cast_forward_inputs=True, ) self.device_mesh = create_fsdp_device_mesh(self.world_size, self.world_size) diff --git a/arealite/engine/ppo/actor.py b/arealite/engine/ppo/actor.py index 802c03e13c..ba40e0ccb9 100644 --- a/arealite/engine/ppo/actor.py +++ b/arealite/engine/ppo/actor.py @@ -128,7 +128,7 @@ def compute_advantages(self, data: TensorDict) -> None: advantages = torch.stack(advantages_reversed[::-1], dim=1) # Optionally perform advantage normalization. - if self.adv_norm: + if self.adv_norm or self.group_adv_norm: if self.group_adv_norm: adv_list = [] for i in range(0, bs, self.group_size): diff --git a/arealite/workflow/multi_turn.py b/arealite/workflow/multi_turn.py index eec8aa1e77..14bd263b49 100644 --- a/arealite/workflow/multi_turn.py +++ b/arealite/workflow/multi_turn.py @@ -31,7 +31,7 @@ async def arun_episode(self, engine: InferenceEngine, data): messages = data["messages"] # Run multi-turn rollout until correct t = reward = 0 - discount = 0 + discount = 1 rid = uuid.uuid4().hex while reward == 0 and t < self.max_turns: # Amend a prompt if the previous answer is incorrect diff --git a/docs/arealite/gsm8k_grpo.md b/docs/arealite/gsm8k_grpo.md index 26125d32ed..be3b9ce92c 100644 --- a/docs/arealite/gsm8k_grpo.md +++ b/docs/arealite/gsm8k_grpo.md @@ -233,7 +233,7 @@ def prepare_batch( workflow: "RolloutWorkflow", ): if not hasattr(self, "data_generator"): - self.data_generator = iter(dataloader) + self.data_generator = itertools.cycle(dataloader) assert dataloader.batch_size is not None while True: # Submit at least two batches to allow maximum overlap @@ -242,11 +242,7 @@ def prepare_batch( and self.input_queue.qsize() + dataloader.batch_size < self.input_queue.maxsize ): - try: - data = next(self.data_generator) - except StopIteration: - self.data_generator = iter(dataloader) - data = next(self.data_generator) + data = next(self.data_generator) for item in data: # submit data into input_queue self.submit(item, workflow=workflow) @@ -264,18 +260,13 @@ rollout = RemoteSGLangEngine(config.rollout) rollout.initialize() eval_rollout = ... -data_generator = iter(train_dataloader) +data_generator = iterools.cycle(train_dataloader) for global_step in range(max_steps): # rollout batched training data for current step 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 = rollout.rollout_batch(next(data_generator), workflow=workflow) ``` If you want to use rollout workflows with custom reward functions or agentic tool @@ -375,17 +366,12 @@ Now a complete GRPO training step in AReaLite is done! The core logic of our exa training script can be summarized as: ```python -data_generator = iter(train_dataloader) +data_generator = itertools.cycle(train_dataloader) for global_step in range(max_steps): 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 = rollout.rollout_batch(next(data_generator), workflow=workflow) logp = actor.compute_logp(batch) batch["prox_logp"] = logp diff --git a/docs/customization/agent.md b/docs/customization/agent.md index c71aeab050..781bee4739 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -77,7 +77,7 @@ and converting it into an `LLMRequest` object for the inference engine: class MultiTurnWorkflow(RolloutWorkflow): # ... __init__ method above ... - async def arun_episode(self, engine: InferenceEngine, data): + async def arun_episode(self, engine: InferenceEngine, data) -> TensorDict: # Initialize result containers seq, logprobs, loss_mask, versions = [], [], [], [] messages = data["messages"] @@ -121,7 +121,7 @@ apply a discount, add feedback to the conversation, and let the model try again: class MultiTurnWorkflow(RolloutWorkflow): # ... previous methods ... - async def arun_episode(self, engine: InferenceEngine, data): + async def arun_episode(self, engine: InferenceEngine, data) -> TensorDict: # ... initialization code ... while reward == 0 and t < self.max_turns: # Add feedback if the previous answer was incorrect @@ -192,7 +192,7 @@ Finally, let's complete the implementation by collecting trajectories in the class MultiTurnWorkflow(RolloutWorkflow): # ... previous methods ... - async def arun_episode(self, engine: InferenceEngine, data): + async def arun_episode(self, engine: InferenceEngine, data) -> TensorDict: # ... episode logic above ... while reward == 0 and t < self.max_turns: @@ -247,18 +247,13 @@ def main(args): ) # Run training—no other changes needed! - data_generator = iter(train_dataloader) + data_generator = itertools.cycle(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) + batch = rollout.rollout_batch(danext(data_generator)ta, workflow=workflow) # ... continue with training loop ... ``` diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index f087f83dbe..04a910031a 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -171,16 +171,11 @@ def main(args): ) # Main training loop + data_generator = itertools.cycle(dataloader) 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 = rollout.rollout_batch(next(data_generator), workflow=workflow) batch = batch.to(actor.device) diff --git a/examples/arealite/configs/gsm8k_grpo.yaml b/examples/arealite/configs/gsm8k_grpo.yaml index 6613784488..4ad2d9cd4a 100644 --- a/examples/arealite/configs/gsm8k_grpo.yaml +++ b/examples/arealite/configs/gsm8k_grpo.yaml @@ -68,6 +68,7 @@ ref: trial_name: ${trial_name} path: ${actor.path} init_from_scratch: false + disable_dropout: true dtype: ${actor.dtype} mb_spec: max_tokens_per_mb: 10240 diff --git a/examples/arealite/gsm8k_grpo.py b/examples/arealite/gsm8k_grpo.py index f44c0d35d0..dbb1ccea31 100644 --- a/examples/arealite/gsm8k_grpo.py +++ b/examples/arealite/gsm8k_grpo.py @@ -1,3 +1,4 @@ +import itertools import os import sys @@ -129,7 +130,7 @@ def main(args): 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) + data_generator = itertools.cycle(train_dataloader) for global_step in range(max_steps): epoch = global_step // steps_per_epoch step = global_step % steps_per_epoch @@ -138,12 +139,7 @@ def main(args): 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 = rollout.rollout_batch(next(data_generator), workflow=workflow) batch = batch.to(actor.device) # Create barrier to synchronize all rollout processes.