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..c10129d1d1 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 for policy training + enable: false + ranks: [0] + save_path: null # required when enable=true; absolute local path + skip_first: 10 + wait: 0 + warmup: 1 + active: 1 + 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 + 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` 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: policy forward/backward/optimizer only. RL critic/ref models and generation are not profiled. + +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. + +- `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 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..1796f9198d 100644 --- a/examples/train/async/async_trainer.py +++ b/examples/train/async/async_trainer.py @@ -51,63 +51,73 @@ 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() + + # One profiler step per global step. + 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: + 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..05a11e5874 100644 --- a/examples/train/megatron/run_megatron.sh +++ b/examples/train/megatron/run_megatron.sh @@ -17,7 +17,7 @@ MEGATRON_TP=2 MEGATRON_PP=2 MEGATRON_CP=1 -# torch profiler config +# 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}" @@ -33,9 +33,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 \ @@ -69,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 1a8bf8c752..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}" @@ -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 \ @@ -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 6022639d9d..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 @@ -26,68 +28,81 @@ 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) + + # One profiler step per global step. + 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: + self._profiler_stop() self.tracker.finish() logger.info("Dummy training completed successfully!") diff --git a/skyrl/backends/skyrl_train/utils/profiler.py b/skyrl/backends/skyrl_train/utils/profiler.py index da01dfbd50..2e3f6a6147 100644 --- a/skyrl/backends/skyrl_train/utils/profiler.py +++ b/skyrl/backends/skyrl_train/utils/profiler.py @@ -4,82 +4,126 @@ import torch.distributed from loguru import logger +# Config string -> torch.profiler activity. +_ACTIVITY_MAP = { + "cpu": torch.profiler.ProfilerActivity.CPU, + "cuda": torch.profiler.ProfilerActivity.CUDA, +} + + +def build_profiler_from_policy_cfg(trainer_cfg): + """Build the policy profiler, or return None when disabled.""" + cfg = trainer_cfg.policy.torch_profiler_config + if not cfg.enable: + return None + return Profiler(cfg) + class Profiler: - """ - A PyTorch profiler wrapper class for collecting performance metrics. - """ + """Thin ``torch.profiler`` wrapper driven by trainer start/step/stop calls.""" def __init__(self, config): - """ - config contains: - - enable: bool - - ranks: list[int] - - save_path: str - """ self.enable = config.enable + self.prof = None + # 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 + # Validated at startup by TorchProfilerConfig. 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.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), ) - - def check(self): + 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: + """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": + 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: + # 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 ``{"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)} + + 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..e80cf65d5e 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) + # 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"): # 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..844d762b37 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 + # 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,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) + # Created only on profiled ranks. + 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..e179dc9e6e 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/skyrl/backends/skyrl_train/workers/worker.py @@ -233,6 +233,8 @@ 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 + # Populated by init_model when torch profiling is enabled. + self.profiler = None if self.cfg.algorithm.temperature is None: raise ValueError("`cfg.algorithm.temperature` must be set") @@ -292,6 +294,29 @@ def set_algorithm_config(self, **kwargs) -> None: for key, value in kwargs.items(): setattr(self.cfg.algorithm, key, value) + # ------------------------------------------------------------------ + # torch.profiler RPCs, dispatched via WorkerDispatch pass_through. + # ------------------------------------------------------------------ + + 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.""" + 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 summary, or None.""" + 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..cf08210d0b 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,48 @@ 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. Avoid _ensure_on_gpu so profiling does not perturb + # the colocation offload state. + # ------------------------------------------------------------------ + + def start_profile(self, model: str) -> None: + """Start profiling on ``model`` workers.""" + 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 profiling by one global step.""" + 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 profiling on ``model`` workers.""" + 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 per-rank last-window kernel summaries for ``model``.""" + 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..7dff56f60e 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -156,11 +156,106 @@ 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): + """``torch.profiler`` config for policy training steps.""" + 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. Required when ``enable=True``; must be a local absolute path.""" + + # torch.profiler.schedule + skip_first: int = 10 + """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 means forever.""" + + # torch.profiler.profile + 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`` or ``stacks``; stacks require ``with_stack=True``.""" + + 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: + """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.") + # 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. " + "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." + ) + # 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] + 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}.") + + # 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: " + "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 @@ -207,7 +302,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 +367,8 @@ 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) + """``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 f2e216992c..107b11095b 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,8 @@ 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 for policy training steps.""" # ---- SFT-specific flat fields ---- strategy: str = "megatron" # "megatron" or "fsdp" @@ -304,6 +307,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 +392,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..e909226119 100644 --- a/skyrl/train/fully_async_trainer.py +++ b/skyrl/train/fully_async_trainer.py @@ -461,185 +461,197 @@ 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 + + # One profiler step per async global step. + 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}" - ) - - 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) - # 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: + 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..1892d147f1 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -715,6 +715,11 @@ def __init__( self._steps_per_epoch: int = 0 self._current_epoch: int = 0 + @property + def _torch_profiler_enabled(self) -> bool: + """Whether to dispatch policy profiler RPCs.""" + 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 +1309,11 @@ def train_step(self, batch: TrainingInputBatch, step: int) -> dict: grad_norm = self.dispatch.optim_step("policy") metrics = output.metrics + + # One profiler step per SFT global step. + 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 +1404,41 @@ 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: + if self._torch_profiler_enabled: + self.dispatch.stop_profile("policy") logger.info("Dummy SFT training complete!") @@ -1553,118 +1569,125 @@ 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: + 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..e23c8dd8cd 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -188,6 +188,26 @@ 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 to dispatch policy profiler RPCs.""" + return self.cfg.trainer.policy.torch_profiler_config.enable + + def _profiler_start(self) -> None: + """Start policy profiling when enabled.""" + if self._torch_profiler_enabled: + self.dispatch.start_profile("policy") + + def _profiler_step(self) -> None: + """Advance policy profiling by one global step.""" + if self._torch_profiler_enabled: + self.dispatch.profile_step("policy") + + def _profiler_stop(self) -> None: + """Stop policy profiling when enabled.""" + 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 +310,235 @@ 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) + + # One profiler step per RL global step. + 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: + 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..ce4eadf693 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -270,6 +270,13 @@ def validate_cfg(cfg: SkyRLTrainConfig): "or negative to keep all checkpoints" ) + 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() 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..003345aa76 --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_profiler.py @@ -0,0 +1,358 @@ +"""CPU tests for the config-driven torch.profiler wrapper.""" + +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: + """TorchProfilerConfig stand-in.""" + + 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, 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): + # 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) + 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, 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 warmup+active cycles. + prof = Profiler( + _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*")) + 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, 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_is_taken_verbatim(tmp_path): + # Path validation lives in TorchProfilerConfig. + 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, 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, 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"], 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) + 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"], save_path=str(tmp_path))) + 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(save_path=str(tmp_path))) + + class _Boom: + """Faulting profiler stub; avoids opening a real kineto session.""" + + def start(self): + pass + + def step(self): + raise RuntimeError("boom") + + def stop(self): + pass + + # Swap before start() so no real profiler session opens. + prof.prof = _Boom() + prof.start() + + prof.step() + assert prof.enable is False + assert prof.prof is None + # Subsequent calls remain safe no-ops. + prof.step() + prof.stop() + + +class TestWorkerProfilerRPCs: + """Worker profiler RPC coverage.""" + + 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 + + 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: + """Coverage for the worker profiler factory.""" + + @staticmethod + def _trainer_cfg(prof_cfg): + from types import SimpleNamespace + + 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, save_path=str(tmp_path))) + assert build_profiler_from_policy_cfg(cfg) is None + + 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, + ) + + explicit = str(tmp_path / "explicit") + 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 + + +class TestWorkerDispatchProfilerRPCs: + """WorkerDispatch profiler RPC coverage.""" + + @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 + 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): + 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 profiler helper coverage.""" + + @staticmethod + def _trainer(enable): + from types import SimpleNamespace + + from skyrl.train.trainer import RayPPOTrainer + + 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..d579841ef2 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -422,3 +422,112 @@ 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 validation coverage.""" + + @staticmethod + def _cfg(**overrides): + from skyrl.train.config.config import TorchProfilerConfig + + # 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 + + 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 + + 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): + 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): + 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() + + def test_empty_activities_rejected(self): + 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 + + 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() + 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): + 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) + + # 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): + 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): + self._cfg().validate(strategy="fsdp", colocate_all=False, colocate_policy_ref=False, fsdp_cpu_offload=False) + + def test_fsdp_native_cpu_offload_allowed(self): + self._cfg().validate(strategy="fsdp", colocate_all=True, colocate_policy_ref=True, fsdp_cpu_offload=True) + + def test_megatron_colocation_allowed(self): + self._cfg().validate(strategy="megatron", colocate_all=True, colocate_policy_ref=True, fsdp_cpu_offload=False) + + def test_offload_check_skipped_without_context(self): + self._cfg().validate() + + def test_validate_cfg_rejects_profiler_under_default_colocation(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" + 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) diff --git a/tests/train/test_sft_config.py b/tests/train/test_sft_config.py index 621b49be72..65bb80ef2c 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: + """SFT profiler config bridge coverage.""" + + 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]", + "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.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): + cfg = _sft_cfg_from_overrides( + [ + "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"): + build_skyrl_config_for_sft(cfg) + + class TestFSDPConfigOverrides: """FSDP config overrides propagate when strategy=fsdp."""