Tensor parallelism

Overview

Tensor parallelism is a way to train or run a neural network that is too large for one GPU by cutting the weight matrices inside each layer into pieces. Each GPU holds one piece, computes its share of the layer, and the partial results are combined with a fast communication step. Several devices therefore compute a single layer together, instead of each device computing different layers or different data.

It is most useful for the wide layers of transformers, where one matrix can be larger than the memory of a single device. Because the devices must exchange data inside every layer, tensor parallelism works best on a small group of GPUs joined by very fast links, typically inside one server. It is usually combined with data parallelism and pipeline parallelism. For tabular foundation models, it matters when the model grows large or when the attention over rows and columns becomes heavy.

What it is

Tensor parallelism is a form of model parallelism. In model parallelism, the model itself is divided among devices. Tensor parallelism divides it at the finest useful level: the individual tensors (weight matrices) within a layer.

Think of a large spreadsheet that has to be summed by hand. One person could do it, but slowly. If you cut the sheet into four vertical strips and give one to each of four people, each person sums their strip, and then the four subtotals are added. The work was done in parallel, and only the small subtotals had to be shared. Tensor parallelism does the same with the big matrix multiplications at the heart of a network.

Intra-layer versus inter-layer splitting

Tensor parallelism sends more data more often than the other two, which is why it needs the fastest links.

Typical group sizes

Tensor parallel groups are small. Common sizes are two, four, or eight devices, matching the number of GPUs in one server. Going beyond one server usually makes the communication too slow, so larger scale is achieved by adding the other kinds of parallelism.

Where it came from

Splitting matrix multiplications across processors is an old idea in scientific computing. Distributed linear algebra libraries have long cut large matrices into blocks so that many processors can work on them at once.

For deep learning, early model parallel work divided networks by layers across devices. As transformers grew, researchers needed a cleaner way to split the biggest layers. A 2019 paper on training multi-billion parameter language models with model parallelism, from a large GPU company's applied research group, described a simple scheme for transformers. It cut the feed-forward block and the attention block so that each needed only two communication points per layer in the forward pass. The scheme became widely known through the library built around it, Megatron-LM.

Later work extended the idea. Sequence parallelism splits the parts of a layer, such as normalization and dropout, that tensor parallelism leaves replicated. Two-dimensional and three-dimensional tensor parallel variants split matrices along more than one axis to reduce communication. Frameworks such as Colossal-AI, DeepSpeed, and PyTorch now offer tensor parallel features, and inference servers use the same idea to serve large models on several GPUs.

How it works

The key fact is that matrix multiplication can be split in two clean ways. Understanding both is enough to understand the whole method.

The basic matrix identity

A layer computes Y = X times W, where X is the input (rows are examples) and W is the weight matrix.

Column split. Cut W into column blocks: W = [W1, W2]. Then Y = [X times W1, X times W2]. Each device needs the full X and produces a slice of the output columns. No communication is needed to compute the pieces. The output is spread across devices.

Row split. Cut W into row blocks, and cut X into matching column blocks. Then Y = X1 times W1 + X2 times W2. Each device produces a partial result of full size, and the final Y is the sum of the partial results. An all-reduce performs the sum.

Everyday analogy: a column split is like giving each of two cooks half of the recipe's dishes to cook from the same ingredients. A row split is like giving each cook half of the ingredients, so each cook produces a partial sauce, and the sauces must be mixed at the end.

The two-matrix trick

Transformer feed-forward blocks have two linear layers with an activation in the middle: Y = GELU(X times A) times B. The scheme splits A by columns and B by rows.

1. Each device computes a slice of the hidden layer: GELU(X times A_i). Since A is split by columns, each device has a whole set of hidden units. The activation is element-wise, so it needs no communication.
2. Each device multiplies its hidden slice by its row block B_i. This gives a partial output of full size.
3. One all-reduce sums the partial outputs across devices.

The result is one all-reduce in the forward pass for the whole block, not one per matrix. The column split followed by a row split is the central design choice.

Splitting attention

Multi-head attention is easy to split because heads are independent.

1. Assign a subset of heads to each device.
2. Split the query, key, and value projections by columns, so each device produces only its own heads.
3. Each device computes attention for its heads and gets partial context vectors.
4. The output projection is split by rows, so each device multiplies its heads' output by a row block.
5. One all-reduce sums the partial outputs.

