nemo_automodel.components.speculative.target_cp

Target-side context parallelism for speculative-decoding draft training.

The draft trainers (EAGLE-3, DFlash, DSpark) shard the FROZEN target’s forward along the sequence and gather its captured hidden states / logits back to the full sequence before handing them to the draft. That flow is specific to speculative decoding — it needs no gradients through the target and no model-owned CP sharder — so it lives here rather than in the shared components/distributed/cp_utils surface.

Module Contents

Functions

Name	Description
`attach_cp_kv_gather_hooks`	Context-parallel self-attention for a FROZEN (forward-only) target.
`gather_cp_seq`	Gather context-parallel sharded `tensors` back to the full sequence.
`make_target_cp_ctx`	Build a context-parallel context for a frozen target forward.
`run_target_cp_forward_and_gather`	Run a frozen target under context parallelism and gather its outputs.

API

nemo_automodel.components.speculative.target_cp.attach_cp_kv_gather_hooks(
    model: torch.nn.Module,
    cp_mesh
) -> None

Context-parallel self-attention for a FROZEN (forward-only) target.

Torch’s context_parallel ring dispatch does not fire for a plain HuggingFace forward — q/k/v reach F.scaled_dot_product_attention as ordinary local tensors, so each rank silently attends only to its own sequence shard and every position past the first shard is wrong. Because the target is frozen (no backward through attention), the correct fix is simple: all-gather K/V across the cp group and attend the local Q against the full K/V with a global causal mask. The O(S^2) attention matrix stays sharded [S/cp, S] per rank (the memory win); only the O(S) K/V is replicated.

Assumes contiguous (non-load-balanced) sharding — rank r holds global positions [r*S_local, (r+1)*S_local) — which is what :func:make_target_cp_ctx produces (it disables load balancing). Q/K/V are [B, nH, S_local, D] at the SDPA call, so the sequence dim is 2.

nemo_automodel.components.speculative.target_cp.gather_cp_seq(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    tensors: typing.List[torch.Tensor],
    seq_dim: int,
    orig_len: int
)

Gather context-parallel sharded tensors back to the full sequence.

Inverse of the sharding done by :func:make_target_cp_ctx. Uses torch’s context_parallel_unshard with load_balancer=None (matching the load-balancing-disabled sharding) and slices the right-pad back off.

Parameters:

cp_mesh

DeviceMesh

The context-parallel device (sub)mesh used to shard.

tensors

List[torch.Tensor]

Local-shard tensors (e.g. captured aux hidden states, logits), each sharded to T/cp along seq_dim.

seq_dim

int

The sequence dimension to gather along.

orig_len

int

The pre-pad sequence length to slice back to.

Returns:

A list of full-sequence tensors of length orig_len along seq_dim.

nemo_automodel.components.speculative.target_cp.make_target_cp_ctx(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    input_ids,
    position_ids = None
)

Build a context-parallel context for a frozen target forward.

Shards input_ids (and position_ids) along the sequence dim across cp_mesh so the target’s self-attention runs as ring attention. Unlike :func:make_cp_batch_and_ctx, this does not require labels and is meant for the EAGLE-3 target wrapper, which gathers the aux/logits back to the full sequence (see :func:gather_cp_seq) before handing them to the draft.

Load balancing is disabled (_cp_options.enable_load_balance = False) so each rank holds a contiguous sequence chunk and the gather is a plain ordered concat (no round-robin un-permute). The sharding is thrown away right after the forward, so load balancing buys nothing here, and the ordered shard makes the gather deterministic. This is a process-global torch flag; the EAGLE-3 recipe is the only context-parallel user in its process.

The sequence is right-padded to a multiple of cp_size; the returned orig_len lets the caller slice the gathered outputs back down.

Parameters:

cp_mesh

DeviceMesh

The context-parallel device (sub)mesh.

input_ids

[B, T] token ids.

position_ids

Defaults to None

Optional [B, T] (or [1, T]) position ids; an arange is injected when omitted.

Returns:

(cp_ctx, sharded_input_ids, sharded_position_ids, orig_len). Enter

nemo_automodel.components.speculative.target_cp.run_target_cp_forward_and_gather(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    model: torch.nn.Module,
    input_ids: torch.Tensor,
    forward_kwargs: dict,
    collect: typing.Callable[[object], typing.List[torch.Tensor]],
    position_ids: typing.Optional[torch.Tensor] = None,
    filter_kwargs: bool = False
) -> tuple

Run a frozen target under context parallelism and gather its outputs.

Shards input_ids (and position_ids) along the sequence via :func:make_target_cp_ctx, runs model as ring attention with attention_mask=None (the self_attn hooks force is_causal), and — still inside the CP context, before any capture hooks are removed — gathers the tensors returned by collect(outputs) back to the full sequence with :func:gather_cp_seq (seq_dim=1, un-padded to orig_len).

Centralizes the gather-inside-context invariant and the seq_dim/orig_len contract shared by the eagle3/dflash/dspark target wrappers, so a fix lands in one place instead of drifting across three copies.

Parameters:

cp_mesh

DeviceMesh

The cp device submesh.

model

torch.nn.Module

The frozen target module.

input_ids

torch.Tensor

Full (unsharded) [B, T] token ids.

forward_kwargs

dict

Extra kwargs for model.forward. Must not include input_ids / attention_mask / position_ids — those are injected here (CP forces attention_mask=None).

collect

Callable[[object], List[torch.Tensor]]

callable(outputs) -> list[Tensor] selecting the tensors to gather; invoked inside the CP context after the forward, so it also sees any tensors captured by forward hooks.

position_ids

Optional[torch.Tensor]Defaults to None

Optional [B, T] / [1, T] positions; an arange is injected by :func:make_target_cp_ctx when omitted.

filter_kwargs

boolDefaults to False

Drop kwargs the model’s forward does not accept (via :func:filter_forward_kwargs) — needed for the VLM/MoE targets.

Returns: tuple

(outputs, gathered) — the raw model outputs and the list of