Introduction
When using PyTorch, sometimes we would like to check if certain variable conditions are met during the forward pass. If a PyTorch program runs on GPU, we would like to have this check to be performed asynchronously on the GPU device stream, without blocking the CPU thread. Otherwise, it is a graph break operation which can disturb the optimization of the computation graph from a neural network compiler, such as torch.compile
. In PyTorch, such asynchronous assertion can be performed using the torch._assert_async
API.
In this blog post, I would like to quickly discuss how to use the torch._assert_async
API and how it is implemented in PyTorch.
PyTorch Asynchronous Assert
The torch._assert_async
is not well documented. What’s different from the torch._assert API is that
torch._assert_async
accepts a boolean tensor whereas torch._assert
accepts a Python boolean value. When the boolean tensor is on GPU, the assertion will be performed asynchronously on the GPU device stream. Since the assertion is performed asynchronously, if the assertion fails, the error will be reported at a later time only when the GPU stream is synchronized with the CPU thread.In the following example, we inserted a torch._assert_async
assertion in a PyTorch model. We will test what happens when the assertion fails and how it is reported asynchronously when torch.compile
is used or not used.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215
|
import argparseimport loggingfrom pathlib import Pathfrom typing import Optionalimport torchimport torch.nn as nnlogging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")class NormalizedLinear(nn.Module): """ Applies a linear transformation followed by Softmax, then validates that output probabilities sum to 1.0 asynchronously on the device stream. """ def __init__(self, in_features: int, out_features: int, trigger_fail: bool = False) -> None: super().__init__() self.linear = nn.Linear(in_features, out_features) self.softmax = nn.Softmax(dim=-1) self.trigger_fail = trigger_fail def forward(self, x: torch.Tensor) -> torch.Tensor: probs = self.softmax(self.linear(x)) if self.trigger_fail: # Force condition to False to demonstrate deferred stream failure cond = torch.tensor(False, device=probs.device) else: prob_sum = probs.sum(dim=-1) # Pure GPU Tensor Ops (abs -> le -> all): Returns a 0-D boolean CUDA Tensor # WITHOUT calling .item() or causing CPU-GPU host synchronization. cond = torch.abs(prob_sum - 1.0).le(1e-5).all() # Non-blocking async check enqueued into the GPU execution stream torch._assert_async( cond, "Assertion Failed: Probabilities must sum to 1.0!") return probs * 10.0def create_profiler( device: torch.device, trace_path: Path, stacks_path: Optional[Path] = None, warmup_steps: int = 3, active_steps: int = 5,) -> torch.profiler.profile: """Configures and returns a PyTorch profiler using torch.profiler.schedule for warmup and active steps.""" # Ensure parent directories exist trace_path.parent.mkdir(parents=True, exist_ok=True) if stacks_path: stacks_path.parent.mkdir(parents=True, exist_ok=True) activities = [torch.profiler.ProfilerActivity.CPU] if device.type == "cuda": activities.append(torch.profiler.ProfilerActivity.CUDA) def on_trace_ready(prof: torch.profiler.profile) -> None: prof.export_chrome_trace(str(trace_path)) logging.info(f"Saved Chrome trace to: {trace_path.resolve()}") if stacks_path: metric = "self_cuda_time_total" if device.type == "cuda" else "self_cpu_time_total" prof.export_stacks(str(stacks_path), metric=metric) logging.info( f"Saved Flamegraph stacks to: {stacks_path.resolve()}") return torch.profiler.profile( activities=activities, schedule=torch.profiler.schedule(wait=0, warmup=warmup_steps, active=active_steps, repeat=1), record_shapes=True, profile_memory=True, with_stack=stacks_path is not None, on_trace_ready=on_trace_ready, )def run_benchmark(args: argparse.Namespace) -> None: device = torch.device(args.device) logging.info("=" * 60) logging.info(f"Device : {device.type.upper()}") logging.info(f"Compile : {args.compile}") logging.info(f"Profile : {args.profile}") logging.info(f"Trigger Fail : {args.trigger_fail}") if args.profile: logging.info( f"Profiler Config: {args.warmup} Warmup Steps | {args.iters} Active Steps" ) logging.info("=" * 60) # Initialize model model = NormalizedLinear(args.in_features, args.out_features, trigger_fail=args.trigger_fail).to(device) x = torch.randn(args.batch_size, args.in_features, device=device) if args.compile: logging.info("Compiling model via torch.compile()...") model = torch.compile(model) try: if args.profile: total_steps = args.warmup + args.iters logging.info( f"Running Profiler ({args.warmup} warmup + {args.iters} active = {total_steps} total steps)..." ) with create_profiler(device=device, trace_path=args.trace_path, stacks_path=args.stacks_path, warmup_steps=args.warmup, active_steps=args.iters) as prof: for step in range(total_steps): _ = model(x) prof.step( ) # Advances schedule: warmup -> active -> trace ready if device.type == "cuda": torch.cuda.synchronize() else: # Standalone Warmup logging.info(f"Warming up ({args.warmup} iterations)...") for _ in range(args.warmup): _ = model(x) if device.type == "cuda": torch.cuda.synchronize() # Standalone Benchmark Loop logging.info(f"Executing {args.iters} benchmark iterations...") for _ in range(args.iters): _ = model(x) if device.type == "cuda": torch.cuda.synchronize() logging.info("Run finished successfully!") except RuntimeError as err: logging.error( "Caught Exception from GPU Stream Sync (Expected if --trigger-fail was set):" ) logging.error(f"--> {err}")def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description= "PyTorch Async Assertion & Compiler Benchmark with Profiler Schedule", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) # Feature Flags parser.add_argument("--compile", action="store_true", help="Compile model graph with torch.compile") parser.add_argument("--profile", action="store_true", help="Enable torch.profiler tracing") parser.add_argument( "--trigger-fail", action="store_true", help="Force condition failure to test deferred stream errors") # File Paths parser.add_argument("--trace-path", type=Path, default=Path("my_trace.json"), help="Output JSON path for Chrome/Perfetto trace") parser.add_argument( "--stacks-path", type=Path, default=None, help="Optional output text path for Flamegraph stack trace") # Model Hyperparameters & Profiler Steps parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Target device") parser.add_argument("--batch-size", type=int, default=128, help="Batch size") parser.add_argument("--in-features", type=int, default=1024, help="Input feature size") parser.add_argument("--out-features", type=int, default=2048, help="Output feature size") parser.add_argument("--warmup", type=int, default=3, help="Warmup iterations before active profiling") parser.add_argument( "--iters", type=int, default=5, help="Number of active benchmark iterations to capture in profile") return parser.parse_args()if __name__ == "__main__": run_benchmark(parse_args())
|
To profile the PyTorch program with and without using torch.compile
, we could run the following commands, which will generate Perfetto profiling traces for both cases.
12
|
$ python assert_async.py --profile --trace-path assert_async_trace.json$ python assert_async.py --compile --profile --trace-path assert_async_compiled_trace.json
|
In the profiling trace of the run that does not use torch.compile
, we could see that the assertion is performed asynchronously on the GPU stream.
In the profiling trace of the run that uses torch.compile
, we could see that the assertion is fused into the compiled graph with other operations, which is friendly to the optimization of the computation graph.
If an assertion fails, the error will be reported asynchronously when the GPU stream is synchronized with the CPU thread. For example, if we run the following command to trigger an assertion failure.
123456789101112131415
|
$ python assert_async.py --trigger-fail[INFO] ============================================================[INFO] Device : CUDA[INFO] Compile : False[INFO] Profile : False[INFO] Trigger Fail : True[INFO] ============================================================[INFO] Warming up (3 iterations).../opt/pytorch/pytorch/aten/src/ATen/native/cuda/TensorCompare.cu:109: _assert_async_cuda_kernel: block: [0,0,0], thread: [0,0,0] Assertion `Assertion Failed: Probabilities must sum to 1.0!` failed.[ERROR] Caught Exception from GPU Stream Sync (Expected if --trigger-fail was set):[ERROR] --> CUDA error: device-side assert triggeredSearch for `cudaErrorAssert' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.For debugging consider passing CUDA_LAUNCH_BLOCKING=1Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
|
On caveat is that an assertion failure will cause CUDA context to be poisoned, which will cause all subsequent CUDA calls to fail. Without restarting the CUDA context, the PyTorch program will not be able to continue on GPU. That is why it is an assertion rather than an exception which can be caught and handled.
PyTorch Asynchronous Assert Implementation
Because the purpose of assertion failure is to terminate the program, the implementation of torch._assert_async
is designed to poison the CUDA context. In the CUDA operation, the __trap
instruction, which translates to asm volatile("trap;")
I believe, is used to terminate the program when an assertion fails. In the Triton compilation, the tl.device_assert instruction is used to terminate the program when an assertion fails.
Conclusions
The torch._assert_async
calls are not completely free and it can poison the CUDA context when an assertion fails. Therefore, ideally it should be used in development, and should not be used in production because assertion should be expected to always pass. In C++, similarly, the assert
will be optimized away by the compiler in release or production builds when the macro NDEBUG
is defined.
PyTorch Asynchronous Assert