Low Level Design

Low-Level Design — Modules, Classes, and Control Flow

Source-level components

The main implementation is in llm_ddp_parallel_training.py. The diagrams below distinguish project classes (IndexedTextDataset, TinyCausalTransformer) from imported PyTorch/PEFT classes.

classDiagram
    class IndexedTextDataset {
      +texts: list~str~
      +__len__() int
      +__getitem__(index) tuple
    }
    class TinyCausalTransformer {
      +token_emb: Embedding
      +pos_emb: Embedding
      +encoder: TransformerEncoder
      +norm: LayerNorm
      +lm_head: Linear
      +forward(input_ids, labels) loss
    }
    class DistributedSampler {
      +set_epoch(epoch)
    }
    class DataLoader {
      +__iter__()
    }
    class DistributedDataParallel {
      +forward()
      +synchronize_gradients()
    }
    class AdamW {
      +zero_grad()
      +step()
    }
    class PEFT_LoRA {
      +get_peft_model()
      +save_pretrained()
    }
    IndexedTextDataset --> DataLoader : source of examples
    DistributedSampler --> DataLoader : per-rank subset
    DataLoader --> TinyCausalTransformer : tiny mode tensors
    TinyCausalTransformer ..> DistributedDataParallel : wrapped if world size > 1
    PEFT_LoRA ..> DistributedDataParallel : qwen mode model wrapped
    DistributedDataParallel --> AdamW : trainable parameters

Loading

LLD function contracts

Function/class	Input	Output / side effect
IndexedTextDataset(texts)	list of non-empty texts	indexed (record_id, text) pairs
tiny_collator(items, seq_len)	batch of indexed texts	padded input_ids, masked labels, record_ids
TinyCausalTransformer.forward(input_ids, labels)	[B,T] integer tensors	scalar shifted next-token cross-entropy
main()	argparse + launcher environment	process initialization, training, rank-0 saving, teardown
Qwen nested collate_fn	indexed strings	token IDs, attention mask, labels with padded positions -100

Important code branches

1. --model tiny: local Embedding(258,64), learned positions, one nn.TransformerEncoderLayer, LayerNorm and LM head. Learning rate defaults to 3e-4.
2. --model qwen: HF model Qwen/Qwen2.5-0.5B-Instruct and PEFT LoRA on q_proj/v_proj; learning rate defaults to 2e-4.
3. WORLD_SIZE > 1: initialize distributed backend; wrap model in DDP and use DistributedSampler.
4. WORLD_SIZE == 1: run without a process group, DDP wrapper or DistributedSampler.
5. rank == 0: log aggregate losses and write artifacts; all ranks synchronize at a final barrier if distributed.

Shapes for tiny mode

• input_ids: [B, T]; B is per-rank batch size; T = seq_len.
• Embeddings: [B, T, 64].
• Causal attention mask: [T, T], true above diagonal.
• LM logits: [B, T, 258].
• Cross-entropy on shifted positions: logits [:, :-1, :] vs labels [:, 1:]; pad targets -100 ignored.

Read Tiny-Transformer, Qwen-LoRA, and Runtime-Sequence.