[Llama] Support attention_bias and mlp_bias config flags

python/mlc_llm/loader/standard_loader.py

@@ -32,6 +32,8 @@ def make_standard_hf_loader(
     gate_up_names: Sequence[str] = ("gate_proj", "up_proj"),
     gate_up_concat_axis: int = 0,
     gate_up_target_name: str = "gate_up_proj",
+    add_gate_up_bias: bool = False,
+    gate_up_bias_optional: bool = False,
     include_qkv: bool = True,
     include_gate_up: bool = True,
     add_unused: Optional[Iterable[str]] = None,  # noqa: UP045
@@ -133,6 +135,24 @@ def huggingface(
                             ),
                         )
 
+                    if add_gate_up_bias and gate_up_names:
+                        mlc_bias_name = f"{mlp}.{gate_up_target_name}.bias"
+                        if (not gate_up_bias_optional) or mlc_bias_name in named_parameters:
+                            mlc_param = named_parameters[mlc_bias_name]
+                            mapping.add_mapping(
+                                mlc_bias_name,
+                                [
+                                    name_transform_fn(f"{mlp}.{name}.bias")
+                                    for name in gate_up_names
+                                ],
+                                functools.partial(
+                                    lambda gate, up, dtype: np.concatenate(
+                                        [gate, up], axis=gate_up_concat_axis
+                                    ).astype(dtype),
+                                    dtype=mlc_param.dtype,
+                                ),
+                            )
+
                 for unused_name in unused_names:
                     mapping.add_unused(name_transform_fn(f"{attn}.{unused_name}"))
 

python/mlc_llm/model/llama/llama_loader.py

@@ -18,6 +18,10 @@
 
 huggingface = make_standard_hf_loader(
     model_cls=LlamaForCausalLM,
+    add_qkv_bias=True,
+    qkv_bias_optional=True,
+    add_gate_up_bias=True,
+    gate_up_bias_optional=True,
     add_unused=["rotary_emb.inv_freq"],
 )
 

python/mlc_llm/model/llama/llama_model.py

@@ -37,6 +37,8 @@ class LlamaConfig(ConfigBase):
     prefill_chunk_size: int = 0
     num_key_value_heads: int = 0
     head_dim: int = 0
+    attention_bias: bool = False
+    mlp_bias: bool = False
     tensor_parallel_shards: int = 1
     pipeline_parallel_stages: int = 1
     max_batch_size: int = 1
@@ -115,9 +117,11 @@ def __init__(self, config: LlamaConfig):
         self.gate_up_proj = nn.Linear(
             in_features=config.hidden_size,
             out_features=2 * self.intermediate_size,
-            bias=False,
+            bias=config.mlp_bias,
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, config.hidden_size, bias=config.mlp_bias
         )
-        self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False)
 
     def forward(self, x: Tensor):
         concat_x1_x2 = self.gate_up_proj(x)
@@ -150,9 +154,11 @@ def __init__(self, config: LlamaConfig):
         self.qkv_proj = nn.Linear(
             in_features=config.hidden_size,
             out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim,
-            bias=False,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_q_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
         )
-        self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False)
 
     def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int):
         d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads
@@ -182,6 +188,13 @@ def _set_tp():
             def _set(layer, hint):
                 layer.weight.attrs["shard_strategy"] = hint
 
+            def _set_bias(layer, hint):
+                # Column-parallel layers shard their bias along the same dim as the weight's
+                # output dim. Row-parallel layers (o_proj, down_proj) replicate the bias
+                # across shards — it is added once to the post-allreduce sum, so no
+                # shard_strategy is required for them.
+                layer.bias.attrs["shard_strategy"] = hint
+
             hd = config.head_dim
             q = self.self_attn.num_q_heads * hd
             k = self.self_attn.num_kv_heads * hd
@@ -197,6 +210,16 @@ def _set(layer, hint):
                 tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0),
             )
             _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1))
+            if config.attention_bias:
+                _set_bias(
+                    self.self_attn.qkv_proj,
+                    tp.ShardSingleDim("_shard_qkv_bias", segs=[q, k, v], dim=0),
+                )
+            if config.mlp_bias:
+                _set_bias(
+                    self.mlp.gate_up_proj,
+                    tp.ShardSingleDim("_shard_mlp_up_bias", segs=[i, i], dim=0),
+                )
 
         self.tensor_parallel_shards = config.tensor_parallel_shards
         _set_tp()