# PyTorch for Control Systems and Decision Making

Vincent Moens, Meta | MLOps Podcast | Episode 276 | 55:26
Hosted by Demetrios Brinkmann

Source: https://www.youtube.com/watch?v=JO8bWCsfPnE
Channel: MLOps Community, now AAIF Live (https://www.youtube.com/@AAIFLive-x1r). Summarised by MLOps Talks.
Page: https://mlopstalks.com/talks/pytorch-for-control-systems-and-decision-making
Published: 2024-12-03
Tags: debugging, reinforcement-learning, testing

## TL;DR
- Benchmark pinned-memory transfers instead of assuming that pinning will make CPU-to-CUDA copies faster.
- Avoid in-place tensor updates when possible because they rarely improve speed and make Torch Compile harder to use.
- TorchRL and TensorDict use flexible dictionary-like data structures to support reinforcement learning workloads and broader PyTorch applications.

## Summary
Vincent Moens discusses practical ways to make PyTorch code faster and easier to maintain, with examples from reinforcement learning and TorchRL. He explains why pinned memory can be slower than a direct CPU-to-CUDA transfer, why in-place operations create problems for Torch Compile, and how small reinforcement-learning models can lose substantial time to Python and module-management overhead. He then describes the design of TorchRL and TensorDict. Reinforcement learning covers many data types and algorithm interfaces, so TorchRL uses dictionary-like inputs and outputs, while TensorDict adds tensor operations such as stacking, reshaping, device transfers, and size inspection. Moens also covers Torch Compile debugging, systematic testing with shared test classes, and how to write useful GitHub issues. He recommends Hydra for configuration sweeps and Submitit for launching experiments on Slurm clusters. His examples are aimed at people who need to understand PyTorch behavior rather than rely on common advice without measuring it.

## Key ideas
### Pinned memory should be benchmarked rather than applied by default
[01:38](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=98s)
Moens explains that copying a tensor from RAM into pinned memory is itself a copy. A CPU-to-CUDA transfer passes through pinned memory whether the user calls pin_memory or PyTorch handles the transition internally. Calling pin_memory explicitly can therefore add work instead of removing it. In his experiments, a direct tensor.to CUDA call, or a list of such calls, could be faster than pinning first. His advice is to benchmark the workload. For many CPU-to-CUDA transfers, using non_blocking=True without explicitly calling pin_memory already provides enough speedup.

### In-place tensor updates usually create compiler problems without saving time
[05:43](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=343s)
Moens advises against operations such as add_ or using modules with inplace=True as a general optimization. He says the performance gain is marginal, if there is any gain at all, while in-place changes make Torch Compile harder to handle because several views may refer to the same storage. A compiled graph has to account for those relationships, which can make the generated graph less efficient. His practical recommendation is simple: benchmark in-place code before keeping it, and prefer out-of-place operations unless there is a clear measured reason to do otherwise.

### Small reinforcement-learning models can be limited by Python overhead
[07:32](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=452s)
Reinforcement-learning models are often small and called very frequently, so time spent in Python and module management can outweigh the matrix operations themselves. Moens points to module.train and module.eval as examples. Each call walks through submodules and changes attributes, and repeated switching can consume a noticeable part of a training run. One workaround is to create two module objects that share the same parameter storage, with one kept in training mode and the other in evaluation mode. The same pattern can apply to temporarily changing requires_grad. Because the models are often small, the extra Python objects need little additional memory.

### TorchRL needs abstractions that cover very different reinforcement-learning algorithms
[11:37](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=697s)
Moens says reinforcement learning has a much wider range of data and algorithm interfaces than areas such as computer vision. A policy for one algorithm may return an action distribution and log probability, while a DQN policy returns values that are later converted into actions. Applications also range from games and robotics to portfolio management and autonomous driving. TorchRL handles this variation by giving policies and replay buffers dictionary-like inputs and outputs. This creates a shared pipeline interface without forcing every algorithm into the same tensor signature.

### TensorDict adds tensor behavior to flexible structured data
[15:44](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=944s)
Dictionary objects are useful for heterogeneous reinforcement-learning data, but ordinary dictionaries are awkward to stack, reshape, move between devices, or store efficiently. TensorDict was created to provide those operations while retaining a dictionary-like structure. Moens describes support for stacking TensorDict objects, moving their tensors between CPU and CUDA, distributed communication, and inspecting the byte size of stored data. TensorDict can also hold non-tensor values, such as headlines alongside financial data. Its users later expanded beyond reinforcement learning, including people working on generative models and other forms of machine learning.

### TensorDict supports functional models and broader PyTorch APIs
[19:40](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=1180s)
Moens describes using TensorDict to hold model parameters and run several parameter configurations through a vectorized map instead of looping over configurations one at a time. The parameters can be gathered and stacked, then passed to a functional version of the model. He says this can be faster because the operations are vectorized. The broader design goal is close correspondence with ordinary PyTorch tensor operations, so code written for one tensor can often work with a TensorDict. He also discusses distributed tensors, where a sharded tensor across multiple nodes can be manipulated through a single tensor-like object.

### Torch Compile debugging depends on finding graph breaks and recompilations
[31:55](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=1915s)
Moens explains that Torch Compile converts Python code into graphs and falls back to Python when it reaches unsupported operations. Each fallback creates a graph break. Users can inspect these breaks with torch._dynamo logging, then rewrite code where practical, such as replacing some conditionals with torch.where. Compile also adds guards for input properties and recompiles when those properties change. Moens recommends examining the guards and recompilations instead of treating compilation as a black box. He reports that, after learning how to use it, his team sometimes made models six to eight times faster than their uncompiled versions.

### Good PyTorch contributions explain intent and reduce the reproduction
[37:31](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=2251s)
When submitting a PyTorch or TorchRL issue, Moens wants to know what the contributor was trying to do, not only that a piece of code failed. A useful report includes a minimal reproducible example with as few dependencies and as little code as possible. Contributors should read the relevant documentation and check for existing issues first. Torch Compile bugs can be difficult to reduce because behavior depends on the surrounding stack, so Moens says it is still fine to submit an issue when a minimal example is not possible. The maintainers can ask for logs and other diagnostic output, and contributors can follow up if an issue has gone quiet.

### Shared test classes let TorchRL check features across implementations
[43:38](https://www.youtube.com/watch?v=JO8bWCsfPnE&t=2618s)
TorchRL uses pytest and shared abstract test classes for transforms. A contributor can apply the common test class to a new transform, checking that it works in different thread settings, with replay buffers, and as a module. TensorDict uses a similar pattern across its different implementations, including a data-class form and storage backed by H5 databases. Running the same tests across these variants exposes edge cases, such as what a device transfer should mean when the structure contains non-tensor data. Moens recommends discussing a proposed feature with maintainers before implementing it, then asking how it should fit the existing test system.

## Notable quotes
- Vincent Moens: "The TLDR here is basically benchmark what you're doing. Pin memory might be working for you for some reason, and if it doesn't, so most of the time you were sending from CPU to CUDA, just doing tensor.to CUDA with non_blocking was true will already give you a pretty decent amount of speedup." (04:58)
- Vincent Moens: "The official guideline is basically try not to do things in place. It might break things. It's not going to be faster." (06:48)
- Vincent Moens: "The game for both the developers and the user is to have as few graph breaks as possible." (32:51)
- Vincent Moens: "The takeaway message here is really if you're dealing with Torch Compile and you have an issue and you cannot find a minimal reproducible example, don't worry about that. Submit that issue, let people know, and start the conversation as early as you can." (40:19)

## Tools & references mentioned
- PyTorch
- Meta
- TorchRL
- TensorDict
- CUDA
- NVIDIA
- Torch Compile
- TorchDynamo
- Hydra
- OmegaConf
- Submitit
- Slurm
- pytest
- TensorDictClass
- Llama 3.2
- AlphaGo
- DQN
- TorchVision
- TorchRec
- distributed tensor
- H5 databases
- Weights & Biases

## Who should watch
- You are moving PyTorch tensors between CPU and CUDA and want to know whether pinned memory is helping your workload.
- Your reinforcement-learning code uses small models at high call frequency, so Python and module-management overhead may be affecting runtime.
- You maintain or contribute to PyTorch libraries and need practical advice on Torch Compile diagnostics, tests, and GitHub issue reports.

## Related talks

- [How to Optimize Large AI Models with PyTorch](https://mlopstalks.com/talks/how-to-optimize-large-ai-models-with-pytorch) (Michael Gschwind, Meta Platforms, 57:44)
- [PyTorch: Bridging AI Research and Production](https://mlopstalks.com/talks/pytorch-bridging-ai-research-and-production) (Dmytro Dzhulgakov, Facebook, 52:55)
- [Performance Optimization and Software/Hardware Co-design across PyTorch, CUDA, and NVIDIA GPUs](https://mlopstalks.com/talks/performance-optimization-and-software-hardware-co-design-across-pytorch-cuda) (Chris Fregly, AI performance engineer, startup founder, and investor, 1:25:50)
- [Testing AI Intelligence: The Benchmarking Battle](https://mlopstalks.com/talks/testing-ai-intelligence-the-benchmarking-battle) (Greg Kamradt, Arc Prize, 48:31)
- [Building RedPajama](https://mlopstalks.com/talks/building-redpajama) (Vipul Ved Prakash, Together, 27:52)