Summary: Fixes #58512

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 2417a47
Pull Request resolved: #60975

Summary: Fixes #58512

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 384fccb
Pull Request resolved: #60975

Summary: Fixes #58512

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 2e6248d
Pull Request resolved: #60975

Summary: Fixes #58512

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: f3c33b6
Pull Request resolved: #60975

Fixes #58659. This PR builds directly on top of [#58512](#58512 (comment)).

~~Creates a context manager `with torch.autograd.graph.saved_tensors_default_hooks(pack, unpack)` that can be used on the Python side.~~

Expose a pair of functions to Python users: `torch.autograd.graph.set_saved_tensors_default_hooks(pack, unpack)` and `torch.autograd.graph.reset_saved_tensors_default_hooks()`.
These functions control the hooks applied to saved tensors: all tensors *saved* in that context will be packed using the `pack` function, then unpacked accordingly when needed.

Currently, this works by simply calling `register_hooks` (cf #60975) directly at the end of the constructor of a SavedVariable. This could be optimized further by not performing the copy before registering default hooks, but this would require a small refactor. Edit: the refactor is done in #61927.

A current limitation is that if users create tensors in this context, they will not be able to register additional hooks on the saved tensor.

For instance, to perform something like #28997, one could define a `pack` function that saves to disk whenever the tensor size is too big and returns a filename, then `unpack` simply reads the content of the file and outputs a tensor, e.g.:

```python
def pack(x):
    name = os.path.join(tmp_dir, str(uuid.uuid4()))
    torch.save(x, name)
    return name

def unpack(name):
    return torch.load(name)
```

Differential Revision: [D29792193](https://our.internmc.facebook.com/intern/diff/D29792193)

[ghstack-poisoned]

Fixes #58659. This PR builds directly on top of [#58512](#58512 (comment)).

~~Creates a context manager `with torch.autograd.graph.saved_tensors_default_hooks(pack, unpack)` that can be used on the Python side.~~

Expose a pair of functions to Python users: `torch.autograd.graph.set_saved_tensors_default_hooks(pack, unpack)` and `torch.autograd.graph.reset_saved_tensors_default_hooks()`.
These functions control the hooks applied to saved tensors: all tensors *saved* in that context will be packed using the `pack` function, then unpacked accordingly when needed.

Currently, this works by simply calling `register_hooks` (cf #60975) directly at the end of the constructor of a SavedVariable. This could be optimized further by not performing the copy before registering default hooks, but this would require a small refactor. Edit: the refactor is done in #61927.

A current limitation is that if users create tensors in this context, they will not be able to register additional hooks on the saved tensor.

For instance, to perform something like #28997, one could define a `pack` function that saves to disk whenever the tensor size is too big and returns a filename, then `unpack` simply reads the content of the file and outputs a tensor, e.g.:

```python
def pack(x):
    name = os.path.join(tmp_dir, str(uuid.uuid4()))
    torch.save(x, name)
    return name

def unpack(name):
    return torch.load(name)
```

Differential Revision: [D29792193](https://our.internmc.facebook.com/intern/diff/D29792193)

[ghstack-poisoned]

Fixes #58659. This PR builds directly on top of [#58512](#58512 (comment)).

~~Creates a context manager `with torch.autograd.graph.saved_tensors_default_hooks(pack, unpack)` that can be used on the Python side.~~

Expose a pair of functions to Python users: `torch.autograd.graph.set_saved_tensors_default_hooks(pack, unpack)` and `torch.autograd.graph.reset_saved_tensors_default_hooks()`.
These functions control the hooks applied to saved tensors: all tensors *saved* in that context will be packed using the `pack` function, then unpacked accordingly when needed.

Currently, this works by simply calling `register_hooks` (cf #60975) directly at the end of the constructor of a SavedVariable. This could be optimized further by not performing the copy before registering default hooks, but this would require a small refactor. Edit: the refactor is done in #61927.

A current limitation is that if users create tensors in this context, they will not be able to register additional hooks on the saved tensor.

For instance, to perform something like #28997, one could define a `pack` function that saves to disk whenever the tensor size is too big and returns a filename, then `unpack` simply reads the content of the file and outputs a tensor, e.g.:

```python
def pack(x):
    name = os.path.join(tmp_dir, str(uuid.uuid4()))
    torch.save(x, name)
    return name

def unpack(name):
    return torch.load(name)
```

Differential Revision: [D29792193](https://our.internmc.facebook.com/intern/diff/D29792193)

[ghstack-poisoned]

Summary:
Pull Request resolved: #61834

Expose a pair of functions to Python users: torch.autograd.graph.set_saved_tensors_default_hooks(pack, unpack) and torch.autograd.graph.reset_saved_tensors_default_hooks().
These functions control the hooks applied to saved tensors: all tensors saved in that context will be packed using the pack function, then unpacked accordingly when needed.

Currently, this works by simply calling register_hooks (cf #60975) directly at the end of the constructor of a SavedVariable. This could be optimized further by not performing the copy before registering default hooks, but this would require a small refactor. Edit: the refactor is done in #61927.

A current limitation is that if users create tensors in this context, they will not be able to register additional hooks on the saved tensor.

For instance, to perform something like #28997, one could define a pack function that saves to disk whenever the tensor size is too big and returns a filename, then unpack simply reads the content of the file and outputs a tensor, e.g.:

```
def pack(x):
    name = os.path.join(tmp_dir, str(uuid.uuid4()))
    torch.save(x, name)
    return name

def unpack(name):
    return torch.load(name)
```

Test Plan: Imported from OSS

Reviewed By: zou3519

Differential Revision: D29792193

Pulled By: Varal7

fbshipit-source-id: 33e931230ef59faa3ec8b5d11ef7c05539bce77c

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (pytorch#52451), how to register packing / unpacking
hooks (pytorch#60975) and how to use default hooks (pytorch#61834)

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Summary:

Expose a pair of functions to Python users: torch.autograd.graph.set_saved_tensors_default_hooks(pack, unpack) and torch.autograd.graph.reset_saved_tensors_default_hooks().
These functions control the hooks applied to saved tensors: all tensors saved in that context will be packed using the pack function, then unpacked accordingly when needed.

Currently, this works by simply calling register_hooks (cf #60975) directly at the end of the constructor of a SavedVariable. This could be optimized further by not performing the copy before registering default hooks, but this would require a small refactor. Edit: the refactor is done in #61927.

A current limitation is that if users create tensors in this context, they will not be able to register additional hooks on the saved tensor.

For instance, to perform something like #28997, one could define a pack function that saves to disk whenever the tensor size is too big and returns a filename, then unpack simply reads the content of the file and outputs a tensor, e.g.:

```
def pack(x):
    name = os.path.join(tmp_dir, str(uuid.uuid4()))
    torch.save(x, name)
    return name

def unpack(name):
    return torch.load(name)
```

Relanding previous PR: #61834

Original PR led to timeout error in: https://www.internalfb.com/mast/job/yuguo-release_canary_offline_training-inlinecvrp_a-canary_offline_train_28a7ecfc

Now passing: https://www.internalfb.com/mast/job/quach-release_canary_offline_training-inlinecvrp_a-canary_offline_train_9bb57e98

The difference with the new version is we don't need to acquire the GIL when calling `PyDefaultSavedVariableHooks::get_hooks`.

[ghstack-poisoned]

Summary:

Expose a pair of functions to Python users: torch.autograd.graph.set_saved_tensors_default_hooks(pack, unpack) and torch.autograd.graph.reset_saved_tensors_default_hooks().
These functions control the hooks applied to saved tensors: all tensors saved in that context will be packed using the pack function, then unpacked accordingly when needed.

Currently, this works by simply calling register_hooks (cf #60975) directly at the end of the constructor of a SavedVariable. This could be optimized further by not performing the copy before registering default hooks, but this would require a small refactor. Edit: the refactor is done in #61927.

A current limitation is that if users create tensors in this context, they will not be able to register additional hooks on the saved tensor.

For instance, to perform something like #28997, one could define a pack function that saves to disk whenever the tensor size is too big and returns a filename, then unpack simply reads the content of the file and outputs a tensor, e.g.:

```
def pack(x):
    name = os.path.join(tmp_dir, str(uuid.uuid4()))
    torch.save(x, name)
    return name

def unpack(name):
    return torch.load(name)
```

Relanding previous PR: #61834

Original PR led to timeout error in: https://www.internalfb.com/mast/job/yuguo-release_canary_offline_training-inlinecvrp_a-canary_offline_train_28a7ecfc

Now passing: https://www.internalfb.com/mast/job/quach-release_canary_offline_training-inlinecvrp_a-canary_offline_train_9bb57e98

The difference with the new version is we don't need to acquire the GIL when calling `PyDefaultSavedVariableHooks::get_hooks`.

[ghstack-poisoned]

Summary:
Pull Request resolved: #62563

Expose a pair of functions to Python users: torch.autograd.graph.set_saved_tensors_default_hooks(pack, unpack) and torch.autograd.graph.reset_saved_tensors_default_hooks().
These functions control the hooks applied to saved tensors: all tensors saved in that context will be packed using the pack function, then unpacked accordingly when needed.

Currently, this works by simply calling register_hooks (cf #60975) directly at the end of the constructor of a SavedVariable. This could be optimized further by not performing the copy before registering default hooks, but this would require a small refactor. Edit: the refactor is done in #61927.

A current limitation is that if users create tensors in this context, they will not be able to register additional hooks on the saved tensor.

For instance, to perform something like #28997, one could define a pack function that saves to disk whenever the tensor size is too big and returns a filename, then unpack simply reads the content of the file and outputs a tensor, e.g.:

```
def pack(x):
    name = os.path.join(tmp_dir, str(uuid.uuid4()))
    torch.save(x, name)
    return name

def unpack(name):
    return torch.load(name)
```

Relanding previous PR: #61834

Original PR led to timeout error in: https://www.internalfb.com/mast/job/yuguo-release_canary_offline_training-inlinecvrp_a-canary_offline_train_28a7ecfc

Now passing: https://www.internalfb.com/mast/job/quach-release_canary_offline_training-inlinecvrp_a-canary_offline_train_9bb57e98

The difference with the new version is we don't need to acquire the GIL when calling `PyDefaultSavedVariableHooks::get_hooks`.

Test Plan: Imported from OSS

Reviewed By: iramazanli

Differential Revision: D30045405

Pulled By: Varal7

fbshipit-source-id: 7f6c07af3a56fe8835d5edcc815c15ea4fb4e332

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

ghstack-source-id: 9d9aa29
Pull Request resolved: #63647

Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

[ghstack-poisoned]

Summary:
Add section to the Autograd mechanics docs to describe the recently
exposed saved tensors (#52451), how to register packing / unpacking
hooks (#60975) and how to use default hooks (#61834)

Sister PR: #62361 (will add a link from autograd.rst to notes/autograd in whatever PR does not land first)

Pull Request resolved: #62362

Reviewed By: soulitzer

Differential Revision: D30453177

Pulled By: Varal7

fbshipit-source-id: f5759977b069ff0ef36a47b08856d297691a6caa