Disable JAX GPU preallocation by default

[YonghaoZhao722]

Summary

Disable JAX GPU memory preallocation by default for MGW.

Motivation

MGW uses both PyTorch and JAX/OTT. By default, JAX preallocates a large fraction of GPU memory on first use, even when the actual working set is much smaller. This can make MGW appear to consume excessive GPU memory and may cause unnecessary conflicts with PyTorch, other processes, or shared-GPU environments.

Disabling preallocation makes GPU memory usage more proportional to the actual workload while preserving the existing MGW algorithm and benchmark behavior.

Changes

• Set XLA_PYTHON_CLIENT_PREALLOCATE=false before JAX/OTT initialization.
• Use os.environ.setdefault(...) so user-provided allocator settings are not overwritten.
• Keep the change limited to runtime memory allocation behavior; no algorithmic logic is changed.
• Allow users to override this behavior through their own environment variables if needed.

Test plan

• Verified that JAX/OTT operations still run correctly without preallocation.
• Verified MGW benchmark scripts run successfully with the new default.
• Tested in limited/shared GPU memory settings to confirm reduced unnecessary GPU reservation.