Fix Metric Reporting Crash (ZeroDivisionError) using Tunix PerfTracer

[igorts-git]

Description

This PR fixes a ZeroDivisionError that occurred in the MaxText MetricLogger during Tunix-based SFT runs.

The Problem

As the result of Tunix PR1289. The Tunix library's internal training loop now hard-codes the step time to zero when calling the on_train_step_end hook.
The MaxText SFT hook passes the zero value to the metric_logger.
The MetricLogger divides FLOPS and tokens values by the passed zero step_time.

The Tunix change was made because the old reported step_time metric was inaccurate due to async execution. The recommended solution was to use Tunix PerfTracer, which correctly tracks the async step time.

The Fix

1. Initialize a PerfTracer in setup_trainer_state and passes it to the PeftTrainer.
2. Add sft_perf_export_fn to extract the peft_train_step span duration from the Tunix timeline.
3. The training hooks now pull the step_time from the tracer.

FIXES: b/497665859.

Tests

Manual invocation of SFT training with and without the fix with Tunix installed from main/HEAD.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/post_train/sft/hooks.py	0.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!