cuda=torch.device('cuda')# Default CUDA devicecuda0=torch.device('cuda:0')cuda2=torch.device('cuda:2')# GPU 2 (these are 0-indexed)x=torch.tensor([1.,2.],device=cuda0)# x.device is device(type='cuda', index=0)y=torch.tensor([1.,2.]).cuda()# y.device is device(type='cuda', index=0)withtorch.cuda.device(1):# allocates a tensor on GPU 1a=torch.tensor([1.,2.],device=cuda)# transfers a tensor from CPU to GPU 1b=torch.tensor([1.,2.]).cuda()# a.device and b.device are device(type='cuda', index=1)# You can also use ``Tensor.to`` to transfer a tensor:b2=torch.tensor([1.,2.]).to(device=cuda)# b.device and b2.device are device(type='cuda', index=1)c=a+b# c.device is device(type='cuda', index=1)z=x+y# z.device is device(type='cuda', index=0)# even within a context, you can specify the device# (or give a GPU index to the .cuda call)d=torch.randn(2,device=cuda2)e=torch.randn(2).to(cuda2)f=torch.randn(2).cuda(cuda2)# d.device, e.device, and f.device are all device(type='cuda', index=2)

Warning

Old settings withallow_tf32 as follows is going to be deprecated. We suggest to use the above new settings forbetter control. And we do not support to use mix of old and new settings.

CUDA streams#

ACUDA stream is a linear sequence of execution that belongs to a specificdevice. You normally do not need to create one explicitly: by default, eachdevice uses its own “default” stream.

Operations inside each stream are serialized in the order they are created,but operations from different streams can execute concurrently in anyrelative order, unless explicit synchronization functions (such assynchronize() orwait_stream()) areused. For example, the following code is incorrect:

cuda=torch.device('cuda')s=torch.cuda.Stream()# Create a new stream.A=torch.empty((100,100),device=cuda).normal_(0.0,1.0)withtorch.cuda.stream(s):# sum() may start execution before normal_() finishes!B=torch.sum(A)

When the “current stream” is the default stream, PyTorch automatically performsnecessary synchronization when data is moved around, as explained above.However, when using non-default streams, it is the user’s responsibility toensure proper synchronization. The fixed version of this example is:

cuda=torch.device('cuda')s=torch.cuda.Stream()# Create a new stream.A=torch.empty((100,100),device=cuda).normal_(0.0,1.0)s.wait_stream(torch.cuda.default_stream(cuda))# NEW!withtorch.cuda.stream(s):B=torch.sum(A)A.record_stream(s)# NEW!

There are two new additions. Thetorch.cuda.Stream.wait_stream() callensures that thenormal_() execution has finished before we start runningsum(A) on a side stream. Thetorch.Tensor.record_stream() (see formore details) ensures that we do not deallocate A beforesum(A) hascompleted. You can also manually wait on the stream at some later point intime withtorch.cuda.default_stream(cuda).wait_stream(s) (note that itis pointless to wait immediately, since that will prevent the stream executionfrom running in parallel with other work on the default stream.) See thedocumentation fortorch.Tensor.record_stream() on more details on whento use one or another.

Note that this synchronization is necessary even when there is noread dependency, e.g., as seen in this example:

cuda=torch.device('cuda')s=torch.cuda.Stream()# Create a new stream.A=torch.empty((100,100),device=cuda)s.wait_stream(torch.cuda.default_stream(cuda))# STILL REQUIRED!withtorch.cuda.stream(s):A.normal_(0.0,1.0)A.record_stream(s)

Despite the computation ons not reading the contents ofA and noother uses ofA, it is still necessary to synchronize, becauseAmay correspond to memory reallocated by the CUDA caching allocator, withpending operations from the old (deallocated) memory.

Stream semantics of backward passes#

Each backward CUDA op runs on the same stream that was used for its corresponding forward op.If your forward pass runs independent ops in parallel on different streams,this helps the backward pass exploit that same parallelism.

