Function profile_model

profile.py:529–648

Profile the model and return a list of KernelRecords sorted by GPU time desc. Returns: Tuple of (records, extras) where extras contains paths to exported artifacts and optional HTA analysis results.

(
    model: nn.Module,
    inputs: Dict[str, Any],
    warmup_iters: int = WARMUP_ITERS,
    profile_iters: int = PROFILE_ITERS,
    export_trace: bool = False,
    memory_snapshot: bool = False,
)

def profile_model(
    model: nn.Module,
    inputs: Dict[str, Any],
    warmup_iters: int = WARMUP_ITERS,
    profile_iters: int = PROFILE_ITERS,
    export_trace: bool = False,
    memory_snapshot: bool = False,
) -> Tuple[List[KernelRecord], Dict[str, Any]]:
    """Profile the model and return a list of KernelRecords sorted by GPU time desc.

    Returns:
        Tuple of (records, extras) where extras contains paths to exported
        artifacts and optional HTA analysis results.
    """
    extras: Dict[str, Any] = {}

    os.makedirs(WORKSPACE_DIR, exist_ok=True)
    trace_path = os.path.join(WORKSPACE_DIR, "trace.json")
    snapshot_path = os.path.join(WORKSPACE_DIR, "memory_snapshot.pickle")

    # --- Warmup ---
    with torch.no_grad():
        for _ in range(warmup_iters):
            _run_forward(model, inputs)

    torch.cuda.synchronize()

    # --- Start memory recording if requested ---
    if memory_snapshot:
        try:
            torch.cuda.memory._record_memory_history(max_entries=100000)
        except Exception as e:
            print(f"  WARNING: Could not start memory history recording: {e}")
            memory_snapshot = False

    # --- Profile ---
    with torch.no_grad():
        with torch.profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            record_shapes=True,
            with_stack=False,
        ) as prof:
            for _ in range(profile_iters):
                _run_forward(model, inputs)
                torch.cuda.synchronize()

    # --- Export Chrome trace ---
    if export_trace:
        try:
            prof.export_chrome_trace(trace_path)
            extras["trace_path"] = trace_path
        except Exception as e:
            print(f"  WARNING: Could not export Chrome trace: {e}")

    # --- Capture memory snapshot ---