Misleading "Uses tensor cores" print and dead TF32 setting in ch05/10_llm-training-speed

[hellozdp]

Bug description

Two small issues in ch05/10_llm-training-speed/01_opt_single_gpu.py (the second one also applies to 02_opt_multi_gpu_ddp.py). Line numbers below are against the current main (01_opt_single_gpu.py blob sha b99c969ad3).

1. Uses tensor cores: reports CUDA availability, not tensor-core support

01_opt_single_gpu.py:380-389

if torch.cuda.is_available():
    capability = torch.cuda.get_device_capability()
    if capability[0] >= 7:
        torch.set_float32_matmul_precision("high")
        print("Uses tensor cores")
    else:
        print("Tensor cores not supported on this GPU. Using default precision.")
print(f"Uses tensor cores: {torch.cuda.is_available()}")   # line 389

On a pre-Volta card (e.g. GTX 1080, compute capability 6.1) the output contradicts itself:

Tensor cores not supported on this GPU. Using default precision.
Uses tensor cores: True

Suggested fix:

has_tensor_cores = (
    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 7
)
print(f"Uses tensor cores: {has_tensor_cores}")

02_opt_multi_gpu_ddp.py does not have this print, so this part is specific to 01.

2. set_float32_matmul_precision("high") has no effect once the model is cast to bf16

Both scripts enable TF32 and then convert the whole model to bfloat16:

file	TF32	bf16 cast
01_opt_single_gpu.py	line 385	line 415 model.to(device).to(torch.bfloat16)
02_opt_multi_gpu_ddp.py	line 455	line 489 model = model.to(torch.bfloat16)

set_float32_matmul_precision only affects fp32 matmuls, and after the cast there are none left. I checked this by hooking every nn.Linear in GPTModel after .to(torch.bfloat16):

parameters: 37x torch.bfloat16
buffers:     2x torch.bfloat16
forward pass, 13 Linear layers:
  input bfloat16, weight bfloat16, output bfloat16
final logits: torch.bfloat16

So as the scripts are currently written, the TF32 branch never does anything.

It is still useful if a reader comments out the bf16 cast to compare "fp32 + TF32" against "pure bf16", but as written it reads as if the two optimizations stack, which they don't. A short comment would make the intent clear, e.g.:

# Only relevant if the bfloat16 cast below is disabled; with a bf16 model
# there are no fp32 matmuls left for this setting to affect.
torch.set_float32_matmul_precision("high")

Happy to open a PR for either or both if that would be useful.

Thanks for the book and the repo — they have been a great resource.

What operating system are you using?

macOS

Where do you run your code?

Local (laptop, desktop)

Environment

{
  "machine": "MacBook Pro (Mac15,6, Apple M3 Pro)",
  "OS": "macOS-15.7.3-arm64-arm-64bit",
  "architecture": "arm64",
  "Python": "3.11.15",
  "PyTorch": "2.13.0",
  "installed_with": "pip",
  "installation_source": null,
  "CUDA_build": null,
  "CUDA_available": false,
  "MPS_available": true
}

Note: this is a code-reading report, not a runtime failure — both issues are visible from the source and neither requires a CUDA device to observe.