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 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] 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))