Again, one all-reduce per block in the forward pass. The number of heads must be divisible by the tensor parallel degree for an even split.

The backward pass

The backward pass mirrors the forward pass. Where the forward pass had an all-reduce, the backward pass has an identity, and where the forward pass had an identity, the backward pass has an all-reduce. In total, a transformer layer needs about two all-reduces in the forward pass and two in the backward pass, over activations of the size of the layer input.

The embedding and output layers

The vocabulary or input embedding table can also be split. For language models the table is split by rows, so each device owns part of the vocabulary. Lookups gather partial results, and the output layer can compute logits in slices. For tabular models, the analogue is the set of per-column or per-type embedding tables, which can be split in the same way if they are large.

Sequence parallelism as a companion

Operations such as layer normalization and dropout act on each token independently. Plain tensor parallelism keeps them replicated, so every device holds the full activations for these steps. Sequence parallelism splits those activations along the sequence axis, which saves memory. It turns each all-reduce into a pair: a reduce-scatter and an all-gather. The total data moved stays about the same, but activation memory drops.

Why fast links matter

Each layer sends activations to the other devices. The amount sent is proportional to batch size times sequence length times hidden size. With dozens or hundreds of layers, this happens very many times per step, and none of it can be easily hidden behind computation because the next step needs the result. That is why tensor parallelism is placed inside a node, connected by high-bandwidth links, and not across slower network links between nodes. See NCCL for the communication library used on GPUs.

Key concepts and terms

A worked example

A numeric walk-through of a row and column split

Let X be a single row [1, 2], and let W be a 2 by 4 matrix:

W = [[1, 0, 2, 1],
     [0, 1, 1, 3]]

The full result is Y = X times W = [11 + 20, 10 + 21, 12 + 21, 11 + 23] = [1, 2, 4, 7].

Column split across two devices. Device 1 holds columns 1 and 2 of W: [[1, 0], [0, 1]]. Its output is [1, 2]. Device 2 holds columns 3 and 4: [[2, 1], [1, 3]]. Its output is [12 + 21, 11 + 23] = [4, 7]. Joining the slices gives [1, 2, 4, 7]. The same answer, and no summation was needed.

Row split across two devices. Device 1 holds row 1 of W and the first element of X: it computes 1 times [1, 0, 2, 1] = [1, 0, 2, 1]. Device 2 holds row 2 and the second element of X: it computes 2 times [0, 1, 1, 3] = [0, 2, 2, 6]. The all-reduce sum is [1, 2, 4, 7]. Again the same answer, and this time a sum was needed.

A code sketch of a column and row parallel pair

The following sketch shows the logic in plain PyTorch, simulating the devices in one process. Real systems place each shard on its own GPU and use a collective library.

import torch

torch.manual_seed(0)
d_model, d_hidden, tp = 8, 16, 2   # tp = tensor parallel degree

x = torch.randn(4, d_model)
A = torch.randn(d_model, d_hidden)   # first linear layer
B = torch.randn(d_hidden, d_model)   # second linear layer

# Reference: no parallelism
ref = torch.nn.functional.gelu(x @ A) @ B

# Shard A by columns and B by rows
A_shards = A.chunk(tp, dim=1)
B_shards = B.chunk(tp, dim=0)

partials = []
for i in range(tp):
    h = torch.nn.functional.gelu(x @ A_shards[i])   # no communication
    partials.append(h @ B_shards[i])                # partial output

out = sum(partials)   # this sum plays the role of the all-reduce

print(torch.allclose(ref, out, atol=1e-5))   # True

The check confirms that the parallel result equals the reference. In a real run, the line that sums the partial results would be an all-reduce across GPUs.

Memory per device

Suppose a transformer layer has weights of total size S. With a tensor parallel degree of 4, each device holds about S divided by 4 for the split matrices. Parts that are not split, such as small bias vectors and norm parameters, stay replicated, but they are tiny.

Activations behave differently. The input to each block is the full size on every device, unless sequence parallelism is used. Thus tensor parallelism cuts weight memory almost in proportion to the degree, but activation memory only partly.

A configuration sketch

Parallel layouts are normally described by three degrees whose product equals the number of devices.

cluster:
  nodes: 4
  gpus_per_node: 8          # 32 devices in total
parallelism:
  tensor_parallel: 8        # inside one node
  pipeline_parallel: 2      # across two nodes
  data_parallel: 2          # 8 * 2 * 2 = 32
  sequence_parallel: true

