Move nf4 op registration from utils to modules (#1035)

pytorch · May 31, 2024 · 0c4056a · 0c4056a
1 parent 3b1bcf9
commit 0c4056a
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/torchtune/modules/low_precision/__init__.py b/torchtune/modules/low_precision/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/...htune/utils/_register_nf4_dispatch_ops.py → ...w_precision/_register_nf4_dispatch_ops.py b/...htune/utils/_register_nf4_dispatch_ops.py → ...w_precision/_register_nf4_dispatch_ops.py
diff --git a/torchtune/modules/peft/lora.py b/torchtune/modules/peft/lora.py
@@ -11,8 +11,8 @@
 from torch import nn, Tensor
 
 from torchao.dtypes.nf4tensor import linear_nf4, to_nf4
+from torchtune.modules.low_precision import _register_nf4_dispatch_ops  # noqa: F401
 from torchtune.modules.peft.peft_utils import AdapterModule
-from torchtune.utils import _register_nf4_dispatch_ops  # noqa: F401
 
 
 class LoRALinear(nn.Module, AdapterModule):