[Feature] Async Tensor Parallelism for Archon Engine

Checklist

• This feature will maintain backward compatibility with the current APIs in areal/api/. If not, please raise a refactor issue first.

Background

Picking up the async TP task from the contributor list.

Noted at qwen2/infra/parallelize.py L286-287:

For PP, by default do not reshard after forward to avoid per-microbatch all-gathers, which can be expensive and non-overlapped

TP Communication Is on Activations, Not Weights

I checked whether FSDP2's set_modules_to_forward_prefetch could be extended to cover TP. It cannot — FSDP prefetch only triggers FSDP unshard (DP-dimension all-gather of weights). TP communication in Megatron-SP is fundamentally different: it redistributes activations (hidden states) between TP ranks, not weights.

Since each all-gather depends on the immediately preceding computation's output, inter-layer activation prefetch is not possible — the data to redistribute does not exist until the previous operation finishes. This is unlike FSDP weight prefetch where weights exist independently of the data flow.

Potential Solution

There are two levels of overlap to consider:

Intra-layer: Within a single linear layer, decompose all-gather + matmul into chunked pipelined operations. PyTorch has a torch.compile graph pass for this, validated in the TorchTitan ICLR 2025 paper (up to 29% forward pass speedup). Since Archon already compiles transformer blocks via apply_compile with fullgraph=True, and _c10d_functional ops are Inductor-compatible, enabling this should be achievable.

Sub-layer boundary overlap: Between the reduce-scatter at the end of one sub-layer (attention or FFN) and the all-gather at the start of the next, there is a small window (residual add + RMSNorm) where compute and communication could overlap. Since RMSNorm reduces over the hidden dimension and the all-gather is on the sequence dimension, norm(all_gather(x)) == all_gather(norm(x)) — reordering the all-gather before the norm is mathematically valid, letting it start earlier. The gain per boundary is small (norm + residual are cheap), but there are 4 boundaries per block (2 intra-block, 2 inter-block).

Both are complementary. My plan:

• Phase 1 — Enable intra-layer async TP via the compile pass for Qwen2. Lower risk, leverages existing PyTorch machinery. Verify correctness and measure speedup against baseline.
• Phase 2 — Investigate sub-layer boundary overlap (reorder all-gather before norm) if Phase 1 alone is insufficient. This touches model forward passes or PrepareModuleInput hooks.
• Phase 3 — Extend to Qwen3 (MoE + EP coexistence) and Qwen3.5 (hybrid attention, once basic TP lands). Verify with PP scheduling (zero-bubble), activation checkpointing (recomputation path), and torch.compile.

Additional Information

• SequenceParallel on norm layers has an "async redistribute which breaks..." comment at qwen3/infra/parallelize.py L581 — directly relevant since sub-layer boundary overlap would interact with these async redistributions. Using raw _c10d_functional ops + wait_tensor instead of DTensor .redistribute(async_op=True) avoids this.
• Known limitation: torch.compile + TP has edge cases with MXFP8 linear layers (PyTorch blog). Worth checking if Archon's FP8 path is affected.
• SAC save lists (_get_op_sac_save_list) already include _c10d_functional.reduce_scatter_tensor.default — new async ops from boundary overlap may need to be added.