Feat/multi gpu fastsac

scripts/train_offpolicy.py

@@ -133,13 +133,79 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--use_amp", action="store_true", help="Enable mixed precision training (FP16)"
     )
+    parser.add_argument(
+        "--num_gpus",
+        type=int,
+        default=1,
+        help="Number of GPUs for multi-GPU training via NCCL all-reduce (default: 1 = single GPU)",
+    )
     return parser
 
 
 def build_runner(algo_name: str, args, cfg):
     """Build algorithm runner from unified config."""
     if algo_name == "sac":
+        from unilab.algos.torch.fast_sac.learner import FastSACLearner
         from unilab.algos.torch.fast_sac.runner import FastSACRunner
+        from unilab.utils.device_utils import get_default_device, get_env_dims
+
+        # Multi-GPU path
+        if getattr(args, "num_gpus", 1) > 1:
+            from unilab.algos.torch.offpolicy.multi_gpu_runner import MultiGPUOffPolicyRunner
+            from unilab.base import registry
+            from unilab.utils.algo_utils import ensure_registries
+
+            ensure_registries()
+            device = args.device or get_default_device()
+            env = registry.make(args.task, num_envs=1, sim_backend=args.sim_backend)
+            assert env.observation_space.shape and env.action_space.shape
+            obs_dim = env.observation_space.shape[0]
+            action_dim = env.action_space.shape[0]
+            env.close()
+
+            # Learner kwargs needed to reconstruct FastSACLearner in each worker
+            learner_kwargs = {
+                "obs_dim": obs_dim,
+                "action_dim": action_dim,
+                "gamma": cfg.gamma,
+                "tau": cfg.tau,
+                "actor_lr": cfg.actor_lr,
+                "critic_lr": cfg.critic_lr,
+                "alpha_lr": cfg.algo_params.alpha_lr,
+                "alpha_init": cfg.algo_params.alpha_init,
+                "target_entropy_ratio": cfg.algo_params.target_entropy_ratio,
+                "actor_hidden_dim": cfg.actor_hidden_dim,
+                "critic_hidden_dim": cfg.critic_hidden_dim,
+                "num_atoms": cfg.num_atoms,
+                "use_layer_norm": cfg.use_layer_norm,
+                "max_grad_norm": cfg.algo_params.max_grad_norm,
+                "use_amp": args.use_amp,
+                # symmetry not supported in multi-GPU mode (mujoco_model not picklable)
+                "use_symmetry": False,
+            }
+            # Main-process learner (used for WeightSync initialisation and single-GPU fallback)
+            main_learner = FastSACLearner(device=device, **learner_kwargs)
+
+            return MultiGPUOffPolicyRunner(
+                learner=main_learner,
+                env_name=args.task,
+                algo_type="sac",
+                learner_kwargs=learner_kwargs,
+                num_gpus=args.num_gpus,
+                num_envs=cfg.num_envs,
+                replay_buffer_n=cfg.replay_buffer_n,
+                batch_size=cfg.batch_size,
+                warmup_steps=cfg.warmup_steps,
+                updates_per_step=cfg.updates_per_step,
+                policy_frequency=cfg.policy_frequency,
+                sync_collection=not args.no_sync_collection,
+                env_steps_per_sync=cfg.env_steps_per_sync,
+                device=device,
+                actor_hidden_dim=cfg.actor_hidden_dim,
+                use_layer_norm=cfg.use_layer_norm,
+                obs_normalization=False,
+                sim_backend=args.sim_backend,
+            )
 
         return FastSACRunner(
             env_name=args.task,

src/unilab/algos/torch/common/shared_buffer.py

@@ -6,14 +6,16 @@
 class SharedBufferBase:
     """Device-adaptive shared memory buffer base class."""
 
-    def __init__(self, capacity: int, device: str):
+    def __init__(self, capacity: int, device: str, defer_gpu: bool = False):
         self.capacity = capacity
         self.device = device
         self.ptr = torch.zeros(1, dtype=torch.int64).share_memory_()
 
         if device == "cuda":
             self._gpu_synced_ptr = 0
-            self._cuda_stream = torch.cuda.Stream()
+            # defer_gpu=True: skip CUDA stream creation so the buffer is picklable
+            # (used in multi-GPU mode where workers call init_local_gpu_cache later)
+            self._cuda_stream = None if defer_gpu else torch.cuda.Stream()  # type: ignore[assignment]
         else:
             self._gpu_synced_ptr = None  # type: ignore[assignment]
             self._cuda_stream = None  # type: ignore[assignment]

src/unilab/algos/torch/fast_sac/learner.py

@@ -14,6 +14,7 @@
 from typing import Any, Dict, Tuple
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
@@ -359,13 +360,15 @@ def __init__(
         use_amp: bool = False,
         mujoco_model=None,
         obs_structure: dict | None = None,
+        world_size: int = 1,
     ):
         self.device = device
         self.gamma = gamma
         self.tau = tau
         self.max_grad_norm = max_grad_norm
         self.use_autotune = use_autotune
         self.use_amp = use_amp and device == "cuda"
+        self.world_size = world_size
 
         # Build actor
         self.actor = SACActor(
@@ -455,6 +458,27 @@ def __init__(
             assert obs_structure is not None
             self.symmetry = G1SymmetryAugmentation(mujoco_model, obs_structure, device=device)
 
+    def _reduce_gradients(self, model: nn.Module) -> None:
+        """All-reduce gradients across all workers and divide by world_size.
+
+        Must be called after ``backward()`` and, when using AMP, after
+        ``scaler.unscale_(optimizer)`` so that gradients are in full precision.
+        """
+        if self.world_size <= 1:
+            return
+        grads = [p.grad.view(-1) for p in model.parameters() if p.grad is not None]
+        if not grads:
+            return
+        flat = torch.cat(grads)
+        dist.all_reduce(flat, op=dist.ReduceOp.SUM)
+        flat /= self.world_size
+        offset = 0
+        for p in model.parameters():
+            if p.grad is not None:
+                n = p.grad.numel()
+                p.grad.copy_(flat[offset : offset + n].view_as(p.grad))
+                offset += n
+
     def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
         """One critic update step."""
         obs = batch["obs"]
@@ -504,8 +528,9 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
             self.q_optimizer.zero_grad(set_to_none=True)
             if self.scaler:
                 self.scaler.scale(qf_loss).backward()
+                self.scaler.unscale_(self.q_optimizer)
+                self._reduce_gradients(self.qnet)
                 if self.max_grad_norm > 0:
-                    self.scaler.unscale_(self.q_optimizer)
                     critic_grad_norm = torch.nn.utils.clip_grad_norm_(
                         self.qnet.parameters(), max_norm=self.max_grad_norm
                     )
@@ -515,6 +540,7 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
                 self.scaler.update()
             else:
                 qf_loss.backward()
+                self._reduce_gradients(self.qnet)
                 if self.max_grad_norm > 0:
                     critic_grad_norm = torch.nn.utils.clip_grad_norm_(
                         self.qnet.parameters(), max_norm=self.max_grad_norm
@@ -536,6 +562,9 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
             ).mean()
             if torch.isfinite(alpha_loss):
                 alpha_loss.backward()
+                if self.world_size > 1 and self.log_alpha.grad is not None:
+                    dist.all_reduce(self.log_alpha.grad, op=dist.ReduceOp.SUM)
+                    self.log_alpha.grad /= self.world_size
                 self.alpha_optimizer.step()
 
         return {
@@ -574,8 +603,9 @@ def update_actor(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
             self.actor_optimizer.zero_grad(set_to_none=True)
             if self.scaler:
                 self.scaler.scale(actor_loss).backward()
+                self.scaler.unscale_(self.actor_optimizer)
+                self._reduce_gradients(self.actor)
                 if self.max_grad_norm > 0:
-                    self.scaler.unscale_(self.actor_optimizer)
                     actor_grad_norm = torch.nn.utils.clip_grad_norm_(
                         self.actor.parameters(), max_norm=self.max_grad_norm
                     )
@@ -585,6 +615,7 @@ def update_actor(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
                 self.scaler.update()
             else:
                 actor_loss.backward()
+                self._reduce_gradients(self.actor)
                 if self.max_grad_norm > 0:
                     actor_grad_norm = torch.nn.utils.clip_grad_norm_(
                         self.actor.parameters(), max_norm=self.max_grad_norm

src/unilab/algos/torch/fast_sac/runner.py

@@ -38,6 +38,7 @@ def __init__(
         use_amp: bool = False,
         sim_backend: str = "mujoco",
         use_symmetry: bool = False,
+        world_size: int = 1,
     ):
         from unilab.base import registry
         from unilab.utils.algo_utils import ensure_registries
@@ -77,6 +78,7 @@ def __init__(
             use_symmetry=use_symmetry,
             mujoco_model=mujoco_model,
             obs_structure=obs_structure,
+            world_size=getattr(self, "world_size", world_size),
         )
 
         # Auto-adjust batch_size when symmetry is enabled

src/unilab/algos/torch/offpolicy/__init__.py

@@ -1,6 +1,7 @@
 """Off-policy RL unified infrastructure."""
 
+from unilab.algos.torch.offpolicy.multi_gpu_runner import MultiGPUOffPolicyRunner
 from unilab.algos.torch.offpolicy.runner import OffPolicyRunner
 from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn
 
-__all__ = ["OffPolicyRunner", "off_policy_collector_fn"]
+__all__ = ["OffPolicyRunner", "MultiGPUOffPolicyRunner", "off_policy_collector_fn"]