DeepSpeed Launcher Runner Implementation for Multi-node Multi-GPU Training

from copy import deepcopy

def filter_host_resources(host_info, include_list="", exclude_list=""):
    '''Process inclusion or exclusion specifications for host resource filtering.

    Format is NODE_SPEC[@NODE_SPEC ...], where
        NODE_SPEC = NAME[:SLOT[,SLOT ...]].
    Omitting :SLOT includes/excludes all slots on that host.

    Examples:
        include_list="worker-0@worker-1:0,2" uses all slots on worker-0 and
          slots [0, 2] on worker-1.
        exclude_list="worker-1:0" uses all available resources except
          slot 0 on worker-1.
    '''

    # Syntax constants
    NODE_DELIMITER = '@'
    SLOT_MARKER = ':'
    SLOT_DELIMITER = ','

    # Mutual exclusivity check
    if include_list and exclude_list:
        raise ValueError('include_list and exclude_list are mutually exclusive.')

    # No filtering needed
    if not include_list and not exclude_list:
        return host_info

    # Initialize result container
    filtered_hosts = dict()
    parse_string = include_list if include_list else exclude_list

    # Process each node specification
    for spec in parse_string.split(NODE_DELIMITER):
        # Check if specific slots are mentioned
        if SLOT_MARKER in spec:
            host_name, slot_data = spec.split(SLOT_MARKER)
            slots = [int(x) for x in slot_data.split(SLOT_DELIMITER)]

            # Validation checks
            if host_name not in host_info:
                raise ValueError(f"Host '{host_name}' not found in hostfile")
            for slot in slots:
                if slot not in host_info[host_name]:
                    raise ValueError(f"Slot '{slot}' not found on host '{host_name}'")

            # Apply filtering based on mode
            if include_list:
                if host_name not in filtered_hosts:
                    filtered_hosts[host_name] = []
                filtered_hosts[host_name].extend(slots)
            else:
                # Exclude mode
                if host_name not in filtered_hosts:
                    filtered_hosts[host_name] = host_info[host_name][:]
                for slot in slots:
                    if slot in filtered_hosts[host_name]:
                        filtered_hosts[host_name].remove(slot)
                # Remove host if no slots remain
                if not filtered_hosts[host_name]:
                    del filtered_hosts[host_name]
        else:
            # All slots requested
            host_name = spec
            if host_name not in host_info:
                raise ValueError(f"Host '{host_name}' not found in hostfile")
            
            if include_list:
                filtered_hosts[host_name] = host_info[host_name][:]
            else:
                # Remove entire host from exclusion
                if host_name in filtered_hosts:
                    del filtered_hosts[host_name]

    return filtered_hosts

The filter_host_resources function handles the parsing and application of resource inclusion and exclusion rules for distributed training environments. It processes host specifications in the format HOSTNAME[:SLOT[,SLOT...]], supporting both explicit slot selection and full host inclusion/exclusion.

The implementation follows these steps:

  1. Defines constants for parsing syntax
  2. Validates mutual exclusivity of include/exclude parameters
  3. Returns input unchanged when no filtering is required
  4. Initializes filtering state based on operation type
  5. Processes each node specification in the input string
  6. Handles slot-specific configurations when colon notation is present
  7. Performs validation checks for host and slot existence
  8. Applies filtering logic according to include/exclude mode
  9. Manages cleanup when hosts become empty after filtering

For multi-node scenarios, DeepSpeed uses hostfiles compatible with OpenMPI and Horovod. These files list machine names (or SSH aliases) accessible via passwordless SSH along with GPU counts per machine. Example format:

worker-1 slots=4
worker-2 slots=4

To launch distributed training, the deepspeed command-line tool is used with options like --hostfile, --num_nodes, and --num_gpus. Additional filtering can be achieved using --include and --exclude flags with the same syntax as the function above.

Environment variables can be propagated across nodes using a .deepspeed_env file containing VAR=VAL entries. This enables configuration of specialized networking parameters required by certain clusters.

DeepSpeed supports MPI-based launching through mpi4py integration. When using MPI, the launcher automatically initializes torch distributed backend with proper world size and rank information.

In single-node setups, hostfiles are optional. DeepSpeed detects local GPU count when no hostfile is provided. Resource filtering works similarly, but requires "localhost" as hostname.

The system handles distributed training initialization through the deepspeed.initialize function which sets up model engine, optimizer, and data loaders with appropriate distributed configurations. Training loops use simple three-step operations: forward pass, backward propagation, and weight updates.

Model checkpoints are managed through save_checkpoint and load_checkpoint APIs that handle model weights, optimizer states, and custom client state dictionaries transparently.

Configuration is defined in JSON files specifying training batch sizes, gradient accumulation, optimizer settings, mixed precision options, and ZeRO optimization levels.

Tags: deepspeed distributed-training multi-gpu resource-management launcher

Posted on Tue, 04 Aug 2026 16:15:09 +0000 by g.grillo