Profile JAX training

Scope

Ingestion sources:

• XPlane protobufs inside Levanter profile directories (source of truth):
  • plugins/profile/<timestamp>/*.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/<timestamp>/perfetto_trace.json.gz
  • plugins/profile/<timestamp>/*.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 <trainer.log_dir>/<run_id>/profiler. Remote Marin runs also upload to MARIN_PREFIX TTL storage and print an XProf link:

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:

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:

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: https://docs.jax.dev/en/latest/gpu_performance_tips.html#code-generation-flags. 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: https://docs.jax.dev/en/latest/gpu_performance_tips.html

Ingest to Structured Summary

Use /tmp for ephemeral downloads. Use scratch/ only when the working tree must retain an uncommitted analysis artifact.

# /tmp (ephemeral)
uv run python lib/marin/tools/profile_summary.py summarize \
  --run-target marin-community/marin/<run_id> \
  --download-root /tmp/marin-profiles \
  --breakdown-mode exclusive_global \
  --output /tmp/profile_summary.json

Option A: From a W&B artifact reference

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

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

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.

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:

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

uv run python lib/marin/tools/profile_summary.py query \
  --summary /tmp/profile_summary.json \
  --question "<top ops, compute vs communication, gap, region, or op context>"

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:

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):

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):

uv run python lib/marin/tools/profile_summary.py history \
  --history /tmp/profile_regression_history.jsonl

7. One-shot compare bundle:

uv run python lib/marin/tools/profile_summary.py bundle \
  --before-run-target marin-community/marin/<baseline_run_id> \
  --after-run-target marin-community/marin/<candidate_run_id> \
  --output-dir /tmp/profile_bundle \
  --history /tmp/profile_regression_history.jsonl

8. Publish summary/report back to W&B:

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.