Skip to content

Introduce Hyrax's own auto_model function to remove DataParallel#996

Open
SamSandwich07 wants to merge 11 commits into
mainfrom
custom-auto-model
Open

Introduce Hyrax's own auto_model function to remove DataParallel#996
SamSandwich07 wants to merge 11 commits into
mainfrom
custom-auto-model

Conversation

@SamSandwich07

Copy link
Copy Markdown
Collaborator

Introduce a helper function _auto_model to pytorch_ignite.py to replace idist.auto_model and remove DataParallel. In other words, models will never be wrapped in DataParallel anymore; they will either be unwrapped or wrapped in DistributedDataParallel.

@SamSandwich07 SamSandwich07 self-assigned this Jul 21, 2026
torch.nn.Module
"""
wrapped_model = idist.auto_model(model, sync_bn=sync_bn, **kwargs)
if type(wrapped_model) is DataParallel:

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.

We could simply make this a RuntimeError. Thoughts?

@SamSandwich07 SamSandwich07 Jul 21, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

No, that wouldn't work. idist.auto_model is going to make it DataParallel whenever we run Hyrax in an environment with multiple GPUs (like Gondor) but doesn't use idist.Parallel (i.e. set distributed to false or run a verb other than train). If we raise an error here then Hyrax will essentially always raise an error on Gondor and the like.

Comment thread src/hyrax/pytorch_ignite.py Outdated

Args:
model: model to adapt.
sync_bn: if True, applies `torch convert_sync_batchnorm`_ to the model for native torch

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.

Do we know we will need to change sync_bn from callers of this function?

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.

I guess I'm unclear why we're promoting it out of **kwargs in this interface.

If we need it set to false consistently, then we can put that, refactor when usage forces it?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

We are the only callers of idist.auto_model, and in our code we never call it with any arguments other than the torch model, so the answer for now is "we never change the argument".

@SamSandwich07 SamSandwich07 Jul 21, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

As for why we promote it out of **kwargs, it's because idist.auto_model keeps it outside of **kwargs too, so I just used the same input signature (below taken from ignite docs):

def auto_model(model: nn.Module, sync_bn: bool = False, **kwargs: Any) -> nn.Module:

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants