refactor: unify AWEX converters and adapters across v1 and v2 by sitabulaixizawaluduo · Pull Request #1684 · areal-project/AReaL

sitabulaixizawaluduo · 2026-09-07T04:45:32Z

Description

Unify the Megatron and SGLang AWEX integrations so v1 colocated weight updates and v2 separated weight updates share adapter implementations and native AWEX model conversion.

Shared adapters

Centralize Megatron, SGLang, and FSDP AWEX adapters and the training/inference Protocols in areal/engine/awex/.
Connect v1 and v2 entry points directly to the shared adapters and remove the legacy reader/writer facades and superseded v2 adapter modules.
Keep NCCL group setup, metadata exchange, DTE configuration/detection, and device utilities alongside the adapters, with no imports from areal.v2 in the shared AWEX package.
Reuse MegatronWeightResidency for weight, optimizer, and gradient memory management.
Retain v1 colocated IPC transfer and v2 separated full-weight and delta-weight transfer paths.

Native conversion and metadata

Use AWEX's native Megatron and SGLang converters, rank extraction, and sharding metadata.
Pass inference converter context through the weight-update gateway to training workers before generating parameter metadata.
Propagate the training parameter layout to SGLang metadata generation, including an explicit HF layout for FSDP.
Replace duplicated conversion and metadata handling with shared adapter logic.

Type of Change

♻️ Refactoring

Le8r0nJames · 2026-09-09T06:29:12Z

AwexFSDPAdapter now imports and calls batch_send_recv_by_peer, so this mock targets a symbol that no longer exists in the FSDP module. The FSDP test case raises AttributeError before reaching the barrier assertion. Please patch batch_send_recv_by_peer instead, and verify that the transfer is invoked before the Gloo completion barrier.

Fixed in a17882b. The parameterized FSDP/Megatron/SGLang test now patches batch_send_recv_by_peer, checks its arguments for both use_group settings, and asserts that the transfer returns before the Gloo completion barrier. The SGLang copy test now asserts transfer -> copy -> synchronize -> barrier. I also updated the two DTE module-loader stubs to expose the new helper.

Validation: full pre-commit, Ruff, and Python compilation passed. Unit tests were not run locally.

Le8r0nJames · 2026-09-10T08:23:54Z

+        model_config, _ = PretrainedConfig.get_config_dict(actor.path)
+        if model_config.get("model_type") == "qwen2" or "Qwen2ForCausalLM" in (
+            model_config.get("architectures") or []
+        ):
+            raise ValueError(
+                "Qwen2/Qwen2.5 text models do not support AWEX weight updates. "
+                "For RL, set actor._version=v1, rollout._version=v1, and "
+                "actor.weight_update_mode=xccl or disk. Controller v2 is not "
+                "supported for these models' RL weight updates."
+            )


This guard also rejects Qwen2/Qwen2.5 with v2 + use_lora=True, but the existing trainer routes that configuration through disk-based adapter updates, bypassing AWEX conversion. Is there a separate incompatibility that requires rejecting this path? If not, could we scope the guard to configurations that actually select AWEX?

Le8r0nJames · 2026-09-10T08:33:12Z

+    for peer in sorted(ops_by_peer):
+        peer_ops = ops_by_peer[peer]
+        batch_send_recv(
+            send_ops=peer_ops if send_ops else [],
+            recv_ops=peer_ops if recv_ops else [],
+            blocking=True,
+            use_group=True,
+        )


Is peer serialization required for the shared adapter/native converter integration to work? It appears independent of the refactor and changes the default transfer scheduling. Unless there is a demonstrated correctness or resource issue that makes it necessary, could we preserve the existing scheduling here and move this change to a separate PR with before/after measurements?

Share Megatron and SGLang weight-update adapters so both versions use native AWEX conversion and sharding metadata. Key changes: - Consolidate v1 colocate and v2 separation adapter implementations - Preserve legacy entry points and shared Megatron weight residency - Forward inference converter context through the weight-update gateway - Propagate parameter layouts between training and inference workers - Retain full-weight and delta-weight separation transfer paths

Allow metadata fetches to honor the gateway initialization timeout instead of always using the aiohttp default. Key changes: - Extend training and inference adapter contracts compatibly - Forward the timeout through FSDP, Megatron, and SGLang adapters - Apply the timeout to metadata HTTP requests

Merge same-name TP shards before the planner indexes parameters, and validate inference replicas before retaining a single layout. This avoids missing TP shards and redundant metadata distribution. Forward the initialization timeout to separation adapters.

Use the actor request timeout for the weight-update controller and its gateway initialization and update requests, rather than silently falling back to the shorter internal defaults.

Bound grouped P2P concurrency to one peer per rank while preserving AWEX's blocking device synchronization and manual-stream fallback.

Keep senders and receivers on the same deterministic peer schedule to limit concurrent NCCL P2P traffic during full weight updates. Key changes: - Use peer-serialized transfers in FSDP and Megatron senders. - Use the matching schedule in the SGLang receiver. - Preserve v1 colocate, DTE paths, and the manual-stream fallback.

Forked workers can inherit a per-task CUDA_VISIBLE_DEVICES namespace. Use logical device zero when that namespace already covers the server.

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

Keep adapter tests aligned with the peer-serialized transfer helper. Assert transfer completion before the Gloo barrier and preserve the SGLang transfer, copy, synchronize, barrier sequence. Update the DTE module-loader stubs to expose the same helper.

Defer peer-serialized transport to a separate PR while keeping the shared adapter and native converter changes. Key changes: - Restore blocking AWEX batch_send_recv on training and inference. - Update transfer mocks while retaining completion-order assertions.

