Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Extend auto shard capabilities to work around torch.fx edge cases. #817

Merged
merged 1 commit into from
Oct 22, 2021
Merged

Extend auto shard capabilities to work around torch.fx edge cases. #817

merged 1 commit into from
Oct 22, 2021

Commits on Oct 21, 2021

  1. Extend auto shard capabilities to work around torch.fx edge cases. 

    auto_shard.py currently uses torch.fx to create a symbolic DAG of
operations and linearizes that DAG into an nn.Sequential so it can later
be used for model offloading. This works in most cases but runs into
issues for certain eager mode features, such as dynamic conditionals,
shape-dependent computation, etc.

This PR extends auto_shard.py to first run a preprocessing step which wraps
any nn.Module which cannot be traced through. It adds a test for dynamic
conditionals and updates existing failing test code.

There are some immediate extensions to this approach which are marked as
TODO in the code.
    Eugen Hotaj committed Oct 21, 2021
    Configuration menu
    Copy the full SHA
    e55fdf6 View commit details
    Browse the repository at this point in the history