get_device() sets new TF32 API on torch >= 2.9, breaking torch.compile in ch02 section 2.9 (InductorError)

[DDRXV]

Summary

get_device() in reasoning_from_scratch/ch02.py sets the new TF32 API on PyTorch >= 2.9. That combination makes torch.compile fail with an InductorError, so section 2.9 ("Faster inference via PyTorch model compilation") cannot run on a CUDA GPU with PyTorch 2.9 or newer.

Eager mode is unaffected, so everything works until the reader reaches the compilation section.

Environment

• PyTorch 2.10.0+cu128
• reasoning_from_scratch 0.1.21
• NVIDIA RTX 4090, driver 550.163.01
• Python 3.12, Linux

Offending code

reasoning_from_scratch/ch02.py, lines 17-19:

if (major, minor) >= (2, 9):
    torch.backends.cuda.matmul.fp32_precision = "tf32"
    torch.backends.cudnn.conv.fp32_precision = "tf32"
else:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

The same lines appear in the get_device() cell in ch02/01_main-chapter-code/ch02_main.ipynb (section 2.5).

Error

InductorError: LoweringException: RuntimeError: PyTorch is checking whether
allow_tf32_new is enabled for cuBlas matmul, Current status indicate that you
have used mix of the legacy and new APIs to set the TF32 status for cublas
matmul. We suggest only using the new API to set the TF32 flag.
  target: aten.bmm.default

Root cause

This is an upstream PyTorch bug, not a mistake in the book's code. Inductor's own compile_fx.py reads the legacy getter torch.backends.cuda.matmul.allow_tf32. Once user code sets the new-style flag, PyTorch sees both API families used in the same process and escalates from a deprecation warning to a hard error.

• torch._dynamo.exc.BackendCompilerFailed about mixed TF32 API usage using torch.compile when only new API is used pytorch/pytorch#166387 — same error with new API only
• unnecessary warning triggered by inductor in pytorch 2.9 pytorch/pytorch#166286 — identifies the legacy read inside compile_fx.py

vLLM hit the same issue (vllm-project/vllm#31579, vllm-project/vllm#29349).

Minimal reproduction

import torch
from pathlib import Path
from reasoning_from_scratch.qwen3 import Qwen3Model, QWEN_CONFIG_06_B, Qwen3Tokenizer

# These two lines are what break compilation:
torch.backends.cuda.matmul.fp32_precision = "tf32"
torch.backends.cudnn.conv.fp32_precision = "tf32"

dev = torch.device("cuda")
tok = Qwen3Tokenizer(tokenizer_file_path=Path("qwen3") / "tokenizer-base.json")
m = Qwen3Model(QWEN_CONFIG_06_B)
m.load_state_dict(torch.load(Path("qwen3") / "qwen3-0.6B-base.pth"))
m.to(dev).eval()

torch._dynamo.config.allow_unspec_int_on_nn_module = True
mc = torch.compile(m)
ids = torch.tensor(tok.encode("Explain large language models in a single sentence."), device=dev).unsqueeze(0)

with torch.inference_mode():
    mc(ids)   # InductorError

Results across TF32 settings

Same script, only the TF32 lines varied:

TF32 setting	torch.compile
none	works
torch.backends.cuda.matmul.allow_tf32 = True (legacy)	works, deprecation warning
torch.backends.fp32_precision = "tf32" (global, new)	works
torch.backends.cuda.matmul.fp32_precision = "tf32" (current code)	InductorError

Suggested fix

Use the legacy API unconditionally, or the global new-style setter, until the upstream issue is resolved:

if enable_tensor_cores:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

Worth noting that for chapter 2 specifically the flag has no effect either way: QWEN_CONFIG_06_B sets "dtype": torch.bfloat16, so every Linear is bf16 and TF32 (which only governs fp32 matmuls) never applies. It becomes relevant in the later training chapters.

Note

One additional detail that may be worth a sentence in the text: the TF32 flags are process-global and sticky, so a reader who edits the cell still hits the error until they restart the kernel.

Happy to open a PR if useful.

[rasbt]

Awesome, thanks for reporting. Some readers had exactly this issue, and I was wondering where it was coming from. Fixing it in #257