# Profile Training > Profile a named JAX, Levanter, or Marin run, or investigate a measured startup, compilation, initialization, or throughput bottleneck. - Skill: `marin-community/profile-training` (Agent Skill) - Install (CLI): `npx skillmds@latest add marin-community/profile-training` - Raw SKILL.md: https://api.skillmd.com/api/skills/marin-community/profile-training/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - Author: marin-community (https://skillmd.com/u/marin-community) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/marin-community/profile-training --- # Profile JAX training ## Scope Ingestion sources: - XPlane protobufs inside Levanter profile directories (source of truth): - `plugins/profile//*.xplane.pb` - explicit local `*.xplane.pb` files via `--xplane-file` - xprof aggregate tables exported from the same XPlane protobuf when the optional `xprof` package is available: step overview timing, kernel stats, collective breakdowns, xprof bottleneck statements. - Perfetto trace JSON as an explicit/fallback source for older profiles: - `plugins/profile//perfetto_trace.json.gz` - `plugins/profile//*.trace.json.gz` Prefer XPlane protobuf for new work. Perfetto trace JSON commonly hits the trace event cap; XPlane contains the uncapped timeline events needed for named-scope regions, pre-op gaps, gap context, process/thread metadata, and xprof aggregate tables. Use `--trace-file` only for a specific Perfetto JSON trace or an older profile with no XPlane protobuf. ## Capture Profiles Use Levanter profiler flags so profiles land under `//profiler`. Remote Marin runs also upload to `MARIN_PREFIX` TTL storage and print an XProf link: ```bash uv run ... \ --trainer.profiler.enabled true \ --trainer.profiler.start_step 5 \ --trainer.profiler.num_steps 10 \ --trainer.profiler.upload.ttl_days 30 ``` For profiles where xprof/HLO protobuf tables matter, enable JAX profile options through the Levanter profiler config: ```bash uv run ... \ --trainer.profiler.enabled true \ --trainer.profiler.start_step 5 \ --trainer.profiler.num_steps 5 \ --trainer.profiler.profile_options.host_tracer_level 1 \ --trainer.profiler.profile_options.python_tracer_level 0 \ --trainer.profiler.profile_options.device_tracer_level 0 \ --trainer.profiler.profile_options.enable_hlo_proto true ``` HLO metadata increases artifact size, so keep these profile windows short. The `XProf profile:` link appears after upload. Set `--trainer.profiler.upload.enabled false` for local-only capture. Do not copy profiles to another GCS region for inspection. Known-good TensorBoard scope recipe from CoreWeave Grug MoE profiling: `trainer.profiler.enabled=true`, `trainer.profiler.start_step=3`, `trainer.profiler.num_steps=2`, `trainer.profiler.perfetto_link=false`, `trainer.profiler.profile_options.host_tracer_level=1`, `trainer.profiler.profile_options.python_tracer_level=0`, and `trainer.profiler.profile_options.enable_hlo_proto=true` preserved useful `jax.named_scope` / `named_call` regions in TensorBoard for `GM2560-MAY-120S4096-W2048-B8-R1-E8M1-FA4PROFILE-S3B-N1-cw-20260617-2353`. Leave `device_tracer_level` unset unless device timelines are specifically needed; this profile retained useful hierarchical host/XLA metadata without it. On GPU, command buffers can collapse or suppress the visible name stack in TensorBoard/Perfetto. For profile-readability runs, disable command buffers: ```bash export XLA_FLAGS="${XLA_FLAGS:-} --xla_gpu_enable_command_buffer=''" ``` This hurts performance, so use it only when the goal is semantic trace attribution; leave it out of throughput comparisons unless command-buffer behavior is the axis being tested. For GPU throughput runs, keep profile-readability flags separate from XLA code generation and scheduling flags. Start from JAX's GPU performance guide, especially the code generation flags section: . The exact set of useful XLA flags is `jaxlib`-version dependent, so record the full `XLA_FLAGS` value with each profile or W&B run. For better profile readability, use `haliax.jax_utils.named_call` and `jax.named_scope` liberally in model code; these names flow into trace annotations and make region-level summaries far more actionable. Reference: - `lib/levanter/docs/Performance-Guide.md` - `.agents/skills/add-pallas-kernel/SKILL.md` - JAX GPU performance tips: ## Ingest to Structured Summary Use `/tmp` for ephemeral downloads. Use `scratch/` only when the working tree must retain an uncommitted analysis artifact. ```bash # /tmp (ephemeral) uv run python lib/marin/tools/profile_summary.py summarize \ --run-target marin-community/marin/ \ --download-root /tmp/marin-profiles \ --breakdown-mode exclusive_global \ --output /tmp/profile_summary.json ``` ### Option A: From a W&B artifact reference ```bash uv run python lib/marin/tools/profile_summary.py summarize \ --artifact marin-community/marin/run-grug-125m-profile-apples-pallas_tpu-20260217-225239-055ab2-profiler:v0 \ --download-root /tmp/marin-profiles \ --output /tmp/profile_summary.json ``` `--run-target` accepts: a bare run id (requires `--entity` and `--project`), `entity/project/run_id`, or a full W&B run URL. The profiler directory is resolved from `trainer.log_dir` in the run config. ### Option B: From a local artifact directory ```bash uv run python lib/marin/tools/profile_summary.py summarize \ --profile-dir /path/to/profiler_dir \ --output /tmp/profile_summary.json ``` If the directory contains `*.xplane.pb`, `--profile-dir` uses the XPlane path automatically. When both `*.xplane.pb` and Perfetto trace JSON are present, `--profile-dir` reads the XPlane protobuf by default (Perfetto exports are often capped). Use `--trace-file` to force a specific Perfetto JSON file. ### Option C: From a specific trace file ```bash uv run python lib/marin/tools/profile_summary.py summarize \ --trace-file /path/to/perfetto_trace.json.gz \ --output /tmp/profile_summary.json ``` ### Option D: From a specific XPlane protobuf Direct XPlane timeline parsing uses `protobuf` and does not require TensorFlow-generated `xplane_pb2` modules. If `xprof` is installed, ingestion also exports compact xprof table JSON and augments the timeline summary with aggregate step, kernel, collective, and bottleneck evidence. ```bash uv run --with xprof --with protobuf python lib/marin/tools/profile_summary.py summarize \ --xplane-file /path/to/profile.xplane.pb \ --xplane-output-dir /tmp/profile_xprof_tables \ --xplane-count-trace-events \ --output /tmp/profile_summary.json ``` Without `--xplane-output-dir` the command still parses XPlane timeline events directly. Add `--with xprof` for xprof aggregate table augmentation; add `--xplane-output-dir` to preserve the exported table JSON (this flag requires the optional `xprof` package). XPlane summaries expose hierarchical named-scope regions, pre-op gaps, gap region context, process/thread/timeline event metadata, step timing (when step markers or xprof overview rows exist), xprof bottleneck statements, kernel stats, collective breakdowns, and optimization candidates. Summary version tag: `profile_summary.v1` Generate a deterministic markdown root-cause report: ```bash uv run python lib/marin/tools/profile_summary.py report \ --summary /tmp/profile_summary.json \ --output /tmp/profile_report.md ``` Trace quality checks are surfaced in `trace_overview`: - `suspected_truncation`: `true` when event counts match a known export cap. - `quality_warnings`: warnings to treat hotspot/gap attribution with caution. ## Query the summary ```bash uv run python lib/marin/tools/profile_summary.py query \ --summary /tmp/profile_summary.json \ --question "" ``` Query top exclusive-time ops, compute/communication balance and collectives, specific pre-op gaps, hierarchical regions, noisy-op context, and suggested optimizations. Useful query forms include: - `What are the top 10 ops by exclusive time?` - `Is comm or compute dominating? Which collective is worst?` - `gap before _linear_softmax_cross_entropy_loss_bwd_pallas_mosaic_tpu_combined.1` - `show hierarchical regions` - `show context for op copy.564` - `What should we try next?` Pre-op gap attribution is marker-aware: - `gap_before_ops[].payload_op`: op where useful work starts after the idle period. - `gap_before_ops[].marker_op`: first op observed after the gap (often lightweight setup like `iota.*`). ## Optimization Workflow Use a strict workflow: 1. **Measure**: generate `before.json`. 2. **Change**: apply one bounded patch/config tweak. 3. **Re-measure**: generate `after.json`. 4. **Compare**: ```bash uv run python lib/marin/tools/profile_summary.py compare \ --before /tmp/profile_before.json \ --after /tmp/profile_after.json \ --strict-provenance ``` 5. **Track** (thresholded pass/warn/fail + history): ```bash uv run python lib/marin/tools/profile_summary.py track \ --before /tmp/profile_before.json \ --after /tmp/profile_after.json \ --label "pallas-kernel-attempt-3" \ --history /tmp/profile_regression_history.jsonl ``` 6. **History summary** (regression trend tracking): ```bash uv run python lib/marin/tools/profile_summary.py history \ --history /tmp/profile_regression_history.jsonl ``` 7. **One-shot compare bundle**: ```bash uv run python lib/marin/tools/profile_summary.py bundle \ --before-run-target marin-community/marin/ \ --after-run-target marin-community/marin/ \ --output-dir /tmp/profile_bundle \ --history /tmp/profile_regression_history.jsonl ``` 8. **Publish summary/report back to W&B**: ```bash uv run python lib/marin/tools/profile_summary.py publish \ --summary /tmp/profile_summary.json \ --report /tmp/profile_report.md \ --alias latest ``` The comparison reports: steady-state step-time delta, step class deltas (light/heavy when detected), compute/comm/host/stall share deltas, semantic family deltas with workload-normalized metrics, provenance checks (trace hash/run identity), and regressed/improved ops by exclusive duration.