diff --git a/examples/models/llama/source_transformation/quantized_kv_cache.py b/examples/models/llama/source_transformation/quantized_kv_cache.py index 023fc6800ff..2320eec529e 100644 --- a/examples/models/llama/source_transformation/quantized_kv_cache.py +++ b/examples/models/llama/source_transformation/quantized_kv_cache.py @@ -202,11 +202,15 @@ def replace_kv_cache_with_quantized_kv_cache(module): torch.ops.load_library(libs[0]) op = torch.ops.quantized_decomposed.quantize_per_token.out assert op is not None + import gc + import resource + # This is needed to ensure that custom ops are registered from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + rusage = resource.getrusage(resource.RUSAGE_SELF) logging.warning( - "Replacing KVCache with QuantizedKVCache. This modifies the model in place." + f"Replacing KVCache with QuantizedKVCache. This modifies the model in place. (HACK: rusage: {rusage} gc stats: {gc.get_stats()})" ) for name, child in module.named_children(): if isinstance(child, KVCache) or isinstance(child, CustomKVCache): @@ -270,8 +274,12 @@ def replace_kv_cache_with_custom_kv_cache(module): This is because the custom op treats second dim as sequence dim. Future work: support [B, H, S, D] """ + import gc + import resource + + rusage = resource.getrusage(resource.RUSAGE_SELF) logging.warning( - "Replacing KVCache with CustomKVCache. This modifies the model in place." + f"Replacing KVCache with CustomKVCache. This modifies the model in place. (HACK: rusage: {rusage} gc stats: {gc.get_stats()})" ) for name, child in module.named_children(): if isinstance(child, KVCache): diff --git a/examples/models/llava/test/test_llava.py b/examples/models/llava/test/test_llava.py index 5fd60399415..aeea0bb964a 100644 --- a/examples/models/llava/test/test_llava.py +++ b/examples/models/llava/test/test_llava.py @@ -24,6 +24,8 @@ logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) +import gc +gc.set_debug(gc.DEBUG_STATS | gc.DEBUG_COLLECTABLE) class TestLlava(unittest.TestCase): def setUp(self): diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py index be6977b639c..296f6a1efd5 100644 --- a/extension/llm/export/builder.py +++ b/extension/llm/export/builder.py @@ -156,7 +156,13 @@ def source_transform( if self.verbose: logging.info(f"Applied source transforms: {self.applied_source_transforms}") - logging.info(f"Model after source transforms: {self.model}") + import gc + import resource + + rusage = resource.getrusage(resource.RUSAGE_SELF) + logging.info( + f"Model after source transforms: {self.model} (HACK: rusage: {rusage} gc_stats: {gc.get_stats()})" + ) return self def _get_dynamic_shape(self) -> Any: