Refactor HF_MODULE_KEY

GitOrigin-RevId: 166f661

Author: changlan

axlearn/common/module.py

@@ -177,6 +177,8 @@ def forward(self, x: Tensor) -> Tensor:
 
 _CallableT = TypeVar("_CallableT", bound=Callable)
 
+HF_MODULE_KEY = "hf_module"
+
 
 def nowrap(fun: _CallableT) -> _CallableT:
     """Marks the specified module method as one that doesn't need to be wrapped.

axlearn/experiments/test_utils.py

@@ -27,6 +27,7 @@
 from axlearn.common.checkpointer import every_n_steps_policy
 from axlearn.common.config import TrainerConfigFn
 from axlearn.common.evaler import every_n_steps_policy as eval_every_n_steps_policy
+from axlearn.common.module import HF_MODULE_KEY
 from axlearn.common.summary_writer import SummaryWriter
 from axlearn.common.test_utils import (
     ParamInitSpec,
@@ -37,7 +38,6 @@
 )
 from axlearn.common.trainer import SpmdTrainer
 from axlearn.common.utils import NestedTree, Tensor, flatten_items, set_data_dir
-from axlearn.huggingface.hf_module import HF_MODULE_KEY
 
 
 def named_parameters(

axlearn/huggingface/hf_module.py

@@ -19,11 +19,9 @@
 from axlearn.common.base_layer import ParameterSpec
 from axlearn.common.base_model import BaseModel
 from axlearn.common.config import REQUIRED, Required, RequiredFieldMissingError, config_class
-from axlearn.common.module import Module, NestedTensor  # pylint: disable=unused-import
+from axlearn.common.module import HF_MODULE_KEY, Module, NestedTensor
 from axlearn.common.utils import Nested, Tensor
 
-HF_MODULE_KEY = "hf_module"
-
 
 def get_hf_models_cache_dir() -> Path:
     """Returns the directory where HF models are cached.