Llama 3.1 8B · Gate projection

32 independent transformer blocks › SwiGLU

Gate projection

learnedself.w1(x)

Construct features that will modulate the parallel up branch.

{\htmlData{light=g gate}{g}} = {\htmlData{light=x}{x}}W_{\mathrm{gate}}

in: $x$ [B batch, T positions, 4096 d]
out: $g$ [B batch, T positions, 14336 d_ff]
parameters: 58,720,256 · × 32 blocks = 1,879,048,192

Parameters

58,720,256 parameters per instance.

Exact per-instance accounting
Component	Calculation	Parameters
Weight matrix	$4{,}096 \times 14{,}336$	58,720,256
Per-instance total		58,720,256

This projection has no bias parameter.

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

Across 32 independent blocks: $58{,}720{,}256 \times 32 = 1{,}879{,}048{,}192$ parameters.

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

Codemodel.py 193–234 · 2 more files

model.py 208–219

        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
        self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )
        self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

model.py 193–206

class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        # custom dim factor multiplier
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

model.py 229–234

        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,
            multiple_of=args.multiple_of,
            ffn_dim_multiplier=args.ffn_dim_multiplier,
        )

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,
        ),

Sources4

Direct evidence

The MLP owns three bias-free projections. Its forward expression applies SiLU to w1(x), multiplies w3(x) element-wise and projects through w2.

mlp · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 208–219

Read code excerpt

        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
        self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )
        self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(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,
        ),

Direct evidence

The block passes four times model width; FeedForward truncates, applies the configured multiplier and rounds to the configured multiple.

mlp-width · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 193–234

Read code excerpt

class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        # custom dim factor multiplier
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

… lines 207–228 omitted …

        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,
            multiple_of=args.multiple_of,
            ffn_dim_multiplier=args.ffn_dim_multiplier,
        )

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)