Disk updates bypass AWEX native conversion, so the model support guard must allow them for both controller versions.

sitabulaixizawaluduo requested review from Le8r0nJames, TaoZex, garrett4wade, geshi001, guozhihao-224, nuzant and rchardx as code owners September 7, 2026 04:45

sitabulaixizawaluduo marked this pull request as draft September 7, 2026 04:58

Le8r0nJames reviewed Sep 9, 2026

View reviewed changes

sitabulaixizawaluduo force-pushed the feature/v2-awex-native-converters branch from b377204 to 4d69307 Compare September 9, 2026 09:42

sitabulaixizawaluduo added the safe-to-test Ready to run unit-tests in a PR. label Sep 9, 2026

sitabulaixizawaluduo had a problem deploying to AReaL-unittests September 9, 2026 09:59 — with GitHub Actions Failure

sitabulaixizawaluduo added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Sep 10, 2026

sitabulaixizawaluduo temporarily deployed to AReaL-unittests September 10, 2026 04:30 — with GitHub Actions Inactive

sitabulaixizawaluduo marked this pull request as ready for review September 10, 2026 06:12

sitabulaixizawaluduo requested review from HwVanICI and fishcrap as code owners September 10, 2026 06:12

sitabulaixizawaluduo added a commit that referenced this pull request Sep 10, 2026

revert(v2): remove AWEX converter changes superseded by #1684

08b3fb5

sitabulaixizawaluduo mentioned this pull request Sep 10, 2026

feat(v2): enable multimodal processor in the data proxy #1694

Merged

Le8r0nJames reviewed Sep 10, 2026

View reviewed changes

Le8r0nJames previously approved these changes Sep 11, 2026

View reviewed changes

sitabulaixizawaluduo added 5 commits September 16, 2026 11:42

refactor(engine): centralize AWEX adapter protocols

2a4bc2b

refactor(engine): centralize AWEX metadata exchange helpers

3ef085a

refactor(engine): centralize AWEX NCCL group utilities

aa9ac93

refactor(engine): centralize AWEX colocate device utilities

3d1f4a0

sitabulaixizawaluduo added 28 commits September 16, 2026 11:43

fix(v2): honor actor timeouts for weight updates

d172ecd

Use the actor request timeout for the weight-update controller and its gateway initialization and update requests, rather than silently falling back to the shorter internal defaults.

fix(engine): add peer-serialized AWEX transfers

4cd03ad

Bound grouped P2P concurrency to one peer per rank while preserving AWEX's blocking device synchronization and manual-stream fallback.

fix(engine): respect isolated GPUs in AWEX colocate launch

c00c2ba

Forked workers can inherit a per-task CUDA_VISIBLE_DEVICES namespace. Use logical device zero when that namespace already covers the server.

refactor(engine): group AWEX adapter protocols

ac57043

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group AWEX metadata transport

d77145e

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group AWEX NCCL transport

82b71fa

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group AWEX DTE configuration

460b6e3

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group AWEX delta detection

4a22eeb

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group the FSDP AWEX adapter

97942b0

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group the Megatron AWEX adapter

7548828

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group the SGLang AWEX adapter

b4c36f1

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

refactor(engine): group AWEX inference plugins

e339e18

Keep the shared AWEX package navigable without changing transfer logic. Key changes: - Move this module into its responsibility-based subpackage - Update imports and affected runtime or test-loader references together

test(engine): migrate DTE overrides to native AWEX conversion

5a5aab9

fix(trainer): limit Qwen2 RL weight updates to v1 XCCL or disk

b07e5b2

docs: clarify Qwen2 RL weight update support

b3b8667

test(trainer): initialize MOPD plan in evaluation fixture

1612ce2

test(infra): align RTensor test storage and fetch backends

a97375a

test(engine): remove unused AWEX peer-transfer stubs

deaca2c

revert(engine): remove peer-serialized AWEX transfer helper

9511e1a

fix(trainer): allow Qwen2 disk weight updates in v2

7c01d0f

Disk updates bypass AWEX native conversion, so the model support guard must allow them for both controller versions.

docs: clarify Qwen2 disk weight updates across versions

32b848a

fix(engine): align Megatron AWEX sends to receiver dtype

84493e0

fix(engine): align FSDP AWEX sends to receiver dtype

599718c

sitabulaixizawaluduo dismissed Le8r0nJames’s stale review via 599718c September 16, 2026 03:47

sitabulaixizawaluduo force-pushed the feature/v2-awex-native-converters branch from 7f09d5a to 599718c Compare September 16, 2026 03:47

Navigation Menu

refactor: unify AWEX converters and adapters across v1 and v2 - #1684

refactor: unify AWEX converters and adapters across v1 and v2#1684
sitabulaixizawaluduo wants to merge 38 commits into
mainfrom
feature/v2-awex-native-converters

sitabulaixizawaluduo commented Sep 7, 2026 •

edited

Loading

Uh oh!

Le8r0nJames Sep 9, 2026

Uh oh!

sitabulaixizawaluduo Sep 9, 2026

Uh oh!

Le8r0nJames Sep 10, 2026

Uh oh!

Le8r0nJames Sep 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

sitabulaixizawaluduo commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Shared adapters

Native conversion and metadata

Type of Change

Uh oh!

Le8r0nJames Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

sitabulaixizawaluduo Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

Le8r0nJames Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Le8r0nJames Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sitabulaixizawaluduo commented Sep 7, 2026 •

edited

Loading