NOCOMMIT: log peak mem usage and swapouts to debug test-llava-runner

Attempting to test my theory that the timestamp gaps in #8180 are caused by swapping.

ghstack-source-id: 45e7e64
ghstack-comment-id: 2634969596
Pull Request resolved: #8192

Author: swolchok

examples/models/llama/source_transformation/quantized_kv_cache.py

@@ -205,8 +205,11 @@ def replace_kv_cache_with_quantized_kv_cache(module):
     # This is needed to ensure that custom ops are registered
     from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
 
+    import resource
+    maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    nswap = resource.getrusage(resource.RUSAGE_SELF).ru_nswap
     logging.warning(
-        "Replacing KVCache with QuantizedKVCache. This modifies the model in place."
+        f"Replacing KVCache with QuantizedKVCache. This modifies the model in place. (HACK: maxrss: {maxrss} nswap: {nswap})"
     )
     for name, child in module.named_children():
         if isinstance(child, KVCache) or isinstance(child, CustomKVCache):
@@ -270,8 +273,11 @@ 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 resource
+    maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    nswap = resource.getrusage(resource.RUSAGE_SELF).ru_nswap
     logging.warning(
-        "Replacing KVCache with CustomKVCache. This modifies the model in place."
+        f"Replacing KVCache with CustomKVCache. This modifies the model in place. (HACK: maxrss: {maxrss} nswap: {nswap})"
     )
     for name, child in module.named_children():
         if isinstance(child, KVCache):

extension/llm/export/builder.py

@@ -156,7 +156,10 @@ 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 resource
+        maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        nswap = resource.getrusage(resource.RUSAGE_SELF).ru_nswap
+        logging.info(f"Model after source transforms: {self.model} (HACK: maxrss: {maxrss} nswap: {nswap})")
         return self
 
     def _get_dynamic_shape(self) -> Any: