Sanket Jayant Purandare 2e5366fbc0 Extended Module Tracker (#128508) ... This is an extension of [ModuleTracker](https://github.com/pytorch/pytorch/blob/main/torch/utils/module_tracker.py) with added features and bug fixes. 1. Allows installing user-defined hooks to be called in pre-fw, post-fw, pre-bw and post-bw hooks of the ``ModTracker``. 2. Adds a function ``get_known_fqn`` that retrieves the fqn of the module as tracked by the ``ModTracker``. 3. Only registers the multi-grad hooks if we are in the forward pass. This is important because, a module's pre-fw and post-fw hooks get called in the backward during AC and we do not want to register multi-grad hooks in this case. 4. Sets the kwarg ``always_call=True`` for post-fw hooks, so that they are called post AC. Pull Request resolved: https://github.com/pytorch/pytorch/pull/128508 Approved by: https://github.com/wanchaol		2024-06-14 19:48:46 +00:00
..
test_memory_tracker.py	Enable UFMT on all of test/distributed (#123539)	2024-04-17 06:46:02 +00:00
test_mod_tracker.py	Extended Module Tracker (#128508)	2024-06-14 19:48:46 +00:00