feat(profiler): reject FSDP profiler configs that crash on swap-based CPU offload

The FSDP2 manual CPU-offload path moves models with model.to("cpu") ->
nn.Module._apply -> torch.utils.swap_tensors, which raises
"RuntimeError: _apply(): Couldn't swap <param>" while torch.profiler holds
weakrefs to those params during an active window (reproduced on CPU with
torch 2.12). That offload only fires mid-loop under colocation.

Map of the crash precondition (all required):
  - strategy == "fsdp"  (Megatron offloads via flat-buffer/.data reassignment,
    never swap_tensors -> immune)
  - fsdp_config.cpu_offload == False  (the default "manual" path; cpu_offload=True
    uses FSDP2-native offload, no manual swap)
  - colocate_all OR colocate_policy_ref  (otherwise no in-loop offload happens)
SFT is single-model and hardcodes colocate_all=False -> never at risk.

Enforce via TorchProfilerConfig.validate(): the RL validator (validate_cfg) now
passes strategy/colocation/cpu_offload context so an incompatible config fails
fast at startup with an actionable message (set cpu_offload=true, disable
colocation, or use Megatron). SFT calls validate() with no context so the
cross-field check is skipped. Added unit + end-to-end tests and documented the
restriction in config.mdx.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: dyurk-lila

docs/content/docs/configuration/config.mdx

@@ -306,6 +306,8 @@ policy:
 
 **Scope:** the profiler captures **only the policy model's training step** (forward/backward + optimizer). In an RL run it does **not** profile the critic or reference models, and it does **not** profile generation/inference — only policy training compute on the configured `ranks`.
 
+**FSDP + colocation restriction:** on the FSDP backend, the default CPU-offload path moves models to CPU with `torch.utils.swap_tensors`, which crashes mid-run (`RuntimeError: _apply(): Couldn't swap <param>`) when the profiler holds references to those parameters. That offload only happens under colocation, so `validate_cfg` **rejects at startup** the combination of `enable: true` + `trainer.strategy=fsdp` + `policy.fsdp_config.cpu_offload: false` (default) + colocation (`placement.colocate_all: true` *or* `placement.colocate_policy_ref: true`). To profile FSDP, either set `policy.fsdp_config.cpu_offload: true` (FSDP2-native offload, no swap) or disable colocation. The **Megatron** backend offloads via a different mechanism and is unaffected.
+
 Which steps are recorded is controlled entirely by [`torch.profiler.schedule`](https://docs.pytorch.org/docs/stable/profiler.html#torch.profiler.schedule): the profiler skips the first `skip_first` steps, then repeats a cycle of `wait` (idle) + `warmup` (tracing discarded) + `active` (tracing recorded) steps `repeat` times (`repeat: 0` profiles every cycle for the whole run). This is how you profile multiple steps, at an interval, repeating.
 
 - `policy.torch_profiler_config.enable`: Master switch. When `false` (default), no profiler RPCs are dispatched and there is zero overhead.

skyrl/train/config/config.py

@@ -206,14 +206,27 @@ class TorchProfilerConfig(BaseConfig):
     """``chrome_trace`` -> ``*.pt.trace.json`` (HTA-friendly) or ``stacks`` ->
     flamegraph-style self-CUDA-time stacks (requires ``with_stack=True``)."""
 
-    def validate(self) -> None:
+    def validate(
+        self,
+        strategy: Optional[str] = None,
+        colocate_all: Optional[bool] = None,
+        colocate_policy_ref: Optional[bool] = None,
+        fsdp_cpu_offload: Optional[bool] = None,
+    ) -> None:
         """Validate the profiler config. No-op when disabled.
 
         Called from both ``validate_cfg`` (RL) and ``validate_sft_cfg`` (SFT) so
         an invalid config fails fast at startup rather than silently degrading
         (e.g. an unknown ``export_type`` would otherwise fall through to the
         chrome-trace branch, and an unknown ``activities`` entry would disable
         profiling mid-run via the wrapper's exception isolation).
+
+        The ``strategy``/``colocate_*``/``fsdp_cpu_offload`` args (passed by the RL
+        validator, which has the full root config) gate an incompatibility check:
+        the FSDP2 manual CPU-offload path moves params with ``torch.utils.swap_tensors``,
+        which fails (``RuntimeError: Couldn't swap <param>``) while ``torch.profiler``
+        holds weakrefs to those params during an active window. SFT calls this with no
+        context (single-model, never colocated) so the check is skipped there.
         """
         if not self.enable:
             return
@@ -262,6 +275,29 @@ def validate(self) -> None:
         if self.active < 1:
             raise ValueError(f"`torch_profiler_config.active` must be >= 1, got {self.active}.")
 
+        # Cross-field: reject configs that would crash mid-run on the FSDP2 swap-based
+        # CPU offload. The manual offload path (fsdp_strategy: `manual_offload`, used
+        # when `fsdp_config.cpu_offload=False`) calls `model.to("cpu")` -> nn.Module._apply
+        # -> torch.utils.swap_tensors, which raises "Couldn't swap <param>" if any weakref
+        # to a param is alive. torch.profiler holds such weakrefs to every param it observes
+        # during an active window. That offload only fires mid-loop under colocation
+        # (colocate_all, or colocate_policy_ref for the policy/ref pair). Megatron offloads
+        # via its own flat-buffer/`.data` reassignment path (no swap_tensors) and is immune;
+        # `fsdp_config.cpu_offload=True` uses FSDP2-native offload (also no manual swap).
+        if strategy == "fsdp" and fsdp_cpu_offload is False and (colocate_all or colocate_policy_ref):
+            raise ValueError(
+                "`torch_profiler_config.enable=true` is incompatible with this FSDP configuration: "
+                "with the manual CPU-offload path (`policy.fsdp_config.cpu_offload=false`, the default) "
+                "under colocation "
+                f"(`placement.colocate_all={colocate_all}`, `placement.colocate_policy_ref={colocate_policy_ref}`), "
+                "the trainer offloads models to CPU via `torch.utils.swap_tensors` while the profiler holds "
+                "references to their parameters, which crashes mid-run with "
+                "`RuntimeError: _apply(): Couldn't swap <param>`. "
+                "To profile: set `policy.fsdp_config.cpu_offload=true` (FSDP2-native offload, no swap), or "
+                "disable colocation (`placement.colocate_all=false` and `placement.colocate_policy_ref=false`), "
+                "or use the Megatron backend (`trainer.strategy=megatron`)."
+            )
+
 
 @dataclass
 class MegatronLoraConfig(BaseConfig):

skyrl/train/utils/utils.py

@@ -271,7 +271,12 @@ def validate_cfg(cfg: SkyRLTrainConfig):
             "or negative to keep all checkpoints"
         )
 
