Hugging Face's logo

Buckets:
hf-doc-build
/
doc-dev

Files
xet
hf-doc-build/doc-dev / transformers /pr_33962 /en /perf_infer_gpu_multi.html
rtrm's picture
rtrm
download
raw
56.1 kB
<meta charset="utf-8" /><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;Distributed inference&quot;,&quot;local&quot;:&quot;distributed-inference&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Partitioning a model&quot;,&quot;local&quot;:&quot;partitioning-a-model&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Partitioning strategies&quot;,&quot;local&quot;:&quot;partitioning-strategies&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Packed strategies&quot;,&quot;local&quot;:&quot;packed-strategies&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Local strategies&quot;,&quot;local&quot;:&quot;local-strategies&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Custom partitioning strategies&quot;,&quot;local&quot;:&quot;custom-partitioning-strategies&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Benchmarks&quot;,&quot;local&quot;:&quot;benchmarks&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Design implementation&quot;,&quot;local&quot;:&quot;design-implementation&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;DeviceMesh&quot;,&quot;local&quot;:&quot;devicemesh&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;DTensor&quot;,&quot;local&quot;:&quot;dtensor&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2}],&quot;depth&quot;:1}">
 <link href="/docs/transformers/pr_33962/en/_app/immutable/assets/0.e3b0c442.css" rel="modulepreload">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/entry/start.1776ac26.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/scheduler.31fdf58d.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/singletons.73ad6261.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/index.252883d5.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/paths.7e4d6052.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/entry/app.d2b9a7d8.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/preload-helper.769e5991.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/index.2f76fdf0.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/nodes/0.29a76cb1.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/each.e59479a4.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/nodes/504.c703f62d.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/CopyLLMTxtMenu.9922688f.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/globals.7f7f1b26.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/IconCopy.ac192424.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/MermaidChart.svelte_svelte_type_style_lang.4975a63c.js">
 <link rel="modulepreload" href="/docs/transformers/pr_33962/en/_app/immutable/chunks/CodeBlock.ab12f8e1.js"><!-- HEAD_svelte-u9bgzb_START --><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;Distributed inference&quot;,&quot;local&quot;:&quot;distributed-inference&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Partitioning a model&quot;,&quot;local&quot;:&quot;partitioning-a-model&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Partitioning strategies&quot;,&quot;local&quot;:&quot;partitioning-strategies&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Packed strategies&quot;,&quot;local&quot;:&quot;packed-strategies&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;Local strategies&quot;,&quot;local&quot;:&quot;local-strategies&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Custom partitioning strategies&quot;,&quot;local&quot;:&quot;custom-partitioning-strategies&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Benchmarks&quot;,&quot;local&quot;:&quot;benchmarks&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Design implementation&quot;,&quot;local&quot;:&quot;design-implementation&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;DeviceMesh&quot;,&quot;local&quot;:&quot;devicemesh&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3},{&quot;title&quot;:&quot;DTensor&quot;,&quot;local&quot;:&quot;dtensor&quot;,&quot;sections&quot;:[],&quot;depth&quot;:3}],&quot;depth&quot;:2}],&quot;depth&quot;:1}"><!-- HEAD_svelte-u9bgzb_END --> <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 max-sm:gap-0.5 h-6 max-sm:h-5 px-2 max-sm:px-1.5 text-[11px] max-sm:text-[9px] 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"><svg class="w-3 h-3 max-sm:w-2.5 max-sm:h-2.5" 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-6 max-sm:h-5 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 w-3 h-3 max-sm:w-2.5 max-sm:h-2.5 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> </div> <h1 class="relative group"><a id="distributed-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="#distributed-inference"><span><svg class="" 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>Distributed inference</span></h1> <p data-svelte-h="svelte-1pjj6ol">When a model doesn’t fit on a single GPU, distributed inference with <a href="./perf_train_gpu_many#tensor-parallelism">tensor parallelism</a> can help. Tensor parallelism shards a model onto multiple accelerators (CUDA GPU, Intel XPU, etc.) and parallelizes computations such as matrix multiplication. It enables fitting larger model sizes into memory and is faster because each accelerator can process a tensor slice.</p> <p data-svelte-h="svelte-jzkds">However, tensor parallelism adds communication overhead and should be used on single machine setups with multiple accelerators to take advantage of fast intra-node communication. For multi-node training, it may be more efficient to use pipeline or data parallelism depending on your use case.</p> <blockquote class="tip" data-svelte-h="svelte-vec04o"><p>Refer to the <a href="https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism" rel="nofollow">Ultra-Scale Playbook</a> section on tensor parallelism to learn more.</p></blockquote> <p data-svelte-h="svelte-1dtx2x6">Check the list below for models that natively support tensor parallelism. Open a GitHub issue or pull request to add support for a model.</p> <details data-svelte-h="svelte-1w7okx8"><summary>Show supported models</summary> <ul><li><a href="./model_doc/cohere">Cohere</a> and <a href="./model_doc/cohere2">Cohere 2</a></li> <li><a href="./model_doc/gemma">Gemma</a> and <a href="./model_doc/gemma2">Gemma 2</a></li> <li><a href="./model_doc/glm">GLM</a></li> <li><a href="./model_doc/granite">Granite</a></li> <li><a href="./model_doc/llama">Llama</a></li> <li><a href="./model_doc/mistral">Mistral</a></li> <li><a href="./model_doc/mixtral">Mixtral</a></li> <li><a href="./model_doc/olmo">OLMo</a> and <a href="./model_doc/olmo2">OLMo2</a></li> <li><a href="./model_doc/phi">Phi</a> and <a href="./model_doc/phi3">Phi-3</a></li> <li><a href="./model_doc/qwen2">Qwen2</a>, <a href="./model_doc/qwen2_moe">Qwen2Moe</a>, and <a href="./model_doc/qwen2_5_vl">Qwen2-VL</a></li> <li><a href="./model_doc/starcoder2">Starcoder2</a></li></ul></details> <p data-svelte-h="svelte-5hwspe">This guide shows how to enable tensor parallelism with Transformers and different partitioning strategies.</p> <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 class="" 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 data-svelte-h="svelte-4k9sxp">Transformers supports tensor parallelism if a model has a <code>tp_plan</code>. Set <code>tp_plan=&quot;auto&quot;</code> to automatically use a tensor parallelism plan based on a model’s predefined configuration.</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 class="" 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=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> os
