[fx/meta/rpc] move _meta_registration.py to fx folder / register fx functions with compatibility checks / remove color debug #1710
Conversation
colossalai/fx/__init__.py
Outdated
| try: | ||
| from . import _meta_registrations | ||
| META_COMPATIBILITY = True | ||
| except: | ||
| import torch | ||
| META_COMPATIBILITY = False | ||
| print(f'_meta_registrations seems to be incompatible with PyTorch {torch.__version__}.') |
There was a problem hiding this comment.
This would still print a statement to the user upon importing colossalai if colossalai.fx is used by any other module in the library.
There was a problem hiding this comment.
Let's not print something like this. Instead, we can offer an API like xxx_is_compatible to check the compatibility.
There was a problem hiding this comment.
Do not expose a constant to the user, no matter he is a user of the colossalai or the developer who wants to use colossalai.fx.
There was a problem hiding this comment.
I see. For the compatibility concerns, probably I have to come up with a global registration for this. Currently in this PR I will just define an API for checking though.
There was a problem hiding this comment.
But is there any existing modules in ColossalAI for compatibility checks?
There was a problem hiding this comment.
this should be only within this fx module as compatibility is specific to this module.
|
The CI failed, please check it out. |
|
Remember to remove the test label when your PR is not completely ready for review as it will incur testing runs for every push. I will help you remove it first. |
super-dainiu
changed the title
super-dainiu
changed the title
super-dainiu
changed the title
| def compatibility(is_backward_compatible: bool = False) -> Callable: | ||
| """A decorator to make a function compatible with different versions of PyTorch. | ||
|
|
||
| Args: | ||
| is_backward_compatible (bool, optional): Whether the function is backward compatible. Defaults to False. | ||
|
|
||
| Returns: | ||
| Callable: The decorated function | ||
| """ | ||
|
|
||
| def decorator(func): | ||
| if META_COMPATIBILITY: | ||
| return func | ||
| else: | ||
| if is_backward_compatible: | ||
| return func | ||
| else: | ||
|
|
||
| def wrapper(*args, **kwargs): | ||
| raise RuntimeError(f'Function `{func.__name__}` is not compatible with PyTorch {torch.__version__}') | ||
|
|
||
| return wrapper | ||
|
|
||
| return decorator |
There was a problem hiding this comment.
This is a decorator that wraps our defined function. If the PyTorch version is not satisfied, we can wrap the function to be a RuntimeError.
There was a problem hiding this comment.
What is the difference between this and the PyTorch one?
colossalai/fx/_compatibility.py
Outdated
| def check_meta_compatibility() -> bool: | ||
| """Check the meta compatibility. Normally it should be called before importing some of the `colossalai.fx` | ||
| modules. If the meta compatibility is not satisfied, the `colossalai.fx` modules will be replaced by its | ||
| experimental counterparts. | ||
|
|
||
| Returns: | ||
| bool: The meta compatibility | ||
| """ | ||
| return META_COMPATIBILITY |
There was a problem hiding this comment.
An API for checking
| @@ -0,0 +1,32 @@ | |||
| import torch | |||
There was a problem hiding this comment.
The file colossalai/fx/profiler/constant.py is deleted and then re-created? You should change the file directly instead of deleting it to minimize git change.
There was a problem hiding this comment.
I split this file into two, so git cannot identify this as a filename change.
| from .memory import calculate_fwd_in, calculate_fwd_out, calculate_fwd_tmp | ||
| from .profiler import profile_function, profile_method, profile_module | ||
| from .profiler_function import * | ||
| from .profiler_module import * | ||
| from .profiler import profile_function, profile_method, profile_module | ||
| from .registry import meta_profiler_function, meta_profiler_module |
There was a problem hiding this comment.
Add some documentation in this file to explain why this folder is experimental.
There was a problem hiding this comment.
I might refactor this experimental/ folder later, it hasn't been updated for a long time.
Why move?
Currently
_meta_registration.pyis only used forMetaTensorinfx.profiler. Users with PyTorch 11 or lower cannot use theMetaTensorand will get an error message. I believe we should show the error message only when they importcolossalai.fx. So move it fromcolossalaitocolossalai.fxis a good choice.Other changes
I removed all the
color_debugbecause it causes conflicts on our CI. (should I?) cc @LSTM-KirigayaTests