DOCSEARCHbyAlgolia
Submit Docs

PyTorch Examples

repository·main·Indexed 12 days ago

https://github.com/pytorch/examples

A curated repository of high-quality PyTorch examples demonstrating various machine learning tasks. Includes implementations for DCGAN, MNIST, linear regression, and transfer learning, with extensive support for the PyTorch C++ frontend (LibTorch) and Distributed Data Parallel (DDP) training across single-GPU, multi-GPU, and multi-node Slurm clusters.

Tokens
41.9K
Snippets
163
Records
201
Agent score
96%

What's inside PyTorch Examples

  1. Overview of Tensor Parallelism examples

    main

    This directory provides demonstrations of SPMD Megatron-LM style Tensor Parallelism using PyTorch native Tensor Parallel APIs. The examples cover three progressive levels of complexity:

    1. Simple module-level Tensor Parallelism: Demonstrated on a dummy MLP model.
    2. Tensor Parallelism with Sequence Parallel: Demonstrates module-level parallelism with Sequence Parallel inputs and outputs on a dummy MLP model.
    3. FSDP + Tensor Parallel: An end-to-end (E2E) demonstration combining Fully Sharded Data Parallel (FSDP) and Tensor Parallel (including Sequence Parallel) using a Llama2 model.

    For detailed technical documentation on the underlying APIs, refer to the official PyTorch distributed.tensor.parallel documentation.

  2. Overview of Distributed DataParallel + Distributed RPC Example

    main

    This example demonstrates a hybrid training architecture that combines Distributed DataParallel (DDP) with the Distributed RPC Framework. It is designed for scenarios where a model has both dense and sparse components.

    Architecture Model

    • Dense Component: Uses nn.Linear and is replicated across trainer nodes using DDP.
    • Sparse Component: Uses nn.EmbeddingBag and resides on a dedicated Parameter Server.

    Workflow

    1. Initialization: The master node initializes an embedding table on the parameter server.
    2. Forward Pass: Each trainer node performs embedding lookups on the parameter server via RPC, then processes the retrieved data through its local nn.Linear module.
    3. Backward Pass:
      • DDP aggregates gradients for the dense (nn.Linear) part using allreduce.
      • The distributed backward pass updates the sparse embedding table parameters directly on the parameter server via RPC.

  3. Overview of minGPT-DDP

    main

    minGPT-DDP is a codebase designed for training GPT-like models using PyTorch's Distributed Data Parallel (DDP) strategy. It is the companion code for the PyTorch tutorial: https://pytorch.org/tutorials/intermediate/ddp_series_minGPT.html.

    Key components include:

    • main.py: The entry point that initializes the DDP process group, loads configurations, and starts the training job.
    • mingpt/trainer.py: Contains the Trainer class responsible for executing distributed training iterations.
    • mingpt/model.py: Defines the GPT model architecture.
    • mingpt/char_dataset.py: Implements a Dataset class for character-level data.
    • mingpt/gpt2_train_cfg.yaml: A configuration file for model, data, optimizer, and training parameters.
    • mingpt/slurm/: Contains scripts for multi-node training on AWS clusters using Slurm.

  4. Overview of PyTorch Examples

    main
    The pytorch/examples repository provides curated, high-quality, and short examples of using PyTorch. These examples are designed to have minimal dependencies and serve as templates that can be emulated in your own work. They cover a wide range of domains including computer vision, natural language processing, reinforcement learning, and distributed training.

  5. Explore torch.fx program transformation examples

    main

    The fx/ directory provides standalone Python examples demonstrating various program transformations implemented using torch.fx. These examples are designed to be runnable as independent scripts to help developers understand how to manipulate and transform PyTorch programs using the FX framework.

    Note: Since torch.fx is in Beta, both the API and these specific examples are subject to change.

  6. Explore PyTorch training examples by domain

    main

    The pytorch/examples repository provides a collection of reference implementations for various machine learning tasks. You can use these examples to learn PyTorch syntax, experiment with specific architectures, or understand how to implement research papers.

    Available domains include:

    • Computer Vision: Image classification (ConvNets, ResNet, AlexNet, VGG), Siamese Networks, DCGAN, Variational Auto-Encoders (VAE), Super-resolution, and Neural Style Transfer.
    • Natural Language Processing: Word-level language modeling using RNNs (GRU, LSTM) and Transformers.
    • Reinforcement Learning: Actor-Critic methods using the Gymnasium toolkit.
    • Time Series: Sequence prediction using LSTMCell.
    • Graph Learning: Graph Convolutional Networks (GCN).
    • Advanced/Research: Forward-Forward algorithm, HOGWILD! training, and torch.fx module transformations.
    • Infrastructure: Distributed training (DDP and RPC) and C++ frontend usage.

  7. Understand the GCN model architecture and dataset

    main

    Model Architecture

    The GCN implementation follows the architecture from the paper "Semi-Supervised Classification with Graph Convolutional Networks". It uses multiple graph convolutional layers with ReLU activation, ending with a softmax layer for classification. Key configurable components include:

    • Number of hidden units
    • Number of layers
    • Dropout rate

    Dataset

    The implementation is pre-configured to use the Cora dataset.

    • Structure: Nodes represent scientific papers; edges represent citation relationships.
    • Task: Semi-supervised node classification (7 binary classes).
    • Handling: The dataset is automatically downloaded and preprocessed.

  8. How the GAT model architecture works

    main

    The Graph Attention Network (GAT) uses multi-head attention to capture information from neighboring nodes.

    Architecture Details:

    • Layers: The implementation uses two graph attention layers.
    • Attention Mechanism: Each layer applies a shared self-attention mechanism to every node to learn importance weights.
    • Activations:
      • ELU (Exponential Linear Unit): Used for hidden layer activations.
      • LeakyReLU: Applied to attention coefficients to ensure non-zero gradients for negative values.
    • Layer Configuration (Official Implementation):
      • First Layer: Uses $K=8$ attention heads computing $F'=8$ features each (total 64 features), followed by ELU.
      • Second Layer: A single attention head computing $C$ features (where $C$ is the number of classes), followed by a log-softmax activation for probabilistic outputs.

    Note on Implementation: This version uses the full dense form of the adjacency matrix rather than a sparse form. While this is less efficient for large graphs, it does not affect accuracy.

  9. Understand Distributed Data Parallel (DDP) Topologies

    main

    In a DDP application, multiple workers train the same global model on different data shards. They compute local gradients and synchronize them using AllReduce. This follows the Single Program Multiple Data (SPMD) model.

    Key concepts for mapping processes to hardware:

    • World Size (W): Total number of processes across all nodes.
    • Local World Size (L): Number of processes per node.
    • Global Rank: The unique ID of a process in the range [0, W-1].
    • Local Rank: The unique ID of a process on a specific node in the range [0, L-1].

    Best Practice: A good rule of thumb is to have one process span a single GPU. This balances I/O and computational costs by providing as many parallel reader streams as there are GPUs.

  11. Train a super-resolution network

    main

    Use main.py to train a super-resolution network using an efficient sub-pixel convolution layer. The training process uses the BSD300 dataset (crops from 200 training images) and saves model snapshots after every epoch with the naming convention model_epoch_<epoch_number>.pth.

    python main.py --upscale_factor 3 --batchSize 4 --testBatchSize 100 --nEpochs 30 --lr 0.001 --accel

  12. Run Multi-processing Distributed Data Parallel (DDP) training

    main

    For optimal performance on CUDA, use the nccl backend with --multiprocessing-distributed.

    Note: XPU multiprocessing is not supported as of PyTorch 2.6.

    Single node, multiple GPUs

    Use --dist-url with a local address and a free port, and set --world-size to 1 and --rank to 0.

    Multiple nodes

    Set --dist-url to the IP address of the master node (Node 0). Each node must have a unique --rank (e.g., Node 0 is --rank 0, Node 1 is --rank 1).

    # Single node, multiple GPUs
python main.py -a resnet50 --dist-url 'tcp://127.0.0.1:FREEPORT' --dist-backend 'nccl' --multiprocessing-distributed --world-size 1 --rank 0 [imagenet-folder]

# Multiple nodes (Node 0)
python main.py -a resnet50 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' --dist-backend 'nccl' --multiprocessing-distributed --world-size 2 --rank 0 [imagenet-folder]

# Multiple nodes (Node 1)
python main.py -a resnet50 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' --dist-backend 'nccl' --multiprocessing-distributed --world-size 2 --rank 1 [imagenet-folder]