tree flatten -> tree_leaves

Summary

Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

Author: HDCharles

src/llmcompressor/modifiers/smoothquant/base.py

@@ -6,7 +6,7 @@
 from loguru import logger
 from pydantic import ConfigDict, Field
 from torch.nn import Module
-from torch.utils._pytree import tree_flatten
+from torch.utils._pytree import tree_leaves
 
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers import Modifier
@@ -202,7 +202,7 @@ def _resolve_mappings(self, model: Module) -> List[SmoothQuantMapping]:
         module_to_name = get_module_to_name_dict(model)
         for mapping in self.mappings:
             for *nested_balance_layers, smooth_layers in match_modules_set(
-                model, tree_flatten(mapping)[0], self.ignore
+                model, tree_leaves(mapping), self.ignore
             ):
                 assert len(smooth_layers) == 1, (
                     "SmoothQuant mappings must match a single smooth layer for each "
@@ -211,7 +211,7 @@ def _resolve_mappings(self, model: Module) -> List[SmoothQuantMapping]:
                 )
                 smooth_layer = smooth_layers[0]
                 smooth_name = module_to_name.get(smooth_layers[0])
-                balance_layers = tree_flatten(nested_balance_layers)[0]
+                balance_layers = tree_leaves(nested_balance_layers)
                 resolved_mappings.append(
                     SmoothQuantMapping(smooth_name, smooth_layer, balance_layers)
                 )

src/llmcompressor/modifiers/transform/spinquant/base.py

@@ -11,7 +11,7 @@
 )
 from compressed_tensors.utils import TorchDtype, get_head_dim
 from pydantic import Field, ValidationInfo, field_validator
-from torch.utils._pytree import tree_flatten
+from torch.utils._pytree import tree_leaves
 from transformers import PreTrainedModel
 
 from llmcompressor.core import Event, EventType, State
@@ -208,7 +208,7 @@ def _fuse_norms(self, model: PreTrainedModel):
             ):
                 # match_modules_set returns a list of lists
                 assert len(norm) == 1
-                fuse_norm_linears(norm[0], tree_flatten(linears)[0])
+                fuse_norm_linears(norm[0], tree_leaves(linears))
 
     def _create_r1_scheme(self) -> TransformScheme:
         return TransformScheme(