Skip to content

Updates - #533

Merged
felix-e-h-p merged 3 commits into
mainfrom
enable-tf32-tensor-cores
Jul 20, 2026
Merged

Updates#533
felix-e-h-p merged 3 commits into
mainfrom
enable-tf32-tensor-cores

Conversation

@felix-e-h-p

@felix-e-h-p felix-e-h-p commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Enable TF32 tensor cores and make validation bf16-safe

  • run.py: add torch.set_float32_matmul_precision("high") — lets fp32 matmuls / convs use tensor cores on Ampere+ GPUs. Free speedup, negligible precision cost.
  • lightning_module.py: y_hat.float().cpu().numpy() in val storage - fixes a crash when training with precision=bf16-mixed (numpy has no bfloat16). No-op under fp32.

Tested on A6000 (DE model): GPU fully utilised, no change to training behaviour.

@felix-e-h-p
felix-e-h-p requested a review from dfulu July 18, 2026 10:43
@dfulu

dfulu commented Jul 20, 2026

Copy link
Copy Markdown
Member

Looks good. What's the reasoning around the precision used? In pvnet-summation it was locked to "medium"

@felix-e-h-p

felix-e-h-p commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Looks good. What's the reasoning around the precision used? In pvnet-summation it was locked to "medium"

Ah yeah, bit more precision than medium but gets same tensor core speedup on A6000s, maybe the safer default. Happy to change to medium to match the summation if you prefer.

@dfulu

dfulu commented Jul 20, 2026

Copy link
Copy Markdown
Member

So would that mean it would be better for the summation model to use "high" precision as well?

@felix-e-h-p

felix-e-h-p commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

So would that mean it would be better for the summation model to use "high" precision as well?

Not sure, but I suppose yeah - speed won't be a bottleneck anyway right. May be worth updating also.

@dfulu

dfulu commented Jul 20, 2026

Copy link
Copy Markdown
Member

Yeh it sounds like it would. It feels like we shouldn't have different standards for the two libraries

dfulu
dfulu reviewed Jul 20, 2026
View reviewed changes
Comment thread run.py Resolved
dfulu
dfulu reviewed Jul 20, 2026
View reviewed changes
Comment thread pvnet/training/lightning_module.py Resolved
@felix-e-h-p

felix-e-h-p commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Couple comments added - also made the relevant update for summation: openclimatefix/pvnet-summation#48

dfulu
dfulu approved these changes Jul 20, 2026
View reviewed changes

@dfulu dfulu left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, thanks

@felix-e-h-p
felix-e-h-p merged commit 8bb428c into main Jul 20, 2026
6 checks passed
@felix-e-h-p
felix-e-h-p deleted the enable-tf32-tensor-cores branch July 20, 2026 12:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@dfulu dfulu dfulu approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@felix-e-h-p @dfulu