From 69dde2ae12c3ddcc454f74e8441c8e6cc64bf6b0 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Thu, 4 Jun 2026 11:16:19 -0700 Subject: [PATCH 1/7] feat(profiler): drive torch.profiler around the training loop SkyRL constructed a Profiler object on the Megatron policy worker but never drove it -- start()/step()/stop() were called nowhere, so torch_profiler_config was dead code. This wires torch.profiler end to end for both Megatron and FSDP, both RL and SFT, with the full torch.profiler surface exposed as config. - start_profile / profile_step / stop_profile RPCs on the shared Worker base, dispatched via pass_through thin wrappers in WorkerDispatch. - Trainers bracket the loop: start before, one profile_step per global step, stop after -- all gated on torch_profiler_config.enable so non-profiling runs dispatch zero extra RPCs. - TorchProfilerConfig hoisted to PolicyConfig (also wired through the SFT config bridge), exposing schedule (skip_first/wait/warmup/active/repeat) and capture (activities/record_shapes/profile_memory/with_stack/with_flops/with_modules/ export_type) knobs. Defaults reproduce prior effective behavior; enable=false by default. Traces written by tensorboard_trace_handler as HTA/Kineto-friendly *.pt.trace.json under {ckpt_path}/profiler_traces by default. - Profiles only the policy model's training step (fwd/bwd + optimizer), not the critic/ref models and not generation/inference. - All profiler paths are exception-isolated: a fault disables profiling for the rest of the run rather than crashing it. - Removes the redundant Megatron-only torch_profiler_config. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/api-pages.yaml | 2 +- docs/content/docs/configuration/config.mdx | 33 ++ examples/train/async/async_trainer.py | 127 ++--- examples/train/megatron/run_megatron.sh | 10 +- .../megatron/run_megatron_nemotron_mini_4b.sh | 6 +- skyrl/backends/skyrl_train/utils/profiler.py | 206 +++++--- .../skyrl_train/workers/fsdp/fsdp_worker.py | 4 + .../workers/megatron/megatron_worker.py | 9 +- skyrl/backends/skyrl_train/workers/worker.py | 54 +++ .../skyrl_train/workers/worker_dispatch.py | 55 +++ skyrl/train/config/__init__.py | 4 +- skyrl/train/config/config.py | 85 +++- skyrl/train/config/sft_config.py | 8 + skyrl/train/fully_async_trainer.py | 341 ++++++------- skyrl/train/sft_trainer.py | 305 ++++++------ skyrl/train/trainer.py | 459 ++++++++++-------- skyrl/train/utils/utils.py | 2 + .../gpu_ci/megatron/test_megatron_worker.py | 6 +- .../skyrl_train/utils/test_profiler.py | 394 +++++++++++++++ tests/train/test_config.py | 58 +++ tests/train/test_sft_config.py | 55 +++ 21 files changed, 1581 insertions(+), 642 deletions(-) create mode 100644 tests/backends/skyrl_train/utils/test_profiler.py diff --git a/docs/api-pages.yaml b/docs/api-pages.yaml index f3b7b7a252..a29d6ed5a2 100644 --- a/docs/api-pages.yaml +++ b/docs/api-pages.yaml @@ -147,7 +147,7 @@ - skyrl.train.config.config.MegatronConfig - skyrl.train.config.config.MegatronDDPConfig - skyrl.train.config.config.MegatronLoraConfig - - skyrl.train.config.config.MegatronTorchProfilerConfig + - skyrl.train.config.config.TorchProfilerConfig - heading: Placement description: "" objects: diff --git a/docs/content/docs/configuration/config.mdx b/docs/content/docs/configuration/config.mdx index b7ed540acc..6ac77e22d4 100644 --- a/docs/content/docs/configuration/config.mdx +++ b/docs/content/docs/configuration/config.mdx @@ -273,6 +273,23 @@ policy: use_torch_compile: false # Enable torch compile for the entropy calculation record_memory: false # Dump memory snapshot for debugging + torch_profiler_config: # torch.profiler-based training-loop profiler (see below) + enable: false + ranks: [0] + save_path: null # defaults to {ckpt_path}/profiler_traces + skip_first: 10 # torch.profiler.schedule + wait: 0 + warmup: 1 + active: 1 + repeat: 1 # 0 = profile for the whole run + activities: ["cpu", "cuda"] + record_shapes: true + profile_memory: false + with_stack: true + with_flops: false + with_modules: false + export_type: "chrome_trace" # chrome_trace | stacks + model_config_kwargs: {} # pass through kwargs to the HuggingFace model config for FSDP training backends (i.e. for overriding vocab size, etc) - for megatron, use policy.megatron_config.transformer_config_kwargs instead ``` @@ -283,6 +300,22 @@ policy: - `policy.use_torch_compile`: Whether to enable torch compile for entropy calculation - `policy.record_memory`: Whether to record memory usage. If `True`, this will use PyTorch's [memory snapshotting utility](https://docs.pytorch.org/docs/stable/torch_cuda_memory.html) to record memory usage and dump memory snapshots after each policy model training step. +### Torch Profiler Configuration + +`policy.torch_profiler_config` enables a [`torch.profiler`](https://docs.pytorch.org/docs/stable/profiler.html)-based profiler that the trainer drives around the training loop (both FSDP and Megatron backends, and both RL and SFT). When enabled, it writes one Kineto/[HTA](https://github.com/facebookresearch/HolisticTraceAnalysis)-friendly `*.pt.trace.json` per active window per profiled rank into `save_path`. + +**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`. + +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. +- `policy.torch_profiler_config.ranks`: List of global ranks to profile (e.g. `[0]`). Add a mid-pipeline-stage rank to diagnose pipeline bubbles. +- `policy.torch_profiler_config.save_path`: Output directory for traces. Defaults to `{ckpt_path}/profiler_traces`. +- `policy.torch_profiler_config.{skip_first,wait,warmup,active,repeat}`: Passed directly to `torch.profiler.schedule`. +- `policy.torch_profiler_config.activities`: Subset of `["cpu", "cuda"]` to record. +- `policy.torch_profiler_config.{record_shapes,profile_memory,with_stack,with_flops,with_modules}`: Passed directly to `torch.profiler.profile`. `with_stack` and `record_shapes` add overhead; `profile_memory` is heavier still and off by default. +- `policy.torch_profiler_config.export_type`: `chrome_trace` writes `*.pt.trace.json` (Kineto/HTA-friendly); `stacks` writes flamegraph-style self-CUDA-time stacks (requires `with_stack: true`). + ### LoRA Configuration LoRA (Low-Rank Adaptation) enables parameter-efficient fine-tuning by training only a small number of additional low-rank matrices instead of the full model weights: diff --git a/examples/train/async/async_trainer.py b/examples/train/async/async_trainer.py index 288812aaa5..8c1d289284 100644 --- a/examples/train/async/async_trainer.py +++ b/examples/train/async/async_trainer.py @@ -51,63 +51,76 @@ async def train(self): start_epoch = self.global_step // len(self.train_dataloader) # Start from step 1 self.global_step += 1 - for epoch in range(start_epoch, self.cfg.trainer.epochs): - # while this is just off by one, you can image a more general queue based approach - # where the generation buffer holds a list of objects that the trainer can read from - # bit by bit. - generation_buffer = asyncio.Queue(maxsize=1) - self.sync_finished = asyncio.Event() - self.generation_ack = asyncio.Event() - - # start generator task - generator_task = asyncio.create_task(self._run_generate_loop(generation_buffer)) - - for idx in range(len(self.train_dataloader)): - with Timer("step", self.all_timings): - status = await self._run_training(generation_buffer) - - # request the generation loop that we should sync sometime soon. - if idx != len(self.train_dataloader) - 1: - await self.generation_ack.wait() - - # sync weights - async with Timer("sync_weights", self.all_timings): - await self.dispatch.save_weights_for_sampler() - - self.sync_finished.set() - self.generation_ack.clear() - - # 5. set logs - logger.info(status) - # log epoch info - self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step}) - self.tracker.log(self.all_metrics, step=self.global_step) - self.all_metrics = {} - pbar.update(1) - - if self.cfg.trainer.eval_interval > 0 and ( - self.global_step % self.cfg.trainer.eval_interval == 0 - or self.global_step == self.total_training_steps - ): - with Timer("eval", self.all_timings): - eval_metrics = await self.eval() - self.all_metrics.update(eval_metrics) - if self.cfg.trainer.ckpt_interval > 0 and self.global_step % self.cfg.trainer.ckpt_interval == 0: - with Timer("save_checkpoints", self.all_timings): - self.save_checkpoints() - if self.cfg.trainer.hf_save_interval > 0 and self.global_step % self.cfg.trainer.hf_save_interval == 0: - with Timer("save_hf_model", self.all_timings): - self.save_models() - self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step) - self.all_timings = {} - self.global_step += 1 - - if self.cfg.trainer.update_ref_every_epoch and self.ref_model is not None: - with Timer("update_ref_with_policy", self.all_timings): - await asyncio.to_thread(self.update_ref_with_policy) - - # cancel generation task for this epoch - generator_task.cancel() + self._profiler_start() + try: + for epoch in range(start_epoch, self.cfg.trainer.epochs): + # while this is just off by one, you can image a more general queue based approach + # where the generation buffer holds a list of objects that the trainer can read from + # bit by bit. + generation_buffer = asyncio.Queue(maxsize=1) + self.sync_finished = asyncio.Event() + self.generation_ack = asyncio.Event() + + # start generator task + generator_task = asyncio.create_task(self._run_generate_loop(generation_buffer)) + + for idx in range(len(self.train_dataloader)): + with Timer("step", self.all_timings): + status = await self._run_training(generation_buffer) + + # request the generation loop that we should sync sometime soon. + if idx != len(self.train_dataloader) - 1: + await self.generation_ack.wait() + + # sync weights + async with Timer("sync_weights", self.all_timings): + await self.dispatch.save_weights_for_sampler() + + self.sync_finished.set() + self.generation_ack.clear() + + # Advance the torch profiler schedule once per global step + # (no-op unless profiling is enabled). + self._profiler_step() + + # 5. set logs + logger.info(status) + # log epoch info + self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step}) + self.tracker.log(self.all_metrics, step=self.global_step) + self.all_metrics = {} + pbar.update(1) + + if self.cfg.trainer.eval_interval > 0 and ( + self.global_step % self.cfg.trainer.eval_interval == 0 + or self.global_step == self.total_training_steps + ): + with Timer("eval", self.all_timings): + eval_metrics = await self.eval() + self.all_metrics.update(eval_metrics) + if self.cfg.trainer.ckpt_interval > 0 and self.global_step % self.cfg.trainer.ckpt_interval == 0: + with Timer("save_checkpoints", self.all_timings): + self.save_checkpoints() + if ( + self.cfg.trainer.hf_save_interval > 0 + and self.global_step % self.cfg.trainer.hf_save_interval == 0 + ): + with Timer("save_hf_model", self.all_timings): + self.save_models() + self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step) + self.all_timings = {} + self.global_step += 1 + + if self.cfg.trainer.update_ref_every_epoch and self.ref_model is not None: + with Timer("update_ref_with_policy", self.all_timings): + await asyncio.to_thread(self.update_ref_with_policy) + + # cancel generation task for this epoch + generator_task.cancel() + finally: + # Always stop/flush the profiler when the loop exits (incl. on error) + # so the open kineto trace window isn't leaked. No-op when disabled. + self._profiler_stop() pbar.close() if self.cfg.trainer.ckpt_interval > 0: diff --git a/examples/train/megatron/run_megatron.sh b/examples/train/megatron/run_megatron.sh index 48dd907466..9f4c594ae3 100644 --- a/examples/train/megatron/run_megatron.sh +++ b/examples/train/megatron/run_megatron.sh @@ -17,7 +17,9 @@ MEGATRON_TP=2 MEGATRON_PP=2 MEGATRON_CP=1 -# torch profiler config +# torch profiler config. Profiles ONLY the policy model's training step +# (forward/backward + optim) -- not the critic or ref models, and not generation/ +# inference. See trainer.policy.torch_profiler_config for schedule/capture knobs. ENABLE_TORCH_PROFILER=false RANKS_TO_PROFILE="[0]" SAVE_PATH="$HOME/megatron_prof/tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_${MODEL_NAME}" @@ -33,9 +35,9 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \ trainer.placement.ref_num_gpus_per_node=$NUM_GPUS \ generator.inference_engine.num_engines=$NUM_GPUS \ generator.inference_engine.tensor_parallel_size=1 \ - trainer.policy.megatron_config.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \ - trainer.policy.megatron_config.torch_profiler_config.ranks=$RANKS_TO_PROFILE \ - trainer.policy.megatron_config.torch_profiler_config.save_path=$SAVE_PATH \ + trainer.policy.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \ + trainer.policy.torch_profiler_config.ranks=$RANKS_TO_PROFILE \ + trainer.policy.torch_profiler_config.save_path=$SAVE_PATH \ trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ diff --git a/examples/train/megatron/run_megatron_nemotron_mini_4b.sh b/examples/train/megatron/run_megatron_nemotron_mini_4b.sh index 1a8bf8c752..6fe0f35288 100755 --- a/examples/train/megatron/run_megatron_nemotron_mini_4b.sh +++ b/examples/train/megatron/run_megatron_nemotron_mini_4b.sh @@ -53,9 +53,9 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \ trainer.placement.ref_num_gpus_per_node=$NUM_GPUS \ generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ generator.inference_engine.tensor_parallel_size=$INFERENCE_TP \ - trainer.policy.megatron_config.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \ - trainer.policy.megatron_config.torch_profiler_config.ranks=$RANKS_TO_PROFILE \ - trainer.policy.megatron_config.torch_profiler_config.save_path=$SAVE_PATH \ + trainer.policy.torch_profiler_config.enable=$ENABLE_TORCH_PROFILER \ + trainer.policy.torch_profiler_config.ranks=$RANKS_TO_PROFILE \ + trainer.policy.torch_profiler_config.save_path=$SAVE_PATH \ trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ diff --git a/skyrl/backends/skyrl_train/utils/profiler.py b/skyrl/backends/skyrl_train/utils/profiler.py index da01dfbd50..23988fa8f2 100644 --- a/skyrl/backends/skyrl_train/utils/profiler.py +++ b/skyrl/backends/skyrl_train/utils/profiler.py @@ -4,82 +4,176 @@ import torch.distributed from loguru import logger +# Map config activity strings to torch ProfilerActivity members. +_ACTIVITY_MAP = { + "cpu": torch.profiler.ProfilerActivity.CPU, + "cuda": torch.profiler.ProfilerActivity.CUDA, +} -class Profiler: + +def build_profiler_from_policy_cfg(trainer_cfg): + """Construct a :class:`Profiler` from ``policy.torch_profiler_config``, or None. + + Returns ``None`` when profiling is disabled, so callers can simply assign the + result to ``self.profiler``. The trace ``save_path`` defaults to + ``{ckpt_path}/profiler_traces`` (mirrors how memory snapshots default under + ``ckpt_path``). """ - A PyTorch profiler wrapper class for collecting performance metrics. + cfg = trainer_cfg.policy.torch_profiler_config + if not cfg.enable: + return None + default_save_path = os.path.join(trainer_cfg.ckpt_path, "profiler_traces") + return Profiler(cfg, default_save_path=default_save_path) + + +class Profiler: + """A configurable ``torch.profiler`` wrapper driven by the training loop. + + The trainer brackets the loop: ``start()`` once before it, ``step()`` once + per global step, ``stop()`` once after. Which steps are actually recorded is + decided by ``torch.profiler.schedule(skip_first, wait, warmup, active, + repeat)`` -- this is the "profile N steps, every M, repeating K times" knob, + so nothing about the window is hardcoded. + + Traces are written by ``torch.profiler.tensorboard_trace_handler`` (one + ``*.pt.trace.json`` per active window, per rank), which is the + Kineto/Holistic-Trace-Analysis-friendly format -- no manual export. + + Every method is exception-isolated: a profiler fault disables profiling for + the rest of the run rather than crashing training. + + ``config`` fields (see :class:`skyrl.train.config.config.TorchProfilerConfig`): + enable, ranks, save_path, + skip_first, wait, warmup, active, repeat, + activities, record_shapes, profile_memory, with_stack, with_flops, + with_modules, export_type. """ - def __init__(self, config): - """ - config contains: - - enable: bool - - ranks: list[int] - - save_path: str - """ + def __init__(self, config, default_save_path: str = None): self.enable = config.enable + self.prof = None + # Per-window self-device-time kernel summary, refreshed by on_trace_ready + # at the close of each active window (the per-kernel attribution denominator). + # Exposed read-only via get_kernel_summary(). + self._last_pairs: list = [] + self._window_count: int = 0 if not config.enable: return self.config = config - self.save_path = config.save_path - self.ranks = config.ranks - self.saved = False - self.prof = None - self.rank = torch.distributed.get_rank() - if self.rank in self.ranks: - logger.info(f"[Profiler] Profiler init for rank {self.rank}") + self.save_path = config.save_path or default_save_path or "./profiler_traces" + self.ranks = list(config.ranks) + self.export_type = getattr(config, "export_type", "chrome_trace") + self.rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + if self.rank not in self.ranks: + return + try: + activities = [_ACTIVITY_MAP[a.lower()] for a in getattr(config, "activities", ["cpu", "cuda"])] + schedule = torch.profiler.schedule( + skip_first=getattr(config, "skip_first", 0), + wait=getattr(config, "wait", 0), + warmup=getattr(config, "warmup", 0), + active=getattr(config, "active", 1), + repeat=getattr(config, "repeat", 1), + ) + logger.info( + f"[Profiler] init rank {self.rank}: schedule(skip_first={getattr(config, 'skip_first', 0)}, " + f"wait={getattr(config, 'wait', 0)}, warmup={getattr(config, 'warmup', 0)}, " + f"active={getattr(config, 'active', 1)}, repeat={getattr(config, 'repeat', 1)}) " + f"-> traces under {self.save_path}" + ) self.prof = torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ], - schedule=torch.profiler.schedule( - wait=0, - warmup=0, - active=1, - repeat=1, - ), - record_shapes=True, - with_stack=True, + activities=activities, + schedule=schedule, + on_trace_ready=self._on_trace_ready, + record_shapes=getattr(config, "record_shapes", True), + profile_memory=getattr(config, "profile_memory", False), + with_stack=getattr(config, "with_stack", True), + with_flops=getattr(config, "with_flops", False), + with_modules=getattr(config, "with_modules", False), ) + except Exception as e: + logger.warning(f"[Profiler] init failed on rank {self.rank}; profiling disabled: {e}") + self.enable = False + self.prof = None + + def _on_trace_ready(self, prof) -> None: + """Fires at the close of each active window. Writes the trace AND stashes + a pickle-safe per-kernel self-device-time summary for the just-closed + window (the per-kernel attribution denominator -- exact, no cross-stream + overlap double-counting). ``rank{N}`` in the worker_name keeps the rank + parseable by HTA and avoids cross-rank filename collisions in a shared + ``save_path``. Best-effort: a fault here must never crash the worker.""" + os.makedirs(self.save_path, exist_ok=True) + worker_name = f"rank{self.rank}" + if self.export_type == "stacks": + # Flamegraph-style self-CUDA-time stacks (requires with_stack=True). + out = os.path.join(self.save_path, f"{worker_name}_stacks.txt") + prof.export_stacks(out, "self_cuda_time_total") + logger.info(f"[Profiler] rank {self.rank}: exported stacks -> {out}") + else: + torch.profiler.tensorboard_trace_handler(self.save_path, worker_name=worker_name)(prof) + logger.info(f"[Profiler] rank {self.rank}: exported chrome trace under {self.save_path}") + + try: + # ``self_device_time_total`` is torch 2.11's field (the older + # ``self_cuda_time_total`` was removed). Microseconds, self time. + self._last_pairs = [(str(e.key), float(e.self_device_time_total)) for e in prof.key_averages()] + self._window_count += 1 + except Exception as e: + logger.warning(f"[Profiler] rank {self.rank}: kernel-summary capture failed: {e}") + + def get_kernel_summary(self): + """Return the last closed window's self-device-time kernel summary, or None. + + Shape (pickle-safe, no tensors):: + + {"window_count": int, "pairs": [(kernel_name, self_device_us), ...]} + + ``None`` when profiling is disabled or no profiler was constructed on + this rank. ``pairs`` is empty until the first active window closes. + + NOTE: SkyRL's own trainers never read this -- the high-level + ``*.pt.trace.json`` files are the deliverable. This (and the + ``_on_trace_ready`` capture that feeds it) is a deliberately-provided + low-level API for downstream consumers that want per-kernel self-time + attribution without re-parsing the trace. Reached via + ``Worker.dump_profiler_summary`` -> ``WorkerDispatch.dump_profiler_summary``. + """ + if not self.enable or self.prof is None: + return None + return {"window_count": self._window_count, "pairs": list(self._last_pairs)} - def check(self): + def check(self) -> bool: return self.prof is not None and self.enable - def start(self): - if self.check(): - logger.info(f"[Profiler] started for rank {self.rank}") - self.prof.start() - - def step(self): - if self.check(): - self.prof.step() + def _disable(self, where: str, err: Exception) -> None: + logger.warning(f"[Profiler] {where} failed on rank {getattr(self, 'rank', '?')}; profiling disabled: {err}") + self.enable = False + self.prof = None - def stop(self): + def start(self) -> None: if self.check(): - logger.info(f"[Profiler] stopped for rank {self.rank}") - self.prof.stop() - - def save(self): - if self.prof is not None and not self.saved: - if not os.path.exists(self.save_path): - os.makedirs(self.save_path) - save_file_name = f"/prof_rank_{self.rank}.json" - logger.info(f"[Profiler] Saving trace to {self.save_path + save_file_name}") - self.prof.export_chrome_trace(self.save_path + save_file_name) - self.enable = False - self.saved = True + try: + logger.info(f"[Profiler] started for rank {self.rank}") + self.prof.start() + except Exception as e: + self._disable("start", e) - def stop_and_save(self): + def step(self) -> None: if self.check(): - self.stop() - self.save() + try: + self.prof.step() + except Exception as e: + self._disable("step", e) - def stop_trace(self): + def stop(self) -> None: if self.check(): - logger.info(f"[Profiler] Trace stopped for rank {self.rank}") - self.enable = False + try: + logger.info(f"[Profiler] stopped for rank {self.rank}") + self.prof.stop() + except Exception as e: + self._disable("stop", e) class CudaTimer: diff --git a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py index d03e583c24..da109ce86d 100644 --- a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py +++ b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py @@ -24,6 +24,7 @@ from skyrl.backends.skyrl_train.training_batch import ( TrainingInputBatch, ) +from skyrl.backends.skyrl_train.utils.profiler import build_profiler_from_policy_cfg from skyrl.backends.skyrl_train.weight_sync import ( LoraLoadRequest, WeightChunk, @@ -195,6 +196,9 @@ def init_model(self, model_path, num_training_steps: int = None): # standard CUDA memory; only subsequent activations use expandable segments. self._set_expandable_segments(True) + # create profiler (driven by the trainer via start/profile_step/stop RPCs) + self.profiler = build_profiler_from_policy_cfg(self.cfg) + async def init_weight_sync_state(self, inference_engine_client, inference_engine_cfg: "InferenceEngineConfig"): # Call super first to set _transfer_strategy_cls and create sender/receivers await super().init_weight_sync_state(inference_engine_client, inference_engine_cfg) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index d2bee26a56..6382e41b14 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -42,7 +42,7 @@ TrainingInputBatch, TrainingOutputBatch, ) -from skyrl.backends.skyrl_train.utils.profiler import Profiler +from skyrl.backends.skyrl_train.utils.profiler import build_profiler_from_policy_cfg from skyrl.backends.skyrl_train.weight_sync import ( LoraLoadRequest, WeightChunk, @@ -704,7 +704,7 @@ def __init__(self, **kwargs): self.actor_module: List[nn.Module] = None self.scheduler: OptimizerParamScheduler = None self.optimizer: DistributedOptimizer = None - self.profiler: Profiler = None + # self.profiler is initialized on the Worker base; populated in init_model. self._is_lora = self.cfg.policy.model.lora.rank > 0 # Per-worker store of LoRA adapter snapshots. Allocated only for the # LoRA path; FFT runs single-tenant exactly as before. @@ -806,9 +806,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): if self._rank == 0: print_model_size(self.actor_module[0]) - # create profiler - if self.cfg.policy.megatron_config.torch_profiler_config.enable: - self.profiler = Profiler(self.cfg.policy.megatron_config.torch_profiler_config) + # create profiler (driven by the trainer via start/profile_step/stop RPCs) + self.profiler = build_profiler_from_policy_cfg(self.cfg) # create optimizer (skipped for inference-only flows; Megatron's # DistributedOptimizer eagerly materializes fp32 master + AdamW state diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py index ed1a3e74b6..7feb3b8569 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/skyrl/backends/skyrl_train/workers/worker.py @@ -233,6 +233,10 @@ def __init__(self, cfg: TrainerConfig, *args, **kwargs): super().__init__(*args, **kwargs) self.cfg = cfg self._transfer_strategy_cls = None # Set in init_weight_transfer_communicator + # torch.profiler wrapper. Constructed in ``init_model`` when + # ``policy.torch_profiler_config.enable`` is set; driven by the trainer + # via the start_profile/profile_step/stop_profile RPCs below. + self.profiler = None if self.cfg.algorithm.temperature is None: raise ValueError("`cfg.algorithm.temperature` must be set") @@ -292,6 +296,56 @@ def set_algorithm_config(self, **kwargs) -> None: for key, value in kwargs.items(): setattr(self.cfg.algorithm, key, value) + # ------------------------------------------------------------------ + # torch.profiler control RPCs (dispatched by the trainer via "pass_through"). + # + # These live on the shared Worker base so they're on the snapshotted Ray + # actor method table of every PolicyWorker (Megatron + FSDP) without a + # subclass. They no-op when ``self.profiler`` is None (profiling disabled + # or rank not selected), so the trainer can call them on every step. The + # ``Profiler`` itself is exception-isolated; these are an extra guard so a + # profiler fault can never abort a training step. + # ------------------------------------------------------------------ + + def start_profile(self) -> None: + """Arm the profiler before the training loop (no-op when disabled).""" + if self.profiler is not None: + self.profiler.start() + + def profile_step(self) -> None: + """Advance the profiler schedule by one global step (no-op when disabled). + + Call exactly once per global step. ``torch.profiler``'s schedule decides + which steps are actually recorded; trace files are written automatically + at the close of each active window. + """ + if self.profiler is not None: + self.profiler.step() + + def stop_profile(self) -> None: + """Stop the profiler after the training loop, flushing any open window.""" + if self.profiler is not None: + self.profiler.stop() + + def dump_profiler_summary(self): + """Return this rank's last-window kernel self-time summary, or None. + + Pickle-safe dict (``{"window_count", "pairs": [(name, self_us), ...]}``) + or None when profiling is disabled / no profiler on this rank. The + low-level per-kernel attribution data path for downstream consumers; the + trace files remain the high-level (HTA) source. + + NOTE: SkyRL itself does not call this RPC (or ``get_kernel_summary``) in + its own training loop -- the high-level trace files are SkyRL's + deliverable. This is a deliberately-provided forward-looking API for + downstream consumers that want per-kernel self-time attribution without + re-parsing the on-disk trace. Keep it wired end-to-end (Profiler -> + Worker -> WorkerDispatch) so such a consumer needs only to call + ``dispatch.dump_profiler_summary("policy")``; do not remove it as + "dead code". + """ + return self.profiler.get_kernel_summary() if self.profiler is not None else None + def _get_module_for_offload(self): """Return the model module(s) to be offloaded/backloaded. Megatron offloads `self.actor_module`. FSDP workers use `self.model` directly.""" return self.model diff --git a/skyrl/backends/skyrl_train/workers/worker_dispatch.py b/skyrl/backends/skyrl_train/workers/worker_dispatch.py index 54c17673bc..7dd8ba9fe1 100644 --- a/skyrl/backends/skyrl_train/workers/worker_dispatch.py +++ b/skyrl/backends/skyrl_train/workers/worker_dispatch.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import ray +from loguru import logger from ray import ObjectRef from skyrl.backends.skyrl_train.distributed.dispatch import ( @@ -419,6 +420,60 @@ def set_algorithm_config(self, model: str, **kwargs) -> None: self._ensure_on_gpu(model, need_optimizer=False, need_model=False) ray.get(self._actor_groups[model].async_run_ray_method("pass_through", "set_algorithm_config", **kwargs)) + # ------------------------------------------------------------------ + # torch.profiler control. Dispatched via "pass_through" to every worker. + # Deliberately do NOT call _ensure_on_gpu so they don't perturb the + # colocation offload state machine. Each is best-effort: a profiler fault + # must never abort the training loop. + # ------------------------------------------------------------------ + + def start_profile(self, model: str) -> None: + """Arm the torch profiler on ``model``'s workers (before the loop).""" + if model not in self._actor_groups: + return + try: + ray.get(self._actor_groups[model].async_run_ray_method("pass_through", "start_profile")) + except Exception as e: + logger.warning(f"[profiler] start_profile dispatch for {model} failed: {e}") + + def profile_step(self, model: str) -> None: + """Advance the torch profiler schedule by one global step on ``model``.""" + if model not in self._actor_groups: + return + try: + ray.get(self._actor_groups[model].async_run_ray_method("pass_through", "profile_step")) + except Exception as e: + logger.warning(f"[profiler] profile_step dispatch for {model} failed: {e}") + + def stop_profile(self, model: str) -> None: + """Stop the torch profiler on ``model``'s workers (after the loop).""" + if model not in self._actor_groups: + return + try: + ray.get(self._actor_groups[model].async_run_ray_method("pass_through", "stop_profile")) + except Exception as e: + logger.warning(f"[profiler] stop_profile dispatch for {model} failed: {e}") + + def dump_profiler_summary(self, model: str) -> Optional[List]: + """Collect the per-rank last-window kernel self-time summaries for ``model``. + + Returns a per-actor list (one entry per worker): profiled ranks return a + dict ``{"window_count", "pairs"}``; other ranks return None. Returns None + when the model is unknown or dispatch fails. + + NOTE: not invoked by SkyRL's own trainers (which rely on the high-level + trace files). This is the dispatch-side entry of a forward-looking + low-level API for downstream consumers that want per-kernel self-time + attribution without re-parsing traces -- see ``Worker.dump_profiler_summary``. + """ + if model not in self._actor_groups: + return None + try: + return ray.get(self._actor_groups[model].async_run_ray_method("pass_through", "dump_profiler_summary")) + except Exception as e: + logger.warning(f"[profiler] dump_profiler_summary dispatch for {model} failed: {e}") + return None + def _save_memory_snapshot(self, model: str, tag: str) -> None: """Save memory snapshot on workers.""" ray.get( diff --git a/skyrl/train/config/__init__.py b/skyrl/train/config/__init__.py index 869a72ea4c..6c99e24d0f 100644 --- a/skyrl/train/config/__init__.py +++ b/skyrl/train/config/__init__.py @@ -20,7 +20,6 @@ MegatronConfig, MegatronDDPConfig, MegatronLoraConfig, - MegatronTorchProfilerConfig, MixedPrecisionConfig, ModelConfig, OffPolicyCorrectionConfig, @@ -33,6 +32,7 @@ SkyRLGymConfig, SkyRLLoraConfig, SkyRLTrainConfig, + TorchProfilerConfig, TrainerConfig, get_config_as_dict, get_config_as_yaml_str, @@ -63,7 +63,7 @@ "OptimizerConfig", "FSDPConfig", "MegatronConfig", - "MegatronTorchProfilerConfig", + "TorchProfilerConfig", "PlacementConfig", "SamplingParams", "ChatTemplateConfig", diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 5bdc74214b..dab4e7683e 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -156,11 +156,87 @@ class MegatronDDPConfig(BaseConfig): average_in_collective: bool = True +TORCH_PROFILER_ACTIVITIES = ("cpu", "cuda") +TORCH_PROFILER_EXPORT_TYPES = ("chrome_trace", "stacks") + + @dataclass -class MegatronTorchProfilerConfig(BaseConfig): +class TorchProfilerConfig(BaseConfig): + """Configuration for the ``torch.profiler``-based training-loop profiler. + + Mirrors ``torch.profiler.profile`` + ``torch.profiler.schedule`` so every + knob is overridable. Defaults reproduce the previous hardcoded behavior + (CPU+CUDA, ``record_shapes``+``with_stack``) but are now fully configurable. + The trainer drives it (``start`` before the loop, one ``step`` per global + step, ``stop`` after); the schedule decides which steps are recorded. + + Scope: this profiles **only the policy model's training step** + (forward/backward + optimizer). In an RL run it does **not** profile the + critic or ref models, and it does **not** profile generation/inference -- + only the policy training compute on the configured ``ranks``. + """ + enable: bool = False - ranks: List[int] = field(default_factory=list) + ranks: List[int] = field(default_factory=lambda: [0]) save_path: Optional[str] = None + """Trace output dir. Defaults to ``{ckpt_path}/profiler_traces`` when None.""" + + # torch.profiler.schedule -- one cycle = wait + warmup + active steps. + skip_first: int = 10 + """Steps to skip before the schedule begins (warmup/steady-state).""" + wait: int = 0 + warmup: int = 1 + active: int = 1 + """Number of steps recorded per cycle.""" + repeat: int = 1 + """Number of cycles. 0 = repeat for the whole run.""" + + # torch.profiler.profile capture knobs. + activities: List[str] = field(default_factory=lambda: ["cpu", "cuda"]) + record_shapes: bool = True + profile_memory: bool = False + with_stack: bool = True + with_flops: bool = False + with_modules: bool = False + export_type: str = "chrome_trace" + """``chrome_trace`` -> ``*.pt.trace.json`` (HTA-friendly) or ``stacks`` -> + flamegraph-style self-CUDA-time stacks (requires ``with_stack=True``).""" + + def validate(self) -> 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). + """ + if not self.enable: + return + if not self.ranks: + raise ValueError("`torch_profiler_config.ranks` must be non-empty when profiling is enabled.") + bad_activities = [a for a in self.activities if a.lower() not in TORCH_PROFILER_ACTIVITIES] + if bad_activities: + raise ValueError( + f"invalid `torch_profiler_config.activities` entries {bad_activities}. " + f"Each must be one of {list(TORCH_PROFILER_ACTIVITIES)}." + ) + if self.export_type not in TORCH_PROFILER_EXPORT_TYPES: + raise ValueError( + f"invalid `torch_profiler_config.export_type`: {self.export_type!r}. " + f"Must be one of {list(TORCH_PROFILER_EXPORT_TYPES)}." + ) + if self.export_type == "stacks" and not self.with_stack: + raise ValueError( + "`torch_profiler_config.export_type='stacks'` requires `with_stack=true` " + "(torch.profiler.export_stacks needs stack records)." + ) + for name in ("skip_first", "wait", "warmup", "repeat"): + value = getattr(self, name) + if value < 0: + raise ValueError(f"`torch_profiler_config.{name}` must be >= 0, got {value}.") + if self.active < 1: + raise ValueError(f"`torch_profiler_config.active` must be >= 1, got {self.active}.") @dataclass @@ -207,7 +283,6 @@ class MegatronConfig(BaseConfig): moe_router_dtype: str = "fp32" """Pass through to Megatron-Bridge - can be set to 'fp64' for additional numerical stability.""" ddp_config: MegatronDDPConfig = field(default_factory=MegatronDDPConfig) - torch_profiler_config: MegatronTorchProfilerConfig = field(default_factory=MegatronTorchProfilerConfig) lora_config: MegatronLoraConfig = field(default_factory=MegatronLoraConfig) optimizer_config_kwargs: Dict[str, Any] = field( default_factory=lambda: copy.deepcopy(DEFAULT_MEGATRON_OPTIMIZER_KWARGS) @@ -273,6 +348,10 @@ class PolicyConfig(BaseConfig): record_memory: bool = False """Save memory snapshots to ``{ckpt_path}/memory_snapshots/``. Visualize by dragging pickle files to https://docs.pytorch.org/memory_viz.""" + torch_profiler_config: TorchProfilerConfig = field(default_factory=TorchProfilerConfig) + """Backend-agnostic ``torch.profiler`` config (FSDP + Megatron). When + ``enable`` is true the trainer drives the profiler around the training loop + and writes traces to ``save_path``. See :class:`TorchProfilerConfig`.""" megatron_config: MegatronConfig = field(default_factory=MegatronConfig) model_config_kwargs: dict = field(default_factory=dict) """Pass-through kwargs for the HuggingFace model config (FSDP backends). diff --git a/skyrl/train/config/sft_config.py b/skyrl/train/config/sft_config.py index f2e216992c..bd5b4bdb80 100644 --- a/skyrl/train/config/sft_config.py +++ b/skyrl/train/config/sft_config.py @@ -21,6 +21,7 @@ ModelConfig, OptimizerConfig, SkyRLTrainConfig, + TorchProfilerConfig, ) # --------------------------------------------------------------------------- @@ -141,6 +142,10 @@ def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig": record_memory: bool = False """Save memory snapshots to ``{ckpt_path}/memory_snapshots/``. Visualize by dragging pickle files to https://docs.pytorch.org/memory_viz.""" + torch_profiler_config: TorchProfilerConfig = field(default_factory=TorchProfilerConfig) + """torch.profiler config (FSDP + Megatron). Set ``enable=true`` to capture + traces of the policy training step around the training loop; see + :class:`TorchProfilerConfig`.""" # ---- SFT-specific flat fields ---- strategy: str = "megatron" # "megatron" or "fsdp" @@ -304,6 +309,8 @@ def validate_sft_cfg(cfg: SFTConfig) -> None: if cfg.max_training_steps is not None and cfg.max_training_steps <= 0: raise ValueError(f"max_training_steps must be > 0, got {cfg.max_training_steps}") + cfg.torch_profiler_config.validate() + # Eval config if cfg.eval_interval < 0: raise ValueError(f"eval_interval must be >= 0, got {cfg.eval_interval}") @@ -387,6 +394,7 @@ def build_skyrl_config_for_sft(sft_cfg: SFTConfig) -> SkyRLTrainConfig: cfg.trainer.policy.model_config_kwargs = sft_cfg.model_config_kwargs cfg.trainer.policy.use_torch_compile = sft_cfg.use_torch_compile cfg.trainer.policy.record_memory = sft_cfg.record_memory + cfg.trainer.policy.torch_profiler_config = sft_cfg.torch_profiler_config # SFT doesn't use KL/ref model cfg.trainer.algorithm.use_kl_loss = False diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py index c6c3081f07..24221f45ea 100644 --- a/skyrl/train/fully_async_trainer.py +++ b/skyrl/train/fully_async_trainer.py @@ -461,185 +461,198 @@ async def train(self): start_epoch = resumed_start_epoch if resumed_start_epoch is not None else 0 self.global_step += 1 # start training at global_step 1 stop_training = False - for epoch in range(start_epoch, self.cfg.trainer.epochs): - self.epoch = epoch - # 0. Per-epoch prologue. Note that we do not do any cross-epoch asynchrony here. + self._profiler_start() + try: + for epoch in range(start_epoch, self.cfg.trainer.epochs): + self.epoch = epoch + # 0. Per-epoch prologue. Note that we do not do any cross-epoch asynchrony here. - # Buffer of completed generation, size bounded by capacity - consumed = B * (max_staleness_steps + 1) - generation_output_group_buffer = asyncio.Queue[GeneratedOutputGroup]( - maxsize=self.mini_batch_size * (self.max_staleness_steps + 1) - ) + # Buffer of completed generation, size bounded by capacity - consumed = B * (max_staleness_steps + 1) + generation_output_group_buffer = asyncio.Queue[GeneratedOutputGroup]( + maxsize=self.mini_batch_size * (self.max_staleness_steps + 1) + ) - # Maintain self.num_parallel_generation_workers concurrent group-generation workers - generator_tasks = [ - asyncio.create_task(self._run_generate_for_a_group_loop(generation_output_group_buffer)) - for _ in range(self.num_parallel_generation_workers) - ] - - # Lets the consumer detect epoch exhaustion (all generators done + buffer empty) instead of - # blocking forever on buffer.get() -- under sample_full_batch, drops can exhaust an epoch - # before num_steps_per_epoch steps complete. - all_generators_done = asyncio.Event() - generators_done_watcher = None - if self.sample_full_batch: - - async def _watch_generators_done(tasks=generator_tasks, event=all_generators_done): - await asyncio.gather(*tasks, return_exceptions=True) - event.set() - - generators_done_watcher = asyncio.create_task(_watch_generators_done()) - - # Steps trained in THIS epoch (not global_step % num_steps_per_epoch: sample_full_batch can - # end an epoch early, drifting global_step out of epoch alignment). On resume the dataloader - # already reflects this epoch's trained steps. The range below is just an upper bound. - trained_steps_this_epoch = self.async_train_dataloader.num_trained() // self.mini_batch_size - for _step_idx in range(self.global_step, (1 + epoch) * self.num_steps_per_epoch + 1): - with Timer("step", self.all_timings): - # 1. Wait until we have a full mini-batch buffered (dropping zero-variance groups if - # sample_full_batch). - ( - cur_generation_group_mini_batch, - cur_dropped_groups, - epoch_exhausted, - ) = await self._collect_generation_mini_batch(generation_output_group_buffer, all_generators_done) - - if epoch_exhausted: - # Exhausted mid mini-batch: discard the partial batch (marked consumed so it - # isn't regenerated on resume) and end the epoch early. - if cur_generation_group_mini_batch: - for _ in cur_generation_group_mini_batch: - await self._staleness_manager.on_rollout_filtered() - await self.async_train_dataloader.mark_filtered_uids( - [g.uid for g in cur_generation_group_mini_batch] + # Maintain self.num_parallel_generation_workers concurrent group-generation workers + generator_tasks = [ + asyncio.create_task(self._run_generate_for_a_group_loop(generation_output_group_buffer)) + for _ in range(self.num_parallel_generation_workers) + ] + + # Lets the consumer detect epoch exhaustion (all generators done + buffer empty) instead of + # blocking forever on buffer.get() -- under sample_full_batch, drops can exhaust an epoch + # before num_steps_per_epoch steps complete. + all_generators_done = asyncio.Event() + generators_done_watcher = None + if self.sample_full_batch: + + async def _watch_generators_done(tasks=generator_tasks, event=all_generators_done): + await asyncio.gather(*tasks, return_exceptions=True) + event.set() + + generators_done_watcher = asyncio.create_task(_watch_generators_done()) + + # Steps trained in THIS epoch (not global_step % num_steps_per_epoch: sample_full_batch can + # end an epoch early, drifting global_step out of epoch alignment). On resume the dataloader + # already reflects this epoch's trained steps. The range below is just an upper bound. + trained_steps_this_epoch = self.async_train_dataloader.num_trained() // self.mini_batch_size + for _step_idx in range(self.global_step, (1 + epoch) * self.num_steps_per_epoch + 1): + with Timer("step", self.all_timings): + # 1. Wait until we have a full mini-batch buffered (dropping zero-variance groups if + # sample_full_batch). + ( + cur_generation_group_mini_batch, + cur_dropped_groups, + epoch_exhausted, + ) = await self._collect_generation_mini_batch(generation_output_group_buffer, all_generators_done) + + if epoch_exhausted: + # Exhausted mid mini-batch: discard the partial batch (marked consumed so it + # isn't regenerated on resume) and end the epoch early. + if cur_generation_group_mini_batch: + for _ in cur_generation_group_mini_batch: + await self._staleness_manager.on_rollout_filtered() + await self.async_train_dataloader.mark_filtered_uids( + [g.uid for g in cur_generation_group_mini_batch] + ) + logger.warning( + f"sample_full_batch: epoch {epoch} exhausted with a partial mini-batch of " + f"{len(cur_generation_group_mini_batch)} group(s); discarding and ending the epoch." + ) + # Save the end-of-epoch checkpoint the normal is_epoch_end path would have, since + # we break before reaching it. + if self.cfg.trainer.ckpt_interval > 0: + with Timer("save_checkpoints", self.all_timings): + await asyncio.to_thread(self.save_checkpoints) + if self.cfg.trainer.hf_save_interval > 0: + with Timer("save_hf_model", self.all_timings): + await asyncio.to_thread(self.save_models) + break + + if self.sample_full_batch: + self.all_metrics["async/num_groups_dropped"] = len(cur_dropped_groups) + + # 2. Post-process the generated groups, aggregating to a single GeneratorOutput, and convert to training format. + with Timer("convert_to_training_input", self.all_timings): + training_input = await asyncio.to_thread( + self.convert_generation_group_mini_batch_to_training_input, + cur_generation_group_mini_batch, + cur_dropped_groups, ) - logger.warning( - f"sample_full_batch: epoch {epoch} exhausted with a partial mini-batch of " - f"{len(cur_generation_group_mini_batch)} group(s); discarding and ending the epoch." + + # 3. Run training and update consumed UIDs. + with Timer("run_training", self.all_timings): + status = await self._run_training(training_input) + await self.async_train_dataloader.mark_consumed_uids( + [g.uid for g in cur_generation_group_mini_batch] ) - # Save the end-of-epoch checkpoint the normal is_epoch_end path would have, since - # we break before reaching it. - if self.cfg.trainer.ckpt_interval > 0: + + # 4. After training: pause generation, sync weights, resume. + with Timer("sync_weights", self.all_timings): + await self.dispatch.save_weights_for_sampler() + + # A training step completed: count it for this epoch's bookkeeping. + trained_steps_this_epoch += 1 + + # Advance the torch profiler schedule once per global step + # (no-op unless profiling is enabled). One schedule step == + # one full async global step; the schedule decides which are recorded. + self._profiler_step() + + # 5. Set logs for this training step. + logger.info(status) + self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step}) + self.tracker.log(self.all_metrics, step=self.global_step, commit=False) + self.all_metrics = {} + pbar.update(1) + + # 6. Eval. At interval and at the last step. + # NOTE(Charlie): eval does not overlap with training, but overlaps with generation. + if self.cfg.trainer.eval_interval > 0 and ( + self.global_step % self.cfg.trainer.eval_interval == 0 + or self.global_step == self.total_training_steps + ): + with Timer("eval", self.all_timings): + eval_metrics = await self.eval() + self.all_metrics.update(eval_metrics) + + # 7. Checkpointing. At interval and at the last step of each epoch. + is_epoch_end = trained_steps_this_epoch == self.num_steps_per_epoch + if self.cfg.trainer.ckpt_interval > 0: + if is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0: with Timer("save_checkpoints", self.all_timings): await asyncio.to_thread(self.save_checkpoints) - if self.cfg.trainer.hf_save_interval > 0: + if self.cfg.trainer.hf_save_interval > 0: + if is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0: with Timer("save_hf_model", self.all_timings): await asyncio.to_thread(self.save_models) - break - if self.sample_full_batch: - self.all_metrics["async/num_groups_dropped"] = len(cur_dropped_groups) - - # 2. Post-process the generated groups, aggregating to a single GeneratorOutput, and convert to training format. - with Timer("convert_to_training_input", self.all_timings): - training_input = await asyncio.to_thread( - self.convert_generation_group_mini_batch_to_training_input, - cur_generation_group_mini_batch, - cur_dropped_groups, - ) + timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()} + if self._vllm_metrics_scraper is not None: + timing_payload.update(await self._vllm_metrics_scraper.sample()) + self.tracker.log(timing_payload, step=self.global_step, commit=True) + self.all_timings = {} + self.global_step += 1 + + if ( + self.cfg.trainer.max_training_steps is not None + and self.global_step > self.cfg.trainer.max_training_steps + ): + logger.info(f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early.") + for t in generator_tasks: + t.cancel() + await asyncio.gather(*generator_tasks, return_exceptions=True) + if generators_done_watcher is not None: + generators_done_watcher.cancel() + await asyncio.gather(generators_done_watcher, return_exceptions=True) + stop_training = True + break - # 3. Run training and update consumed UIDs. - with Timer("run_training", self.all_timings): - status = await self._run_training(training_input) - await self.async_train_dataloader.mark_consumed_uids( - [g.uid for g in cur_generation_group_mini_batch] - ) + # 8. Notify generation workers that the capacity has increased, unblocking them. + await self._staleness_manager.notify_capacity_change(self.global_step) + # Only trained UIDs map to completed steps; filtered/dropped UIDs are extra consumption. + expected_trained_in_epoch = self.mini_batch_size * trained_steps_this_epoch + actual_trained_in_epoch = self.async_train_dataloader.num_trained() + assert actual_trained_in_epoch == expected_trained_in_epoch, ( + "Unexpected number of trained (consumed minus filtered) data UIDs. Got: " + f"{actual_trained_in_epoch} != {expected_trained_in_epoch}" + ) - # 4. After training: pause generation, sync weights, resume. - with Timer("sync_weights", self.all_timings): - await self.dispatch.save_weights_for_sampler() - - # A training step completed: count it for this epoch's bookkeeping. - trained_steps_this_epoch += 1 - - # 5. Set logs for this training step. - logger.info(status) - self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step}) - self.tracker.log(self.all_metrics, step=self.global_step, commit=False) - self.all_metrics = {} - pbar.update(1) - - # 6. Eval. At interval and at the last step. - # NOTE(Charlie): eval does not overlap with training, but overlaps with generation. - if self.cfg.trainer.eval_interval > 0 and ( - self.global_step % self.cfg.trainer.eval_interval == 0 - or self.global_step == self.total_training_steps - ): - with Timer("eval", self.all_timings): - eval_metrics = await self.eval() - self.all_metrics.update(eval_metrics) - - # 7. Checkpointing. At interval and at the last step of each epoch. - is_epoch_end = trained_steps_this_epoch == self.num_steps_per_epoch - if self.cfg.trainer.ckpt_interval > 0: - if is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0: - with Timer("save_checkpoints", self.all_timings): - await asyncio.to_thread(self.save_checkpoints) - if self.cfg.trainer.hf_save_interval > 0: - if is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0: - with Timer("save_hf_model", self.all_timings): - await asyncio.to_thread(self.save_models) - - timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()} - if self._vllm_metrics_scraper is not None: - timing_payload.update(await self._vllm_metrics_scraper.sample()) - self.tracker.log(timing_payload, step=self.global_step, commit=True) - self.all_timings = {} - self.global_step += 1 - - if ( - self.cfg.trainer.max_training_steps is not None - and self.global_step > self.cfg.trainer.max_training_steps - ): - logger.info(f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early.") - for t in generator_tasks: - t.cancel() - await asyncio.gather(*generator_tasks, return_exceptions=True) - if generators_done_watcher is not None: - generators_done_watcher.cancel() - await asyncio.gather(generators_done_watcher, return_exceptions=True) - stop_training = True + if stop_training: break - # 8. Notify generation workers that the capacity has increased, unblocking them. - await self._staleness_manager.notify_capacity_change(self.global_step) - # Only trained UIDs map to completed steps; filtered/dropped UIDs are extra consumption. - expected_trained_in_epoch = self.mini_batch_size * trained_steps_this_epoch - actual_trained_in_epoch = self.async_train_dataloader.num_trained() - assert actual_trained_in_epoch == expected_trained_in_epoch, ( - "Unexpected number of trained (consumed minus filtered) data UIDs. Got: " - f"{actual_trained_in_epoch} != {expected_trained_in_epoch}" - ) + # 9. Per-epoch epilogue. + if self.cfg.trainer.update_ref_every_epoch and self.ref_model is not None: + with Timer("update_ref_with_policy", self.all_timings): + await asyncio.to_thread(self.update_ref_with_policy) - if stop_training: - break - - # 9. Per-epoch epilogue. - if self.cfg.trainer.update_ref_every_epoch and self.ref_model is not None: - with Timer("update_ref_with_policy", self.all_timings): - await asyncio.to_thread(self.update_ref_with_policy) - - # Cancel generator tasks for this epoch - for t in generator_tasks: - t.cancel() - try: - await asyncio.gather(*generator_tasks, return_exceptions=True) - except Exception: - pass - if generators_done_watcher is not None: - generators_done_watcher.cancel() - await asyncio.gather(generators_done_watcher, return_exceptions=True) - - # Per-epoch reset/validation for data loading and staleness management - assert all( - t.done() for t in generator_tasks - ), "Generator tasks must be done before resetting the dataloader manager and validating the staleness manager." - assert ( - generation_output_group_buffer.qsize() == 0 - ), f"We expect all generation output to be consumed by the training worker at end of an epoch, got {generation_output_group_buffer.qsize()}." - await self.async_train_dataloader.reset_at_epoch_end() - await self._staleness_manager.validate_state_at_epoch_end(self.global_step) + # Cancel generator tasks for this epoch + for t in generator_tasks: + t.cancel() + try: + await asyncio.gather(*generator_tasks, return_exceptions=True) + except Exception: + pass + if generators_done_watcher is not None: + generators_done_watcher.cancel() + await asyncio.gather(generators_done_watcher, return_exceptions=True) + + # Per-epoch reset/validation for data loading and staleness management + assert all( + t.done() for t in generator_tasks + ), "Generator tasks must be done before resetting the dataloader manager and validating the staleness manager." + assert ( + generation_output_group_buffer.qsize() == 0 + ), f"We expect all generation output to be consumed by the training worker at end of an epoch, got {generation_output_group_buffer.qsize()}." + await self.async_train_dataloader.reset_at_epoch_end() + await self._staleness_manager.validate_state_at_epoch_end(self.global_step) + + # End of an epoch. + finally: + # Always stop/flush the profiler when the loop exits -- including + # via an exception -- so the open kineto trace window isn't leaked. + # No-op when profiling is disabled. + self._profiler_stop() - # End of an epoch. pbar.close() if not stop_training: diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 9ac815db79..8646663b7b 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -715,6 +715,12 @@ def __init__( self._steps_per_epoch: int = 0 self._current_epoch: int = 0 + @property + def _torch_profiler_enabled(self) -> bool: + """Whether the trainer should drive the torch profiler. Gates all + profiler RPC dispatch so non-profiling runs pay zero extra round-trips.""" + return self.cfg.trainer.policy.torch_profiler_config.enable + def _build_collator(self, tokenizer): """Select the batch collator from the configured packing mode. @@ -1304,6 +1310,12 @@ def train_step(self, batch: TrainingInputBatch, step: int) -> dict: grad_norm = self.dispatch.optim_step("policy") metrics = output.metrics + + # Advance the torch profiler schedule once per global step (no-op unless + # profiling is enabled; the schedule decides which steps are recorded). + if self._torch_profiler_enabled: + self.dispatch.profile_step("policy") + loss_val = metrics.get("final_loss", metrics.get("loss", float("nan"))) return { "loss": loss_val, @@ -1394,35 +1406,43 @@ def _train_dummy(self): if self._ray_gpu_monitor is not None: self._ray_gpu_monitor.start() - for step in range(num_steps): - all_timings: dict[str, float] = {} - - with Timer("step", all_timings): - step_result = self.train_step(batch, step) - all_timings.update(step_result["timings"]) - - actual_num_tokens = batch["attention_mask"].sum().item() - self._total_tokens_processed += actual_num_tokens - tokens_per_second = actual_num_tokens / all_timings["step"] - - log_dict = { - "train/loss": step_result["loss"], - "train/grad_norm": step_result["grad_norm"], - "train/tokens_per_second": tokens_per_second, - "train/tokens_per_second_per_gpu": tokens_per_second / self._num_training_gpus, - "train/actual_num_tokens": actual_num_tokens, - "train/total_tokens_processed": self._total_tokens_processed, - } - log_dict.update({f"timing/{k}": v for k, v in all_timings.items()}) - if self._ray_gpu_monitor is not None: - log_dict.update(self._ray_gpu_monitor.flush()) - - self.tracker.log(log_dict, step=step, commit=True) - logger.info( - f"Step {step}: loss={step_result['loss']:.4f}, " - f"grad_norm={step_result['grad_norm']}, " - f"tokens_per_second={tokens_per_second:.0f}" - ) + if self._torch_profiler_enabled: + self.dispatch.start_profile("policy") + try: + for step in range(num_steps): + all_timings: dict[str, float] = {} + + with Timer("step", all_timings): + step_result = self.train_step(batch, step) + all_timings.update(step_result["timings"]) + + actual_num_tokens = batch["attention_mask"].sum().item() + self._total_tokens_processed += actual_num_tokens + tokens_per_second = actual_num_tokens / all_timings["step"] + + log_dict = { + "train/loss": step_result["loss"], + "train/grad_norm": step_result["grad_norm"], + "train/tokens_per_second": tokens_per_second, + "train/tokens_per_second_per_gpu": tokens_per_second / self._num_training_gpus, + "train/actual_num_tokens": actual_num_tokens, + "train/total_tokens_processed": self._total_tokens_processed, + } + log_dict.update({f"timing/{k}": v for k, v in all_timings.items()}) + if self._ray_gpu_monitor is not None: + log_dict.update(self._ray_gpu_monitor.flush()) + + self.tracker.log(log_dict, step=step, commit=True) + logger.info( + f"Step {step}: loss={step_result['loss']:.4f}, " + f"grad_norm={step_result['grad_norm']}, " + f"tokens_per_second={tokens_per_second:.0f}" + ) + finally: + # Always stop/flush the profiler when the loop exits (including via + # an exception) so the open trace window isn't leaked. No-op when off. + if self._torch_profiler_enabled: + self.dispatch.stop_profile("policy") logger.info("Dummy SFT training complete!") @@ -1553,118 +1573,127 @@ def train(self): self._fire("on_epoch_start") epoch_in_progress = True - while self.global_step <= num_steps: - all_timings: dict[str, float] = {} - - with Timer("step", all_timings): - - # Data loading with wrap-around - with Timer("data_loading", all_timings): - start_idx = (self.global_step * batch_size) % len(tokenized) - end_idx = start_idx + batch_size - if end_idx > len(tokenized): - batch_examples = tokenized[start_idx:] + tokenized[: end_idx - len(tokenized)] - else: - batch_examples = tokenized[start_idx:end_idx] - batch = self.collator(batch_examples, batch_size=batch_size) - - self._fire("on_step_start", batch=batch) - - # Training step - step_result = self.train_step(batch, self.global_step) - all_timings.update(step_result["timings"]) - - # Compute throughput using actual (non-padding) tokens - batch_padded_seq_len = batch["sequences"].shape[1] - actual_num_tokens = batch["attention_mask"].sum().item() - self._total_tokens_processed += actual_num_tokens - tokens_per_second = actual_num_tokens / all_timings["step"] - - # Build log dict - log_dict = { - "train/loss": step_result["loss"], - "train/grad_norm": step_result["grad_norm"], - "train/tokens_per_second": tokens_per_second, - "train/tokens_per_second_per_gpu": tokens_per_second / self._num_training_gpus, - "train/actual_num_tokens": actual_num_tokens, - "train/batch_padded_seq_len": batch_padded_seq_len, - "train/total_tokens_processed": self._total_tokens_processed, - } - log_dict.update({f"timing/{k}": v for k, v in all_timings.items()}) - if self._ray_gpu_monitor is not None: - log_dict.update(self._ray_gpu_monitor.flush()) - - self._fire("on_step_end", batch=batch, metrics=step_result) - - # Capture callback-driven triggers, then reset so they only fire once. - force_save = self._training_control.should_save - force_eval = self._training_control.should_evaluate - self._training_control.should_save = False - self._training_control.should_evaluate = False - - # Checkpoint: interval-driven or callback-requested. - interval_save = ( - self.sft_cfg.ckpt_interval > 0 - and self.global_step > 0 - and self.global_step % self.sft_cfg.ckpt_interval == 0 - ) - did_save_last_step = force_save or interval_save - if did_save_last_step: - with Timer("save_checkpoint", all_timings): - ckpt_path = self.save_checkpoint() - log_dict["timing/save_checkpoint"] = all_timings["save_checkpoint"] - self._fire("on_save", ckpt_path=ckpt_path) - - # HF export at regular intervals - if self.sft_cfg.hf_save_interval > 0 and self.global_step % self.sft_cfg.hf_save_interval == 0: - with Timer("save_hf_model", all_timings): - self.save_hf_model() - log_dict["timing/save_hf_model"] = all_timings["save_hf_model"] - - eval_metrics = None - num_eval_batches: int | None = None - # Eval fires at step N where N % eval_interval == 0 and N > 0, OR - # whenever a callback set ``control.should_evaluate``. - interval_eval = self.sft_cfg.eval_interval > 0 and self.global_step % self.sft_cfg.eval_interval == 0 - if eval_tokenized is not None and (force_eval or interval_eval): - self._fire("on_eval_start") - with Timer("eval", all_timings): - eval_metrics, num_eval_batches = self.run_eval(eval_tokenized) - self._fire("on_eval_end", metrics=eval_metrics) - if eval_metrics: - log_dict.update({f"eval/{k}": v for k, v in eval_metrics.items()}) - log_dict["timing/eval"] = all_timings["eval"] - - log_dict.update({"train/epoch": current_epoch, "train/global_step": self.global_step}) - # Callbacks may mutate log_dict in place via on_log. - self._fire("on_log", logs=log_dict) - self.tracker.log(log_dict, step=self.global_step, commit=True) - - if self.global_step % 5 == 0: - logger.info( - f"Step {self.global_step}: loss={step_result['loss']:.4f}, " f"grad_norm={step_result['grad_norm']}" + if self._torch_profiler_enabled: + self.dispatch.start_profile("policy") + try: + while self.global_step <= num_steps: + all_timings: dict[str, float] = {} + + with Timer("step", all_timings): + + # Data loading with wrap-around + with Timer("data_loading", all_timings): + start_idx = (self.global_step * batch_size) % len(tokenized) + end_idx = start_idx + batch_size + if end_idx > len(tokenized): + batch_examples = tokenized[start_idx:] + tokenized[: end_idx - len(tokenized)] + else: + batch_examples = tokenized[start_idx:end_idx] + batch = self.collator(batch_examples, batch_size=batch_size) + + self._fire("on_step_start", batch=batch) + + # Training step + step_result = self.train_step(batch, self.global_step) + all_timings.update(step_result["timings"]) + + # Compute throughput using actual (non-padding) tokens + batch_padded_seq_len = batch["sequences"].shape[1] + actual_num_tokens = batch["attention_mask"].sum().item() + self._total_tokens_processed += actual_num_tokens + tokens_per_second = actual_num_tokens / all_timings["step"] + + # Build log dict + log_dict = { + "train/loss": step_result["loss"], + "train/grad_norm": step_result["grad_norm"], + "train/tokens_per_second": tokens_per_second, + "train/tokens_per_second_per_gpu": tokens_per_second / self._num_training_gpus, + "train/actual_num_tokens": actual_num_tokens, + "train/batch_padded_seq_len": batch_padded_seq_len, + "train/total_tokens_processed": self._total_tokens_processed, + } + log_dict.update({f"timing/{k}": v for k, v in all_timings.items()}) + if self._ray_gpu_monitor is not None: + log_dict.update(self._ray_gpu_monitor.flush()) + + self._fire("on_step_end", batch=batch, metrics=step_result) + + # Capture callback-driven triggers, then reset so they only fire once. + force_save = self._training_control.should_save + force_eval = self._training_control.should_evaluate + self._training_control.should_save = False + self._training_control.should_evaluate = False + + # Checkpoint: interval-driven or callback-requested. + interval_save = ( + self.sft_cfg.ckpt_interval > 0 + and self.global_step > 0 + and self.global_step % self.sft_cfg.ckpt_interval == 0 ) + did_save_last_step = force_save or interval_save + if did_save_last_step: + with Timer("save_checkpoint", all_timings): + ckpt_path = self.save_checkpoint() + log_dict["timing/save_checkpoint"] = all_timings["save_checkpoint"] + self._fire("on_save", ckpt_path=ckpt_path) + + # HF export at regular intervals + if self.sft_cfg.hf_save_interval > 0 and self.global_step % self.sft_cfg.hf_save_interval == 0: + with Timer("save_hf_model", all_timings): + self.save_hf_model() + log_dict["timing/save_hf_model"] = all_timings["save_hf_model"] + + eval_metrics = None + num_eval_batches: int | None = None + # Eval fires at step N where N % eval_interval == 0 and N > 0, OR + # whenever a callback set ``control.should_evaluate``. + interval_eval = self.sft_cfg.eval_interval > 0 and self.global_step % self.sft_cfg.eval_interval == 0 + if eval_tokenized is not None and (force_eval or interval_eval): + self._fire("on_eval_start") + with Timer("eval", all_timings): + eval_metrics, num_eval_batches = self.run_eval(eval_tokenized) + self._fire("on_eval_end", metrics=eval_metrics) + if eval_metrics: + log_dict.update({f"eval/{k}": v for k, v in eval_metrics.items()}) + log_dict["timing/eval"] = all_timings["eval"] + + log_dict.update({"train/epoch": current_epoch, "train/global_step": self.global_step}) + # Callbacks may mutate log_dict in place via on_log. + self._fire("on_log", logs=log_dict) + self.tracker.log(log_dict, step=self.global_step, commit=True) + + if self.global_step % 5 == 0: + logger.info( + f"Step {self.global_step}: loss={step_result['loss']:.4f}, " + f"grad_norm={step_result['grad_norm']}" + ) - if eval_metrics: - logger.info( - f"Step {self.global_step}: eval_loss={eval_metrics.get('eval_loss', float('nan')):.4f} " - f"over {num_eval_batches} batches" - ) + if eval_metrics: + logger.info( + f"Step {self.global_step}: eval_loss={eval_metrics.get('eval_loss', float('nan')):.4f} " + f"over {num_eval_batches} batches" + ) - # Check for epoch boundary and reshuffle - epoch = (self.global_step * batch_size) // len(tokenized) - if epoch > current_epoch: - self._fire("on_epoch_end") - epoch_in_progress = False - for _ in range(epoch - current_epoch): - rng.shuffle(tokenized) - current_epoch = epoch - self._current_epoch = epoch - if self.global_step + 1 <= num_steps: - self._fire("on_epoch_start") - epoch_in_progress = True - - self.global_step += 1 + # Check for epoch boundary and reshuffle + epoch = (self.global_step * batch_size) // len(tokenized) + if epoch > current_epoch: + self._fire("on_epoch_end") + epoch_in_progress = False + for _ in range(epoch - current_epoch): + rng.shuffle(tokenized) + current_epoch = epoch + self._current_epoch = epoch + if self.global_step + 1 <= num_steps: + self._fire("on_epoch_start") + epoch_in_progress = True + + self.global_step += 1 + finally: + # Always stop/flush the profiler when the loop exits (including via + # an exception) so the open trace window isn't leaked. No-op when off. + if self._torch_profiler_enabled: + self.dispatch.stop_profile("policy") self.global_step = min(self.global_step, num_steps) # Pair the leading on_epoch_start: fire on_epoch_end if we exited the diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index b70ca0701d..119256ad2f 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -188,6 +188,41 @@ def has_critic(self) -> bool: """Check if critic model is configured.""" return bool(self.cfg.trainer.critic.model.path) + @property + def _torch_profiler_enabled(self) -> bool: + """Whether the trainer should drive the torch profiler on policy workers. + Gates all profiler RPC dispatch so non-profiling runs pay nothing.""" + return self.cfg.trainer.policy.torch_profiler_config.enable + + def _profiler_start(self) -> None: + """Arm the torch profiler on the policy workers before the training loop. + + No-op unless profiling is enabled. Shared by every trainer flavor + (sync / fully-async / one-step async) so the driving logic lives in one + place; subclasses that override ``train()`` call this at loop start. + """ + if self._torch_profiler_enabled: + self.dispatch.start_profile("policy") + + def _profiler_step(self) -> None: + """Advance the torch profiler schedule once per global step. + + No-op unless profiling is enabled. One call == one full global step + (not per minibatch); the torch schedule decides which steps are + recorded. Call exactly once per global step from every ``train()``. + """ + if self._torch_profiler_enabled: + self.dispatch.profile_step("policy") + + def _profiler_stop(self) -> None: + """Stop/flush the torch profiler after the training loop. + + No-op unless profiling is enabled. Call from a ``finally`` so an open + kineto trace window isn't leaked when the loop raises. + """ + if self._torch_profiler_enabled: + self.dispatch.stop_profile("policy") + def _build_train_dataloader_and_compute_training_steps(self): """ Hook for constructing the training dataloader. Subclasses can override @@ -290,224 +325,236 @@ async def train(self): # as well as hf model at step end will_save_ckpts = False hf_model_save = False - for epoch in range(start_epoch, self.cfg.trainer.epochs): - self._current_epoch = epoch - self._fire("on_epoch_start") - # ``step_started`` tracks the on_step_start/on_step_end pairing taking - # dynamic-sampling into account (which span multiple inner iterations - # before completing a logical step). - step_started = False - for _, rand_prompts in enumerate(self.train_dataloader): - if not step_started: - self._fire("on_step_start") - step_started = True - # Open the train-rollout metrics window once per logical - # step; paused so only the generation spans count toward the - # throughput denominator (dynamic sampling may generate more - # than once before the step completes). - if self._vllm_metrics_scraper is not None: - await self._vllm_metrics_scraper.start("vllm/train") - self._vllm_metrics_scraper.pause() - with Timer("step", self.all_timings): - # for colocate_all=true, inference engine is always on GPU when starting the training step - - # 0. truncate data to have even shards - rand_prompts = self._remove_tail_data(rand_prompts) - generator_input, uids = prepare_generator_input( - rand_prompts, - self.cfg.generator.n_samples_per_prompt, - get_sampling_params_for_backend( - self.cfg.generator.inference_engine.backend, self.cfg.generator.sampling_params - ), - self.cfg.environment.env_class, - "train", - self.global_step, - ) - - # 1.1. generation phase - if self._vllm_metrics_scraper is not None: - self._vllm_metrics_scraper.resume() - with Timer("generate", self.all_timings): - generator_output: GeneratorOutput = await self.generate(generator_input) - if self._vllm_metrics_scraper is not None: - self._vllm_metrics_scraper.pause() - - if self.cfg.generator.step_wise_trajectories: - # NOTE: We use instance_ids from `trajectory_ids` here instead of re-using `uids` - # this is because in step-wise training, len(uids) != len(generator_output["response_ids"]) - uids = [trajectory_id.instance_id for trajectory_id in generator_output["trajectory_ids"]] - - # dynamic sampling - if self.cfg.trainer.algorithm.dynamic_sampling.type is not None: - generator_output, uids, keep_sampling = self.handle_dynamic_sampling(generator_output, uids) - if keep_sampling: # continue sampling - # update progress bar for current batch (but not global step) - pbar.update(1) - continue - - if self.colocate_all: - # if we are not continuing sampling, we sleep the inference engine - await self.inference_engine_client.sleep() - - # The train rollout for this step is done generating; close - # its metrics window. ``vllm/eval/*`` is collected separately - # around eval below. - vllm_metrics: Dict[str, float] = {} - if self._vllm_metrics_scraper is not None: - vllm_metrics = await self._vllm_metrics_scraper.stop() - - # 1.2 postprocess rewards (and merge step-wise turns if enabled) - with Timer("postprocess_generator_output", self.all_timings): - generator_output, uids = self.postprocess_generator_output(generator_output, uids) - - # 2. print example just for debugging - log_interval = self.cfg.trainer.log_example_interval - if log_interval > 0 and self.global_step % log_interval == 0: - vis = self.tokenizer.decode(generator_output["response_ids"][0]) - log_example( - logger, - prompt=generator_input["prompts"][0], - response=vis, - reward=generator_output["rewards"][0], + self._profiler_start() + try: + for epoch in range(start_epoch, self.cfg.trainer.epochs): + self._current_epoch = epoch + self._fire("on_epoch_start") + # ``step_started`` tracks the on_step_start/on_step_end pairing taking + # dynamic-sampling into account (which span multiple inner iterations + # before completing a logical step). + step_started = False + for _, rand_prompts in enumerate(self.train_dataloader): + if not step_started: + self._fire("on_step_start") + step_started = True + # Open the train-rollout metrics window once per logical + # step; paused so only the generation spans count toward the + # throughput denominator (dynamic sampling may generate more + # than once before the step completes). + if self._vllm_metrics_scraper is not None: + await self._vllm_metrics_scraper.start("vllm/train") + self._vllm_metrics_scraper.pause() + with Timer("step", self.all_timings): + # for colocate_all=true, inference engine is always on GPU when starting the training step + + # 0. truncate data to have even shards + rand_prompts = self._remove_tail_data(rand_prompts) + generator_input, uids = prepare_generator_input( + rand_prompts, + self.cfg.generator.n_samples_per_prompt, + get_sampling_params_for_backend( + self.cfg.generator.inference_engine.backend, self.cfg.generator.sampling_params + ), + self.cfg.environment.env_class, + "train", + self.global_step, ) - # 3. Convert GeneratorOutput to TrainingInputBatch - with Timer("convert_to_training_input", self.all_timings): - training_input: TrainingInputBatch = self.convert_to_training_input(generator_output, uids) - - # 4. Inference and calculate values, log probs, rewards, kl divergence - with Timer("fwd_logprobs_values_reward", self.all_timings): - training_input = self.fwd_logprobs_values_reward(training_input) - - # 5. apply kl divergence penalty to rewards - if self.cfg.trainer.algorithm.use_kl_in_reward: - with Timer("apply_reward_kl_penalty", self.all_timings): - training_input = self.apply_reward_kl_penalty(training_input) - - # 6. calculate advantages and returns - with Timer("compute_advantages_and_returns", self.all_timings): - training_input = self.compute_advantages_and_returns(training_input) - # remove some unwanted keys - for key in ["rewards"]: - training_input.pop(key) - training_input.metadata.pop("uids") - training_input.metadata.pop("is_last_step", None) - - if self.cfg.trainer.dump_data_batch: - # dump data to file - with Timer("dump_data_batch"): - self.dump_data(training_input, file_name=f"global_step_{self.global_step}_training_input") - - # 7. train policy/critic model - # Policy model is backloaded to GPU during training - with Timer("train_critic_and_policy", self.all_timings): - status = self.train_critic_and_policy(training_input) - - self._fire("on_step_end", batch=training_input, metrics=status) - step_started = False - - # Capture callback-driven triggers, then reset. - force_save = self._training_control.should_save - force_eval = self._training_control.should_evaluate - self._training_control.should_save = False - self._training_control.should_evaluate = False - - # 8. conditionally save checkpoints and hf model - is_epoch_end = self.global_step % len(self.train_dataloader) == 0 - hf_model_save = self.cfg.trainer.hf_save_interval > 0 and ( - is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0 - ) - ckpt_interval_save = self.cfg.trainer.ckpt_interval > 0 and ( - is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0 - ) - will_save_ckpts = force_save or ckpt_interval_save - if will_save_ckpts: - with Timer("save_checkpoints", self.all_timings): - ckpt_path = self.save_checkpoints() - self._fire("on_save", ckpt_path=ckpt_path) - if hf_model_save: - with Timer("save_hf_model", self.all_timings): - self.save_models() - - # 9. conditionally sync policy and ref at the end of the epoch - if ( - self.cfg.trainer.update_ref_every_epoch - and self.ref_model is not None - and is_epoch_end - and epoch != self.cfg.trainer.epochs - 1 # skip updating ref at the end of the last epoch - ): - with Timer("update_ref_with_policy", self.all_timings): - self.update_ref_with_policy() - - # 10. Prepare weights for sampling - with Timer("sync_weights", self.all_timings): - await self.dispatch.save_weights_for_sampler() - - # 11. set logs - logger.info(status) - # Throughput metrics - train_time = self.all_timings.get("train_critic_and_policy", 0.0) - if train_time > 0 and training_input.get("attention_mask") is not None: - total_tokens = int(training_input["attention_mask"].sum().item()) - self.all_metrics["trainer/tokens_per_second_per_gpu"] = total_tokens / ( - train_time * self._num_training_gpus + # 1.1. generation phase + if self._vllm_metrics_scraper is not None: + self._vllm_metrics_scraper.resume() + with Timer("generate", self.all_timings): + generator_output: GeneratorOutput = await self.generate(generator_input) + if self._vllm_metrics_scraper is not None: + self._vllm_metrics_scraper.pause() + + if self.cfg.generator.step_wise_trajectories: + # NOTE: We use instance_ids from `trajectory_ids` here instead of re-using `uids` + # this is because in step-wise training, len(uids) != len(generator_output["response_ids"]) + uids = [trajectory_id.instance_id for trajectory_id in generator_output["trajectory_ids"]] + + # dynamic sampling + if self.cfg.trainer.algorithm.dynamic_sampling.type is not None: + generator_output, uids, keep_sampling = self.handle_dynamic_sampling(generator_output, uids) + if keep_sampling: # continue sampling + # update progress bar for current batch (but not global step) + pbar.update(1) + continue + + if self.colocate_all: + # if we are not continuing sampling, we sleep the inference engine + await self.inference_engine_client.sleep() + + # The train rollout for this step is done generating; close + # its metrics window. ``vllm/eval/*`` is collected separately + # around eval below. + vllm_metrics: Dict[str, float] = {} + if self._vllm_metrics_scraper is not None: + vllm_metrics = await self._vllm_metrics_scraper.stop() + + # 1.2 postprocess rewards (and merge step-wise turns if enabled) + with Timer("postprocess_generator_output", self.all_timings): + generator_output, uids = self.postprocess_generator_output(generator_output, uids) + + # 2. print example just for debugging + log_interval = self.cfg.trainer.log_example_interval + if log_interval > 0 and self.global_step % log_interval == 0: + vis = self.tokenizer.decode(generator_output["response_ids"][0]) + log_example( + logger, + prompt=generator_input["prompts"][0], + response=vis, + reward=generator_output["rewards"][0], + ) + + # 3. Convert GeneratorOutput to TrainingInputBatch + with Timer("convert_to_training_input", self.all_timings): + training_input: TrainingInputBatch = self.convert_to_training_input(generator_output, uids) + + # 4. Inference and calculate values, log probs, rewards, kl divergence + with Timer("fwd_logprobs_values_reward", self.all_timings): + training_input = self.fwd_logprobs_values_reward(training_input) + + # 5. apply kl divergence penalty to rewards + if self.cfg.trainer.algorithm.use_kl_in_reward: + with Timer("apply_reward_kl_penalty", self.all_timings): + training_input = self.apply_reward_kl_penalty(training_input) + + # 6. calculate advantages and returns + with Timer("compute_advantages_and_returns", self.all_timings): + training_input = self.compute_advantages_and_returns(training_input) + # remove some unwanted keys + for key in ["rewards"]: + training_input.pop(key) + training_input.metadata.pop("uids") + training_input.metadata.pop("is_last_step", None) + + if self.cfg.trainer.dump_data_batch: + # dump data to file + with Timer("dump_data_batch"): + self.dump_data(training_input, file_name=f"global_step_{self.global_step}_training_input") + + # 7. train policy/critic model + # Policy model is backloaded to GPU during training + with Timer("train_critic_and_policy", self.all_timings): + status = self.train_critic_and_policy(training_input) + + # Advance the torch profiler schedule once per global step + # (no-op unless profiling is enabled). One schedule step == + # one full RL global step; the schedule decides which are recorded. + self._profiler_step() + + self._fire("on_step_end", batch=training_input, metrics=status) + step_started = False + + # Capture callback-driven triggers, then reset. + force_save = self._training_control.should_save + force_eval = self._training_control.should_evaluate + self._training_control.should_save = False + self._training_control.should_evaluate = False + + # 8. conditionally save checkpoints and hf model + is_epoch_end = self.global_step % len(self.train_dataloader) == 0 + hf_model_save = self.cfg.trainer.hf_save_interval > 0 and ( + is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0 + ) + ckpt_interval_save = self.cfg.trainer.ckpt_interval > 0 and ( + is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0 + ) + will_save_ckpts = force_save or ckpt_interval_save + if will_save_ckpts: + with Timer("save_checkpoints", self.all_timings): + ckpt_path = self.save_checkpoints() + self._fire("on_save", ckpt_path=ckpt_path) + if hf_model_save: + with Timer("save_hf_model", self.all_timings): + self.save_models() + + # 9. conditionally sync policy and ref at the end of the epoch + if ( + self.cfg.trainer.update_ref_every_epoch + and self.ref_model is not None + and is_epoch_end + and epoch != self.cfg.trainer.epochs - 1 # skip updating ref at the end of the last epoch + ): + with Timer("update_ref_with_policy", self.all_timings): + self.update_ref_with_policy() + + # 10. Prepare weights for sampling + with Timer("sync_weights", self.all_timings): + await self.dispatch.save_weights_for_sampler() + + # 11. set logs + logger.info(status) + # Throughput metrics + train_time = self.all_timings.get("train_critic_and_policy", 0.0) + if train_time > 0 and training_input.get("attention_mask") is not None: + total_tokens = int(training_input["attention_mask"].sum().item()) + self.all_metrics["trainer/tokens_per_second_per_gpu"] = total_tokens / ( + train_time * self._num_training_gpus + ) + # log epoch info + self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step}) + interval_eval = self.cfg.trainer.eval_interval > 0 and ( + self.global_step % self.cfg.trainer.eval_interval == 0 + or self.global_step == self.total_training_steps ) - # log epoch info - self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step}) - interval_eval = self.cfg.trainer.eval_interval > 0 and ( - self.global_step % self.cfg.trainer.eval_interval == 0 - or self.global_step == self.total_training_steps - ) - if force_eval or interval_eval: - # Open the eval-rollout window; the scraper itself measures - # the generation spans via resume()/pause() inside eval(). - if self._vllm_metrics_scraper is not None: - await self._vllm_metrics_scraper.start("vllm/eval") - self._vllm_metrics_scraper.pause() - self._fire("on_eval_start") - with Timer("eval", self.all_timings): - eval_metrics = await self.eval(vllm_metrics_scraper=self._vllm_metrics_scraper) - self.all_metrics.update(eval_metrics) - self._fire("on_eval_end", metrics=eval_metrics) - if self._vllm_metrics_scraper is not None: - vllm_metrics.update(await self._vllm_metrics_scraper.stop()) - - log_payload = { - **self.all_metrics, - **{f"timing/{k}": v for k, v in self.all_timings.items()}, - # vllm/train/* = train rollout, vllm/eval/* = eval rollout, - # each over its own generation time (owned by the scraper). - **vllm_metrics, - } + if force_eval or interval_eval: + # Open the eval-rollout window; the scraper itself measures + # the generation spans via resume()/pause() inside eval(). + if self._vllm_metrics_scraper is not None: + await self._vllm_metrics_scraper.start("vllm/eval") + self._vllm_metrics_scraper.pause() + self._fire("on_eval_start") + with Timer("eval", self.all_timings): + eval_metrics = await self.eval(vllm_metrics_scraper=self._vllm_metrics_scraper) + self.all_metrics.update(eval_metrics) + self._fire("on_eval_end", metrics=eval_metrics) + if self._vllm_metrics_scraper is not None: + vllm_metrics.update(await self._vllm_metrics_scraper.stop()) + + log_payload = { + **self.all_metrics, + **{f"timing/{k}": v for k, v in self.all_timings.items()}, + # vllm/train/* = train rollout, vllm/eval/* = eval rollout, + # each over its own generation time (owned by the scraper). + **vllm_metrics, + } - if self._ray_gpu_monitor is not None: - log_payload.update(self._ray_gpu_monitor.flush()) + if self._ray_gpu_monitor is not None: + log_payload.update(self._ray_gpu_monitor.flush()) - self._fire("on_log", logs=log_payload) + self._fire("on_log", logs=log_payload) - self.tracker.log(log_payload, step=self.global_step, commit=True) - self.all_metrics = {} - self.all_timings = {} + self.tracker.log(log_payload, step=self.global_step, commit=True) + self.all_metrics = {} + self.all_timings = {} - # update progress bar after logging - pbar.update(1) + # update progress bar after logging + pbar.update(1) - self.global_step += 1 + self.global_step += 1 - if ( - self.cfg.trainer.max_training_steps is not None - and self.global_step > self.cfg.trainer.max_training_steps - ): - logger.info(f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early.") - stop_training = True - break + if ( + self.cfg.trainer.max_training_steps is not None + and self.global_step > self.cfg.trainer.max_training_steps + ): + logger.info(f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early.") + stop_training = True + break - del training_input, generator_output + del training_input, generator_output - self._fire("on_epoch_end") + self._fire("on_epoch_end") - if stop_training: - break + if stop_training: + break + finally: + # Always stop/flush the profiler when the training loop exits -- + # including via an exception -- so the open kineto trace window + # isn't leaked. No-op when profiling is disabled. + self._profiler_stop() pbar.close() if self.colocate_all: diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 8abbe05db0..342ae3a8c8 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -270,6 +270,8 @@ def validate_cfg(cfg: SkyRLTrainConfig): "or negative to keep all checkpoints" ) + cfg.trainer.policy.torch_profiler_config.validate() + # TODO (devpatel): move to initializing ray and syncing registries codepath at startup repopulate_all_registries() available_policy_losses = PolicyLossRegistry.list_available() diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py index 8ce0ce6c2a..2b5a1bf571 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_worker.py @@ -18,9 +18,9 @@ from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch from skyrl.backends.skyrl_train.utils.torch_utils import logprobs_from_logits from skyrl.train.config import ( - MegatronTorchProfilerConfig, SkyRLLoraConfig, SkyRLTrainConfig, + TorchProfilerConfig, ) from skyrl.train.utils.utils import print_mem, validate_cfg from tests.backends.skyrl_train.gpu.utils import ( @@ -684,7 +684,7 @@ async def test_megatron_dp(ray_init_fixture, worker_type, tp, pp, gpus_per_node) cfg.trainer.micro_train_batch_size_per_gpu = 4 # set torch profiler config - cfg.trainer.policy.megatron_config.torch_profiler_config = MegatronTorchProfilerConfig( + cfg.trainer.policy.torch_profiler_config = TorchProfilerConfig( enable=False, ranks=[0], save_path=f"/home/ray/megatron_prof/tp{tp}_pp{pp}/" ) @@ -726,7 +726,7 @@ async def test_megatron_dp(ray_init_fixture, worker_type, tp, pp, gpus_per_node) cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 # set torch profiler config - cfg.trainer.policy.megatron_config.torch_profiler_config = MegatronTorchProfilerConfig( + cfg.trainer.policy.torch_profiler_config = TorchProfilerConfig( enable=False, ranks=[0], save_path="/home/ray/megatron_prof/dp4/" ) diff --git a/tests/backends/skyrl_train/utils/test_profiler.py b/tests/backends/skyrl_train/utils/test_profiler.py new file mode 100644 index 0000000000..2c0df310ee --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_profiler.py @@ -0,0 +1,394 @@ +"""CPU unit tests for the config-driven torch.profiler wrapper. + +These run without a GPU or torch.distributed: ``torch.profiler`` works on CPU, +and ``Profiler`` falls back to rank 0 when distributed is not initialized. + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/utils/test_profiler.py -v +""" + +import glob +import os +from dataclasses import dataclass, field +from typing import List, Optional +from unittest.mock import patch + +from skyrl.backends.skyrl_train.utils.profiler import Profiler + + +@dataclass +class _ProfCfg: + """Minimal stand-in for TorchProfilerConfig (avoids importing skyrl_gym).""" + + enable: bool = True + ranks: List[int] = field(default_factory=lambda: [0]) + save_path: Optional[str] = None + skip_first: int = 0 + wait: int = 0 + warmup: int = 0 + active: int = 1 + repeat: int = 1 + activities: List[str] = field(default_factory=lambda: ["cpu"]) + record_shapes: bool = False + profile_memory: bool = False + with_stack: bool = False + with_flops: bool = False + with_modules: bool = False + export_type: str = "chrome_trace" + + +def _run_loop(prof: Profiler, n_steps: int) -> None: + prof.start() + for _ in range(n_steps): + prof.step() + prof.stop() + + +def test_disabled_is_noop(tmp_path): + prof = Profiler(_ProfCfg(enable=False), default_save_path=str(tmp_path)) + assert prof.prof is None + assert prof.check() is False + _run_loop(prof, 5) # must not raise + assert glob.glob(os.path.join(str(tmp_path), "*")) == [] + + +def test_rank_not_selected_is_noop(tmp_path): + # is_initialized() is False in this process -> rank resolves to 0, which is + # not in ranks=[1], so the profiler must not arm. + prof = Profiler(_ProfCfg(ranks=[1]), default_save_path=str(tmp_path)) + assert prof.prof is None + _run_loop(prof, 3) + assert glob.glob(os.path.join(str(tmp_path), "*")) == [] + + +def test_single_window_writes_one_trace(tmp_path): + prof = Profiler( + _ProfCfg(skip_first=0, wait=0, warmup=0, active=1, repeat=1), + default_save_path=str(tmp_path), + ) + assert prof.check() is True + _run_loop(prof, 3) + traces = glob.glob(os.path.join(str(tmp_path), "*.pt.trace.json*")) + assert len(traces) == 1, f"expected exactly one trace, got {traces}" + assert "rank0" in os.path.basename(traces[0]) + + +def test_repeat_writes_multiple_windows(tmp_path): + # Two cycles of (warmup=1 + active=1); needs >= 2*(1+1) steps to fire twice. + prof = Profiler( + _ProfCfg(skip_first=0, wait=0, warmup=1, active=1, repeat=2), + default_save_path=str(tmp_path), + ) + _run_loop(prof, 8) + traces = glob.glob(os.path.join(str(tmp_path), "*.pt.trace.json*")) + assert len(traces) == 2, f"expected two traces for repeat=2, got {traces}" + + +def test_skip_first_defers_recording(tmp_path): + # With skip_first larger than the loop, no window should ever open. + prof = Profiler( + _ProfCfg(skip_first=100, wait=0, warmup=0, active=1, repeat=1), + default_save_path=str(tmp_path), + ) + _run_loop(prof, 5) + assert glob.glob(os.path.join(str(tmp_path), "*.pt.trace.json*")) == [] + + +def test_save_path_overrides_default(tmp_path): + explicit = tmp_path / "explicit" + prof = Profiler( + _ProfCfg(save_path=str(explicit)), + default_save_path=str(tmp_path / "default"), + ) + assert prof.save_path == str(explicit) + + +def test_default_save_path_used_when_none(tmp_path): + default = str(tmp_path / "default") + prof = Profiler(_ProfCfg(save_path=None), default_save_path=default) + assert prof.save_path == default + + +def test_kernel_summary_none_when_disabled(tmp_path): + prof = Profiler(_ProfCfg(enable=False), default_save_path=str(tmp_path)) + assert prof.get_kernel_summary() is None + + +def test_kernel_summary_empty_before_first_window(tmp_path): + prof = Profiler(_ProfCfg(skip_first=0, wait=0, warmup=0, active=1), default_save_path=str(tmp_path)) + summary = prof.get_kernel_summary() + assert summary == {"window_count": 0, "pairs": []} + + +def test_kernel_summary_populated_after_window(tmp_path): + prof = Profiler( + _ProfCfg(skip_first=0, wait=0, warmup=0, active=1, repeat=1, activities=["cpu"]), + default_save_path=str(tmp_path), + ) + prof.start() + for _ in range(3): + # Do a little CPU work so the profiler records some ops. + import torch + + _ = torch.randn(64, 64) @ torch.randn(64, 64) + prof.step() + prof.stop() + + summary = prof.get_kernel_summary() + assert summary is not None + assert summary["window_count"] == 1 + assert isinstance(summary["pairs"], list) + # Each pair is (name:str, self_us:float); summary is pickle-safe. + import pickle + + pickle.dumps(summary) + for name, self_us in summary["pairs"]: + assert isinstance(name, str) + assert isinstance(self_us, float) + + +def test_activities_threaded_to_torch(tmp_path): + import torch + + prof = Profiler(_ProfCfg(activities=["cpu"]), default_save_path=str(tmp_path)) + # The underlying torch profiler was constructed with the CPU activity only. + assert torch.profiler.ProfilerActivity.CPU in prof.prof.activities + assert torch.profiler.ProfilerActivity.CUDA not in prof.prof.activities + + +def test_step_failure_disables_without_raising(tmp_path): + prof = Profiler(_ProfCfg(), default_save_path=str(tmp_path)) + prof.start() + + class _Boom: + def step(self): + raise RuntimeError("boom") + + # Simulate an internal profiler fault mid-loop; the wrapper must swallow it. + prof.prof = _Boom() + prof.step() + assert prof.enable is False + assert prof.prof is None + # Subsequent calls remain safe no-ops. + prof.step() + prof.stop() + + +class TestWorkerProfilerRPCs: + """The Worker base exposes the profiler-control RPCs and they no-op when + ``self.profiler`` is None (so the trainer can dispatch them unconditionally, + and so they're on the snapshotted Ray actor method table without a subclass).""" + + def test_methods_exist_on_worker_base(self): + from skyrl.backends.skyrl_train.workers.worker import Worker + + for name in ("start_profile", "profile_step", "stop_profile", "dump_profiler_summary"): + assert callable(getattr(Worker, name)), f"Worker.{name} missing" + + def test_dump_profiler_summary_none_when_profiler_none(self): + from types import SimpleNamespace + + from skyrl.backends.skyrl_train.workers.worker import Worker + + stub = SimpleNamespace(profiler=None) + assert Worker.dump_profiler_summary(stub) is None + + def test_dump_profiler_summary_delegates_to_profiler(self): + from types import SimpleNamespace + + from skyrl.backends.skyrl_train.workers.worker import Worker + + expected = {"window_count": 2, "pairs": [("gemm", 1.5)]} + stub = SimpleNamespace(profiler=SimpleNamespace(get_kernel_summary=lambda: expected)) + assert Worker.dump_profiler_summary(stub) == expected + + def test_rpcs_noop_when_profiler_none(self): + from types import SimpleNamespace + + from skyrl.backends.skyrl_train.workers.worker import Worker + + # Call the unbound methods against a stub whose profiler is None; they + # must early-return without touching anything. + stub = SimpleNamespace(profiler=None) + Worker.start_profile(stub) + Worker.profile_step(stub) + Worker.stop_profile(stub) + + def test_rpcs_drive_profiler_when_present(self): + from types import SimpleNamespace + + from skyrl.backends.skyrl_train.workers.worker import Worker + + calls = [] + fake_profiler = SimpleNamespace( + start=lambda: calls.append("start"), + step=lambda: calls.append("step"), + stop=lambda: calls.append("stop"), + ) + stub = SimpleNamespace(profiler=fake_profiler) + Worker.start_profile(stub) + Worker.profile_step(stub) + Worker.stop_profile(stub) + assert calls == ["start", "step", "stop"] + + +class TestBuildProfilerFromPolicyCfg: + """``build_profiler_from_policy_cfg`` is the production entry both worker + backends call in ``init_model``. It gates on ``enable`` and composes the + default ``{ckpt_path}/profiler_traces`` save path when none is set.""" + + @staticmethod + def _trainer_cfg(prof_cfg, ckpt_path): + from types import SimpleNamespace + + return SimpleNamespace(ckpt_path=ckpt_path, policy=SimpleNamespace(torch_profiler_config=prof_cfg)) + + def test_returns_none_when_disabled(self, tmp_path): + from skyrl.backends.skyrl_train.utils.profiler import ( + build_profiler_from_policy_cfg, + ) + + cfg = self._trainer_cfg(_ProfCfg(enable=False), str(tmp_path)) + assert build_profiler_from_policy_cfg(cfg) is None + + def test_builds_profiler_with_ckpt_default_save_path(self, tmp_path): + from skyrl.backends.skyrl_train.utils.profiler import ( + Profiler, + build_profiler_from_policy_cfg, + ) + + cfg = self._trainer_cfg(_ProfCfg(enable=True, save_path=None), str(tmp_path)) + prof = build_profiler_from_policy_cfg(cfg) + assert isinstance(prof, Profiler) + assert prof.save_path == os.path.join(str(tmp_path), "profiler_traces") + + def test_explicit_save_path_wins_over_ckpt_default(self, tmp_path): + from skyrl.backends.skyrl_train.utils.profiler import ( + build_profiler_from_policy_cfg, + ) + + explicit = str(tmp_path / "explicit") + cfg = self._trainer_cfg(_ProfCfg(enable=True, save_path=explicit), str(tmp_path)) + prof = build_profiler_from_policy_cfg(cfg) + assert prof.save_path == explicit + + +class TestWorkerDispatchProfilerRPCs: + """The WorkerDispatch profiler wrappers dispatch via ``pass_through`` to the + named model's actor group, no-op for an unknown model, and swallow dispatch + faults so a profiler error can never abort the training loop. Tested against + the unbound methods with a stub ``self`` (no live Ray actors required).""" + + @staticmethod + def _stub(actor_groups): + from types import SimpleNamespace + + return SimpleNamespace(_actor_groups=actor_groups) + + def _fake_group(self, calls, raises=False): + from types import SimpleNamespace + + def async_run_ray_method(mode, method, *args, **kwargs): + calls.append((mode, method)) + if raises: + raise RuntimeError("dispatch boom") + return ["sentinel"] + + return SimpleNamespace(async_run_ray_method=async_run_ray_method) + + def test_unknown_model_is_noop(self): + from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch + + stub = self._stub({}) # no "policy" group + # None of these should raise or dispatch. + WorkerDispatch.start_profile(stub, "policy") + WorkerDispatch.profile_step(stub, "policy") + WorkerDispatch.stop_profile(stub, "policy") + assert WorkerDispatch.dump_profiler_summary(stub, "policy") is None + + def test_control_rpcs_dispatch_pass_through(self): + from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch + + calls = [] + stub = self._stub({"policy": self._fake_group(calls)}) + with patch("skyrl.backends.skyrl_train.workers.worker_dispatch.ray.get", side_effect=lambda x: x): + WorkerDispatch.start_profile(stub, "policy") + WorkerDispatch.profile_step(stub, "policy") + WorkerDispatch.stop_profile(stub, "policy") + assert calls == [ + ("pass_through", "start_profile"), + ("pass_through", "profile_step"), + ("pass_through", "stop_profile"), + ] + + def test_dump_profiler_summary_returns_per_rank_payload(self): + from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch + + calls = [] + payload = [{"window_count": 1, "pairs": [("gemm", 2.0)]}, None] + stub = self._stub({"policy": self._fake_group(calls)}) + with patch("skyrl.backends.skyrl_train.workers.worker_dispatch.ray.get", side_effect=lambda x: payload): + out = WorkerDispatch.dump_profiler_summary(stub, "policy") + assert out == payload + assert calls == [("pass_through", "dump_profiler_summary")] + + def test_dispatch_fault_is_swallowed(self): + from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch + + calls = [] + stub = self._stub({"policy": self._fake_group(calls, raises=True)}) + + def boom(_): + raise RuntimeError("ray.get boom") + + with patch("skyrl.backends.skyrl_train.workers.worker_dispatch.ray.get", side_effect=boom): + # Must not propagate -- a profiler fault cannot abort training. + WorkerDispatch.start_profile(stub, "policy") + WorkerDispatch.profile_step(stub, "policy") + WorkerDispatch.stop_profile(stub, "policy") + assert WorkerDispatch.dump_profiler_summary(stub, "policy") is None + + +class TestTrainerProfilerHelpers: + """``RayPPOTrainer`` exposes gated ``_profiler_{start,step,stop}`` helpers so + every trainer flavor (sync / fully-async / one-step async) drives the + profiler through one code path. The helpers dispatch only when + ``_torch_profiler_enabled`` is True, and always target the policy model.""" + + @staticmethod + def _trainer(enable): + from types import SimpleNamespace + + from skyrl.train.trainer import RayPPOTrainer + + # Build a bare instance without running __init__ (which needs Ray/cfg). + trainer = object.__new__(RayPPOTrainer) + calls = [] + trainer.dispatch = SimpleNamespace( + start_profile=lambda m: calls.append(("start", m)), + profile_step=lambda m: calls.append(("step", m)), + stop_profile=lambda m: calls.append(("stop", m)), + ) + trainer.cfg = SimpleNamespace( + trainer=SimpleNamespace(policy=SimpleNamespace(torch_profiler_config=SimpleNamespace(enable=enable))) + ) + return trainer, calls + + def test_helpers_exist(self): + from skyrl.train.trainer import RayPPOTrainer + + for name in ("_profiler_start", "_profiler_step", "_profiler_stop"): + assert callable(getattr(RayPPOTrainer, name)), f"RayPPOTrainer.{name} missing" + + def test_noop_when_disabled(self): + trainer, calls = self._trainer(enable=False) + trainer._profiler_start() + trainer._profiler_step() + trainer._profiler_stop() + assert calls == [] + + def test_dispatch_to_policy_when_enabled(self): + trainer, calls = self._trainer(enable=True) + trainer._profiler_start() + trainer._profiler_step() + trainer._profiler_stop() + assert calls == [("start", "policy"), ("step", "policy"), ("stop", "policy")] diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 2c6853c79b..6961c37fd2 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -422,3 +422,61 @@ def test_validate_cfg_allows_explicit_max_seq_len_for_seq_mean_token_sum_norm(se cfg.trainer.algorithm.max_seq_len = 4096 validate_cfg(cfg) + + +class TestTorchProfilerConfigValidation: + """``TorchProfilerConfig.validate()`` rejects unusable profiler settings up + front (so an enabled run fails fast instead of silently degrading), and is a + no-op when profiling is disabled.""" + + @staticmethod + def _cfg(**overrides): + from skyrl.train.config.config import TorchProfilerConfig + + return TorchProfilerConfig(enable=True, **overrides) + + def test_disabled_skips_all_checks(self): + from skyrl.train.config.config import TorchProfilerConfig + + # Garbage values must be tolerated while disabled (the default state). + TorchProfilerConfig(enable=False, export_type="bogus", activities=["gpu"], ranks=[], active=0).validate() + + def test_defaults_are_valid_when_enabled(self): + self._cfg().validate() # must not raise + + def test_empty_ranks_rejected(self): + with pytest.raises(ValueError, match=r"ranks.*non-empty"): + self._cfg(ranks=[]).validate() + + def test_unknown_activity_rejected(self): + with pytest.raises(ValueError, match=r"activities"): + self._cfg(activities=["cpu", "gpu"]).validate() + + def test_activities_case_insensitive(self): + self._cfg(activities=["CPU", "CUDA"]).validate() # must not raise + + def test_unknown_export_type_rejected(self): + with pytest.raises(ValueError, match=r"export_type"): + self._cfg(export_type="bogus").validate() + + def test_stacks_requires_with_stack(self): + with pytest.raises(ValueError, match=r"with_stack"): + self._cfg(export_type="stacks", with_stack=False).validate() + # With with_stack=True it is accepted. + self._cfg(export_type="stacks", with_stack=True).validate() + + def test_negative_schedule_field_rejected(self): + with pytest.raises(ValueError, match=r"skip_first"): + self._cfg(skip_first=-1).validate() + + def test_active_must_be_at_least_one(self): + with pytest.raises(ValueError, match=r"active"): + self._cfg(active=0).validate() + + def test_validate_cfg_invokes_profiler_validation(self): + # The RL entrypoint validator must surface profiler config errors. + cfg = _make_validated_test_config() + cfg.trainer.policy.torch_profiler_config.enable = True + cfg.trainer.policy.torch_profiler_config.export_type = "bogus" + with pytest.raises(ValueError, match=r"export_type"): + validate_cfg(cfg) diff --git a/tests/train/test_sft_config.py b/tests/train/test_sft_config.py index 621b49be72..3aa3312646 100644 --- a/tests/train/test_sft_config.py +++ b/tests/train/test_sft_config.py @@ -134,6 +134,61 @@ def test_lora_disabled_by_default(self): assert skyrl_cfg.trainer.policy.model.lora.rank == 0 +class TestTorchProfilerConfigOverrides: + """torch_profiler_config bridges to policy.torch_profiler_config.""" + + def test_disabled_by_default(self): + cfg = _sft_cfg_from_overrides([]) + skyrl_cfg = build_skyrl_config_for_sft(cfg) + assert skyrl_cfg.trainer.policy.torch_profiler_config.enable is False + + def test_enable_and_schedule_propagate(self): + cfg = _sft_cfg_from_overrides( + [ + "torch_profiler_config.enable=true", + "torch_profiler_config.skip_first=3", + "torch_profiler_config.active=2", + "torch_profiler_config.repeat=0", + "torch_profiler_config.save_path=/tmp/sft_prof", + ] + ) + skyrl_cfg = build_skyrl_config_for_sft(cfg) + prof = skyrl_cfg.trainer.policy.torch_profiler_config + assert prof.enable is True + assert prof.skip_first == 3 + assert prof.active == 2 + assert prof.repeat == 0 + assert prof.save_path == "/tmp/sft_prof" + + def test_capture_flags_propagate(self): + cfg = _sft_cfg_from_overrides( + [ + "torch_profiler_config.enable=true", + "torch_profiler_config.profile_memory=true", + "torch_profiler_config.with_stack=false", + "torch_profiler_config.activities=[cuda]", + ] + ) + skyrl_cfg = build_skyrl_config_for_sft(cfg) + prof = skyrl_cfg.trainer.policy.torch_profiler_config + assert prof.profile_memory is True + assert prof.with_stack is False + assert list(prof.activities) == ["cuda"] + + def test_invalid_profiler_config_rejected_by_sft_validation(self): + # validate_sft_cfg (run by build_skyrl_config_for_sft) must surface + # profiler config errors on the SFT path, not just the RL path. + cfg = _sft_cfg_from_overrides( + [ + "model.path=test/my-model", + "torch_profiler_config.enable=true", + "torch_profiler_config.export_type=bogus", + ] + ) + with pytest.raises(ValueError, match=r"export_type"): + build_skyrl_config_for_sft(cfg) + + class TestFSDPConfigOverrides: """FSDP config overrides propagate when strategy=fsdp.""" From 067baaa47320951696c67d31e8bbce6aec32bd44 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Mon, 8 Jun 2026 14:48:51 -0500 Subject: [PATCH 2/7] feat(profiler): drive torch.profiler in the full-context trainer Bracket FullCtxTrainer.train() with _profiler_start / _profiler_step / _profiler_stop, matching the wiring in RayPPOTrainer.train(). One profile_step per dummy global step; stop runs in a finally so the open kineto trace window isn't leaked if a step raises. No-op unless torch_profiler_config.enable is set, so non-profiling runs pay nothing. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../full_context/trainer_full_ctx.py | 141 ++++++++++-------- 1 file changed, 79 insertions(+), 62 deletions(-) diff --git a/examples/train_scripts/full_context/trainer_full_ctx.py b/examples/train_scripts/full_context/trainer_full_ctx.py index 6022639d9d..3160c2358c 100644 --- a/examples/train_scripts/full_context/trainer_full_ctx.py +++ b/examples/train_scripts/full_context/trainer_full_ctx.py @@ -26,68 +26,85 @@ async def train(self): # Run a few training steps self.global_step += 1 # start from 1 - for step in range(self.cfg.trainer.num_dummy_steps): - logger.info(f"Running dummy training step {step + 1}/{self.cfg.trainer.num_dummy_steps}") - - # Run a single training step - with Timer("step", self.all_timings): - # Create training input directly with max length sequences - num_samples = self.cfg.trainer.train_batch_size * self.cfg.generator.n_samples_per_prompt - uids = [str(i) for i in range(self.cfg.trainer.train_batch_size)] - prompt_token_ids = [ - [random.randint(0, self.tokenizer.vocab_size - 1)] * self.cfg.generator.max_input_length - ] * self.cfg.trainer.train_batch_size - prompt_token_ids = sum( - [ - [prompt_token_id] * self.cfg.generator.n_samples_per_prompt - for prompt_token_id in prompt_token_ids - ], - [], - ) - response_ids = [ - [random.randint(0, self.tokenizer.vocab_size - 1)] - * self.cfg.generator.sampling_params.max_generate_length - ] * num_samples - uids = sum([[uid] * self.cfg.generator.n_samples_per_prompt for uid in uids], []) - - dummy_generator_output = { - "prompt_token_ids": prompt_token_ids, - "response_ids": response_ids, - "rewards": [ - [0] * (self.cfg.generator.sampling_params.max_generate_length - 1) + [random.randint(0, 1)] - ] - * num_samples, - "loss_masks": [[1] * self.cfg.generator.sampling_params.max_generate_length] * num_samples, - } - training_input = self.convert_to_training_input(dummy_generator_output, uids) - - with Timer("fwd_logprobs_values_reward", self.all_timings): - training_input = self.fwd_logprobs_values_reward(training_input) - - # 1.5 apply kl divergence penalty to rewards - if self.cfg.trainer.algorithm.use_kl_in_reward: - with Timer("apply_reward_kl_penalty", self.all_timings): - training_input = self.apply_reward_kl_penalty(training_input) - - # 3. calculate advantages and returns - with Timer("compute_advantages_and_returns", self.all_timings): - training_input = self.compute_advantages_and_returns(training_input) - # remove some unwanted keys - for key in ["rewards"]: - training_input.pop(key) - training_input.metadata.pop("uids") - - # 4. train policy/critic model - with Timer("train_critic_and_policy", self.all_timings): - status = self.train_critic_and_policy(training_input) - - self.tracker.log(self.all_metrics, step=self.global_step) - self.all_metrics = {} - self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step) - self.all_timings = {} - self.global_step += 1 - - logger.info(f"Step {step + 1} completed. Status: {status}") + self._profiler_start() + try: + for step in range(self.cfg.trainer.num_dummy_steps): + logger.info(f"Running dummy training step {step + 1}/{self.cfg.trainer.num_dummy_steps}") + + # Run a single training step + with Timer("step", self.all_timings): + # Create training input directly with max length sequences + num_samples = self.cfg.trainer.train_batch_size * self.cfg.generator.n_samples_per_prompt + uids = [str(i) for i in range(self.cfg.trainer.train_batch_size)] + prompt_token_ids = [ + [random.randint(0, self.tokenizer.vocab_size - 1)] * self.cfg.generator.max_input_length + ] * self.cfg.trainer.train_batch_size + prompt_token_ids = sum( + [ + [prompt_token_id] * self.cfg.generator.n_samples_per_prompt + for prompt_token_id in prompt_token_ids + ], + [], + ) + response_ids = [ + [random.randint(0, self.tokenizer.vocab_size - 1)] + * self.cfg.generator.sampling_params.max_generate_length + ] * num_samples + uids = sum( + [[uid] * self.cfg.generator.n_samples_per_prompt for uid in uids], + [], + ) + + dummy_generator_output = { + "prompt_token_ids": prompt_token_ids, + "response_ids": response_ids, + "rewards": [ + [0] * (self.cfg.generator.sampling_params.max_generate_length - 1) + [random.randint(0, 1)] + ] + * num_samples, + "loss_masks": [[1] * self.cfg.generator.sampling_params.max_generate_length] * num_samples, + } + training_input = self.convert_to_training_input(dummy_generator_output, uids) + + with Timer("fwd_logprobs_values_reward", self.all_timings): + training_input = self.fwd_logprobs_values_reward(training_input) + + # 1.5 apply kl divergence penalty to rewards + if self.cfg.trainer.algorithm.use_kl_in_reward: + with Timer("apply_reward_kl_penalty", self.all_timings): + training_input = self.apply_reward_kl_penalty(training_input) + + # 3. calculate advantages and returns + with Timer("compute_advantages_and_returns", self.all_timings): + training_input = self.compute_advantages_and_returns(training_input) + # remove some unwanted keys + for key in ["rewards"]: + training_input.pop(key) + training_input.metadata.pop("uids") + + # 4. train policy/critic model + with Timer("train_critic_and_policy", self.all_timings): + status = self.train_critic_and_policy(training_input) + + # Advance the torch profiler schedule once per global step + # (no-op unless profiling is enabled). + self._profiler_step() + + self.tracker.log(self.all_metrics, step=self.global_step) + self.all_metrics = {} + self.tracker.log( + {"timing/" + k: v for k, v in self.all_timings.items()}, + step=self.global_step, + ) + self.all_timings = {} + self.global_step += 1 + + logger.info(f"Step {step + 1} completed. Status: {status}") + finally: + # Always stop/flush the profiler when the loop exits -- including via + # an exception -- so the open kineto trace window isn't leaked. No-op + # when profiling is disabled. + self._profiler_stop() self.tracker.finish() logger.info("Dummy training completed successfully!") From e458d560ecb6b0b641b6312950b75afbcf28a75a Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Mon, 8 Jun 2026 15:59:43 -0500 Subject: [PATCH 3/7] feat(profiler): fall back to local save_path for cloud ckpt_path; reject empty activities Two robustness fixes from PR review: - Profiler.__init__ detects a cloud save_path (s3://, gs://, gcs:// via the existing is_cloud_path helper) and falls back to ./profiler_traces with a warning. save_path commonly defaults to {ckpt_path}/profiler_traces, and ckpt_path can be a cloud URI -- which torch.profiler can't write to, so the trace would otherwise be silently lost. - TorchProfilerConfig.validate rejects an empty `activities` list. It passed the membership check vacuously, but torch.profiler.profile(activities=[]) records nothing -- now it fails fast at startup. Co-Authored-By: Claude Opus 4.8 (1M context) --- skyrl/backends/skyrl_train/utils/profiler.py | 12 ++++++++++++ skyrl/train/config/config.py | 4 ++++ tests/backends/skyrl_train/utils/test_profiler.py | 8 ++++++++ tests/train/test_config.py | 6 ++++++ 4 files changed, 30 insertions(+) diff --git a/skyrl/backends/skyrl_train/utils/profiler.py b/skyrl/backends/skyrl_train/utils/profiler.py index 23988fa8f2..14130e3e81 100644 --- a/skyrl/backends/skyrl_train/utils/profiler.py +++ b/skyrl/backends/skyrl_train/utils/profiler.py @@ -4,6 +4,8 @@ import torch.distributed from loguru import logger +from skyrl.backends.skyrl_train.utils.io.io import is_cloud_path + # Map config activity strings to torch ProfilerActivity members. _ACTIVITY_MAP = { "cpu": torch.profiler.ProfilerActivity.CPU, @@ -61,6 +63,16 @@ def __init__(self, config, default_save_path: str = None): return self.config = config self.save_path = config.save_path or default_save_path or "./profiler_traces" + # torch.profiler writes traces with the local filesystem only. ``save_path`` + # commonly defaults to ``{ckpt_path}/profiler_traces``, and ckpt_path can be a + # cloud URI (s3://, gs://) -- which torch can't write to, so the trace would be + # silently lost. Fall back to a local dir so profiling still produces output. + if is_cloud_path(self.save_path): + logger.warning( + f"[Profiler] cloud save_path {self.save_path!r} is not writable by torch.profiler; " + f"falling back to local './profiler_traces'." + ) + self.save_path = "./profiler_traces" self.ranks = list(config.ranks) self.export_type = getattr(config, "export_type", "chrome_trace") self.rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index dab4e7683e..3853bf67bc 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -215,6 +215,10 @@ def validate(self) -> None: return if not self.ranks: raise ValueError("`torch_profiler_config.ranks` must be non-empty when profiling is enabled.") + # An empty `activities` passes the membership check below vacuously, but + # `torch.profiler.profile(activities=[])` records nothing -- fail fast instead. + if not self.activities: + raise ValueError("`torch_profiler_config.activities` must be non-empty when profiling is enabled.") bad_activities = [a for a in self.activities if a.lower() not in TORCH_PROFILER_ACTIVITIES] if bad_activities: raise ValueError( diff --git a/tests/backends/skyrl_train/utils/test_profiler.py b/tests/backends/skyrl_train/utils/test_profiler.py index 2c0df310ee..e80757a1fc 100644 --- a/tests/backends/skyrl_train/utils/test_profiler.py +++ b/tests/backends/skyrl_train/utils/test_profiler.py @@ -108,6 +108,14 @@ def test_default_save_path_used_when_none(tmp_path): assert prof.save_path == default +def test_cloud_save_path_falls_back_to_local(tmp_path): + # ckpt_path (and thus the derived default save_path) can be a cloud URI, which + # torch.profiler can't write to. The wrapper must fall back to a local dir so + # the trace isn't silently lost. + prof = Profiler(_ProfCfg(save_path="s3://my-bucket/run/profiler_traces"), default_save_path=str(tmp_path)) + assert prof.save_path == "./profiler_traces" + + def test_kernel_summary_none_when_disabled(tmp_path): prof = Profiler(_ProfCfg(enable=False), default_save_path=str(tmp_path)) assert prof.get_kernel_summary() is None diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 6961c37fd2..f6e8fced0a 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -452,6 +452,12 @@ def test_unknown_activity_rejected(self): with pytest.raises(ValueError, match=r"activities"): self._cfg(activities=["cpu", "gpu"]).validate() + def test_empty_activities_rejected(self): + # An empty list profiles nothing; the membership check would pass it + # vacuously, so it needs its own guard. + with pytest.raises(ValueError, match=r"activities.*non-empty"): + self._cfg(activities=[]).validate() + def test_activities_case_insensitive(self): self._cfg(activities=["CPU", "CUDA"]).validate() # must not raise From ce2588e3abf9977b7c5cf107e9cf7c930f9f19c3 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Mon, 15 Jun 2026 12:42:50 -0500 Subject: [PATCH 4/7] feat(profiler): require explicit local save_path; drop implicit ckpt-derived default Per review: remove the non-obvious {ckpt_path}/profiler_traces default and the silent cloud-URI -> ./profiler_traces fallback. The relative fallback would land traces in the Ray runtime working dir under /tmp/ray, and the ckpt-derived default was surprising. save_path is now a required, explicit local path, validated fail-fast at startup (non-empty + not a cloud URI). Updated docs, config defaults comment, and tests accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/content/docs/configuration/config.mdx | 4 +- skyrl/backends/skyrl_train/utils/profiler.py | 26 ++----- skyrl/train/config/config.py | 21 ++++- .../skyrl_train/utils/test_profiler.py | 77 ++++++------------- tests/train/test_config.py | 21 ++++- tests/train/test_sft_config.py | 2 + 6 files changed, 76 insertions(+), 75 deletions(-) diff --git a/docs/content/docs/configuration/config.mdx b/docs/content/docs/configuration/config.mdx index 6ac77e22d4..4c29535151 100644 --- a/docs/content/docs/configuration/config.mdx +++ b/docs/content/docs/configuration/config.mdx @@ -276,7 +276,7 @@ policy: torch_profiler_config: # torch.profiler-based training-loop profiler (see below) enable: false ranks: [0] - save_path: null # defaults to {ckpt_path}/profiler_traces + save_path: null # REQUIRED when enable=true; absolute local path skip_first: 10 # torch.profiler.schedule wait: 0 warmup: 1 @@ -310,7 +310,7 @@ Which steps are recorded is controlled entirely by [`torch.profiler.schedule`](h - `policy.torch_profiler_config.enable`: Master switch. When `false` (default), no profiler RPCs are dispatched and there is zero overhead. - `policy.torch_profiler_config.ranks`: List of global ranks to profile (e.g. `[0]`). Add a mid-pipeline-stage rank to diagnose pipeline bubbles. -- `policy.torch_profiler_config.save_path`: Output directory for traces. Defaults to `{ckpt_path}/profiler_traces`. +- `policy.torch_profiler_config.save_path`: Output directory for traces. **Required** when `enable: true` (there is no implicit default). Use an absolute local path: Ray workers run from a `/tmp/ray/.../working_dir_files` runtime working dir, so a relative path would scatter traces there, and `torch.profiler` cannot write to cloud URIs (`s3://`/`gs://`) — both are rejected at startup. - `policy.torch_profiler_config.{skip_first,wait,warmup,active,repeat}`: Passed directly to `torch.profiler.schedule`. - `policy.torch_profiler_config.activities`: Subset of `["cpu", "cuda"]` to record. - `policy.torch_profiler_config.{record_shapes,profile_memory,with_stack,with_flops,with_modules}`: Passed directly to `torch.profiler.profile`. `with_stack` and `record_shapes` add overhead; `profile_memory` is heavier still and off by default. diff --git a/skyrl/backends/skyrl_train/utils/profiler.py b/skyrl/backends/skyrl_train/utils/profiler.py index 14130e3e81..8e23b800d4 100644 --- a/skyrl/backends/skyrl_train/utils/profiler.py +++ b/skyrl/backends/skyrl_train/utils/profiler.py @@ -4,8 +4,6 @@ import torch.distributed from loguru import logger -from skyrl.backends.skyrl_train.utils.io.io import is_cloud_path - # Map config activity strings to torch ProfilerActivity members. _ACTIVITY_MAP = { "cpu": torch.profiler.ProfilerActivity.CPU, @@ -17,15 +15,13 @@ def build_profiler_from_policy_cfg(trainer_cfg): """Construct a :class:`Profiler` from ``policy.torch_profiler_config``, or None. Returns ``None`` when profiling is disabled, so callers can simply assign the - result to ``self.profiler``. The trace ``save_path`` defaults to - ``{ckpt_path}/profiler_traces`` (mirrors how memory snapshots default under - ``ckpt_path``). + result to ``self.profiler``. ``save_path`` is an explicit, required local path + (validated at startup) -- there is no implicit ``ckpt_path``-derived default. """ cfg = trainer_cfg.policy.torch_profiler_config if not cfg.enable: return None - default_save_path = os.path.join(trainer_cfg.ckpt_path, "profiler_traces") - return Profiler(cfg, default_save_path=default_save_path) + return Profiler(cfg) class Profiler: @@ -51,7 +47,7 @@ class Profiler: with_modules, export_type. """ - def __init__(self, config, default_save_path: str = None): + def __init__(self, config): self.enable = config.enable self.prof = None # Per-window self-device-time kernel summary, refreshed by on_trace_ready @@ -62,17 +58,9 @@ def __init__(self, config, default_save_path: str = None): if not config.enable: return self.config = config - self.save_path = config.save_path or default_save_path or "./profiler_traces" - # torch.profiler writes traces with the local filesystem only. ``save_path`` - # commonly defaults to ``{ckpt_path}/profiler_traces``, and ckpt_path can be a - # cloud URI (s3://, gs://) -- which torch can't write to, so the trace would be - # silently lost. Fall back to a local dir so profiling still produces output. - if is_cloud_path(self.save_path): - logger.warning( - f"[Profiler] cloud save_path {self.save_path!r} is not writable by torch.profiler; " - f"falling back to local './profiler_traces'." - ) - self.save_path = "./profiler_traces" + # `save_path` is a required, explicit local path -- validated by + # TorchProfilerConfig.validate() at startup (non-empty + not a cloud URI). + self.save_path = config.save_path self.ranks = list(config.ranks) self.export_type = getattr(config, "export_type", "chrome_trace") self.rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 3853bf67bc..eda3fc0965 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -179,7 +179,10 @@ class TorchProfilerConfig(BaseConfig): enable: bool = False ranks: List[int] = field(default_factory=lambda: [0]) save_path: Optional[str] = None - """Trace output dir. Defaults to ``{ckpt_path}/profiler_traces`` when None.""" + """Trace output dir. Required when ``enable=True``; use an absolute local path. + Ray workers run from a ``/tmp/ray/.../working_dir_files`` runtime dir, so a + relative path would scatter traces there -- and ``torch.profiler`` cannot write + to cloud URIs (``s3://``/``gs://``).""" # torch.profiler.schedule -- one cycle = wait + warmup + active steps. skip_first: int = 10 @@ -215,6 +218,22 @@ def validate(self) -> None: return if not self.ranks: raise ValueError("`torch_profiler_config.ranks` must be non-empty when profiling is enabled.") + # `save_path` is a required, explicit, local path -- no implicit default. Ray + # workers run from a /tmp/ray runtime working dir, so a relative path would + # scatter traces there, and torch.profiler can't write cloud URIs. + if not self.save_path: + raise ValueError( + "`torch_profiler_config.save_path` must be set when profiling is enabled. " + "Use an absolute local path -- Ray workers run from a /tmp/ray runtime " + "working dir, so a relative path would write traces there." + ) + from skyrl.backends.skyrl_train.utils.io.io import is_cloud_path + + if is_cloud_path(self.save_path): + raise ValueError( + f"`torch_profiler_config.save_path` must be a local path; got cloud URI " + f"{self.save_path!r}. torch.profiler cannot write to cloud storage." + ) # An empty `activities` passes the membership check below vacuously, but # `torch.profiler.profile(activities=[])` records nothing -- fail fast instead. if not self.activities: diff --git a/tests/backends/skyrl_train/utils/test_profiler.py b/tests/backends/skyrl_train/utils/test_profiler.py index e80757a1fc..a53ddef065 100644 --- a/tests/backends/skyrl_train/utils/test_profiler.py +++ b/tests/backends/skyrl_train/utils/test_profiler.py @@ -44,7 +44,7 @@ def _run_loop(prof: Profiler, n_steps: int) -> None: def test_disabled_is_noop(tmp_path): - prof = Profiler(_ProfCfg(enable=False), default_save_path=str(tmp_path)) + prof = Profiler(_ProfCfg(enable=False, save_path=str(tmp_path))) assert prof.prof is None assert prof.check() is False _run_loop(prof, 5) # must not raise @@ -54,7 +54,7 @@ def test_disabled_is_noop(tmp_path): def test_rank_not_selected_is_noop(tmp_path): # is_initialized() is False in this process -> rank resolves to 0, which is # not in ranks=[1], so the profiler must not arm. - prof = Profiler(_ProfCfg(ranks=[1]), default_save_path=str(tmp_path)) + prof = Profiler(_ProfCfg(ranks=[1], save_path=str(tmp_path))) assert prof.prof is None _run_loop(prof, 3) assert glob.glob(os.path.join(str(tmp_path), "*")) == [] @@ -62,8 +62,7 @@ def test_rank_not_selected_is_noop(tmp_path): def test_single_window_writes_one_trace(tmp_path): prof = Profiler( - _ProfCfg(skip_first=0, wait=0, warmup=0, active=1, repeat=1), - default_save_path=str(tmp_path), + _ProfCfg(skip_first=0, wait=0, warmup=0, active=1, repeat=1, save_path=str(tmp_path)), ) assert prof.check() is True _run_loop(prof, 3) @@ -75,8 +74,7 @@ def test_single_window_writes_one_trace(tmp_path): def test_repeat_writes_multiple_windows(tmp_path): # Two cycles of (warmup=1 + active=1); needs >= 2*(1+1) steps to fire twice. prof = Profiler( - _ProfCfg(skip_first=0, wait=0, warmup=1, active=1, repeat=2), - default_save_path=str(tmp_path), + _ProfCfg(skip_first=0, wait=0, warmup=1, active=1, repeat=2, save_path=str(tmp_path)), ) _run_loop(prof, 8) traces = glob.glob(os.path.join(str(tmp_path), "*.pt.trace.json*")) @@ -86,51 +84,35 @@ def test_repeat_writes_multiple_windows(tmp_path): def test_skip_first_defers_recording(tmp_path): # With skip_first larger than the loop, no window should ever open. prof = Profiler( - _ProfCfg(skip_first=100, wait=0, warmup=0, active=1, repeat=1), - default_save_path=str(tmp_path), + _ProfCfg(skip_first=100, wait=0, warmup=0, active=1, repeat=1, save_path=str(tmp_path)), ) _run_loop(prof, 5) assert glob.glob(os.path.join(str(tmp_path), "*.pt.trace.json*")) == [] -def test_save_path_overrides_default(tmp_path): - explicit = tmp_path / "explicit" - prof = Profiler( - _ProfCfg(save_path=str(explicit)), - default_save_path=str(tmp_path / "default"), - ) - assert prof.save_path == str(explicit) - - -def test_default_save_path_used_when_none(tmp_path): - default = str(tmp_path / "default") - prof = Profiler(_ProfCfg(save_path=None), default_save_path=default) - assert prof.save_path == default - - -def test_cloud_save_path_falls_back_to_local(tmp_path): - # ckpt_path (and thus the derived default save_path) can be a cloud URI, which - # torch.profiler can't write to. The wrapper must fall back to a local dir so - # the trace isn't silently lost. - prof = Profiler(_ProfCfg(save_path="s3://my-bucket/run/profiler_traces"), default_save_path=str(tmp_path)) - assert prof.save_path == "./profiler_traces" +def test_save_path_is_taken_verbatim(tmp_path): + # `save_path` is an explicit, required local path -- the wrapper uses it as-is + # (no implicit default, no rewriting). Cloud-URI/empty rejection lives in + # TorchProfilerConfig.validate(), exercised in tests/train/test_config.py. + explicit = str(tmp_path / "explicit") + prof = Profiler(_ProfCfg(save_path=explicit)) + assert prof.save_path == explicit def test_kernel_summary_none_when_disabled(tmp_path): - prof = Profiler(_ProfCfg(enable=False), default_save_path=str(tmp_path)) + prof = Profiler(_ProfCfg(enable=False, save_path=str(tmp_path))) assert prof.get_kernel_summary() is None def test_kernel_summary_empty_before_first_window(tmp_path): - prof = Profiler(_ProfCfg(skip_first=0, wait=0, warmup=0, active=1), default_save_path=str(tmp_path)) + prof = Profiler(_ProfCfg(skip_first=0, wait=0, warmup=0, active=1, save_path=str(tmp_path))) summary = prof.get_kernel_summary() assert summary == {"window_count": 0, "pairs": []} def test_kernel_summary_populated_after_window(tmp_path): prof = Profiler( - _ProfCfg(skip_first=0, wait=0, warmup=0, active=1, repeat=1, activities=["cpu"]), - default_save_path=str(tmp_path), + _ProfCfg(skip_first=0, wait=0, warmup=0, active=1, repeat=1, activities=["cpu"], save_path=str(tmp_path)), ) prof.start() for _ in range(3): @@ -157,14 +139,14 @@ def test_kernel_summary_populated_after_window(tmp_path): def test_activities_threaded_to_torch(tmp_path): import torch - prof = Profiler(_ProfCfg(activities=["cpu"]), default_save_path=str(tmp_path)) + prof = Profiler(_ProfCfg(activities=["cpu"], save_path=str(tmp_path))) # The underlying torch profiler was constructed with the CPU activity only. assert torch.profiler.ProfilerActivity.CPU in prof.prof.activities assert torch.profiler.ProfilerActivity.CUDA not in prof.prof.activities def test_step_failure_disables_without_raising(tmp_path): - prof = Profiler(_ProfCfg(), default_save_path=str(tmp_path)) + prof = Profiler(_ProfCfg(save_path=str(tmp_path))) prof.start() class _Boom: @@ -241,42 +223,33 @@ def test_rpcs_drive_profiler_when_present(self): class TestBuildProfilerFromPolicyCfg: """``build_profiler_from_policy_cfg`` is the production entry both worker - backends call in ``init_model``. It gates on ``enable`` and composes the - default ``{ckpt_path}/profiler_traces`` save path when none is set.""" + backends call in ``init_model``. It gates on ``enable`` and passes the + explicit, validated ``save_path`` straight through (no implicit default).""" @staticmethod - def _trainer_cfg(prof_cfg, ckpt_path): + def _trainer_cfg(prof_cfg): from types import SimpleNamespace - return SimpleNamespace(ckpt_path=ckpt_path, policy=SimpleNamespace(torch_profiler_config=prof_cfg)) + return SimpleNamespace(policy=SimpleNamespace(torch_profiler_config=prof_cfg)) def test_returns_none_when_disabled(self, tmp_path): from skyrl.backends.skyrl_train.utils.profiler import ( build_profiler_from_policy_cfg, ) - cfg = self._trainer_cfg(_ProfCfg(enable=False), str(tmp_path)) + cfg = self._trainer_cfg(_ProfCfg(enable=False, save_path=str(tmp_path))) assert build_profiler_from_policy_cfg(cfg) is None - def test_builds_profiler_with_ckpt_default_save_path(self, tmp_path): + def test_builds_profiler_with_explicit_save_path(self, tmp_path): from skyrl.backends.skyrl_train.utils.profiler import ( Profiler, build_profiler_from_policy_cfg, ) - cfg = self._trainer_cfg(_ProfCfg(enable=True, save_path=None), str(tmp_path)) - prof = build_profiler_from_policy_cfg(cfg) - assert isinstance(prof, Profiler) - assert prof.save_path == os.path.join(str(tmp_path), "profiler_traces") - - def test_explicit_save_path_wins_over_ckpt_default(self, tmp_path): - from skyrl.backends.skyrl_train.utils.profiler import ( - build_profiler_from_policy_cfg, - ) - explicit = str(tmp_path / "explicit") - cfg = self._trainer_cfg(_ProfCfg(enable=True, save_path=explicit), str(tmp_path)) + cfg = self._trainer_cfg(_ProfCfg(enable=True, save_path=explicit)) prof = build_profiler_from_policy_cfg(cfg) + assert isinstance(prof, Profiler) assert prof.save_path == explicit diff --git a/tests/train/test_config.py b/tests/train/test_config.py index f6e8fced0a..62c0e5ea91 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -433,13 +433,18 @@ class TestTorchProfilerConfigValidation: def _cfg(**overrides): from skyrl.train.config.config import TorchProfilerConfig + # `save_path` is required when enabled; supply a valid local default so + # tests targeting *other* fields don't trip the save_path guard. + overrides.setdefault("save_path", "/tmp/skyrl_prof_test") return TorchProfilerConfig(enable=True, **overrides) def test_disabled_skips_all_checks(self): from skyrl.train.config.config import TorchProfilerConfig # Garbage values must be tolerated while disabled (the default state). - TorchProfilerConfig(enable=False, export_type="bogus", activities=["gpu"], ranks=[], active=0).validate() + TorchProfilerConfig( + enable=False, export_type="bogus", activities=["gpu"], ranks=[], active=0, save_path=None + ).validate() def test_defaults_are_valid_when_enabled(self): self._cfg().validate() # must not raise @@ -448,6 +453,19 @@ def test_empty_ranks_rejected(self): with pytest.raises(ValueError, match=r"ranks.*non-empty"): self._cfg(ranks=[]).validate() + def test_missing_save_path_rejected(self): + # save_path has no implicit default; an enabled run must set it explicitly. + with pytest.raises(ValueError, match=r"save_path.*must be set"): + self._cfg(save_path=None).validate() + with pytest.raises(ValueError, match=r"save_path.*must be set"): + self._cfg(save_path="").validate() + + def test_cloud_save_path_rejected(self): + # torch.profiler can only write the local filesystem; cloud URIs fail fast. + for uri in ("s3://bucket/run/traces", "gs://bucket/run/traces", "gcs://bucket/run/traces"): + with pytest.raises(ValueError, match=r"save_path.*local path"): + self._cfg(save_path=uri).validate() + def test_unknown_activity_rejected(self): with pytest.raises(ValueError, match=r"activities"): self._cfg(activities=["cpu", "gpu"]).validate() @@ -483,6 +501,7 @@ def test_validate_cfg_invokes_profiler_validation(self): # The RL entrypoint validator must surface profiler config errors. 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.torch_profiler_config.export_type = "bogus" with pytest.raises(ValueError, match=r"export_type"): validate_cfg(cfg) diff --git a/tests/train/test_sft_config.py b/tests/train/test_sft_config.py index 3aa3312646..86c4e39fc6 100644 --- a/tests/train/test_sft_config.py +++ b/tests/train/test_sft_config.py @@ -167,6 +167,7 @@ def test_capture_flags_propagate(self): "torch_profiler_config.profile_memory=true", "torch_profiler_config.with_stack=false", "torch_profiler_config.activities=[cuda]", + "torch_profiler_config.save_path=/tmp/sft_prof", ] ) skyrl_cfg = build_skyrl_config_for_sft(cfg) @@ -183,6 +184,7 @@ def test_invalid_profiler_config_rejected_by_sft_validation(self): "model.path=test/my-model", "torch_profiler_config.enable=true", "torch_profiler_config.export_type=bogus", + "torch_profiler_config.save_path=/tmp/sft_prof", ] ) with pytest.raises(ValueError, match=r"export_type"): From dcd4a2d22437d5c5f2bf1666808f27767f861c37 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Mon, 15 Jun 2026 13:08:02 -0500 Subject: [PATCH 5/7] 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 " 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) --- docs/content/docs/configuration/config.mdx | 2 + skyrl/train/config/config.py | 38 ++++++++++++++++- skyrl/train/utils/utils.py | 7 +++- tests/train/test_config.py | 47 ++++++++++++++++++++++ 4 files changed, 92 insertions(+), 2 deletions(-) diff --git a/docs/content/docs/configuration/config.mdx b/docs/content/docs/configuration/config.mdx index 4c29535151..d2733b8b72 100644 --- a/docs/content/docs/configuration/config.mdx +++ b/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 `) 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. diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index eda3fc0965..870944b191 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -205,7 +205,13 @@ 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 @@ -213,6 +219,13 @@ def validate(self) -> None: (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 ``) 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 @@ -261,6 +274,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 " 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 `. " + "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): diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 342ae3a8c8..ce4eadf693 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -270,7 +270,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() diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 62c0e5ea91..6e3425f217 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -505,3 +505,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 ") 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 From 54f71b74528ad615563405b5462d95724b6d50d0 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Tue, 16 Jun 2026 10:56:05 -0500 Subject: [PATCH 6/7] test(profiler): don't start real kineto in step-fault test (fixes teardown segfault) test_step_failure_disables_without_raising started a real torch.profiler session, then nulled self.prof via the fault path without stopping it. The orphaned libkineto session crashed at interpreter shutdown (exit 139), surfacing as a CPU-CI teardown segfault after all tests passed. Swap the faulting stub in before start() so no real profiler is ever opened; the swallow-and-disable path is still exercised. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../skyrl_train/utils/test_profiler.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/tests/backends/skyrl_train/utils/test_profiler.py b/tests/backends/skyrl_train/utils/test_profiler.py index a53ddef065..3b59a00863 100644 --- a/tests/backends/skyrl_train/utils/test_profiler.py +++ b/tests/backends/skyrl_train/utils/test_profiler.py @@ -147,14 +147,29 @@ def test_activities_threaded_to_torch(tmp_path): def test_step_failure_disables_without_raising(tmp_path): prof = Profiler(_ProfCfg(save_path=str(tmp_path))) - prof.start() class _Boom: + """Stands in for a live torch profiler whose ``step`` faults. ``start``/ + ``stop`` are no-ops: we must NOT start the real kineto profiler here, or + the wrapper's fault path nulls ``self.prof`` while that session is still + running, leaving it unstopped -> libkineto segfaults at interpreter exit + (process-wide, surfaces as a teardown crash long after this test passes).""" + + def start(self): + pass + def step(self): raise RuntimeError("boom") - # Simulate an internal profiler fault mid-loop; the wrapper must swallow it. + def stop(self): + pass + + # Swap in the faulting stub BEFORE start() so no real profiler session is ever + # opened; this still exercises the wrapper's swallow-and-disable path. prof.prof = _Boom() + prof.start() + + # Simulate an internal profiler fault mid-loop; the wrapper must swallow it. prof.step() assert prof.enable is False assert prof.prof is None From 604f8383ba0553c4ebe8ee921fd4d5b43b23ac02 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Wed, 1 Jul 2026 14:09:50 -0700 Subject: [PATCH 7/7] docs: tighten profiler comments --- docs/content/docs/configuration/config.mdx | 32 ++++----- examples/train/async/async_trainer.py | 5 +- examples/train/megatron/run_megatron.sh | 6 +- .../megatron/run_megatron_nemotron_mini_4b.sh | 4 +- .../full_context/trainer_full_ctx.py | 12 ++-- skyrl/backends/skyrl_train/utils/profiler.py | 66 +++--------------- .../skyrl_train/workers/fsdp/fsdp_worker.py | 2 +- .../workers/megatron/megatron_worker.py | 4 +- skyrl/backends/skyrl_train/workers/worker.py | 37 ++-------- .../skyrl_train/workers/worker_dispatch.py | 24 ++----- skyrl/train/config/config.py | 68 ++++--------------- skyrl/train/config/sft_config.py | 4 +- skyrl/train/fully_async_trainer.py | 15 ++-- skyrl/train/sft_trainer.py | 10 +-- skyrl/train/trainer.py | 38 +++-------- .../skyrl_train/utils/test_profiler.py | 54 +++------------ tests/train/test_config.py | 29 ++------ tests/train/test_sft_config.py | 4 +- 18 files changed, 96 insertions(+), 318 deletions(-) diff --git a/docs/content/docs/configuration/config.mdx b/docs/content/docs/configuration/config.mdx index d2733b8b72..c10129d1d1 100644 --- a/docs/content/docs/configuration/config.mdx +++ b/docs/content/docs/configuration/config.mdx @@ -273,22 +273,22 @@ policy: use_torch_compile: false # Enable torch compile for the entropy calculation record_memory: false # Dump memory snapshot for debugging - torch_profiler_config: # torch.profiler-based training-loop profiler (see below) + torch_profiler_config: # torch.profiler for policy training enable: false ranks: [0] - save_path: null # REQUIRED when enable=true; absolute local path - skip_first: 10 # torch.profiler.schedule + save_path: null # required when enable=true; absolute local path + skip_first: 10 wait: 0 warmup: 1 active: 1 - repeat: 1 # 0 = profile for the whole run + repeat: 1 # 0 = repeat forever activities: ["cpu", "cuda"] record_shapes: true profile_memory: false with_stack: true with_flops: false with_modules: false - export_type: "chrome_trace" # chrome_trace | stacks + export_type: "chrome_trace" # chrome_trace | stacks model_config_kwargs: {} # pass through kwargs to the HuggingFace model config for FSDP training backends (i.e. for overriding vocab size, etc) - for megatron, use policy.megatron_config.transformer_config_kwargs instead @@ -302,21 +302,19 @@ policy: ### Torch Profiler Configuration -`policy.torch_profiler_config` enables a [`torch.profiler`](https://docs.pytorch.org/docs/stable/profiler.html)-based profiler that the trainer drives around the training loop (both FSDP and Megatron backends, and both RL and SFT). When enabled, it writes one Kineto/[HTA](https://github.com/facebookresearch/HolisticTraceAnalysis)-friendly `*.pt.trace.json` per active window per profiled rank into `save_path`. +`policy.torch_profiler_config` profiles policy training steps for FSDP and Megatron in RL and SFT. It writes one Kineto/[HTA](https://github.com/facebookresearch/HolisticTraceAnalysis)-friendly `*.pt.trace.json` per active window and profiled rank. -**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`. +Scope: policy forward/backward/optimizer only. RL critic/ref models and generation are not profiled. -**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 `) 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. +FSDP restriction: profiling is rejected when manual CPU offload and colocation are both active (`policy.fsdp_config.cpu_offload: false` with `placement.colocate_all` or `placement.colocate_policy_ref`). Use FSDP native CPU offload, disable colocation, or use Megatron. -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. -- `policy.torch_profiler_config.ranks`: List of global ranks to profile (e.g. `[0]`). Add a mid-pipeline-stage rank to diagnose pipeline bubbles. -- `policy.torch_profiler_config.save_path`: Output directory for traces. **Required** when `enable: true` (there is no implicit default). Use an absolute local path: Ray workers run from a `/tmp/ray/.../working_dir_files` runtime working dir, so a relative path would scatter traces there, and `torch.profiler` cannot write to cloud URIs (`s3://`/`gs://`) — both are rejected at startup. -- `policy.torch_profiler_config.{skip_first,wait,warmup,active,repeat}`: Passed directly to `torch.profiler.schedule`. -- `policy.torch_profiler_config.activities`: Subset of `["cpu", "cuda"]` to record. -- `policy.torch_profiler_config.{record_shapes,profile_memory,with_stack,with_flops,with_modules}`: Passed directly to `torch.profiler.profile`. `with_stack` and `record_shapes` add overhead; `profile_memory` is heavier still and off by default. -- `policy.torch_profiler_config.export_type`: `chrome_trace` writes `*.pt.trace.json` (Kineto/HTA-friendly); `stacks` writes flamegraph-style self-CUDA-time stacks (requires `with_stack: true`). +- `policy.torch_profiler_config.enable`: Enables profiling. +- `policy.torch_profiler_config.ranks`: Global ranks to profile, for example `[0]`. +- `policy.torch_profiler_config.save_path`: Required absolute local trace directory. Relative paths land under Ray's `/tmp/ray/.../working_dir_files`; cloud URIs are rejected. +- `policy.torch_profiler_config.{skip_first,wait,warmup,active,repeat}`: Passed to [`torch.profiler.schedule`](https://docs.pytorch.org/docs/stable/profiler.html#torch.profiler.schedule). +- `policy.torch_profiler_config.activities`: Subset of `["cpu", "cuda"]`. +- `policy.torch_profiler_config.{record_shapes,profile_memory,with_stack,with_flops,with_modules}`: Passed to `torch.profiler.profile`. +- `policy.torch_profiler_config.export_type`: `chrome_trace` for `*.pt.trace.json`; `stacks` for self-CUDA-time stacks (`with_stack: true` required). ### LoRA Configuration diff --git a/examples/train/async/async_trainer.py b/examples/train/async/async_trainer.py index 8c1d289284..1796f9198d 100644 --- a/examples/train/async/async_trainer.py +++ b/examples/train/async/async_trainer.py @@ -79,8 +79,7 @@ async def train(self): self.sync_finished.set() self.generation_ack.clear() - # Advance the torch profiler schedule once per global step - # (no-op unless profiling is enabled). + # One profiler step per global step. self._profiler_step() # 5. set logs @@ -118,8 +117,6 @@ async def train(self): # cancel generation task for this epoch generator_task.cancel() finally: - # Always stop/flush the profiler when the loop exits (incl. on error) - # so the open kineto trace window isn't leaked. No-op when disabled. self._profiler_stop() pbar.close() diff --git a/examples/train/megatron/run_megatron.sh b/examples/train/megatron/run_megatron.sh index 9f4c594ae3..05a11e5874 100644 --- a/examples/train/megatron/run_megatron.sh +++ b/examples/train/megatron/run_megatron.sh @@ -17,9 +17,7 @@ MEGATRON_TP=2 MEGATRON_PP=2 MEGATRON_CP=1 -# torch profiler config. Profiles ONLY the policy model's training step -# (forward/backward + optim) -- not the critic or ref models, and not generation/ -# inference. See trainer.policy.torch_profiler_config for schedule/capture knobs. +# torch.profiler for policy training steps. ENABLE_TORCH_PROFILER=false RANKS_TO_PROFILE="[0]" SAVE_PATH="$HOME/megatron_prof/tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_${MODEL_NAME}" @@ -71,4 +69,4 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \ trainer.run_name="gsm8k_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_${MODEL_NAME}" \ trainer.resume_mode=null \ trainer.ckpt_path="$HOME/ckpts/gsm8k_megatron_ckpt" \ - $@ \ No newline at end of file + $@ diff --git a/examples/train/megatron/run_megatron_nemotron_mini_4b.sh b/examples/train/megatron/run_megatron_nemotron_mini_4b.sh index 6fe0f35288..3eb302a498 100755 --- a/examples/train/megatron/run_megatron_nemotron_mini_4b.sh +++ b/examples/train/megatron/run_megatron_nemotron_mini_4b.sh @@ -37,7 +37,7 @@ MEGATRON_CP=1 NUM_INFERENCE_ENGINES=8 INFERENCE_TP=1 -# torch profiler config +# torch.profiler config ENABLE_TORCH_PROFILER=false RANKS_TO_PROFILE="[0]" SAVE_PATH="$HOME/megatron_prof/tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_${MODEL_NAME}" @@ -89,4 +89,4 @@ uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \ trainer.run_name="gsm8k_megatron_nemotron_mini_4b_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.resume_mode=null \ trainer.ckpt_path="$HOME/ckpts/gsm8k_megatron_nemotron_ckpt" \ - $@ \ No newline at end of file + $@ diff --git a/examples/train_scripts/full_context/trainer_full_ctx.py b/examples/train_scripts/full_context/trainer_full_ctx.py index 3160c2358c..6598d75ba1 100644 --- a/examples/train_scripts/full_context/trainer_full_ctx.py +++ b/examples/train_scripts/full_context/trainer_full_ctx.py @@ -1,6 +1,8 @@ -from skyrl.train.trainer import RayPPOTrainer -from loguru import logger import random + +from loguru import logger + +from skyrl.train.trainer import RayPPOTrainer from skyrl.train.utils.utils import Timer @@ -86,8 +88,7 @@ async def train(self): with Timer("train_critic_and_policy", self.all_timings): status = self.train_critic_and_policy(training_input) - # Advance the torch profiler schedule once per global step - # (no-op unless profiling is enabled). + # One profiler step per global step. self._profiler_step() self.tracker.log(self.all_metrics, step=self.global_step) @@ -101,9 +102,6 @@ async def train(self): logger.info(f"Step {step + 1} completed. Status: {status}") finally: - # Always stop/flush the profiler when the loop exits -- including via - # an exception -- so the open kineto trace window isn't leaked. No-op - # when profiling is disabled. self._profiler_stop() self.tracker.finish() diff --git a/skyrl/backends/skyrl_train/utils/profiler.py b/skyrl/backends/skyrl_train/utils/profiler.py index 8e23b800d4..2e3f6a6147 100644 --- a/skyrl/backends/skyrl_train/utils/profiler.py +++ b/skyrl/backends/skyrl_train/utils/profiler.py @@ -4,7 +4,7 @@ import torch.distributed from loguru import logger -# Map config activity strings to torch ProfilerActivity members. +# Config string -> torch.profiler activity. _ACTIVITY_MAP = { "cpu": torch.profiler.ProfilerActivity.CPU, "cuda": torch.profiler.ProfilerActivity.CUDA, @@ -12,12 +12,7 @@ def build_profiler_from_policy_cfg(trainer_cfg): - """Construct a :class:`Profiler` from ``policy.torch_profiler_config``, or None. - - Returns ``None`` when profiling is disabled, so callers can simply assign the - result to ``self.profiler``. ``save_path`` is an explicit, required local path - (validated at startup) -- there is no implicit ``ckpt_path``-derived default. - """ + """Build the policy profiler, or return None when disabled.""" cfg = trainer_cfg.policy.torch_profiler_config if not cfg.enable: return None @@ -25,41 +20,18 @@ def build_profiler_from_policy_cfg(trainer_cfg): class Profiler: - """A configurable ``torch.profiler`` wrapper driven by the training loop. - - The trainer brackets the loop: ``start()`` once before it, ``step()`` once - per global step, ``stop()`` once after. Which steps are actually recorded is - decided by ``torch.profiler.schedule(skip_first, wait, warmup, active, - repeat)`` -- this is the "profile N steps, every M, repeating K times" knob, - so nothing about the window is hardcoded. - - Traces are written by ``torch.profiler.tensorboard_trace_handler`` (one - ``*.pt.trace.json`` per active window, per rank), which is the - Kineto/Holistic-Trace-Analysis-friendly format -- no manual export. - - Every method is exception-isolated: a profiler fault disables profiling for - the rest of the run rather than crashing training. - - ``config`` fields (see :class:`skyrl.train.config.config.TorchProfilerConfig`): - enable, ranks, save_path, - skip_first, wait, warmup, active, repeat, - activities, record_shapes, profile_memory, with_stack, with_flops, - with_modules, export_type. - """ + """Thin ``torch.profiler`` wrapper driven by trainer start/step/stop calls.""" def __init__(self, config): self.enable = config.enable self.prof = None - # Per-window self-device-time kernel summary, refreshed by on_trace_ready - # at the close of each active window (the per-kernel attribution denominator). - # Exposed read-only via get_kernel_summary(). + # Last closed-window kernel self time, exposed via get_kernel_summary(). self._last_pairs: list = [] self._window_count: int = 0 if not config.enable: return self.config = config - # `save_path` is a required, explicit local path -- validated by - # TorchProfilerConfig.validate() at startup (non-empty + not a cloud URI). + # Validated at startup by TorchProfilerConfig. self.save_path = config.save_path self.ranks = list(config.ranks) self.export_type = getattr(config, "export_type", "chrome_trace") @@ -98,16 +70,10 @@ def __init__(self, config): self.prof = None def _on_trace_ready(self, prof) -> None: - """Fires at the close of each active window. Writes the trace AND stashes - a pickle-safe per-kernel self-device-time summary for the just-closed - window (the per-kernel attribution denominator -- exact, no cross-stream - overlap double-counting). ``rank{N}`` in the worker_name keeps the rank - parseable by HTA and avoids cross-rank filename collisions in a shared - ``save_path``. Best-effort: a fault here must never crash the worker.""" + """Write a trace and cache the last-window kernel self-time summary.""" os.makedirs(self.save_path, exist_ok=True) worker_name = f"rank{self.rank}" if self.export_type == "stacks": - # Flamegraph-style self-CUDA-time stacks (requires with_stack=True). out = os.path.join(self.save_path, f"{worker_name}_stacks.txt") prof.export_stacks(out, "self_cuda_time_total") logger.info(f"[Profiler] rank {self.rank}: exported stacks -> {out}") @@ -116,30 +82,14 @@ def _on_trace_ready(self, prof) -> None: logger.info(f"[Profiler] rank {self.rank}: exported chrome trace under {self.save_path}") try: - # ``self_device_time_total`` is torch 2.11's field (the older - # ``self_cuda_time_total`` was removed). Microseconds, self time. + # Microseconds, self time. self._last_pairs = [(str(e.key), float(e.self_device_time_total)) for e in prof.key_averages()] self._window_count += 1 except Exception as e: logger.warning(f"[Profiler] rank {self.rank}: kernel-summary capture failed: {e}") def get_kernel_summary(self): - """Return the last closed window's self-device-time kernel summary, or None. - - Shape (pickle-safe, no tensors):: - - {"window_count": int, "pairs": [(kernel_name, self_device_us), ...]} - - ``None`` when profiling is disabled or no profiler was constructed on - this rank. ``pairs`` is empty until the first active window closes. - - NOTE: SkyRL's own trainers never read this -- the high-level - ``*.pt.trace.json`` files are the deliverable. This (and the - ``_on_trace_ready`` capture that feeds it) is a deliberately-provided - low-level API for downstream consumers that want per-kernel self-time - attribution without re-parsing the trace. Reached via - ``Worker.dump_profiler_summary`` -> ``WorkerDispatch.dump_profiler_summary``. - """ + """Return ``{"window_count": int, "pairs": [(name, self_us), ...]}`` or None.""" if not self.enable or self.prof is None: return None return {"window_count": self._window_count, "pairs": list(self._last_pairs)} diff --git a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py index da109ce86d..e80cf65d5e 100644 --- a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py +++ b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py @@ -196,7 +196,7 @@ def init_model(self, model_path, num_training_steps: int = None): # standard CUDA memory; only subsequent activations use expandable segments. self._set_expandable_segments(True) - # create profiler (driven by the trainer via start/profile_step/stop RPCs) + # Created only on profiled ranks. self.profiler = build_profiler_from_policy_cfg(self.cfg) async def init_weight_sync_state(self, inference_engine_client, inference_engine_cfg: "InferenceEngineConfig"): diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 6382e41b14..844d762b37 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -704,7 +704,7 @@ def __init__(self, **kwargs): self.actor_module: List[nn.Module] = None self.scheduler: OptimizerParamScheduler = None self.optimizer: DistributedOptimizer = None - # self.profiler is initialized on the Worker base; populated in init_model. + # Worker base owns self.profiler; init_model may populate it. self._is_lora = self.cfg.policy.model.lora.rank > 0 # Per-worker store of LoRA adapter snapshots. Allocated only for the # LoRA path; FFT runs single-tenant exactly as before. @@ -806,7 +806,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9): if self._rank == 0: print_model_size(self.actor_module[0]) - # create profiler (driven by the trainer via start/profile_step/stop RPCs) + # Created only on profiled ranks. self.profiler = build_profiler_from_policy_cfg(self.cfg) # create optimizer (skipped for inference-only flows; Megatron's diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py index 7feb3b8569..e179dc9e6e 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/skyrl/backends/skyrl_train/workers/worker.py @@ -233,9 +233,7 @@ def __init__(self, cfg: TrainerConfig, *args, **kwargs): super().__init__(*args, **kwargs) self.cfg = cfg self._transfer_strategy_cls = None # Set in init_weight_transfer_communicator - # torch.profiler wrapper. Constructed in ``init_model`` when - # ``policy.torch_profiler_config.enable`` is set; driven by the trainer - # via the start_profile/profile_step/stop_profile RPCs below. + # Populated by init_model when torch profiling is enabled. self.profiler = None if self.cfg.algorithm.temperature is None: @@ -297,14 +295,7 @@ def set_algorithm_config(self, **kwargs) -> None: setattr(self.cfg.algorithm, key, value) # ------------------------------------------------------------------ - # torch.profiler control RPCs (dispatched by the trainer via "pass_through"). - # - # These live on the shared Worker base so they're on the snapshotted Ray - # actor method table of every PolicyWorker (Megatron + FSDP) without a - # subclass. They no-op when ``self.profiler`` is None (profiling disabled - # or rank not selected), so the trainer can call them on every step. The - # ``Profiler`` itself is exception-isolated; these are an extra guard so a - # profiler fault can never abort a training step. + # torch.profiler RPCs, dispatched via WorkerDispatch pass_through. # ------------------------------------------------------------------ def start_profile(self) -> None: @@ -313,12 +304,7 @@ def start_profile(self) -> None: self.profiler.start() def profile_step(self) -> None: - """Advance the profiler schedule by one global step (no-op when disabled). - - Call exactly once per global step. ``torch.profiler``'s schedule decides - which steps are actually recorded; trace files are written automatically - at the close of each active window. - """ + """Advance the profiler schedule by one global step.""" if self.profiler is not None: self.profiler.step() @@ -328,22 +314,7 @@ def stop_profile(self) -> None: self.profiler.stop() def dump_profiler_summary(self): - """Return this rank's last-window kernel self-time summary, or None. - - Pickle-safe dict (``{"window_count", "pairs": [(name, self_us), ...]}``) - or None when profiling is disabled / no profiler on this rank. The - low-level per-kernel attribution data path for downstream consumers; the - trace files remain the high-level (HTA) source. - - NOTE: SkyRL itself does not call this RPC (or ``get_kernel_summary``) in - its own training loop -- the high-level trace files are SkyRL's - deliverable. This is a deliberately-provided forward-looking API for - downstream consumers that want per-kernel self-time attribution without - re-parsing the on-disk trace. Keep it wired end-to-end (Profiler -> - Worker -> WorkerDispatch) so such a consumer needs only to call - ``dispatch.dump_profiler_summary("policy")``; do not remove it as - "dead code". - """ + """Return this rank's last-window kernel summary, or None.""" return self.profiler.get_kernel_summary() if self.profiler is not None else None def _get_module_for_offload(self): diff --git a/skyrl/backends/skyrl_train/workers/worker_dispatch.py b/skyrl/backends/skyrl_train/workers/worker_dispatch.py index 7dd8ba9fe1..cf08210d0b 100644 --- a/skyrl/backends/skyrl_train/workers/worker_dispatch.py +++ b/skyrl/backends/skyrl_train/workers/worker_dispatch.py @@ -421,14 +421,12 @@ def set_algorithm_config(self, model: str, **kwargs) -> None: ray.get(self._actor_groups[model].async_run_ray_method("pass_through", "set_algorithm_config", **kwargs)) # ------------------------------------------------------------------ - # torch.profiler control. Dispatched via "pass_through" to every worker. - # Deliberately do NOT call _ensure_on_gpu so they don't perturb the - # colocation offload state machine. Each is best-effort: a profiler fault - # must never abort the training loop. + # torch.profiler control. Avoid _ensure_on_gpu so profiling does not perturb + # the colocation offload state. # ------------------------------------------------------------------ def start_profile(self, model: str) -> None: - """Arm the torch profiler on ``model``'s workers (before the loop).""" + """Start profiling on ``model`` workers.""" if model not in self._actor_groups: return try: @@ -437,7 +435,7 @@ def start_profile(self, model: str) -> None: logger.warning(f"[profiler] start_profile dispatch for {model} failed: {e}") def profile_step(self, model: str) -> None: - """Advance the torch profiler schedule by one global step on ``model``.""" + """Advance profiling by one global step.""" if model not in self._actor_groups: return try: @@ -446,7 +444,7 @@ def profile_step(self, model: str) -> None: logger.warning(f"[profiler] profile_step dispatch for {model} failed: {e}") def stop_profile(self, model: str) -> None: - """Stop the torch profiler on ``model``'s workers (after the loop).""" + """Stop profiling on ``model`` workers.""" if model not in self._actor_groups: return try: @@ -455,17 +453,7 @@ def stop_profile(self, model: str) -> None: logger.warning(f"[profiler] stop_profile dispatch for {model} failed: {e}") def dump_profiler_summary(self, model: str) -> Optional[List]: - """Collect the per-rank last-window kernel self-time summaries for ``model``. - - Returns a per-actor list (one entry per worker): profiled ranks return a - dict ``{"window_count", "pairs"}``; other ranks return None. Returns None - when the model is unknown or dispatch fails. - - NOTE: not invoked by SkyRL's own trainers (which rely on the high-level - trace files). This is the dispatch-side entry of a forward-looking - low-level API for downstream consumers that want per-kernel self-time - attribution without re-parsing traces -- see ``Worker.dump_profiler_summary``. - """ + """Collect per-rank last-window kernel summaries for ``model``.""" if model not in self._actor_groups: return None try: diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 870944b191..7dff56f60e 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -162,39 +162,24 @@ class MegatronDDPConfig(BaseConfig): @dataclass class TorchProfilerConfig(BaseConfig): - """Configuration for the ``torch.profiler``-based training-loop profiler. - - Mirrors ``torch.profiler.profile`` + ``torch.profiler.schedule`` so every - knob is overridable. Defaults reproduce the previous hardcoded behavior - (CPU+CUDA, ``record_shapes``+``with_stack``) but are now fully configurable. - The trainer drives it (``start`` before the loop, one ``step`` per global - step, ``stop`` after); the schedule decides which steps are recorded. - - Scope: this profiles **only the policy model's training step** - (forward/backward + optimizer). In an RL run it does **not** profile the - critic or ref models, and it does **not** profile generation/inference -- - only the policy training compute on the configured ``ranks``. - """ + """``torch.profiler`` config for policy training steps.""" enable: bool = False ranks: List[int] = field(default_factory=lambda: [0]) save_path: Optional[str] = None - """Trace output dir. Required when ``enable=True``; use an absolute local path. - Ray workers run from a ``/tmp/ray/.../working_dir_files`` runtime dir, so a - relative path would scatter traces there -- and ``torch.profiler`` cannot write - to cloud URIs (``s3://``/``gs://``).""" + """Trace output dir. Required when ``enable=True``; must be a local absolute path.""" - # torch.profiler.schedule -- one cycle = wait + warmup + active steps. + # torch.profiler.schedule skip_first: int = 10 - """Steps to skip before the schedule begins (warmup/steady-state).""" + """Steps to skip before scheduling begins.""" wait: int = 0 warmup: int = 1 active: int = 1 """Number of steps recorded per cycle.""" repeat: int = 1 - """Number of cycles. 0 = repeat for the whole run.""" + """Number of cycles. 0 means forever.""" - # torch.profiler.profile capture knobs. + # torch.profiler.profile activities: List[str] = field(default_factory=lambda: ["cpu", "cuda"]) record_shapes: bool = True profile_memory: bool = False @@ -202,8 +187,7 @@ class TorchProfilerConfig(BaseConfig): with_flops: bool = False with_modules: bool = False export_type: str = "chrome_trace" - """``chrome_trace`` -> ``*.pt.trace.json`` (HTA-friendly) or ``stacks`` -> - flamegraph-style self-CUDA-time stacks (requires ``with_stack=True``).""" + """``chrome_trace`` or ``stacks``; stacks require ``with_stack=True``.""" def validate( self, @@ -212,28 +196,12 @@ def validate( 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 ``) 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. - """ + """Fail fast on invalid or known-incompatible profiler settings.""" if not self.enable: return if not self.ranks: raise ValueError("`torch_profiler_config.ranks` must be non-empty when profiling is enabled.") - # `save_path` is a required, explicit, local path -- no implicit default. Ray - # workers run from a /tmp/ray runtime working dir, so a relative path would - # scatter traces there, and torch.profiler can't write cloud URIs. + # Avoid implicit relative paths in Ray runtime working dirs. if not self.save_path: raise ValueError( "`torch_profiler_config.save_path` must be set when profiling is enabled. " @@ -247,8 +215,7 @@ def validate( f"`torch_profiler_config.save_path` must be a local path; got cloud URI " f"{self.save_path!r}. torch.profiler cannot write to cloud storage." ) - # An empty `activities` passes the membership check below vacuously, but - # `torch.profiler.profile(activities=[])` records nothing -- fail fast instead. + # Empty activities record nothing. if not self.activities: raise ValueError("`torch_profiler_config.activities` must be non-empty when profiling is enabled.") bad_activities = [a for a in self.activities if a.lower() not in TORCH_PROFILER_ACTIVITIES] @@ -274,15 +241,8 @@ def validate( 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 " 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). + # FSDP manual CPU offload uses swap_tensors, which conflicts with profiler-held + # parameter refs during colocated runs. 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: " @@ -408,9 +368,7 @@ class PolicyConfig(BaseConfig): """Save memory snapshots to ``{ckpt_path}/memory_snapshots/``. Visualize by dragging pickle files to https://docs.pytorch.org/memory_viz.""" torch_profiler_config: TorchProfilerConfig = field(default_factory=TorchProfilerConfig) - """Backend-agnostic ``torch.profiler`` config (FSDP + Megatron). When - ``enable`` is true the trainer drives the profiler around the training loop - and writes traces to ``save_path``. See :class:`TorchProfilerConfig`.""" + """``torch.profiler`` config for policy training steps.""" megatron_config: MegatronConfig = field(default_factory=MegatronConfig) model_config_kwargs: dict = field(default_factory=dict) """Pass-through kwargs for the HuggingFace model config (FSDP backends). diff --git a/skyrl/train/config/sft_config.py b/skyrl/train/config/sft_config.py index bd5b4bdb80..107b11095b 100644 --- a/skyrl/train/config/sft_config.py +++ b/skyrl/train/config/sft_config.py @@ -143,9 +143,7 @@ def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig": """Save memory snapshots to ``{ckpt_path}/memory_snapshots/``. Visualize by dragging pickle files to https://docs.pytorch.org/memory_viz.""" torch_profiler_config: TorchProfilerConfig = field(default_factory=TorchProfilerConfig) - """torch.profiler config (FSDP + Megatron). Set ``enable=true`` to capture - traces of the policy training step around the training loop; see - :class:`TorchProfilerConfig`.""" + """torch.profiler config for policy training steps.""" # ---- SFT-specific flat fields ---- strategy: str = "megatron" # "megatron" or "fsdp" diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py index 24221f45ea..e909226119 100644 --- a/skyrl/train/fully_async_trainer.py +++ b/skyrl/train/fully_async_trainer.py @@ -503,7 +503,9 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don cur_generation_group_mini_batch, cur_dropped_groups, epoch_exhausted, - ) = await self._collect_generation_mini_batch(generation_output_group_buffer, all_generators_done) + ) = await self._collect_generation_mini_batch( + generation_output_group_buffer, all_generators_done + ) if epoch_exhausted: # Exhausted mid mini-batch: discard the partial batch (marked consumed so it @@ -553,9 +555,7 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don # A training step completed: count it for this epoch's bookkeeping. trained_steps_this_epoch += 1 - # Advance the torch profiler schedule once per global step - # (no-op unless profiling is enabled). One schedule step == - # one full async global step; the schedule decides which are recorded. + # One profiler step per async global step. self._profiler_step() # 5. Set logs for this training step. @@ -597,7 +597,9 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don self.cfg.trainer.max_training_steps is not None and self.global_step > self.cfg.trainer.max_training_steps ): - logger.info(f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early.") + logger.info( + f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early." + ) for t in generator_tasks: t.cancel() await asyncio.gather(*generator_tasks, return_exceptions=True) @@ -648,9 +650,6 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don # End of an epoch. finally: - # Always stop/flush the profiler when the loop exits -- including - # via an exception -- so the open kineto trace window isn't leaked. - # No-op when profiling is disabled. self._profiler_stop() pbar.close() diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 8646663b7b..1892d147f1 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -717,8 +717,7 @@ def __init__( @property def _torch_profiler_enabled(self) -> bool: - """Whether the trainer should drive the torch profiler. Gates all - profiler RPC dispatch so non-profiling runs pay zero extra round-trips.""" + """Whether to dispatch policy profiler RPCs.""" return self.cfg.trainer.policy.torch_profiler_config.enable def _build_collator(self, tokenizer): @@ -1311,8 +1310,7 @@ def train_step(self, batch: TrainingInputBatch, step: int) -> dict: metrics = output.metrics - # Advance the torch profiler schedule once per global step (no-op unless - # profiling is enabled; the schedule decides which steps are recorded). + # One profiler step per SFT global step. if self._torch_profiler_enabled: self.dispatch.profile_step("policy") @@ -1439,8 +1437,6 @@ def _train_dummy(self): f"tokens_per_second={tokens_per_second:.0f}" ) finally: - # Always stop/flush the profiler when the loop exits (including via - # an exception) so the open trace window isn't leaked. No-op when off. if self._torch_profiler_enabled: self.dispatch.stop_profile("policy") @@ -1690,8 +1686,6 @@ def train(self): self.global_step += 1 finally: - # Always stop/flush the profiler when the loop exits (including via - # an exception) so the open trace window isn't leaked. No-op when off. if self._torch_profiler_enabled: self.dispatch.stop_profile("policy") self.global_step = min(self.global_step, num_steps) diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 119256ad2f..e23c8dd8cd 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -190,36 +190,21 @@ def has_critic(self) -> bool: @property def _torch_profiler_enabled(self) -> bool: - """Whether the trainer should drive the torch profiler on policy workers. - Gates all profiler RPC dispatch so non-profiling runs pay nothing.""" + """Whether to dispatch policy profiler RPCs.""" return self.cfg.trainer.policy.torch_profiler_config.enable def _profiler_start(self) -> None: - """Arm the torch profiler on the policy workers before the training loop. - - No-op unless profiling is enabled. Shared by every trainer flavor - (sync / fully-async / one-step async) so the driving logic lives in one - place; subclasses that override ``train()`` call this at loop start. - """ + """Start policy profiling when enabled.""" if self._torch_profiler_enabled: self.dispatch.start_profile("policy") def _profiler_step(self) -> None: - """Advance the torch profiler schedule once per global step. - - No-op unless profiling is enabled. One call == one full global step - (not per minibatch); the torch schedule decides which steps are - recorded. Call exactly once per global step from every ``train()``. - """ + """Advance policy profiling by one global step.""" if self._torch_profiler_enabled: self.dispatch.profile_step("policy") def _profiler_stop(self) -> None: - """Stop/flush the torch profiler after the training loop. - - No-op unless profiling is enabled. Call from a ``finally`` so an open - kineto trace window isn't leaked when the loop raises. - """ + """Stop policy profiling when enabled.""" if self._torch_profiler_enabled: self.dispatch.stop_profile("policy") @@ -433,16 +418,16 @@ async def train(self): if self.cfg.trainer.dump_data_batch: # dump data to file with Timer("dump_data_batch"): - self.dump_data(training_input, file_name=f"global_step_{self.global_step}_training_input") + self.dump_data( + training_input, file_name=f"global_step_{self.global_step}_training_input" + ) # 7. train policy/critic model # Policy model is backloaded to GPU during training with Timer("train_critic_and_policy", self.all_timings): status = self.train_critic_and_policy(training_input) - # Advance the torch profiler schedule once per global step - # (no-op unless profiling is enabled). One schedule step == - # one full RL global step; the schedule decides which are recorded. + # One profiler step per RL global step. self._profiler_step() self._fire("on_step_end", batch=training_input, metrics=status) @@ -540,7 +525,9 @@ async def train(self): self.cfg.trainer.max_training_steps is not None and self.global_step > self.cfg.trainer.max_training_steps ): - logger.info(f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early.") + logger.info( + f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early." + ) stop_training = True break @@ -551,9 +538,6 @@ async def train(self): if stop_training: break finally: - # Always stop/flush the profiler when the training loop exits -- - # including via an exception -- so the open kineto trace window - # isn't leaked. No-op when profiling is disabled. self._profiler_stop() pbar.close() diff --git a/tests/backends/skyrl_train/utils/test_profiler.py b/tests/backends/skyrl_train/utils/test_profiler.py index 3b59a00863..003345aa76 100644 --- a/tests/backends/skyrl_train/utils/test_profiler.py +++ b/tests/backends/skyrl_train/utils/test_profiler.py @@ -1,10 +1,4 @@ -"""CPU unit tests for the config-driven torch.profiler wrapper. - -These run without a GPU or torch.distributed: ``torch.profiler`` works on CPU, -and ``Profiler`` falls back to rank 0 when distributed is not initialized. - -uv run --isolated --extra dev pytest tests/backends/skyrl_train/utils/test_profiler.py -v -""" +"""CPU tests for the config-driven torch.profiler wrapper.""" import glob import os @@ -17,7 +11,7 @@ @dataclass class _ProfCfg: - """Minimal stand-in for TorchProfilerConfig (avoids importing skyrl_gym).""" + """TorchProfilerConfig stand-in.""" enable: bool = True ranks: List[int] = field(default_factory=lambda: [0]) @@ -52,8 +46,7 @@ def test_disabled_is_noop(tmp_path): def test_rank_not_selected_is_noop(tmp_path): - # is_initialized() is False in this process -> rank resolves to 0, which is - # not in ranks=[1], so the profiler must not arm. + # Local tests resolve to rank 0. prof = Profiler(_ProfCfg(ranks=[1], save_path=str(tmp_path))) assert prof.prof is None _run_loop(prof, 3) @@ -72,7 +65,7 @@ def test_single_window_writes_one_trace(tmp_path): def test_repeat_writes_multiple_windows(tmp_path): - # Two cycles of (warmup=1 + active=1); needs >= 2*(1+1) steps to fire twice. + # Two warmup+active cycles. prof = Profiler( _ProfCfg(skip_first=0, wait=0, warmup=1, active=1, repeat=2, save_path=str(tmp_path)), ) @@ -91,9 +84,7 @@ def test_skip_first_defers_recording(tmp_path): def test_save_path_is_taken_verbatim(tmp_path): - # `save_path` is an explicit, required local path -- the wrapper uses it as-is - # (no implicit default, no rewriting). Cloud-URI/empty rejection lives in - # TorchProfilerConfig.validate(), exercised in tests/train/test_config.py. + # Path validation lives in TorchProfilerConfig. explicit = str(tmp_path / "explicit") prof = Profiler(_ProfCfg(save_path=explicit)) assert prof.save_path == explicit @@ -127,7 +118,6 @@ def test_kernel_summary_populated_after_window(tmp_path): assert summary is not None assert summary["window_count"] == 1 assert isinstance(summary["pairs"], list) - # Each pair is (name:str, self_us:float); summary is pickle-safe. import pickle pickle.dumps(summary) @@ -140,7 +130,6 @@ def test_activities_threaded_to_torch(tmp_path): import torch prof = Profiler(_ProfCfg(activities=["cpu"], save_path=str(tmp_path))) - # The underlying torch profiler was constructed with the CPU activity only. assert torch.profiler.ProfilerActivity.CPU in prof.prof.activities assert torch.profiler.ProfilerActivity.CUDA not in prof.prof.activities @@ -149,11 +138,7 @@ def test_step_failure_disables_without_raising(tmp_path): prof = Profiler(_ProfCfg(save_path=str(tmp_path))) class _Boom: - """Stands in for a live torch profiler whose ``step`` faults. ``start``/ - ``stop`` are no-ops: we must NOT start the real kineto profiler here, or - the wrapper's fault path nulls ``self.prof`` while that session is still - running, leaving it unstopped -> libkineto segfaults at interpreter exit - (process-wide, surfaces as a teardown crash long after this test passes).""" + """Faulting profiler stub; avoids opening a real kineto session.""" def start(self): pass @@ -164,12 +149,10 @@ def step(self): def stop(self): pass - # Swap in the faulting stub BEFORE start() so no real profiler session is ever - # opened; this still exercises the wrapper's swallow-and-disable path. + # Swap before start() so no real profiler session opens. prof.prof = _Boom() prof.start() - # Simulate an internal profiler fault mid-loop; the wrapper must swallow it. prof.step() assert prof.enable is False assert prof.prof is None @@ -179,9 +162,7 @@ def stop(self): class TestWorkerProfilerRPCs: - """The Worker base exposes the profiler-control RPCs and they no-op when - ``self.profiler`` is None (so the trainer can dispatch them unconditionally, - and so they're on the snapshotted Ray actor method table without a subclass).""" + """Worker profiler RPC coverage.""" def test_methods_exist_on_worker_base(self): from skyrl.backends.skyrl_train.workers.worker import Worker @@ -211,8 +192,6 @@ def test_rpcs_noop_when_profiler_none(self): from skyrl.backends.skyrl_train.workers.worker import Worker - # Call the unbound methods against a stub whose profiler is None; they - # must early-return without touching anything. stub = SimpleNamespace(profiler=None) Worker.start_profile(stub) Worker.profile_step(stub) @@ -237,9 +216,7 @@ def test_rpcs_drive_profiler_when_present(self): class TestBuildProfilerFromPolicyCfg: - """``build_profiler_from_policy_cfg`` is the production entry both worker - backends call in ``init_model``. It gates on ``enable`` and passes the - explicit, validated ``save_path`` straight through (no implicit default).""" + """Coverage for the worker profiler factory.""" @staticmethod def _trainer_cfg(prof_cfg): @@ -269,10 +246,7 @@ def test_builds_profiler_with_explicit_save_path(self, tmp_path): class TestWorkerDispatchProfilerRPCs: - """The WorkerDispatch profiler wrappers dispatch via ``pass_through`` to the - named model's actor group, no-op for an unknown model, and swallow dispatch - faults so a profiler error can never abort the training loop. Tested against - the unbound methods with a stub ``self`` (no live Ray actors required).""" + """WorkerDispatch profiler RPC coverage.""" @staticmethod def _stub(actor_groups): @@ -295,7 +269,6 @@ def test_unknown_model_is_noop(self): from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch stub = self._stub({}) # no "policy" group - # None of these should raise or dispatch. WorkerDispatch.start_profile(stub, "policy") WorkerDispatch.profile_step(stub, "policy") WorkerDispatch.stop_profile(stub, "policy") @@ -337,7 +310,6 @@ def boom(_): raise RuntimeError("ray.get boom") with patch("skyrl.backends.skyrl_train.workers.worker_dispatch.ray.get", side_effect=boom): - # Must not propagate -- a profiler fault cannot abort training. WorkerDispatch.start_profile(stub, "policy") WorkerDispatch.profile_step(stub, "policy") WorkerDispatch.stop_profile(stub, "policy") @@ -345,10 +317,7 @@ def boom(_): class TestTrainerProfilerHelpers: - """``RayPPOTrainer`` exposes gated ``_profiler_{start,step,stop}`` helpers so - every trainer flavor (sync / fully-async / one-step async) drives the - profiler through one code path. The helpers dispatch only when - ``_torch_profiler_enabled`` is True, and always target the policy model.""" + """RayPPOTrainer profiler helper coverage.""" @staticmethod def _trainer(enable): @@ -356,7 +325,6 @@ def _trainer(enable): from skyrl.train.trainer import RayPPOTrainer - # Build a bare instance without running __init__ (which needs Ray/cfg). trainer = object.__new__(RayPPOTrainer) calls = [] trainer.dispatch = SimpleNamespace( diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 6e3425f217..d579841ef2 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -425,23 +425,19 @@ def test_validate_cfg_allows_explicit_max_seq_len_for_seq_mean_token_sum_norm(se class TestTorchProfilerConfigValidation: - """``TorchProfilerConfig.validate()`` rejects unusable profiler settings up - front (so an enabled run fails fast instead of silently degrading), and is a - no-op when profiling is disabled.""" + """TorchProfilerConfig validation coverage.""" @staticmethod def _cfg(**overrides): from skyrl.train.config.config import TorchProfilerConfig - # `save_path` is required when enabled; supply a valid local default so - # tests targeting *other* fields don't trip the save_path guard. + # Valid default for tests targeting other fields. overrides.setdefault("save_path", "/tmp/skyrl_prof_test") return TorchProfilerConfig(enable=True, **overrides) def test_disabled_skips_all_checks(self): from skyrl.train.config.config import TorchProfilerConfig - # Garbage values must be tolerated while disabled (the default state). TorchProfilerConfig( enable=False, export_type="bogus", activities=["gpu"], ranks=[], active=0, save_path=None ).validate() @@ -454,14 +450,12 @@ def test_empty_ranks_rejected(self): self._cfg(ranks=[]).validate() def test_missing_save_path_rejected(self): - # save_path has no implicit default; an enabled run must set it explicitly. with pytest.raises(ValueError, match=r"save_path.*must be set"): self._cfg(save_path=None).validate() with pytest.raises(ValueError, match=r"save_path.*must be set"): self._cfg(save_path="").validate() def test_cloud_save_path_rejected(self): - # torch.profiler can only write the local filesystem; cloud URIs fail fast. for uri in ("s3://bucket/run/traces", "gs://bucket/run/traces", "gcs://bucket/run/traces"): with pytest.raises(ValueError, match=r"save_path.*local path"): self._cfg(save_path=uri).validate() @@ -471,8 +465,6 @@ def test_unknown_activity_rejected(self): self._cfg(activities=["cpu", "gpu"]).validate() def test_empty_activities_rejected(self): - # An empty list profiles nothing; the membership check would pass it - # vacuously, so it needs its own guard. with pytest.raises(ValueError, match=r"activities.*non-empty"): self._cfg(activities=[]).validate() @@ -486,7 +478,6 @@ def test_unknown_export_type_rejected(self): def test_stacks_requires_with_stack(self): with pytest.raises(ValueError, match=r"with_stack"): self._cfg(export_type="stacks", with_stack=False).validate() - # With with_stack=True it is accepted. self._cfg(export_type="stacks", with_stack=True).validate() def test_negative_schedule_field_rejected(self): @@ -498,7 +489,6 @@ def test_active_must_be_at_least_one(self): self._cfg(active=0).validate() def test_validate_cfg_invokes_profiler_validation(self): - # The RL entrypoint validator must surface profiler config errors. 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" @@ -506,40 +496,29 @@ def test_validate_cfg_invokes_profiler_validation(self): 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 ") 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. + # FSDP manual-offload incompatibility. 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" @@ -551,4 +530,4 @@ def test_validate_cfg_allows_profiler_with_native_offload(self): 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 + validate_cfg(cfg) diff --git a/tests/train/test_sft_config.py b/tests/train/test_sft_config.py index 86c4e39fc6..65bb80ef2c 100644 --- a/tests/train/test_sft_config.py +++ b/tests/train/test_sft_config.py @@ -135,7 +135,7 @@ def test_lora_disabled_by_default(self): class TestTorchProfilerConfigOverrides: - """torch_profiler_config bridges to policy.torch_profiler_config.""" + """SFT profiler config bridge coverage.""" def test_disabled_by_default(self): cfg = _sft_cfg_from_overrides([]) @@ -177,8 +177,6 @@ def test_capture_flags_propagate(self): assert list(prof.activities) == ["cuda"] def test_invalid_profiler_config_rejected_by_sft_validation(self): - # validate_sft_cfg (run by build_skyrl_config_for_sft) must surface - # profiler config errors on the SFT path, not just the RL path. cfg = _sft_cfg_from_overrides( [ "model.path=test/my-model",