Buckets:
| <meta charset="utf-8" /><meta name="hf:doc:metadata" content="{"title":"Tensor parallelism for inference","local":"tensor-parallelism-for-inference","sections":[{"title":"Partitioning a model","local":"partitioning-a-model","sections":[],"depth":2},{"title":"Partitioning strategies","local":"partitioning-strategies","sections":[{"title":"Packed strategies","local":"packed-strategies","sections":[],"depth":3}],"depth":2},{"title":"Custom partitioning strategies","local":"custom-partitioning-strategies","sections":[],"depth":2},{"title":"Benchmarks","local":"benchmarks","sections":[],"depth":2},{"title":"Design implementation","local":"design-implementation","sections":[{"title":"DeviceMesh","local":"devicemesh","sections":[],"depth":3},{"title":"DTensor","local":"dtensor","sections":[],"depth":3}],"depth":2},{"title":"Resources","local":"resources","sections":[],"depth":2}],"depth":1}"/> | |
| <link href="/docs/transformers/main/en/_app/immutable/entry/start._QoxNuaL.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/CDai3N9n.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/BhmhGQv5.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/entry/app.BAbB5buu.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/qMuWdE-Y.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/CfMEBQAN.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/DOLBYqcr.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/DsnmJJEf.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/nodes/0.BbJ0x0zO.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/B76T2TVW.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/nodes/655.CVxjarbh.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/vlAHDplH.js" rel="modulepreload"> | |
| <link href="/docs/transformers/main/en/_app/immutable/chunks/BpA0a5UK.js" rel="modulepreload"> | |
| <!--1skuau8--><meta name="hf:doc:metadata" content="{"title":"Tensor parallelism for inference","local":"tensor-parallelism-for-inference","sections":[{"title":"Partitioning a model","local":"partitioning-a-model","sections":[],"depth":2},{"title":"Partitioning strategies","local":"partitioning-strategies","sections":[{"title":"Packed strategies","local":"packed-strategies","sections":[],"depth":3}],"depth":2},{"title":"Custom partitioning strategies","local":"custom-partitioning-strategies","sections":[],"depth":2},{"title":"Benchmarks","local":"benchmarks","sections":[],"depth":2},{"title":"Design implementation","local":"design-implementation","sections":[{"title":"DeviceMesh","local":"devicemesh","sections":[],"depth":3},{"title":"DTensor","local":"dtensor","sections":[],"depth":3}],"depth":2},{"title":"Resources","local":"resources","sections":[],"depth":2}],"depth":1}"/><!----> | |
| <link href="/docs/transformers/main/en/_app/immutable/assets/0.tn0RQdqM.css" rel="modulepreload"> <!--[--><!--[0--><!--[--><!--[0--><!--[--><p></p> <div class="items-center shrink-0 min-w-[100px] max-sm:min-w-[50px] justify-end ml-auto flex" style="float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"><div class="inline-flex rounded-md max-sm:rounded-sm"><button class="inline-flex items-center gap-1 h-7 max-sm:h-7 px-2 max-sm:px-1.5 text-sm font-medium text-gray-800 border border-r-0 rounded-l-md max-sm:rounded-l-sm border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-live="polite"><span class="inline-flex items-center justify-center rounded-md p-0.5 max-sm:p-0 hover:text-gray-800 dark:hover:text-gray-200"><svg class="sm:size-3.5 size-3" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----></span> <span>Copy page</span></button> <button class="inline-flex items-center justify-center w-6 max-sm:w-5 h-7 max-sm:h-7 disabled:pointer-events-none text-sm text-gray-500 hover:text-gray-700 dark:hover:text-white rounded-r-md max-sm:rounded-r-sm border border-l transition border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-haspopup="menu" aria-expanded="false" aria-label="Open copy menu"><svg class="transition-transform text-gray-400 overflow-visible sm:size-3.5 size-3 rotate-0" width="1em" height="1em" viewBox="0 0 12 7" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M1 1L6 6L11 1" stroke="currentColor"></path></svg><!----></button></div> <!--[-1--><!--]--></div><!----> <!--[0--><h1 class="relative group"><a id="tensor-parallelism-for-inference" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#tensor-parallelism-for-inference"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Tensor parallelism for inference</span></h1><!--]--><!----> <p><a href="./perf_train_gpu_many#tensor-parallelism">Tensor parallelism</a> slices a model layer into pieces so multiple hardware accelerators work on it simultaneously. This lets you run models that exceed a single GPU’s memory capacity and achieve higher throughput. You’ll need fast intra-node communication because GPUs exchange partial results at each layer.</p> <p>A model supports tensor parallelism if its config defines <code>base_model_tp_plan</code>. Check a loaded model with the <code>supports_tp_plan</code> property.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoModelForCausalLM | |
| model = AutoModelForCausalLM.from_pretrained(<span class="hljs-string">"meta-llama/Meta-Llama-3-8B-Instruct"</span>) | |
| <span class="hljs-built_in">print</span>(model.supports_tp_plan)<!----></pre></div><!----> <p>This guide covers enabling tensor parallelism in Transformers and the available partitioning strategies.</p> <!--[1--><h2 class="relative group"><a id="partitioning-a-model" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#partitioning-a-model"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Partitioning a model</span></h2><!--]--><!----> <p>Configure the number of tensor parallel devices with <code>tp_size</code> in <a href="/docs/transformers/pr_36708/en/expert_parallelism#transformers.DistributedConfig">DistributedConfig</a>.</p> <ul><li>Set <code>DistributedConfig(tp_size=N)</code> to use the model’s predefined plan.</li> <li>Define a manual <code>tp_plan</code> and pass it to <a href="/docs/transformers/pr_36708/en/expert_parallelism#transformers.DistributedConfig">DistributedConfig</a> along with <code>tp_size</code>.</li></ul> <p>You can also set <code>tp_plan="auto"</code> to request the predefined plan explicitly. When <code>tp_size</code> is omitted and a <code>tp_plan</code> is set, <code>tp_size</code> is derived from <code>WORLD_SIZE</code> divided by the other parallel sizes. Passing <code>tp_plan</code> directly to <a href="/docs/transformers/pr_36708/en/main_classes/model#transformers.PreTrainedModel.from_pretrained">from_pretrained()</a> is deprecated and will be removed in v5.18.</p> <div class="flex space-x-2 items-center my-1.5 mr-8 h-7 !pl-0 -mx-3 md:mx-0"><!--[--><div class="flex items-center border rounded-lg px-1.5 py-1 leading-none select-none text-smd border-gray-800 bg-black dark:bg-gray-700 text-white">auto plan</div><div class="flex items-center border rounded-lg px-1.5 py-1 leading-none select-none text-smd text-gray-500 cursor-pointer opacity-90 hover:text-gray-700 dark:hover:text-gray-200 hover:shadow-sm">manual plan</div><!--]--></div> <div class="language-select"><!--[0--><div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoModelForCausalLM, AutoTokenizer, DistributedConfig | |
| <span class="hljs-comment"># model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies</span> | |
| distributed_config = DistributedConfig(tp_size=<span class="hljs-number">4</span>) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| <span class="hljs-string">"meta-llama/Meta-Llama-3-8B-Instruct"</span>, | |
| dtype=torch.bfloat16, | |
| distributed_config=distributed_config, | |
| ) | |
| <span class="hljs-built_in">print</span>(model.tp_plan) | |
| tokenizer = AutoTokenizer.from_pretrained(<span class="hljs-string">"meta-llama/Meta-Llama-3-8B-Instruct"</span>) | |
| prompt = <span class="hljs-string">"Can I help"</span> | |
| inputs = tokenizer(prompt, return_tensors=<span class="hljs-string">"pt"</span>).input_ids.to(model.device) | |
| <span class="hljs-comment"># distributed run</span> | |
| outputs = model(inputs)<!----></pre></div><!----> <p>Launch the inference script with <a href="https://pytorch.org/docs/stable/elastic/run.html" rel="nofollow">torchrun</a>. Use one process per GPU.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-bash "><!---->torchrun --nproc-per-node 4 demo.py<!----></pre></div><!----><!----><!--]--><!----> <!--[-1--><!--]--><!----><!----></div><!----> <!--[1--><h2 class="relative group"><a id="partitioning-strategies" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#partitioning-strategies"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Partitioning strategies</span></h2><!--]--><!----> <p>The <code>ParallelInterface</code> class maps each strategy name you can use in a <code>tp_plan</code> to a configured strategy instance. You don’t interact with it directly to shard a model, but it’s the authoritative list of available names.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">class</span> <span class="hljs-title class_">ParallelInterface</span>(<span class="hljs-title class_ inherited__">GeneralInterface</span>): | |
| _global_mapping = { | |
| <span class="hljs-string">"embedding_rowwise"</span>: RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), | |
| <span class="hljs-string">"colwise_gather_output"</span>: ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), | |
| <span class="hljs-string">"colwise_rep"</span>: ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), | |
| <span class="hljs-string">"colwise"</span>: ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-<span class="hljs-number">1</span>)), | |
| <span class="hljs-string">"rowwise"</span>: RowwiseParallel(input_layouts=Shard(-<span class="hljs-number">1</span>), output_layouts=Replicate()), | |
| <span class="hljs-string">"rowwise_split_input"</span>: RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), | |
| <span class="hljs-string">"rowwise_rep"</span>: RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), | |
| <span class="hljs-string">"packed_colwise"</span>: PackedColwiseParallel(), | |
| <span class="hljs-string">"packed_rowwise"</span>: PackedRowwiseParallel(), | |
| <span class="hljs-string">"sequence_parallel"</span>: SequenceParallel(use_local_output=<span class="hljs-literal">True</span>), | |
| <span class="hljs-string">"grouped_gemm"</span>: MoEParamShard(Shard(<span class="hljs-number">0</span>), shards_expert_dim=<span class="hljs-literal">True</span>), | |
| <span class="hljs-string">"ep_router"</span>: EpRouterParallel(), | |
| <span class="hljs-string">"megamoe_router"</span>: RouterParallelMegaMoe(), | |
| <span class="hljs-string">"moe_tp_experts"</span>: MoeExpertsParallel(), | |
| <span class="hljs-string">"megamoe_experts"</span>: MoeTensorParalellMegaMoeExperts(), | |
| <span class="hljs-string">"moe_identity_expert"</span>: MoeIdentityParallel(), | |
| <span class="hljs-string">"replicated_with_grad_allreduce"</span>: ReplicatedWithGradAllReduce(), | |
| <span class="hljs-string">"mla_kv_a_proj"</span>: MlaKvAProjParallel(), | |
| <span class="hljs-string">"all_reduce"</span>: AllReduceParallel(), | |
| }<!----></pre></div><!----> <p>Every strategy is a subclass of <code>TensorParallelLayer</code> in <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py" rel="nofollow">distributed/tensor_parallel.py</a>. The two you’ll reach for most, <code>ColwiseParallel</code> and <code>RowwiseParallel</code>, take <code>input_layouts</code> and <code>output_layouts</code> placements, which is how one class covers several plan names. <code>colwise</code> leaves its output sharded on the last dim for a following <code>rowwise</code> layer to consume, while <code>colwise_gather_output</code> all-gathers it back to a full tensor.</p> <p>The table below describes each strategy.</p> <table><thead><tr><th>Strategy</th><th>Description</th></tr></thead><tbody><tr><td><code>ColwiseParallel</code></td><td>Shards a weight on its output-feature dim (<code>Shard(0)</code> for a 2D <code>nn.Linear</code> weight, <code>Shard(1)</code> for <code>nn.Embedding</code>) and shards a 1D bias. Redistributes the input to <code>Replicate()</code>, then places the output at <code>output_layouts</code>.</td></tr><tr><td><code>RowwiseParallel</code></td><td>Shards a weight on its input-feature dim (<code>Shard(-1)</code>, or <code>Shard(0)</code> for <code>nn.Embedding</code>) and replicates the bias, which is added after the reduction. Each rank’s forward produces a <code>Partial()</code> output holding its share of the sum, and <code>transform_output_post_forward</code> redistributes that to <code>output_layouts</code>. Reducing to <code>Replicate()</code> issues an all-reduce, and reducing to <code>Shard(1)</code> issues a reduce-scatter.</td></tr><tr><td><code>PackedColwiseParallel</code></td><td>A variant of <code>ColwiseParallel</code> for fused weights, for example <code>up_proj</code> and <code>gate_proj</code> packed into <code>gate_up_proj</code>. Use <code>split_factor</code> when a weight packs a number of layers other than two.</td></tr><tr><td><code>PackedRowwiseParallel</code></td><td>The row-wise counterpart, for weights packed along the final dim. Replicates 1D parameters.</td></tr><tr><td><code>SequenceParallel</code></td><td>Replicates the module’s parameters and shards its input on <code>sequence_dim</code> (defaults to <code>1</code>). Used for norms that operate per-token, such as <code>LayerNorm</code> and <code>RMSNorm</code>.</td></tr><tr><td><code>ReplicatedWithGradAllReduce</code></td><td>Replicates a parameter but all-reduces its gradient. Needed for norms that sit between a column-wise and a row-wise layer and normalize along a sharded axis, where each rank only sees its own heads.</td></tr><tr><td><code>AllReduceParallel</code></td><td>All-reduces a module’s <code>Partial()</code> forward output to <code>Replicate()</code>. Use it as a sync point for a module whose compute ends in a partial sum.</td></tr><tr><td><code>MlaKvAProjParallel</code></td><td>Splits the <code>kv_a_proj_with_mqa</code> output of DeepSeek-V2 style MLA attention and all-reduces the gradient of the RoPE half, which bypasses <code>kv_b_proj</code> and would otherwise keep a partial gradient. Requires <code>qk_rope_head_dim</code> in the model config.</td></tr><tr><td><code>MoEParamShard</code></td><td>Shards MoE expert weights on a given placement. Backs the <code>grouped_gemm</code> name, where <code>shards_expert_dim=True</code> also rewrites <code>module.num_experts</code> to the per-rank expert count.</td></tr><tr><td><code>EpRouterParallel</code></td><td>Masks router scores for non-local experts and remaps global expert IDs to local ones, so each rank runs only the experts it owns. Requires <code>num_experts</code> to be divisible by the mesh size.</td></tr><tr><td><code>RouterParallelMegaMoe</code></td><td>Router variant for DeepGEMM Mega MoE, which dispatches experts inside the kernel which wants the router output untouched.</td></tr><tr><td><code>MoeExpertsParallel</code></td><td>Tensor parallel MoE experts. All-reduces the expert output forward and adds the backward all-reduces for hidden states and routing weights.</td></tr><tr><td><code>MoeTensorParalellMegaMoeExperts</code></td><td>Inference-only experts layer for DeepGEMM Mega MoE. Skips the gradient syncs and passes the process group into the module so the kernel can set up its shared buffers on the first forward.</td></tr><tr><td><code>MoeIdentityParallel</code></td><td>Pre-divides the input of a zero or identity expert by the mesh size, cancelling the all-reduce that <code>moe_tp_experts</code> applies downstream.</td></tr></tbody></table> <!--[2--><h3 class="relative group"><a id="packed-strategies" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#packed-strategies"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Packed strategies</span></h3><!--]--><!----> <p>Weight packing combines multiple linear layers into a single, larger layer. The <code>PackedColwiseParallel</code> and <code>PackedRowwiseParallel</code> strategies shard packed weights correctly. Basic <code>ColwiseParallel</code> or <code>RowwiseParallel</code> strategies shard packed weights incorrectly.</p> <p>The example below packs <code>up_proj</code> and <code>gate_proj</code> into a single <code>gate_up_proj</code> module and requires the <code>packed_rowwise</code> strategy to shard <code>gate_up_proj</code>.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">class</span> <span class="hljs-title class_">Llama4TextExperts</span>(nn.Module): | |
| ... | |
| <span class="hljs-variable language_">self</span>.gate_up_proj = nn.Parameter(torch.zeros(<span class="hljs-variable language_">self</span>.num_experts, <span class="hljs-variable language_">self</span>.hidden_size, <span class="hljs-number">2</span> * <span class="hljs-variable language_">self</span>.expert_dim))<!----></pre></div><!----> <p>Use batch matrix multiplication in the <code>forward</code> pass to compute the output of the <code>gate_up_proj</code> module.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, hidden_states</span>): | |
| ... | |
| gate_up = torch.bmm(hidden_states, <span class="hljs-variable language_">self</span>.gate_up_proj) <span class="hljs-comment"># Compute the output of the gate_up_proj module</span> | |
| gate, up = gate_up.chunk(<span class="hljs-number">2</span>, dim=-<span class="hljs-number">1</span>) <span class="hljs-comment"># Split the output into gate and up</span><!----></pre></div><!----> <p>A plain <code>Shard</code> splits that dimension into contiguous blocks, so a rank would receive the tail of <code>gate</code> and the head of <code>up</code> instead of a slice of each. The packed strategies use <code>_StridedShard</code> with <code>split_factor</code> (<code>2</code> by default) to interleave the split, giving every rank a matching slice of both halves so <code>chunk</code> still lines up after sharding.</p> <!--[1--><h2 class="relative group"><a id="custom-partitioning-strategies" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#custom-partitioning-strategies"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Custom partitioning strategies</span></h2><!--]--><!----> <p>Inherit from <code>TensorParallelLayer</code> in <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py" rel="nofollow">distributed/tensor_parallel.py</a> to create a custom partitioning strategy. Override only the hooks your strategy needs, since every one has a no-op default.</p> <table><thead><tr><th>Hook</th><th>Purpose</th></tr></thead><tbody><tr><td><code>validate_param</code></td><td>Reject a parameter this strategy can’t shard, before any weights load.</td></tr><tr><td><code>shard_param</code></td><td>Replace one parameter with a <code>DTensor</code> placeholder so the loader knows which shard belongs to this rank.</td></tr><tr><td><code>transform_inputs_pre_forward</code></td><td>Redistribute the module’s inputs to the layout its forward expects.</td></tr><tr><td><code>context_around_forward</code></td><td>Wrap the forward in a context manager, for example to expose local tensors to a kernel.</td></tr><tr><td><code>transform_output_post_forward</code></td><td>Redistribute or reduce the module’s output.</td></tr><tr><td><code>should_use_local_tensors</code></td><td>Report that this module’s forward needs plain tensors rather than <code>DTensor</code>s.</td></tr><tr><td><code>install_forward</code></td><td>Replace <code>module.forward</code> outright. Override this only when the hooks above aren’t enough, as <code>ReplicatedWithGradAllReduce</code> does to register a backward hook.</td></tr></tbody></table> <p>The example below walks through a trimmed version of <code>ColwiseParallel</code>.</p> <ol><li><p>Inherit from <code>TensorParallelLayer</code> and store the placements the strategy works with. The base class defines no <code>__init__</code>, so there’s nothing to call <code>super()</code> on.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">class</span> <span class="hljs-title class_">ColwiseParallel</span>(<span class="hljs-title class_ inherited__">TensorParallelLayer</span>): | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, *, input_layouts=<span class="hljs-literal">None</span>, output_layouts=<span class="hljs-literal">None</span>, use_local_output: <span class="hljs-built_in">bool</span> = <span class="hljs-literal">True</span></span>): | |
| <span class="hljs-variable language_">self</span>.input_layouts = input_layouts <span class="hljs-keyword">or</span> Replicate() | |
| <span class="hljs-variable language_">self</span>.output_layouts = output_layouts <span class="hljs-keyword">if</span> output_layouts <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-literal">None</span> <span class="hljs-keyword">else</span> Shard(-<span class="hljs-number">1</span>) | |
| <span class="hljs-variable language_">self</span>.use_local_output = use_local_output<!----></pre></div><!----></li> <li><p>Implement <code>shard_param</code> to wrap one parameter as a <code>DTensor</code> placeholder. It runs on meta tensors, so <code>distribute_tensor</code> only builds metadata and moves no data. Pass <code>src_data_rank=None</code> because there’s no full tensor to scatter from yet.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">shard_param</span>(<span class="hljs-params">self, module, param, mesh</span>): | |
| meta = module._parameters.get(param) | |
| <span class="hljs-keyword">if</span> meta <span class="hljs-keyword">is</span> <span class="hljs-literal">None</span>: | |
| <span class="hljs-keyword">return</span> | |
| <span class="hljs-comment"># Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias</span> | |
| placement = Shard(<span class="hljs-number">1</span>) <span class="hljs-keyword">if</span> <span class="hljs-built_in">isinstance</span>(module, torch.nn.Embedding) <span class="hljs-keyword">else</span> Shard(meta.ndim - <span class="hljs-number">2</span>) | |
| module._parameters[param] = torch.nn.Parameter( | |
| distribute_tensor(meta, mesh, [placement], src_data_rank=<span class="hljs-literal">None</span>), | |
| requires_grad=meta.requires_grad, | |
| )<!----></pre></div><!----></li> <li><p>Implement the input and output transforms. <code>install_forward</code> calls them around the module’s original forward, so they only need to move tensors between layouts. Column-wise partitioning expects a replicated input and produces an output sharded on the last dim.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">transform_inputs_pre_forward</span>(<span class="hljs-params">self, module, args, kwargs, mesh</span>): | |
| x = args[<span class="hljs-number">0</span>] | |
| <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> <span class="hljs-built_in">isinstance</span>(x, DTensor): | |
| x = DTensor.from_local(x, mesh, [<span class="hljs-variable language_">self</span>.input_layouts], run_check=<span class="hljs-literal">False</span>) | |
| <span class="hljs-keyword">if</span> x.placements != (Replicate(),): | |
| x = x.redistribute(placements=[Replicate()]) | |
| <span class="hljs-keyword">return</span> (x,) + args[<span class="hljs-number">1</span>:], kwargs | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">transform_output_post_forward</span>(<span class="hljs-params">self, module, output, mesh</span>): | |
| <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> <span class="hljs-built_in">isinstance</span>(output, DTensor): | |
| output = DTensor.from_local(output, mesh, [Shard(-<span class="hljs-number">1</span>)], run_check=<span class="hljs-literal">False</span>) | |
| <span class="hljs-keyword">if</span> output.placements != (<span class="hljs-variable language_">self</span>.output_layouts,): | |
| output = output.redistribute(placements=[<span class="hljs-variable language_">self</span>.output_layouts]) | |
| <span class="hljs-keyword">return</span> output.to_local() <span class="hljs-keyword">if</span> <span class="hljs-variable language_">self</span>.use_local_output <span class="hljs-keyword">else</span> output<!----></pre></div><!----> <p>The shipped <code>ColwiseParallel</code> adds fast paths on top of this that skip the <code>DTensor</code> round trip for plain <code>nn.Linear</code> inference and for quantized modules that need local tensors. Read the source before copying it if your strategy needs the same treatment.</p></li> <li><p>Register the strategy so a <code>tp_plan</code> can name it. Registration takes a strategy instance, so the placements you pass to <code>__init__</code> are stored with the name. Note the <code>()</code> below.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoModelForCausalLM, DistributedConfig | |
| <span class="hljs-keyword">from</span> transformers.distributed.tensor_parallel <span class="hljs-keyword">import</span> ParallelInterface | |
| ParallelInterface.register(<span class="hljs-string">"colwise_custom"</span>, ColwiseParallel()) | |
| tp_plan = { | |
| <span class="hljs-string">"model.layers.*.self_attn.q_proj"</span>: <span class="hljs-string">"colwise_custom"</span>, | |
| ... | |
| } | |
| distributed_config = DistributedConfig(tp_size=<span class="hljs-number">4</span>, tp_plan=tp_plan) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| <span class="hljs-string">"meta-llama/Meta-Llama-3-8B-Instruct"</span>, | |
| dtype=torch.bfloat16, | |
| distributed_config=distributed_config, | |
| )<!----></pre></div><!----></li></ol> <!--[1--><h2 class="relative group"><a id="benchmarks" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#benchmarks"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Benchmarks</span></h2><!--]--><!----> <p>Tensor parallelism significantly speeds up inference, especially for large batch sizes or long sequences.</p> <p>This chart shows the expected speedup for a single forward pass on <a href="./model_doc/llama">Llama</a> with a sequence length of 512.</p> <div style="text-align: center"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/Meta-Llama-3-8B-Instruct%2C%20seqlen%20%3D%20512%2C%20python%2C%20w_%20compile.png"/></div> <!--[1--><h2 class="relative group"><a id="design-implementation" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#design-implementation"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Design implementation</span></h2><!--]--><!----> <p>Transformers implements tensor parallelism in a framework-agnostic way. It relies on <a href="https://docs.pytorch.org/tutorials/recipes/distributed_device_mesh.html" rel="nofollow">DeviceMesh</a> and <a href="https://docs.pytorch.org/docs/stable/distributed.tensor.html" rel="nofollow">DTensor</a> from <a href="https://docs.pytorch.org/tutorials/beginner/dist_overview.html" rel="nofollow">torch.distributed</a> to provide a simple, extensible interface.</p> <!--[2--><h3 class="relative group"><a id="devicemesh" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#devicemesh"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>DeviceMesh</span></h3><!--]--><!----> <p><code>DeviceMesh</code> creates a multi-dimensional grid of devices that communicate together. Different parallelization strategies require different communication patterns. Create a <code>DeviceMesh</code> with multiple sub-meshes to handle these patterns.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> torch.distributed.device_mesh <span class="hljs-keyword">import</span> init_device_mesh | |
| <span class="hljs-comment"># Create a 1D mesh of 4 accelerators</span> | |
| device_type = torch.accelerator.current_accelerator().<span class="hljs-built_in">type</span> | |
| device_mesh = init_device_mesh(device_type, (<span class="hljs-number">4</span>,), mesh_dim_names=[<span class="hljs-string">"tp"</span>])<!----></pre></div><!----> <p>Most <code>torch.distributed</code> parallelization strategies apply to the mesh itself or its sub-mesh. The mesh automatically handles communication patterns.</p> <!--[2--><h3 class="relative group"><a id="dtensor" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#dtensor"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>DTensor</span></h3><!--]--><!----> <p><code>DTensor</code> (Distributed Tensor) handles distributed logic on top of usual tensor operations. Model weights under tensor parallelism are stored as <code>DTensor</code>s, which is what lets a strategy describe communication as a change of layout instead of an explicit collective.</p> <p>The <code>placements</code> attribute tells PyTorch how a tensor is laid out across the devices in a <code>DeviceMesh</code>. It accepts the following values:</p> <ul><li><p><code>Shard(dimension)</code> splits a <code>DTensor</code> across a given dimension over the <code>DeviceMesh</code> it was constructed under. Column-wise partitioning shards the output-feature dim of the weight and the only dim of the bias.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->weight = DTensor.from_local(weight, device_mesh[<span class="hljs-string">"tp"</span>], placements=[Shard(<span class="hljs-number">0</span>)]) <span class="hljs-comment"># Shard the output features</span> | |
| bias = DTensor.from_local(bias, device_mesh[<span class="hljs-string">"tp"</span>], placements=[Shard(-<span class="hljs-number">1</span>)]) <span class="hljs-comment"># Shard the ONLY dimension</span><!----></pre></div><!----> <p>Row-wise partitioning shards the input-feature dim instead and replicates the bias, because the bias is added once after the reduction rather than on every rank.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->weight = DTensor.from_local(weight, device_mesh[<span class="hljs-string">"tp"</span>], placements=[Shard(-<span class="hljs-number">1</span>)]) <span class="hljs-comment"># Shard the input features</span> | |
| bias = DTensor.from_local(bias, device_mesh[<span class="hljs-string">"tp"</span>], placements=[Replicate()]) <span class="hljs-comment"># Replicate bias across all GPUs</span><!----></pre></div><!----></li> <li><p><code>Replicate()</code> replicates a <code>DTensor</code> across the <code>DeviceMesh</code>, creating a full copy of the tensor on each device.</p></li> <li><p><code>Partial()</code> marks a tensor as pending a reduction. A row-wise layer’s forward output is <code>Partial()</code>, and redistributing it to <code>Replicate()</code> is what issues the all-reduce.</p></li></ul> <!--[1--><h2 class="relative group"><a id="resources" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#resources"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Resources</span></h2><!--]--><!----> <ul><li><p>The <a href="https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism" rel="nofollow">Ultra-Scale Playbook</a> section on tensor parallelism provides more details.</p></li> <li><p>Check the <a href="./expert_parallelism">expert parallelism</a> guide if you’re using a mixture-of-experts (MoE) model. These models support tensor parallelism and expert parallelism.</p></li> <li><p>Read the <a href="https://huggingface.co/blog/qgallouedec/tp" rel="nofollow">Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand</a> blog post for a quick overview of tensor parallelism and learn how column and row parallel setups differ.</p></li> <li><p>See the <a href="./tensor_parallelism">Tensor parallelism</a> training guide to learn how to use it in a training setting.</p></li></ul> <a class="!text-gray-400 !no-underline text-sm flex items-center not-prose mt-4" href="https://github.com/huggingface/transformers/blob/main/docs/source/en/perf_infer_gpu_multi.md" target="_blank"><svg class="mr-1" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M31,16l-7,7l-1.41-1.41L28.17,16l-5.58-5.59L24,9l7,7z"></path><path d="M1,16l7-7l1.41,1.41L3.83,16l5.58,5.59L8,23l-7-7z"></path><path d="M12.419,25.484L17.639,6.552l1.932,0.518L14.351,26.002z"></path></svg><!----> <span><span class="underline">Update</span> on GitHub</span></a><!----> <p></p><!--]--><!----><!--]--><!--]--><!--]--> <!--[-1--><!--]--><!--]--> | |
| <script> | |
| { | |
| __sveltekit_5tx6kx = { | |
| base: "/docs/transformers/main/en", | |
| assets: "/docs/transformers/main/en" | |
| }; | |
| const element = document.currentScript.parentElement; | |
| Promise.all([ | |
| import("/docs/transformers/main/en/_app/immutable/entry/start._QoxNuaL.js"), | |
| import("/docs/transformers/main/en/_app/immutable/entry/app.BAbB5buu.js") | |
| ]).then(([kit, app]) => { | |
| kit.start(app, element, { | |
| node_ids: [0, 655], | |
| data: [null,null], | |
| form: null, | |
| error: null | |
| }); | |
| }); | |
| } | |
| </script> | |
Xet Storage Details
- Size:
- 60.3 kB
- Xet hash:
- faf640f7de025520a9f3e4514cb17b4a133bbb717feb98ca9d1b182092d321a1
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.