<span class="hljs-keyword">import</span> torch
<span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoModelForCausalLM, AutoTokenizer
<span class="hljs-comment"># model_id = &quot;meta-llama/Llama-4-Scout-17B-16E-Instruct&quot; # better to visualize all the possible strategies</span>
model_id = <span class="hljs-string">&quot;meta-llama/Meta-Llama-3-8B-Instruct&quot;</span> <span class="hljs-comment"># better for smaller number of GPUs</span>
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, tp_plan=<span class="hljs-string">&quot;auto&quot;</span>)
<span class="hljs-built_in">print</span>(model._tp_plan)
tokenizer = AutoTokenizer.from_pretrained(<span class="hljs-string">&quot;meta-llama/Meta-Llama-3-8B-Instruct&quot;</span>)
prompt = <span class="hljs-string">&quot;Can I help&quot;</span>
inputs = tokenizer(prompt, return_tensors=<span class="hljs-string">&quot;pt&quot;</span>).input_ids.to(model.device)
<span class="hljs-comment"># distributed run</span>
outputs = model(inputs)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-968uev">Launch the inference script above on <a href="https://pytorch.org/docs/stable/elastic/run.html" rel="nofollow">torchrun</a> with 4 processes 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 class="" 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=""><!-- HTML_TAG_START -->torchrun --nproc-per-node 4 demo.py<!-- HTML_TAG_END --></pre></div> <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 class="" 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 data-svelte-h="svelte-12qtzv2">All partitioning strategies are defined in the <code>ParallelInterface</code> class which maps a string to the strategy implementation. You don’t need to interact with this class directly since all the strategies are set with <code>tp_plan</code> in <a href="/docs/transformers/pr_33962/en/main_classes/model#transformers.PreTrainedModel.from_pretrained">from_pretrained()</a>, but it is useful for checking what strategies are available.</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 class="" 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=""><!-- HTML_TAG_START --><span class="hljs-keyword">class</span> <span class="hljs-title class_">ParallelInterface</span>(<span class="hljs-title class_ inherited__">MutableMapping</span>):
 <span class="hljs-string">&quot;&quot;&quot;
 Dict-like object keeping track of allowed attention functions. You can easily add a new attention function
 with a call to `register()`. If a model needs to locally overwrite an existing attention function, say `sdpa`,
 it needs to declare a new instance of this class inside the `modeling_&lt;model&gt;.py`, and declare it on that instance.
 &quot;&quot;&quot;</span>
 _global_mapping = {
 <span class="hljs-string">&quot;colwise&quot;</span>: ColwiseParallel(),
 <span class="hljs-string">&quot;rowwise&quot;</span>: RowwiseParallel(),
 <span class="hljs-string">&quot;colwise_rep&quot;</span>: ColwiseParallel(output_layouts=Replicate()),
 <span class="hljs-string">&quot;rowwise_rep&quot;</span>: RowwiseParallel(input_layouts=Replicate()),
 <span class="hljs-string">&quot;local_colwise&quot;</span>: ColwiseParallel(use_dtensor=<span class="hljs-literal">False</span>),
 <span class="hljs-string">&quot;local_rowwise&quot;</span>: RowwiseParallel(use_dtensor=<span class="hljs-literal">False</span>),
 <span class="hljs-string">&quot;local&quot;</span>: IsolatedParallel(),
 <span class="hljs-string">&quot;gather&quot;</span>: GatherParallel(),
 <span class="hljs-string">&quot;local_packed_rowwise&quot;</span>: PackedRowwiseParallel(use_dtensor=<span class="hljs-literal">False</span>),
 <span class="hljs-string">&quot;sequence_parallel&quot;</span>: SequenceParallel(),
 <span class="hljs-string">&quot;replicate&quot;</span>: ReplicateParallel(),
 }<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1tu5la0">Refer to the table below to learn more about each strategy.</p> <table data-svelte-h="svelte-3u90w6"><thead><tr><th>Strategy</th> <th>Description</th></tr></thead> <tbody><tr><td><code>ColwiseParallel</code></td> <td>Column-wise partitioning of weights and biases.</td></tr> <tr><td><code>RowwiseParallel</code></td> <td>Row-wise partitioning of weights and biases. Also supports partitioning <code>nn.Embedding</code> modules.</td></tr> <tr><td><code>SequenceParallel</code></td> <td>Sequence parallel implementation to support <code>LayerNorm</code> and <code>Dropout</code> layers. Also supports Python implementation of <a href="https://github.com/facebookresearch/llama/blob/main/llama/model.py#L34" rel="nofollow">RMSNorm</a>.</td></tr> <tr><td><code>PackedColwiseParallel</code></td> <td>Variant of <code>ColwiseParallel</code> to support packed weights (for example, packing <code>up_proj</code> and <code>gate_proj</code> together). Refer to the <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/tensor_parallel.py#L79-#L108" rel="nofollow">code</a> for more details.</td></tr> <tr><td><code>PackedRowwiseParallel</code></td> <td>Variant of <code>RowwiseParallel</code> to support packed weights (refer to the <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/tensor_parallel.py#L79-#L108" rel="nofollow">code</a> for more details).</td></tr> <tr><td><code>GatherParallel</code></td> <td>Gather outputs of the module across devices.</td></tr> <tr><td><code>IsolatedParallel</code></td> <td>Used for Experts in Mixture-of-Experts (MoE) layers to isolates module from other devices.</td></tr> <tr><td><code>ReplicateParallel</code></td> <td>Replicate modules across all devices to prevent <code>torch.distributed</code> APIs from breaking due to a partially sharded model.</td></tr></tbody></table> <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 class="" 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 data-svelte-h="svelte-dw1gnc">Weight packing packs multiple linear layers into a single, bigger layer. Packed strategies, <code>PackedColwiseParallel</code> and <code>PackedRowwiseParallel</code>, are used to shard packed weights. The more basic <code>ColwiseParallel</code> or <code>RowwiseParallel</code> will incorrectly shard the packed weights.</p> <p data-svelte-h="svelte-1swuft2">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>PackedRowwiseParallel</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 class="" 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=""><!-- HTML_TAG_START --><span class="hljs-keyword">class</span> <span class="hljs-title class_">Llama4TextExperts</span>(nn.Module):
 ...
 self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, <span class="hljs-number">2</span> * self.expert_dim))<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-ppwgaw">Batch matrix multiplication can be used 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 class="" 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=""><!-- HTML_TAG_START --><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, self.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><!-- HTML_TAG_END --></pre></div> <blockquote class="tip" data-svelte-h="svelte-ji3s4j"><p>Refer to <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/tensor_parallel.py#L79-#L108" rel="nofollow">this comment</a> for an visual representation of why <code>Packed*</code> needs to be used.</p></blockquote> <h3 class="relative group"><a id="local-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="#local-strategies"><span><svg class="" 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>Local strategies</span></h3> <p data-svelte-h="svelte-py6c6p">Local strategies (<code>local_colwise</code>, <code>local_rowwise</code>, <code>local_packed_rowwise</code>) don’t use <a href="https://docs.pytorch.org/docs/stable/distributed.tensor.html" rel="nofollow">DTensor</a> because it isn’t supported for some operations such as <a href="https://docs.pytorch.org/docs/stable/generated/torch.chunk.html" rel="nofollow">torch.chunk</a>. Instead, local strategies use the basic <a href="https://docs.pytorch.org/docs/stable/tensors.html" rel="nofollow">torch.Tensor</a> and performs some of the distributed logic manually.</p> <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 class="" 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 data-svelte-h="svelte-1y826iv">A custom partitioning strategy should inherit from <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/tensor_parallel.py" rel="nofollow"><code>TensorParallelLayer</code></a> and implement <code>partition_tensor</code>, <code>_prepare_input_fn</code> and <code>_prepare_output_fn</code>.</p> <p data-svelte-h="svelte-1ql3jzh">Then it needs to be registered in the <code>ParallelInterface</code> mapping so the dispatching logic can find it when specified in <code>tp_plan</code>.</p> <p data-svelte-h="svelte-okfrkp">The example below shows how to implement <code>ColwiseParallel</code> with this workflow.</p> <ol><li><p data-svelte-h="svelte-d16j6n">Inherit from <code>TensorParallelLayer</code>. In the <code>__init__</code> method, define <code>input_layouts</code> and <code>output_layouts</code> to describe how the input and output tensors should be placed on devices. The <code>desired_input_layouts</code> attribute is used to specify how the input <em>should</em> be placed on devices.</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 class="" 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=""><!-- HTML_TAG_START --><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-type">Optional</span>[Placement] = <span class="hljs-literal">None</span>, <span class="hljs-comment"># The input layout coming from the previous layer</span>
 output_layouts: <span class="hljs-type">Optional</span>[Placement] = <span class="hljs-literal">None</span>, <span class="hljs-comment"># The output layout we want to achieve</span>
 use_local_output: <span class="hljs-built_in">bool</span> = <span class="hljs-literal">True</span>, <span class="hljs-comment"># Whether to use local output or not</span>
 use_dtensor=<span class="hljs-literal">True</span>, <span class="hljs-comment"># Whether to use DTensor or not</span>
 </span>):
 self.input_layouts = (input_layouts <span class="hljs-keyword">or</span> Replicate(),) <span class="hljs-comment"># The input sharding coming from the previous layer</span>
 self.output_layouts = (output_layouts <span class="hljs-keyword">or</span> Shard(-<span class="hljs-number">1</span>),) <span class="hljs-comment"># Desired output sharding</span>
 self.desired_input_layouts = (Replicate(),) <span class="hljs-comment"># Desired input sharding, inputs should be replicated across GPUs</span>
 self.use_local_output = use_local_output
 self.use_dtensor = use_dtensor<!-- HTML_TAG_END --></pre></div></li> <li><p data-svelte-h="svelte-1tvq20x">Implement the <code>partition_tensor</code>, <code>_prepare_input_fn</code> and <code>_prepare_output_fn</code> methods.</p> <p data-svelte-h="svelte-io1t4t">The <code>partition_tensor</code> method partitions the tensor and fills <code>empty_param</code> with the partitioned tensor. Use the utility function <code>get_tensor_shard</code> to help you get the correct shard of the original parameter for a given rank and <code>get_packed_weights</code> to help with packed weights.</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 class="" 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=""><!-- HTML_TAG_START --><span class="hljs-keyword">def</span> <span class="hljs-title function_">partition_tensor</span>(<span class="hljs-params">
 self,
 param, <span class="hljs-comment"># Full tensor of the parameter</span>
 empty_param, <span class="hljs-comment"># Empty tensor of the parameter, will be filled with the partitioned tensor</span>
 param_type, <span class="hljs-comment"># Type of the parameter, `bias` or `weight`</span>
 param_casting_dtype, <span class="hljs-comment"># The type to cast the parameter to</span>
 to_contiguous, <span class="hljs-comment"># Whether to convert the tensor to a contiguous memory layout</span>
 rank, <span class="hljs-comment"># The rank of the current device</span>
 device_mesh, <span class="hljs-comment"># The device mesh</span>
</span>) -&gt; nn.Parameter: <span class="hljs-comment"># Return the partitioned parameter</span>
 ...<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-493bdp">The <code>_prepare_input_fn</code> and <code>_prepare_output_fn</code> methods are used in the <a href="https://docs.pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_forward_pre_hook.html" rel="nofollow">pre-forward</a> and <a href="https://docs.pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_forward_hook.html" rel="nofollow">forward</a> hooks. They redistribute the inputs and outputs to the desired layout as specified in the <code>__init__</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 class="" 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=""><!-- HTML_TAG_START --><span class="hljs-keyword">def</span> <span class="hljs-title function_">_prepare_input_fn</span>(<span class="hljs-params">input_layouts, desired_input_layouts, mod, inputs, device_mesh</span>):
 ...
 <span class="hljs-comment"># Do some custom logic, cast to DTensor etc.</span>
 ...
 <span class="hljs-keyword">return</span> inputs.redistribute(placements=desired_input_layouts, device_mesh=device_mesh)