The layout above keeps the chatty tensor parallel traffic inside each node, uses pipeline stages across nodes, and replicates the whole arrangement twice for data parallelism.

Following one transformer block step by step

It helps to trace one block with a tensor parallel degree of two. Call the devices D0 and D1.

1. The block input X is present in full on both devices.
2. Each device applies layer normalization to X. This is replicated work.
3. Each device projects X into queries, keys, and values for its own half of the heads.
4. Each device runs attention for its own heads. No messages are sent.
5. Each device multiplies its attention output by its row block of the output projection.
6. An all-reduce sums the two partial results. Now both devices hold the same full output.
7. Each device adds the residual connection and applies the second layer normalization.
8. Each device multiplies by its column block of the first feed-forward matrix and applies the activation.
9. Each device multiplies by its row block of the second feed-forward matrix.
10. A second all-reduce sums the partial results.
11. Both devices add the residual and pass the result to the next block.

Two all-reduces per block, and everything else is local. This regularity is what makes the method easy to reason about and to optimize.

Estimating communication volume

The volume of each all-reduce is set by the size of the activation tensor.

• Per all-reduce, the tensor holds batch size times sequence length times hidden size values.
• Per block in the forward pass there are two all-reduces.
• Per block in the backward pass there are two more.
• A ring all-reduce moves close to twice the tensor size per device.

Multiply by the number of blocks and by the bytes per value to get the traffic per step. Because this scales with the activation size, halving the precision halves the traffic. It also explains why long contexts and large batches raise the communication load, while a larger tensor parallel degree does not shrink it. The volume per device stays about the same, and the compute per device falls, so the share of time spent communicating grows with the degree.

Choosing the degree: a worked reasoning

Imagine a layer whose compute takes 8 milliseconds on one device and whose all-reduce takes 1 millisecond on a fast link.

These numbers are illustrative, not measurements. They show the pattern: each doubling of the degree halves compute but leaves communication roughly unchanged, so the efficiency falls. The right degree is the smallest one that fits the model in memory or meets a latency goal.

Two-dimensional and higher variants

The simple scheme splits each matrix along one axis. Two-dimensional schemes split along both axes, arranging devices in a grid. Each device holds a block of the matrix, and multiplication proceeds in rounds in which blocks are broadcast along rows and columns of the grid. The benefit is that communication per device can shrink as the grid grows. The cost is more complex code and more rounds of messages. Higher-dimensional variants extend the same idea. These are mainly of interest at very large scale.

Tensor parallelism in inference servers

Serving a large model is a different problem from training it.

• There are no gradients or optimizer state, so memory pressure comes from weights and from the key-value cache.
• Latency for a single request matters, not only throughput.
• Splitting each layer across devices shortens the time per token, because each device multiplies by a smaller matrix.
• The key-value cache can also be split by heads, so each device stores only its own heads.

For tabular in-context learning, the analogue of the cache is the set of encoded context rows. If the same labeled rows are reused for many query rows, the encoded context can be computed once and kept, with each device holding the part for its heads.

Interaction with quantization

Weights of a tensor parallel model can be compressed after training with methods such as GPTQ or AWQ. Compression is applied per shard or before sharding. Either way, the quantization groups should not cross shard boundaries, or the shards cannot be loaded independently. Check that the group size divides the shard size.

Troubleshooting a slow tensor parallel run

If throughput is below expectation, check the following in order.

1. Confirm that all devices in the group are in the same server and that the fast links are active.
2. Confirm that the collective library selected the right transport and is not falling back to a slow path.
3. Profile a step and look at the share of time spent in collectives.
4. Check that matrix dimensions after sharding are still multiples of the sizes that tensor cores like.
5. Check for load imbalance, for example when heads do not divide evenly.
6. Try a smaller degree. A smaller group may be faster if the model still fits.
7. Try mixed precision if not already enabled, to cut the communication volume.

Strengths and limitations

Strengths

• Fits layers that are too large for a single device.
• Cuts weight and optimizer memory almost in proportion to the degree.
• Reduces the per-device compute for each layer, which lowers latency in inference.
• Does not create pipeline bubbles, because all devices work on the same layer at once.
• Keeps the batch size unchanged, so the optimization behavior is unaffected.
• Combines with data, pipeline, and sequence parallelism.

Limitations

