src/fragments/tp-plan.html

<pre><code class="language-python"># In the model's config (example: ERNIE 4.5-style decoder blocks)
base_model_tp_plan = {
    "layers.*.self_attn.q_proj": "colwise",
    "layers.*.self_attn.k_proj": "colwise",
    "layers.*.self_attn.v_proj": "colwise",
    "layers.*.self_attn.o_proj": "rowwise",
    "layers.*.mlp.gate_proj": "colwise",
    "layers.*.mlp.up_proj":   "colwise",
    "layers.*.mlp.down_proj": "rowwise",
}

# Runtime
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "your/model-or-local-checkpoint"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    tp_plan=base_model_tp_plan,   # <-- plan defined above
)
tok = AutoTokenizer.from_pretrained(model_id)
inputs = tok("Hello", return_tensors="pt").to(model.device)
out = model(**inputs)</code></pre>