refactor: unify AWEX converters and adapters across v1 and v2

README.md

@@ -245,6 +245,12 @@ All RL algorithms support both asynchronous and synchronous versions by setting
 | **Gemma 3**                | ❌       | ✅           | ❌             | Vision-language model                                    |
 | **Other Hugging Face LLM** | ❌       | ✅           | ❌             | Compatibility depending on the version of `transformers` |
 
+Qwen2 and Qwen2.5 **text models** support RL weight updates with
+`actor.weight_update_mode=disk` in both v1 and v2. XCCL weight updates require
+`actor._version=v1`, `rollout._version=v1`, and `actor.weight_update_mode=xccl`. AWEX
+weight updates are not supported for these models. This restriction does not apply to
+Qwen2.5-VL.
+
 Check the [AI Coding Assistant Guide](docs/en/reference/ai_assisted_dev.md) and
 [Archon Reference](docs/en/tutorial/archon.md) for how to integrate new models into
 AReaL.

areal/engine/awex/__init__.py

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

areal/engine/awex/adapters/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared AWEX training and inference adapters."""

areal/v2/weight_update/awex/fsdp_adapter.py → areal/engine/awex/adapters/fsdp_adapter.py

@@ -1,4 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
+"""AWEX training adapter for FSDP engines."""
+
 from __future__ import annotations
 
 # pyright: reportMissingImports=false
@@ -18,19 +20,20 @@
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.placement_types import Shard
 
-from areal.engine.core.model import is_qwen_vl_model
-from areal.utils import logging
-from areal.v2.weight_update.awex import (
+from areal.engine.awex.adapters.training_adapter import (
+    AwexTrainingAdapter,
+)
+from areal.engine.awex.transport.dtype import align_send_ops_to_recv_dtype
+from areal.engine.awex.transport.metadata import (
     awex_wu_use_group,
     fetch_kv_metadata,
 )
-from areal.v2.weight_update.nccl_group import (
+from areal.engine.awex.transport.nccl_group import (
     init_weights_update_group,
     setup_batch_isend_irecv,
 )
-from areal.v2.weight_update.training_adapter import (
-    AwexTrainingAdapter,
-)
+from areal.engine.core.model import is_qwen_vl_model
+from areal.utils import logging
 
 if TYPE_CHECKING:
     from areal.engine.fsdp_engine import FSDPEngine
@@ -61,6 +64,7 @@ def parallelism_strategy(self) -> dict:
             "dp_size": self._engine.data_parallel_world_size,
             "ep_size": 1,
             "dp_replicated": False,
+            "parameter_layout": "hf",
         }
 
     @property
@@ -138,10 +142,11 @@ def init_weight_update_group(
         infer_world_size: int,
         train_world_size: int,
         num_engines: int,
+        timeout_s: float = 300.0,
     ) -> None:
         self._transfer_rank = transfer_rank
 
-        infer_meta, train_meta = fetch_kv_metadata(kv_store_url, pair_name)
+        infer_meta, train_meta = fetch_kv_metadata(kv_store_url, pair_name, timeout_s)
 
         builder = TransferPlanBuilder(
             infer_world_size=infer_world_size,
@@ -198,6 +203,7 @@ def execute_weight_update(self, version: int) -> None:
             self._weights_update_group,
             copy_rank=self._transfer_rank,
         )
+        align_send_ops_to_recv_dtype(send_ops, self._transfer_plan)
         batch_send_recv(
             send_ops=send_ops,
             recv_ops=[],

areal/v2/weight_update/inference_adapter.py → areal/engine/awex/adapters/inference_adapter.py

@@ -1,4 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
+"""Shared contract for AWEX inference adapters."""
+
 from __future__ import annotations
 
 from typing import Protocol, runtime_checkable
@@ -42,6 +44,7 @@ def init_weight_update_group(
         infer_world_size: int,
         train_world_size: int,
         num_engines: int,
+        timeout_s: float = 300.0,
     ) -> None:
         """Pull peer meta from KV store, build local recv plan, join NCCL group."""
         ...