feat(awex): add separation AdamW delta weight transfer

areal/api/cli_args.py

@@ -1138,6 +1138,34 @@ class SchedulingStrategy:
     )
 
 
+@dataclass
+class DTEConfig:
+    """Configuration for separation AdamW delta weight updates."""
+
+    enabled: bool = field(
+        default=False,
+        metadata={"help": "Enable DTE-backed separation delta weight updates."},
+    )
+    transfer: str = field(
+        default="delta",
+        metadata={
+            "help": "Weight transfer type. The supported DTE mode is 'delta'.",
+            "choices": ["delta"],
+        },
+    )
+    delta_method: str = field(
+        default="adamw",
+        metadata={
+            "help": "How DTE finds changed weights.",
+            "choices": ["adamw"],
+        },
+    )
+    anchor_interval: int = field(
+        default=0,
+        metadata={"help": "Force a full sync every N deltas. 0 means never."},
+    )
+
+
 @dataclass
 class SchedulingSpec:
     cpu: int = field(
@@ -1646,6 +1674,8 @@ def __post_init__(self):
 class PPOActorConfig(TrainEngineConfig):
     """Configuration for PPO actor model, a subclass of a TrainEngine."""
 
+    dte: DTEConfig = field(default_factory=DTEConfig)
+
     # Core PPO/GRPO Parameters
     ppo_n_minibatches: int = field(
         default=4, metadata={"help": "Number of minibatches for each PPO update"}

areal/trainer/rl_trainer.py

@@ -48,6 +48,7 @@
 from areal.infra.utils.concurrent import call_maybe_async
 from areal.utils import logging, perf_tracer, seeding, stats_tracker
 from areal.utils.dataloader import create_dataloader
+from areal.utils.dte import apply_dte_config_envvars
 from areal.utils.environ import is_single_controller
 from areal.utils.evaluator import Evaluator
 from areal.utils.hf_utils import load_hf_processor_and_tokenizer
@@ -132,6 +133,7 @@ def _init_impl(
             logging.setup_file_logging(StatsLogger.get_log_path(config.stats_logger))
 
         self.config = config
+        self._apply_dte_config_envvars()
         self.processor, self.tokenizer = load_hf_processor_and_tokenizer(
             config.tokenizer_path
         )
@@ -1056,6 +1058,10 @@ def _init_scheduler(self) -> Scheduler:
             return SlurmScheduler(exp_config=self.config)
         raise NotImplementedError(f"Unknown scheduler type: {cfg.type}")
 
+    def _apply_dte_config_envvars(self) -> None:
+        """Export ``actor.dte`` config to worker-visible runtime switches."""
+        apply_dte_config_envvars(self.config)
+
     def _create_dataloader(
         self,
         dataset: Dataset,

areal/utils/dte.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dependency-free configuration helpers for separation DTE updates."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import MutableMapping
+from enum import Enum
+from typing import Any
+
+
+def apply_dte_config_envvars(
+    config: Any,
+    environ: MutableMapping[str, str] | None = None,
+) -> dict[str, str]:
+    """Validate ``actor.dte`` and propagate its switches to GPU workers."""
+    dte_cfg = getattr(config.actor, "dte", None)
+    if dte_cfg is None or not dte_cfg.enabled:
+        return {}
+
+    topology = config.rollout.scheduling_strategy.type
+    topology = topology.value if isinstance(topology, Enum) else topology
+    if topology != "separation":
+        raise ValueError(
+            "actor.dte is currently supported only with rollout scheduling "
+            f"strategy 'separation', got {topology!r}"
+        )
+    if dte_cfg.transfer != "delta":
+        raise ValueError("actor.dte.transfer must be 'delta' when DTE is enabled")
+    if dte_cfg.delta_method != "adamw":
+        raise ValueError("actor.dte.delta_method must be 'adamw'")
+    if dte_cfg.anchor_interval < 0:
+        raise ValueError("actor.dte.anchor_interval must be non-negative")
+
+    exported_env = {
+        "DTE_SEPARATION_WEIGHT_UPDATE": "1",
+        "DTE_DELTA_TRANSFER": "1",
+        "DTE_DELTA_ANCHOR_INTERVAL": str(dte_cfg.anchor_interval),
+        "DTE_STREAMING_RECONSTRUCT": "1",
+    }
+    environ = os.environ if environ is None else environ
+    environ.update(exported_env)
+
+    for cfg_part in (config.actor, config.rollout):
+        for spec in getattr(cfg_part, "scheduling_spec", ()) or ():
+            if isinstance(spec, dict):
+                spec.setdefault("env_vars", {}).update(exported_env)
+                continue
+            env_vars = getattr(spec, "env_vars", None)
+            if env_vars is None:
+                env_vars = {}
+                setattr(spec, "env_vars", env_vars)
+            env_vars.update(exported_env)
+
+    return exported_env

areal/v2/weight_update/awex/delta_config.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Runtime gates and factories for separation AdamW delta transfer.
+
+The delta algorithm itself lives in the standalone ``dte`` package; this module
+only decides whether the separated-card AWEX adapters invoke it and constructs
+the writer tracker (with a clear error if DTE is not installed). It
+accepts the new ``DTE_*`` runtime environment emitted from ``actor.dte.*`` CLI
+config, while preserving the old ``AWEX_*`` names as a compatibility fallback.
+
+DTE is imported lazily: a default AReaL install does not require it unless the
+separation delta path is explicitly enabled.
+
+Switches:
+    DTE_DELTA_TRANSFER        enable sparse incremental transfer (default off)
+    DTE_SEPARATION_WEIGHT_UPDATE
+                              allow the separation-only sparse P2P path
+    DTE_DELTA_ANCHOR_INTERVAL force a full sync every N deltas (0 = never)
+"""
+
+from __future__ import annotations
+
+import os
+
+_DTE_MISSING_MSG = (
+    "DTE separation delta transfer requires the 'dte' package "
+    "(delta-transfer-engine), which the AWEX adapters import "
+    "lazily. Install it with `pip install -e <path>/delta-transfer-engine` "
+    "(local dev) or add DTE_SRC to PYTHONPATH."
+)
+
+
+def _env_value(name: str, legacy_name: str, default: str | None = None) -> str | None:
+    value = os.environ.get(name)
+    if value is not None and value.strip() != "":
+        return value
+    value = os.environ.get(legacy_name)
+    if value is not None and value.strip() != "":
+        return value
+    return default
+
+
+def _env_bool(name: str, legacy_name: str, default: bool = False) -> bool:
+    value = _env_value(name, legacy_name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def delta_transfer_enabled() -> bool:
+    """Master switch for sparse incremental transfer."""
+    return _env_bool("DTE_DELTA_TRANSFER", "AWEX_DELTA_TRANSFER")
+
+
+def separation_weight_update_enabled() -> bool:
+    """Whether the configured topology permits separation-only DTE code."""
+    return _env_bool("DTE_SEPARATION_WEIGHT_UPDATE", "AWEX_SEPARATION_WEIGHT_UPDATE")
+
+
+def separation_delta_transfer_enabled() -> bool:
+    """Whether sparse separated-card transfer is explicitly enabled."""
+    return separation_weight_update_enabled() and delta_transfer_enabled()
+
+
+def delta_anchor_interval() -> int:
+    """Force a full sync every N deltas (0 = never; rely on seed + chain-break)."""
+    return int(
+        _env_value(
+            "DTE_DELTA_ANCHOR_INTERVAL",
+            "AWEX_DELTA_ANCHOR_INTERVAL",
+            "0",
+        )
+    )
+
+
+def make_delta_tracker():
+    """Sender-side dte ``DeltaTracker``, configured from env.
+
+    Raises a clear ``ImportError`` if dte is not installed.
+    """
+    try:
+        from dte.core import DeltaTracker
+    except ImportError as e:  # pragma: no cover - exercised only without dte
+        raise ImportError(_DTE_MISSING_MSG) from e
+    return DeltaTracker(anchor_interval=delta_anchor_interval())
+
+
+def cuda_mem_stats_mb(reset_peak: bool = True) -> tuple[float, float]:
+    """Return ``(allocated_mb, peak_mb)`` for the current CUDA device.
+
+    ``peak_mb`` is the high-water mark since the previous call (the peak
+    counter is reset afterwards by default), which lets [dte-perf] stage
+    marks attribute allocation spikes to individual weight-sync stages.
+    Returns ``(-1.0, -1.0)`` when CUDA is unavailable (CPU tests).
+    """
+    import torch
+
+    if not torch.cuda.is_available():
+        return -1.0, -1.0
+    allocated = torch.cuda.memory_allocated() / (1024.0 * 1024.0)
+    peak = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
+    if reset_peak:
+        torch.cuda.reset_peak_memory_stats()
+    return allocated, peak