PyTorchModelMemoryProfiler Module

Helpers to profile PyTorch model memory usage during training and inference.

class HMB.PyTorchModelMemoryProfiler.PyTorchModelMemoryProfiler(model, inputShape, batchSize=1, precision='FP32', device='cpu')[source]

Bases: object

Profiler for estimating PyTorch model memory usage (parameters, gradients, activations, optimizer state, attention matrices) and compute (FLOPs/GFLOPs) for training and inference. Uses a dummy forward pass with hooks to gather activation shapes and sizes.

Parameters:

  • model (nn.Module) – The PyTorch model to profile.

  • inputShape (Tuple[int, ...]) – Shape of a single input sample (channels, H, W, …).

  • batchSize (int) – Batch size used for the dummy forward pass. Defaults to 1.

  • precision (str) – “FP32” or “FP16” – affects bytes-per-parameter calculations.

  • device (str) – Target device string (e.g., “cpu” or “cuda”). Used for the dummy forward pass attempt; profiling logic falls back to CPU if unavailable.

ProfileModelMemory(optimizerType='Adam', optimizerKwargs=None, isTransformer=False, sequenceLength=None, checkpointing=False, checkpointSavingsFactor=0.5, deviceFLOPSGFLOPS=None, datasetSize=None, trainingMultiplier=3.0, runMicroBenchmark=False)[source]

Perform a full memory and compute profile for the configured model. This runs a dummy forward pass (with hooks) to collect activation shapes, estimates parameter/buffer/optimizer/attention memory, and computes FLOPs-based performance estimates for training and inference.

Parameters:

  • optimizerType (str) – Optimizer used for estimating optimizer state memory.

  • optimizerKwargs (Optional[Dict[str, Any]]) – Options passed to optimizer estimation (e.g., {“amsgrad”: True}).

  • isTransformer (bool) – If True, uses transformer-specific attention memory estimation; otherwise auto-detection may enable it.

  • sequenceLength (int) – Required when isTransformer is True; sequence length used for attention memory estimation.

  • checkpointing (bool) – Whether gradient checkpointing is enabled (reduces retained activation memory estimate).

  • checkpointSavingsFactor (float) – Fraction of activation memory saved by checkpointing (0..1). Defaults to 0.5.

  • deviceFLOPSGFLOPS (Optional[float]) – Optional device GFLOPS peak to use for performance estimates. If None a heuristic/default is chosen.

  • datasetSize (Optional[int]) – Dataset size to estimate steps per epoch.

  • trainingMultiplier (float) – Factor to scale forward GFLOPs to training GFLOPs (includes backward and optimizer work). Defaults to 3.0.

  • runMicroBenchmark (bool) – If True attempts a small GEMM on the target device to empirically measure GFLOPS and refine timing estimates.

Returns:

A comprehensive dictionary containing memory breakdowns (bytes and MB), layer-wise activations, top-K lists, FLOPs estimates, and performance estimates.

Return type:

Dict[str, Any]

SaveProfileToJSON(memoryProfile, path)[source]

Persist a memory profile dictionary to a JSON file path using utf-8 encoding and pretty-print indentation for readability.

Parameters:

  • memoryProfile (Dict[str, Any]) – The profile produced by ProfileModelMemory.

  • path (str) – Filesystem path where the JSON will be written.

Return type:

None

PrintMemoryReport(memoryProfile)[source]

Nicely format and print a human-readable memory and performance report to standard output based on a profile produced by ProfileModelMemory.

Parameters:

memoryProfile (Dict[str, Any]) – The profile produced by ProfileModelMemory.

Return type:

None