profile-model

Files (3)

• SKILL.md
• eval.yaml · evals/
• simple_trace.json · evals/fixtures/analyze_trace_basic/

SKILL.md (first lines)

# Profile Model

Analyze PyTorch profiler traces (`.json`/`.json.gz` files) to diagnose GPU performance issues, find slow kernels, and identify bottlenecks.

---

## Step 1: Get Trace Path

Ask the user for the trace file path. Traces are typically:
- Single file: `/path/to/trace.json` or `/path/to/trace.json.gz`
- Directory (multi-rank): `/path/to/traces/` containing `rank0.json`, `rank1.json`, etc.

## Step 2: Load Trace as DataFrame

Load the trace into a pandas DataFrame for flexible querying.

**⚠️ IMPORTANT: Ask user before caching.** Do NOT cache automatically. Ask:

> "Would you like me to cache this trace as Parquet for faster future loads?
> This saves to `/tmp/trace_cache/` and makes subsequent loads faster.
> (Recommended for large traces or if you plan to ask follow-up questions)"

Only proceed with caching if the user agrees.

**DataFrame structure:**
- Standard columns: `name`, `cat`, `ph`, `ts`, `dur`, `pid`, `tid`, `id`
- Args flattened to: `args_device`, `args_stream`, `args_Input Dims`, etc.

**Loading approach:**
1. Parse JSON/JSON.gz trace file
2. Flatten each event into a row, with `args` dict expanded to `args_*` columns
3. Convert `ts` and `dur` to numeric

**Parquet caching notes (if user agrees):**
- Cache to `/tmp/trace_cache/{basename}.parquet`
- Check cache freshness by comparing mtime with source file
- **Handle mixed types**: `args_*` columns often have mixed types (int, str, None). Before saving to Parquet, detect mixed-type columns and convert them to strings. Replace `'nan'`/`'None'` strings back to `None`.
- If Parquet save fails, retry with all `args_*` columns as strings
- If still fails, continue without cache

## Step 3: Quick Summary

After loading, show a quick summary:
- Total events and breakdown by category (`df["cat"].value_counts()`)
- Kernel count and total kernel time
- Top kernels by duration

## Step 4: Analyze with DataFrame Queries

Key column values for filtering:
- GPU kernels: `cat == "kernel"`
- CPU ops: `cat == "cpu_op"`
- NCCL/communication: `name.str.contains("nccl", case=False)`
- Memory ops: `cat` contains "mem" or "memcpy"

Timestamps (`ts`) and durations (`dur`) are in **microseconds (μs)**.

Use standard pandas operations (groupby, filter, sort, aggregate) to answer user questions.

## Step 5: Interpret Results

| Finding | Likely Issue | Recommendation |
|---------|--------------|----------------|
| Many small kernels (avg < 10μs) | Kernel launch overhead | Use `torch.compile`, fuse operations |
| One kernel dominates (>50%) | Expected for compute-bound | Check if kernel is optimized |
| Many memory operations | Memory-bound | Reduce copies, use pinned memory |
| High communication time | Distributed overhead | Overlap comm with compute |

## Step 6: Deep Analysis with HTA (Optional)

If the user wants deeper analysis, use HTA (Holistic Trace Analysis):

```python
# Check if HTA is available
try:
    from hta.trace_analysis import TraceAnalysis
except ImportError:
    print("HTA not installed. Install with: pip install HolisticTraceAnalysis")
```

#### HTA Analysis Functions

| User Question | HTA Function | What It Returns |
|---------------|--------------|-----------------|
| "What's taking GPU time?" | `get_gpu_kernel_breakdown()` | Kernel time distribution |
| "Where's the critical path?" | `critical_path_analysis()` | Bottleneck path through execution |
| "Why is GPU idle?" | `get_idle_time_breakdown()` | Host wait vs kernel wait breakdown |
| "Comm/compute overlap?" | `get_comm_comp_overlap()` | Overlap percentage per rank |
| "Any slow ranks?" | `get_potential_stragglers()` | Straggler rank identification |
| "Overall time breakdown?" | `get_temporal_breakdown()` | Idle/compute/non-compute time |

#### HTA Usage Examples

```python
from hta.trace_analysis import TraceAnalysis

# Load traces
analyzer = TraceAnalysis(trace_dir="/path/to/traces/")
# OR: analyzer = TraceAnalysis(trace_files={0: "/path/to/trace.json.gz"})

# GPU Kernel Breakdown
kernel_type_df, kernel_df = analyzer.get_gpu_kernel_breakdown(visualize=False)

# Temporal Breakdown (idle vs compute vs non-compute)
temporal_df = analyzer.get_temporal_breakdown(visualize=False)

# Idle Time Breakdown (why is GPU idle?)
idle_df, _ = analyzer.get_idle_time_breakdown(ranks=[0], visualize=False)

# Critical Path Analysis (for distributed training)
cp_graph, success = analyzer.critical_path_analysis(
    rank=0, annotation="ProfilerStep", instance_id=1
)
if success:
    analyzer.overlay_critical_path_analysis(
        rank=0, critical_path_graph=cp_graph, output_dir="/tmp/cp_output"
    )

# Communication-Computation Overlap

From meta-pytorch/skills, BSD-3-Clause.

More from meta-pytorch/skills