`nemo_rl.models.megatron.draft.dspark`#

DSpark heads and their checkpoint key contract.

This module defines the Markov and confidence heads. The DSpark draft provider in PR #3726 constructs both heads. This module does not connect them to a trainer. Attach the modules under markov_head and confidence_head to create the checkpoint key prefixes. Keep markov_w1, markov_w2, and proj. The markov_w2.weight shape is [draft_vocab_size, markov_rank] before TP sharding.

The pinned schema for deepseek-ai/dspark_qwen3_8b_block7 is in tests/unit/models/megatron/fixtures/dspark_qwen3_8b_block7_03326e50.json. See docs/design-docs/dspark-auxiliary-heads.md for inputs, TP behavior, and the limits of the checkpoint test.

Module Contents#

Classes#

`_CopyToTensorParallelRegion`	Keep the forward value local and SUM its gradient across TP ranks.
`DSparkMarkovHead`	Add a previous-token low-rank bias to caller-owned base logits.
`DSparkConfidenceHead`	Predict per-slot acceptance logits from draft and optional Markov state.

Data#

__all__

API#

class nemo_rl.models.megatron.draft.dspark._CopyToTensorParallelRegion#

Bases: torch.autograd.Function

Keep the forward value local and SUM its gradient across TP ranks.

static forward( ctx: Any, tensor: torch.Tensor, group: torch.distributed.ProcessGroup, ) → torch.Tensor#

static backward( ctx: Any, grad_output: torch.Tensor, ) → tuple[torch.Tensor, None]#

class nemo_rl.models.megatron.draft.dspark.DSparkMarkovHead( *, target_vocab_size: int, draft_vocab_size: int, markov_rank: int, draft_vocab_start_index: int = 0, draft_vocab_end_index: int | None = None, tensor_parallel_group: torch.distributed.ProcessGroup | None = None, device: torch.device | str | None = None, dtype: torch.dtype | None = None, )#

Bases: torch.nn.Module

Add a previous-token low-rank bias to caller-owned base logits.

markov_w1 remains a replicated target-vocabulary embedding. The output rows of markov_w2 cover either the full draft vocabulary or one explicit tensor-parallel draft-vocabulary shard. This module owns neither embeddings used by the draft backbone nor an LM head.

Initialization

embed_previous_tokens( *, previous_token_ids: torch.Tensor, slot_valid: torch.Tensor, reduce_across_tensor_parallel: bool, ) → torch.Tensor#

Embed previous-token IDs without invalid lookups or replicated TP drift.

Invalid slots use token ID zero for the lookup; consumers must still apply slot_valid to their outputs.

reduce_across_tensor_parallel must reflect the CONSUMER, not this head: pass True when the embeddings feed a vocabulary-sharded weight (each rank then holds a partial gradient that must be summed into the replicated markov_w1, as forward does for markov_w2); pass False when the consumer is TP-replicated (e.g. the confidence head’s proj) – every rank already computes the identical, complete gradient there, and summing would over-count it by the TP world size.

forward( base_logits: torch.Tensor, *, previous_token_ids: torch.Tensor, slot_valid: torch.Tensor, ) → torch.Tensor#: Return local-vocabulary logits with the Markov bias applied.

sharded_state_dict( prefix: str = '', sharded_offsets: tuple[tuple[int, int, int], ...] = (), metadata: dict[str, Any] | None = None, ) → dict[str, Any]#: Return Megatron checkpoint metadata for the TP-local output rows.

class nemo_rl.models.megatron.draft.dspark.DSparkConfidenceHead( *, hidden_size: int, markov_rank: int, with_markov: bool, device: torch.device | str | None = None, dtype: torch.dtype | None = None, )#