Skip to main content

PyTorch Elastic Distributed Training

When configuring PyTorch distributed training on Kubernetes, jobs frequently hang or fail quietly when a single worker encounters an out-of-memory error (OOM) or rank desynchronization. With default NCCL and TorchElastic configurations, surviving workers can block for up to 90 minutes before timing out.

In flyte-sdk, the flyteplugins.pytorch.task module provides Elastic and RunPolicy configuration classes alongside the TorchFunctionTask execution wrapper. These classes integrate Flyte with TorchElastic (torch.distributed.launcher.api.elastic_launch) and the Kubeflow PyTorch Operator, handling multi-process launching, rank coordination, rendezvous configuration, failure detection, and NCCL tuning.

Setting Up an Elastic Task​

To run distributed PyTorch training, configure a flyte.TaskEnvironment with plugin_config=Elastic(...) and specify resource requirements such as GPU counts.

import typing
import flyte
import torch
import torch.nn as nn
import torch.optim as optim
from flyteplugins.pytorch.task import Elastic
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler

image = (
    flyte.Image.from_debian_base(name="pytorch_image")
    .with_pip_packages("torch", "flyteplugins-pytorch")
)

torch_env = flyte.TaskEnvironment(
    name="torch_env",
    image=image,
    plugin_config=Elastic(),
)

class ToyDataset(Dataset):
    def __len__(self):
        return 100

    def __getitem__(self, idx):
        return torch.randn(10), torch.randn(1)

@torch_env.task
def torch_distributed_train(epochs: int = 3) -> typing.Optional[float]:
    torch.distributed.init_process_group("gloo")
    
    rank = torch.distributed.get_rank()
    world_size = torch.distributed.get_world_size()

    dataset = ToyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=16, sampler=sampler)

    model = DDP(nn.Linear(10, 1))
    criterion = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    final_loss = 0.0
    for _ in range(epochs):
        sampler.set_epoch(_)
        for x, y in dataloader:
            optimizer.zero_grad()
            output = model(x)
            loss = criterion(output, y)
            loss.backward()
            optimizer.step()
            final_loss = loss.item()

    if rank == 0:
        print(f"Final training loss: {final_loss}")

    return final_loss

Single-Node vs. Multi-Node Execution​

When TorchFunctionTask initializes in flyteplugins/pytorch/task.py, it inspects plugin_config.nnodes:

  • If nnodes == 1: self.task_type is set to "python-task". The task executes locally or in a standard container without provisioning a Kubeflow PyTorchJob custom resource.
  • If nnodes > 1 or an elastic range string is provided (such as nnodes="2:4"): self.task_type is set to "pytorch". The task is submitted to the backend as a multi-replica Kubeflow PyTorchJob.

The helper run.parse_min_max_nnodes(str(self.plugin_config.nnodes)) parses nnodes into min_nodes and max_nodes.

Return Values and Rank Coordination​

When elastic_launch spawns worker processes via launcher_entrypoint, each process executes your task function. However, Flyte captures return values based on distributed rank:

out = elastic_launch(config=config, entrypoint=launcher_entrypoint)(
    tctx,
    cloudpickle.dumps(self.func),
    kwargs,
)
result = out[0] if 0 in out else None
await self.post(result)
return result
  • Rank 0 Return: Only rank 0 (out[0]) returns output to the Flyte workflow engine. Outputs returned by non-zero ranks are ignored.
  • Worker Logging & Tracking: When integrating experiment trackers such as flyteplugins.wandb or flyteplugins.mlflow, guard logging calls to run only on rank 0 unless tracking distributed metrics across all processes.

Process Initialization and Subprocess Management​

Before launching worker processes, TorchFunctionTask.pre prepares the process environment:

  1. Unbuffered Output: Sets PYTHONUNBUFFERED="1" so standard output from all ranks is visible in logs before process exit.
  2. OpenMP Threads: If nproc_per_node > 1 and OMP_NUM_THREADS is not set, Flyte defaults OMP_NUM_THREADS="1" to prevent CPU core oversubscription across local worker processes.
  3. Main Thread Signal Registration: elastic_launch executes directly on the main thread in TorchFunctionTask.execute so it can register signal handlers (SIGINT, SIGTERM) for subprocess cleanup without warning or leaving orphaned worker processes.

Each worker subprocess enters launcher_entrypoint(tctx, fn, kwargs), which deserializes the task function with cloudpickle, calls flyte.init(...) with task context metadata, and executes the user function inside the configured internal_ctx().

Rendezvous Backends and Kubeflow RunPolicy​

You can configure Kubeflow pod cleanup policies, retry limits, and TorchElastic rendezvous parameters directly on Elastic:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyteplugins.pytorch.task.RunPolicy' does not accept keyword(s): active_deadline_seconds, backoff_limit, clean_pod_policy, ttl_seconds_after_finished 'flyteplugins.pytorch.task.Elastic' does not accept keyword(s): max_restarts, monitor_interval, nnodes, nproc_per_node, rdzv_backend, rdzv_configs, …

import flyte
from flyteplugins.pytorch.task import Elastic, RunPolicy

image = flyte.Image.from_debian_base(name="pytorch_image")

policy = RunPolicy(
    clean_pod_policy="Running",
    backoff_limit=2,
    ttl_seconds_after_finished=300,
    active_deadline_seconds=3600,
)

