accelerator: bind CUDA before init_process_group; compile_dynamic -> compile_fused

The CUDA device is bound before init_process_group again, with main's exact
torch.cuda.set_device("cuda:<local_rank>") call, so the GPU path keeps main's
call order. TPU still binds after init: touching its one visible device earlier
breaks the rendezvous. The two phases are two named calls in _init_common.

compile_dynamic becomes compile_fused, whose docstring names its three callers:
each is elementwise/reduction math, where compiling only fuses kernels, so the
uncompiled function TPU gets computes the same values.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: HJSang

miles/backends/training_utils/loss/hub/math_utils.py

@@ -136,7 +136,7 @@ def hook(grad: torch.Tensor) -> torch.Tensor:
     tensor.register_hook(hook)
 
 
-@accelerator.compile_dynamic
+@accelerator.compile_fused
 def compute_approx_kl(
     log_probs: torch.Tensor,
     log_probs_base: torch.Tensor,
@@ -251,7 +251,7 @@ def compute_gspo_kl(
     return ppo_kl
 
 
-@accelerator.compile_dynamic
+@accelerator.compile_fused
 def compute_policy_loss(
     ppo_kl: torch.Tensor,
     advantages: torch.Tensor,
@@ -426,7 +426,7 @@ class _VocabParallelEntropy(torch.autograd.Function):
     @staticmethod
     def forward(ctx, vocab_parallel_logits: torch.Tensor, process_group: dist.ProcessGroup) -> torch.Tensor:
 
-        @accelerator.compile_dynamic
+        @accelerator.compile_fused
         def mul_reduce(a, b):
             return (a * b).sum(dim=-1, keepdim=True)
 

miles/ray/train_actor.py

@@ -153,6 +153,8 @@ def _init_common(self, args: Namespace, role: str, with_ref: bool = False, with_
 
         torch.serialization.add_safe_globals([miles.utils.eval_config.EvalDatasetConfig])
 
+        accelerator.bind_device_before_process_group(int(os.environ.get("LOCAL_RANK", 0)))
+
         if args.debug_deterministic_collective:
             register_det_nccl_backend()
             args.distributed_backend = DET_NCCL_BACKEND_NAME
@@ -169,7 +171,7 @@ def _init_common(self, args: Namespace, role: str, with_ref: bool = False, with_
             backend=backend,
             timeout=timedelta(minutes=args.distributed_timeout_minutes),
         )
-        accelerator.bind_local_device(int(os.environ.get("LOCAL_RANK", 0)))
+        accelerator.bind_device_after_process_group()
         init_gloo_group()
 
         args.rank = dist.get_rank()

miles/utils/accelerator.py

@@ -18,10 +18,16 @@
     import torch_tpu  # noqa: F401  registers the tpu device with torch
 
 
-def bind_local_device(local_rank: int) -> None:
-    """Bind this process to its device. Call after init_process_group: TPU binds its one visible device at
-    process-group init (touching it earlier breaks the rendezvous), CUDA binds lazily."""
-    torch.accelerator.set_device_index(0 if IS_TPU else local_rank)
+def bind_device_before_process_group(local_rank: int) -> None:
+    # CUDA binds before init_process_group, as it always has; a TPU touched this early breaks the rendezvous
+    if not IS_TPU:
+        torch.cuda.set_device(f"cuda:{local_rank}")
+
+
+def bind_device_after_process_group() -> None:
+    # TPU binds its one visible device once the process group is up
+    if IS_TPU:
+        torch.accelerator.set_device_index(0)
 
 
 def current_device() -> torch.device:
@@ -36,6 +42,8 @@ def device_name() -> str | None:
     return torch.cuda.get_device_name()
 
 
-def compile_dynamic(fn):
-    # the TPU compile backend rejects dynamic shapes; these helpers only fuse kernels
+def compile_fused(fn):
+    """`torch.compile(dynamic=True)` on CUDA, the function itself on TPU, whose compile backend rejects dynamic
+    shapes. Only for elementwise/reduction helpers, where compiling only fuses kernels and eager gives the same
+    values: `compute_approx_kl`, `compute_policy_loss` and the entropy's `mul_reduce`."""
     return fn if IS_TPU else torch.compile(dynamic=True)(fn)