facebookresearch / facebookresearch/fairscale
Fairscale checkpoint wrapper does not support deepcopy
- Dominant language
- Python
- Stars
- 3.4k
- Forks
- 293
- PR merge metrics
- No merged PRs in 30d
Description
Came across a very tricky issue with fairscale checkpoint wrapper.
Checkpoint wrapper wraps the forward function in https://github.com/facebookresearch/fairscale/blob/8acbec718f3c70a6b9785470bb9e05cd84fc3f8e/fairscale/nn/checkpoint/checkpoint_activations.py#L154-L156
copy.deepcopy(module).forward will be calling module.forward due to the weakref being passed to through the partial function. This is not what the user would intend as the deepcopy of a module is expected work as a completely separate copy. For e.g. even when the deepcopy module parameters are updated, the forward function would still call the forward of the original module using the original module's parameters.
The workaround I tried is to define __deepcopy__ method that deepcopy the module correctly by unwrapping the checkpoint wrapper and wrapping it again like:
```
def _add_deepcopy_method(module):
def deepcopy_module(m, memodict):
"""deepcopy the module correctly by unwrapping the checkpoint
wrapper and wrapping it again. Ensures that correct forward
function is called in the deepcopy.
"""
deepcopy_method = m.__deepcopy__
forward_method = m.forward
m.__deepcopy__ = m.old_deepcopy_method
m.forward = m.precheckpoint_forward
mc = copy.deepcopy(m, memodict)
delattr(mc, "precheckpoint_forward")
m.forward = forward_method
m.__deepcopy__ = deepcopy_method
return checkpoint_wrapper(mc, offload_to_cpu=offload_to_cpu)
module.old_deepcopy_method = getattr(module, "__deepcopy__", None)
module.__deepcopy__ = deepcopy_module.__get__(module)
```
Contributor guide
Research direction
Start in fairscale/nn/checkpoint/checkpoint_activations.py at the wrapper around lines 154-156, then reproduce the issue with copy.deepcopy(module). Check how the weakref and partial affect the copied module's forward call. Done means the deep-copied module invokes forward with its own parameters rather than the original module's parameters.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100