[askhat-g]

Summary

Part 1 of 2 (Platform & Runtime Patches):

1. TPU Distributed & Runtime Patches (verl_hardware_plugin/accelerators/tpu/patches/reduce_avg_allreduce_patch.py):
   • Rewrites torch.distributed.all_reduce(op=ReduceOp.AVG) into ReduceOp.SUM + tensor.div_(world_size) on tpu_dist process groups (matching the XPU reduce_avg_allreduce_patch pattern).
   • Routes torch.distributed.all_gather_object(..., group=dp_group) on tpu_dist device-mesh subgroups through either the default Gloo group (when group_size == world_size) or fixed-shape int32 byte tensors via dist.all_gather (avoiding tensor.resize_() crashes on PJRT/XLA), and detaches any TPU metric tensors to CPU.
   • Replaces verl.utils.distributed.set_numa_affinity with a no-op on TPU and routes verl.utils.torch_functional.logprobs_from_logits to vectorized logprobs_from_logits_naive for TPU tensors.
2. torchrun Multi-Controller PJRT Environment Setup (PlatformTPU):
   • Adds configure_torchrun_tpu_env() in platform_tpu.py (invoked from PlatformTPU.__init__) so standalone torchrun -m verl.trainer.sft_trainer automatically populates PJRT_DEVICE, TORCH_TPU_SLICEBUILDER_ADDRESSES, TPU_PROCESS_PORT, TPU_VISIBLE_CHIPS, and topology bounds from LOCAL_RANK / RANK / WORLD_SIZE.

Test Plan

• pytest tests/test_plugin_registration.py -v