• Needs communication inside every layer, so it depends on fast links.
• Rarely scales well across servers.
• The degree must divide the number of heads and the hidden size.
• Adds code complexity compared with plain data parallelism.
• Small layers do not gain much, because communication cost dominates.
• Replicated parts, such as layer norm, waste memory unless sequence parallelism is added.
• Debugging is harder, because a bug can appear as slightly wrong numbers instead of an error.

How it compares with alternatives

Tensor parallelism versus FSDP

Both reduce the weight memory per device. FSDP stores shards and gathers each layer's full weights just in time, so every device still performs the whole layer computation on its own data. Tensor parallelism never gathers weights. Each device computes with only its shard and exchanges activations instead. Tensor parallelism shortens the per-layer compute on each device. FSDP is easier to apply, while tensor parallelism can give better efficiency for very large layers on fast links.

Combining methods

Large training runs use all of them. A typical recipe: tensor parallelism inside a node, pipeline parallelism across a few nodes, data parallelism (often sharded) across the remaining devices, mixed precision everywhere, and gradient checkpointing when activations still do not fit.

A decision guide

Use these questions in order to decide whether tensor parallelism belongs in a plan.

1. Does the model, with its optimizer state, fit on one device? If yes, use data parallelism and stop.
2. Does it fit if the state is sharded across data parallel workers? If yes, use FSDP or DeepSpeed ZeRO.
3. Does a single layer fit on one device after sharding? If not, tensor parallelism is needed.
4. Is the model so deep that a few devices cannot hold it even with tensor parallelism? Add pipeline parallelism.
5. Is single-request latency the goal? Tensor parallelism is a good choice, up to the size of one server.

Typical layouts by scale

The table is a guide, not a rule. Hardware, sequence length, and batch size shift the boundaries. Always measure on the actual cluster.

Common patterns in code

Most libraries expose tensor parallelism as drop-in layer types. The idea is to replace a standard linear layer with a column-parallel or row-parallel version, and to leave the rest of the model alone.

class ParallelMLP(torch.nn.Module):
    def __init__(self, d_model, d_hidden, tp_group):
        super().__init__()
        self.up = ColumnParallelLinear(d_model, d_hidden, group=tp_group)
        self.down = RowParallelLinear(d_hidden, d_model, group=tp_group)

    def forward(self, x):
        h = torch.nn.functional.gelu(self.up(x))
        return self.down(h)   # the all-reduce happens inside RowParallelLinear

The layer names above are placeholders for whatever the chosen library calls them. The structure is the important part: column first, row second, one all-reduce hidden inside the second layer.

Role in tabular foundation models

Tabular foundation models learn from raw tables: rows and columns, mixed types, missing values, and often no useful column names. They are pre-trained on large numbers of synthetic or real tables and then used through in-context learning, where labeled rows are shown at inference time and the model predicts new rows without weight updates. For the overall design, see /model/architecture and /model/research.

When tabular models need tensor parallelism

Many tabular models are small compared with large language models, and a single table fits on one GPU. For such models, data parallelism is enough. Tensor parallelism becomes relevant in four situations.

1. The model has grown to billions of parameters, as foundation models do when they are trained for broad coverage of table types.
2. The hidden size is large, so single feed-forward and attention matrices are heavy.
3. The context holds many example rows, so attention activations become large.
4. Latency matters at inference, and splitting each layer across devices lowers response time.

Attention over rows and columns

A tabular transformer often attends over two axes: across columns within a row, and across rows within a column. See axial attention and Perceiver IO. Each of these attention blocks has the same structure as in a language model: query, key, and value projections, several heads, and an output projection. Head parallelism applies directly. Giving each GPU a subset of heads splits the attention work for both axes.

When the table is wide, the column axis has many tokens, and when it is tall, the row axis has many. Either way, splitting heads across devices lowers the per-device load for each attention pass.

Mixed types and per-type embeddings

Tabular inputs combine numeric, categorical, date, and text-derived features. Models encode each type with its own embedding or projection. If the categorical embedding tables are large, they can be split by rows across devices in the same way as vocabulary tables. Numeric encoders are small, so they are usually replicated. The key point is that the splitting happens inside the model and is invisible to the data format.

Missing values and masks

Missing cells are represented by masks or special tokens. Because these operate element by element, they are compatible with tensor parallelism with no extra communication. The mask tensors are replicated on each device in the group, and each device applies them to its slice of the activations.

In-context learning with long contexts