<span class="hljs-keyword">def</span> <span class="hljs-title function_">_prepare_output_fn</span>(<span class="hljs-params">output_layouts, use_local_output, mod, outputs, device_mesh</span>):
 ...
 <span class="hljs-comment"># Do some custom logic, cast to DTensor etc.</span>
 ...
 <span class="hljs-keyword">return</span> outputs.redistribute(placements=output_layouts, device_mesh=device_mesh)<!-- HTML_TAG_END --></pre></div></li> <li><p data-svelte-h="svelte-fkwwr3">Register the strategy to <code>ParallelInterface</code> to enable it for use with <code>tp_plan</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 class="" 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=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers.integrations.tensor_parallel <span class="hljs-keyword">import</span> ParallelInterface
ParallelInterface.register_strategy(<span class="hljs-string">&quot;colwise_custom&quot;</span>, ColwiseParallel)
tp_plan = {
 <span class="hljs-string">&quot;model.layers.*.self_attn.q_proj&quot;</span>: <span class="hljs-string">&quot;colwise_custom&quot;</span>,
 ...
}
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, tp_plan=tp_plan)<!-- HTML_TAG_END --></pre></div></li></ol> <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 class="" 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 data-svelte-h="svelte-15hmkdk">Tensor parallelism can considerably speedup inference, especially for inputs with large batch sizes or long sequences.</p> <p data-svelte-h="svelte-18qebw6">Refer to the chart below for 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" data-svelte-h="svelte-1tp7cj2"><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> <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 class="" 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 data-svelte-h="svelte-1s3zgaq">The Transformers tensor parallelism implementation is framework-agnostic, but for specific implementations, we rely 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 and extensible interface.</p> <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 class="" 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 data-svelte-h="svelte-evldcy">Imagine <code>DeviceMesh</code> as a multi-dimensional grid of devices that communicate together. Different parallelization strategies require different types of communication patterns, so you can create a <code>DeviceMesh</code> with multiple sub-meshes.</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 class="" 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=""><!-- HTML_TAG_START --><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 GPUs</span>
device_mesh = init_device_mesh(<span class="hljs-string">&quot;cuda&quot;</span>, (<span class="hljs-number">4</span>,), mesh_dim_names=[<span class="hljs-string">&quot;tp&quot;</span>])<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-17a74sp">Most of the <code>torch.distributed</code> defined parallelization strategies can be applied to the mesh itself, or its sub-mesh, and it automatically handles the communication patterns.</p> <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 class="" 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 data-svelte-h="svelte-1pz9ttn"><code>DTensor</code> (Distributed Tensor) is a tensor subclass that handles the distributed logic on top of the usual tensor operations. Most of the model weights in tensor parallelism are stored as <code>DTensor</code>s.</p> <p data-svelte-h="svelte-axo0c9">The most important part of DTensor is the <code>placement</code> attribute because it tells PyTorch how a tensor is placed on the devices in <code>DeviceMesh</code>. The <code>placement</code> attribute can take the following values.</p> <ul><li><p data-svelte-h="svelte-165m90y"><code>Shard(dimension)</code> - Indicates how a <code>DTensor</code> is sharded across a given dimension, over the <code>DeviceMesh</code> it was constructed under. The example below demonstrates how to shard weights over different dimensions for column-wise partitioning.</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 class="" 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=""><!-- HTML_TAG_START -->weight = ...
weight = DTensor.from_local(weight, device_mesh[<span class="hljs-string">&quot;tp&quot;</span>], placements=[Shard(<span class="hljs-number">0</span>)]) <span class="hljs-comment"># Shard across the 1st (column-wise) dimension</span>
bias = ...
bias = DTensor.from_local(bias, device_mesh[<span class="hljs-string">&quot;tp&quot;</span>], placements=[Shard(-<span class="hljs-number">1</span>)]) <span class="hljs-comment"># Shard across the ONLY dimension</span><!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-2fow1t">This example demonstrates how to shard weights over different dimensions for row-wise partitioning.</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 class="" 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=""><!-- HTML_TAG_START -->weight = ...
weight = DTensor.from_local(weight, device_mesh[<span class="hljs-string">&quot;tp&quot;</span>], placements=[Shard(<span class="hljs-number">1</span>)]) <span class="hljs-comment"># Shard across the 2nd (row-wise) dimension</span>
bias = ...
bias = DTensor.from_local(bias, device_mesh[<span class="hljs-string">&quot;tp&quot;</span>], placements=[Replicate()]) <span class="hljs-comment"># Replicate bias across all GPUs</span><!-- HTML_TAG_END --></pre></div></li> <li><p data-svelte-h="svelte-ios0p0"><code>Replicate()</code> - Indicates a <code>DTensor</code> is replicated across the <code>DeviceMesh</code>. It only creates a full copy of the tensor on each device.</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 class="" 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=""><!-- HTML_TAG_START -->bias = ...
bias = DTensor.from_local(bias, device_mesh[<span class="hljs-string">&quot;tp&quot;</span>], placements=[Replicate()]) <span class="hljs-comment"># Replicate bias across all GPUs</span><!-- HTML_TAG_END --></pre></div></li> <li data-svelte-h="svelte-5v07kk"><p><code>Partial()</code> - Indicates a tensor is pending a reduction operation (not typically relevant for usage in Transformers).</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 data-svelte-h="svelte-zjs2n5"><span class="underline">Update</span> on GitHub</span></a> <p></p>
<script>
 {
 __sveltekit_1yuubte = {
 assets: "/docs/transformers/pr_33962/en",
 base: "/docs/transformers/pr_33962/en",
 env: {}
 };
 const element = document.currentScript.parentElement;
 const data = [null,null];
 Promise.all([
 import("/docs/transformers/pr_33962/en/_app/immutable/entry/start.1776ac26.js"),
 import("/docs/transformers/pr_33962/en/_app/immutable/entry/app.d2b9a7d8.js")
 ]).then(([kit, app]) => {
 kit.start(app, element, {
 node_ids: [0, 504],
 data,
 form: null,
 error: null
 });
 });
 }
 </script>

Xet Storage Details

Size:
56.1 kB
·
Xet hash:
5a9531ad0e6909a4f0b6dfe9f0d51f876d7aa73cbe0e41cd83a63ce31ef4cd40

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.