mapping_networks.distributed.profile_peak_memory

mapping_networks.distributed.profile_peak_memory(fn: Callable[[...], Any], *args: Any, device: str = 'cpu', **kwargs: Any) → tuple[Any, int]

Measure the peak memory used by fn during its execution.

On CPU the measurement uses tracemalloc to track Python/C-level allocations. On CUDA it resets torch.cuda.max_memory_allocated before the call and reads it afterwards.

Parameters:

  • fn – Callable to profile.

  • *args – Positional arguments forwarded to fn.

  • device – "cpu" (default) or "cuda". Only affects which counter is used; the function itself must have already placed tensors on the correct device.

  • **kwargs – Keyword arguments forwarded to fn.

Returns:

A (result, peak_bytes) tuple where result is the return value of fn and peak_bytes is the peak allocation in bytes.

Example:

result, peak = profile_peak_memory(model.generator.generate_parameters)
print(f"Peak: {peak / 1e6:.2f} MB")