set_layer -> Module.set_submodule

Author: ved1beta

src/llmcompressor/modifiers/distillation/output/base.py

@@ -11,7 +11,7 @@
 )
 from llmcompressor.utils.fsdp.context import summon_full_params_context
 from llmcompressor.utils.fsdp.helpers import maybe_get_wrapped, set_wrapped_model
-from llmcompressor.utils.pytorch.module import get_layers, set_layer
+from llmcompressor.utils.pytorch.module import get_layers
 
 __all__ = ["OutputDistillationModifier"]
 
@@ -85,8 +85,8 @@ def on_initialize(self, state: State, **kwargs) -> bool:
 
         with summon_full_params_context(state.teacher_model, offload_to_cpu=True):
             for key, (student_wrapper, teacher_wrapper) in self.wrappers_.items():
-                set_layer(key, student_wrapper, state.model)
-                set_layer(key, teacher_wrapper, state.teacher_model)
+                Module.set_submodule(key, student_wrapper, state.model)
+                Module.set_submodule(key, teacher_wrapper, state.teacher_model)
 
         self.wrapped_kd_model_ = self._create_model_wrapper(
             student_model=maybe_get_wrapped(state.model),
@@ -109,8 +109,8 @@ def on_finalize(self, state: State, **kwargs) -> bool:
 
         with summon_full_params_context(state.teacher_model, offload_to_cpu=True):
             for key, (student_wrapper, teacher_wrapper) in self.wrappers_.items():
-                set_layer(key, student_wrapper.layer, state.model)
-                set_layer(key, teacher_wrapper.layer, state.teacher_model)
+                Module.set_submodule(key, student_wrapper.layer, state.model)
+                Module.set_submodule(key, teacher_wrapper.layer, state.teacher_model)
                 del student_wrapper
                 del teacher_wrapper
 

src/llmcompressor/utils/pytorch/module.py

@@ -51,7 +51,6 @@
     "match_layers_params",
     "get_layers",
     "get_layer",
-    "set_layer",
     "get_params",
     "get_param",
     "get_terminal_layers",
@@ -197,22 +196,6 @@ def get_layer(target: str, module: Module) -> Tuple[str, Module]:
     return name, layer
 
 
-def set_layer(target: str, layer: Module, module: Module) -> Module:
-    with summon_full_params_context(module):
-        # importing here to avoid circular import
-        from llmcompressor.utils.fsdp.helpers import maybe_get_wrapped
-
-        parent_target = ".".join(target.split(".")[:-1])
-        if parent_target != "":
-            parent_layer = get_layer(parent_target, module)[1]
-        else:
-            parent_layer = maybe_get_wrapped(module)
-        old_layer = getattr(parent_layer, target.split(".")[-1])
-        setattr(parent_layer, target.split(".")[-1], layer)
-
-    return old_layer
-
-
 def get_params(targets: Union[str, List[str]], module: Module) -> Dict[str, Parameter]:
     return match_layers_params(targets, module, params=True)