Reading a TPU profile

The on-demand route starts a server inside the training process and samples it from outside. xp.start_server(port) returns an object whose lifetime is the server's, and xp.trace(service_addr, logdir, duration_ms) blocks while it collects. Its defaults are worth reading once: host tracing at level 2, device tracing on, and a retry loop that keeps re-sending the request every interval_s until timeout_s runs out, because the thing being profiled may be busy compiling when the first request lands.

The in-process route skips the server. xp.start_trace(log_dir) opens a profiler session, xp.stop_trace() stops it and exports to that directory, and a second start_trace while one is running raises rather than quietly nesting. The implementation says it is based on the jax profiler, and the shape of the API is the same shape jax users already know.

Both routes produce the same kind of capture, and neither one knows anything about torch. What makes the capture legible is the third surface, the annotations you put in the loop, which is the rest of this lesson.

  options = {
      'host_tracer_level': host_tracer_level,
      'device_tracer_level': device_tracer_level,
      'delay_ms': delay_ms,
  }
  torch_xla._XLAC.profiler.trace(
      service_addr,
      logdir,
      duration_ms=duration_ms,
      num_tracing_attempts=num_tracing_attempts,
      timeout_s=timeout_s,
      interval_s=interval_s,
      options=options)

Wrap a region in xp.Trace('fwd') and two things happen at once. A host trace event opens, which is the part you expected. The context manager also pushes a lazy scope, and every IR node built while that scope is open records the name in its metadata.

That recorded scope is what makes a device timeline readable. When the lowering writes HLO metadata, a node whose metadata carries a scope gets that scope as its op name prefix, so an instruction that would have been called xla__add is called fwd/xla__add instead. The comment above that code says why the string manipulation is there: the xprof backend groups and nests traces by op name and type patterns, so the naming is the grouping.

One flag gates the whole thing. The metadata is only populated when XLA_HLO_DEBUG is set or the lazy IR debug flag is on, which means a profile captured without either shows your host annotations and shows unnamed device work. Nothing errors. The names simply do not arrive on the far side.

Trace also enters a jax named scope when jax is importable, so the same annotation carries through a torchax lowering rather than needing a second, jax-shaped version of itself.

    std::string op_name_prefix;
    size_t max_stack_depth = nmeta.frame_info.size();

    if (custom_opname_meta != nullptr) {
      op_name_prefix = custom_opname_meta->op_name_prefix;
      max_stack_depth = custom_opname_meta->max_stack_depth;
    }

    else if (!nmeta.scope.empty()) {
      op_name_prefix =
          absl::StrCat(absl::StrReplaceAll(nmeta.scope, {{":", "_"}}), "/");
    }
    metadata.set_op_name(absl::StrCat(op_name_prefix, op_type));

xp.StepTrace looks like xp.Trace with a step number attached, and its exit does something the plain version does not. It deletes the scope first, then calls torch_xla.sync(), then closes the host event. The comment explains the ordering: the step marker checks that no scope is still open, so the scope has to go before the barrier does.

That makes the step annotation a cut point as well as a label. Whatever the bridge arc's second lesson taught about lines that force materialization applies to this one, with the difference that here the cut is the intended behaviour rather than an accident you are hunting.

The profiling guide puts one restriction on the plain version, and it follows from the same ordering. A region wrapped in xp.Trace() must not contain a call to torch_xla.sync(). Reach for StepTrace when a barrier belongs inside the region, and for Trace when it does not.

  def __init__(self, name: str, **kwargs):
    super().__init__(name, _r=1, **kwargs)

  def __enter__(self):
    set_tracer_marked_step(True)
    super().__enter__()

  def __exit__(self, type, value, traceback):
    if getattr(self, 'scope', None):
      # In ir.cpp ResetScopeContext we ensure that we have no remaining scope
      # before marking step.
      del self.scope
    torch_xla.sync()
    super().__exit__(type, value, traceback)

The repo ships a worked reading of two real profiles, and it is the closest thing to an answer key this course can point at while it has no chip of its own. A Stable Diffusion 2.1 pipeline on a v4-8 profiles as busy with one small gap in the middle, and without annotations there is no way to find out what the host was doing during it. Adding xp.Trace calls to the pipeline and the U-Net turns that gap into a named Python region.

The XL version of the same model shows two shapes at once. There is one large gap caused by watermarking, which the walkthrough diagnoses by noticing that the gap is preceded by a TransferFromDevice, and which turns out to be tensors moving to CPU and becoming numpy arrays for two libraries that need them there. There are also many small gaps inside the denoising loop, traced to .item() and .nonzero() calls inside the scheduler's step, each one cutting the loop graph into smaller pieces.

If the gaps in the profile are due to Python code tracing that happens on the host, then this might be a bottleneck and there is no further straightforward optimization that can be done.

That quote is the guide's own conclusion, and it is the honest half of profiling a lazy program. Some gaps are your program asking for a value it did not need. Others are tracing itself, and no amount of rearranging the loop removes them.

A trace tells you where time sat. It does not tell you how many times the frontend crossed into the runtime, or which of the seam's named paths it took, and those are usually the first questions worth asking. The metrics report answers them by name, counting calls raised inside PjRtComputationClient itself.

LAB·P5 does that reading end to end against the pinned line numbers, and its reference run on a Colab TPU v6e-1 is the number to compare your own against: the same two-step loop reports one compile and two executes on the lazy path, and 23 compiles and 58 executes under the eager names. The lab hangs off chapter 10 with the rest of the runnable work, and this lesson does not repeat its table; the point of naming it here is the ordering.

Read the counters first, because they are exact and cost nothing. Capture a trace when the counts already look the way you expected and the wall clock still does not. A profile of a program that is compiling four times per step is a picture of the wrong problem.

None of the captures above ran here. The machine this course was written on has no TPU, xp.trace needs a profiler server on a machine that has one, and a device plane without a device is nothing. Everything in this lesson is read from the pinned source and the pinned docs.

The pending half is small and specific, which is how a pending marker should read. Start a server in a Colab TPU runtime, wrap one step in StepTrace with XLA_HLO_DEBUG=1 set, capture, and check two things: that your step name appears on device instructions and not only on host events, and that the gaps you predicted are the gaps you got. LAB·P5 already holds the counter half of that same run.