From 6b4dd44fb18c0bbeac59584d4f5a55d922a2c449 Mon Sep 17 00:00:00 2001 From: pyq623 <12421170@zju.edu.cn> Date: Fri, 14 Aug 2026 12:18:51 +0800 Subject: [PATCH 1/5] feat(config): add separation DTE topology gates --- areal/api/cli_args.py | 30 ++++++++++++ areal/trainer/rl_trainer.py | 6 +++ areal/utils/dte.py | 55 +++++++++++++++++++++ docs/en/cli_reference.md | 15 ++++++ docs/zh/cli_reference.md | 15 ++++++ tests/test_dte_topology_gating.py | 81 +++++++++++++++++++++++++++++++ 6 files changed, 202 insertions(+) create mode 100644 areal/utils/dte.py create mode 100644 tests/test_dte_topology_gating.py diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 903e94458b..c8786f0561 100644 --- a/areal/api/cli_args.py +++ b/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"} diff --git a/areal/trainer/rl_trainer.py b/areal/trainer/rl_trainer.py index 117ff61c8c..b9afe254ec 100644 --- a/areal/trainer/rl_trainer.py +++ b/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, diff --git a/areal/utils/dte.py b/areal/utils/dte.py new file mode 100644 index 0000000000..3729742da9 --- /dev/null +++ b/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 diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index ef9a88e079..f134f2608a 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -76,6 +76,7 @@ For detailed examples, see the experiment configurations in the `examples/` dire - [ArchonFP8 Configuration](section-archon-fp8) - [DPO Configuration](section-dpo) - [DPOEngine Configuration](section-dpo-engine) +- [DTE Configuration](section-dte) - [DistributedDataParallel Configuration](section-distributed-data-parallel) - [FP8Engine Configuration](section-fp8-engine) - [MegatronEngine Configuration](section-megatron-engine) @@ -387,6 +388,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `setup_timeout` | float | `3600.0` | Gateway setup timeout in seconds for controller v2. | | `workers_ready_timeout` | float | `30.0` | Timeout (seconds) for initialize() to wait for guards to be ready. | | `scheduling_strategy` | [`SchedulingStrategy`](section-scheduling-strategy) | **Required** | The scheduling strategy of this TrainEngine, either separation or colocation. Currently only used by the TrainController. | +| `dte` | [`DTEConfig`](section-dte) | **Required** | - | | `ppo_n_minibatches` | integer | `4` | Number of minibatches for each PPO update | | `eps_clip` | float | `0.2` | Clipping factor for policy ratio | | `eps_clip_higher` | float \| None | `None` | Clipping factor (higher value) for policy ratio. Default is None. When eps_clip_higher is set (decoupled), eps_clip will be used as the lower value. | @@ -1033,6 +1035,19 @@ fields. | `beta` | float | `0.1` | KL penalty coefficient for DPO loss. | | `loss_type` | string | `"sigmoid"` | DPO loss variant. 'sigmoid': original DPO loss (Rafailov et al. 2023). 'ipo': Identity Preference Optimization with per-token length normalization (Azar et al. 2023). **Choices:** `sigmoid`, `ipo` | +(section-dte)= + +## DTE Configuration + +Configuration for separation AdamW delta weight updates. + +| Parameter | Type | Default | Description | +| ----------------- | ------- | --------- | ----------------------------------------------------------------------------- | +| `enabled` | boolean | `False` | Enable DTE-backed separation delta weight updates. | +| `transfer` | string | `"delta"` | Weight transfer type. The supported DTE mode is 'delta'. **Choices:** `delta` | +| `delta_method` | string | `"adamw"` | How DTE finds changed weights. **Choices:** `adamw` | +| `anchor_interval` | integer | `0` | Force a full sync every N deltas. 0 means never. | + (section-distributed-data-parallel)= ## DistributedDataParallel Configuration diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 42233dbe18..b9860386ad 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -74,6 +74,7 @@ python3 train.py --config path/to/config.yaml actor.lr=1e-4 seed=42 - [ArchonFP8 Configuration](section-archon-fp8) - [DPO Configuration](section-dpo) - [DPOEngine Configuration](section-dpo-engine) +- [DTE Configuration](section-dte) - [DistributedDataParallel Configuration](section-distributed-data-parallel) - [FP8Engine Configuration](section-fp8-engine) - [MegatronEngine Configuration](section-megatron-engine) @@ -385,6 +386,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `setup_timeout` | float | `3600.0` | Gateway setup timeout in seconds for controller v2. | | `workers_ready_timeout` | float | `30.0` | Timeout (seconds) for initialize() to wait for guards to be ready. | | `scheduling_strategy` | [`SchedulingStrategy`](section-scheduling-strategy) | **Required** | The scheduling strategy of this TrainEngine, either separation or colocation. Currently only used by the TrainController. | +| `dte` | [`DTEConfig`](section-dte) | **Required** | - | | `ppo_n_minibatches` | integer | `4` | Number of minibatches for each PPO update | | `eps_clip` | float | `0.2` | Clipping factor for policy ratio | | `eps_clip_higher` | float \| None | `None` | Clipping factor (higher value) for policy ratio. Default is None. When eps_clip_higher is set (decoupled), eps_clip will be used as the lower value. | @@ -1031,6 +1033,19 @@ fields. | `beta` | float | `0.1` | KL penalty coefficient for DPO loss. | | `loss_type` | string | `"sigmoid"` | DPO loss variant. 'sigmoid': original DPO loss (Rafailov et al. 2023). 'ipo': Identity Preference Optimization with per-token length normalization (Azar et al. 2023). **Choices:** `sigmoid`, `ipo` | +(section-dte)= + +## DTE Configuration + +Configuration for separation AdamW delta weight updates. + +| Parameter | Type | Default | Description | +| ----------------- | ------- | --------- | ----------------------------------------------------------------------------- | +| `enabled` | boolean | `False` | Enable DTE-backed separation delta weight updates. | +| `transfer` | string | `"delta"` | Weight transfer type. The supported DTE mode is 'delta'. **Choices:** `delta` | +| `delta_method` | string | `"adamw"` | How DTE finds changed weights. **Choices:** `adamw` | +| `anchor_interval` | integer | `0` | Force a full sync every N deltas. 0 means never. | + (section-distributed-data-parallel)= ## DistributedDataParallel Configuration diff --git a/tests/test_dte_topology_gating.py b/tests/test_dte_topology_gating.py new file mode 100644 index 0000000000..14912825aa --- /dev/null +++ b/tests/test_dte_topology_gating.py @@ -0,0 +1,81 @@ +"""Tests for the separation-only DTE configuration boundary.""" + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_HELPER_PATH = Path(__file__).parents[1] / "areal/utils/dte.py" + + +def _load_helpers(): + spec = importlib.util.spec_from_file_location("areal_utils_dte", _HELPER_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _config(*, enabled=False, topology="separation", **overrides): + dte = { + "enabled": enabled, + "transfer": "delta", + "delta_method": "adamw", + "anchor_interval": 20, + **overrides, + } + actor_spec = SimpleNamespace(env_vars={}) + rollout_spec = SimpleNamespace(env_vars={}) + return SimpleNamespace( + actor=SimpleNamespace( + dte=SimpleNamespace(**dte), scheduling_spec=(actor_spec,) + ), + rollout=SimpleNamespace( + scheduling_strategy=SimpleNamespace(type=topology), + scheduling_spec=(rollout_spec,), + ), + ) + + +def test_dte_disabled_does_not_change_environment(): + config = _config() + + exported = _load_helpers().apply_dte_config_envvars(config, environ={}) + + assert exported == {} + assert config.actor.scheduling_spec[0].env_vars == {} + assert config.rollout.scheduling_spec[0].env_vars == {} + + +def test_dte_enabled_exports_only_separation_adamw_switches(): + config = _config(enabled=True) + + exported = _load_helpers().apply_dte_config_envvars(config, environ={}) + + assert exported == { + "DTE_SEPARATION_WEIGHT_UPDATE": "1", + "DTE_DELTA_TRANSFER": "1", + "DTE_DELTA_ANCHOR_INTERVAL": "20", + "DTE_STREAMING_RECONSTRUCT": "1", + } + assert config.actor.scheduling_spec[0].env_vars == exported + assert config.rollout.scheduling_spec[0].env_vars == exported + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"topology": "colocation"}, "only with.*separation"), + ({"transfer": "full"}, "transfer must be 'delta'"), + ({"delta_method": "snapshot"}, "delta_method must be 'adamw'"), + ({"anchor_interval": -1}, "anchor_interval must be non-negative"), + ], +) +def test_dte_rejects_out_of_scope_modes(kwargs, match): + kwargs = dict(kwargs) + topology = kwargs.pop("topology", "separation") + config = _config(enabled=True, topology=topology, **kwargs) + + with pytest.raises(ValueError, match=match): + _load_helpers().apply_dte_config_envvars(config, environ={}) From c3ea90dc84b598a8412007e90933743aa13fcf5c Mon Sep 17 00:00:00 2001 From: pyq623 <12421170@zju.edu.cn> Date: Fri, 14 Aug 2026 12:20:16 +0800 Subject: [PATCH 2/5] feat(awex): add separation AdamW delta weight transfer --- areal/v2/weight_update/awex/delta_config.py | 103 + areal/v2/weight_update/awex/delta_detect.py | 1652 +++++++++++++++++ .../v2/weight_update/awex/megatron_adapter.py | 283 ++- areal/v2/weight_update/awex/sglang_adapter.py | 114 +- tests/test_awex_delta_common.py | 313 ++++ tests/test_awex_separation_delta.py | 375 ++++ 6 files changed, 2834 insertions(+), 6 deletions(-) create mode 100644 areal/v2/weight_update/awex/delta_config.py create mode 100644 areal/v2/weight_update/awex/delta_detect.py create mode 100644 tests/test_awex_delta_common.py create mode 100644 tests/test_awex_separation_delta.py diff --git a/areal/v2/weight_update/awex/delta_config.py b/areal/v2/weight_update/awex/delta_config.py new file mode 100644 index 0000000000..2f20209ea5 --- /dev/null +++ b/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 /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 diff --git a/areal/v2/weight_update/awex/delta_detect.py b/areal/v2/weight_update/awex/delta_detect.py new file mode 100644 index 0000000000..5012b9242a --- /dev/null +++ b/areal/v2/weight_update/awex/delta_detect.py @@ -0,0 +1,1652 @@ +# SPDX-License-Identifier: Apache-2.0 +"""AdamW inversion detector for AWEX separation delta weight transfer. + +The detector answers one question each weight-update step: which bf16 elements +of the converted HF payload changed since the previous version? It returns +``{hf_name: bool mask | int32/int64 flat indices}`` consumed by +``DeltaTracker.encode(masks=...)``. + +It reconstructs the *pre-step* weights from the optimizer's resident AdamW +moments (``exp_avg`` / ``exp_avg_sq``), converts them through the same HF path +as the live payload, and bitwise-compares both versions. Unsupported optimizer +states return ``None`` so every rank falls back to a full transfer. + +AdamW inversion (decoupled AdamW, the mcore default): + + theta_t = theta_{t-1}·(1 - lr·wd) - (lr/bc1)·m / (sqrt(v)/sqrt(bc2) + eps) + +is element-wise invertible because ``m=exp_avg`` / ``v=exp_avg_sq`` stay resident +after ``optimizer.step()``: + + theta_{t-1} = (theta_t + (lr/bc1)·m / (sqrt(v)/sqrt(bc2) + eps)) / (1 - lr·wd) + +with ``bc1 = 1 - beta1^step``, ``bc2 = 1 - beta2^step``. + +Megatron specifics (confirmed against mcore distrib_optimizer): + +- Moments live ONLY in mcore main-param space (HF space has QKV already split, + gate/expert converted) so the mask MUST be computed AFTER convert. +- The distributed optimizer shards the fp32 main param + moments across DP, but + the FULL bf16 model param (theta_t) is resident on every DP rank (post-step + all-gather). So inversion is a per-rank, shard-local fp32 compute scattered + into a full-param buffer, then ONE DP all-reduce(SUM) of the *correction* + (zero outside the owned slice) assembles the full pre-step param. + +Hard gates (any failure -> return None -> writer ships dense this step): + +- ``use_precision_aware_optimizer`` (adam_bf16): moments are bf16/TE-fused, not + plain fp32 ``exp_avg`` -> inversion infeasible. +- non-decoupled AdamW (L2-Adam folds wd into grad): the ``/(1-lr·wd)`` form is + wrong. +- no reconstructable moment state for the whole step, or ``step < 1`` on every + usable shard (first step, recover): division blow-up. Per-param/per-rank + missing state contributes zeros to DP reduction so ranks do not deadlock. +- compact per-shard fingerprints are only a guard for ``step`` unchanged or + missing-watermark recovery. A known one-step AdamW transition is always + reconstructed; the fingerprint is not trusted as proof of no payload change. +- more than one distributed-optimizer instance: the DP group identity differs + from the all-gather group. Multiple optimizer DP topologies in one model are + allowed; MoE expert params can use a smaller expert-DP group while dense + params use the regular DP group, so reconstruction reduces each param on its + own optimizer DP group. + +Migration note (origin/gh AWEX): the detector now drives the new adapter's +convert path — ``adapter._get_inner_optimizers()`` for the moments and +``adapter._convert_hf_with_overrides(theta_by_id)`` to push reconstructed +pre-step weights through the same all_gather + convert_to_hf as the live +payload (the old awex ``_make_weight_converter`` / ``_convert_parameters_with`` +are gone). The mcore reconstruction below is GPU-only and must be validated on +the cluster; ``dte.core.invert_adamw`` itself is CPU-unit-tested. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Iterator + +import torch + +from areal.utils import logging +from areal.v2.weight_update.awex.delta_config import cuda_mem_stats_mb + +logger = logging.getLogger("AwexDeltaDetect") + + +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 _inversion_debug_enabled() -> bool: + return _env_bool("DTE_DELTA_INVERSION_DEBUG", "AWEX_DELTA_INVERSION_DEBUG") + + +def _streaming_reconstruct_enabled() -> bool: + """Whether to reconstruct pre-step tensors lazily during HF conversion.""" + return _env_bool("DTE_STREAMING_RECONSTRUCT", "PYQ_STREAMING_RECONSTRUCT", False) + + +def _inversion_allreduce_window_bytes() -> int: + """In-flight byte budget for pipelined inversion all-reduces. + + The reconstruct loop posts per-param correction/status all-reduces with + ``async_op=True`` and only waits once this many correction bytes are in + flight. ``0`` restores the fully synchronous per-param path (A/B knob). + """ + mb = float( + _env_value( + "DTE_INVERSION_ALLREDUCE_WINDOW_MB", + "AWEX_INVERSION_ALLREDUCE_WINDOW_MB", + "512", + ) + ) + return max(int(mb * 1024 * 1024), 0) + + +def _inversion_compute_device(flat_t: torch.Tensor) -> torch.device: + """Device used for the AdamW-inversion elementwise math. + + Colocate offloads the fp32 main shards and AdamW moments to CPU before the + weight sync runs, so deriving the compute device from ``main_shard`` makes + the whole inversion (invert_adamw + mask probes) run on the CPU: ~18s/step + for Qwen3-30B versus a few seconds of staged H2D copies. Default to the + visible payload's device (GPU under colocate) and stage the offloaded + operands there. Set DTE_INVERSION_COMPUTE_ON_CPU=1 to restore the legacy + behaviour of computing wherever the fp32 main shard lives. + """ + if _env_bool("DTE_INVERSION_COMPUTE_ON_CPU", "AWEX_INVERSION_COMPUTE_ON_CPU"): + return torch.device("cpu") + return flat_t.device + + +def _inversion_bf16_margin_rel() -> float: + return float( + _env_value( + "DTE_DELTA_INVERSION_BF16_MARGIN_REL", + "AWEX_DELTA_INVERSION_BF16_MARGIN_REL", + "1e-4", + ) + ) + + +def _inversion_dense_param_suffixes() -> tuple[str, ...]: + """HF parameter suffixes to send dense under AdamW inversion. + + Qwen3 MoE router/gate weights have repeatedly shown one-element BF16 false + negatives in snapshot verification. They are small relative to full model + payloads, so sending them dense is the conservative default. Set the env to + an empty string to disable, or to a comma-separated suffix list to override. + """ + if "DTE_DELTA_INVERSION_DENSE_PARAM_SUFFIXES" in os.environ: + value = os.environ["DTE_DELTA_INVERSION_DENSE_PARAM_SUFFIXES"] + elif "AWEX_DELTA_INVERSION_DENSE_PARAM_SUFFIXES" in os.environ: + value = os.environ["AWEX_DELTA_INVERSION_DENSE_PARAM_SUFFIXES"] + else: + value = ".mlp.gate.weight" + return tuple(item.strip() for item in value.split(",") if item.strip()) + + +def _inversion_force_dense_param(name: str) -> bool: + return any(name.endswith(suffix) for suffix in _inversion_dense_param_suffixes()) + + +def _dist_rank(group=None) -> int: + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return -1 + try: + return torch.distributed.get_rank(group=group) + except Exception: # pragma: no cover - debug-only best effort + return -1 + + +def _dist_world_size(group=None) -> int: + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return 1 + try: + return torch.distributed.get_world_size(group=group) + except Exception: # pragma: no cover - debug-only best effort + return 1 + + +def _adamw_hparams(param_group: dict) -> tuple[float, float, float, float, float]: + """Extract (lr, weight_decay, beta1, beta2, eps) from a torch param_group.""" + betas = param_group.get("betas", (0.9, 0.999)) + beta1 = float(betas[0]) + beta2 = float(betas[1]) + if not param_group.get("bias_correction", True): + # Apex FusedAdam stores step on the param group and can disable bias + # correction there. invert_adamw uses beta**step only to compute those + # correction factors, so beta=0 gives bc1=bc2=1. + beta1 = 0.0 + beta2 = 0.0 + return ( + float(param_group.get("_areal_last_step_lr", param_group["lr"])), + float(param_group.get("weight_decay", 0.0)), + beta1, + beta2, + float(param_group.get("eps", 1e-8)), + ) + + +def _resolve_offloaded_state_value( + state: dict, offloaded_state: dict, key: str +) -> object | None: + """Return optimizer state value, following AwexMegatronAdapter offload slots.""" + val = state.get(key) + if isinstance(val, torch.Tensor) and val.numel() == 0 and key in offloaded_state: + return offloaded_state[key] + return val + + +def _state_step_value(state: dict, offloaded_state: dict) -> float | None: + step_val = _resolve_offloaded_state_value(state, offloaded_state, "step") + return _step_to_float(step_val) + + +def _param_group_step_value(param_group: dict) -> float | None: + return _step_to_float(param_group.get("step")) + + +def _step_to_float(step_val: object | None) -> float | None: + if step_val is None: + return None + if isinstance(step_val, torch.Tensor): + if step_val.numel() == 0: + return None + return float(step_val.item()) + return float(step_val) + + +@torch.no_grad() +def _tensor_fingerprint(tensor: torch.Tensor) -> tuple: + """Small value fingerprint for a tensor slice, not a stored snapshot. + + The detector uses this as a no-full-copy guard to decide whether a local + model-visible shard changed since the last successful sync. The actual delta + mask still comes from AdamW inversion; this fingerprint only prevents + replaying moments for a shard whose synced bf16 payload did not move. + """ + data = tensor.detach().contiguous().view(torch.uint8).reshape(-1) + n = data.numel() + if n == 0: + return (str(tensor.dtype), 0, 0, 0, 0, 0) + + bins = min(4096, n) + width = max(n // bins, 1) + head_n = bins * width + head = data[:head_n].reshape(bins, width) + bin_sums = head.sum(dim=1, dtype=torch.int64) + weights = torch.arange(1, bins + 1, device=data.device, dtype=torch.int64) + total = bin_sums.sum() + weighted = (bin_sums * weights).sum() + if head_n < n: + tail = data[head_n:].sum(dtype=torch.int64) + total = total + tail + weighted = weighted + tail * (bins + 1) + + sample_count = min(16, n) + if sample_count == 1: + sample_idx = torch.zeros(1, device=data.device, dtype=torch.long) + else: + sample_idx = ( + torch.arange(sample_count, device=data.device, dtype=torch.long) + * (n - 1) + // (sample_count - 1) + ) + samples = data.index_select(0, sample_idx).to(torch.int64) + sample_weights = torch.arange( + 1, sample_count + 1, device=data.device, dtype=torch.int64 + ) + sample_hash = (samples * sample_weights).sum() + return ( + str(tensor.dtype), + n, + int(total.item()), + int(weighted.item()), + int(sample_hash.item()), + int(data[0].item()), + int(data[-1].item()), + ) + + +@torch.no_grad() +def _bf16_rounding_boundary_mask( + old_fp32: torch.Tensor, cur_bf16: torch.Tensor +) -> torch.Tensor: + """Conservatively include values near a BF16 rounding boundary.""" + cur_fp32 = cur_bf16.to(torch.float32) + old_fp32 = old_fp32.to(torch.float32) + + # BF16 bins are not symmetric at exponent boundaries: for example the lower + # neighbor of 1.0 is 0.99609375 while the upper neighbor is 1.0078125. Using + # one spacing for both sides misses values close to the narrower boundary. + neg_inf = torch.full_like(cur_bf16, float("-inf")) + pos_inf = torch.full_like(cur_bf16, float("inf")) + lower = torch.nextafter(cur_bf16, neg_inf).to(torch.float32) + upper = torch.nextafter(cur_bf16, pos_inf).to(torch.float32) + lower_half_ulp = (cur_fp32 - lower).abs() * 0.5 + upper_half_ulp = (upper - cur_fp32).abs() * 0.5 + half_ulp = torch.where(old_fp32 < cur_fp32, lower_half_ulp, upper_half_ulp) + threshold = half_ulp * (1.0 - _inversion_bf16_margin_rel()) + dist = (old_fp32 - cur_fp32).abs() + return torch.isfinite(old_fp32) & torch.isfinite(cur_fp32) & (dist >= threshold) + + +@torch.no_grad() +def _bf16_rounding_boundary_mask_chunked( + old_fp32: torch.Tensor, cur_bf16: torch.Tensor +) -> torch.Tensor: + """Bound boundary-check temporaries by processing large tensors in chunks.""" + chunk_elems = int( + _env_value( + "DTE_BOUNDARY_CHUNK_ELEMS", + "PYQ_BOUNDARY_CHUNK_ELEMS", + str(50_000_000), + ) + ) + if chunk_elems <= 0 or cur_bf16.numel() <= chunk_elems: + return _bf16_rounding_boundary_mask(old_fp32, cur_bf16) + + flat_old = old_fp32.reshape(-1) + flat_cur = cur_bf16.reshape(-1) + result = torch.empty(flat_cur.numel(), dtype=torch.bool, device=flat_cur.device) + for start in range(0, flat_cur.numel(), chunk_elems): + end = min(start + chunk_elems, flat_cur.numel()) + result[start:end] = _bf16_rounding_boundary_mask( + flat_old[start:end], flat_cur[start:end] + ).reshape(-1) + return result.reshape(cur_bf16.shape) + + +class _ReconstructionAborted(Exception): + """Lazy reconstruction reached a rank-consistent dense fallback verdict.""" + + +_NO_RECONSTRUCTION = object() + + +class _LazyReconstructDict: + """Dict-like lazy view over reconstructed pre-step tensors. + + Current Megatron converters call ``get(id(param), param)`` and, after + conversion, ``pop(id(param), None)``. ``get`` may advance the producer; + the matching ``pop`` must only discard the value already returned and must + never advance again. The producer emits a sentinel for params without an + override so one missing value cannot force buffering the rest of the model. + """ + + def __init__(self, producer: Iterator[tuple[int, torch.Tensor | object]]) -> None: + self._producer = producer + self._buffer: dict[int, torch.Tensor | object] = {} + self._consumed: set[int] = set() + self._exhausted = False + self._aborted = False + self._producer_seconds = 0.0 + self._reconstructed_count = 0 + + def _pull_one(self) -> tuple[int, torch.Tensor | object] | None: + if self._exhausted: + return None + started = time.monotonic() + try: + item = next(self._producer) + except _ReconstructionAborted: + self._aborted = True + self._exhausted = True + return None + except StopIteration: + self._exhausted = True + return None + finally: + self._producer_seconds += time.monotonic() - started + if item[1] is not _NO_RECONSTRUCTION: + self._reconstructed_count += 1 + return item + + def _fill(self, key: int) -> None: + while key not in self._buffer and not self._exhausted: + item = self._pull_one() + if item is None: + break + item_key, value = item + self._buffer[item_key] = value + + def get(self, key: int, default=None): + if key in self._consumed: + return default + self._fill(key) + value = self._buffer.get(key, _NO_RECONSTRUCTION) + self._consumed.add(key) + if value is _NO_RECONSTRUCTION: + return default + return value + + def pop(self, key: int, default=None): + if key in self._consumed: + value = self._buffer.pop(key, _NO_RECONSTRUCTION) + return default if value is _NO_RECONSTRUCTION else value + self._fill(key) + self._consumed.add(key) + value = self._buffer.pop(key, _NO_RECONSTRUCTION) + return default if value is _NO_RECONSTRUCTION else value + + def finish(self) -> None: + """Drain the producer so trailing collectives and aborts always run.""" + while not self._exhausted: + item = self._pull_one() + if item is None: + break + # Conversion has ended, so unrequested overrides can be released + # immediately. Only their collectives and trailing verdict matter. + del item + + def clear(self) -> None: + self._buffer.clear() + self._consumed.clear() + + @property + def aborted(self) -> bool: + return self._aborted + + @property + def producer_seconds(self) -> float: + return self._producer_seconds + + @property + def reconstructed_count(self) -> int: + return self._reconstructed_count + + +class AdamWInversionDetector: + """Compute change masks by inverting the resident AdamW moments. + + Bound to ``AwexMegatronAdapter`` (needs ``_get_inner_optimizers`` for the + moments and ``_convert_hf_with_overrides`` to convert reconstructed pre-step + params through the live path). All mcore/DP/convert work is GPU-only; any + unmet precondition returns ``None`` so the writer falls back to dense for + that step (and the snapshot tracker re-seeds), never silently corrupting. + """ + + name = "inversion" + + def __init__(self, adapter): + self._adapter = adapter + # All watermark/fingerprint maps are keyed by the mcore param NAME, not + # id(param): on EP topologies the optimizer can hold different param + # OBJECTS than the live module traversal (observed d2p8e4: one DP + # replica's ids never match), and ids are also not stable across + # offload/reload cycles. Names are rank- and version-invariant. + self._last_synced_steps: dict[str, float | None] = {} + self._last_synced_fingerprints: dict[str, tuple] = {} + self._last_synced_model_fingerprints: dict[str, tuple] = {} + self._last_synced_payload_fingerprints: dict[str, tuple] = {} + # Single-slot cache filled by precompute_masks() before the colocate + # orchestration offloads the optimizer to CPU; consumed (or discarded) + # by the next compute_masks()/mark_synced() call. + self._precomputed_masks: ( + tuple[int, tuple[str, ...], dict[str, torch.Tensor] | None] | None + ) = None + + def has_synced_watermark(self) -> bool: + """Whether a successful full/delta payload sync recorded optimizer steps.""" + return bool(self._last_synced_steps) + + # -- gating ------------------------------------------------------------ + def _inversion_feasible(self, inner_optimizers) -> bool: + """All hard gates from the design; any failure -> dense fallback.""" + if not inner_optimizers: + logger.warning("Inversion: no inner optimizers; falling back to dense.") + return False + for opt in inner_optimizers: + base_opt = getattr(opt, "optimizer", opt) + adam_w_mode = getattr(base_opt, "adam_w_mode", None) + if adam_w_mode is not None and int(adam_w_mode) == 0: + logger.warning("Inversion: non-decoupled FusedAdam -> dense.") + return False + cfg = getattr(opt, "config", None) + if cfg is not None: + # adam_bf16 / precision-aware: moments are bf16 / TE-fused, not + # plain fp32 exp_avg -> inversion infeasible. + pa = getattr( + cfg, "use_precision_aware_optimizer_no_fp8_or_ds_fp8", None + ) + if pa is None: + pa = getattr(cfg, "use_precision_aware_optimizer", False) + if pa: + logger.warning("Inversion: use_precision_aware_optimizer -> dense.") + return False + # decoupled AdamW: the (1-lr*wd) form requires decoupled wd. + if not getattr(cfg, "decoupled_weight_decay", True): + logger.warning("Inversion: non-decoupled AdamW -> dense.") + return False + # Single distributed-optimizer instance: otherwise the DP group + # identity differs from the param all-gather group (design risk #3). + n_inst = getattr(opt, "num_distributed_optimizer_instances", 1) + if n_inst and n_inst > 1: + logger.warning( + "Inversion: %d distributed-optimizer instances -> dense.", + n_inst, + ) + return False + return True + + def _module_param_key_maps( + self, + ) -> tuple[dict[int, str], dict[tuple[int, tuple], str]]: + """Map live module params to stable name keys, by id and (data_ptr, shape). + + The (data_ptr, shape) index bridges optimizer-held param objects to the + module params when object identities diverge (observed on d2p8e4: one DP + replica's optimizer groups hold different objects than the module + traversal). Both views alias the same storage while weights are + resident, so the pointer identifies the tensor. Pointers are only valid + within one call (offload/reload reallocates storage) — never persist + them, only the derived names. + """ + id2key: dict[int, str] = {} + ptr2key: dict[tuple[int, tuple], str] = {} + engine = getattr(self._adapter, "_engine", None) + if engine is None or getattr(engine, "model", None) is None: + return id2key, ptr2key + from areal.engine.megatron_utils.megatron import get_named_parameters + + num_moe_experts = getattr(engine.tf_config, "num_moe_experts", None) + for name, param in get_named_parameters(engine.model, num_moe_experts): + id2key.setdefault(id(param), name) + try: + if param.data.numel() > 0: + ptr2key.setdefault( + (param.data.data_ptr(), tuple(param.shape)), name + ) + except RuntimeError: + pass # released storage has no data_ptr + return id2key, ptr2key + + @staticmethod + def _param_key(param, id2key, ptr2key) -> str | int | None: + key = id2key.get(id(param)) + if key is not None: + return key + try: + if param.data.numel() > 0: + key = ptr2key.get((param.data.data_ptr(), tuple(param.shape))) + if key is not None: + return key + except RuntimeError: + pass + if not id2key and not ptr2key: + return id(param) + return None + + def _collect_current_steps(self, inner_optimizers) -> dict[str, float | None]: + """Return optimizer step watermarks keyed by mcore param name. + + A value of ``None`` means this param had no local optimizer state when the + payload was synced. On the next step, ``None -> 1`` is still a valid + one-step transition. + """ + steps: dict[str, float | None] = {} + id2key, ptr2key = self._module_param_key_maps() + offloaded_states = getattr(self._adapter, "_offloaded_optimizer_states", {}) + for opt in inner_optimizers: + base_opt = getattr(opt, "optimizer", opt) + state = getattr(base_opt, "state", None) + group_index_map = getattr(opt, "model_param_group_index_map", None) + fp32_groups = getattr(opt, "shard_fp32_from_float16_groups", None) + model_groups = getattr(opt, "model_float16_groups", None) + if state is None or fp32_groups is None or model_groups is None: + continue + for g, (fp32_group, model_group) in enumerate( + zip(fp32_groups, model_groups) + ): + for main_shard, model_param in zip(fp32_group, model_group): + if main_shard is None: + continue + key = self._param_key(model_param, id2key, ptr2key) + if key is None or key in steps: + continue + gi = ( + group_index_map[model_param][0] + if group_index_map is not None + and model_param in group_index_map + else g + ) + st = state.get(main_shard) + step = None + if st: + offloaded_state = offloaded_states.get(main_shard, {}) + step = _state_step_value(st, offloaded_state) + if step is None: + step = _param_group_step_value(base_opt.param_groups[gi]) + steps[key] = step + return steps + + def _collect_current_fingerprints(self, inner_optimizers) -> dict[str, tuple]: + """Record compact fingerprints of model-visible local shards. + + This is intentionally not a snapshot: it stores a few scalar values per + optimizer-owned shard, not tensor contents. The fingerprint gates AdamW + replay to slices whose synced payload actually changed. + """ + fingerprints: dict[str, tuple] = {} + id2key, ptr2key = self._module_param_key_maps() + for opt in inner_optimizers: + fp32_groups = getattr(opt, "shard_fp32_from_float16_groups", None) + model_groups = getattr(opt, "model_float16_groups", None) + if fp32_groups is None or model_groups is None: + continue + for fp32_group, model_group in zip(fp32_groups, model_groups): + for main_shard, model_param in zip(fp32_group, model_group): + if main_shard is None: + continue + key = self._param_key(model_param, id2key, ptr2key) + if key is None or key in fingerprints: + continue + try: + rng = opt._get_model_param_range_map(model_param)["param"] + except Exception: + continue + visible_slice = model_param.detach().reshape(-1)[ + rng.start : rng.end + ] + fingerprints[key] = _tensor_fingerprint(visible_slice) + return fingerprints + + def _ordered_model_params(self, fallback_order: list[torch.Tensor]): + iter_fn = getattr(self._adapter, "_iter_model_params_for_delta", None) + if iter_fn is not None: + params = list(iter_fn()) + if params: + return params + + ordered_model_params: list[torch.Tensor] = [] + engine = getattr(self._adapter, "_engine", None) + if engine is not None and getattr(engine, "model", None) is not None: + from areal.engine.megatron_utils.megatron import get_named_parameters + + num_moe_experts = getattr(engine.tf_config, "num_moe_experts", None) + seen_order: set[int] = set() + for _mcore_name, model_param in get_named_parameters( + engine.model, num_moe_experts + ): + pid = id(model_param) + if pid in seen_order: + continue + ordered_model_params.append(model_param) + seen_order.add(pid) + else: + ordered_model_params = fallback_order + return ordered_model_params + + def _collect_current_model_fingerprints( + self, fallback_order: list[torch.Tensor] | None = None + ) -> dict[str, tuple]: + fingerprints: dict[str, tuple] = {} + id2key, ptr2key = self._module_param_key_maps() + for model_param in self._ordered_model_params(fallback_order or []): + key = self._param_key(model_param, id2key, ptr2key) + if key is None or key in fingerprints: + continue + fingerprints[key] = _tensor_fingerprint(model_param.detach().reshape(-1)) + return fingerprints + + def _collect_payload_fingerprints( + self, payload_params: dict[str, torch.Tensor] | None = None + ) -> dict[str, tuple]: + if not payload_params: + return {} + return { + name: _tensor_fingerprint(tensor) for name, tensor in payload_params.items() + } + + @staticmethod + def _mask_to_index_mask(mask: torch.Tensor) -> torch.Tensor: + """Convert a temporary bool mask into compact external flat indices.""" + indices = mask.nonzero(as_tuple=False).squeeze(1) + if mask.numel() <= torch.iinfo(torch.int32).max: + return indices.to(torch.int32) + return indices + + def capture_synced_state( + self, payload_params: dict[str, torch.Tensor] | None = None + ): + """Capture post-step watermarks while model weights are still resident.""" + inner = self._adapter._get_inner_optimizers() + if not self._inversion_feasible(inner): + return None + return ( + self._collect_current_steps(inner), + self._collect_current_fingerprints(inner), + self._collect_current_model_fingerprints(), + self._collect_payload_fingerprints(payload_params), + ) + + def mark_synced(self, version: int, captured_state=None) -> None: + """Record optimizer step watermarks after a successful payload sync. + + AdamW inversion only reconstructs the most recent optimizer step. Without + this watermark, a later weight sync with no new optimizer step would replay + the previous step's moments and report false positives. + """ + # Drop any unconsumed precompute cache (e.g. the sync took the + # full-sync branch and never called compute_masks). + self._precomputed_masks = None + if captured_state is None: + captured_state = self.capture_synced_state() + if captured_state is None: + self._last_synced_steps.clear() + self._last_synced_fingerprints.clear() + self._last_synced_model_fingerprints.clear() + self._last_synced_payload_fingerprints.clear() + return + if len(captured_state) == 3: + ( + self._last_synced_steps, + self._last_synced_fingerprints, + self._last_synced_model_fingerprints, + ) = captured_state + self._last_synced_payload_fingerprints = {} + else: + ( + self._last_synced_steps, + self._last_synced_fingerprints, + self._last_synced_model_fingerprints, + self._last_synced_payload_fingerprints, + ) = captured_state + known = len(self._last_synced_steps) + stepped = sum( + 1 for step in self._last_synced_steps.values() if step is not None + ) + max_step = max( + (step for step in self._last_synced_steps.values() if step is not None), + default=0.0, + ) + logger.info( + "Inversion: recorded optimizer step watermark for %d params " + "at version %d (stepped=%d shard_fingerprints=%d " + "model_fingerprints=%d payload_fingerprints=%d max_step=%.0f)", + known, + version, + stepped, + len(self._last_synced_fingerprints), + len(self._last_synced_model_fingerprints), + len(self._last_synced_payload_fingerprints), + max_step, + ) + + # -- reconstruction (GPU-only) ---------------------------------------- + @torch.no_grad() + def _iter_reconstruct_pre_step_mcore( + self, inner_optimizers + ) -> Iterator[tuple[int, torch.Tensor | object]]: + """Yield theta_{t-1} overrides in model-parameter order. + + Every model parameter produces exactly one item. Parameters without a + usable override produce ``_NO_RECONSTRUCTION`` so an on-demand + consumer never has to drain the rest of the model just to establish a + missing key. Ambiguous replay is reported after all rank-aligned + collectives finish by raising ``_ReconstructionAborted``. + """ + from dte.core import invert_adamw + + reconstructed_count = 0 + skipped_no_state = 0 + skipped_bad_state = 0 + skipped_global_no_state = 0 + skipped_step_unchanged = 0 + skipped_missing_watermark = 0 + skipped_step_jump = 0 + skipped_missing_fingerprint = 0 + skipped_tracked_unchanged = 0 + skipped_payload_changed_without_step = 0 + skipped_untracked_non_optimizer = 0 + skipped_changed_non_optimizer = 0 + skipped_unchanged_non_optimizer = 0 + skipped_partial_state = 0 + unkeyed_opt_params = 0 + force_dense = False + debug = _inversion_debug_enabled() + global_rank = _dist_rank() + id2key, ptr2key = self._module_param_key_maps() + # opt_entries is keyed by the stable mcore name (see _param_key): on EP + # topologies id(optimizer param) may never match id(module param), and + # any rank-local miss must not change which collectives this rank + # enters (P-deadlock 2026-07-04: half the ranks fell back to the + # default group while their peers used the expert group -> crossed + # communicators, all ranks spinning in the first expert-param + # all_reduce). + opt_entries: dict[object, tuple] = {} + fallback_order: list[torch.Tensor] = [] + seen_fallback: set[int] = set() + default_dp_group = None + group_by_ranks: dict[tuple, object] = {} + for opt_idx, opt in enumerate(inner_optimizers): + base_opt = getattr(opt, "optimizer", opt) + state = getattr(base_opt, "state", None) + if state is None: + raise _ReconstructionAborted + group_index_map = getattr(opt, "model_param_group_index_map", None) + fp32_groups = getattr(opt, "shard_fp32_from_float16_groups", None) + model_groups = getattr(opt, "model_float16_groups", None) + if fp32_groups is None or model_groups is None: + raise _ReconstructionAborted + dp_group = getattr(opt, "data_parallel_group", None) + if default_dp_group is None: + default_dp_group = dp_group + if dp_group is not None: + # Canonicalize group OBJECTS by their actual global-rank + # membership: distinct communicators with identical members + # must collapse to one object, or peers end up posting the + # same logical reduce on different NCCL comms. + try: + ranks_key = tuple( + torch.distributed.get_process_group_ranks(dp_group) + ) + except Exception: + ranks_key = ( + "sig", + _dist_rank(dp_group), + _dist_world_size(dp_group), + ) + group_by_ranks.setdefault(ranks_key, dp_group) + if debug: + logger.info( + "Inversion debug: rank=%d opt=%d dp_rank=%d/%d " + "fp32_groups=%d model_groups=%d state=%d", + global_rank, + opt_idx, + _dist_rank(dp_group), + _dist_world_size(dp_group), + len(fp32_groups), + len(model_groups), + len(state), + ) + offloaded_states = getattr(self._adapter, "_offloaded_optimizer_states", {}) + for g, (fp32_group, model_group) in enumerate( + zip(fp32_groups, model_groups) + ): + if debug: + logger.info( + "Inversion debug: rank=%d opt=%d group=%d " + "fp32_len=%d model_len=%d", + global_rank, + opt_idx, + g, + len(fp32_group), + len(model_group), + ) + for i, (main_shard, model_param) in enumerate( + zip(fp32_group, model_group) + ): + if main_shard is None: + # precision-aware sentinel (should be gated out already) + raise _ReconstructionAborted + pid = id(model_param) + if pid not in seen_fallback: + fallback_order.append(model_param) + seen_fallback.add(pid) + key = self._param_key(model_param, id2key, ptr2key) + if key is None: + unkeyed_opt_params += 1 + continue + if key in opt_entries: + continue + gi = ( + group_index_map[model_param][0] + if group_index_map is not None + and model_param in group_index_map + else g + ) + opt_entries[key] = ( + opt_idx, + opt, + base_opt, + state, + main_shard, + gi, + model_param, + ) + + # Resolve the reduce group(s) once, rank-invariantly. A param's group + # must NEVER depend on whether THIS rank found an optimizer entry. + expert_group = None + if len(group_by_ranks) <= 1: + canonical_group = ( + next(iter(group_by_ranks.values())) + if group_by_ranks + else default_dp_group + ) + elif len(group_by_ranks) == 2: + dense_ranks = None + try: + from megatron.core import parallel_state as mpu + + try: + dense_group = mpu.get_data_parallel_group( + with_context_parallel=True + ) + except TypeError: + dense_group = mpu.get_data_parallel_group() + dense_ranks = tuple( + torch.distributed.get_process_group_ranks(dense_group) + ) + except Exception: + dense_ranks = None + if dense_ranks is None or dense_ranks not in group_by_ranks: + try: + dense_ranks = tuple( + torch.distributed.get_process_group_ranks(default_dp_group) + ) + except Exception: + dense_ranks = next(iter(group_by_ranks.keys())) + canonical_group = group_by_ranks.get(dense_ranks, default_dp_group) + expert_group = next( + (grp for k, grp in group_by_ranks.items() if k != dense_ranks), + None, + ) + else: + logger.warning( + "Inversion: %d distinct DP replica sets; only dense+expert " + "are supported -> dense.", + len(group_by_ranks), + ) + raise _ReconstructionAborted + logger.info( + "Inversion: canonical DP groups resolved (replica_sets=%d " + "expert_split=%s entries=%d unkeyed_opt_params=%d)", + len(group_by_ranks), + expert_group is not None, + len(opt_entries), + unkeyed_opt_params, + ) + + ordered_model_params = self._ordered_model_params(fallback_order) + + if debug: + logger.info( + "Inversion debug: rank=%d ordered_params=%d opt_entries=%d", + global_rank, + len(ordered_model_params), + len(opt_entries), + ) + + if not ordered_model_params: + raise _ReconstructionAborted + + offloaded_states = getattr(self._adapter, "_offloaded_optimizer_states", {}) + + # Post-reduce finalization, shared by the synchronous and pipelined + # paths. Runs strictly in parameter order on every rank; only local + # classification happens here so draining order can never diverge + # across ranks. + def _finalize_reduced(ctx) -> tuple[int, torch.Tensor | object]: + nonlocal force_dense + nonlocal reconstructed_count + nonlocal skipped_no_state, skipped_global_no_state + nonlocal skipped_partial_state, skipped_untracked_non_optimizer + nonlocal skipped_changed_non_optimizer + nonlocal skipped_unchanged_non_optimizer + ( + param_idx, + model_param, + param_key, + entry, + rng, + expected_one_step, + has_local_update, + flat_t, + theta_t_full, + correction, + status, + dp_group, + ) = ctx + # Single host sync for all three counters (was 3x .item()). + status_vals = status.tolist() + if debug: + logger.info( + "Inversion debug: rank=%d param_idx=%d post_reduce " + "contrib=%d bad=%d", + global_rank, + param_idx, + int(status_vals[0]), + int(status_vals[1]), + ) + if int(status_vals[1]) != 0: + force_dense = True + return id(model_param), _NO_RECONSTRUCTION + if int(status_vals[0]) == 0: + # No DP rank could reconstruct this param. For a known one-step + # AdamW transition that is unsafe: current-vs-current would hide + # a real update, so fall back to dense. If this rank had no + # optimizer entry at all, only now (after the DP all-reduce) + # classify it as a true non-AdamW payload tensor; a different DP + # rank may have owned and contributed the optimizer shard. + if entry is None: + last_model_fingerprint = ( + self._last_synced_model_fingerprints.get(param_key) + if param_key is not None + else None + ) + if last_model_fingerprint is None: + skipped_untracked_non_optimizer += 1 + elif _tensor_fingerprint(flat_t) != last_model_fingerprint: + skipped_changed_non_optimizer += 1 + force_dense = True + else: + skipped_unchanged_non_optimizer += 1 + else: + skipped_global_no_state += 1 + if expected_one_step: + force_dense = True + return id(model_param), _NO_RECONSTRUCTION + covered = int(status_vals[2]) + if covered != flat_t.numel(): + # Some DP rank's owned state slice is missing (unkeyed entry or + # unusable state on that rank). theta_old would silently equal + # theta_t on the uncovered slice and hide real updates there — + # never ship a sparse delta built from partial coverage. + skipped_partial_state += 1 + if debug and skipped_partial_state <= 8: + slice_start = -1 if rng is None else rng.start + slice_end = -1 if rng is None else rng.end + logger.warning( + "Inversion debug: partial optimizer coverage " + "rank=%d param_idx=%d key=%s shape=%s covered=%d " + "numel=%d local_slice=%d:%d dp_rank=%d/%d " + "local_update=%d force_dense=%d", + global_rank, + param_idx, + param_key, + tuple(theta_t_full.shape), + covered, + flat_t.numel(), + slice_start, + slice_end, + _dist_rank(dp_group), + _dist_world_size(dp_group), + int(has_local_update), + int(force_dense), + ) + force_dense = True + return id(model_param), _NO_RECONSTRUCTION + theta_old_full = (flat_t.to(torch.float32) + correction).reshape( + theta_t_full.shape + ) + reconstructed_count += 1 + return id(model_param), theta_old_full + + # Pipelined all-reduce window: post async reduces and only drain once + # this many correction bytes are in flight. Collective ORDER is + # identical to the synchronous path on every rank; only completion is + # overlapped ([dte-perf] showed ~18s of reconstruct dominated by + # serialized per-param reduce + 3x .item() sync round-trips). + window_bytes = _inversion_allreduce_window_bytes() + pending: list[tuple] = [] + pending_bytes = 0 + + def _drain_one_pending() -> tuple[int, torch.Tensor | object]: + nonlocal pending_bytes + ctx, work_c, work_s = pending.pop(0) + correction = ctx[9] + pending_bytes -= correction.numel() * correction.element_size() + work_c.wait() + work_s.wait() + return _finalize_reduced(ctx) + + def _drain_pending_batch() -> list[tuple[int, torch.Tensor | object]]: + # Never yield into the HF converter while a DP async collective is + # still in flight. The converter may issue all-gathers on another + # communicator, and interleaving those with this batch can deadlock + # when ranks advance the lazy producer at slightly different times. + ready = [] + while pending: + ready.append(_drain_one_pending()) + return ready + + for param_idx, model_param in enumerate(ordered_model_params): + param_key = self._param_key(model_param, id2key, ptr2key) + entry = opt_entries.get(param_key) if param_key is not None else None + # Group choice is a pure function of the (rank-invariant) param + # name — never of local entry state. + if ( + expert_group is not None + and isinstance(param_key, str) + and (".experts." in param_key) + ): + dp_group = expert_group + else: + dp_group = canonical_group + # theta_t for the full param is resident on every DP rank. + theta_t_full = model_param.detach() + flat_t = theta_t_full.reshape(-1) + # Every DP rank must enter the same collectives for every model + # param in the same order. Ranks without a usable local optimizer + # shard contribute an all-zero correction and a zero contribution + # count; ranks with moments fill only their owned slice. + correction = torch.zeros_like(flat_t, dtype=torch.float32) + has_local_update = False + expected_one_step = False + rng = None + opt_idx = -1 + if entry is None: + skipped_no_state += 1 + else: + opt_idx, opt, base_opt, state, main_shard, gi, opt_model_param = entry + rng = opt._get_model_param_range_map(opt_model_param)["param"] + st = state.get(main_shard) + offloaded_state = offloaded_states.get(main_shard, {}) if st else {} + step = _state_step_value(st, offloaded_state) if st else None + if step is None: + step = _param_group_step_value(base_opt.param_groups[gi]) + if step is None or step < 1: + skipped_bad_state += 1 + elif param_key is None or param_key not in self._last_synced_steps: + skipped_missing_watermark += 1 + force_dense = True + else: + last_step = self._last_synced_steps[param_key] + baseline_step = 0.0 if last_step is None else last_step + step_delta = step - baseline_step + if step_delta == 0: + last_fingerprint = self._last_synced_fingerprints.get(param_key) + if last_fingerprint is not None and ( + _tensor_fingerprint(flat_t[rng.start : rng.end]) + != last_fingerprint + ): + skipped_payload_changed_without_step += 1 + force_dense = True + else: + skipped_step_unchanged += 1 + elif step_delta != 1: + skipped_step_jump += 1 + force_dense = True + else: + expected_one_step = True + last_fingerprint = self._last_synced_fingerprints.get(param_key) + if last_fingerprint is None: + skipped_missing_fingerprint += 1 + force_dense = True + elif not st or "exp_avg" not in st or "exp_avg_sq" not in st: + # A rank may not own this shard in Megatron's + # distributed optimizer; do not force dense locally. + # If no DP rank contributes below, the global status + # check will fall back to dense. + skipped_bad_state += 1 + else: + exp_avg = _resolve_offloaded_state_value( + st, offloaded_state, "exp_avg" + ) + exp_avg_sq = _resolve_offloaded_state_value( + st, offloaded_state, "exp_avg_sq" + ) + main_flat = main_shard.detach().reshape(-1) + if ( + not isinstance(exp_avg, torch.Tensor) + or not isinstance(exp_avg_sq, torch.Tensor) + or exp_avg.numel() == 0 + or exp_avg_sq.numel() == 0 + or exp_avg.numel() != main_flat.numel() + or exp_avg_sq.numel() != main_flat.numel() + ): + skipped_bad_state += 1 + else: + lr, wd, b1, b2, eps = _adamw_hparams( + base_opt.param_groups[gi] + ) + # The fp32 main shard and moments may be + # offloaded to CPU under colocate; stage them on + # the inversion compute device (the visible + # payload's device by default) so invert_adamw + # runs on GPU instead of the CPU. Inversion + # still uses the exact post-step optimizer + # weight, not the bf16 model copy. Keep the + # reconstructed previous value in fp32 until + # mask creation so boundary-near BF16 roundoff + # is handled conservatively. + dev = _inversion_compute_device(flat_t) + theta_t_slice = main_flat.to( + device=dev, dtype=torch.float32 + ) + theta_old_slice = invert_adamw( + theta_t_slice, + exp_avg.to(device=dev, dtype=torch.float32).reshape( + -1 + ), + exp_avg_sq.to( + device=dev, dtype=torch.float32 + ).reshape(-1), + step, + lr, + wd, + b1, + b2, + eps, + ) + visible_slice = flat_t[rng.start : rng.end] + visible_fp32 = visible_slice.to( + device=theta_old_slice.device, + dtype=torch.float32, + ) + correction[rng.start : rng.end] = ( + theta_old_slice - visible_fp32 + ).to( + device=correction.device, + dtype=correction.dtype, + ) + has_local_update = True + + owned_numel = ( + int(rng.end - rng.start) if has_local_update and rng is not None else 0 + ) + status = torch.tensor( + [ + 1 if has_local_update else 0, + 1 if force_dense else 0, + owned_numel, + ], + dtype=torch.int64, + device=flat_t.device, + ) + if debug: + slice_start = -1 if rng is None else rng.start + slice_end = -1 if rng is None else rng.end + logger.info( + "Inversion debug: rank=%d param_idx=%d opt=%d " + "pre_reduce numel=%d slice=%d:%d local=%d", + global_rank, + param_idx, + opt_idx, + flat_t.numel(), + slice_start, + slice_end, + int(has_local_update), + ) + ctx = ( + param_idx, + model_param, + param_key, + entry, + rng, + expected_one_step, + has_local_update, + flat_t, + theta_t_full, + correction, + status, + dp_group, + ) + if dp_group is None: + yield _finalize_reduced(ctx) + elif window_bytes <= 0: + torch.distributed.all_reduce( + correction, + op=torch.distributed.ReduceOp.SUM, + group=dp_group, + ) + torch.distributed.all_reduce( + status, + op=torch.distributed.ReduceOp.SUM, + group=dp_group, + ) + yield _finalize_reduced(ctx) + else: + work_c = torch.distributed.all_reduce( + correction, + op=torch.distributed.ReduceOp.SUM, + group=dp_group, + async_op=True, + ) + work_s = torch.distributed.all_reduce( + status, + op=torch.distributed.ReduceOp.SUM, + group=dp_group, + async_op=True, + ) + pending.append((ctx, work_c, work_s)) + pending_bytes += correction.numel() * correction.element_size() + if pending_bytes > window_bytes: + yield from _drain_pending_batch() + yield from _drain_pending_batch() + if force_dense: + logger.warning( + "Inversion: optimizer step replay is ambiguous " + "(missing_watermark=%d step_jump=%d missing_fingerprint=%d " + "payload_changed_without_step=%d " + "untracked_non_optimizer=%d changed_non_optimizer=%d " + "partial_state=%d unkeyed_opt_params=%d) " + "-> dense.", + skipped_missing_watermark, + skipped_step_jump, + skipped_missing_fingerprint, + skipped_payload_changed_without_step, + skipped_untracked_non_optimizer, + skipped_changed_non_optimizer, + skipped_partial_state, + unkeyed_opt_params, + ) + raise _ReconstructionAborted + if reconstructed_count == 0: + logger.info( + "Inversion: no BF16 payload shard changed since last sync " + "(step_unchanged=%d tracked_unchanged=%d no_state=%d " + "bad_state=%d global_no_state=%d " + "unchanged_non_optimizer=%d)", + skipped_step_unchanged, + skipped_tracked_unchanged, + skipped_no_state, + skipped_bad_state, + skipped_global_no_state, + skipped_unchanged_non_optimizer, + ) + return + if ( + skipped_no_state + or skipped_bad_state + or skipped_global_no_state + or skipped_step_unchanged + or skipped_tracked_unchanged + or skipped_missing_fingerprint + or skipped_payload_changed_without_step + or skipped_unchanged_non_optimizer + ): + logger.info( + "Inversion: reconstructed %d params, skipped no_state=%d " + "bad_state=%d global_no_state=%d step_unchanged=%d " + "tracked_unchanged=%d missing_fingerprint=%d " + "payload_changed_without_step=%d " + "unchanged_non_optimizer=%d", + reconstructed_count, + skipped_no_state, + skipped_bad_state, + skipped_global_no_state, + skipped_step_unchanged, + skipped_tracked_unchanged, + skipped_missing_fingerprint, + skipped_payload_changed_without_step, + skipped_unchanged_non_optimizer, + ) + return + + @torch.no_grad() + def _reconstruct_pre_step_mcore( + self, inner_optimizers + ) -> dict[int, torch.Tensor] | None: + """Eager compatibility wrapper around streaming reconstruction.""" + theta_old_by_id: dict[int, torch.Tensor] = {} + try: + for param_id, theta_old in self._iter_reconstruct_pre_step_mcore( + inner_optimizers + ): + if theta_old is not _NO_RECONSTRUCTION: + theta_old_by_id[param_id] = theta_old + except _ReconstructionAborted: + return None + return theta_old_by_id + + # -- entry point ------------------------------------------------------- + @torch.no_grad() + def precompute_masks(self, names, tensors, version) -> bool: + """Compute and cache masks while optimizer state is still GPU-resident. + + The colocate gateway offloads the optimizer (fp32 main shards + AdamW + moments) to CPU right before the weight sync to make room for the + inference-side apply. Calling this beforehand lets the inversion run + with all operands on GPU and moves the mask computation off the sync + critical path. The cached result (including ``None`` = infeasible) is + consumed by the next ``compute_masks`` call for the same version and + parameter list; any mismatch discards the cache and recomputes. + + All training ranks must call this together: the mask computation runs + the same DP/TP collectives as ``compute_masks``. + """ + self._precomputed_masks = None + masks = self.compute_masks(names, tensors, version) + self._precomputed_masks = (int(version), tuple(names), masks) + return masks is not None + + def _pop_precomputed_masks(self, names, version): + """Return ``(hit, masks)`` consuming the precompute cache if valid. + + The hit decision must be rank-invariant: a cache hit skips the DP/TP + collectives inside the recompute path, so if any training rank misses + (e.g. a restarted worker lost its cache, or a stale version), every + rank must recompute or the collectives would mismatch and hang. A + cheap global MIN-reduce promotes any local miss to a global miss. + """ + cached = self._precomputed_masks + self._precomputed_masks = None + local_hit = False + masks = None + if cached is not None: + c_version, c_names, c_masks = cached + if c_version == int(version) and c_names == tuple(names): + local_hit = True + masks = c_masks + else: + logger.warning( + "Inversion: discarding stale precomputed masks " + "(cached v%d/%d params vs requested v%d/%d params)", + c_version, + len(c_names), + int(version), + len(names), + ) + if ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and _dist_world_size() > 1 + ): + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + hit_t = torch.tensor( + [1 if local_hit else 0], dtype=torch.int32, device=device + ) + torch.distributed.all_reduce(hit_t, op=torch.distributed.ReduceOp.MIN) + if int(hit_t.item()) == 0: + if local_hit: + logger.warning( + "Inversion: dropping precomputed masks for v%d; " + "another rank missed its cache, recomputing on all " + "ranks to keep collectives aligned", + int(version), + ) + return False, None + return (local_hit, masks) if local_hit else (False, None) + + def compute_masks(self, names, tensors, version): + """Return sparse index masks or ``None`` to force a dense step.""" + hit, cached_masks = self._pop_precomputed_masks(names, version) + if hit: + logger.info( + "Inversion: using precomputed masks for v%d (%d params, %s)", + int(version), + len(names), + "feasible" if cached_masks is not None else "dense_infeasible", + ) + return cached_masks + perf_start = time.monotonic() + # Reset the peak counter so peak_mb below covers inversion only. + cuda_mem_stats_mb() + adapter = self._adapter + t_inner = time.monotonic() + inner = adapter._get_inner_optimizers() + inner_ms = (time.monotonic() - t_inner) * 1000 + t_feasible = time.monotonic() + if not self._inversion_feasible(inner): + logger.info( + "[dte-perf][inversion] v%d result=dense_infeasible " + "get_optimizers_ms=%.1f feasible_ms=%.1f total_ms=%.1f", + version, + inner_ms, + (time.monotonic() - t_feasible) * 1000, + (time.monotonic() - perf_start) * 1000, + ) + return None + feasible_ms = (time.monotonic() - t_feasible) * 1000 + + cur = dict(zip(names, tensors)) + masks: dict[str, torch.Tensor] = {} + from dte.core import bitwise_changed_mask + + # Convert reconstructed pre-step params through the SAME all_gather + + # convert_to_hf path as the live payload. Prefer the streaming adapter + # interface so AdamW inversion does not materialize a second full HF + # payload (current HF payload + old HF payload) before mask creation. + iter_hf = getattr(adapter, "_iter_hf_with_overrides", None) + lazy: _LazyReconstructDict | None = None + streaming_reconstruct = _streaming_reconstruct_enabled() and iter_hf is not None + t_reconstruct = time.monotonic() + if streaming_reconstruct: + theta_old_by_id = _LazyReconstructDict( + self._iter_reconstruct_pre_step_mcore(inner) + ) + reconstruct_ms = 0.0 + theta_old_param_count = 0 + else: + theta_old_by_id = self._reconstruct_pre_step_mcore(inner) + reconstruct_ms = (time.monotonic() - t_reconstruct) * 1000 + if theta_old_by_id is None: + logger.warning("Inversion: reconstruction unavailable -> dense step.") + logger.info( + "[dte-perf][inversion] v%d " + "result=dense_reconstruct_unavailable " + "get_optimizers_ms=%.1f feasible_ms=%.1f " + "reconstruct_mcore_ms=%.1f total_ms=%.1f", + version, + inner_ms, + feasible_ms, + reconstruct_ms, + (time.monotonic() - perf_start) * 1000, + ) + return None + theta_old_param_count = len(theta_old_by_id) + + old_items = None + stream_old_hf = iter_hf is not None + convert_hf_ms = 0.0 + if stream_old_hf: + if streaming_reconstruct: + lazy = theta_old_by_id + old_items = iter(iter_hf(theta_old_by_id)) + else: + t_convert = time.monotonic() + hf_old = adapter._convert_hf_with_overrides(theta_old_by_id) + convert_hf_ms = (time.monotonic() - t_convert) * 1000 + old_items = iter(hf_old.items()) + + mask_loop_ms = 0.0 + full_changed_masks = 0 + shape_mismatch_masks = 0 + forced_dense_params = 0 + forced_dense_elements = 0 + seen_names: set[str] = set() + dense_verdict: str | None = None + while True: + if stream_old_hf: + t_next = time.monotonic() + try: + name, old_t = next(old_items) + except StopIteration: + break + if stream_old_hf: + convert_hf_ms += (time.monotonic() - t_next) * 1000 + + cur_t = cur.get(name) + if cur_t is None: + continue + seen_names.add(name) + t_mask = time.monotonic() + if _inversion_force_dense_param(name): + # Known-unsupported inversion cases (e.g. tiny Qwen3 MoE router + # gates) stay dense so snapshot verification cannot miss + # changed BF16 values (c3c80a935 false-negative fix). + masks[name] = torch.ones( + cur_t.numel(), dtype=torch.bool, device=cur_t.device + ) + forced_dense_params += 1 + forced_dense_elements += cur_t.numel() + mask_loop_ms += (time.monotonic() - t_mask) * 1000 + continue + if old_t is None or old_t.shape != cur_t.shape: + # No pre-step counterpart -> treat as fully changed (dense). + masks[name] = torch.ones( + cur_t.numel(), dtype=torch.bool, device=cur_t.device + ) + full_changed_masks += 1 + if old_t is not None: + shape_mismatch_masks += 1 + mask_loop_ms += (time.monotonic() - t_mask) * 1000 + continue + old_payload = old_t.to(cur_t.dtype) + mask = bitwise_changed_mask(cur_t, old_payload).reshape(-1) + if cur_t.dtype == torch.bfloat16 and old_t.dtype != cur_t.dtype: + same_payload = ~mask + boundary = _bf16_rounding_boundary_mask_chunked(old_t, cur_t).reshape( + -1 + ) + mask = mask | (same_payload & boundary) + if not bool(mask.any().item()) and self._last_synced_payload_fingerprints: + last_payload_fingerprint = self._last_synced_payload_fingerprints.get( + name + ) + if last_payload_fingerprint is None: + if dense_verdict is None: + logger.warning( + "Inversion: payload %s has no synced fingerprint while " + "detector mask is empty -> dense step.", + name, + ) + dense_verdict = "missing_payload_fingerprint" + cur_payload_fingerprint = _tensor_fingerprint(cur_t) + if ( + last_payload_fingerprint is not None + and cur_payload_fingerprint != last_payload_fingerprint + ): + if dense_verdict is None: + logger.warning( + "Inversion: payload %s changed since last sync but " + "detector mask is empty -> dense step.", + name, + ) + dense_verdict = "payload_fingerprint_changed" + masks[name] = self._mask_to_index_mask(mask) + mask_loop_ms += (time.monotonic() - t_mask) * 1000 + + if lazy is not None: + # The converter may not request every model parameter. Draining is + # still mandatory so trailing collectives and a late dense verdict + # execute on every rank. + producer_during_convert_ms = lazy.producer_seconds * 1000 + lazy.finish() + reconstruct_ms = lazy.producer_seconds * 1000 + theta_old_param_count = lazy.reconstructed_count + convert_hf_ms = max(0.0, convert_hf_ms - producer_during_convert_ms) + if lazy.aborted: + dense_verdict = dense_verdict or "reconstruction_aborted" + lazy.clear() + + for name, cur_t in cur.items(): + if name in seen_names: + continue + t_mask = time.monotonic() + masks[name] = torch.ones( + cur_t.numel(), dtype=torch.bool, device=cur_t.device + ) + full_changed_masks += 1 + mask_loop_ms += (time.monotonic() - t_mask) * 1000 + if dense_verdict is not None: + logger.warning( + "Inversion: completed HF conversion before dense fallback (reason=%s).", + dense_verdict, + ) + return None + logger.info( + "Inversion: computed masks for %d HF params at version %d " + "(forced_dense_params=%d forced_dense_elements=%d)", + len(masks), + version, + forced_dense_params, + forced_dense_elements, + ) + alloc_mb, peak_mb = cuda_mem_stats_mb() + logger.info( + "[dte-perf][inversion] v%d result=sparse params=%d " + "theta_old_params=%d full_changed_masks=%d " + "shape_mismatch_masks=%d get_optimizers_ms=%.1f " + "feasible_ms=%.1f reconstruct_mcore_ms=%.1f " + "convert_hf_ms=%.1f mask_loop_ms=%.1f stream_old_hf=%d " + "streaming_reconstruct=%d " + "total_ms=%.1f alloc_mb=%.0f peak_mb=%.0f", + version, + len(masks), + theta_old_param_count, + full_changed_masks, + shape_mismatch_masks, + inner_ms, + feasible_ms, + reconstruct_ms, + convert_hf_ms, + mask_loop_ms, + int(stream_old_hf), + int(streaming_reconstruct), + (time.monotonic() - perf_start) * 1000, + alloc_mb, + peak_mb, + ) + return masks diff --git a/areal/v2/weight_update/awex/megatron_adapter.py b/areal/v2/weight_update/awex/megatron_adapter.py index cec905ceee..ed3cc2f988 100644 --- a/areal/v2/weight_update/awex/megatron_adapter.py +++ b/areal/v2/weight_update/awex/megatron_adapter.py @@ -18,7 +18,8 @@ from awex.sharding.param_sharding import ShardingType from awex.sharding.rank_info import RankInfo from awex.transfer.nccl_comm import batch_send_recv, nccl_build_send_ops -from awex.transfer.transfer_plan import TransferPlan, TransferPlanBuilder +from awex.transfer.nccl_stream_batch import NcclColocateStreamBatchTransport +from awex.transfer.transfer_plan import TransferPlan, TransferPlanBuilder, slice_tensor from awex.util.tensor_util import ( cuda_ipc_serialize, group_tensors_by_shape_and_dtype, @@ -29,6 +30,11 @@ awex_wu_use_group, fetch_kv_metadata, ) +from areal.v2.weight_update.awex.delta_config import ( + make_delta_tracker, + separation_delta_transfer_enabled, +) +from areal.v2.weight_update.awex.delta_detect import AdamWInversionDetector from areal.v2.weight_update.nccl_group import ( init_weights_update_group, setup_batch_isend_irecv, @@ -62,6 +68,8 @@ def __init__(self, engine: MegatronEngine): self._transfer_plan: TransferPlan | None = None self._weights_update_group = None self._weights_update_group_gloo = None + self._world_size: int | None = None + self._separation_delta_transport: NcclColocateStreamBatchTransport | None = None self._transfer_rank: int | None = None self._offloaded_optimizer_states: dict = {} self._offloaded_weights: dict[str, torch.Tensor] = {} @@ -70,6 +78,8 @@ def __init__(self, engine: MegatronEngine): self._colocate_admin_api_key: str = "areal-admin-key" self._colocate_http_client: httpx.Client | None = None self._colocate_timeout_s: float = 120.0 + self._delta_tracker = None + self._delta_detector = None @property def parallelism_strategy(self) -> dict: @@ -164,6 +174,7 @@ def init_weight_update_group( num_engines: int, ) -> None: self._transfer_rank = transfer_rank + self._world_size = world_size infer_meta, train_meta = fetch_kv_metadata(kv_store_url, pair_name) @@ -205,7 +216,14 @@ def init_weight_update_group( ) def execute_weight_update(self, version: int) -> None: - del version + if separation_delta_transfer_enabled(): + self._release_grad_buffers_for_separation_sync() + try: + self._execute_separation_weight_update(version) + finally: + self._restore_grad_buffers_after_separation_sync() + return + if self._transfer_plan is None: raise RuntimeError("Transfer plan is not initialized") if self._weights_update_group is None: @@ -230,6 +248,204 @@ def execute_weight_update(self, version: int) -> None: ) dist.barrier(group=self._weights_update_group_gloo) + def _release_grad_buffers_for_separation_sync(self) -> None: + """Temporarily release Megatron DDP grad buffers during transfer.""" + model = getattr(self._engine, "model", None) + if model is None: + return + modules = model if isinstance(model, (list, tuple)) else [model] + for module in modules: + release = getattr(module, "offload_grad_buffers", None) + if release is not None: + release(synchronize=False, empty_cache=False) + + def _restore_grad_buffers_after_separation_sync(self) -> None: + """Restore Megatron DDP grad buffers even when transfer fails.""" + model = getattr(self._engine, "model", None) + if model is None: + return + modules = model if isinstance(model, (list, tuple)) else [model] + for module in modules: + restore = getattr(module, "restore_grad_buffers", None) + if restore is not None: + restore(synchronize=False) + + def _execute_separation_weight_update(self, version: int) -> None: + """Send an AdamW-derived sparse update, with a dense fallback.""" + if self._transfer_plan is None: + raise RuntimeError("Transfer plan is not initialized") + if self._weights_update_group is None: + raise RuntimeError("Weight update group is not initialized") + if self._weights_update_group_gloo is None: + raise RuntimeError("Gloo weight update group is not initialized") + if self._transfer_rank is None or self._world_size is None: + raise RuntimeError("Transfer rank/world size is not initialized") + + params = self.get_local_shard_parameters() + self._ensure_delta_components() + synced_state = self._delta_detector.capture_synced_state(params) + masks, local_is_delta = self._delta_prepare_masks(params, version) + + decision = torch.tensor([int(local_is_delta)], dtype=torch.int64) + dist.all_reduce( + decision, op=dist.ReduceOp.MIN, group=self._weights_update_group_gloo + ) + use_delta = bool(decision.item()) and masks is not None + + if use_delta: + self._execute_separation_delta_send(params, masks, version) + else: + send_ops, _, _ = nccl_build_send_ops( + params, + self._transfer_plan, + self._weights_update_group, + copy_rank=self._transfer_rank, + ) + batch_send_recv( + send_ops=send_ops, + recv_ops=[], + blocking=True, + use_group=awex_wu_use_group(), + ) + + # The receiver joins this barrier after applying the payload. Only then + # may the sender advance its version/watermark state. + dist.barrier(group=self._weights_update_group_gloo) + if use_delta: + self._delta_tracker.mark_delta_committed(version) + self._delta_detector.mark_synced(version, synced_state) + + def _delta_prepare_masks( + self, + params: dict[str, torch.Tensor], + version: int, + ) -> tuple[dict[str, torch.Tensor] | None, bool]: + reason = self._delta_tracker.full_sync_reason(version) + if reason is None and not self._delta_detector.has_synced_watermark(): + reason = "initial_full" + reason = self._sync_full_reason(reason, version) + + masks = None + if reason is None: + try: + masks = self._delta_detector.compute_masks( + list(params), list(params.values()), version + ) + except Exception: + logger.exception( + "separation delta v%d: AdamW inversion failed; using full sync", + version, + ) + reason = "adamw_inversion_error" + if masks is None and reason is None: + reason = "adamw_inversion_infeasible" + + reason = self._sync_full_reason(reason, version) + if reason is not None: + self._delta_tracker.seed(params.items(), version, store_snapshot=False) + logger.info( + "separation delta v%d: FULL sync fallback (%s)", version, reason + ) + return None, False + + logger.info( + "separation delta v%d: sparse AdamW path (%d params)", + version, + len(params), + ) + return masks, True + + def _sync_full_reason(self, reason: str | None, version: int) -> str | None: + """Promote a rank-local dense fallback to every training rank.""" + if not dist.is_available() or not dist.is_initialized(): + return reason + try: + world_size = dist.get_world_size() + except RuntimeError: + return reason + if world_size <= 1: + return reason + + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + needs_full = torch.tensor( + [1 if reason is not None else 0], dtype=torch.int32, device=device + ) + dist.all_reduce(needs_full, op=dist.ReduceOp.MAX) + if int(needs_full.item()) == 0: + return None + if reason is not None: + return reason + logger.warning("separation delta v%d: peer rank requires a full sync", version) + return "peer_rank_fallback" + + def _execute_separation_delta_send( + self, + params: dict[str, torch.Tensor], + masks: dict[str, torch.Tensor], + version: int, + ) -> None: + from dte.core.colocate_protocol import ( + _filter_plan_by_dtype, + _ops_by_recv_dtype, + _PlanView, + two_round_delta_exchange, + ) + from dte.core.delta_p2p import build_send_payloads_by_op + + assert self._transfer_plan is not None + assert self._weights_update_group is not None + assert self._transfer_rank is not None + assert self._world_size is not None + + operations = [ + op for ops in self._transfer_plan.operations.values() for op in ops + ] + operations_by_dtype = _ops_by_recv_dtype(operations) + identity_mapping = {rank: rank for rank in range(self._world_size)} + empty_plan = _PlanView({}) + device = torch.device(f"cuda:{torch.cuda.current_device()}") + + if self._separation_delta_transport is None: + self._separation_delta_transport = NcclColocateStreamBatchTransport( + self._transfer_rank, self._world_size + ) + schedule_fn = ( + self._separation_delta_transport.execute_recursive_partition_stream_transfer + ) + + payload_count = 0 + for dtype, ops in operations_by_dtype.items(): + payloads = build_send_payloads_by_op(ops, masks, params) + send_plan = _filter_plan_by_dtype(self._transfer_plan, dtype, is_send=True) + two_round_delta_exchange( + transfer_rank=self._transfer_rank, + world_size=self._world_size, + send_plan=send_plan, + recv_plan=empty_plan, + train_to_infer_device_mapping=identity_mapping, + weights_update_group=self._weights_update_group, + send_payloads_by_op=payloads, + recv_params={}, + value_dtype=dtype, + device=device, + schedule_fn=schedule_fn, + slice_fn=slice_tensor, + rank_coordinate=f"train-{self._transfer_rank}", + step_id=version, + ) + payload_count += len(payloads) + + logger.info( + "separation delta v%d sent %d payload ops across %d dtypes", + version, + payload_count, + len(operations_by_dtype), + ) + def batch_isend_irecv(self, **kwargs) -> None: if self._weights_update_group_gloo is None: raise RuntimeError("Gloo weight update group is not initialized") @@ -253,6 +469,10 @@ def teardown_weight_update_group(self) -> None: self._weights_update_group_gloo = None self._transfer_plan = None self._transfer_rank = None + self._world_size = None + self._separation_delta_transport = None + self._delta_tracker = None + self._delta_detector = None if self._colocate_http_client is not None: self._colocate_http_client.close() self._colocate_http_client = None @@ -296,7 +516,11 @@ def _build_rank_info(self) -> RankInfo: cp_mode="ring" if cp_size > 1 else "none", ) - def _iter_hf_params(self): + def _iter_hf_params( + self, + theta_by_id: dict[int, torch.Tensor] | None = None, + consume_overrides: bool = False, + ): """Yield (hf_name, tensor) for every parameter on this rank. Uses get_named_parameters + all_gather_param + convert_to_hf to produce @@ -315,13 +539,23 @@ def _iter_hf_params(self): tie_word_embeddings = getattr( self._engine.hf_config, "tie_word_embeddings", False ) + overrides = theta_by_id if theta_by_id is not None else {} for mcore_name, param in get_named_parameters( self._engine.model, num_moe_experts ): + src = overrides.get(id(param), param) + if src is not param: + for attr in ( + "tensor_model_parallel", + "partition_dim", + "partition_stride", + ): + if hasattr(param, attr): + setattr(src, attr, getattr(param, attr)) gathered = all_gather_param( mcore_name, - param, + src, fp8_direct_convert=False, quantization_config=None, duplicated_param_names=self._engine._duplicated_param_names, @@ -338,6 +572,47 @@ def _iter_hf_params(self): if tie_word_embeddings and hf_name == "lm_head.weight": continue yield hf_name, tensor.detach() + if consume_overrides: + overrides.pop(id(param), None) + + def _iter_model_params_for_delta(self): + """Yield model tensors in the same order used by the HF converter.""" + from areal.engine.megatron_utils.megatron import get_named_parameters + + num_moe_experts = getattr(self._engine.tf_config, "num_moe_experts", None) + seen: set[int] = set() + for _mcore_name, param in get_named_parameters( + self._engine.model, num_moe_experts + ): + if id(param) in seen: + continue + seen.add(id(param)) + yield param + + def _convert_hf_with_overrides( + self, theta_by_id: dict[int, torch.Tensor] + ) -> dict[str, torch.Tensor]: + return dict(self._iter_hf_params(theta_by_id)) + + @torch.no_grad() + def _iter_hf_with_overrides(self, theta_by_id: dict[int, torch.Tensor]): + yield from self._iter_hf_params(theta_by_id, consume_overrides=True) + + def _ensure_delta_components(self) -> None: + if self._delta_tracker is None: + self._delta_tracker = make_delta_tracker() + if self._delta_detector is None: + self._delta_detector = AdamWInversionDetector(self) + + def _get_inner_optimizers(self): + optimizer = self._engine.optimizer + if optimizer is None: + return [] + if hasattr(optimizer, "chained_optimizers"): + return optimizer.chained_optimizers + if hasattr(optimizer, "optimizers"): + return optimizer.optimizers + return [optimizer] # ── Colocated weight transfer methods ───────────────────────────────── diff --git a/areal/v2/weight_update/awex/sglang_adapter.py b/areal/v2/weight_update/awex/sglang_adapter.py index 674f84a21a..6d4155593c 100644 --- a/areal/v2/weight_update/awex/sglang_adapter.py +++ b/areal/v2/weight_update/awex/sglang_adapter.py @@ -24,7 +24,7 @@ ) from awex.transfer.nccl_comm import batch_send_recv, nccl_build_recv_ops from awex.transfer.nccl_stream_batch import NcclColocateStreamBatchTransport -from awex.transfer.transfer_plan import TransferPlan, TransferPlanBuilder +from awex.transfer.transfer_plan import TransferPlan, TransferPlanBuilder, slice_tensor from awex.util.tensor_util import ( cuda_ipc_deserialize, reconstruct_tensors_from_groups, @@ -36,6 +36,9 @@ awex_wu_use_group, fetch_kv_metadata, ) +from areal.v2.weight_update.awex.delta_config import ( + separation_delta_transfer_enabled, +) from areal.v2.weight_update.inference_adapter import ( AwexInferenceAdapter, ) @@ -55,6 +58,8 @@ def __init__(self, scheduler: Any): self._transfer_plan: TransferPlan | None = None self._weights_update_group = None self._weights_update_group_gloo = None + self._world_size: int | None = None + self._separation_delta_transport: NcclColocateStreamBatchTransport | None = None self._transfer_rank: int | None = None self._rank_info: RankInfo | None = None self._parameters: dict[str, torch.Tensor] | None = None @@ -379,6 +384,7 @@ def init_weight_update_group( engine_local_rank = pp_rank * tp_size + tp_rank global_rank = transfer_rank * per_engine_world + engine_local_rank self._transfer_rank = global_rank + self._world_size = world_size infer_meta, train_meta = fetch_kv_metadata(kv_store_url, pair_name) @@ -420,7 +426,10 @@ def init_weight_update_group( ) def execute_weight_update(self, version: int) -> None: - del version + if separation_delta_transfer_enabled(): + self._execute_separation_weight_update(version) + return + if self._transfer_plan is None: raise RuntimeError("Transfer plan is not initialized") if self._weights_update_group is None: @@ -447,6 +456,105 @@ def execute_weight_update(self, version: int) -> None: current_platform.synchronize() dist.barrier(group=self._weights_update_group_gloo) + def _execute_separation_weight_update(self, version: int) -> None: + """Receive either a sparse AdamW update or its dense fallback.""" + if self._transfer_plan is None: + raise RuntimeError("Transfer plan is not initialized") + if self._weights_update_group is None: + raise RuntimeError("Weight update group is not initialized") + if self._weights_update_group_gloo is None: + raise RuntimeError("Gloo weight update group is not initialized") + + decision = torch.tensor([1], dtype=torch.int64) + dist.all_reduce( + decision, op=dist.ReduceOp.MIN, group=self._weights_update_group_gloo + ) + use_delta = bool(decision.item()) + params = self.get_local_shard_parameters() + + if use_delta: + self._execute_separation_delta_recv(params, version) + else: + recv_ops, non_contiguous_pairs, _ = nccl_build_recv_ops( + params, + self._transfer_plan, + self._weights_update_group, + ) + batch_send_recv( + send_ops=[], + recv_ops=recv_ops, + blocking=True, + use_group=awex_wu_use_group(), + ) + for original, contiguous in non_contiguous_pairs: + original.copy_(contiguous) + + current_platform.synchronize() + dist.barrier(group=self._weights_update_group_gloo) + + def _execute_separation_delta_recv( + self, + recv_params: dict[str, torch.Tensor], + version: int, + ) -> None: + from dte.core.colocate_protocol import ( + _filter_plan_by_dtype, + _ops_by_recv_dtype, + _PlanView, + two_round_delta_exchange, + ) + + if self._transfer_plan is None: + raise RuntimeError("Transfer plan is not initialized") + if self._weights_update_group is None: + raise RuntimeError("Weight update group is not initialized") + if self._transfer_rank is None or self._world_size is None: + raise RuntimeError("Transfer rank/world size is not initialized") + + operations = [ + op for ops in self._transfer_plan.operations.values() for op in ops + ] + operations_by_dtype = _ops_by_recv_dtype(operations) + identity_mapping = {rank: rank for rank in range(self._world_size)} + empty_plan = _PlanView({}) + device = torch.device(f"cuda:{torch.cuda.current_device()}") + + if self._separation_delta_transport is None: + self._separation_delta_transport = NcclColocateStreamBatchTransport( + self._transfer_rank, self._world_size + ) + schedule_fn = ( + self._separation_delta_transport.execute_recursive_partition_stream_transfer + ) + + operation_count = 0 + for dtype, ops in operations_by_dtype.items(): + recv_plan = _filter_plan_by_dtype(self._transfer_plan, dtype, is_send=False) + two_round_delta_exchange( + transfer_rank=self._transfer_rank, + world_size=self._world_size, + send_plan=empty_plan, + recv_plan=recv_plan, + train_to_infer_device_mapping=identity_mapping, + weights_update_group=self._weights_update_group, + send_payloads_by_op={}, + recv_params=recv_params, + value_dtype=dtype, + device=device, + schedule_fn=schedule_fn, + slice_fn=slice_tensor, + rank_coordinate=f"infer-{self._transfer_rank}", + step_id=version, + ) + operation_count += len(ops) + + logger.info( + "separation delta v%d received %d ops across %d dtypes", + version, + operation_count, + len(operations_by_dtype), + ) + def batch_isend_irecv(self, **kwargs) -> None: if self._weights_update_group_gloo is None: raise RuntimeError("Gloo weight update group is not initialized") @@ -470,6 +578,8 @@ def teardown_weight_update_group(self) -> None: self._weights_update_group_gloo = None self._transfer_plan = None self._transfer_rank = None + self._world_size = None + self._separation_delta_transport = None self._rank_info = None self._parameters = None if self._colocate_http_client is not None: diff --git a/tests/test_awex_delta_common.py b/tests/test_awex_delta_common.py new file mode 100644 index 0000000000..4a72e99652 --- /dev/null +++ b/tests/test_awex_delta_common.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared Stage-3 tests and loaders for AWEX DTE delta transfer. + +The runtime adapters import optional packages such as AWEX, DTE, and httpx. +These tests file-load the target modules with narrow stubs so separation logic +can be unit-tested without importing the full AReaL runtime. +""" + +from __future__ import annotations + +import importlib.util +import logging as stdlib_logging +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +_ROOT = Path(__file__).resolve().parent.parent +_DC_PATH = _ROOT / "areal/v2/weight_update/awex/delta_config.py" + + +def _load_delta_config(): + spec = importlib.util.spec_from_file_location("awex_delta_config", _DC_PATH) + mod = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(mod) + return mod + + +@pytest.fixture +def dc(): + return _load_delta_config() + + +def _encode_like_adapter(tracker, params, version): + """Mirror the sender-side DeltaTracker contract used by adapters.""" + params_list = list(params.items()) + reason = tracker.full_sync_reason(version) + if reason is not None: + tracker.seed(params_list, version) + return list(params.keys()), list(params.values()) + encoded = tracker.encode(params_list, version) + return encoded.names, encoded.tensors + + +def _over_the_wire(names, tensors): + """Stand in for cuda_ipc serialize/deserialize: clone to detach storage.""" + return dict(zip(names, [t.clone() for t in tensors])) + + +def _weights(seed: int): + generator = torch.Generator().manual_seed(seed) + return { + "embed.weight": torch.randn(32, 16, generator=generator), + "layer.weight": torch.randn(16, 16, generator=generator), + "layer.bias": torch.randn(16, generator=generator), + } + + +def test_delta_config_env_gates(dc, monkeypatch): + """DTE env vars must honor DTE_* overrides over legacy AWEX_* names.""" + monkeypatch.delenv("DTE_DELTA_TRANSFER", raising=False) + monkeypatch.delenv("AWEX_DELTA_TRANSFER", raising=False) + assert dc.delta_transfer_enabled() is False + monkeypatch.setenv("AWEX_DELTA_TRANSFER", "1") + assert dc.delta_transfer_enabled() is True + monkeypatch.setenv("DTE_DELTA_TRANSFER", "0") + assert dc.delta_transfer_enabled() is False + monkeypatch.setenv("DTE_DELTA_TRANSFER", "1") + assert dc.delta_transfer_enabled() is True + + monkeypatch.setenv("AWEX_DELTA_ANCHOR_INTERVAL", "5") + assert dc.delta_anchor_interval() == 5 + monkeypatch.setenv("DTE_DELTA_ANCHOR_INTERVAL", "7") + assert dc.delta_anchor_interval() == 7 + + +def test_cuda_mem_stats_mb_without_cuda_returns_sentinel(dc): + """CPU-only hosts must not fail while reporting DTE memory telemetry.""" + alloc_mb, peak_mb = dc.cuda_mem_stats_mb() + if torch.cuda.is_available(): + assert alloc_mb >= 0.0 + assert peak_mb >= 0.0 + else: + assert (alloc_mb, peak_mb) == (-1.0, -1.0) + assert dc.cuda_mem_stats_mb(reset_peak=False) == (-1.0, -1.0) + + +def test_factory_builds_dte_tracker(dc, monkeypatch): + """The lazy factory should build a DTE tracker when DTE is available.""" + pytest.importorskip("dte") + monkeypatch.setenv("AWEX_DELTA_ANCHOR_INTERVAL", "0") + tracker = dc.make_delta_tracker() + assert hasattr(tracker, "encode") and hasattr(tracker, "seed") + + +def test_invert_adamw_roundtrip(): + """DTE AdamW inversion recovers pre-step weights on CPU.""" + pytest.importorskip("dte") + from dte.core import invert_adamw + + torch.manual_seed(0) + theta_prev = torch.randn(512, dtype=torch.float32) + param = theta_prev.clone().requires_grad_(True) + lr, wd, b1, b2, eps = 1e-3, 0.01, 0.9, 0.999, 1e-8 + opt = torch.optim.AdamW([param], lr=lr, betas=(b1, b2), eps=eps, weight_decay=wd) + param.grad = torch.randn_like(param) + opt.step() + state = opt.state[param] + recovered = invert_adamw( + param.detach().clone(), + state["exp_avg"], + state["exp_avg_sq"], + float(state["step"]), + lr, + wd, + b1, + b2, + eps, + ) + torch.testing.assert_close(recovered, theta_prev, rtol=1e-3, atol=1e-4) + + +def _stub_colocate_device(monkeypatch): + colocate_device_mod = types.ModuleType( + "areal.v2.weight_update.awex.colocate_device" + ) + colocate_device_mod.device_mapping_key = lambda ip, device: f"{ip}_{device}" + colocate_device_mod.get_colocate_ip_address = lambda: "127.0.0.1" + colocate_device_mod.get_physical_cuda_device_id = lambda local_index=None: str( + local_index or 0 + ) + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.awex.colocate_device", + colocate_device_mod, + ) + + +def _stub_areal_packages(monkeypatch): + monkeypatch.setitem(sys.modules, "httpx", types.ModuleType("httpx")) + monkeypatch.setitem(sys.modules, "areal", types.ModuleType("areal")) + monkeypatch.setitem(sys.modules, "areal.v2", types.ModuleType("areal.v2")) + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update", + types.ModuleType("areal.v2.weight_update"), + ) + + awex_mod = types.ModuleType("areal.v2.weight_update.awex") + awex_mod.awex_wu_use_group = lambda: False + awex_mod.fetch_kv_metadata = lambda *args, **kwargs: ([], []) + awex_mod.load_kv_metadata_file = lambda *args, **kwargs: None + awex_mod.resolve_physical_gpu_id = lambda *args, **kwargs: 0 + awex_mod.__path__ = [] + monkeypatch.setitem(sys.modules, "areal.v2.weight_update.awex", awex_mod) + _stub_colocate_device(monkeypatch) + + logging_mod = types.ModuleType("areal.utils.logging") + logging_mod.getLogger = stdlib_logging.getLogger + utils_mod = types.ModuleType("areal.utils") + utils_mod.logging = logging_mod + monkeypatch.setitem(sys.modules, "areal.utils", utils_mod) + monkeypatch.setitem(sys.modules, "areal.utils.logging", logging_mod) + + infra_mod = types.ModuleType("areal.infra") + platforms_mod = types.ModuleType("areal.infra.platforms") + platforms_mod.current_platform = SimpleNamespace(synchronize=lambda: None) + monkeypatch.setitem(sys.modules, "areal.infra", infra_mod) + monkeypatch.setitem(sys.modules, "areal.infra.platforms", platforms_mod) + + weight_digest_mod = types.ModuleType("areal.v2.weight_update.awex.weight_digest") + weight_digest_mod.log_tensor_digest = lambda *args, **kwargs: None + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.awex.weight_digest", + weight_digest_mod, + ) + + +def _load_delta_detect(monkeypatch): + """Load delta_detect.py without the full AReaL runtime.""" + fake = types.ModuleType("areal.utils.logging") + fake.getLogger = stdlib_logging.getLogger + monkeypatch.setitem(sys.modules, "areal", types.ModuleType("areal")) + monkeypatch.setitem(sys.modules, "areal.utils", types.ModuleType("areal.utils")) + monkeypatch.setitem(sys.modules, "areal.utils.logging", fake) + for package in ( + "areal.v2", + "areal.v2.weight_update", + "areal.v2.weight_update.awex", + ): + monkeypatch.setitem(sys.modules, package, types.ModuleType(package)) + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.awex.delta_config", + _load_delta_config(), + ) + path = _ROOT / "areal/v2/weight_update/awex/delta_detect.py" + spec = importlib.util.spec_from_file_location("awex_delta_detect", path) + mod = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(mod) + return mod + + +def _bind_inversion_param_names(inv, *named_params): + """Bind fake mcore parameter names for CPU-only inversion tests.""" + id2key = {id(param): name for name, param in named_params} + inv._module_param_key_maps = lambda: (id2key, {}) + + +def _load_sglang_adapter(monkeypatch): + """Load sglang_adapter.py with only its AReaL imports stubbed.""" + pytest.importorskip("awex.meta.weight_meta") + pytest.importorskip("awex.sharding.sglang_sharding") + + _stub_areal_packages(monkeypatch) + + delta_config_mod = types.ModuleType("areal.v2.weight_update.awex.delta_config") + delta_config_mod.separation_delta_transfer_enabled = lambda: False + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.awex.delta_config", + delta_config_mod, + ) + + inference_adapter_mod = types.ModuleType("areal.v2.weight_update.inference_adapter") + + class _AwexInferenceAdapter: + pass + + inference_adapter_mod.AwexInferenceAdapter = _AwexInferenceAdapter + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.inference_adapter", + inference_adapter_mod, + ) + + nccl_group_mod = types.ModuleType("areal.v2.weight_update.nccl_group") + nccl_group_mod.init_weights_update_group = lambda *args, **kwargs: None + nccl_group_mod.setup_batch_isend_irecv = lambda *args, **kwargs: None + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.nccl_group", + nccl_group_mod, + ) + + path = _ROOT / "areal/v2/weight_update/awex/sglang_adapter.py" + spec = importlib.util.spec_from_file_location("awex_sglang_adapter_test", path) + mod = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(mod) + return mod + + +def _load_megatron_adapter(monkeypatch): + """Load megatron_adapter.py with runtime-heavy AReaL imports stubbed.""" + pytest.importorskip("awex.meta.weight_meta") + pytest.importorskip("awex.sharding.param_sharding") + pytest.importorskip("awex.transfer.transfer_plan") + pytest.importorskip("awex.util.tensor_util") + + _stub_areal_packages(monkeypatch) + + delta_config_mod = types.ModuleType("areal.v2.weight_update.awex.delta_config") + delta_config_mod.separation_delta_transfer_enabled = lambda: True + delta_config_mod.make_delta_tracker = lambda *args, **kwargs: None + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.awex.delta_config", + delta_config_mod, + ) + + delta_detect_mod = types.ModuleType("areal.v2.weight_update.awex.delta_detect") + delta_detect_mod.AdamWInversionDetector = lambda *args, **kwargs: None + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.awex.delta_detect", + delta_detect_mod, + ) + + nccl_group_mod = types.ModuleType("areal.v2.weight_update.nccl_group") + nccl_group_mod.init_weights_update_group = lambda *args, **kwargs: None + nccl_group_mod.setup_batch_isend_irecv = lambda *args, **kwargs: None + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.nccl_group", + nccl_group_mod, + ) + + training_adapter_mod = types.ModuleType("areal.v2.weight_update.training_adapter") + + class _AwexTrainingAdapter: + pass + + training_adapter_mod.AwexTrainingAdapter = _AwexTrainingAdapter + monkeypatch.setitem( + sys.modules, + "areal.v2.weight_update.training_adapter", + training_adapter_mod, + ) + + path = _ROOT / "areal/v2/weight_update/awex/megatron_adapter.py" + spec = importlib.util.spec_from_file_location("awex_megatron_adapter_test", path) + mod = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(mod) + return mod diff --git a/tests/test_awex_separation_delta.py b/tests/test_awex_separation_delta.py new file mode 100644 index 0000000000..a2b01478a9 --- /dev/null +++ b/tests/test_awex_separation_delta.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Separation-specific distributed ordering tests for AWEX delta transfer.""" + +import importlib.util +import sys +import types +from types import SimpleNamespace + +import pytest + +from tests import test_awex_delta_common as common + +torch = common.torch + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("dte") is None or importlib.util.find_spec("awex") is None, + reason="DTE/AWEX source path is not available", +) + + +class _FakeTransferPlanBuilder: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def build_local_transfer_plan(self, *args, **kwargs): + del args, kwargs + return object() + + +def _capture_weight_update_groups(monkeypatch, mod): + calls = [] + + def _init_group(**kwargs): + calls.append(kwargs) + return kwargs.get("backend", "nccl") + + monkeypatch.setattr(mod, "init_weights_update_group", _init_group) + monkeypatch.setattr(mod, "fetch_kv_metadata", lambda *args, **kwargs: ([], [])) + monkeypatch.setattr(mod, "TransferPlanBuilder", _FakeTransferPlanBuilder) + return calls + + +def test_megatron_full_transfer_initializes_gloo_control_group(monkeypatch): + """Separated Full sync must not fall back to a NCCL control barrier.""" + mod = common._load_megatron_adapter(monkeypatch) + monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) + calls = _capture_weight_update_groups(monkeypatch, mod) + adapter = object.__new__(mod.AwexMegatronAdapter) + + adapter.init_weight_update_group( + pair_name="pair", + master_addr="127.0.0.1", + master_port=23456, + transfer_rank=8, + world_size=16, + kv_store_url="http://127.0.0.1:9999", + infer_world_size=8, + train_world_size=8, + num_engines=2, + ) + + assert [call.get("backend", "nccl") for call in calls] == ["nccl", "gloo"] + assert adapter._weights_update_group == "nccl" + assert adapter._weights_update_group_gloo == "gloo" + + +def test_sglang_full_transfer_initializes_gloo_control_group(monkeypatch): + """SGLang Full sync creates the same CPU control group as training.""" + mod = common._load_sglang_adapter(monkeypatch) + monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) + calls = _capture_weight_update_groups(monkeypatch, mod) + adapter = object.__new__(mod.AwexSGLangAdapter) + adapter._get_model_context = lambda: { + "tp_size": 4, + "tp_rank": 0, + "pp_size": 1, + "pp_rank": 0, + } + + adapter.init_weight_update_group( + pair_name="pair", + master_addr="127.0.0.1", + master_port=23456, + transfer_rank=0, + world_size=16, + kv_store_url="http://127.0.0.1:9999", + infer_world_size=8, + train_world_size=8, + num_engines=2, + ) + + assert [call.get("backend", "nccl") for call in calls] == ["nccl", "gloo"] + assert adapter._weights_update_group == "nccl" + assert adapter._weights_update_group_gloo == "gloo" + + +def test_megatron_full_transfer_uses_gloo_completion_barrier(monkeypatch): + """Full payload P2P completes through the sideband control group.""" + mod = common._load_megatron_adapter(monkeypatch) + monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) + monkeypatch.setattr( + mod, "nccl_build_send_ops", lambda *args, **kwargs: ([], [], []) + ) + monkeypatch.setattr(mod, "batch_send_recv", lambda **kwargs: None) + barriers = [] + monkeypatch.setattr(mod.dist, "barrier", lambda *, group: barriers.append(group)) + adapter = object.__new__(mod.AwexMegatronAdapter) + adapter._transfer_plan = object() + adapter._weights_update_group = "nccl" + adapter._weights_update_group_gloo = "gloo" + adapter._transfer_rank = 8 + adapter._world_size = 16 + adapter.get_local_shard_parameters = lambda: {} + + adapter.execute_weight_update(version=1) + + assert barriers == ["gloo"] + + +def test_sglang_full_transfer_uses_gloo_completion_barrier(monkeypatch): + """Receiver Full payload completion avoids the NCCL data group.""" + mod = common._load_sglang_adapter(monkeypatch) + monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) + monkeypatch.setattr( + mod, "nccl_build_recv_ops", lambda *args, **kwargs: ([], [], []) + ) + monkeypatch.setattr(mod, "batch_send_recv", lambda **kwargs: None) + barriers = [] + monkeypatch.setattr(mod.dist, "barrier", lambda *, group: barriers.append(group)) + adapter = object.__new__(mod.AwexSGLangAdapter) + adapter._transfer_plan = object() + adapter._weights_update_group = "nccl" + adapter._weights_update_group_gloo = "gloo" + adapter._transfer_rank = 0 + adapter.get_local_shard_parameters = lambda: {} + + adapter.execute_weight_update(version=1) + + assert barriers == ["gloo"] + + +def test_megatron_separation_delta_commits_tracker_after_receiver_barrier( + monkeypatch, +): + """A successful direct sparse transfer advances the periodic anchor.""" + mod = common._load_megatron_adapter(monkeypatch) + events = [] + + class _Tracker: + def mark_delta_committed(self, version): + events.append(("commit", version)) + + monkeypatch.setattr(mod.dist, "all_reduce", lambda *args, **kwargs: None) + monkeypatch.setattr( + mod.dist, + "barrier", + lambda *, group: events.append(("barrier", group)), + ) + adapter = object.__new__(mod.AwexMegatronAdapter) + adapter._transfer_plan = object() + adapter._weights_update_group = "nccl" + adapter._weights_update_group_gloo = "gloo" + adapter._transfer_rank = 8 + adapter._world_size = 16 + adapter._delta_tracker = _Tracker() + adapter.get_local_shard_parameters = lambda: {"w": torch.ones(1)} + adapter._ensure_delta_components = lambda: None + adapter._delta_detector = SimpleNamespace( + capture_synced_state=lambda params: {}, + mark_synced=lambda version, state: events.append(("synced", version)), + ) + adapter._delta_prepare_masks = lambda params, version: ( + {"w": torch.ones(1, dtype=torch.bool)}, + True, + ) + adapter._execute_separation_delta_send = ( + lambda params, masks, version: events.append(("send", version)) + ) + + adapter._execute_separation_weight_update(version=2) + + assert events == [ + ("send", 2), + ("barrier", "gloo"), + ("commit", 2), + ("synced", 2), + ] + + +def test_megatron_separation_failed_delta_does_not_advance_tracker(monkeypatch): + """A failed sparse transfer must leave the anchor counter unchanged.""" + mod = common._load_megatron_adapter(monkeypatch) + commits = [] + + class _Tracker: + def mark_delta_committed(self, version): + commits.append(version) + + monkeypatch.setattr(mod.dist, "all_reduce", lambda *args, **kwargs: None) + monkeypatch.setattr( + mod.dist, + "barrier", + lambda **kwargs: pytest.fail("failed transfer must not reach barrier"), + ) + adapter = object.__new__(mod.AwexMegatronAdapter) + adapter._transfer_plan = object() + adapter._weights_update_group = "nccl" + adapter._weights_update_group_gloo = "gloo" + adapter._transfer_rank = 8 + adapter._world_size = 16 + adapter._delta_tracker = _Tracker() + adapter.get_local_shard_parameters = lambda: {"w": torch.ones(1)} + adapter._ensure_delta_components = lambda: None + adapter._delta_detector = SimpleNamespace( + capture_synced_state=lambda params: {}, mark_synced=lambda *args: None + ) + adapter._delta_prepare_masks = lambda params, version: ( + {"w": torch.ones(1, dtype=torch.bool)}, + True, + ) + adapter._execute_separation_delta_send = lambda *args: (_ for _ in ()).throw( + RuntimeError("transfer failed") + ) + + with pytest.raises(RuntimeError, match="transfer failed"): + adapter._execute_separation_weight_update(version=2) + + assert commits == [] + + +def test_reconstructed_override_preserves_tensor_parallel_metadata(monkeypatch): + """A plain theta_old tensor must gather like its live TP parameter.""" + mod = common._load_megatron_adapter(monkeypatch) + param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + param.tensor_model_parallel = True + param.partition_dim = 0 + param.partition_stride = 1 + override = param.detach().clone() + overrides = {id(param): override} + captured = {} + + megatron_mod = types.ModuleType("areal.engine.megatron_utils.megatron") + megatron_mod.get_named_parameters = lambda model, experts: [("qkv", param)] + + def _all_gather_param(name, tensor, **kwargs): + del name, kwargs + captured["uses_override"] = tensor is override + captured["metadata"] = ( + tensor.tensor_model_parallel, + tensor.partition_dim, + tensor.partition_stride, + ) + return tensor + + megatron_mod.all_gather_param = _all_gather_param + megatron_mod.convert_to_hf = lambda config, model, name, tensor: [("w", tensor)] + monkeypatch.setitem( + sys.modules, "areal.engine.megatron_utils.megatron", megatron_mod + ) + + adapter = object.__new__(mod.AwexMegatronAdapter) + adapter._engine = SimpleNamespace( + model=object(), + tf_config=SimpleNamespace(num_moe_experts=None), + hf_config=SimpleNamespace(model_type="qwen3", tie_word_embeddings=False), + _duplicated_param_names=set(), + ) + + items = list(adapter._iter_hf_params(overrides, consume_overrides=True)) + + assert len(items) == 1 + assert items[0][0] == "w" + torch.testing.assert_close(items[0][1], override, rtol=0, atol=0) + assert captured["uses_override"] is True + assert captured["metadata"] == (True, 0, 1) + assert overrides == {} + + +def test_streaming_generator_waits_async_batch_before_yield(monkeypatch): + """HF conversion never starts with inversion all-reduces in flight.""" + monkeypatch.setenv("DTE_INVERSION_ALLREDUCE_WINDOW_MB", "0.000001") + dd = common._load_delta_detect(monkeypatch) + param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + base_opt = torch.optim.AdamW([param], lr=1e-3) + dp_group = object() + in_flight = 0 + + class _FakeWork: + def wait(self): + nonlocal in_flight + in_flight -= 1 + + def _fake_all_reduce(tensor, op=None, group=None, async_op=False): + nonlocal in_flight + del op + assert group is dp_group + if tensor.dtype == torch.int64: + tensor[:] = torch.tensor([1, 0, param.numel()]) + if async_op: + in_flight += 1 + return _FakeWork() + return None + + monkeypatch.setattr(dd.torch.distributed, "all_reduce", _fake_all_reduce) + + class _FakeDistOpt: + optimizer = base_opt + shard_fp32_from_float16_groups = [[param]] + model_float16_groups = [[param]] + model_param_group_index_map = None + data_parallel_group = dp_group + + def _get_model_param_range_map(self, model_param): + assert model_param is param + return {"param": SimpleNamespace(start=0, end=param.numel())} + + inv = dd.AdamWInversionDetector(SimpleNamespace(_offloaded_optimizer_states={})) + common._bind_inversion_param_names(inv, ("w", param)) + producer = inv._iter_reconstruct_pre_step_mcore([_FakeDistOpt()]) + + param_id, old = next(producer) + + assert param_id == id(param) + assert old is not dd._NO_RECONSTRUCTION + assert in_flight == 0 + + +def test_missing_watermark_finishes_all_collectives_before_abort(monkeypatch): + """A dense verdict cannot reduce the per-parameter collective count.""" + monkeypatch.setenv("DTE_INVERSION_ALLREDUCE_WINDOW_MB", "0.000001") + dd = common._load_delta_detect(monkeypatch) + params = [ + torch.nn.Parameter(torch.tensor([1.0, 2.0])), + torch.nn.Parameter(torch.tensor([3.0, 4.0])), + ] + base_opt = torch.optim.AdamW(params, lr=1e-3) + for param in params: + param.grad = torch.ones_like(param) + base_opt.step() + dp_group = object() + calls = 0 + + class _FakeWork: + def wait(self): + return None + + def _fake_all_reduce(tensor, op=None, group=None, async_op=False): + nonlocal calls + del tensor, op + assert group is dp_group + calls += 1 + return _FakeWork() if async_op else None + + monkeypatch.setattr(dd.torch.distributed, "all_reduce", _fake_all_reduce) + + class _FakeDistOpt: + optimizer = base_opt + shard_fp32_from_float16_groups = [params] + model_float16_groups = [params] + model_param_group_index_map = None + data_parallel_group = dp_group + + def _get_model_param_range_map(self, model_param): + return {"param": SimpleNamespace(start=0, end=model_param.numel())} + + inv = dd.AdamWInversionDetector(SimpleNamespace(_offloaded_optimizer_states={})) + common._bind_inversion_param_names(inv, ("w0", params[0]), ("w1", params[1])) + lazy = dd._LazyReconstructDict( + inv._iter_reconstruct_pre_step_mcore([_FakeDistOpt()]) + ) + + lazy.finish() + + assert lazy.aborted is True + assert calls == 2 * len(params) From e90eca68c37fa58029149f54b556838ab9b2a1c4 Mon Sep 17 00:00:00 2001 From: pyq623 <12421170@zju.edu.cn> Date: Fri, 14 Aug 2026 14:30:03 +0800 Subject: [PATCH 3/5] docs(examples): add DTE separation GSM8K example --- examples/dte/README.md | 45 ++++++++++++++++++++ examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml | 49 ++++++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 examples/dte/README.md create mode 100644 examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml diff --git a/examples/dte/README.md b/examples/dte/README.md new file mode 100644 index 0000000000..1b665520f3 --- /dev/null +++ b/examples/dte/README.md @@ -0,0 +1,45 @@ +# DTE separation example + +This directory contains an opt-in example for transferring AdamW weight updates from a +Megatron actor to a separately scheduled SGLang rollout worker with +[AReaL-DTE](https://github.com/areal-project/AReaL-DTE). + +Install AReaL-DTE in the controller and worker environments before enabling the example. +For local development, use an editable checkout: + +```bash +pip install -e /path/to/AReaL-DTE +``` + +The Qwen3-30B-A3B GSM8K example uses two 8-GPU allocations: one for the actor and one +for rollout. Launch it with the scheduler used by your cluster, for example: + +```bash +python3 examples/math/gsm8k_rl.py \ + --config examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml \ + scheduler.type=slurm \ + actor.scheduling_spec.0.image=/path/to/areal.sif +``` + +The first weight update is a full synchronization. Later contiguous versions use sparse +AdamW deltas, with a periodic full-weight anchor after every 20 successfully committed +deltas. Set `actor.dte.enabled=false` to return to the existing AWEX full-weight +behavior. + +## Effective performance defaults + +The example keeps its public YAML surface small. The following optimizations are already +active without adding site-specific environment overrides: + +| Optimization | Effective setting | Source | +| ------------------------------- | --------------------------------------- | ------------------------------------------------- | +| Streaming AdamW reconstruction | `DTE_STREAMING_RECONSTRUCT=1` | Propagated by AReaL when `actor.dte.enabled=true` | +| Coalesced two-round sparse P2P | `DTE_DELTA_P2P_COALESCE=1` | AReaL-DTE default | +| Pipelined inversion collectives | 512 MiB in-flight window | AReaL default | +| Inversion compute device | Payload device (GPU for this example) | AReaL default | +| Compact change indices | `int32` when the parameter size permits | Standard AReaL detector path | +| Batched operation remapping | Once per parameter | Standard AReaL-DTE payload path | + +Snapshot verification, weight digests, phase timing, recovery, and deterministic rollout +diagnostics are intentionally not enabled here because they are validation or experiment +controls rather than requirements of the separation delta path. diff --git a/examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml b/examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml new file mode 100644 index 0000000000..3a062cfcc7 --- /dev/null +++ b/examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml @@ -0,0 +1,49 @@ +defaults: + - ../math/gsm8k_grpo_megatron@_global_ + - _self_ + +experiment_name: gsm8k-grpo-qwen3-30b-a3b-dte +trial_name: trial0 + +# This example uses separate 8-GPU allocations for training and inference. +cluster: + n_nodes: 2 + n_gpus_per_node: 8 + +total_train_steps: 20 + +rollout: + _version: v2 + backend: "sglang:d2t4p1" + scheduling_strategy: + type: separation + +actor: + _version: v2 + backend: "megatron:(attn:d2p1t4|ffn:d1p1e8)" + path: Qwen/Qwen3-30B-A3B + weight_update_mode: awex + dte: + enabled: true + transfer: delta + delta_method: adamw + anchor_interval: 20 + # Effective performance defaults do not need duplicate YAML switches: + # - AReaL propagates DTE_STREAMING_RECONSTRUCT=1 to actor/rollout workers. + # - AReaL-DTE defaults DTE_DELTA_P2P_COALESCE to 1. + # - compact integer indices and batched per-parameter remapping are the + # standard payload-construction path. + +ref: + path: ${actor.path} + +sglang: + model_path: ${actor.path} + +train_dataset: + batch_size: 8 + path: openai/gsm8k + +valid_dataset: + batch_size: 8 + path: openai/gsm8k From 5669c633e469fb99eda7132f39c39badb2b4674f Mon Sep 17 00:00:00 2001 From: pyq623 <12421170@zju.edu.cn> Date: Fri, 14 Aug 2026 14:59:03 +0800 Subject: [PATCH 4/5] docs(examples): expand DTE GSM8K configuration --- examples/dte/README.md | 22 +-- examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml | 176 ++++++++++++++++++++- 2 files changed, 180 insertions(+), 18 deletions(-) diff --git a/examples/dte/README.md b/examples/dte/README.md index 1b665520f3..6fdec9d0f7 100644 --- a/examples/dte/README.md +++ b/examples/dte/README.md @@ -28,17 +28,17 @@ behavior. ## Effective performance defaults -The example keeps its public YAML surface small. The following optimizations are already -active without adding site-specific environment overrides: - -| Optimization | Effective setting | Source | -| ------------------------------- | --------------------------------------- | ------------------------------------------------- | -| Streaming AdamW reconstruction | `DTE_STREAMING_RECONSTRUCT=1` | Propagated by AReaL when `actor.dte.enabled=true` | -| Coalesced two-round sparse P2P | `DTE_DELTA_P2P_COALESCE=1` | AReaL-DTE default | -| Pipelined inversion collectives | 512 MiB in-flight window | AReaL default | -| Inversion compute device | Payload device (GPU for this example) | AReaL default | -| Compact change indices | `int32` when the parameter size permits | Standard AReaL detector path | -| Batched operation remapping | Once per parameter | Standard AReaL-DTE payload path | +The example spells out the portable performance settings in the worker environment; none +of them depends on a site-specific path or cluster: + +| Optimization | Effective setting | Source | +| ------------------------------- | --------------------------------------- | ---------------------------------------------- | +| Streaming AdamW reconstruction | `DTE_STREAMING_RECONSTRUCT=1` | Explicit; also enforced by `actor.dte.enabled` | +| Coalesced two-round sparse P2P | `DTE_DELTA_P2P_COALESCE=1` | Explicit; also the AReaL-DTE default | +| Pipelined inversion collectives | 512 MiB in-flight window | Explicit AReaL setting | +| Inversion compute device | Payload device (GPU for this example) | Explicit AReaL setting | +| Compact change indices | `int32` when the parameter size permits | Standard AReaL detector path | +| Batched operation remapping | Once per parameter | Standard AReaL-DTE payload path | Snapshot verification, weight digests, phase timing, recovery, and deterministic rollout diagnostics are intentionally not enabled here because they are validation or experiment diff --git a/examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml b/examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml index 3a062cfcc7..c5c947bd79 100644 --- a/examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml +++ b/examples/dte/gsm8k_grpo_qwen3_30b_a3b.yaml @@ -5,45 +5,207 @@ defaults: experiment_name: gsm8k-grpo-qwen3-30b-a3b-dte trial_name: trial0 -# This example uses separate 8-GPU allocations for training and inference. +seed: 42 +# Separation keeps the actor parameters and AdamW moments resident so DTE can +# reconstruct the pre-step weights. +enable_offload: false +total_train_epochs: 10 +total_train_steps: 20 +tokenizer_path: ${actor.path} + +# Training and rollout are scheduled on separate 8-GPU allocations. cluster: n_nodes: 2 n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve -total_train_steps: 20 +scheduler: + type: null + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 2048 + max_tokens: 3072 + greedy: false + temperature: 1.0 + top_p: 1.0 + top_k: 1000000 rollout: _version: v2 backend: "sglang:d2t4p1" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + setup_timeout: 43200 + request_timeout: 7200 + workers_ready_timeout: 43200 + pause_grace_period: 30 + max_concurrent_rollouts: 64 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 8 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} scheduling_strategy: type: separation + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: false + agent: + mode: inline + export_style: individual + turn_discount: 1.0 actor: _version: v2 backend: "megatron:(attn:d2p1t4|ffn:d1p1e8)" + experiment_name: ${experiment_name} + trial_name: ${trial_name} path: Qwen/Qwen3-30B-A3B + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + grad_reduce_dtype: float32 + mb_spec: + n_mbs: 1 + granularity: 1 + max_tokens_per_mb: 6144 + n_mbs_divisor: 1 + optimizer: + type: adam + lr: 3e-6 + weight_decay: 0.003 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0 + megatron: + wrap_with_ddp: true + ddp: + grad_reduce_in_fp32: true + overlap_grad_reduce: true + overlap_param_gather: false + use_distributed_optimizer: true + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + moe_router_dtype: fp32 + moe_token_dispatcher_type: alltoall + moe_router_bias_update_rate: 0.0 + eps_clip: 0.4 + eps_clip_higher: null + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: false + prox_logp_method: reuse_train_logp + use_decoupled_loss: true + rejection_sampling: + level: token + action: mask + 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} dte: enabled: true transfer: delta delta_method: adamw anchor_interval: 20 - # Effective performance defaults do not need duplicate YAML switches: - # - AReaL propagates DTE_STREAMING_RECONSTRUCT=1 to actor/rollout workers. - # - AReaL-DTE defaults DTE_DELTA_P2P_COALESCE to 1. - # - compact integer indices and batched per-parameter remapping are the - # standard payload-construction path. + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + cpu: 4 + mem: 64 + image: /path/to/areal.sif + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: + # These values match the effective defaults and make the optimized + # separation path visible to readers of this standalone example. + DTE_STREAMING_RECONSTRUCT: "1" + DTE_DELTA_P2P_COALESCE: "1" + DTE_INVERSION_ALLREDUCE_WINDOW_MB: "512" + DTE_INVERSION_COMPUTE_ON_CPU: "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: ${actor.mb_spec.max_tokens_per_mb} + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} 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.65 + cuda_graph_max_bs: 16 train_dataset: batch_size: 8 + shuffle: true + pin_memory: true + num_workers: 4 path: openai/gsm8k + type: rl + max_length: 1024 valid_dataset: batch_size: 8 + pin_memory: true + num_workers: 4 path: openai/gsm8k + type: rl + +saver: + mode: auto + freq_epochs: null + freq_steps: null + freq_secs: null + +recover: + mode: disabled + freq_epochs: null + freq_steps: null + freq_secs: null + +evaluator: + freq_epochs: null + freq_steps: null + freq_secs: null + +stats_logger: + wandb: + mode: disabled + +perf_tracer: + enabled: false From c04658784a128d85ffa5595287f69878273868cb Mon Sep 17 00:00:00 2001 From: pyq623 <12421170@zju.edu.cn> Date: Fri, 14 Aug 2026 16:48:13 +0800 Subject: [PATCH 5/5] fix(awex): address separation DTE review feedback --- areal/api/cli_args.py | 13 +- areal/engine/megatron_engine.py | 8 + areal/utils/environ.py | 52 ++++- areal/v2/weight_update/awex/delta_config.py | 132 +++++-------- areal/v2/weight_update/awex/delta_detect.py | 179 ++++++++++-------- .../v2/weight_update/awex/megatron_adapter.py | 10 +- areal/v2/weight_update/awex/sglang_adapter.py | 7 +- docs/en/cli_reference.md | 2 +- docs/zh/cli_reference.md | 2 +- examples/dte/README.md | 16 ++ tests/test_awex_delta_common.py | 168 ++++++++++++++-- tests/test_awex_separation_delta.py | 8 +- tests/test_dte_topology_gating.py | 14 ++ tests/test_environ.py | 101 ++++++++++ tests/test_megatron_optimizer_config.py | 33 ++++ 15 files changed, 542 insertions(+), 203 deletions(-) create mode 100644 tests/test_environ.py diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index c8786f0561..f3cff6ad3c 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -1678,7 +1678,11 @@ class PPOActorConfig(TrainEngineConfig): # 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"} @@ -1879,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( diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index a28a60ac4c..90fae868e7 100644 --- a/areal/engine/megatron_engine.py +++ b/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"] diff --git a/areal/utils/environ.py b/areal/utils/environ.py index 2a444a650d..6cfbf07058 100644 --- a/areal/utils/environ.py +++ b/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: diff --git a/areal/v2/weight_update/awex/delta_config.py b/areal/v2/weight_update/awex/delta_config.py index 2f20209ea5..bf61cada49 100644 --- a/areal/v2/weight_update/awex/delta_config.py +++ b/areal/v2/weight_update/awex/delta_config.py @@ -1,25 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -"""Runtime gates and factories for separation AdamW delta transfer. +"""Runtime configuration 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) +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 -import os +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 " @@ -28,76 +19,55 @@ "(local dev) or add DTE_SRC to PYTHONPATH." ) +_DTE_TRUTHY_VALUES = ("1", "true", "yes", "on") +_DTE_FALSY_VALUES = ("0", "false", "no", "off") -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") +@dataclass(frozen=True) +class DTERuntimeConfig: + """Worker-local runtime settings for separation delta transfer.""" -def separation_delta_transfer_enabled() -> bool: - """Whether sparse separated-card transfer is explicitly enabled.""" - return separation_weight_update_enabled() and delta_transfer_enabled() + delta_transfer: bool + separation_weight_update: bool + anchor_interval: int - -def delta_anchor_interval() -> int: - """Force a full sync every N deltas (0 = never; rely on seed + chain-break).""" - return int( - _env_value( + @classmethod + def from_env(cls) -> DTERuntimeConfig: + """Snapshot DTE runtime settings from the worker environment.""" + anchor_value = get_env_var( "DTE_DELTA_ANCHOR_INTERVAL", - "AWEX_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, ) - ) - - -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 + @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) diff --git a/areal/v2/weight_update/awex/delta_detect.py b/areal/v2/weight_update/awex/delta_detect.py index 5012b9242a..348700e4d0 100644 --- a/areal/v2/weight_update/awex/delta_detect.py +++ b/areal/v2/weight_update/awex/delta_detect.py @@ -61,42 +61,40 @@ from __future__ import annotations -import os import time from collections.abc import Iterator import torch from areal.utils import logging -from areal.v2.weight_update.awex.delta_config import cuda_mem_stats_mb +from areal.utils.environ import get_bool_env_var, get_env_var logger = logging.getLogger("AwexDeltaDetect") - -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 +_DTE_TRUTHY_VALUES = ("1", "true", "yes", "on") +_DTE_FALSY_VALUES = ("0", "false", "no", "off") -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 _dte_bool_env(name: str, legacy_name: str, default: bool = False) -> bool: + return get_bool_env_var( + name, + "true" if default else "false", + fallback_names=(legacy_name,), + truthy_values=_DTE_TRUTHY_VALUES, + falsy_values=_DTE_FALSY_VALUES, + strip_value=True, + ) def _inversion_debug_enabled() -> bool: - return _env_bool("DTE_DELTA_INVERSION_DEBUG", "AWEX_DELTA_INVERSION_DEBUG") + return _dte_bool_env("DTE_DELTA_INVERSION_DEBUG", "AWEX_DELTA_INVERSION_DEBUG") def _streaming_reconstruct_enabled() -> bool: """Whether to reconstruct pre-step tensors lazily during HF conversion.""" - return _env_bool("DTE_STREAMING_RECONSTRUCT", "PYQ_STREAMING_RECONSTRUCT", False) + return _dte_bool_env( + "DTE_STREAMING_RECONSTRUCT", "PYQ_STREAMING_RECONSTRUCT", False + ) def _inversion_allreduce_window_bytes() -> int: @@ -107,10 +105,10 @@ def _inversion_allreduce_window_bytes() -> int: flight. ``0`` restores the fully synchronous per-param path (A/B knob). """ mb = float( - _env_value( + get_env_var( "DTE_INVERSION_ALLREDUCE_WINDOW_MB", - "AWEX_INVERSION_ALLREDUCE_WINDOW_MB", "512", + fallback_names=("AWEX_INVERSION_ALLREDUCE_WINDOW_MB",), ) ) return max(int(mb * 1024 * 1024), 0) @@ -127,17 +125,17 @@ def _inversion_compute_device(flat_t: torch.Tensor) -> torch.device: operands there. Set DTE_INVERSION_COMPUTE_ON_CPU=1 to restore the legacy behaviour of computing wherever the fp32 main shard lives. """ - if _env_bool("DTE_INVERSION_COMPUTE_ON_CPU", "AWEX_INVERSION_COMPUTE_ON_CPU"): + if _dte_bool_env("DTE_INVERSION_COMPUTE_ON_CPU", "AWEX_INVERSION_COMPUTE_ON_CPU"): return torch.device("cpu") return flat_t.device def _inversion_bf16_margin_rel() -> float: return float( - _env_value( + get_env_var( "DTE_DELTA_INVERSION_BF16_MARGIN_REL", - "AWEX_DELTA_INVERSION_BF16_MARGIN_REL", "1e-4", + fallback_names=("AWEX_DELTA_INVERSION_BF16_MARGIN_REL",), ) ) @@ -150,15 +148,27 @@ def _inversion_dense_param_suffixes() -> tuple[str, ...]: payloads, so sending them dense is the conservative default. Set the env to an empty string to disable, or to a comma-separated suffix list to override. """ - if "DTE_DELTA_INVERSION_DENSE_PARAM_SUFFIXES" in os.environ: - value = os.environ["DTE_DELTA_INVERSION_DENSE_PARAM_SUFFIXES"] - elif "AWEX_DELTA_INVERSION_DENSE_PARAM_SUFFIXES" in os.environ: - value = os.environ["AWEX_DELTA_INVERSION_DENSE_PARAM_SUFFIXES"] - else: - value = ".mlp.gate.weight" + value = get_env_var( + "DTE_DELTA_INVERSION_DENSE_PARAM_SUFFIXES", + ".mlp.gate.weight", + fallback_names=("AWEX_DELTA_INVERSION_DENSE_PARAM_SUFFIXES",), + allow_empty=True, + ) + assert value is not None return tuple(item.strip() for item in value.split(",") if item.strip()) +def _cuda_mem_stats_mb(reset_peak: bool = True) -> tuple[float, float]: + """Return current and peak CUDA allocation, or CPU sentinels.""" + 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 + + def _inversion_force_dense_param(name: str) -> bool: return any(name.endswith(suffix) for suffix in _inversion_dense_param_suffixes()) @@ -181,8 +191,13 @@ def _dist_world_size(group=None) -> int: return 1 -def _adamw_hparams(param_group: dict) -> tuple[float, float, float, float, float]: +def _adamw_hparams( + param_group: dict, +) -> tuple[float, float, float, float, float] | None: """Extract (lr, weight_decay, beta1, beta2, eps) from a torch param_group.""" + last_step_lr = param_group.get("_areal_last_step_lr") + if last_step_lr is None: + return None betas = param_group.get("betas", (0.9, 0.999)) beta1 = float(betas[0]) beta2 = float(betas[1]) @@ -193,7 +208,7 @@ def _adamw_hparams(param_group: dict) -> tuple[float, float, float, float, float beta1 = 0.0 beta2 = 0.0 return ( - float(param_group.get("_areal_last_step_lr", param_group["lr"])), + float(last_step_lr), float(param_group.get("weight_decay", 0.0)), beta1, beta2, @@ -311,10 +326,10 @@ def _bf16_rounding_boundary_mask_chunked( ) -> torch.Tensor: """Bound boundary-check temporaries by processing large tensors in chunks.""" chunk_elems = int( - _env_value( + get_env_var( "DTE_BOUNDARY_CHUNK_ELEMS", - "PYQ_BOUNDARY_CHUNK_ELEMS", str(50_000_000), + fallback_names=("PYQ_BOUNDARY_CHUNK_ELEMS",), ) ) if chunk_elems <= 0 or cur_bf16.numel() <= chunk_elems: @@ -769,6 +784,7 @@ def _iter_reconstruct_pre_step_mcore( skipped_step_unchanged = 0 skipped_missing_watermark = 0 skipped_step_jump = 0 + skipped_missing_last_step_lr = 0 skipped_missing_fingerprint = 0 skipped_tracked_unchanged = 0 skipped_payload_changed_without_step = 0 @@ -1163,50 +1179,53 @@ def _drain_pending_batch() -> list[tuple[int, torch.Tensor | object]]: ): skipped_bad_state += 1 else: - lr, wd, b1, b2, eps = _adamw_hparams( - base_opt.param_groups[gi] - ) - # The fp32 main shard and moments may be - # offloaded to CPU under colocate; stage them on - # the inversion compute device (the visible - # payload's device by default) so invert_adamw - # runs on GPU instead of the CPU. Inversion - # still uses the exact post-step optimizer - # weight, not the bf16 model copy. Keep the - # reconstructed previous value in fp32 until - # mask creation so boundary-near BF16 roundoff - # is handled conservatively. - dev = _inversion_compute_device(flat_t) - theta_t_slice = main_flat.to( - device=dev, dtype=torch.float32 - ) - theta_old_slice = invert_adamw( - theta_t_slice, - exp_avg.to(device=dev, dtype=torch.float32).reshape( - -1 - ), - exp_avg_sq.to( + hparams = _adamw_hparams(base_opt.param_groups[gi]) + if hparams is None: + skipped_missing_last_step_lr += 1 + force_dense = True + else: + lr, wd, b1, b2, eps = hparams + # The fp32 main shard and moments may be + # offloaded to CPU under colocate; stage them on + # the inversion compute device (the visible + # payload's device by default) so invert_adamw + # runs on GPU instead of the CPU. Inversion + # still uses the exact post-step optimizer + # weight, not the bf16 model copy. Keep the + # reconstructed previous value in fp32 until + # mask creation so boundary-near BF16 roundoff + # is handled conservatively. + dev = _inversion_compute_device(flat_t) + theta_t_slice = main_flat.to( device=dev, dtype=torch.float32 - ).reshape(-1), - step, - lr, - wd, - b1, - b2, - eps, - ) - visible_slice = flat_t[rng.start : rng.end] - visible_fp32 = visible_slice.to( - device=theta_old_slice.device, - dtype=torch.float32, - ) - correction[rng.start : rng.end] = ( - theta_old_slice - visible_fp32 - ).to( - device=correction.device, - dtype=correction.dtype, - ) - has_local_update = True + ) + theta_old_slice = invert_adamw( + theta_t_slice, + exp_avg.to( + device=dev, dtype=torch.float32 + ).reshape(-1), + exp_avg_sq.to( + device=dev, dtype=torch.float32 + ).reshape(-1), + step, + lr, + wd, + b1, + b2, + eps, + ) + visible_slice = flat_t[rng.start : rng.end] + visible_fp32 = visible_slice.to( + device=theta_old_slice.device, + dtype=torch.float32, + ) + correction[rng.start : rng.end] = ( + theta_old_slice - visible_fp32 + ).to( + device=correction.device, + dtype=correction.dtype, + ) + has_local_update = True owned_numel = ( int(rng.end - rng.start) if has_local_update and rng is not None else 0 @@ -1283,13 +1302,15 @@ def _drain_pending_batch() -> list[tuple[int, torch.Tensor | object]]: if force_dense: logger.warning( "Inversion: optimizer step replay is ambiguous " - "(missing_watermark=%d step_jump=%d missing_fingerprint=%d " + "(missing_watermark=%d step_jump=%d missing_last_step_lr=%d " + "missing_fingerprint=%d " "payload_changed_without_step=%d " "untracked_non_optimizer=%d changed_non_optimizer=%d " "partial_state=%d unkeyed_opt_params=%d) " "-> dense.", skipped_missing_watermark, skipped_step_jump, + skipped_missing_last_step_lr, skipped_missing_fingerprint, skipped_payload_changed_without_step, skipped_untracked_non_optimizer, @@ -1442,7 +1463,7 @@ def compute_masks(self, names, tensors, version): return cached_masks perf_start = time.monotonic() # Reset the peak counter so peak_mb below covers inversion only. - cuda_mem_stats_mb() + _cuda_mem_stats_mb() adapter = self._adapter t_inner = time.monotonic() inner = adapter._get_inner_optimizers() @@ -1624,7 +1645,7 @@ def compute_masks(self, names, tensors, version): forced_dense_params, forced_dense_elements, ) - alloc_mb, peak_mb = cuda_mem_stats_mb() + alloc_mb, peak_mb = _cuda_mem_stats_mb() logger.info( "[dte-perf][inversion] v%d result=sparse params=%d " "theta_old_params=%d full_changed_masks=%d " diff --git a/areal/v2/weight_update/awex/megatron_adapter.py b/areal/v2/weight_update/awex/megatron_adapter.py index ed3cc2f988..2f24dfbe88 100644 --- a/areal/v2/weight_update/awex/megatron_adapter.py +++ b/areal/v2/weight_update/awex/megatron_adapter.py @@ -30,10 +30,7 @@ awex_wu_use_group, fetch_kv_metadata, ) -from areal.v2.weight_update.awex.delta_config import ( - make_delta_tracker, - separation_delta_transfer_enabled, -) +from areal.v2.weight_update.awex.delta_config import DTERuntimeConfig from areal.v2.weight_update.awex.delta_detect import AdamWInversionDetector from areal.v2.weight_update.nccl_group import ( init_weights_update_group, @@ -78,6 +75,7 @@ def __init__(self, engine: MegatronEngine): self._colocate_admin_api_key: str = "areal-admin-key" self._colocate_http_client: httpx.Client | None = None self._colocate_timeout_s: float = 120.0 + self._dte_config = DTERuntimeConfig.from_env() self._delta_tracker = None self._delta_detector = None @@ -216,7 +214,7 @@ def init_weight_update_group( ) def execute_weight_update(self, version: int) -> None: - if separation_delta_transfer_enabled(): + if self._dte_config.enabled: self._release_grad_buffers_for_separation_sync() try: self._execute_separation_weight_update(version) @@ -600,7 +598,7 @@ def _iter_hf_with_overrides(self, theta_by_id: dict[int, torch.Tensor]): def _ensure_delta_components(self) -> None: if self._delta_tracker is None: - self._delta_tracker = make_delta_tracker() + self._delta_tracker = self._dte_config.create_delta_tracker() if self._delta_detector is None: self._delta_detector = AdamWInversionDetector(self) diff --git a/areal/v2/weight_update/awex/sglang_adapter.py b/areal/v2/weight_update/awex/sglang_adapter.py index 6d4155593c..208180bc82 100644 --- a/areal/v2/weight_update/awex/sglang_adapter.py +++ b/areal/v2/weight_update/awex/sglang_adapter.py @@ -36,9 +36,7 @@ awex_wu_use_group, fetch_kv_metadata, ) -from areal.v2.weight_update.awex.delta_config import ( - separation_delta_transfer_enabled, -) +from areal.v2.weight_update.awex.delta_config import DTERuntimeConfig from areal.v2.weight_update.inference_adapter import ( AwexInferenceAdapter, ) @@ -70,6 +68,7 @@ def __init__(self, scheduler: Any): self._colocate_transport = None self._train_to_infer_device_mapping: dict | None = None self._infer_to_train_device_mapping: dict | None = None + self._dte_config = DTERuntimeConfig.from_env() def _get_model(self) -> torch.nn.Module: return self._scheduler.tp_worker.model_runner.model @@ -426,7 +425,7 @@ def init_weight_update_group( ) def execute_weight_update(self, version: int) -> None: - if separation_delta_transfer_enabled(): + if self._dte_config.enabled: self._execute_separation_weight_update(version) return diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index f134f2608a..8b4bbd36af 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -389,7 +389,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `workers_ready_timeout` | float | `30.0` | Timeout (seconds) for initialize() to wait for guards to be ready. | | `scheduling_strategy` | [`SchedulingStrategy`](section-scheduling-strategy) | **Required** | The scheduling strategy of this TrainEngine, either separation or colocation. Currently only used by the TrainController. | | `dte` | [`DTEConfig`](section-dte) | **Required** | - | -| `ppo_n_minibatches` | integer | `4` | Number of minibatches for each PPO update | +| `ppo_n_minibatches` | integer | `4` | Number of minibatches for each PPO update. Separation DTE AdamW delta transfer currently requires 1. | | `eps_clip` | float | `0.2` | Clipping factor for policy ratio | | `eps_clip_higher` | float \| None | `None` | Clipping factor (higher value) for policy ratio. Default is None. When eps_clip_higher is set (decoupled), eps_clip will be used as the lower value. | | `c_clip` | float \| None | `None` | Dual clipping factor for policy ratio, must be > 1.0. None disables dual clipping. | diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index b9860386ad..47ab88753e 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -387,7 +387,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `workers_ready_timeout` | float | `30.0` | Timeout (seconds) for initialize() to wait for guards to be ready. | | `scheduling_strategy` | [`SchedulingStrategy`](section-scheduling-strategy) | **Required** | The scheduling strategy of this TrainEngine, either separation or colocation. Currently only used by the TrainController. | | `dte` | [`DTEConfig`](section-dte) | **Required** | - | -| `ppo_n_minibatches` | integer | `4` | Number of minibatches for each PPO update | +| `ppo_n_minibatches` | integer | `4` | Number of minibatches for each PPO update. Separation DTE AdamW delta transfer currently requires 1. | | `eps_clip` | float | `0.2` | Clipping factor for policy ratio | | `eps_clip_higher` | float \| None | `None` | Clipping factor (higher value) for policy ratio. Default is None. When eps_clip_higher is set (decoupled), eps_clip will be used as the lower value. | | `c_clip` | float \| None | `None` | Dual clipping factor for policy ratio, must be > 1.0. None disables dual clipping. | diff --git a/examples/dte/README.md b/examples/dte/README.md index 6fdec9d0f7..d499b57e3c 100644 --- a/examples/dte/README.md +++ b/examples/dte/README.md @@ -26,6 +26,22 @@ AdamW deltas, with a periodic full-weight anchor after every 20 successfully com deltas. Set `actor.dte.enabled=false` to return to the existing AWEX full-weight behavior. +## Optimizer-step boundary + +The current separation AdamW detector reconstructs exactly one optimizer step between +weight updates. Therefore, enabling DTE requires: + +```yaml +actor: + ppo_n_minibatches: 1 +``` + +Each PPO minibatch performs an optimizer step, while weight synchronization happens only +after all PPO minibatches finish. Values greater than one would make the optimizer step +advance by more than the detector can invert and would force a full-weight fallback. +AReaL rejects that configuration at startup instead of silently running full +synchronization while DTE appears enabled. + ## Effective performance defaults The example spells out the portable performance settings in the worker environment; none diff --git a/tests/test_awex_delta_common.py b/tests/test_awex_delta_common.py index 4a72e99652..0bea4790fa 100644 --- a/tests/test_awex_delta_common.py +++ b/tests/test_awex_delta_common.py @@ -10,6 +10,7 @@ import importlib.util import logging as stdlib_logging +import os import sys import types from pathlib import Path @@ -23,11 +24,53 @@ _DC_PATH = _ROOT / "areal/v2/weight_update/awex/delta_config.py" +def _make_environ_stub(): + environ_mod = types.ModuleType("areal.utils.environ") + + def get_env_var(name, default=None, *, fallback_names=(), allow_empty=False): + for candidate in (name, *fallback_names): + value = os.environ.get(candidate) + if value is not None and (allow_empty or value.strip() != ""): + return value + return default + + def get_bool_env_var( + name, + default="false", + *, + fallback_names=(), + truthy_values=("true", "1"), + falsy_values=("false", "0"), + strip_value=False, + ): + del falsy_values + value = get_env_var(name, default, fallback_names=fallback_names) + value = value.strip() if strip_value else value + return value.lower() in truthy_values + + environ_mod.get_env_var = get_env_var + environ_mod.get_bool_env_var = get_bool_env_var + return environ_mod + + def _load_delta_config(): spec = importlib.util.spec_from_file_location("awex_delta_config", _DC_PATH) mod = importlib.util.module_from_spec(spec) assert spec.loader is not None - spec.loader.exec_module(mod) + module_names = ("areal", "areal.utils", "areal.utils.environ") + saved_modules = {name: sys.modules.get(name) for name in module_names} + sys.modules["areal"] = types.ModuleType("areal") + sys.modules["areal.utils"] = types.ModuleType("areal.utils") + sys.modules["areal.utils.environ"] = _make_environ_stub() + sys.modules[spec.name] = mod + try: + spec.loader.exec_module(mod) + finally: + for name, saved in saved_modules.items(): + if saved is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = saved return mod @@ -65,36 +108,41 @@ def test_delta_config_env_gates(dc, monkeypatch): """DTE env vars must honor DTE_* overrides over legacy AWEX_* names.""" monkeypatch.delenv("DTE_DELTA_TRANSFER", raising=False) monkeypatch.delenv("AWEX_DELTA_TRANSFER", raising=False) - assert dc.delta_transfer_enabled() is False + assert dc.DTERuntimeConfig.from_env().delta_transfer is False monkeypatch.setenv("AWEX_DELTA_TRANSFER", "1") - assert dc.delta_transfer_enabled() is True + assert dc.DTERuntimeConfig.from_env().delta_transfer is True monkeypatch.setenv("DTE_DELTA_TRANSFER", "0") - assert dc.delta_transfer_enabled() is False + assert dc.DTERuntimeConfig.from_env().delta_transfer is False monkeypatch.setenv("DTE_DELTA_TRANSFER", "1") - assert dc.delta_transfer_enabled() is True + assert dc.DTERuntimeConfig.from_env().delta_transfer is True monkeypatch.setenv("AWEX_DELTA_ANCHOR_INTERVAL", "5") - assert dc.delta_anchor_interval() == 5 + assert dc.DTERuntimeConfig.from_env().anchor_interval == 5 monkeypatch.setenv("DTE_DELTA_ANCHOR_INTERVAL", "7") - assert dc.delta_anchor_interval() == 7 + assert dc.DTERuntimeConfig.from_env().anchor_interval == 7 -def test_cuda_mem_stats_mb_without_cuda_returns_sentinel(dc): - """CPU-only hosts must not fail while reporting DTE memory telemetry.""" - alloc_mb, peak_mb = dc.cuda_mem_stats_mb() - if torch.cuda.is_available(): - assert alloc_mb >= 0.0 - assert peak_mb >= 0.0 - else: - assert (alloc_mb, peak_mb) == (-1.0, -1.0) - assert dc.cuda_mem_stats_mb(reset_peak=False) == (-1.0, -1.0) +def test_delta_runtime_config_preserves_legacy_values_and_snapshots(dc, monkeypatch): + monkeypatch.setenv("DTE_DELTA_TRANSFER", "") + monkeypatch.setenv("AWEX_DELTA_TRANSFER", " yes ") + monkeypatch.setenv("DTE_SEPARATION_WEIGHT_UPDATE", "on") + + config = dc.DTERuntimeConfig.from_env() + + assert config.delta_transfer is True + assert config.separation_weight_update is True + assert config.enabled is True + + monkeypatch.setenv("DTE_DELTA_TRANSFER", "0") + assert config.enabled is True + assert dc.DTERuntimeConfig.from_env().enabled is False def test_factory_builds_dte_tracker(dc, monkeypatch): """The lazy factory should build a DTE tracker when DTE is available.""" pytest.importorskip("dte") monkeypatch.setenv("AWEX_DELTA_ANCHOR_INTERVAL", "0") - tracker = dc.make_delta_tracker() + tracker = dc.DTERuntimeConfig.from_env().create_delta_tracker() assert hasattr(tracker, "encode") and hasattr(tracker, "seed") @@ -189,6 +237,7 @@ def _load_delta_detect(monkeypatch): monkeypatch.setitem(sys.modules, "areal", types.ModuleType("areal")) monkeypatch.setitem(sys.modules, "areal.utils", types.ModuleType("areal.utils")) monkeypatch.setitem(sys.modules, "areal.utils.logging", fake) + monkeypatch.setitem(sys.modules, "areal.utils.environ", _make_environ_stub()) for package in ( "areal.v2", "areal.v2.weight_update", @@ -208,6 +257,71 @@ def _load_delta_detect(monkeypatch): return mod +def test_cuda_mem_stats_mb_without_cuda_returns_sentinel(monkeypatch): + """CPU-only hosts must not fail while reporting DTE memory telemetry.""" + mod = _load_delta_detect(monkeypatch) + alloc_mb, peak_mb = mod._cuda_mem_stats_mb() + if torch.cuda.is_available(): + assert alloc_mb >= 0.0 + assert peak_mb >= 0.0 + else: + assert (alloc_mb, peak_mb) == (-1.0, -1.0) + assert mod._cuda_mem_stats_mb(reset_peak=False) == (-1.0, -1.0) + + +def test_adamw_hparams_require_and_prefer_recorded_step_lr(monkeypatch): + """A scheduler's next-step LR cannot silently drive AdamW inversion.""" + mod = _load_delta_detect(monkeypatch) + param_group = { + "lr": 2e-6, + "_areal_last_step_lr": 3e-6, + "weight_decay": 0.1, + "betas": (0.9, 0.95), + "eps": 1e-8, + } + + hparams = mod._adamw_hparams(param_group) + + assert hparams is not None + assert hparams[0] == 3e-6 + del param_group["_areal_last_step_lr"] + assert mod._adamw_hparams(param_group) is None + + +def test_missing_recorded_step_lr_forces_dense_reconstruction(monkeypatch): + """Missing LR metadata must not fall back to the scheduler's current LR.""" + mod = _load_delta_detect(monkeypatch) + param = torch.nn.Parameter(torch.tensor([1.0, 2.0], dtype=torch.float32)) + optimizer = torch.optim.AdamW([param], lr=1e-3) + before = param.detach().clone() + + inversion = mod.AdamWInversionDetector( + SimpleNamespace(_offloaded_optimizer_states={}) + ) + _bind_inversion_param_names(inversion, ("w", param)) + inversion._last_synced_steps = {"w": None} + inversion._last_synced_fingerprints = {"w": mod._tensor_fingerprint(before)} + + param.grad = torch.ones_like(param) + optimizer.step() + assert "_areal_last_step_lr" not in optimizer.param_groups[0] + + class _FakeDistOpt: + shard_fp32_from_float16_groups = [[param]] + model_float16_groups = [[param]] + model_param_group_index_map = None + data_parallel_group = None + + def __init__(self): + self.optimizer = optimizer + + def _get_model_param_range_map(self, model_param): + assert model_param is param + return {"param": SimpleNamespace(start=0, end=param.numel())} + + assert inversion._reconstruct_pre_step_mcore([_FakeDistOpt()]) is None + + def _bind_inversion_param_names(inv, *named_params): """Bind fake mcore parameter names for CPU-only inversion tests.""" id2key = {id(param): name for name, param in named_params} @@ -222,7 +336,13 @@ def _load_sglang_adapter(monkeypatch): _stub_areal_packages(monkeypatch) delta_config_mod = types.ModuleType("areal.v2.weight_update.awex.delta_config") - delta_config_mod.separation_delta_transfer_enabled = lambda: False + + class _DTERuntimeConfig: + @classmethod + def from_env(cls): + return SimpleNamespace(enabled=False) + + delta_config_mod.DTERuntimeConfig = _DTERuntimeConfig monkeypatch.setitem( sys.modules, "areal.v2.weight_update.awex.delta_config", @@ -268,8 +388,16 @@ def _load_megatron_adapter(monkeypatch): _stub_areal_packages(monkeypatch) delta_config_mod = types.ModuleType("areal.v2.weight_update.awex.delta_config") - delta_config_mod.separation_delta_transfer_enabled = lambda: True - delta_config_mod.make_delta_tracker = lambda *args, **kwargs: None + + class _DTERuntimeConfig: + @classmethod + def from_env(cls): + return SimpleNamespace( + enabled=True, + create_delta_tracker=lambda: None, + ) + + delta_config_mod.DTERuntimeConfig = _DTERuntimeConfig monkeypatch.setitem( sys.modules, "areal.v2.weight_update.awex.delta_config", diff --git a/tests/test_awex_separation_delta.py b/tests/test_awex_separation_delta.py index a2b01478a9..bcf2ff8a6a 100644 --- a/tests/test_awex_separation_delta.py +++ b/tests/test_awex_separation_delta.py @@ -43,9 +43,9 @@ def _init_group(**kwargs): def test_megatron_full_transfer_initializes_gloo_control_group(monkeypatch): """Separated Full sync must not fall back to a NCCL control barrier.""" mod = common._load_megatron_adapter(monkeypatch) - monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) calls = _capture_weight_update_groups(monkeypatch, mod) adapter = object.__new__(mod.AwexMegatronAdapter) + adapter._dte_config = SimpleNamespace(enabled=False) adapter.init_weight_update_group( pair_name="pair", @@ -67,9 +67,9 @@ def test_megatron_full_transfer_initializes_gloo_control_group(monkeypatch): def test_sglang_full_transfer_initializes_gloo_control_group(monkeypatch): """SGLang Full sync creates the same CPU control group as training.""" mod = common._load_sglang_adapter(monkeypatch) - monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) calls = _capture_weight_update_groups(monkeypatch, mod) adapter = object.__new__(mod.AwexSGLangAdapter) + adapter._dte_config = SimpleNamespace(enabled=False) adapter._get_model_context = lambda: { "tp_size": 4, "tp_rank": 0, @@ -97,7 +97,6 @@ def test_sglang_full_transfer_initializes_gloo_control_group(monkeypatch): def test_megatron_full_transfer_uses_gloo_completion_barrier(monkeypatch): """Full payload P2P completes through the sideband control group.""" mod = common._load_megatron_adapter(monkeypatch) - monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) monkeypatch.setattr( mod, "nccl_build_send_ops", lambda *args, **kwargs: ([], [], []) ) @@ -105,6 +104,7 @@ def test_megatron_full_transfer_uses_gloo_completion_barrier(monkeypatch): barriers = [] monkeypatch.setattr(mod.dist, "barrier", lambda *, group: barriers.append(group)) adapter = object.__new__(mod.AwexMegatronAdapter) + adapter._dte_config = SimpleNamespace(enabled=False) adapter._transfer_plan = object() adapter._weights_update_group = "nccl" adapter._weights_update_group_gloo = "gloo" @@ -120,7 +120,6 @@ def test_megatron_full_transfer_uses_gloo_completion_barrier(monkeypatch): def test_sglang_full_transfer_uses_gloo_completion_barrier(monkeypatch): """Receiver Full payload completion avoids the NCCL data group.""" mod = common._load_sglang_adapter(monkeypatch) - monkeypatch.setattr(mod, "separation_delta_transfer_enabled", lambda: False) monkeypatch.setattr( mod, "nccl_build_recv_ops", lambda *args, **kwargs: ([], [], []) ) @@ -128,6 +127,7 @@ def test_sglang_full_transfer_uses_gloo_completion_barrier(monkeypatch): barriers = [] monkeypatch.setattr(mod.dist, "barrier", lambda *, group: barriers.append(group)) adapter = object.__new__(mod.AwexSGLangAdapter) + adapter._dte_config = SimpleNamespace(enabled=False) adapter._transfer_plan = object() adapter._weights_update_group = "nccl" adapter._weights_update_group_gloo = "gloo" diff --git a/tests/test_dte_topology_gating.py b/tests/test_dte_topology_gating.py index 14912825aa..44ad3638c1 100644 --- a/tests/test_dte_topology_gating.py +++ b/tests/test_dte_topology_gating.py @@ -79,3 +79,17 @@ def test_dte_rejects_out_of_scope_modes(kwargs, match): with pytest.raises(ValueError, match=match): _load_helpers().apply_dte_config_envvars(config, environ={}) + + +def test_dte_requires_one_ppo_minibatch_per_weight_update(): + pytest.importorskip("httpx") + from areal.api.cli_args import DTEConfig, PPOActorConfig + + enabled = DTEConfig(enabled=True) + assert PPOActorConfig(dte=enabled, ppo_n_minibatches=1).ppo_n_minibatches == 1 + + with pytest.raises(ValueError, match="requires ppo_n_minibatches=1"): + PPOActorConfig(dte=enabled, ppo_n_minibatches=2) + + disabled = DTEConfig(enabled=False) + assert PPOActorConfig(dte=disabled).ppo_n_minibatches == 4 diff --git a/tests/test_environ.py b/tests/test_environ.py new file mode 100644 index 0000000000..edf016c008 --- /dev/null +++ b/tests/test_environ.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for shared environment-variable parsing helpers.""" + +import importlib.util +import logging as stdlib_logging +import sys +import types +from pathlib import Path + +_ENVIRON_PATH = Path(__file__).parents[1] / "areal/utils/environ.py" + + +def _load_environ(monkeypatch): + logging_mod = types.ModuleType("areal.utils.logging") + logging_mod.getLogger = stdlib_logging.getLogger + utils_mod = types.ModuleType("areal.utils") + utils_mod.logging = logging_mod + monkeypatch.setitem(sys.modules, "areal", types.ModuleType("areal")) + monkeypatch.setitem(sys.modules, "areal.utils", utils_mod) + monkeypatch.setitem(sys.modules, "areal.utils.logging", logging_mod) + + spec = importlib.util.spec_from_file_location("areal_test_environ", _ENVIRON_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_get_env_var_uses_primary_then_fallback(monkeypatch): + environ = _load_environ(monkeypatch) + monkeypatch.delenv("AREAL_TEST_PRIMARY", raising=False) + monkeypatch.setenv("AREAL_TEST_LEGACY", "legacy") + + assert ( + environ.get_env_var( + "AREAL_TEST_PRIMARY", + fallback_names=("AREAL_TEST_LEGACY",), + ) + == "legacy" + ) + + monkeypatch.setenv("AREAL_TEST_PRIMARY", "primary") + assert ( + environ.get_env_var( + "AREAL_TEST_PRIMARY", + fallback_names=("AREAL_TEST_LEGACY",), + ) + == "primary" + ) + + +def test_get_env_var_empty_value_fallback_is_configurable(monkeypatch): + environ = _load_environ(monkeypatch) + monkeypatch.setenv("AREAL_TEST_PRIMARY", "") + monkeypatch.setenv("AREAL_TEST_LEGACY", "legacy") + + assert ( + environ.get_env_var( + "AREAL_TEST_PRIMARY", + fallback_names=("AREAL_TEST_LEGACY",), + ) + == "legacy" + ) + assert ( + environ.get_env_var( + "AREAL_TEST_PRIMARY", + fallback_names=("AREAL_TEST_LEGACY",), + allow_empty=True, + ) + == "" + ) + + +def test_get_bool_env_var_preserves_defaults_and_supports_opt_in_values(monkeypatch): + environ = _load_environ(monkeypatch) + monkeypatch.setenv("AREAL_TEST_BOOL", "yes") + + assert environ.get_bool_env_var("AREAL_TEST_BOOL") is False + assert ( + environ.get_bool_env_var( + "AREAL_TEST_BOOL", + truthy_values=("true", "1", "yes", "on"), + falsy_values=("false", "0", "no", "off"), + ) + is True + ) + + +def test_get_bool_env_var_can_strip_legacy_dte_values(monkeypatch): + environ = _load_environ(monkeypatch) + monkeypatch.setenv("AREAL_TEST_BOOL", " on ") + + assert ( + environ.get_bool_env_var( + "AREAL_TEST_BOOL", + truthy_values=("true", "1", "yes", "on"), + falsy_values=("false", "0", "no", "off"), + strip_value=True, + ) + is True + ) diff --git a/tests/test_megatron_optimizer_config.py b/tests/test_megatron_optimizer_config.py index 6563b7899e..c829a0e6fd 100644 --- a/tests/test_megatron_optimizer_config.py +++ b/tests/test_megatron_optimizer_config.py @@ -116,6 +116,39 @@ def capture_scheduler(*args, **kwargs): assert captured["wd_incr_steps"] == 10 +def test_dte_records_optimizer_step_lr_before_scheduler_advances() -> None: + """AdamW inversion must retain the LR consumed by the completed step.""" + engine = megatron_engine_module.MegatronEngine.__new__( + megatron_engine_module.MegatronEngine + ) + param_groups = [{"lr": 3e-6}, {"lr": 4e-6}] + + class _Optimizer: + def __init__(self): + self.param_groups = param_groups + + def step(self): + return True, torch.tensor(1.0), None + + class _Scheduler: + def step(self, increment): + assert increment == 1 + param_groups[0]["lr"] = 2e-6 + param_groups[1]["lr"] = 1e-6 + + engine.optimizer = _Optimizer() + engine.lr_scheduler = _Scheduler() + engine._dte_runtime_config = SimpleNamespace(enabled=True) + + engine.optimizer_step() + engine.lr_scheduler_step() + + assert param_groups[0]["lr"] == 2e-6 + assert param_groups[1]["lr"] == 1e-6 + assert param_groups[0]["_areal_last_step_lr"] == 3e-6 + assert param_groups[1]["_areal_last_step_lr"] == 4e-6 + + @pytest.mark.parametrize( ("optimizer_config", "ft_spec", "match"), [