Implementation:Mlfoundations Open flamingo Init distributed device

Overview

Concrete tool for initializing distributed process groups across multiple backends provided by the OpenFlamingo training module.

Description

The init_distributed_device() function auto-detects the distributed launch method (Horovod, SLURM, or torchrun) and initializes the appropriate process group. It sets the following attributes on the args namespace:

args.distributed — boolean flag indicating whether distributed mode is active.
args.world_size — total number of processes across all nodes.
args.rank — global rank of the current process.
args.local_rank — rank of the current process on its node.
args.device — the torch.device assigned to this process.

For GPU training, the function calls torch.cuda.set_device() to bind the process to its local GPU and initializes the NCCL backend via torch.distributed.init_process_group. It returns the torch.device to use for all subsequent training operations.

Usage

Call at the beginning of training before model creation or data loading. The function must be invoked in every process spawned by the distributed launcher.

Code Reference

Source: Repository: https://github.com/mlfoundations/open_flamingo; File: open_flamingo/train/distributed.py, Lines L73–132

Signature

def init_distributed_device(args) -> torch.device: """ Sets: args.distributed, args.world_size, args.rank, args.local_rank, args.device Returns: torch.device for this process """

Import: from open_flamingo.train.distributed import init_distributed_device

I/O Contract

Inputs

Name	Type	Required	Description
`args`	`argparse.Namespace`	Yes	Training arguments with `dist_backend` and `dist_url` fields

Outputs

Name	Type	Description
`device`	`torch.device`	The device assigned to this process
`args` modifications	in-place	Sets `args.distributed`, `args.world_size`, `args.rank`, `args.local_rank`, `args.device`

Usage Examples

Initializing distributed training with torchrun:

import argparse
import torch
from open_flamingo.train.distributed import init_distributed_device

args = argparse.Namespace(
    dist_backend="nccl",
    dist_url="env://",
    no_set_device_rank=False,
    horovod=False,
)

device = init_distributed_device(args)

print(f"Process rank {args.rank}/{args.world_size} using device: {device}")
# Build model after initialization
model = build_model(args).to(device)

Launch the above script with:

torchrun --nproc_per_node=4 train.py

Environment:Mlfoundations_Open_flamingo_PyTorch_CUDA_Distributed

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment