Skip to content

feat: add bounded colocate transfer and Qwen4Exp converters - #121

Merged
chaokunyang merged 2 commits into
inclusionAI:mainfrom
dingzhiqiang:fix/bounded-colocate-transfer
Sep 22, 2026
Merged

chaokunyang merged 2 commits into
inclusionAI:mainfrom
dingzhiqiang:fix/bounded-colocate-transfer

Conversation

@dingzhiqiang

@dingzhiqiang dingzhiqiang commented Sep 22, 2026 •

Copy link
Copy Markdown
Contributor

Description

Keep transfer planning as tensor views and allocate contiguous send/receive buffers only for the active batch, including strided column slices. Copy receives back before releasing that batch. Pack compatible expert operations within configured bounds, check packing limits across inference ranks before P2P starts, and provide CUDA IPC staging allocation that restores the allocator settings used for training. The packing byte limit is not a hard cap for a single oversized tensor.

Add a reader factory so integrations can explicitly choose the bounded transport; the AWEX default remains unchanged. All participating readers must choose the same transport. This extracts the shared transport implementation from areal-project/AReaL#1731.

Own Qwen4Exp weight converters, GDN/gated-QKV layouts and frozen-parameter declarations in AWEX. Registration is explicit and accepts caller-supplied binders; checkpoint evidence and live engine ownership stay in the integration. AReaL consumes these APIs in areal-project/AReaL#1751 without a second converter implementation.

Validation

  • 41 targeted CPU tests passed, including real two-rank Gloo agreement/disagreement, strided buffer allocation/release and independent GDN/QKV coordinate checks across TP sizes.
  • AReaL integration regression suite: 557 passed, 9 optional-runtime skips; separate Transformers 5.16.1 vision run: 31 passed.
  • Ruff check and format passed for changed files.
  • GPU NCCL transfers, repeated weight equality, and cross-image CUDA IPC remain unverified and require validation before merge.

@dingzhiqiang dingzhiqiang changed the title feat(transfer): add bounded colocate transport and portable IPC allocation feat: add bounded colocate transfer and Qwen4Exp converters Sep 22, 2026
@dingzhiqiang
dingzhiqiang marked this pull request as ready for review September 22, 2026 11:11

@chaokunyang chaokunyang left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@chaokunyang
chaokunyang merged commit b6c24aa into inclusionAI:main Sep 22, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants