[Llama] Support attention_bias and mlp_bias config flags

[przemek-arch]

Motivation

HuggingFace Transformers added attention_bias and mlp_bias flags to LlamaConfig in PR #30031 (May 2024) to support llama-variant models that include bias terms in attention and MLP linear layers. Examples in the wild:

• IBM Calico models — used mlp_bias=True for stability and performance
• SpeakLeash Bielik v3 (speakleash/Bielik-4.5B-v3.0-Instruct, Apache 2.0) — uses both attention_bias=True and mlp_bias=True as a deliberate architectural choice for the Polish-language fine-tune
• Future llama variants will likely follow this pattern

MLC LLM's current llama_model.py hardcodes bias=False on the QKV, output, gate-up, and down projections, so converting bias-laden llama variants fails at weight load. The common workaround (using --model-type qwen2, which natively carries a fused QKV bias) silently drops 4 of 7 bias vectors per layer (o_proj.bias, gate_proj.bias, up_proj.bias, down_proj.bias) — for Bielik 4.5B v3 (60 layers) that is 240 of 420 HF bias vectors lost.

This PR adds optional bias support to MLC's Llama implementation, defaulting to False for full backward compatibility with existing Llama 2/3/3.1/3.2 conversions.

Changes

python/mlc_llm/model/llama/llama_model.py

• LlamaConfig dataclass — added attention_bias: bool = False and mlp_bias: bool = False fields
• LlamaAttention.__init__ — qkv_proj and o_proj now use bias=config.attention_bias
• LlamaFFN.__init__ — gate_up_proj and down_proj now use bias=config.mlp_bias
• LlamaDecoderLayer._set_tp — added a _set_bias helper that shards column-parallel biases (qkv, gate_up) along the same dim as the corresponding weight. Row-parallel biases (o_proj, down_proj) remain replicated across shards (added once after allreduce).

python/mlc_llm/loader/standard_loader.py

• make_standard_hf_loader — added add_gate_up_bias: bool = False and gate_up_bias_optional: bool = False parameters. The fusion logic mirrors the existing add_qkv_bias path: HF gate_proj.bias + up_proj.bias are concatenated along gate_up_concat_axis into MLC's fused gate_up_proj.bias.

python/mlc_llm/model/llama/llama_loader.py

• The default Llama loader now passes add_qkv_bias=True, qkv_bias_optional=True, add_gate_up_bias=True, gate_up_bias_optional=True. With *_optional=True, the loader picks up bias parameters when the model was instantiated with bias=True (so the bias keys are present in named_parameters), and skips them otherwise — backward compatibility with bias-free Llama configs is preserved.

The o_proj.bias and down_proj.bias parameters are passed through the existing 1:1 generic mapping at the end of the loader (no fusion needed, names match HF directly).

Testing

Backward compatibility — bit-identical regression

Converted unsloth/Llama-3.2-1B-Instruct (a bias-free Llama 3.2 mirror) twice with --quantization q4f16_1:

1. Once on main HEAD (baseline)
2. Once with this branch applied (patched)

Then byte-compared every params_shard_*.bin:

22 / 22 shards bit-identical
tensor-cache.json identical

The attention_bias=False, mlp_bias=False default path is provably unchanged.

Forward path — Bielik 4.5B v3 conversion

Converted speakleash/Bielik-4.5B-v3.0-Instruct (60 layers, attention_bias=True, mlp_bias=True, 32K context) with --quantization q4f16_1 --model-type llama:

4,757,260,288 parameters → 100 shards, 2.5 GB on disk, 4.505 bits/param

Tensor-cache analysis vs the previous qwen2-workaround conversion:

Conversion	Bias tensors
--model-type qwen2 (workaround)	60 fused QKV biases (= 180 HF biases)
--model-type llama (this PR)	60 fused QKV + 60 o_proj + 60 fused gate/up + 60 down_proj = all 420 HF biases

Compile to Metal succeeds; the runtime config log confirms attention_bias=True, mlp_bias=True propagating from config.json through to TVM compilation:

LlamaConfig(..., attention_bias=True, mlp_bias=True, ...)

Functional smoke test (3 prompts in Polish, including domain-specific legal Q&A) — all produced coherent, well-structured Polish output with appropriate legal terminology.

Empirical quality note — important caveat

We ran a head-to-head A/B against the qwen2-workaround conversion and an FP16/bfloat16 reference (HF transformers, Apple Metal MPS) on:

• 5 short discrete prompts (block-type detection, NER, structured JSON, terminology) — all three close, within q4 quantization noise
• 3 long-form prompts (rękojmia explanation, NDA draft, multistep labor-law reasoning)
• 5 multistep reasoning prompts spanning Polish labor / consumer / civil / criminal / commercial law, requiring specific article citations

