CS336 Assignments Implementation

This repository contains the implementation of assignments for the CS336 course (Spring 2025). The project covers various aspects of Large Language Models (LLMs), from basic Transformer architecture to systems, scaling laws, data processing, and alignment.

Project Structure

The repository is organized into five main assignments:

• Assignment 1: Basics

  • Implementation of core Transformer components (Attention, BPE, etc.).
  • Basics of language modeling.

• Assignment 2: Systems

  • Distributed training techniques.
  • Data Distributed Parallelism (DDP) and optimizer sharding.

• Assignment 3: Scaling

  • Investigation of scaling laws.
  • Compute-optimal training (IsoFLOPs).

• Assignment 4: Data

  • Data processing pipelines for LLM pre-training.
  • Deduplication, PII redaction, and quality filtering.

• Assignment 5: Alignment

  • Techniques for aligning LLMs with human intent.
  • Supervised Fine-Tuning (SFT), DPO, and RLHF (GRPO).

Getting Started

Prerequisites

This project uses uv for dependency management. Ensure you have uv installed.

Installation & Usage

Each assignment is a self-contained Python project. To work on a specific assignment, navigate to its directory and sync dependencies.

Example for Assignment 1:

cd assignment1-basics
uv sync
uv run pytest

Please refer to the README.md within each assignment directory for specific instructions and details.