penr-oz-neural-network-v4-torch-fsdp

Version 4 Implementation of Neural Network service leveraging pytorch fully sharded data parallel (fsdp)

This repository demonstrates same key concepts in neural networks as in penr-oz-neural-network-v3-torch-ddp with automatic gradient descent calculations relying on PyTorch library leveraging Neural Network (nn) package, Fully Sharded Data Parallel (FSDP2) feature to shard model parameters and optimizer state across multiple GPU (CUDA) devices for memory-efficient scaling, and changes to API to support downloading/sharding training data for local read instead of an API payload.

Implementation follows:

• nn-zero-to-hero lectures
• makemore
• build-nanogpt
• nanoGPT

Backpropagation: Auto Gradient Calculation

The gradients are automatically computed using PyTorch and the PyTorch Neural Network package

Scaling Calculation Speed and Concurrency

This is done by leveraging PyTorch Fully Sharded Data Parallel (FSDP2), which shards parameters, gradients and optimizer state across devices via fully_shard, lowering per-GPU memory versus DDP's full replication.

Quickstart Guide

1. Clone the Repository:

   git clone https://github.com/derinworks/penr-oz-neural-network-v4-torch-fsdp.git
   cd penr-oz-neural-network-v4-torch-fsdp

2. Create and Activate a Virtual Environment:

   • Install python 3.10

     $ python3 --version
     Python 3.10.18

   • Create:

     python3 -m venv venv

   • Activate:
     • On Unix or macOS:

       source venv/bin/activate

     • On Windows:

       venv\Scripts\activate

3. Install Dependencies:

   pip install -r requirements.txt

4. Run the Service:

   python main.py

   or

   uvicorn main:app --log-config log_config.json

5. Interact with the Service Test the endpoints using Swagger at http://127.0.0.1:8000/docs.

6. Interact with the Dashboard Diagnose model training at http://127.0.0.1:8000/dashboard.

Import from Hugging Face

Import a pretrained GPT-2 family model directly from the Hugging Face Hub and make it immediately usable with all existing endpoints (/generate/, /train/, /output/, etc.):

curl -X POST "http://localhost:8000/import/" \
  -H "Content-Type: application/json" \
  -d '{"hf_repo_id": "gpt2", "model_id": "gpt2-imported"}'

Optional fields:

• revision – specific HuggingFace branch / tag / commit (default: null)
• device – PyTorch device, e.g. "cuda" (default: "cpu")

Supported architectures: GPT-2 family (gpt2, gpt2-medium, gpt2-large, gpt2-xl, openai-community/gpt2-*).

7. Quickly spin up the Service in a brand-new Linux VM

   ./run-in-vm.sh

---

Testing and Coverage

To ensure code quality and maintainability, follow these steps to run tests and check code coverage:

1. Run All Tests:

   python -m pytest -v

   The test suite includes 136 tests across 7 test files:

   • test_main.py - API endpoint tests (36 tests)
   • test_neural_net_model.py - Model implementation tests (65 tests)
   • test_neural_net_layers.py - Custom layer tests (12 tests)
   • test_loaders.py - Dataset loader/downloader tests (7 tests)
   • test_mappers.py - Layer/optimizer mapper tests (12 tests)
   • test_gpt_tokenizers.py - Tokenization tests (4 tests)
   • test_fsdp.py - Distributed (FSDP) training tests (29 tests)

2. Run Tests with Coverage: Execute the following commands to run tests and generate a coverage report:

   coverage run -m pytest
   coverage report

3. Generate HTML Coverage Report (Optional): For a detailed coverage report in HTML format:

   coverage html

   Open the htmlcov/index.html file in a web browser to view the report.

Platform-Specific Tests

Some tests require Linux-specific features (e.g., /dev/shm for shared memory caching) and will be automatically skipped on macOS/Windows:

• test_train_* - Full training integration tests with model persistence
• test_cache_miss - Shared memory cache behavior
• test_delete - Model deletion with shared memory cleanup

These tests will run automatically on Linux systems where /dev/shm is available.