---
title: Train Your Large Model on Multiple GPUs with Tensor Parallelism
url: https://www.dataloco.com/en/train-your-large-model-on-multiple-gpus-with-tensor-parallelism
published: 2026-09-15T16:11:08+00:00
language: en
section: Developers
source: https://machinelearningmastery.com/train-your-large-model-on-multiple-gpus-with-tensor-parallelism/
publisher: Dataloco
---

# Train Your Large Model on Multiple GPUs with Tensor Parallelism

Tensor parallelism is a technique that supports model parallelism by dividing a tensor along a specific dimension. This method aims to distribute tensor computation across multiple devices with minimal communication overhead. It is particularly useful for models with large parameter tensors where a single matrix multiplication cannot be accommodated on one GPU.

An example of tensor parallelism can be found in matrix multiplication. In a column-wise tensor parallel approach, the weight tensor is divided into columns, enabling the matrix multiplication operation to produce sharded outputs that must be concatenated. For instance, a $3\times 4$ matrix multiplied by a $4\times 6$ matrix results in a $3\times 6$ matrix. This operation can be decomposed into smaller multiplications, allowing for memory savings by not requiring the entire weight matrix on a single device.

Alternatively, row-wise tensor parallelism involves sharding the weight tensor into rows and producing partial outputs that are summed element-wise. This allows the workload to be lighter than conducting full matrix multiplications, although it requires greater bandwidth for result communication.

Notably, tensor parallelism does not apply to every operation within deep learning models, as some functions such as activation functions and normalization layers may require different parallelization techniques. For operations that do not permit parallelization, computation must occur in their original forms.

Beyond memory efficiency, tensor parallelism allows for detailed control over computation and communication patterns. The method involves sharding matrix multiplications, thus providing the option to manage whether or not to combine results, which minimizes communication overhead.

In PyTorch, tensor parallelism is integrated within the distributed framework. The execution script is initiated using the torchrun command, similar to distributed data parallelism and other parallelism techniques. It is essential to set up the distributed environment and establish the device mesh, which serves as a high-level abstraction of the process group, essential for wrapping a model into a tensor parallel model.

Preparing a model for tensor parallelism entails recognizing the fully-qualified names of each module and submodule, identical to keys in the model's state_dict(). This information is used to construct a parallelization plan, a Python dictionary that links module names to ParallelStyle objects, facilitating the model's adaptation for parallel execution.