diff --git a/python/mlc_llm/loader/standard_loader.py b/python/mlc_llm/loader/standard_loader.py index 5444c1e6c8..8944985192 100644 --- a/python/mlc_llm/loader/standard_loader.py +++ b/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}")) diff --git a/python/mlc_llm/model/llama/llama_loader.py b/python/mlc_llm/model/llama/llama_loader.py index 6931593c11..a3734f91ad 100644 --- a/python/mlc_llm/model/llama/llama_loader.py +++ b/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"], ) diff --git a/python/mlc_llm/model/llama/llama_model.py b/python/mlc_llm/model/llama/llama_model.py index 4659a2497d..5a1330c146 100644 --- a/python/mlc_llm/model/llama/llama_model.py +++ b/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,14 @@ def _set_tp(): def _set(layer, hint): layer.weight.attrs["shard_strategy"] = hint + def _set_bias(layer, hint): + # Column-parallel layers (qkv_proj, gate_up_proj) shard their bias along the + # same dim as the weight's output dim. Row-parallel layers (o_proj, down_proj) + # keep a full (unsharded) bias and instead divide it by the shard count at + # forward time via tp.shard_bias, so the subsequent allreduce(sum) restores + # the original value — they don't get a shard_strategy here. + 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,14 +211,29 @@ 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() def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): - out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + # o_proj and down_proj are row-parallel: their bias must be divided by the shard + # count so the following ccl_allreduce(sum) in _apply_residual restores it (no-op + # when there is no bias or on a single device). + with tp.shard_bias(self.self_attn.o_proj, self.tensor_parallel_shards): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) hidden_states = self._apply_residual(out, residual=hidden_states) - out = self.mlp(self.post_attention_layernorm(hidden_states)) + with tp.shard_bias(self.mlp.down_proj, self.tensor_parallel_shards): + out = self.mlp(self.post_attention_layernorm(hidden_states)) hidden_states = self._apply_residual(out, residual=hidden_states) return hidden_states diff --git a/python/mlc_llm/support/tensor_parallel.py b/python/mlc_llm/support/tensor_parallel.py index 90ce7f7888..cff9eaa893 100644 --- a/python/mlc_llm/support/tensor_parallel.py +++ b/python/mlc_llm/support/tensor_parallel.py @@ -102,6 +102,10 @@ def shard_bias(linear: nn.Linear, tensor_parallel_shards: int): """ A context manager to shard the bias of a linear into `tensor_parallel_shards` shards. + Divides the bias by the shard count so that the subsequent ``ccl_allreduce(sum)`` + over the row-parallel output restores it to its original value. No-op when the + linear has no bias or when running on a single device. + Parameters ---------- @@ -112,7 +116,7 @@ def shard_bias(linear: nn.Linear, tensor_parallel_shards: int): The number of shards. """ original_bias = linear.bias - if tensor_parallel_shards > 1: + if tensor_parallel_shards > 1 and linear.bias is not None: linear.bias = linear.bias / tensor_parallel_shards yield linear.bias = original_bias