Implementation:OpenGVLab InternVL DDP Hooks

Knowledge Sources	OpenGVLab_InternVL
Domains	Distributed Training, Gradient Compression, Segmentation
Last Updated	2026-02-07 14:00 GMT

Overview

Custom DDP (Distributed Data Parallel) communication hooks that reduce gradient communication bandwidth by compressing gradients to half-precision formats during distributed training.

Description

This module provides five DDP communication hooks and wrappers for PyTorch's DistributedDataParallel:

• allreduce_hook -- Standard gradient averaging via all_reduce, dividing by world size before the reduction to avoid overflow (especially for FP16). Returns a future that resolves to the averaged gradient.

• fp16_compress_hook -- Casts the GradBucket tensor to torch.float16, divides by world size, performs all_reduce, then decompresses back to the original dtype via an in-place copy to reduce peak memory.

• bf16_compress_hook -- Identical to fp16_compress_hook but uses torch.bfloat16 (requires NCCL > 2.9.6).

• fp16_compress_wrapper -- A higher-order function that wraps any existing communication hook with FP16 compression. The bucket buffer is cast to float16 before the inner hook runs, then decompressed afterward. Equivalent to composing fp16 compression with any hook (e.g., PowerSGD).

• bf16_compress_wrapper -- Same as fp16_compress_wrapper but for BFloat16.

All hooks perform in-place decompression to minimize peak memory usage, following the pattern recommended in PyTorch issue #45968.

Usage

Register these hooks with DDP models during segmentation training to reduce communication overhead in multi-GPU setups, particularly beneficial for training large InternViT backbone models.

Code Reference

Source Location

• Repository: OpenGVLab_InternVL
• File: segmentation/mmcv_custom/ddp_hooks.py
• Lines: 1-180

Signature

def allreduce_hook(process_group, bucket) -> torch.futures.Future[torch.Tensor]: ...
def fp16_compress_hook(process_group, bucket) -> torch.futures.Future[torch.Tensor]: ...
def bf16_compress_hook(process_group, bucket) -> torch.futures.Future[torch.Tensor]: ...
def fp16_compress_wrapper(hook) -> Callable: ...
def bf16_compress_wrapper(hook) -> Callable: ...

Import

from mmcv_custom.ddp_hooks import fp16_compress_hook, bf16_compress_hook

I/O Contract

Inputs

Name	Type	Required	Description
process_group	dist.ProcessGroup	Yes	The distributed process group (or None for WORLD)
bucket	dist.GradBucket	Yes	The gradient bucket from DDP containing tensors to communicate

Outputs

Name	Type	Description
future	torch.futures.Future[torch.Tensor]	Future resolving to the averaged (and decompressed) gradient tensor

Usage Examples

Basic Usage

import torch.distributed as dist
from mmcv_custom.ddp_hooks import fp16_compress_hook

# Register with a DDP model
ddp_model = torch.nn.parallel.DistributedDataParallel(model)
ddp_model.register_comm_hook(dist.group.WORLD, fp16_compress_hook)

Related Pages

• Principle:OpenGVLab_InternVL_DDP_Gradient_Compression