diff --git a/areal/api/alloc_mode.py b/areal/api/alloc_mode.py index 96f23d8047..0c0ece14aa 100644 --- a/areal/api/alloc_mode.py +++ b/areal/api/alloc_mode.py @@ -603,7 +603,8 @@ def gen_instance_size(self) -> int: inf_para: modern_inf_para modern_inf_para: INFER_BACKEND ("[" NAME "]")? ":" inf_dim+ - train_para: train_backend_with_name | train_backend_hybrid | train_backend_only | train_name_only | train_dims_only | hybrid_moe_syntax + train_para: train_backend_name_hybrid | train_backend_with_name | train_backend_hybrid | train_backend_only | train_name_only | train_dims_only | hybrid_moe_syntax + train_backend_name_hybrid: TRAIN_BACKEND "[" NAME "]" ":" hybrid_moe_syntax train_backend_with_name: TRAIN_BACKEND "[" NAME "]" ":" common_dim+ train_backend_hybrid: TRAIN_BACKEND ":" hybrid_moe_syntax train_backend_only: TRAIN_BACKEND ":" common_dim+ @@ -853,6 +854,19 @@ def train_backend_with_name(self, items): SchedulingStrategy(type=SchedulingStrategyType.separation, target=None), ) + def train_backend_name_hybrid(self, items): + """Handle: TRAIN_BACKEND [ NAME ] : hybrid_moe_syntax""" + backend = str(items[0]) + name = str(items[1]) + strategy = items[2] # ParallelStrategy from hybrid_moe_syntax + + return self._build_model_allocation( + backend, + name, + strategy, + SchedulingStrategy(type=SchedulingStrategyType.separation, target=None), + ) + def train_backend_hybrid(self, items): """Handle: TRAIN_BACKEND : hybrid_moe_syntax""" backend = str(items[0]) diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 82ad59c0dc..ff30f0a850 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -1237,7 +1237,11 @@ class TrainEngineConfig: weight_update_mode: str = field( default="xccl", - metadata={"help": "Weight update backend type.", "choices": ["disk", "xccl"]}, + metadata={ + "help": "Weight update backend type. 'awex' requires a Megatron actor " + "and an SGLang rollout, and targets colocated actor-rollout setups.", + "choices": ["disk", "xccl", "awex"], + }, ) fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) archon: ArchonEngineConfig = field(default_factory=ArchonEngineConfig) @@ -1997,6 +2001,7 @@ class SGLangConfig: triton_attention_reduce_in_fp32: bool = False triton_attention_num_kv_splits: int = 8 num_continuous_decode_steps: int = 1 + load_format: str = "auto" enable_memory_saver: bool = False allow_auto_truncate: bool = False attention_backend: str | None = "fa3" @@ -2098,7 +2103,6 @@ def build_args( # Model and tokenizer tokenizer_path=sglang_config.model_path, tokenizer_mode="auto", - load_format="auto", trust_remote_code=True, is_embedding=False, # Other runtime options diff --git a/areal/api/io_struct.py b/areal/api/io_struct.py index 5fa7a4a8ac..abb640787f 100644 --- a/areal/api/io_struct.py +++ b/areal/api/io_struct.py @@ -292,6 +292,7 @@ def from_fsdp_xccl( @classmethod def from_awex( cls, + meta_server_addr: str | None = None, use_lora: bool = False, lora_name: str = "", lora_int_id: int = 1, @@ -299,6 +300,7 @@ def from_awex( ): return cls( type="awex", + nccl_master_address=meta_server_addr, use_lora=use_lora, lora_name=lora_name, lora_int_id=lora_int_id, @@ -440,9 +442,25 @@ def log(self, head: str = "", rank: int = 0, precision: int = 2): mem_used = f"{self.mem_used:.{precision}f}" mem_total = f"{self.mem_total:.{precision}f}" if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank): + # Append host RssAnon/RssShmem (from /proc/self/status) to every + # device-stats log point so per-step host memory growth can be + # attributed to a phase boundary. Near-zero cost; rank0-only. + host_str = "" + try: + anon = shmem = 0 + with open("/proc/self/status") as f: + for line in f: + if line.startswith("RssAnon:"): + anon = int(line.split()[1]) // 1024 + elif line.startswith("RssShmem:"): + shmem = int(line.split()[1]) // 1024 + host_str = f" | host RssAnon: {anon}MB, RssShmem: {shmem}MB" + except Exception: + pass logger.info( f"Memory-Usage {head}: " f"memory allocated ({self.unit}): {mem_allocated}, " f"memory reserved ({self.unit}): {mem_reserved}, " f"device memory used/total ({self.unit}): {mem_used}/{mem_total}" + f"{host_str}" ) diff --git a/areal/engine/awex/__init__.py b/areal/engine/awex/__init__.py new file mode 100644 index 0000000000..84b818accf --- /dev/null +++ b/areal/engine/awex/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""AWEX colocated weight-transfer integration (writer, reader, SGLang plugin).""" diff --git a/areal/engine/awex/colocate_reader.py b/areal/engine/awex/colocate_reader.py new file mode 100644 index 0000000000..13025394d2 --- /dev/null +++ b/areal/engine/awex/colocate_reader.py @@ -0,0 +1,593 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""AWEX colocate weight reader (native awex worker-reader adapter). + +Runs inside the SGLang scheduler process. This is a thin shell around awex's +native ``NCCLWorkerWeightsReader`` that: + +1. Eager-registers the inference-side metadata the train writer waits for + (``infer_conf`` + ``num_infer_engines``), computed via awex's own + ``InferParamMetaResolver._get_model_param_info`` + ``_build_params_meta`` + (no hand-rolled name normalization or shard merging). +2. Lazily constructs the awex ``NCCLWorkerWeightsReader`` on the first weight + update (it needs ``training_params_meta``, which only appears after the + first training step) and delegates the whole IPC-collect + StreamBatch + transport + writer handshake to it. + +Why the awex-native reader instead of a hand-rolled receiver: the community +SGLang scheduler has no ``execute_task_in_model_worker`` driver layer, so we +build the awex *worker* reader directly in-process. The native worker reader +uses ``NcclColocateStreamBatchTransport`` (recursive partition), the transport +AWEX ships -- a hand-rolled ring-shift transport deadlocks on mismatched +train/infer pipeline layouts (e.g. train PP=4 vs infer PP=1). + +The plugin shell still owns the steps awex's *driver* would normally do +(``_pre_update_weights`` wait-for-offload + resume weights, ``_resume_kvcache`` +signal-finished); see ``awex_sglang_plugin.process_awex_queue``. +""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def _patch_tms_hook_mode() -> None: + """Make ``torch_memory_saver.hook_mode`` setter a no-op once initialized. + + ``megatron.core.inference.contexts.dynamic_context`` (pulled in transitively + by ``awex.converter.mcore_converter`` -> ``megatron.core``) runs a + module-level ``torch_memory_saver.hook_mode = "torch"``. In the SGLang + scheduler process the memory_saver singleton is already initialized (sglang + ran ``_ensure_initialized``, which ``del``s ``_impl_ctor_kwargs``), so that + late assignment raises ``AttributeError``. awex's model registry swallows the + import error, the BailingMoe converter never registers, and weight transfer + later dies with ``Unsupported attention parameter name: attention.g_proj``. + The singleton's own assert already declares post-init configuration + unsupported, so dropping the late set is the intended behavior. + """ + try: + import torch_memory_saver as _tms + except Exception: + return + inst = getattr(_tms, "torch_memory_saver", None) + if inst is None: + return + cls = type(inst) + prop = cls.hook_mode + if getattr(prop.fset, "_awex_safe", False): + return + + def _safe_setter(self, value): + if not hasattr(self, "_impl_ctor_kwargs"): + return # singleton already initialized; late set is a design no-op + prop.fset(self, value) + + _safe_setter._awex_safe = True + cls.hook_mode = property(prop.fget, _safe_setter) + + +# Must run before any awex import: awex.models.registry auto-imports model +# modules at module load, and the BailingMoe module's transitive megatron import +# trips the hook_mode race above. +_patch_tms_hook_mode() + +from awex.meta.infer_meta_resolver import InferParamMetaResolver # noqa: E402 +from awex.meta.meta_resolver import ParamMetaResolver # noqa: E402 +from awex.reader.nccl_reader import NCCLWorkerWeightsReader # noqa: E402 +from awex.sharding import get_sharding_strategy_builder # noqa: E402 +from awex.util.common import simple_hf_config # noqa: E402 + +from areal.utils.logging import getLogger # noqa: E402 + +logger = getLogger("AwexColocateReader") + + +def _ensure_awex_models_registered() -> None: + """Rebuild awex's model registry in case it cached a failed auto-import. + + ``import_model_configs`` is ``lru_cache``-d and ``ModelRegistry`` is built + once at module load. If anything imported the registry before our hook_mode + patch took effect, the BailingMoe converter would be silently missing. Clear + the cache and rebuild now that the patch is in place. + """ + try: + from awex.models import registry as _reg + + _reg.import_model_configs.cache_clear() + _reg.ModelRegistry.models = _reg.import_model_configs() + missing = [ + m + for m in ("BailingMoeV2_5ForCausalLM", "BailingMoeV2ForCausalLM") + if m not in _reg.ModelRegistry.models + ] + if missing: + logger.warning(f"awex model registry still missing converters: {missing}") + except Exception as e: # pragma: no cover - diagnostics only + logger.warning(f"Failed to rebuild awex model registry: {e}") + + +_ensure_awex_models_registered() + + +class _SingleInstanceMetaResolver(ParamMetaResolver): + """Aggregate per-rank raw meta of ONE inference instance into ParameterMeta. + + awex's ``InferParamMetaResolver`` normally drives this via + ``execute_task_in_model_worker`` (a driver fan-out we do not have). We + instead exchange the per-rank raw meta dicts through the MetaServer + (see ``_build_instance_params_meta``) and reuse awex's ``_build_params_meta`` + for the aggregation, plus awex's own sharding strategy builder for + ``_get_sharding_info``. This yields the exact same ``parameters_meta`` the + native reader expects, with awex converter parameter names (no hand-rolled + normalization). + """ + + def __init__(self, hf_config, engine_name, infer_engine_config, raw_meta_list): + super().__init__(hf_config) + self._raw_meta_list = raw_meta_list + rank0 = self._select_rank0(raw_meta_list) + self._model_arch_name = rank0["model_arch_name"] + self._sharding_strategy = get_sharding_strategy_builder(engine_name)( + self._model_arch_name, + infer_engine_config, + rank0["rank_info"], + ) + + @staticmethod + def _select_rank0(raw_meta_list): + for info in raw_meta_list: + if info["rank_info"].global_rank == 0: + return info + return raw_meta_list[0] + + def get_model_arch_name(self) -> str: + return self._model_arch_name + + def get_parameters_meta(self): + return self._build_params_meta() + + def _get_params_raw_meta(self): + return self._raw_meta_list + + def _get_sharding_info(self, name, rank_info, param_meta): + return self._sharding_strategy.get_sharding_strategy( + name, rank_info=rank_info, param_meta=param_meta + ) + + +class AwexColocateReader: + """Thin adapter binding awex's native worker reader into a SGLang scheduler.""" + + def __init__(self, scheduler: Any): + self._scheduler = scheduler + self._meta_server_client = None + self._reader: NCCLWorkerWeightsReader | None = None + self._released_tags: set[str] = set() + + self._transfer_rank: int | None = None + self._local_gpu_id: int | None = None + self._infer_world_size: int | None = None + self._train_world_size: int | None = None + self._meta_server_addr: str | None = None + + # External-instance decomposition (computed in initialize()). + self._infer_instance_world_size: int | None = None + self._num_infer_engines: int | None = None + self._engine_rank: int | None = None + self._instance_local_rank: int | None = None + + # Inference-side parameters_meta for ONE engine instance, computed via + # awex resolver + MetaServer raw-meta exchange. Reused as the native + # reader's ``parameters_meta`` constructor arg. + self._infer_params_meta = None + self._infer_conf: dict | None = None + self._initialized = False + + # ── model / context helpers ─────────────────────────────────────── + + def _get_model(self) -> torch.nn.Module: + return self._scheduler.tp_worker.model_runner.model + + def _build_model_context(self) -> dict[str, Any]: + """awex model_context describing ONE inference engine instance. + + ``world_size`` is the single-server tp*pp; ``global_rank`` is the + instance-local rank (= tp_rank for pp=1). The cross-server NCCL identity + (engine_rank / global transfer_rank) is tracked separately by the awex + reader. ``infer_engine_config`` (== server_args) is required by + ``WorkerWeightsReader.__init__`` and the backport's model_context omits + it, so we add it here. + """ + scheduler = self._scheduler + server_args = scheduler.server_args + tp_size = int(getattr(server_args, "tp_size", 1)) + pp_size = int(getattr(server_args, "pp_size", 1)) + dp_size = int(getattr(server_args, "dp_size", 1)) + tp_rank = int(getattr(scheduler, "tp_rank", 0)) + + if self._infer_instance_world_size is not None: + world_size = self._infer_instance_world_size + global_rank = self._instance_local_rank + else: + world_size = tp_size * pp_size + global_rank = tp_rank + + return { + "scheduler": scheduler, + "infer_engine_config": server_args, + "tp_rank": tp_rank, + "tp_size": tp_size, + "pp_rank": int(getattr(scheduler, "pp_rank", 0)), + "pp_size": pp_size, + "dp_size": dp_size, + "world_size": world_size, + "global_rank": global_rank, + "local_rank": tp_rank, + "attn_tp_rank": int(getattr(scheduler, "attn_tp_rank", tp_rank)), + "attn_tp_size": int(getattr(scheduler, "attn_tp_size", tp_size)), + "attn_dp_rank": int(getattr(scheduler, "attn_dp_rank", 0)), + } + + def get_parallelism(self) -> dict: + ctx = self._build_model_context() + server_args = self._scheduler.server_args + return { + "world_size": ctx["world_size"], + "tp_size": int(getattr(server_args, "tp_size", ctx["tp_size"])), + "pp_size": int(getattr(server_args, "pp_size", ctx["pp_size"])), + "dp_size": int(getattr(server_args, "dp_size", ctx["dp_size"])), + "ep_size": int(getattr(server_args, "ep_size", 1)), + "num_engines": self._num_infer_engines or 1, + } + + # ── metadata (awex-native, no hand-rolled normalization) ────────── + + def _compute_local_raw_meta(self) -> dict: + """Per-rank raw meta via awex's own staticmethod (HF-converted names).""" + server_args = self._scheduler.server_args + model_context = self._build_model_context() + return InferParamMetaResolver._get_model_param_info( + "sglang", + server_args, + convert_params=True, + engine_rank=self._engine_rank or 0, + model=self._get_model(), + model_context=model_context, + ) + + def _build_instance_params_meta(self): + """Gather single-instance raw meta via the MetaServer, then aggregate. + + Returns the awex ``parameters_meta`` (list[ParameterMeta]) for ONE + inference engine instance (the ``instance_world`` instance-local ranks). + + We exchange per-rank raw meta through the MetaServer instead of an + ``all_gather`` over ``tp_cpu_group``: that group is sglang's TP + request-broadcast group, driven by the scheduler MainThread's + ``recv_requests`` -> ``broadcast_pyobj``. This method runs on the + plugin's background thread, so a collective on the shared group races + the MainThread broadcast and deadlocks (two ops in flight on one + non-thread-safe group). The MetaServer exchange needs no process-group + collective, is isolated per engine instance by ``engine_rank``, and also + sidesteps the ``dist.new_group`` collective-ordering trap (train + infer + share the default world in colocate mode). + """ + local_raw = self._compute_local_raw_meta() + + instance_world = self._infer_instance_world_size or 1 + if instance_world > 1: + client = self._meta_server_client + prefix = f"infer_instance_raw_meta_{self._engine_rank}" + client.put_object(f"{prefix}_{self._instance_local_rank}", local_raw) + raw_meta_list = [ + client.get_object(f"{prefix}_{r}", timeout=300.0) + for r in range(instance_world) + ] + else: + raw_meta_list = [local_raw] + + # MetaServer serializes RankInfo to a dict on the wire (as did the + # legacy all_gather); rebuild the object before awex's resolver reads it. + from awex.sharding.rank_info import RankInfo + + for info in raw_meta_list: + ri = info.get("rank_info") + if isinstance(ri, dict): + info["rank_info"] = RankInfo(**ri) + + resolver = _SingleInstanceMetaResolver( + self._get_model().config, + "sglang", + self._scheduler.server_args, + raw_meta_list, + ) + return resolver.get_parameters_meta() + + def get_weight_metadata(self): + """Inference-side parameters_meta for ONE engine instance.""" + if self._engine_rank is None: + raise RuntimeError( + "AwexColocateReader must be initialized before getting weight metadata" + ) + if self._infer_params_meta is None: + self._infer_params_meta = self._build_instance_params_meta() + return self._infer_params_meta + + # ── eager init: register infer_conf + num_infer_engines ─────────── + + def initialize( + self, + meta_server_addr: str, + transfer_rank: int, + infer_world_size: int, + train_world_size: int, + local_gpu_id: int, + timeout_s: float = 300.0, + ) -> None: + """Eager init: publish the metadata the train writer waits for. + + Must NOT block on the training side (runs before the first training step + finishes). The native ``NCCLWorkerWeightsReader`` is built lazily in + ``update_weights`` once ``training_params_meta`` is available. Device + entry registration (``inference_device_rank_entries``) is left to the + native reader's ``_init_reader_in_colocate_mode``. + """ + from awex.meta.meta_server import MetaServerClient + + if infer_world_size != train_world_size: + raise ValueError( + f"Colocate mode requires equal total rank counts " + f"(same physical GPUs), got infer={infer_world_size} " + f"vs train={train_world_size}" + ) + + self._transfer_rank = transfer_rank + self._local_gpu_id = local_gpu_id + self._infer_world_size = infer_world_size + self._train_world_size = train_world_size + self._meta_server_addr = meta_server_addr + + server_args = self._scheduler.server_args + tp_size = int(getattr(server_args, "tp_size", 1)) + pp_size = int(getattr(server_args, "pp_size", 1)) + instance_world = max(1, tp_size * pp_size) + if infer_world_size % instance_world != 0: + raise ValueError( + f"infer_world_size ({infer_world_size}) must be divisible by the " + f"per-instance world tp*pp ({instance_world})" + ) + self._infer_instance_world_size = instance_world + self._num_infer_engines = infer_world_size // instance_world + self._engine_rank = transfer_rank // instance_world + self._instance_local_rank = transfer_rank % instance_world + logger.info( + "AWEX instance decomposition: transfer_rank=%d -> engine_rank=%d, " + "instance_local_rank=%d (instance_world=%d, num_engines=%d)", + transfer_rank, + self._engine_rank, + self._instance_local_rank, + instance_world, + self._num_infer_engines, + ) + + host, port = meta_server_addr.rsplit(":", 1) + self._meta_server_client = MetaServerClient(host, int(port)) + + # Compute single-instance parameters_meta (also reused as the native + # reader's constructor arg later). + self.get_weight_metadata() + + par = self.get_parallelism() + infer_conf = { + "engine_name": "sglang", + "infer_atten_tp_size": par["tp_size"], + "infer_world_size": infer_world_size, + "hf_config": simple_hf_config(self._get_model().config), + # AWEX's native reader publishes router_dtype so the train-side + # converter casts mlp.gate.weight to the dtype the inference + # engine actually holds (fp32 for BailingMoe). Omitting it makes + # the converter fall back to its bf16 default: gate shards go out + # as 2N bytes against a 4N irecv and the transfer wedges + # deterministically. The wire-level dtype reconciliation below + # papers over any such mismatch generically, but keep the + # semantic path whole so new models behave identically to native + # awex. + "router_dtype": getattr(self._get_model().config, "router_dtype", "bf16"), + } + self._infer_conf = infer_conf + + # Only one rank publishes the engine-instance-wide info the writer waits + # for. transfer_rank 0 is engine_rank 0, instance_local_rank 0. + if transfer_rank == 0: + self._meta_server_client.put_object("infer_conf", infer_conf) + self._meta_server_client.put_object( + "num_infer_engines", self._num_infer_engines + ) + logger.info( + "Registered infer_conf + num_infer_engines=%d with MetaServer", + self._num_infer_engines, + ) + + self._initialized = True + logger.info( + "Eager init done: transfer_rank=%d, local_gpu_id=%d, infer_world_size=%d " + "(native worker reader construction deferred to first update_weights)", + transfer_rank, + local_gpu_id, + infer_world_size, + ) + + # ── lazy native-reader construction + weight update ─────────────── + + def _ensure_reader(self) -> NCCLWorkerWeightsReader: + if self._reader is not None: + return self._reader + + client = self._meta_server_client + training_params_meta = client.get_object( + "training_params_meta", timeout=10000.0 + ) + logger.info("Got training_params_meta from MetaServer") + + model_context = self._build_model_context() + reader = NCCLWorkerWeightsReader( + engine_name="sglang", + model=self._get_model(), + model_context=model_context, + infer_conf=self._infer_conf, + engine_rank=self._engine_rank, + num_engines=self._num_infer_engines, + meta_server_addr=self._meta_server_addr, + parameters_meta=self._infer_params_meta, + training_params_meta=training_params_meta, + enable_colocate_mode=True, + ipc_backend="cuda", + enable_debug_mode=False, + ) + reader.initialize() + self._reader = reader + logger.info( + "Constructed native NCCLWorkerWeightsReader (transfer_rank=%d, " + "engine_rank=%d, num_engines=%d)", + reader.transfer_rank, + self._engine_rank, + self._num_infer_engines, + ) + return reader + + def update_weights(self, version: int) -> None: + """Run one colocate weight update via the native awex worker reader. + + The native reader internally does: IPC collect -> StreamBatch transport + -> put ``weights_update_finished`` -> barrier -> get_then_delete + ``write_finished`` -> flush_cache. The plugin only needs to wrap this + with the driver-equivalent wait-for-offload + resume + signal steps. + """ + if not self._initialized: + raise RuntimeError("AwexColocateReader not initialized") + reader = self._ensure_reader() + reader.update_weights(step_id=version) + self._rebuild_derived_weights() + logger.info("Colocate weight update completed: version=%d", version) + + def _rebuild_derived_weights(self) -> None: + """Re-derive non-parameter tensors after an in-place AWEX weight write. + + Root cause: sglang's ``load_model`` ends with + ``post_load_weights()``, which splits each MLA layer's + ``kv_b_proj.weight`` into the absorbed-path tensors ``w_kc``/``w_vc`` + — ``.contiguous()`` copies stored as plain attributes, in neither + ``named_parameters`` nor ``named_buffers``. The memory-saver + release/resume cycle remaps their pages to zeros, and the AWEX reader + rewrites only named parameters via in-place ``copy_`` (bypassing + ``model.load_weights``), so nothing ever rebuilds them: decode's + forward_absorb then consumes zeros and the 4 MLA layers degenerate + while the 28 Lightning layers stay healthy (reward 0.77 -> ~0 within + 5 steps). Rebuild after EVERY transfer — train weights move each + version, so a one-time fix would go stale. ``bind_or_assign`` copies + into the existing tensors in place, which keeps captured CUDA-graph + addresses valid. + """ + model = self._get_model() + fn = getattr(model, "post_load_weights", None) + if fn is None: + return + fn() + torch.cuda.synchronize() + logger.info("post_load_weights() re-derived absorbed MLA weights") + + # ── memory release/resume (delegate to SGLang native) ───────────── + + def release_memory(self, tags: list[str] | None = None) -> None: + from sglang.srt.managers.io_struct import ReleaseMemoryOccupationReqInput + + tags = tags or ["kv_cache"] + native_tags = [t for t in tags if t not in self._released_tags] + if native_tags: + req = ReleaseMemoryOccupationReqInput(tags=native_tags) + self._scheduler.release_memory_occupation(req) + self._released_tags.update(native_tags) + logger.info("release_memory: tags=%s", tags) + + def resume_memory(self, tags: list[str] | None = None) -> None: + from sglang.srt.managers.io_struct import ResumeMemoryOccupationReqInput + + tags = tags or ["kv_cache"] + resume_tags = [t for t in tags if t in self._released_tags] + if resume_tags: + req = ResumeMemoryOccupationReqInput(tags=resume_tags) + self._scheduler.resume_memory_occupation(req) + self._released_tags.difference_update(resume_tags) + logger.info("resume_memory: tags=%s", tags) + + # ── writer-coordination handshake (driver-equivalent shell steps) ── + + def wait_for_training_offloaded(self, version: int) -> None: + """Wait for the writer to offload its model weights (avoid 2x weights). + + Equivalent to awex driver ``_pre_update_weights``'s wait on + ``all_training_offloaded_weights``. + """ + from areal.engine.awex.colocate_writer import awex_colocate_timeout_s + + self._meta_server_client.wait_set_until_size( + "all_training_offloaded_weights", + self._train_world_size, + timeout=awex_colocate_timeout_s(), + ) + + def wait_for_weights_ready( + self, version: int, timeout_s: float | None = None + ) -> None: + """Block until the writer has published THIS version's IPC handles. + + Used by the plugin's background thread as the per-version trigger to + enqueue a weight-update marker. We probe the per-version + ``training_serialized_weights_{ip}_{gpu}_{version}`` key with MetaServer + ``wait_key`` (existence-only, NO deserialization), for two reasons: + + 1. Per-version gating. The unversioned ``all_training_offloaded_weights`` + set is only deleted by the writer's rank0 in ``finish_colocate_weight_update`` + (a later phase than the engine's signal_finished), so gating on it + lets the background thread fire v+1 off a *stale* satisfied set while + the writer is still in v's finish phase. The collected v+1 IPC then + blocks waiting for a not-yet-published key, hogging the scheduler main + loop so it cannot serve rollout -> train waits on rollout -> deadlock. + The writer only puts the v+1 serialized key in the NEXT training cycle, + so gating on it cannot fire early. + 2. No double-attach. ``get_object`` would deserialize the CUDA IPC handle + in the background thread, racing the worker reader's own collect inside + update_weights. ``wait_key`` only checks presence (``_has_key``). + """ + from awex.util.common import get_ip_address + + from areal.engine.awex.colocate_writer import awex_colocate_timeout_s + + ip = get_ip_address() + key = f"training_serialized_weights_{ip}_{self._local_gpu_id}_{version}" + self._meta_server_client.wait_key( + key, + timeout=awex_colocate_timeout_s() if timeout_s is None else timeout_s, + ) + + def signal_finished_weights_update(self) -> None: + """Signal this engine finished, so the writer can resume kv_cache. + + Equivalent to awex driver ``_resume_kvcache``'s add to + ``finished_weights_update_engines``. Only one rank per engine instance + (instance_local_rank == 0) signals, with its real engine_rank, so the + set collects exactly num_infer_engines unique entries. + """ + if self._instance_local_rank != 0: + return + self._meta_server_client.add_object_to_set( + "finished_weights_update_engines", self._engine_rank + ) + + def teardown(self) -> None: + self._reader = None + + +__all__ = ["AwexColocateReader"] diff --git a/areal/engine/awex/colocate_writer.py b/areal/engine/awex/colocate_writer.py new file mode 100644 index 0000000000..0a32c8c07c --- /dev/null +++ b/areal/engine/awex/colocate_writer.py @@ -0,0 +1,788 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Licensed under the Apache License, Version 2.0 +"""AWEX colocate adapter for MegatronEngine (training side). + +Provides: +- Manual GPU→CPU offload for model weights and optimizer states +- CUDA IPC weight transfer to colocated SGLang (same GPU, via MetaServer) +- Coordinates with SGLang inference via MetaServer signals + +Weight transfer flow (mirrors the AWEX reference nccl_writer colocate mode): + 1. Convert Megatron params → HF format + 2. Group tensors by shape/dtype → share_memory_() → cuda_ipc_serialize + 3. Put serialized IPC handles to MetaServer + 4. Infer side (same GPU) deserializes via CUDA IPC (zero-copy) + 5. Infer-only NCCL group handles redistribution among infer ranks + 6. Infer signals done → train cleans up shared tensors + +This adapter is used when weight_update_type == "awex" in colocate mode. +""" + +from __future__ import annotations + +import gc +import os +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist + +if TYPE_CHECKING: + from areal.engine.megatron_engine import MegatronEngine + +from areal.utils.logging import getLogger + +logger = getLogger("AwexColocate") + + +def resolve_physical_gpu_id(relative_gpu_id: int) -> int: + """Map a CUDA-masked relative device index to its physical GPU id. + + CUDA IPC keys must be unique per node, so both sides of a colocated + transfer have to agree on physical GPU ids. Inside a process that was + given a device mask, ``torch.cuda.current_device()`` and SGLang's + ``gpu_id`` are indices into that mask rather than physical ids, so the + mask itself is the only ground truth. Falls back to the relative index + when the mask is absent or holds GPU UUIDs. + """ + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not cuda_visible: + return relative_gpu_id + try: + gpu_ids = [int(x) for x in cuda_visible.split(",") if x.strip()] + return gpu_ids[relative_gpu_id] + except (ValueError, IndexError): + return relative_gpu_id + + +def awex_colocate_timeout_s(default: float = 1800.0) -> float: + value = os.environ.get("AWEX_COLOCATE_TIMEOUT_S", "").strip() + if not value: + return default + try: + return float(value) + except ValueError: + logger.warning( + "Invalid AWEX_COLOCATE_TIMEOUT_S=%r; using default %.1fs", + value, + default, + ) + return default + + +class AwexMegatronAdapter: + """Training-side adapter for AWEX colocated weight transfer. + + Uses CUDA IPC (share_memory + ForkingPickler serialization) for zero-copy + weight transfer to the colocated SGLang process on the same GPU. The infer + side handles redistribution among infer ranks via its own NCCL group. + """ + + def __init__(self, engine: MegatronEngine): + self._engine = engine + self._offloaded_weights: dict[str, torch.Tensor] = {} + self._released_tags: set[str] = set() + self._meta_server_addr: str | None = None + self._meta_server_client = None + self._transfer_rank: int | None = None + self._weight_converter = None + self._initialized = False + self._rank_info = None + self._ip_address: str | None = None + self._infer_world_size: int | None = None + self._num_infer_engines: int | None = None + self._logical_train_rank: int | None = None + + def init_colocate_weight_update( + self, + meta_server_addr: str | None = None, + pair_name: str = "default", + transfer_rank: int = 0, + timeout_s: float | None = None, + ) -> None: + """Initialize MetaServer connection. NCCL group creation is deferred + to the first weight update (lazy init) to allow SGLang to start first. + """ + from awex.meta.meta_server import MetaServerClient, start_meta_server + + if not meta_server_addr: + meta_server_addr = os.environ.get("AWEX_META_SERVER_ADDR", "") + if not meta_server_addr: + host, port = start_meta_server() + meta_server_addr = f"{host}:{port}" + os.environ["AWEX_META_SERVER_ADDR"] = meta_server_addr + logger.info("Started MetaServer at %s", meta_server_addr) + + host, port = meta_server_addr.rsplit(":", 1) + self._meta_server_client = MetaServerClient(host, int(port)) + self._meta_server_addr = meta_server_addr + self._transfer_rank = transfer_rank + # Train-side wait budget for infer's weights_update_finished signal. + # Keep the writer/reader/plugin on one env-controlled timeout to avoid + # split-brain diagnostics. + self._timeout_s = awex_colocate_timeout_s() if timeout_s is None else timeout_s + if dist.get_rank() == 0: + self._meta_server_client.put_object( + "awex_train_info", {"train_world_size": dist.get_world_size()} + ) + logger.info( + "Registered awex_train_info (train_world_size=%d) with MetaServer", + dist.get_world_size(), + ) + + logger.info( + "AwexMegatronAdapter initialized: meta_server=%s, transfer_rank=%d", + meta_server_addr, + transfer_rank, + ) + + def _lazy_initialize(self) -> None: + """Perform deferred initialization: metadata exchange and weight converter setup. + + In colocate mode, train side does NOT join any NCCL group. + Weight transfer uses CUDA IPC (share_memory + serialize via MetaServer). + The infer side creates its own infer-only NCCL group for redistribution. + """ + if self._initialized: + return + + from awex.models.registry import get_train_weights_converter + from awex.sharding.param_sharding import get_rank_info_extractor + from awex.util.common import get_ip_address + + rank = dist.get_rank() + + self._rank_info = get_rank_info_extractor("mcore")() + training_world_size = self._rank_info.world_size + self._ip_address = get_ip_address() + self._physical_gpu_id = resolve_physical_gpu_id(torch.cuda.current_device()) + + from awex.meta.train_meta_resolver import McoreParamMetaResolver + + class _EngineShim: + def __init__(self, engine): + self.model = engine.model + if not isinstance(self.model, (list, tuple)): + self.model = [self.model] + self.hf_config = engine.hf_config + self.enable_debug_mode = False + self.enable_colocate_mode = False + self.engine_name = "mcore" + self.config = {} + self.meta_server_addr = "" + + def release_memory_occupation(self, tags=None): + pass + + def resume_memory_occupation(self, tags=None): + pass + + shim = _EngineShim(self._engine) + + infer_conf = self._meta_server_client.get_object( + "infer_conf", timeout=self._timeout_s + ) + logger.info("Got infer_conf from MetaServer: %s", infer_conf) + + if isinstance(infer_conf.get("hf_config"), dict): + from types import SimpleNamespace + + infer_conf["hf_config"] = SimpleNamespace(**infer_conf["hf_config"]) + + meta_resolver = McoreParamMetaResolver(shim, self._engine.hf_config, infer_conf) + parameters_meta = meta_resolver.get_parameters_meta() + logger.info( + "Collected training parameters metadata: %d params", len(parameters_meta) + ) + + if rank == 0: + self._meta_server_client.put_object("training_params_meta", parameters_meta) + logger.info("Registered training_params_meta with MetaServer") + + self._infer_world_size = infer_conf["infer_world_size"] + self._logical_train_rank = self._infer_world_size + self._rank_info.global_rank + + # Register physical device entry for (ip, node_local_gpu_id) -> rank + # pairing on the infer side (AWEX reader._init_reader_in_colocate_mode). + # device_id must be the node-local physical GPU id (matching the infer + # side and the CUDA IPC key), NOT a global rank. CUDA_VISIBLE_DEVICES is + # the ground truth since torch.cuda.current_device() is always 0 here. + self._meta_server_client.add_object_to_set( + "training_device_rank_entries", + (self._ip_address, self._physical_gpu_id, self._logical_train_rank), + ) + logger.info( + "Registered training_device_rank_entries: (ip=%s, gpu=%d, rank=%d)", + self._ip_address, + self._physical_gpu_id, + self._logical_train_rank, + ) + self._num_infer_engines = self._meta_server_client.get_object( + "num_infer_engines", timeout=self._timeout_s + ) + logger.info("Got num_infer_engines=%d from MetaServer", self._num_infer_engines) + + self._weight_converter = get_train_weights_converter( + "mcore", + self._engine.hf_config.architectures[0], + self._engine.hf_config, + self._rank_info, + { + **infer_conf, + "train_pp_stage_layer_id_map": ( + meta_resolver.get_pp_stage_layer_id_map() + ), + }, + tf_config=_get_tf_config(self._engine.model), + ) + + self._initialized = True + logger.info( + "Colocate train side initialized: logical_train_rank=%d, " + "infer_world_size=%d, train_world_size=%d", + self._logical_train_rank, + self._infer_world_size, + training_world_size, + ) + + def _release_grad_memory(self) -> None: + """Release gradient buffers to free GPU memory before weight conversion. + + Mirrors the AWEX reference release_grad_memory(). + Saves original sizes to buffer.grad_data_size for later restoration. + """ + from megatron.core.distributed import DistributedDataParallel as DDP + + model = self._engine.model + if model is None: + return + if not isinstance(model, (list, tuple)): + model = [model] + count = 0 + for chunk in model: + if isinstance(chunk, DDP): + for buffers in [chunk.buffers, chunk.expert_parallel_buffers]: + for buf in buffers: + if buf.grad_data.storage().size() > 0: + buf.grad_data_size = buf.grad_data.storage().size() + buf.grad_data.storage().resize_(0) + count += 1 + if count > 0: + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + logger.info("Released %d grad buffers", count) + + @torch.no_grad() + def execute_colocate_weight_update(self, version: int) -> None: + """Send weights to colocated inference via CUDA IPC through MetaServer. + + Flow (mirrors AWEX nccl_writer._write_weights_in_colocate_mode): + 1. Release optimizer states + grad memory to free GPU space + 2. Convert Megatron params to HF format + 3. Group tensors by shape/dtype → release originals → offload weights + 4. Signal all_training_offloaded_weights (reader waits for this) + 5. share_memory_() + cuda_ipc_serialize → put to MetaServer + 6. Wait for weights_update_finished (reader done copying) + 7. Clean up shared tensors → signal write_finished + 8. Wait for all infer engines to finish (finished_weights_update_engines) + """ + from awex.util.tensor_util import ( + cuda_ipc_serialize, + group_tensors_by_shape_and_dtype, + release_tensors, + ) + + weights_were_offloaded = "weights" in self._released_tags + + # Reclaim any IPC-exported blocks from the previous version whose + # peer mappings closed after our last collect (belt-and-braces with the + # collect in this method's finally block). + torch.cuda.ipc_collect() + + # Free optimizer states + grad buffers BEFORE reloading the train + # weights, not after. The colocated inference engine has already + # resumed its full PP=1 weight set on this same physical GPU, so + # reloading the train param buffers (resize_ allocates GiB-scale + # storage per buffer) OOMs on the very first buffer unless optimizer + + # grad memory is freed first. This matches the AWEX weights_writer + # order (release_memory(optimizer) THEN resume_memory(weights), see + # _release_memory_for_weights_exchange). + # Optimizer/grad offload operate on independent Megatron buffers and do + # not require the weights to be resumed, so reordering is safe. + self.release_memory(tags=["optimizer"]) + + self._release_grad_memory() + + if weights_were_offloaded: + self.resume_memory(tags=["weights"]) + + # _lazy_initialize AFTER the weights resume — its meta resolver + # runs convert_param over live params, which dies with CUDA invalid + # argument on resize_(0)-ed storages. The recover path is the only + # one that reaches the first transfer with weights offloaded (see + # re-releases them right after the recover load), which is why the + # normal step-1 path never hit this. + self._lazy_initialize() + + parameters = self._convert_parameters() + tensors = list(parameters.values()) + names = list(parameters.keys()) + logger.info( + "Converted %d params for colocate IPC transfer (version=%d)", + len(tensors), + version, + ) + + group_tensors, metadata = group_tensors_by_shape_and_dtype(tensors) + torch.cuda.synchronize() + logger.info( + "Grouped into %d tensor groups for IPC serialization", len(group_tensors) + ) + + # convert_param returns the ORIGINAL tensor (shared storage) + # whenever the conversion is an identity — direct_name_mapping + # (embed_tokens / final_layernorm) returns `parameter` as-is, and + # `.to(dtype)` (gate.weight, expert_bias) is a no-op when the dtype + # already matches. release_tensors() does untyped_storage().resize_(0), + # so releasing those entries frees the LIVE module storage. Params are + # later rebuilt by the weights offload/reload round-trip, but buffers + # (e.g. the fp32 router expert_bias, a register_buffer) are not part of + # offload/reload and stay dangling forever -> the post-transfer IMA in + # the router/EP path. Only release tensors we actually own (copies). + live_storages = set() + model = self._engine.model + for chunk in model if isinstance(model, (list, tuple)) else [model]: + # NB: Megatron DDP shadows nn.Module.buffers with a LIST attribute + # (ParamAndGradBuffer); named_parameters/named_buffers stay methods. + for _, p in chunk.named_parameters(): + live_storages.add(p.untyped_storage().data_ptr()) + for _, b in chunk.named_buffers(): + live_storages.add(b.untyped_storage().data_ptr()) + owned = [ + t for t in tensors if t.untyped_storage().data_ptr() not in live_storages + ] + logger.info( + "Releasing %d/%d converted tensors (%d alias live module storage, " + "left to the weights offload path)", + len(owned), + len(tensors), + len(tensors) - len(owned), + ) + release_tensors(owned) + del tensors, owned + parameters.clear() + + self.release_memory(tags=["weights"]) + + ip_address = self._ip_address + device_id = self._physical_gpu_id + key_suffix = f"_{ip_address}_{device_id}_{version}" + + self._meta_server_client.add_object_to_set( + "all_training_offloaded_weights", self._logical_train_rank + ) + logger.info( + "Signaled all_training_offloaded_weights (rank=%d)", + self._logical_train_rank, + ) + + group_shared = [t.share_memory_() for t in group_tensors] + serialized_weights = cuda_ipc_serialize((group_shared, metadata, names)) + torch.cuda.synchronize() + logger.info("CUDA IPC serialization complete, putting to MetaServer") + + serialized_weights_key = f"training_serialized_weights{key_suffix}" + # Tell the reader which version we are publishing. The plugin's + # background worker used to assume the stream starts at v1, which + # deadlocks recover runs (writer resumes at v=global_step, e.g. 9). + writer_version_key = f"awex_writer_version_{ip_address}_{device_id}" + self._meta_server_client.put_object(writer_version_key, version) + logger.info( + "Put writer version to MetaServer: key=%s version=%s", + writer_version_key, + version, + ) + self._meta_server_client.put_object( + serialized_weights_key, + (self._logical_train_rank, self._rank_info, serialized_weights), + ) + logger.info("Put IPC weights to MetaServer: key=%s", serialized_weights_key) + + update_finished_key = f"weights_update_finished{key_suffix}" + try: + try: + self._meta_server_client.get_object( + update_finished_key, timeout=self._timeout_s + ) + except Exception: + logger.error( + "Timed out or failed after %ss waiting for the inference " + "side to consume published weights (key=%s). The reader " + "likely died or never entered the colocate update; the " + "transfer cannot complete.", + self._timeout_s, + update_finished_key, + ) + raise + self._meta_server_client.delete_if_exists(update_finished_key) + self._meta_server_client.delete_if_exists(serialized_weights_key) + logger.info("Got done signal from infer side: %s", update_finished_key) + finally: + release_tensors(group_tensors) + release_tensors(group_shared) + del group_tensors, group_shared + torch.cuda.synchronize() + gc.collect() + # Storages exported via cudaIpcGetMemHandle park in PyTorch's + # CudaIPCSentDataLimbo when freed and are NOT returned to the + # allocator until ipc_collect() confirms the peer closed its + # mapping. GiB-scale group tensors are exported every version; + # without collection the train-process residual grows by GBs per + # version, eating the colocated rollout's prefill headroom until + # its logits all-gather OOMs. + torch.cuda.ipc_collect() + torch.cuda.empty_cache() + + write_finished_key = f"write_finished{key_suffix}" + self._meta_server_client.put_object(write_finished_key, True) + logger.info("Signaled write_finished: %s", write_finished_key) + + logger.info("Colocate weight update completed: version=%d", version) + + def finish_colocate_weight_update(self, training_world_size: int) -> None: + """Wait for all inference engines to finish weight update, then clean up. + + Mirrors AWEX _finish_weights_update(). + Called from megatron_engine.update_weights() after barrier. + """ + num_infer_engines = self._num_infer_engines + logger.info( + "Waiting for %d inference engine(s) to signal finished_weights_update_engines", + num_infer_engines, + ) + self._meta_server_client.wait_set_until_size( + "finished_weights_update_engines", + num_infer_engines, + timeout=self._timeout_s, + ) + logger.info("All inference engines finished weights update") + + dist.barrier(group=self._engine.cpu_group) + + if dist.get_rank() == 0: + self._meta_server_client.delete_if_exists("finished_weights_update_engines") + self._meta_server_client.delete_if_exists("all_training_offloaded_weights") + logger.info("Cleaned up MetaServer coordination keys") + + @torch.no_grad() + def _convert_parameters(self) -> dict[str, torch.Tensor]: + """Convert Megatron parameters to HF format for IPC transfer.""" + from awex.converter.mcore_converter import get_mcore_model_parameters + + model = self._engine.model + if not isinstance(model, (list, tuple)): + model = [model] + + converted = {} + for vp_stage, m in enumerate(model): + for name, param in get_mcore_model_parameters(m).items(): + for hf_name, hf_param in self._weight_converter.convert_param( + name, param.detach(), vp_stage=vp_stage + ): + converted[hf_name] = hf_param + + hf_config = self._engine.hf_config + if ( + getattr(hf_config, "tie_word_embeddings", False) + and self._rank_info.pp_rank == self._rank_info.pp_size - 1 + and "lm_head.weight" not in converted + and "model.embed_tokens.weight" in converted + ): + converted["lm_head.weight"] = converted["model.embed_tokens.weight"] + + logger.info("Converted %d parameters for IPC transfer", len(converted)) + return converted + + # ── Memory management (manual offload) ──────────────────────────────── + + def release_memory(self, tags: list[str] | None = None) -> None: + tags = tags or ["optimizer", "weights"] + tags_to_release = [t for t in tags if t not in self._released_tags] + if not tags_to_release: + return + + if "optimizer" in tags_to_release: + self._offload_optimizer_states() + self._released_tags.add("optimizer") + + if "weights" in tags_to_release: + self._offload_model_weights() + self._released_tags.add("weights") + + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + logger.info("release_memory done: tags=%s", tags_to_release) + + def resume_memory(self, tags: list[str] | None = None) -> None: + tags = tags or ["optimizer", "weights"] + tags_to_resume = [t for t in tags if t in self._released_tags] + if not tags_to_resume: + return + + if "weights" in tags_to_resume: + self._reload_model_weights(load_grad=False) + self._released_tags.discard("weights") + + if "optimizer" in tags_to_resume: + self._reload_optimizer_states() + self._released_tags.discard("optimizer") + + torch.cuda.synchronize() + logger.info("resume_memory done: tags=%s", tags_to_resume) + + def _offload_model_weights(self) -> None: + from megatron.core.distributed import DistributedDataParallel as DDP + + model = self._engine.model + if model is None: + return + if not isinstance(model, (list, tuple)): + model = [model] + count = 0 + for chunk in model: + if isinstance(chunk, DDP): + for buffers in [chunk.buffers, chunk.expert_parallel_buffers]: + for buf in buffers: + if hasattr(buf, "offload_to_cpu"): + buf.offload_to_cpu() + count += 1 + continue + if buf.param_data.storage().size() > 0: + if not hasattr(buf.param_data, "cpu_data"): + buf.param_data.cpu_data = torch.zeros( + buf.param_data.data.shape, + dtype=buf.param_data.data.dtype, + pin_memory=True, + device="cpu", + ) + buf.param_data.cpu_data.copy_(buf.param_data.data) + buf.param_data_size = buf.param_data.storage().size() + buf.param_data.storage().resize_(0) + count += 1 + if buf.grad_data.storage().size() > 0: + buf.grad_data_size = buf.grad_data.storage().size() + buf.grad_data.storage().resize_(0) + else: + for name, param in chunk.named_parameters(): + if param.data.is_cuda: + self._offloaded_weights[name] = param.data.detach().to( + "cpu", non_blocking=True + ) + param.data = torch.empty(0, device="cpu") + count += 1 + torch.cuda.synchronize() + logger.info("Offloaded %d weight buffers to CPU", count) + + def _reload_model_weights(self, load_grad: bool = False) -> None: + from megatron.core.distributed import DistributedDataParallel as DDP + + model = self._engine.model + if model is None: + return + if not isinstance(model, (list, tuple)): + model = [model] + device = self._engine.device + for chunk in model: + if isinstance(chunk, DDP): + for buffers in [chunk.buffers, chunk.expert_parallel_buffers]: + for buf in buffers: + if hasattr(buf, "reload_from_cpu"): + buf.reload_from_cpu(move_grads=load_grad) + continue + if buf.param_data.storage().size() == 0: + buf.param_data.storage().resize_(buf.param_data_size) + buf.param_data.copy_(buf.param_data.cpu_data, non_blocking=True) + if ( + load_grad + and hasattr(buf, "grad_data_size") + and buf.grad_data.storage().size() == 0 + ): + buf.grad_data.storage().resize_(buf.grad_data_size) + buf.grad_data.zero_() + else: + for name, param in chunk.named_parameters(): + if name in self._offloaded_weights: + param.data = self._offloaded_weights[name].to( + device, non_blocking=True + ) + self._offloaded_weights.clear() + torch.cuda.synchronize() + logger.info("Reloaded model weights to GPU (load_grad=%s)", load_grad) + + def ensure_grad_buffers(self) -> None: + """Allocate grad buffers if they were freed during offload. + + Called before forward_backward (training) to ensure grad storage + is available for backward pass. Separate from _reload_model_weights + because compute_logp (inference-only) should not allocate grad buffers. + """ + from megatron.core.distributed import DistributedDataParallel as DDP + + model = self._engine.model + if model is None: + return + if not isinstance(model, (list, tuple)): + model = [model] + count = 0 + for chunk in model: + if isinstance(chunk, DDP): + for buffers in [chunk.buffers, chunk.expert_parallel_buffers]: + for buf in buffers: + if ( + hasattr(buf, "grad_data_size") + and buf.grad_data.storage().size() == 0 + ): + buf.grad_data.storage().resize_(buf.grad_data_size) + buf.grad_data.zero_() + count += 1 + if count > 0: + torch.cuda.synchronize() + logger.info("Allocated %d grad buffers for training", count) + + def _get_inner_optimizers(self): + optimizer = self._engine.optimizer + if optimizer is None: + return [] + if hasattr(optimizer, "chained_optimizers"): + inner_optimizers = optimizer.chained_optimizers + elif hasattr(optimizer, "optimizers"): + inner_optimizers = optimizer.optimizers + else: + inner_optimizers = [optimizer] + return inner_optimizers + + def _offload_optimizer_states(self) -> None: + optimizer = self._engine.optimizer + if optimizer is None: + return + # Default path mirrors the AWEX reference optimizer offload + # (megatron_util.offload_megatron_optimizer): swap .data / state-dict + # references to CPU, never resize_ storages, then purge TE's global + # _dummy_wgrads cache and synchronize. Megatron HybridDeviceOptimizer's + # offload_to_cpu/restore_from_cpu is kept only as an opt-in fallback — + # its internal pointer bookkeeping is hard to validate and the AWEX + # reference integration deliberately avoids it. + if os.environ.get("AWEX_OPT_OFFLOAD_VIA_HDO", "").strip() == "1" and hasattr( + optimizer, "offload_to_cpu" + ): + optimizer.offload_to_cpu() + logger.info("Offloaded optimizer via offload_to_cpu()") + return + + inner_optimizers = self._get_inner_optimizers() + if not inner_optimizers: + return + + count = 0 + for opt in inner_optimizers: + # Offload FP32 main parameter copies (shard_fp32_from_float16_groups) + if hasattr(opt, "shard_fp32_from_float16_groups"): + for group in opt.shard_fp32_from_float16_groups: + if isinstance(group, list): + for t in group: + if t is not None and t.data.is_cuda: + t.data = t.data.to("cpu", non_blocking=True) + count += 1 + elif group is not None and group.data.is_cuda: + group.data = group.data.to("cpu", non_blocking=True) + count += 1 + + # Offload Adam states (exp_avg, exp_avg_sq) + base_opt = getattr(opt, "optimizer", opt) + if not hasattr(base_opt, "state") or base_opt.state is None: + continue + for state in base_opt.state.values(): + for key in ("exp_avg", "exp_avg_sq"): + if ( + key in state + and isinstance(state[key], torch.Tensor) + and state[key].is_cuda + ): + state[key] = state[key].to("cpu", non_blocking=True) + count += 1 + + # Targeted fix from the AWEX reference: transformer_engine caches dummy wgrad + # tensors in a module-global dict; without purging it the GPU memory + # is never actually freed and stale references survive the offload. + try: + from transformer_engine.pytorch.module.base import _dummy_wgrads + + purged = len(_dummy_wgrads) + for k in list(_dummy_wgrads): + del _dummy_wgrads[k] + if purged: + logger.info("Purged %d TE _dummy_wgrads cache entries", purged) + except ImportError: + pass + torch.cuda.synchronize() + logger.info("Offloaded %d optimizer state tensors to CPU", count) + + def _reload_optimizer_states(self) -> None: + optimizer = self._engine.optimizer + if optimizer is None: + return + if os.environ.get("AWEX_OPT_OFFLOAD_VIA_HDO", "").strip() == "1" and hasattr( + optimizer, "restore_from_cpu" + ): + optimizer.restore_from_cpu() + logger.info("Reloaded optimizer via restore_from_cpu()") + return + + inner_optimizers = self._get_inner_optimizers() + if not inner_optimizers: + return + + device = self._engine.device + count = 0 + for opt in inner_optimizers: + # Reload FP32 main parameter copies + if hasattr(opt, "shard_fp32_from_float16_groups"): + for group in opt.shard_fp32_from_float16_groups: + if isinstance(group, list): + for t in group: + if t is not None and not t.data.is_cuda: + t.data = t.data.to(device, non_blocking=True) + count += 1 + elif group is not None and not group.data.is_cuda: + group.data = group.data.to(device, non_blocking=True) + count += 1 + + # Reload Adam states + base_opt = getattr(opt, "optimizer", opt) + if not hasattr(base_opt, "state") or base_opt.state is None: + continue + for state in base_opt.state.values(): + for key in ("exp_avg", "exp_avg_sq"): + if ( + key in state + and isinstance(state[key], torch.Tensor) + and not state[key].is_cuda + ): + state[key] = state[key].to(device, non_blocking=True) + count += 1 + torch.cuda.synchronize() + logger.info("Reloaded %d optimizer state tensors to GPU", count) + + +def _get_tf_config(models): + if not isinstance(models, (list, tuple)): + models = [models] + for model in models: + for attr in ("transformer_config", "config"): + cfg = getattr(model, attr, None) + if cfg is not None: + return cfg + return None diff --git a/areal/engine/awex/sglang_plugin.py b/areal/engine/awex/sglang_plugin.py new file mode 100644 index 0000000000..4766423f03 --- /dev/null +++ b/areal/engine/awex/sglang_plugin.py @@ -0,0 +1,753 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""AWEX SGLang scheduler plugin for colocated weight transfer. + +Patches SGLang's scheduler to inject CUDA IPC weight receiving capabilities. +When AWEX_META_SERVER_ADDR env var is set, starts a background thread that +fetches IPC handles from MetaServer (CPU I/O) and queues them for the +scheduler's main loop to process (CUDA copy on main thread). + +Weight transfer flow (mirrors the AWEX reference colocate mode): + 1. Training side: convert params → cuda_ipc_serialize → MetaServer put + 2. Background thread: MetaServer get → queue IPC data (CPU only) + 3. Scheduler main loop: release_memory → deserialize + copy → resume_memory + 4. Main loop: signal done → train side releases shared tensors + +Usage: + # Option 1: Register plugin then launch SGLang + from areal.engine.awex.sglang_plugin import register_awex_plugin + register_awex_plugin() + + # Option 2: Run as entry module (replaces sglang.launch_server) + # python3 -m areal.engine.awex.sglang_plugin --model-path ... +""" + +from __future__ import annotations + +import os +import queue +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + + +def assert_alloc_conf_supports_memory_saver(conf: str) -> None: + """Reject allocator configs that silently disable SGLang's memory saver. + + torch_memory_saver disables itself when it sees expandable_segments, so + release/resume becomes a no-op: the rollout never hands its GPU back and + weight pages stay mapped, which surfaces much later as a colocate OOM or an + invalid CUDA IPC target. Colocated roles must therefore carry their own + scheduling_spec env_vars rather than share the actor's. + """ + if "expandable_segments:true" in conf.lower().replace(" ", ""): + raise RuntimeError( + "SGLang's memory saver cannot unmap/remap expandable segments, so " + f"it would disable itself (PYTORCH_CUDA_ALLOC_CONF={conf!r}). Give " + "the rollout role its own scheduling_spec env_vars without " + "expandable_segments instead of sharing the actor's." + ) + + +assert_alloc_conf_supports_memory_saver(os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")) + + +from areal.utils import pkg_version # noqa: E402 +from areal.utils.logging import getLogger # noqa: E402 + +logger = getLogger("AwexSGLangPlugin") + + +SUPPORTED_SGLANG_VERSIONS = ("0.5.9", "0.5.10.post1") + + +def assert_supported_sglang_version() -> None: + """Refuse to patch a SGLang build whose internals were not verified. + + This plugin reaches into scheduler internals: it wraps Scheduler.__init__, + replaces both event loops, and backports a model-worker task dispatcher. + Those touch points move between releases, so a silent mismatch surfaces as + a hang or a corrupt transfer rather than an import error. Fail loudly + instead, and extend the tuple only after re-checking the patched surfaces. + """ + installed = pkg_version.get_version("sglang") + if installed not in SUPPORTED_SGLANG_VERSIONS: + raise RuntimeError( + f"AWEX colocate patches SGLang internals and was verified against " + f"{', '.join(SUPPORTED_SGLANG_VERSIONS)}, but found {installed}. " + f"Re-check Scheduler.__init__, the event loops, and " + f"execute_task_in_model_worker before allowing this version." + ) + + +def _float_env(name: str, default: float) -> float: + value = os.environ.get(name, "") + if not value: + return default + try: + return float(value) + except ValueError: + logger.warning("Invalid %s=%r; using %.3f", name, value, default) + return default + + +class AwexSchedulerPlugin: + """Binds awex weight-receive to a SGLang Scheduler instance. + + Architecture: background thread handles MetaServer I/O (CPU only), + scheduler main loop handles CUDA weight copy (via process_awex_queue). + """ + + def __init__(self, scheduler: Any) -> None: + self._scheduler = scheduler + self._receiver = None + self._bg_thread: threading.Thread | None = None + self._weight_queue: queue.Queue = queue.Queue() + self._version = 0 + self._paused_poll_interval_s = max( + 0.0, _float_env("AWEX_PAUSED_POLL_INTERVAL_S", 0.01) + ) + + def bind(self) -> None: + methods = [ + "awex_init_receiver", + "awex_receive_weights", + "awex_release_memory", + "awex_resume_memory", + "awex_get_weight_metadata", + "awex_get_parallelism", + "process_awex_queue", + ] + for name in methods: + setattr(self._scheduler, name, getattr(self, name)) + logger.info( + f"[AWEX] AwexSchedulerPlugin bound {len(methods)} methods to scheduler", + ) + + meta_server_addr = os.environ.get("AWEX_META_SERVER_ADDR") + if meta_server_addr: + self._start_background_worker(meta_server_addr) + self._patch_event_loop() + + def _require_receiver(self): + if self._receiver is None: + from areal.engine.awex.colocate_reader import AwexColocateReader + + self._receiver = AwexColocateReader(self._scheduler) + return self._receiver + + def awex_init_receiver(self, **kwargs: Any) -> None: + self._require_receiver().initialize(**kwargs) + + def awex_receive_weights(self, version: int = 0) -> None: + self._require_receiver().update_weights(version) + + def awex_release_memory(self, tags: list[str] | None = None) -> None: + self._require_receiver().release_memory(tags) + + def awex_resume_memory(self, tags: list[str] | None = None) -> None: + self._require_receiver().resume_memory(tags) + + def awex_get_weight_metadata(self) -> list: + return self._require_receiver().get_weight_metadata() + + def awex_get_parallelism(self) -> dict: + return self._require_receiver().get_parallelism() + + # ── Main loop hook: process queued weight updates ───────────────── + + def process_awex_queue(self) -> None: + """Called from scheduler main loop. Processes pending weight updates. + + This is a TP-collective operation: ALL TP ranks must call it together + (since it's called between recv_requests() calls which use broadcast_pyobj). + + Uses all_reduce(MIN) to check if all TP ranks have a pending update. + Only proceeds when ALL ranks have queued an update, preventing the deadlock + where one rank blocks in CUDA ops while others wait in broadcast_pyobj. + + We act as the awex *driver* layer (the community SGLang scheduler has no + ``execute_task_in_model_worker`` driver). The collect-IPC + StreamBatch + transport + writer handshake is delegated to the awex-native worker reader + (``AwexColocateReader.update_weights`` -> ``NCCLWorkerWeightsReader``). We + only own the driver-equivalent steps around it: + 1. Wait for all_training_offloaded_weights (= driver _pre_update_weights) + 2. resume_memory_occupation(weights) — re-allocate infer weight buffers + 3. reader.update_weights(version) — awex worker reader does the rest: + collect IPC + StreamBatch transport + put weights_update_finished + + barrier + get_then_delete write_finished + flush_cache + 4. signal_finished_weights_update (= driver _resume_kvcache) + """ + import torch + import torch.distributed + + tp_cpu_group = self._scheduler.tp_cpu_group + tp_size = self._scheduler.tp_size + + has_item = 1 if not self._weight_queue.empty() else 0 + + if tp_size > 1: + has_item_tensor = torch.tensor([has_item], dtype=torch.int32) + torch.distributed.all_reduce( + has_item_tensor, + op=torch.distributed.ReduceOp.MIN, + group=tp_cpu_group, + ) + all_ready = has_item_tensor.item() == 1 + else: + all_ready = has_item == 1 + + if not all_ready: + return + + item = self._weight_queue.get_nowait() + version = item["version"] + gpu_id = getattr(self._scheduler, "gpu_id", "?") + logger.info( + f"[AWEX] main loop: processing weight update v{version} (gpu_id={gpu_id})", + ) + + from sglang.srt.managers.io_struct import ResumeMemoryOccupationReqInput + + receiver = self._require_receiver() + + # Step 1: Wait for writer to offload its model weights first (= awex driver + # _pre_update_weights). Ensures no 2x model weights on GPU simultaneously. + # The background thread already gated on this, so this returns immediately; + # kept for driver-equivalent clarity. + logger.info( + f"[AWEX] main loop: waiting for all_training_offloaded_weights (gpu_id={gpu_id})", + ) + receiver.wait_for_training_offloaded(version) + logger.info( + f"[AWEX] main loop: writer offloaded weights confirmed (gpu_id={gpu_id})", + ) + + # Step 2: Resume weight memory (memory_saver re-allocates buffers). + resume_req = ResumeMemoryOccupationReqInput(tags=["weights"]) + self._scheduler.resume_memory_occupation(resume_req) + logger.info( + f"[AWEX] main loop: resumed weight memory for v{version} (gpu_id={gpu_id})", + ) + + # Step 3: Delegate the whole collect-IPC + StreamBatch transport + writer + # handshake (put weights_update_finished + barrier + get_then_delete + # write_finished + flush_cache) to the awex-native worker reader. + try: + receiver.update_weights(version) + logger.info( + f"[AWEX] main loop: weight update done for v{version} (gpu_id={gpu_id})", + ) + except Exception: + logger.exception( + "AWEX main loop failed to update weights v%s on gpu_id=%s", + version, + gpu_id, + ) + raise + + # Step 4: Signal that this infer engine finished weight update, so the + # writer can resume kv_cache (= awex driver _resume_kvcache). + receiver.signal_finished_weights_update() + self._version = version + + # ── Patch scheduler event loop to call process_awex_queue ───────── + + def _patch_event_loop(self) -> None: + """Inject process_awex_queue into scheduler's event loops. + + SGLang uses event_loop_overlap by default. Patch both for safety. + Weight updates process when engine is paused (no ongoing inference). + """ + scheduler = self._scheduler + plugin = self + + _decode_hooks_available = hasattr(scheduler, "log_decode_stats") and hasattr( + scheduler, "log_decode_stats_every_iteration" + ) + if _decode_hooks_available: + _orig_log_decode_stats = scheduler.log_decode_stats + _orig_log_decode_stats_every_iteration = ( + scheduler.log_decode_stats_every_iteration + ) + + def _tracked_log_decode_stats(*args, **kwargs): + scheduler._areal_awex_last_decode_stats_ct = getattr( + scheduler, "forward_ct_decode", None + ) + return _orig_log_decode_stats(*args, **kwargs) + + def _tracked_log_decode_stats_every_iteration(*args, **kwargs): + scheduler._areal_awex_last_decode_stats_every_iter_ct = getattr( + scheduler, "forward_ct_decode", None + ) + return _orig_log_decode_stats_every_iteration(*args, **kwargs) + + scheduler.log_decode_stats = _tracked_log_decode_stats + scheduler.log_decode_stats_every_iteration = ( + _tracked_log_decode_stats_every_iteration + ) + else: + logger.warning( + "[AWEX] sglang scheduler has no log_decode_stats hooks " + "(removed in sglang>=0.5.10); skipping decode-stats tracking" + ) + + def _maybe_restore_decode_metrics(stage, batch, result): + if not _decode_hooks_available: + return + if os.environ.get("AREAL_AWEX_FORCE_SGLANG_METRICS", "1") != "1": + return + if stage != "after_process_batch_result" or batch is None: + return + mode = getattr(getattr(batch, "forward_mode", None), "name", None) + if mode != "DECODE": + return + if not getattr(scheduler, "current_scheduler_metrics_enabled", False): + return + + current_ct = getattr(scheduler, "forward_ct_decode", None) + interval = ( + getattr( + getattr(scheduler, "server_args", None), "decode_log_interval", 1 + ) + or 1 + ) + should_log_decode = current_ct is not None and current_ct % interval == 0 + + if ( + should_log_decode + and getattr(scheduler, "_areal_awex_last_decode_stats_ct", None) + != current_ct + ): + can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) + logger.debug( + f"[AWEX-METRICS] restoring native log_decode_stats " + f"gpu_id={getattr(scheduler, 'gpu_id', '?')} " + f"forward_ct_decode={current_ct}", + ) + scheduler.log_decode_stats(can_run_cuda_graph, running_batch=batch) + + if ( + getattr(scheduler, "_areal_awex_last_decode_stats_every_iter_ct", None) + != current_ct + ): + scheduler.log_decode_stats_every_iteration( + batch, + num_accepted_tokens=getattr(result, "num_accepted_tokens", 0), + ) + + # Patch event_loop_overlap (the one actually used by SGLang) + _orig_overlap = scheduler.event_loop_overlap + + def _patched_overlap(): + """Patched overlap loop that checks awex queue when paused.""" + from collections import deque + + scheduler.result_queue = deque() + _loop_count = 0 + _paused_reported = False + + def pop_and_process(): + tmp_batch, tmp_result = scheduler.result_queue.popleft() + scheduler.process_batch_result(tmp_batch, tmp_result) + _maybe_restore_decode_metrics( + "after_process_batch_result", tmp_batch, tmp_result + ) + + logger.info( + f"[AWEX] _patched_overlap STARTING (gpu_id={getattr(scheduler, 'gpu_id', '?')})", + ) + + while True: + recv_reqs = scheduler.recv_requests() + if recv_reqs: + req_types = [type(r).__name__ for r in recv_reqs] + has_control = any( + t + not in ( + "TokenizedGenerateReqInput", + "TokenizedEmbeddingReqInput", + ) + for t in req_types + ) + if has_control or _loop_count % 500 == 0: + logger.info( + f"[AWEX] loop gpu_id={getattr(scheduler, 'gpu_id', '?')}: " + f"recv {len(recv_reqs)} reqs, types={req_types[:5]}, " + f"_engine_paused={scheduler._engine_paused}, loop={_loop_count}", + ) + + was_paused = scheduler._engine_paused + scheduler.process_input_requests(recv_reqs) + if scheduler._engine_paused != was_paused: + logger.info( + f"[AWEX] _engine_paused CHANGED: {was_paused} → {scheduler._engine_paused} " + f"(gpu_id={getattr(scheduler, 'gpu_id', '?')}, loop={_loop_count})", + ) + + if scheduler._engine_paused: + if not _paused_reported: + logger.info( + f"[AWEX] _patched_overlap: _engine_paused=True detected! " + f"(gpu_id={getattr(scheduler, 'gpu_id', '?')}, loop_count={_loop_count})", + ) + _paused_reported = True + plugin.process_awex_queue() + time.sleep(plugin._paused_poll_interval_s) + continue + + _loop_count += 1 + batch = scheduler.get_next_batch_to_run() + scheduler.cur_batch = batch + disable_overlap_for_batch = scheduler.is_disable_overlap_for_batch( + batch + ) + + if disable_overlap_for_batch: + pop_and_process() + + if batch: + batch_result = scheduler.run_batch(batch) + scheduler.result_queue.append((batch.copy(), batch_result)) + else: + batch_result = None + + if scheduler.last_batch: + if not disable_overlap_for_batch: + pop_and_process() + elif batch is None: + scheduler.self_check_during_idle() + + if scheduler.is_generation: + scheduler.launch_batch_sample_if_needed(batch_result) + + scheduler.last_batch = batch + + scheduler.event_loop_overlap = _patched_overlap + + # Also patch event_loop_normal as fallback + _orig_normal = scheduler.event_loop_normal + + def _patched_normal(): + logger.info( + f"[AWEX] _patched_normal STARTING (gpu_id={getattr(scheduler, 'gpu_id', '?')})", + ) + while True: + recv_reqs = scheduler.recv_requests() + scheduler.process_input_requests(recv_reqs) + if scheduler._engine_paused: + plugin.process_awex_queue() + time.sleep(plugin._paused_poll_interval_s) + continue + batch = scheduler.get_next_batch_to_run() + scheduler.cur_batch = batch + if batch: + result = scheduler.run_batch(batch) + scheduler.process_batch_result(batch, result) + _maybe_restore_decode_metrics( + "after_process_batch_result", batch, result + ) + else: + scheduler.self_check_during_idle() + scheduler.last_batch = batch + + scheduler.event_loop_normal = _patched_normal + logger.info( + "[AWEX] Patched event_loop_overlap + event_loop_normal with awex queue", + ) + + # ── Background thread: MetaServer I/O only (no CUDA ops) ───────── + + def _start_background_worker(self, meta_server_addr: str) -> None: + self._bg_thread = threading.Thread( + target=self._background_worker, + args=(meta_server_addr,), + daemon=True, + ) + self._bg_thread.start() + gpu_id = int(getattr(self._scheduler, "gpu_id", -1)) + logger.info( + f"[AWEX] Started background worker thread " + f"(gpu_id={gpu_id}, meta_server={meta_server_addr})", + ) + + def _background_worker(self, meta_server_addr: str) -> None: + """Initialize the reader, then gate weight-update triggers to the main loop. + + This thread does NOT perform any CUDA memory writes. It only: + 1. Connects to MetaServer and initializes the awex worker reader + 2. Blocks on the per-version writer-offload signal (a set-size wait) + 3. Enqueues a version marker so the TP-collective main-loop gate fires + (the awex worker reader collects the IPC handles itself inside + update_weights, so no large payload is prefetched here) + """ + import torch + + gpu_id = int(getattr(self._scheduler, "gpu_id", 0)) + torch.cuda.set_device(gpu_id) + logger.info(f"[AWEX] background worker: set CUDA device to {gpu_id}") + + try: + self._init_receiver_from_meta_server(meta_server_addr) + except Exception: + logger.exception("AWEX background worker initialization failed") + return + + logger.info( + "[AWEX] background worker: initialization complete, entering fetch loop", + ) + # Recover can resume weight-transfer versions from the checkpoint step + # instead of v1, so sync the first version from the writer. + from awex.meta.meta_server import MetaServerClient as _MSC + from awex.util.common import get_ip_address as _get_ip + + _host, _port = meta_server_addr.rsplit(":", 1) + _ver_client = _MSC(_host, int(_port)) + from areal.engine.awex.colocate_writer import ( + awex_colocate_timeout_s, + resolve_physical_gpu_id, + ) + + # Shares a key namespace with the training writer, so it needs the + # physical GPU id rather than the mask-relative index. + _ver_key = f"awex_writer_version_{_get_ip()}_{resolve_physical_gpu_id(gpu_id)}" + + version = int( + _ver_client.get_object( + _ver_key, + timeout=awex_colocate_timeout_s(), + ) + ) + logger.info( + f"[AWEX] background worker: writer stream starts at v{version}", + ) + retries = 0 + # Slow online environments can spend tens of minutes between updates. + # Keep the reader alive across transient wait timeouts. + max_retries = int(os.environ.get("AWEX_READER_MAX_RETRIES", "1000")) + + while True: + try: + # Block on THIS version's writer-published IPC handles + # (existence-only probe, no deserialization). This is the + # per-version trigger: the writer only publishes v+1's key in the + # next training cycle, so the background thread cannot fire early + # off a stale unversioned set and dead-lock the main loop. See + # AwexColocateReader.wait_for_weights_ready for the full rationale. + logger.info( + f"[AWEX] background worker: waiting for writer weights v{version}", + ) + receiver = self._require_receiver() + receiver.wait_for_weights_ready(version) + logger.info( + f"[AWEX] background worker: writer published v{version}, " + f"queuing for main loop", + ) + + # Queue a version marker for the main loop (no CUDA ops here). + self._weight_queue.put({"version": version}) + + # Wait for main loop to finish processing before gating the next. + while self._version < version: + time.sleep(0.1) + + version += 1 + retries = 0 + except Exception as e: + retries += 1 + logger.exception( + "AWEX background worker failed while waiting for writer " + "weights v%s (attempt %s/%s): %s", + version, + retries, + max_retries, + e, + ) + if retries >= max_retries: + logger.info( + f"[AWEX] background worker: giving up after {max_retries} failures", + ) + break + time.sleep(min(2**retries, 30)) + + def _init_receiver_from_meta_server(self, meta_server_addr: str): + """Connect to MetaServer, get train info, initialize colocate receiver.""" + from awex.meta.meta_server import MetaServerClient + + host, port = meta_server_addr.rsplit(":", 1) + + client = None + for attempt in range(60): + try: + client = MetaServerClient(host, int(port)) + break + except Exception: + if attempt % 10 == 0: + logger.info( + f"[AWEX] background worker: MetaServer not ready, retrying... " + f"(attempt {attempt + 1}, addr={meta_server_addr})", + ) + time.sleep(5) + if client is None: + raise RuntimeError( + f"Failed to connect to MetaServer at {meta_server_addr} after 60 attempts" + ) + + logger.info( + f"[AWEX] background worker: connected to MetaServer at {meta_server_addr}", + ) + + receiver = self._require_receiver() + + # `gpu_id` is node-local. Multi-node colocate needs a globally unique + # transfer rank that stays physically paired with the training process. + gpu_id = int(getattr(self._scheduler, "gpu_id", 0)) + node_id = int(os.environ.get("SLURM_NODEID", "0")) + nnodes = int(os.environ.get("SLURM_NNODES", "1")) + + logger.info( + f"[AWEX] background worker: waiting for awex_train_info " + f"(gpu_id={gpu_id}, node_id={node_id}, nnodes={nnodes})", + ) + # The driver publishes awex_train_info only after rollout init finishes, + # so large models need the same timeout budget as the weight path. + from areal.engine.awex.colocate_writer import ( + awex_colocate_timeout_s, + resolve_physical_gpu_id, + ) + + train_info = client.get_object( + "awex_train_info", + timeout=awex_colocate_timeout_s(), + ) + train_world_size = train_info["train_world_size"] + # In colocate mode train and infer share the same N physical GPUs, so the + # global infer NCCL world spans the same N ranks (numerically == train + # world). This is a *physical* coincidence (same GPUs), NOT a requirement + # that train/infer parallel topologies match: the infer side decomposes + # into num_infer_engines DP replicas inside receiver.initialize(). + infer_world_size = train_world_size + + n_gpus_per_node = max(1, infer_world_size // nnodes) + transfer_rank = node_id * n_gpus_per_node + gpu_id + + logger.info( + f"[AWEX] background worker: got train_world_size={train_world_size}, " + f"infer_world_size={infer_world_size}, n_gpus_per_node={n_gpus_per_node}, " + f"transfer_rank={transfer_rank}", + ) + + # transfer_rank stays a relative index so it lines up with the infer + # NCCL world, but the CUDA IPC keys are shared with the training writer + # and must therefore use physical GPU ids. + receiver.initialize( + meta_server_addr=meta_server_addr, + transfer_rank=transfer_rank, + infer_world_size=infer_world_size, + train_world_size=train_world_size, + local_gpu_id=resolve_physical_gpu_id(gpu_id), + ) + logger.info( + f"[AWEX] background worker: receiver initialized " + f"(transfer_rank={transfer_rank}, infer_world_size={infer_world_size})", + ) + + +@dataclass +class ModelWorkerTask: + """Task for execute_task_in_model_worker (PR #13595 backport for SGLang 0.5.9).""" + + task_func: Callable + kwargs: dict = field(default_factory=dict) + + +def register_awex_plugin() -> None: + """Patch Scheduler.__init__ to inject awex plugin after construction. + + Must be called INSIDE the scheduler child process (not the parent), + because SGLang spawns scheduler processes via mp.Process with "spawn" + start method, which doesn't inherit parent-process monkey-patches. + """ + assert_supported_sglang_version() + + from sglang.srt.managers.scheduler import Scheduler + + _orig_init = Scheduler.__init__ + + def _patched_init(self, *args, **kwargs): + _orig_init(self, *args, **kwargs) + AwexSchedulerPlugin(self).bind() + _patch_execute_task_in_model_worker(self) + + Scheduler.__init__ = _patched_init + logger.info("[AWEX] Patched Scheduler.__init__ with awex plugin") + + +def _patch_execute_task_in_model_worker(scheduler) -> None: + """Add execute_task_in_model_worker to Scheduler (backport from PR #13595).""" + + def execute_task_in_model_worker(task_spec: ModelWorkerTask): + model_context = dict( + tp_rank=scheduler.tp_rank, + tp_size=scheduler.tp_size, + server_args=scheduler.server_args, + scheduler=scheduler, + ) + kwargs = dict(task_spec.kwargs) + kwargs["model_context"] = model_context + kwargs["model"] = scheduler.tp_worker.model_runner.model + kwargs["model_runner"] = scheduler.tp_worker.model_runner + return task_spec.task_func(**kwargs) + + scheduler.execute_task_in_model_worker = execute_task_in_model_worker + + if hasattr(scheduler, "_request_dispatcher"): + scheduler._request_dispatcher._mapping[ModelWorkerTask] = ( + execute_task_in_model_worker + ) + logger.info("[AWEX] Registered execute_task_in_model_worker in dispatcher") + + +def awex_run_scheduler_process(*args, **kwargs): + """Scheduler process entry point that registers awex plugin. + + Memory management (pause/resume weights, KV cache, CUDA graphs) is handled + at runtime by AWEX's release_memory/resume_memory, matching the AWEX + reference integration. + No init-time memory patching needed. + """ + import os + + meta_addr = os.environ.get("AWEX_META_SERVER_ADDR") + if meta_addr: + register_awex_plugin() + else: + logger.info( + "[AWEX] No AWEX_META_SERVER_ADDR, skipping plugin registration", + ) + from sglang.srt.managers.scheduler import run_scheduler_process + + return run_scheduler_process(*args, **kwargs) + + +if __name__ == "__main__": + import os + import sys + + logger.info("[AWEX] awex_sglang_plugin __main__ starting") + + from sglang.srt.entrypoints.http_server import launch_server + from sglang.srt.server_args import prepare_server_args + from sglang.srt.utils import kill_process_tree + + server_args = prepare_server_args(sys.argv[1:]) + try: + launch_server( + server_args, + run_scheduler_process_func=awex_run_scheduler_process, + ) + finally: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 984d9143b5..a8b75ad017 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -202,6 +202,7 @@ def __init__(self, config: TrainEngineConfig): self.own_global_group: bool = False self.is_offload: bool = False self._offload_depth: int = 0 + self._awex_adapter = None # AwexMegatronAdapter for colocate mode self.enable_tree_training: bool = self.config.enable_tree_training # FP8 configuration self.fp8_config = self.mcore_config.fp8_config @@ -248,6 +249,7 @@ def create_process_group(self, parallel_strategy: ParallelStrategy | None = None self.own_global_group = True self.logger = logging.getLogger(f"[MegatronEngine Rank {dist.get_rank()}]") self._context_and_model_parallel_group = None + self._cpu_model_parallel_group = None self._init_context_and_model_parallel_group() # This is needed for barrier synchronization when models are moved to CPU self._cpu_group = dist.new_group( @@ -682,6 +684,11 @@ def context_and_model_parallel_group(self) -> dist.ProcessGroup: assert self.process_group_initialized return self._context_and_model_parallel_group + @property + def cpu_model_parallel_group(self) -> dist.ProcessGroup: + assert self.process_group_initialized + return self._cpu_model_parallel_group + @property def cpu_group(self) -> dist.ProcessGroup: assert self.process_group_initialized @@ -743,6 +750,17 @@ def connect_engine(self, engine: InferenceEngine, meta: WeightUpdateMeta): if meta.type == "xccl" and not self.weight_update_group_initialized: self._init_weight_update_from_distributed(meta) self.weight_update_group_initialized = True + elif meta.type == "awex": + from areal.engine.awex.colocate_writer import AwexMegatronAdapter + + if self._awex_adapter is None: + self._awex_adapter = AwexMegatronAdapter(self) + self._awex_adapter.init_colocate_weight_update( + meta_server_addr=meta.nccl_master_address, + pair_name=meta.nccl_group_name or "default", + transfer_rank=self.rank or 0, + ) + self.logger.info("Initialized AWEX colocate adapter") current_platform.synchronize() dist.barrier(group=self.cpu_group) @@ -791,6 +809,31 @@ def prepare_batch( def update_weights(self, meta: WeightUpdateMeta): self._check_rollout_engine_connected() + if meta.type == "awex": + # Colocate mode flow (mirrors the AWEX reference integration): + # 1. execute_colocate_weight_update: release grad → convert → offload + # weights → signal offloaded → IPC serialize → wait reader done → + # cleanup shared → signal write_finished + # 2. finish: wait all infer engines done → cleanup MetaServer keys + # 3. resume kv_cache + continue generation + self._awex_adapter.execute_colocate_weight_update(meta.version or 0) + # Do NOT flip is_offload here: the AWEX adapter tracks released + # memory via _released_tags, and the trainer onloads explicitly + # at the next train phase. Marking is_offload would make every + # _offload_aware_context RPC (e.g. export_stats) reload optimizer + # states onto a GPU already fully occupied by the resumed rollout. + + dist.barrier(group=self.cpu_group) + + self._awex_adapter.finish_colocate_weight_update( + training_world_size=dist.get_world_size(self.cpu_group) + ) + + if dist.get_rank() == 0: + self.rollout_engine.onload(tags=["kv_cache"]) + self.rollout_engine.continue_generation() + dist.barrier(group=self.cpu_group) + return with self._offload_aware_context(): if meta.type == "xccl": assert self.weight_update_group_initialized @@ -807,6 +850,14 @@ def get_version(self) -> int: return self._version def save(self, meta: SaveLoadMeta): + if self._awex_adapter is not None: + # Post-ppo_update the fp32 grad buffers (~2x param bytes) are + # dead weight until the next train_batch (which rebuilds them via + # ensure_grad_buffers); drop them here to fund the HF saver's TP + # coalesced all-gather transient. + self._awex_adapter._release_grad_memory() + gc.collect() + torch.cuda.empty_cache() with self._offload_aware_context(): if meta.weight_format == "hf": if meta.with_optim: @@ -1010,6 +1061,8 @@ def train_batch( loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], ) -> dict[str, float]: self._ensure_ready() + if self._awex_adapter is not None: + self._awex_adapter.ensure_grad_buffers() self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -1187,11 +1240,58 @@ def export_stats(self) -> dict[str, float]: data.update(data_list[0]) return data + def init_awex_adapter(self, meta_server_addr: str | None = None) -> None: + """Create awex adapter early for selective memory management. + + Must be called before offload() in colocate mode so that offload uses + the adapter's tag-based mechanism instead of TMS (which is all-or-nothing + and causes OOM on resume when SGLang occupies GPU memory). + """ + if self._awex_adapter is None: + from areal.engine.awex.colocate_writer import AwexMegatronAdapter + + self._awex_adapter = AwexMegatronAdapter(self) + self.logger.info("Created AWEX adapter for memory management") + + self._eager_publish_awex_train_info(meta_server_addr) + + def _eager_publish_awex_train_info(self, meta_server_addr: str | None) -> None: + addr = meta_server_addr or os.environ.get("AWEX_META_SERVER_ADDR", "") + if not addr or (dist.is_initialized() and dist.get_rank() != 0): + return + try: + from awex.meta.meta_server import MetaServerClient + + host, port = addr.rsplit(":", 1) + client = MetaServerClient(host, int(port)) + world = dist.get_world_size() if dist.is_initialized() else 1 + client.put_object("awex_train_info", {"train_world_size": world}) + self.logger.info( + "[AWEX] eager-published awex_train_info (train_world_size=%d) to %s", + world, + addr, + ) + except Exception as e: + self.logger.warning("[AWEX] eager publish awex_train_info failed: %s", e) + def offload(self) -> None: - """Offload model memory to CPU using torch_memory_saver. + """Offload model memory to CPU. + + In colocate mode (awex adapter active): manual tag-based offload via + storage resize. Otherwise: torch_memory_saver pause. Ref: https://github.com/THUDM/slime/blob/main/slime/backends/megatron_utils/actor.py """ + if self._awex_adapter is not None: + self.get_device_stats().log("before offload model") + current_platform.clear_memory() + self._awex_adapter.release_memory(tags=["optimizer", "weights"]) + current_platform.synchronize() + dist.barrier(group=self.cpu_group) + self.get_device_stats().log("after offload model") + self.is_offload = True + return + if not is_tms_enabled(): raise RuntimeError( "torch_memory_saver requires `enable_offload=True` in yaml config." @@ -1220,10 +1320,21 @@ def offload(self) -> None: self.is_offload = True def onload(self) -> None: - """Onload model memory from CPU back to GPU using torch_memory_saver. + """Onload model memory from CPU back to GPU. + + Uses awex adapter (selective, tag-based) when available, otherwise + torch_memory_saver. Ref: https://github.com/THUDM/slime/blob/main/slime/backends/megatron_utils/actor.py """ + if self._awex_adapter is not None: + self._awex_adapter.resume_memory(tags=["optimizer", "weights"]) + current_platform.clear_memory() + current_platform.synchronize() + dist.barrier(group=self.cpu_group) + self.get_device_stats().log("after onload model") + self.is_offload = False + return torch_memory_saver.resume() @@ -1401,6 +1512,18 @@ def _init_context_and_model_parallel_group(self) -> None: ) if dp_rank == mpu.get_data_parallel_rank(): self._context_and_model_parallel_group = group + # The gloo mirror is only read once an offloaded engine has handed the + # accelerator to rollout and device collectives are unusable (see + # resolve_broadcast_target). Building it otherwise costs one collective + # per data-parallel group at startup for a group nothing ever uses; + # resolve_broadcast_target falls back to the device group when it is None. + if self.config.offload: + for dp_rank, ranks in enumerate(context_and_model_parallel_ranks): + cpu_group = dist.new_group( + ranks, timeout=DIST_GROUP_DEFAULT_TIMEOUT, backend="gloo" + ) + if dp_rank == mpu.get_data_parallel_rank(): + self._cpu_model_parallel_group = cpu_group def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: if self.optimizer_config is None: diff --git a/areal/engine/sglang_remote.py b/areal/engine/sglang_remote.py index 5787be3b9d..cd54bc8a27 100644 --- a/areal/engine/sglang_remote.py +++ b/areal/engine/sglang_remote.py @@ -35,8 +35,11 @@ from areal.infra.platforms import current_platform from areal.infra.utils.launcher import TRITON_CACHE_PATH from areal.utils import perf_tracer, stats_tracker +from areal.utils.logging import getLogger from areal.utils.network import format_host_for_url +logger = getLogger("SGLangRemote") + class SGLangBackend: """SGLang-specific backend implementation for remote inference.""" @@ -331,21 +334,49 @@ def build_init_weights_group_request( return HttpRequest(endpoint="/init_weights_update_group", payload=payload) - def get_pause_request(self) -> HttpRequest: + def get_pause_request(self, mode: str | None = None) -> HttpRequest: """Get SGLang pause request.""" - return HttpRequest(endpoint="/pause_generation", payload={}) + payload = {} if mode is None else {"mode": mode} + return HttpRequest(endpoint="/pause_generation", payload=payload) + + def get_pause_requests(self) -> list[HttpRequest]: + """Pause in two steps so memory can be released safely. + + The first request keeps SGLang's default mode, which aborts in-flight + requests and returns their partial output so the client resumes them by + extending the prompt. That also leaves the scheduler fully idle, which + SGLang requires before releasing memory. + + The second request looks redundant because Scheduler.pause_generation + sets its paused flag unconditionally, but the abort path never reaches + the scheduler: TokenizerManager forwards the request only for non-abort + modes, and otherwise just drains via abort_request(). Abort therefore + raises the tokenizer's own gate while the scheduler keeps scheduling. + Only an in-place pause raises the scheduler flag that the colocate loop + watches before it services awex work, and by then the abort has already + left nothing for that mode to retain. + """ + return [ + self.get_pause_request(), + self.get_pause_request(mode="in_place"), + ] def get_resume_request(self) -> HttpRequest: """Get SGLang resume request.""" return HttpRequest(endpoint="/continue_generation", payload={}) + def get_abort_all_request(self) -> HttpRequest: + """Get SGLang abort all requests.""" + return HttpRequest(endpoint="/abort_request", payload={"abort_all": True}) + def get_health_check_request(self) -> HttpRequest: """Get SGLang health check request.""" return HttpRequest(endpoint="/health", payload={}, method="GET") - def get_offload_request(self) -> HttpRequest: + def get_offload_request(self, tags: list[str] | None = None) -> HttpRequest: """Get SGLang offload request.""" - return HttpRequest(endpoint="/release_memory_occupation", payload={}) + payload = {"tags": tags} if tags is not None else {} + return HttpRequest(endpoint="/release_memory_occupation", payload=payload) def get_onload_request(self, tags: list[str] | None = None) -> HttpRequest: """Get SGLang onload request. @@ -360,9 +391,48 @@ def get_onload_request(self, tags: list[str] | None = None) -> HttpRequest: def launch_server(self, server_args: dict[str, Any]) -> subprocess.Popen: """Launch SGLang server subprocess.""" + awex_meta_addr = server_args.pop("awex_meta_server_addr", None) + awex_colocate = server_args.pop("awex_colocate_mode", False) + # Colocate placement: derive base_gpu_id from SLURM_LOCALID so two SGLang + # servers sharing a node never claim the same GPU range. The controller + # cannot do this reliably because its global rank -> node-slot mapping is + # not guaranteed by SLURM task dispatch (a collision degrades the + # TP group into an unsharded single-GPU load -> OOM). SLURM_LOCALID is the + # only id guaranteed unique per node-slot, and only the worker sees it at + # runtime. `_awex_gpus_per_server` is injected by the controller exclusively + # for real colocation, so its presence doubles as the colocate gate; it is + # absent for separated mode (where CVD isolation keeps base_gpu_id at 0). + awex_gpus_per_server = server_args.pop("_awex_gpus_per_server", None) + if awex_gpus_per_server is not None: + slurm_localid = os.environ.get("SLURM_LOCALID") + if slurm_localid is not None: + base_gpu_id = int(slurm_localid) * int(awex_gpus_per_server) + server_args["base_gpu_id"] = base_gpu_id + logger.info( + "AWEX colocate base_gpu_id override: SLURM_LOCALID=%s x " + "gpus_per_server=%s -> base_gpu_id=%s", + slurm_localid, + awex_gpus_per_server, + base_gpu_id, + ) cmd = SGLangConfig.build_cmd_from_args(server_args) _env = self.build_server_env(os.environ) + if not awex_meta_addr: + awex_meta_addr = os.environ.get("AWEX_META_SERVER_ADDR") + if awex_colocate or awex_meta_addr: + sglang_entrypoints = ( + "sglang.launch_server", + "areal.v2.inference_service.sglang.launch_server", + ) + cmd = [ + "areal.engine.awex.sglang_plugin" if c in sglang_entrypoints else c + for c in cmd + ] + if awex_meta_addr: + _env["AWEX_META_SERVER_ADDR"] = awex_meta_addr + logger.info("AWEX mode: using awex_sglang_plugin entry, cmd=%s", cmd[:4]) + return subprocess.Popen( cmd, env=_env, @@ -594,8 +664,14 @@ def launch_server(self, server_args: dict[str, Any]) -> LocalInfServerInfo: def teardown_server(self): return self._engine.teardown_server() - def offload(self): - return self._engine.offload() + def offload(self, tags: list[str] | None = None): + logger.info("RemoteSGLangEngine.offload(tags=%s) called", tags) + result = self._engine.offload(tags=tags) + logger.info("RemoteSGLangEngine.offload(tags=%s) done", tags) + return result + + def abort_all_requests(self): + return self._engine.abort_all_requests() def onload(self, tags: list[str] | None = None): return self._engine.onload(tags=tags) diff --git a/areal/engine/vllm_remote.py b/areal/engine/vllm_remote.py index 40dcb39444..2d76930a9d 100644 --- a/areal/engine/vllm_remote.py +++ b/areal/engine/vllm_remote.py @@ -257,11 +257,20 @@ def get_health_check_request(self) -> HttpRequest: """Get vLLM health check request.""" return HttpRequest(endpoint="/health", payload={}, method="GET") - def get_offload_request(self) -> HttpRequest: + def get_abort_all_request(self) -> HttpRequest: + raise NotImplementedError("vLLM does not support abort_all_requests") + + def get_offload_request(self, tags: list[str] | None = None) -> HttpRequest: """Get vLLM offload request. Uses vLLM's /sleep endpoint to offload model memory to CPU. Default level is 1. + + Parameters + ---------- + tags : list[str], optional + Accepted for RemoteInfEngine API compatibility. vLLM sleep does not + support component-specific tags, so this value is ignored. """ return HttpRequest(endpoint="/sleep", payload={}, method="POST") @@ -522,8 +531,8 @@ def launch_server(self, server_args: dict[str, Any]) -> LocalInfServerInfo: def teardown_server(self): return self._engine.teardown_server() - def offload(self): - return self._engine.offload() + def offload(self, tags: list[str] | None = None): + return self._engine.offload(tags=tags) def onload(self, tags: list[str] | None = None): return self._engine.onload(tags=tags) diff --git a/areal/infra/controller/rollout_callback.py b/areal/infra/controller/rollout_callback.py index 1f106a5373..fc00cef15a 100644 --- a/areal/infra/controller/rollout_callback.py +++ b/areal/infra/controller/rollout_callback.py @@ -180,3 +180,13 @@ def continue_generation(self) -> None: This is synchronous as it should complete before returning control. """ self._post("/callback/continue_generation") + + def onload(self, tags: list[str] | None = None) -> None: + """Callback to controller to resume memory occupation on inference side.""" + payload = {"tags": tags} if tags else {} + self._post("/callback/onload", payload) + + def offload(self, tags: list[str] | None = None) -> None: + """Callback to controller to release memory occupation on inference side.""" + payload = {"tags": tags} if tags else {} + self._post("/callback/offload", payload) diff --git a/areal/infra/controller/rollout_controller.py b/areal/infra/controller/rollout_controller.py index be3bd09751..8878358dfe 100644 --- a/areal/infra/controller/rollout_controller.py +++ b/areal/infra/controller/rollout_controller.py @@ -34,6 +34,7 @@ InferenceEngineConfig, PerfTracerConfig, SchedulingSpec, + SchedulingStrategyType, ) from areal.infra.rpc.serialization import deserialize_value from areal.infra.utils.concurrent import run_async_task @@ -89,6 +90,7 @@ def __init__( self.workers: list[Worker] = [] # List of Worker objects from scheduler self.server_infos: list[LocalInfServerInfo] = [] self._worker_role: str + self._gpus_per_server: int = 1 # Round-robin scheduling self._current_worker_idx = 0 @@ -174,6 +176,7 @@ def initialize( self.rollout_alloc.parallel.tp_size * self.rollout_alloc.parallel.pp_size ) dp_size = self.rollout_alloc.parallel.dp_size + self._gpus_per_server = instance_size # The first element of `self.config.scheduling_spec` is the resource spec # of workers, aka the RPC server process. Since a worker exactly matches @@ -292,7 +295,61 @@ async def _async_initialize( ) ] await asyncio.gather(*tasks) + elif ( + self.config.scheduling_strategy.type + == SchedulingStrategyType.colocation.value + ): + # Colocation (AWEX) path: multiple servers share a node, so SLURM does + # NOT isolate GPUs per worker. We must compute base_gpu_id explicitly and + # inject `_awex_gpus_per_server` so the worker recomputes base_gpu_id from + # its own SLURM_LOCALID at runtime (the only value guaranteed unique per + # node-slot). See SGLangBackend.launch_server. + # + # NOTE: this assumes a server fits within one node + # (gpus_per_server <= n_gpus_per_node). Cross-node TP servers would + # collapse slots_per_node to 1 and collide; only the SLURM_LOCALID + # path is collision-safe in that case. + slots_per_node = max( + 1, + getattr(self.scheduler, "n_gpus_per_node", 8) // self._gpus_per_server, + ) + launch_tasks = [] + for rank, worker in enumerate(self.workers): + per_worker_args = { + **server_args, + "base_gpu_id": (rank % slots_per_node) * self._gpus_per_server, + "_awex_gpus_per_server": self._gpus_per_server, + } + launch_tasks.append( + self.scheduler.async_call_engine( + worker_id=worker.id, + method="launch_server", + engine_name=self._engine_name(rank), + server_args=per_worker_args, + ) + ) + self.server_infos = await asyncio.gather(*launch_tasks) + tasks = [ + self.scheduler.async_call_engine( + worker_id=worker.id, + method="initialize", + engine_name=self._engine_name(rank), + engine_id=str(rank), + addr=f"{info.host}:{info.port}", + *args, + **kwargs, + ) + for rank, (worker, info) in enumerate( + zip(self.workers, self.server_infos) + ) + ] + await asyncio.gather(*tasks) else: + # Separation path: each rollout server gets its own SLURM-isolated GPUs, + # so we must NOT override base_gpu_id (SLURM already sets + # CUDA_VISIBLE_DEVICES per worker). Use the collective launch + addr-less + # initialize that the verified disaggregated baseline relies on; the + # worker discovers its server address via name_resolve. self.server_infos = await self._collective_rpc_async( "launch_server", server_args=server_args ) @@ -301,7 +358,6 @@ async def _async_initialize( worker_id=worker.id, method="initialize", engine_name=self._engine_name(rank), - # args in `engine_api` engine_id=str(rank), *args, **kwargs, @@ -588,6 +644,20 @@ def continue_generation(): self._callback_loop.run_until_complete(self.continue_generation()) return jsonify({"status": "ok"}) + @app.route("/callback/onload", methods=["POST"]) + def onload(): + payload = request.get_json() or {} + tags = payload.get("tags") + self.onload(tags=tags) + return jsonify({"status": "ok"}) + + @app.route("/callback/offload", methods=["POST"]) + def offload(): + payload = request.get_json() or {} + tags = payload.get("tags") + self.offload(tags=tags) + return jsonify({"status": "ok"}) + @app.route("/callback/rollout_complete", methods=["POST"]) def rollout_complete(): payload = request.get_json() or {} @@ -1094,17 +1164,12 @@ async def update_weights_from_disk(self, meta: WeightUpdateMeta): async def pause_generation(self): await self._collective_rpc_async("pause_generation") + def pause_generation_sync(self): + self._collective_rpc("pause_generation", http_timeout=120.0) + async def continue_generation(self): await self._collective_rpc_async("continue_generation") - def offload(self) -> None: - """Offload rollout model memory on all inference workers.""" - self._collective_rpc("offload") - - def onload(self, tags: list[str] | None = None) -> None: - """Onload rollout model memory on all inference workers.""" - self._collective_rpc("onload", tags=tags) - def set_version(self, version: int) -> None: with self._version_lock: self._version = version @@ -1126,6 +1191,15 @@ def resume(self): self._collective_rpc("resume", http_timeout=60.0) self.dispatcher.resume() + def offload(self, tags: list[str] | None = None): + self._collective_rpc("offload", tags=tags, http_timeout=120.0) + + def abort_all_requests(self): + self._collective_rpc("abort_all_requests", http_timeout=60.0) + + def onload(self, tags: list[str] | None = None): + self._collective_rpc("onload", tags=tags, http_timeout=120.0) + def export_stats(self) -> dict[str, float]: all_raw_stats = self._collective_rpc(method="export_stats", http_timeout=60.0) stats = defaultdict(float) diff --git a/areal/infra/controller/train_controller.py b/areal/infra/controller/train_controller.py index fe37d1ed8e..45f1c4b340 100644 --- a/areal/infra/controller/train_controller.py +++ b/areal/infra/controller/train_controller.py @@ -707,6 +707,12 @@ def load(self, meta: SaveLoadMeta): """ self._custom_function_call("load", meta) + def init_awex_adapter(self, meta_server_addr: str | None = None): + """Create awex adapter early for selective memory management.""" + self._custom_function_call( + "init_awex_adapter", meta_server_addr=meta_server_addr + ) + def step_lr_scheduler(self): """Step the learning rate scheduler. diff --git a/areal/infra/launcher/sglang_server.py b/areal/infra/launcher/sglang_server.py index 07b76344ca..f7d437a70d 100644 --- a/areal/infra/launcher/sglang_server.py +++ b/areal/infra/launcher/sglang_server.py @@ -203,7 +203,15 @@ def run(self): self._monitor_server_processes(server_addresses) def launch_one_server(self, cmd, host_ip, server_port, node_rank): - server_process = launch_server_cmd(cmd) + custom_env = None + awex_meta_addr = os.environ.get("AWEX_META_SERVER_ADDR") + if awex_meta_addr: + custom_env = {"AWEX_META_SERVER_ADDR": awex_meta_addr} + cmd = [ + "areal.engine.awex.sglang_plugin" if c == "sglang.launch_server" else c + for c in cmd + ] + server_process = launch_server_cmd(cmd, custom_env=custom_env) wait_for_server(f"http://{format_hostport(host_ip, server_port)}") if node_rank == 0: name = names.gen_servers(self.experiment_name, self.trial_name) diff --git a/areal/infra/remote_inf_engine.py b/areal/infra/remote_inf_engine.py index d5c440b286..e634a97ab0 100644 --- a/areal/infra/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -353,7 +353,11 @@ def get_health_check_request(self) -> HttpRequest: """ ... - def get_offload_request(self) -> HttpRequest: + def get_abort_all_request(self) -> HttpRequest: + """Get request to abort all in-flight requests.""" + ... + + def get_offload_request(self, tags: list[str] | None = None) -> HttpRequest: """Get request to offload model memory. Returns @@ -1373,8 +1377,14 @@ def prepare_batch( @trace_perf("remote_inf_engine.pause_generation", category="misc") def pause_generation(self): """Pause request submission for async rollout.""" - pause_req = self.backend.get_pause_request() - self._run_request_on_all_servers(pause_req) + get_pause_requests = getattr(self.backend, "get_pause_requests", None) + pause_requests = ( + get_pause_requests() + if get_pause_requests is not None + else [self.backend.get_pause_request()] + ) + for pause_req in pause_requests: + self._run_request_on_all_servers(pause_req) # The above http request may require some time to be scheduled and executed. # The following line waits until all requests are indeed dropped. @@ -1396,10 +1406,22 @@ def resume(self): """Resume request submission for async rollout.""" return self.workflow_executor.resume() - def offload(self) -> None: + def offload(self, tags: list[str] | None = None) -> None: """Offload model memory on all servers.""" - offload_req = self.backend.get_offload_request() + offload_req = self.backend.get_offload_request(tags=tags) + self.logger.info( + "RemoteInfEngine.offload(tags=%s) sending to %s: endpoint=%s", + tags, + self.addresses, + offload_req.endpoint, + ) self._run_request_on_all_servers(offload_req) + self.logger.info("RemoteInfEngine.offload(tags=%s) completed", tags) + + def abort_all_requests(self) -> None: + """Abort all in-flight requests on all servers.""" + abort_req = self.backend.get_abort_all_request() + self._run_request_on_all_servers(abort_req) def onload(self, tags: list[str] | None = None) -> None: """Onload model memory on all servers.""" diff --git a/areal/infra/rpc/guard/engine_blueprint.py b/areal/infra/rpc/guard/engine_blueprint.py index 0c4fdf0a6e..e481e97981 100644 --- a/areal/infra/rpc/guard/engine_blueprint.py +++ b/areal/infra/rpc/guard/engine_blueprint.py @@ -89,6 +89,20 @@ def _should_broadcast_payload( return broadcast +def resolve_broadcast_target( + engine: TrainEngine | InferenceEngine, device: Any +) -> tuple[Any, Any]: + """Offloaded engines gave the accelerator to the inference engine, so device + collectives are unusable and a gloo mirror is used. The mirror must span the + ranks of ``context_and_model_parallel_group`` or the broadcast deadlocks; + engines without one keep the pre-offload device broadcast. + """ + cpu_mirror_group = getattr(engine, "cpu_model_parallel_group", None) + if getattr(engine, "is_offload", False) and cpu_mirror_group is not None: + return cpu_mirror_group, "cpu" + return engine.context_and_model_parallel_group, device + + engine_bp = Blueprint("engine", __name__) # --------------------------------------------------------------------------- @@ -471,21 +485,20 @@ def execute_in_engine_thread(): ) if should_broadcast: logger.debug(f"Broadcasting RPC payload for method: {method_name}") - args_bcast = tensor_container_to( - args, current_platform.current_device() + bcast_group, bcast_device = resolve_broadcast_target( + engine, current_platform.current_device() ) + args_bcast = tensor_container_to(args, bcast_device) args_bcast = broadcast_tensor_container( args_bcast, src_rank=engine.current_data_parallel_head(), - group=engine.context_and_model_parallel_group, - ) - kwargs_bcast = tensor_container_to( - kwargs, current_platform.current_device() + group=bcast_group, ) + kwargs_bcast = tensor_container_to(kwargs, bcast_device) kwargs_bcast = broadcast_tensor_container( kwargs_bcast, src_rank=engine.current_data_parallel_head(), - group=engine.context_and_model_parallel_group, + group=bcast_group, ) logger.debug("Broadcasting RPC payload done.") diff --git a/areal/infra/scheduler/slurm.py b/areal/infra/scheduler/slurm.py index 54696ccd07..9edde79549 100644 --- a/areal/infra/scheduler/slurm.py +++ b/areal/infra/scheduler/slurm.py @@ -995,56 +995,79 @@ def create_workers(self, job: Job, *args, **kwargs) -> list[str]: ) target_workers = self._workers[colocate_role] - if num_workers != len(target_workers): - raise WorkerCreationError( - role, - "Replica count mismatch", - f"Colocated role must have same replica count as target " - f"({num_workers} != {len(target_workers)})", - ) + if num_workers == len(target_workers): + # Check if fork mode is enabled + if strategy.fork: + # Fork mode: spawn new processes on same nodes via /fork endpoint + return self.fork_workers(role, colocate_role) - # Check if fork mode is enabled - if strategy.fork: - # Fork mode: spawn new processes on same nodes via /fork endpoint - return self.fork_workers(role, colocate_role) + # Reuse existing workers - no new Slurm job submitted + worker_ids = [w.worker.id for w in target_workers] + self._colocated_roles[role] = colocate_role - # Reuse existing workers - no new Slurm job submitted - worker_ids = [w.worker.id for w in target_workers] - self._colocated_roles[role] = colocate_role + logger.info( + f"Role '{role}' colocated with '{colocate_role}': " + f"reusing workers {worker_ids}" + ) + return worker_ids + # Different worker counts: submit new job on the same nodes + # (e.g., AWEX colocation where rollout has TP-grouped instances) + target_job_id = self._jobs[colocate_role] + job_infos = query_jobs(slurm_ids=[target_job_id]) + if not job_infos: + raise WorkerCreationError( + role, f"Target job {target_job_id} not found in queue" + ) + colocation_nodelist = job_infos[0].host + spec = schedulings[0] + total_gpus = spec.gpu * replicas + nodes = max( + 1, + (total_gpus + self.n_gpus_per_node - 1) // self.n_gpus_per_node, + ) + cpus_per_task = spec.cpu + mem_per_task = spec.mem * 1024 logger.info( - f"Role '{role}' colocated with '{colocate_role}': " - f"reusing workers {worker_ids}" + f"Creating {replicas} workers for role '{role}' colocated with " + f"'{colocate_role}' on nodes {colocation_nodelist}: " + f"nodes={nodes}, cpus={cpus_per_task}, mem={mem_per_task}MB" ) - return worker_ids + nodelist = colocation_nodelist + elif strategy_type == SchedulingStrategyType.separation: + # Non-colocated: calculate nodes needed and submit new Slurm job + spec = schedulings[0] + total_gpus = spec.gpu * replicas + nodes = max( + 1, (total_gpus + self.n_gpus_per_node - 1) // self.n_gpus_per_node + ) + nodelist = spec.nodelist + cpus_per_task = spec.cpu + mem_per_task = spec.mem * 1024 # Convert GB to MB - if strategy_type != SchedulingStrategyType.separation: + logger.info( + f"Creating {replicas} workers for role '{role}': " + f"nodes={nodes}, gpus_per_node={self.n_gpus_per_node}, " + f"cpus={cpus_per_task}, mem={mem_per_task}MB" + ) + else: raise ValueError(f"Unknown scheduling strategy type: {strategy_type}") - # Non-colocated: calculate nodes needed and submit new Slurm job - spec = schedulings[0] - total_gpus = spec.gpu * replicas - nodes = max(1, (total_gpus + self.n_gpus_per_node - 1) // self.n_gpus_per_node) - nodelist = spec.nodelist - - # Calculate resource requirements - n_gpus_per_node = min( - self.n_gpus_per_node, (spec.gpu * replicas + nodes - 1) // nodes - ) - cpus_per_task = spec.cpu - mem_per_task = spec.mem * 1024 # Convert GB to MB - - logger.info( - f"Creating {replicas} workers for role '{role}': " - f"nodes={nodes}, gpus_per_node={n_gpus_per_node}, " - f"cpus={cpus_per_task}, mem={mem_per_task}MB" - ) # Generate sbatch script + # Colocated roles must not request GPU gres: the target role's job + # already holds the nodes' GPUs, so a second gres request would + # deadlock in the queue. Colocated workers address GPUs directly + # via base_gpu_id / CUDA_VISIBLE_DEVICES instead. + request_gpus = ( + 0 + if strategy_type == SchedulingStrategyType.colocation + else spec.gpu * replicas + ) sbatch_script = self._generate_sbatch_script( role=role, replicas=replicas, nodes=nodes, - total_gpus=spec.gpu * replicas, + total_gpus=request_gpus, cpus_per_task=cpus_per_task, mem_per_task=mem_per_task, schedulings=schedulings, diff --git a/areal/models/mcore/hf_save.py b/areal/models/mcore/hf_save.py index 5945c3d98a..f2423569e4 100644 --- a/areal/models/mcore/hf_save.py +++ b/areal/models/mcore/hf_save.py @@ -416,7 +416,7 @@ def _emit_per_expert_flat_expert_sd( for s in expert_specs: param = s.param if etp_size > 1: - params = all_gather_outputs[s.global_name].chunk(etp_size, dim=0) + params = all_gather_outputs.pop(s.global_name).chunk(etp_size, dim=0) else: params = [param] @@ -436,7 +436,7 @@ def _emit_per_expert_flat_expert_sd( ) for n, p in zip(converted_names, converted_params): assert n not in expert_sd, n - expert_sd[n] = p + expert_sd[n] = p.cpu() return expert_sd @@ -613,6 +613,10 @@ def save_weights_to_hf_with_mbridge_fast( # all-gather weights across the TP group and converts to HF format # Optimized via a single `all_gather_into_tensor_coalesced` call, which should be # faster than plain all_gather. + # + # Converted tensors are stashed on CPU and each gather buffer is popped as + # it is consumed, so the save never holds the gather transient plus the + # full HF state dict on GPU (OOMs memory-tight, e.g. colocated, setups). non_expert_sd = {} _all_gather_specs = [] all_gather_outputs = {} @@ -629,6 +633,7 @@ def save_weights_to_hf_with_mbridge_fast( ) for s, gathered_param in zip(_all_gather_specs, _all_gather_outputs): all_gather_outputs[s.global_name] = gathered_param + del _all_gather_outputs for s in non_expert_specs: param = s.param @@ -637,7 +642,7 @@ def save_weights_to_hf_with_mbridge_fast( if mpu.get_tensor_model_parallel_world_size() <= 1: infer_params = [param] else: - infer_params = all_gather_outputs[s.global_name].chunk( + infer_params = all_gather_outputs.pop(s.global_name).chunk( mpu.get_tensor_model_parallel_world_size(), dim=0 ) # Convert TE FP8 -> torch bf16 -> torch FP8 and finally save the native torch FP8 model @@ -666,7 +671,8 @@ def save_weights_to_hf_with_mbridge_fast( ) for n, p in zip(converted_names, converted_params): assert n not in non_expert_sd, n - non_expert_sd[n] = p + non_expert_sd[n] = p.cpu() + torch.cuda.empty_cache() # Split the state dict into shards and save the process's own shard. shards = split_state_dict_into_shards(non_expert_sd, n_shards) @@ -757,6 +763,7 @@ def _save_one_shard(x): ) for s, gathered_param in zip(_all_gather_specs, _all_gather_outputs): all_gather_outputs[s.global_name] = gathered_param + del _all_gather_outputs if stacked_experts and ep_size > 1: expert_sd = _emit_stacked_moe_expert_sd( bridge, @@ -817,6 +824,10 @@ def _save_one_shard(x): shard_idx = shard_offset + i for k in shard: weight_map[k] = output_filename.format(shard=shard_idx + 1) + # Free per-rank tensors before the NCCL metadata gather: comm setup + # for ep_pp_group needs scratch GPU memory. + del shards, expert_sd, all_gather_outputs + torch.cuda.empty_cache() ep_pp_group = mpu.get_expert_tensor_model_pipeline_parallel_group() weight_map_list = [None for _ in range(dist.get_world_size(ep_pp_group))] dist.all_gather_object( diff --git a/areal/trainer/rl_trainer.py b/areal/trainer/rl_trainer.py index bdb6b22a72..117ff61c8c 100644 --- a/areal/trainer/rl_trainer.py +++ b/areal/trainer/rl_trainer.py @@ -108,6 +108,23 @@ def __init__( config: PPOConfig, train_dataset: Dataset | None = None, valid_dataset: Dataset | None = None, + ): + try: + self._init_impl(config, train_dataset, valid_dataset) + except Exception: + logger.error( + "PPOTrainer construction failed; tearing down partially " + "created workers", + exc_info=True, + ) + self.close() + raise + + def _init_impl( + self, + config: PPOConfig, + train_dataset: Dataset | None = None, + valid_dataset: Dataset | None = None, ): rank = int(os.getenv("RANK", "0")) if is_single_controller(): @@ -144,6 +161,12 @@ def __init__( self._should_offload_teacher = ( config.teacher is not None and config.teacher.offload ) + # In colocate (awex) mode the GPU switch between rollout and training + # is managed by the AWEX adapter (manual offload/onload + tagged SGLang + # release), not by the TMS-based offload machinery below. + if self._is_v1_awex_colocate(config): + self._should_offload_rollout = False + self._should_offload_actor = False # Validate config before proceeding with weight initialization self._validate_cfg() @@ -309,6 +332,26 @@ def __init__( if self._should_offload_actor: self._offload_model(self.actor, role="actor") + # In colocate (awex) mode, offload training weights before SGLang starts + # so that GPU memory is available for inference engine allocation. + # Uses adapter-based manual offload (not TMS), so enable_offload is not required. + self._awex_meta_server_addr: str | None = None + if self._is_v1_awex_colocate(config): + from awex.meta.meta_server import start_meta_server + + from areal.utils.network import gethostip + + host, port = start_meta_server() + if host in ("0.0.0.0", ""): + host = gethostip() + self._awex_meta_server_addr = f"{host}:{port}" + logger.info( + "Started MetaServer on controller at %s", + self._awex_meta_server_addr, + ) + self.actor.init_awex_adapter(meta_server_addr=self._awex_meta_server_addr) + self.actor.offload() + # Initialize inference with LoRA path self.rollout = self._init_rollout( config.rollout, is_eval=False, lora_path=initial_lora_path @@ -395,6 +438,10 @@ def __init__( ) else: self.weight_update_meta = WeightUpdateMeta.from_fsdp_xccl(**xccl_kwargs) + elif self.config.actor.weight_update_mode == "awex": + self.weight_update_meta = WeightUpdateMeta.from_awex( + meta_server_addr=self._awex_meta_server_addr, + ) else: raise ValueError( f"Invalid weight update mode: {self.config.actor.weight_update_mode}" @@ -421,6 +468,9 @@ def __init__( self.train_dataloader, inference_engine=self.rollout, weight_update_meta=self.weight_update_meta, + # Recompute placement instead of reusing _should_offload_rollout: + # AWEX clears that flag because it drives the handover itself. + colocated_rollout=self._is_actor_rollout_colocated(config), ) # After recovery, sync the staleness manager so its capacity formula @@ -454,6 +504,25 @@ def _is_actor_rollout_colocated(self, config: PPOConfig) -> bool: self._is_colocation(rollout_s) and rollout_s.target == "actor" ) + def _is_v1_awex_colocate(self, config: PPOConfig) -> bool: + """Whether this run is the v1 AWEX colocated actor-rollout setup. + + ``weight_update_mode`` alone is not enough: controller v2 selects AWEX + from ``use_lora`` and never reads that field, so a v2 separation run may + legitimately carry ``weight_update_mode="awex"`` and would otherwise + take the v1 colocation handover. + + The scheduling strategy is deliberately not part of this check. v1 AWEX + only exists colocated and runs opt in through ``weight_update_mode`` + while leaving actor and rollout on the default (separation) strategy; + requiring colocation here skips the meta-server handoff, every training + worker then starts its own server, and the run waits on ``infer_conf`` + forever. + """ + return ( + config.actor._version == "v1" and config.actor.weight_update_mode == "awex" + ) + def _onload_model(self, engine, role: str) -> None: with ( stats_tracker.record_timing(f"{role}_onload"), @@ -694,6 +763,20 @@ def train( if self._should_offload_teacher: self._offload_model(self.teacher, role="teacher") + # In colocate (awex) mode: switch GPU from inference to training. + # Release SGLang KV cache + weights to free GPU for actor. + if self._is_v1_awex_colocate(self.config): + logger.info("[AWEX] colocate: pausing rollout...") + self.rollout.pause() + logger.info("[AWEX] colocate: pause_generation_sync...") + self.rollout.pause_generation_sync() + logger.info("[AWEX] colocate: offload kv_cache...") + self.rollout.offload(tags=["kv_cache"]) + logger.info("[AWEX] colocate: offload weights...") + self.rollout.offload(tags=["weights"]) + logger.info("[AWEX] colocate: offload done, onloading actor...") + self.actor.onload() + if self._should_offload_actor: self._onload_model(self.actor, role="actor") if config.actor.should_compute_prox_logp(): @@ -769,6 +852,22 @@ def train( if self._should_offload_critic: self._offload_model(self.critic, role="critic") + # Save BEFORE update_weights. In AWEX colocate mode the + # transfer ends with actor weights offloaded, so saving afterwards + # would resume weights onto a card already crowded by the + # fully-resumed rollout plus transfer staging leftovers and the HF + # saver's TP coalesced all-gather transient can OOM. Here the + # actor weights are still onloaded from ppo_update (no resume + # needed) and MegatronEngine.save() drops the dead fp32 grad + # buffers to fund the transient. Weights are identical on both + # sides of the transfer, so the checkpoint content is unchanged. + if self._is_v1_awex_colocate(config): + self._save_training_state( + epoch=epoch, + epoch_step=step, + global_step=global_step, + ) + # pause inference for updating weights, save, and evaluation self.rollout.pause() @@ -795,26 +894,11 @@ def train( if self.eval_rollout is not None: self.eval_rollout.set_version(new_version) - with ( - stats_tracker.record_timing("save"), - perf_tracer.trace_scope( - "train.save", - category=Category.IO, - args={"global_step": global_step}, - ), - ): - self._save_hf(epoch=epoch, epoch_step=step, global_step=global_step) - - with ( - stats_tracker.record_timing("checkpoint_for_recover"), - perf_tracer.trace_scope( - "train.checkpoint", - category=Category.IO, - args={"global_step": global_step}, - ), - ): - self._save_recover_checkpoint( - epoch=epoch, epoch_step=step, global_step=global_step + if not self._is_v1_awex_colocate(config): + self._save_training_state( + epoch=epoch, + epoch_step=step, + global_step=global_step, ) # Offload actor before eval @@ -879,6 +963,37 @@ def train( self._save_perf_tracer(step=global_step) + def _save_training_state( + self, + *, + epoch: int, + epoch_step: int, + global_step: int, + ) -> None: + with ( + stats_tracker.record_timing("save"), + perf_tracer.trace_scope( + "train.save", + category=Category.IO, + args={"global_step": global_step}, + ), + ): + self._save_hf(epoch=epoch, epoch_step=epoch_step, global_step=global_step) + + with ( + stats_tracker.record_timing("checkpoint_for_recover"), + perf_tracer.trace_scope( + "train.checkpoint", + category=Category.IO, + args={"global_step": global_step}, + ), + ): + self._save_recover_checkpoint( + epoch=epoch, + epoch_step=epoch_step, + global_step=global_step, + ) + def close(self): self.saver.finalize() if hasattr(self, "_train_rdataset") and self._train_rdataset is not None: @@ -1093,6 +1208,9 @@ def _init_rollout( pp_size=self.rollout_alloc.parallel.pp_size, base_gpu_id=0, ) + if self._is_v1_awex_colocate(self.config): + server_args["awex_colocate_mode"] = True + server_args["awex_meta_server_addr"] = self._awex_meta_server_addr elif rollout_backend == "vllm": if self.config.rollout.return_routed_experts: raise ValueError( @@ -1358,15 +1476,27 @@ def _validate_cfg(self): "offload is enabled. Please set enable_offload=True." ) - if ( - self._is_actor_rollout_colocated(self.config) - and self.config.actor.weight_update_mode != "disk" - ): + if self._is_actor_rollout_colocated( + self.config + ) and self.config.actor.weight_update_mode not in ("disk", "awex"): raise ValueError( - "weight_update_mode must be 'disk' when colocation scheduling is enabled. " - "Please set actor.weight_update_mode=disk." + "weight_update_mode must be 'disk' or 'awex' when colocation " + "scheduling is enabled. Please set actor.weight_update_mode " + "to one of them." ) + if self._is_v1_awex_colocate(self.config): + if actor_backend != "megatron": + raise ValueError( + "weight_update_mode='awex' requires Megatron actor training " + f"backend, got {actor_backend!r}." + ) + if rollout_backend != "sglang": + raise ValueError( + "weight_update_mode='awex' requires SGLang rollout backend, " + f"got {rollout_backend!r}." + ) + if rollout_backend == "vllm" and self.config.rollout.return_routed_experts: raise ValueError( "return_routed_experts is only supported with SGLang backend. " diff --git a/areal/utils/logging.py b/areal/utils/logging.py index 89d1966d2b..f54a384e1f 100644 --- a/areal/utils/logging.py +++ b/areal/utils/logging.py @@ -126,6 +126,10 @@ "InferenceRouter": "white", "InferenceGateway": "white", "RPCGuard": "white", + # AWEX weight exchange - cyan (compute backend) + "AwexColocate": "light_cyan", + "AwexColocateReader": "light_cyan", + "AwexSGLangPlugin": "light_cyan", } # Prefix patterns checked in order (first match wins) diff --git a/areal/utils/recover.py b/areal/utils/recover.py index 9af81b7efb..2ff898ea16 100644 --- a/areal/utils/recover.py +++ b/areal/utils/recover.py @@ -2,6 +2,7 @@ from __future__ import annotations import dataclasses +import inspect import json import os import pickle @@ -217,6 +218,51 @@ def _normalize_recover_engines( return engine return {"default": engine} + @staticmethod + def _should_run_awex_colocate_transfer( + inference_engine: InferenceEngine | None, + weight_update_meta: WeightUpdateMeta | None, + colocated_rollout: bool, + ) -> bool: + """Whether recovery must drive the AWEX colocate pre-transfer sequence. + + The transport type alone is not enough: v2 selects AWEX for every + non-LoRA run regardless of placement, so the caller has to state whether + actor and rollout physically share devices. + """ + return ( + inference_engine is not None + and getattr(weight_update_meta, "type", None) == "awex" + and colocated_rollout + ) + + @staticmethod + def _require_colocate_rollout_protocol( + inference_engine: InferenceEngine, + ) -> None: + missing = [] + if not callable(getattr(inference_engine, "pause_generation_sync", None)): + missing.append("pause_generation_sync()") + + offload = getattr(inference_engine, "offload", None) + if not callable(offload): + missing.append("offload(tags=...)") + else: + try: + accepts_tags = "tags" in inspect.signature(offload).parameters + except (TypeError, ValueError): + accepts_tags = True + if not accepts_tags: + missing.append("offload(tags=...)") + + if missing: + raise NotImplementedError( + "Colocated AWEX recovery needs a rollout engine implementing " + f"{', '.join(missing)}, which {type(inference_engine).__name__} " + "does not provide. Disable `recover.mode` or run this " + "configuration without actor-rollout colocation." + ) + def dump( self, engine: TrainEngine @@ -279,6 +325,7 @@ def load( inference_engine: InferenceEngine | None = None, weight_update_meta: WeightUpdateMeta | None = None, inference_engine_update_from: str = "default", + colocated_rollout: bool = False, ) -> RecoverInfo | None: if self.config.mode in ("disabled", "off"): return @@ -309,19 +356,54 @@ def load( stats_logger.load_state_dict(recover_info.stats_logger_info) dataloader.load_state_dict(recover_info.dataloader_info) - for name, engine_ in normalized_engine.items(): - self._load_checkpoint(engine_, name=name) global_step = recover_info.last_step_info.global_step + recovery_version = global_step + 1 + + is_awex_colocate = self._should_run_awex_colocate_transfer( + inference_engine=inference_engine, + weight_update_meta=weight_update_meta, + colocated_rollout=colocated_rollout, + ) + if is_awex_colocate: + self._require_colocate_rollout_protocol(inference_engine) + + if not is_awex_colocate: + for name, engine_ in normalized_engine.items(): + self._load_checkpoint(engine_, name=name) if inference_engine is not None: assert weight_update_meta is not None update_engine = normalized_engine[inference_engine_update_from] - recovery_version = global_step + 1 versioned_meta = weight_update_meta.with_version(recovery_version) update_engine.connect_engine(inference_engine, versioned_meta) inference_engine.pause() - update_engine.update_weights(versioned_meta) - inference_engine.resume() + try: + # AWEX colocate transfer requires the full engine-level + # pause/offload protocol, not just the controller pause. The + # sglang plugin's patched event loop only drains the weight- + # update queue while scheduler._engine_paused is True (set by + # pause_generation), and the reader-side protocol expects the + # engine's kv/weights released before the writer publishes. + # Without this the recover-path transfer deadlocks: reader + # never consumes the queued version marker, writer blocks on + # weights_update_finished forever. + # Mirror of the trainer's pre-update sequence; the reverse + # side (kv_cache onload) happens inside update_weights. + if is_awex_colocate: + inference_engine.pause_generation_sync() + inference_engine.offload(tags=["kv_cache"]) + inference_engine.offload(tags=["weights"]) + # Load the actor checkpoint only after the colocated + # rollout engine has released its GPU memory; loading + # first would stack DCP weights/optimizer on top of the + # still-resident sglang allocation and risk OOM. + for name, engine_ in normalized_engine.items(): + self._load_checkpoint(engine_, name=name) + update_engine.update_weights(versioned_meta) + finally: + # Always resume: leaving rollout paused after a failed + # checkpoint load or transfer would hang every later step. + inference_engine.resume() update_engine.set_version(recovery_version) inference_engine.set_version(recovery_version) return recover_info diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 89fbe1b31f..e32922380d 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -365,7 +365,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -440,7 +440,7 @@ Configuration for PPO critic model, a subclass of a TrainEngine. | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -488,7 +488,7 @@ Core configuration for model training, including optimization and backend settin | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -607,6 +607,7 @@ https://github.com/sgl-project/sglang for detailed documentation. | `triton_attention_reduce_in_fp32` | boolean | `False` | - | | `triton_attention_num_kv_splits` | integer | `8` | - | | `num_continuous_decode_steps` | integer | `1` | - | +| `load_format` | string | `"auto"` | - | | `enable_memory_saver` | boolean | `False` | - | | `allow_auto_truncate` | boolean | `False` | - | | `attention_backend` | string \| None | `"fa3"` | - | @@ -1004,7 +1005,7 @@ fields. | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index eb7d182f42..91c5e8a99c 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -363,7 +363,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -438,7 +438,7 @@ Configuration for PPO critic model, a subclass of a TrainEngine. | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -486,7 +486,7 @@ Core configuration for model training, including optimization and backend settin | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -605,6 +605,7 @@ https://github.com/sgl-project/sglang for detailed documentation. | `triton_attention_reduce_in_fp32` | boolean | `False` | - | | `triton_attention_num_kv_splits` | integer | `8` | - | | `num_continuous_decode_steps` | integer | `1` | - | +| `load_format` | string | `"auto"` | - | | `enable_memory_saver` | boolean | `False` | - | | `allow_auto_truncate` | boolean | `False` | - | | `attention_backend` | string \| None | `"fa3"` | - | @@ -1002,7 +1003,7 @@ fields. | `grad_reduce_dtype` | string | `"float32"` | Gradient reduction data type. | | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | -| `weight_update_mode` | string | `"xccl"` | Weight update backend type. **Choices:** `disk`, `xccl` | +| `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | diff --git a/examples/math/gsm8k_grpo_awex_colocate.yaml b/examples/math/gsm8k_grpo_awex_colocate.yaml new file mode 100644 index 0000000000..3306ae442d --- /dev/null +++ b/examples/math/gsm8k_grpo_awex_colocate.yaml @@ -0,0 +1,194 @@ +experiment_name: gsm8k-grpo-awex-colocate +trial_name: trial0 + +seed: 1 +enable_offload: true +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 1 + n_gpus_per_node: 2 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + + +scheduler: + type: null + +rollout: + backend: "sglang:d1t2p1" + 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: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: false + agent: + mode: inline + export_style: individual + turn_discount: 1.0 + +gconfig: + n_samples: 2 + min_new_tokens: 0 + max_new_tokens: 256 + max_tokens: 2048 + greedy: false + temperature: 1.0 + +actor: + backend: "megatron:d2" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen3-0.6B + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 1024 + packing_algorithm: ffd + optimizer: + type: adam + lr: 1.70e-5 + weight_decay: 0.017 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + 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 + rejection_sampling: + metric: ratio + upper: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + weight_update_mode: awex + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: + # SGLang opens its own memory-saver regions for the colocated rollout. + # An auto-opened region on the training side would nest inside them and + # trip the memory saver's assertion, so keep the hook loadable but idle. + TMS_INIT_ENABLE: "0" + TMS_INIT_ENABLE_CPU_BACKUP: "0" + +ref: + backend: ${actor.backend} + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + packing_algorithm: ffd + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SGLang +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.3 + enable_memory_saver: true + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 32768 + gpu_memory_utilization: 0.8 + +# datasets +train_dataset: + batch_size: 2 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + max_length: 1024 + +valid_dataset: + batch_size: 2 + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + 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 + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false diff --git a/tests/test_alloc_conf_import_side_effects.py b/tests/test_alloc_conf_import_side_effects.py new file mode 100644 index 0000000000..ff306f388d --- /dev/null +++ b/tests/test_alloc_conf_import_side_effects.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Guards that the CUDA allocator config is not rewritten at import time.""" + +import os +import subprocess +import sys + +import pytest + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def _run(code: str, env_extra: dict[str, str]) -> subprocess.CompletedProcess: + env = {**os.environ, "PYTHONPATH": REPO_ROOT, **env_extra} + return subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, env=env + ) + + +class TestNoImportTimeAllocConfRewrite: + @pytest.mark.parametrize("role", ["actor", "ref", "rollout"]) + def test_importing_areal_leaves_the_allocator_config_alone(self, role): + code = ( + "import sys, os, json;" + f" sys.argv = ['prog', '--role', {role!r}];" + " import areal;" + " print(json.dumps(os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '')))" + ) + result = _run( + code, + { + "PYTORCH_CUDA_ALLOC_CONF": "", + "AWEX_ACTOR_ALLOC_CONF": "expandable_segments:True", + }, + ) + + assert result.returncode == 0, result.stderr[-1500:] + assert result.stdout.strip().splitlines()[-1] == '""', ( + "importing areal rewrote PYTORCH_CUDA_ALLOC_CONF; the per-role " + "allocator config belongs in each role's scheduling_spec.env_vars" + ) + + def test_awex_actor_alloc_conf_is_gone(self): + hits = subprocess.run( + ["grep", "-rIl", "AWEX_ACTOR_ALLOC_CONF", "areal"], + cwd=REPO_ROOT, + capture_output=True, + text=True, + ) + assert hits.stdout.strip() == "", ( + f"AWEX_ACTOR_ALLOC_CONF still referenced in: {hits.stdout.strip()}" + ) + + +class TestSglangPluginRejectsExpandableSegments: + def test_import_fails_loudly_when_expandable_segments_is_enabled(self): + result = _run( + "import areal.engine.awex.sglang_plugin", + {"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}, + ) + + assert result.returncode != 0 + assert "expandable segments" in result.stderr.lower(), result.stderr[-1500:] + + def test_import_is_allowed_without_expandable_segments(self): + result = _run( + "import areal.engine.awex.sglang_plugin", + {"PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:128"}, + ) + + assert "expandable segments" not in result.stderr.lower() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_awex_colocate_device_id.py b/tests/test_awex_colocate_device_id.py new file mode 100644 index 0000000000..b8c83f2982 --- /dev/null +++ b/tests/test_awex_colocate_device_id.py @@ -0,0 +1,26 @@ +from areal.engine.awex.colocate_writer import resolve_physical_gpu_id + + +def test_physical_gpu_id_maps_through_visible_devices(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7") + + assert resolve_physical_gpu_id(0) == 4 + assert resolve_physical_gpu_id(3) == 7 + + +def test_physical_gpu_id_is_identity_without_visible_devices(monkeypatch): + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + + assert resolve_physical_gpu_id(2) == 2 + + +def test_physical_gpu_id_falls_back_on_non_integer_entries(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-abc,GPU-def") + + assert resolve_physical_gpu_id(1) == 1 + + +def test_physical_gpu_id_falls_back_when_index_out_of_range(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4") + + assert resolve_physical_gpu_id(2) == 2 diff --git a/tests/test_awex_sglang_version_guard.py b/tests/test_awex_sglang_version_guard.py new file mode 100644 index 0000000000..404237316c --- /dev/null +++ b/tests/test_awex_sglang_version_guard.py @@ -0,0 +1,34 @@ +"""The colocate plugin patches SGLang internals, so it pins the versions it knows.""" + +from unittest import mock + +import pytest + +from areal.engine.awex import sglang_plugin + + +def test_register_rejects_unverified_sglang_version(): + with mock.patch.object( + sglang_plugin.pkg_version, "get_version", return_value="0.5.11" + ): + with pytest.raises(RuntimeError, match="0.5.11"): + sglang_plugin.assert_supported_sglang_version() + + +def test_register_accepts_verified_sglang_versions(): + for version in sglang_plugin.SUPPORTED_SGLANG_VERSIONS: + with mock.patch.object( + sglang_plugin.pkg_version, "get_version", return_value=version + ): + sglang_plugin.assert_supported_sglang_version() + + +def test_error_names_the_supported_versions(): + with mock.patch.object( + sglang_plugin.pkg_version, "get_version", return_value="0.4.0" + ): + with pytest.raises(RuntimeError) as excinfo: + sglang_plugin.assert_supported_sglang_version() + message = str(excinfo.value) + for version in sglang_plugin.SUPPORTED_SGLANG_VERSIONS: + assert version in message diff --git a/tests/test_engine_blueprint_offload_broadcast.py b/tests/test_engine_blueprint_offload_broadcast.py new file mode 100644 index 0000000000..74cac05f66 --- /dev/null +++ b/tests/test_engine_blueprint_offload_broadcast.py @@ -0,0 +1,40 @@ +"""Offloaded engines broadcast over CPU only when they expose a CPU mirror group.""" + +from types import SimpleNamespace + +from areal.infra.rpc.guard.engine_blueprint import resolve_broadcast_target + + +def _engine(is_offload: bool, with_cpu_group: bool): + fields = { + "is_offload": is_offload, + "context_and_model_parallel_group": "device-group", + } + if with_cpu_group: + fields["cpu_model_parallel_group"] = "cpu-group" + return SimpleNamespace(**fields) + + +def test_offloaded_engine_uses_cpu_mirror_group(): + group, device = resolve_broadcast_target( + _engine(is_offload=True, with_cpu_group=True), device="cuda:0" + ) + assert group == "cpu-group" + assert device == "cpu" + + +def test_engine_without_cpu_group_keeps_device_broadcast(): + """FSDP tracks is_offload but has no CPU mirror group.""" + group, device = resolve_broadcast_target( + _engine(is_offload=True, with_cpu_group=False), device="cuda:0" + ) + assert group == "device-group" + assert device == "cuda:0" + + +def test_resident_engine_keeps_device_broadcast(): + group, device = resolve_broadcast_target( + _engine(is_offload=False, with_cpu_group=True), device="cuda:0" + ) + assert group == "device-group" + assert device == "cuda:0" diff --git a/tests/test_examples.py b/tests/test_examples.py index 7b8cf41805..9c5fe27d70 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -439,6 +439,55 @@ def test_gsm8k_ppo_colocate(tmp_path_factory): assert success, "GSM8K PPO colocated example failed" +@pytest.mark.sglang +@pytest.mark.multi_gpu +def test_gsm8k_grpo_awex_colocate(tmp_path_factory, monkeypatch): + """Actor and rollout time-share the GPUs, syncing weights through AWEX.""" + experiments_path = tmp_path_factory.mktemp("experiments") + name_resolve_path = tmp_path_factory.mktemp("name_resolve") + model_path = get_model_path( + "/storage/openpsi/models/Qwen__Qwen3-0.6B", "Qwen/Qwen3-0.6B" + ) + dataset_path = get_dataset_path("/storage/openpsi/data/gsm8k", "openai/gsm8k") + + example_file = "examples/math/gsm8k_rl.py" + config_name = "examples/math/gsm8k_grpo.yaml" + + monkeypatch.setenv("AREAL_ALLOW_DEFAULT_ADMIN_KEY", "1") + + success = run_async_task( + run_example, + example_file, + config_name, + # One inference server keeps every rank in a single NCCL group, and two + # GPUs give the ranks a peer to talk to. + "rollout.backend=sglang:d1t2p1", + "actor.backend=megatron:d2", + "actor.weight_update_mode=awex", + "enable_offload=True", + "gconfig.n_samples=2", + "gconfig.max_new_tokens=256", + "sglang.mem_fraction_static=0.3", + "+sglang.enable_memory_saver=True", + # Keep the memory-saver hook available but let SGLang open its own + # regions; an auto-opened region would nest and trip TMS's assertion. + "+actor.scheduling_spec.0.env_vars.TMS_INIT_ENABLE=0", + "+actor.scheduling_spec.0.env_vars.TMS_INIT_ENABLE_CPU_BACKUP=0", + "actor.mb_spec.max_tokens_per_mb=1024", + "train_dataset.batch_size=2", + "valid_dataset.batch_size=2", + f"train_dataset.path={dataset_path}", + f"valid_dataset.path={dataset_path}", + "cluster.n_gpus_per_node=2", + f"cluster.fileroot={str(experiments_path)}", + f"cluster.name_resolve.nfs_record_root={str(name_resolve_path)}", + f"actor.path={model_path}", + "scheduler.type=local", + timeout=900, + ) + assert success, "GSM8K GRPO AWEX colocated example failed" + + @pytest.mark.ci @pytest.mark.parametrize( "rollout_backend,actor_backend", diff --git a/tests/test_gloo_mirror_group_gating.py b/tests/test_gloo_mirror_group_gating.py new file mode 100644 index 0000000000..b95f5830d8 --- /dev/null +++ b/tests/test_gloo_mirror_group_gating.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The gloo mirror group is only built when the engine can offload.""" + +import ast +import pathlib + +import pytest + +import areal.engine.megatron_engine as megatron_engine_module + + +def _init_group_method() -> ast.FunctionDef: + source = pathlib.Path(megatron_engine_module.__file__).read_text() + for node in ast.walk(ast.parse(source)): + if ( + isinstance(node, ast.FunctionDef) + and node.name == "_init_context_and_model_parallel_group" + ): + return node + raise AssertionError("_init_context_and_model_parallel_group not found") + + +def _gloo_group_calls(node: ast.AST) -> list[ast.Call]: + calls = [] + for sub in ast.walk(node): + if not isinstance(sub, ast.Call): + continue + if not (isinstance(sub.func, ast.Attribute) and sub.func.attr == "new_group"): + continue + for kw in sub.keywords: + if ( + kw.arg == "backend" + and isinstance(kw.value, ast.Constant) + and kw.value.value == "gloo" + ): + calls.append(sub) + return calls + + +class TestGlooMirrorGroupIsGated: + def test_the_gloo_group_is_built_only_under_an_offload_guard(self): + method = _init_group_method() + gloo_calls = _gloo_group_calls(method) + + assert len(gloo_calls) == 1, ( + f"expected exactly one gloo new_group call, found {len(gloo_calls)}" + ) + target = gloo_calls[0] + + guarded = False + for node in ast.walk(method): + if not isinstance(node, ast.If): + continue + if any(c is target for c in _gloo_group_calls(node)): + test_src = ast.dump(node.test) + if "offload" in test_src: + guarded = True + assert guarded, ( + "the gloo mirror group is created unconditionally; a separation run " + "pays one collective per data-parallel group for a group that " + "resolve_broadcast_target never uses" + ) + + def test_the_mirror_group_defaults_to_none(self): + source = pathlib.Path(megatron_engine_module.__file__).read_text() + + assert "self._cpu_model_parallel_group = None" in source, ( + "resolve_broadcast_target falls back on a None mirror group, so the " + "attribute must exist even when the group is not built" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_recover.py b/tests/test_recover.py index 653e63a004..7e4c80cc20 100644 --- a/tests/test_recover.py +++ b/tests/test_recover.py @@ -7,7 +7,11 @@ from areal.api.cli_args import RecoverConfig from areal.api.io_struct import FinetuneSpec, StepInfo -from areal.utils.recover import RecoverHandler, check_if_auto_recover, check_if_recover +from areal.utils.recover import ( + RecoverHandler, + check_if_auto_recover, + check_if_recover, +) from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) @@ -234,3 +238,75 @@ def test_dump_rejects_gateway_train_controller(self, mode): assert "GatewayTrainController" in str(exc_info.value) assert "recover.mode" in str(exc_info.value) + + +class TestAwexColocateGate: + """The AWEX pre-transfer sequence must run only for colocated rollouts.""" + + @staticmethod + def _awex_meta(): + return Mock(type="awex") + + def test_awex_transport_without_colocation_is_not_colocate(self): + assert not RecoverHandler._should_run_awex_colocate_transfer( + inference_engine=Mock(), + weight_update_meta=self._awex_meta(), + colocated_rollout=False, + ) + + def test_awex_transport_with_colocation_is_colocate(self): + assert RecoverHandler._should_run_awex_colocate_transfer( + inference_engine=Mock(), + weight_update_meta=self._awex_meta(), + colocated_rollout=True, + ) + + @pytest.mark.parametrize("meta_type", ["disk", "xccl"]) + def test_non_awex_transport_is_never_colocate(self, meta_type): + assert not RecoverHandler._should_run_awex_colocate_transfer( + inference_engine=Mock(), + weight_update_meta=Mock(type=meta_type), + colocated_rollout=True, + ) + + def test_missing_inference_engine_is_not_colocate(self): + assert not RecoverHandler._should_run_awex_colocate_transfer( + inference_engine=None, + weight_update_meta=self._awex_meta(), + colocated_rollout=True, + ) + + def test_meta_without_type_attribute_is_not_colocate(self): + assert not RecoverHandler._should_run_awex_colocate_transfer( + inference_engine=Mock(), + weight_update_meta=None, + colocated_rollout=True, + ) + + +class TestColocateRolloutProtocol: + """Engines lacking the colocate protocol must fail before any side effect.""" + + def test_engine_with_full_protocol_is_accepted(self): + engine = Mock(spec=["pause_generation_sync", "offload"]) + engine.offload = lambda tags=None: None + + RecoverHandler._require_colocate_rollout_protocol(engine) + + def test_engine_without_pause_generation_sync_is_rejected(self): + engine = Mock(spec=["offload"]) + engine.offload = lambda tags=None: None + + with pytest.raises(NotImplementedError) as exc_info: + RecoverHandler._require_colocate_rollout_protocol(engine) + + assert "pause_generation_sync" in str(exc_info.value) + + def test_engine_with_untagged_offload_is_rejected(self): + engine = Mock(spec=["pause_generation_sync", "offload"]) + engine.offload = lambda: None + + with pytest.raises(NotImplementedError) as exc_info: + RecoverHandler._require_colocate_rollout_protocol(engine) + + assert "tags" in str(exc_info.value) diff --git a/tests/test_sglang_pause_for_offload.py b/tests/test_sglang_pause_for_offload.py new file mode 100644 index 0000000000..458db49421 --- /dev/null +++ b/tests/test_sglang_pause_for_offload.py @@ -0,0 +1,50 @@ +from types import SimpleNamespace + +from areal.engine.sglang_remote import SGLangBackend +from areal.infra.remote_inf_engine import RemoteInfEngine + + +def test_pause_generation_drains_then_pauses_scheduler(): + requests = [] + engine = RemoteInfEngine.__new__(RemoteInfEngine) + engine.backend = SGLangBackend() + engine.config = SimpleNamespace(pause_grace_period=0) + engine._run_request_on_all_servers = requests.append + + pause = getattr( + RemoteInfEngine.pause_generation, + "__wrapped__", + RemoteInfEngine.pause_generation, + ) + pause(engine) + + assert [request.payload for request in requests] == [ + {}, + {"mode": "in_place"}, + ] + + +def test_default_pause_request_matches_upstream_payload(): + assert SGLangBackend().get_pause_request().payload == {} + + +def test_pause_generation_keeps_single_request_backends_compatible(): + requests = [] + + class _SingleRequestBackend: + def get_pause_request(self): + return SimpleNamespace(payload={"backend": "single"}) + + engine = RemoteInfEngine.__new__(RemoteInfEngine) + engine.backend = _SingleRequestBackend() + engine.config = SimpleNamespace(pause_grace_period=0) + engine._run_request_on_all_servers = requests.append + + pause = getattr( + RemoteInfEngine.pause_generation, + "__wrapped__", + RemoteInfEngine.pause_generation, + ) + pause(engine) + + assert [request.payload for request in requests] == [{"backend": "single"}] diff --git a/tests/test_v1_awex_colocate_gate.py b/tests/test_v1_awex_colocate_gate.py new file mode 100644 index 0000000000..0f52001fc1 --- /dev/null +++ b/tests/test_v1_awex_colocate_gate.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Guards that v1 AWEX colocation never activates for controller v2.""" + +import ast +import pathlib + +import pytest + +import areal.trainer.rl_trainer as rl_trainer_module +from areal.api.cli_args import SchedulingStrategy +from areal.trainer.rl_trainer import PPOTrainer + + +def _config(version: str, mode: str, colocated: bool): + actor = type( + "Actor", + (), + { + "_version": version, + "weight_update_mode": mode, + "scheduling_strategy": SchedulingStrategy( + type="colocation" if colocated else "separation", + target="rollout" if colocated else "", + ), + }, + )() + rollout = type("Rollout", (), {"scheduling_strategy": None})() + return type("Cfg", (), {"actor": actor, "rollout": rollout})() + + +class TestV1AwexColocateGate: + @pytest.mark.parametrize( + "version,mode,colocated,expected", + [ + ("v1", "awex", True, True), + ("v1", "awex", False, True), + ("v2", "awex", True, False), + ("v2", "awex", False, False), + ("v1", "xccl", True, False), + ("v1", "disk", True, False), + ], + ) + def test_gate_selects_v1_awex_whatever_the_strategy( + self, version, mode, colocated, expected + ): + trainer = object.__new__(PPOTrainer) + cfg = _config(version, mode, colocated) + + assert PPOTrainer._is_v1_awex_colocate(trainer, cfg) is expected + + def test_the_default_separation_strategy_still_selects_v1_awex(self): + """AWEX runs opt in with weight_update_mode, not with a strategy. + + examples/math/gsm8k_grpo.yaml leaves actor and rollout on the dataclass + default, so requiring colocation here skips the meta-server handoff and + every worker starts its own; the run then waits on 'infer_conf' forever. + """ + trainer = object.__new__(PPOTrainer) + default = SchedulingStrategy() + + assert default.type == "separation" + assert ( + PPOTrainer._is_v1_awex_colocate( + trainer, _config("v1", "awex", colocated=False) + ) + is True + ) + + +class TestNoUngatedAwexChecks: + def test_weight_update_mode_awex_is_only_compared_in_the_meta_dispatch(self): + source = pathlib.Path(rl_trainer_module.__file__).read_text() + tree = ast.parse(source) + + gate = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_is_v1_awex_colocate" + ) + gate_lines = range(gate.lineno, gate.end_lineno + 1) + + bare = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Compare): + continue + if not ( + isinstance(node.left, ast.Attribute) + and node.left.attr == "weight_update_mode" + ): + continue + if node.lineno in gate_lines: + continue + for comparator in node.comparators: + if isinstance(comparator, ast.Constant) and comparator.value == "awex": + bare.append(node.lineno) + + # The meta dispatch keeps its comparison: it sits in an elif chain that + # controller v2 already short-circuits. + assert len(bare) <= 1, ( + 'ungated `weight_update_mode == "awex"` checks at lines ' + f"{bare}; use _is_v1_awex_colocate so controller v2 separation " + "runs are unaffected" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])