Migrate to strict mode export (dynamo) to support AC tags and HOPs

[xmfan]

Depends on: pytorch/pytorch#161479

[ezyang]

Checking with @avikchaudhuri @tugsbayasgalan to make sure this doesn't make the eventual pre compile rework harder

[ezyang]

I'm OK for the use of export to torch ir. I'll leave it up to fmassa when we should merge this (e.g., does deepcopy need to be fixed first)

[tugsbayasgalan]

Checking with @avikchaudhuri @tugsbayasgalan to make sure this doesn't make the eventual pre compile rework harder

Since we just plan to swap out export_to_torch_ir with the eventual correct API that also works with precompile , i think this change is fine.

[fmassa]

@xmfan I was looking into implementing the deepcopy for our parametrization so that we can get this merged, but I realized that there might be other issues in here.

The captured graph from _export_to_torch_ir gives us a graph like the following for the example_autoparallel.py code, when adding back mixed precision

graph():
    %l_args_0_ : [num_users=1] = placeholder[target=arg0]
    %to : [num_users=2] = call_method[target=to](args = (%l_args_0_, torch.bfloat16), kwargs = {})
    %l__self___wq_weight : [num_users=1] = get_attr[target=L__self___wq_weight]
    %l__self___wk_weight : [num_users=1] = get_attr[target=L__self___wk_weight]
    %l__self___wv_weight : [num_users=1] = get_attr[target=L__self___wv_weight]
    %l__self___wo_weight : [num_users=1] = get_attr[target=L__self___wo_weight]
    %wrap_body_0 : [num_users=1] = get_attr[target=wrap_body_0]
    %tag_activation_checkpoint : [num_users=1] = call_function[target=torch.ops.higher_order.tag_activation_checkpoint](args = (%wrap_body_0, %to, %l__self___wq_weight, %l__self___wk_weight, %l__self___wv_weight, %l__self___wo_weight), kwargs = {use_reentrant: False})
    %o : [num_users=1] = call_function[target=operator.getitem](args = (%tag_activation_checkpoint, 0), kwargs = {})
    %o0 : [num_users=2] = call_function[target=operator.add](args = (%o, %to), kwargs = {})
    %l__self___w1 : [num_users=1] = get_attr[target=L__self___w1]
    %o_1 : [num_users=1] = call_method[target=forward](args = (%l__self___w1, %o0), kwargs = {})
    %o_2 : [num_users=1] = call_function[target=torch.nn.functional.relu](args = (%o_1,), kwargs = {})
    %l__self___w2 : [num_users=1] = get_attr[target=L__self___w2]
    %o_3 : [num_users=1] = call_method[target=forward](args = (%l__self___w2, %o_2), kwargs = {})
    %o_4 : [num_users=1] = call_function[target=operator.add](args = (%o0, %o_3), kwargs = {})
    %output : [num_users=1] = call_method[target=to](args = (%o_4, torch.bfloat16), kwargs = {})
    return [output]

You can see that the input and output cast get properly called, but the weight parametrization disappear.
My guess is that this is because the weight names have been renamed (like L__self__wq_weight), so we never call into the parametrization.

Am I reading this right?

[xmfan]

@fmassa I see the casts properly if I enable torch._dynamo.config.install_free_tensors=True: https://gist.github.com/xmfan/5176d488358e77943dbca0ecd6fe0005. This was caused by a divergence between torch.export dynamo and torch.compile dynamo: https://github.com/pytorch/pytorch/blob/eba20d2d748cb17dce9aa26e5513e4567bfd8282/torch/_dynamo/variables/builder.py#L1881-L1901

[fmassa]

I validated that with this change we get back the dtype cast hooks.

But looks like test failures are related, and this seems to be breaking init_weights for now.

Also, am I understanding this right that this PR will also make AutoParallel only support tuple inputs?

[xmfan]

For the init_weights, dynamo encodes the fqn a certain way, self.linear.weight -> self____modules__linear____parameters__weight, iirc we want the fqns to match eager exactly for model saving/loading purposes? @fmassa

[fmassa]

For the init_weights, dynamo encodes the fqn a certain way, self.linear.weight -> self____modules__linear____parameters__weight, iirc we want the fqns to match eager exactly for model saving/loading purposes? @fmassa

@xmfan yes, we would want to keep the same FQN for saving / loading.

[xmfan]

Scrap scrap that. It's too brittle to manually restore FQNs. Keep using the export stuff.

Scrap this. Apparently it's okay, people are working on AOT precompile and we'll just move to that.

I synced with @tugsbayasgalan, there's some refactors upcoming to export so it'd be best if we stuck to public APIs. Given that strict=False works for trunk, there's probably not that much gap.

So I'm breaking this PR up:
1- I'll ensure that torch.export(strict=True) works for what we currently have landed in autoparallel: #104
2- Tugsuu is looking at getting the AC HOP to proxy itself during pre-dispatch, the tests in this PR should pass after

[fmassa]

LGTM once tests pass, thanks for working on it!

Also, could you add some more comments on the need of the monkey_patch?

[fmassa]

Does the code also works with the reduce_dtype=torch.float32? If yes, can you set it the default? That is the setup we mostly use so it would be good to have it as the default

[xmfan]

yes works, i'll update these comments

[fmassa]

Can you explain a bit (in code as well) why we currently need this?

Is this something that you expect we will remove in the future or is it something that is meant to stay?

[xmfan]

added comment, it's something we will remove in the future, either when:

• export offers a mode to drop the serializability constraints
• precompile frontend