nemo_rl.models.automodel.draft.draft_qwen3#

Module Contents#

Classes#

Qwen3DSparkAttention

Qwen3DSparkDecoderLayer

Qwen3DSparkModel

Functions#

apply_rotary_pos_emb

Data#

__all__

API#

nemo_rl.models.automodel.draft.draft_qwen3.apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1)#
class nemo_rl.models.automodel.draft.draft_qwen3.Qwen3DSparkAttention(config, layer_idx: int)#

Bases: torch.nn.Module

Initialization

forward(
hidden_states: torch.Tensor,
target_hidden_states: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor],
attention_mask: Optional[torch.Tensor],
past_key_values: Optional[transformers.cache_utils.Cache] = None,
cache_position: Optional[torch.LongTensor] = None,
**kwargs: typing_extensions.Unpack[transformers.models.qwen3.modeling_qwen3.FlashAttentionKwargs],
) → tuple[torch.Tensor, Optional[torch.Tensor]]#
class nemo_rl.models.automodel.draft.draft_qwen3.Qwen3DSparkDecoderLayer(config, layer_idx: int)#

Bases: transformers.models.qwen3.modeling_qwen3.GradientCheckpointingLayer

Initialization

forward(
target_hidden_states: Optional[torch.Tensor] = None,
hidden_states: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[transformers.cache_utils.Cache] = None,
output_attentions: Optional[bool] = False,
use_cache: Optional[bool] = False,
cache_position: Optional[torch.LongTensor] = None,
position_embeddings: Optional[typing_extensions.Tuple[torch.Tensor, torch.Tensor]] = None,
**kwargs: typing_extensions.Unpack[transformers.models.qwen3.modeling_qwen3.FlashAttentionKwargs],
) → typing_extensions.Tuple[torch.FloatTensor, Optional[typing_extensions.Tuple[torch.FloatTensor, torch.FloatTensor]]]#
class nemo_rl.models.automodel.draft.draft_qwen3.Qwen3DSparkModel(config)#

Bases: transformers.models.qwen3.modeling_qwen3.Qwen3PreTrainedModel

Initialization

_no_split_modules#

[‘Qwen3DSparkDecoderLayer’]

_get_d2t_target_ids() → torch.Tensor#

Absolute target token id for each draft index.

The checkpoint’s d2t buffer stores OFFSETS (eagle3/speculators convention): target_id = draft_idx + d2t[draft_idx] — the same interpretation as vLLM’s arange(draft_vocab) + d2t scatter index.

_get_t2d_index() → torch.Tensor#

Target token id -> draft index, -1 when absent.

Built lazily so it reflects the checkpoint-loaded d2t buffer (from_pretrained fills buffers after __init__).

_apply(fn, recurse: bool = True)#

Keep the RoPE inv_freq buffer in fp32 across dtype casts.

model.to(bfloat16) (the training build path) would otherwise round inv_freq to bf16 and dephase RoPE with absolute position, eroding draft acceptance (see pin_rope_inv_freq_fp32).

initialize_embeddings_and_head(
*,
embed_tokens: torch.nn.Module,
lm_head: torch.nn.Module,
freeze: bool = True,
)#
set_embedding_head_trainable(trainable: bool)#
compute_logits(hidden_states: torch.Tensor) → torch.Tensor#
predict_confidence_step(
hidden_states: torch.Tensor,
prev_token_ids: Optional[torch.Tensor] = None,
) → Optional[torch.Tensor]#
sample_draft_tokens(
base_logits: torch.Tensor,
*,
first_prev_token_ids: torch.Tensor,
temperature: float = 0.0,
hidden_states: Optional[torch.Tensor] = None,
) → tuple[torch.Tensor, torch.Tensor]#
sample_draft_token_step(
base_logits: torch.Tensor,
*,
prev_token_ids: torch.Tensor,
temperature: float = 0.0,
hidden_states: Optional[torch.Tensor] = None,
) → tuple[torch.Tensor, torch.Tensor]#
_forward_backbone(
*,
position_ids: torch.LongTensor,
attention_mask: Optional[torch.Tensor] = None,
noise_embedding: Optional[torch.Tensor] = None,
target_hidden_states: Optional[torch.Tensor] = None,
past_key_values: Optional[transformers.cache_utils.Cache] = None,
use_cache: bool = False,
**kwargs,
) → torch.Tensor#
forward(
input_ids: torch.Tensor,
target_hidden_states: torch.Tensor,
loss_mask: torch.Tensor,
target_last_hidden_states: Optional[torch.Tensor] = None,
position_ids: Optional[torch.Tensor] = None,
seq_lens: Optional[torch.Tensor] = None,
doc_remaining: Optional[torch.Tensor] = None,
teacher_logits: Optional[torch.Tensor] = None,
anchor_generator: Optional[torch.Generator] = None,
) → nemo_rl.models.automodel.draft.common.DSparkForwardOutput#

Run one DSpark training forward.

Sequence packing (position_ids [B, S] per-document reset positions, seq_lens [B, max_docs], doc_remaining [B, S]) keeps every block inside its anchor’s document: the anchor’s first target must be in-document, the block’s context prefix and supervision are restricted to that document, and the draft’s RoPE uses the per-document positions.

nemo_rl.models.automodel.draft.draft_qwen3.__all__#

[‘Qwen3DSparkModel’, ‘Qwen3DSparkAttention’, ‘Qwen3DSparkDecoderLayer’]