NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Add Conv2d tensor parallelism support in Megatron-LM

Open
#1,694 7 comments 0 reactions 1 assignee Claimed by @Phlip79 View on GitHub
community-request enhancement waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

## 🚀 Feature Request

### Example

I have been working with **extremely large convolutional models on ultra-high-resolution images** (e.g., 10,000 × 10,000 pixels).

We have a single Conv2d layer configured as follows:

- Input channels: **1024**
- Output channels: **2048**
- Kernel size: **3 × 3**
- Spatial resolution: **10,000 × 10,000**
- Batch size: **1**

This one layer alone requires:

- ~18.9 million parameters (~38 MB in FP16 — not the main problem).
- **~400 GB** activation memory (FP16) **per forward pass**.
- **~400 GB** gradient memory (FP16) **per backward pass**.

> **Total memory footprint per GPU for this single layer alone: >800 GB (FP16).**

It is completely impossible to fit on even the largest GPUs available today (e.g., A100 80GB or H100 120GB).

### Motivation

Megatron-LM currently supports tensor parallelism for large linear (dense) layers, which is extremely effective for transformer-based models. However, many emerging large-scale vision and multi-modal architectures (e.g., ConvNets with huge conv layers, hybrid conv-transformer backbones) involve extremely large convolutional layers that cannot fit into a single GPU, even with pipeline or data parallelism.

Adding tensor parallelism support for `Conv2d` layers would allow distributing these large convolutional weights and activations across multiple GPUs, enabling the training of much larger purely convolutional or hybrid models.

---

### Proposed solution

Implement a `ColumnParallelConv2d` module similar to `ColumnParallelLinear`, which:

- **Splits output channels** across GPUs (i.e., each rank stores only a shard of the output channels).
- Performs partial forward computation on each GPU.
- Uses `torch.distributed.all_gather` (or `reduce_scatter` if applicable) to assemble the final output feature maps.
- Automatically handles gradient synchronization through PyTorch autograd.

This design matches the current tensor parallel style in Megatron-LM and can leverage existing collective communication utilities in `megatron.core`.

---

### Potential API design

```python
class ColumnParallelConv2d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, ...):
# Initializes and shards conv weights across GPUs (output channels split)

def forward(self, x):
# Local conv computation
# All-gather outputs along channel dimension
return y

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.