qwen3

class Qwen3MLP(nn.Module):
    """Qwen3 MLP with SwiGLU activation and tensor parallelism."""
def __init__(
    self,
    hidden_size: int,
    intermediate_size: int,
    hidden_act: str,
    quant_config: QuantizationConfig | None = None,
    bias: bool = False,
    prefix: str = "",
) -> None:
    super().__init__()
    self.gate_up_proj = MergedColumnParallelLinear(
        input_size=hidden_size,
        output_sizes=[intermediate_size] * 2,
        bias=bias,
        quant_config=quant_config,
        prefix=f"{prefix}.gate_up_proj",
    )
    self.down_proj = RowParallelLinear(
        input_size=intermediate_size,
        output_size=hidden_size,
        bias=bias,
        quant_config=quant_config,
        prefix=f"{prefix}.down_proj",
    )
    if hidden_act != "silu":
        raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.")
    self.act_fn = SiluAndMul()

def forward(self, x: torch.Tensor) -> torch.Tensor:
    x, _ = self.gate_up_proj(x)
    x = self.act_fn(x)
    x, _ = self.down_proj(x)
    return x</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3MLP.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x, _ = self.gate_up_proj(x)
    x = self.act_fn(x)
    x, _ = self.down_proj(x)
    return x
fastvideo.models.encoders.qwen3.Qwen3Attention
Qwen3Attention(config: Qwen3TextConfig, hidden_size: int, num_heads: int, num_kv_heads: int, rope_theta: float = 1000000.0, rope_scaling: dict[str, Any] | None = None, max_position_embeddings: int = 40960, quant_config: QuantizationConfig | None = None, bias: bool = False, prefix: str = '')
Qwen3 attention with QK-Norm and tensor parallelism.
Key difference from LLaMA: RMSNorm is applied to Q and K before attention.
Source
class Qwen3Attention(nn.Module):
    """Qwen3 attention with QK-Norm and tensor parallelism.
Key difference from LLaMA: RMSNorm is applied to Q and K before attention.
&quot;&quot;&quot;

def __init__(
    self,
    config: Qwen3TextConfig,
    hidden_size: int,
    num_heads: int,
    num_kv_heads: int,
    rope_theta: float = 1000000.0,
    rope_scaling: dict&#91;str, Any&#93; | None = None,
    max_position_embeddings: int = 40960,
    quant_config: QuantizationConfig | None = None,
    bias: bool = False,
    prefix: str = &quot;&quot;,
) -&gt; None:
    super().__init__()
    self.hidden_size = hidden_size
    tp_size = get_tp_world_size()
    self.total_num_heads = num_heads
    assert self.total_num_heads % tp_size == 0
    self.num_heads = self.total_num_heads // tp_size
    self.total_num_kv_heads = num_kv_heads
    if self.total_num_kv_heads &gt;= tp_size:
        assert self.total_num_kv_heads % tp_size == 0
    else:
        assert tp_size % self.total_num_kv_heads == 0
    self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

    self.head_dim = getattr(config, &quot;head_dim&quot;, self.hidden_size // self.total_num_heads)
    self.rotary_dim = self.head_dim
    self.q_size = self.num_heads * self.head_dim
    self.kv_size = self.num_kv_heads * self.head_dim
    self.scaling = self.head_dim**-0.5
    self.rope_theta = rope_theta
    self.max_position_embeddings = max_position_embeddings

    self.qkv_proj = QKVParallelLinear(
        hidden_size=hidden_size,
        head_size=self.head_dim,
        total_num_heads=self.total_num_heads,
        total_num_kv_heads=self.total_num_kv_heads,
        bias=bias,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.qkv_proj&quot;,
    )
    self.o_proj = RowParallelLinear(
        input_size=self.total_num_heads * self.head_dim,
        output_size=hidden_size,
        bias=bias,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.o_proj&quot;,
    )

    rms_norm_eps = getattr(config, &quot;rms_norm_eps&quot;, 1e-6)
    self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
    self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

    self.rotary_emb = get_rope(
        self.head_dim,
        rotary_dim=self.rotary_dim,
        max_position=max_position_embeddings,
        base=int(rope_theta),
        rope_scaling=rope_scaling,
        is_neox_style=True,
    )

    self.attn = LocalAttention(
        self.num_heads,
        self.head_dim,
        self.num_kv_heads,
        softmax_scale=self.scaling,
        causal=True,
        supported_attention_backends=config._supported_attention_backends,
    )

def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -&gt; torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states)
    q, k, v = qkv.split(&#91;self.q_size, self.kv_size, self.kv_size&#93;, dim=-1)

    batch_size, seq_len = q.shape&#91;0&#93;, q.shape&#91;1&#93;
    q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
    k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
    v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

    q = self.q_norm(q)
    k = self.k_norm(k)

    q = q.reshape(batch_size, seq_len, -1)
    k = k.reshape(batch_size, seq_len, -1)

    q, k = self.rotary_emb(positions, q, k)

    q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
    k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

    if attention_mask is None:
        attn_output = self.attn(q, k, v)
    else:
        q_sdpa = q.transpose(1, 2)
        k_sdpa = k.transpose(1, 2)
        v_sdpa = v.transpose(1, 2)
        causal_mask = torch.ones(
            seq_len,
            seq_len,
            device=q.device,
            dtype=torch.bool,
        ).tril()
        key_mask = attention_mask.to(device=q.device, dtype=torch.bool)
        attn_mask = causal_mask&#91;None, None, :, :&#93; &amp; key_mask&#91;:, None, None, :&#93;
        attn_output = torch.nn.functional.scaled_dot_product_attention(
            q_sdpa,
            k_sdpa,
            v_sdpa,
            attn_mask=attn_mask,
            dropout_p=0.0,
            is_causal=False,
            scale=self.scaling,
            enable_gqa=self.num_heads != self.num_kv_heads,
        ).transpose(1, 2)

    attn_output = attn_output.reshape(batch_size, seq_len, -1)

    output, _ = self.o_proj(attn_output)
    return output</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3Attention.forward
forward(positions: torch.Tensor, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor
Source
def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
batch_size, seq_len = q.shape&#91;0&#93;, q.shape&#91;1&#93;
q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

q = self.q_norm(q)
k = self.k_norm(k)

q = q.reshape(batch_size, seq_len, -1)
k = k.reshape(batch_size, seq_len, -1)

q, k = self.rotary_emb(positions, q, k)

q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

if attention_mask is None:
    attn_output = self.attn(q, k, v)
else:
    q_sdpa = q.transpose(1, 2)
    k_sdpa = k.transpose(1, 2)
    v_sdpa = v.transpose(1, 2)
    causal_mask = torch.ones(
        seq_len,
        seq_len,
        device=q.device,
        dtype=torch.bool,
    ).tril()
    key_mask = attention_mask.to(device=q.device, dtype=torch.bool)
    attn_mask = causal_mask&#91;None, None, :, :&#93; &amp; key_mask&#91;:, None, None, :&#93;
    attn_output = torch.nn.functional.scaled_dot_product_attention(
        q_sdpa,
        k_sdpa,
        v_sdpa,
        attn_mask=attn_mask,
        dropout_p=0.0,
        is_causal=False,
        scale=self.scaling,
        enable_gqa=self.num_heads != self.num_kv_heads,
    ).transpose(1, 2)

attn_output = attn_output.reshape(batch_size, seq_len, -1)

output, _ = self.o_proj(attn_output)
return output</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3DecoderLayer
Qwen3DecoderLayer(config: Qwen3TextConfig, quant_config: QuantizationConfig | None = None, prefix: str = '')
Qwen3 transformer decoder layer.
Source
class Qwen3DecoderLayer(nn.Module):
    """Qwen3 transformer decoder layer."""
def __init__(
    self,
    config: Qwen3TextConfig,
    quant_config: QuantizationConfig | None = None,
    prefix: str = &quot;&quot;,
) -&gt; None:
    super().__init__()
    self.hidden_size = config.hidden_size
    rope_theta = getattr(config, &quot;rope_theta&quot;, 1000000.0)
    rope_scaling = getattr(config, &quot;rope_scaling&quot;, None)
    max_position_embeddings = getattr(config, &quot;max_position_embeddings&quot;, 40960)
    attention_bias = getattr(config, &quot;attention_bias&quot;, False)

    self.self_attn = Qwen3Attention(
        config=config,
        hidden_size=self.hidden_size,
        num_heads=config.num_attention_heads,
        num_kv_heads=getattr(config, &quot;num_key_value_heads&quot;, config.num_attention_heads),
        rope_theta=rope_theta,
        rope_scaling=rope_scaling,
        max_position_embeddings=max_position_embeddings,
        quant_config=quant_config,
        bias=attention_bias,
        prefix=f&quot;&#123;prefix&#125;.self_attn&quot;,
    )
    self.mlp = Qwen3MLP(
        hidden_size=self.hidden_size,
        intermediate_size=config.intermediate_size,
        hidden_act=config.hidden_act,
        quant_config=quant_config,
        bias=getattr(config, &quot;mlp_bias&quot;, False),
        prefix=f&quot;&#123;prefix&#125;.mlp&quot;,
    )
    self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
    self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    residual: torch.Tensor | None,
    attention_mask: torch.Tensor | None = None,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor&#93;:
    if residual is None:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
    else:
        hidden_states, residual = self.input_layernorm(hidden_states, residual)

    hidden_states = self.self_attn(
        positions=positions,
        hidden_states=hidden_states,
        attention_mask=attention_mask,
    )

    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
    hidden_states = self.mlp(hidden_states)
    return hidden_states, residual</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3DecoderLayer.forward
forward(positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, attention_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    residual: torch.Tensor | None,
    attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    if residual is None:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
    else:
        hidden_states, residual = self.input_layernorm(hidden_states, residual)
hidden_states = self.self_attn(
    positions=positions,
    hidden_states=hidden_states,
    attention_mask=attention_mask,
)

hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
hidden_states = self.mlp(hidden_states)
return hidden_states, residual</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3ForCausalLM
Qwen3ForCausalLM(config: Qwen3TextConfig)
Qwen3 causal language model for text encoding in diffusion models (e.g. Flux2 Klein).
Features:
- Tensor parallelism support
- FlashAttention/SDPA support via LocalAttention
- QK-Norm for better training stability
- output_hidden_states for Klein (layers 9, 18, 27)
Source
class Qwen3ForCausalLM(TextEncoder):
    """Qwen3 causal language model for text encoding in diffusion models (e.g. Flux2 Klein).
