fix(peft_trainer): make scalar-metric host-fetch multi-process safe

[lokic233]

What

Makes the scalar-metric host-fetch in PeftTrainer safe under a multi-process / multi-host mesh.

Why

Loss, perplexity and learning_rate are replicated scalars. Under a multi-process / multi-host mesh the underlying jax.Array spans non-process-local (non-addressable) devices, so jax.device_get / np.asarray on it raises:

RuntimeError: Fetching value for `jax.Array` that spans non-addressable
(non process local) devices is not possible.

This crashes the train→metrics logging path on any multi-host / bare-SPMD run — the training step itself succeeds, but the metrics flush at the end of the step throws.

What changes

Adds a small _addressable_scalar helper (following the existing is_fully_addressable idiom already used in _shard) that, for a replicated scalar jax.Array, returns this process's local addressable shard before the host fetch.

It only reduces when the local shard shape equals the global shape, so a genuinely sharded, non-scalar array is never silently truncated to a partial shard. Fully-addressable arrays and plain Python/NumPy scalars pass through unchanged — single-process behavior is identical.

Applied at the four scalar host-fetch sites:

• MetricsBuffer.loss (buffered-loss mean)
• _log_metrics: jax.device_get(loss) and jax.device_get(learning_rate)
• _write_metrics._to_np_array (additional scalar metrics: reward, kl, etc.)

Test

Validated on a 16-process bare-SPMD RL run on TPU v7x: the metrics flush that previously crashed at jax.device_get(learning_rate) now completes, and training proceeds across many steps. peft_trainer.py py_compile clean and 80-col compliant.