From b8ae38702cc252ee99129f05ca44e9097d015915 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Tue, 10 Dec 2024 18:59:59 +0000 Subject: [PATCH] add lifecycle docstring --- src/compressed_tensors/utils/offload.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/compressed_tensors/utils/offload.py b/src/compressed_tensors/utils/offload.py index 6c20da08..f99e2b91 100644 --- a/src/compressed_tensors/utils/offload.py +++ b/src/compressed_tensors/utils/offload.py @@ -11,6 +11,19 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +""" +Utilities associated with offloading functionality provided by `accelerate`. + +| ----------------------------------------------------------------------------------------------------- | # noqa: E501 +| Operation | Without offloading support | With offloading support | # noqa: E501 +| --------- | -------------------------------------- | ------------------------------------------------ | # noqa: E501 +| Add | module.register_parameter(name, param) | register_offload_parameter(module, name, param) | # noqa: E501 +| Check | N/A | has_offloaded_params(module) | # noqa: E501 +| Onload | N/A | with align_module_device(module) | # noqa: E501 +| Update | module.name.data.copy_(new_data) | update_offload_parameter(module, name, new_data) | # noqa: E501 +| Delete | del module.name | delete_offload_parameter(module, name) | # noqa: E501 +| ----------------------------------------------------------------------------------------------------- | # noqa: E501 +""" import contextlib from functools import wraps