32 independent transformer blocks › Attention
Query projection
×learnedself.wq(x)
Construct query features from the normalized input.
- in
- [B batch, T positions, 4096 d]
- out
- [B batch, T positions, 4096 d]
- parameters
- 16,777,216 · × 32 blocks = 536,870,912
Parameters
16,777,216 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Weight matrix | 16,777,216 | |
| Per-instance total | 16,777,216 | |
This projection has no bias parameter.
The reference module holds the weight output by input, [out, in].
Across 32 independent blocks: parameters.
Derived from the configured matrix dimensions; checkpoint tensors have not been independently counted.
Codemodel.py 93–106, 153–154 · 2 more files
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
self.wq = ColumnParallelLinear(
args.dim,
args.n_heads * self.head_dim,
bias=False,
gather_output=False,
init_method=lambda x: x,
) bsz, seqlen, _ = x.shape
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
# 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)) 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)
# 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)) # 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)LLAMA3_VOCAB_SIZE = 128256def llama3_base_models() -> list[Model]:
return [
Model(
core_model_id=CoreModelId.llama3_8b,
description="Llama 3 8b model",
huggingface_repo="meta-llama/Llama-3-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": False,
},
pth_file_count=1,
),Sources3
Direct evidence
The query projection is bias-free and produces query-head count times head width features.
attn-q · 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
self.wq = ColumnParallelLinear(
args.dim,
args.n_heads * self.head_dim,
bias=False,
gather_output=False,
init_method=lambda x: x,
)… lines 107–152 omitted …
bsz, seqlen, _ = x.shape
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)Direct evidence
Meta’s original Llama 3 8B base catalogue entry supplies dimensions, head counts, vocabulary, MLP rounding settings, epsilon and unscaled RoPE settings.
config · meta-llama/llama-models / models/sku_list.py
Pinned revision 0e0b8c519242 ↗ · lines 19–207
Read code excerpt
LLAMA3_VOCAB_SIZE = 128256… lines 20–187 omitted …
def llama3_base_models() -> list[Model]:
return [
Model(
core_model_id=CoreModelId.llama3_8b,
description="Llama 3 8b model",
huggingface_repo="meta-llama/Llama-3-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": False,
},
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)