[Distributed] Fix FrozenLake distributed launcher and multi-core TPU mesh handling

[andytwigg]

Summary of Changes

This PR resolves several issues when running distributed FrozenLake (frozenlake_dist) on TPUs, specifically supporting multi-core-per-chip hardware architectures (such as TPU 7x) as well as small batch/smoke test configurations:

1. Micro-batch size default in frozenlake_dist/launcher.sh:

   • Defaults TRAIN_MICRO_BATCH_SIZE to 1 (aligning with math_gsm8k_dist/launcher.sh and CI smoke recipes). This prevents validation failure (MINI_BATCH_SIZE * NUM_GENERATIONS must be divisible by TRAIN_MICRO_BATCH_SIZE) for lean test runs like BATCH_SIZE=1 MINI_BATCH_SIZE=1 NUM_GENERATIONS=2.

2. Libtpu co-location support:

   • Exports ALLOW_MULTIPLE_LIBTPU_LOAD=1 in launcher.sh so co-located trainer and rollout processes on the same host can initialize libtpu without collision.

3. Dynamic TPU_CHIPS_PER_HOST_BOUNDS derivation:

   • Auto-derives host bounds from TRAINER_TPU_CHIPS count when unset (e.g. 1,2,1 for 2 chips instead of hardcoding 1,4,1).

4. Multi-core TPU chip mesh compatibility:

   • In run_trainer_node.py and run_rollout_node.py, auto-adjust mesh_tp and tensor_parallel_size when mesh_fsdp == 1 and jax.device_count() is a multiple of args.mesh_tp (e.g. on TPU architectures with 2 cores per chip where TP is specified in physical chips).

5. Tooling & Dependency isolation in scripts/install_raiden.sh:

   • Uses uv pip install grpcio-tools if uv is available during protobuf compilation to avoid unintended dependency resolution and setuptools version conflicts with vllm.

6. Repository hygiene (.gitignore):

   • Ignores .venv/, myenv/, logs (*.log), artifacts, checkpoints, and generated protobuf stubs (*_pb2.py, *_pb2_grpc.py).

Verification

• Ran full distributed FrozenLake with Qwen3-1.7B on TPU:

  MODEL_NAME=Qwen3-1.7B MODEL_ID=Qwen/Qwen3-1.7B \
  TRAINER_TPU_CHIPS=0,1 TRAINER_TP=2 \
  ROLLOUT_TPU_CHIPS=2,3 ROLLOUT_TP=2 \
  BATCH_SIZE=1 MINI_BATCH_SIZE=1 NUM_GENERATIONS=2 \
  MAX_STEPS=1 MAX_TURNS=3 WEIGHT_SYNC_MODE=none ./launcher.sh

  Verified mesh creation, model loading, rollout trajectory generation, token logp evaluation, gradient accumulation, optimizer step, and clean teardown.
• Verified test suite passes: pytest tests/experimental/examples/frozenlake_dist/ tests/experimental/examples/common/run_trainer_node_test.py (58 passed).