fix: support run_on_local_node in gemm_throttling

[simrankaurb]

Summary

Support configurable mesh locality (run_on_local_node) in gemm_throttling to allow running GEMM thermal soak microbenchmarks on single host devices across multi-node clusters.

Rationale & Context

Following the design in #125, run_on_local_node allows running microbenchmarks per single host (jax.local_devices()) rather than over the entire global cluster slice (jax.devices()).

Changes Included

1. gemm_throttling(): Added run_on_local_node: bool = False parameter and passed local_mesh=run_on_local_node to create_mesh().
2. gemm_throttling_calculate_metrics(): Added run_on_local_node: bool = False parameter and set device_count = jax.local_device_count() if run_on_local_node else jax.device_count() when calling handle_based_on_sharding(). This prevents inflated total FLOPs and throughput metrics in multi-node environments.

[gemini-code-assist]

Code Review

This pull request introduces a run_on_local_node parameter to the gemm_throttling benchmark function to allow running on local node devices. However, the reviewer pointed out a critical issue: introducing this parameter without updating gemm_throttling_calculate_metrics will cause incorrect metrics calculations in multi-node environments, as the global device count will be used instead of the local device count. The reviewer provided a code suggestion to fix this by passing the correct device count based on run_on_local_node.