Skip to content

CS336-RUST-EQUIVALENT.md

Latest commit

 

History

History
100 lines (77 loc) · 8.24 KB

CS336-RUST-EQUIVALENT.md

File metadata and controls

100 lines (77 loc) · 8.24 KB
Raw

CS336 Rust Equivalent

Status: original Rust curriculum design.

This track is based on the public structure of Stanford CS336: Language Modeling from Scratch. It is not Stanford coursework, not a mirror, and not a copy of the lectures, handouts, tests, or code.

The goal is to build an equivalent learning journey in Rust:

language modeling from scratch
  -> typed Rust implementation
  -> newtype-protected meaning
  -> category-theory intuition about maps and composition

For the scraped public course map, see CS336 source map. For the object/map/newtype bridge that connects the core lessons to this track, see Concept Atlas and the executable primer in code/category_lens, including the first ReviewedCompositionTrace -> PublicCompositionTrace release boundary.

Audience

This track is for learners who have finished the core repo path through 06 Attention and want a serious language-modeling systems path.

Expected background:

  • enough Rust to read structs, enums, traits, Result, tests, and iterators
  • basic linear algebra from the earlier lessons
  • patience for systems details such as memory, shapes, batching, and measurement

Learning Objectives

After this track, the learner should be able to:

  • implement the pieces of a tiny language model in Rust
  • explain tokenizer, dataset, Transformer, optimizer, inference, evaluation, and alignment boundaries
  • use newtypes to distinguish tokens, positions, logits, losses, metrics, and resource budgets
  • read model training as composition plus feedback
  • reason about scaling, data quality, and systems constraints without treating them as separate from ML

Lecture-To-Rust Map

Unit Public CS336 topic Original Rust lab direction
1 Overview, tokenization build typed RawText -> TokenSequence -> TokenIds maps with public text review in code/lm_basics
2 Resource accounting count FLOPs, bytes moved, arithmetic intensity, and public measurement reports for tiny layers
3 Architectures, hyperparameters model TransformerConfig with validated dimensions, head-width checks, parameter budgets, and public encoder trace review in code/transformer
4 Attention alternatives, mixture of experts compare attention maps and route tokens through typed experts in code/transformer
5 GPUs, TPUs model accelerator memory hierarchy as typed byte, bandwidth, memory-tier, transfer-time, and public-report maps in code/systems
6 Kernels and compilation write tiled CPU kernels, typed resource traces, and public report reviews in code/kernels before discussing GPU lowering
7 Parallelism I split batches and parameters while preserving ownership, shape contracts, and public trace boundaries in code/parallelism
8 Parallelism II reason about communication, synchronization, failure boundaries, and public report review in code/parallelism
9 Scaling laws I log small experiments and fit a simple power-law curve
10 Inference implement decoding, typed sampling controls, KV-cache traces, latency budgets, and public trace reviews in code/inference
11 Scaling laws II use fitted curves and public-report review to compare data, model, and compute tradeoffs in code/scaling
12 Evaluation build metric newtypes, deterministic evaluation records, and public report reviews in code/evaluation
13 Data sources stream documents through DocumentId, Source, CorpusShard, and public manifest types in code/data
14 Data filtering implement filters, dedup keys, mixture weights, provenance checks, and public/private source boundaries in code/data
15 SFT and RLHF represent instruction examples and preference pairs with explicit roles
16 RLVR model verifier feedback and rollout traces as typed learning signals
17 Alignment systems design auditable alignment pipelines with checked state transitions and public-release review in code/alignment
18 Guest systems synthesis write a design review of the full Rust language-modeling stack
19 Final synthesis present a tested Rust artifact and explain the invariants it protects

Assignment Sequence

The Rust sequence mirrors the public assignment themes while staying original:

  1. R1 Basics: tokenizer, checked token IDs, next-token batches, loss, public text review, and a tiny trainable language-modeling core in code/lm_basics.
  2. R2 Systems: profiling, memory accounting, attention FLOP estimates, median timings, arithmetic intensity, and public report review in code/systems.
  3. R3 Scaling: experiment logs, scaling curves, component ablations, forecast limitations, and public report review in code/scaling.
  4. R4 Data: corpus ingestion, filtering, deduplication, sampling mixtures, and public corpus manifests in code/data.
  5. R5 Alignment: supervised finetuning data, preference pairs, verifier feedback, public release review, and safety notes in code/alignment.

Repository Integration

The current repo already provides the first conceptual bridge:

  • 03 Neuron: typed prediction, update roles, and public training-step review
  • 04 Learning: feedback loops and epochs
  • 05 MLP: hidden representations, shape flow, and public trace review
  • 06 Attention: query/key/value roles, scores, weights, value mixing, and public trace review
  • 07 Transformer: typed encoder architecture and public encoder trace review
  • code/category_lens: executable object/map/composition primer with rust_ml_category_lens and public composition-trace review
  • code/lm_basics: first executable CS336 R1 artifact for text-to-loss language modeling and public text review
  • code/systems: first executable CS336 R2 artifact for typed resource accounting and public report review
  • code/kernels: first executable kernels artifact for typed tiling, reductions, matrix-vector traces, resource estimates, and public report review
  • code/scaling: first executable CS336 R3 artifact for typed experiment evidence, power-law fitting, and public report review
  • code/data: first executable CS336 R4 artifact for typed corpus preparation and public manifest boundaries
  • code/evaluation: first executable evaluation artifact for typed examples, predictions, reports, run comparison, and public report review
  • code/inference: first executable inference artifact for typed decoding controls, KV-cache traces, latency budgets, and public trace review
  • code/parallelism: first executable parallelism artifact for typed ranks, sharding plans, collective traces, pipeline schedules, and public report review
  • code/alignment: first executable CS336 R5 artifact for typed post-training signals and public release review

Public-Safe Use

Use the Stanford page as a public reference for topic order and assignment themes. Do not copy handout wording, tests, starter code, private course communication, or non-public materials.

When adding Rust labs, write original prompts, original tests, and original explanations. The learning target can be equivalent; the artifact must be this repo's own work.