feat(colocate): support AWEX colocated actor-rollout training

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])

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

areal/api/io_struct.py

@@ -292,13 +292,15 @@ 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,
         base_model_name: str = "",
     ):
         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}"
             )

areal/engine/awex/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""AWEX colocated weight-transfer integration (writer, reader, SGLang plugin)."""