elastic_config = Elastic(
    nnodes="2:4",  # Dynamic elastic scaling between 2 and 4 nodes
    nproc_per_node=2,
    rdzv_backend="c10d",
    rdzv_configs={"timeout": 900, "join_timeout": 900},
    max_restarts=3,
    monitor_interval=3,
    run_policy=policy,
)

torch_env = flyte.TaskEnvironment(
    name="elastic_resilient_env",
    image=image,
    plugin_config=elastic_config,
)

Serialization to Kubeflow Custom Resources​

TorchFunctionTask.custom_config serializes these parameters to a DistributedPyTorchTrainingTask protobuf message:

  • elasticConfig: Serializes rdzv_backend, min_replicas, max_replicas, nproc_per_node, and max_restarts.
  • runPolicy: Maps clean_pod_policy to common_pb2.CleanPodPolicy ("CLEANPOD_POLICY_ALL", "CLEANPOD_POLICY_RUNNING", or "CLEANPOD_POLICY_NONE"), alongside backoff_limit, ttl_seconds_after_finished, and active_deadline_seconds.
  • worker_replicas: Sets DistributedPyTorchTrainingReplicaSpec(replicas=self.max_nodes).

NCCL Tuning, Timeout Controls, and Deadlock Prevention​

During distributed GPU training using NCCL, a failure on one worker (such as CUDA OOM during backward pass) causes that worker to skip collectives. Surviving ranks block inside collective operations (like all-reduce). With PyTorch default settings (600s collective timeout, 1800s heartbeat timeout, 3 restarts), resolving a deterministic failure can take up to:

$$\text{Total Wait} = (\text{max_restarts} + 1) \times (\text{nccl_collective_timeout_sec} + \text{nccl_heartbeat_timeout_sec})$$

$$4 \times (600\text{s} + 300\text{s}) = 3600\text{s (60 minutes)}$$

To fail fast and prevent cluster resource starvation, tune the NCCL timeouts on Elastic:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyteplugins.pytorch.task.Elastic' does not accept keyword(s): max_restarts, nccl_async_error_handling, nccl_collective_timeout_sec, nccl_enable_monitoring, nccl_heartbeat_timeout_sec, nnodes, nproc_per_node

import flyte
from flyteplugins.pytorch.task import Elastic

image = flyte.Image.from_debian_base(name="pytorch_image")

elastic_oom_config = Elastic(
    nnodes=1,
    nproc_per_node=2,
    max_restarts=0,  # Fail immediately on deterministic errors like OOM
    nccl_collective_timeout_sec=60,  # Reduces collective timeout from 600s to 60s
    nccl_heartbeat_timeout_sec=60,   # Reduces heartbeat timeout from 300s to 60s
    nccl_async_error_handling=True,  # Aborts stuck collectives asynchronously
    nccl_enable_monitoring=True,     # Keeps NCCL heartbeat monitor active
)

torch_env = flyte.TaskEnvironment(
    name="fast_fail_env",
    image=image,
    plugin_config=elastic_oom_config,
)

NCCL Configuration Mechanics​

TorchFunctionTask.pre and launcher_entrypoint apply these settings:

  • nccl_heartbeat_timeout_sec: Sets os.environ["TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC"]. Defaults to 300 seconds (compared to PyTorch's default of 1800s).
  • nccl_async_error_handling: Sets os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1". NCCL aborts stuck collectives asynchronously, causing the worker process to crash-exit so TorchElastic detects it within monitor_interval seconds (~3s).
  • nccl_collective_timeout_sec: Sets os.environ["FLYTE_NCCL_COLLECTIVE_TIMEOUT_SEC"]. In launcher_entrypoint, Flyte imports and patches torch.distributed.constants.default_pg_nccl_timeout and torch.distributed.distributed_c10d.default_pg_nccl_timeout to the specified timedelta before user code invokes torch.distributed.init_process_group().
  • nccl_enable_monitoring: Sets os.environ["TORCH_NCCL_ENABLE_MONITORING"] = "1", enabling the NCCL monitor thread that checks heartbeat counters and issues SIGABRT if stalled.

Zombie Watchdog Thread​

When all workers abort simultaneously due to SIGABRT, PyTorch's TorchElastic agent contains a known condition where multiprocessing.Event.set() deadlocks inside Condition.notify() waiting on an unreleased semaphore from dead workers.

TorchFunctionTask.execute invokes _start_zombie_watchdog(nproc=self.plugin_config.nproc_per_node) before elastic_launch. This daemon thread periodically scans /proc for child processes. When it detects that at least nproc direct child processes have entered the zombie state (State: Z), it triggers an immediate os._exit(1):

if len(zombie_pids) >= nproc:
    logger.error(
        "Zombie watchdog: %d worker processes are zombies (PIDs %s). "
        "This indicates a PyTorch elastic agent deadlock in "
        "multiprocessing.Event.set(). Force-exiting.",
        len(zombie_pids),
        zombie_pids,
    )
    os._exit(1)

Diagnostic Failure Formatting​

If a worker process fails, _format_worker_failure translates the resulting ChildFailedError into actionable diagnostics:

  • -signal.SIGABRT: Identifies collective timeout or rank desynchronization typically caused by CUDA OOM.
  • -signal.SIGKILL: Identifies process termination from the host or container OOM killer exceeding memory limits.