Runtime Sequence

Runtime Sequence and Process Lifecycle

Detailed two-rank training step

sequenceDiagram
    autonumber
    participant L as torchrun
    participant A as Rank 0
    participant B as Rank 1
    participant PG as Process group / collective
    participant FS as Checkpoint store
    L->>A: Launch RANK=0, LOCAL_RANK=0
    L->>B: Launch RANK=1, LOCAL_RANK=1
    A->>PG: init_process_group(env://)
    B->>PG: init_process_group(env://)
    Note over A,B: Initialize model replicas and per-rank DistributedSamplers
    A->>A: Forward and local loss
    B->>B: Forward and local loss
    A->>PG: backward(): gradient collectives
    B->>PG: backward(): gradient collectives
    PG-->>A: synchronized averaged gradients
    PG-->>B: synchronized averaged gradients
    A->>A: clip gradients + optimizer.step()
    B->>B: clip gradients + optimizer.step()
    A->>PG: all_reduce(loss) for logging
    B->>PG: all_reduce(loss) for logging
    PG-->>A: global mean loss
    A->>FS: Save checkpoint or LoRA adapter
    A->>PG: barrier()
    B->>PG: barrier()
    A->>PG: destroy_process_group()
    B->>PG: destroy_process_group()

Loading

Process state machine

stateDiagram-v2
    [*] --> ParseCLI
    ParseCLI --> ReadRankEnvironment
    ReadRankEnvironment --> ChooseDevice
    ChooseDevice --> InitGroup: world_size > 1
    ChooseDevice --> PrepareData: world_size = 1
    InitGroup --> PrepareData
    PrepareData --> BuildModel
    BuildModel --> WrapDDP: world_size > 1
    BuildModel --> Train: world_size = 1
    WrapDDP --> Train
    Train --> SaveRankZero
    SaveRankZero --> Barrier: world_size > 1
    SaveRankZero --> End: world_size = 1
    Barrier --> DestroyGroup
    DestroyGroup --> End
    End --> [*]

Loading

torchrun sets environment variables per subprocess. Both ranks must run compatible collective operations in the same order; otherwise the distributed job may hang or fail. The script initializes the process group with a ten-minute timeout, then destroys it in finally.

The diagram shows two ranks for clarity. With one rank the script bypasses DDP and distributed collectives.