Profile Training

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.