AI Distributed Job Initialization

 

Job Orchestration and Resource Allocation

When a distributed training job is launched, the cluster scheduler or orchestration layer (such as SLURM or Kubernetes) allocates the required compute nodes and provides job-specific information to the job launcher. This information includes the JobID, the allocated nodes, and the resources assigned to the job.

The launcher provides the configuration required by the Torchrun agents, such as the number of nodes, the number of local workers, and the rendezvous configuration. Torchrun uses this configuration and the rendezvous mechanism to establish the worker-group membership and determine the distributed identities of the worker processes. These include the global rank (RANK), local rank (LOCAL_RANK), total number of workers (WORLD_SIZE), and number of local workers (LOCAL_WORLD_SIZE).

In a typical multi-node PyTorch deployment, one Torchrun agent runs on each allocated node. The agent is a host-side process responsible for launching and managing the local worker group. It starts the configured number of worker processes, typically one per GPU. Each worker is an independent operating-system process that executes the training application and uses its assigned GPU for computation.

The Torchrun agent monitors the worker processes and manages their lifecycle, including restarting the worker group if a failure occurs and a restart is permitted by the configured policy.

Each Torchrun agent manages a local worker group on its node. In a multi-node training job, the local worker groups on all participating nodes together form one distributed worker group (called a WorkerGroup in PyTorch terminology). The agents participate in a shared rendezvous mechanism to discover the participating nodes, synchronize worker-group initialization, and establish the global worker-group configuration. This process provides the information needed to assign global ranks and determine the total number of workers.

Rendezvous Configuration and Store Endpoint

Before the worker processes initialize their distributed communication process group, the rendezvous configuration and the endpoint used for distributed initialization must be established.

The rendezvous configuration identifies the participating worker group and specifies how the Torchrun agents discover and coordinate with each other. It includes a rendezvous ID, a rendezvous backend, and a rendezvous endpoint. The rendezvous endpoint identifies where the rendezvous backend can be accessed.

The rendezvous configuration is provided to the Torchrun agents using one of two common approaches:

Dynamic Discovery via Cluster Scheduler (e.g., SLURM): In large-scale production environments, the job launch script uses the cluster manager's allocation information to identify the participating nodes and select a node for the rendezvous service. It retrieves a hostname or IP address that can be reached by the other nodes and determines the required TCP port.

Static Configuration in Launch Scripts: For smaller test beds, static labs, or deterministic environments (such as the two-node topology depicted in Figure 2-1), the rendezvous configuration is explicitly defined in the launch script or its configuration.

After rendezvous, Torchrun provides the worker processes with the MASTER_ADDR and MASTER_PORT environment variables. These identify the host and TCP port used for the C10d TCPStore, which provides the key-value store used to bootstrap distributed communication. Depending on the rendezvous backend and Store configuration, this endpoint may be established as part of the rendezvous process rather than being identical to the rendezvous endpoint.

Environment Variable Propagation

When the Torchrun agents launch the worker processes, they provide each process with a set of environment variables containing shared job information and process-specific information.

In Figure 2-1, each Torchrun agent launches two worker processes on its node. Each worker receives its own environment variable set, which is available in the process environment in the host's memory (CPU DRAM). These variables contain information such as the global rank and local rank. The JobID is typically inherited from the scheduler-provided environment, if it has been exported to the worker processes. Each worker process also receives a Process ID (PID), which is assigned by the operating system when the process is created. The PID identifies the process on its host and is not an environment variable.

Torchrun provides the distributed process coordination variables, including:

·       WORLD_SIZE: Total number of worker processes participating in the worker group.

·       LOCAL_WORLD_SIZE: Number of worker processes running in each local worker group.

·       RANK: Global rank of the worker process, unique across all nodes.

·       LOCAL_RANK: Local rank of the worker process within its node.

·       GROUP_RANK: Rank of the worker group. When there is one worker group per node, this identifies the node's position in the distributed job.

·       MASTER_ADDR: Hostname or IP address used to access the C10d TCPStore.

·       MASTER_PORT: TCP port used to access the C10d TCPStore.

The node rank identifies a node's position within the distributed job. Depending on the launch configuration, it is supplied as a node-rank parameter or represented by the worker-group rank (GROUP_RANK). When there is one worker group per node, the group rank identifies the node's position in the job.

Environment Variable Interpretation

The environment variables provide the information PyTorch needs to identify each worker process, determine the process-group membership, and select the GPU used by each local process.

The WORLD_SIZE variable specifies the total number of participating worker processes. It allows the distributed framework to determine the expected process-group membership and coordinate the initialization of all participating processes.

Each process has a global rank (RANK), which uniquely identifies it across the worker group. In a homogeneous cluster, where every node runs the same number of worker processes, the global rank can be expressed using the worker-group rank, the number of local processes, and the local rank:

 RANK = GROUP_RANK × LOCAL_WORLD_SIZE + LOCAL_RANK 

This formula illustrates how global ranks are mapped across nodes. In practice, Torchrun assigns the RANK value to each worker process as part of its distributed worker-group configuration.

