Phase 01 — PyTorch & TensorFlow Framework Internals

Difficulty: ⭐⭐⭐⭐☆
Estimated Time: 4 weeks (80–100 hours)
Roles Supported: All — this is the bedrock phase. Every downstream topic (quantization, torch.compile, NPU deployment) requires internalizing what happens below the Python API.

Why This Phase Exists

Every Senior Staff ML engineer who influences Qualcomm's framework stack will be asked — in an interview or on the first week of the job — to trace execution from loss.backward() down to the CUDA kernel. If you can't do that, you cannot meaningfully review PRs touching the dispatcher, write a custom operator that behaves correctly under torch.compile, or diagnose why a quantized model produces NaN after a specific graph transformation.

PyTorch's execution model has three distinct layers that every practitioner conflates: the Python eager frontend, the dispatcher (C++ c10), and the backend kernels. TensorFlow's execution model has evolved through three paradigms (eager, tf.function + AutoGraph, tf.experimental.mlir_bridge) that coexist in production code. Understanding both frameworks at this depth is what separates framework users from framework engineers.

This phase will make you dangerous enough to write test cases that expose correctness bugs in framework patches, write custom ops that are safe under all dispatch keys (CPU, CUDA, autograd, torch.compile, functional), and reason about TF's tracing semantics when a colleague reports that "the model gives different results in graph mode."

Concepts

  • PyTorch dispatch stack: Python → C++ dispatcher (c10::Dispatcher) → dispatch keys (CPU, CUDA, Autograd, Functionalize, Meta) → kernel registration via TORCH_LIBRARY / TORCH_LIBRARY_IMPL
  • Autograd engine: DAG construction during forward pass, AccumulateGrad, Engine::execute(), gradient accumulation, retain_graph, create_graph for higher-order gradients
  • torch.autograd.Function: custom forward/backward, save_for_backward, composite functions, double backward
  • Eager vs compiled mode: what torch.compile changes at the boundary — TorchDynamo captures bytecode into FX IR, TorchInductor lowers to Triton/C++
  • torch.fx symbolic tracing: Interpreter, Transformer, Graph, Node, proxy mechanism, limitations (control flow, data-dependent shapes)
  • Custom C++ operators: TORCH_LIBRARY schema, TORCH_LIBRARY_IMPL, pybind11 vs torch::Library, OpCheck, structured kernels
  • TF graph execution: Session.run() (TF1) vs tf.function + tf.Graph (TF2), ConcreteFunction, FuncGraph, AutoGraph transformation rules
  • tf.function tracing semantics: Python-side effects, tf.Variable vs Python variables, retracing triggers (new input signatures, Python closures)
  • TF SavedModel format: tf.saved_model.save, SignatureDef, asset handling, interop with TFLite/ONNX
  • Memory management: PyTorch caching allocator, CUDA unified memory, TF tensor lifetimes in sessions vs eager
  • Numerical precision: FP32 vs FP16 vs BF16 accumulation, catastrophic cancellation, Kahan summation, determinism flags
  • Operator decomposition: how PyTorch's _decomp registry breaks complex ops into primitive ops for portability

Labs

Lab 01 — Custom Autograd Engine from Scratch

FieldValue
GoalImplement a mini automatic differentiation library (minigrad) that supports reverse-mode autodiff over a computation graph, passing numerical gradient checks against PyTorch on 15 standard test cases
ConceptsReverse-mode AD, topological sort, Function subclassing, Tensor wrapper, backward() dispatch
Steps1. Implement Tensor class wrapping np.ndarray with grad, grad_fn, requires_grad; 2. Implement Function base class with forward, backward, save_for_backward; 3. Implement add, mul, matmul, relu, softmax, log_softmax, cross_entropy as Function subclasses; 4. Implement backward() via topological sort + reverse traversal; 5. Implement no_grad() context manager; 6. Run gradcheck style numerical verification vs PyTorch
StackPython 3.10, NumPy 1.26, PyTorch (for verification only)
DatasetsSynthetic: random MLPs on MNIST features (loaded from sklearn)
Outputminigrad/ package; all 15 gradient checks pass within 1e-5 relative error
How to Testpython test_lab.py — runs gradcheck on all Function subclasses; trains a 2-layer MLP for 100 steps and validates loss decreases monotonically
Talking PointsWhy topological sort? What breaks with cycles? How does PyTorch's C++ engine differ from this pure Python version? When would you use create_graph=True?
Resume BulletImplemented reverse-mode AD engine from scratch in Python; passed numerical gradient checks for 15 op types within 1e-5 relative tolerance, demonstrating mastery of PyTorch autograd internals
ExtensionsAdd second-order gradient support; implement vmap style batched AD; add XLA-style HLO emission

Lab 02 — Custom C++ PyTorch Operator

FieldValue
GoalWrite a custom warp_softmax C++ operator registered via TORCH_LIBRARY, including a CPU kernel (SIMD-friendly), a CUDA kernel (one warp per row), and metadata for torch.compile compatibility
ConceptsTORCH_LIBRARY, TORCH_LIBRARY_IMPL, dispatch keys, at::Tensor, structured kernels, OpCheck, pybind11, CMake build system, torch.compile compatibility rules
Steps1. Define operator schema in TORCH_LIBRARY; 2. Implement CPU kernel with vectorized exp+sum; 3. Implement CUDA kernel (32 threads/warp, __shfl_xor_sync reduction); 4. Register both impls via TORCH_LIBRARY_IMPL; 5. Write pybind11 PYBIND11_MODULE wrapper; 6. Add Abstract (Meta) kernel for torch.compile shape propagation; 7. Run torch.library.opcheck(); 8. Benchmark vs torch.nn.functional.softmax
StackC++17, CUDA 12.x, CMake 3.20+, PyTorch 2.3+ C++ extension API, torch.utils.cpp_extension
DatasetsSynthetic: torch.randn(batch, seq_len, vocab_size) tensors
Outputwarp_softmax.so loadable via torch.ops.load_library(); 10–20% latency improvement vs baseline on sequence length ≥ 4096
How to Testpytest test_lab.py — tests: (1) output matches F.softmax within 1e-6, (2) OpCheck passes all dispatch keys, (3) works under torch.compile, (4) gradients match via gradcheck
Talking PointsWhat is a dispatch key and why does order matter? What is the Meta kernel and why does torch.compile require it? How does TORCH_LIBRARY_IMPL(..., CUDA) interact with device placement?
Resume BulletWrote production-ready custom CUDA softmax op in C++/CUDA registered via TORCH_LIBRARY; passed opcheck on all dispatch keys and achieved 1.18x speedup over F.softmax at vocab=32K
ExtensionsAdd torch.amp (AMP) support; add FP8 kernel path; add torch.vmap batch support via vmap dispatch key

Lab 03 — TF tf.function Tracing & Graph Analysis

FieldValue
GoalBuild a tool that (a) traces a TF2 model under tf.function, (b) visualizes the concrete FuncGraph node-by-node, (c) identifies retracing triggers, and (d) exports to TFLite and validates accuracy delta
Conceptstf.function, AutoGraph, ConcreteFunction, FuncGraph, tracing vs execution, Python side effects, input_signature, reduce_retracing, SavedModel, TFLite conversion
Steps1. Implement a trace_and_inspect(fn, inputs) utility that captures the FuncGraph, counts nodes by op type, and prints a summary; 2. Demonstrate 3 retracing trigger scenarios (new shape, new Python value, new closure); 3. Build a custom BERT-tiny model with tf.function on call(); 4. Use tf.saved_model.save then reload and verify outputs match; 5. Convert to TFLite (dynamic range quant) and measure accuracy delta on MNLI dev set; 6. Annotate the FuncGraph with compute cost estimates
StackTensorFlow 2.15, TFLite 2.15, Python 3.10, tensorflow_datasets
DatasetsMNLI dev set (from tensorflow_datasets); BERT-tiny weights from HuggingFace
Outputgraph_inspector.py tool; FuncGraph visualization (JSON + text); MNLI accuracy: FP32 baseline vs TFLite dynamic-range within 0.5%
How to Testpython test_lab.py — checks: retracing count matches expected, FuncGraph node count is deterministic, TFLite output matches TF output within 1e-4 for 100 test samples
Talking PointsWhen does tf.function retrace vs reuse? What are tf.Variable constraints inside traced code? How does AutoGraph transform Python if/for? When is TFLite conversion not lossless?
Resume BulletBuilt TF2 graph analysis tool that identifies retracing triggers and validates TFLite conversion accuracy; demonstrated <0.3% MNLI accuracy degradation under dynamic-range quantization
ExtensionsAdd MLIR bridge analysis (tf.experimental.mlir_bridge); profile XLA compilation cost; instrument with TF Profiler

Deliverables Checklist

  • minigrad package with all 15 gradient checks passing
  • warp_softmax.so with OpCheck passing and benchmarks documented
  • TF graph inspector tool with FuncGraph visualization output
  • All test_lab.py suites pass: pytest phase-01-*/lab-0*/test_lab.py
  • Personal notes in each HITCHHIKERS-GUIDE.md section annotated with your insights

Interview Relevance

  • "Explain how PyTorch's backward() traverses the computation graph." (you can draw the DAG, explain topological sort, explain AccumulateGrad nodes)
  • "What is a dispatch key in PyTorch? How does torch.compile use it?" (you can explain Meta kernels, Functionalize key, and why OpCheck tests all of them)
  • "When does tf.function retrace?" (you can enumerate all 5 retracing scenarios from memory after Lab 03)
  • "How would you debug a NaN appearing after quantization in only graph mode?" (you have the mental model to distinguish Python-side bugs from kernel-side bugs)
  • "Write a custom backward for a fused attention operation." (Lab 01 gives you the foundations to extend on a whiteboard)