Skip to content

Refactor throughput dtype inference into Precision plugin API#21537

Open
KAVYANSHTYAGI wants to merge 10 commits intoLightning-AI:masterfrom
KAVYANSHTYAGI:prc-dtype-fix
Open

Refactor throughput dtype inference into Precision plugin API#21537
KAVYANSHTYAGI wants to merge 10 commits intoLightning-AI:masterfrom
KAVYANSHTYAGI:prc-dtype-fix

Conversation

@KAVYANSHTYAGI
Copy link
Contributor

@KAVYANSHTYAGI KAVYANSHTYAGI commented Feb 15, 2026

Summary

This PR moves throughput dtype inference into the Precision plugin itself by introducing a new method:

compute_dtype()

Instead of using isinstance(...) checks inside throughput utilities, we now simply call:

plugin.compute_dtype()

What Changed

  • Added compute_dtype() to the base Precision class (default: torch.float32)
  • Implemented compute_dtype() across all built-in Fabric and PyTorch precision plugins
  • Removed hard-coded type checks and internal attribute access in throughput utilities
  • Simplified _plugin_to_compute_dtype() to delegate directly to the plugin API
  • Added tests to verify correct behavior and integration with ThroughputMonitor

Why

Previously, throughput relied on chained isinstance(...) checks and accessed internal plugin attributes.
This tightly coupled throughput logic to specific precision implementations.

By moving dtype responsibility into the plugin:

  • We improve encapsulation
  • Reduce maintenance overhead
  • Automatically support custom precision plugins
  • Keep throughput utilities clean and extensible

fix #21535

@github-actions github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package labels Feb 15, 2026
@codecov
Copy link

codecov bot commented Feb 15, 2026

Codecov Report

❌ Patch coverage is 84.74576% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 87%. Comparing base (7d2de87) to head (93dcbd5).
⚠️ Report is 1 commits behind head on master.
✅ All tests successful. No failed tests found.

Additional details and impacted files
@@           Coverage Diff           @@
##           master   #21537   +/-   ##
=======================================
  Coverage      87%      87%           
=======================================
  Files         270      264    -6     
  Lines       24071    23980   -91     
=======================================
- Hits        20867    20861    -6     
+ Misses       3204     3119   -85     

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Refactor throughput dtype inference into the Precision plugin API

1 participant