Factor out dispatch and layout registration table

[jerryzh168]

Summary:
att, after the refactor we can use common utils for new dtypes as well

# added for adding ops to op dispatch table, works for both aten and torch function
def _implements(cls, aten_ops_or_torch_fns):
   ...

# added for registering new layout class
def _register_layout_cls(cls, extended_layout: str):
 ...

# get layout tensor constructor
def _get_layout_tensor_constructor(cls, extended_layout: str) -> Callable:
 ...

Test Plan:
python test/dtypes/test_nf4.py
python test/dtypes/test_aqt.py
python test/integration/test_integration.py

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/360

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8929358 with merge base 8841094 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[drisspg]

Nit: is this typing right?

[jerryzh168]

yeah, but I can change Any to Callable as well to be more specific I think

[drisspg]

Part of me still doesn't like having the function table be shared among subclasses.... I don't have a great reason against though lol so that's not blocking

[msaroufim]

Actually curious to hear more. I want us to get the aesthetics right

[drisspg]

I actually dont have any logically good reason lol it is more just gut check I dont love global things that tie a bunch of potentially disperate ideas/implementations together

The dispatcher in PyTorch is a giant global vtable that ties everything together so there are obvious counter points thats why this super nit and not blocking

[jerryzh168]

OK yeah, I think we can follow what pytorch is already doing

[weifengpy]

This breaks QLora in TorchTune. I plannning to add a FSDP unit test in TorchAO. Probably need to revert NF4 changes. #380

[rank0]:   File "/data/users/weif/pytorch-official/pytorch/torch/distributed/_tensor/dispatch.py", line 205, in dispatch
[rank0]:     local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
[rank0]:   File "/data/users/weif/pytorch-official/pytorch/torch/_ops.py", line 666, in __call__
[rank0]:     return self_._op(*args, **kwargs)
[rank0]:   File "/data/users/weif/ao/torchao/dtypes/nf4tensor.py", line 541, in __torch_function__
[rank0]:     return _ATEN_OP_OR_TORCH_FN_TABLE[cls][func](*args, **kwargs)
[rank0]:   File "/data/users/weif/ao/torchao/dtypes/utils.py", line 25, in wrapper
[rank0]:     return func(*args, **kwargs)
[rank0]: TypeError: nf4_detach() missing 1 required positional argument: 'args'