-    cfg.trainer.policy.torch_profiler_config.validate()
+    cfg.trainer.policy.torch_profiler_config.validate(
+        strategy=cfg.trainer.strategy,
+        colocate_all=cfg.trainer.placement.colocate_all,
+        colocate_policy_ref=cfg.trainer.placement.colocate_policy_ref,
+        fsdp_cpu_offload=cfg.trainer.policy.fsdp_config.cpu_offload,
+    )
 
     # TODO (devpatel): move to initializing ray and syncing registries codepath at startup
     repopulate_all_registries()

tests/train/test_config.py

@@ -339,3 +339,50 @@ def test_validate_cfg_invokes_profiler_validation(self):
         cfg.trainer.policy.torch_profiler_config.export_type = "bogus"
         with pytest.raises(ValueError, match=r"export_type"):
             validate_cfg(cfg)
+
+    # -- FSDP swap-offload incompatibility (cross-field) --------------------------
+    # The FSDP2 manual CPU-offload path moves params via torch.utils.swap_tensors,
+    # which crashes mid-run ("Couldn't swap <param>") while the profiler holds
+    # weakrefs to those params. That offload only fires under colocation; Megatron
+    # and fsdp cpu_offload=true use non-swap paths and are safe.
+
+    def test_fsdp_colocate_all_manual_offload_rejected(self):
+        with pytest.raises(ValueError, match=r"Couldn't swap"):
+            self._cfg().validate(strategy="fsdp", colocate_all=True, colocate_policy_ref=True, fsdp_cpu_offload=False)
+
+    def test_fsdp_colocate_policy_ref_only_rejected(self):
+        # colocate_policy_ref alone still offloads the policy/ref pair mid-loop.
+        with pytest.raises(ValueError, match=r"Couldn't swap"):
+            self._cfg().validate(strategy="fsdp", colocate_all=False, colocate_policy_ref=True, fsdp_cpu_offload=False)
+
+    def test_fsdp_no_colocation_allowed(self):
+        # No colocation -> no in-loop offload -> safe.
+        self._cfg().validate(strategy="fsdp", colocate_all=False, colocate_policy_ref=False, fsdp_cpu_offload=False)
+
+    def test_fsdp_native_cpu_offload_allowed(self):
+        # cpu_offload=true uses FSDP2-native offload (no manual swap_tensors) -> safe.
+        self._cfg().validate(strategy="fsdp", colocate_all=True, colocate_policy_ref=True, fsdp_cpu_offload=True)
+
+    def test_megatron_colocation_allowed(self):
+        # Megatron offloads via flat-buffer/.data reassignment (no swap_tensors) -> safe.
+        self._cfg().validate(strategy="megatron", colocate_all=True, colocate_policy_ref=True, fsdp_cpu_offload=False)
+
+    def test_offload_check_skipped_without_context(self):
+        # Called with no context (e.g. the SFT path), the cross-field check is skipped.
+        self._cfg().validate()
+
+    def test_validate_cfg_rejects_profiler_under_default_colocation(self):
+        # End-to-end: the default config is fsdp + colocate_all + cpu_offload=false,
+        # so simply enabling the profiler must fail fast through validate_cfg.
+        cfg = _make_validated_test_config()
+        cfg.trainer.policy.torch_profiler_config.enable = True
+        cfg.trainer.policy.torch_profiler_config.save_path = "/tmp/skyrl_prof_test"
+        with pytest.raises(ValueError, match=r"Couldn't swap"):
+            validate_cfg(cfg)
+
+    def test_validate_cfg_allows_profiler_with_native_offload(self):
+        cfg = _make_validated_test_config()
+        cfg.trainer.policy.torch_profiler_config.enable = True
+        cfg.trainer.policy.torch_profiler_config.save_path = "/tmp/skyrl_prof_test"
+        cfg.trainer.policy.fsdp_config.cpu_offload = True
+        validate_cfg(cfg)  # must not raise on the profiler/offload check