System Architecture

High-Level Design — System Architecture

Purpose and boundary

The application is an offline command-line training demonstration, not an API or distributed training service. Python or torchrun launches ranks. Each rank owns a full model and a subset of examples; DDP synchronizes gradients so their optimizer steps stay consistent.

System context

flowchart LR
    Developer["Developer / Data Scientist"] --> CLI["Python CLI or Jupyter notebook"]
    Dataset["Built-in text corpus or local UTF-8 corpus"] --> Run["Training run"]
    CLI --> Launcher["python or torchrun launcher"]
    Launcher --> Run
    Run --> Devices["CPU ranks or CUDA devices"]
    Devices --> Artifacts["Rank 0 output: tiny checkpoint or LoRA adapter"]
    Pretrained["Hugging Face Qwen weights and tokenizer (Qwen mode only)"] --> Run

Loading

HLD — two-rank execution

flowchart TB
    User["CLI flags / notebook"] --> Torchrun["torchrun: RANK, LOCAL_RANK, WORLD_SIZE"]
    Corpus["Local corpus / built-in examples"] --> Dataset["IndexedTextDataset"]
    Torchrun --> R0["Rank 0 process"]
    Torchrun --> R1["Rank 1 process (when world_size >= 2)"]
    Dataset --> S0["DistributedSampler rank 0"]
    Dataset --> S1["DistributedSampler rank 1"]
    S0 --> R0
    S1 --> R1
    subgraph Replica0["Rank 0: full model replica"]
      R0 --> C0["Collate and tokenize"] --> F0["Forward + local loss"] --> B0["Backward"] --> O0["AdamW update"]
    end
    subgraph Replica1["Rank 1: full model replica"]
      R1 --> C1["Collate and tokenize"] --> F1["Forward + local loss"] --> B1["Backward"] --> O1["AdamW update"]
    end
    B0 <-->|"DDP gradient AllReduce"| B1
    F0 -. "Separate loss averaging" .-> Metrics["AllReduce reported mean loss"]
    F1 -. "Separate loss averaging" .-> Metrics
    O0 --> Save["Rank 0 saves checkpoint / adapters"]
    O1 -. "No save" .-> Done["Barrier and cleanup"]
    Save --> Done

Loading

Subsystems

Subsystem	Actual implementation
Launcher	python or python -m torch.distributed.run (torchrun)
Control-plane inputs	CLI args; RANK, LOCAL_RANK, WORLD_SIZE environment vars
Dataset	IndexedTextDataset, built-in examples or --corpus-file
Partitioning	DistributedSampler, one shard of indices per rank, epoch reseeding
Model	TinyCausalTransformer or pretrained Qwen + LoRA
Communication	PyTorch process group: Gloo on CPU / Windows; NCCL for Linux CUDA
Training	Per-rank DataLoader, AdamW, gradient clipping; Qwen CUDA uses GradScaler
Artifact persistence	Rank 0 alone writes tiny checkpoint or PEFT adapters + tokenizer

Important: DDP does not shard model weights across GPUs. A two-GPU setup does not combine device memory for fitting one model. Qwen is downloaded from Hugging Face only in Qwen mode.

See Low-Level-Design, GPU-and-Distributed-Backends, and Checkpoints.