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,9 +1674,15 @@ 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"}
+        default=4,
+        metadata={
+            "help": "Number of minibatches for each PPO update. Separation DTE "
+            "AdamW delta transfer currently requires 1."
+        },
     )
     eps_clip: float = field(
         default=0.2, metadata={"help": "Clipping factor for policy ratio"}
@@ -1849,6 +1883,13 @@ def __post_init__(self):
                 "to a single optimizer step per PPO update."
             )
             self.ppo_n_minibatches = 1
+        if self.dte.enabled and self.ppo_n_minibatches != 1:
+            raise ValueError(
+                "actor.dte.enabled=true currently requires "
+                "ppo_n_minibatches=1 because separation AdamW inversion "
+                "supports exactly one optimizer step between weight updates; "
+                f"got ppo_n_minibatches={self.ppo_n_minibatches}"
+            )
         # Warn if rejection_sampling is configured but use_decoupled_loss is False
         if not self.use_decoupled_loss and self.rejection_sampling is not None:
             logger.warning(

areal/engine/megatron_engine.py

@@ -137,6 +137,7 @@
 from areal.utils.offload import is_tms_enabled, torch_memory_saver
 from areal.utils.perf_tracer import trace_perf, trace_scope
 from areal.utils.seeding import get_seed
+from areal.v2.weight_update.awex.delta_config import DTERuntimeConfig
 
 if TYPE_CHECKING:
     from areal.api import Scheduler
@@ -336,6 +337,7 @@ def __init__(self, config: TrainEngineConfig):
         self.is_offload: bool = False
         self._offload_depth: int = 0
         self._awex_adapter = None  # AwexMegatronAdapter for colocate mode
+        self._dte_runtime_config = DTERuntimeConfig.from_env()
         self.enable_tree_training: bool = self.config.enable_tree_training
         _validate_areal_lm_head_compatibility(
             self.mcore_config.enable_chunked_logits,
@@ -1077,6 +1079,12 @@ def optimizer_zero_grad(self):
             model.zero_grad_buffer()
 
     def optimizer_step(self):
+        if self._dte_runtime_config.enabled:
+            # The LR scheduler advances before the subsequent weight update.
+            # Preserve the LR consumed by this optimizer step for AdamW
+            # inversion instead of reading the next-step LR later.
+            for param_group in self.optimizer.param_groups:
+                param_group["_areal_last_step_lr"] = float(param_group["lr"])
         with trace_scope("megatron_engine.step"):
             update_successful, grad_norm, _ = self.optimizer.step()
         current_lr = self.optimizer.param_groups[0]["lr"]

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/utils/environ.py

@@ -10,12 +10,52 @@
 _warned_rank_env_var_values = set()
 
 
-def get_bool_env_var(name: str, default: str = "false") -> bool:
-    value = os.getenv(name, default)
-    value = value.lower()
-
-    truthy_values = ("true", "1")
-    falsy_values = ("false", "0")
+def get_env_var(
+    name: str,
+    default: str | None = None,
+    *,
+    fallback_names: tuple[str, ...] = (),
+    allow_empty: bool = False,
+) -> str | None:
+    """Read an environment variable with ordered fallback names.
+
+    Empty values are skipped by default so a legacy name can supply the value.
+    Set ``allow_empty=True`` when an explicitly empty value has domain meaning.
+    """
+    for candidate in (name, *fallback_names):
+        value = os.getenv(candidate)
+        if value is None:
+            continue
+        if allow_empty or value.strip() != "":
+            return value
+    return default
+
+
+def get_bool_env_var(
+    name: str,
+    default: str = "false",
+    *,
+    fallback_names: tuple[str, ...] = (),
+    truthy_values: tuple[str, ...] = ("true", "1"),
+    falsy_values: tuple[str, ...] = ("false", "0"),
+    strip_value: bool = False,
+) -> bool:
+    """Read a boolean environment variable.
+
+    The default accepted values remain backward compatible. Callers that
+    historically accepted additional spellings can opt in through
+    ``truthy_values`` and ``falsy_values``.
+    """
+    value = get_env_var(
+        name,
+        default,
+        fallback_names=fallback_names,
+        # Preserve the original single-name behavior: an explicitly empty
+        # boolean is invalid rather than silently replaced by the default.
+        allow_empty=not fallback_names,
+    )
+    assert value is not None
+    value = (value.strip() if strip_value else value).lower()
 
     if (value not in truthy_values) and (value not in falsy_values):
         if value not in _warned_bool_env_var_keys:

areal/v2/weight_update/awex/delta_config.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Runtime configuration for separation AdamW delta transfer.
+
+The delta algorithm itself lives in the standalone ``dte`` package. This
+module snapshots the environment propagated to a GPU worker and lazily creates
+the sender-side tracker only when the opt-in separation path needs it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from areal.utils.environ import get_bool_env_var, get_env_var
+
+_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."
+)
+
+_DTE_TRUTHY_VALUES = ("1", "true", "yes", "on")
+_DTE_FALSY_VALUES = ("0", "false", "no", "off")
+
+
+@dataclass(frozen=True)
+class DTERuntimeConfig:
+    """Worker-local runtime settings for separation delta transfer."""
+
+    delta_transfer: bool
+    separation_weight_update: bool
+    anchor_interval: int
+
+    @classmethod
+    def from_env(cls) -> DTERuntimeConfig:
+        """Snapshot DTE runtime settings from the worker environment."""
+        anchor_value = get_env_var(
+            "DTE_DELTA_ANCHOR_INTERVAL",
+            "0",
+            fallback_names=("AWEX_DELTA_ANCHOR_INTERVAL",),
+        )
+        assert anchor_value is not None
+        anchor_interval = int(anchor_value)
+        return cls(
+            delta_transfer=get_bool_env_var(
+                "DTE_DELTA_TRANSFER",
+                fallback_names=("AWEX_DELTA_TRANSFER",),
+                truthy_values=_DTE_TRUTHY_VALUES,
+                falsy_values=_DTE_FALSY_VALUES,
+                strip_value=True,
+            ),
+            separation_weight_update=get_bool_env_var(
+                "DTE_SEPARATION_WEIGHT_UPDATE",
+                fallback_names=("AWEX_SEPARATION_WEIGHT_UPDATE",),
+                truthy_values=_DTE_TRUTHY_VALUES,
+                falsy_values=_DTE_FALSY_VALUES,
+                strip_value=True,
+            ),
+            anchor_interval=anchor_interval,
+        )
+
+    @property
+    def enabled(self) -> bool:
+        """Whether sparse separated-card transfer is explicitly enabled."""
+        return self.separation_weight_update and self.delta_transfer
+
+    def create_delta_tracker(self):
+        """Create the sender-side tracker without making DTE a base dependency."""
+        try:
+            from dte.core import DeltaTracker
+        except ImportError as e:  # pragma: no cover - only without optional DTE
+            raise ImportError(_DTE_MISSING_MSG) from e
+        return DeltaTracker(anchor_interval=self.anchor_interval)