Fix ZeroDivisionError in SFT during record_train_metrics

[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 Fix

This PR modifies SFTTrainingHooks in src/maxtext/trainers/post_train/sft/hooks.py to:

1. Independently track step time: Records time.time() at on_train_step_start.
2. Bypass the provided value: Calculates the actual_step_time at on_train_step_end and passes it to the logger, ensuring a valid, non-zero duration is always reported.

Alternatives

Alternatively, we can ask Tunix to restore the code that reports step time.

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 20.00000% with 4 lines in your changes missing coverage. Please review.

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

📢 Thoughts on this report? Let us know!