Llama 3.1 8B · Key projection

32 independent transformer blocks › Attention

Key projection

learnedself.wk(x)

Construct key features from the normalized input.

{\htmlData{light=k}{{\htmlClass{vl-stream-k}{K}}}}={\htmlData{light=attention-input}{x}}W_K

in: $x$ [B batch, T positions, 4096 d]
out: $K$ [B batch, T positions, 1024 n_kv·d_h]
parameters: 4,194,304 · × 32 blocks = 134,217,728

Parameters

4,194,304 parameters per instance.

Exact per-instance accounting
Component	Calculation	Parameters
Weight matrix	4,096 × (8 × 128)	4,194,304
Per-instance total		4,194,304

Key/value width, derived from the configuration

$d_h=d/n_h=4{,}096/32=128$ config attn-k
$n_{kv}\,d_h=8\times 128=1{,}024$ config attn-k

This projection has no bias parameter.

The reference module holds the weight output by input, [out, in].

Across 32 independent blocks: $4{,}194{,}304 \times 32 = 134{,}217{,}728$ parameters.

Derived from the configured matrix dimensions; checkpoint tensors have not been independently counted.

Codemodel.py 93–154 · 2 more files

model.py 93–98

        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

model.py 107–113

        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )

model.py 153–154

        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

layers.py 258–262


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.output_size_per_partition, self.in_features))

layers.py 286–290

    def forward(self, input_: torch.Tensor) -> torch.Tensor:  # type: ignore
        # Set up backprop all-reduce.
        input_parallel = copy_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight, self.bias)

layers.py 345–349


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.out_features, self.input_size_per_partition))

layers.py 374–380

        # Set up backprop all-reduce.
        if self.input_is_parallel:
            input_parallel = input_
        else:
            input_parallel = scatter_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight)

sku_list.py 19

LLAMA3_VOCAB_SIZE = 128256

sku_list.py 231–248

        Model(
            core_model_id=CoreModelId.llama3_1_8b,
            description="Llama 3.1 8b model",
            huggingface_repo="meta-llama/Llama-3.1-8B",
            arch_args={
                "dim": 4096,
                "n_layers": 32,
                "n_heads": 32,
                "n_kv_heads": 8,
                "vocab_size": LLAMA3_VOCAB_SIZE,
                "ffn_dim_multiplier": 1.3,
                "multiple_of": 1024,
                "norm_eps": 1e-05,
                "rope_theta": 500000.0,
                "use_scaled_rope": True,
            },
            pth_file_count=1,
        ),

Sources3

Direct evidence

The key projection is bias-free and produces KV-head count times head width features.

attn-k · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 93–154

Read code excerpt

        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

… lines 99–106 omitted …

        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )

… lines 114–152 omitted …

        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

Direct evidence

Meta's Llama 3.1 8B base catalogue entry supplies dimensions, head counts, the vocabulary (LLAMA3_VOCAB_SIZE, line 19), MLP rounding, epsilon, rope_theta and use_scaled_rope: True.

config · meta-llama/llama-models / models/sku_list.py
Pinned revision 0e0b8c519242 ↗ · lines 19–248

Read code excerpt

LLAMA3_VOCAB_SIZE = 128256

… lines 20–230 omitted …

        Model(
            core_model_id=CoreModelId.llama3_1_8b,
            description="Llama 3.1 8b model",
            huggingface_repo="meta-llama/Llama-3.1-8B",
            arch_args={
                "dim": 4096,
                "n_layers": 32,
                "n_heads": 32,
                "n_kv_heads": 8,
                "vocab_size": LLAMA3_VOCAB_SIZE,
                "ffn_dim_multiplier": 1.3,
                "multiple_of": 1024,
                "norm_eps": 1e-05,
                "rope_theta": 500000.0,
                "use_scaled_rope": True,
            },
            pth_file_count=1,
        ),

Dependency operation context

FairScale 0.4.13 stores linear weights with output features first and calls F.linear. This is operation-context evidence; the Meta requirements leave the FairScale version unpinned.

linear-storage · facebookresearch/fairscale / fairscale/nn/model_parallel/layers.py
Pinned revision 50b06d251730 ↗ · lines 258–380

Read code excerpt


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.output_size_per_partition, self.in_features))

… lines 263–285 omitted …

    def forward(self, input_: torch.Tensor) -> torch.Tensor:  # type: ignore
        # Set up backprop all-reduce.
        input_parallel = copy_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight, self.bias)

… lines 291–344 omitted …


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.out_features, self.input_size_per_partition))

… lines 350–373 omitted …

        # Set up backprop all-reduce.
        if self.input_is_parallel:
            input_parallel = input_
        else:
            input_parallel = scatter_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight)