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

[EugenHotaj]

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.

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

[anj-s]

Can you print out the original model and sharded model?

[EugenHotaj]

Here is the full model:

BranchedNetwork(
  (net): ModuleList(
    (0): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (1): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (2): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (3): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (4): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (5): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (6): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (7): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (8): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
    (9): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
)

And the sharded model:

[GraphModule(
  (net): Module(
    (0): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (1): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (2): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (3): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (4): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (5): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (6): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (7): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (8): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
), GraphModule(
  (net): Module(
    (9): Branch(
      (left): Linear(in_features=10, out_features=10, bias=True)
      (right): Linear(in_features=10, out_features=10, bias=True)
    )
  )
)]

Looks reasonable to me but I'm not sure why the shard_count is not being respected.

[EugenHotaj]

I'm not sure why the shard_count is not being respected.

BTW, this also happens without the changes in this PR, i.e. if we just use torch.fx.symbolic_trace directly, so I think something might be going wrong in the _split_nodes logic.

(I could try looking into it in a follow up PR).