Fea nn refactor #1265
Fea nn refactor #1265
Conversation
loliverhennigh
requested review from
Alexey-Kamenev,
CharlelieLrt,
NickGeneva,
RishikeshRanade,
coreyjadams,
dallasfoster,
ktangsali,
mnabian,
peterdsharpe and
pzharrington
as code owners
|
Skipped: This PR changes more files than the configured file change limit: ( |
coreyjadams
left a comment
coreyjadams
left a comment
There was a problem hiding this comment.
A lot of the pieces you've put in module might be better suited in functional, I think. For example, the neighbors utilities are pretty much pure functions. Similar for many activation functions.
Is ASV supported with this too or just reorganization? Is that benchmark mixing meant to use ASV or something else?
| class Function(TorchAutogradFunction, BenchmarkMixin): | ||
| """Base class for PhysicsNeMo custom autograd functions.""" | ||
| # Placeholder for utilities to bring in warp, fuser, etc. No newline at end of file |
There was a problem hiding this comment.
This locks us in to the inherit-from-TorchAutogradFunction interface and away from the custom op solution that pytorch seems to prefer here:
https://docs.pytorch.org/tutorials/advanced/python_custom_ops.html
I favor the custom op decorators for several reasons:
- better integration with torch.compile and avoiding compile breaks
- Registration with pytorch dispatch tools, used to enable ShardTensor's distributed dispatch
- built-in testing utilities from pytorch to validate compile support, etc.
|
|
||
|
|
||
| class Module(torch.nn.Module): | ||
| class Module(torch.nn.Module, BenchmarkMixin): |
There was a problem hiding this comment.
This will automatically expose our models to this BenchmarkMixin functionality. Not necessarily a bad thing, though I wonder if we'll hit unforeseen consequences in save / load / restore / broken checkpoints, etc.
2 participants
PhysicsNeMo Pull Request
Description
Refactor physicsnemo
nnto more closely align with pytorch.