facebookresearch / facebookresearch/fairscale

Fairscale checkpoint wrapper does not support deepcopy

Open
#816 2 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.