[TPU] Correctly profile peak memory usage & Upgrade PyTorch XLA

[WoosukKwon]

Should be merged after #9437 and after the 10/17 version of PyTorch XLA nightly is available.

This PR upgrades the PyTorch XLA, and uses the peak_bytes_used to correctly profile the peak HBM usage during the dummy profile run.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[robertgshaw2-redhat]

Hey @WoosukKwon - I have been using this PR for the compressed-tensors loading as well as #9437.

I noticed that even with this PR and #9437, value of m["peak_bytes_used"] >> m["bytes_used"] (it can be 2-4GB more) before we do the dummy run. This is often bigger than the activation size, which means we are wasting space for KV cache allocation. This extra memory occurs when we actually load the weights:

  def load_model(self, *, model_config: ModelConfig,
                 device_config: DeviceConfig,
                 lora_config: Optional[LoRAConfig],
                 parallel_config: ParallelConfig,
                 scheduler_config: SchedulerConfig,
                 cache_config: CacheConfig) -> nn.Module:
      target_device = torch.device(device_config.device)
      with set_default_torch_dtype(model_config.dtype):
          with target_device:
              model = _initialize_model(model_config, self.load_config,
                                        lora_config, cache_config,
                                        scheduler_config) # << here, value of peak_bytes_used == bytes_used == weight_size

          model.load_weights(self._get_all_weights(model_config, model)) # here, peak_bytes_used > bytes_used

So, I was thinking we might want to reset peak_bytes_used prior to running the dummy pass