Features:
- Tensor parallelism support
- FlashAttention/SDPA support via LocalAttention
- QK-Norm for better training stability
- output_hidden_states for Klein (layers 9, 18, 27)
&quot;&quot;&quot;

supports_hf_from_pretrained = True

def __init__(self, config: Qwen3TextConfig) -&gt; None:
    super().__init__(config)

    self.config = config
    self.quant_config = getattr(config, &quot;quant_config&quot;, None)

    if getattr(config, &quot;lora_config&quot;, None) is not None:
        max_loras = getattr(config.lora_config, &quot;max_loras&quot;, 1)
        lora_vocab_size = getattr(config.lora_config, &quot;lora_extra_vocab_size&quot;, 1)
        lora_vocab = lora_vocab_size * max_loras
    else:
        lora_vocab = 0
    self.vocab_size = config.vocab_size + lora_vocab
    self.org_vocab_size = config.vocab_size

    self.embed_tokens = VocabParallelEmbedding(
        self.vocab_size,
        config.hidden_size,
        org_num_embeddings=config.vocab_size,
        quant_config=self.quant_config,
    )

    self.layers = nn.ModuleList(&#91;
        Qwen3DecoderLayer(
            config=config,
            quant_config=self.quant_config,
            prefix=f&quot;&#123;config.prefix&#125;.layers.&#123;i&#125;&quot;,
        ) for i in range(config.num_hidden_layers)
    &#93;)

    self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

@classmethod
def from_pretrained_local(
    cls,
    model_path: str,
    model_config: Qwen3TextConfig,
    dtype: torch.dtype,
    device: torch.device,
) -&gt; nn.Module:
    from transformers import AutoModel

    # FastVideo uses Qwen3 only as a text encoder. Loading the body avoids
    # materializing an unused LM head and full-vocabulary logits, including
    # for checkpoints whose metadata names Qwen3ForCausalLM.
    return AutoModel.from_pretrained(
        model_path,
        local_files_only=True,
        dtype=dtype,
        low_cpu_mem_usage=True,
    ).eval().to(device)

def get_input_embeddings(self, input_ids: torch.Tensor) -&gt; torch.Tensor:
    return self.embed_tokens(input_ids)

def forward(
    self,
    input_ids: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs: Any,
) -&gt; BaseEncoderOutput:
    output_hidden_states = (output_hidden_states
                            if output_hidden_states is not None else self.config.output_hidden_states)

    if inputs_embeds is not None:
        hidden_states = inputs_embeds
    else:
        assert input_ids is not None
        hidden_states = self.get_input_embeddings(input_ids)

    residual = None

    if position_ids is None:
        # Expand to &#91;batch_size, seq_len&#93;: the rotary layer flattens
        # positions to ``num_tokens`` and reshapes q/k to
        # ``(num_tokens, -1, head_dim)``. A bare &#91;1, seq_len&#93; only matches
        # ``num_tokens`` when batch_size == 1; for batched inputs it folds
        # the batch dim into the head dim and misaligns RoPE. Expanding to
        # ``batch_size * seq_len`` tokens keeps the layout correct.
        position_ids = torch.arange(
            0,
            hidden_states.shape&#91;1&#93;,
            device=hidden_states.device,
            dtype=torch.long,
        ).unsqueeze(0).expand(hidden_states.shape&#91;0&#93;, -1)

    all_hidden_states: tuple&#91;Any, ...&#93; | None = (() if output_hidden_states else None)

    for layer in self.layers:
        if all_hidden_states is not None:
            all_hidden_states += ((hidden_states, ) if residual is None else (hidden_states + residual, ))
        hidden_states, residual = layer(
            position_ids,
            hidden_states,
            residual,
            attention_mask=attention_mask,
        )

    hidden_states, _ = self.norm(hidden_states, residual)

    if all_hidden_states is not None:
        all_hidden_states += (hidden_states, )

    return BaseEncoderOutput(
        last_hidden_state=hidden_states,
        hidden_states=all_hidden_states,
    )

def load_weights(self, weights: Iterable&#91;tuple&#91;str, torch.Tensor&#93;&#93;) -&gt; set&#91;str&#93;:
    params_dict = dict(self.named_parameters())
    loaded_params: set&#91;str&#93; = set()
    stacked_params_mapping = self.config.arch_config.stacked_params_mapping
    # A fused destination is initialized either by one already-fused tensor
    # or after every split source projection has been loaded. Include
    # auxiliary quantization parameters (for example scale_weight) rather
    # than limiting completeness checks to weight/bias tensors.
    expected_stacked_shards = &#123;(name, shard_id)
                               for name in params_dict
                               for param_name, _, shard_id in stacked_params_mapping if param_name in name&#125;
    fused_param_names = &#123;name for name, _ in expected_stacked_shards&#125;
    loaded_stacked_shards: set&#91;tuple&#91;str, str | int&#93;&#93; = set()
    loaded_fused_params: set&#91;str&#93; = set()

    for name, loaded_weight in weights:
        if name.startswith(&quot;model.&quot;):
            name = name&#91;6:&#93;

        if &quot;rotary_emb.inv_freq&quot; in name:
            continue
        if &quot;rotary_emb.cos_cached&quot; in name or &quot;rotary_emb.sin_cached&quot; in name:
            continue

        if &quot;scale&quot; in name:
            kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
            if kv_scale_name is None:
                continue
            name = kv_scale_name

        matched_stacked_param = False
        for param_name, weight_name, shard_id in stacked_params_mapping:
            if weight_name not in name:
                continue
            matched_stacked_param = True
            target_name = name.replace(weight_name, param_name)

            if target_name.endswith(&quot;.bias&quot;) and target_name not in params_dict:
                break

            if target_name not in params_dict:
                break

            param = params_dict&#91;target_name&#93;
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            loaded_params.add(target_name)
            loaded_stacked_shards.add((target_name, shard_id))
            break

        if matched_stacked_param:
            continue

        if name.endswith(&quot;.bias&quot;) and name not in params_dict:
            continue

        if name not in params_dict:
            continue

        param = params_dict&#91;name&#93;
        if name in fused_param_names and name.endswith(&quot;.scale_weight&quot;):
            # Merged scale loaders interpret a missing shard id as shard 0.
            # An exact fused key is already a complete vector, so copy it
            # atomically and retain the default loader&#x27;s shape validation.
            default_weight_loader(param, loaded_weight)
        else:
            weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
            weight_loader(param, loaded_weight)
        loaded_params.add(name)
        if name in fused_param_names:
            loaded_fused_params.add(name)

    required_split_shards = &#123;(name, shard_id)
                             for name, shard_id in expected_stacked_shards if name not in loaded_fused_params&#125;
    missing_stacked_shards = required_split_shards - loaded_stacked_shards
    if missing_stacked_shards:
        formatted_missing = &quot;, &quot;.join(f&quot;&#123;name&#125;&#91;&#123;shard_id&#125;&#93;&quot; for name, shard_id in sorted(
            missing_stacked_shards,
            key=lambda item: (item&#91;0&#93;, str(item&#91;1&#93;)),
        ))
        raise ValueError(&quot;Missing required stacked checkpoint shards: &quot;
                         f&quot;&#123;formatted_missing&#125;&quot;)

    return loaded_params</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.from_pretrained_local
from_pretrained_local(model_path: str, model_config: Qwen3TextConfig, dtype: torch.dtype, device: torch.device) -> nn.Module
Source
@classmethod
def from_pretrained_local(
    cls,
    model_path: str,
    model_config: Qwen3TextConfig,
    dtype: torch.dtype,
    device: torch.device,
) -> nn.Module:
    from transformers import AutoModel
# FastVideo uses Qwen3 only as a text encoder. Loading the body avoids
# materializing an unused LM head and full-vocabulary logits, including
# for checkpoints whose metadata names Qwen3ForCausalLM.
return AutoModel.from_pretrained(
    model_path,
    local_files_only=True,
    dtype=dtype,
    low_cpu_mem_usage=True,
).eval().to(device)</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.get_input_embeddings
get_input_embeddings(input_ids: torch.Tensor) -> torch.Tensor
Source
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
    return self.embed_tokens(input_ids)
fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.forward
forward(input_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs: Any) -> BaseEncoderOutput
Source
def forward(
    self,
    input_ids: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs: Any,
) -> BaseEncoderOutput:
    output_hidden_states = (output_hidden_states
                            if output_hidden_states is not None else self.config.output_hidden_states)
if inputs_embeds is not None:
    hidden_states = inputs_embeds
else:
    assert input_ids is not None
    hidden_states = self.get_input_embeddings(input_ids)

residual = None

if position_ids is None:
    # Expand to &#91;batch_size, seq_len&#93;: the rotary layer flattens
    # positions to ``num_tokens`` and reshapes q/k to
    # ``(num_tokens, -1, head_dim)``. A bare &#91;1, seq_len&#93; only matches
    # ``num_tokens`` when batch_size == 1; for batched inputs it folds
    # the batch dim into the head dim and misaligns RoPE. Expanding to
    # ``batch_size * seq_len`` tokens keeps the layout correct.
    position_ids = torch.arange(
        0,
        hidden_states.shape&#91;1&#93;,
        device=hidden_states.device,
        dtype=torch.long,
    ).unsqueeze(0).expand(hidden_states.shape&#91;0&#93;, -1)

all_hidden_states: tuple&#91;Any, ...&#93; | None = (() if output_hidden_states else None)

for layer in self.layers:
    if all_hidden_states is not None:
        all_hidden_states += ((hidden_states, ) if residual is None else (hidden_states + residual, ))
    hidden_states, residual = layer(
        position_ids,
        hidden_states,
        residual,
        attention_mask=attention_mask,
    )

hidden_states, _ = self.norm(hidden_states, residual)

if all_hidden_states is not None:
    all_hidden_states += (hidden_states, )

return BaseEncoderOutput(
    last_hidden_state=hidden_states,
    hidden_states=all_hidden_states,
)</code></pre></details>
fastvideo.models.encoders.qwen3.Qwen3ForCausalLM.load_weights
load_weights(weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]
Source
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
    params_dict = dict(self.named_parameters())
    loaded_params: set[str] = set()
    stacked_params_mapping = self.config.arch_config.stacked_params_mapping
    # A fused destination is initialized either by one already-fused tensor
    # or after every split source projection has been loaded. Include
    # auxiliary quantization parameters (for example scale_weight) rather
    # than limiting completeness checks to weight/bias tensors.
    expected_stacked_shards = {(name, shard_id)
                               for name in params_dict
                               for param_name, _, shard_id in stacked_params_mapping if param_name in name}
    fused_param_names = {name for name, _ in expected_stacked_shards}
    loaded_stacked_shards: set[tuple[str, str | int]] = set()
    loaded_fused_params: set[str] = set()
for name, loaded_weight in weights:
    if name.startswith(&quot;model.&quot;):
        name = name&#91;6:&#93;

    if &quot;rotary_emb.inv_freq&quot; in name:
        continue
    if &quot;rotary_emb.cos_cached&quot; in name or &quot;rotary_emb.sin_cached&quot; in name:
        continue

    if &quot;scale&quot; in name:
        kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
        if kv_scale_name is None:
            continue
        name = kv_scale_name

    matched_stacked_param = False
    for param_name, weight_name, shard_id in stacked_params_mapping:
        if weight_name not in name:
            continue
        matched_stacked_param = True
        target_name = name.replace(weight_name, param_name)

        if target_name.endswith(&quot;.bias&quot;) and target_name not in params_dict:
            break

        if target_name not in params_dict:
            break

        param = params_dict&#91;target_name&#93;
        weight_loader = param.weight_loader
        weight_loader(param, loaded_weight, shard_id)
        loaded_params.add(target_name)
        loaded_stacked_shards.add((target_name, shard_id))
        break

    if matched_stacked_param:
        continue

    if name.endswith(&quot;.bias&quot;) and name not in params_dict:
        continue

    if name not in params_dict:
        continue

    param = params_dict&#91;name&#93;
    if name in fused_param_names and name.endswith(&quot;.scale_weight&quot;):
        # Merged scale loaders interpret a missing shard id as shard 0.
        # An exact fused key is already a complete vector, so copy it
        # atomically and retain the default loader&#x27;s shape validation.
        default_weight_loader(param, loaded_weight)
    else:
        weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
        weight_loader(param, loaded_weight)
    loaded_params.add(name)
    if name in fused_param_names:
        loaded_fused_params.add(name)

required_split_shards = &#123;(name, shard_id)
                         for name, shard_id in expected_stacked_shards if name not in loaded_fused_params&#125;
missing_stacked_shards = required_split_shards - loaded_stacked_shards
if missing_stacked_shards:
    formatted_missing = &quot;, &quot;.join(f&quot;&#123;name&#125;&#91;&#123;shard_id&#125;&#93;&quot; for name, shard_id in sorted(
        missing_stacked_shards,
        key=lambda item: (item&#91;0&#93;, str(item&#91;1&#93;)),
    ))
    raise ValueError(&quot;Missing required stacked checkpoint shards: &quot;
                     f&quot;&#123;formatted_missing&#125;&quot;)

return loaded_params</code></pre></details>