[Design proposal]: shared residual-connection integration for mHC and Wide Residual

[kwangjunahn]

Proposal

Share the layer-integration machinery for mHC and Wide Residual while preserving their separate algorithms, kernels, and supported behavior. The goal is to reduce duplicated code and make future integration changes easier to review and maintain. This shared structure could also make it easier to implement other wide-residual architectures from the literature without duplicating the layer-integration code.

Motivation

Both implementations read an ordinary-width input from a wide residual stream, execute a branch, and write its update back to the stream. WR uses a ResidualConnection contract through shared Transformer hooks. mHC has separate layer overrides that repeat related normalization, checkpoint, and offload bookkeeping.

HybridStack also coordinates separate mHC and WR replay paths, although both use CheckpointWithoutOutputManager. Sharing these mechanics could reduce duplicated maintenance without changing the mathematical operations.

Proposed Structure

Use a common connection lifecycle, conceptually:

branch_input, state = connection.read(wide_stream)
branch_update = branch(norm(branch_input))
next_stream = connection.write(branch_update, state)

The connection owns its maps, carried tensor state, bias/dropout semantics, and residual update. The layer owns branch execution. Replay coordination owns block boundaries and manager lifetime.

WR will retain its token-independent streamwise maps and optional discounting/retention. mHC will retain its input-dependent maps, Sinkhorn normalization, and cross-stream mixing. Sharing the layer-integration structure should also make it easier to add other wide-residual architectures emerging in the literature while keeping their algorithm-specific behavior explicit.

Start from the existing connection contract, extending it only where necessary to preserve mHC's state and execution requirements.

Plans for Incremental PRs

Stage	Scope	Expected Result
1. Connection contract	Expose mHC's existing operations through a compatible interface	mHC caller uses the contract while retaining its execution schedule.
2. Transformer execution	Reuse shared read, normalization, and writeback hooks, retaining necessary architecture-specific policies.	Less duplicated layer code, with existing mHC and WR modes preserved.
3. Replay coordination	Share block planning, manager creation, and finalization; preserve each algorithm's operation-level checkpoint policy.	Common coordination with unchanged replay semantics.
4. HybridStack orchestration	Simplify dispatch and expansion/readout integration.	A more generic stack without changing model boundaries.

[Connor-XY]

Thanks for writing this up. +1 on sharing the integration code. A few thoughts:

1. Start with stage 3 (replay coordination). mHC and WR already group layers into recompute blocks the same way. We just have three copies of the planner. Unifying them changes no math and is easy to verify bit-exactly. One thing to preserve: mHC passes its recompute manager outside the forward kwargs so that it works with CUDA graphs. The shared version should do the same.

2. ResidualConnection needs a few extensions before mHC fits. The connection should be able to:

