[Feature] Add patch to accelerate SGLang weight loading

areal/api/cli_args.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import os
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
@@ -599,6 +600,11 @@ class SGLangConfig:
     # The interval (in decoding iterations) to log throughput
     # and update prometheus metrics
     decode_log_interval: int = 1
+    # Extra loader arguments
+    # NOTE: These arguments will be parsed into a dict json-string
+    # and passed as `model_loader_extra_config` to SGLang.
+    enable_multithread_load: bool = False
+    enable_fast_load: bool = False
 
     # Use staticmethod to make OmegaConf happy.
     @staticmethod
@@ -649,6 +655,19 @@ def build_args(
     ):
         # Map "all-linear" to "all"
         args: Dict = conf_as_dict(sglang_config)
+        if sglang_config.enable_multithread_load or sglang_config.enable_fast_load:
+            assert pkg_version.is_version_equal(
+                "sglang", "0.5.2"
+            ), f"Customized model loading requires exact SGLang version 0.5.2"
+            model_loader_extra_config = dict(
+                enable_multithread_load=sglang_config.enable_multithread_load,
+                enable_fast_load=sglang_config.enable_fast_load,
+            )
+            args.pop("enable_multithread_load", None)
+            args.pop("enable_fast_load", None)
+            args["model_loader_extra_config"] = json.dumps(
+                model_loader_extra_config, separators=(",", ":")
+            )
         # Map "all-linear" to "all"
         if "lora_target_modules" in args and args["lora_target_modules"]:
             args["lora_target_modules"] = [

areal/experimental/megatron_engine.py

@@ -612,8 +612,6 @@ def _update_weights_from_disk(self, meta: WeightUpdateMeta):
         # dist.barrier() are called when _save_model_to_hf finished
 
         if dist.get_rank() == 0:
-            fut.result()
-
             update_name = names.update_weights_from_disk(
                 self.config.experiment_name,
                 self.config.trial_name,
@@ -623,6 +621,8 @@ def _update_weights_from_disk(self, meta: WeightUpdateMeta):
                 update_name, str(datetime.now().timestamp()), keepalive_ttl=120
             )
 
+            fut.result()
+
         dist.barrier(device_ids=[self.device.index])
         current_platform.synchronize()
 
@@ -642,7 +642,10 @@ def connect_engine(self, engine: InferenceEngine, meta: WeightUpdateMeta):
             )
         self.rollout_engine = engine
 
-        if not self.weight_update_group_initialized:
+        if (
+            meta.type == current_platform.communication_backend
+            and not self.weight_update_group_initialized
+        ):
             self._init_weight_update_from_distributed(meta)
             self.weight_update_group_initialized = True
 

areal/launcher/sglang_server.py

@@ -22,7 +22,7 @@
 )
 from areal.platforms import current_platform
 from areal.utils import logging, name_resolve, names
-from areal.utils.launcher import TRITON_CACHE_PATH
+from areal.utils.launcher import TRITON_CACHE_PATH, apply_sglang_patch
 from areal.utils.network import find_free_ports, gethostip
 
 logger = logging.getLogger("SGLangServer Wrapper")
@@ -130,6 +130,9 @@ def __init__(
         self.server_process = None
         self.n_gpus_per_node = n_gpus_per_node
 
+        if self.config.enable_fast_load or self.config.enable_multithread_load:
+            apply_sglang_patch()
+
     def run(self):
         gpus_per_server = self.allocation_mode.gen_instance_size
         cross_nodes = False

areal/utils/launcher.py

@@ -3,11 +3,15 @@
 import getpass
 import os
 import pathlib
+import shutil
+import subprocess
+import sys
 import time
+from pathlib import Path
 from typing import Dict, Optional
 
 from areal.api.alloc_mode import AllocationMode, AllocationType
-from areal.utils import logging, name_resolve, names
+from areal.utils import logging, name_resolve, names, pkg_version
 
 logger = logging.getLogger("Launcher Utils")
 
@@ -154,3 +158,66 @@ def validate_config_for_distributed_launcher(config):
         assert (
             allocation_mode.gen.tp_size <= config.cluster.n_gpus_per_node
         ), "Currently only support vLLM TP size less <= #GPUs per node."
+
+
+def apply_sglang_patch():
+    p = Path(os.path.dirname(__file__))
+    patch_path = str(
+        p.parent.parent
+        / "patch"
+        / "sglang"
+        / f"v{pkg_version.get_version('sglang')}.patch"
+    )
+    target_path = None
+    sglang_meta = subprocess.check_output(
+        [sys.executable, "-m", "pip", "show", "sglang"]
+    ).decode("utf-8")
+    # Prioritize editable install location, since pip show lists both locations
+    # if installed in editable mode.
+    for line in sglang_meta.split("\n"):
+        line = line.strip()
+        if line.startswith("Editable project location: "):
+            target_path = str(Path(line.split(": ")[1]) / "sglang")
+            break
+    else:
+        for line in sglang_meta.split("\n"):
+            line = line.strip()
+            if line.startswith("Location: "):
+                target_path = str(Path(line.split(": ")[1]) / "sglang")
+                break
+
+    if not target_path or not os.path.exists(target_path):
+        raise RuntimeError("Could not determine the installation path of SGLang.")
+
+    patch_binary = shutil.which("patch")
+    if not patch_binary:
+        raise RuntimeError(
+            "Could not locate the `patch` command; SGLang patch application failed."
+        )
+    result = subprocess.run(
+        [patch_binary, "-p1", "-N", "-i", patch_path],
+        cwd=target_path,
+        capture_output=True,
+        text=True,
+    )
+
+    output = (result.stdout or "") + (result.stderr or "")
+    if result.returncode == 0:
+        logger.info(f"Applied SGLang patch {patch_path} to {target_path}")
+    elif (
+        "Reversed (or previously applied) patch detected" in output
+        or "Skipping patch." in output
+    ):
+        logger.warning(
+            f"SGLang patch {patch_path} appears to be already applied for {target_path}."
+        )
+    else:
+        logger.error(
+            "Failed to apply SGLang patch %s to %s. Output:\n%s",
+            patch_path,
+            target_path,
+            output.strip(),
+        )
+        raise RuntimeError(
+            f"SGLang patch {patch_path} failed with exit code {result.returncode}."
+        )

areal/utils/pkg_version.py

@@ -51,3 +51,15 @@ def is_version_less(package_name: str, target_version: str) -> bool:
     """
     installed_version = get_version(package_name)
     return compare_versions(installed_version, target_version) < 0
+
+
+def is_version_equal(package_name: str, target_version: str) -> bool:
+    """
+    Check if the installed version of a package is equal to the target version.
+
+    :param package_name: Name of the package.
+    :param target_version: Target version to compare against.
+    :return: True if the installed version is equal to the target version, False otherwise.
+    """
+    installed_version = get_version(package_name)
+    return compare_versions(installed_version, target_version) == 0

docs/cli_reference.md

@@ -527,6 +527,8 @@ https://github.com/sgl-project/sglang for detailed documentation.
 | `show_time_cost`                  | boolean                 | `False`      | -           |
 | `enable_metrics`                  | boolean                 | `True`       | -           |
 | `decode_log_interval`             | integer                 | `1`          | -           |
+| `enable_multithread_load`         | boolean                 | `False`      | -           |
+| `enable_fast_load`                | boolean                 | `False`      | -           |
 
 (section-v-llm)=