Using the LOCAL_RANK value, the PyTorch application selects the GPU used by each local worker process. In Figure 2-1, for example, the first worker on Node 0 (LOCAL_RANK: 0, RANK: 0) selects GPU 0, while the second worker (LOCAL_RANK: 1, RANK: 1) selects GPU 1. The mapping between local ranks and physical GPUs depends on the GPU devices made visible to the worker processes.

The worker with global rank 0 is commonly responsible for certain initialization tasks and is often informally referred to as the master rank. The MASTER_ADDR and MASTER_PORT values identify the host and TCP port used for the C10d TCPStore, which provides a coordination and bootstrapping mechanism for distributed initialization. Depending on the rendezvous and Store configuration, the TCPStore server can be hosted by the Torchrun agent on the rank-0 node or by the worker-group rank 0 worker process. The other worker processes use these values to establish TCP connections to the service and exchange the information required for distributed initialization.

Together, the JobID, worker-group membership, local ranks, global ranks, and master endpoint information allow the distributed training framework to organize the participating processes and establish the foundation for distributed communication.

Once the processes have selected their respective GPUs, the application can allocate GPU memory (VRAM) for model parameters, optimizer states, gradients, and activation buffers as required by the training workload.

NCCL Unique ID Distribution and Communication Groups

During distributed job initialization, the Torchrun Agents participate in a C10d Rendezvous process to establish the participating worker group. In the configuration shown in Figure 2-1, the rendezvous process also provides access to a shared key-value Store implemented by C10d TCPStore. The TCPStore server can be hosted on Node 0 and accessed by workers on both nodes through its network endpoint.

For example, in Figure 2-1, the TCPStore endpoint is configured using MASTER_ADDR and MASTER_PORT. These environment variables identify the address and port used by workers to access the Store.

The C10d Store allows workers to publish and retrieve initialization information. It uses a key-value mechanism: one worker publishes information under a specific key, while other workers wait for that key to become available and retrieve its value.

Let's consider a training job running on two compute nodes, each with two GPU workers. The model uses tensor parallelism (TP=2) and data parallelism (DP=2). Each node hosts one data-parallel model replica. Within each replica, tensor parallelism divides the model computation between the two local workers. Therefore, Node 0 contains the tensor-parallel group consisting of Ranks 0 and 1, while Node 1 contains the corresponding tensor-parallel group consisting of Ranks 2 and 3. The two data-parallel replicas process different portions of the training dataset. Workers within the same tensor-parallel group cooperate on the same input batch or micro-batch.

The workers within each tensor-parallel group must establish an NCCL communicator to exchange intermediate tensors during model computation. To initialize this communicator, one participating worker generates an NCCL Unique ID and publishes it in the C10d Store. The other participating worker retrieves the same ID and uses it to initialize its local NCCL communicator instance.

For example, Rank 0 generates UID-A and publishes it in the C10d Store. Rank 1 retrieves UID-A and uses it to initialize the tensor-parallel communicator shared by Ranks 0 and 1. The same process takes place on Node 1, where Rank 2 generates UID-B and publishes it in the Store. Rank 3 retrieves UID-B to initialize the tensor-parallel communicator shared by Ranks 2 and 3.

Data parallelism requires additional communication groups to synchronize gradients between corresponding model partitions across the two nodes. For this purpose, Rank 0 generates UID-C and publishes it in the Store for Rank 2. Similarly, Rank 1 generates UID-D and publishes it for Rank 3. The corresponding workers retrieve these UIDs and initialize their data-parallel NCCL communicators.

The resulting NCCL communicators are reused for the communication operations required by their respective parallelization dimensions. Tensor-parallel communicators support intermediate tensor exchanges during forward and backward computation, while data-parallel communicators support gradient synchronization, typically during the backward pass.

Consequently, TCPStore provides a shared mechanism for distributing NCCL UIDs and initializing different GPU communication groups according to the training job's parallelization strategy. The NCCL UID identifies a communicator initialization, not an individual NCCL operation.

NCCL Topology Discovery and Communication Path Selection

After the NCCL communicators have been initialized, NCCL establishes the communication paths required between the participating GPUs. During communicator initialization, NCCL discovers the available GPU and network topology, including the GPUs, PCIe connections, NVLink and NVSwitch interconnects, network adapters, and their associated communication paths.

 Based on the discovered topology, NCCL establishes communication channels and selects suitable transport mechanisms for communication between participating GPUs. Within a compute node, communication can use GPU-to-GPU paths such as NVLink or PCIe. Between compute nodes, communication uses network adapters and the available cluster network fabric.

In an AI cluster, these communication paths correspond to two major scaling dimensions:

Scale-up communication: Communication between GPUs within the same compute node, typically using NVLink, NVSwitch, or PCIe.

Scale-out communication: Communication between GPUs located on different compute nodes, using network adapters and the cluster network fabric, such as InfiniBand, RoCE, or Ultra Ethernet.

NCCL also supports different collective communication algorithms, such as Ring and Tree. These algorithms define how participating GPUs exchange data during collective operations. NCCL selects suitable algorithms and communication protocols based on factors such as the collective operation, message size, and available hardware topology.

The combination of topology discovery, transport selection, and communication algorithms allows NCCL to use the available GPU interconnects and network infrastructure efficiently.

Figure 2-1: Job Orchestration and Resource Allocation.