• choose what it carries as the residual (mHC uses its own view to keep its fused gradient add)
• choose what gets recomputed (mHC keeps its mapping coefficients and replays only the aggregation; replaying the whole read would redo Sinkhorn in backward)
• pass state from one sublayer to the next (single-pass mHC in #7500 needs this)

Calling the connection module for both read and write also has side effects. For WR this needed the FSDP fixes in #7757/#7819. mHC's write is a plain method call today, so it's worth planning for in stage 1.

3. HybridStack: wrapper or per-layer subclass? mHC wraps any hybrid layer, while WR uses a subclass per layer type. Please say which design we end up with. Moving mHC to subclasses changes checkpoint keys, which needs a key mapping and a Megatron-Bridge update. It also slightly changes numerics for wrapped Mamba/GDN layers.

4. Check the design against the variants already in flight: gated-residual mHC (#7553) and single-pass mHC (#7500). If the contract covers mHC, WR and those two, it's probably right. Several other open PRs touch the same files (Add mHC precision controls, mean output contract, and full activation recompute with mHC #7521, Add mHC pipeline and virtual-pipeline support #7336, Enable TE and full-iteration CUDA graphs for mHC #7337, Fix deterministic mHC backward reduction selection #7704, Wrap residual-stream replay children as MFSDP v2 units #7819, Enable inference CUDA graphs with Shortcut MoE and wide residuals #7772), so it would help to agree on ordering with their authors.

[kwangjunahn]

Thanks @Connor-XY for the detailed feedback. To address your questions and concerns directly:

1. Replay coordination: Agreed that this should come first. We will preserve existing graph-compatible manager handling and each algorithm’s replay boundaries. The shared planner also needs to preserve WR’s ShortcutMoE pairs coordination.
2. Connection contract and FSDP: Agreed that the current WR contract is insufficient as-is. The connection must choose its carried residual, including mHC’s autograd aliases. We should also retain mapping coefficients while replaying aggregation, and support explicit forward-local carry between sublayers (for single-pass mHC). A common read/write interface will not require both operations to enter nn.Module.__call__. We should preserve mHC’s current call boundaries initially and validate hook/parameter-lifetime behavior before changing them.
3. Wrapper versus subclass: My preferred end state is explicit subclasses for standard Hybrid recipes. They can access branch updates directly (no "subtraction" required) and fit the existing two-stage execution used by ShortcutMoE. However, wrapper migration should be separate from the initial mechanical refactors. Where parameter paths change, it needs checkpoint key mapping and Megatron-Bridge coordination. That said, the generic wrapper can remain for custom layouts.
4. Variants in flight: I reviewed both proposals/PRs. Gated residual requires normalized, channel-wise reads and different norm ownership. Single-pass mHC requires incoming coefficients, outgoing carry, and different entry/exit behavior. The proposed preparation/read/transport/injection components (see below) accommodate these differences without forcing separate kernel launches. We should confirm these boundaries with the authors and coordinate overlapping changes before implementation.

With that all said, the incremental PR plan below follows those decisions.

(Updated) Plans for Incremental PRs

1. Share replay coordination. (Unify mHC and WR replay coordination #7923) Consolidate block partitioning and per-forward manager creation while preserving algorithm-specific checkpoint boundaries, existing graph handling, and WR’s ShortcutMoE atomic grouping. (Also keep manager outside the forward kwargs as you suggested.)
2. Extend the residual-connection contract. Allow each implementation to choose its carried residual, saved write state, and replay policy, with separate forward-local carry for single-pass mHC. Adapt a real classic-mHC caller without changing its registered modules. I'll also check the contract against the single-pass and gated-residual proposals with their authors.
3. Share compatible Transformer-layer scheduling. Route read, normalization, branch execution, and write through the agreed contract, retaining explicit layer classes and static specs.
4. Simplify HybridStack orchestration. Share compatible expansion, carry handoff, replay finalization, and exit-contraction bookkeeping without changing parameter paths. Keep this step only where it genuinely simplifies the existing code.
5. Migrate standard Hybrid mHC recipes to explicit subclasses. Integrate directly with Transformer/GDN and Mamba/GDP branch updates, avoiding subtraction-based recovery on those paths. Review checkpoint conversion and numerical differences separately; retain the wrapper for transition and custom layouts. New mHC + ShortcutMoE support may warrant a separate feature PR.

Each PR should remain independently functional and be validated against its parent implementation for outputs and gradients, supported replay/distributed modes, and runtime and memory overhead. Subclass migration that removes add/subtract recovery will require separate numerical evaluation.

Proposed Modular Structure

The connection would compose four responsibilities:

• Map preparation: static parameters, joint token-dependent prediction, or normalized low-rank gate prediction.
• Read: scalar-weighted stream reduction or channel-wise gated reduction.
• Transport: identity, bounded discounting/retention, or constrained cross-stream mixing.
• Injection: distribute the branch update using write coefficients.

I expect this modular structure will help unify the existing recipes. Crucially, these are semantic boundaries, not mandatory execution boundaries. Specialized backends can fuse multiple components.