How to Profile ML Workloads to Find Performance Bottlenecks in cs249r_book

The Profiler class in the cs249r_book repository provides lightweight instrumentation to quantify model parameters, estimate floating-point operations (FLOPs), and measure wall-clock latency, enabling systematic identification of compute and memory bottlenecks in machine learning models.

The cs249r_book educational framework includes TinyTorch, a minimal deep learning library designed for systems-level understanding. When you need to profile ML workloads within this ecosystem, the core implementation resides in tinytorch/src/14_profiling/14_profiling.py, which exposes a unified interface for static analysis and dynamic benchmarking without requiring external profiling tools.

The Profiler Architecture

The Profiler class serves as the central instrumentation engine. Upon instantiation, it initializes a measurements dictionary and a defaultdict(int) for operation counters, creating a stateful context for accumulating metrics across multiple runs. The class introspects models through their public APIs—checking for layers, parameters, or weight attributes—to gather structural information without modifying the underlying layer implementations. All metric results are stored in self.measurements for programmatic access or downstream visualization.

Measuring Model Complexity

Static analysis methods compute theoretical costs before execution, helping you identify architectural inefficiencies early in the design cycle.

Parameter Counting

The Profiler.count_parameters() method traverses the model hierarchy to tally trainable parameters. It first attempts to iterate over a layers attribute, then falls back to a generic parameters() iterator. For single-layer models, it delegates to _count_layer_parameters(), which inspects the weight and bias tensors directly. This approach works for both sequential containers and individual modules like Linear or Conv2d.

FLOP Estimation

To estimate computational cost, the profiler implements analytical FLOP counters for specific layer types. The _count_linear_flops() helper calculates operations as 2 * input_features * output_features * batch_size (accounting for multiply-accumulate operations), while _count_conv_flops() computes 2 * kernel_h * kernel_w * in_channels * out_channels * out_h * out_w * batch_size. For sequential models, _count_sequential_flops() aggregates these values across the layer stack, returning the total FLOP count for a single forward pass.

Measuring Runtime Performance

While static metrics indicate theoretical load, wall-clock time reveals real-world execution bottlenecks including kernel launch overhead and memory bandwidth limitations.

The Profiler.measure_latency() method implements a robust timing protocol. It executes a configurable number of warm-up iterations to stabilize CPU caches and GPU contexts, then runs the target model for a specified number of iterations, recording the average execution time in seconds. This method accepts any TinyTorch module and input tensor, making it suitable for benchmarking individual layers or complete networks.

Practical Profiling Workflow

Combine these utilities to diagnose a model's performance profile. The following example demonstrates profiling a simple linear classifier:

from tinytorch.perf.profiling import Profiler
from tinytorch.core.layers import Linear
from tinytorch.core.tensor import Tensor
import numpy as np

# Initialize model and profiler

model = Linear(1024, 512)
profiler = Profiler()

# 1. Static analysis: parameter count and FLOPs

params = profiler.count_parameters(model)
print(f"Total parameters: {params:,}")

# For FLOPs, we use the internal helper (batch_size=1)

flops = profiler._count_linear_flops(model, input_shape=(1, 1024))
print(f"FLOPs per forward pass: {flops:,}")

# 2. Dynamic analysis: latency measurement

input_tensor = Tensor(np.random.randn(1, 1024))
avg_latency = profiler.measure_latency(
    model, 
    input_tensor, 
    warmup=5, 
    iterations=100
)
print(f"Average latency: {avg_latency*1000:.3f} ms")

Interpretation guidelines:

• High parameter count with low FLOPs indicates memory-bound workloads where weight loading dominates execution time.
• High FLOPs relative to latency suggests compute-bound operations that benefit from hardware acceleration.
• Disproportionate latency compared to FLOP estimates signals inefficiencies such as excessive Python overhead or poor memory access patterns.

Key Source Files

The profiling system spans several locations within the repository:

• tinytorch/src/14_profiling/14_profiling.py — Contains the complete Profiler implementation including parameter counting, FLOP estimation, and latency measurement logic.
• tinytorch/perf/__init__.py — Serves as the public entry point, exposing the profiling utilities through the tinytorch.perf namespace.
• tinytorch/tests/14_profiling/test_profiler_core.py — Validates all public methods including edge cases for parameter counting and latency measurement accuracy.
• tinytorch/src/20_capstone/20_capstone.py — Demonstrates integration of measure_latency within a full training pipeline benchmark.

Summary

• The Profiler class in tinytorch/src/14_profiling/14_profiling.py unifies static analysis (parameters, FLOPs) and dynamic benchmarking (latency) in a single interface.
• Parameter counting supports both sequential containers and individual layers through introspection of layers, parameters, or weight attributes.
• FLOP estimation uses analytical formulas for Linear and Conv2d layers, aggregating costs across sequential stacks.
• Latency measurement implements warm-up and statistical averaging to provide reproducible wall-clock timings.
• The profiling API integrates seamlessly with TinyTorch tensors and layers, requiring no external dependencies.

Frequently Asked Questions

How do I access the Profiler in my cs249r_book project?

Import the Profiler class from the tinytorch.perf.profiling module. The tinytorch/perf/__init__.py file exposes this interface, allowing you to instantiate Profiler() and immediately begin calling count_parameters() or measure_latency() on your models.

Can I profile custom layers that are not Linear or Conv2d?

Yes. For parameter counting, the profiler falls back to iterating over parameters() or inspecting weight attributes directly. For FLOP estimation, you may need to extend the Profiler class with custom logic for your specific layer's mathematical operations, as the built-in FLOP counters specialize in matrix multiplication and convolution dimensions.

What is the difference between FLOPs and latency in this context?

FLOPs (floating-point operations) represent a static, hardware-agnostic count of arithmetic operations required for one forward pass, calculated analytically from layer dimensions. Latency measures the actual wall-clock time elapsed during execution, which includes not only arithmetic but also memory transfers, kernel launch overhead, and Python interpreter delays. A model with low FLOPs but high latency is likely memory-bound or suffering from overhead inefficiencies.

How accurate are the latency measurements compared to production profilers?

The measure_latency() method provides millisecond-level precision suitable for educational and relative benchmarking purposes. It uses standard Python timing utilities with warm-up iterations to minimize jitter. For production-grade nanosecond precision or hardware-level counter data, you would need to integrate platform-specific tools such as NVIDIA Nsight or Intel VTune, though the TinyTorch profiler offers sufficient accuracy for identifying architectural bottlenecks during model development.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}