PyTorch Developer Podcast

Today I'm going to talk about a famous issue in PyTorch, DataLoader with num_workers > 0 causes memory leak (https://github.com/pytorch/pytorch/issues/13246). This bug is a good opportunity to talk about DataSet/DataLoader design in PyTorch, fork and copy-on-write memory in Linux and Python reference counting; you have to know about all of these things to understand why this bug occurs, but once you do, it also explains why the workarounds help.

Further reading.

• A nice summary of the full issue https://github.com/pytorch/pytorch/issues/13246#issuecomment-905703662
• DataLoader architecture RFC https://github.com/pytorch/pytorch/issues/49440
• Cinder Python https://github.com/facebookincubator/cinder

PyTorch operates on its input data in a batched manner, typically processing multiple batches of an input at once (rather than once at a time, as would be the case in typical programming). In this podcast, we talk a little about the implications of batching operations in this way, and then also about how PyTorch's API is structured for batching (hint: poorly) and how Numpy introduced a concept of ufunc/gufuncs to standardize over broadcasting and batching behavior. There is some overlap between this podcast and previous podcasts about TensorIterator and vmap; you may also be interested in those episodes.

Further reading.

• ufuncs and gufuncs https://numpy.org/doc/stable/reference/ufuncs.html and https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html
• A brief taxonomy of PyTorch operators by shape behavior http://blog.ezyang.com/2020/05/a-brief-taxonomy-of-pytorch-operators-by-shape-behavior/
• Related episodes on TensorIterator and vmap https://pytorch-dev-podcast.simplecast.com/episodes/tensoriterator and https://pytorch-dev-podcast.simplecast.com/episodes/vmap

Python is a single dispatch OO language, but there are some operations such as binary magic methods which implement a simple form of multiple dispatch. torch_function__ (through its Numpy predecessor __array_function) generalizes this mechanism so that invocations of torch.add with different subclasses work properly. This podcast describes how this mechanism works and how it can be used (in an unconventional way) to build composable subclasses ala JAX in functorch.

Further reading:

• This podcast in written form https://dev-discuss.pytorch.org/t/functorch-levels-as-dynamically-allocated-classes/294
• Multiple dispatch resolution rules in the RFC https://github.com/pytorch/rfcs/blob/master/RFC-0001-torch-function-for-methods.md#process-followed-during-a-functionmethod-call

Writing multithreading code has always been a pain, and in PyTorch there are buckets and buckets of multithreading related issues you have to be aware about and deal with when writing code that makes use of it. We'll cover how you interface with multithreading in PyTorch, what goes into implementing those interfaces (thread pools!) and also some miscellaneous stuff like TLS, forks and data structure thread safety that is also relevant.

Further reading:

• TorchScript CPU inference threading documentation https://github.com/pytorch/pytorch/blob/master/docs/source/notes/cpu_threading_torchscript_inference.rst
• c10 thread pool https://github.com/pytorch/pytorch/blob/master/c10/core/thread_pool.h and autograd thread pool https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/engine.cpp
• Tracking issue for TLS propagation across threads https://github.com/pytorch/pytorch/issues/28520

CUDA is asynchronous, CPU is synchronous. Making them play well together can be one of the more thorny and easy to get wrong aspects of the PyTorch API. I talk about why non_blocking is difficult to use correctly, a hypothetical "asynchronous CPU" device which would help smooth over some of the API problems and also why it used to be difficult to implement async CPU (but it's not hard anymore!) At the end, I also briefly talk about how async/sync impedance can also show up in unusual places, namely the CUDA caching allocator.

Further reading.

• CUDA semantics which discuss non_blocking somewhat https://pytorch.org/docs/stable/notes/cuda.html
• Issue requesting async cpu https://github.com/pytorch/pytorch/issues/44343

We talk about gradcheck, the property based testing mechanism that we use to verify the correctness of analytic gradient formulas in PyTorch. I'll talk a bit about testing in general, property based testing and why gradcheck is a particularly useful property based test. There will be some calculus, although I've tried to keep the math mostly to intuitions and pointers on what to read up on elsewhere.

Further reading.

• Gradcheck mechanics, a detailed mathematical explanation of how it works https://pytorch.org/docs/stable/notes/gradcheck.html In particular, it also explains how gradcheck extends to complex numbers
• JAX has a pretty good explanation about vjp and jvp at https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html
• Fast gradcheck tracking issue https://github.com/pytorch/pytorch/issues/53876

torch.use_deterministic_algorithms lets you force PyTorch to use deterministic algorithms. It's very useful for debugging!

There are some errors in the recording: the feature is called torch.use_deterministic_algorithms, and there is not actually a capability to warn (this was in an old version of the PR but taken out), we just error if you hit nondeterministic code.

Docs: https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html#torch.use_deterministic_algorithms

Reference counting is a common memory management technique in C++ but PyTorch does its reference counting in a slightly idiosyncratic way using intrusive_ptr. We'll talk about why intrusive_ptr exists, the reason why refcount bumps are slow in C++ (but not in Python), what's up with const Tensor& everywhere, why the const is a lie and how TensorRef lets you create a const Tensor& from a TensorImpl* without needing to bump your reference count.

Further reading.

• Why you shouldn't feel bad about passing tensor by reference https://dev-discuss.pytorch.org/t/we-shouldnt-feel-bad-about-passing-tensor-by-reference/85
• Const correctness in PyTorch https://github.com/zdevito/ATen/issues/27
• TensorRef RFC https://github.com/pytorch/rfcs/pull/16

Memory layout specifies how the logical multi-dimensional tensor maps its elements onto physical linear memory. Some layouts admit more efficient implementations, e.g., NCHW versus NHWC. Memory layout makes use of striding to allow users to conveniently represent their tensors with different physical layouts without having to explicitly tell every operator what to do.

Further reading.

• Tutorial https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html
• Memory format RFC https://github.com/pytorch/pytorch/issues/19092
• Layout permutation proposal (not implemented) https://github.com/pytorch/pytorch/issues/32078

pytorch-probot is a GitHub application that we use to automate common tasks in GitHub. I talk about what it does and some design philosophy for it. Repo is at: https://github.com/pytorch/pytorch-probot