Measure memory

Author: ciaranbor

src/exo/utils/memory.py

@@ -0,0 +1,17 @@
+from time import sleep
+
+import mlx.core as mx
+import psutil
+
+from exo.worker.runner.bootstrap import logger
+
+
+def log_memory(label: str) -> None:
+    """Log current and peak GPU memory usage."""
+    active_mb = mx.get_active_memory() / (1024 * 1024)
+    peak_mb = mx.get_peak_memory() / (1024 * 1024)
+    res_mem = psutil.Process().memory_info().rss / (1024 * 1024)
+    logger.info(
+        f"[MEMORY] {label}: active={active_mb:.1f}MB, peak={peak_mb:.1f}MB, res={res_mem:.1f}MB"
+    )
+    sleep(3)

src/exo/worker/engines/image/distributed_model.py

@@ -7,6 +7,7 @@
 
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.utils.memory import log_memory
 from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models import (
@@ -43,8 +44,11 @@ def __init__(
         quantize: int | None = None,
     ):
         # Get model config and create adapter (adapter owns the model)
+        mx.metal.reset_peak_memory()
+
         config = get_config_for_model(model_id)
         adapter = create_adapter_for_model(config, model_id, local_path, quantize)
+        log_memory("After model load (adapter created)")
 
         if group is not None:
             adapter.slice_transformer_blocks(
@@ -53,6 +57,7 @@ def __init__(
                 total_joint_blocks=config.joint_block_count,
                 total_single_blocks=config.single_block_count,
             )
+            log_memory("After block slicing")
 
         # Create diffusion runner (handles both single-node and distributed modes)
         num_sync_steps = config.get_num_sync_steps("medium") if group else 0
@@ -63,20 +68,23 @@ def __init__(
             shard_metadata=shard_metadata,
             num_sync_steps=num_sync_steps,
         )
+        log_memory("After DiffusionRunner creation")
 
         if group is not None:
             logger.info("Initialized distributed diffusion runner")
 
             mx.eval(adapter.model.parameters())
+            log_memory("After mx.eval(parameters)")
 
-            # TODO(ciaran): Do we need this?
             mx.eval(adapter.model)
+            log_memory("After mx.eval(model)")
 
             # Synchronize processes before generation to avoid timeout
             mx_barrier(group)
             logger.info(f"Transformer sharded for rank {group.rank()}")
         else:
             logger.info("Single-node initialization")
+            log_memory("Single-node init complete")
 
         object.__setattr__(self, "_config", config)
         object.__setattr__(self, "_adapter", adapter)

src/exo/worker/engines/image/generate.py

@@ -1,11 +1,14 @@
 import io
 from typing import Generator, Literal
 
+import mlx.core as mx
 from PIL import Image
 
 from exo.shared.types.api import ImageGenerationTaskParams
 from exo.shared.types.worker.runner_response import ImageGenerationResponse
+from exo.utils.memory import log_memory
 from exo.worker.engines.image.base import ImageGenerator
+from exo.worker.runner.bootstrap import logger
 
 
 def parse_size(size_str: str | None) -> tuple[int, int]:
@@ -26,13 +29,16 @@ def parse_size(size_str: str | None) -> tuple[int, int]:
 
 
 def warmup_image_generator(model: ImageGenerator) -> Image.Image | None:
-    return model.generate(
+    log_memory("Before warmup generation")
+    result = model.generate(
         prompt="Warmup",
         height=256,
         width=256,
         quality="low",
         seed=2,
     )
+    log_memory("After warmup generation")
+    return result
 
 
 def generate_image(

src/exo/worker/engines/image/models/qwen/adapter.py

@@ -15,6 +15,7 @@
 )
 from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
 
+from exo.utils.memory import log_memory
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models.base import BaseModelAdapter
 from exo.worker.engines.image.pipeline.adapter import (
@@ -307,7 +308,11 @@ def final_projection(
     ) -> mx.array:
         """Apply final normalization and projection."""
         hidden_states = self._transformer.norm_out(hidden_states, text_embeddings)
-        return self._transformer.proj_out(hidden_states)
+        mx.eval(hidden_states)
+        log_memory("after norm out")
+        hidden_states = self._transformer.proj_out(hidden_states)
+        mx.eval(hidden_states)
+        log_memory("after proj_out")
 
     def get_joint_blocks(self) -> list[JointBlockInterface]:
         """Return all 60 transformer blocks."""

src/exo/worker/engines/image/pipeline/runner.py

@@ -9,6 +9,7 @@
 from tqdm import tqdm
 
 from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.utils.memory import log_memory
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.pipeline.adapter import (
     BlockWrapperMode,
@@ -194,8 +195,14 @@ def generate_image(
             GeneratedImage result
         """
         runtime_config = RuntimeConfig(settings, self.adapter.model.model_config)
+
         latents = self.adapter.create_latents(seed, runtime_config)
+        mx.eval(latents)
+        log_memory("generate_image: after create_latents")
+
         prompt_data = self.adapter.encode_prompt(prompt)
+        mx.eval(prompt_data)
+        log_memory("generate_image: after encode_prompt")
 
         latents = self._run_diffusion_loop(
             latents=latents,
@@ -204,9 +211,14 @@ def generate_image(
             seed=seed,
             prompt=prompt,
         )
+        mx.eval(latents)
+        log_memory("generate_image: after diffusion_loop")
 
         if self.is_last_stage:
-            return self.adapter.decode_latents(latents, runtime_config, seed, prompt)
+            result = self.adapter.decode_latents(latents, runtime_config, seed, prompt)
+            mx.eval(result)
+            log_memory("generate_image: after decode_latents")
+            return result
 
     def _run_diffusion_loop(
         self,
@@ -260,6 +272,7 @@ def _run_diffusion_loop(
                 )
 
                 mx.eval(latents)
+                log_memory("after diffusion step")
 
             except KeyboardInterrupt:  # noqa: PERF203
                 Callbacks.interruption(
@@ -310,16 +323,24 @@ def _forward_pass(
         if config is None:
             raise ValueError("config must be provided in kwargs")
         scaled_latents = config.scheduler.scale_model_input(latents, t)
+        mx.eval(scaled_latents)
+        log_memory("after scaling model input")
 
         hidden_states, encoder_hidden_states = self.adapter.compute_embeddings(
             scaled_latents, prompt_embeds
         )
+        mx.eval(hidden_states, encoder_hidden_states)
+        log_memory("after computing embeddings")
         text_embeddings = self.adapter.compute_text_embeddings(
             t, config, pooled_prompt_embeds, hidden_states=hidden_states
         )
+        mx.eval(text_embeddings)
+        log_memory("after computing text embeddings")
         rotary_embeddings = self.adapter.compute_rotary_embeddings(
             prompt_embeds, config, **kwargs
         )
+        mx.eval(rotary_embeddings)
+        log_memory("after computing rotary embeddings")
 
         text_seq_len = prompt_embeds.shape[1]
 
@@ -337,6 +358,8 @@ def _forward_pass(
                 **kwargs,
             )
 
+        mx.eval(hidden_states, encoder_hidden_states)
+        log_memory("after joint blocks")
         # Merge streams
         if self.joint_block_wrappers:
             hidden_states = self.adapter.merge_streams(
@@ -354,9 +377,15 @@ def _forward_pass(
                 mode=BlockWrapperMode.CACHING,
             )
 
+        mx.eval(hidden_states)
+        log_memory("after single blocks")
+
         # Extract image portion and project
         hidden_states = hidden_states[:, text_seq_len:, ...]
-        return self.adapter.final_projection(hidden_states, text_embeddings)
+        hidden_states = self.adapter.final_projection(hidden_states, text_embeddings)
+        mx.eval(hidden_states)
+        log_memory("after final projection")
+        return hidden_states
 
     def _diffusion_step(
         self,
@@ -371,7 +400,10 @@ def _diffusion_step(
         configuration and current timestep.
         """
         if self.group is None:
-            return self._single_node_step(t, config, latents, prompt_data)
+            latents = self._single_node_step(t, config, latents, prompt_data)
+            mx.eval(latents)
+            log_memory("single node step complete")
+            return latents
         elif t < self.num_sync_steps:
             return self._sync_pipeline(
                 t,
@@ -435,6 +467,7 @@ def _single_node_step(
                 kwargs,
             )
 
+        log_memory("scheduler step")
         return config.scheduler.step(model_output=noise, timestep=t, sample=latents)
 
     def _initialize_kv_caches(