The stream semantics of a backward call with respect to surrounding ops are the sameas for any other call. The backward pass inserts internal syncs to ensure this even whenbackward ops run on multiple streams as described in the previous paragraph.More concretely, when callingautograd.backward,autograd.grad, ortensor.backward,and optionally supplying CUDA tensor(s) as the initial gradient(s) (e.g.,autograd.backward(...,grad_tensors=initial_grads),autograd.grad(...,grad_outputs=initial_grads), ortensor.backward(...,gradient=initial_grad)),the acts of

1. optionally populating initial gradient(s),

2. invoking the backward pass, and

3. using the gradients

have the same stream-semantics relationship as any group of ops:

s=torch.cuda.Stream()# Safe, grads are used in the same stream context as backward()withtorch.cuda.stream(s):loss.backward()usegrads# Unsafewithtorch.cuda.stream(s):loss.backward()usegrads# Safe, with synchronizationwithtorch.cuda.stream(s):loss.backward()torch.cuda.current_stream().wait_stream(s)usegrads# Safe, populating initial grad and invoking backward are in the same stream contextwithtorch.cuda.stream(s):loss.backward(gradient=torch.ones_like(loss))# Unsafe, populating initial_grad and invoking backward are in different stream contexts,# without synchronizationinitial_grad=torch.ones_like(loss)withtorch.cuda.stream(s):loss.backward(gradient=initial_grad)# Safe, with synchronizationinitial_grad=torch.ones_like(loss)s.wait_stream(torch.cuda.current_stream())withtorch.cuda.stream(s):initial_grad.record_stream(s)loss.backward(gradient=initial_grad)

withtorch.cuda.stream(s):loss.backward()usegrads

withtorch.cuda.stream(s):loss.backward()torch.cuda.current_stream().wait_stream(s)usegrads

Optimizing memory usage withPYTORCH_CUDA_ALLOC_CONF#

Use of a caching allocator can interfere with memory checking tools such ascuda-memcheck. To debug memory errors usingcuda-memcheck, setPYTORCH_NO_CUDA_MEMORY_CACHING=1 in your environment to disable caching.

The behavior of the caching allocator can be controlled via the environment variablePYTORCH_CUDA_ALLOC_CONF.The format isPYTORCH_CUDA_ALLOC_CONF=<option>:<value>,<option2>:<value2>...Available options:

• backend allows selecting the underlying allocator implementation.Currently, valid options arenative, which uses PyTorch’s nativeimplementation, andcudaMallocAsync, which usesCUDA’s built-in asynchronous allocator.cudaMallocAsync requires CUDA 11.4 or newer. The default isnative.backend applies to all devices used by the process, and can’t bespecified on a per-device basis.

• max_split_size_mb prevents the native allocatorfrom splitting blocks larger than this size (in MB). This can reducefragmentation and may allow some borderline workloads to complete withoutrunning out of memory. Performance cost can range from ‘zero’ to ‘substantial’depending on allocation patterns. Default value is unlimited, i.e. all blockscan be split. Thememory_stats() andmemory_summary() methods are useful for tuning. Thisoption should be used as a last resort for a workload that is abortingdue to ‘out of memory’ and showing a large amount of inactive split blocks.max_split_size_mb is only meaningful withbackend:native.Withbackend:cudaMallocAsync,max_split_size_mb is ignored.

• roundup_power2_divisions helps with rounding the requested allocationsize to nearest power-2 division and making better use of the blocks. Inthe native CUDACachingAllocator, the sizes are rounded up in multipleof blocks size of 512, so this works fine for smaller sizes. However, thiscan be inefficient for large near-by allocations as each will go to differentsize of blocks and reuse of those blocks are minimized. This might createlots of unused blocks and will waste GPU memory capacity. This option enablesthe rounding of allocation size to nearest power-2 division. For example, ifwe need to round-up size of 1200 and if number of divisions is 4,the size 1200 lies between 1024 and 2048 and if we do 4 divisions betweenthem, the values are 1024, 1280, 1536, and 1792. So, allocation size of 1200will be rounded to 1280 as the nearest ceiling of power-2 division.Specify a single value to apply for all allocation sizes or specify anarray of key value pairs to set power-2 division individually for eachpower of two interval. For example to set 1 division for all allocationsunder 256MB, 2 division for allocations between 256MB and 512MB, 4 divisionsfor allocations between 512MB and 1GB and 8 divisions for any larger allocations,set the knob value to: [256:1,512:2,1024:4,>:8].roundup_power2_divisions is only meaningful withbackend:native.Withbackend:cudaMallocAsync,roundup_power2_divisions is ignored.

• max_non_split_rounding_mb will allow non-split blocks for better reuse, eg,

  a 1024MB cached block can be reused for a 512MB allocation request. In the defaultcase, we only allow up to 20MB of rounding of non-split blocks, so a 512MB blockcan only be served with between 512-532 MB size block. If we set the value of thisoption to 1024, it will allow 512-1536 MB size blocks to be used for a 512MB blockwhich increases reuse of larger blocks. This will also help in reducing the stallsin avoiding expensive cudaMalloc calls.

• garbage_collection_threshold helps actively reclaiming unused GPU memory toavoid triggering expensive sync-and-reclaim-all operation (release_cached_blocks),which can be unfavorable to latency-critical GPU applications (e.g., servers).Upon setting this threshold (e.g., 0.8), the allocator will start reclaimingGPU memory blocks if the GPU memory capacity usage exceeds the threshold (i.e.,80% of the total memory allocated to the GPU application). The algorithm prefersto free old & unused blocks first to avoid freeing blocks that are actively beingreused. The threshold value should be between greater than 0.0 and less than 1.0.The default value is set at 1.0.

  garbage_collection_threshold is only meaningful withbackend:native.Withbackend:cudaMallocAsync,garbage_collection_threshold is ignored.

• expandable_segments (experimental, default:False) If set toTrue, this setting instructsthe allocator to create CUDA allocations that can later be expanded to better handle caseswhere a job changing allocation sizes frequently, such as having a changing batch size.Normally for large (>2MB) allocations, the allocator calls cudaMalloc to get allocationsthat are the same size as what the user requests. In the future, parts of theseallocations can be reused for other requests if they are free. This works wellwhen the program makes many requests of exactly the same size or of sizes thateven multiples of that size. Many deep learning models follow this behavior.However, one common exception is when the batch size changes slightly from oneiteration to the next, e.g. in batched inference. When the program runsinitially with batch sizeN, it will make allocations appropriate for that size.If in the future, it runs at sizeN - 1, the existing allocations will still bebig enough. However, if it runs at sizeN + 1, then it will have to make newallocations that are slightly larger. Not all the tensors are the same size.Some might be(N + 1)*A and others(N + 1)*A*B whereA andB are some non-batchdimensions in the model. Because the allocator reuses existing allocations whenthey are big enough, some number of(N + 1)*A allocations will actually fit inthe already existingN*B*A segments, though not perfectly. As the model runs itwill partially fill up all of these segments leaving unusable free slices ofmemory at the end of these segments. The allocator at some point will need tocudaMalloc a new(N + 1)*A*B segment. If there is not enough memory, there isnow no way to recover the slices of memory that are free at the end of existingsegments. With models 50+ layers deep, this pattern might repeat 50+ timescreating many slivers.

  expandable_segments allows the allocator to create a segment initially and thenexpand its size later when more memory is needed. Instead of making one segmentper allocation, it tries to make one segment (per stream) that grows asnecessary. Now when theN + 1 case runs, the allocations will tile nicely intothe one large segment until it fills up. Then more memory is requested andappended to the end of the segment. This process does not create as many sliversof unusable memory, so it is more likely to succeed at finding this memory.

• pinned_use_cuda_host_register option is a boolean flag that determines whether touse the CUDA API’s cudaHostRegister function for allocating pinned memory insteadof the default cudaHostAlloc. When set to True, the memory is allocated using regularmalloc and then pages are mapped to the memory before calling cudaHostRegister.This pre-mapping of pages helps reduce the lock time during the executionof cudaHostRegister.

• pinned_num_register_threads option is only valid when pinned_use_cuda_host_registeris set to True. By default, one thread is used to map the pages. This option allowsusing more threads to parallelize the page mapping operations to reduce the overallallocation time of pinned memory. A good value for this option is 8 based onbenchmarking results.

• pinned_use_background_threads option is a boolean flag to enable background threadfor processing events. This avoids any slow path associated with querying/processing ofevents in the fast allocation path. This feature is disabled by default.

• graph_capture_record_stream_reuse (experimental, default:False)If set toTrue, the CUDA caching allocator will attempt to reclaim device memory duringCUDA Graph capture by using the graph topology (instead of CUDA events) to determinewhen a freed block is safe to reuse. This can reduce peak memory during long captures that freeand reallocate buffers across multiple streams, especially when the capture DAG frequentlyreaches joined frontiers. Note: Enabling this option can significantly increase the time spentcapturing the graph.

Note

WhilecudaMallocManaged offers convenient automatic memory management using CUDA Unified Virtual Memory (UVM),it is not recommended for DL workloads. For DL workloads that fit in GPU memory, explicit placement consistentlyoutperforms UVM, since there are no page faults and access patterns remain predictable. When GPU memory getssaturated, UVM has to perform costly double transfers, evicting pages to CPU before bringing in new ones.

Note

• torch.cuda.MemPool holds a reference to the pool. When you use thetorch.cuda.use_mem_pool context manager, it will also acquire another referenceto the pool. On exit of the context manager, it will release its reference. After that,ideally it should only be tensors holding references to the pool. Once the tensors releasetheir references, the use count of the pool will be 1, reflecting that only thetorch.cuda.MemPool object is holding a reference. Only at that point, can the memoryheld by the pool be returned to the system when the pool’s destructor is called usingdel.

• torch.cuda.MemPool doesn’t currently supportexpandable_segments mode ofCUDACachingAllocator.

• NCCL has specific requirements for a buffer to be compatible with NVLS reductions.These requirements can be broken in a dynamic workload, for instance, the buffer beingsent to NCCL by the CUDACachingAllocator might be split and hence, not correctly aligned.In those cases, NCCL can use a fallback algorithm instead of NVLS.

• Allocators likencclMemAlloc can use more memory than requested, due to alignmentrequirements (CU_MULTICAST_GRANULARITY_RECOMMENDED,CU_MULTICAST_GRANULARITY_MINIMUM),and can cause your workload to run out of memory.

Device-agnostic code#

Due to the structure of PyTorch, you may need to explicitly writedevice-agnostic (CPU or GPU) code; an example may be creating a new tensor asthe initial hidden state of a recurrent neural network.

The first step is to determine whether the GPU should be used or not. A commonpattern is to use Python’sargparse module to read in user arguments, andhave a flag that can be used to disable CUDA, in combination withis_available(). In the following,args.device results in atorch.device object that can be used to move tensors to CPU or CUDA.

importargparseimporttorchparser=argparse.ArgumentParser(description='PyTorch Example')parser.add_argument('--disable-cuda',action='store_true',help='Disable CUDA')args=parser.parse_args()args.device=Noneifnotargs.disable_cudaandtorch.cuda.is_available():args.device=torch.device('cuda')else:args.device=torch.device('cpu')

Now that we haveargs.device, we can use it to create a Tensor on thedesired device.

x=torch.empty((8,42),device=args.device)net=Network().to(device=args.device)

This can be used in a number of cases to produce device agnostic code. Belowis an example when using a dataloader:

cuda0=torch.device('cuda:0')# CUDA GPU 0fori,xinenumerate(train_loader):x=x.to(cuda0)

When working with multiple GPUs on a system, you can use theCUDA_VISIBLE_DEVICES environment flag to manage which GPUs are available toPyTorch. As mentioned above, to manually control which GPU a tensor is createdon, the best practice is to use atorch.cuda.device context manager.

print("Outside device is 0")# On device 0 (default in most scenarios)withtorch.cuda.device(1):print("Inside device is 1")# On device 1print("Outside device is still 0")# On device 0

If you have a tensor and would like to create a new tensor of the same type onthe same device, then you can use atorch.Tensor.new_* method(seetorch.Tensor).Whilst the previously mentionedtorch.* factory functions(Creation Ops) depend on the current GPU context andthe attributes arguments you pass in,torch.Tensor.new_* methods preservethe device and other attributes of the tensor.

This is the recommended practice when creating modules in which newtensors need to be created internally during the forward pass.

cuda=torch.device('cuda')x_cpu=torch.empty(2)x_gpu=torch.empty(2,device=cuda)x_cpu_long=torch.empty(2,dtype=torch.int64)y_cpu=x_cpu.new_full([3,2],fill_value=0.3)print(y_cpu)tensor([[0.3000,0.3000],[0.3000,0.3000],[0.3000,0.3000]])y_gpu=x_gpu.new_full([3,2],fill_value=-5)print(y_gpu)tensor([[-5.0000,-5.0000],[-5.0000,-5.0000],[-5.0000,-5.0000]],device='cuda:0')y_cpu_long=x_cpu_long.new_tensor([[1,2,3]])print(y_cpu_long)tensor([[1,2,3]])

If you want to create a tensor of the same type and size of another tensor, andfill it with either ones or zeros,ones_like() orzeros_like() are provided as convenient helper functions (whichalso preservetorch.device andtorch.dtype of a Tensor).

x_cpu=torch.empty(2,3)x_gpu=torch.empty(2,3)y_cpu=torch.ones_like(x_cpu)y_gpu=torch.zeros_like(x_gpu)

Use pinned memory buffers#

Host to GPU copies are much faster when they originate from pinned (page-locked)memory. CPU tensors and storages expose apin_memory()method, that returns a copy of the object, with data put in a pinned region.

Also, once you pin a tensor or storage, you can use asynchronous GPU copies.Just pass an additionalnon_blocking=True argument to ato() or acuda() call. This can be usedto overlap data transfers with computation.

You can make theDataLoader return batches placed inpinned memory by passingpin_memory=True to its constructor.

Use nn.parallel.DistributedDataParallel instead of multiprocessing or nn.DataParallel#

Most use cases involving batched inputs and multiple GPUs should default tousingDistributedDataParallel to utilize morethan one GPU.

There are significant caveats to using CUDA models withmultiprocessing; unless care is taken to meet the data handlingrequirements exactly, it is likely that your program will have incorrect orundefined behavior.

It is recommended to useDistributedDataParallel,instead ofDataParallel to do multi-GPU training, even ifthere is only a single node.

The difference betweenDistributedDataParallel andDataParallel is:DistributedDataParalleluses multiprocessing where a process is created for each GPU, whileDataParallel uses multithreading. By using multiprocessing,each GPU has its dedicated process, this avoids the performance overhead causedby GIL of Python interpreter.

If you useDistributedDataParallel, you could usetorch.distributed.launch utility to launch your program, seeLaunch utility.

Why CUDA Graphs?#

Replaying a graph sacrifices the dynamic flexibility of typical eager execution in exchange forgreatly reduced CPU overhead. A graph’s arguments and kernels are fixed, so a graph replayskips all layers of argument setup and kernel dispatch, including Python, C++, and CUDA driveroverheads. Under the hood, a replay submits the entire graph’s work to the GPU witha single call tocudaGraphLaunch. Kernels in a replay also execute slightly fasteron the GPU, but eliding CPU overhead is the main benefit.

You should try CUDA graphs if all or part of your network is graph-safe (usually this meansstatic shapes and static control flow, but see the otherconstraints)and you suspect its runtime is at least somewhat CPU-limited.

PyTorch API#

PyTorch exposes graphs via a rawtorch.cuda.CUDAGraph classand two convenience wrappers,torch.cuda.graph andtorch.cuda.make_graphed_callables.

torch.cuda.graph is a simple, versatile context manager thatcaptures CUDA work in its context.Before capture, warm up the workload to be captured by runninga few eager iterations. Warmup must occur on a side stream.Because the graph reads from and writes to the same memory addresses in everyreplay, you must maintain long-lived references to tensors that holdinput and output data during capture.To run the graph on new input data, copy new data to the capture’s input tensor(s),replay the graph, then read the new output from the capture’s output tensor(s).Example:

g=torch.cuda.CUDAGraph()# Placeholder input used for capturestatic_input=torch.empty((5,),device="cuda")# Warmup before captures=torch.cuda.Stream()s.wait_stream(torch.cuda.current_stream())withtorch.cuda.stream(s):for_inrange(3):static_output=static_input*2torch.cuda.current_stream().wait_stream(s)# Captures the graph# To allow capture, automatically sets a side stream as the current stream in the contextwithtorch.cuda.graph(g):static_output=static_input*2# Fills the graph's input memory with new data to compute onstatic_input.copy_(torch.full((5,),3,device="cuda"))g.replay()# static_output holds the resultsprint(static_output)# full of 3 * 2 = 6# Fills the graph's input memory with more data to compute onstatic_input.copy_(torch.full((5,),4,device="cuda"))g.replay()print(static_output)# full of 4 * 2 = 8

SeeWhole-network capture,Usage with torch.cuda.amp, andUsage with multiple streamsfor realistic and advanced patterns.

make_graphed_callables is more sophisticated.make_graphed_callables accepts Python functions andtorch.nn.Modules. For each passed function or Module,it creates separate graphs of the forward-pass and backward-pass work. SeePartial-network capture.

Whole-network capture#

If your entire network is capturable, you can capture and replay an entire iteration:

N,D_in,H,D_out=640,4096,2048,1024model=torch.nn.Sequential(torch.nn.Linear(D_in,H),torch.nn.Dropout(p=0.2),torch.nn.Linear(H,D_out),torch.nn.Dropout(p=0.1)).cuda()loss_fn=torch.nn.MSELoss()optimizer=torch.optim.SGD(model.parameters(),lr=0.1)# Placeholders used for capturestatic_input=torch.randn(N,D_in,device='cuda')static_target=torch.randn(N,D_out,device='cuda')# warmup# Uses static_input and static_target here for convenience,# but in a real setting, because the warmup includes optimizer.step()# you must use a few batches of real data.s=torch.cuda.Stream()s.wait_stream(torch.cuda.current_stream())withtorch.cuda.stream(s):foriinrange(3):optimizer.zero_grad(set_to_none=True)y_pred=model(static_input)loss=loss_fn(y_pred,static_target)loss.backward()optimizer.step()torch.cuda.current_stream().wait_stream(s)# captureg=torch.cuda.CUDAGraph()# Sets grads to None before capture, so backward() will create# .grad attributes with allocations from the graph's private pooloptimizer.zero_grad(set_to_none=True)withtorch.cuda.graph(g):static_y_pred=model(static_input)static_loss=loss_fn(static_y_pred,static_target)static_loss.backward()optimizer.step()real_inputs=[torch.rand_like(static_input)for_inrange(10)]real_targets=[torch.rand_like(static_target)for_inrange(10)]fordata,targetinzip(real_inputs,real_targets):# Fills the graph's input memory with new data to compute onstatic_input.copy_(data)static_target.copy_(target)# replay() includes forward, backward, and step.# You don't even need to call optimizer.zero_grad() between iterations# because the captured backward refills static .grad tensors in place.g.replay()# Params have been updated. static_y_pred, static_loss, and .grad# attributes hold values from computing on this iteration's data.

Partial-network capture#

If some of your network is unsafe to capture (e.g., due to dynamic control flow,dynamic shapes, CPU syncs, or essential CPU-side logic), you can run the unsafepart(s) eagerly and usetorch.cuda.make_graphed_callables() to graph onlythe capture-safe part(s).

By default, callables returned bymake_graphed_callables()are autograd-aware, and can be used in the training loop as direct replacementsfor the functions ornn.Modules you passed.

make_graphed_callables() internally createsCUDAGraph objects, runs warmup iterations, and maintainsstatic inputs and outputs as needed. Therefore (unlike withtorch.cuda.graph) you don’t need to handle those manually.

In the following example, data-dependent dynamic control flow means thenetwork isn’t capturable end-to-end, butmake_graphed_callables()lets us capture and run graph-safe sections as graphs regardless:

N,D_in,H,D_out=640,4096,2048,1024module1=torch.nn.Linear(D_in,H).cuda()module2=torch.nn.Linear(H,D_out).cuda()module3=torch.nn.Linear(H,D_out).cuda()loss_fn=torch.nn.MSELoss()optimizer=torch.optim.SGD(chain(module1.parameters(),module2.parameters(),module3.parameters()),lr=0.1)# Sample inputs used for capture# requires_grad state of sample inputs must match# requires_grad state of real inputs each callable will see.x=torch.randn(N,D_in,device='cuda')h=torch.randn(N,H,device='cuda',requires_grad=True)module1=torch.cuda.make_graphed_callables(module1,(x,))module2=torch.cuda.make_graphed_callables(module2,(h,))module3=torch.cuda.make_graphed_callables(module3,(h,))real_inputs=[torch.rand_like(x)for_inrange(10)]real_targets=[torch.randn(N,D_out,device="cuda")for_inrange(10)]fordata,targetinzip(real_inputs,real_targets):optimizer.zero_grad(set_to_none=True)tmp=module1(data)# forward ops run as a graphiftmp.sum().item()>0:tmp=module2(tmp)# forward ops run as a graphelse:tmp=module3(tmp)# forward ops run as a graphloss=loss_fn(tmp,target)# module2's or module3's (whichever was chosen) backward ops,# as well as module1's backward ops, run as graphsloss.backward()optimizer.step()

Usage with torch.cuda.amp#

For typical optimizers,GradScaler.step syncsthe CPU with the GPU, which is prohibited during capture. To avoid errors, either usepartial-network capture, or (if forward, loss,and backward are capture-safe) capture forward, loss, and backward but not theoptimizer step:

# warmup# In a real setting, use a few batches of real data.s=torch.cuda.Stream()s.wait_stream(torch.cuda.current_stream())withtorch.cuda.stream(s):foriinrange(3):optimizer.zero_grad(set_to_none=True)withtorch.cuda.amp.autocast():y_pred=model(static_input)loss=loss_fn(y_pred,static_target)scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()torch.cuda.current_stream().wait_stream(s)# captureg=torch.cuda.CUDAGraph()optimizer.zero_grad(set_to_none=True)withtorch.cuda.graph(g):withtorch.cuda.amp.autocast():static_y_pred=model(static_input)static_loss=loss_fn(static_y_pred,static_target)scaler.scale(static_loss).backward()# don't capture scaler.step(optimizer) or scaler.update()real_inputs=[torch.rand_like(static_input)for_inrange(10)]real_targets=[torch.rand_like(static_target)for_inrange(10)]fordata,targetinzip(real_inputs,real_targets):static_input.copy_(data)static_target.copy_(target)g.replay()# Runs scaler.step and scaler.update eagerlyscaler.step(optimizer)scaler.update()

Usage with multiple streams#

Capture mode automatically propagates to any streams that sync with a capturing stream.Within capture, you may expose parallelism by issuing calls to different streams,but the overall stream dependency DAG must branch out from theinitial capturing stream after capture begins and rejoin the initial streambefore capture ends:

withtorch.cuda.graph(g):# at context manager entrance, torch.cuda.current_stream()# is the initial capturing stream# INCORRECT (does not branch out from or rejoin initial stream)withtorch.cuda.stream(s):cuda_work()# CORRECT:# branches out from initial streams.wait_stream(torch.cuda.current_stream())withtorch.cuda.stream(s):cuda_work()# rejoins initial stream before capture endstorch.cuda.current_stream().wait_stream(s)

Usage with DistributedDataParallel#

Graph memory management#

A captured graph acts on the same virtual addresses every time it replays.If PyTorch frees the memory, a later replay can hit an illegal memory access.If PyTorch reassigns the memory to new tensors, the replay can corrupt the valuesseen by those tensors. Therefore, the virtual addresses used by the graph must bereserved for the graph across replays. The PyTorch caching allocator achieves thisby detecting when capture is underway and satisfying the capture’s allocationsfrom a graph-private memory pool. The private pool stays alive until itsCUDAGraph object and all tensors created during capturego out of scope.

Private pools are maintained automatically. By default, the allocator creates aseparate private pool for each capture. If you capture multiple graphs,this conservative approach ensures graph replays never corrupt each other’s values,but sometimes needlessly wastes memory.

g1=torch.cuda.CUDAGraph()g2=torch.cuda.CUDAGraph()# (create static inputs for g1 and g2, run warmups of their workloads...)# Captures g1withtorch.cuda.graph(g1):static_out_1=g1_workload(static_in_1)# Captures g2, hinting that g2 may share a memory pool with g1withtorch.cuda.graph(g2,pool=g1.pool()):static_out_2=g2_workload(static_in_2)static_in_1.copy_(real_data_1)static_in_2.copy_(real_data_2)g1.replay()g2.replay()