[ENHANCEMENT] Add Conv2d tensor parallelism support in Megatron-LM
- 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
Assessment
This issue has not been assessed yet.