diff --git a/.gitignore b/.gitignore index 558ef0b4a..448caf14e 100644 --- a/.gitignore +++ b/.gitignore @@ -64,6 +64,11 @@ src/unilab/algos/torch/rsl_rl third-party temp/ agent/ +memory/ +tmp/ +issues/ +issue2/ +docs/superpowers/ # Motion assets (downloaded from HF at runtime) src/unilab/assets/motions/g1/*.npz diff --git a/conf/offpolicy/algo/flashsac.yaml b/conf/offpolicy/algo/flashsac.yaml index 40034306d..e4e38f94e 100644 --- a/conf/offpolicy/algo/flashsac.yaml +++ b/conf/offpolicy/algo/flashsac.yaml @@ -42,3 +42,7 @@ algo_params: n_step: 1 amp_dtype: auto use_compile: true + use_cuda_graph_critic: false + use_cuda_graph_actor: false + use_cuda_graph_critic_packed_staging: false + use_cuda_graph_actor_packed_staging: false diff --git a/conf/offpolicy/algo/sac.yaml b/conf/offpolicy/algo/sac.yaml index 18a8d65db..ced2f104a 100644 --- a/conf/offpolicy/algo/sac.yaml +++ b/conf/offpolicy/algo/sac.yaml @@ -35,3 +35,7 @@ algo_params: max_grad_norm: 0.0 amp_dtype: auto use_compile: true + use_cuda_graph_critic: false + use_cuda_graph_actor: false + use_cuda_graph_critic_packed_staging: false + use_cuda_graph_actor_packed_staging: false diff --git a/conf/offpolicy/config.yaml b/conf/offpolicy/config.yaml index 74984f226..5c300048d 100644 --- a/conf/offpolicy/config.yaml +++ b/conf/offpolicy/config.yaml @@ -43,6 +43,7 @@ training: trace_output_dir: null trace_thread_time: false trace_cuda_events: true + nvtx_profile_ranges: false replay_prefetch_mode: one_tick verbose_metrics: false diff --git a/docs/superpowers/specs/2026-07-02-issue662-pr-markdown-design.md b/docs/superpowers/specs/2026-07-02-issue662-pr-markdown-design.md new file mode 100644 index 000000000..da4c97922 --- /dev/null +++ b/docs/superpowers/specs/2026-07-02-issue662-pr-markdown-design.md @@ -0,0 +1,94 @@ +# Issue662 PR Markdown Design + +## Goal + +Create a Chinese GitHub PR-description markdown document at: + +```text +issues/pr/issue662_fast_sac_gpu_utilization.md +``` + +The document should explain how the current `issue662` work addresses the Issue 1 FastSAC GPU-utilization problem. It should read like an engineering PR, not like an LLM conversation or a chat transcript. + +## Source Material + +The PR document must be grounded in repository artifacts only: + +- `issues/issue1.md` +- `issues/why_issue1/README.md` +- recursively referenced markdown from `issues/why_issue1/README.md` +- `issues/6_main_vs_issue662_ab_profile/reports/report.md` +- `issues/8_issue_node_ablation/reports/report.md` +- `memory/tools.md` +- `memory/hardware.md` + +SVG figures must be read as source and rendered to PNG for visual inspection before being used. Report figures and selected root-cause figures should be copied into `issues/pr/figures/` so the PR document is self-contained. + +## Recommended Narrative + +Use a problem-attribution-optimization-validation narrative: + +1. State the issue correspondence: FastSAC on `sac/g1_walk_flat/mujoco` exposes low GPU duty cycle caused by learner tiny-kernel / CUDA launch fragmentation. +2. Establish the environment boundary from local hardware/tooling records. +3. Summarize root-cause evidence from `why_issue1`: trace excludes steady replay/H2D starvation, Nsight Systems shows launch API time dominating kernel time, NCU shows underfilled kernels, `torch.compile` helps but is incomplete, and NVTX ranges localize tiny kernels to learner subranges. +4. Present `main` vs `issue662` evidence from report 6 across `g1_walk_flat_mujoco` and `g1_walk_rough_mujoco`. +5. Present issue-node ablation evidence from report 8 to show that the major improvement comes from CUDA graph and graph-boundary launch reduction, while earlier local cleanup stages are not the main E2E driver. +6. Close with limits and next steps: reward observations are smoke-level evidence, stream-sync increases must be acknowledged, and the next work is to evaluate whether other algorithms show similar launch-bound fragmentation. + +## Document Structure + +The PR markdown should contain these sections: + +- `## 背景与问题对应` +- `## 当前实验环境` +- `## 根因证据` +- `## 优化效果：main vs issue662` +- `## 优化来源：issue node ablation` +- `## 结论与边界` +- `## 后续工作` +- `## Validation` + +## Figure Policy + +Copy all figures from the two final reports into `issues/pr/figures/`: + +- 9 figures from `issues/6_main_vs_issue662_ab_profile/figures/` +- 4 figures from `issues/8_issue_node_ablation/figures/` + +Also copy selected root-cause figures from `issues/why_issue1/` when they are necessary to make the issue correspondence self-contained. Each copied image must be referenced in the PR body. + +Every figure reference must include a paragraph-length caption in Chinese that explains: + +- what experiment or measurement the figure represents +- what each important numeric axis or metric means +- what conclusion can and cannot be drawn from the figure + +Captions should be academic-paper-like but concise. They should not exaggerate beyond the measured artifact. + +## Claim Boundaries + +The PR must not claim: + +- that replay/H2D was the main bottleneck for this issue +- that `torch.compile` alone solved the problem +- that reward curves prove full long-run learning equivalence +- that `make test-all` passed unless there is explicit evidence +- that improvements generalize to all algorithms or all hardware + +The PR should explicitly state: + +- the measured environment is RTX 5090 D / 170 SM, EPYC 7K62, Ubuntu 22.04.5, PyTorch 2.7.0+cu128, CUDA 12.8 runtime, Nsight Systems 2024.6.2, Nsight Compute 2025.1.1 +- report 6 used 3 repeats for each task/variant +- report 8 used a stage-level ablation with repeat count 1 +- Nsight metrics and E2E wall time come from separate measurement contexts where appropriate +- stream synchronization time increased in `issue662`, so it should be treated as a graph-boundary/runtime-shape tradeoff to keep watching + +## Validation For This Documentation Task + +Before completion: + +- Verify the PR file exists. +- Verify every image reference in the PR points to an existing file. +- Verify all report 6 and report 8 figures are copied into `issues/pr/figures/`. +- Verify the PR includes environment information from `memory/tools.md` and `memory/hardware.md`. +- Verify the PR states the next step: evaluate other algorithms for similar launch-bound fragmentation. diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index daf7d98de..8f5f456ce 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -263,6 +263,27 @@ def build_runner(algo_name: str, cfg: DictConfig): "amp_dtype": cfg.algo.algo_params.amp_dtype, "use_compile": cfg.algo.algo_params.use_compile, "obs_normalization": cfg.algo.obs_normalization, + "use_cuda_graph_critic": bool( + getattr(cfg.algo.algo_params, "use_cuda_graph_critic", False) + ), + "use_cuda_graph_actor": bool( + getattr(cfg.algo.algo_params, "use_cuda_graph_actor", False) + ), + "use_cuda_graph_critic_packed_staging": bool( + getattr( + cfg.algo.algo_params, + "use_cuda_graph_critic_packed_staging", + False, + ) + ), + "use_cuda_graph_actor_packed_staging": bool( + getattr( + cfg.algo.algo_params, + "use_cuda_graph_actor_packed_staging", + False, + ) + ), + "nvtx_profile_ranges": bool(getattr(cfg.training, "nvtx_profile_ranges", False)), "use_symmetry": cfg.algo.use_symmetry, "symmetry_augmentation": _symmetry_aug, "critic_obs_dim": _critic_dim, diff --git a/src/unilab/algos/torch/fast_sac/learner.py b/src/unilab/algos/torch/fast_sac/learner.py index 9841eab07..30460ed48 100644 --- a/src/unilab/algos/torch/fast_sac/learner.py +++ b/src/unilab/algos/torch/fast_sac/learner.py @@ -11,6 +11,8 @@ from __future__ import annotations import math +from collections.abc import Iterator +from contextlib import contextmanager from typing import Any, Dict, Tuple, cast import torch @@ -35,6 +37,19 @@ def normalize_fast_sac_distributed_sync_mode(mode: str) -> str: return normalized +@contextmanager +def _cuda_nvtx_range(name: str, enabled: bool) -> Iterator[None]: + if not enabled: + yield + return + + torch.cuda.nvtx.range_push(name) + try: + yield + finally: + torch.cuda.nvtx.range_pop() + + # --------------------------------------------------------------------------- # Actor Network (holosoma-style: SiLU + LayerNorm + Tanh squashing) # --------------------------------------------------------------------------- @@ -140,26 +155,38 @@ def forward(self, obs: torch.Tensor) -> torch.Tensor: return _Wrapper() def get_actions_and_log_probs( - self, obs: torch.Tensor + self, + obs: torch.Tensor, + eps: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Sample actions and compute log probabilities. Returns (action, log_prob, log_std).""" _, mean, log_std = self(obs) + action, log_prob = self._sample_action_and_log_prob(mean, log_std, eps=eps) + return action, log_prob, log_std + + def _sample_action_and_log_prob( + self, + mean: torch.Tensor, + log_std: torch.Tensor, + eps: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: std = log_std.exp() - dist = torch.distributions.Normal(mean, std) - raw_action = dist.rsample() + if eps is None: + eps = torch.randn_like(mean) + raw_action = mean + std * eps + log_prob = -0.5 * ( + ((raw_action - mean) / std).pow(2) + 2.0 * log_std + math.log(2.0 * math.pi) + ) if self.use_tanh: tanh_action = torch.tanh(raw_action) action = tanh_action * self.action_scale + self.action_bias - log_prob = dist.log_prob(raw_action) log_prob -= torch.log(1 - tanh_action.pow(2) + 1e-6) log_prob -= torch.log(self.action_scale + 1e-6) else: action = raw_action - log_prob = dist.log_prob(raw_action) - log_prob = log_prob.sum(1) - return action, log_prob, log_std + return action, log_prob.sum(1) @torch.no_grad() def explore( @@ -187,9 +214,7 @@ def explore( return torch.tanh(mean) * self.action_scale + self.action_bias return mean - std = log_std.exp() - dist = torch.distributions.Normal(mean, std) - raw_action = dist.rsample() + raw_action = mean + log_std.exp() * torch.randn_like(mean) if self.use_tanh: return torch.tanh(raw_action) * self.action_scale + self.action_bias @@ -414,6 +439,11 @@ def __init__( amp_dtype: str = "auto", use_compile: bool = False, obs_normalization: bool = False, + use_cuda_graph_critic: bool = False, + use_cuda_graph_actor: bool = False, + use_cuda_graph_critic_packed_staging: bool = False, + use_cuda_graph_actor_packed_staging: bool = False, + nvtx_profile_ranges: bool = False, symmetry_augmentation: SymmetryAugmentation | None = None, world_size: int = 1, distributed_sync_mode: str = "sync_sgd", @@ -428,6 +458,15 @@ def __init__( self.use_compile = ( bool(use_compile) and get_torch_compile_for_cuda(self.device, warn=True) is not None ) + self.use_cuda_graph_critic = ( + bool(use_cuda_graph_critic) and self._device_type == "cuda" and world_size <= 1 + ) + requested_cuda_graph_critic_packed_staging = bool(use_cuda_graph_critic_packed_staging) + requested_cuda_graph_actor_packed_staging = bool(use_cuda_graph_actor_packed_staging) + self.use_cuda_graph_actor = ( + bool(use_cuda_graph_actor) and self._device_type == "cuda" and world_size <= 1 + ) + self.nvtx_profile_ranges = bool(nvtx_profile_ranges) and self._device_type == "cuda" self.amp_dtype = amp_dtype self._amp_dtype = self._resolve_amp_dtype(amp_dtype, self._device_type) self.world_size = world_size @@ -475,6 +514,7 @@ def __init__( # Entropy coefficient self.log_alpha = torch.tensor([math.log(alpha_init)], requires_grad=True, device=device) self.target_entropy = -action_dim * target_entropy_ratio + self._zero_metric = torch.zeros((), device=device) self.obs_normalizer: EmpiricalNormalization | nn.Identity if obs_normalization: @@ -484,6 +524,7 @@ def __init__( # fused AdamW requires CUDA; MPS and CPU do not support it _fused = isinstance(device, str) and device.startswith("cuda") + _optimizer_cuda_kwargs = {"capturable": True} if _fused else {} # Optimizers (AdamW with holosoma betas) self.q_optimizer = optim.AdamW( @@ -492,6 +533,7 @@ def __init__( weight_decay=weight_decay, fused=_fused, betas=(0.9, 0.95), + **_optimizer_cuda_kwargs, ) self.actor_optimizer = optim.AdamW( self.actor.parameters(), @@ -499,6 +541,7 @@ def __init__( weight_decay=weight_decay, fused=_fused, betas=(0.9, 0.95), + **_optimizer_cuda_kwargs, ) self.alpha_optimizer = optim.AdamW( [self.log_alpha], @@ -506,6 +549,7 @@ def __init__( fused=_fused, betas=(0.9, 0.95), weight_decay=0.0, + **_optimizer_cuda_kwargs, ) # Step counter @@ -517,6 +561,17 @@ def __init__( if self._should_use_grad_scaler(self.use_amp, self._device_type, self._amp_dtype) else None ) + self.use_cuda_graph_critic_packed_staging = ( + requested_cuda_graph_critic_packed_staging + and self.use_cuda_graph_critic + and self.scaler is None + ) + self.use_cuda_graph_actor_packed_staging = ( + requested_cuda_graph_actor_packed_staging + and self.use_cuda_graph_actor + and self.scaler is None + and not use_symmetry + ) self.symmetry = symmetry_augmentation if use_symmetry and symmetry_augmentation is None: @@ -524,6 +579,38 @@ def __init__( "FastSACLearner use_symmetry=True requires a symmetry_augmentation contract" ) self.use_symmetry = use_symmetry + self._cuda_graph_critic: torch.cuda.CUDAGraph | None = None + self._cuda_graph_critic_static_inputs: dict[str, torch.Tensor] | None = None + self._cuda_graph_critic_static_packed_input: torch.Tensor | None = None + self._cuda_graph_sac_static_packed_input: torch.Tensor | None = None + self._cuda_graph_sac_static_source_ptr: int | None = None + self._cuda_graph_critic_action_noise: torch.Tensor | None = None + self._cuda_graph_critic_outputs: ( + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ] + | None + ) = None + self._cuda_graph_critic_shapes: dict[str, torch.Size] | None = None + self._cuda_graph_actor: torch.cuda.CUDAGraph | None = None + self._cuda_graph_actor_static_inputs: dict[str, torch.Tensor] | None = None + self._cuda_graph_actor_static_packed_input: torch.Tensor | None = None + self._cuda_graph_actor_action_noise: torch.Tensor | None = None + self._cuda_graph_actor_outputs: ( + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ] + | None + ) = None + self._cuda_graph_actor_shapes: dict[str, torch.Size] | None = None if self.use_compile: self._compile_training_methods() @@ -552,12 +639,14 @@ def _compile_training_methods(self) -> None: return compile_kwargs = {"options": {"triton.cudagraphs": False}} - self._critic_loss_tensors = compile_fn( # type: ignore[method-assign] - self._critic_loss_tensors, **compile_kwargs - ) - self._actor_loss_tensors = compile_fn( # type: ignore[method-assign] - self._actor_loss_tensors, **compile_kwargs - ) + if not self.use_cuda_graph_critic: + self._critic_loss_tensors = compile_fn( # type: ignore[method-assign] + self._critic_loss_tensors, **compile_kwargs + ) + if not self.use_cuda_graph_actor: + self._actor_loss_tensors = compile_fn( # type: ignore[method-assign] + self._actor_loss_tensors, **compile_kwargs + ) def _autocast(self): return torch.amp.autocast( # pyright: ignore[reportPrivateImportUsage] @@ -681,6 +770,7 @@ def _get_actions_and_log_probs_for_critic( self, actor_obs: torch.Tensor, critic_obs: torch.Tensor, + eps: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Sample actor actions for critic targets. @@ -688,16 +778,17 @@ def _get_actions_and_log_probs_for_critic( preserving the standard SAC update path. """ del critic_obs - return self.actor.get_actions_and_log_probs(actor_obs) + return self.actor.get_actions_and_log_probs(actor_obs, eps=eps) def _get_actions_and_log_probs_for_actor( self, actor_obs: torch.Tensor, critic_obs: torch.Tensor, + eps: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Sample actor actions for the actor loss update.""" del critic_obs - return self.actor.get_actions_and_log_probs(actor_obs) + return self.actor.get_actions_and_log_probs(actor_obs, eps=eps) def _critic_loss_tensors( self, @@ -708,6 +799,7 @@ def _critic_loss_tensors( critic_next_obs: torch.Tensor, dones: torch.Tensor, truncated: torch.Tensor, + next_action_eps: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: bootstrap = torch.clamp(1.0 - dones.float() + truncated.float(), 0.0, 1.0) discount = torch.full_like(dones, self.gamma) @@ -717,6 +809,7 @@ def _critic_loss_tensors( next_actions, next_log_probs, _ = self._get_actions_and_log_probs_for_critic( next_obs, critic_next_obs, + eps=next_action_eps, ) adjusted_rewards = ( rewards - discount * bootstrap * self.log_alpha.exp() * next_log_probs @@ -738,15 +831,21 @@ def _critic_loss_tensors( return qf_loss, target_q_max, target_q_min, next_log_probs.detach() + def _alpha_loss_tensor(self, next_log_probs: torch.Tensor) -> torch.Tensor: + entropy_error_mean = (next_log_probs + self.target_entropy).detach().mean() + return -(self.log_alpha.exp() * entropy_error_mean) + def _actor_loss_tensors( self, obs: torch.Tensor, critic_obs: torch.Tensor, + action_eps: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: with self._autocast(): actions, log_probs, log_std = self._get_actions_and_log_probs_for_actor( obs, critic_obs, + eps=action_eps, ) with torch.no_grad(): @@ -762,6 +861,591 @@ def _actor_loss_tensors( return actor_loss, policy_entropy, action_std + def _update_actor_capture_candidate( + self, + obs: torch.Tensor, + critic_obs: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + actor_loss, policy_entropy, action_std = self._actor_loss_tensors( + obs, + critic_obs, + action_eps=self._cuda_graph_actor_action_noise, + ) + + 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: + actor_grad_norm = torch.nn.utils.clip_grad_norm_( + self.actor.parameters(), max_norm=self.max_grad_norm + ) + else: + actor_grad_norm = self._zero_metric + self.scaler.step(self.actor_optimizer) + 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 + ) + else: + actor_grad_norm = self._zero_metric + self.actor_optimizer.step() + + return actor_loss, actor_grad_norm, policy_entropy, action_std + + def _update_critic_capture_candidate( + self, + critic_obs: torch.Tensor, + actions: torch.Tensor, + rewards: torch.Tensor, + next_obs: torch.Tensor, + critic_next_obs: torch.Tensor, + dones: torch.Tensor, + truncated: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + qf_loss, target_q_max, target_q_min, next_log_probs = self._critic_loss_tensors( + critic_obs, + actions, + rewards, + next_obs, + critic_next_obs, + dones, + truncated, + next_action_eps=self._cuda_graph_critic_action_noise, + ) + + 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: + critic_grad_norm = torch.nn.utils.clip_grad_norm_( + self.qnet.parameters(), max_norm=self.max_grad_norm + ) + else: + critic_grad_norm = self._zero_metric + self.scaler.step(self.q_optimizer) + 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 + ) + else: + critic_grad_norm = self._zero_metric + self.q_optimizer.step() + + alpha_loss = self._zero_metric + if self.use_autotune: + self.alpha_optimizer.zero_grad(set_to_none=True) + alpha_loss = self._alpha_loss_tensor(next_log_probs) + alpha_loss.backward() + if ( + self.world_size > 1 + and self.distributed_sync_mode == "sync_sgd" + 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 ( + qf_loss, + critic_grad_norm, + target_q_max, + target_q_min, + alpha_loss, + self.log_alpha.exp(), + ) + + def _critic_graph_input_shapes(self, batch: Dict[str, torch.Tensor]) -> dict[str, torch.Size]: + return { + "critic": batch["critic"].shape, + "actions": batch["actions"].shape, + "rewards": batch["rewards"].shape, + "next_obs": batch["next_obs"].shape, + "next_critic": batch["next_critic"].shape, + "dones": batch["dones"].shape, + "truncated": batch["truncated"].shape, + } + + @staticmethod + def _critic_graph_input_keys() -> tuple[str, ...]: + return ( + "critic", + "actions", + "rewards", + "next_obs", + "next_critic", + "dones", + "truncated", + ) + + @classmethod + def _critic_graph_packed_width(cls, batch: Dict[str, torch.Tensor]) -> int: + width = 0 + for key in cls._critic_graph_input_keys(): + tensor = batch[key] + width += int(tensor.reshape(tensor.shape[0], -1).shape[1]) + return width + + @classmethod + def _critic_graph_static_views_from_packed( + cls, + packed: torch.Tensor, + shapes: dict[str, torch.Size], + ) -> dict[str, torch.Tensor]: + views: dict[str, torch.Tensor] = {} + offset = 0 + for key in cls._critic_graph_input_keys(): + shape = shapes[key] + width = 1 + for dim in shape[1:]: + width *= int(dim) + view = packed.narrow(1, offset, width).view(shape) + views[key] = view + offset += width + return views + + @staticmethod + def _sac_graph_offsets( + actor_shapes: dict[str, torch.Size], + critic_shapes: dict[str, torch.Size], + ) -> dict[str, tuple[int, int]]: + def width(shape: torch.Size) -> int: + value = 1 + for dim in shape[1:]: + value *= int(dim) + return value + + widths = { + "obs": width(actor_shapes["obs"]), + "critic": width(critic_shapes["critic"]), + "actions": width(critic_shapes["actions"]), + "rewards": width(critic_shapes["rewards"]), + "next_obs": width(critic_shapes["next_obs"]), + "next_critic": width(critic_shapes["next_critic"]), + "dones": width(critic_shapes["dones"]), + "truncated": width(critic_shapes["truncated"]), + } + offsets: dict[str, tuple[int, int]] = {} + offset = 0 + for key in ( + "obs", + "critic", + "actions", + "rewards", + "next_obs", + "next_critic", + "dones", + "truncated", + ): + key_width = widths[key] + offsets[key] = (offset, key_width) + offset += key_width + return offsets + + @classmethod + def _critic_graph_static_views_from_sac_packed( + cls, + packed: torch.Tensor, + critic_shapes: dict[str, torch.Size], + actor_shapes: dict[str, torch.Size], + ) -> dict[str, torch.Tensor]: + offsets = cls._sac_graph_offsets(actor_shapes, critic_shapes) + views: dict[str, torch.Tensor] = {} + for key in cls._critic_graph_input_keys(): + offset, width = offsets[key] + views[key] = packed.narrow(1, offset, width).view(critic_shapes[key]) + return views + + @classmethod + def _actor_graph_static_views_from_sac_packed( + cls, + packed: torch.Tensor, + actor_shapes: dict[str, torch.Size], + ) -> dict[str, torch.Tensor]: + batch_size = int(actor_shapes["obs"][0]) + critic_shapes = { + "critic": actor_shapes["critic"], + "actions": torch.Size((batch_size, 0)), + "rewards": torch.Size((batch_size,)), + "next_obs": actor_shapes["obs"], + "next_critic": actor_shapes["critic"], + "dones": torch.Size((batch_size,)), + "truncated": torch.Size((batch_size,)), + } + offsets = cls._sac_graph_offsets(actor_shapes, critic_shapes) + views: dict[str, torch.Tensor] = {} + for key in ("obs", "critic"): + offset, width = offsets[key] + views[key] = packed.narrow(1, offset, width).view(actor_shapes[key]) + return views + + def _copy_critic_graph_inputs(self, batch: Dict[str, torch.Tensor]) -> None: + assert self._cuda_graph_critic_static_inputs is not None + packed_source = batch.get("sac_graph_packed_source") + if packed_source is not None and self._cuda_graph_sac_static_packed_input is not None: + self._cuda_graph_sac_static_packed_input.copy_(packed_source) + self._cuda_graph_sac_static_source_ptr = int(packed_source.data_ptr()) + else: + self._cuda_graph_sac_static_source_ptr = None + packed_source = batch.get("critic_graph_packed_source") + if ( + packed_source is not None + and self._cuda_graph_critic_static_packed_input is not None + ): + self._cuda_graph_critic_static_packed_input.copy_(packed_source) + else: + for key, tensor in self._cuda_graph_critic_static_inputs.items(): + tensor.copy_(batch[key]) + assert self._cuda_graph_critic_action_noise is not None + self._cuda_graph_critic_action_noise.normal_() + + def _copy_actor_graph_inputs(self, batch: Dict[str, torch.Tensor]) -> None: + assert self._cuda_graph_actor_static_inputs is not None + packed_source = batch.get("sac_graph_packed_source") + if packed_source is not None: + static_packed = self._cuda_graph_actor_static_packed_input + if static_packed is None: + static_packed = self._cuda_graph_sac_static_packed_input + if static_packed is not None: + source_ptr = int(packed_source.data_ptr()) + if ( + static_packed is not self._cuda_graph_sac_static_packed_input + or self._cuda_graph_sac_static_source_ptr != source_ptr + ): + static_packed.copy_(packed_source) + if static_packed is self._cuda_graph_sac_static_packed_input: + self._cuda_graph_sac_static_source_ptr = source_ptr + else: + for key, tensor in self._cuda_graph_actor_static_inputs.items(): + tensor.copy_(batch[key]) + else: + for key, tensor in self._cuda_graph_actor_static_inputs.items(): + tensor.copy_(batch[key]) + assert self._cuda_graph_actor_action_noise is not None + self._cuda_graph_actor_action_noise.normal_() + + def _actor_graph_input_shapes(self, batch: Dict[str, torch.Tensor]) -> dict[str, torch.Size]: + return { + "obs": batch["obs"].shape, + "critic": batch["critic"].shape, + } + + def _materialize_capturable_critic_optimizer_state( + self, + batch: Dict[str, torch.Tensor], + ) -> None: + optimizer_lrs = [group["lr"] for group in self.q_optimizer.param_groups] + optimizer_weight_decays = [group["weight_decay"] for group in self.q_optimizer.param_groups] + alpha_lrs = [group["lr"] for group in self.alpha_optimizer.param_groups] + cpu_rng_state = torch.random.get_rng_state() + cuda_rng_state = torch.cuda.get_rng_state() if self._device_type == "cuda" else None + try: + for group in self.q_optimizer.param_groups: + group["lr"] = 0.0 + group["weight_decay"] = 0.0 + for group in self.alpha_optimizer.param_groups: + group["lr"] = 0.0 + self._update_critic_capture_candidate( + batch["critic"], + batch["actions"], + batch["rewards"], + batch["next_obs"], + batch["next_critic"], + batch["dones"], + batch["truncated"], + ) + finally: + torch.random.set_rng_state(cpu_rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) + for group, lr, weight_decay in zip( + self.q_optimizer.param_groups, + optimizer_lrs, + optimizer_weight_decays, + strict=True, + ): + group["lr"] = lr + group["weight_decay"] = weight_decay + for group, lr in zip(self.alpha_optimizer.param_groups, alpha_lrs, strict=True): + group["lr"] = lr + + for optimizer in (self.q_optimizer, self.alpha_optimizer): + optimizer.zero_grad(set_to_none=True) + for state in optimizer.state.values(): + step = state.get("step") + if isinstance(step, torch.Tensor): + step.zero_() + elif step is not None: + state["step"] = 0 + for name in ("exp_avg", "exp_avg_sq", "max_exp_avg_sq"): + tensor = state.get(name) + if isinstance(tensor, torch.Tensor): + tensor.zero_() + + def _reset_critic_cuda_graph(self) -> None: + self._cuda_graph_critic = None + self._cuda_graph_critic_static_inputs = None + self._cuda_graph_critic_static_packed_input = None + self._cuda_graph_sac_static_packed_input = None + self._cuda_graph_sac_static_source_ptr = None + self._cuda_graph_critic_action_noise = None + self._cuda_graph_critic_outputs = None + self._cuda_graph_critic_shapes = None + + def _materialize_capturable_actor_optimizer_state( + self, + batch: Dict[str, torch.Tensor], + ) -> None: + optimizer_lrs = [group["lr"] for group in self.actor_optimizer.param_groups] + optimizer_weight_decays = [ + group["weight_decay"] for group in self.actor_optimizer.param_groups + ] + cpu_rng_state = torch.random.get_rng_state() + cuda_rng_state = torch.cuda.get_rng_state() if self._device_type == "cuda" else None + try: + for group in self.actor_optimizer.param_groups: + group["lr"] = 0.0 + group["weight_decay"] = 0.0 + self._update_actor_capture_candidate( + batch["obs"], + batch["critic"], + ) + finally: + torch.random.set_rng_state(cpu_rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) + for group, lr, weight_decay in zip( + self.actor_optimizer.param_groups, + optimizer_lrs, + optimizer_weight_decays, + strict=True, + ): + group["lr"] = lr + group["weight_decay"] = weight_decay + + self.actor_optimizer.zero_grad(set_to_none=True) + for state in self.actor_optimizer.state.values(): + step = state.get("step") + if isinstance(step, torch.Tensor): + step.zero_() + elif step is not None: + state["step"] = 0 + for name in ("exp_avg", "exp_avg_sq", "max_exp_avg_sq"): + tensor = state.get(name) + if isinstance(tensor, torch.Tensor): + tensor.zero_() + + def _reset_actor_cuda_graph(self) -> None: + self._cuda_graph_actor = None + self._cuda_graph_actor_static_inputs = None + self._cuda_graph_actor_static_packed_input = None + self._cuda_graph_actor_action_noise = None + self._cuda_graph_actor_outputs = None + self._cuda_graph_actor_shapes = None + + def _capture_actor_cuda_graph(self, batch: Dict[str, torch.Tensor]) -> None: + if not self.use_cuda_graph_actor: + return + self._cuda_graph_actor_shapes = self._actor_graph_input_shapes(batch) + if self.use_cuda_graph_actor_packed_staging and "sac_graph_packed_source" in batch: + packed_source = batch["sac_graph_packed_source"] + if ( + self._cuda_graph_sac_static_packed_input is not None + and self._cuda_graph_sac_static_packed_input.shape == packed_source.shape + ): + self._cuda_graph_actor_static_packed_input = ( + self._cuda_graph_sac_static_packed_input + ) + else: + self._cuda_graph_sac_static_packed_input = packed_source.detach().clone() + self._cuda_graph_sac_static_source_ptr = None + self._cuda_graph_actor_static_packed_input = ( + self._cuda_graph_sac_static_packed_input + ) + self._cuda_graph_actor_static_inputs = self._actor_graph_static_views_from_sac_packed( + self._cuda_graph_actor_static_packed_input, + self._cuda_graph_actor_shapes, + ) + else: + self._cuda_graph_actor_static_packed_input = None + self._cuda_graph_actor_static_inputs = { + "obs": batch["obs"].detach().clone(), + "critic": batch["critic"].detach().clone(), + } + self._cuda_graph_actor_action_noise = torch.empty( + batch["obs"].shape[:-1] + (self.actor.action_dim,), + device=batch["obs"].device, + dtype=batch["obs"].dtype, + ) + self._copy_actor_graph_inputs(batch) + + graph = torch.cuda.CUDAGraph() + capture_stream = cast(torch.cuda.Stream, torch.cuda.Stream()) + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream), torch.cuda.graph(graph): + self._cuda_graph_actor_outputs = self._update_actor_capture_candidate( + self._cuda_graph_actor_static_inputs["obs"], + self._cuda_graph_actor_static_inputs["critic"], + ) + torch.cuda.current_stream().wait_stream(capture_stream) + torch.cuda.synchronize() + self._cuda_graph_actor = graph + + def _actor_graph_output_metrics(self, *, read_items: bool = True) -> Dict[str, float]: + if not read_items: + return {} + assert self._cuda_graph_actor_outputs is not None + actor_loss, actor_grad_norm, policy_entropy, action_std = self._cuda_graph_actor_outputs + return { + "actor_loss": actor_loss.item(), + "actor_grad_norm": actor_grad_norm.item(), + "policy_entropy": policy_entropy.item(), + "action_std": action_std.item(), + } + + def _capture_critic_cuda_graph(self, batch: Dict[str, torch.Tensor]) -> None: + if not self.use_cuda_graph_critic: + return + self._cuda_graph_critic_shapes = self._critic_graph_input_shapes(batch) + if self.use_cuda_graph_critic_packed_staging and "sac_graph_packed_source" in batch: + packed_source = batch["sac_graph_packed_source"] + self._cuda_graph_sac_static_packed_input = packed_source.detach().clone() + self._cuda_graph_critic_static_packed_input = None + actor_shapes = self._actor_graph_input_shapes(batch) + self._cuda_graph_critic_static_inputs = self._critic_graph_static_views_from_sac_packed( + self._cuda_graph_sac_static_packed_input, + self._cuda_graph_critic_shapes, + actor_shapes, + ) + elif self.use_cuda_graph_critic_packed_staging and "critic_graph_packed_source" in batch: + packed_source = batch["critic_graph_packed_source"] + self._cuda_graph_sac_static_packed_input = None + self._cuda_graph_critic_static_packed_input = packed_source.detach().clone() + self._cuda_graph_critic_static_inputs = self._critic_graph_static_views_from_packed( + self._cuda_graph_critic_static_packed_input, + self._cuda_graph_critic_shapes, + ) + else: + self._cuda_graph_sac_static_packed_input = None + self._cuda_graph_critic_static_packed_input = None + self._cuda_graph_critic_static_inputs = { + "critic": batch["critic"].detach().clone(), + "actions": batch["actions"].detach().clone(), + "rewards": batch["rewards"].detach().clone(), + "next_obs": batch["next_obs"].detach().clone(), + "next_critic": batch["next_critic"].detach().clone(), + "dones": batch["dones"].detach().clone(), + "truncated": batch["truncated"].detach().clone(), + } + self._cuda_graph_critic_action_noise = torch.empty( + batch["next_obs"].shape[:-1] + (batch["actions"].shape[-1],), + device=batch["next_obs"].device, + dtype=batch["next_obs"].dtype, + ) + self._copy_critic_graph_inputs(batch) + + graph = torch.cuda.CUDAGraph() + capture_stream = cast(torch.cuda.Stream, torch.cuda.Stream()) + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream), torch.cuda.graph(graph): + self._cuda_graph_critic_outputs = self._update_critic_capture_candidate( + self._cuda_graph_critic_static_inputs["critic"], + self._cuda_graph_critic_static_inputs["actions"], + self._cuda_graph_critic_static_inputs["rewards"], + self._cuda_graph_critic_static_inputs["next_obs"], + self._cuda_graph_critic_static_inputs["next_critic"], + self._cuda_graph_critic_static_inputs["dones"], + self._cuda_graph_critic_static_inputs["truncated"], + ) + torch.cuda.current_stream().wait_stream(capture_stream) + torch.cuda.synchronize() + self._cuda_graph_critic = graph + + def _critic_graph_output_metrics(self, *, read_items: bool = True) -> Dict[str, float]: + if not read_items: + return {} + assert self._cuda_graph_critic_outputs is not None + qf_loss, critic_grad_norm, target_q_max, target_q_min, alpha_loss, alpha = ( + self._cuda_graph_critic_outputs + ) + return { + "qf_loss": qf_loss.item(), + "critic_grad_norm": critic_grad_norm.item(), + "target_q_max": target_q_max.item(), + "target_q_min": target_q_min.item(), + "alpha_loss": alpha_loss.item(), + "alpha": alpha.item(), + } + + def update_critic_cuda_graph( + self, + batch: Dict[str, torch.Tensor], + *, + read_metrics: bool = True, + ) -> Dict[str, float]: + if not self.use_cuda_graph_critic: + return self.update_critic(batch) + if self._device_type != "cuda": + return self.update_critic(batch) + if self.scaler is not None or self.world_size > 1: + return self.update_critic(batch) + if self._cuda_graph_critic_shapes != self._critic_graph_input_shapes(batch): + self._reset_critic_cuda_graph() + self._materialize_capturable_critic_optimizer_state(batch) + self._capture_critic_cuda_graph(batch) + with _cuda_nvtx_range( + "critic_graph/output_metrics_item", + self.nvtx_profile_ranges, + ): + return self._critic_graph_output_metrics(read_items=read_metrics) + assert self._cuda_graph_critic is not None + with _cuda_nvtx_range("critic_graph/copy_inputs", self.nvtx_profile_ranges): + self._copy_critic_graph_inputs(batch) + with _cuda_nvtx_range("critic_graph/replay", self.nvtx_profile_ranges): + self._cuda_graph_critic.replay() + with _cuda_nvtx_range("critic_graph/output_metrics_item", self.nvtx_profile_ranges): + return self._critic_graph_output_metrics(read_items=read_metrics) + + def update_actor_cuda_graph( + self, + batch: Dict[str, torch.Tensor], + *, + read_metrics: bool = True, + ) -> Dict[str, float]: + if not self.use_cuda_graph_actor: + return self.update_actor(batch) + if self._device_type != "cuda": + return self.update_actor(batch) + if self.scaler is not None or self.world_size > 1 or self.use_symmetry: + return self.update_actor(batch) + if self._cuda_graph_actor_shapes != self._actor_graph_input_shapes(batch): + self._reset_actor_cuda_graph() + self._materialize_capturable_actor_optimizer_state(batch) + self._capture_actor_cuda_graph(batch) + with _cuda_nvtx_range( + "actor_graph/output_metrics_item", + self.nvtx_profile_ranges, + ): + return self._actor_graph_output_metrics(read_items=read_metrics) + assert self._cuda_graph_actor is not None + with _cuda_nvtx_range("actor_graph/copy_inputs", self.nvtx_profile_ranges): + self._copy_actor_graph_inputs(batch) + with _cuda_nvtx_range("actor_graph/replay", self.nvtx_profile_ranges): + self._cuda_graph_actor.replay() + with _cuda_nvtx_range("actor_graph/output_metrics_item", self.nvtx_profile_ranges): + return self._actor_graph_output_metrics(read_items=read_metrics) + def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: """One critic update step.""" obs = batch["obs"] @@ -775,86 +1459,108 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: # Apply symmetry augmentation if self.use_symmetry: - orig_actions = actions - - assert self.symmetry is not None - obs, actions = self.symmetry.augment_obs_and_actions(obs, actions, obs_group="obs") - next_obs, _ = self.symmetry.augment_obs_and_actions( - next_obs, orig_actions, obs_group="obs" - ) + with _cuda_nvtx_range("critic/symmetry_augment", self.nvtx_profile_ranges): + assert self.symmetry is not None + with _cuda_nvtx_range("critic/symmetry_obs_actions", self.nvtx_profile_ranges): + obs, actions = self.symmetry.augment_obs_and_actions( + obs, + actions, + obs_group="obs", + ) + with _cuda_nvtx_range("critic/symmetry_next_obs", self.nvtx_profile_ranges): + next_obs = self.symmetry.augment_obs( + next_obs, + obs_group="obs", + ) - critic_obs, _ = self.symmetry.augment_obs_and_actions( - critic_obs, - orig_actions, - obs_group="critic", - ) - critic_next_obs, _ = self.symmetry.augment_obs_and_actions( - critic_next_obs, - orig_actions, - obs_group="critic", - ) + with _cuda_nvtx_range("critic/symmetry_critic_obs", self.nvtx_profile_ranges): + critic_obs = self.symmetry.augment_obs( + critic_obs, + obs_group="critic", + ) + with _cuda_nvtx_range("critic/symmetry_critic_next_obs", self.nvtx_profile_ranges): + critic_next_obs = self.symmetry.augment_obs( + critic_next_obs, + obs_group="critic", + ) - # Double the batch size for other tensors - rewards = rewards.repeat(2) - dones = dones.repeat(2) - truncated = truncated.repeat(2) + # Double the batch size for other tensors + with _cuda_nvtx_range("critic/symmetry_aux_repeat", self.nvtx_profile_ranges): + rewards = rewards.repeat(2) + dones = dones.repeat(2) + truncated = truncated.repeat(2) self.normalize_obs(obs, update=True) next_obs = self.normalize_obs(next_obs, update=False) - qf_loss, target_q_max, target_q_min, next_log_probs = self._critic_loss_tensors( - critic_obs, - actions, - rewards, - next_obs, - critic_next_obs, - dones, - truncated, - ) + with _cuda_nvtx_range("critic/loss_compiled", self.nvtx_profile_ranges): + qf_loss, target_q_max, target_q_min, next_log_probs = self._critic_loss_tensors( + critic_obs, + actions, + rewards, + next_obs, + critic_next_obs, + dones, + truncated, + ) # Skip if NaN if torch.isfinite(qf_loss): self.q_optimizer.zero_grad(set_to_none=True) if self.scaler: - self.scaler.scale(qf_loss).backward() + with _cuda_nvtx_range("critic/backward", self.nvtx_profile_ranges): + self.scaler.scale(qf_loss).backward() self.scaler.unscale_(self.q_optimizer) 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 - ) + with _cuda_nvtx_range("critic/grad_clip", self.nvtx_profile_ranges): + critic_grad_norm = torch.nn.utils.clip_grad_norm_( + self.qnet.parameters(), max_norm=self.max_grad_norm + ) else: - critic_grad_norm = torch.tensor(0.0, device=self.device) - self.scaler.step(self.q_optimizer) + critic_grad_norm = self._zero_metric + with _cuda_nvtx_range("critic/q_optimizer_step", self.nvtx_profile_ranges): + self.scaler.step(self.q_optimizer) self.scaler.update() else: - qf_loss.backward() + with _cuda_nvtx_range("critic/backward", self.nvtx_profile_ranges): + 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 - ) + with _cuda_nvtx_range("critic/grad_clip", self.nvtx_profile_ranges): + critic_grad_norm = torch.nn.utils.clip_grad_norm_( + self.qnet.parameters(), max_norm=self.max_grad_norm + ) else: - critic_grad_norm = torch.tensor(0.0, device=self.device) - self.q_optimizer.step() + critic_grad_norm = self._zero_metric + with _cuda_nvtx_range("critic/q_optimizer_step", self.nvtx_profile_ranges): + self.q_optimizer.step() else: - critic_grad_norm = torch.tensor(0.0, device=self.device) + critic_grad_norm = self._zero_metric # Alpha loss (temperature update) - matching holosoma - alpha_loss = torch.tensor(0.0, device=self.device) + alpha_loss = self._zero_metric if self.use_autotune: - self.alpha_optimizer.zero_grad(set_to_none=True) - alpha_loss = (-self.log_alpha.exp() * (next_log_probs + self.target_entropy)).mean() - if torch.isfinite(alpha_loss): - alpha_loss.backward() - if ( - self.world_size > 1 - and self.distributed_sync_mode == "sync_sgd" - 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() + with _cuda_nvtx_range("critic/alpha_update", self.nvtx_profile_ranges): + self.alpha_optimizer.zero_grad(set_to_none=True) + with _cuda_nvtx_range("critic/alpha_loss", self.nvtx_profile_ranges): + alpha_loss = self._alpha_loss_tensor(next_log_probs) + if torch.isfinite(alpha_loss): + with _cuda_nvtx_range("critic/alpha_backward", self.nvtx_profile_ranges): + alpha_loss.backward() + if ( + self.world_size > 1 + and self.distributed_sync_mode == "sync_sgd" + and self.log_alpha.grad is not None + ): + with _cuda_nvtx_range( + "critic/alpha_distributed_reduce", + self.nvtx_profile_ranges, + ): + dist.all_reduce(self.log_alpha.grad, op=dist.ReduceOp.SUM) + self.log_alpha.grad /= self.world_size + with _cuda_nvtx_range("critic/alpha_optimizer_step", self.nvtx_profile_ranges): + self.alpha_optimizer.step() return { "qf_loss": qf_loss.item(), @@ -872,43 +1578,50 @@ def update_actor(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: # Apply symmetry augmentation if self.use_symmetry: - assert self.symmetry is not None - obs = torch.cat([obs, self.symmetry.mirror_obs(obs, obs_group="obs")], dim=0) - critic_obs = torch.cat( - [critic_obs, self.symmetry.mirror_obs(critic_obs, obs_group="critic")], - dim=0, - ) + with _cuda_nvtx_range("actor/symmetry_augment", self.nvtx_profile_ranges): + assert self.symmetry is not None + with _cuda_nvtx_range("actor/symmetry_obs", self.nvtx_profile_ranges): + obs = self.symmetry.augment_obs(obs, obs_group="obs") + with _cuda_nvtx_range("actor/symmetry_critic_obs", self.nvtx_profile_ranges): + critic_obs = self.symmetry.augment_obs(critic_obs, obs_group="critic") obs = self.normalize_obs(obs, update=False) - actor_loss, policy_entropy, action_std = self._actor_loss_tensors(obs, critic_obs) + with _cuda_nvtx_range("actor/loss_compiled", self.nvtx_profile_ranges): + actor_loss, policy_entropy, action_std = self._actor_loss_tensors(obs, critic_obs) # Skip if NaN if torch.isfinite(actor_loss): self.actor_optimizer.zero_grad(set_to_none=True) if self.scaler: - self.scaler.scale(actor_loss).backward() + with _cuda_nvtx_range("actor/backward", self.nvtx_profile_ranges): + self.scaler.scale(actor_loss).backward() self.scaler.unscale_(self.actor_optimizer) 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 - ) + with _cuda_nvtx_range("actor/grad_clip", self.nvtx_profile_ranges): + actor_grad_norm = torch.nn.utils.clip_grad_norm_( + self.actor.parameters(), max_norm=self.max_grad_norm + ) else: - actor_grad_norm = torch.tensor(0.0, device=self.device) - self.scaler.step(self.actor_optimizer) + actor_grad_norm = self._zero_metric + with _cuda_nvtx_range("actor/optimizer_step", self.nvtx_profile_ranges): + self.scaler.step(self.actor_optimizer) self.scaler.update() else: - actor_loss.backward() + with _cuda_nvtx_range("actor/backward", self.nvtx_profile_ranges): + 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 - ) + with _cuda_nvtx_range("actor/grad_clip", self.nvtx_profile_ranges): + actor_grad_norm = torch.nn.utils.clip_grad_norm_( + self.actor.parameters(), max_norm=self.max_grad_norm + ) else: - actor_grad_norm = torch.tensor(0.0, device=self.device) - self.actor_optimizer.step() + actor_grad_norm = self._zero_metric + with _cuda_nvtx_range("actor/optimizer_step", self.nvtx_profile_ranges): + self.actor_optimizer.step() else: - actor_grad_norm = torch.tensor(0.0, device=self.device) + actor_grad_norm = self._zero_metric return { "actor_loss": actor_loss.item(), @@ -920,8 +1633,15 @@ def update_actor(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: def soft_update_target(self) -> None: """Polyak-average update of the target Q-network.""" with torch.no_grad(): - for tgt, src in zip(self.qnet_target.parameters(), self.qnet.parameters()): - tgt.data.mul_(1.0 - self.tau).add_(src.data, alpha=self.tau) + with _cuda_nvtx_range("target/soft_update_loop", self.nvtx_profile_ranges): + target_params = cast(list[torch.Tensor], list(self.qnet_target.parameters())) + source_params = cast(list[torch.Tensor], list(self.qnet.parameters())) + try: + torch._foreach_mul_(target_params, 1.0 - self.tau) + torch._foreach_add_(target_params, source_params, alpha=self.tau) + except RuntimeError: + for tgt, src in zip(target_params, source_params): + tgt.mul_(1.0 - self.tau).add_(src, alpha=self.tau) def get_state_dict(self) -> Dict[str, Any]: """Save all components.""" @@ -953,6 +1673,8 @@ def load_state_dict(self, state_dict: Dict) -> None: if state_dict.get("obs_normalizer") and hasattr(self.obs_normalizer, "load_state_dict"): self.obs_normalizer.load_state_dict(state_dict["obs_normalizer"]) self.update_count = state_dict.get("update_count", 0) + self._reset_critic_cuda_graph() + self._reset_actor_cuda_graph() # --------------------------------------------------------------------------- diff --git a/src/unilab/algos/torch/fast_sac/runner.py b/src/unilab/algos/torch/fast_sac/runner.py index 8f5a6fb48..7b0cc08f9 100644 --- a/src/unilab/algos/torch/fast_sac/runner.py +++ b/src/unilab/algos/torch/fast_sac/runner.py @@ -39,6 +39,8 @@ def __init__( max_grad_norm: float = 0.0, use_amp: bool = False, amp_dtype: str = "auto", + use_cuda_graph_critic: bool = False, + use_cuda_graph_actor: bool = False, sim_backend: str = "mujoco", use_symmetry: bool = False, world_size: int = 1, @@ -93,6 +95,8 @@ def __init__( use_amp=use_amp, amp_dtype=amp_dtype, obs_normalization=obs_normalization, + use_cuda_graph_critic=use_cuda_graph_critic, + use_cuda_graph_actor=use_cuda_graph_actor, use_symmetry=use_symmetry, symmetry_augmentation=symmetry_augmentation, world_size=getattr(self, "world_size", world_size), diff --git a/src/unilab/algos/torch/flash_sac/double_buffer.py b/src/unilab/algos/torch/flash_sac/double_buffer.py index 178906698..b93d2e3b7 100644 --- a/src/unilab/algos/torch/flash_sac/double_buffer.py +++ b/src/unilab/algos/torch/flash_sac/double_buffer.py @@ -103,6 +103,14 @@ def build_flashsac_double_buffer_runner( "use_amp": cfg.training.use_amp, "amp_dtype": cfg.algo.algo_params.amp_dtype, "use_compile": cfg.algo.algo_params.use_compile, + "use_cuda_graph_critic": cfg.algo.algo_params.use_cuda_graph_critic, + "use_cuda_graph_actor": cfg.algo.algo_params.use_cuda_graph_actor, + "use_cuda_graph_critic_packed_staging": ( + cfg.algo.algo_params.use_cuda_graph_critic_packed_staging + ), + "use_cuda_graph_actor_packed_staging": ( + cfg.algo.algo_params.use_cuda_graph_actor_packed_staging + ), } actor_kwargs = { "actor_num_blocks": cfg.algo.algo_params.actor_num_blocks, diff --git a/src/unilab/algos/torch/flash_sac/learner.py b/src/unilab/algos/torch/flash_sac/learner.py index 744bde3f0..05f7be0a0 100644 --- a/src/unilab/algos/torch/flash_sac/learner.py +++ b/src/unilab/algos/torch/flash_sac/learner.py @@ -156,6 +156,7 @@ class FlashSACLearner: supports_multi_gpu = True supports_multi_gpu_symmetry = False supported_multi_gpu_sync_modes = frozenset({"sync_sgd", "local_sgd"}) + supports_cuda_graph_packed_staging = True def __init__( self, @@ -192,6 +193,10 @@ def __init__( use_amp: bool = False, amp_dtype: str = "auto", use_compile: bool = False, + use_cuda_graph_critic: bool = False, + use_cuda_graph_actor: bool = False, + use_cuda_graph_critic_packed_staging: bool = False, + use_cuda_graph_actor_packed_staging: bool = False, world_size: int = 1, distributed_sync_mode: str = "sync_sgd", ): @@ -210,6 +215,14 @@ def __init__( self.use_compile = bool( use_compile and get_torch_compile_for_cuda(self.device, warn=True) is not None ) + self.use_cuda_graph_critic = bool(use_cuda_graph_critic) + self.use_cuda_graph_actor = bool(use_cuda_graph_actor) + self.use_cuda_graph_critic_packed_staging = bool( + use_cuda_graph_critic_packed_staging and self.use_cuda_graph_critic + ) + self.use_cuda_graph_actor_packed_staging = bool( + use_cuda_graph_actor_packed_staging and self.use_cuda_graph_actor + ) self.world_size = int(world_size) self.distributed_sync_mode = normalize_distributed_sync_mode(distributed_sync_mode) @@ -260,12 +273,27 @@ def __init__( else None ) lr_peak = learning_rate_peak if learning_rate_peak > 0 else actor_lr - fused = self.device.type == "cuda" - self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr_peak, fused=fused) - self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr_peak, fused=fused) + optimizer_kwargs: dict[str, Any] = {"fused": self.device.type == "cuda"} + if self.device.type == "cuda": + optimizer_kwargs["capturable"] = True + self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr_peak, **optimizer_kwargs) + self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr_peak, **optimizer_kwargs) self.temperature_optimizer = optim.Adam( - self.temperature.parameters(), lr=lr_peak, fused=fused + self.temperature.parameters(), lr=lr_peak, **optimizer_kwargs ) + self._cuda_graph_critic: torch.cuda.CUDAGraph | None = None + self._cuda_graph_critic_static_inputs: dict[str, torch.Tensor] | None = None + self._cuda_graph_sac_static_packed_input: torch.Tensor | None = None + self._cuda_graph_sac_static_source_ptr: int | None = None + self._cuda_graph_critic_outputs: tuple[torch.Tensor, torch.Tensor] | None = None + self._cuda_graph_critic_shapes: dict[str, torch.Size] | None = None + self._cuda_graph_actor: torch.cuda.CUDAGraph | None = None + self._cuda_graph_actor_static_inputs: dict[str, torch.Tensor] | None = None + self._cuda_graph_actor_static_packed_input: torch.Tensor | None = None + self._cuda_graph_actor_outputs: ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None + ) = None + self._cuda_graph_actor_shapes: dict[str, torch.Size] | None = None scheduler_fn = build_lr_lambda( init_lr=learning_rate_init, @@ -292,12 +320,14 @@ def _compile_training_methods(self) -> None: self.actor.get_mean_and_std = compile_fn( # type: ignore[method-assign] self.actor.get_mean_and_std, **compile_kwargs ) - self._critic_loss_tensors = compile_fn( # type: ignore[method-assign] - self._critic_loss_tensors, **compile_kwargs - ) - self._actor_loss_tensors = compile_fn( # type: ignore[method-assign] - self._actor_loss_tensors, **compile_kwargs - ) + if not self.use_cuda_graph_critic: + self._critic_loss_tensors = compile_fn( # type: ignore[method-assign] + self._critic_loss_tensors, **compile_kwargs + ) + if not self.use_cuda_graph_actor: + self._actor_loss_tensors = compile_fn( # type: ignore[method-assign] + self._actor_loss_tensors, **compile_kwargs + ) @staticmethod def _resolve_amp_dtype(amp_dtype: str, device_type: str) -> torch.dtype: @@ -509,6 +539,526 @@ def _actor_loss_tensors( entropy = -log_probs.detach().mean() return actor_loss, entropy + @staticmethod + def _critic_graph_input_keys() -> tuple[str, ...]: + return ( + "obs", + "actions", + "rewards", + "next_obs", + "dones", + "truncated", + "critic", + "next_critic", + ) + + def _critic_graph_input_shapes(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Size]: + return {key: inputs[key].shape for key in self._critic_graph_input_keys()} + + @staticmethod + def _graph_width(shape: torch.Size) -> int: + value = 1 + for dim in shape[1:]: + value *= int(dim) + return value + + @classmethod + def _sac_graph_offsets( + cls, + actor_shapes: dict[str, torch.Size], + critic_shapes: dict[str, torch.Size], + ) -> dict[str, tuple[int, int]]: + widths = { + "obs": cls._graph_width(actor_shapes["obs"]), + "critic": cls._graph_width(critic_shapes["critic"]), + "actions": cls._graph_width(critic_shapes["actions"]), + "rewards": cls._graph_width(critic_shapes["rewards"]), + "next_obs": cls._graph_width(critic_shapes["next_obs"]), + "next_critic": cls._graph_width(critic_shapes["next_critic"]), + "dones": cls._graph_width(critic_shapes["dones"]), + "truncated": cls._graph_width(critic_shapes["truncated"]), + } + offsets: dict[str, tuple[int, int]] = {} + offset = 0 + for key in ( + "obs", + "critic", + "actions", + "rewards", + "next_obs", + "next_critic", + "dones", + "truncated", + ): + key_width = widths[key] + offsets[key] = (offset, key_width) + offset += key_width + return offsets + + @classmethod + def _critic_graph_static_views_from_sac_packed( + cls, + packed: torch.Tensor, + critic_shapes: dict[str, torch.Size], + actor_shapes: dict[str, torch.Size], + ) -> dict[str, torch.Tensor]: + offsets = cls._sac_graph_offsets(actor_shapes, critic_shapes) + views: dict[str, torch.Tensor] = {} + for key in cls._critic_graph_input_keys(): + offset, width = offsets[key] + views[key] = packed.narrow(1, offset, width).view(critic_shapes[key]) + return views + + @classmethod + def _actor_graph_static_views_from_sac_packed( + cls, + packed: torch.Tensor, + actor_shapes: dict[str, torch.Size], + ) -> dict[str, torch.Tensor]: + batch_size = int(actor_shapes["obs"][0]) + critic_shapes = { + "critic": actor_shapes["critic"], + "actions": actor_shapes["actions"], + "rewards": torch.Size((batch_size,)), + "next_obs": actor_shapes["next_obs"], + "next_critic": actor_shapes["critic"], + "dones": torch.Size((batch_size,)), + "truncated": torch.Size((batch_size,)), + } + offsets = cls._sac_graph_offsets(actor_shapes, critic_shapes) + views: dict[str, torch.Tensor] = {} + for key in cls._actor_graph_input_keys(): + source_key = "actions" if key == "actions" else key + offset, width = offsets[source_key] + views[key] = packed.narrow(1, offset, width).view(actor_shapes[key]) + return views + + def _prepare_critic_graph_inputs( + self, + batch: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + obs = batch["obs"].to(self.device) + actions = batch["actions"].to(self.device) + rewards = batch["rewards"].to(self.device) + next_obs = batch["next_obs"].to(self.device) + dones = batch["dones"].to(self.device) + truncated = batch["truncated"].to(self.device) + critic_obs = batch["critic"].to(self.device) + critic_next_obs = batch["next_critic"].to(self.device) + + obs = self._maybe_normalize_obs(obs, update=True) + next_obs = self._maybe_normalize_obs(next_obs, update=False) + if self.reward_normalizer is not None: + rewards = self.reward_normalizer.normalize(rewards) + + prepared = { + "obs": obs, + "actions": actions, + "rewards": rewards, + "next_obs": next_obs, + "dones": dones, + "truncated": truncated, + "critic": critic_obs, + "next_critic": critic_next_obs, + } + if "sac_graph_packed_source" in batch: + prepared["sac_graph_packed_source"] = batch["sac_graph_packed_source"].to(self.device) + return prepared + + def _copy_critic_graph_inputs(self, inputs: dict[str, torch.Tensor]) -> None: + assert self._cuda_graph_critic_static_inputs is not None + packed_source = inputs.get("sac_graph_packed_source") + if packed_source is not None and self._cuda_graph_sac_static_packed_input is not None: + self._cuda_graph_sac_static_packed_input.copy_(packed_source) + self._cuda_graph_sac_static_source_ptr = int(packed_source.data_ptr()) + return + for key, tensor in self._cuda_graph_critic_static_inputs.items(): + tensor.copy_(inputs[key]) + + def _update_critic_capture_candidate( + self, + inputs: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + actions = inputs["actions"] + rewards = inputs["rewards"] + next_obs = inputs["next_obs"] + dones = inputs["dones"] + truncated = inputs["truncated"] + critic_obs = inputs["critic"] + critic_next_obs = inputs["next_critic"] + + gamma = self.gamma**self.n_step + obs_all = torch.cat([critic_obs, critic_next_obs], dim=0) + + with torch.no_grad(): + with self._autocast(): + next_actions, actor_info = self.actor(next_obs, training=False) + actor_entropy = self.temperature().detach() * actor_info["log_prob"] + act_all = torch.cat([actions, next_actions], dim=0) + qs_all, q_info_all = self.target_critic(obs_all, act_all, training=True) + next_q_values = qs_all.chunk(2, dim=1)[1] + next_q_log_probs_full = q_info_all["log_prob"].chunk(2, dim=1)[1] + support = cast(torch.Tensor, self.target_critic.predictor.support) + + with self._autocast(): + _, pred_info_all = self.critic(obs_all, act_all, training=True) + pred_log_probs = pred_info_all["log_prob"].chunk(2, dim=1)[0] + critic_loss = self._critic_loss_tensors( + next_q_values, + next_q_log_probs_full, + support, + rewards, + dones, + truncated, + actor_entropy, + pred_log_probs, + gamma, + ) + + self.critic_optimizer.zero_grad(set_to_none=True) + critic_loss.backward() + self.critic_optimizer.step() + reward_scale_std = ( + torch.sqrt(self.reward_normalizer.rms.var) + if self.reward_normalizer is not None + else torch.ones((), device=self.device) + ) + return critic_loss, reward_scale_std + + def _reset_critic_cuda_graph(self) -> None: + self._cuda_graph_critic = None + self._cuda_graph_critic_static_inputs = None + self._cuda_graph_sac_static_packed_input = None + self._cuda_graph_sac_static_source_ptr = None + self._cuda_graph_critic_outputs = None + self._cuda_graph_critic_shapes = None + + def _materialize_capturable_critic_optimizer_state( + self, + inputs: dict[str, torch.Tensor], + ) -> None: + optimizer_lrs = [group["lr"] for group in self.critic_optimizer.param_groups] + optimizer_weight_decays = [ + group["weight_decay"] for group in self.critic_optimizer.param_groups + ] + cpu_rng_state = torch.random.get_rng_state() + cuda_rng_state = torch.cuda.get_rng_state() if self.device.type == "cuda" else None + try: + for group in self.critic_optimizer.param_groups: + group["lr"] = 0.0 + group["weight_decay"] = 0.0 + self._update_critic_capture_candidate(inputs) + finally: + torch.random.set_rng_state(cpu_rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) + for group, lr, weight_decay in zip( + self.critic_optimizer.param_groups, + optimizer_lrs, + optimizer_weight_decays, + strict=True, + ): + group["lr"] = lr + group["weight_decay"] = weight_decay + + self.critic_optimizer.zero_grad(set_to_none=True) + for state in self.critic_optimizer.state.values(): + step = state.get("step") + if isinstance(step, torch.Tensor): + step.zero_() + elif step is not None: + state["step"] = 0 + for name in ("exp_avg", "exp_avg_sq", "max_exp_avg_sq"): + tensor = state.get(name) + if isinstance(tensor, torch.Tensor): + tensor.zero_() + + def _capture_critic_cuda_graph(self, inputs: dict[str, torch.Tensor]) -> None: + self._cuda_graph_critic_shapes = self._critic_graph_input_shapes(inputs) + packed_source = inputs.get("sac_graph_packed_source") + if self.use_cuda_graph_critic_packed_staging and packed_source is not None: + self._cuda_graph_sac_static_packed_input = packed_source.detach().clone() + actor_shapes = self._actor_graph_input_shapes(inputs) + self._cuda_graph_critic_static_inputs = self._critic_graph_static_views_from_sac_packed( + self._cuda_graph_sac_static_packed_input, + self._cuda_graph_critic_shapes, + actor_shapes, + ) + else: + self._cuda_graph_sac_static_packed_input = None + self._cuda_graph_critic_static_inputs = { + key: inputs[key].detach().clone() for key in self._critic_graph_input_keys() + } + self._copy_critic_graph_inputs(inputs) + + graph = torch.cuda.CUDAGraph() + capture_stream = cast(torch.cuda.Stream, torch.cuda.Stream()) + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream), torch.cuda.graph(graph): + self._cuda_graph_critic_outputs = self._update_critic_capture_candidate( + self._cuda_graph_critic_static_inputs + ) + torch.cuda.current_stream().wait_stream(capture_stream) + torch.cuda.synchronize() + self._cuda_graph_critic = graph + + def _critic_graph_output_metrics(self, *, read_items: bool = True) -> dict[str, float]: + if not read_items: + return {} + assert self._cuda_graph_critic_outputs is not None + critic_loss, reward_scale_std = self._cuda_graph_critic_outputs + return { + "critic_loss": float(critic_loss.detach().cpu()), + "reward_scale_std": float(reward_scale_std.detach().cpu()), + } + + def update_critic_cuda_graph( + self, + batch: dict[str, torch.Tensor], + *, + read_metrics: bool = True, + ) -> dict[str, float]: + if not self.use_cuda_graph_critic: + return self.update_critic(batch) + if self.device.type != "cuda": + return self.update_critic(batch) + if self.scaler is not None or self.world_size > 1: + return self.update_critic(batch) + if not isinstance(self.obs_normalizer, nn.Identity): + return self.update_critic(batch) + + inputs = self._prepare_critic_graph_inputs(batch) + if self._cuda_graph_critic_shapes != self._critic_graph_input_shapes(inputs): + self._reset_critic_cuda_graph() + self._materialize_capturable_critic_optimizer_state(inputs) + self._capture_critic_cuda_graph(inputs) + self.critic_scheduler.step() + self.critic.normalize_parameters() + return self._critic_graph_output_metrics(read_items=read_metrics) + + assert self._cuda_graph_critic is not None + self._copy_critic_graph_inputs(inputs) + self._cuda_graph_critic.replay() + self.critic_scheduler.step() + self.critic.normalize_parameters() + return self._critic_graph_output_metrics(read_items=read_metrics) + + @staticmethod + def _actor_graph_input_keys() -> tuple[str, ...]: + return ("obs", "next_obs", "actions", "critic") + + def _actor_graph_input_shapes(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Size]: + return {key: inputs[key].shape for key in self._actor_graph_input_keys()} + + def _prepare_actor_graph_inputs( + self, + batch: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + obs = batch["obs"].to(self.device) + next_obs = batch["next_obs"].to(self.device) + expert_actions = batch["actions"].to(self.device) + critic_obs = batch["critic"].to(self.device) + prepared = { + "obs": self._maybe_normalize_obs(obs, update=False), + "next_obs": self._maybe_normalize_obs(next_obs, update=False), + "actions": expert_actions, + "critic": critic_obs, + } + if "sac_graph_packed_source" in batch: + prepared["sac_graph_packed_source"] = batch["sac_graph_packed_source"].to(self.device) + return prepared + + def _copy_actor_graph_inputs(self, inputs: dict[str, torch.Tensor]) -> None: + assert self._cuda_graph_actor_static_inputs is not None + packed_source = inputs.get("sac_graph_packed_source") + if packed_source is not None: + static_packed = self._cuda_graph_actor_static_packed_input + if static_packed is None: + static_packed = self._cuda_graph_sac_static_packed_input + if static_packed is not None: + source_ptr = int(packed_source.data_ptr()) + if ( + static_packed is not self._cuda_graph_sac_static_packed_input + or self._cuda_graph_sac_static_source_ptr != source_ptr + ): + static_packed.copy_(packed_source) + if static_packed is self._cuda_graph_sac_static_packed_input: + self._cuda_graph_sac_static_source_ptr = source_ptr + return + for key, tensor in self._cuda_graph_actor_static_inputs.items(): + tensor.copy_(inputs[key]) + + def _update_actor_capture_candidate( + self, + inputs: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + obs = inputs["obs"] + next_obs = inputs["next_obs"] + expert_actions = inputs["actions"] + critic_obs = inputs["critic"] + obs_all = torch.cat([obs, next_obs], dim=0) + + with self._autocast(): + actions_all, actor_info_all = self.actor(obs_all, training=True) + actions = actions_all.chunk(2, dim=0)[0] + log_probs = actor_info_all["log_prob"].chunk(2, dim=0)[0] + + self._set_requires_grad(self.critic, False) + q_values, _ = self.critic(critic_obs, actions, training=False) + self._set_requires_grad(self.critic, True) + temp_value = self.temperature() + actor_loss, entropy = self._actor_loss_tensors( + log_probs, q_values, actions, expert_actions, temp_value + ) + + self.actor_optimizer.zero_grad(set_to_none=True) + actor_loss.backward() + self.actor_optimizer.step() + + temp_loss = temp_value * (entropy - self.target_entropy) + self.temperature_optimizer.zero_grad(set_to_none=True) + temp_loss.backward() + self.temperature_optimizer.step() + return actor_loss, entropy, temp_value, temp_loss + + def _reset_actor_cuda_graph(self) -> None: + self._cuda_graph_actor = None + self._cuda_graph_actor_static_inputs = None + self._cuda_graph_actor_static_packed_input = None + self._cuda_graph_actor_outputs = None + self._cuda_graph_actor_shapes = None + + def _materialize_capturable_actor_optimizer_state( + self, + inputs: dict[str, torch.Tensor], + ) -> None: + optimizers = (self.actor_optimizer, self.temperature_optimizer) + optimizer_lrs = [ + [group["lr"] for group in optimizer.param_groups] for optimizer in optimizers + ] + optimizer_weight_decays = [ + [group["weight_decay"] for group in optimizer.param_groups] for optimizer in optimizers + ] + cpu_rng_state = torch.random.get_rng_state() + cuda_rng_state = torch.cuda.get_rng_state() if self.device.type == "cuda" else None + try: + for optimizer in optimizers: + for group in optimizer.param_groups: + group["lr"] = 0.0 + group["weight_decay"] = 0.0 + self._update_actor_capture_candidate(inputs) + finally: + torch.random.set_rng_state(cpu_rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) + for optimizer, lrs, weight_decays in zip( + optimizers, + optimizer_lrs, + optimizer_weight_decays, + strict=True, + ): + for group, lr, weight_decay in zip( + optimizer.param_groups, + lrs, + weight_decays, + strict=True, + ): + group["lr"] = lr + group["weight_decay"] = weight_decay + + for optimizer in optimizers: + optimizer.zero_grad(set_to_none=True) + for state in optimizer.state.values(): + step = state.get("step") + if isinstance(step, torch.Tensor): + step.zero_() + elif step is not None: + state["step"] = 0 + for name in ("exp_avg", "exp_avg_sq", "max_exp_avg_sq"): + tensor = state.get(name) + if isinstance(tensor, torch.Tensor): + tensor.zero_() + + def _capture_actor_cuda_graph(self, inputs: dict[str, torch.Tensor]) -> None: + self._cuda_graph_actor_shapes = self._actor_graph_input_shapes(inputs) + packed_source = inputs.get("sac_graph_packed_source") + if self.use_cuda_graph_actor_packed_staging and packed_source is not None: + if ( + self._cuda_graph_sac_static_packed_input is not None + and self._cuda_graph_sac_static_packed_input.shape == packed_source.shape + ): + self._cuda_graph_actor_static_packed_input = ( + self._cuda_graph_sac_static_packed_input + ) + else: + self._cuda_graph_actor_static_packed_input = packed_source.detach().clone() + self._cuda_graph_actor_static_inputs = self._actor_graph_static_views_from_sac_packed( + self._cuda_graph_actor_static_packed_input, + self._cuda_graph_actor_shapes, + ) + else: + self._cuda_graph_actor_static_packed_input = None + self._cuda_graph_actor_static_inputs = { + key: inputs[key].detach().clone() for key in self._actor_graph_input_keys() + } + self._copy_actor_graph_inputs(inputs) + + graph = torch.cuda.CUDAGraph() + capture_stream = cast(torch.cuda.Stream, torch.cuda.Stream()) + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream), torch.cuda.graph(graph): + self._cuda_graph_actor_outputs = self._update_actor_capture_candidate( + self._cuda_graph_actor_static_inputs + ) + torch.cuda.current_stream().wait_stream(capture_stream) + torch.cuda.synchronize() + self._cuda_graph_actor = graph + + def _actor_graph_output_metrics(self, *, read_items: bool = True) -> dict[str, float]: + if not read_items: + return {} + assert self._cuda_graph_actor_outputs is not None + actor_loss, entropy, temp_value, temp_loss = self._cuda_graph_actor_outputs + return { + "actor_loss": float(actor_loss.detach().cpu()), + "actor_entropy": float(entropy.detach().cpu()), + "temperature": float(temp_value.detach().cpu()), + "temperature_loss": float(temp_loss.detach().cpu()), + } + + def update_actor_cuda_graph( + self, + batch: dict[str, torch.Tensor], + *, + read_metrics: bool = True, + ) -> dict[str, float]: + if not self.use_cuda_graph_actor: + return self.update_actor(batch) + if self.device.type != "cuda": + return self.update_actor(batch) + if self.scaler is not None or self.world_size > 1: + return self.update_actor(batch) + if not isinstance(self.obs_normalizer, nn.Identity): + return self.update_actor(batch) + + inputs = self._prepare_actor_graph_inputs(batch) + if self._cuda_graph_actor_shapes != self._actor_graph_input_shapes(inputs): + self._reset_actor_cuda_graph() + self._materialize_capturable_actor_optimizer_state(inputs) + self._capture_actor_cuda_graph(inputs) + self.actor_scheduler.step() + self.temperature_scheduler.step() + self.actor.normalize_parameters() + return self._actor_graph_output_metrics(read_items=read_metrics) + + assert self._cuda_graph_actor is not None + self._copy_actor_graph_inputs(inputs) + self._cuda_graph_actor.replay() + self.actor_scheduler.step() + self.temperature_scheduler.step() + self.actor.normalize_parameters() + return self._actor_graph_output_metrics(read_items=read_metrics) + def update_critic(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: obs = batch["obs"].to(self.device) actions = batch["actions"].to(self.device) diff --git a/src/unilab/algos/torch/flash_sac/runner.py b/src/unilab/algos/torch/flash_sac/runner.py index 165257edf..ea7a189f7 100644 --- a/src/unilab/algos/torch/flash_sac/runner.py +++ b/src/unilab/algos/torch/flash_sac/runner.py @@ -53,6 +53,10 @@ def __init__( n_step: int = 1, amp_dtype: str = "auto", use_compile: bool = False, + use_cuda_graph_critic: bool = False, + use_cuda_graph_actor: bool = False, + use_cuda_graph_critic_packed_staging: bool = False, + use_cuda_graph_actor_packed_staging: bool = False, seed: int | None = None, trace_enabled: bool = False, trace_output_dir: str | None = None, @@ -110,6 +114,10 @@ def __init__( use_amp=use_amp, amp_dtype=amp_dtype, use_compile=use_compile, + use_cuda_graph_critic=use_cuda_graph_critic, + use_cuda_graph_actor=use_cuda_graph_actor, + use_cuda_graph_critic_packed_staging=use_cuda_graph_critic_packed_staging, + use_cuda_graph_actor_packed_staging=use_cuda_graph_actor_packed_staging, ) super().__init__( diff --git a/src/unilab/algos/torch/hora/sac_learner.py b/src/unilab/algos/torch/hora/sac_learner.py index 2968a1e9e..4275dab78 100644 --- a/src/unilab/algos/torch/hora/sac_learner.py +++ b/src/unilab/algos/torch/hora/sac_learner.py @@ -76,6 +76,8 @@ def __init__( symmetry_augmentation=None, **kwargs, ) + self.use_cuda_graph_critic = False + self.use_cuda_graph_actor = False self.priv_info_dim = int(priv_info_dim) self.actor = HoraSACActor( obs_dim=obs_dim, @@ -103,7 +105,9 @@ def _get_actions_and_log_probs_for_critic( self, actor_obs: torch.Tensor, critic_obs: torch.Tensor, + eps: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del eps priv_info = derive_priv_info_from_critic_obs( actor_obs, critic_obs, @@ -116,7 +120,9 @@ def _get_actions_and_log_probs_for_actor( self, actor_obs: torch.Tensor, critic_obs: torch.Tensor, + eps: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del eps priv_info = derive_priv_info_from_critic_obs( actor_obs, critic_obs, diff --git a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py b/src/unilab/algos/torch/offpolicy/double_buffer_runner.py index 73194c1dd..f68e0dc9f 100644 --- a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py +++ b/src/unilab/algos/torch/offpolicy/double_buffer_runner.py @@ -186,6 +186,26 @@ def learn( warn_if_over_budget, ) + graph_packed_staging_supported = self.algo_type == "sac" or ( + self.algo_type == "flashsac" + and bool(getattr(self.learner, "supports_cuda_graph_packed_staging", False)) + ) + use_critic_graph_packed_source = ( + graph_packed_staging_supported + and bool(getattr(self.learner, "use_cuda_graph_critic_packed_staging", False)) + and self.critic_obs_dim > 0 + ) + use_sac_graph_pack_layout = use_critic_graph_packed_source and bool( + getattr(self.learner, "use_cuda_graph_actor_packed_staging", False) + ) + use_critic_graph_packed_source = ( + use_critic_graph_packed_source and not use_sac_graph_pack_layout + ) + critic_graph_staging_width = ( + self.critic_obs_dim + self.action_dim + 1 + self.obs_dim + self.critic_obs_dim + 1 + 1 + if use_critic_graph_packed_source + else 0 + ) mem_est = estimate_offpolicy_bytes( num_envs=self.num_envs, replay_buffer_n=self.replay_buffer_n, @@ -194,6 +214,7 @@ def learn( critic_dim=self.critic_obs_dim, batch_size=self.batch_size, updates_per_step=self.updates_per_step, + critic_graph_staging_width=critic_graph_staging_width, ) warn_if_over_budget(mem_est, label=f"Off-policy ({self.algo_type})") raise_if_shared_memory_over_budget(mem_est, label=f"Off-policy ({self.algo_type})") @@ -219,10 +240,22 @@ def learn( collector_pack_request_queue = _SPAWN_CTX.Queue(maxsize=1) collector_pack_ready_queue = _SPAWN_CTX.Queue(maxsize=1) packed_width = int(replay_buffer._storage.shape[1]) + if use_sac_graph_pack_layout: + packed_width = int(replay_buffer.sac_graph_packed_width()) collector_pack_shared_slots = [ torch.empty((sample_count, packed_width), dtype=torch.float32).share_memory_() for _ in range(2) ] + collector_pack_critic_graph_shared_slots = None + if use_critic_graph_packed_source: + critic_graph_width = int(replay_buffer.critic_graph_packed_width()) + collector_pack_critic_graph_shared_slots = [ + torch.empty( + (sample_count, critic_graph_width), + dtype=torch.float32, + ).share_memory_() + for _ in range(2) + ] _verbose_output_dir: str | None = None if self.verbose_metrics: _vroot = Path(self.trace_output_dir) if self.trace_output_dir else Path(log_dir) @@ -239,6 +272,9 @@ def learn( collector_pack_request_queue=collector_pack_request_queue, collector_pack_ready_queue=collector_pack_ready_queue, collector_pack_shared_slots=collector_pack_shared_slots, + pack_layout="sac_graph" if use_sac_graph_pack_layout else "packed", + use_critic_graph_packed_source=use_critic_graph_packed_source, + collector_pack_critic_graph_shared_slots=collector_pack_critic_graph_shared_slots, ) self.replay_h2d_submitter = getattr( replay_pipeline, @@ -340,6 +376,10 @@ def learn( "collector_pack_ready_queue": collector_pack_ready_queue, "collector_pack_shared_slots": collector_pack_shared_slots, } + if use_critic_graph_packed_source: + collector_kwargs["collector_pack_critic_graph_shared_slots"] = ( + collector_pack_critic_graph_shared_slots + ) self._start_collector( target_fn=off_policy_collector_fn, kwargs={"stop_event": self._stop_event, **collector_kwargs}, @@ -575,9 +615,18 @@ def learn( s = update_idx * self.batch_size e = s + self.batch_size batch = {k: v[s:e] for k, v in large_batch.items()} + read_critic_graph_metrics = update_idx == self.updates_per_step - 1 _critic_ns = time.perf_counter_ns() - critic_metrics = learner.update_critic(batch) + if getattr(learner, "use_cuda_graph_critic", False) and hasattr( + learner, "update_critic_cuda_graph" + ): + critic_metrics = learner.update_critic_cuda_graph( + batch, + read_metrics=read_critic_graph_metrics, + ) + else: + critic_metrics = learner.update_critic(batch) if trace_recorder: trace_recorder.add_slice( "learner/update_critic", @@ -590,8 +639,18 @@ def learn( iter_metrics[k].append(v) if update_idx % self.policy_frequency == 0: + next_actor_update = update_idx + self.policy_frequency + read_actor_graph_metrics = next_actor_update >= self.updates_per_step _actor_ns = time.perf_counter_ns() - actor_metrics = learner.update_actor(batch) + if getattr(learner, "use_cuda_graph_actor", False) and hasattr( + learner, "update_actor_cuda_graph" + ): + actor_metrics = learner.update_actor_cuda_graph( + batch, + read_metrics=read_actor_graph_metrics, + ) + else: + actor_metrics = learner.update_actor(batch) if trace_recorder: trace_recorder.add_slice( "learner/update_actor", diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/torch/offpolicy/runner.py index e9bd7a5fa..737a9797a 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/torch/offpolicy/runner.py @@ -542,9 +542,18 @@ def learn( s = update_idx * self.batch_size e = s + self.batch_size batch = {k: v[s:e] for k, v in large_batch.items()} + read_critic_graph_metrics = update_idx == self.updates_per_step - 1 _critic_ns = time.perf_counter_ns() if trace_recorder else 0 - critic_metrics = learner.update_critic(batch) + if getattr(learner, "use_cuda_graph_critic", False) and hasattr( + learner, "update_critic_cuda_graph" + ): + critic_metrics = learner.update_critic_cuda_graph( + batch, + read_metrics=read_critic_graph_metrics, + ) + else: + critic_metrics = learner.update_critic(batch) if trace_recorder: trace_recorder.add_slice( "learner/update_critic", @@ -557,8 +566,18 @@ def learn( iter_metrics[k].append(v) if update_idx % self.policy_frequency == 0: + next_actor_update = update_idx + self.policy_frequency + read_actor_graph_metrics = next_actor_update >= self.updates_per_step _actor_ns = time.perf_counter_ns() if trace_recorder else 0 - actor_metrics = learner.update_actor(batch) + if getattr(learner, "use_cuda_graph_actor", False) and hasattr( + learner, "update_actor_cuda_graph" + ): + actor_metrics = learner.update_actor_cuda_graph( + batch, + read_metrics=read_actor_graph_metrics, + ) + else: + actor_metrics = learner.update_actor(batch) if trace_recorder: trace_recorder.add_slice( "learner/update_actor", diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index e64b9d699..5c53810c6 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -239,7 +239,12 @@ def _replay_write_exclude_ranges( return ranges -def _collector_pack_shared_batch(replay_buffer, request: dict, shared_slots) -> dict: +def _collector_pack_shared_batch( + replay_buffer, + request: dict, + shared_slots, + critic_graph_shared_slots=None, +) -> dict: tick_id = int(request["tick_id"]) rank = int(request.get("rank", 0)) world_size = int(request.get("world_size", 1)) @@ -278,7 +283,20 @@ def _collector_pack_shared_batch(replay_buffer, request: dict, shared_slots) -> if rank_shared_slots is None: raise RuntimeError("collector replay pack request is missing shared slots") dst = rank_shared_slots[shared_slot] - torch.index_select(replay_buffer._storage, 0, indices, out=dst) + pack_layout = str(request.get("pack_layout", "packed")) + if pack_layout == "sac_graph": + sampled = torch.index_select(replay_buffer._storage, 0, indices) + replay_buffer.pack_sac_graph_source(sampled, out=dst) + else: + torch.index_select(replay_buffer._storage, 0, indices, out=dst) + if bool(request.get("use_critic_graph_packed_source", False)): + rank_critic_slots = _ranked_entry(critic_graph_shared_slots, rank, world_size) + if rank_critic_slots is None: + raise RuntimeError("collector replay pack request is missing critic graph slots") + replay_buffer.pack_critic_graph_source( + dst, + out=rank_critic_slots[shared_slot], + ) pack_end_ns = time.perf_counter_ns() return { "tick_id": tick_id, @@ -291,7 +309,7 @@ def _collector_pack_shared_batch(replay_buffer, request: dict, shared_slots) -> "shared_slot": shared_slot, "target_gpu_slot": target_gpu_slot, "learner_hot_gpu_slot": learner_hot_gpu_slot, - "pack_layout": "packed", + "pack_layout": pack_layout, "pack_executor": "collector_thread", "pack_begin_ns": pack_begin_ns, "pack_end_ns": pack_end_ns, @@ -303,6 +321,7 @@ def _service_collector_pack_requests( request_queue, ready_queue, shared_slots, + critic_graph_shared_slots=None, trace_recorder=None, *, block_timeout: float = 0.0, @@ -325,7 +344,12 @@ def _service_collector_pack_requests( min_snapshot_ptr = int(request.get("min_snapshot_ptr", 0)) if int(replay_buffer.ptr[0]) < min_snapshot_ptr: return False, request - ready = _collector_pack_shared_batch(replay_buffer, request, shared_slots) + ready = _collector_pack_shared_batch( + replay_buffer, + request, + shared_slots, + critic_graph_shared_slots, + ) if trace_recorder: trace_recorder.add_slice( "collector/cpu_pack_sample_batch", @@ -362,6 +386,7 @@ def _drain_collector_pack_requests( request_queue, ready_queue, shared_slots, + critic_graph_shared_slots=None, trace_recorder=None, *, pending_request: dict | None = None, @@ -378,6 +403,7 @@ def _drain_collector_pack_requests( request_queue, ready_queue, shared_slots, + critic_graph_shared_slots, trace_recorder, block_timeout=0.0, pending_request=pending, @@ -396,6 +422,7 @@ def __init__( request_queue, ready_queue, shared_slots, + critic_graph_shared_slots=None, trace_recorder=None, *, stop_event=None, @@ -404,13 +431,17 @@ def __init__( self._request_queue = request_queue self._ready_queue = ready_queue self._shared_slots = shared_slots + self._critic_graph_shared_slots = critic_graph_shared_slots self._trace_recorder = trace_recorder self._stop_event = stop_event self._threads: list[threading.Thread] = [] self._started = False @staticmethod - def should_start(request_queue, ready_queue, shared_slots) -> bool: + def should_start( + request_queue, ready_queue, shared_slots, critic_graph_shared_slots=None + ) -> bool: + del critic_graph_shared_slots return ( isinstance(request_queue, list) and isinstance(ready_queue, list) @@ -437,6 +468,7 @@ def _rank_worker(self, rank: int, world_size: int) -> None: request_queue = self._request_queue[rank] ready_queue = self._ready_queue shared_slots = self._shared_slots + critic_graph_shared_slots = self._critic_graph_shared_slots pending_request = None while True: if self._stop_event is not None and self._stop_event.is_set(): @@ -446,6 +478,7 @@ def _rank_worker(self, rank: int, world_size: int) -> None: request_queue, ready_queue, shared_slots, + critic_graph_shared_slots, self._trace_recorder, block_timeout=0.001, pending_request=pending_request, @@ -489,6 +522,7 @@ def off_policy_collector_fn( collector_pack_request_queue=None, collector_pack_ready_queue=None, collector_pack_shared_slots=None, + collector_pack_critic_graph_shared_slots=None, nan_guard_cfg=None, **kwargs, ): @@ -528,6 +562,7 @@ def off_policy_collector_fn( collector_pack_request_queue=collector_pack_request_queue, collector_pack_ready_queue=collector_pack_ready_queue, collector_pack_shared_slots=collector_pack_shared_slots, + collector_pack_critic_graph_shared_slots=collector_pack_critic_graph_shared_slots, nan_guard_cfg=nan_guard_cfg, ) @@ -562,6 +597,7 @@ def _run_collector( collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + collector_pack_critic_graph_shared_slots=None, nan_guard_cfg=None, ): del learning_starts @@ -661,12 +697,14 @@ def _run_collector( collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + collector_pack_critic_graph_shared_slots, ): collector_pack_service = _CollectorPackService( replay_buffer, collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + collector_pack_critic_graph_shared_slots, trace_recorder, stop_event=stop_event, ) @@ -820,6 +858,7 @@ def _run_collector( collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + collector_pack_critic_graph_shared_slots, trace_recorder, pending_request=pending_collector_pack_request, ) @@ -872,6 +911,7 @@ def _run_collector( collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + collector_pack_critic_graph_shared_slots, trace_recorder, pending_request=pending_collector_pack_request, ) @@ -889,6 +929,7 @@ def _run_collector( collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + collector_pack_critic_graph_shared_slots, trace_recorder, pending_request=pending_collector_pack_request, ) diff --git a/src/unilab/base/augmentation.py b/src/unilab/base/augmentation.py index 4fdad70ad..b71bd958d 100644 --- a/src/unilab/base/augmentation.py +++ b/src/unilab/base/augmentation.py @@ -22,6 +22,13 @@ def augment_obs_and_actions( obs_group: str = "obs", ) -> tuple[torch.Tensor, torch.Tensor]: ... + def augment_obs( + self, + obs: torch.Tensor, + *, + obs_group: str = "obs", + ) -> torch.Tensor: ... + def mirror_obs( self, obs: torch.Tensor, diff --git a/src/unilab/demo.py b/src/unilab/demo.py index d341b3efd..c4fdc6ffc 100644 --- a/src/unilab/demo.py +++ b/src/unilab/demo.py @@ -300,7 +300,7 @@ def _run_teaser_demo() -> int: str(_repo_root() / "src" / "unilab" / "tools" / "render_teaser.py"), ] env = os.environ.copy() - env.setdefault("UV_PROJECT_ENVIRONMENT", str(_repo_root() / ".venv")) + env["UV_PROJECT_ENVIRONMENT"] = str(_repo_root() / ".venv") return subprocess.run(command, check=False, env=env).returncode from unilab.tools.render_teaser import main as render_teaser_main @@ -325,7 +325,7 @@ def run_demo(*, demo_name: str, refresh: bool = False, device: str | None = None demo_name=demo_name, checkpoint_path=checkpoint_path, device=device ) env = os.environ.copy() - env.setdefault("UV_PROJECT_ENVIRONMENT", str(_repo_root() / ".venv")) + env["UV_PROJECT_ENVIRONMENT"] = str(_repo_root() / ".venv") returncode = subprocess.run(command, check=False, env=env).returncode if returncode == 0: print(f"Demo finished: {demo_name} (algo={spec.algo}, task={spec.task}, sim={spec.sim})") diff --git a/src/unilab/envs/locomotion/g1/symmetry.py b/src/unilab/envs/locomotion/g1/symmetry.py index 9e1c06647..e1fa33cc7 100644 --- a/src/unilab/envs/locomotion/g1/symmetry.py +++ b/src/unilab/envs/locomotion/g1/symmetry.py @@ -12,9 +12,8 @@ @dataclass(frozen=True) class _ObsGroupTransform: dim: int - flip_mask: torch.Tensor joint_map: torch.Tensor - joint_sign: torch.Tensor + sign_mask: torch.Tensor class G1SymmetryAugmentation(SymmetryAugmentation): @@ -120,9 +119,8 @@ def _build_obs_group_transform( return _ObsGroupTransform( dim=obs_dim, - flip_mask=flip_mask, joint_map=joint_map, - joint_sign=joint_sign, + sign_mask=flip_mask * joint_sign, ) @staticmethod @@ -141,7 +139,10 @@ def mirror_obs(self, obs: torch.Tensor, *, obs_group: str = "obs") -> torch.Tens raise ValueError( f"Symmetry obs group {obs_group!r} expects dim {transform.dim}, got {obs.shape[-1]}" ) - return obs[..., transform.joint_map] * transform.flip_mask * transform.joint_sign + return obs[..., transform.joint_map] * transform.sign_mask + + def augment_obs(self, obs: torch.Tensor, *, obs_group: str = "obs") -> torch.Tensor: + return torch.cat([obs, self.mirror_obs(obs, obs_group=obs_group)], dim=0) def augment_obs_and_actions( self, @@ -150,7 +151,7 @@ def augment_obs_and_actions( *, obs_group: str = "obs", ) -> tuple[torch.Tensor, torch.Tensor]: - return torch.cat([obs, self.mirror_obs(obs, obs_group=obs_group)], dim=0), torch.cat( + return self.augment_obs(obs, obs_group=obs_group), torch.cat( [actions, self.mirror_action(actions)], dim=0, ) diff --git a/src/unilab/ipc/memory_budget.py b/src/unilab/ipc/memory_budget.py index 9a23b78b1..80e567867 100644 --- a/src/unilab/ipc/memory_budget.py +++ b/src/unilab/ipc/memory_budget.py @@ -19,6 +19,7 @@ def estimate_offpolicy_bytes( critic_dim: int, batch_size: int, updates_per_step: int, + critic_graph_staging_width: int = 0, ) -> dict[str, int | str]: """Estimate memory for off-policy replay buffer + double-buffer slots.""" row_width = 2 * obs_dim + action_dim + 3 + 2 * critic_dim @@ -27,17 +28,21 @@ def estimate_offpolicy_bytes( sample_count = batch_size * updates_per_step slot_bytes = sample_count * row_width * 4 * 2 + critic_graph_staging_bytes = sample_count * int(critic_graph_staging_width) * 4 * 2 - total = replay_bytes + slot_bytes + total = replay_bytes + slot_bytes + critic_graph_staging_bytes return { "replay_buffer": replay_bytes, "double_buffer_slots": slot_bytes, + "critic_graph_staging_slots": critic_graph_staging_bytes, "total": total, "breakdown": ( f"Replay: {replay_bytes / 1024**2:.0f} MB " f"({num_envs} envs × {replay_buffer_n} steps × {row_width} cols × 4B)\n" f" Double-buffer: {slot_bytes / 1024**2:.0f} MB " f"({sample_count} samples × {row_width} cols × 4B × 2 slots)\n" + f" Critic graph staging: {critic_graph_staging_bytes / 1024**2:.0f} MB " + f"({sample_count} samples × {int(critic_graph_staging_width)} cols × 4B × 2 slots)\n" " Excludes MuJoCo BatchEnvPool/native allocations, CUDA pinned/shared " "registration, and driver memory." ), diff --git a/src/unilab/ipc/replay_buffer.py b/src/unilab/ipc/replay_buffer.py index 9aa861019..a18e0143b 100644 --- a/src/unilab/ipc/replay_buffer.py +++ b/src/unilab/ipc/replay_buffer.py @@ -66,6 +66,81 @@ def _init_packed_storage( self._ncritic_sl = slice(c, c + critic_dim) c += critic_dim + def critic_graph_packed_width(self) -> int: + """Return graph-order packed width for FastSAC critic graph inputs.""" + if self._critic_dim <= 0: + raise RuntimeError("critic_graph_packed_source requires critic replay storage") + return self._critic_dim + self._action_dim + 1 + self._obs_dim + self._critic_dim + 1 + 1 + + def sac_graph_packed_width(self) -> int: + """Return graph-union packed width for one-H2D FastSAC graph staging.""" + if self._critic_dim <= 0: + raise RuntimeError("sac_graph_packed_source requires critic replay storage") + return int(self._storage.shape[1]) + + def pack_critic_graph_source( + self, + packed_batch: torch.Tensor, + *, + out: torch.Tensor, + ) -> torch.Tensor: + """Pack replay rows in the exact input order consumed by critic CUDA graph.""" + if self._critic_dim <= 0: + raise RuntimeError("critic_graph_packed_source requires critic replay storage") + expected_shape = (packed_batch.shape[0], self.critic_graph_packed_width()) + if tuple(out.shape) != expected_shape: + raise ValueError( + "critic_graph_packed_source output shape mismatch: " + f"expected {expected_shape}, got {tuple(out.shape)}" + ) + offset = 0 + fields = ( + packed_batch[:, self._critic_sl], + packed_batch[:, self._act_sl], + packed_batch[:, self._rew_col : self._rew_col + 1], + packed_batch[:, self._nobs_sl], + packed_batch[:, self._ncritic_sl], + packed_batch[:, self._done_col : self._done_col + 1], + packed_batch[:, self._trunc_col : self._trunc_col + 1], + ) + for field in fields: + width = int(field.shape[1]) + out[:, offset : offset + width].copy_(field) + offset += width + return out + + def pack_sac_graph_source( + self, + packed_batch: torch.Tensor, + *, + out: torch.Tensor, + ) -> torch.Tensor: + """Pack replay rows once in a layout friendly to SAC actor/critic graphs.""" + if self._critic_dim <= 0: + raise RuntimeError("sac_graph_packed_source requires critic replay storage") + expected_shape = (packed_batch.shape[0], self.sac_graph_packed_width()) + if tuple(out.shape) != expected_shape: + raise ValueError( + "sac_graph_packed_source output shape mismatch: " + f"expected {expected_shape}, got {tuple(out.shape)}" + ) + offset = 0 + fields = ( + packed_batch[:, self._obs_sl], + packed_batch[:, self._critic_sl], + packed_batch[:, self._act_sl], + packed_batch[:, self._rew_col : self._rew_col + 1], + packed_batch[:, self._nobs_sl], + packed_batch[:, self._ncritic_sl], + packed_batch[:, self._done_col : self._done_col + 1], + packed_batch[:, self._trunc_col : self._trunc_col + 1], + ) + for field in fields: + width = int(field.shape[1]) + out[:, offset : offset + width].copy_(field) + offset += width + return out + def __getstate__(self) -> dict: """Custom pickle support. diff --git a/src/unilab/ipc/replay_pipelines/cpu_pinned_double_buffer.py b/src/unilab/ipc/replay_pipelines/cpu_pinned_double_buffer.py index 09488f04a..26db1a77a 100644 --- a/src/unilab/ipc/replay_pipelines/cpu_pinned_double_buffer.py +++ b/src/unilab/ipc/replay_pipelines/cpu_pinned_double_buffer.py @@ -43,6 +43,9 @@ def __init__( collector_pack_request_queue=None, collector_pack_ready_queue=None, collector_pack_shared_slots=None, + pack_layout: str = "packed", + use_critic_graph_packed_source: bool = False, + collector_pack_critic_graph_shared_slots=None, ) -> None: self._replay_buffer = replay_buffer self._device = torch.device(device) @@ -53,7 +56,7 @@ def __init__( self._trace_cuda_events = bool(trace_cuda_events) and self._device_type == "cuda" self._verbose = bool(verbose) self._verbose_output_dir = verbose_output_dir if self._verbose else None - self._pack_layout = "packed" + self._pack_layout = str(pack_layout) self._pack_executor = "collector_thread" self._ring_depth = 2 self._transfer_backend = build_replay_transfer_backend( @@ -85,8 +88,32 @@ def __init__( self._collector_pack_request_queue = collector_pack_request_queue self._collector_pack_ready_queue = collector_pack_ready_queue self._collector_pack_shared_slots = collector_pack_shared_slots + if self._pack_layout not in {"packed", "sac_graph"}: + raise ValueError( + "CPUPinnedDoubleBufferReplayPipeline pack_layout must be packed or sac_graph" + ) + self._use_critic_graph_packed_source = ( + bool(use_critic_graph_packed_source) and self._pack_layout != "sac_graph" + ) + self._collector_pack_critic_graph_shared_slots = collector_pack_critic_graph_shared_slots self._fields: Dict[str, tuple[torch.Tensor, int]] = {} - self._packed_width = int(replay_buffer._storage.shape[1]) + self._packed_width = ( + int(replay_buffer.sac_graph_packed_width()) + if self._pack_layout == "sac_graph" + else int(replay_buffer._storage.shape[1]) + ) + self._critic_graph_packed_width = ( + int(replay_buffer.critic_graph_packed_width()) + if self._use_critic_graph_packed_source + else 0 + ) + if ( + self._use_critic_graph_packed_source + and collector_pack_critic_graph_shared_slots is None + ): + raise ValueError( + "critic graph packed source requires collector critic graph shared slots" + ) self._host_packed: list[torch.Tensor] = [] self._register_collector_shared_slots() self._gpu_packed = self._transfer_backend.allocate_device_slots( @@ -94,6 +121,15 @@ def __init__( shape=(self._sample_count, self._packed_width), dtype=torch.float32, ) + self._gpu_critic_graph_packed: list[torch.Tensor] = ( + self._transfer_backend.allocate_device_slots( + count=self._ring_depth, + shape=(self._sample_count, self._critic_graph_packed_width), + dtype=torch.float32, + ) + if self._use_critic_graph_packed_source + else [] + ) self._host: list[Dict[str, torch.Tensor]] = [] self._gpu: list[Dict[str, torch.Tensor]] = [] @@ -139,6 +175,11 @@ def transfer_manifest(self) -> dict[str, object]: def _register_collector_shared_slots(self) -> None: assert self._collector_pack_shared_slots is not None self._transfer_backend.register_host_slots(self._collector_pack_shared_slots) + if self._use_critic_graph_packed_source: + assert self._collector_pack_critic_graph_shared_slots is not None + self._transfer_backend.register_host_slots( + self._collector_pack_critic_graph_shared_slots + ) self._host_pinned = self._transfer_backend.host_pinned self._direct_pinned_shared = self._transfer_backend.direct_pinned_shared @@ -151,8 +192,41 @@ def _packed_h2d_source(self, slot: int) -> torch.Tensor: assert self._collector_pack_shared_slots is not None return self._collector_pack_shared_slots[slot] + def _critic_graph_h2d_source(self, slot: int) -> torch.Tensor: + assert self._collector_pack_critic_graph_shared_slots is not None + return self._collector_pack_critic_graph_shared_slots[slot] + def _packed_batch_view(self, packed: torch.Tensor) -> Dict[str, torch.Tensor]: rb = self._replay_buffer + if self._pack_layout == "sac_graph": + c = 0 + obs_sl = slice(c, c + rb._obs_dim) + c += rb._obs_dim + critic_sl = slice(c, c + rb._critic_dim) + c += rb._critic_dim + act_sl = slice(c, c + rb._action_dim) + c += rb._action_dim + rew_col = c + c += 1 + nobs_sl = slice(c, c + rb._obs_dim) + c += rb._obs_dim + ncritic_sl = slice(c, c + rb._critic_dim) + c += rb._critic_dim + done_col = c + c += 1 + trunc_col = c + batch = { + "obs": packed[:, obs_sl], + "next_obs": packed[:, nobs_sl], + "actions": packed[:, act_sl], + "rewards": packed[:, rew_col], + "dones": packed[:, done_col], + "truncated": packed[:, trunc_col], + "critic": packed[:, critic_sl], + "next_critic": packed[:, ncritic_sl], + "sac_graph_packed_source": packed, + } + return batch batch = { "obs": packed[:, rb._obs_sl], "next_obs": packed[:, rb._nobs_sl], @@ -166,6 +240,12 @@ def _packed_batch_view(self, packed: torch.Tensor) -> Dict[str, torch.Tensor]: batch["next_critic"] = packed[:, rb._ncritic_sl] return batch + def _large_batch_view(self, slot: int) -> Dict[str, torch.Tensor]: + batch = self._packed_batch_view(self._gpu_packed[slot]) + if self._use_critic_graph_packed_source: + batch["critic_graph_packed_source"] = self._gpu_critic_graph_packed[slot] + return batch + # -- snapshot / H2D -------------------------------------------------------- def _snapshot(self) -> tuple[int, int]: @@ -185,6 +265,23 @@ def _submit_h2d(self, slot: int, metadata: ReplayTickMetadata | None = None) -> pack_executor=self._pack_executor, ) + def _submit_critic_graph_h2d( + self, + slot: int, + metadata: ReplayTickMetadata | None = None, + ) -> float: + return self._transfer_backend.submit_h2d( + slot=slot, + dst=self._gpu_critic_graph_packed[slot], + src=self._critic_graph_h2d_source(slot), + metadata=metadata, + trace_recorder=self._trace_recorder, + trace_cuda_events=False, + h2d_bytes=self._critic_graph_h2d_bytes(), + pack_layout="critic_graph_packed", + pack_executor=self._pack_executor, + ) + def _submit_collector_packed_h2d(self, ready: dict) -> ReplayTickMetadata: metadata = ReplayTickMetadata( tick_id=int(ready["tick_id"]), @@ -210,6 +307,8 @@ def _submit_device_transfer(self, metadata: ReplayTickMetadata) -> None: assert shared_slot is not None h2d_submit_ns = time.perf_counter_ns() self.last_incremental_h2d_time_s = self._submit_h2d(slot, metadata) + if self._use_critic_graph_packed_source: + self.last_incremental_h2d_time_s += self._submit_critic_graph_h2d(slot, metadata) if self._trace_recorder is not None: self._trace_recorder.add_slice( "replay_pipeline/batch_h2d_submit", @@ -224,6 +323,11 @@ def _submit_device_transfer(self, metadata: ReplayTickMetadata) -> None: "pack_executor": self._pack_executor, "h2d_submitter": self._h2d_submitter, "h2d_bytes": self._h2d_bytes(), + "critic_graph_h2d_bytes": ( + self._critic_graph_h2d_bytes() + if self._use_critic_graph_packed_source + else 0 + ), "h2d_submitted": True, "pinned_memory": self._host_pinned, "direct_pinned_shared": self._direct_pinned_shared, @@ -274,6 +378,10 @@ def _h2d_bytes(self) -> int: source = self._packed_h2d_source(0) return int(source.numel() * source.element_size()) + def _critic_graph_h2d_bytes(self) -> int: + source = self._critic_graph_h2d_source(0) + return int(source.numel() * source.element_size()) + def _clear_ready(self, slot: int) -> None: self._transfer_backend.clear_ready(slot) @@ -349,6 +457,7 @@ def start_prepare( "target_gpu_slot": slot, "pack_layout": self._pack_layout, "pack_executor": self._pack_executor, + "use_critic_graph_packed_source": self._use_critic_graph_packed_source, } if self._trace_recorder is not None: _req_ns = time.perf_counter_ns() @@ -416,7 +525,7 @@ def sample_large_batch(self, tick_id: int, sample_count: int) -> Dict[str, torch f"Hot batch tick {self._hot_metadata.tick_id} does not match " f"requested tick {tick_id}" ) - return self._packed_batch_view(self._gpu_packed[self._hot]) + return self._large_batch_view(self._hot) if not self._has_hot_batch: if not self.batch_ready(tick_id, sample_count): self.wait_until_ready(tick_id, sample_count) @@ -466,7 +575,7 @@ def sample_large_batch(self, tick_id: int, sample_count: int) -> Dict[str, torch self._prepared_metadata = None self._prepare_tick_id = None self._prepare_state = "idle" - return self._packed_batch_view(self._gpu_packed[self._hot]) + return self._large_batch_view(self._hot) def after_tick(self) -> None: self._has_hot_batch = False @@ -504,3 +613,5 @@ def close(self) -> None: self._host_packed.clear() if hasattr(self, "_gpu_packed"): self._gpu_packed.clear() + if hasattr(self, "_gpu_critic_graph_packed"): + self._gpu_critic_graph_packed.clear() diff --git a/src/unilab/structured_configs.py b/src/unilab/structured_configs.py index 8e8d5a9f7..b50ed5f98 100644 --- a/src/unilab/structured_configs.py +++ b/src/unilab/structured_configs.py @@ -27,6 +27,10 @@ class SACAlgoParams: max_grad_norm: float = 0.0 amp_dtype: str = "auto" use_compile: bool = True + use_cuda_graph_critic: bool = False + use_cuda_graph_actor: bool = False + use_cuda_graph_critic_packed_staging: bool = False + use_cuda_graph_actor_packed_staging: bool = False @dataclass @@ -128,6 +132,10 @@ class FlashSACAlgoParams: n_step: int = 1 amp_dtype: str = "auto" use_compile: bool = True + use_cuda_graph_critic: bool = False + use_cuda_graph_actor: bool = False + use_cuda_graph_critic_packed_staging: bool = False + use_cuda_graph_actor_packed_staging: bool = False @dataclass diff --git a/tests/algos/test_fast_sac_compile.py b/tests/algos/test_fast_sac_compile.py index 4141939a2..298a6ff85 100644 --- a/tests/algos/test_fast_sac_compile.py +++ b/tests/algos/test_fast_sac_compile.py @@ -6,7 +6,67 @@ import pytest import torch -from unilab.algos.torch.fast_sac.learner import FastSACLearner +from unilab.algos.torch.fast_sac.learner import FastSACLearner, SACActor + + +def _small_fast_sac_learner(*, use_autotune: bool = True) -> FastSACLearner: + return FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=use_autotune, + max_grad_norm=0.0, + ) + + +def _small_offpolicy_batch(batch_size: int = 4) -> dict[str, torch.Tensor]: + return { + "obs": torch.linspace(-0.4, 0.7, steps=batch_size * 4).view(batch_size, 4), + "critic": torch.linspace(-0.2, 0.9, steps=batch_size * 5).view(batch_size, 5), + "actions": torch.linspace(-0.5, 0.5, steps=batch_size * 2).view(batch_size, 2), + "rewards": torch.linspace(-0.3, 0.6, steps=batch_size), + "next_obs": torch.linspace(0.1, 1.2, steps=batch_size * 4).view(batch_size, 4), + "next_critic": torch.linspace(-0.7, 0.4, steps=batch_size * 5).view(batch_size, 5), + "dones": torch.tensor([0.0, 1.0, 0.0, 1.0]), + "truncated": torch.tensor([0.0, 1.0, 0.0, 0.0]), + } + + +_CRITIC_GRAPH_INPUT_KEYS = ( + "critic", + "actions", + "rewards", + "next_obs", + "next_critic", + "dones", + "truncated", +) +_ACTOR_GRAPH_INPUT_KEYS = ("obs", "critic") + + +def _critic_graph_static_inputs(batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {key: batch[key].detach().clone() for key in _CRITIC_GRAPH_INPUT_KEYS} + + +def _actor_graph_static_inputs(batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {key: batch[key].detach().clone() for key in _ACTOR_GRAPH_INPUT_KEYS} + + +def _critic_graph_packed_source(batch: dict[str, torch.Tensor]) -> torch.Tensor: + return torch.cat( + [batch[key].reshape(batch[key].shape[0], -1) for key in _CRITIC_GRAPH_INPUT_KEYS], dim=1 + ) + + +def _sac_graph_packed_source(batch: dict[str, torch.Tensor]) -> torch.Tensor: + keys = ("obs", "critic", "actions", "rewards", "next_obs", "next_critic", "dones", "truncated") + return torch.cat([batch[key].reshape(batch[key].shape[0], -1) for key in keys], dim=1) def test_fast_sac_compile_targets_training_hot_paths(monkeypatch) -> None: @@ -45,6 +105,216 @@ def fake_compile(fn: Callable, **kwargs): ] +def test_fast_sac_graph_critic_skips_compiling_critic_loss(monkeypatch) -> None: + calls: list[tuple[str, dict[str, Any]]] = [] + + def fake_compile(fn: Callable, **kwargs): + calls.append((fn.__qualname__, kwargs)) + return fn + + learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_cuda_graph_critic=True, + ) + learner.device = "cuda" + learner.use_cuda_graph_critic = True + monkeypatch.setattr(torch, "compile", fake_compile) + + learner._compile_training_methods() + + assert calls == [ + ( + "FastSACLearner._actor_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ] + + +def test_fast_sac_graph_actor_skips_compiling_actor_loss(monkeypatch) -> None: + calls: list[tuple[str, dict[str, Any]]] = [] + + def fake_compile(fn: Callable, **kwargs): + calls.append((fn.__qualname__, kwargs)) + return fn + + learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_cuda_graph_actor=True, + ) + learner.device = "cuda" + learner.use_cuda_graph_actor = True + monkeypatch.setattr(torch, "compile", fake_compile) + + learner._compile_training_methods() + + assert calls == [ + ( + "FastSACLearner._critic_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ] + + +def test_fast_sac_cuda_adamw_optimizers_are_capture_ready(monkeypatch) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA-only optimizer kwargs require a CUDA-enabled torch build") + + calls: list[dict[str, Any]] = [] + + class _FakeAdamW: + def __init__(self, _params, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(torch.optim, "AdamW", _FakeAdamW) + + FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cuda", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_compile=False, + ) + + assert len(calls) == 3 + assert all(call["fused"] for call in calls) + assert all(call["capturable"] for call in calls) + + +def test_fast_sac_cpu_adamw_optimizers_keep_default_capturability(monkeypatch) -> None: + calls: list[dict[str, Any]] = [] + + class _FakeAdamW: + def __init__(self, _params, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(torch.optim, "AdamW", _FakeAdamW) + + FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_compile=False, + ) + + assert len(calls) == 3 + assert not any(call["fused"] for call in calls) + assert all("capturable" not in call for call in calls) + + +def test_fast_sac_cuda_graph_critic_is_opt_in_and_cuda_only() -> None: + cpu_learner = _small_fast_sac_learner() + assert not cpu_learner.use_cuda_graph_critic + + if not torch.cuda.is_available(): + pytest.skip("CUDA graph opt-in requires a CUDA-enabled torch build") + + cuda_learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cuda", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_compile=False, + use_amp=False, + use_cuda_graph_critic=True, + ) + assert cuda_learner.use_cuda_graph_critic + + +def test_fast_sac_critic_packed_staging_disabled_when_fp16_scaler_active( + monkeypatch, +) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA graph staging with fp16 scaler requires a CUDA-enabled torch build") + + monkeypatch.setattr( + FastSACLearner, + "_should_use_grad_scaler", + staticmethod(lambda use_amp, device_type, amp_dtype: True), + ) + learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cuda", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_amp=True, + amp_dtype="fp16", + use_cuda_graph_critic=True, + use_cuda_graph_critic_packed_staging=True, + ) + + assert learner.use_cuda_graph_critic + assert learner.scaler is not None + assert not learner.use_cuda_graph_critic_packed_staging + + +def test_fast_sac_cuda_graph_actor_is_opt_in_and_cuda_only() -> None: + cpu_learner = _small_fast_sac_learner() + assert not cpu_learner.use_cuda_graph_actor + + if not torch.cuda.is_available(): + pytest.skip("CUDA graph opt-in requires a CUDA-enabled torch build") + + cuda_learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cuda", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_compile=False, + use_amp=False, + use_cuda_graph_actor=True, + ) + assert cuda_learner.use_cuda_graph_actor + + def test_fast_sac_amp_dtype_resolution_and_scaler_rules() -> None: assert FastSACLearner._resolve_amp_dtype("auto", "cuda") is torch.bfloat16 assert FastSACLearner._resolve_amp_dtype("auto", "xpu") is torch.bfloat16 @@ -58,3 +328,621 @@ def test_fast_sac_amp_dtype_resolution_and_scaler_rules() -> None: with pytest.raises(ValueError, match="amp_dtype"): FastSACLearner._resolve_amp_dtype("tf32", "cuda") + + +def test_fast_sac_alpha_loss_helper_matches_reference_value_and_grad() -> None: + learner = FastSACLearner( + obs_dim=4, + action_dim=3, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=True, + ) + next_log_probs = torch.tensor([-1.25, -0.5, 0.25, 1.5], dtype=torch.float32) + learner.target_entropy = -1.75 + learner.log_alpha.data.fill_(-2.0) + + reference_log_alpha = learner.log_alpha.detach().clone().requires_grad_(True) + reference_loss = (-reference_log_alpha.exp() * (next_log_probs + learner.target_entropy)).mean() + reference_loss.backward() + + learner.log_alpha.grad = None + alpha_loss = learner._alpha_loss_tensor(next_log_probs) + alpha_loss.backward() + + assert torch.allclose(alpha_loss.detach(), reference_loss.detach()) + assert learner.log_alpha.grad is not None + assert reference_log_alpha.grad is not None + assert torch.allclose(learner.log_alpha.grad, reference_log_alpha.grad) + assert not next_log_probs.requires_grad + + +def test_sac_actor_tensor_gaussian_sampling_matches_normal_reference() -> None: + actor = SACActor( + obs_dim=4, + action_dim=3, + hidden_dim=12, + use_layer_norm=False, + action_scale=torch.tensor([0.5, 1.5, 2.0]), + action_bias=torch.tensor([-0.25, 0.0, 0.75]), + ) + obs = torch.tensor( + [ + [-1.0, -0.25, 0.5, 1.25], + [0.25, 0.5, -0.75, 1.0], + ], + dtype=torch.float32, + ) + + _, mean, log_std = actor(obs) + std = log_std.exp() + eps = torch.tensor( + [ + [-0.5, 0.25, 1.0], + [1.5, -1.0, 0.0], + ], + dtype=torch.float32, + ) + raw_action = mean + std * eps + dist = torch.distributions.Normal(mean, std) + tanh_action = torch.tanh(raw_action) + expected_action = tanh_action * actor.action_scale + actor.action_bias + expected_log_prob = dist.log_prob(raw_action) + expected_log_prob -= torch.log(1 - tanh_action.pow(2) + 1e-6) + expected_log_prob -= torch.log(actor.action_scale + 1e-6) + expected_log_prob = expected_log_prob.sum(1) + + action, log_prob = actor._sample_action_and_log_prob(mean, log_std, eps=eps) + + torch.testing.assert_close(action, expected_action) + torch.testing.assert_close(log_prob, expected_log_prob) + + +def test_sac_actor_tensor_gaussian_sampling_matches_normal_without_tanh() -> None: + actor = SACActor(obs_dim=2, action_dim=3, hidden_dim=12, use_layer_norm=False, use_tanh=False) + mean = torch.tensor( + [[-0.5, 0.25, 1.0], [1.5, -1.0, 0.0]], + dtype=torch.float32, + requires_grad=True, + ) + log_std = torch.tensor( + [[-1.0, -0.25, 0.5], [0.0, -0.75, 0.25]], + dtype=torch.float32, + requires_grad=True, + ) + eps = torch.tensor( + [[0.25, -1.5, 0.75], [-0.5, 1.0, 1.5]], + dtype=torch.float32, + ) + + action, log_prob = actor._sample_action_and_log_prob(mean, log_std, eps=eps) + + reference_mean = mean.detach().clone().requires_grad_(True) + reference_log_std = log_std.detach().clone().requires_grad_(True) + reference_std = reference_log_std.exp() + reference_raw_action = reference_mean + reference_std * eps + reference_dist = torch.distributions.Normal(reference_mean, reference_std) + expected_log_prob = reference_dist.log_prob(reference_raw_action).sum(1) + + torch.testing.assert_close(action, reference_raw_action) + torch.testing.assert_close(log_prob, expected_log_prob) + + loss = (action + log_prob.unsqueeze(1)).sum() + reference_loss = (reference_raw_action + expected_log_prob.unsqueeze(1)).sum() + loss.backward() + reference_loss.backward() + + assert mean.grad is not None + assert log_std.grad is not None + assert reference_mean.grad is not None + assert reference_log_std.grad is not None + torch.testing.assert_close(mean.grad, reference_mean.grad) + torch.testing.assert_close(log_std.grad, reference_log_std.grad) + + +def test_fast_sac_capture_candidate_matches_public_critic_update_for_finite_loss() -> None: + public_learner = _small_fast_sac_learner() + capture_learner = _small_fast_sac_learner() + capture_learner.actor.load_state_dict(public_learner.actor.state_dict()) + capture_learner.qnet.load_state_dict(public_learner.qnet.state_dict()) + capture_learner.qnet_target.load_state_dict(public_learner.qnet_target.state_dict()) + capture_learner.log_alpha.data.copy_(public_learner.log_alpha.data) + batch = _small_offpolicy_batch() + + torch.manual_seed(2024) + public_metrics = public_learner.update_critic(batch) + + torch.manual_seed(2024) + capture_outputs = capture_learner._update_critic_capture_candidate( + batch["critic"], + batch["actions"], + batch["rewards"], + batch["next_obs"], + batch["next_critic"], + batch["dones"], + batch["truncated"], + ) + capture_metrics = { + "qf_loss": capture_outputs[0].item(), + "critic_grad_norm": capture_outputs[1].item(), + "target_q_max": capture_outputs[2].item(), + "target_q_min": capture_outputs[3].item(), + "alpha_loss": capture_outputs[4].item(), + "alpha": capture_outputs[5].item(), + } + + assert public_metrics.keys() == capture_metrics.keys() + for key, public_value in public_metrics.items(): + assert capture_metrics[key] == pytest.approx(public_value) + for public_param, capture_param in zip( + public_learner.qnet.parameters(), + capture_learner.qnet.parameters(), + strict=True, + ): + torch.testing.assert_close(capture_param, public_param) + torch.testing.assert_close(capture_learner.log_alpha, public_learner.log_alpha) + + +def test_fast_sac_capture_candidate_matches_public_actor_update_for_finite_loss() -> None: + public_learner = _small_fast_sac_learner() + capture_learner = _small_fast_sac_learner() + capture_learner.actor.load_state_dict(public_learner.actor.state_dict()) + capture_learner.qnet.load_state_dict(public_learner.qnet.state_dict()) + capture_learner.qnet_target.load_state_dict(public_learner.qnet_target.state_dict()) + capture_learner.log_alpha.data.copy_(public_learner.log_alpha.data) + batch = _small_offpolicy_batch() + + torch.manual_seed(2025) + public_metrics = public_learner.update_actor(batch) + + torch.manual_seed(2025) + capture_outputs = capture_learner._update_actor_capture_candidate( + batch["obs"], + batch["critic"], + ) + capture_metrics = { + "actor_loss": capture_outputs[0].item(), + "actor_grad_norm": capture_outputs[1].item(), + "policy_entropy": capture_outputs[2].item(), + "action_std": capture_outputs[3].item(), + } + + assert public_metrics.keys() == capture_metrics.keys() + for key, public_value in public_metrics.items(): + assert capture_metrics[key] == pytest.approx(public_value) + for public_param, capture_param in zip( + public_learner.actor.parameters(), + capture_learner.actor.parameters(), + strict=True, + ): + torch.testing.assert_close(capture_param, public_param) + + +def test_fast_sac_cuda_graph_state_materialization_preserves_cpu_rng_state() -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + torch.manual_seed(12345) + expected_rng_state = torch.random.get_rng_state() + + learner._materialize_capturable_critic_optimizer_state(batch) + + torch.testing.assert_close(torch.random.get_rng_state(), expected_rng_state) + + +def test_fast_sac_load_state_dict_resets_cuda_graph_caches() -> None: + learner = _small_fast_sac_learner() + state_dict = learner.get_state_dict() + marker = object() + learner._cuda_graph_critic = marker # type: ignore[assignment] + learner._cuda_graph_critic_static_inputs = {"obs": torch.zeros(1)} + learner._cuda_graph_critic_action_noise = torch.zeros(1) + learner._cuda_graph_critic_outputs = (torch.zeros(()),) * 6 # type: ignore[assignment] + learner._cuda_graph_critic_shapes = {"obs": torch.Size([1])} + learner._cuda_graph_actor = marker # type: ignore[assignment] + learner._cuda_graph_actor_static_inputs = {"obs": torch.zeros(1)} + learner._cuda_graph_actor_action_noise = torch.zeros(1) + learner._cuda_graph_actor_outputs = (torch.zeros(()),) * 4 # type: ignore[assignment] + learner._cuda_graph_actor_shapes = {"obs": torch.Size([1])} + + learner.load_state_dict(state_dict) + + assert learner._cuda_graph_critic is None + assert learner._cuda_graph_critic_static_inputs is None + assert learner._cuda_graph_critic_action_noise is None + assert learner._cuda_graph_critic_outputs is None + assert learner._cuda_graph_critic_shapes is None + assert learner._cuda_graph_actor is None + assert learner._cuda_graph_actor_static_inputs is None + assert learner._cuda_graph_actor_action_noise is None + assert learner._cuda_graph_actor_outputs is None + assert learner._cuda_graph_actor_shapes is None + + +def test_fast_sac_cuda_graph_static_inputs_keep_addresses_after_copy() -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + learner._cuda_graph_critic_static_inputs = _critic_graph_static_inputs(batch) + learner._cuda_graph_actor_static_inputs = _actor_graph_static_inputs(batch) + learner._cuda_graph_critic_action_noise = torch.zeros_like(batch["actions"]) + learner._cuda_graph_actor_action_noise = torch.zeros_like(batch["actions"]) + + critic_static_inputs = learner._cuda_graph_critic_static_inputs + actor_static_inputs = learner._cuda_graph_actor_static_inputs + assert critic_static_inputs is not None + assert actor_static_inputs is not None + critic_ptrs_before = {name: tensor.data_ptr() for name, tensor in critic_static_inputs.items()} + actor_ptrs_before = {name: tensor.data_ptr() for name, tensor in actor_static_inputs.items()} + + learner._copy_critic_graph_inputs(batch) + learner._copy_actor_graph_inputs(batch) + + critic_ptrs_after = {name: tensor.data_ptr() for name, tensor in critic_static_inputs.items()} + actor_ptrs_after = {name: tensor.data_ptr() for name, tensor in actor_static_inputs.items()} + assert critic_ptrs_after == critic_ptrs_before + assert actor_ptrs_after == actor_ptrs_before + + +def test_fast_sac_cuda_graph_static_inputs_match_batch_values_after_copy() -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + learner._cuda_graph_critic_static_inputs = { + key: torch.empty_like(batch[key]) for key in _CRITIC_GRAPH_INPUT_KEYS + } + learner._cuda_graph_actor_static_inputs = { + key: torch.empty_like(batch[key]) for key in _ACTOR_GRAPH_INPUT_KEYS + } + learner._cuda_graph_critic_action_noise = torch.zeros_like(batch["actions"]) + learner._cuda_graph_actor_action_noise = torch.zeros_like(batch["actions"]) + + learner._copy_critic_graph_inputs(batch) + learner._copy_actor_graph_inputs(batch) + + critic_static_inputs = learner._cuda_graph_critic_static_inputs + actor_static_inputs = learner._cuda_graph_actor_static_inputs + assert critic_static_inputs is not None + assert actor_static_inputs is not None + for key in _CRITIC_GRAPH_INPUT_KEYS: + torch.testing.assert_close(critic_static_inputs[key], batch[key]) + for key in _ACTOR_GRAPH_INPUT_KEYS: + torch.testing.assert_close(actor_static_inputs[key], batch[key]) + + +def test_fast_sac_cuda_graph_critic_packed_source_updates_static_views() -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + packed_source = _critic_graph_packed_source(batch) + batch_from_packed = {"critic_graph_packed_source": packed_source} + learner._cuda_graph_critic_shapes = learner._critic_graph_input_shapes(batch) + learner._cuda_graph_critic_static_packed_input = torch.empty_like(packed_source) + learner._cuda_graph_critic_static_inputs = learner._critic_graph_static_views_from_packed( + learner._cuda_graph_critic_static_packed_input, + learner._cuda_graph_critic_shapes, + ) + learner._cuda_graph_critic_action_noise = torch.zeros_like(batch["actions"]) + + learner._copy_critic_graph_inputs(batch_from_packed) + + critic_static_inputs = learner._cuda_graph_critic_static_inputs + assert critic_static_inputs is not None + for key in _CRITIC_GRAPH_INPUT_KEYS: + torch.testing.assert_close(critic_static_inputs[key], batch[key]) + + +def test_fast_sac_cuda_graph_sac_graph_packed_source_updates_critic_and_actor_views() -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + packed_source = _sac_graph_packed_source(batch) + batch_from_packed = {"sac_graph_packed_source": packed_source} + learner._cuda_graph_critic_shapes = learner._critic_graph_input_shapes(batch) + learner._cuda_graph_actor_shapes = learner._actor_graph_input_shapes(batch) + learner._cuda_graph_sac_static_packed_input = torch.empty_like(packed_source) + learner._cuda_graph_critic_static_inputs = learner._critic_graph_static_views_from_sac_packed( + learner._cuda_graph_sac_static_packed_input, + learner._cuda_graph_critic_shapes, + learner._cuda_graph_actor_shapes, + ) + learner._cuda_graph_actor_static_inputs = learner._actor_graph_static_views_from_sac_packed( + learner._cuda_graph_sac_static_packed_input, + learner._cuda_graph_actor_shapes, + ) + learner._cuda_graph_critic_action_noise = torch.zeros_like(batch["actions"]) + learner._cuda_graph_actor_action_noise = torch.zeros_like(batch["actions"]) + + learner._copy_critic_graph_inputs(batch_from_packed) + learner._copy_actor_graph_inputs(batch_from_packed) + + critic_static_inputs = learner._cuda_graph_critic_static_inputs + actor_static_inputs = learner._cuda_graph_actor_static_inputs + assert critic_static_inputs is not None + assert actor_static_inputs is not None + for key in _CRITIC_GRAPH_INPUT_KEYS: + torch.testing.assert_close(critic_static_inputs[key], batch[key]) + for key in _ACTOR_GRAPH_INPUT_KEYS: + torch.testing.assert_close(actor_static_inputs[key], batch[key]) + + +def test_fast_sac_actor_reuses_sac_graph_static_buffer_after_critic_copy() -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + packed_source = _sac_graph_packed_source(batch) + learner._cuda_graph_critic_shapes = learner._critic_graph_input_shapes(batch) + learner._cuda_graph_actor_shapes = learner._actor_graph_input_shapes(batch) + learner._cuda_graph_sac_static_packed_input = torch.empty_like(packed_source) + learner._cuda_graph_critic_static_inputs = learner._critic_graph_static_views_from_sac_packed( + learner._cuda_graph_sac_static_packed_input, + learner._cuda_graph_critic_shapes, + learner._cuda_graph_actor_shapes, + ) + learner._cuda_graph_actor_static_packed_input = learner._cuda_graph_sac_static_packed_input + learner._cuda_graph_actor_static_inputs = learner._actor_graph_static_views_from_sac_packed( + learner._cuda_graph_actor_static_packed_input, + learner._cuda_graph_actor_shapes, + ) + learner._cuda_graph_critic_action_noise = torch.zeros_like(batch["actions"]) + learner._cuda_graph_actor_action_noise = torch.zeros_like(batch["actions"]) + + learner._copy_critic_graph_inputs({"sac_graph_packed_source": packed_source}) + source_ptr_after_critic = learner._cuda_graph_sac_static_source_ptr + static_after_critic = learner._cuda_graph_sac_static_packed_input.clone() + + learner._copy_actor_graph_inputs({"sac_graph_packed_source": packed_source}) + + assert learner._cuda_graph_sac_static_source_ptr == source_ptr_after_critic + torch.testing.assert_close(learner._cuda_graph_sac_static_packed_input, static_after_critic) + + +def test_fast_sac_cuda_graph_shape_change_resets_cache_before_new_capture(monkeypatch) -> None: + learner = _small_fast_sac_learner() + old_batch = _small_offpolicy_batch() + new_batch = _small_offpolicy_batch(batch_size=5) + new_batch["dones"] = torch.tensor([0.0, 1.0, 0.0, 1.0, 0.0]) + new_batch["truncated"] = torch.tensor([0.0, 0.0, 1.0, 0.0, 1.0]) + marker = object() + events: list[str] = [] + + learner.use_cuda_graph_critic = True + learner.use_cuda_graph_actor = True + learner._device_type = "cuda" + learner._cuda_graph_critic_shapes = learner._critic_graph_input_shapes(old_batch) + learner._cuda_graph_critic_static_inputs = _critic_graph_static_inputs(old_batch) + learner._cuda_graph_critic_action_noise = torch.zeros_like(old_batch["actions"]) + learner._cuda_graph_critic_outputs = (torch.zeros(()),) * 6 + learner._cuda_graph_critic = marker # type: ignore[assignment] + learner._cuda_graph_actor_shapes = learner._actor_graph_input_shapes(old_batch) + learner._cuda_graph_actor_static_inputs = _actor_graph_static_inputs(old_batch) + learner._cuda_graph_actor_action_noise = torch.zeros_like(old_batch["actions"]) + learner._cuda_graph_actor_outputs = (torch.zeros(()),) * 4 + learner._cuda_graph_actor = marker # type: ignore[assignment] + + def assert_critic_cache_reset(prefix: str) -> None: + assert learner._cuda_graph_critic is None, prefix + assert learner._cuda_graph_critic_static_inputs is None, prefix + assert learner._cuda_graph_critic_action_noise is None, prefix + assert learner._cuda_graph_critic_outputs is None, prefix + assert learner._cuda_graph_critic_shapes is None, prefix + + def assert_actor_cache_reset(prefix: str) -> None: + assert learner._cuda_graph_actor is None, prefix + assert learner._cuda_graph_actor_static_inputs is None, prefix + assert learner._cuda_graph_actor_action_noise is None, prefix + assert learner._cuda_graph_actor_outputs is None, prefix + assert learner._cuda_graph_actor_shapes is None, prefix + + def fake_materialize_critic(_batch: dict[str, torch.Tensor]) -> None: + events.append("critic_materialize") + assert_critic_cache_reset("critic cache should reset before materialization") + + def fake_capture_critic(capture_batch: dict[str, torch.Tensor]) -> None: + events.append("critic_capture") + assert_critic_cache_reset("critic cache should reset before capture") + learner._cuda_graph_critic_shapes = learner._critic_graph_input_shapes(capture_batch) + learner._cuda_graph_critic_static_inputs = _critic_graph_static_inputs(capture_batch) + learner._cuda_graph_critic_action_noise = torch.zeros_like(capture_batch["actions"]) + learner._cuda_graph_critic_outputs = (torch.zeros(()),) * 6 + learner._cuda_graph_critic = object() # type: ignore[assignment] + + def fake_materialize_actor(_batch: dict[str, torch.Tensor]) -> None: + events.append("actor_materialize") + assert_actor_cache_reset("actor cache should reset before materialization") + + def fake_capture_actor(capture_batch: dict[str, torch.Tensor]) -> None: + events.append("actor_capture") + assert_actor_cache_reset("actor cache should reset before capture") + learner._cuda_graph_actor_shapes = learner._actor_graph_input_shapes(capture_batch) + learner._cuda_graph_actor_static_inputs = _actor_graph_static_inputs(capture_batch) + learner._cuda_graph_actor_action_noise = torch.zeros_like(capture_batch["actions"]) + learner._cuda_graph_actor_outputs = (torch.zeros(()),) * 4 + learner._cuda_graph_actor = object() # type: ignore[assignment] + + monkeypatch.setattr( + learner, + "_materialize_capturable_critic_optimizer_state", + fake_materialize_critic, + ) + monkeypatch.setattr(learner, "_capture_critic_cuda_graph", fake_capture_critic) + monkeypatch.setattr( + learner, + "_materialize_capturable_actor_optimizer_state", + fake_materialize_actor, + ) + monkeypatch.setattr(learner, "_capture_actor_cuda_graph", fake_capture_actor) + + assert learner.update_critic_cuda_graph(new_batch, read_metrics=False) == {} + assert learner.update_actor_cuda_graph(new_batch, read_metrics=False) == {} + + assert events == [ + "critic_materialize", + "critic_capture", + "actor_materialize", + "actor_capture", + ] + assert learner._cuda_graph_critic_shapes == learner._critic_graph_input_shapes(new_batch) + assert learner._cuda_graph_actor_shapes == learner._actor_graph_input_shapes(new_batch) + + +def test_fast_sac_cuda_graph_replay_paths_emit_outer_nvtx_ranges(monkeypatch) -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + pushed: list[str] = [] + popped = 0 + + class _FakeGraph: + def replay(self) -> None: + pushed.append("graph.replay.called") + + def fake_push(name: str) -> None: + pushed.append(name) + + def fake_pop() -> None: + nonlocal popped + popped += 1 + + monkeypatch.setattr(torch.cuda.nvtx, "range_push", fake_push) + monkeypatch.setattr(torch.cuda.nvtx, "range_pop", fake_pop) + + learner.use_cuda_graph_critic = True + learner.use_cuda_graph_actor = True + learner._device_type = "cuda" + learner.nvtx_profile_ranges = True + learner._cuda_graph_critic_shapes = learner._critic_graph_input_shapes(batch) + learner._cuda_graph_critic_static_inputs = { + key: tensor.clone() + for key, tensor in batch.items() + if key in {"critic", "actions", "rewards", "next_obs", "next_critic", "dones", "truncated"} + } + learner._cuda_graph_critic_action_noise = torch.empty(batch["actions"].shape) + learner._cuda_graph_critic_outputs = (torch.zeros(()),) * 6 + learner._cuda_graph_critic = _FakeGraph() # type: ignore[assignment] + learner._cuda_graph_actor_shapes = learner._actor_graph_input_shapes(batch) + learner._cuda_graph_actor_static_inputs = { + "obs": batch["obs"].clone(), + "critic": batch["critic"].clone(), + } + learner._cuda_graph_actor_action_noise = torch.empty(batch["actions"].shape) + learner._cuda_graph_actor_outputs = (torch.zeros(()),) * 4 + learner._cuda_graph_actor = _FakeGraph() # type: ignore[assignment] + + learner.update_critic_cuda_graph(batch) + learner.update_actor_cuda_graph(batch) + + assert pushed == [ + "critic_graph/copy_inputs", + "critic_graph/replay", + "graph.replay.called", + "critic_graph/output_metrics_item", + "actor_graph/copy_inputs", + "actor_graph/replay", + "graph.replay.called", + "actor_graph/output_metrics_item", + ] + assert popped == 6 + + +def test_fast_sac_cuda_graph_metrics_can_skip_item_reads(monkeypatch) -> None: + learner = _small_fast_sac_learner() + calls = 0 + + class _Metric: + def __init__(self, value: float) -> None: + self.value = value + + def item(self) -> float: + nonlocal calls + calls += 1 + return self.value + + learner._cuda_graph_critic_outputs = tuple(_Metric(float(i)) for i in range(6)) # type: ignore[assignment] + learner._cuda_graph_actor_outputs = tuple(_Metric(float(i)) for i in range(4)) # type: ignore[assignment] + + assert learner._critic_graph_output_metrics(read_items=False) == {} + assert learner._actor_graph_output_metrics(read_items=False) == {} + assert calls == 0 + + assert learner._critic_graph_output_metrics(read_items=True)["qf_loss"] == 0.0 + assert learner._actor_graph_output_metrics(read_items=True)["actor_loss"] == 0.0 + assert calls == 10 + + +def test_fast_sac_cuda_graph_metric_skip_preserves_replay_updates(monkeypatch) -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + replay_calls = 0 + metric_calls = 0 + + class _FakeGraph: + def replay(self) -> None: + nonlocal replay_calls + replay_calls += 1 + + class _Metric: + def item(self) -> float: + nonlocal metric_calls + metric_calls += 1 + return 0.0 + + learner.use_cuda_graph_critic = True + learner.use_cuda_graph_actor = True + learner._device_type = "cuda" + learner._cuda_graph_critic_shapes = learner._critic_graph_input_shapes(batch) + learner._cuda_graph_critic_static_inputs = { + key: tensor.clone() + for key, tensor in batch.items() + if key in {"critic", "actions", "rewards", "next_obs", "next_critic", "dones", "truncated"} + } + learner._cuda_graph_critic_action_noise = torch.zeros_like(batch["actions"]) + learner._cuda_graph_critic_outputs = tuple(_Metric() for _ in range(6)) # type: ignore[assignment] + learner._cuda_graph_critic = _FakeGraph() # type: ignore[assignment] + learner._cuda_graph_actor_shapes = learner._actor_graph_input_shapes(batch) + learner._cuda_graph_actor_static_inputs = { + "obs": batch["obs"].clone(), + "critic": batch["critic"].clone(), + } + learner._cuda_graph_actor_action_noise = torch.zeros_like(batch["actions"]) + learner._cuda_graph_actor_outputs = tuple(_Metric() for _ in range(4)) # type: ignore[assignment] + learner._cuda_graph_actor = _FakeGraph() # type: ignore[assignment] + + assert learner.update_critic_cuda_graph(batch, read_metrics=False) == {} + assert learner.update_actor_cuda_graph(batch, read_metrics=False) == {} + assert replay_calls == 2 + assert metric_calls == 0 + + learner.update_critic_cuda_graph(batch, read_metrics=True) + learner.update_actor_cuda_graph(batch, read_metrics=True) + assert replay_calls == 4 + assert metric_calls == 10 + + +def test_fast_sac_cuda_graph_input_copy_fills_static_noise_in_place(monkeypatch) -> None: + learner = _small_fast_sac_learner() + batch = _small_offpolicy_batch() + learner._cuda_graph_critic_static_inputs = { + key: tensor.clone() + for key, tensor in batch.items() + if key in {"critic", "actions", "rewards", "next_obs", "next_critic", "dones", "truncated"} + } + learner._cuda_graph_actor_static_inputs = { + "obs": batch["obs"].clone(), + "critic": batch["critic"].clone(), + } + learner._cuda_graph_critic_action_noise = torch.zeros_like(batch["actions"]) + learner._cuda_graph_actor_action_noise = torch.zeros_like(batch["actions"]) + + def fail_randn_like(_tensor: torch.Tensor) -> torch.Tensor: + raise AssertionError("graph input copy should fill static noise buffers in place") + + monkeypatch.setattr(torch, "randn_like", fail_randn_like) + + learner._copy_critic_graph_inputs(batch) + learner._copy_actor_graph_inputs(batch) + + assert not torch.equal( + learner._cuda_graph_critic_action_noise, torch.zeros_like(batch["actions"]) + ) + assert not torch.equal( + learner._cuda_graph_actor_action_noise, torch.zeros_like(batch["actions"]) + ) diff --git a/tests/algos/test_fast_sac_symmetry_contract.py b/tests/algos/test_fast_sac_symmetry_contract.py index 344daf5b0..d176ecad7 100644 --- a/tests/algos/test_fast_sac_symmetry_contract.py +++ b/tests/algos/test_fast_sac_symmetry_contract.py @@ -1,5 +1,6 @@ from __future__ import annotations +from contextlib import contextmanager from typing import Any import gymnasium as gym @@ -12,10 +13,21 @@ class _FakeSymmetryAugmentation: batch_multiplier = 2 + def __init__(self): + self.augment_obs_calls: list[str] = [] + self.augment_obs_and_actions_calls: list[str] = [] + self.mirror_obs_calls: list[str] = [] + + def augment_obs(self, obs, *, obs_group: str = "obs"): + self.augment_obs_calls.append(obs_group) + return torch.cat([obs, obs], dim=0) + def augment_obs_and_actions(self, obs, actions, *, obs_group: str = "obs"): - return obs, actions + self.augment_obs_and_actions_calls.append(obs_group) + return torch.cat([obs, obs], dim=0), torch.cat([actions, actions], dim=0) def mirror_obs(self, obs, *, obs_group: str = "obs"): + self.mirror_obs_calls.append(obs_group) return obs @@ -362,3 +374,143 @@ def fake_all_reduce(tensor, op=None): assert len(seen_sizes) == 1 assert seen_sizes[0] > 0 assert torch.allclose(learner.log_alpha, before) + + +def test_fast_sac_symmetry_augmentation_emits_fine_grained_nvtx_ranges( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unilab.algos.torch.fast_sac.learner import FastSACLearner + + seen_ranges: list[str] = [] + + @contextmanager + def record_range(name: str, enabled: bool): + if enabled: + seen_ranges.append(name) + yield + + monkeypatch.setattr(learner_module, "_cuda_nvtx_range", record_range) + + learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_symmetry=True, + symmetry_augmentation=_FakeSymmetryAugmentation(), + ) + learner.nvtx_profile_ranges = True + + def fake_critic_loss_tensors(*args, **kwargs): + del args, kwargs + return ( + torch.tensor(0.0, requires_grad=True), + torch.tensor(0.0), + torch.tensor(0.0), + torch.zeros(4), + ) + + def fake_actor_loss_tensors(*args, **kwargs): + del args, kwargs + return ( + torch.tensor(0.0, requires_grad=True), + torch.tensor(0.0), + torch.tensor(0.0), + ) + + monkeypatch.setattr(learner, "_critic_loss_tensors", fake_critic_loss_tensors) + monkeypatch.setattr(learner, "_actor_loss_tensors", fake_actor_loss_tensors) + batch = { + "obs": torch.randn(4, 4), + "critic": torch.randn(4, 5), + "actions": torch.randn(4, 2), + "rewards": torch.randn(4), + "next_obs": torch.randn(4, 4), + "next_critic": torch.randn(4, 5), + "dones": torch.zeros(4), + "truncated": torch.zeros(4), + } + + learner.update_critic(batch) + learner.update_actor(batch) + + for expected in { + "critic/symmetry_obs_actions", + "critic/symmetry_next_obs", + "critic/symmetry_critic_obs", + "critic/symmetry_critic_next_obs", + "critic/symmetry_aux_repeat", + "actor/symmetry_obs", + "actor/symmetry_critic_obs", + }: + assert expected in seen_ranges + + +def test_fast_sac_symmetry_uses_obs_only_augmentation_for_obs_only_paths( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unilab.algos.torch.fast_sac.learner import FastSACLearner + + symmetry = _FakeSymmetryAugmentation() + learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + use_symmetry=True, + symmetry_augmentation=symmetry, + ) + + def fake_critic_loss_tensors(critic_obs, actions, rewards, next_obs, critic_next_obs, *args): + assert critic_obs.shape[0] == 8 + assert actions.shape[0] == 8 + assert rewards.shape[0] == 8 + assert next_obs.shape[0] == 8 + assert critic_next_obs.shape[0] == 8 + return ( + torch.tensor(0.0, requires_grad=True), + torch.tensor(0.0), + torch.tensor(0.0), + torch.zeros(8), + ) + + def fake_actor_loss_tensors(obs, critic_obs): + assert obs.shape[0] == 8 + assert critic_obs.shape[0] == 8 + return ( + torch.tensor(0.0, requires_grad=True), + torch.tensor(0.0), + torch.tensor(0.0), + ) + + monkeypatch.setattr(learner, "_critic_loss_tensors", fake_critic_loss_tensors) + monkeypatch.setattr(learner, "_actor_loss_tensors", fake_actor_loss_tensors) + batch = { + "obs": torch.randn(4, 4), + "critic": torch.randn(4, 5), + "actions": torch.randn(4, 2), + "rewards": torch.randn(4), + "next_obs": torch.randn(4, 4), + "next_critic": torch.randn(4, 5), + "dones": torch.zeros(4), + "truncated": torch.zeros(4), + } + + learner.update_critic(batch) + learner.update_actor(batch) + + assert symmetry.augment_obs_and_actions_calls == ["obs"] + assert symmetry.augment_obs_calls == ["obs", "critic", "critic", "obs", "critic"] + assert symmetry.mirror_obs_calls == [] diff --git a/tests/algos/test_flash_sac_learner.py b/tests/algos/test_flash_sac_learner.py index 1df91d949..dab4d1330 100644 --- a/tests/algos/test_flash_sac_learner.py +++ b/tests/algos/test_flash_sac_learner.py @@ -59,6 +59,136 @@ def test_flashsac_learner_exposes_expected_dims(): assert learner.action_dim == 29 +def test_flashsac_cuda_graph_options_are_opt_in() -> None: + default_learner = _make_small_learner() + + assert default_learner.supports_cuda_graph_packed_staging is True + assert default_learner.use_cuda_graph_critic is False + assert default_learner.use_cuda_graph_actor is False + assert default_learner.use_cuda_graph_critic_packed_staging is False + assert default_learner.use_cuda_graph_actor_packed_staging is False + + graph_learner = _make_small_learner( + use_cuda_graph_critic=True, + use_cuda_graph_actor=True, + use_cuda_graph_critic_packed_staging=True, + use_cuda_graph_actor_packed_staging=True, + ) + + assert graph_learner.use_cuda_graph_critic is True + assert graph_learner.use_cuda_graph_actor is True + assert graph_learner.use_cuda_graph_critic_packed_staging is True + assert graph_learner.use_cuda_graph_actor_packed_staging is True + + +def test_flashsac_update_critic_cuda_graph_falls_back_on_cpu() -> None: + learner = FlashSACLearner( + obs_dim=98, + action_dim=29, + critic_obs_dim=101, + device="cpu", + use_cuda_graph_critic=True, + ) + batch = _make_batch(batch_size=4) + + metrics = learner.update_critic_cuda_graph(batch) + + assert "critic_loss" in metrics + assert "reward_scale_std" in metrics + + +def test_flashsac_update_actor_cuda_graph_falls_back_on_cpu() -> None: + learner = FlashSACLearner( + obs_dim=98, + action_dim=29, + critic_obs_dim=101, + device="cpu", + use_cuda_graph_actor=True, + ) + batch = _make_batch(batch_size=4) + + metrics = learner.update_actor_cuda_graph(batch) + + assert "actor_loss" in metrics + assert "temperature" in metrics + + +def test_flashsac_sac_graph_packed_source_updates_critic_and_actor_views() -> None: + learner = _make_small_learner( + use_cuda_graph_critic=True, + use_cuda_graph_actor=True, + use_cuda_graph_critic_packed_staging=True, + use_cuda_graph_actor_packed_staging=True, + ) + batch = { + "obs": torch.randn(4, 4), + "critic": torch.randn(4, 6), + "actions": torch.randn(4, 2), + "rewards": torch.randn(4), + "next_obs": torch.randn(4, 4), + "next_critic": torch.randn(4, 6), + "dones": torch.zeros(4), + "truncated": torch.zeros(4), + } + packed = torch.cat( + [ + batch["obs"], + batch["critic"], + batch["actions"], + batch["rewards"].view(4, 1), + batch["next_obs"], + batch["next_critic"], + batch["dones"].view(4, 1), + batch["truncated"].view(4, 1), + ], + dim=1, + ) + critic_shapes = learner._critic_graph_input_shapes(batch) + actor_shapes = learner._actor_graph_input_shapes(batch) + + critic_views = learner._critic_graph_static_views_from_sac_packed( + packed, + critic_shapes, + actor_shapes, + ) + actor_views = learner._actor_graph_static_views_from_sac_packed(packed, actor_shapes) + + torch.testing.assert_close(critic_views["obs"], batch["obs"]) + torch.testing.assert_close(critic_views["critic"], batch["critic"]) + torch.testing.assert_close(critic_views["actions"], batch["actions"]) + torch.testing.assert_close(critic_views["rewards"], batch["rewards"]) + torch.testing.assert_close(critic_views["next_obs"], batch["next_obs"]) + torch.testing.assert_close(critic_views["next_critic"], batch["next_critic"]) + torch.testing.assert_close(actor_views["obs"], batch["obs"]) + torch.testing.assert_close(actor_views["critic"], batch["critic"]) + + +def test_flashsac_cuda_adam_optimizers_are_capture_ready(monkeypatch) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA-only optimizer kwargs require a CUDA-enabled torch build") + + calls: list[dict[str, Any]] = [] + + class _FakeAdam: + def __init__(self, _params, **kwargs): + calls.append(kwargs) + self.param_groups = [{"lr": kwargs["lr"]}] + + class _FakeLambdaLR: + def __init__(self, optimizer, lr_lambda): + self.optimizer = optimizer + self.lr_lambda = lr_lambda + + monkeypatch.setattr(torch.optim, "Adam", _FakeAdam) + monkeypatch.setattr(torch.optim.lr_scheduler, "LambdaLR", _FakeLambdaLR) + + FlashSACLearner(obs_dim=98, action_dim=29, critic_obs_dim=101, device="cuda") + + assert len(calls) == 3 + assert all(call["fused"] for call in calls) + assert all(call["capturable"] for call in calls) + + def test_flashsac_learner_declares_multi_gpu_contract() -> None: validate_distributed_learner_capability( learner_cls=FlashSACLearner, @@ -188,6 +318,68 @@ def fake_compile(fn: Callable, **kwargs): ] +def test_flashsac_graph_critic_skips_compiling_critic_loss(monkeypatch) -> None: + calls: list[tuple[str, dict[str, Any]]] = [] + + def fake_compile(fn: Callable, **kwargs): + calls.append((fn.__qualname__, kwargs)) + return fn + + learner = FlashSACLearner( + obs_dim=98, + action_dim=29, + critic_obs_dim=101, + device="cpu", + use_cuda_graph_critic=True, + ) + learner.device = torch.device("cuda") + monkeypatch.setattr(torch, "compile", fake_compile) + + learner._compile_training_methods() + + assert calls == [ + ( + "FlashSACActor.get_mean_and_std", + {"options": {"triton.cudagraphs": False}}, + ), + ( + "FlashSACLearner._actor_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ] + + +def test_flashsac_graph_actor_skips_compiling_actor_loss(monkeypatch) -> None: + calls: list[tuple[str, dict[str, Any]]] = [] + + def fake_compile(fn: Callable, **kwargs): + calls.append((fn.__qualname__, kwargs)) + return fn + + learner = FlashSACLearner( + obs_dim=98, + action_dim=29, + critic_obs_dim=101, + device="cpu", + use_cuda_graph_actor=True, + ) + learner.device = torch.device("cuda") + monkeypatch.setattr(torch, "compile", fake_compile) + + learner._compile_training_methods() + + assert calls == [ + ( + "FlashSACActor.get_mean_and_std", + {"options": {"triton.cudagraphs": False}}, + ), + ( + "FlashSACLearner._critic_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ] + + def test_flashsac_amp_dtype_resolution_and_scaler_rules() -> None: assert FlashSACLearner._resolve_amp_dtype("auto", "cuda") is torch.bfloat16 assert FlashSACLearner._resolve_amp_dtype("auto", "xpu") is torch.bfloat16 diff --git a/tests/algos/test_hora_contract.py b/tests/algos/test_hora_contract.py index 63b3ec105..affcb89ac 100644 --- a/tests/algos/test_hora_contract.py +++ b/tests/algos/test_hora_contract.py @@ -94,6 +94,27 @@ def test_hora_sac_learner_updates_with_privileged_tail() -> None: assert torch.isfinite(torch.tensor(list(actor_metrics.values()))).all() +def test_hora_sac_disables_cuda_graph_critic_path() -> None: + from unilab.algos.torch.hora.sac_learner import HoraSACLearner + + learner = HoraSACLearner( + obs_dim=5, + critic_obs_dim=8, + priv_info_dim=3, + action_dim=2, + device="cpu", + actor_hidden_dim=16, + critic_hidden_dim=16, + num_atoms=11, + use_layer_norm=False, + use_cuda_graph_critic=True, + use_cuda_graph_actor=True, + ) + + assert not learner.use_cuda_graph_critic + assert not learner.use_cuda_graph_actor + + def test_hora_sac_distilled_student_forward_does_not_require_priv_info() -> None: from unilab.algos.torch.hora.distill import HoraSACDistillActor, HoraSACDistillShared diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index ad8b538e0..dc22b13a0 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -373,6 +373,244 @@ def __init__(self, *args, **kwargs): assert runner.kwargs["learner"].kwargs["use_compile"] is True +def test_sac_cuda_graph_critic_override_is_passed_to_learner( + monkeypatch: pytest.MonkeyPatch, +): + import gymnasium as gym + + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "algo=sac", + "training.device=cuda", + "algo.use_symmetry=false", + "algo.algo_params.use_cuda_graph_critic=true", + ] + ) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def build_symmetry_augmentation(self, device=None): + return None + + def close(self): + pass + + class _FakeLearner: + class actor: + @staticmethod + def state_dict(): + return {"w": MagicMock(shape=(4,))} + + update_count = 0 + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(mod, "ensure_registries", lambda: None) + monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + + import unilab.algos.torch.fast_sac.learner as learner_mod + + monkeypatch.setattr(learner_mod, "FastSACLearner", _FakeLearner) + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + + monkeypatch.setattr(db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = mod.build_runner("sac", cfg) + + assert runner.kwargs["learner"].kwargs["use_cuda_graph_critic"] is True + + +def test_sac_cuda_graph_actor_override_is_passed_to_learner( + monkeypatch: pytest.MonkeyPatch, +): + import gymnasium as gym + + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "algo=sac", + "training.device=cuda", + "algo.use_symmetry=false", + "algo.algo_params.use_cuda_graph_actor=true", + ] + ) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def build_symmetry_augmentation(self, device=None): + return None + + def close(self): + pass + + class _FakeLearner: + class actor: + @staticmethod + def state_dict(): + return {"w": MagicMock(shape=(4,))} + + update_count = 0 + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(mod, "ensure_registries", lambda: None) + monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + + import unilab.algos.torch.fast_sac.learner as learner_mod + + monkeypatch.setattr(learner_mod, "FastSACLearner", _FakeLearner) + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + + monkeypatch.setattr(db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = mod.build_runner("sac", cfg) + + assert runner.kwargs["learner"].kwargs["use_cuda_graph_actor"] is True + + +def test_sac_cuda_graph_critic_packed_staging_override_is_passed_to_learner( + monkeypatch: pytest.MonkeyPatch, +): + import gymnasium as gym + + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "algo=sac", + "training.device=cuda", + "algo.use_symmetry=false", + "algo.algo_params.use_cuda_graph_critic=true", + "algo.algo_params.use_cuda_graph_critic_packed_staging=true", + ] + ) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def build_symmetry_augmentation(self, device=None): + return None + + def close(self): + pass + + class _FakeLearner: + class actor: + @staticmethod + def state_dict(): + return {"w": MagicMock(shape=(4,))} + + update_count = 0 + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(mod, "ensure_registries", lambda: None) + monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + + import unilab.algos.torch.fast_sac.learner as learner_mod + + monkeypatch.setattr(learner_mod, "FastSACLearner", _FakeLearner) + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + + monkeypatch.setattr(db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = mod.build_runner("sac", cfg) + + assert runner.kwargs["learner"].kwargs["use_cuda_graph_critic_packed_staging"] is True + + +def test_sac_cuda_graph_actor_packed_staging_override_is_passed_to_learner( + monkeypatch: pytest.MonkeyPatch, +): + import gymnasium as gym + + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "algo=sac", + "training.device=cuda", + "algo.use_symmetry=false", + "algo.algo_params.use_cuda_graph_actor=true", + "algo.algo_params.use_cuda_graph_actor_packed_staging=true", + ] + ) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def build_symmetry_augmentation(self, device=None): + return None + + def close(self): + pass + + class _FakeLearner: + class actor: + @staticmethod + def state_dict(): + return {"w": MagicMock(shape=(4,))} + + update_count = 0 + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(mod, "ensure_registries", lambda: None) + monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + + import unilab.algos.torch.fast_sac.learner as learner_mod + + monkeypatch.setattr(learner_mod, "FastSACLearner", _FakeLearner) + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + + monkeypatch.setattr(db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = mod.build_runner("sac", cfg) + + assert runner.kwargs["learner"].kwargs["use_cuda_graph_actor_packed_staging"] is True + + +def test_td3_config_does_not_expose_sac_critic_packed_staging(): + cfg = _offpolicy_cfg(["algo=td3"]) + + assert "use_cuda_graph_critic_packed_staging" not in cfg.algo.algo_params + + +def test_td3_config_does_not_expose_sac_actor_packed_staging(): + cfg = _offpolicy_cfg(["algo=td3"]) + + assert "use_cuda_graph_actor_packed_staging" not in cfg.algo.algo_params + + def test_sac_async_collection_passes_sync_collection_false( monkeypatch: pytest.MonkeyPatch, ): @@ -889,6 +1127,60 @@ def __init__(self, *args, **kwargs): assert runner.kwargs["replay_prefetch_mode"] == "one_tick" +def test_sac_nvtx_profile_ranges_passed_to_learner(monkeypatch: pytest.MonkeyPatch): + import gymnasium as gym + + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "algo=sac", + "training.device=cuda", + "training.nvtx_profile_ranges=true", + "algo.use_symmetry=false", + ] + ) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def build_symmetry_augmentation(self, device=None): + return None + + def close(self): + pass + + class _FakeLearner: + class actor: + @staticmethod + def state_dict(): + return {"w": MagicMock(shape=(4,))} + + update_count = 0 + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(mod, "ensure_registries", lambda: None) + monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + + import unilab.algos.torch.fast_sac.learner as learner_mod + + monkeypatch.setattr(learner_mod, "FastSACLearner", _FakeLearner) + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + + monkeypatch.setattr(db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = mod.build_runner("sac", cfg) + + assert runner.kwargs["learner"].kwargs["nvtx_profile_ranges"] is True + + def test_flashsac_double_buffer_dispatches_to_correct_runner(monkeypatch: pytest.MonkeyPatch): import gymnasium as gym @@ -1134,6 +1426,183 @@ def test_one_tick_loop_waits_when_prefetch_not_ready(): assert names.count("wait_until_ready") == 2 +def test_cpu_pinned_pipeline_omits_critic_graph_packed_source_by_default(): + import queue + + import torch + + from unilab.ipc.replay_buffer import ReplayBuffer + from unilab.ipc.replay_pipelines.cpu_pinned_double_buffer import ( + CPUPinnedDoubleBufferReplayPipeline, + ) + + replay_buffer = ReplayBuffer( + capacity=8, + obs_dim=4, + action_dim=2, + device="cpu", + defer_gpu=True, + critic_dim=5, + packed_cpu_storage=True, + ) + sample_count = 4 + slots = [ + torch.empty((sample_count, replay_buffer._storage.shape[1]), dtype=torch.float32) + for _ in range(2) + ] + pipeline = CPUPinnedDoubleBufferReplayPipeline( + replay_buffer, + device="cpu", + sample_count=sample_count, + collector_pack_request_queue=queue.Queue(), + collector_pack_ready_queue=queue.Queue(), + collector_pack_shared_slots=slots, + ) + try: + pipeline._gpu_packed[0].copy_( + torch.arange(sample_count * replay_buffer._storage.shape[1], dtype=torch.float32).view( + sample_count, replay_buffer._storage.shape[1] + ) + ) + pipeline._hot = 0 + pipeline._has_hot_batch = True + + batch = pipeline.sample_large_batch(tick_id=1, sample_count=sample_count) + + assert "critic_graph_packed_source" not in batch + finally: + pipeline.close() + + +def test_cpu_pinned_pipeline_returns_critic_graph_packed_source_when_enabled(): + import queue + + import torch + + from unilab.ipc.replay_buffer import ReplayBuffer + from unilab.ipc.replay_pipelines.cpu_pinned_double_buffer import ( + CPUPinnedDoubleBufferReplayPipeline, + ) + + replay_buffer = ReplayBuffer( + capacity=8, + obs_dim=4, + action_dim=2, + device="cpu", + defer_gpu=True, + critic_dim=5, + packed_cpu_storage=True, + ) + sample_count = 4 + packed_width = int(replay_buffer._storage.shape[1]) + critic_graph_width = replay_buffer.critic_graph_packed_width() + slots = [torch.empty((sample_count, packed_width), dtype=torch.float32) for _ in range(2)] + critic_graph_slots = [ + torch.empty((sample_count, critic_graph_width), dtype=torch.float32) for _ in range(2) + ] + pipeline = CPUPinnedDoubleBufferReplayPipeline( + replay_buffer, + device="cpu", + sample_count=sample_count, + collector_pack_request_queue=queue.Queue(), + collector_pack_ready_queue=queue.Queue(), + collector_pack_shared_slots=slots, + use_critic_graph_packed_source=True, + collector_pack_critic_graph_shared_slots=critic_graph_slots, + ) + try: + packed = torch.arange(sample_count * packed_width, dtype=torch.float32).view( + sample_count, packed_width + ) + expected = torch.empty((sample_count, critic_graph_width), dtype=torch.float32) + replay_buffer.pack_critic_graph_source(packed, out=expected) + pipeline._gpu_packed[0].copy_(packed) + pipeline._gpu_critic_graph_packed[0].copy_(expected) + pipeline._hot = 0 + pipeline._has_hot_batch = True + + batch = pipeline.sample_large_batch(tick_id=1, sample_count=sample_count) + + source = batch["critic_graph_packed_source"] + assert source.is_contiguous() + torch.testing.assert_close(source, expected) + roundtrip = torch.cat( + [ + batch["critic"].reshape(sample_count, -1), + batch["actions"].reshape(sample_count, -1), + batch["rewards"].reshape(sample_count, -1), + batch["next_obs"].reshape(sample_count, -1), + batch["next_critic"].reshape(sample_count, -1), + batch["dones"].reshape(sample_count, -1), + batch["truncated"].reshape(sample_count, -1), + ], + dim=1, + ) + torch.testing.assert_close(source, roundtrip) + finally: + pipeline.close() + + +def test_cpu_pinned_pipeline_sac_graph_layout_uses_primary_h2d_buffer_for_graph_sources(): + import queue + + import torch + + from unilab.ipc.replay_buffer import ReplayBuffer + from unilab.ipc.replay_pipelines.cpu_pinned_double_buffer import ( + CPUPinnedDoubleBufferReplayPipeline, + ) + + replay_buffer = ReplayBuffer( + capacity=8, + obs_dim=4, + action_dim=2, + device="cpu", + defer_gpu=True, + critic_dim=5, + packed_cpu_storage=True, + ) + sample_count = 4 + graph_width = replay_buffer.sac_graph_packed_width() + slots = [torch.empty((sample_count, graph_width), dtype=torch.float32) for _ in range(2)] + pipeline = CPUPinnedDoubleBufferReplayPipeline( + replay_buffer, + device="cpu", + sample_count=sample_count, + collector_pack_request_queue=queue.Queue(), + collector_pack_ready_queue=queue.Queue(), + collector_pack_shared_slots=slots, + pack_layout="sac_graph", + ) + try: + standard = torch.arange( + sample_count * replay_buffer._storage.shape[1], + dtype=torch.float32, + ).view(sample_count, replay_buffer._storage.shape[1]) + graph_packed = torch.empty((sample_count, graph_width), dtype=torch.float32) + replay_buffer.pack_sac_graph_source(standard, out=graph_packed) + pipeline._gpu_packed[0].copy_(graph_packed) + pipeline._hot = 0 + pipeline._has_hot_batch = True + + batch = pipeline.sample_large_batch(tick_id=1, sample_count=sample_count) + + assert batch["sac_graph_packed_source"].data_ptr() == pipeline._gpu_packed[0].data_ptr() + assert "critic_graph_packed_source" not in batch + assert not pipeline._gpu_critic_graph_packed + torch.testing.assert_close(batch["obs"], standard[:, replay_buffer._obs_sl]) + torch.testing.assert_close(batch["critic"], standard[:, replay_buffer._critic_sl]) + torch.testing.assert_close(batch["actions"], standard[:, replay_buffer._act_sl]) + torch.testing.assert_close(batch["rewards"], standard[:, replay_buffer._rew_col]) + torch.testing.assert_close(batch["next_obs"], standard[:, replay_buffer._nobs_sl]) + torch.testing.assert_close(batch["next_critic"], standard[:, replay_buffer._ncritic_sl]) + torch.testing.assert_close(batch["dones"], standard[:, replay_buffer._done_col]) + torch.testing.assert_close(batch["truncated"], standard[:, replay_buffer._trunc_col]) + assert pipeline._h2d_bytes() == graph_packed.numel() * graph_packed.element_size() + finally: + pipeline.close() + + def test_one_tick_loop_skips_wait_when_prefetch_ready(): log, _, _ = _run_fake_loop(max_iterations=2, batch_ready=True) names = log.names() @@ -1338,6 +1807,222 @@ def capture_start_collector(*, target_fn, kwargs): assert captured["nan_guard_cfg"].enabled is True +@pytest.mark.parametrize( + ("algo_type", "supports_graph_staging", "packed_staging_enabled", "expected_enabled"), + [ + ("sac", True, True, True), + ("sac", True, False, False), + ("flashsac", True, True, True), + ("flashsac", False, True, False), + ("td3", True, True, False), + ], +) +def test_double_buffer_runner_enables_critic_graph_source_only_for_supported_opt_in( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + algo_type, + supports_graph_staging, + packed_staging_enabled, + expected_enabled, +): + import queue + + import torch + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + import unilab.algos.torch.offpolicy.runner as runner_mod + + class _FakeActor: + def state_dict(self): + return {"w": torch.zeros(1)} + + class _FakeLearner: + def __init__(self): + self.actor = _FakeActor() + self.update_count = 0 + self.supports_cuda_graph_packed_staging = supports_graph_staging + self.use_cuda_graph_critic_packed_staging = packed_staging_enabled + + def get_state_dict(self): + return {"update_count": self.update_count} + + class _FakeReplayBuffer: + def __init__(self, **kwargs): + self.size = torch.zeros(1, dtype=torch.int64) + self.ptr = torch.zeros(1, dtype=torch.int64) + self._storage = torch.zeros(16, 16) + self._critic_dim = 5 + + def critic_graph_packed_width(self): + return 16 + + def close(self): + pass + + class _FakeWeightSync: + def __init__(self): + self.name = "fake-ws" + self._lock = None + + @classmethod + def from_state_dict(cls, state_dict, create=True): + return cls() + + def close(self): + pass + + captured_pipeline = {} + + class _FakePipeline: + def __init__(self, *args, **kwargs): + captured_pipeline.update(kwargs) + + def close(self): + pass + + monkeypatch.setattr(db_mod, "ReplayBuffer", _FakeReplayBuffer) + monkeypatch.setattr(db_mod, "SharedWeightSync", _FakeWeightSync) + monkeypatch.setattr(db_mod, "CPUPinnedDoubleBufferReplayPipeline", _FakePipeline) + monkeypatch.setattr(runner_mod, "get_env_dims", lambda *args, **kwargs: (4, 2, 5)) + monkeypatch.setattr(db_mod.torch, "save", lambda *args, **kwargs: None) + + learner = _FakeLearner() + runner = db_mod.DoubleBufferOffPolicyRunner( + learner=learner, + env_name="DummyEnv", + algo_type=algo_type, + num_envs=2, + replay_buffer_n=8, + batch_size=8, + learning_starts=6, + updates_per_step=1, + policy_frequency=1, + sync_collection=False, + env_steps_per_sync=1, + device="cpu", + ) + + captured_collector = {} + + def capture_start_collector(*, target_fn, kwargs): + del target_fn + captured_collector.update(kwargs) + + monkeypatch.setattr(runner, "_start_collector", capture_start_collector) + monkeypatch.setattr(db_mod._SPAWN_CTX, "Queue", lambda maxsize=0: queue.Queue()) + monkeypatch.setattr(db_mod.time, "sleep", lambda seconds: None) + + runner.learn(max_iterations=0, save_interval=0, log_dir=str(tmp_path)) + + assert captured_pipeline.get("use_critic_graph_packed_source", False) is expected_enabled + assert ( + captured_pipeline.get("collector_pack_critic_graph_shared_slots") is not None + ) is expected_enabled + assert ( + captured_collector.get("collector_pack_critic_graph_shared_slots") is not None + ) is expected_enabled + + +def test_double_buffer_runner_uses_sac_graph_primary_layout_without_extra_graph_slots( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +): + import queue + + import torch + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + import unilab.algos.torch.offpolicy.runner as runner_mod + + class _FakeActor: + def state_dict(self): + return {"w": torch.zeros(1)} + + class _FakeLearner: + def __init__(self): + self.actor = _FakeActor() + self.update_count = 0 + self.use_cuda_graph_critic_packed_staging = True + self.use_cuda_graph_actor_packed_staging = True + + def get_state_dict(self): + return {"update_count": self.update_count} + + class _FakeReplayBuffer: + def __init__(self, **kwargs): + self.size = torch.zeros(1, dtype=torch.int64) + self.ptr = torch.zeros(1, dtype=torch.int64) + self._storage = torch.zeros(16, 16) + self._critic_dim = 5 + + def critic_graph_packed_width(self): + return 16 + + def sac_graph_packed_width(self): + return 16 + + def close(self): + pass + + class _FakeWeightSync: + name = "fake-ws" + _lock = None + + @classmethod + def from_state_dict(cls, state_dict, create=True): + return cls() + + def close(self): + pass + + captured_pipeline = {} + + class _FakePipeline: + def __init__(self, *args, **kwargs): + captured_pipeline.update(kwargs) + + def close(self): + pass + + monkeypatch.setattr(db_mod, "ReplayBuffer", _FakeReplayBuffer) + monkeypatch.setattr(db_mod, "SharedWeightSync", _FakeWeightSync) + monkeypatch.setattr(db_mod, "CPUPinnedDoubleBufferReplayPipeline", _FakePipeline) + monkeypatch.setattr(runner_mod, "get_env_dims", lambda *args, **kwargs: (4, 2, 5)) + monkeypatch.setattr(db_mod.torch, "save", lambda *args, **kwargs: None) + + runner = db_mod.DoubleBufferOffPolicyRunner( + learner=_FakeLearner(), + env_name="DummyEnv", + algo_type="sac", + num_envs=2, + replay_buffer_n=8, + batch_size=8, + learning_starts=6, + updates_per_step=1, + policy_frequency=1, + sync_collection=False, + env_steps_per_sync=1, + device="cpu", + ) + + captured_collector = {} + + def capture_start_collector(*, target_fn, kwargs): + del target_fn + captured_collector.update(kwargs) + + monkeypatch.setattr(runner, "_start_collector", capture_start_collector) + monkeypatch.setattr(db_mod._SPAWN_CTX, "Queue", lambda maxsize=0: queue.Queue()) + monkeypatch.setattr(db_mod.time, "sleep", lambda seconds: None) + + runner.learn(max_iterations=0, save_interval=0, log_dir=str(tmp_path)) + + assert captured_pipeline["pack_layout"] == "sac_graph" + assert captured_pipeline.get("use_critic_graph_packed_source", False) is False + assert captured_pipeline.get("collector_pack_critic_graph_shared_slots") is None + assert captured_collector.get("collector_pack_critic_graph_shared_slots") is None + + # --------------------------------------------------------------------------- # _safe_put_trainer_done helper (issue #594 sync deadlock fix) # --------------------------------------------------------------------------- diff --git a/tests/algos/test_offpolicy_runner_unit.py b/tests/algos/test_offpolicy_runner_unit.py index 952fd653c..5d07ed140 100644 --- a/tests/algos/test_offpolicy_runner_unit.py +++ b/tests/algos/test_offpolicy_runner_unit.py @@ -46,7 +46,11 @@ def __init__(self, *args, **kwargs) -> None: self.reward_normalizer = None self.update_count = 0 self.critic_updates = 0 + self.graph_critic_updates = 0 self.actor_updates = 0 + self.graph_actor_updates = 0 + self.graph_critic_read_metrics: list[bool] = [] + self.graph_actor_read_metrics: list[bool] = [] self.target_updates = 0 self.average_parameter_calls = 0 self.kwargs = dict(kwargs) @@ -56,10 +60,34 @@ def update_critic(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: self.critic_updates += 1 return {"critic_loss": float(batch["obs"].shape[0])} + def update_critic_cuda_graph( + self, + batch: dict[str, torch.Tensor], + *, + read_metrics: bool = True, + ) -> dict[str, float]: + self.graph_critic_updates += 1 + self.graph_critic_read_metrics.append(read_metrics) + if not read_metrics: + return {} + return {"critic_loss": float(batch["obs"].shape[0])} + def update_actor(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: self.actor_updates += 1 return {"actor_loss": float(batch["obs"].shape[0])} + def update_actor_cuda_graph( + self, + batch: dict[str, torch.Tensor], + *, + read_metrics: bool = True, + ) -> dict[str, float]: + self.graph_actor_updates += 1 + self.graph_actor_read_metrics.append(read_metrics) + if not read_metrics: + return {} + return {"actor_loss": float(batch["obs"].shape[0])} + def soft_update_target(self) -> None: self.target_updates += 1 @@ -810,6 +838,73 @@ def fake_sleep(seconds: float) -> None: assert learner.target_updates == 3 +def test_offpolicy_runner_uses_cuda_graph_critic_when_learner_opts_in( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + runner = _make_runner( + monkeypatch, + sync_collection=False, + updates_per_step=3, + policy_frequency=2, + ) + runner.learner.use_cuda_graph_critic = True + threshold = runner.train_start_threshold + sleep_sizes = iter([4, 8, threshold]) + + def fake_sleep(seconds: float) -> None: + if seconds < 0.5: + next_size = next(sleep_sizes, threshold) + replay_buffer = _FakeReplayBuffer.last_instance + assert replay_buffer is not None + replay_buffer.size[0] = next_size + replay_buffer.ptr[0] = next_size + + monkeypatch.setattr(runner_module._SPAWN_CTX, "Queue", lambda maxsize=0: queue.Queue()) + monkeypatch.setattr(runner_module.time, "sleep", fake_sleep) + + runner.learn(max_iterations=1, save_interval=0, log_dir=str(tmp_path / "logs")) + + learner = cast(_FakeLearner, runner.learner) + assert learner.critic_updates == 0 + assert learner.graph_critic_updates == 3 + assert learner.graph_critic_read_metrics == [False, False, True] + assert learner.actor_updates == 2 + assert learner.target_updates == 3 + + +def test_offpolicy_runner_uses_cuda_graph_actor_when_learner_opts_in( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + runner = _make_runner( + monkeypatch, + sync_collection=False, + updates_per_step=3, + policy_frequency=2, + ) + runner.learner.use_cuda_graph_actor = True + threshold = runner.train_start_threshold + sleep_sizes = iter([4, 8, threshold]) + + def fake_sleep(seconds: float) -> None: + if seconds < 0.5: + next_size = next(sleep_sizes, threshold) + replay_buffer = _FakeReplayBuffer.last_instance + assert replay_buffer is not None + replay_buffer.size[0] = next_size + replay_buffer.ptr[0] = next_size + + monkeypatch.setattr(runner_module._SPAWN_CTX, "Queue", lambda maxsize=0: queue.Queue()) + monkeypatch.setattr(runner_module.time, "sleep", fake_sleep) + + runner.learn(max_iterations=1, save_interval=0, log_dir=str(tmp_path / "logs")) + + learner = cast(_FakeLearner, runner.learner) + assert learner.graph_actor_updates == 2 + assert learner.graph_actor_read_metrics == [False, True] + assert learner.actor_updates == 0 + assert learner.target_updates == 3 + + class _FakeDoubleBufferPipeline: def __init__( self, @@ -825,6 +920,9 @@ def __init__( collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + pack_layout, + use_critic_graph_packed_source, + collector_pack_critic_graph_shared_slots, ) -> None: del ( replay_buffer, @@ -836,6 +934,9 @@ def __init__( collector_pack_request_queue, collector_pack_ready_queue, collector_pack_shared_slots, + pack_layout, + use_critic_graph_packed_source, + collector_pack_critic_graph_shared_slots, ) self.sample_count = sample_count self.trace_recorder = trace_recorder diff --git a/tests/envs/locomotion/g1/test_symmetry_contract.py b/tests/envs/locomotion/g1/test_symmetry_contract.py index b0666d917..fd16ad1ed 100644 --- a/tests/envs/locomotion/g1/test_symmetry_contract.py +++ b/tests/envs/locomotion/g1/test_symmetry_contract.py @@ -75,10 +75,14 @@ def test_g1_walk_flat_symmetry_can_augment_critic_group(): actions, obs_group="critic", ) + actor_obs_aug = augmentation.augment_obs(obs, obs_group="obs") + critic_obs_aug = augmentation.augment_obs(critic, obs_group="critic") assert actor_aug.shape == (2, env.obs_groups_spec["obs"]) assert critic_aug.shape == (2, env.obs_groups_spec["critic"]) assert action_aug.shape == (2, action_dim) assert critic_action_aug.shape == (2, action_dim) + assert torch.equal(actor_obs_aug, actor_aug) + assert torch.equal(critic_obs_aug, critic_aug) finally: env.close() diff --git a/tests/ipc/test_memory_budget.py b/tests/ipc/test_memory_budget.py index fbd57b6e3..c994a9ab2 100644 --- a/tests/ipc/test_memory_budget.py +++ b/tests/ipc/test_memory_budget.py @@ -27,6 +27,23 @@ def test_offpolicy_memory_budget_notes_native_exclusions() -> None: assert "driver memory" in breakdown +def test_offpolicy_memory_budget_includes_opt_in_critic_graph_staging() -> None: + estimate = estimate_offpolicy_bytes( + num_envs=10, + replay_buffer_n=4, + obs_dim=2, + action_dim=1, + critic_dim=3, + batch_size=5, + updates_per_step=2, + critic_graph_staging_width=9, + ) + + assert estimate["critic_graph_staging_slots"] == 5 * 2 * 9 * 4 * 2 + assert int(estimate["total"]) >= int(estimate["critic_graph_staging_slots"]) + assert "Critic graph staging" in str(estimate["breakdown"]) + + def test_shared_memory_budget_unknown_available_is_noop(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(shutil, "disk_usage", lambda path: (_ for _ in ()).throw(OSError())) estimate = {"total": 1024, "breakdown": "test"} diff --git a/tests/ipc/test_replay_pipeline_double_buffer.py b/tests/ipc/test_replay_pipeline_double_buffer.py index 64350e704..9d01b3aa4 100644 --- a/tests/ipc/test_replay_pipeline_double_buffer.py +++ b/tests/ipc/test_replay_pipeline_double_buffer.py @@ -27,6 +27,34 @@ def _make_cpu_replay(capacity: int = 64, obs_dim: int = 4, action_dim: int = 2) return rb +def _make_cpu_critic_replay( + capacity: int = 64, + obs_dim: int = 4, + action_dim: int = 2, + critic_dim: int = 5, +) -> ReplayBuffer: + rb = ReplayBuffer( + capacity=capacity, + obs_dim=obs_dim, + action_dim=action_dim, + device="cpu", + critic_dim=critic_dim, + packed_cpu_storage=True, + ) + n = min(32, capacity) + rb.add( + obs=torch.arange(n * obs_dim, dtype=torch.float32).view(n, obs_dim), + actions=torch.arange(n * action_dim, dtype=torch.float32).view(n, action_dim) + 100, + rewards=torch.arange(n, dtype=torch.float32) + 200, + next_obs=torch.arange(n * obs_dim, dtype=torch.float32).view(n, obs_dim) + 300, + dones=torch.zeros(n), + truncated=torch.ones(n), + critic=torch.arange(n * critic_dim, dtype=torch.float32).view(n, critic_dim) + 400, + next_critic=torch.arange(n * critic_dim, dtype=torch.float32).view(n, critic_dim) + 500, + ) + return rb + + def test_collector_pack_shared_batch_writes_expected_packed_rows(): from unilab.algos.torch.offpolicy.worker import _collector_pack_shared_batch @@ -58,6 +86,102 @@ def test_collector_pack_shared_batch_writes_expected_packed_rows(): assert ready["learner_hot_gpu_slot"] == 0 +def test_collector_pack_shared_batch_writes_critic_graph_source_from_same_rows(): + from unilab.algos.torch.offpolicy.worker import _collector_pack_shared_batch + + rb = _make_cpu_critic_replay(capacity=64, obs_dim=5, action_dim=3, critic_dim=7) + sample_count = 8 + shared_slots = [ + torch.empty((sample_count, rb._storage.shape[1])).share_memory_() for _ in range(2) + ] + critic_graph_slots = [ + torch.empty((sample_count, rb.critic_graph_packed_width())).share_memory_() + for _ in range(2) + ] + request = { + "tick_id": 8, + "snapshot_ptr": int(rb.ptr[0]), + "snapshot_size": int(rb.size[0]), + "sample_seed": 321, + "sample_count": sample_count, + "shared_slot": 1, + "learner_hot_gpu_slot": 0, + "target_gpu_slot": 1, + "use_critic_graph_packed_source": True, + } + + ready = _collector_pack_shared_batch(rb, request, shared_slots, critic_graph_slots) + + gen = torch.Generator(device="cpu") + gen.manual_seed(321) + expected_indices = torch.randint(0, int(rb.size[0]), (sample_count,), generator=gen) + expected_packed = rb._storage[expected_indices] + expected_graph_source = torch.empty_like(critic_graph_slots[1]) + rb.pack_critic_graph_source(expected_packed, out=expected_graph_source) + torch.testing.assert_close(shared_slots[1], expected_packed) + torch.testing.assert_close(critic_graph_slots[1], expected_graph_source) + assert ready["tick_id"] == 8 + assert ready["shared_slot"] == 1 + assert "critic_graph_shared_slots" not in ready + + +def test_replay_buffer_packs_sac_graph_source_without_changing_width(): + rb = _make_cpu_critic_replay(capacity=64, obs_dim=5, action_dim=3, critic_dim=7) + sample_count = 8 + packed = rb._storage[:sample_count].clone() + out = torch.empty((sample_count, rb.sac_graph_packed_width())) + + result = rb.pack_sac_graph_source(packed, out=out) + + assert result is out + assert rb.sac_graph_packed_width() == int(rb._storage.shape[1]) + expected = torch.cat( + [ + packed[:, rb._obs_sl], + packed[:, rb._critic_sl], + packed[:, rb._act_sl], + packed[:, rb._rew_col : rb._rew_col + 1], + packed[:, rb._nobs_sl], + packed[:, rb._ncritic_sl], + packed[:, rb._done_col : rb._done_col + 1], + packed[:, rb._trunc_col : rb._trunc_col + 1], + ], + dim=1, + ) + torch.testing.assert_close(out, expected) + + +def test_collector_pack_shared_batch_writes_sac_graph_layout_to_primary_slot(): + from unilab.algos.torch.offpolicy.worker import _collector_pack_shared_batch + + rb = _make_cpu_critic_replay(capacity=64, obs_dim=5, action_dim=3, critic_dim=7) + sample_count = 8 + shared_slots = [ + torch.empty((sample_count, rb.sac_graph_packed_width())).share_memory_() for _ in range(2) + ] + request = { + "tick_id": 11, + "snapshot_ptr": int(rb.ptr[0]), + "snapshot_size": int(rb.size[0]), + "sample_seed": 456, + "sample_count": sample_count, + "shared_slot": 1, + "learner_hot_gpu_slot": 0, + "target_gpu_slot": 1, + "pack_layout": "sac_graph", + } + + ready = _collector_pack_shared_batch(rb, request, shared_slots) + + gen = torch.Generator(device="cpu") + gen.manual_seed(456) + expected_indices = torch.randint(0, int(rb.size[0]), (sample_count,), generator=gen) + expected = torch.empty_like(shared_slots[1]) + rb.pack_sac_graph_source(rb._storage[expected_indices], out=expected) + torch.testing.assert_close(shared_slots[1], expected) + assert ready["pack_layout"] == "sac_graph" + + def test_collector_pack_shared_batch_rejects_hot_gpu_slot_target(): from unilab.algos.torch.offpolicy.worker import _collector_pack_shared_batch diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index fa53c1f45..d9dec2f54 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -327,6 +327,7 @@ def test_offpolicy_hydra_default_trace_flags(): assert cfg.training.trace_output_dir is None assert cfg.training.trace_thread_time is False assert cfg.training.trace_cuda_events is True + assert cfg.training.nvtx_profile_ranges is False assert cfg.training.verbose_metrics is False assert cfg.training.multi_gpu_sync_mode == "local_sgd" assert cfg.training.multi_gpu_sync_interval == 1