In-context learning for tables places labeled training rows and unlabeled query rows in one sequence. The more rows shown, the better the prediction can be, up to a point. Long contexts mean heavy activations. Tensor parallelism together with sequence parallelism and FlashAttention allows longer contexts to fit in memory. This is especially useful at inference, where a large context must be processed in a single forward pass, as in TabPFN and TabICL.

Pre-training objectives

Pre-training on synthetic tables, masked cell modeling, and imputation objectives all use the same network body. See masked cell modeling and imputation networks. Tensor parallelism changes how the body is laid out across devices, but not what the loss is. The loss heads are small and are replicated, or split if they output a very large vocabulary of discrete bins.

Continual learning

In continual learning the model is updated on new tables while keeping old skills. Methods like elastic weight consolidation and synaptic intelligence keep importance values for every weight. With tensor parallelism, those importance values must be sharded in the same way as the weights. Each device stores and uses the importance values for its own shard, so no extra communication is needed. Checkpoint tools must save and reload the shards consistently, and converting between different tensor parallel degrees requires re-slicing the weights.

Anomaly detection and sector classification

Task heads such as anomaly scores or semantic sector heads sit on top of the shared representation. They are small. In a tensor parallel model, the final hidden state is gathered or already replicated after the last all-reduce, so heads can run on any device. Because these outputs often feed thresholds, tiny numeric differences between parallel layouts should be tested for stability.

Why it is often not the first choice

For small and medium tabular models, the simplest scaling path wins: data parallelism with mixed precision. Tensor parallelism adds code, communication, and checkpoint complexity. It should be adopted when memory or latency truly requires it. See /model/benchmarks for measured trade-offs in a tabular context.

Practical guidance and best practices

Pick the degree wisely

• Use the smallest tensor parallel degree that fits the model.
• Keep the degree within one server so that fast links carry the traffic.
• Make sure the degree divides the number of attention heads and the hidden size.
• Measure throughput at degrees 1, 2, 4, and 8 before choosing.

Use proven libraries

• Prefer well-tested implementations such as Megatron-LM, Colossal-AI, or the tensor parallel features in PyTorch.
• Avoid writing your own collective logic unless you must.
• Use NCCL for collectives on GPUs, and verify the topology of the links.

Control memory

• Turn on sequence parallelism to shard the activations of layer norm and dropout.
• Add gradient checkpointing when activations are still too large.
• Use mixed precision to halve activation and communication volume.
• Consider FlashAttention to avoid storing full attention matrices.

Initialize and seed correctly

• Replicated parameters must start identical on all devices in the group.
• Sharded parameters should be initialized so the combined matrix has the intended distribution, not a scaled-down version of it.
• Random number generators for dropout need care so that replicated activations use the same mask and split activations use different ones.

Test for equivalence

• Compare the loss of a tensor parallel model against a single-device model on the same data for a few steps.
• Compare gradients after one step. They should match closely.
• Re-run the check after any change to the layout.

Plan checkpoints

• Save shards with metadata describing the parallel layout.
• Provide a tool to merge shards into a single file for sharing and for serving.
• Provide a tool to re-split shards for a different degree.

Common mistakes

• Spreading a tensor parallel group across servers. The slower network turns communication into the bottleneck and can make the run slower than with no splitting.
• Choosing a degree that does not divide the heads. Some libraries fail, others pad and waste compute.
• Using tensor parallelism for small layers. The communication costs more than the computation saved.
• Forgetting that some parameters are replicated. Gradients of replicated parameters must be synchronized inside the group.
• Mismatched random seeds. Different dropout masks on replicated activations quietly break equivalence.
• Skipping the equivalence test. A split that is off by a scale factor still trains, but it trains worse.
• Ignoring checkpoint layout. Shards saved with one degree cannot be loaded with another without conversion.
• Expecting memory savings for activations. Without sequence parallelism, activation memory drops only partly.
• Treating tensor parallelism as a replacement for data parallelism. It adds to it. It does not increase the number of examples processed per step.
• Overlooking the embedding and output layers. Large vocabularies or bin sets need their own split strategy.

Checklist before a long run

• The degree divides the number of heads and the hidden size.
• All devices in a group share one server.
• A short equivalence test against one device passed.
• Sequence parallelism is on if activations are heavy.
• Checkpoints record the layout, and a merge tool is ready.
• Throughput was measured at two or three degrees.

Frequently asked questions

What is tensor parallelism in simple terms?

