nemo_automodel.components.distributed.blockdiag_cp.exchange

Bases: Function

All-gather a per-rank sequence shard along seq_dim over group.

Forward concatenates every rank’s shard [B, H, L, D] in rank order (sequential sharding), producing the full sequence [B, H, L*world, D]. Backward reduce-scatters the incoming gradient: each K/V position is read by every rank’s queries, so its gradient is the sum over ranks of the per-rank grad slice; reduce-scatter(SUM) hands each rank the summed gradient for the shard it owns. All shards have equal length L (the sequence is padded to a multiple of the CP world size before sharding).

nemo_automodel.components.distributed.blockdiag_cp.exchange._AllGatherSeqDiff.backward(
    ctx,
    grad_out: torch.Tensor
)

nemo_automodel.components.distributed.blockdiag_cp.exchange._AllGatherSeqDiff.forward(
    ctx,
    x: torch.Tensor,
    group,
    seq_dim: int
) -> torch.Tensor

Bases: Function

Uniform-size single-step neighbor exchange for node-local halo mode.

Every rank sends the last halo_size tokens of its chunk [B, C, L, D] to next_peer = (rank+1)%world and receives halo_size tokens from prev_peer = (rank-1)%world, send-before-recv, in one batch_isend_irecv. The caller slices the last recv_count of the received block as its real straddle (rank 0’s wraparound block and any over-send are ignored). Differentiable: backward is the reverse exchange (send grad to prev, recv from next).

nemo_automodel.components.distributed.blockdiag_cp.exchange._LeftHaloExchange.backward(
    ctx,
    grad_recv
)

nemo_automodel.components.distributed.blockdiag_cp.exchange._LeftHaloExchange.forward(
    ctx,
    x_local,
    group,
    halo_size,
    prev_peer,
    next_peer
)

Bases: Function

Differentiable multi-cast all-to-all-v K/V delivery.

Delivers each rank exactly the K/V range it attends to, zero-redundantly (a source token is sent once per rank that needs it). Forward maps the local chunk [B, C, L, D] to the received needed range [B, C, R_recv, D]. Backward scatter-ADDs the per-destination grads back to the source token (index_add), matching what the all-gather’s reduce_scatter would sum.

nemo_automodel.components.distributed.blockdiag_cp.exchange._NeededKVExchange.backward(
    ctx,
    grad_out
)

nemo_automodel.components.distributed.blockdiag_cp.exchange._NeededKVExchange.forward(
    ctx,
    x_local,
    group,
    in_splits,
    out_splits,
    send_index
)

Needed-only block-diagonal CP attention via the general all-to-all-v exchange.

Handles documents spanning >2 ranks (single long doc / unpacked long context) that the left-halo can’t, still zero-redundantly (vs the full all-gather fallback). Same argument contract as :func:_blockdiag_halo_attention.

Parameters:

Returns:

Attention output [B, Hq, L, D], or None when the kernel is

Needed-only block-diagonal CP attention via the left-halo exchange.

Each rank attends its local queries against [boundary-doc halo from rank-1] + [local chunk] instead of the full all-gathered sequence. The caller must fail fast when this returns None: the halo collective has already executed, so a rank-local all-gather fallback would fork the CP group’s collective order.

Parameters:

Returns:

Attention output [B, Hq, L, D], or None when the kernel is

Global left-halo exchange plan, computed from replicated doc_ids.

For each rank r the block-diagonal path needs the contiguous K/V range [s_first_r, real_end_r) = its boundary document’s straddle (owned by ranks < r) plus its own local chunk. back_r = r*local_len - s_first_r is the straddle length. If back_r <= local_len for ALL ranks, every straddle fits entirely in the LEFT neighbor (rank r-1), so a single neighbor p2p suffices (use_halo). If any document spans >2 ranks (back_r > local_len — e.g. one doc across the whole sequence) the caller falls back to the general all-to-all-v. Every rank computes this from the same replicated doc_ids and therefore agrees on use_halo and on all send/recv counts without communication.

Parameters:

Returns: dict

A plan dict with per-rank lists recv (tokens received from r-1 ==

Per-(src,dst) split sizes + local gather index for the general needed-only all-to-all-v exchange (the >2-rank / single-doc case the halo can’t cover).

Each rank r needs the contiguous global range [s_first_r, real_end_r). As a SENDER, this rank s sends to each dst r the intersection of its owned chunk [s*L,(s+1)*L) with r’s needed range (a token may go to several dsts, so the send buffer duplicates it once per destination). As a receiver, the pieces arrive in source order and concatenate into the contiguous needed range. All counts derive from the replicated plan, so every rank agrees.

Parameters:

Returns:

(in_splits, out_splits, send_index): per-destination send counts,