Result on multistep reasoning (n=5): the qwen2-workaround conversion tracked the FP16 reference more faithfully than the patched fullbias conversion — qwen2 won 2/5 prompts on FP16-faithfulness, fullbias won 0/5, 3/5 ties. This PR brings the Llama architecture to parity with HF transformers (PR #30031), but on this particular model/quantization the extra biases did not improve — and on multistep reasoning slightly hurt — output quality.

We do not yet have a clean explanation. Plausible hypotheses:

1. The current group-quantization (q4f16_1, group_size=32, 4-bit weights, fp16 activations) may handle 1D bias vectors sub-optimally — biases have a different statistical distribution than 2D weight matrices and may benefit from a separate precision path.
2. q4f32_1 (4-bit weights, fp32 activations) recovered some quality on a single multistep prompt but did not generalize across the n=5 set, so the bottleneck is bias precision in q4 group quant rather than activation precision.
3. Bielik's training data has independent inconsistencies (the FP16 reference itself cited a wrong KPK article in one of our tests), so it is not a clean ground truth.

So this PR is offered as an architectural correctness contribution (Llama parity with HF transformers), not as a measured quality improvement for Bielik q4f16_1. Downstream users converting bias-laden models should benchmark output quality against alternative conversion paths for their specific model + quantization rather than assuming full bias coverage is strictly better. Use cases most likely to benefit:

• Bias-laden llama variants other than Bielik (Calico, future variants)
• Higher-precision quantizations (q5f16, q8f16, no-quant) where bias precision noise is smaller relative to signal
• Anyone writing tests against HF transformers parity

Happy to take reviewer guidance on whether this should land as-is, or whether bias precision in group quantization deserves a follow-up patch first. (No existing MLC issue/PR found for attention_bias/mlp_bias Llama support at the time of writing — if there is a tracking issue I missed, please point me to it.)

References

• HuggingFace Transformers PR #30031 — added attention_bias / mlp_bias to LlamaConfig: add mlp bias for llama models huggingface/transformers#30031
• Bielik v3 technical report (arXiv:2505.02550): https://arxiv.org/abs/2505.02550
• Bielik v3 model card (Apache 2.0): https://huggingface.co/speakleash/Bielik-4.5B-v3.0-Instruct
• Background work: Pelnora project (privilege-preserving legal AI for Polish lawyers, https://pelnora.app)

Side benefit

Enables MLC LLM browser-native deployment of:

• Bielik 4.5B v3 / 11B v3 (Polish, Apache 2.0)
• Future llama variants with bias support (Calico, etc.)

---

Diff summary

 python/mlc_llm/loader/standard_loader.py   | 20 ++++++++++++++++++++
 python/mlc_llm/model/llama/llama_loader.py |  4 ++++
 python/mlc_llm/model/llama/llama_model.py  | 31 ++++++++++++++++++++++++++----
 3 files changed, 51 insertions(+), 4 deletions(-)

All changes default to existing behavior; the only behavioral change is when config.attention_bias or config.mlp_bias are True (previously a weight-loading failure, now succeeds with proper bias handling).

Acknowledgements

• SpeakLeash team — for the Bielik model series that motivated this investigation
• Pelnora project — for funding and attributing the integration work
• Implementation co-authored with Claude Opus 4.7

[gemini-code-assist]

Code Review

This pull request introduces support for attention and MLP biases in the Llama model and its standard loader. It updates the configuration, model layers, and loading logic to handle these biases, including initial support for sharding them during Tensor Parallelism. However, a significant issue was identified in the Tensor Parallelism implementation: row-parallel biases for o_proj and down_proj are added before the all-reduce operation, which will cause the bias to be multiplied by the number of shards. This needs to be corrected by either scaling the bias or adding it after the reduction.

[gemini-code-assist]

The comment stating that row-parallel biases (for o_proj and down_proj) are "added once to the post-allreduce sum" is currently not reflected in the implementation.

In the current code, these biases are part of the nn.Linear layers, which means they are added to the local result on each shard before the op.ccl_allreduce(out, "sum") operation in _apply_residual (line 236). Consequently, the allreduce will sum the bias across all shards, effectively multiplying it by the number of shards (e.g., in a 2-GPU setup, the bias will be doubled).

To support Tensor Parallelism correctly with biases, you should either:

1. Divide the row-parallel biases by the number of shards during loading or in _set_tp.
2. Set bias=False for o_proj and down_proj and manually add the bias parameter after the allreduce in _apply_residual.

[przemek-arch]

Good catch — you're right, the o_proj / down_proj biases were being added before the all-reduce. Pushed 3356155.

I went with option 1, using the existing tp.shard_bias context manager (the same idiom gpt_neox and starcoder2 already use for their row-parallel dense / dense_4h_to_h biases): the attention and MLP calls in LlamaDecoderLayer.forward are now wrapped in tp.shard_bias(self.self_attn.o_proj, self.tensor_parallel_shards) / tp.shard_bias(self.mlp.down_proj, ...), so the row-parallel bias is divided by the shard count and the following ccl_allreduce(sum) restores it. qkv_proj / gate_up_proj are column-parallel and keep their ShardSingleDim shard_strategy.

I also made tp.shard_bias a no-op when linear.bias is None, so the default Llama path (attention_bias=False, mlp_bias=False) is unaffected at any shard count.

Backward compat re-verified: Llama-3.2-1B reconversion is bit-identical to the pre-patch baseline (22/22 params_shard_*.bin + tensor-cache.json); Bielik 4.5B v3 still captures all 420 HF biases and the single-device smoke test is unchanged. Also updated the _set_bias docstring that prompted the comment. Let me know if you'd prefer the explicit "separate nn.Parameter + post-allreduce add" form instead.