Tensor parallelism cuts the big weight matrices inside a layer into pieces and gives each piece to a different GPU. All the GPUs work on the same layer at the same time, each with its own piece, and then they combine their partial answers. It lets a layer that is too big for one GPU run across several.

How is tensor parallelism different from pipeline parallelism?

Tensor parallelism splits inside each layer, so every device takes part in every layer. Pipeline parallelism splits between layers, so each device owns a group of layers and passes activations to the next device. Tensor parallelism needs very frequent communication, while pipeline parallelism needs less but suffers idle time called bubbles. See pipeline parallelism.

Why must tensor parallelism stay inside one server?

Every layer requires an all-reduce or a similar exchange of activations, and these exchanges are on the critical path. Links inside a server are much faster than links between servers. Across servers, communication time can grow larger than computation time, and the speedup disappears.

What is a column-parallel and a row-parallel layer?

A column-parallel layer splits its weight matrix by output columns, so each device produces a slice of the output without any communication. A row-parallel layer splits the weight by input rows, so each device produces a partial sum that must be added with an all-reduce. Pairing a column-parallel layer with a row-parallel layer needs only one all-reduce per block.

Does tensor parallelism change the model results?

In exact arithmetic it does not, because the split computes the same matrix product. In floating-point arithmetic, the order of additions differs, so tiny numeric differences appear. Well-implemented systems match a single-device run closely, and equivalence tests are the way to confirm this.

Can I use tensor parallelism together with data parallelism?

Yes, and most large runs do. Tensor parallelism handles the model size inside a group of devices, and data parallelism replicates that group to process more data. The total number of devices equals the tensor parallel degree times the pipeline degree times the data parallel degree. See data parallelism.

Is tensor parallelism useful for inference?

Yes. Splitting each layer across several GPUs reduces the memory needed per device and can reduce latency for a single request, since each device does less work per layer. It is widely used to serve models that do not fit on a single GPU. Servers pay for it with communication on every layer.

Do tabular foundation models need tensor parallelism?

Often not. Many tabular models are small enough to train with data parallelism alone, since one table fits on a single device. Tensor parallelism becomes useful when the model is very large, when hidden sizes are big, or when in-context learning uses long contexts of example rows that stress memory and latency.

How does tensor parallelism interact with mixed precision?

They work well together. Lower precision halves the size of activations and of the data exchanged between devices, which eases the communication load. Accumulation in the all-reduce should be handled carefully so that partial sums do not lose accuracy. See mixed precision.

What is sequence parallelism and why is it paired with tensor parallelism?

Sequence parallelism splits the parts of a layer that act on each token on its own, such as layer normalization and dropout, along the sequence axis. Plain tensor parallelism leaves these parts replicated, which wastes activation memory. Adding sequence parallelism removes that waste at almost no extra communication cost.

Glossary

• All-gather: A collective that gathers slices from all devices so each has the full tensor.
• All-reduce: A collective that sums values across devices and shares the result.
• Column-parallel: A linear layer split by output columns.
• Head parallelism: Assigning different attention heads to different devices.
• Interconnect: The physical link between devices, such as a high-bandwidth bus.
• Model parallelism: Dividing a model among several devices.
• Reduce-scatter: A collective that sums values and gives each device one slice of the result.
• Row-parallel: A linear layer split by input rows.
• Sequence parallelism: Splitting per-token operations along the sequence axis.
• Shard: One piece of a tensor held by one device.
• Tensor parallel degree: The number of devices sharing each layer.

Summary

Tensor parallelism splits the weight matrices inside each layer across several GPUs, using a column split followed by a row split so that one all-reduce per block is enough. It cuts weight memory, shortens per-layer compute, and avoids pipeline bubbles, at the price of frequent communication that demands fast links inside a single server. It is almost always combined with data parallelism, pipeline parallelism, sequence parallelism, and mixed precision.

For tabular foundation models, tensor parallelism is a tool for the large end of the scale: big hidden sizes, many attention heads over rows and columns, long in-context sets of example rows, and low-latency serving. Smaller tabular models are usually served well by data parallelism alone. When tensor parallelism is used, shard the importance values and checkpoints consistently, test equivalence against a single-device run, and keep the group inside one node. Related reading: Megatron-LM, FSDP, pipeline parallelism, and /model/architecture.

Related

It belongs to Foundation models from scratch in my stack, under Foundation models and applied ML. Others in the same group: