skills.rest
marin-community
marin-communityOFFICIAL
Explore all skills by this author →
WebsiteGitHub

profile-training

Profile JAX, Levanter, and Marin training runs to diagnose startup, compilation, and throughput bottlenecks.

3.2k|266|Updated Mar 22, 2024
Search Tags:#performance analysis#WandB#xplane#JAX#Profiling#training bottleneck#levanter
One-click install
npx skills add https://github.com/marin-community/marin --skill profile-training
Or copy as Structured Prompt for Agent▼
Please help me install this Agent Skill.
Skill: profile-training
Source: https://github.com/marin-community/marin/tree/main/.agents/skills/profile-training
Command: npx skills add https://github.com/marin-community/marin --skill profile-training

SYSTEM DOCUMENTATION & REQUIREMENTS

💡 This Skill requires protobuf, xprof.

What problem does it solve?

Training large JAX models often suffers from slow startup, long compilation, or poor throughput, and raw profiler output is hard to interpret. This Skill ingests XPlane protobufs and Perfetto traces from Levanter profiler directories and turns them into structured summaries, root-cause reports, and before/after comparisons.

Core Features & Use Cases

  • Profile Capture Guidance: Configure Levanter profiler flags, JAX profile options, and XLA flags so traces retain named-scope regions and HLO metadata.
  • Structured Summaries: Parse XPlane protobufs (with optional xprof aggregate tables) into JSON summaries exposing top ops, pre-op gaps, collective breakdowns, and bottleneck statements.
  • Optimization Workflow: Measure, patch, re-measure, then compare runs with provenance checks, regression tracking history, and publish reports back to W&B.
  • Use Case: A run shows a large idle gap before the loss backward kernel; query the summary for gap attribution, apply a bounded fix, and use the compare command to verify the step-time improvement.

Quick Start

Ask the assistant to summarize the profiler output for a given W&B run id and report the top ops by exclusive time and any pre-op gaps.

Frequently Asked Questions about profile-training

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I profile a JAX training run with Levanter?▼

Enable the Levanter profiler with trainer.profiler.enabled=true plus start_step and num_steps flags so profiles land under the trainer log directory. Remote Marin runs also upload to TTL storage and print an XProf link.

How do I analyze an XPlane protobuf profile from a training run?▼

Run the profile_summary.py summarize command with --xplane-file pointing at the .xplane.pb file, optionally adding --with xprof for aggregate tables. It produces a JSON summary with top ops, gaps, collectives, and bottleneck statements.

Should I use XPlane protobuf or Perfetto trace JSON for profiling?▼

Prefer XPlane protobuf because Perfetto trace JSON commonly hits the trace event cap and loses timeline events. Use --trace-file only for older profiles that have no XPlane protobuf available.

Why do GPU profiles lose named-scope regions in TensorBoard?▼

CUDA command buffers can collapse the visible name stack on GPU. Disable them with --xla_gpu_enable_command_buffer='' in XLA_FLAGS for profile-readability runs, accepting the performance cost.

How do I compare two training profiles to detect regressions?▼

Generate before and after summaries, then run the compare command with --strict-provenance to get step-time deltas, compute/comm share changes, and regressed ops. Use track and history commands for thresholded regression trends.

Related Skills

Explore More Software Engineering Skills→
sqjiansqjian

generate-profile

Launch an SGLang server, validate accuracy, and generate a Chrome/Perfetto trace.

Community
Advanced
VCERSVCERS
7

nsys-optimizer

Identify CUDA performance bottlenecks from Nsight Systems profiling data.

Official
Advanced
igamenovoerigamenovoer

krnopt-cuda-profiling

Identify CUDA bottlenecks from NSYS traces and NCU reports and attribute them to source regions.

Community
Advanced
onejsonejs
4.5k

ios-simulator-skill

Automate iOS simulator testing, building, and UI interaction using accessibility-driven scripts.

Official
Advanced